diff --git a/BUILD b/BUILD index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..4bf647e47aa56cff0b3fd5af7d5df99d8b70549b 100644 --- a/BUILD +++ b/BUILD @@ -0,0 +1,6 @@ +exports_files( + [ + "LICENSE", + "ACKNOWLEDGEMENTS", + ], +) diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md index ff11d131409b65880f16b80f9fe38dc39ac0e5fa..5fff9d05a1c589636bc9c711e6eb7cc4aba86b2f 100644 --- a/CODE_OF_CONDUCT.md +++ b/CODE_OF_CONDUCT.md @@ -67,4 +67,4 @@ If the Project Stewards receive a report alleging a violation of the Code of Con ## Attribution -This Code of Conduct is adapted from the Contributor Covenant, version 1.4, available at http://contributor-covenant.org/version/1/4, and includes some aspects of the Geek Feminism Code of Conduct and the Drupal Code of Conduct. +This Code of Conduct is adapted from the Contributor Covenant, version 1.4, available at https://contributor-covenant.org/version/1/4, and includes some aspects of the Geek Feminism Code of Conduct and the Drupal Code of Conduct. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index dc96bc2e3d3960827efd109551f8eaa78a6cfb48..3dad41a88c8212b7445c32f241d887306d3c19ad 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -8,8 +8,8 @@ We'd love to accept your patches! Before we can take them, we have to jump a cou Please fill out either the individual or corporate Contributor License Agreement (CLA). - * If you are an individual writing original source code and you're sure you own the intellectual property, then you'll need to sign an [individual CLA](http://code.google.com/legal/individual-cla-v1.0.html). - * If you work for a company that wants to allow you to contribute your work, then you'll need to sign a [corporate CLA](http://code.google.com/legal/corporate-cla-v1.0.html). + * If you are an individual writing original source code and you're sure you own the intellectual property, then you'll need to sign an [individual CLA](https://code.google.com/legal/individual-cla-v1.0.html). + * If you work for a company that wants to allow you to contribute your work, then you'll need to sign a [corporate CLA](https://code.google.com/legal/corporate-cla-v1.0.html). Follow either of the two links above to access the appropriate CLA and instructions for how to sign and return it. Once we receive it, we'll be able to accept your pull requests. @@ -41,7 +41,7 @@ TensorFlow coding style. #### General guidelines and philosophy for contribution * Include unit tests when you contribute new features, as they help to - a) prove that your code works correctly, b) guard against future breaking + a) prove that your code works correctly, and b) guard against future breaking changes to lower the maintenance cost. * Bug fixes also generally require unit tests, because the presence of bugs usually indicates insufficient test coverage. @@ -51,7 +51,7 @@ TensorFlow coding style. non-backward-compatible API changes without a major release. Reviewers of your pull request will comment on any API compatibility issues. * When you contribute a new feature to TensorFlow, the maintenance burden is (by - default) transferred to the TensorFlow team. This means that benefit of + default) transferred to the TensorFlow team. This means that benefit of the contribution must be compared against the cost of maintaining the feature. * Full new features (e.g., a new op implementing a cutting-edge algorithm) typically will live in @@ -68,8 +68,8 @@ Include a license at the top of new files. * [Java license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/java/src/main/java/org/tensorflow/Graph.java#L1) * [Go license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/go/operation.go#L1) * [Bash license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/ci_build/ci_sanity.sh#L2) -* [HTML license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorboard/dist/index.html#L2) -* [JavaScript/TypeScript license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorboard/components/tf_backend/backend.ts#L1) +* [HTML license example](https://github.com/tensorflow/tensorboard/blob/master/tensorboard/components/tf_backend/tf-backend.html#L2) +* [JavaScript/TypeScript license example](https://github.com/tensorflow/tensorboard/blob/master/tensorboard/components/tf_backend/backend.ts#L1) Bazel BUILD files also need to include a license section, e.g., [BUILD example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/BUILD#L61). @@ -117,7 +117,7 @@ pylint --rcfile=/tmp/pylintrc myfile.py * [Google Java Style Guide](https://google.github.io/styleguide/javaguide.html) * [Google JavaScript Style Guide](https://google.github.io/styleguide/jsguide.html) * [Google Shell Style Guide](https://google.github.io/styleguide/shell.xml) -* [Google Objective-C Style Guide](http://google.github.io/styleguide/objcguide.html) +* [Google Objective-C Style Guide](https://google.github.io/styleguide/objcguide.html) #### Running sanity check @@ -163,7 +163,7 @@ There are two ways to run TensorFlow unit tests. bazel test ${flags} //tensorflow/python/... ``` -2. Using [Docker](www.docker.com) and TensorFlow's CI scripts. +2. Using [Docker](https://www.docker.com) and TensorFlow's CI scripts. ```bash # Install Docker first, then this will build and run cpu tests diff --git a/ISSUE_TEMPLATE.md b/ISSUE_TEMPLATE.md index 1a401997c649518766acb2ebb0dea1c128bd0ba4..2f3df7cda9cec29ed0c2266629022f0a22b37df9 100644 --- a/ISSUE_TEMPLATE.md +++ b/ISSUE_TEMPLATE.md @@ -4,7 +4,7 @@ https://stackoverflow.com/questions/tagged/tensorflow If you open a GitHub issue, here is our policy: -1. It must be a bug or a feature request. +1. It must be a bug, a feature request, or a significant problem with documentation (for small docs fixes please send a PR instead). 2. The form below must be filled out. 3. It shouldn't be a TensorBoard issue. Those go [here](https://github.com/tensorflow/tensorboard/issues). diff --git a/LICENSE b/LICENSE index 15ae42140452d32ccf929f59f7eca01a3c7b555f..4862420c0234f7542d4fe8f3520516b484a64aed 100644 --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright 2017 The TensorFlow Authors. All rights reserved. +Copyright 2018 The TensorFlow Authors. All rights reserved. Apache License Version 2.0, January 2004 diff --git a/README.md b/README.md index aff3427bddb307aea6d6c2466eac14c9edffcc32..916e5200b29841028652c861c49dbb3650baea3c 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ | **`Linux CPU`** | **`Linux GPU`** | **`Mac OS CPU`** | **`Windows CPU`** | **`Android`** | |-----------------|---------------------|------------------|-------------------|---------------| -| [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-cpu)](https://ci.tensorflow.org/job/tensorflow-master-cpu) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-linux-gpu)](https://ci.tensorflow.org/job/tensorflow-master-linux-gpu) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-mac)](https://ci.tensorflow.org/job/tensorflow-master-mac) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-win-cmake-py)](https://ci.tensorflow.org/job/tensorflow-master-win-cmake-py) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-android)](https://ci.tensorflow.org/job/tensorflow-master-android) | +| [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-cpu)](https://ci.tensorflow.org/job/tensorflow-master-cpu) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-linux-gpu)](https://ci.tensorflow.org/job/tensorflow-master-linux-gpu) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-mac)](https://ci.tensorflow.org/job/tensorflow-master-mac) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-win-cmake-py)](https://ci.tensorflow.org/job/tensorflow-master-win-cmake-py) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-android)](https://ci.tensorflow.org/job/tensorflow-master-android) [ ![Download](https://api.bintray.com/packages/google/tensorflow/tensorflow/images/download.svg) ](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 @@ -27,10 +27,14 @@ guidelines](CONTRIBUTING.md). This project adheres to TensorFlow's uphold this code.** **We use [GitHub issues](https://github.com/tensorflow/tensorflow/issues) for -tracking requests and bugs. So please see +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: + +[![CII Best Practices](https://bestpractices.coreinfrastructure.org/projects/1486/badge)](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.* @@ -46,11 +50,11 @@ packages on Linux, Mac, and Windows. **Individual whl files** -* Linux CPU-only: [Python 2](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-cp27-none-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave/)) / [Python 3.4](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-cp34-cp34m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=cpu-slave/)) / [Python 3.5](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-cp35-cp35m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=cpu-slave/)) -* Linux GPU: [Python 2](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-linux/42/artifact/pip_test/whl/tf_nightly_gpu-1.head-cp27-none-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-linux/)) / [Python 3.4](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly_gpu-1.head-cp34-cp34m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-linux/)) / [Python 3.5](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly_gpu-1.head-cp35-cp35m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=gpu-linux/)) +* Linux CPU-only: [Python 2](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-cp27-none-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave/)) / [Python 3.4](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-cp34-cp34m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=cpu-slave/)) / [Python 3.5](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-cp35-cp35m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=cpu-slave/)) / [Python 3.6](http://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.6,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-cp36-cp36m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.6,label=cpu-slave/)) +* Linux GPU: [Python 2](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-linux/42/artifact/pip_test/whl/tf_nightly_gpu-1.head-cp27-none-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-linux/)) / [Python 3.4](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly_gpu-1.head-cp34-cp34m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-linux/)) / [Python 3.5](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly_gpu-1.head-cp35-cp35m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=gpu-linux/)) / [Python 3.6](http://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.6,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly_gpu-1.head-cp36-cp36m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-linux/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.6,label=gpu-linux/)) * Mac CPU-only: [Python 2](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-mac/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=mac-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-py2-none-any.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-mac/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=mac-slave/)) / [Python 3](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-mac/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=mac-slave/lastSuccessfulBuild/artifact/pip_test/whl/tf_nightly-1.head-py3-none-any.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-mac/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=mac-slave/)) -* Windows CPU-only: [Python 3.5 64-bit](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows,PY=35/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tf_nightly-1.head-cp35-cp35m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows,PY=35/)) / [Python 3.6 64-bit](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows,PY=36/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tf_nightly-1.head-cp36-cp36m-win_amd64.whl) ([build history](http://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows,PY=36/)) -* Windows GPU: [Python 3.5 64-bit](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows-gpu,PY=35/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tf_nightly_gpu-1.head-cp35-cp35m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows-gpu,PY=35/)) / [Python 3.6 64-bit](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows-gpu,PY=36/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tf_nightly_gpu-1.head-cp36-cp36m-win_amd64.whl) ([build history](http://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows-gpu,PY=36/)) +* Windows CPU-only: [Python 3.5 64-bit](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows,PY=35/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tf_nightly-1.head-cp35-cp35m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows,PY=35/)) / [Python 3.6 64-bit](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows,PY=36/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tf_nightly-1.head-cp36-cp36m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows,PY=36/)) +* Windows GPU: [Python 3.5 64-bit](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows-gpu,PY=35/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tf_nightly_gpu-1.head-cp35-cp35m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows-gpu,PY=35/)) / [Python 3.6 64-bit](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows-gpu,PY=36/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tf_nightly_gpu-1.head-cp36-cp36m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/tf-nightly/job/tf-nightly-windows/M=windows-gpu,PY=36/)) * Android: [demo APK](https://ci.tensorflow.org/view/Nightly/job/nightly-android/lastSuccessfulBuild/artifact/out/tensorflow_demo.apk), [native libs](https://ci.tensorflow.org/view/Nightly/job/nightly-android/lastSuccessfulBuild/artifact/out/native/) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-android/)) diff --git a/RELEASE.md b/RELEASE.md index e04bd3fc505d51ade9e9fa12c822cb695e90b4f3..6f54dee58f75c29a16545ba25de12fe059baf1eb 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,3 +1,285 @@ +# Release 1.6.0 + +## Breaking Changes +* Prebuilt binaries are now built against CUDA 9.0 and cuDNN 7. +* Prebuilt binaries will use AVX instructions. This may break TF on older CPUs. + +## Major Features And Improvements +* New Optimizer internal API for non-slot variables. Descendants of AdamOptimizer that access _beta[12]_power will need to be updated. +* `tf.estimator.{FinalExporter,LatestExporter}` now export stripped SavedModels. This improves forward compatibility of the SavedModel. +* FFT support added to XLA CPU/GPU. + +## Bug Fixes and Other Changes +* Documentation updates: + * Added a second version of Getting Started, which is aimed at ML +newcomers. + * Clarified documentation on `resize_images.align_corners` parameter. + * Additional documentation for TPUs. +* Google Cloud Storage (GCS): + * Add client-side throttle. + * Add a `FlushCaches()` method to the FileSystem interface, with an implementation for GcsFileSystem. +* Other: + * Add `tf.contrib.distributions.Kumaraswamy`. + * `RetryingFileSystem::FlushCaches()` calls the base FileSystem's `FlushCaches()`. + * Add `auto_correlation` to distributions. + * Add `tf.contrib.distributions.Autoregressive`. + * Add SeparableConv1D layer. + * Add convolutional Flipout layers. + * When both inputs of `tf.matmul` are bfloat16, it returns bfloat16, instead of float32. + * Added `tf.contrib.image.connected_components`. + * Add `tf.contrib.framework.CriticalSection` that allows atomic variable access. + * Output variance over trees predictions for classifications tasks. + * For `pt` and `eval` commands, allow writing tensor values to filesystem as numpy files. + * gRPC: Propagate truncated errors (instead of returning gRPC internal error). + * Augment `parallel_interleave` to support 2 kinds of prefetching. + * Improved XLA support for C64-related ops log, pow, atan2, tanh. + * Add probabilistic convolutional layers. + +## API Changes +* Introducing `prepare_variance` boolean with default setting to False for backward compatibility. +* Move `layers_dense_variational_impl.py` to `layers_dense_variational.py`. + +## Known Bugs +* Using XLA:GPU with CUDA 9 and CUDA 9.1 results in garbage results and/or + `CUDA_ILLEGAL_ADDRESS` failures. + + Google discovered in mid-December 2017 that the PTX-to-SASS compiler in CUDA 9 + and CUDA 9.1 sometimes does not properly compute the carry bit when + decomposing 64-bit address calculations with large offsets (e.g. `load [x + + large_constant]`) into 32-bit arithmetic in SASS. + + As a result, these versions of `ptxas` miscompile most XLA programs which use + more than 4GB of temp memory. This results in garbage results and/or + `CUDA_ERROR_ILLEGAL_ADDRESS` failures. + + A fix in CUDA 9.1.121 is expected in late February 2018. We do not expect a + fix for CUDA 9.0.x. Until the fix is available, the only workaround is to + [downgrade](https://developer.nvidia.com/cuda-toolkit-archive) to CUDA 8.0.x + or disable XLA:GPU. + + TensorFlow will print a warning if you use XLA:GPU with a known-bad version of + CUDA; see e00ba24c4038e7644da417ddc639169b6ea59122. + +## Thanks to our Contributors + +This release contains contributions from many people at Google, as well as: + +4d55397500, Ag Ramesh, Aiden Scandella, Akimasa Kimura, Alex Rothberg, Allen Goodman, +amilioto, Andrei Costinescu, Andrei Nigmatulin, Anjum Sayed, Anthony Platanios, +Anush Elangovan, Armando Fandango, Ashish Kumar Ram, Ashwini Shukla, Ben, Bhavani Subramanian, +Brett Koonce, Carl Thomé, cclauss, Cesc, Changming Sun, Christoph Boeddeker, Clayne Robison, +Clemens Schulz, Clint (Woonhyuk Baek), codrut3, Cole Gerdemann, Colin Raffel, Daniel Trebbien, +Daniel Ylitalo, Daniel Zhang, Daniyar, Darjan Salaj, Dave Maclachlan, David Norman, Dong--Jian, +dongsamb, dssgsra, Edward H, eladweiss, elilienstein, Eric Lilienstein, error.d, Eunji Jeong, fanlu, +Florian Courtial, fo40225, Fred, Gregg Helt, Guozhong Zhuang, Hanchen Li, hsm207, hyunyoung2, +ImSheridan, Ishant Mrinal Haloi, Jacky Ko, Jay Young, Jean Flaherty, Jerome, JerrikEph, Jesse +Kinkead, jfaath, Jian Lin, jinghuangintel, Jiongyan Zhang, Joel Hestness, Joel Shor, Johnny Chan, +Julian Niedermeier, Julian Wolff, JxKing, K-W-W, Karl Lessard, Kasper Marstal, Keiji Ariyama, +Koan-Sin Tan, Loki Der Quaeler, Loo Rong Jie, Luke Schaefer, Lynn Jackson, ManHyuk, Matt Basta, +Matt Smith, Matthew Schulkind, Michael, michaelkhan3, Miguel Piedrafita, Mikalai Drabovich, +Mike Knapp, mjwen, mktozk, Mohamed Aly, Mohammad Ashraf Bhuiyan, Myungjoo Ham, Naman Bhalla, +Namrata-Ibm, Nathan Luehr, nathansilberman, Netzeband, Niranjan Hasabnis, Omar Aflak, Ozge +Yalcinkaya, Parth P Panchal, patrickzzy, Patryk Chrabaszcz, Paul Van Eck, Paweł Kapica, Peng Yu, +Philip Yang, Pierre Blondeau, Po-Hsien Chu, powderluv, Puyu Wang, Rajendra Arora, Rasmus, Renat +Idrisov, resec, Robin Richtsfeld, Ronald Eddy Jr, Sahil Singh, Sam Matzek, Sami Kama, sandipmgiri, +Santiago Castro, Sayed Hadi Hashemi, Scott Tseng, Sergii Khomenko, Shahid, Shengpeng Liu, Shreyash +Sharma, Shrinidhi Kl, Simone Cirillo, simsicon, Stanislav Levental, starsblinking, Stephen Lumenta, +Steven Hickson, Su Tang, Taehoon Lee, Takuya Wakisaka, Ted Chang, Ted Ying, Tijmen Verhulsdonck, +Timofey Kondrashov, vade, vaibhav, Valentin Khrulkov, vchigrin, Victor Costan, Viraj Navkal, +Vivek Rane, wagonhelm, Yan Facai (颜发才), Yanbo Liang, Yaroslav Bulatov, yegord, Yong Tang, +Yoni Tsafir, yordun, Yuan (Terry) Tang, Yuxin Wu, zhengdi, Zhengsheng Wei, 田传武 + +# Release 1.5.0 + +## Breaking Changes +* Prebuilt binaries are now built against CUDA 9.0 and cuDNN 7. +* Starting from 1.6 release, our prebuilt binaries will use AVX instructions. + This may break TF on older CPUs. + +## Major Features And Improvements +* [Eager execution](https://github.com/tensorflow/tensorflow/tree/r1.5/tensorflow/contrib/eager) + preview version is now available. +* [TensorFlow Lite](https://github.com/tensorflow/tensorflow/tree/r1.5/tensorflow/contrib/lite) + dev preview is now available. +* CUDA 9.0 and cuDNN 7 support. +* Accelerated Linear Algebra (XLA): + * Add `complex64` support to XLA compiler. + * `bfloat` support is now added to XLA infrastructure. + * Make `ClusterSpec` propagation work with XLA devices. + * Use a determinisitic executor to generate XLA graph. +* `tf.contrib`: + * `tf.contrib.distributions`: + * Add `tf.contrib.distributions.Autoregressive`. + * Make `tf.contrib.distributions` QuadratureCompound classes support batch + * Infer `tf.contrib.distributions.RelaxedOneHotCategorical` `dtype` from arguments. + * Make `tf.contrib.distributions` quadrature family parameterized by + `quadrature_grid_and_prob` vs `quadrature_degree`. + * `auto_correlation` added to `tf.contrib.distributions` + * Add `tf.contrib.bayesflow.layers`, a collection of probabilistic (neural) layers. + * Add `tf.contrib.bayesflow.halton_sequence`. + * Add `tf.contrib.data.make_saveable_from_iterator.` + * Add `tf.contrib.data.shuffle_and_repeat`. + * Add new custom transformation: `tf.contrib.data.scan()`. + * `tf.contrib.distributions.bijectors`: + * Add `tf.contrib.distributions.bijectors.MaskedAutoregressiveFlow`. + * Add `tf.contrib.distributions.bijectors.Permute`. + * Add `tf.contrib.distributions.bijectors.Gumbel`. + * Add `tf.contrib.distributions.bijectors.Reshape`. + * Support shape inference (i.e., shapes containing -1) in the Reshape bijector. +* Add `streaming_precision_recall_at_equal_thresholds,` a method for computing + streaming precision and recall with `O(num_thresholds + size of predictions)` + time and space complexity. +* Change `RunConfig` default behavior to not set a random seed, making random + behavior independently random on distributed workers. We expect this to + generally improve training performance. Models that do rely on determinism + should set a random seed explicitly. +* Replaced the implementation of `tf.flags` with `absl.flags`. +* Add support for `CUBLAS_TENSOR_OP_MATH` in fp16 GEMM +* Add support for CUDA on NVIDIA Tegra devices + +## Bug Fixes and Other Changes +* Documentation updates: + * Clarified that you can only install TensorFlow on 64-bit machines. + * Added a short doc explaining how `Estimator`s save checkpoints. + * Add documentation for ops supported by the `tf2xla` bridge. + * Fix minor typos in the doc of `SpaceToDepth` and `DepthToSpace`. + * Updated documentation comments in `mfcc_mel_filterbank.h` and `mfcc.h` to + clarify that the input domain is squared magnitude spectra and the weighting + is done on linear magnitude spectra (sqrt of inputs). + * Change `tf.contrib.distributions` docstring examples to use `tfd` alias + rather than `ds`, `bs`. + * Fix docstring typos in `tf.distributions.bijectors.Bijector`. + * `tf.assert_equal` no longer raises `ValueError.` It now raises + `InvalidArgumentError,` as documented. + * Update Getting Started docs and API intro. +* Google Cloud Storage (GCS): + * Add userspace DNS caching for the GCS client. + * Customize request timeouts for the GCS filesystem. + * Improve GCS filesystem caching. +* Bug Fixes: + * Fix bug where partitioned integer variables got their wrong shapes. Before + * Fix correctness bug in CPU and GPU implementations of Adadelta. + * Fix a bug in `import_meta_graph`'s handling of partitioned variables when + importing into a scope. WARNING: This may break loading checkpoints of + graphs with partitioned variables saved after using `import_meta_graph` with + a non-empty `import_scope` argument. + * Fix bug in offline debugger which prevented viewing events. + * Added the `WorkerService.DeleteWorkerSession` method to the gRPC interface, + to fix a memory leak. Ensure that your master and worker servers are running + the same version of TensorFlow to avoid compatibility issues. + * Fix bug in peephole implementation of BlockLSTM cell. + * Fix bug by casting dtype of `log_det_jacobian` to match `log_prob` in + `TransformedDistribution`. + * Fix a bug in `import_meta_graph`'s handling of partitioned variables when + * Ensure `tf.distributions.Multinomial` doesn't underflow in `log_prob`. + Before this change, all partitions of an integer variable were initialized + with the shape of the unpartitioned variable; after this change they are + initialized correctly. +* Other: + * Add necessary shape util support for bfloat16. + * Add a way to run ops using a step function to MonitoredSession. + * Add `DenseFlipout` probabilistic layer. + * A new flag `ignore_live_threads` is available on train. If set to `True`, it + will ignore threads that remain running when tearing down infrastructure + after successfully completing training, instead of throwing a RuntimeError. + * Restandardize `DenseVariational` as simpler template for other probabilistic + layers. + * `tf.data` now supports `tf.SparseTensor` components in dataset elements. + * It is now possible to iterate over `Tensor`s. + * Allow `SparseSegmentReduction` ops to have missing segment IDs. + * Modify custom export strategy to account for multidimensional sparse float + splits. + * `Conv2D`, `Conv2DBackpropInput`, `Conv2DBackpropFilter` now supports arbitrary + dilations with GPU and cuDNNv6 support. + * `Estimator` now supports `Dataset`: `input_fn` can return a `Dataset` + instead of `Tensor`s. + * Add `RevBlock`, a memory-efficient implementation of reversible residual layers. + * Reduce BFCAllocator internal fragmentation. + * Add `cross_entropy` and `kl_divergence` to `tf.distributions.Distribution`. + * Add `tf.nn.softmax_cross_entropy_with_logits_v2` which enables backprop + w.r.t. the labels. + * GPU back-end now uses `ptxas` to compile generated PTX. + * `BufferAssignment`'s protocol buffer dump is now deterministic. + * Change embedding op to use parallel version of `DynamicStitch`. + * Add support for sparse multidimensional feature columns. + * Speed up the case for sparse float columns that have only 1 value. + * Allow sparse float splits to support multivalent feature columns. + * Add `quantile` to `tf.distributions.TransformedDistribution`. + * Add `NCHW_VECT_C` support for `tf.depth_to_space` on GPU. + * Add `NCHW_VECT_C` support for `tf.space_to_depth` on GPU. + +## API Changes +* Rename `SqueezeDims` attribute to `Axis` in C++ API for Squeeze op. +* `Stream::BlockHostUntilDone` now returns Status rather than bool. +* Minor refactor: move stats files from `stochastic` to `common` and remove + `stochastic`. + +## Known Bugs +* Using XLA:GPU with CUDA 9 and CUDA 9.1 results in garbage results and/or + `CUDA_ILLEGAL_ADDRESS` failures. + + Google discovered in mid-December 2017 that the PTX-to-SASS compiler in CUDA 9 + and CUDA 9.1 sometimes does not properly compute the carry bit when + decomposing 64-bit address calculations with large offsets (e.g. `load [x + + large_constant]`) into 32-bit arithmetic in SASS. + + As a result, these versions of `ptxas` miscompile most XLA programs which use + more than 4GB of temp memory. This results in garbage results and/or + `CUDA_ERROR_ILLEGAL_ADDRESS` failures. + + A fix in CUDA 9.1.121 is expected in late February 2018. We do not expect a + fix for CUDA 9.0.x. Until the fix is available, the only workaround is to + [downgrade](https://developer.nvidia.com/cuda-toolkit-archive) to CUDA 8.0.x + or disable XLA:GPU. + + TensorFlow will print a warning if you use XLA:GPU with a known-bad version of + CUDA; see e00ba24c4038e7644da417ddc639169b6ea59122. + +## Thanks to our Contributors + +This release contains contributions from many people at Google, as well as: + +Adam Zahran, Ag Ramesh, Alan Lee, Alan Yee, Alex Sergeev, Alexander, Amir H. Jadidinejad, +Amy, Anastasios Doumoulakis, Andrei Costinescu, Andrei Nigmatulin, Anthony Platanios, +Anush Elangovan, arixlin, Armen Donigian, ArtëM Sobolev, Atlas7, Ben Barsdell, Bill Prin, +Bo Wang, Brett Koonce, Cameron Thomas, Carl Thomé, Cem Eteke, cglewis, Changming Sun, +Charles Shenton, Chi-Hung, Chris Donahue, Chris Filo Gorgolewski, Chris Hoyean Song, +Chris Tava, Christian Grail, Christoph Boeddeker, cinqS, Clayne Robison, codrut3, concerttttt, +CQY, Dan Becker, Dan Jarvis, Daniel Zhang, David Norman, dmaclach, Dmitry Trifonov, +Donggeon Lim, dongpilYu, Dr. Kashif Rasul, Edd Wilder-James, Eric Lv, fcharras, Felix Abecassis, +FirefoxMetzger, formath, FredZhang, Gaojin Cao, Gary Deer, Guenther Schmuelling, Hanchen Li, +Hanmin Qin, hannesa2, hyunyoung2, Ilya Edrenkin, Jackson Kontny, Jan, Javier Luraschi, +Jay Young, Jayaram Bobba, Jeff, Jeff Carpenter, Jeremy Sharpe, Jeroen BéDorf, Jimmy Jia, +Jinze Bai, Jiongyan Zhang, Joe Castagneri, Johan Ju, Josh Varty, Julian Niedermeier, +JxKing, Karl Lessard, Kb Sriram, Keven Wang, Koan-Sin Tan, Kyle Mills, lanhin, LevineHuang, +Loki Der Quaeler, Loo Rong Jie, Luke Iwanski, LáSzló Csomor, Mahdi Abavisani, Mahmoud Abuzaina, +ManHyuk, Marek ŠUppa, MathSquared, Mats Linander, Matt Wytock, Matthew Daley, Maximilian Bachl, +mdymczyk, melvyniandrag, Michael Case, Mike Traynor, miqlas, Namrata-Ibm, Nathan Luehr, +Nathan Van Doorn, Noa Ezra, Nolan Liu, Oleg Zabluda, opensourcemattress, Ouwen Huang, +Paul Van Eck, peisong, Peng Yu, PinkySan, pks, powderluv, Qiao Hai-Jun, Qiao Longfei, +Rajendra Arora, Ralph Tang, resec, Robin Richtsfeld, Rohan Varma, Ryohei Kuroki, SaintNazaire, +Samuel He, Sandeep Dcunha, sandipmgiri, Sang Han, scott, Scott Mudge, Se-Won Kim, Simon Perkins, +Simone Cirillo, Steffen Schmitz, Suvojit Manna, Sylvus, Taehoon Lee, Ted Chang, Thomas Deegan, +Till Hoffmann, Tim, Toni Kunic, Toon Verstraelen, Tristan Rice, Urs KöSter, Utkarsh Upadhyay, +Vish (Ishaya) Abrams, Winnie Tsang, Yan Chen, Yan Facai (颜发才), Yi Yang, Yong Tang, +Youssef Hesham, Yuan (Terry) Tang, Zhengsheng Wei, zxcqwe4906, 张志豪, 田传武 + +We are also grateful to all who filed issues or helped resolve them, asked and +answered questions, and were part of inspiring discussions. + +# Release 1.4.1 + +## Bug Fixes and Other Changes +* `LinearClassifier` fix. + +# Release 1.4.0 + +## Major Features And Improvements +* `tf.keras` is now part of the core TensorFlow API. +* [`tf.data`](http://tensorflow.org/programmers_guide/datasets) is now part of + the core TensorFlow API. + * The API is now subject to backwards compatibility guarantees. + # Release 1.4.0 ## Major Features And Improvements @@ -351,7 +633,7 @@ answered questions, and were part of inspiring discussions. * Fixed LIBXSMM integration. * Make decode_jpeg/decode_png/decode_gif handle all formats, since users frequently try to decode an image as the wrong type. * Improve implicit broadcasting lowering. -* Improving stability of GCS/Bigquery clients by a faster retrying of stale transmissions. +* Improving stability of GCS/BigQuery clients by a faster retrying of stale transmissions. * Remove OpKernelConstruction::op_def() as part of minimizing proto dependencies. * VectorLaplaceDiag distribution added. * Android demo no longer requires libtensorflow_demo.so to run (libtensorflow_inference.so still required) diff --git a/WORKSPACE b/WORKSPACE index b40913801ba8e3c8ee73f7ba69540b520ad698a6..1e38a9a8cd754886fc5232531816b875de0879a3 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -2,11 +2,11 @@ workspace(name = "org_tensorflow") http_archive( name = "io_bazel_rules_closure", - sha256 = "110fe68753413777944b473c25eed6368c4a0487cee23a7bac1b13cc49d3e257", - strip_prefix = "rules_closure-4af89ef1db659eb41f110df189b67d4cf14073e1", + sha256 = "6691c58a2cd30a86776dd9bb34898b041e37136f2dc7e24cadaeaf599c95c657", + strip_prefix = "rules_closure-08039ba8ca59f64248bb3b6ae016460fe9c9914f", urls = [ - "https://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/4af89ef1db659eb41f110df189b67d4cf14073e1.tar.gz", - "https://github.com/bazelbuild/rules_closure/archive/4af89ef1db659eb41f110df189b67d4cf14073e1.tar.gz", # 2017-08-28 + "https://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/08039ba8ca59f64248bb3b6ae016460fe9c9914f.tar.gz", + "https://github.com/bazelbuild/rules_closure/archive/08039ba8ca59f64248bb3b6ae016460fe9c9914f.tar.gz", # 2018-01-16 ], ) @@ -41,12 +41,12 @@ load("//tensorflow:workspace.bzl", "tf_workspace") tf_workspace() new_http_archive( - name = "inception5h", + name = "inception_v1", build_file = "models.BUILD", - sha256 = "d13569f6a98159de37e92e9c8ec4dae8f674fbf475f69fe6199b514f756d4364", + sha256 = "7efe12a8363f09bc24d7b7a450304a15655a57a7751929b2c1593a71183bb105", urls = [ - "http://storage.googleapis.com/download.tensorflow.org/models/inception5h.zip", - "http://download.tensorflow.org/models/inception5h.zip", + "http://storage.googleapis.com/download.tensorflow.org/models/inception_v1.zip", + "http://download.tensorflow.org/models/inception_v1.zip", ], ) diff --git a/configure.py b/configure.py index e4218b56518a3d76564533df9b31d514a7313c49..3aa1a3e956c6a559b89cdeb593a96a95188c32ae 100644 --- a/configure.py +++ b/configure.py @@ -43,6 +43,7 @@ _DEFAULT_CUDA_PATH = '/usr/local/cuda' _DEFAULT_CUDA_PATH_LINUX = '/opt/cuda' _DEFAULT_CUDA_PATH_WIN = ('C:/Program Files/NVIDIA GPU Computing ' 'Toolkit/CUDA/v%s' % _DEFAULT_CUDA_VERSION) +_DEFAULT_TENSORRT_PATH_LINUX = '/usr/lib/x86_64-linux-gnu' _TF_OPENCL_VERSION = '1.2' _DEFAULT_COMPUTECPP_TOOLKIT_PATH = '/usr/local/computecpp' _DEFAULT_TRISYCL_INCLUDE_DIR = '/usr/local/triSYCL/include' @@ -297,11 +298,17 @@ def get_var(environ_cp, System". enabled_by_default: boolean for default behavior. question: optional string for how to ask for user input. - yes_reply: optionanl string for reply when feature is enabled. + yes_reply: optional string for reply when feature is enabled. no_reply: optional string for reply when feature is disabled. Returns: boolean value of the variable. + + Raises: + UserInputError: if an environment variable is set, but it cannot be + interpreted as a boolean indicator, assume that the user has made a + scripting error, and will continue to provide invalid input. + Raise the error to avoid infinitely looping. """ if not question: question = 'Do you wish to build TensorFlow with %s support?' % query_item @@ -319,6 +326,23 @@ def get_var(environ_cp, question += ' [y/N]: ' var = environ_cp.get(var_name) + if var is not None: + var_content = var.strip().lower() + true_strings = ('1', 't', 'true', 'y', 'yes') + false_strings = ('0', 'f', 'false', 'n', 'no') + if var_content in true_strings: + var = True + elif var_content in false_strings: + var = False + else: + raise UserInputError( + 'Environment variable %s must be set as a boolean indicator.\n' + 'The following are accepted as TRUE : %s.\n' + 'The following are accepted as FALSE: %s.\n' + 'Current value is %s.' % ( + var_name, ', '.join(true_strings), ', '.join(false_strings), + var)) + while var is None: user_input_origin = get_input(question) user_input = user_input_origin.strip().lower() @@ -387,7 +411,7 @@ def set_action_env_var(environ_cp, System". enabled_by_default: boolean for default behavior. question: optional string for how to ask for user input. - yes_reply: optionanl string for reply when feature is enabled. + yes_reply: optional string for reply when feature is enabled. no_reply: optional string for reply when feature is disabled. """ var = int( @@ -421,7 +445,7 @@ def convert_version_to_int(version): def check_bazel_version(min_version): - """Check installed bezel version is at least min_version. + """Check installed bazel version is at least min_version. Args: min_version: string for minimum bazel version. @@ -605,8 +629,9 @@ def prompt_loop_or_load_from_env( Raises: UserInputError: if a query has been attempted n_ask_attempts times without - success, assume that the user has made a scripting error, and will continue - to provide invalid input. Raise the error to avoid infinitely looping. + success, assume that the user has made a scripting error, and will + continue to provide invalid input. Raise the error to avoid infinitely + looping. """ default = environ_cp.get(var_name) or var_default full_query = '%s [Default is %s]: ' % ( @@ -802,6 +827,28 @@ def set_gcc_host_compiler_path(environ_cp): write_action_env_to_bazelrc('GCC_HOST_COMPILER_PATH', gcc_host_compiler_path) +def reformat_version_sequence(version_str, sequence_count): + """Reformat the version string to have the given number of sequences. + + For example: + Given (7, 2) -> 7.0 + (7.0.1, 2) -> 7.0 + (5, 1) -> 5 + (5.0.3.2, 1) -> 5 + + Args: + version_str: String, the version string. + sequence_count: int, an integer. + Returns: + string, reformatted version string. + """ + v = version_str.split('.') + if len(v) < sequence_count: + v = v + (['0'] * (sequence_count - len(v))) + + return '.'.join(v[:sequence_count]) + + def set_tf_cuda_version(environ_cp): """Set CUDA_TOOLKIT_PATH and TF_CUDA_VERSION.""" ask_cuda_version = ( @@ -812,6 +859,7 @@ def set_tf_cuda_version(environ_cp): # Configure the Cuda SDK version to use. tf_cuda_version = get_from_env_or_user_or_default( environ_cp, 'TF_CUDA_VERSION', ask_cuda_version, _DEFAULT_CUDA_VERSION) + tf_cuda_version = reformat_version_sequence(str(tf_cuda_version), 2) # Find out where the CUDA toolkit is installed default_cuda_path = _DEFAULT_CUDA_PATH @@ -868,6 +916,7 @@ def set_tf_cudnn_version(environ_cp): tf_cudnn_version = get_from_env_or_user_or_default( environ_cp, 'TF_CUDNN_VERSION', ask_cudnn_version, _DEFAULT_CUDNN_VERSION) + tf_cudnn_version = reformat_version_sequence(str(tf_cudnn_version) ,1) default_cudnn_path = environ_cp.get('CUDA_TOOLKIT_PATH') ask_cudnn_path = (r'Please specify the location where cuDNN %s library is ' @@ -935,6 +984,128 @@ def set_tf_cudnn_version(environ_cp): write_action_env_to_bazelrc('TF_CUDNN_VERSION', tf_cudnn_version) +def set_tf_tensorrt_install_path(environ_cp): + """Set TENSORRT_INSTALL_PATH and TF_TENSORRT_VERSION. + + Adapted from code contributed by Sami Kama (https://github.com/samikama). + + Args: + environ_cp: copy of the os.environ. + + Raises: + ValueError: if this method was called under non-Linux platform. + UserInputError: if user has provided invalid input multiple times. + """ + if not is_linux(): + raise ValueError('Currently TensorRT is only supported on Linux platform.') + + # Ask user whether to add TensorRT support. + if str(int(get_var( + environ_cp, 'TF_NEED_TENSORRT', 'TensorRT', False))) != '1': + return + + for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS): + ask_tensorrt_path = (r'Please specify the location where TensorRT is ' + 'installed. [Default is %s]:') % ( + _DEFAULT_TENSORRT_PATH_LINUX) + trt_install_path = get_from_env_or_user_or_default( + environ_cp, 'TENSORRT_INSTALL_PATH', ask_tensorrt_path, + _DEFAULT_TENSORRT_PATH_LINUX) + + # Result returned from "read" will be used unexpanded. That make "~" + # unusable. Going through one more level of expansion to handle that. + trt_install_path = os.path.realpath( + os.path.expanduser(trt_install_path)) + + def find_libs(search_path): + """Search for libnvinfer.so in "search_path".""" + fl = set() + if os.path.exists(search_path) and os.path.isdir(search_path): + fl.update([os.path.realpath(os.path.join(search_path, x)) + for x in os.listdir(search_path) if 'libnvinfer.so' in x]) + return fl + + possible_files = find_libs(trt_install_path) + possible_files.update(find_libs(os.path.join(trt_install_path, 'lib'))) + possible_files.update(find_libs(os.path.join(trt_install_path, 'lib64'))) + + def is_compatible(tensorrt_lib, cuda_ver, cudnn_ver): + """Check the compatibility between tensorrt and cudnn/cudart libraries.""" + ldd_bin = which('ldd') or '/usr/bin/ldd' + ldd_out = run_shell([ldd_bin, tensorrt_lib]).split(os.linesep) + cudnn_pattern = re.compile('.*libcudnn.so\\.?(.*) =>.*$') + cuda_pattern = re.compile('.*libcudart.so\\.?(.*) =>.*$') + cudnn = None + cudart = None + for line in ldd_out: + if 'libcudnn.so' in line: + cudnn = cudnn_pattern.search(line) + elif 'libcudart.so' in line: + cudart = cuda_pattern.search(line) + if cudnn and len(cudnn.group(1)): + cudnn = convert_version_to_int(cudnn.group(1)) + if cudart and len(cudart.group(1)): + cudart = convert_version_to_int(cudart.group(1)) + return (cudnn == cudnn_ver) and (cudart == cuda_ver) + + cuda_ver = convert_version_to_int(environ_cp['TF_CUDA_VERSION']) + cudnn_ver = convert_version_to_int(environ_cp['TF_CUDNN_VERSION']) + nvinfer_pattern = re.compile('.*libnvinfer.so.?(.*)$') + highest_ver = [0, None, None] + + for lib_file in possible_files: + if is_compatible(lib_file, cuda_ver, cudnn_ver): + ver_str = nvinfer_pattern.search(lib_file).group(1) + ver = convert_version_to_int(ver_str) if len(ver_str) else 0 + if ver > highest_ver[0]: + highest_ver = [ver, ver_str, lib_file] + if highest_ver[1] is not None: + trt_install_path = os.path.dirname(highest_ver[2]) + tf_tensorrt_version = highest_ver[1] + break + + # Try another alternative from ldconfig. + ldconfig_bin = which('ldconfig') or '/sbin/ldconfig' + ldconfig_output = run_shell([ldconfig_bin, '-p']) + search_result = re.search( + '.*libnvinfer.so\\.?([0-9.]*).* => (.*)', ldconfig_output) + if search_result: + libnvinfer_path_from_ldconfig = search_result.group(2) + if os.path.exists(libnvinfer_path_from_ldconfig): + if is_compatible(libnvinfer_path_from_ldconfig, cuda_ver, cudnn_ver): + trt_install_path = os.path.dirname(libnvinfer_path_from_ldconfig) + tf_tensorrt_version = search_result.group(1) + break + + # Reset and Retry + if len(possible_files): + print('TensorRT libraries found in one the following directories', + 'are not compatible with selected cuda and cudnn installations') + print(trt_install_path) + print(os.path.join(trt_install_path, 'lib')) + print(os.path.join(trt_install_path, 'lib64')) + if search_result: + print(libnvinfer_path_from_ldconfig) + else: + print('Invalid path to TensorRT. None of the following files can be found:') + print(trt_install_path) + print(os.path.join(trt_install_path, 'lib')) + print(os.path.join(trt_install_path, 'lib64')) + if search_result: + print(libnvinfer_path_from_ldconfig) + + else: + raise UserInputError('Invalid TF_TENSORRT setting was provided %d ' + 'times in a row. Assuming to be a scripting mistake.' % + _DEFAULT_PROMPT_ASK_ATTEMPTS) + + # Set TENSORRT_INSTALL_PATH and TF_TENSORRT_VERSION + environ_cp['TENSORRT_INSTALL_PATH'] = trt_install_path + write_action_env_to_bazelrc('TENSORRT_INSTALL_PATH', trt_install_path) + environ_cp['TF_TENSORRT_VERSION'] = tf_tensorrt_version + write_action_env_to_bazelrc('TF_TENSORRT_VERSION', tf_tensorrt_version) + + def get_native_cuda_compute_capabilities(environ_cp): """Get native cuda compute capabilities. @@ -1098,17 +1269,20 @@ def set_computecpp_toolkit_path(environ_cp): write_action_env_to_bazelrc('COMPUTECPP_TOOLKIT_PATH', computecpp_toolkit_path) + def set_trisycl_include_dir(environ_cp): - """Set TRISYCL_INCLUDE_DIR""" + """Set TRISYCL_INCLUDE_DIR.""" + ask_trisycl_include_dir = ('Please specify the location of the triSYCL ' 'include directory. (Use --config=sycl_trisycl ' 'when building with Bazel) ' '[Default is %s]: ' - ) % (_DEFAULT_TRISYCL_INCLUDE_DIR) + ) % (_DEFAULT_TRISYCL_INCLUDE_DIR) + while True: trisycl_include_dir = get_from_env_or_user_or_default( - environ_cp, 'TRISYCL_INCLUDE_DIR', ask_trisycl_include_dir, - _DEFAULT_TRISYCL_INCLUDE_DIR) + environ_cp, 'TRISYCL_INCLUDE_DIR', ask_trisycl_include_dir, + _DEFAULT_TRISYCL_INCLUDE_DIR) if os.path.exists(trisycl_include_dir): break @@ -1178,46 +1352,10 @@ def set_other_mpi_vars(environ_cp): raise ValueError('Cannot find the MPI library file in %s/lib' % mpi_home) -def set_mkl(): - write_to_bazelrc('build:mkl --define using_mkl=true') - write_to_bazelrc('build:mkl -c opt') - print( - 'Add "--config=mkl" to your bazel command to build with MKL ' - 'support.\nPlease note that MKL on MacOS or windows is still not ' - 'supported.\nIf you would like to use a local MKL instead of ' - 'downloading, please set the environment variable \"TF_MKL_ROOT\" every ' - 'time before build.\n') - - -def set_monolithic(): - # Add --config=monolithic to your bazel command to use a mostly-static - # build and disable modular op registration support (this will revert to - # loading TensorFlow with RTLD_GLOBAL in Python). By default (without - # --config=monolithic), TensorFlow will build with a dependence on - # //tensorflow:libtensorflow_framework.so. - write_to_bazelrc('build:monolithic --define framework_shared_object=false') - # For projects which use TensorFlow as part of a Bazel build process, putting - # nothing in a bazelrc will default to a monolithic build. The following line - # opts in to modular op registration support by default: - write_to_bazelrc('build --define framework_shared_object=true') - - -def create_android_bazelrc_configs(): - # Flags for --config=android - write_to_bazelrc('build:android --crosstool_top=//external:android/crosstool') - write_to_bazelrc( - 'build:android --host_crosstool_top=@bazel_tools//tools/cpp:toolchain') - # Flags for --config=android_arm - write_to_bazelrc('build:android_arm --config=android') - write_to_bazelrc('build:android_arm --cpu=armeabi-v7a') - # Flags for --config=android_arm64 - write_to_bazelrc('build:android_arm64 --config=android') - write_to_bazelrc('build:android_arm64 --cpu=arm64-v8a') - - def set_grpc_build_flags(): write_to_bazelrc('build --define grpc_no_ares=true') + def set_windows_build_flags(): if is_windows(): # The non-monolithic build is not supported yet @@ -1228,6 +1366,11 @@ def set_windows_build_flags(): write_to_bazelrc('build --verbose_failures') +def config_info_line(name, help_text): + """Helper function to print formatted help text for Bazel config options.""" + print('\t--config=%-12s\t# %s' % (name, help_text)) + + def main(): # Make a copy of os.environ to be clear when functions and getting and setting # environment variables. @@ -1244,13 +1387,16 @@ def main(): environ_cp['TF_NEED_GCP'] = '0' environ_cp['TF_NEED_HDFS'] = '0' environ_cp['TF_NEED_JEMALLOC'] = '0' + environ_cp['TF_NEED_KAFKA'] = '0' environ_cp['TF_NEED_OPENCL_SYCL'] = '0' environ_cp['TF_NEED_COMPUTECPP'] = '0' environ_cp['TF_NEED_OPENCL'] = '0' environ_cp['TF_CUDA_CLANG'] = '0' + environ_cp['TF_NEED_TENSORRT'] = '0' if is_macos(): environ_cp['TF_NEED_JEMALLOC'] = '0' + environ_cp['TF_NEED_TENSORRT'] = '0' set_build_var(environ_cp, 'TF_NEED_JEMALLOC', 'jemalloc as malloc', 'with_jemalloc', True) @@ -1260,6 +1406,8 @@ def main(): 'with_hdfs_support', True, 'hdfs') set_build_var(environ_cp, 'TF_NEED_S3', 'Amazon S3 File System', 'with_s3_support', True, 's3') + set_build_var(environ_cp, 'TF_NEED_KAFKA', 'Apache Kafka Platform', + 'with_kafka_support', False, 'kafka') set_build_var(environ_cp, 'TF_ENABLE_XLA', 'XLA JIT', 'with_xla_support', False, 'xla') set_build_var(environ_cp, 'TF_NEED_GDR', 'GDR', 'with_gdr_support', @@ -1282,7 +1430,11 @@ def main(): 'TF_CUDA_CONFIG_REPO' not in environ_cp): set_tf_cuda_version(environ_cp) set_tf_cudnn_version(environ_cp) + if is_linux(): + set_tf_tensorrt_install_path(environ_cp) set_tf_cuda_compute_capabilities(environ_cp) + if 'LD_LIBRARY_PATH' in environ_cp and environ_cp.get('LD_LIBRARY_PATH') != '1': + write_action_env_to_bazelrc('LD_LIBRARY_PATH', environ_cp.get('LD_LIBRARY_PATH')) set_tf_cuda_clang(environ_cp) if environ_cp.get('TF_CUDA_CLANG') == '1': @@ -1313,10 +1465,7 @@ def main(): set_grpc_build_flags() set_cc_opt_flags(environ_cp) - set_mkl() - set_monolithic() set_windows_build_flags() - create_android_bazelrc_configs() if workspace_has_any_android_rule(): print('The WORKSPACE file has at least one of ["android_sdk_repository", ' @@ -1334,6 +1483,11 @@ def main(): create_android_ndk_rule(environ_cp) create_android_sdk_rule(environ_cp) + print('Preconfigured Bazel build configs. You can use any of the below by ' + 'adding "--config=<>" to your build command. See tools/bazel.rc for ' + 'more details.') + config_info_line('mkl', 'Build with MKL support.') + config_info_line('monolithic', 'Config for mostly static monolithic build.') if __name__ == '__main__': main() diff --git a/tensorflow/BUILD b/tensorflow/BUILD index 6f7e54a90a624d371f1311fcef1807eb90ed36ff..dc995d231d3e591771f801e28024a76610cdba26 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -211,6 +211,12 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "with_kafka_support", + define_values = {"with_kafka_support": "true"}, + visibility = ["//visibility:public"], +) + # Crosses between platforms and file system libraries not supported on those # platforms due to limitations in nested select() statements. config_setting( @@ -364,12 +370,27 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "override_eigen_strong_inline", + values = {"define": "override_eigen_strong_inline=true"}, + visibility = ["//visibility:public"], +) + +# TODO(laigd): consider removing this option and make TensorRT enabled +# automatically when CUDA is enabled. +config_setting( + name = "with_tensorrt_support", + values = {"define": "with_tensorrt_support=true"}, + visibility = ["//visibility:public"], +) + package_group( name = "internal", packages = [ "//learning/meta_rank/...", "//tensorflow/...", "//tensorflow_fold/llgtm/...", + "//third_party/py/tensor2tensor/...", ], ) @@ -401,6 +422,8 @@ filegroup( "//tensorflow/c:all_files", "//tensorflow/cc:all_files", "//tensorflow/cc/saved_model:all_files", + "//tensorflow/cc/saved_model/python:all_files", + "//tensorflow/cc/tools:all_files", "//tensorflow/compiler/aot:all_files", "//tensorflow/compiler/aot/tests:all_files", "//tensorflow/compiler/jit:all_files", @@ -433,9 +456,6 @@ filegroup( "//tensorflow/contrib/all_reduce:all_files", "//tensorflow/contrib/android:all_files", "//tensorflow/contrib/batching:all_files", - "//tensorflow/contrib/batching/kernels:all_files", - "//tensorflow/contrib/batching/test_util:all_files", - "//tensorflow/contrib/batching/util:all_files", "//tensorflow/contrib/bayesflow:all_files", "//tensorflow/contrib/boosted_trees:all_files", "//tensorflow/contrib/boosted_trees/estimator_batch:all_files", @@ -445,6 +465,7 @@ filegroup( "//tensorflow/contrib/cloud:all_files", "//tensorflow/contrib/cloud/kernels:all_files", "//tensorflow/contrib/cluster_resolver:all_files", + "//tensorflow/contrib/coder:all_files", "//tensorflow/contrib/compiler:all_files", "//tensorflow/contrib/copy_graph:all_files", "//tensorflow/contrib/crf:all_files", @@ -462,6 +483,7 @@ filegroup( "//tensorflow/contrib/factorization:all_files", "//tensorflow/contrib/factorization/examples:all_files", "//tensorflow/contrib/factorization/kernels:all_files", + "//tensorflow/contrib/feature_column:all_files", "//tensorflow/contrib/ffmpeg:all_files", "//tensorflow/contrib/ffmpeg/default:all_files", "//tensorflow/contrib/framework:all_files", @@ -520,16 +542,18 @@ filegroup( "//tensorflow/contrib/metrics:all_files", "//tensorflow/contrib/model_pruning:all_files", "//tensorflow/contrib/model_pruning/examples/cifar10:all_files", - "//tensorflow/contrib/mpi:all_files", - "//tensorflow/contrib/mpi_collectives:all_files", "//tensorflow/contrib/nccl:all_files", - "//tensorflow/contrib/ndlstm:all_files", "//tensorflow/contrib/nearest_neighbor:all_files", "//tensorflow/contrib/nn:all_files", "//tensorflow/contrib/opt:all_files", "//tensorflow/contrib/periodic_resample:all_files", "//tensorflow/contrib/predictor:all_files", - "//tensorflow/contrib/quantization:all_files", + "//tensorflow/contrib/py2tf:all_files", + "//tensorflow/contrib/py2tf/converters:all_files", + "//tensorflow/contrib/py2tf/impl:all_files", + "//tensorflow/contrib/py2tf/pyct:all_files", + "//tensorflow/contrib/py2tf/pyct/static_analysis:all_files", + "//tensorflow/contrib/py2tf/utils:all_files", "//tensorflow/contrib/quantize:all_files", "//tensorflow/contrib/receptive_field:all_files", "//tensorflow/contrib/reduce_slice_ops:all_files", @@ -558,6 +582,7 @@ filegroup( "//tensorflow/contrib/tensor_forest/proto:all_files", "//tensorflow/contrib/tensorboard:all_files", "//tensorflow/contrib/tensorboard/db:all_files", + "//tensorflow/contrib/tensorrt:all_files", "//tensorflow/contrib/testing:all_files", "//tensorflow/contrib/text:all_files", "//tensorflow/contrib/tfprof:all_files", @@ -572,6 +597,7 @@ filegroup( "//tensorflow/contrib/util:all_files", "//tensorflow/contrib/verbs:all_files", "//tensorflow/core:all_files", + "//tensorflow/core/api_def:all_files", "//tensorflow/core/debug:all_files", "//tensorflow/core/distributed_runtime:all_files", "//tensorflow/core/distributed_runtime/rpc:all_files", @@ -582,6 +608,7 @@ filegroup( "//tensorflow/core/grappler/optimizers:all_files", "//tensorflow/core/grappler/utils:all_files", "//tensorflow/core/kernels:all_files", + "//tensorflow/core/kernels/batching_util:all_files", "//tensorflow/core/kernels/data:all_files", "//tensorflow/core/kernels/data/sql:all_files", "//tensorflow/core/kernels/fuzzing:all_files", @@ -664,8 +691,6 @@ filegroup( "//third_party/fft2d:all_files", "//third_party/flatbuffers:all_files", "//third_party/hadoop:all_files", - "//third_party/mpi:all_files", - "//third_party/mpi_collectives:all_files", "//third_party/sycl:all_files", "//third_party/sycl/sycl:all_files", ], @@ -793,6 +818,7 @@ tf_cc_shared_object( "//tensorflow/cc:cc_ops", "//tensorflow/cc:client_session", "//tensorflow/cc:scope", + "//tensorflow/cc/profiler", "//tensorflow/core:tensorflow", ], ) diff --git a/tensorflow/SECURITY.md b/tensorflow/SECURITY.md new file mode 100644 index 0000000000000000000000000000000000000000..6ddac1f964dfba3afd240441e2a036bc24ee6d91 --- /dev/null +++ b/tensorflow/SECURITY.md @@ -0,0 +1,239 @@ +# Using TensorFlow Securely + +This document discusses how to safely deal with untrusted programs (models or +model parameters), and input data. Below, we also provide guidelines on how to +report vulnerabilities in TensorFlow. + +## TensorFlow models are programs + +TensorFlow's runtime system interprets and executes programs. What machine +learning practitioners term +[**models**](https://developers.google.com/machine-learning/glossary/#model) are +expressed as programs that TensorFlow executes. TensorFlow programs are encoded +as computation +[**graphs**](https://developers.google.com/machine-learning/glossary/#graph). +The model's parameters are often stored separately in **checkpoints**. + +At runtime, TensorFlow executes the computation graph using the parameters +provided. Note that the behavior of the computation graph may change +depending on the parameters provided. TensorFlow itself is not a sandbox. When +executing the computation graph, TensorFlow may read and write files, send and +receive data over the network, and even spawn additional processes. All these +tasks are performed with the permissions of the TensorFlow process. Allowing +for this flexibility makes for a powerful machine learning platform, +but it has implications for security. + +The computation graph may also accept **inputs**. Those inputs are the +data you supply to TensorFlow to train a model, or to use a model to run +inference on the data. + +**TensorFlow models are programs, and need to be treated as such from a security +perspective.** + +## Running untrusted models + +As a general rule: **Always** execute untrusted models inside a sandbox (e.g., +[nsjail](https://github.com/google/nsjail)). + +There are several ways in which a model could become untrusted. Obviously, if an +untrusted party supplies TensorFlow kernels, arbitrary code may be executed. +The same is true if the untrusted party provides Python code, such as the +Python code that generates TensorFlow graphs. + +Even if the untrusted party only supplies the serialized computation +graph (in form of a `GraphDef`, `SavedModel`, or equivalent on-disk format), the +set of computation primitives available to TensorFlow is powerful enough that +you should assume that the TensorFlow process effectively executes arbitrary +code. One common solution is to whitelist only a few safe Ops. While this is +possible in theory, we still recommend you sandbox the execution. + +It depends on the computation graph whether a user provided checkpoint is safe. +It is easily possible to create computation graphs in which malicious +checkpoints can trigger unsafe behavior. For example, consider a graph that +contains a `tf.cond` depending on the value of a `tf.Variable`. One branch of +the `tf.cond` is harmless, but the other is unsafe. Since the `tf.Variable` is +stored in the checkpoint, whoever provides the checkpoint now has the ability to +trigger unsafe behavior, even though the graph is not under their control. + +In other words, graphs can contain vulnerabilities of their own. To allow users +to provide checkpoints to a model you run on their behalf (e.g., in order to +compare model quality for a fixed model architecture), you must carefully audit +your model, and we recommend you run the TensorFlow process in a sandbox. + +## Accepting untrusted Inputs + +It is possible to write models that are secure in a sense that they can safely +process untrusted inputs assuming there are no bugs. There are two main reasons +to not rely on this: first, it is easy to write models which must not be exposed +to untrusted inputs, and second, there are bugs in any software system of +sufficient complexity. Letting users control inputs could allow them to trigger +bugs either in TensorFlow or in dependent libraries. + +In general, it is good practice to isolate parts of any system which is exposed +to untrusted (e.g., user-provided) inputs in a sandbox. + +A useful analogy to how any TensorFlow graph is executed is any interpreted +programming language, such as Python. While it is possible to write secure +Python code which can be exposed to user supplied inputs (by, e.g., carefully +quoting and sanitizing input strings, size-checking input blobs, etc.), it is +very easy to write Python programs which are insecure. Even secure Python code +could be rendered insecure by a bug in the Python interpreter, or in a bug in a +Python library used (e.g., +[this one](https://www.cvedetails.com/cve/CVE-2017-12852/)). + +## Running a TensorFlow server + +TensorFlow is a platform for distributed computing, and as such there is a +TensorFlow server (`tf.train.Server`). **The TensorFlow server is meant for +internal communication only. It is not built for use in an untrusted network.** + +For performance reasons, the default TensorFlow server does not include any +authorization protocol and sends messages unencrypted. It accepts connections +from anywhere, and executes the graphs it is sent without performing any checks. +Therefore, if you run a `tf.train.Server` in your network, anybody with +access to the network can execute what you should consider arbitrary code with +the privileges of the process running the `tf.train.Server`. + +When running distributed TensorFlow, you must isolate the network in which the +cluster lives. Cloud providers provide instructions for setting up isolated +networks, which are sometimes branded as "virtual private cloud." Refer to the +instructions for +[GCP](https://cloud.google.com/compute/docs/networks-and-firewalls) and +[AWS](https://aws.amazon.com/vpc/)) for details. + +Note that `tf.train.Server` is different from the server created by +`tensorflow/serving` (the default binary for which is called `ModelServer`). +By default, `ModelServer` also has no built-in mechanism for authentication. +Connecting it to an untrusted network allows anyone on this network to run the +graphs known to the `ModelServer`. This means that an attacker may run +graphs using untrusted inputs as described above, but they would not be able to +execute arbitrary graphs. It is possible to safely expose a `ModelServer` +directly to an untrusted network, **but only if the graphs it is configured to +use have been carefully audited to be safe**. + +Similar to best practices for other servers, we recommend running any +`ModelServer` with appropriate privileges (i.e., using a separate user with +reduced permisisons). In the spirit of defense in depth, we recommend +authenticating requests to any TensorFlow server connected to an untrusted +network, as well as sandboxing the server to minimize the adverse effects of +any breach. + +## Vulnerabilities in TensorFlow + +TensorFlow is a large and complex system. It also depends on a large set of +third party libraries (e.g., `numpy`, `libjpeg-turbo`, PNG parsers, `protobuf`). +It is possible that TensorFlow or its dependent libraries contain +vulnerabilities that would allow triggering unexpected or dangerous behavior +with specially crafted inputs. + +### What is a vulnerability? + +Given TensorFlow's flexibility, it is possible to specify computation graphs +which exhibit unexpected or unwanted behaviors. The fact that TensorFlow models +can perform arbitrary computations means that they may read and write files, +communicate via the network, produce deadlocks and infinite loops, or run out +of memory. It is only when these behaviors are outside the specifications of the +operations involved that such behavior is a vulnerability. + +A `FileWriter` writing a file is not unexpected behavior and therefore is not a +vulnerability in TensorFlow. A `MatMul` allowing arbitrary binary code execution +**is** a vulnerability. + +This is more subtle from a system perspective. For example, it is easy to cause +a TensorFlow process to try to allocate more memory than available by specifying +a computation graph containing an ill-considered `tf.tile` operation. TensorFlow +should exit cleanly in this case (it would raise an exception in Python, or +return an error `Status` in C++). However, if the surrounding system is not +expecting the possibility, such behavior could be used in a denial of service +attack (or worse). Because TensorFlow behaves correctly, this is not a +vulnerability in TensorFlow (although it would be a vulnerability of this +hypothetical system). + +As a general rule, it is incorrect behavior for Tensorflow to access memory it +does not own, or to terminate in an unclean way. Bugs in TensorFlow that lead to +such behaviors constitute a vulnerability. + +One of the most critical parts of any system is input handling. If malicious +input can trigger side effects or incorrect behavior, this is a bug, and likely +a vulnerability. + +### Reporting vulnerabilities + +Please email reports about any security related issues you find to +`security@tensorflow.org`. This mail is delivered to a small security team. Your +email will be acknowledged within one business day, and you'll receive a more +detailed response to your email within 7 days indicating the next steps in +handling your report. For critical problems, you may encrypt your report (see +below). + +Please use a descriptive subject line for your report email. After the initial +reply to your report, the security team will endeavor to keep you informed of +the progress being made towards a fix and announcement. + +If you believe that an existing (public) issue is security-related, please send +an email to `security@tensorflow.org`. The email should include the issue ID and +a short description of why it should be handled according to this security +policy. + +Once an issue is reported, TensorFlow uses the following disclosure process: + +* When a report is received, we confirm the issue and determine its severity. +* If we know of specific third-party services or software based on TensorFlow + that require mitigation before publication, those projects will be notified. +* An advisory is prepared (but not published) which details the problem and + steps for mitigation. +* Wherever possible, fixes are prepared for the last minor release of the two + latest major releases, as well as the master branch. We will attempt to + commit these fixes as soon as possible, and as close together as + possible. +* Patch releases are published for all fixed released versions, a + notification is sent to discuss@tensorflow.org, and the advisory is published. + +Past security advisories are listed below. We credit reporters for identifying +security issues, although we keep your name confidential if you request it. + +#### Encryption key for `security@tensorflow.org` + +If your disclosure is extremely sensitive, you may choose to encrypt your +report using the key below. Please only use this for critical security +reports. + +``` +-----BEGIN PGP PUBLIC KEY BLOCK----- + +mQENBFpqdzwBCADTeAHLNEe9Vm77AxhmGP+CdjlY84O6DouOCDSq00zFYdIU/7aI +LjYwhEmDEvLnRCYeFGdIHVtW9YrVktqYE9HXVQC7nULU6U6cvkQbwHCdrjaDaylP +aJUXkNrrxibhx9YYdy465CfusAaZ0aM+T9DpcZg98SmsSml/HAiiY4mbg/yNVdPs +SEp/Ui4zdIBNNs6at2gGZrd4qWhdM0MqGJlehqdeUKRICE/mdedXwsWLM8AfEA0e +OeTVhZ+EtYCypiF4fVl/NsqJ/zhBJpCx/1FBI1Uf/lu2TE4eOS1FgmIqb2j4T+jY +e+4C8kGB405PAC0n50YpOrOs6k7fiQDjYmbNABEBAAG0LVRlbnNvckZsb3cgU2Vj +dXJpdHkgPHNlY3VyaXR5QHRlbnNvcmZsb3cub3JnPokBTgQTAQgAOBYhBEkvXzHm +gOJBnwP4Wxnef3wVoM2yBQJaanc8AhsDBQsJCAcCBhUKCQgLAgQWAgMBAh4BAheA +AAoJEBnef3wVoM2yNlkIAICqetv33MD9W6mPAXH3eon+KJoeHQHYOuwWfYkUF6CC +o+X2dlPqBSqMG3bFuTrrcwjr9w1V8HkNuzzOJvCm1CJVKaxMzPuXhBq5+DeT67+a +T/wK1L2R1bF0gs7Pp40W3np8iAFEh8sgqtxXvLGJLGDZ1Lnfdprg3HciqaVAiTum +HBFwszszZZ1wAnKJs5KVteFN7GSSng3qBcj0E0ql2nPGEqCVh+6RG/TU5C8gEsEf +3DX768M4okmFDKTzLNBm+l08kkBFt+P43rNK8dyC4PXk7yJa93SmS/dlK6DZ16Yw +2FS1StiZSVqygTW59rM5XNwdhKVXy2mf/RtNSr84gSi5AQ0EWmp3PAEIALInfBLR +N6fAUGPFj+K3za3PeD0fWDijlC9f4Ety/icwWPkOBdYVBn0atzI21thPRbfuUxfe +zr76xNNrtRRlbDSAChA1J5T86EflowcQor8dNC6fS+oHFCGeUjfEAm16P6mGTo0p +osdG2XnnTHOOEFbEUeWOwR/zT0QRaGGknoy2pc4doWcJptqJIdTl1K8xyBieik/b +nSoClqQdZJa4XA3H9G+F4NmoZGEguC5GGb2P9NHYAJ3MLHBHywZip8g9oojIwda+ +OCLL4UPEZ89cl0EyhXM0nIAmGn3Chdjfu3ebF0SeuToGN8E1goUs3qSE77ZdzIsR +BzZSDFrgmZH+uP0AEQEAAYkBNgQYAQgAIBYhBEkvXzHmgOJBnwP4Wxnef3wVoM2y +BQJaanc8AhsMAAoJEBnef3wVoM2yX4wIALcYZbQhSEzCsTl56UHofze6C3QuFQIH +J4MIKrkTfwiHlCujv7GASGU2Vtis5YEyOoMidUVLlwnebE388MmaJYRm0fhYq6lP +A3vnOCcczy1tbo846bRdv012zdUA+wY+mOITdOoUjAhYulUR0kiA2UdLSfYzbWwy +7Obq96Jb/cPRxk8jKUu2rqC/KDrkFDtAtjdIHh6nbbQhFuaRuWntISZgpIJxd8Bt +Gwi0imUVd9m9wZGuTbDGi6YTNk0GPpX5OMF5hjtM/objzTihSw9UN+65Y/oSQM81 +v//Fw6ZeY+HmRDFdirjD7wXtIuER4vqCryIqR6Xe9X8oJXz9L/Jhslc= +=CDME +-----END PGP PUBLIC KEY BLOCK----- +``` + +### 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) | + diff --git a/tensorflow/__init__.py b/tensorflow/__init__.py index 083634bd7964b0c12e10a1f3c71be5eab597a6c4..78ad6aec19f3bbbfcb389012ac1577573b3e4901 100644 --- a/tensorflow/__init__.py +++ b/tensorflow/__init__.py @@ -21,7 +21,7 @@ from __future__ import division from __future__ import print_function # pylint: disable=wildcard-import -from tensorflow.python import * +from tensorflow.python import * # pylint: disable=redefined-builtin # pylint: enable=wildcard-import from tensorflow.python.util.lazy_loader import LazyLoader diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD index f258bcd95684cc58c2ead3886b3ce74e4af6c5aa..9060c58c1395f07eff0ccef7bd430b3402f8c826 100644 --- a/tensorflow/c/BUILD +++ b/tensorflow/c/BUILD @@ -6,6 +6,7 @@ licenses(["notice"]) # Apache 2.0 load( "//tensorflow:tensorflow.bzl", "tf_cc_test", + "tf_cuda_cc_test", "tf_copts", "tf_cuda_library", "tf_custom_op_library", @@ -26,6 +27,18 @@ filegroup( visibility = ["//tensorflow:__subpackages__"], ) +filegroup( + name = "srcs", + srcs = glob( + [ + "*.cc", + "*.h", + ], + exclude = ["*test*"], + ), + visibility = ["//visibility:public"], +) + tf_cuda_library( name = "c_api_internal", srcs = ["c_api.h"], @@ -79,6 +92,12 @@ tf_cuda_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", ], + }) + select({ + "//tensorflow:with_xla_support": [ + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/jit", + ], + "//conditions:default": [], }), ) @@ -123,15 +142,21 @@ tf_cuda_library( testonly = 1, srcs = ["c_test_util.cc"], hdrs = ["c_test_util.h"], + visibility = [ + "//learning/brain:__subpackages__", + "//tensorflow:__subpackages__", + ], deps = [ ":c_api", + "//tensorflow/compiler/jit/legacy_flags:mark_for_compilation_pass_flags", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:session_options", "//tensorflow/core:test", ], ) -tf_cc_test( +tf_cuda_cc_test( name = "c_api_test", size = "small", srcs = ["c_api_test.cc"], diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc index bc19044fa27549fde7c66ef249a0743136d555bd..85f1d1639b4d09f2de77d326481a86ec246270d0 100644 --- a/tensorflow/c/c_api.cc +++ b/tensorflow/c/c_api.cc @@ -64,6 +64,7 @@ using tensorflow::AllocationDescription; using tensorflow::DataType; using tensorflow::Graph; using tensorflow::GraphDef; +using tensorflow::mutex_lock; using tensorflow::NameRangeMap; using tensorflow::NameRangesForNode; using tensorflow::NewSession; @@ -77,6 +78,7 @@ using tensorflow::RunMetadata; using tensorflow::RunOptions; using tensorflow::Session; using tensorflow::Status; +using tensorflow::string; using tensorflow::Tensor; using tensorflow::TensorBuffer; using tensorflow::TensorId; @@ -87,8 +89,6 @@ using tensorflow::error::Code; using tensorflow::errors::FailedPrecondition; using tensorflow::errors::InvalidArgument; using tensorflow::gtl::ArraySlice; -using tensorflow::mutex_lock; -using tensorflow::string; using tensorflow::strings::StrCat; extern "C" { @@ -109,6 +109,10 @@ TF_Status* TF_NewStatus() { return new TF_Status; } void TF_DeleteStatus(TF_Status* s) { delete s; } void TF_SetStatus(TF_Status* s, TF_Code code, const char* msg) { + if (code == TF_OK) { + s->status = Status::OK(); + return; + } s->status = Status(static_cast(code), tensorflow::StringPiece(msg)); } @@ -195,11 +199,11 @@ TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims, reinterpret_cast(data) % EIGEN_MAX_ALIGN_BYTES != 0) { // TF_STRING and TF_RESOURCE tensors have a different representation in // TF_Tensor than they do in tensorflow::Tensor. So a copy here is a waste - // (any alignement requirements will be taken care of by TF_TensorToTensor + // (any alignment requirements will be taken care of by TF_TensorToTensor // and TF_TensorFromTensor). // - // Other types have the same represntation, so copy only if it is safe to do - // so. + // Other types have the same representation, so copy only if it is safe to + // do so. buf->data_ = allocate_tensor("TF_NewTensor", len); std::memcpy(buf->data_, data, len); buf->deallocator_ = deallocate_buffer; @@ -211,7 +215,13 @@ TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims, buf->deallocator_ = deallocator; buf->deallocator_arg_ = deallocator_arg; } - return new TF_Tensor{dtype, TensorShape(dimvec), buf}; + TF_Tensor* ret = new TF_Tensor{dtype, TensorShape(dimvec), buf}; + size_t elem_size = TF_DataTypeSize(dtype); + if (elem_size > 0 && len < (elem_size * ret->shape.num_elements())) { + delete ret; + return nullptr; + } + return ret; } TF_Tensor* TF_TensorMaybeMove(TF_Tensor* tensor) { @@ -927,6 +937,7 @@ int TF_DeviceListCount(const TF_DeviceList* list) { status->status = InvalidArgument("index out of bounds"); \ return err_val; \ } \ + status->status = Status::OK(); \ return list->response[index].accessor; \ } @@ -1201,6 +1212,13 @@ void TF_SetAttrTypeList(TF_OperationDescription* desc, const char* attr_name, reinterpret_cast(values), num_values)); } +void TF_SetAttrFuncName(TF_OperationDescription* desc, const char* attr_name, + const char* value, size_t length) { + tensorflow::NameAttrList func_name; + func_name.set_name(std::string(value, value + length)); + desc->node_builder.Attr(attr_name, func_name); +} + void TF_SetAttrShape(TF_OperationDescription* desc, const char* attr_name, const int64_t* dims, int num_dims) { PartialTensorShape shape; @@ -1462,7 +1480,13 @@ int TF_OperationOutputConsumers(TF_Output oper_out, TF_Input* consumers, } int TF_OperationNumControlInputs(TF_Operation* oper) { - return oper->node.in_edges().size() - oper->node.num_inputs(); + int count = 0; + for (const auto* edge : oper->node.in_edges()) { + if (edge->IsControlEdge() && !edge->src()->IsSource()) { + ++count; + } + } + return count; } int TF_OperationGetControlInputs(TF_Operation* oper, @@ -1470,7 +1494,7 @@ int TF_OperationGetControlInputs(TF_Operation* oper, int max_control_inputs) { int count = 0; for (const auto* edge : oper->node.in_edges()) { - if (edge->IsControlEdge()) { + if (edge->IsControlEdge() && !edge->src()->IsSource()) { if (count < max_control_inputs) { control_inputs[count] = ToOperation(edge->src()); } @@ -1483,7 +1507,7 @@ int TF_OperationGetControlInputs(TF_Operation* oper, int TF_OperationNumControlOutputs(TF_Operation* oper) { int count = 0; for (const auto* edge : oper->node.out_edges()) { - if (edge->IsControlEdge()) { + if (edge->IsControlEdge() && !edge->dst()->IsSink()) { ++count; } } @@ -1495,7 +1519,7 @@ int TF_OperationGetControlOutputs(TF_Operation* oper, int max_control_outputs) { int count = 0; for (const auto* edge : oper->node.out_edges()) { - if (edge->IsControlEdge()) { + if (edge->IsControlEdge() && !edge->dst()->IsSink()) { if (count < max_control_outputs) { control_outputs[count] = ToOperation(edge->dst()); } @@ -2130,7 +2154,7 @@ Status CopyGraph(Graph* src_graph, Graph* dst_graph, opts.return_tensors.push_back(ToTensorId(nodes_to_return[i])); } - // TOOD(skyewm): change to OutputTensor + // TODO(skyewm): change to OutputTensor tensorflow::ImportGraphDefResults results; TF_RETURN_IF_ERROR( ImportGraphDef(opts, gdef, dst_graph, dst_refiner, &results)); diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h index 6f1c0606c11c0e102a8b5dcab4ee61fcc36d8b16..ad592ef70961ef427bfe9fd322a82bd64df7f9f1 100644 --- a/tensorflow/c/c_api.h +++ b/tensorflow/c/c_api.h @@ -226,6 +226,10 @@ typedef struct TF_Tensor TF_Tensor; // (*deallocator)(data, len, deallocator_arg) // Clients must provide a custom deallocator function so they can pass in // memory managed by something like numpy. +// +// May return NULL (and invoke the deallocator) if the provided data buffer +// (data, len) is inconsistent with a tensor of the given TF_DataType +// and the shape specified by (dima, num_dims). TF_CAPI_EXPORT extern TF_Tensor* TF_NewTensor( TF_DataType, const int64_t* dims, int num_dims, void* data, size_t len, void (*deallocator)(void* data, size_t len, void* arg), @@ -511,6 +515,11 @@ TF_CAPI_EXPORT extern void TF_SetAttrTypeList(TF_OperationDescription* desc, const char* attr_name, const TF_DataType* values, int num_values); +// Set a 'func' attribute to the specified name. +// `value` must point to a string of length `length` bytes. +TF_CAPI_EXPORT extern void TF_SetAttrFuncName(TF_OperationDescription* desc, + const char* attr_name, + const char* value, size_t length); // Set `num_dims` to -1 to represent "unknown rank". Otherwise, // `dims` points to an array of length `num_dims`. `dims[i]` must be @@ -1029,6 +1038,23 @@ TF_CAPI_EXPORT extern void TF_GraphCopyFunction(TF_Graph* g, const TF_Function* grad, TF_Status* status); +// Returns the number of TF_Functions registered in `g`. +TF_CAPI_EXPORT extern int TF_GraphNumFunctions(TF_Graph* g); + +// Fills in `funcs` with the TF_Function* registered in `g`. +// `funcs` must point to an array of TF_Function* of length at least +// `max_func`. In usual usage, max_func should be set to the result of +// TF_GraphNumFunctions(g). In this case, all the functions registered in +// `g` will be returned. Else, an unspecified subset. +// +// If successful, returns the number of TF_Function* successfully set in +// `funcs` and sets status to OK. The caller takes ownership of +// all the returned TF_Functions. They must be deleted with TF_DeleteFunction. +// On error, returns 0, sets status to the encountered error, and the contents +// of funcs will be undefined. +TF_CAPI_EXPORT extern int TF_GraphGetFunctions(TF_Graph* g, TF_Function** funcs, + int max_func, TF_Status* status); + // Note: The following function may fail on very large protos in the future. TF_CAPI_EXPORT extern void TF_OperationToNodeDef(TF_Operation* oper, @@ -1261,11 +1287,12 @@ TF_CAPI_EXPORT extern void TF_DeleteFunction(TF_Function* func); typedef struct TF_Session TF_Session; -// Return a new execution session with the associated graph, or NULL on error. +// Return a new execution session with the associated graph, or NULL on +// error. Does not take ownership of any input parameters. // -// *graph must be a valid graph (not deleted or nullptr). This function will -// prevent the graph from being deleted until TF_DeleteSession() is called. -// Does not take ownership of opts. +// *`graph` must be a valid graph (not deleted or nullptr). `graph` will be be +// kept alive for the lifetime of the returned TF_Session. New nodes can still +// be added to `graph` after this call. TF_CAPI_EXPORT extern TF_Session* TF_NewSession(TF_Graph* graph, const TF_SessionOptions* opts, TF_Status* status); diff --git a/tensorflow/c/c_api_function.cc b/tensorflow/c/c_api_function.cc index d60d1de315ed37a327bd036ddb914a3c32413f65..384e6c8cb97022264c5327da5ca5861057608fbe 100644 --- a/tensorflow/c/c_api_function.cc +++ b/tensorflow/c/c_api_function.cc @@ -44,8 +44,12 @@ class NodeNameMapping { public: NodeNameMapping() = default; - // Normalize the input/output name and make it unique. - string GetIOName(const string& name); + // Normalize the input name and make it unique. This is the same as the + // function for output, expect that it adds a name mapping for the name. + string GetInputName(const string& name); + + // Normalize the output name and make it unique. + string GetOutputName(const string& name); // Make the node name unique. string Uniquify(const string& name); @@ -107,7 +111,13 @@ string NodeNameMapping::UniquifyHelper(const string& name) const { } } -string NodeNameMapping::GetIOName(const string& name) { +string NodeNameMapping::GetInputName(const string& name) { + const string& input_name = GetOutputName(name); + name_mapping_[name] = input_name; + return input_name; +} + +string NodeNameMapping::GetOutputName(const string& name) { const string& input_name = UniquifyHelper(Normalize(name)); // Record that we used this name, but don't add it to name_mapping_ // since this name is not for a node. @@ -214,10 +224,11 @@ Status FillFunctionBody( // Add control inputs. for (const Edge* edge : control_edges) { - // Add this control input only if the src node is in the body. + // Add this control input only if the src node is in the body or a part of + // the inputs. const string normalized = node_names.Lookup(edge->src()->name()); // If we did not find a name for the source of control edge, this - // source must be outside of the body. Raise an error. + // source must be outside of the body, and not an input. Raise an error. if (normalized.empty()) { return InvalidArgument( "The source of control edge ", edge->DebugString(), @@ -279,7 +290,7 @@ Status GraphToFunctionDef(const Graph& fn_body, const string& fn_name, TF_RETURN_IF_ERROR(node_names.UseOutputName(output_names[i])); argdef->set_name(output_names[i]); } else { - argdef->set_name(node_names.GetIOName(node->name())); + argdef->set_name(node_names.GetOutputName(node->name())); } } @@ -289,7 +300,7 @@ Status GraphToFunctionDef(const Graph& fn_body, const string& fn_name, int idx = inputs[i].index; OpDef::ArgDef* argdef = fdef->mutable_signature()->add_input_arg(); argdef->set_type(node->output_type(idx)); - const string& input_name = node_names.GetIOName(node->name()); + const string& input_name = node_names.GetInputName(node->name()); argdef->set_name(input_name); tensor_renaming[strings::StrCat(node->name(), ":", idx)] = input_name; } @@ -312,7 +323,7 @@ Status GraphToFunctionDef(const Graph& fn_body, const string& fn_name, TF_RETURN_IF_ERROR( NameRangesForNode(*node, node->op_def(), nullptr, &output_ranges)); for (const auto& output : output_ranges) { - const string& output_name = output.first; + const StringPiece& output_name = output.first; int index_start = output.second.first; int index_end = output.second.second; for (int i = index_start; i < index_end; ++i) { @@ -467,7 +478,7 @@ Status ComputeBodyNodes( return Status::OK(); } -} // anonymous namespace +} // namespace } // namespace tensorflow using tensorflow::Node; @@ -548,6 +559,28 @@ void TF_GraphCopyFunction(TF_Graph* g, const TF_Function* func, status->status = g->graph.AddFunctionLibrary(fdef_lib); } +int TF_GraphNumFunctions(TF_Graph* g) { + tensorflow::mutex_lock l(g->mu); + return g->graph.flib_def().num_functions(); +} + +int TF_GraphGetFunctions(TF_Graph* g, TF_Function** funcs, int max_func, + TF_Status* status) { + tensorflow::FunctionDefLibrary lib; + { + tensorflow::mutex_lock l(g->mu); + lib = g->graph.flib_def().ToProto(); + } + const auto len = std::min(max_func, static_cast(lib.function_size())); + for (int i = 0; i < len; ++i) { + TF_Function* func = new TF_Function(); + func->fdef = lib.function(i); + funcs[i] = func; + } + status->status = tensorflow::Status::OK(); + return len; +} + void TF_FunctionToFunctionDef(TF_Function* func, TF_Buffer* output_func_def, TF_Status* status) { status->status = MessageToBuffer(func->fdef, output_func_def); diff --git a/tensorflow/c/c_api_function_test.cc b/tensorflow/c/c_api_function_test.cc index 2e2293ca85175009d1bcc8db5c830789e4701c1d..7ca50119eafe299b307f06c555aec1388e7e82e2 100644 --- a/tensorflow/c/c_api_function_test.cc +++ b/tensorflow/c/c_api_function_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include "tensorflow/core/framework/function.pb.h" #include "tensorflow/core/framework/op_def.pb.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" @@ -330,6 +331,11 @@ class CApiFunctionTest : public ::testing::Test { << "Failed to find expected edge " << e.ToString() << " in fdef: " << fdef.DebugString(); } + for (const EdgeSpec& e : c_edges) { + ASSERT_TRUE(a_edges.find(e) != a_edges.end()) + << "Failed to find expected control edge " << e.ToString() + << " in fdef: " << fdef.DebugString(); + } // If caller specified all edges, check that we have seen all if (is_exact_edges) { @@ -979,7 +985,7 @@ TEST_F(CApiFunctionTest, ControlDependency) { VerifyFDef( {"add_0", "scalar"}, M({{"feed1"}, {"feed2"}}), M({{"add"}}), {{"feed1", "add_0:0"}, {"feed2", "add_0:1"}, {"add_0:sum:0", "add"}}, - {{"scalar", "add_0"}}); + {{"^scalar", "add_0:2"}}); } TEST_F(CApiFunctionTest, ControlDependencyOutsideOfBody) { @@ -1022,12 +1028,17 @@ TEST_F(CApiFunctionTest, ControlDependencyOutsideOfBody_FromInputNode) { TF_Operation* add = AddWithCtrlDependency(feed1, feed2, func_graph_, feed1, s_); EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - Define(-1, {}, {feed1, feed2}, {add}, {}, true); - EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); - EXPECT_EQ(string("The source of control edge [id=3 feed1:-1 -> add:-1] " - "is not in the body. Encountered while creating " - "function 'MyFunc'"), - string(TF_Message(s_))); + Define(-1, {}, {feed1, feed2}, {add}, {}); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, func_op, 2 + 3); + VerifyFDef( + {"add_0"}, M({{"feed1"}, {"feed2"}}), M({{"add"}}), + {{"feed1", "add_0:0"}, {"feed2", "add_0:1"}, {"add_0:sum:0", "add"}}, + {{"^feed1", "add_0:2"}}); } TEST_F(CApiFunctionTest, DuplicateInputsAreNotAllowed) { @@ -1462,7 +1473,11 @@ TEST_F(CApiFunctionTest, AppendHash) { /*append_hash=*/true); tensorflow::FunctionDef fdef; ASSERT_TRUE(GetFunctionDef(func_, &fdef)); +#if (__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) + ASSERT_EQ(string("func_name_base_ZpgUD4x8oqk"), fdef.signature().name()); +#else ASSERT_EQ(string("func_name_base_qaJ8jA8UmGY"), fdef.signature().name()); +#endif } TEST_F(CApiFunctionTest, GetOpDef) { @@ -1531,5 +1546,75 @@ TEST_F(CApiFunctionTest, StatefulOpDef) { TF_DeleteBuffer(buffer); } +void AssertEqual(TF_Function* f1, TF_Function* f2) { + string s1, s2; + tensorflow::FunctionDef fdef1, fdef2; + ASSERT_TRUE(GetFunctionDef(f1, &fdef1)); + ASSERT_TRUE(GetFunctionDef(f2, &fdef2)); + SerializeToStringDeterministic(fdef1, &s1); + SerializeToStringDeterministic(fdef2, &s2); + ASSERT_EQ(s1, s2); +} + +string GetName(TF_Function* func) { + tensorflow::FunctionDef fdef; + GetFunctionDef(func, &fdef); + return fdef.signature().name(); +} + +TEST_F(CApiFunctionTest, GetFunctionsFromGraph) { + TF_Function* funcs[2]; + + // Get functions from empty graph + EXPECT_EQ(TF_GraphNumFunctions(host_graph_), 0); + TF_GraphGetFunctions(host_graph_, nullptr, 0, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + // Define a function and add it to host_graph_ + TF_Function* func0; + DefineFunction("FooFunc0", &func0); + TF_GraphCopyFunction(host_graph_, func0, nullptr, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + // Get this function from host_graph_ + EXPECT_EQ(TF_GraphNumFunctions(host_graph_), 1); + EXPECT_EQ(TF_GraphGetFunctions(host_graph_, funcs, 0, s_), 0); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + EXPECT_EQ(TF_GraphGetFunctions(host_graph_, funcs, 1, s_), 1); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + AssertEqual(func0, funcs[0]); + TF_DeleteFunction(funcs[0]); + EXPECT_EQ(TF_GraphGetFunctions(host_graph_, funcs, 2, s_), 1); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + AssertEqual(func0, funcs[0]); + TF_DeleteFunction(funcs[0]); + + // Define a second function + TF_Function* func1; + DefineFunction("FooFunc1", &func1); + TF_GraphCopyFunction(host_graph_, func1, nullptr, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + // Get both function from host_graph_ + EXPECT_EQ(TF_GraphNumFunctions(host_graph_), 2); + EXPECT_EQ(TF_GraphGetFunctions(host_graph_, funcs, 0, s_), 0); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + EXPECT_EQ(TF_GraphGetFunctions(host_graph_, funcs, 2, s_), 2); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + if (GetName(funcs[0]) == GetName(func0)) { + AssertEqual(func0, funcs[0]); + AssertEqual(func1, funcs[1]); + } else { + AssertEqual(func0, funcs[1]); + AssertEqual(func1, funcs[0]); + } + + TF_DeleteFunction(funcs[0]); + TF_DeleteFunction(funcs[1]); + + TF_DeleteFunction(func0); + TF_DeleteFunction(func1); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/c/c_api_test.cc b/tensorflow/c/c_api_test.cc index df697e16d3d3fcaac66f967c0d3938450f0b0be6..028f146be31790b211e546978302e81afe26b231 100644 --- a/tensorflow/c/c_api_test.cc +++ b/tensorflow/c/c_api_test.cc @@ -57,6 +57,52 @@ static void ExpectHasSubstr(StringPiece s, StringPiece expected) { << "'" << s << "' does not contain '" << expected << "'"; } +// Returns the GPU device name if there is one (with arbitrary tie breaking if +// there are more than one), or "" otherwise. +string GPUDeviceName(TF_Session* session) { + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + TF_Status* s = status.get(); + std::unique_ptr list( + TF_SessionListDevices(session, s), TF_DeleteDeviceList); + TF_DeviceList* device_list = list.get(); + + CHECK_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + + const int num_devices = TF_DeviceListCount(device_list); + LOG(INFO) << "There are " << num_devices << " devices."; + for (int i = 0; i < num_devices; ++i) { + const char* device_name = TF_DeviceListName(device_list, i, s); + CHECK_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + const char* device_type = TF_DeviceListType(device_list, i, s); + CHECK_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + LOG(INFO) << "Device " << i << " has name " << device_name << ", type " + << device_type; + if (string(device_type) == DEVICE_GPU) { + return device_name; + } + } + // No GPU device found. + return ""; +} + +string GPUDeviceName() { + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + TF_Status* s = status.get(); + std::unique_ptr graph(TF_NewGraph(), + TF_DeleteGraph); + + TF_SessionOptions* opts = TF_NewSessionOptions(); + TF_Session* sess = TF_NewSession(graph.get(), opts, s); + TF_DeleteSessionOptions(opts); + + const string gpu_device_name = GPUDeviceName(sess); + TF_DeleteSession(sess, s); + CHECK_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + return gpu_device_name; +} + TEST(CAPI, Version) { EXPECT_STRNE("", TF_Version()); } TEST(CAPI, Status) { @@ -94,6 +140,17 @@ TEST(CAPI, Tensor) { EXPECT_TRUE(deallocator_called); } +void NoOpDeallocator(void* data, size_t, void*) {} + +TEST(CAPI, MalformedTensor) { + // See https://github.com/tensorflow/tensorflow/issues/7394 + // num_dims = 0 implies a scalar, so should be backed by at least 4 bytes of + // data. + TF_Tensor* t = + TF_NewTensor(TF_FLOAT, nullptr, 0, nullptr, 0, &NoOpDeallocator, nullptr); + ASSERT_TRUE(t == nullptr); +} + TEST(CAPI, AllocateTensor) { const int num_bytes = 6 * sizeof(float); int64_t dims[] = {2, 3}; @@ -123,6 +180,10 @@ TEST(CAPI, MaybeMove) { } TEST(CAPI, LibraryLoadFunctions) { + // TODO(b/73318067): Fix linking for the GPU test generated by the + // tf_cuda_cc_test() bazel rule and remove the next line. + if (!GPUDeviceName().empty()) return; + // Load the library. TF_Status* status = TF_NewStatus(); TF_Library* lib = @@ -575,7 +636,7 @@ TEST(CAPI, ImportGraphDef) { TF_Status* s = TF_NewStatus(); TF_Graph* graph = TF_NewGraph(); - // Create a graph with two nodes: x and 3 + // Create a simple graph. Placeholder(graph, s); ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); ASSERT_TRUE(TF_GraphOperationByName(graph, "feed") != nullptr); @@ -586,7 +647,7 @@ TEST(CAPI, ImportGraphDef) { ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); ASSERT_TRUE(TF_GraphOperationByName(graph, "neg") != nullptr); - // Export to a GraphDef + // Export to a GraphDef. TF_Buffer* graph_def = TF_NewBuffer(); TF_GraphToGraphDef(graph, graph_def, s); ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); @@ -606,6 +667,31 @@ TEST(CAPI, ImportGraphDef) { ASSERT_TRUE(feed != nullptr); ASSERT_TRUE(neg != nullptr); + // Test basic structure of the imported graph. + EXPECT_EQ(0, TF_OperationNumInputs(scalar)); + EXPECT_EQ(0, TF_OperationNumInputs(feed)); + ASSERT_EQ(1, TF_OperationNumInputs(neg)); + TF_Output neg_input = TF_OperationInput({neg, 0}); + EXPECT_EQ(scalar, neg_input.oper); + EXPECT_EQ(0, neg_input.index); + + // Test that we can't see control edges involving the source and sink nodes. + TF_Operation* control_ops[100]; + EXPECT_EQ(0, TF_OperationNumControlInputs(scalar)); + EXPECT_EQ(0, TF_OperationGetControlInputs(scalar, control_ops, 100)); + EXPECT_EQ(0, TF_OperationNumControlOutputs(scalar)); + EXPECT_EQ(0, TF_OperationGetControlOutputs(scalar, control_ops, 100)); + + EXPECT_EQ(0, TF_OperationNumControlInputs(feed)); + EXPECT_EQ(0, TF_OperationGetControlInputs(feed, control_ops, 100)); + EXPECT_EQ(0, TF_OperationNumControlOutputs(feed)); + EXPECT_EQ(0, TF_OperationGetControlOutputs(feed, control_ops, 100)); + + EXPECT_EQ(0, TF_OperationNumControlInputs(neg)); + EXPECT_EQ(0, TF_OperationGetControlInputs(neg, control_ops, 100)); + EXPECT_EQ(0, TF_OperationNumControlOutputs(neg)); + EXPECT_EQ(0, TF_OperationGetControlOutputs(neg, control_ops, 100)); + // Import it again, with an input mapping, return outputs, and a return // operation, into the same graph. TF_DeleteImportGraphDefOptions(opts); @@ -629,7 +715,7 @@ TEST(CAPI, ImportGraphDef) { ASSERT_TRUE(neg2 != nullptr); // Check input mapping - TF_Output neg_input = TF_OperationInput({neg, 0}); + neg_input = TF_OperationInput({neg, 0}); EXPECT_EQ(scalar, neg_input.oper); EXPECT_EQ(0, neg_input.index); @@ -887,6 +973,70 @@ TEST(CAPI, Session) { TF_DeleteStatus(s); } +// If `device` is non-empty, run Min op on that device. +// Otherwise run it on the default device (CPU). +void RunMinTest(const string& device, bool use_XLA) { + TF_Status* s = TF_NewStatus(); + TF_Graph* graph = TF_NewGraph(); + + // Make a placeholder operation. + TF_Operation* feed = Placeholder(graph, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + + // Make a constant operation with the scalar "0", for axis. + TF_Operation* one = ScalarConst(0, graph, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + + // Create a session for this graph. + CSession csession(graph, s, use_XLA); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + + if (!device.empty()) { + LOG(INFO) << "Setting op Min on device " << device; + } + TF_Operation* min = MinWithDevice(feed, one, graph, device, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + + // Run the graph. + csession.SetInputs({{feed, Int32Tensor({3, 2, 5})}}); + csession.SetOutputs({min}); + csession.Run(s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Tensor* out = csession.output_tensor(0); + ASSERT_TRUE(out != nullptr); + EXPECT_EQ(TF_INT32, TF_TensorType(out)); + EXPECT_EQ(0, TF_NumDims(out)); // scalar + ASSERT_EQ(sizeof(int32), TF_TensorByteSize(out)); + int32* output_contents = static_cast(TF_TensorData(out)); + EXPECT_EQ(2, *output_contents); + + // Clean up + csession.CloseAndDelete(s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_DeleteGraph(graph); + TF_DeleteStatus(s); +} + +TEST(CAPI, Session_Min_CPU) { RunMinTest(/*device=*/"", /*use_XLA=*/false); } + +TEST(CAPI, Session_Min_XLA_CPU) { RunMinTest(/*device=*/"", /*use_XLA=*/true); } + +TEST(CAPI, Session_Min_GPU) { + const string gpu_device = GPUDeviceName(); + // Skip this test if no GPU is available. + if (gpu_device.empty()) return; + + RunMinTest(gpu_device, /*use_XLA=*/false); +} + +TEST(CAPI, Session_Min_XLA_GPU) { + const string gpu_device = GPUDeviceName(); + // Skip this test if no GPU is available. + if (gpu_device.empty()) return; + + RunMinTest(gpu_device, /*use_XLA=*/true); +} + TEST(CAPI, SessionPRun) { TF_Status* s = TF_NewStatus(); TF_Graph* graph = TF_NewGraph(); @@ -1931,7 +2081,7 @@ TEST_F(CApiAttributesTest, Tensor) { } TEST_F(CApiAttributesTest, StringTensor) { - // Create the string-Tensor "atttribute" value. + // Create the string-Tensor "attribute" value. char encoded[] = { 0, 0, 0, 0, 0, 0, 0, 0, // array[uint64] offsets 1, // varint encoded string length @@ -2029,6 +2179,10 @@ TEST_F(CApiAttributesTest, Errors) { } TEST(TestApiDef, TestCreateApiDef) { + // TODO(b/73318067): Fix linking for the GPU test generated by the + // tf_cuda_cc_test() bazel rule and remove the next line. + if (!GPUDeviceName().empty()) return; + TF_Status* status = TF_NewStatus(); TF_Library* lib = TF_LoadLibrary("tensorflow/c/test_op.so", status); @@ -2059,6 +2213,10 @@ TEST(TestApiDef, TestCreateApiDef) { } TEST(TestApiDef, TestCreateApiDefWithOverwrites) { + // TODO(b/73318067): Fix linking for the GPU test generated by the + // tf_cuda_cc_test() bazel rule and remove the next line. + if (!GPUDeviceName().empty()) return; + TF_Status* status = TF_NewStatus(); TF_Library* lib = TF_LoadLibrary("tensorflow/c/test_op.so", status); diff --git a/tensorflow/c/c_test_util.cc b/tensorflow/c/c_test_util.cc index 37439ff0beac5a5220460465e954b6c093ee1ba9..a55af46ae2baef1cd4f55f478ec234551f370503 100644 --- a/tensorflow/c/c_test_util.cc +++ b/tensorflow/c/c_test_util.cc @@ -15,11 +15,13 @@ limitations under the License. #include "tensorflow/c/c_test_util.h" +#include "tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h" #include "tensorflow/core/framework/function.pb.h" #include "tensorflow/core/framework/op_def.pb.h" #include "tensorflow/core/framework/tensor.pb.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/public/session_options.h" using tensorflow::GraphDef; using tensorflow::NodeDef; @@ -124,8 +126,9 @@ TF_Operation* ScalarConst(double v, TF_Graph* graph, TF_Status* s, return Const(tensor.get(), graph, s, name); } -void AddHelper(TF_Operation* l, TF_Operation* r, TF_Graph* graph, TF_Status* s, - const char* name, TF_Operation** op, bool check) { +void AddOpHelper(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name, TF_Operation** op, + bool check) { TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name); TF_Output add_inputs[2] = {{l, 0}, {r, 0}}; TF_AddInputList(desc, add_inputs, 2); @@ -139,14 +142,14 @@ void AddHelper(TF_Operation* l, TF_Operation* r, TF_Graph* graph, TF_Status* s, TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph, TF_Status* s, const char* name) { TF_Operation* op; - AddHelper(l, r, graph, s, name, &op, true); + AddOpHelper(l, r, graph, s, name, &op, true); return op; } TF_Operation* AddNoCheck(TF_Operation* l, TF_Operation* r, TF_Graph* graph, TF_Status* s, const char* name) { TF_Operation* op; - AddHelper(l, r, graph, s, name, &op, false); + AddOpHelper(l, r, graph, s, name, &op, false); return op; } @@ -160,6 +163,36 @@ TF_Operation* AddWithCtrlDependency(TF_Operation* l, TF_Operation* r, return TF_FinishOperation(desc, s); } +// If `op_device` is non-empty, set the created op on that device. +void BinaryOpHelper(const char* op_name, TF_Operation* l, TF_Operation* r, + TF_Graph* graph, TF_Status* s, const char* name, + TF_Operation** op, const string& op_device, bool check) { + TF_OperationDescription* desc = TF_NewOperation(graph, op_name, name); + if (!op_device.empty()) { + TF_SetDevice(desc, op_device.c_str()); + } + TF_AddInput(desc, {l, 0}); + TF_AddInput(desc, {r, 0}); + *op = TF_FinishOperation(desc, s); + if (check) { + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); + } +} + +TF_Operation* MinWithDevice(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + const string& op_device, TF_Status* s, + const char* name) { + TF_Operation* op; + BinaryOpHelper("Min", l, r, graph, s, name, &op, op_device, true); + return op; +} + +TF_Operation* Min(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name) { + return MinWithDevice(l, r, graph, /*op_device=*/"", s, name); +} + TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s, const char* name) { TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name); @@ -369,8 +402,21 @@ std::vector GetFuncNames(const tensorflow::GraphDef& graph_def) { return names; } -CSession::CSession(TF_Graph* graph, TF_Status* s) { +CSession::CSession(TF_Graph* graph, TF_Status* s, bool use_XLA) { TF_SessionOptions* opts = TF_NewSessionOptions(); + tensorflow::legacy_flags::MarkForCompilationPassFlags* flags = + tensorflow::legacy_flags::GetMarkForCompilationPassFlags(); + flags->tf_xla_cpu_global_jit = use_XLA; + if (use_XLA) { + tensorflow::ConfigProto config; + config.mutable_graph_options() + ->mutable_optimizer_options() + ->set_global_jit_level(tensorflow::OptimizerOptions::ON_1); + std::string contents; + contents.resize(config.ByteSizeLong()); + config.SerializeToArray(&contents[0], contents.size()); + TF_SetConfig(opts, contents.data(), contents.size(), s); + } session_ = TF_NewSession(graph, opts, s); TF_DeleteSessionOptions(opts); } diff --git a/tensorflow/c/c_test_util.h b/tensorflow/c/c_test_util.h index 3429009a71a863ae6b69b5cd29ace3c7fd078f4c..2a70177c724c569844a5d8ad42b99bed20209946 100644 --- a/tensorflow/c/c_test_util.h +++ b/tensorflow/c/c_test_util.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_C_C_TEST_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_C_C_TEST_UTIL_H_ +#ifndef TENSORFLOW_C_C_TEST_UTIL_H_ +#define TENSORFLOW_C_C_TEST_UTIL_H_ #include "tensorflow/c/c_api.h" @@ -69,6 +69,14 @@ TF_Operation* AddWithCtrlDependency(TF_Operation* l, TF_Operation* r, TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s, const char* name = "add"); +TF_Operation* Min(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name = "min"); + +// If `op_device` is non-empty, set the created op on that device. +TF_Operation* MinWithDevice(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + const string& op_device, TF_Status* s, + const char* name = "min"); + TF_Operation* Neg(TF_Operation* n, TF_Graph* graph, TF_Status* s, const char* name = "neg"); @@ -108,7 +116,7 @@ std::vector GetFuncNames(const tensorflow::GraphDef& graph_def); class CSession { public: - CSession(TF_Graph* graph, TF_Status* s); + CSession(TF_Graph* graph, TF_Status* s, bool use_XLA = false); explicit CSession(TF_Session* session); ~CSession(); @@ -124,6 +132,8 @@ class CSession { TF_Tensor* output_tensor(int i) { return output_values_[i]; } + TF_Session* mutable_session() { return session_; } + private: void DeleteInputValues(); void ResetOutputValues(); @@ -136,4 +146,4 @@ class CSession { std::vector targets_; }; -#endif // THIRD_PARTY_TENSORFLOW_C_C_TEST_UTIL_H_ +#endif // TENSORFLOW_C_C_TEST_UTIL_H_ diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD index 53df884e7ca3765541423d3340f35eae1779aee8..e55cb672e97e1403a3dd864c91c176426eb3f067 100644 --- a/tensorflow/c/eager/BUILD +++ b/tensorflow/c/eager/BUILD @@ -6,6 +6,7 @@ load( "tf_cuda_cc_test", "tf_cc_test", "tf_copts", + "tfe_xla_copts", "tf_cuda_library", ) @@ -16,7 +17,7 @@ tf_cuda_library( "c_api_internal.h", ], hdrs = ["c_api.h"], - copts = tf_copts(), + copts = tf_copts() + tfe_xla_copts(), visibility = ["//visibility:public"], deps = select({ "//tensorflow:android": [ @@ -33,7 +34,15 @@ tf_cuda_library( "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", ], - }) + ["//tensorflow/core:gpu_runtime"], + }) + select({ + "//tensorflow:with_xla_support": [ + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/jit", + ], + "//conditions:default": [], + }) + [ + "//tensorflow/core:gpu_runtime", + ], ) tf_cuda_library( @@ -46,6 +55,7 @@ tf_cuda_library( "//tensorflow/c:c_api", "//tensorflow/c:c_api_internal", "//tensorflow/core:core_cpu_lib", + "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:framework_lite", "//tensorflow/core:lib_internal", @@ -55,12 +65,14 @@ tf_cuda_library( tf_cuda_cc_test( name = "c_api_test", srcs = ["c_api_test.cc"], + extra_copts = tfe_xla_copts(), tags = [ "guitar", "multi_gpu", ], deps = [ ":c_api", + "//tensorflow/c:c_test_util", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", @@ -117,3 +129,9 @@ cc_library( "//tensorflow/core:lib", ], ) + +filegroup( + name = "headers", + srcs = ["c_api.h"], + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index 589afb9031b5fd9e3d9a6b5172ba3e9bccd1fe97..8e834eb99c13d1f26da9f0860897267efc2fd01c 100644 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -25,6 +25,9 @@ limitations under the License. #include "tensorflow/c/c_api_internal.h" #include "tensorflow/c/eager/c_api_internal.h" #include "tensorflow/c/eager/runtime.h" +#ifdef TENSORFLOW_EAGER_USE_XLA +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#endif // TENSORFLOW_EAGER_USE_XLA #include "tensorflow/core/common_runtime/copy_tensor.h" #include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/common_runtime/device_mgr.h" @@ -34,6 +37,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_shape.pb.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/core/refcount.h" +#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/map_util.h" #include "tensorflow/core/lib/gtl/stl_util.h" #include "tensorflow/core/platform/mutex.h" @@ -44,13 +48,23 @@ using tensorflow::int64; using tensorflow::string; namespace { -bool IsCPU(tensorflow::Device* d) { +bool IsCPU(const tensorflow::Device* d) { return d == nullptr || d->tensorflow_gpu_device_info() == nullptr; } -string DeviceName(tensorflow::Device* d) { +bool IsXLA(const tensorflow::Device* d) { + if (d == nullptr) return false; + const auto& device_type = d->attributes().device_type(); + return device_type.find("XLA") != std::string::npos; +} + +string DeviceName(const tensorflow::Device* d) { return (d == nullptr) ? "cpu:0" : d->name(); } + +#ifdef TENSORFLOW_EAGER_USE_XLA +std::atomic_int_fast64_t func_id_generator(0); +#endif // TENSORFLOW_EAGER_USE_XLA } // namespace extern "C" { @@ -85,20 +99,15 @@ TFE_Context* TFE_NewContext(const TFE_ContextOptions* opts, TF_Status* status) { return nullptr; } - TFE_Context* ret = new TFE_Context(session); - ret->policy = opts->policy; - ret->pflr.reset(new tensorflow::ProcessFunctionLibraryRuntime( - ret->session->device_mgr, opts->session_options.options.env, - TF_GRAPH_DEF_VERSION, &ret->func_lib_def, {})); - ret->rendezvous = - new tensorflow::IntraProcessRendezvous(ret->session->device_mgr); - - return ret; + return new TFE_Context(*opts, session); } void TFE_DeleteContext(TFE_Context* ctx, TF_Status* status) { status->status = tensorflow::Status::OK(); - tensorflow::gtl::STLDeleteValues(&ctx->kernel_cache); + { + tensorflow::mutex_lock ml(ctx->cache_mu); + tensorflow::gtl::STLDeleteValues(&ctx->kernel_cache); + } TF_Graph* graph = ctx->session->graph; TF_DeleteSession(ctx->session, status); TF_DeleteGraph(graph); @@ -110,6 +119,28 @@ TF_DeviceList* TFE_ContextListDevices(TFE_Context* ctx, TF_Status* status) { return TF_SessionListDevices(ctx->session, status); } +void TFE_ContextClearCaches(TFE_Context* ctx) { + tensorflow::mutex_lock ml(ctx->cache_mu); + tensorflow::gtl::STLDeleteValues(&ctx->kernel_cache); +} + +void TFE_ContextSetThreadLocalDevicePlacementPolicy( + TFE_Context* ctx, TFE_ContextDevicePlacementPolicy policy) { + tensorflow::mutex_lock ml(ctx->policy_map_mu); + ctx->thread_local_policies[std::this_thread::get_id()] = policy; +} + +extern TFE_ContextDevicePlacementPolicy TFE_ContextGetDevicePlacementPolicy( + TFE_Context* ctx) { + tensorflow::mutex_lock ml(ctx->policy_map_mu); + auto policy_map_it = + ctx->thread_local_policies.find(std::this_thread::get_id()); + if (policy_map_it != ctx->thread_local_policies.end()) { + return policy_map_it->second; + } + return ctx->policy; +} + TFE_TensorHandle* TFE_NewTensorHandle(TF_Tensor* t, TF_Status* status) { tensorflow::Tensor tensor; status->status = tensorflow::TF_TensorToTensor(t, &tensor); @@ -165,11 +196,16 @@ TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, bool is_same_device = (srcd == dstd) || (DeviceName(srcd) == DeviceName(dstd)); const bool dst_cpu = IsCPU(dstd); - if (is_same_device) { + const bool src_cpu = IsCPU(srcd); + // both_on_cpu can be true and yet is_same_device is false, if one of src/dst + // has device type XLA_CPU, and the other CPU. + const bool both_on_cpu = src_cpu && dst_cpu; + if (is_same_device || both_on_cpu) { return new TFE_TensorHandle(h->t, dst_cpu ? nullptr : dstd); } tensorflow::Tensor* src = &(h->t); - if (!dst_cpu && !tensorflow::DataTypeCanUseMemcpy(src->dtype())) { + if (!dst_cpu && (src->dtype() != tensorflow::DT_VARIANT && + !tensorflow::DataTypeCanUseMemcpy(src->dtype()))) { TF_SetStatus( status, TF_INVALID_ARGUMENT, tensorflow::strings::StrCat("Can't copy Tensor with type ", @@ -178,13 +214,16 @@ TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, .c_str()); return nullptr; } - tensorflow::Tensor dst(dstd->GetAllocator(tensorflow::AllocatorAttributes()), - src->dtype(), src->shape()); + tensorflow::AllocatorAttributes attr; + if (src->dtype() == tensorflow::DT_VARIANT) { + attr.set_on_host(true); + } + tensorflow::Tensor dst(dstd->GetAllocator(attr), src->dtype(), src->shape()); if (src->shape().num_elements() == 0) { return new TFE_TensorHandle(dst, dst_cpu ? nullptr : dstd); } tensorflow::DeviceContext* src_device_context = nullptr; - if (!IsCPU(srcd)) { + if (!src_cpu) { src_device_context = srcd->tensorflow_gpu_device_info()->default_context; } tensorflow::DeviceContext* dst_device_context = nullptr; @@ -231,15 +270,6 @@ TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, void TFE_DeleteOp(TFE_Op* op) { delete op; } -static void TFE_OpSetDeviceHelper(TFE_Op* op, tensorflow::Device* device, - TF_Status* status) { - // Questionable heuristic: Place the op on the same device as the first input - // placed outside of host memory? - if (IsCPU(op->device) && !IsCPU(device)) { - op->device = device; - } -} - void TFE_OpSetDevice(TFE_Op* op, const char* device_name, TF_Status* status) { tensorflow::Device* d = nullptr; if (device_name != nullptr && strlen(device_name) > 0) { @@ -247,11 +277,32 @@ void TFE_OpSetDevice(TFE_Op* op, const char* device_name, TF_Status* status) { op->ctx->session->device_mgr->LookupDevice(device_name, &d); if (!status->status.ok()) return; } - TFE_OpSetDeviceHelper(op, d, status); + op->device = d; +} + +const char* TFE_OpGetDevice(TFE_Op* op, TF_Status* status) { + tensorflow::Device* device = + (op->device == nullptr) ? op->ctx->devices()[0] : op->device; + return device->name().c_str(); +} + +void TFE_OpSetXLACompilation(TFE_Op* op, unsigned char enable) { + op->use_xla = enable; +#ifndef TENSORFLOW_EAGER_USE_XLA + LOG(WARNING) << "This call is a no-op, as the TensorFlow library is not " + "built with XLA support."; +#endif // TENSORFLOW_EAGER_USE_XLA } void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) { - TFE_OpSetDeviceHelper(op, h->d, status); + // Questionable heuristic ... + // + // Motivation: After an 'op' is placed on GPU because some of its earlier + // inputs are on GPU, we want to keep the 'op' there, even if some later + // inputs of it are not on GPU. + if (IsCPU(op->device) && !IsCPU(h->d)) { + op->device = h->d; + } if (!status->status.ok()) return; op->inputs.push_back(h->t); op->input_devices.push_back(h->d); @@ -268,7 +319,7 @@ TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name, return TF_ATTR_INT; // The compiler requires that we return something. } status->status = - tensorflow::AttrTypeByName(op->attr_types, attr_name, &ret, is_list); + tensorflow::AttrTypeByName(*op->attr_types, attr_name, &ret, is_list); return ret; } @@ -404,6 +455,19 @@ void TFE_OpSetAttrShapeList(TFE_Op* op, const char* attr_name, proto.get(), num_values)); } +void TFE_OpSetAttrFunctionList(TFE_Op* op, const char* attr_name, + const TFE_Op** value, int num_values) { + std::unique_ptr funcs( + new tensorflow::NameAttrList[num_values]); + for (int i = 0; i < num_values; i++) { + funcs[i].set_name(value[i]->name); + value[i]->attrs.FillAttrValueMap(funcs[i].mutable_attr()); + } + op->attrs.Set(attr_name, + tensorflow::gtl::ArraySlice( + funcs.get(), num_values)); +} + namespace { tensorflow::Status ValidateInputTypeAndPlacement( @@ -422,10 +486,17 @@ tensorflow::Status ValidateInputTypeAndPlacement( const tensorflow::Device* actual_device = op->input_devices[i] == nullptr ? host_device : op->input_devices[i]; if (expected_device != actual_device) { - switch (ctx->policy) { - case TFE_DEVICE_PLACEMENT_EXPLICIT: + switch (TFE_ContextGetDevicePlacementPolicy(ctx)) { + case TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32: // TODO(xpan): See if we could bubble python related error up // to python level. + if (op->inputs[i].dtype() == tensorflow::DT_INT32) { + // Note: enabling silent copies of int32 tensors to match behavior + // of graph mode. + break; + } + TF_FALLTHROUGH_INTENDED; + case TFE_DEVICE_PLACEMENT_EXPLICIT: return tensorflow::errors::InvalidArgument( "Tensors on conflicting devices:" " cannot compute ", @@ -478,6 +549,228 @@ tensorflow::Status ValidateInputTypeAndPlacement( } return tensorflow::Status::OK(); } + +#ifdef TENSORFLOW_EAGER_USE_XLA +// Synthesizes and returns a wrapper function over `op`, which must be a +// primitive op (e.g. matmul). +// +// The wrapper function conforms to the function signature expected by +// _XlaLaunchOp, with input params ordered by . For example, if the op has input params , they will be reordered to as the input params to the synthesized function. +// +// It populates `const_input_types`, `arg_input_types` and +// `op_input_to_func_input` based on the reordering results, that the caller can +// use them to build an _XlaLaunchOp. On error, it returns NULL, and sets +// `status` accordingly. +const tensorflow::FunctionDef* OpToFunction( + TFE_Op* op, std::vector* const_input_types, + std::vector* arg_input_types, + tensorflow::gtl::FlatMap* op_input_to_func_input, + TF_Status* status) { + DCHECK(!op->is_function()); + + tensorflow::FunctionDef fdef; + + // Get the OpDef of the op we are trying to encapsulate. + TFE_Context* ctx = op->ctx; + const tensorflow::OpRegistrationData* op_data; + { + tensorflow::tf_shared_lock l(ctx->functions_mu); + status->status = ctx->func_lib_def.LookUp(op->name, &op_data); + if (!status->status.ok()) { + return nullptr; + } + } + const tensorflow::OpDef& op_def = op_data->op_def; + + tensorflow::OpDef* signature = fdef.mutable_signature(); + + // Handle constant inputs. + const std::unordered_set const_inputs( + *tensorflow::XlaOpRegistry::CompileTimeConstantInputs(op->name)); + + // First add place holders for the input args, so that we can refer to them by + // position in the next loop. Also tally up the resource inputs. + int num_resource_inputs = 0; + for (int i = 0; i < op_def.input_arg_size(); ++i) { + if (op_def.input_arg(i).type() == tensorflow::DT_RESOURCE) { + ++num_resource_inputs; + } + signature->add_input_arg(); + } + + // Now we map the input params from `op_def` to `signature`, where the param + // ordering for `signature` is: . + int const_index = 0; + int arg_index = const_inputs.size(); + int resource_index = op_def.input_arg_size() - num_resource_inputs; + for (int i = 0; i < op_def.input_arg_size(); ++i) { + const tensorflow::OpDef::ArgDef& op_input_arg = op_def.input_arg(i); + tensorflow::OpDef::ArgDef* func_input_arg = nullptr; + if (const_inputs.find(op_input_arg.name()) != const_inputs.end()) { + VLOG(1) << "For const input, mapping op input " << i << " to func input " + << const_index; + (*op_input_to_func_input)[i] = const_index; + func_input_arg = signature->mutable_input_arg(const_index++); + const_input_types->push_back( + static_cast(op->inputs[i].dtype())); + } else if (op_input_arg.type() == tensorflow::DT_RESOURCE) { + VLOG(1) << "For resource input, mapping op input " << i + << " to func input " << resource_index; + (*op_input_to_func_input)[i] = resource_index; + func_input_arg = signature->mutable_input_arg(resource_index++); + } else { + VLOG(1) << "For arg input, mapping op input " << i << " to func input " + << arg_index; + (*op_input_to_func_input)[i] = arg_index; + func_input_arg = signature->mutable_input_arg(arg_index++); + arg_input_types->push_back( + static_cast(op->inputs[i].dtype())); + } + + func_input_arg->set_name(op_input_arg.name()); + func_input_arg->set_type(op->inputs[i].dtype()); + } + VLOG(1) << "Added OpDef Inputs: " << fdef.DebugString(); + + // Resources args are at the end of the function input params, and we should + // have iterated over all of them. + DCHECK_EQ(signature->input_arg_size(), resource_index); + + // Make the synthesized function's name unique. + signature->set_name(tensorflow::strings::StrCat( + op_def.name(), func_id_generator.fetch_add(1))); + + // Add the node def and set its input names to match op_def's names. + const tensorflow::NodeDef& ndef = op->attrs.BuildNodeDef(); + DCHECK_EQ(signature->input_arg_size(), ndef.input_size()); + *fdef.add_node_def() = ndef; + for (int i = 0; i < op_def.input_arg_size(); ++i) { + fdef.mutable_node_def(0)->set_input(i, op_def.input_arg(i).name()); + } + VLOG(1) << "Added NodeDef: " << fdef.DebugString(); + + // Fix the output names and set output types. + for (int i = 0; i < op_def.output_arg_size(); ++i) { + tensorflow::OpDef::ArgDef* arg = signature->add_output_arg(); + const tensorflow::OpDef::ArgDef& op_def_arg = op_def.output_arg(i); + const string& out_tensor_name = tensorflow::strings::StrCat( + ndef.name(), ":", op_def_arg.name(), ":", 0); + arg->set_name(op_def_arg.name()); + (*fdef.mutable_ret())[op_def_arg.name()] = out_tensor_name; + const string& type_attr = op_def_arg.type_attr(); + if (!type_attr.empty()) { + auto i = ndef.attr().find(type_attr); + if (i == ndef.attr().end()) { + status->status = tensorflow::errors::InvalidArgument( + tensorflow::strings::StrCat("Could not find attr ", type_attr, + " in NodeDef ", ndef.DebugString())); + return nullptr; + } + arg->set_type(i->second.type()); + } + } + VLOG(1) << "Fixed Output names and all types: " << fdef.DebugString(); + + tensorflow::mutex_lock l(ctx->functions_mu); + status->status = ctx->func_lib_def.AddFunctionDef(fdef); + if (!status->status.ok()) return nullptr; + const auto ret = ctx->func_lib_def.Find(signature->name()); + DCHECK(ret != nullptr); + return ret; +} + +// Builds an _XLALaunchOp as a wrapper over 'op', so that 'op' can be executed +// via XLA. +std::unique_ptr BuildXlaLaunch(TFE_Op* op, TF_Status* status) { + VLOG(1) << "Creating _XlaLaunchOp for TFE_Op " << op->name; + auto launch_op = + std::unique_ptr(TFE_NewOp(op->ctx, "_XlaLaunch", status)); + if (TF_GetCode(status) != TF_OK) return nullptr; + if (op->device) { + TFE_OpSetDevice(launch_op.get(), op->device->name().c_str(), status); + if (TF_GetCode(status) != TF_OK) return nullptr; + } + + const tensorflow::FunctionDef* fdef; + { + tensorflow::tf_shared_lock l(op->ctx->functions_mu); + fdef = op->ctx->func_lib_def.Find(op->name); + } + std::vector const_input_types; + std::vector arg_input_types; + tensorflow::gtl::FlatMap op_input_to_func_input; + if (fdef == nullptr) { + // See if this is a primitive op, and if so create a function for it, so + // that _XlaLaunchOp can access it. + fdef = OpToFunction(op, &const_input_types, &arg_input_types, + &op_input_to_func_input, status); + if (!status->status.ok()) return nullptr; + } else { + // TODO(hongm): XlaOpRegistry::CompileTimeConstantInputs() does not work for + // functions, so we need to find another way to handle constant inputs. + for (int i = const_input_types.size(); + i < fdef->signature().input_arg_size(); ++i) { + VLOG(1) << "Adding Targs from input arg " << i; + const tensorflow::OpDef::ArgDef& arg = fdef->signature().input_arg(i); + arg_input_types.push_back(static_cast(arg.type())); + } + } + DCHECK(fdef != nullptr); + + // Copy inputs and their devices. + // Since input param reordering may have occurred between `op` and `launch_op` + // via `op_input_to_func_input`, adjust the actual inputs accordingly. + launch_op->inputs = op->inputs; + launch_op->input_devices = op->input_devices; + if (!op_input_to_func_input.empty()) { + DCHECK_EQ(op->inputs.size(), op_input_to_func_input.size()); + if (!op->input_devices.empty()) { + DCHECK_EQ(op->input_devices.size(), op_input_to_func_input.size()); + } + for (int i = 0; i < op_input_to_func_input.size(); ++i) { + VLOG(1) << "mapping op input " << i << " to func input " + << op_input_to_func_input[i]; + + launch_op->inputs[op_input_to_func_input[i]] = op->inputs[i]; + if (!op->input_devices.empty()) { + launch_op->input_devices[op_input_to_func_input[i]] = + op->input_devices[i]; + } + } + } + launch_op->attrs.NumInputs(op->inputs.size()); + + TFE_OpSetAttrTypeList(launch_op.get(), "Tconstants", const_input_types.data(), + const_input_types.size()); + + // Set Targs and Nresources attrs. + TFE_OpSetAttrTypeList(launch_op.get(), "Targs", arg_input_types.data(), + arg_input_types.size()); + const int num_resource_inputs = fdef->signature().input_arg_size() - + const_input_types.size() - + arg_input_types.size(); + TFE_OpSetAttrInt(launch_op.get(), "Nresources", num_resource_inputs); + + // Set Tresults attr. + std::vector tresults; + for (const tensorflow::OpDef::ArgDef& arg : fdef->signature().output_arg()) { + tresults.push_back(static_cast(arg.type())); + } + TFE_OpSetAttrTypeList(launch_op.get(), "Tresults", tresults.data(), + tresults.size()); + + // Set function attr. + tensorflow::AttrValue attr_value; + tensorflow::NameAttrList* func = attr_value.mutable_func(); + func->set_name(fdef->signature().name()); + launch_op->attrs.Set("function", attr_value); + + return launch_op; +} +#endif // TENSORFLOW_EAGER_USE_XLA } // namespace void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, @@ -486,11 +779,26 @@ void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, // TODO(ashankar): ASSUMPTION: ctx->devices()[0] is always CPU tensorflow::Device* device = (op->device == nullptr) ? ctx->devices()[0] : op->device; + +#ifdef TENSORFLOW_EAGER_USE_XLA + std::unique_ptr xla_launch_op; + if (op->use_xla && op->name != "_XlaLaunch") { + xla_launch_op = BuildXlaLaunch(op, status); + if (!status->status.ok()) { + return; + } + op = xla_launch_op.get(); + } +#endif // TENSORFLOW_EAGER_USE_XLA + std::vector outputs(1); const tensorflow::MemoryTypeVector* output_memory_types = nullptr; tensorflow::Fprint128 cache_key = op->attrs.CacheKey(device->name()); - tensorflow::KernelAndDevice* kernel = - tensorflow::gtl::FindPtrOrNull(ctx->kernel_cache, cache_key); + tensorflow::KernelAndDevice* kernel; + { + tensorflow::tf_shared_lock l(ctx->cache_mu); + kernel = tensorflow::gtl::FindPtrOrNull(ctx->kernel_cache, cache_key); + } if (kernel == nullptr) { const tensorflow::NodeDef& ndef = op->attrs.BuildNodeDef(); kernel = new tensorflow::KernelAndDevice(ctx->rendezvous); @@ -506,6 +814,7 @@ void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, delete kernel; return; } + tensorflow::mutex_lock ml(ctx->cache_mu); tensorflow::gtl::InsertOrUpdate(&(ctx->kernel_cache), cache_key, kernel); } std::vector copied_tensors; @@ -530,7 +839,7 @@ void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, // WARNING: kernel->Run utilizes the FunctionLibraryRuntime // (ctx->func_lib(device)), which in turn holds a pointer to func_lib_def, // which is GUARDED_BY(ctx->functions_mu). But knowledge of the implementation - // of FunctionLibraryRuntime tells use that func_lib_def is not accessed by + // of FunctionLibraryRuntime tells us that func_lib_def is not accessed by // FunctionLibraryRuntime::Run(), so there is no thread-safety concern here. // This is quite subtle. Re-work things to make this better? (Would it make // sense for FunctionLibraryRuntime to ensure thread-safe access to diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h index 7caab43d00301b9e0df64cbc3e13123bfc3cc427..7a321b54da343fd2b8912187bc620c1e7456db0c 100644 --- a/tensorflow/c/eager/c_api.h +++ b/tensorflow/c/eager/c_api.h @@ -17,6 +17,8 @@ limitations under the License. #define TENSORFLOW_C_EAGER_C_API_H_ // C API extensions to experiment with eager execution of kernels. +// WARNING: Unlike tensorflow/c/c_api.h, the API here is not guaranteed to be +// stable and can change without notice. #include "tensorflow/c/c_api.h" @@ -59,14 +61,16 @@ TF_CAPI_EXPORT extern void TFE_ContextOptionsSetConfig( // Controls how to act when we try to run an operation on a given device but // some input tensors are not on that device. typedef enum TFE_ContextDevicePlacementPolicy { - // The default: running operations with input tensors on the wrong device will - // fail. + // Running operations with input tensors on the wrong device will fail. TFE_DEVICE_PLACEMENT_EXPLICIT = 0, // Copy the tensor to the right device but log a warning. TFE_DEVICE_PLACEMENT_WARN = 1, // Silently copy the tensor, which has a performance cost since the // operation will be blocked till the copy completes. TFE_DEVICE_PLACEMENT_SILENT = 2, + // Default placement policy which silently copies int32 tensors but not other + // dtypes. + TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32 = 3, } TFE_ContextDevicePlacementPolicy; TF_CAPI_EXPORT extern void TFE_ContextOptionsSetDevicePlacementPolicy( @@ -83,10 +87,27 @@ typedef struct TFE_Context TFE_Context; TF_CAPI_EXPORT extern TFE_Context* TFE_NewContext( const TFE_ContextOptions* opts, TF_Status* status); -TF_CAPI_EXPORT extern void TFE_DeleteContext(TFE_Context* ctx, TF_Status* status); +TF_CAPI_EXPORT extern void TFE_DeleteContext(TFE_Context* ctx, + TF_Status* status); TF_CAPI_EXPORT extern TF_DeviceList* TFE_ContextListDevices(TFE_Context* ctx, TF_Status* status); +// Clears the internal caches in the TFE context. Useful when reseeding random +// ops. +TF_CAPI_EXPORT extern void TFE_ContextClearCaches(TFE_Context* ctx); + +// Sets a thread-local device placement policy. After this call, other calls to +// TFE_Execute in the same thread will use the device policy specified here +// instead of the device policy used to construct the context. This has no +// effect on the device policy used by other program threads. +TF_CAPI_EXPORT extern void TFE_ContextSetThreadLocalDevicePlacementPolicy( + TFE_Context*, TFE_ContextDevicePlacementPolicy); + +// Returns the device placement policy to be used by this context in the current +// thread. +TF_CAPI_EXPORT extern TFE_ContextDevicePlacementPolicy +TFE_ContextGetDevicePlacementPolicy(TFE_Context*); + // A handle to a tensor on a device. // // Like a TF_Tensor, a TFE_TensorHandle refers to a tensor with a value, shape, @@ -99,8 +120,10 @@ TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_NewTensorHandle(TF_Tensor* t, TF_CAPI_EXPORT extern void TFE_DeleteTensorHandle(TFE_TensorHandle* h); TF_CAPI_EXPORT extern TF_DataType TFE_TensorHandleDataType(TFE_TensorHandle* h); TF_CAPI_EXPORT extern int TFE_TensorHandleNumDims(TFE_TensorHandle* h); -TF_CAPI_EXPORT extern int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index); -TF_CAPI_EXPORT extern const char* TFE_TensorHandleDeviceName(TFE_TensorHandle* h); +TF_CAPI_EXPORT extern int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, + int dim_index); +TF_CAPI_EXPORT extern const char* TFE_TensorHandleDeviceName( + TFE_TensorHandle* h); TF_CAPI_EXPORT extern TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, TF_Status* status); @@ -110,10 +133,9 @@ TF_CAPI_EXPORT extern TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, // that shares the underlying buffer. Otherwise, it currently requires at least // one of the source or destination devices to be CPU (i.e., for the source or // destination tensor to be placed in host memory). -TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, - TFE_Context* ctx, - const char* device_name, - TF_Status* status); +TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_TensorHandleCopyToDevice( + TFE_TensorHandle* h, TFE_Context* ctx, const char* device_name, + TF_Status* status); // Description of the TensorFlow op to execute. // @@ -128,17 +150,31 @@ TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorH // the additional sanity checks there seem unnecessary; typedef struct TFE_Op TFE_Op; -TF_CAPI_EXPORT extern TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, +TF_CAPI_EXPORT extern TFE_Op* TFE_NewOp(TFE_Context* ctx, + const char* op_or_function_name, TF_Status* status); TF_CAPI_EXPORT extern void TFE_DeleteOp(TFE_Op* op); TF_CAPI_EXPORT extern void TFE_OpSetDevice(TFE_Op* op, const char* device_name, TF_Status* status); +// The returned string remains valid throughout the lifetime of 'op'. +TF_CAPI_EXPORT extern const char* TFE_OpGetDevice(TFE_Op* op, + TF_Status* status); -TF_CAPI_EXPORT extern void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status); +// When 'enable' is set to 1, and if TensorFlow library is built with XLA +// support, a subsequent TFE_Execute() call on `op` will run the op via XLA. +// +// If the library is not built with XLA support, this call would be a no-op. +TF_CAPI_EXPORT extern void TFE_OpSetXLACompilation(TFE_Op* op, + unsigned char enable); + +TF_CAPI_EXPORT extern void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, + TF_Status* status); -TF_CAPI_EXPORT extern TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name, - unsigned char* is_list, TF_Status* status); +TF_CAPI_EXPORT extern TF_AttrType TFE_OpGetAttrType(TFE_Op* op, + const char* attr_name, + unsigned char* is_list, + TF_Status* status); // Get an attribute type given an op name; a fusion of TFE_NewOp and // TFE_OpGetAttrType for use from Python without the overhead of the individual // calls and memory management of TFE_Op. @@ -146,10 +182,13 @@ TF_CAPI_EXPORT extern TF_AttrType TFE_OpNameGetAttrType( TFE_Context* ctx, const char* op_or_function_name, const char* attr_name, unsigned char* is_list, TF_Status* status); -TF_CAPI_EXPORT extern void TFE_OpSetAttrString(TFE_Op* op, const char* attr_name, +TF_CAPI_EXPORT extern void TFE_OpSetAttrString(TFE_Op* op, + const char* attr_name, const char* value); -TF_CAPI_EXPORT extern void TFE_OpSetAttrInt(TFE_Op* op, const char* attr_name, int64_t value); -TF_CAPI_EXPORT extern void TFE_OpSetAttrFloat(TFE_Op* op, const char* attr_name, float value); +TF_CAPI_EXPORT extern void TFE_OpSetAttrInt(TFE_Op* op, const char* attr_name, + int64_t value); +TF_CAPI_EXPORT extern void TFE_OpSetAttrFloat(TFE_Op* op, const char* attr_name, + float value); TF_CAPI_EXPORT extern void TFE_OpSetAttrBool(TFE_Op* op, const char* attr_name, unsigned char value); TF_CAPI_EXPORT extern void TFE_OpSetAttrType(TFE_Op* op, const char* attr_name, @@ -158,7 +197,8 @@ TF_CAPI_EXPORT extern void TFE_OpSetAttrType(TFE_Op* op, const char* attr_name, // -1 and `dims` can be null. If a dimension is unknown, the // corresponding entry in the `dims` array must be -1. TF_CAPI_EXPORT extern void TFE_OpSetAttrShape(TFE_Op* op, const char* attr_name, - const int64_t* dims, const int num_dims, + const int64_t* dims, + const int num_dims, TF_Status* out_status); // Sets the attribute attr_name to be a function specified by 'function'. @@ -169,19 +209,33 @@ TF_CAPI_EXPORT extern void TFE_OpSetAttrFunction(TFE_Op* op, const char* attr_name, const TFE_Op* value); -TF_CAPI_EXPORT extern void TFE_OpSetAttrStringList(TFE_Op* op, const char* attr_name, - const char** value, int num_values); -TF_CAPI_EXPORT extern void TFE_OpSetAttrIntList(TFE_Op* op, const char* attr_name, - const int64_t* values, int num_values); -TF_CAPI_EXPORT extern void TFE_OpSetAttrFloatList(TFE_Op* op, const char* attr_name, - const float* values, int num_values); -TF_CAPI_EXPORT extern void TFE_OpSetAttrBoolList(TFE_Op* op, const char* attr_name, - const unsigned char* values, int num_values); -TF_CAPI_EXPORT extern void TFE_OpSetAttrTypeList(TFE_Op* op, const char* attr_name, - const TF_DataType* values, int num_values); -TF_CAPI_EXPORT extern void TFE_OpSetAttrShapeList(TFE_Op* op, const char* attr_name, - const int64_t** dims, const int* num_dims, - int num_values, TF_Status* out_status); +TF_CAPI_EXPORT extern void TFE_OpSetAttrStringList(TFE_Op* op, + const char* attr_name, + const char** value, + int num_values); +TF_CAPI_EXPORT extern void TFE_OpSetAttrIntList(TFE_Op* op, + const char* attr_name, + const int64_t* values, + int num_values); +TF_CAPI_EXPORT extern void TFE_OpSetAttrFloatList(TFE_Op* op, + const char* attr_name, + const float* values, + int num_values); +TF_CAPI_EXPORT extern void TFE_OpSetAttrBoolList(TFE_Op* op, + const char* attr_name, + const unsigned char* values, + int num_values); +TF_CAPI_EXPORT extern void TFE_OpSetAttrTypeList(TFE_Op* op, + const char* attr_name, + const TF_DataType* values, + int num_values); +TF_CAPI_EXPORT extern void TFE_OpSetAttrShapeList( + TFE_Op* op, const char* attr_name, const int64_t** dims, + const int* num_dims, int num_values, TF_Status* out_status); +TF_CAPI_EXPORT extern void TFE_OpSetAttrFunctionList(TFE_Op* op, + const char* attr_name, + const TFE_Op** value, + int num_values); // Execute the operation defined by 'op' and return handles to computed // tensors in 'retvals'. @@ -196,9 +250,9 @@ TF_CAPI_EXPORT extern void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, // Add a function (serialized FunctionDef protocol buffer) to ctx so // that it can be invoked using TFE_Execute. -TF_CAPI_EXPORT extern void TFE_ContextAddFunctionDef(TFE_Context* ctx, - const char* serialized_function_def, - size_t size, TF_Status* status); +TF_CAPI_EXPORT extern void TFE_ContextAddFunctionDef( + TFE_Context* ctx, const char* serialized_function_def, size_t size, + TF_Status* status); // Adds a function (created from TF_GraphToFunction or // TF_FunctionImportFunctionDef) to the context, allowing it to be executed with diff --git a/tensorflow/c/eager/c_api_internal.h b/tensorflow/c/eager/c_api_internal.h index 11b7a519c3a0632809843a434d5bd863bbeccdee..7b9f1db02ed9c53a280c7bd1284165cac4fb6353 100644 --- a/tensorflow/c/eager/c_api_internal.h +++ b/tensorflow/c/eager/c_api_internal.h @@ -21,6 +21,7 @@ limitations under the License. #include #include #include +#include #include #include "tensorflow/c/c_api.h" @@ -34,20 +35,34 @@ limitations under the License. #include "tensorflow/core/lib/gtl/stl_util.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/thread_annotations.h" +#include "tensorflow/core/public/version.h" struct TFE_ContextOptions { TF_SessionOptions session_options; - TFE_ContextDevicePlacementPolicy policy{TFE_DEVICE_PLACEMENT_EXPLICIT}; + TFE_ContextDevicePlacementPolicy policy{ + TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32}; }; struct TFE_Context { - explicit TFE_Context(TF_Session* s) : session(s) {} - - TFE_ContextDevicePlacementPolicy policy; + explicit TFE_Context(const TFE_ContextOptions& opts, TF_Session* s) + : policy(opts.policy), + session(s), + rendezvous(new tensorflow::IntraProcessRendezvous(s->device_mgr)), + pflr(new tensorflow::ProcessFunctionLibraryRuntime( + session->device_mgr, opts.session_options.options.env, + TF_GRAPH_DEF_VERSION, &func_lib_def, {})) {} + + const TFE_ContextDevicePlacementPolicy policy; + + // Note: we cannot use C++11 thread_local here as there is no concept of a + // thread-local-object-local variable in C++11. + tensorflow::mutex policy_map_mu; + std::unordered_map + thread_local_policies GUARDED_BY(policy_map_mu); // TFE_Context is an extension of TF_Session. And TF_Session needs a TF_Graph. - TF_Session* session; - tensorflow::Rendezvous* rendezvous; + TF_Session* const session; + tensorflow::Rendezvous* const rendezvous; tensorflow::mutex functions_mu; tensorflow::FunctionLibraryDefinition func_lib_def GUARDED_BY(functions_mu){ @@ -56,13 +71,14 @@ struct TFE_Context { // One FunctionLibraryRuntime per device. // func_libs[i] is the FunctionLibraryRuntime corresponding to // session->devices[i]. - std::unique_ptr pflr; + const std::unique_ptr pflr; + tensorflow::mutex cache_mu; std::unordered_map - kernel_cache; + kernel_cache GUARDED_BY(cache_mu); - tensorflow::FunctionLibraryRuntime* func_lib(tensorflow::Device* d) { + tensorflow::FunctionLibraryRuntime* func_lib(tensorflow::Device* d) const { return pflr->GetFLR(d->name()); } @@ -91,6 +107,8 @@ struct TFE_TensorHandle { }; struct TFE_Op { + // t is NULL iff the TFE_Op corresponds to a TensorFlow function instead of a + // primitive operation. TFE_Op(TFE_Context* ctx, const char* op, const tensorflow::AttrTypeMap* t) : ctx(ctx), name(op), attrs(op), attr_types(t), device(nullptr) {} @@ -103,6 +121,7 @@ struct TFE_Op { std::vector inputs; std::vector input_devices; tensorflow::Device* device; + bool use_xla = false; }; #endif // TENSORFLOW_C_EAGER_C_API_INTERNAL_H_ diff --git a/tensorflow/c/eager/c_api_test.cc b/tensorflow/c/eager/c_api_test.cc index 423a7e1ff71bfdc5f51e36ae63359869ea079ddc..4a3ecbc0abb16296a84c0d2184dc3fc9f7f3ebb4 100644 --- a/tensorflow/c/eager/c_api_test.cc +++ b/tensorflow/c/eager/c_api_test.cc @@ -60,6 +60,63 @@ TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) { return op; } +TFE_TensorHandle* TestAxisTensorHandle() { + int64_t dims[] = {1}; + int data[] = {1}; + TF_Tensor* t = TF_AllocateTensor( + TF_INT32, &dims[0], sizeof(dims) / sizeof(int64_t), sizeof(data)); + memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t)); + TF_Status* status = TF_NewStatus(); + TFE_TensorHandle* th = TFE_NewTensorHandle(t, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteTensor(t); + TF_DeleteStatus(status); + return th; +} + +TFE_Op* MinOp(TFE_Context* ctx, TFE_TensorHandle* input, + TFE_TensorHandle* axis) { + TF_Status* status = TF_NewStatus(); + + TFE_Op* op = TFE_NewOp(ctx, "Min", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(op, input, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(op, axis, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpSetAttrBool(op, "keep_dims", 1); + TFE_OpSetAttrType(op, "Tidx", TF_INT32); + TF_DeleteStatus(status); + TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(input)); + + return op; +} + +// If there is a GPU device, returns true and sets 'gpu_device_name' +// accordingly. +bool GetGPUDeviceName(TFE_Context* ctx, string* gpu_device_name) { + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get()); + CHECK_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + + const int num_devices = TF_DeviceListCount(devices); + for (int i = 0; i < num_devices; ++i) { + const string device_type(TF_DeviceListType(devices, i, status.get())); + CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get()); + const string device_name(TF_DeviceListName(devices, i, status.get())); + CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get()); + if (device_type == "GPU") { + *gpu_device_name = device_name; + LOG(INFO) << "Found GPU device " << device_name; + TF_DeleteDeviceList(devices); + return true; + } + } + TF_DeleteDeviceList(devices); + return false; +} + void BM_InitOp(int iters) { tensorflow::testing::StopTiming(); TF_Status* status = TF_NewStatus(); @@ -288,22 +345,56 @@ TEST(CAPI, TensorHandleSilentCopy) { TF_Tensor* t = TFE_TensorHandleResolve(hcpu, status.get()); ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); - TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get()); - ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); - const int num_devices = TF_DeviceListCount(devices); - // Disable the test if no GPU is present. - if (num_devices > 1) { - const int device_to_use = 1; - const string name(TF_DeviceListName(devices, device_to_use, status.get())); + string gpu_device_name; + if (GetGPUDeviceName(ctx, &gpu_device_name)) { + TFE_TensorHandle* hgpu = TFE_TensorHandleCopyToDevice( + hcpu, ctx, gpu_device_name.c_str(), status.get()); ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); - TFE_TensorHandle* hgpu = - TFE_TensorHandleCopyToDevice(hcpu, ctx, name.c_str(), status.get()); + TFE_Op* matmul = MatMulOp(ctx, hcpu, hgpu); + TFE_OpSetDevice(matmul, gpu_device_name.c_str(), status.get()); + ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); + TFE_TensorHandle* retvals[1]; + int num_retvals = 1; + TFE_Execute(matmul, &retvals[0], &num_retvals, status.get()); + ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); + TFE_DeleteOp(matmul); + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteTensorHandle(hgpu); + } + + TF_DeleteTensor(t); + TFE_DeleteTensorHandle(hcpu); + TFE_DeleteContext(ctx, status.get()); + EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); +} + +TEST(CAPI, TensorHandleSilentCopyLocal) { + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetDevicePlacementPolicy(opts, + TFE_DEVICE_PLACEMENT_EXPLICIT); + TFE_Context* ctx = TFE_NewContext(opts, status.get()); + TFE_ContextSetThreadLocalDevicePlacementPolicy(ctx, + TFE_DEVICE_PLACEMENT_SILENT); + TFE_DeleteContextOptions(opts); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + + TFE_TensorHandle* hcpu = TestMatrixTensorHandle(); + TF_Tensor* t = TFE_TensorHandleResolve(hcpu, status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + + // Disable the test if no GPU is present. + string gpu_device_name; + if (GetGPUDeviceName(ctx, &gpu_device_name)) { + TFE_TensorHandle* hgpu = TFE_TensorHandleCopyToDevice( + hcpu, ctx, gpu_device_name.c_str(), status.get()); ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); TFE_Op* matmul = MatMulOp(ctx, hcpu, hgpu); - TFE_OpSetDevice(matmul, name.c_str(), status.get()); + TFE_OpSetDevice(matmul, gpu_device_name.c_str(), status.get()); ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); TFE_TensorHandle* retvals[1]; int num_retvals = 1; @@ -314,14 +405,44 @@ TEST(CAPI, TensorHandleSilentCopy) { TFE_DeleteTensorHandle(hgpu); } - TF_DeleteDeviceList(devices); TF_DeleteTensor(t); TFE_DeleteTensorHandle(hcpu); TFE_DeleteContext(ctx, status.get()); EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); } -TEST(CAPI, Execute) { +TEST(CAPI, SetAndGetOpDevices) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_Op* matmul = MatMulOp(ctx, m, m); + + // Disable the test if no GPU is present. + string gpu_device_name; + if (GetGPUDeviceName(ctx, &gpu_device_name)) { + TFE_OpSetDevice(matmul, "GPU:0", status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + const char* device_name = TFE_OpGetDevice(matmul, status); + ASSERT_TRUE(strstr(device_name, "GPU:0") != nullptr); + + TFE_OpSetDevice(matmul, "CPU:0", status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + device_name = TFE_OpGetDevice(matmul, status); + ASSERT_TRUE(strstr(device_name, "CPU:0") != nullptr); + } + + TFE_DeleteOp(matmul); + TFE_DeleteTensorHandle(m); + TFE_DeleteContext(ctx, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); +} + +TEST(CAPI, Execute_MatMul_CPU) { TF_Status* status = TF_NewStatus(); TFE_ContextOptions* opts = TFE_NewContextOptions(); TFE_Context* ctx = TFE_NewContext(opts, status); @@ -354,6 +475,117 @@ TEST(CAPI, Execute) { TF_DeleteStatus(status); } +TEST(CAPI, Execute_Min_CPU) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* input = TestMatrixTensorHandle(); + TFE_TensorHandle* axis = TestAxisTensorHandle(); + TFE_Op* minOp = MinOp(ctx, input, axis); + TFE_TensorHandle* retvals[2] = {nullptr}; + int num_retvals = 2; // Should be reduced to 1 by the TFE_Execute call. + TFE_Execute(minOp, &retvals[0], &num_retvals, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteOp(minOp); + TFE_DeleteTensorHandle(input); + TFE_DeleteTensorHandle(axis); + TFE_DeleteContext(ctx, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + ASSERT_EQ(1, num_retvals); + + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + TFE_DeleteTensorHandle(retvals[0]); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + float output[2] = {0}; + EXPECT_EQ(sizeof(output), TF_TensorByteSize(t)); + memcpy(&output[0], TF_TensorData(t), TF_TensorByteSize(t)); + TF_DeleteTensor(t); + EXPECT_EQ(1, output[0]); + EXPECT_EQ(3, output[1]); + TF_DeleteStatus(status); +} + +#ifdef TENSORFLOW_EAGER_USE_XLA +TEST(CAPI, Execute_MatMul_XLA_CPU) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_Op* matmul = MatMulOp(ctx, m, m); + + TFE_OpSetXLACompilation(matmul, true); + + TFE_TensorHandle* retvals[2] = {nullptr}; + int num_retvals = 2; // Should be reduced to 1 by the TFE_Execute call. + TFE_Execute(matmul, &retvals[0], &num_retvals, status); + // Running a primitive TF operator via XLA is not yet supported. + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TFE_DeleteOp(matmul); + TFE_DeleteTensorHandle(m); + TFE_DeleteContext(ctx, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + EXPECT_EQ(1, num_retvals); + + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + TFE_DeleteTensorHandle(retvals[0]); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + float product[4] = {0}; + EXPECT_EQ(sizeof(product), TF_TensorByteSize(t)); + memcpy(&product[0], TF_TensorData(t), TF_TensorByteSize(t)); + TF_DeleteTensor(t); + EXPECT_EQ(7, product[0]); + EXPECT_EQ(10, product[1]); + EXPECT_EQ(15, product[2]); + EXPECT_EQ(22, product[3]); + + TF_DeleteStatus(status); +} + +TEST(CAPI, Execute_Min_XLA_CPU) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* input = TestMatrixTensorHandle(); + TFE_TensorHandle* axis = TestAxisTensorHandle(); + TFE_Op* minOp = MinOp(ctx, input, axis); + + TFE_OpSetXLACompilation(minOp, true); + + TFE_TensorHandle* retvals[2] = {nullptr}; + int num_retvals = 2; // Should be reduced to 1 by the TFE_Execute call. + TFE_Execute(minOp, &retvals[0], &num_retvals, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteOp(minOp); + TFE_DeleteTensorHandle(input); + TFE_DeleteTensorHandle(axis); + TFE_DeleteContext(ctx, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + ASSERT_EQ(1, num_retvals); + + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + TFE_DeleteTensorHandle(retvals[0]); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + float output[2] = {0}; + EXPECT_EQ(sizeof(output), TF_TensorByteSize(t)); + memcpy(&output[0], TF_TensorData(t), TF_TensorByteSize(t)); + TF_DeleteTensor(t); + EXPECT_EQ(1, output[0]); + EXPECT_EQ(3, output[1]); + TF_DeleteStatus(status); +} +#endif // TENSORFLOW_EAGER_USE_XLA + TEST(CAPI, ExecuteWithTracing) { TF_Status* status = TF_NewStatus(); TFE_ContextOptions* opts = TFE_NewContextOptions(); @@ -395,7 +627,69 @@ TEST(CAPI, ExecuteWithTracing) { TF_DeleteStatus(status); } -TEST(CAPI, Function) { +TEST(CAPI, Function_ident_CPU) { + // First create a simple identity function. + TF_Graph* function_graph = TF_NewGraph(); + TF_OperationDescription* arg_descr = + TF_NewOperation(function_graph, "Placeholder", "arg"); + TF_SetAttrType(arg_descr, "dtype", TF_INT32); + TF_Status* status = TF_NewStatus(); + TF_Operation* arg = TF_FinishOperation(arg_descr, status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TF_OperationDescription* id_descr = + TF_NewOperation(function_graph, "Identity", "id"); + TF_SetAttrType(id_descr, "T", TF_INT32); + TF_AddInput(id_descr, {arg, 0}); + TF_Operation* id = TF_FinishOperation(id_descr, status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TF_Output input{arg, 0}; + TF_Output output{id, 0}; + TF_Function* fn = + TF_GraphToFunction(function_graph, "ident", 0, 1, &id, 1, &input, 1, + &output, nullptr, nullptr, "test", status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TF_DeleteGraph(function_graph); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TFE_DeleteContextOptions(opts); + TFE_ContextAddFunction(ctx, fn, status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TF_DeleteFunction(fn); + + TF_Tensor* t = + TF_AllocateTensor(TF_INT32, nullptr, 0, 1 * sizeof(tensorflow::int32)); + *reinterpret_cast(TF_TensorData(t)) = 42; + TFE_TensorHandle* h = TFE_NewTensorHandle(t, status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TF_DeleteTensor(t); + + TFE_Op* op = TFE_NewOp(ctx, "ident", status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TFE_OpAddInput(op, h, status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + + std::vector result; + result.push_back(nullptr); + int num_retvals = 1; + TFE_Execute(op, result.data(), &num_retvals, status); + TFE_DeleteOp(op); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + ASSERT_EQ(num_retvals, 1); + + TF_Tensor* r = TFE_TensorHandleResolve(result[0], status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + EXPECT_EQ(*reinterpret_cast(TF_TensorData(r)), 42); + TFE_DeleteTensorHandle(h); + TF_DeleteTensor(r); + TFE_DeleteTensorHandle(result[0]); + TFE_DeleteContext(ctx, status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + TF_DeleteStatus(status); +} + +#ifdef TENSORFLOW_EAGER_USE_XLA +TEST(CAPI, Function_ident_XLA_CPU) { // First create a simple identity function. TF_Graph* function_graph = TF_NewGraph(); TF_OperationDescription* arg_descr = @@ -437,6 +731,9 @@ TEST(CAPI, Function) { TFE_OpAddInput(op, h, status); ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + // Now run it via XLA. + TFE_OpSetXLACompilation(op, true); + std::vector result; result.push_back(nullptr); int num_retvals = 1; @@ -455,6 +752,7 @@ TEST(CAPI, Function) { ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); TF_DeleteStatus(status); } +#endif // TENSORFLOW_EAGER_USE_XLA string MatMulFunction() { tensorflow::FunctionDef def; diff --git a/tensorflow/c/eager/runtime.cc b/tensorflow/c/eager/runtime.cc index ec34b0ea77cfca8b7964c59f79aeeb87ec7ddf52..f77a937f1ffc2d146224cb3191a5ca127daefc22 100644 --- a/tensorflow/c/eager/runtime.cc +++ b/tensorflow/c/eager/runtime.cc @@ -86,10 +86,9 @@ Status AttrTypeMapForOp(const char* op_name, const AttrTypeMap** out) { return Status::OK(); } -Status AttrTypeByName(const AttrTypeMap* m, const string& attr_name, +Status AttrTypeByName(const AttrTypeMap& m, const string& attr_name, TF_AttrType* out, unsigned char* is_list) { - CHECK(m); - auto* t = gtl::FindOrNull(*m, attr_name); + auto* t = gtl::FindOrNull(m, attr_name); if (t == nullptr) { return errors::InvalidArgument("Attribute '", attr_name, "' does not exist for this operation"); @@ -173,14 +172,14 @@ void CombineUnordered(const tensorflow::Fprint128& a, b->high64 += a.high64; } -inline tensorflow::Fprint128 CacheKeyHelper(const StringPiece& s, +inline tensorflow::Fprint128 CacheKeyHelper(StringPiece s, const tensorflow::Fprint128& b) { // TODO(agarwal): avoid ToString(). tensorflow::Fprint128 a = tensorflow::Fingerprint128(s.ToString()); return FingerprintCat128(a, b); } -inline tensorflow::Fprint128 CacheKeyHelper(const StringPiece& s, uint64 b) { +inline tensorflow::Fprint128 CacheKeyHelper(StringPiece s, uint64 b) { return CacheKeyHelper(s, {b, b}); } @@ -316,18 +315,12 @@ Status KernelAndDevice::Run(std::vector* input_tensors, allocator_pair.second->GetRecordsAndUnRef(); } auto* ms = stats->mutable_memory_stats(); - ms->set_host_temp_memory_size(context.host_temp_memory_size()); - ms->set_device_temp_memory_size(context.device_temp_memory_size()); - for (const auto& alloc_id : context.host_persistent_alloc_ids()) { - ms->mutable_host_persistent_tensor_alloc_ids()->Add(alloc_id); + ms->set_temp_memory_size(context.temp_memory_allocated()); + for (const auto& alloc_id : context.persistent_alloc_ids()) { + ms->mutable_persistent_tensor_alloc_ids()->Add(alloc_id); } - for (const auto& alloc_id : context.device_persistent_alloc_ids()) { - ms->mutable_device_persistent_tensor_alloc_ids()->Add(alloc_id); - } - ms->set_host_persistent_memory_size( - context.host_persistent_memory_allocated()); - ms->set_device_persistent_memory_size( - context.device_persistent_memory_allocated()); + + ms->set_persistent_memory_size(context.persistent_memory_allocated()); } return Status::OK(); } diff --git a/tensorflow/c/eager/runtime.h b/tensorflow/c/eager/runtime.h index e28a416e67f8382dbd490648106a7eb6e5fcfd13..4d20b5244a46fcde2eed0a429dced2a77b86aedd 100644 --- a/tensorflow/c/eager/runtime.h +++ b/tensorflow/c/eager/runtime.h @@ -43,7 +43,7 @@ typedef std::unordered_map AttrTypeMap; Status AttrTypeMapForOp(const char* op_name, const AttrTypeMap** out); // Looks for 'attr_name' in 'm' and sets 'out' and 'is_list'. -Status AttrTypeByName(const AttrTypeMap* m, const string& attr_name, +Status AttrTypeByName(const AttrTypeMap& m, const string& attr_name, TF_AttrType* out, unsigned char* is_list); // KernelAndDevice::Init needs a NodeDef only to pass the attribute map through. diff --git a/tensorflow/c/eager/runtime_test.cc b/tensorflow/c/eager/runtime_test.cc index 2ccca66f672b96b3c782ddbfc828eeda270cebee..643153058ce3d6f0c88dd23a0dec4c6eff060319 100644 --- a/tensorflow/c/eager/runtime_test.cc +++ b/tensorflow/c/eager/runtime_test.cc @@ -63,17 +63,17 @@ TEST(AttrTypeMap, Lookup) { TF_AttrType t; unsigned char is_list = 1; - s = AttrTypeByName(m, "ThisAttribyteCannotPossiblyExist", &t, &is_list); + s = AttrTypeByName(*m, "ThisAttribyteCannotPossiblyExist", &t, &is_list); EXPECT_FALSE(s.ok()); EXPECT_NE(is_list, 0); - s = AttrTypeByName(m, "transpose_a", &t, &is_list); + s = AttrTypeByName(*m, "transpose_a", &t, &is_list); ASSERT_TRUE(s.ok()) << s; EXPECT_EQ(TF_ATTR_BOOL, t); EXPECT_EQ(is_list, 0); s = AttrTypeMapForOp("Squeeze", &m); ASSERT_TRUE(s.ok()) << s; - s = AttrTypeByName(m, "squeeze_dims", &t, &is_list); + s = AttrTypeByName(*m, "squeeze_dims", &t, &is_list); ASSERT_TRUE(s.ok()) << s; EXPECT_EQ(TF_ATTR_INT, t); EXPECT_NE(is_list, 0); diff --git a/tensorflow/c/eager/tape.h b/tensorflow/c/eager/tape.h index 2b65e38f54090af6731685f78d5f7f914a875e3c..bdb0815d6b68444ec1c89b835d563db20ce4d8a1 100644 --- a/tensorflow/c/eager/tape.h +++ b/tensorflow/c/eager/tape.h @@ -18,12 +18,12 @@ limitations under the License. // Language-agnostic gradient tape. Does not perform backpropagation, just // maintains the data structures required to do so. -#include -#include #include #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/types.h" namespace tensorflow { @@ -54,11 +54,11 @@ struct OpTapeEntry { // Map from tensor_id to internally-defined operation-id of the operation which // produced this tensor. A value of -1 means that the tensor was directly // watched and not the result of any operation in the tape. -using TensorTape = std::unordered_map; +using TensorTape = gtl::FlatMap; // Map from operation-id to tape entry. template -using OpTape = std::unordered_map>; +using OpTape = gtl::FlatMap>; // Operations the tape needs to perform on tensors to do backpropagation. Named // "vspace" because a subset of these are related to a vector space, such as @@ -159,7 +159,7 @@ class GradientTape { // Map from tensor id to number of remaining usages (i.e. how many entries in // the tape refer to it); to aid in tape garbage collection. - std::unordered_map tensor_usage_; + gtl::FlatMap tensor_usage_; // If false, all activations are deleted in the first call to ComputeGradient. // Else, only when this is destructed. @@ -286,11 +286,11 @@ struct BackpropInitialState { // Map from tensor ID to how many references still exist for this tensor in // the tape. - std::unordered_map tensor_usage_counts; + gtl::FlatMap tensor_usage_counts; // Maps from op ID to how many output tensors of this op still need to have // their gradients computed. - std::unordered_map op_missing_tensor; + gtl::FlatMap op_missing_tensor; }; // If `persistent_tape` is true, op_tape is not changed and none of the @@ -301,8 +301,8 @@ struct BackpropInitialState { template BackpropInitialState PrepareBackprop( gtl::ArraySlice target, const TensorTape& tensor_tape, - OpTape* op_tape, - const std::unordered_set& sources_set, bool persistent_tape) { + OpTape* op_tape, const gtl::FlatSet& sources_set, + bool persistent_tape) { std::vector tensor_stack; tensor_stack.reserve(target.size()); for (auto t : target) { @@ -362,7 +362,7 @@ BackpropInitialState PrepareBackprop( template std::vector InitialStack( const OpTape& op_tape, - const std::unordered_map& op_missing_tensor) { + const gtl::FlatMap& op_missing_tensor) { std::vector result; for (auto& op_entry : op_tape) { if (op_missing_tensor.find(op_entry.first) == op_missing_tensor.end()) { @@ -373,13 +373,13 @@ std::vector InitialStack( } template -Status InitialGradients( - const VSpace& vspace, - gtl::ArraySlice target_tensor_ids, - gtl::ArraySlice output_gradients, const TensorTape& tensor_tape, - const OpTape& op_tape, - const std::unordered_map& tensor_usage_counts, - std::unordered_map>* result) { +Status InitialGradients(const VSpace& vspace, + gtl::ArraySlice target_tensor_ids, + gtl::ArraySlice output_gradients, + const TensorTape& tensor_tape, + const OpTape& op_tape, + const gtl::FlatMap& tensor_usage_counts, + gtl::FlatMap>* result) { for (int i = 0; i < target_tensor_ids.size(); ++i) { const int64 id = target_tensor_ids[i]; if (tensor_usage_counts.find(id) != tensor_usage_counts.end()) { @@ -441,13 +441,13 @@ Status GradientTape::ComputeGradient( gtl::ArraySlice source_tensor_ids, gtl::ArraySlice output_gradients, std::vector* result) { - std::unordered_set sources_set(source_tensor_ids.begin(), - source_tensor_ids.end()); + gtl::FlatSet sources_set(source_tensor_ids.begin(), + source_tensor_ids.end()); BackpropInitialState state = PrepareBackprop( target_tensor_ids, tensor_tape_, &op_tape_, sources_set, persistent_); std::vector op_stack = InitialStack(state.op_tape, state.op_missing_tensor); - std::unordered_map> gradients; + gtl::FlatMap> gradients; Status s = InitialGradients(vspace, target_tensor_ids, output_gradients, tensor_tape_, state.op_tape, state.tensor_usage_counts, &gradients); @@ -463,7 +463,7 @@ Status GradientTape::ComputeGradient( cleanup(); return s; } - std::unordered_map gradients_size; + gtl::FlatMap gradients_size; // TODO(apassos) multiple threads could be dequeuing from op_stack at the same // time, for better CPU backprop performance. VLOG(1) << "Initial stack:"; @@ -472,11 +472,10 @@ Status GradientTape::ComputeGradient( VLOG(1) << " " << t; } } - std::unordered_map> - functions_accept_none_for_indices({ - {"SoftmaxCrossEntropyWithLogits", {1}}, - {"FusedBatchNorm", {1, 2, 3, 4}}, - }); + gtl::FlatMap> functions_accept_none_for_indices({ + {"SoftmaxCrossEntropyWithLogits", {1}}, + {"FusedBatchNorm", {1, 2, 3, 4}}, + }); while (!op_stack.empty()) { const int64 op = op_stack.back(); VLOG(1) << "Popped " << op; diff --git a/tensorflow/c/python_api.h b/tensorflow/c/python_api.h index b51ef2b53122802fef598a26bd6f1843976f11b0..aa9d9e06b28c54cb8869eb547d36ee3cb0d4e6b8 100644 --- a/tensorflow/c/python_api.h +++ b/tensorflow/c/python_api.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_C_PYTHON_API_H_ -#define THIRD_PARTY_TENSORFLOW_C_PYTHON_API_H_ +#ifndef TENSORFLOW_C_PYTHON_API_H_ +#define TENSORFLOW_C_PYTHON_API_H_ #include "tensorflow/c/c_api.h" @@ -39,4 +39,4 @@ void RemoveAllControlInputs(TF_Graph* graph, TF_Operation* op); } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_C_PYTHON_API_H_ +#endif // TENSORFLOW_C_PYTHON_API_H_ diff --git a/tensorflow/cc/BUILD b/tensorflow/cc/BUILD index e354831d7d25af83c068a68a4f844056263a598c..9060c19e9d2cf965c2b9be07be07c42017da45a8 100644 --- a/tensorflow/cc/BUILD +++ b/tensorflow/cc/BUILD @@ -421,7 +421,7 @@ tf_cc_test( tf_gen_op_wrappers_cc( name = "cc_ops", - api_def_srcs = ["//tensorflow/core:base_api_def"], + api_def_srcs = ["//tensorflow/core/api_def:base_api_def"], op_lib_names = [ "array_ops", "audio_ops", @@ -433,6 +433,7 @@ tf_gen_op_wrappers_cc( "linalg_ops", "logging_ops", "lookup_ops", + "manip_ops", "math_ops", "nn_ops", "no_op", @@ -448,7 +449,6 @@ tf_gen_op_wrappers_cc( "ops/const_op.h", "ops/standard_ops.h", ], - override_file = "ops/op_gen_overrides.pbtxt", pkg = "//tensorflow/core", ) @@ -527,14 +527,13 @@ cc_library_with_android_deps( ], copts = tf_copts(), data = [ - "//tensorflow/core:base_api_def", + "//tensorflow/core/api_def:base_api_def", ], deps = [ "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:op_gen_lib", - "//tensorflow/core:op_gen_overrides_proto_cc", "//tensorflow/core:proto_text", "//tensorflow/core:protos_all_cc", ], @@ -547,15 +546,11 @@ tf_cc_test( "framework/cc_op_gen.h", "framework/cc_op_gen_test.cc", ], - data = [ - "//tensorflow/cc:ops/op_gen_overrides.pbtxt", - ], deps = [ "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:op_gen_lib", - "//tensorflow/core:op_gen_overrides_proto_cc", "//tensorflow/core:proto_text", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", @@ -679,7 +674,6 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", - "//tensorflow/core:tensorflow", ], ) diff --git a/tensorflow/cc/client/client_session_test.cc b/tensorflow/cc/client/client_session_test.cc index dfbac9788e16e9c7c65abcd1ea213b51d5d5d060..ea5cf5a1f12be316cc6e0d0a02cd3caf4d177400 100644 --- a/tensorflow/cc/client/client_session_test.cc +++ b/tensorflow/cc/client/client_session_test.cc @@ -23,7 +23,13 @@ limitations under the License. #include "tensorflow/core/platform/test.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) +namespace { + +using ops::Add; +using ops::Const; +using ops::Mul; +using ops::Placeholder; +using ops::Sub; TEST(ClientSessionTest, Basic) { Scope root = Scope::NewRootScope(); @@ -89,4 +95,5 @@ TEST(ClientSessionTest, MultiThreaded) { test::ExpectTensorEqual(outputs[0], test::AsTensor({-1, 2}, {2})); } -} // end namespace tensorflow +} // namespace +} // namespace tensorflow diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc index d889c518f9c38a9f070970b37a2ad4b1fc26671b..a40ad1ffc3b262840e6ca0043139b1b61e04510d 100644 --- a/tensorflow/cc/framework/cc_op_gen.cc +++ b/tensorflow/cc/framework/cc_op_gen.cc @@ -1057,16 +1057,9 @@ string MakeInternal(const string& fname) { } // namespace void WriteCCOps(const OpList& ops, const ApiDefMap& api_def_map, - const string& dot_h_fname, const string& dot_cc_fname, - const string& overrides_fnames) { + const string& dot_h_fname, const string& dot_cc_fname) { Env* env = Env::Default(); - // Load the override map. - OpGenOverrideMap override_map; - if (!overrides_fnames.empty()) { - TF_CHECK_OK(override_map.LoadFileList(env, overrides_fnames)); - } - // Write the initial boilerplate to the .h and .cc files. std::unique_ptr h = nullptr; std::unique_ptr cc = nullptr; diff --git a/tensorflow/cc/framework/cc_op_gen.h b/tensorflow/cc/framework/cc_op_gen.h index cea28990144b9371e8009ce13f912b44044f9aac..c7256a7dc384e652fa1bddfe3aa9893491c2b14c 100644 --- a/tensorflow/cc/framework/cc_op_gen.h +++ b/tensorflow/cc/framework/cc_op_gen.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_ +#define TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_ #include "tensorflow/core/framework/op_def.pb.h" #include "tensorflow/core/framework/op_gen_lib.h" @@ -24,9 +24,8 @@ namespace tensorflow { /// Result is written to files dot_h and dot_cc. void WriteCCOps(const OpList& ops, const ApiDefMap& api_def_map, - const string& dot_h_fname, const string& dot_cc_fname, - const string& overrides_fnames); + const string& dot_h_fname, const string& dot_cc_fname); } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_ diff --git a/tensorflow/cc/framework/cc_op_gen_main.cc b/tensorflow/cc/framework/cc_op_gen_main.cc index 326d5668b8803ee39ffe24900c92e1db87b93601..3157792e15a006555e4924eea3c72ea643e79c1c 100644 --- a/tensorflow/cc/framework/cc_op_gen_main.cc +++ b/tensorflow/cc/framework/cc_op_gen_main.cc @@ -28,7 +28,7 @@ namespace tensorflow { namespace { void PrintAllCCOps(const std::string& dot_h, const std::string& dot_cc, - const std::string& overrides_fnames, bool include_internal, + bool include_internal, const std::vector& api_def_dirs) { OpList ops; OpRegistry::Global()->Export(include_internal, &ops); @@ -49,7 +49,7 @@ void PrintAllCCOps(const std::string& dot_h, const std::string& dot_cc, api_def_map.UpdateDocs(); - WriteCCOps(ops, api_def_map, dot_h, dot_cc, overrides_fnames); + WriteCCOps(ops, api_def_map, dot_h, dot_cc); } } // namespace @@ -57,24 +57,21 @@ void PrintAllCCOps(const std::string& dot_h, const std::string& dot_cc, int main(int argc, char* argv[]) { tensorflow::port::InitMain(argv[0], &argc, &argv); - // TODO(annarev): Update this file to no longer take op_gen_overrides.pbtxt - // as an argument. - if (argc != 6) { + if (argc != 5) { for (int i = 1; i < argc; ++i) { fprintf(stderr, "Arg %d = %s\n", i, argv[i]); } fprintf(stderr, - "Usage: %s out.h out.cc overrides1.pbtxt,2.pbtxt include_internal " + "Usage: %s out.h out.cc include_internal " "api_def_dirs1,api_def_dir2 ...\n" " include_internal: 1 means include internal ops\n", argv[0]); exit(1); } - bool include_internal = tensorflow::StringPiece("1") == argv[4]; + bool include_internal = tensorflow::StringPiece("1") == argv[3]; std::vector api_def_dirs = tensorflow::str_util::Split( - argv[5], ",", tensorflow::str_util::SkipEmpty()); - tensorflow::PrintAllCCOps(argv[1], argv[2], argv[3], include_internal, - api_def_dirs); + argv[4], ",", tensorflow::str_util::SkipEmpty()); + tensorflow::PrintAllCCOps(argv[1], argv[2], include_internal, api_def_dirs); return 0; } diff --git a/tensorflow/cc/framework/cc_op_gen_test.cc b/tensorflow/cc/framework/cc_op_gen_test.cc index 0b7e720a5c7b343415eee1aa157b8de755a1e1a5..1e0f2d241bb350897a840dda90d6d0c009b1daad 100644 --- a/tensorflow/cc/framework/cc_op_gen_test.cc +++ b/tensorflow/cc/framework/cc_op_gen_test.cc @@ -24,10 +24,6 @@ limitations under the License. namespace tensorflow { namespace { -// TODO(annarev): Remove this op_gen_overrides.pbtxt reference. -// It is needed only because WriteCCOps takes it as an argument. -constexpr char kOverridesFnames[] = - "tensorflow/cc/ops/op_gen_overrides.pbtxt"; constexpr char kBaseOpDef[] = R"( op { name: "Foo" @@ -96,7 +92,7 @@ void GenerateCcOpFiles(Env* env, const OpList& ops, const auto internal_h_file_path = io::JoinPath(tmpdir, "test_internal.h"); const auto internal_cc_file_path = io::JoinPath(tmpdir, "test_internal.cc"); - WriteCCOps(ops, api_def_map, h_file_path, cc_file_path, kOverridesFnames); + WriteCCOps(ops, api_def_map, h_file_path, cc_file_path); TF_ASSERT_OK(ReadFileToString(env, h_file_path, h_file_text)); TF_ASSERT_OK( diff --git a/tensorflow/cc/framework/cc_ops_test.cc b/tensorflow/cc/framework/cc_ops_test.cc index 5da23036eaadbef270ba839357dc4613bf3bf490..ac05e3cf95b1ce4009ee1424713baf2d34902a94 100644 --- a/tensorflow/cc/framework/cc_ops_test.cc +++ b/tensorflow/cc/framework/cc_ops_test.cc @@ -22,8 +22,7 @@ limitations under the License. #include "tensorflow/core/lib/core/status_test_util.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) - +namespace ops { namespace { Output Linear(const Scope& scope, Input x, Input w, Input b) { @@ -39,8 +38,6 @@ void GetColocationConstraints(const Output& tensor, constraints)); } -} // namespace - TEST(CCOpTest, Basic) { Scope root = Scope::NewRootScope(); auto c = Const(root, {{1, 1}}); @@ -249,4 +246,6 @@ TEST(CCOpTest, InvalidFinalize) { string::npos); } +} // namespace +} // namespace ops } // namespace tensorflow diff --git a/tensorflow/cc/framework/grad_op_registry.h b/tensorflow/cc/framework/grad_op_registry.h index 190b96f68506c6b5252d6c0184f1712310477a8a..0fc5abb20c884a66539682099497e2c8511a620f 100644 --- a/tensorflow/cc/framework/grad_op_registry.h +++ b/tensorflow/cc/framework/grad_op_registry.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_ +#define TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_ #include @@ -72,4 +72,4 @@ class GradOpRegistry { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_ diff --git a/tensorflow/cc/framework/gradient_checker.h b/tensorflow/cc/framework/gradient_checker.h index d055c60d09c2f33fb1f61165f75b2d04618620b7..1aa215a9088335580667e0c23c7244e6e5047f1a 100644 --- a/tensorflow/cc/framework/gradient_checker.h +++ b/tensorflow/cc/framework/gradient_checker.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_ +#define TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_ #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" @@ -60,4 +60,4 @@ Status ComputeGradientError(const Scope& scope, const Output& x, } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_ diff --git a/tensorflow/cc/framework/gradient_checker_test.cc b/tensorflow/cc/framework/gradient_checker_test.cc index fdc457f40af875d7c0c243246755d0cb87c44a62..d4f0a7f5ab3716be41e22c02a21aca028f76fb88 100644 --- a/tensorflow/cc/framework/gradient_checker_test.cc +++ b/tensorflow/cc/framework/gradient_checker_test.cc @@ -24,10 +24,18 @@ limitations under the License. #include "tensorflow/core/util/equal_graph_def.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) - namespace { +using ops::Complex; +using ops::Const; +using ops::MatMul; +using ops::Placeholder; +using ops::Real; +using ops::Split; +using ops::Square; +using ops::Stack; +using ops::Unstack; + TEST(GradientCheckerTest, BasicFloat) { Scope scope = Scope::NewRootScope(); TensorShape shape({2, 4, 3}); diff --git a/tensorflow/cc/framework/gradients.h b/tensorflow/cc/framework/gradients.h index 717f6f0636d3dd1a546ef7477b100bbfc86ba13d..0a377ad56d139a6ec26ea97b4e1e43495d0b3165 100644 --- a/tensorflow/cc/framework/gradients.h +++ b/tensorflow/cc/framework/gradients.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_ +#define TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_ #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" @@ -49,4 +49,4 @@ Output NoGradient(); } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_ diff --git a/tensorflow/cc/framework/gradients_test.cc b/tensorflow/cc/framework/gradients_test.cc index 07a062e704ed6ffc6389b5897309957a1bfcd1c2..26e3170ad8e4f4fba1c2dc014086acf24d949f72 100644 --- a/tensorflow/cc/framework/gradients_test.cc +++ b/tensorflow/cc/framework/gradients_test.cc @@ -26,10 +26,20 @@ limitations under the License. #include "tensorflow/core/util/equal_graph_def.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) - namespace { +using ops::Assign; +using ops::Const; +using ops::Identity; +using ops::MatMul; +using ops::OnesLike; +using ops::Placeholder; +using ops::Square; +using ops::Stack; +using ops::StopGradient; +using ops::Unstack; +using ops::Variable; + // TODO(andydavis) Add more unit tests once more gradient functions are ported. class GradientsTest : public ::testing::Test { protected: diff --git a/tensorflow/cc/framework/ops.h b/tensorflow/cc/framework/ops.h index 8d4154220c4b18f9286094b10c1b1e96eb4e31e7..a085e1d6e2de5ad63d11eb8979ae64c26b91366f 100644 --- a/tensorflow/cc/framework/ops.h +++ b/tensorflow/cc/framework/ops.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_OPS_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_OPS_H_ +#define TENSORFLOW_CC_FRAMEWORK_OPS_H_ #include @@ -296,4 +296,4 @@ class InputList { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_OPS_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_OPS_H_ diff --git a/tensorflow/cc/framework/scope.h b/tensorflow/cc/framework/scope.h index 0225ac047291d6297af558fddad6e5315389ff40..30c32bd44b0f22d6b29dd3836d431807d0216818 100644 --- a/tensorflow/cc/framework/scope.h +++ b/tensorflow/cc/framework/scope.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_SCOPE_H_ +#define TENSORFLOW_CC_FRAMEWORK_SCOPE_H_ #include #include @@ -242,4 +242,4 @@ struct CompositeOpScopes { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_SCOPE_H_ diff --git a/tensorflow/cc/framework/scope_internal.h b/tensorflow/cc/framework/scope_internal.h index 968c366550ef6f46557cd9b5662d9d0719b31531..8efcfed20d0b86d86d8c20a3d8630c7c6bc909c3 100644 --- a/tensorflow/cc/framework/scope_internal.h +++ b/tensorflow/cc/framework/scope_internal.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ +#define TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ #include "tensorflow/cc/framework/scope.h" @@ -117,4 +117,4 @@ class Scope::Impl { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ diff --git a/tensorflow/cc/framework/testutil.h b/tensorflow/cc/framework/testutil.h index a3e19870ec847bcd4f0e0bf0e71dda724024d5d2..7ad6fb4a676639f5d6d3da6a7c08de1894162f0c 100644 --- a/tensorflow/cc/framework/testutil.h +++ b/tensorflow/cc/framework/testutil.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_ +#define TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_ #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" @@ -44,4 +44,4 @@ void GetTensor(const Scope& scope, const std::vector& assign_vars, } // namespace test } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_ diff --git a/tensorflow/cc/framework/while_gradients.h b/tensorflow/cc/framework/while_gradients.h index 8f592accc93573cb8953a5ab25c04881ca0c2333..cb4e579c8548294ec45b0c3f42cb844e0b87c390 100644 --- a/tensorflow/cc/framework/while_gradients.h +++ b/tensorflow/cc/framework/while_gradients.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_ -#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_ +#ifndef TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_ +#define TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_ #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" @@ -37,4 +37,4 @@ Status AddWhileLoopGradient(WhileContext* while_ctx, const Scope& scope, } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_ +#endif // TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_ diff --git a/tensorflow/cc/gradients/array_grad_test.cc b/tensorflow/cc/gradients/array_grad_test.cc index 455d7330c10cf230462869475f25a1f1b9bf9e9e..4a215fcc9299cf8b8da04cbf151640631ed0d449 100644 --- a/tensorflow/cc/gradients/array_grad_test.cc +++ b/tensorflow/cc/gradients/array_grad_test.cc @@ -23,11 +23,11 @@ limitations under the License. #include "tensorflow/core/lib/core/status_test_util.h" namespace tensorflow { +namespace { + using namespace ops; // NOLINT(build/namespaces) using ops::internal::MirrorPadGrad; -namespace { - class ArrayGradTest : public ::testing::Test { protected: ArrayGradTest() : scope_(Scope::NewRootScope()) {} diff --git a/tensorflow/cc/gradients/data_flow_grad_test.cc b/tensorflow/cc/gradients/data_flow_grad_test.cc index 734dfd3af97b856a7c8c4894c4a6d1a3ade10992..0ba3c0e27b1e545a30925ea3ef9e2c54dc9d0ae9 100644 --- a/tensorflow/cc/gradients/data_flow_grad_test.cc +++ b/tensorflow/cc/gradients/data_flow_grad_test.cc @@ -23,10 +23,13 @@ limitations under the License. #include "tensorflow/core/lib/random/random.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) - namespace { +using ops::Const; +using ops::DynamicPartition; +using ops::DynamicStitch; +using ops::Placeholder; + class DataFlowGradTest : public ::testing::Test { protected: DataFlowGradTest() : scope_(Scope::NewRootScope()) {} diff --git a/tensorflow/cc/gradients/grad_testutil.cc b/tensorflow/cc/gradients/grad_testutil.cc index 04b29d4e8b21eeee200d9e7390868d701eda3c22..304117d3719346202d3a8a18637f7c915d4a47f9 100644 --- a/tensorflow/cc/gradients/grad_testutil.cc +++ b/tensorflow/cc/gradients/grad_testutil.cc @@ -18,16 +18,14 @@ limitations under the License. #include "tensorflow/cc/framework/grad_op_registry.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) - namespace test { Status CallGradFunction(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - GradFunc grad_fn; - TF_RETURN_IF_ERROR( - GradOpRegistry::Global()->Lookup(op.node()->type_string(), &grad_fn)); + ops::GradFunc grad_fn; + TF_RETURN_IF_ERROR(ops::GradOpRegistry::Global()->Lookup( + op.node()->type_string(), &grad_fn)); TF_RETURN_IF_ERROR(grad_fn(scope, op, grad_inputs, grad_outputs)); TF_RETURN_IF_ERROR(scope.status()); return Status::OK(); diff --git a/tensorflow/cc/gradients/grad_testutil.h b/tensorflow/cc/gradients/grad_testutil.h index d31f412754ff59cc7782b14e285071a8d4218d08..70c81f1a73a394322c602a5c51e3c2a40aca2397 100644 --- a/tensorflow/cc/gradients/grad_testutil.h +++ b/tensorflow/cc/gradients/grad_testutil.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_ +#ifndef TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_ +#define TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_ #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" @@ -32,4 +32,4 @@ Status CallGradFunction(const Scope& scope, const Operation& op, } // namespace test } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_ +#endif // TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_ diff --git a/tensorflow/cc/gradients/math_grad.cc b/tensorflow/cc/gradients/math_grad.cc index afd92fbf488c014083284012f01de7c7ac2f090d..52c177212a8c88f1857defcc38de4a01ac47dab0 100644 --- a/tensorflow/cc/gradients/math_grad.cc +++ b/tensorflow/cc/gradients/math_grad.cc @@ -763,24 +763,6 @@ Status LgammaGrad(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("Lgamma", LgammaGrad); -Status SelectGrad(const Scope& scope, const Operation& op, - const std::vector& grad_inputs, - std::vector* grad_outputs) { - auto comparator = op.input(0); - auto x = op.input(1); - auto zeros = ZerosLike(scope, x); - auto grad = grad_inputs[0]; - - auto gx_1 = Where3(scope, comparator, grad, zeros); - auto gx_2 = Where3(scope, comparator, zeros, grad); - - grad_outputs->push_back(NoGradient()); - grad_outputs->push_back(gx_1); - grad_outputs->push_back(gx_2); - return scope.status(); -} -REGISTER_GRADIENT_OP("Select", SelectGrad); - Status MinOrMaxGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { diff --git a/tensorflow/cc/gradients/math_grad_test.cc b/tensorflow/cc/gradients/math_grad_test.cc index b94d797711c9cabbe14f370a7d26c22f71c20410..1b4c7c2688083e74433da3dce2849b8c37443684 100644 --- a/tensorflow/cc/gradients/math_grad_test.cc +++ b/tensorflow/cc/gradients/math_grad_test.cc @@ -23,10 +23,31 @@ limitations under the License. #include "tensorflow/core/lib/random/random.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) - namespace { +using ops::Abs; +using ops::Add; +using ops::AddN; +using ops::BatchMatMul; +using ops::Const; +using ops::Div; +using ops::Greater; +using ops::MatMul; +using ops::Max; +using ops::Maximum; +using ops::Mean; +using ops::Min; +using ops::Minimum; +using ops::Mul; +using ops::Placeholder; +using ops::Pow; +using ops::Prod; +using ops::RealDiv; +using ops::SquaredDifference; +using ops::Sub; +using ops::Sum; +using ops::Where3; + // TODO(andydavis) Test gradient function against numeric gradients output. // TODO(andydavis) As more gradients are added move common test functions // to a testutil library. @@ -83,6 +104,7 @@ class CWiseUnaryGradTest : public ::testing::Test { Output y; switch (op_type) { + using namespace ops; // NOLINT(build/namespaces) case ABS: y = Abs(scope_, x); break; @@ -882,13 +904,5 @@ TEST_F(NaryGradTest, Prod) { RunTest({x}, {x_shape}, {y}, {y_shape}); } -TEST_F(NaryGradTest, Select) { - TensorShape shape({3, 4}); - auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)); - auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)); - auto y = Where3(scope_, Greater(scope_, x1, x2), x1, x2); - RunTest({x1, x2}, {shape, shape}, {y}, {shape}); -} - } // namespace } // namespace tensorflow diff --git a/tensorflow/cc/gradients/nn_grad_test.cc b/tensorflow/cc/gradients/nn_grad_test.cc index f9063e836509669d81d03b1d2f0d32d1166b6eca..0cfe5f6e3c49f7c4a3cafbf48ff4e54a0ffd0d47 100644 --- a/tensorflow/cc/gradients/nn_grad_test.cc +++ b/tensorflow/cc/gradients/nn_grad_test.cc @@ -23,10 +23,22 @@ limitations under the License. #include "tensorflow/core/lib/random/random.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) - namespace { +using ops::BiasAdd; +using ops::Conv2D; +using ops::Elu; +using ops::L2Loss; +using ops::LogSoftmax; +using ops::LRN; +using ops::MaxPool; +using ops::MaxPoolV2; +using ops::Placeholder; +using ops::Relu; +using ops::Relu6; +using ops::Selu; +using ops::Softmax; + class NNGradTest : public ::testing::Test { protected: NNGradTest() : scope_(Scope::NewRootScope()) {} diff --git a/tensorflow/cc/ops/const_op.h b/tensorflow/cc/ops/const_op.h index d11fda475b3db58bf83cdb94079c8fde8d1170f7..424a683665f31b5e25eeceeb40477fc31640ce90 100644 --- a/tensorflow/cc/ops/const_op.h +++ b/tensorflow/cc/ops/const_op.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_OPS_CONST_OP_H_ -#define THIRD_PARTY_TENSORFLOW_CC_OPS_CONST_OP_H_ +#ifndef TENSORFLOW_CC_OPS_CONST_OP_H_ +#define TENSORFLOW_CC_OPS_CONST_OP_H_ #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" @@ -82,4 +82,4 @@ std::vector AsNodeOutList(const Scope& scope, } // namespace ops } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_OPS_CONST_OP_H_ +#endif // TENSORFLOW_CC_OPS_CONST_OP_H_ diff --git a/tensorflow/cc/ops/op_gen_overrides.pbtxt b/tensorflow/cc/ops/op_gen_overrides.pbtxt deleted file mode 100644 index 4aac990e748b0a79cbc3b353b4121a582b0883b0..0000000000000000000000000000000000000000 --- a/tensorflow/cc/ops/op_gen_overrides.pbtxt +++ /dev/null @@ -1,238 +0,0 @@ -# array_ops -op { name: "BroadcastArgs" rename_to: "BroadcastDynamicShape" } -op { name: "BroadcastGradientArgs" hide: true } -op { name: "ConcatOffset" skip: true } # Maybe should just be hidden? -op { name: "Concat" skip: true } -op { name: "ConcatV2" rename_to: "Concat" } -op { name: "ExpandDims" input_rename: { from: "dim" to: "axis" } } -op { name: "ListDiff" rename_to: "SetDiff1D" } -op { name: "MirrorPadGrad" hide: true } -op { name: "Reverse" skip: true } -op { name: "ReverseV2" rename_to: "Reverse" } -op { name: "Split" input_rename: { from: "split_dim" to: "axis" } } -op { name: "SplitV" input_rename: { from: "split_dim" to: "axis" } } -op { name: "Squeeze" attr_rename: { from: "squeeze_dims" to: "axis" } } -op { name: "Pack" rename_to: "Stack" } -op { name: "Unpack" rename_to: "Unstack" } -op { name: "Select" rename_to: "Where3" input_rename: { from: "t" to: "x" } input_rename: { from: "e" to: "y" } } -op { name: "Where" input_rename: { from: "input" to: "condition" } } - - -# candidate_sampling_ops -op { name: "ThreadUnsafeUnigramCandidateSampler", skip: true } - -# control_flow_ops -# TODO(joshl): Hide Switch and Merge once we write and migrate users to -# a Cond() API. -#op { name: "Switch" hide: true } -#op { name: "Merge" hide: true } -op { name: "RefMerge" hide: true } -op { name: "Exit" hide: true } -op { name: "RefExit" hide: true } -op { name: "Enter" hide: true } -op { name: "RefEnter" hide: true } -op { name: "RefIdentity" hide: true } - -# ctc_ops - -# data_flow_ops -op { name: "FakeQueue" skip: true } -op { name: "FIFOQueue" skip: true} -op { name: "FIFOQueueV2" rename_to: "FIFOQueue" } -op { name: "PaddingFIFOQueue" skip: true } -op { name: "PaddingFIFOQueueV2" rename_to: "PaddingFIFOQueue" } -op { name: "PriorityQueue" skip: true } -op { name: "PriorityQueueV2" rename_to: "PriorityQueue" } -op { name: "QueueClose" skip: true } -op { name: "QueueCloseV2" rename_to: "QueueClose" } -op { name: "QueueDequeue" skip: true } -op { name: "QueueDequeueV2" rename_to: "QueueDequeue" } -op { name: "QueueDequeueMany" skip: true } -op { name: "QueueDequeueManyV2" rename_to: "QueueDequeueMany" } -op { name: "QueueDequeueUpTo" skip: true } -op { name: "QueueDequeueUpToV2" rename_to: "QueueDequeueUpTo" } -op { name: "QueueEnqueue" skip: true } -op { name: "QueueEnqueueV2" rename_to: "QueueEnqueue" } -op { name: "QueueEnqueueMany" skip: true } -op { name: "QueueEnqueueManyV2" rename_to: "QueueEnqueueMany" } -op { name: "QueueSize" skip: true } -op { name: "QueueSizeV2" rename_to: "QueueSize" } -op { name: "RandomShuffleQueue" skip: true } -op { name: "RandomShuffleQueueV2" rename_to: "RandomShuffleQueue" } -op { name: "ReaderNumRecordsProduced" skip: true } -op { name: "ReaderNumRecordsProducedV2" rename_to: "ReaderNumRecordsProduced" } -op { name: "ReaderNumWorkUnitsCompleted" skip: true } -op { name: "ReaderNumWorkUnitsCompletedV2" rename_to: "ReaderNumWorkUnitsCompleted" } -op { name: "ReaderRead" skip: true } -op { name: "ReaderReadUpTo" skip: true } -op { name: "ReaderReadUpToV2" rename_to: "ReaderReadUpTo" } -op { name: "ReaderReadV2" rename_to: "ReaderRead" } -op { name: "ReaderReset" skip: true } -op { name: "ReaderResetV2" rename_to: "ReaderReset" } -op { name: "ReaderRestoreState" skip: true } -op { name: "ReaderRestoreStateV2" rename_to: "ReaderRestoreState" } -op { name: "ReaderSerializeState" skip: true } -op { name: "ReaderSerializeStateV2" rename_to: "ReaderSerializeState" } -op { name: "FixedLengthRecordReader" skip: true } -op { name: "FixedLengthRecordReaderV2" rename_to: "FixedLengthRecordReader" } -op { name: "IdentityReader" skip: true } -op { name: "IdentityReaderV2" rename_to: "IdentityReader" } -op { name: "TFRecordReader" skip: true } -op { name: "TFRecordReaderV2" rename_to: "TFRecordReader" } -op { name: "TextLineReader" skip: true } -op { name: "TextLineReaderV2" rename_to: "TextLineReader" } - -# Skip hash table ops until we have better support in C++ (ops are currently -# only used in contrib) -op { name: "HashTable" skip: true } -op { name: "InitializeTable" skip: true } -op { name: "InitializeTableFromTextFile" skip: true } -op { name: "LookupTableFind" skip: true } -op { name: "LookupTableImport" skip: true } -op { name: "LookupTableInsert" skip: true } -op { name: "LookupTableSize" skip: true } -op { name: "MutableDenseHashTable" skip: true } -op { name: "MutableHashTable" skip: true } -op { name: "MutableHashTableOfTensors" skip: true } - -# Stack ops are internal to control flow gradients (not yet implemented in C++) -op { name: "Stack" skip: true } -op { name: "StackClose" skip: true } -op { name: "StackPop" skip: true } -op { name: "StackPush" skip: true } -op { name: "StackV2" skip: true } -op { name: "StackCloseV2" skip: true } -op { name: "StackPopV2" skip: true } -op { name: "StackPushV2" skip: true } - -op { name: "TensorArrayCloseV2" skip: true } -op { name: "TensorArrayCloseV3" rename_to: "TensorArrayClose" } -op { name: "TensorArrayConcatV2" skip: true } -op { name: "TensorArrayConcatV3" rename_to: "TensorArrayConcat" } -op { name: "TensorArrayGatherV2" skip: true } -op { name: "TensorArrayGatherV3" rename_to: "TensorArrayGather" } -op { name: "TensorArrayGradV2" skip: true } -op { name: "TensorArrayGradV3" rename_to: "TensorArrayGrad" } -op { name: "TensorArrayReadV2" skip: true } -op { name: "TensorArrayReadV3" rename_to: "TensorArrayRead" } -op { name: "TensorArrayScatterV2" skip: true } -op { name: "TensorArrayScatterV3" rename_to: "TensorArrayScatter" } -op { name: "TensorArraySizeV2" skip: true } -op { name: "TensorArraySizeV3" rename_to: "TensorArraySize" } -op { name: "TensorArraySplitV2" skip: true } -op { name: "TensorArraySplitV3" rename_to: "TensorArraySplit" } -op { name: "TensorArrayV2" skip: true } -op { name: "TensorArrayV3" rename_to: "TensorArray" } -op { name: "TensorArrayWriteV2" skip: true } -op { name: "TensorArrayWriteV3" rename_to: "TensorArrayWrite" } - -op { name: "WholeFileReader" skip: true } -op { name: "WholeFileReaderV2" rename_to: "WholeFileReader" } - -# functional_ops - -# image_ops -op { name: "AdjustContrastv2" rename_to: "AdjustContrast" } -op { name: "ResizeBilinearGrad" hide: true } -op { name: "ResizeBicubicGrad" hide: true } -op { name: "ResizeNearestNeighborGrad" hide: true } - -# io_ops - -# linalg_ops -op { name: "SelfAdjointEigV2" rename_to: "SelfAdjointEig" } - -# logging_ops -op { name: "AudioSummaryV2" rename_to: "AudioSummary" } - -# lookup_ops -op { name: "LookupTableFind" skip: true } -op { name: "LookupTableFindV2" rename_to: "LookupTableFind" } -op { name: "LookupTableInsert" skip: true } -op { name: "LookupTableInsertV2" rename_to: "LookupTableInsert" } -op { name: "LookupTableSize" skip: true } -op { name: "LookupTableSizeV2" rename_to: "LookupTableSize" } -op { name: "LookupTableExport" skip: true } -op { name: "LookupTableExportV2" rename_to: "LookupTableExport" } -op { name: "LookupTableImport" skip: true } -op { name: "LookupTableImportV2" rename_to: "LookupTableImport" } -op { name: "HashTable" skip: true } -op { name: "HashTableV2" rename_to: "HashTable" } -op { name: "MutableHashTable" skip: true } -op { name: "MutableHashTableV2" rename_to: "MutableHashTable" } -op { name: "MutableHashTableOfTensors" skip: true } -op { name: "MutableHashTableOfTensorsV2" rename_to: "MutableHashTableOfTensors" } -op { name: "MutableDenseHashTable" skip: true } -op { name: "MutableDenseHashTableV2" rename_to: "MutableDenseHashTable" } -op { name: "InitializeTable" skip: true } -op { name: "InitializeTableV2" rename_to: "InitializeTable" } -op { name: "InitializeTableFromTextFile" skip: true } -op { name: "InitializeTableFromTextFileV2" rename_to: "InitializeTableFromTextFile" } - -# math_ops -op { name: "All" alias: "ReduceAll" input_rename: { from: "reduction_indices" to: "axis" } } -op { name: "Any" alias: "ReduceAny" input_rename: { from: "reduction_indices" to: "axis" } } -op { name: "Max" alias: "ReduceMax" input_rename: { from: "reduction_indices" to: "axis" } } -op { name: "Mean" alias: "ReduceMean" input_rename: { from: "reduction_indices" to: "axis" } } -op { name: "Min" alias: "ReduceMin" input_rename: { from: "reduction_indices" to: "axis" } } -op { name: "Mul" rename_to: "Multiply" alias: "Mul" } -op { name: "Neg" rename_to: "Negate" alias: "Neg" } -op { name: "Prod" alias: "ReduceProd" input_rename: { from: "reduction_indices" to: "axis" } } -op { name: "Sub" rename_to: "Subtract" alias: "Sub" } -op { name: "Sum" alias: "ReduceSum" input_rename: { from: "reduction_indices" to: "axis" } } -op { name: "SigmoidGrad" hide: true } -op { name: "TanhGrad" hide: true } -op { name: "InvGrad" hide: true } -op { name: "ReciprocalGrad" hide: true } -op { name: "SqrtGrad" hide: true } -op { name: "RsqrtGrad" hide: true } - -# *Grad ops get hidden, only for use by the gradient code. -op { name: "SigmoidGrad" hide: true } -op { name: "TanhGrad" hide: true } -op { name: "InvGrad" hide: true } -op { name: "ReciprocalGrad" hide: true } -op { name: "SqrtGrad" hide: true } -op { name: "RsqrtGrad" hide: true } - -# nn_ops -op { name: "AvgPoolGrad" hide: true } -op { name: "LRNGrad" hide: true } -op { name: "MaxPoolGrad" hide: true } -op { name: "MaxPoolGradWithArgmax" hide: true } -op { name: "ReluGrad" hide: true } -op { name: "Relu6Grad" hide: true } -op { name: "EluGrad" hide: true } -op { name: "SeluGrad" hide: true } -op { name: "SoftplusGrad" hide: true } -op { name: "SoftsignGrad" hide: true } -op { name: "FractionalAvgPoolGrad" hide: true } -op { name: "FractionalMaxPoolGrad" hide: true } -op { name: "TopKV2" rename_to: "TopK" } -op { name: "BiasAddV1" skip: true } # Use BiasAdd instead - -# parsing_ops - -# random_ops - -op { name: "RandomStandardNormal" rename_to: "RandomNormal" } -# script_ops -# Calling Python functions from a C++ program isn't supported -op { name: "PyFunc" skip: true } -op { name: "PyFuncStateless" skip: true} - -# sdca_ops - -# state_ops - -op { name: "Variable" skip: true } -op { name: "VariableV2" rename_to: "Variable" } - -# sparse_ops - -# string_ops - -# user_ops - -# training_ops - diff --git a/tensorflow/cc/ops/standard_ops.h b/tensorflow/cc/ops/standard_ops.h index 0c021f0b3ac02c596e0511e650a3caa0002c25d1..98f53010ecf78f769c7d89d6aafc48fdb772f42e 100644 --- a/tensorflow/cc/ops/standard_ops.h +++ b/tensorflow/cc/ops/standard_ops.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_OPS_STANDARD_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CC_OPS_STANDARD_OPS_H_ +#ifndef TENSORFLOW_CC_OPS_STANDARD_OPS_H_ +#define TENSORFLOW_CC_OPS_STANDARD_OPS_H_ #include "tensorflow/cc/ops/array_ops.h" #include "tensorflow/cc/ops/candidate_sampling_ops.h" @@ -37,4 +37,4 @@ limitations under the License. #include "tensorflow/cc/ops/training_ops.h" #include "tensorflow/cc/ops/user_ops.h" -#endif // THIRD_PARTY_TENSORFLOW_CC_OPS_STANDARD_OPS_H_ +#endif // TENSORFLOW_CC_OPS_STANDARD_OPS_H_ diff --git a/tensorflow/cc/ops/while_loop.cc b/tensorflow/cc/ops/while_loop.cc index e0251efb2a424f86bd5a4885ef22d1928e04bd3e..d1c918d464bc9684b0db6dade2fb80cb2bd6691a 100644 --- a/tensorflow/cc/ops/while_loop.cc +++ b/tensorflow/cc/ops/while_loop.cc @@ -116,7 +116,7 @@ Status CreateCond(const Scope& scope, const CondGraphBuilderFn& cond, return Status::OK(); } -// Create the bdoy subgraph defined by `body`. `outputs` must be non-null and +// Create the body subgraph defined by `body`. `outputs` must be non-null and // empty. Status CreateBody(const Scope& scope, const BodyGraphBuilderFn& body, const std::vector& inputs, diff --git a/tensorflow/cc/ops/while_loop.h b/tensorflow/cc/ops/while_loop.h index a04476056a058ff0951a6347e8ffc05bc5ff5023..727237b5c7ad4d31dba1aaaf6d5600773d69223e 100644 --- a/tensorflow/cc/ops/while_loop.h +++ b/tensorflow/cc/ops/while_loop.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_OPS_WHILE_LOOP_H_ -#define THIRD_PARTY_TENSORFLOW_CC_OPS_WHILE_LOOP_H_ +#ifndef TENSORFLOW_CC_OPS_WHILE_LOOP_H_ +#define TENSORFLOW_CC_OPS_WHILE_LOOP_H_ #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" @@ -71,4 +71,4 @@ Status BuildWhileLoop(const Scope& scope, const std::vector& inputs, } // namespace ops } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_OPS_WHILE_LOOP_H_ +#endif // TENSORFLOW_CC_OPS_WHILE_LOOP_H_ diff --git a/tensorflow/cc/profiler/BUILD b/tensorflow/cc/profiler/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..00799526fce572e7bb80199ccb8ce1cc89874031 --- /dev/null +++ b/tensorflow/cc/profiler/BUILD @@ -0,0 +1,36 @@ +package( + default_visibility = ["//visibility:public"], +) + +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test") + +tf_cuda_cc_test( + name = "profiler_test", + srcs = ["profiler_test.cc"], + deps = [ + ":profiler", + "//tensorflow/cc:cc_ops", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:tensorflow", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +cc_library( + name = "profiler", + srcs = ["profiler.cc"], + hdrs = ["profiler.h"], + deps = [ + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core/profiler:protos_all_cc", + "//tensorflow/core/profiler:tfprof_options", + "//tensorflow/core/profiler/internal:tfprof_stats", + ], +) diff --git a/tensorflow/cc/profiler/profiler.cc b/tensorflow/cc/profiler/profiler.cc new file mode 100644 index 0000000000000000000000000000000000000000..3e55bac73e6d32a1fa5ddcc1937744e2cf56657d --- /dev/null +++ b/tensorflow/cc/profiler/profiler.cc @@ -0,0 +1,57 @@ +/* Copyright 2017 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/cc/profiler/profiler.h" + +namespace tensorflow { +namespace tfprof { + +Profiler::Profiler(const GraphDef& graph) { + std::unique_ptr graph_ptr(new GraphDef()); + *graph_ptr = graph; + stats_.reset(new TFStats(std::move(graph_ptr), nullptr, nullptr, nullptr)); +} + +void Profiler::AddStep(int64 step, const RunMetadata& run_meta) { + std::unique_ptr run_meta_ptr(new RunMetadata()); + *run_meta_ptr = run_meta; + stats_->AddRunMeta(step, std::move(run_meta_ptr)); +} + +GraphNodeProto Profiler::ProfileGraph(const Options& options) { + stats_->BuildView(kCmds[1]); + return stats_->ShowGraphNode(kCmds[1], options); +} + +GraphNodeProto Profiler::ProfileNameScope(const Options& options) { + stats_->BuildView(kCmds[0]); + return stats_->ShowGraphNode(kCmds[0], options); +} + +MultiGraphNodeProto Profiler::ProfileOperations(const Options& options) { + stats_->BuildView(kCmds[3]); + return stats_->ShowMultiGraphNode(kCmds[3], options); +} + +Status Profiler::SerializeToString(string* content) { + if (!content) { + return Status(error::Code::INVALID_ARGUMENT, + "Cannot use null string pointer for SerializeToString."); + } + stats_->SerializeToString(content); + return Status::OK(); +} + +} // namespace tfprof +} // namespace tensorflow diff --git a/tensorflow/cc/profiler/profiler.h b/tensorflow/cc/profiler/profiler.h new file mode 100644 index 0000000000000000000000000000000000000000..6077c45c5854fd5812ccb7c91522f93ed4e54883 --- /dev/null +++ b/tensorflow/cc/profiler/profiler.h @@ -0,0 +1,97 @@ +/* Copyright 2017 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_CC_PROFILER_PROFILER_H_ +#define TENSORFLOW_CC_PROFILER_PROFILER_H_ + +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/profiler/internal/tfprof_stats.h" +#include "tensorflow/core/profiler/tfprof_options.h" +#include "tensorflow/core/profiler/tfprof_output.pb.h" + +namespace tensorflow { +namespace tfprof { + +/// @addtogroup core +/// @{ + +/// A `Profiler` object lets the caller profile the execution of a graph. +/// +/// Example: +/// // First build a graph and run tracing. +/// Scope root = Scope::NewRootScope(); +/// auto a = Placeholder(root, DT_INT32); +/// auto c = Add(root, a, {41}); +/// +/// ClientSession session(root); +/// std::vector outputs; +/// RunOptions run_options; +/// run_options.set_trace_level(RunOptions::FULL_TRACE); +/// RunMetadata run_meta; +/// Status s = session.Run(run_options, { {a, {1}} }, {c}, &outputs, +/// &run_meta); +/// if (!s.ok()) { ... } +/// +/// // Then create profiler to do profiling. +/// GraphDef graph; +/// root.ToGraphDef(&graph); +/// Profiler profiler(graph); +/// profiler.AddStep(0, run_meta); +/// Options opts = ... // TODO(xpan): Support option building API. +/// MultiGraphNodeProto r = profiler.ProfileOperations(opts); +/// +class Profiler { + public: + /// `graph` is the model's GraphDef. + Profiler(const GraphDef& graph); + + /// 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`. + /// 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. + 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. + /// E.g. /foo is the root of operation /foo/matmul_1 and foo/conv_2. + GraphNodeProto ProfileNameScope(const Options& options); + + /// Profiles the model by organizing nodes by operation types. + /// Each node is an operation type (e.g. Conv2D or MatMul), containing all + /// ops belonging to that type in the model. + MultiGraphNodeProto ProfileOperations(const Options& options); + + /// Serialize the profile content (ProfileProto) into a binary string, + /// User can write the string to file for offline analysis by + /// tfprof command-line tools or graphical user interface. + Status SerializeToString(string* content); + + private: + std::unique_ptr stats_; +}; +/// @} + +} // namespace tfprof +} // namespace tensorflow + +#endif // TENSORFLOW_CC_PROFILER_PROFILER_H_ diff --git a/tensorflow/cc/profiler/profiler_test.cc b/tensorflow/cc/profiler/profiler_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..280cd74827fc8ae80737eaf61286535fec959aa8 --- /dev/null +++ b/tensorflow/cc/profiler/profiler_test.cc @@ -0,0 +1,177 @@ +/* Copyright 2017 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/platform/test.h" + +#include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/cc/profiler/profiler.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/graph/default_device.h" +#include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/public/session.h" + +namespace tensorflow { +namespace tfprof { + +class ProfilerTest : public ::testing::Test { + protected: + ProfilerTest() {} +}; + +GraphDef CreateGraphDef() { + Scope root = Scope::NewRootScope(); + + auto a = ops::Const(root, {{3, 2}, {-1, 0}}); + + auto x = ops::Const(root.WithOpName("x"), {{1.f}, {1.f}}); + + auto y = ops::MatMul(root.WithOpName("y"), a, x); + + auto y2 = ops::Square(root, y); + + auto y2_sum = ops::Sum(root, y2, 0); + + auto y_norm = ops::Sqrt(root, y2_sum); + + auto y_div = ops::Div(root.WithOpName("y_normalized"), y, y_norm); + + GraphDef def; + TF_CHECK_OK(root.ToGraphDef(&def)); + + return def; +} + +Options Default() { + Options opts(1000, /* max_depth */ + 0, /* min_bytes */ + 0, /* min_peak_bytes */ + 0, /* min_residual_bytes */ + 0, /* min_output_bytes */ + 0, /* min_micros */ + 0, /* min_accelerator_micros */ + 0, /* min_cpu_micros */ + 0, /* min_params */ + 0, /* min_float_ops */ + 0, /* min_occurrence */ + 0, /* step */ + "name", /* order_by */ + {".*"}, /* account_type_regexes */ + {".*"}, /* start_name_regexes */ + {}, /* trim_name_regexes */ + {".*"}, {}, /* hide_name_regexes */ + false, /* account_displayed_op_only */ + {"micros"}, /* select */ + {"none"}, /* output_type */ + {}); + return opts; +} + +template +const T* ExtractNode(const T& pb, const string& name) { + if (pb.name() == name) { + return &pb; + } + for (const T& c : pb.children()) { + const T* ret = ExtractNode(c, name); + if (ret) return ret; + } + return nullptr; +} + +TEST_F(ProfilerTest, Basics) { + SessionOptions options; + options.config.set_allow_soft_placement(true); + std::unique_ptr session(NewSession(options)); + GraphDef def = CreateGraphDef(); + if (options.target.empty()) { + graph::SetDefaultDevice("/gpu:0", &def); + } + + TF_CHECK_OK(session->Create(def)); + + Tensor x(DT_FLOAT, TensorShape({2, 1})); + auto x_flat = x.flat(); + x_flat.setRandom(); + Eigen::Tensor inv_norm = + x_flat.square().sum().sqrt().inverse(); + x_flat = x_flat * inv_norm(); + + std::vector outputs; + RunOptions run_options; + run_options.set_trace_level(RunOptions::FULL_TRACE); + RunMetadata run_metadata; + outputs.clear(); + + Profiler profiler(def); + for (int i = 0; i < 2; ++i) { + TF_CHECK_OK(session->Run(run_options, {{"x", x}}, {"y:0", "y_normalized:0"}, + {}, &outputs, &run_metadata)); + profiler.AddStep(i, run_metadata); + CHECK_EQ(size_t{2}, outputs.size()); + } + + std::vector resp; + TF_CHECK_OK(session->ListDevices(&resp)); + bool has_gpu = false; + for (const auto& dev : resp) { + if (dev.device_type() == "GPU") { + has_gpu = true; + } + } + + GraphNodeProto ret = profiler.ProfileNameScope(Default()); + const GraphNodeProto* matmul = ExtractNode(ret, "y"); + EXPECT_TRUE(matmul); + EXPECT_GT(matmul->exec_micros(), 0); + if (has_gpu) { + EXPECT_GT(matmul->accelerator_exec_micros(), 0); + } else { + EXPECT_EQ(matmul->accelerator_exec_micros(), 0); + } + const GraphNodeProto* square = ExtractNode(ret, "Square"); + EXPECT_TRUE(square); + EXPECT_GT(square->exec_micros(), 0); + if (has_gpu) { + EXPECT_GT(square->accelerator_exec_micros(), 0); + } else { + EXPECT_EQ(square->accelerator_exec_micros(), 0); + } + + Options opts2 = Default(); + opts2.output_type = "timeline"; + string timeline_file = io::JoinPath(testing::TmpDir(), "timeline"); + opts2.output_options["outfile"] = timeline_file; + GraphNodeProto ret2 = profiler.ProfileGraph(opts2); + string s; + TF_CHECK_OK(ReadFileToString(Env::Default(), timeline_file + "_0", &s)); + EXPECT_TRUE(s.find("Square") != s.npos); + + MultiGraphNodeProto ret3 = profiler.ProfileOperations(Default()); + const MultiGraphNodeProto* matmul2 = ExtractNode(ret3, "MatMul"); + EXPECT_TRUE(matmul2); + EXPECT_GT(matmul2->exec_micros(), 0); + if (has_gpu) { + EXPECT_GT(matmul2->accelerator_exec_micros(), 0); + } else { + EXPECT_EQ(matmul2->accelerator_exec_micros(), 0); + } + + TF_CHECK_OK(session->Close()); +} + +} // namespace tfprof +} // namespace tensorflow diff --git a/tensorflow/cc/saved_model/constants.h b/tensorflow/cc/saved_model/constants.h index c940df8a8761d97a859be3af30980ff79ca3577a..645a3f101d1ae7dda88ec4ca622c694dc5a7a919 100644 --- a/tensorflow/cc/saved_model/constants.h +++ b/tensorflow/cc/saved_model/constants.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_ -#define THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_ +#ifndef TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_ +#define TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_ namespace tensorflow { @@ -47,4 +47,4 @@ constexpr char kSavedModelVariablesFilename[] = "variables"; } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_ +#endif // TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_ diff --git a/tensorflow/cc/saved_model/loader.cc b/tensorflow/cc/saved_model/loader.cc index acef098c7d07f45d171679bff7c41e13ef0424f1..faa1e378d07ea94ad08ee084d18bf6a113f054af 100644 --- a/tensorflow/cc/saved_model/loader.cc +++ b/tensorflow/cc/saved_model/loader.cc @@ -96,7 +96,9 @@ Status FindMetaGraphDefToLoad(const SavedModel& saved_model_proto, Status LoadMetaGraphIntoSession(const MetaGraphDef& meta_graph_def, const SessionOptions& session_options, std::unique_ptr* session) { - session->reset(NewSession(session_options)); + Session* session_p = nullptr; + TF_RETURN_IF_ERROR(NewSession(session_options, &session_p)); + session->reset(session_p); return (*session)->Create(meta_graph_def.graph_def()); } diff --git a/tensorflow/cc/saved_model/loader.h b/tensorflow/cc/saved_model/loader.h index 3d634dd51543bed8d3c074bdc56c251f97d56976..a8e098fa5440e7a8f72fd0b52737dcb06435b908 100644 --- a/tensorflow/cc/saved_model/loader.h +++ b/tensorflow/cc/saved_model/loader.h @@ -15,8 +15,8 @@ limitations under the License. /// SavedModel loading functions and SavedModelBundle struct. -#ifndef THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ -#define THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ +#ifndef TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ +#define TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ #include #include @@ -61,4 +61,4 @@ bool MaybeSavedModelDirectory(const string& export_dir); } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ +#endif // TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ diff --git a/tensorflow/cc/saved_model/loader_test.cc b/tensorflow/cc/saved_model/loader_test.cc index 0ad6b33bba5fcceaca68e2f179cef2232c689a80..4c64d2cfe3c10e6c7ed82a2d72460a0b34283bb2 100644 --- a/tensorflow/cc/saved_model/loader_test.cc +++ b/tensorflow/cc/saved_model/loader_test.cc @@ -155,6 +155,24 @@ TEST_F(LoaderTest, NoTagMatchMultiple) { << st.error_message(); } +TEST_F(LoaderTest, SessionCreationFailure) { + SavedModelBundle bundle; + // Use invalid SessionOptions to cause session creation to fail. Default + // options work, so provide an invalid value for the target field. + SessionOptions session_options; + constexpr char kInvalidTarget[] = "invalid target"; + session_options.target = kInvalidTarget; + RunOptions run_options; + + const string export_dir = + io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded); + Status st = LoadSavedModel(session_options, run_options, export_dir, + {kSavedModelTagServe}, &bundle); + EXPECT_FALSE(st.ok()); + EXPECT_TRUE(StringPiece(st.error_message()).contains(kInvalidTarget)) + << st.error_message(); +} + TEST_F(LoaderTest, PbtxtFormat) { SavedModelBundle bundle; SessionOptions session_options; diff --git a/tensorflow/cc/saved_model/python/BUILD b/tensorflow/cc/saved_model/python/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..f5fbc75edcba9d5ae9ef7432de224df766bcab9e --- /dev/null +++ b/tensorflow/cc/saved_model/python/BUILD @@ -0,0 +1,30 @@ +# Description: +# CLIF wrappers for TensorFlow SavedModels. + +licenses(["notice"]) # Apache 2.0 + +package( + default_visibility = ["//visibility:public"], +) + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +load("//tensorflow/core:platform/default/build_config.bzl", "tf_py_clif_cc") + +tf_py_clif_cc( + name = "loader", + srcs = ["loader.clif"], + deps = [ + "//tensorflow/cc/saved_model:loader", + ], +) diff --git a/tensorflow/cc/saved_model/python/loader.clif b/tensorflow/cc/saved_model/python/loader.clif new file mode 100644 index 0000000000000000000000000000000000000000..b102757d2eeb46ee713d8ed0d0c3d66b58740ee0 --- /dev/null +++ b/tensorflow/cc/saved_model/python/loader.clif @@ -0,0 +1,4 @@ +from "third_party/tensorflow/cc/saved_model/loader.h": + namespace `tensorflow`: + class SavedModelBundle: + def __init__(self) diff --git a/tensorflow/cc/saved_model/signature_constants.h b/tensorflow/cc/saved_model/signature_constants.h index b2d39bd55beb48a05489236395a208e41deb9c8f..7d8c07f5cf0a310c20193469cb6d18664f738d96 100644 --- a/tensorflow/cc/saved_model/signature_constants.h +++ b/tensorflow/cc/saved_model/signature_constants.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_ -#define THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_ +#ifndef TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_ +#define TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_ namespace tensorflow { @@ -66,4 +66,4 @@ static constexpr char kRegressOutputs[] = "outputs"; } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_ +#endif // TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_ diff --git a/tensorflow/cc/saved_model/tag_constants.h b/tensorflow/cc/saved_model/tag_constants.h index b71cb263ca42dab7e830c1880ec4b311bc272f82..68a090e0c4cf79cfa87771a80447b8112fc37fb9 100644 --- a/tensorflow/cc/saved_model/tag_constants.h +++ b/tensorflow/cc/saved_model/tag_constants.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_ -#define THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_ +#ifndef TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_ +#define TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_ namespace tensorflow { @@ -32,4 +32,4 @@ constexpr char kSavedModelTagTrain[] = "train"; } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_ +#endif // TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_ diff --git a/tensorflow/cc/tools/BUILD b/tensorflow/cc/tools/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..97f66e79b8ad9f383b22f56e9385fc6d2080e1f8 --- /dev/null +++ b/tensorflow/cc/tools/BUILD @@ -0,0 +1,57 @@ +# Description: +# TensorFlow cc tools. + +package( + default_visibility = ["//visibility:public"], +) + +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_test", +) + +cc_library( + name = "freeze_saved_model", + srcs = ["freeze_saved_model.cc"], + hdrs = ["freeze_saved_model.h"], + deps = [ + "//tensorflow/cc/saved_model:loader", + "//tensorflow/core:core_cpu", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], +) + +tf_cc_test( + name = "freeze_saved_model_test", + srcs = ["freeze_saved_model_test.cc"], + deps = [ + ":freeze_saved_model", + "//tensorflow/cc:cc_ops", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + ], +) + +# ----------------------------------------------------------------------------- +# Google-internal targets. + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/cc/tools/freeze_saved_model.cc b/tensorflow/cc/tools/freeze_saved_model.cc new file mode 100644 index 0000000000000000000000000000000000000000..ddf372cdef21e1b3892c9a03714478d5a5785517 --- /dev/null +++ b/tensorflow/cc/tools/freeze_saved_model.cc @@ -0,0 +1,194 @@ +/* Copyright 2017 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/cc/tools/freeze_saved_model.h" + +#include + +#include "tensorflow/core/framework/attr_value.pb.h" +#include "tensorflow/core/framework/function.pb.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/versions.pb.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/protobuf/meta_graph.pb.h" + +namespace tensorflow { + +namespace { + +// Gets tensor names from tensor_info and inserts them into the set of tensor +// names. +void GetTensorNamesFromTensorInfo(const TensorInfo& tensor_info, + std::unordered_set* tensor_names) { + if (tensor_info.has_coo_sparse()) { + // If the tensor is sparse we have to add all three tensors of the sparse + // representations. + const TensorInfo_CooSparse& coo_sparse = tensor_info.coo_sparse(); + tensor_names->insert(coo_sparse.values_tensor_name()); + tensor_names->insert(coo_sparse.indices_tensor_name()); + tensor_names->insert(coo_sparse.dense_shape_tensor_name()); + } else { + tensor_names->insert(tensor_info.name()); + } +} + +// Gets the union of all inputs and outputs of all SignatureDefs in the bundle +void GetSignatureDefsInputsAndOutputs( + const SavedModelBundle& saved_model_bundle, + std::unordered_set* inputs, std::unordered_set* outputs) { + for (auto& sigdef_elem : saved_model_bundle.meta_graph_def.signature_def()) { + const SignatureDef& signature_def = sigdef_elem.second; + for (auto& input_elem : signature_def.inputs()) { + GetTensorNamesFromTensorInfo(input_elem.second, inputs); + } + for (auto& output_elem : signature_def.outputs()) { + GetTensorNamesFromTensorInfo(output_elem.second, outputs); + } + } +} + +// Gets a map from string node name to NodeDef. +void GetNodeNameToNodeDefMap( + GraphDef* graph_def, + std::unordered_map* name_to_node_map) { + for (size_t i = 0; i < graph_def->node_size(); i++) { + NodeDef* node = graph_def->mutable_node(i); + (*name_to_node_map)[node->name()] = node; + } +} + +// Gets the set of node names needed by `outputs` and the corresponding set of +// variable nodes to convert. +void GetReachableNodesAndVariables( + GraphDef* graph_def, const std::unordered_set& outputs, + std::unordered_set* reachable_node_names, + std::unordered_set* variable_node_names) { + // TODO(suharshs): Add support for ResourceVariables. + static const std::unordered_set* kVariableTypes = + new std::unordered_set({"Variable", "VariableV2"}); + // name_to_node_map is needed to get the inputs from the NodeDef corresponding + // the a string node name. These inputs are used when doing our backwards + // traversal. + std::unordered_map name_to_node_map; + GetNodeNameToNodeDefMap(graph_def, &name_to_node_map); + std::queue nodes_to_visit; + for (const string& tensor_name : outputs) { + // We need to strip off the tensor part to get the node name. + std::vector tensor_name_parts = str_util::Split(tensor_name, ':'); + nodes_to_visit.push(tensor_name_parts[0]); + } + // We do a traversal backwards from the outputs specified in the MetaGraphDef. + while (!nodes_to_visit.empty()) { + const string node_name = nodes_to_visit.front(); + nodes_to_visit.pop(); + if (reachable_node_names->find(node_name) != reachable_node_names->end()) { + continue; + } + reachable_node_names->insert(node_name); + NodeDef* node = name_to_node_map[node_name]; + if (kVariableTypes->find(node->op()) != kVariableTypes->end()) { + variable_node_names->insert(node->name()); + } + for (const string& input : node->input()) { + nodes_to_visit.push(input); + } + } +} + +// Gets a map from variable name to variable value. +Status GetVariableNameToTensorMap( + Session* session, std::unordered_set variable_names_set, + std::unordered_map* variable_name_to_value_map) { + if (variable_names_set.empty()) { + return Status::OK(); + } + std::vector variable_names; + std::vector tensor_names; + for (const string& node_name : variable_names_set) { + variable_names.push_back(node_name); + // We need to run tensors, so append ":0". + tensor_names.push_back(node_name + ":0"); + } + std::vector outputs; + TF_RETURN_IF_ERROR( + session->Run(/* inputs */ {}, tensor_names, /* targets */ {}, &outputs)); + for (size_t i = 0; i < variable_names.size(); i++) { + (*variable_name_to_value_map)[variable_names[i]] = outputs[i]; + } + return Status::OK(); +} + +// Converts a Variable NodeDef into a Constant NodeDef. +void ConvertVariableToConstant(const NodeDef& variable_node, + const Tensor& variable_value, + NodeDef* const_node) { + const_node->set_name(variable_node.name()); + const_node->set_op("Const"); + (*const_node->mutable_attr())["dtype"] = variable_node.attr().at("dtype"); + variable_value.AsProtoTensorContent( + (*const_node->mutable_attr())["value"].mutable_tensor()); +} + +// Freezes the subgraph of all nodes needed by `outputs`. +Status FreezeGraphDef(const SavedModelBundle& saved_model_bundle, + const std::unordered_set& outputs, + GraphDef* frozen_graph_def) { + GraphDef graph_def = saved_model_bundle.meta_graph_def.graph_def(); + // Copy versions and library as-is from original graph. + *frozen_graph_def->mutable_versions() = graph_def.versions(); + *frozen_graph_def->mutable_library() = graph_def.library(); + // If the graph is empty there is nothing left to do. + if (graph_def.node_size() == 0) { + return Status::OK(); + } + std::unordered_set reachable_node_names; + std::unordered_set variable_node_names; + GetReachableNodesAndVariables(&graph_def, outputs, &reachable_node_names, + &variable_node_names); + std::unordered_map variable_to_value_map; + TF_RETURN_IF_ERROR( + GetVariableNameToTensorMap(saved_model_bundle.session.get(), + variable_node_names, &variable_to_value_map)); + // We copy the nodes in the same order they were in the original graph_def. + for (const NodeDef& node : graph_def.node()) { + if (reachable_node_names.find(node.name()) == reachable_node_names.end()) { + continue; + } + if (variable_node_names.find(node.name()) != variable_node_names.end()) { + ConvertVariableToConstant(node, variable_to_value_map[node.name()], + frozen_graph_def->add_node()); + } else { + // If the node isn't a variable, just copy the node as-is. + *frozen_graph_def->add_node() = node; + } + } + return Status::OK(); +} + +} // namespace + +Status FreezeSavedModel(const SavedModelBundle& saved_model_bundle, + GraphDef* frozen_graph_def, + std::unordered_set* inputs, + std::unordered_set* outputs) { + GetSignatureDefsInputsAndOutputs(saved_model_bundle, inputs, outputs); + TF_RETURN_IF_ERROR( + FreezeGraphDef(saved_model_bundle, *outputs, frozen_graph_def)); + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/cc/tools/freeze_saved_model.h b/tensorflow/cc/tools/freeze_saved_model.h new file mode 100644 index 0000000000000000000000000000000000000000..b10f29805a4515f9d49426cc41e0d375cd32b072 --- /dev/null +++ b/tensorflow/cc/tools/freeze_saved_model.h @@ -0,0 +1,43 @@ +/* Copyright 2017 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_CC_TOOLS_FREEZE_SAVED_MODEL_H_ +#define TENSORFLOW_CC_TOOLS_FREEZE_SAVED_MODEL_H_ + +#include + +#include "tensorflow/cc/saved_model/loader.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { + +// Returns a frozen GraphDef, input tensors, and output tensors from the loaded +// SavedModelBundle. +// `inputs` and `outputs` consist of the union of all inputs and outputs in the +// SignatureDefs in the SavedModelBundle. +// FreezeSavedModel sets `frozen_graph_def` to a GraphDef of all nodes needed by +// `outputs`. All variables in the supplied SavedModelBundle are converted to +// constants, set to the value of the variables, by running the restored Session +// in the SavedModelBundle. +// WARNING: Only the variable checkpoints will be reflected in the frozen +// graph_def. All saved_model assets will be ignored. +Status FreezeSavedModel(const SavedModelBundle& saved_model_bundle, + GraphDef* frozen_graph_def, + std::unordered_set* inputs, + std::unordered_set* outputs); + +} // namespace tensorflow + +#endif // TENSORFLOW_CC_TOOLS_FREEZE_SAVED_MODEL_H_ diff --git a/tensorflow/cc/tools/freeze_saved_model_test.cc b/tensorflow/cc/tools/freeze_saved_model_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..52a81a50284aec36bba4e56a0232c886cb0cb6cf --- /dev/null +++ b/tensorflow/cc/tools/freeze_saved_model_test.cc @@ -0,0 +1,307 @@ +/* Copyright 2017 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/cc/tools/freeze_saved_model.h" + +#include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/core/framework/function_testlib.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/tensor_testutil.h" +#include "tensorflow/core/framework/versions.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/public/session.h" +#include "tensorflow/core/public/session_options.h" + +namespace tensorflow { +namespace { + +class FreezeTest : public ::testing::Test { + protected: + void GraphDefEqual(const GraphDef& actual, const GraphDef& expected) { + EXPECT_EQ(actual.ShortDebugString(), expected.ShortDebugString()); + } + + // Builds a SignatureDef with the provided `inputs` and `outputs`. + SignatureDef BuildSignatureDef(const std::unordered_set& inputs, + const std::unordered_set& outputs) { + SignatureDef signature_def; + for (const string& input : inputs) { + (*signature_def.mutable_inputs())[input].set_name(input); + } + for (const string& output : outputs) { + (*signature_def.mutable_outputs())[output].set_name(output); + } + return signature_def; + } + + // Adds `signature_def` to `saved_model_bundle` under `key`. + void AddSignatureDefToSavedModelBundle(const SignatureDef& signature_def, + const string& key, + SavedModelBundle* saved_model_bundle) { + MetaGraphDef* meta_graph_def = &saved_model_bundle->meta_graph_def; + (*meta_graph_def->mutable_signature_def())[key] = signature_def; + } + + // Adds an initialized session to `saved_model_bundle` using `graph_def` and + // initializing with `init_node`. + Status InitializeSavedModelBundleSession( + const GraphDef& graph_def, const string& init_node, + SavedModelBundle* saved_model_bundle) { + SessionOptions session_options; + saved_model_bundle->session.reset(NewSession(session_options)); + TF_RETURN_IF_ERROR(saved_model_bundle->session->Create(graph_def)); + if (!init_node.empty()) { + std::vector outputs; + return saved_model_bundle->session->Run( + /* inputs */ {}, /* output_tensors */ {}, {init_node}, &outputs); + } + return Status::OK(); + } + + // Adds `graph_def` to `saved_model_bundle` and initializes a session with + // `init_node`. + Status AddGraphDefToSavedModelBundle(const GraphDef& graph_def, + const string& init_node, + SavedModelBundle* saved_model_bundle) { + MetaGraphDef* meta_graph_def = &saved_model_bundle->meta_graph_def; + *meta_graph_def->mutable_graph_def() = graph_def; + return InitializeSavedModelBundleSession(graph_def, init_node, + saved_model_bundle); + } + + // Adds `graph_def` and `outputs` as the GraphDef and SignatureDef in + // `saved_model_bundle` and initializes a session with `init_node`. + Status AddGraphDefWithOutputsToSavedModelBundle( + const GraphDef& graph_def, const std::unordered_set& outputs, + const string& init_node, SavedModelBundle* saved_model_bundle) { + SignatureDef signature_def = + BuildSignatureDef(std::unordered_set(), outputs); + AddSignatureDefToSavedModelBundle(signature_def, "signature_def", + saved_model_bundle); + return AddGraphDefToSavedModelBundle(graph_def, init_node, + saved_model_bundle); + } + + // Runs and compares the outputs of `tensor_name` on both the + // `unfrozen_session` and the `frozen_graph_def. + void RunAndCompareFrozenAndUnfrozenGraphs(Session* unfrozen_session, + const GraphDef& frozen_graph_def, + const string& tensor_name) { + std::vector unfrozen_outputs; + TF_ASSERT_OK(unfrozen_session->Run(/* inputs */ {}, {tensor_name}, + /* targets */ {}, &unfrozen_outputs)); + + SessionOptions session_options; + std::unique_ptr frozen_session(NewSession(session_options)); + TF_ASSERT_OK(frozen_session->Create(frozen_graph_def)); + std::vector frozen_outputs; + TF_ASSERT_OK(frozen_session->Run(/* inputs */ {}, {tensor_name}, + /* targets */ {}, &frozen_outputs)); + + test::ExpectTensorEqual(unfrozen_outputs[0], frozen_outputs[0]); + } +}; + +TEST_F(FreezeTest, InputsAndOutputsSingleSignatureDef) { + // Test that inputs and outputs get correctly populated for a single + // SignatureDef. + SavedModelBundle saved_model_bundle; + std::unordered_set expected_inputs = {"input0:0", "input1:0"}; + std::unordered_set expected_outputs = {"output0:0", "output1:0"}; + SignatureDef signature_def = + BuildSignatureDef(expected_inputs, expected_outputs); + AddSignatureDefToSavedModelBundle(signature_def, "signature_def", + &saved_model_bundle); + GraphDef frozen_graph_def; + std::unordered_set inputs; + std::unordered_set outputs; + TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs, + &outputs)); + EXPECT_EQ(expected_inputs, inputs); + EXPECT_EQ(expected_outputs, outputs); +} + +TEST_F(FreezeTest, InputsAndOutputsMultipleSignatureDefs) { + // Test that inputs and outputs get correctly merged and populated when + // multiple SignatureDefs are provided. + SavedModelBundle saved_model_bundle; + SignatureDef signature_def_0 = BuildSignatureDef({"input0:0"}, {"output0:0"}); + SignatureDef signature_def_1 = BuildSignatureDef({"input1:0"}, {"output1:0"}); + AddSignatureDefToSavedModelBundle(signature_def_0, "signature_def_0", + &saved_model_bundle); + AddSignatureDefToSavedModelBundle(signature_def_1, "signature_def_1", + &saved_model_bundle); + GraphDef frozen_graph_def; + std::unordered_set inputs; + std::unordered_set outputs; + TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs, + &outputs)); + std::unordered_set expected_inputs = {"input0:0", "input1:0"}; + std::unordered_set expected_outputs = {"output0:0", "output1:0"}; + EXPECT_EQ(expected_inputs, inputs); + EXPECT_EQ(expected_outputs, outputs); +} + +TEST_F(FreezeTest, GraphDefVersionsAndLibrary) { + // Test that GraphDef versions and library are copied correctly into the + // frozen graph. + SavedModelBundle saved_model_bundle; + GraphDef graph_def; + graph_def.mutable_versions()->set_producer(1234); + graph_def.mutable_versions()->set_min_consumer(1234); + *graph_def.mutable_library()->add_function() = test::function::NonZero(); + TF_ASSERT_OK( + AddGraphDefToSavedModelBundle(graph_def, "", &saved_model_bundle)); + + GraphDef frozen_graph_def; + std::unordered_set inputs; + std::unordered_set outputs; + TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs, + &outputs)); + + GraphDefEqual(frozen_graph_def, graph_def); +} + +TEST_F(FreezeTest, GraphDefWithNoVariables) { + // Test freezing a graph with no variables. + SavedModelBundle saved_model_bundle; + GraphDef graph_def; + Scope scope = Scope::NewRootScope(); + Output a = ops::Const(scope.WithOpName("a"), 10.0f, {}); + Output b = ops::Const(scope.WithOpName("b"), 10.0f, {}); + Output c = ops::Mul(scope.WithOpName("c"), a, b); + TF_ASSERT_OK(scope.ToGraphDef(&graph_def)); + TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle(graph_def, {"c:0"}, "", + &saved_model_bundle)); + + GraphDef frozen_graph_def; + std::unordered_set inputs; + std::unordered_set outputs; + TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs, + &outputs)); + + GraphDefEqual(frozen_graph_def, graph_def); +} + +TEST_F(FreezeTest, GraphDefWithVariablesNotNeededByOutputs) { + // Test freezing a graph with variables that are not needed by the outputs in + // the SignatureDef. The resulting graph shouldn't be frozen, but + // non-dependent nodes should be pruned. + SavedModelBundle saved_model_bundle; + GraphDef graph_def; + Scope scope = Scope::NewRootScope(); + Output a = ops::Const(scope.WithOpName("a"), 10.0f, {}); + Output b = ops::Const(scope.WithOpName("b"), 10.0f, {}); + Output c = ops::Mul(scope.WithOpName("c"), a, b); + Output var = ops::Variable(scope.WithOpName("var"), {}, DataType::DT_FLOAT); + Output assign = ops::Assign(scope.WithOpName("assign"), var, a); + TF_ASSERT_OK(scope.ToGraphDef(&graph_def)); + // "c" isnt dependent on the variable, so nothing should be frozen. + TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle( + graph_def, {"c:0"}, assign.name(), &saved_model_bundle)); + + GraphDef frozen_graph_def; + std::unordered_set inputs; + std::unordered_set outputs; + TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs, + &outputs)); + + GraphDef expected_graph_def; + Scope expected_scope = Scope::NewRootScope(); + Output expected_a = ops::Const(expected_scope.WithOpName("a"), 10.0f, {}); + Output expected_b = ops::Const(expected_scope.WithOpName("b"), 10.0f, {}); + Output expected_c = + ops::Mul(expected_scope.WithOpName("c"), expected_a, expected_b); + TF_ASSERT_OK(expected_scope.ToGraphDef(&expected_graph_def)); + + GraphDefEqual(frozen_graph_def, expected_graph_def); + + RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(), + frozen_graph_def, "c:0"); +} + +TEST_F(FreezeTest, GraphDefWithVariablesNeededByOutputs) { + // Test freezing a graph with variables that are needed by outputs in the + // SignatureDef. The variables should be frozen. + SavedModelBundle saved_model_bundle; + GraphDef graph_def; + Scope scope = Scope::NewRootScope(); + Output a = ops::Const(scope.WithOpName("a"), 10.0f, {}); + Output var = ops::Variable(scope.WithOpName("var"), {}, DataType::DT_FLOAT); + Output c = ops::Mul(scope.WithOpName("c"), a, var); + Output assign = ops::Assign(scope.WithOpName("assign"), var, a); + TF_ASSERT_OK(scope.ToGraphDef(&graph_def)); + // "c" isnt dependent on the variable, so nothing should be frozen. + TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle( + graph_def, {"c:0"}, assign.name(), &saved_model_bundle)); + + GraphDef frozen_graph_def; + std::unordered_set inputs; + std::unordered_set outputs; + TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs, + &outputs)); + + // There should be 3 nodes in the resulting graph_def, and none should be + // variables. + EXPECT_EQ(frozen_graph_def.node_size(), 3); + for (const NodeDef& node : frozen_graph_def.node()) { + EXPECT_NE(node.op(), "Variable") << node.name(); + EXPECT_NE(node.op(), "VariableV2") << node.name(); + } + + RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(), + frozen_graph_def, "c:0"); +} + +TEST_F(FreezeTest, GraphDefWithVariablesNeededAndNotNeededByOutputs) { + // Test freezing a graph with some variables that are needed and not needed by + // the outputs in the SignatureDef. The resulting graph should only freeze + // dependent variables. + SavedModelBundle saved_model_bundle; + GraphDef graph_def; + Scope scope = Scope::NewRootScope(); + Output a = ops::Const(scope.WithOpName("a"), 10.0f, {}); + Output var = ops::Variable(scope.WithOpName("var"), {}, DataType::DT_FLOAT); + Output c = ops::Mul(scope.WithOpName("c"), a, var); + Output assign = ops::Assign(scope.WithOpName("assign"), var, a); + Output var_1 = + ops::Variable(scope.WithOpName("var_1"), {}, DataType::DT_FLOAT); + Output assign_1 = ops::Assign(scope.WithOpName("assign_1"), var, a); + TF_ASSERT_OK(scope.ToGraphDef(&graph_def)); + // "c" isnt dependent on the variable, so nothing should be frozen. + TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle( + graph_def, {"c:0"}, assign.name(), &saved_model_bundle)); + + GraphDef frozen_graph_def; + std::unordered_set inputs; + std::unordered_set outputs; + TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs, + &outputs)); + + // There should be 3 nodes in the resulting graph_def, and none should be + // variables. + EXPECT_EQ(frozen_graph_def.node_size(), 3); + for (const NodeDef& node : frozen_graph_def.node()) { + EXPECT_NE(node.op(), "Variable") << node.name(); + EXPECT_NE(node.op(), "VariableV2") << node.name(); + } + + RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(), + frozen_graph_def, "c:0"); +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/cc/training/coordinator.h b/tensorflow/cc/training/coordinator.h index 0e01b19cd98bc797b7bb25da55c05d96f3eb93c7..7168b775251d38687d604b5294405389a8c1b04f 100644 --- a/tensorflow/cc/training/coordinator.h +++ b/tensorflow/cc/training/coordinator.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_TRAINING_COORDINATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CC_TRAINING_COORDINATOR_H_ +#ifndef TENSORFLOW_CC_TRAINING_COORDINATOR_H_ +#define TENSORFLOW_CC_TRAINING_COORDINATOR_H_ #include #include @@ -128,4 +128,4 @@ class Coordinator { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_TRAINING_COORDINATOR_H_ +#endif // TENSORFLOW_CC_TRAINING_COORDINATOR_H_ diff --git a/tensorflow/cc/training/queue_runner.h b/tensorflow/cc/training/queue_runner.h index 2d3450032388bfee96055f23cf621af0fa4731ae..21189b4b046b87b8609483109096fda6144681b8 100644 --- a/tensorflow/cc/training/queue_runner.h +++ b/tensorflow/cc/training/queue_runner.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_ -#define THIRD_PARTY_TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_ +#ifndef TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_ +#define TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_ #include #include @@ -137,4 +137,4 @@ class QueueRunner : public RunnerInterface { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_ +#endif // TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_ diff --git a/tensorflow/compiler/aot/BUILD b/tensorflow/compiler/aot/BUILD index 5740c040e309bad8d7e3bdc468c09a3323fb99e0..0900e87ebabd378e6237b77ca0ef01677c07c244 100644 --- a/tensorflow/compiler/aot/BUILD +++ b/tensorflow/compiler/aot/BUILD @@ -52,6 +52,7 @@ cc_library( "flags.h", ], deps = [ + ":embedded_protocol_buffers", ":runtime", # needed by codegen to print aligned_buffer_bytes "//tensorflow/compiler/tf2xla", "//tensorflow/compiler/tf2xla:common", @@ -68,9 +69,7 @@ cc_library( "//tensorflow/compiler/xla/client:compile_only_client", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service/cpu:cpu_compiler", - "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", - "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", @@ -80,13 +79,18 @@ cc_library( tf_cc_test( name = "codegen_test", srcs = ["codegen_test.cc"], - data = ["codegen_test_h.golden"], + data = [ + "codegen_test_h.golden", + "codegen_test_o.golden", + ], deps = [ ":tfcompile_lib", "//tensorflow/compiler/xla:shape_util", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", + "@llvm//:support", # fixdeps: keep + "@llvm//:x86_code_gen", # fixdeps: keep ], ) @@ -128,7 +132,9 @@ tf_library( config = "test_graph_tfadd.config.pbtxt", cpp_class = "AddComp", graph = "test_graph_tfadd.pbtxt", - tags = ["manual"], + tags = [ + "manual", + ], ) # A test of tf_library that includes a graph with an unknown op, but where @@ -139,7 +145,9 @@ tf_library( config = "test_graph_tfunknownop.config.pbtxt", cpp_class = "UnknownOpAddComp", graph = "test_graph_tfunknownop.pbtxt", - tags = ["manual"], + tags = [ + "manual", + ], ) # A test of tf_library that includes a graph with an unknown op, but where @@ -151,7 +159,9 @@ tf_library( config = "test_graph_tfunknownop2.config.pbtxt", cpp_class = "UnknownOpAddComp", graph = "test_graph_tfunknownop.pbtxt", - tags = ["manual"], + tags = [ + "manual", + ], ) # A test of tf_library that includes a graph with an unknown op, but where @@ -162,7 +172,9 @@ tf_library( config = "test_graph_tfunknownop3.config.pbtxt", cpp_class = "UnknownOpAddComp", graph = "test_graph_tfunknownop.pbtxt", - tags = ["manual"], + tags = [ + "manual", + ], ) # Utility library for benchmark binaries, used by the *_benchmark rules that are @@ -185,11 +197,27 @@ cc_library( name = "benchmark_extra_android", tags = [ "manual", - "notap", ], visibility = ["//visibility:public"], ) +cc_library( + name = "embedded_protocol_buffers", + srcs = ["embedded_protocol_buffers.cc"], + hdrs = ["embedded_protocol_buffers.h"], + deps = [ + "//tensorflow/compiler/tf2xla:common", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", + "//tensorflow/core:lib", + "@llvm//:core", + "@llvm//:execution_engine", + "@llvm//:support", + "@llvm//:target", + ], +) + tf_cc_test( name = "benchmark_test", srcs = ["benchmark_test.cc"], diff --git a/tensorflow/compiler/aot/codegen.cc b/tensorflow/compiler/aot/codegen.cc index 53da2881b60db9ad39565567623eb86f754559af..2cae85e8965216eaaee4d3032015d0016258a5c1 100644 --- a/tensorflow/compiler/aot/codegen.cc +++ b/tensorflow/compiler/aot/codegen.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "tensorflow/compiler/aot/embedded_protocol_buffers.h" #include "tensorflow/compiler/aot/runtime.h" #include "tensorflow/compiler/tf2xla/str_util.h" #include "tensorflow/compiler/tf2xla/tf2xla_util.h" @@ -263,49 +264,6 @@ string GenNameToIndexCode(const T& entries, bool generate) { return code; } -// Converts the given `str` into a comma-separated list of per-character values. -string StringToCharList(const string& str) { - string list; - for (const char c : str) { - if (!list.empty()) { - list += ","; - } - list += strings::StrCat(static_cast(c)); - } - return list; -} - -string GenProgramShapeCode(xla::ProgramShape program_shape, bool generate) { - // No need for any static magic if we're not supposed to generate the data. - if (!generate) { - return "{\n return nullptr;\n }"; - } - // The parameter names are currently meaningless, and redundant with the rest - // of our metadata, so clear them out to avoid confusion and save space. - program_shape.clear_parameter_names(); - const string proto_str = program_shape.SerializeAsString(); - // Embed the program shape as a serialized protobuf in the header file. - // - // TODO(toddw): This strategy will likely fail for larger protobufs, depending - // on the C++ compiler that is used. Figure out another solution if necessary. - string code = R"({ - static const xla::ProgramShape* kShape = []() { - static const char kProto[] = {{{PROTO_LIST}}}; - static constexpr int kProtoSize = {{PROTO_SIZE}}; - xla::ProgramShape* shape = new xla::ProgramShape; - shape->ParseFromArray(kProto, kProtoSize); - return shape; - }(); - return kShape; - })"; - str_util::ReplaceAllPairs( - &code, { - {"{{PROTO_LIST}}", StringToCharList(proto_str)}, - {"{{PROTO_SIZE}}", strings::StrCat(proto_str.size())}, - }); - return code; -} - Status ValidateFeedFetchCppNames(const tf2xla::Config& config) { for (const tf2xla::Feed& feed : config.feed()) { if (!feed.name().empty()) { @@ -322,8 +280,9 @@ Status ValidateFeedFetchCppNames(const tf2xla::Config& config) { } // namespace -Status GenerateHeader(const HeaderOpts& opts, const tf2xla::Config& config, - const CompileResult& compile_result, string* header) { +Status GenerateHeader(const CodegenOpts& opts, const tf2xla::Config& config, + const CompileResult& compile_result, + const MetadataResult& metadata_result, string* header) { TF_RETURN_IF_ERROR(ValidateConfig(config)); TF_RETURN_IF_ERROR(ValidateFeedFetchCppNames(config)); const int64 result_index = compile_result.aot->result_buffer_index(); @@ -373,8 +332,6 @@ Status GenerateHeader(const HeaderOpts& opts, const tf2xla::Config& config, ? R"(#include "tensorflow/compiler/xla/xla_data.pb.h")" : ""; - const string program_shape_code = - GenProgramShapeCode(ps, opts.gen_program_shape); // Use a poor-man's text templating mechanism; first populate the full header // with placeholder tokens, and then rewrite the tokens with real values. @@ -402,6 +359,8 @@ extern "C" void {{ENTRY}}( void* result, const xla::ExecutableRunOptions* run_options, const void** args, void** temps, tensorflow::int64* profile_counters); +{{DECLS_FROM_OBJ_FILE}} + {{NS_START}} // {{CLASS}} represents a computation previously specified in a // TensorFlow graph, now compiled into executable code. This extends the generic @@ -524,7 +483,10 @@ class {{CLASS}} : public tensorflow::XlaCompiledCpuFunction { static const char** StaticResultNames() {{RESULT_NAMES_CODE}} // Shape of the args and results. - static const xla::ProgramShape* StaticProgramShape() {{PROGRAM_SHAPE_CODE}} + static const xla::ProgramShape* StaticProgramShape() { + static const xla::ProgramShape* kShape = {{PROGRAM_SHAPE_SHIM_EXPRESSION}}; + return kShape; + } }; {{NS_END}} @@ -547,18 +509,62 @@ class {{CLASS}} : public tensorflow::XlaCompiledCpuFunction { {"{{NS_END}}\n", ns_end}, {"{{NS_START}}\n", ns_start}, {"{{PROGRAM_SHAPE}}", xla::ShapeUtil::HumanString(ps)}, - {"{{PROGRAM_SHAPE_CODE}}", program_shape_code}, {"{{RESULT_INDEX}}", strings::StrCat(result_index)}, {"{{RESULT_NAMES_CODE}}", result_names_code}, {"{{TEMP_BYTES_ALIGNED}}", strings::StrCat(temp_bytes_aligned)}, {"{{TEMP_BYTES_TOTAL}}", strings::StrCat(temp_bytes_total)}, {"{{TEMP_NUM}}", strings::StrCat(temp_sizes.size())}, {"{{TEMP_SIZES}}", str_util::Join(temp_sizes, ", ")}, - }; + {"{{DECLS_FROM_OBJ_FILE}}", + str_util::Join(metadata_result.header_variable_decls, "\n")}, + {"{{PROGRAM_SHAPE_SHIM_EXPRESSION}}", + metadata_result.program_shape_access_shim}}; str_util::ReplaceAllPairs(header, rewrites); return Status::OK(); } +static string CreateUniqueIdentifierForProgramShape(const CodegenOpts& opts) { + string result = "__tfcompile"; + for (const string& n : opts.namespaces) { + strings::StrAppend(&result, "_", n); + } + + strings::StrAppend(&result, "_", opts.class_name, "_ProgramShape"); + return result; +} + +Status GenerateMetadata(const CodegenOpts& opts, + const CompileResult& compile_result, + MetadataResult* metadata_result) { + std::unique_ptr program_shape; + + if (opts.gen_program_shape) { + program_shape = + tensorflow::MakeUnique(compile_result.program_shape); + // The parameter names are currently meaningless, and redundant with the + // rest of our metadata, so clear them out to avoid confusion and save + // space. + program_shape->clear_parameter_names(); + } + + // When asked to serialize a null protobuf, CreateEmbeddedProtocolBuffer gives + // a shim that evaluates to nullptr, which is what we want. + + TF_ASSIGN_OR_RETURN( + EmbeddedProtocolBuffer embedded_program_shape, + CreateEmbeddedProtocolBuffer(opts.target_triple, + CreateUniqueIdentifierForProgramShape(opts), + "xla::ProgramShape", program_shape.get())); + + metadata_result->program_shape_access_shim = + std::move(embedded_program_shape.cpp_shim_expression); + metadata_result->header_variable_decls.emplace_back( + std::move(embedded_program_shape.cpp_variable_decl)); + metadata_result->object_file_data = + std::move(embedded_program_shape.object_file_data); + return Status::OK(); +} + Status ParseCppClass(const string& cpp_class, string* class_name, std::vector* namespaces) { class_name->clear(); diff --git a/tensorflow/compiler/aot/codegen.h b/tensorflow/compiler/aot/codegen.h index 76dd0cc3cf9470a1beb2a4725724f640aecfec7f..3430b1f96cf4d3c035b76c77ccf124c5d164751e 100644 --- a/tensorflow/compiler/aot/codegen.h +++ b/tensorflow/compiler/aot/codegen.h @@ -26,11 +26,15 @@ limitations under the License. namespace tensorflow { namespace tfcompile { -// HeaderOpts specifies options for header-file generation. -struct HeaderOpts { +// CodegenOpts specifies code generation options for the generated header file +// and the generated metadata object file. +struct CodegenOpts { // The name of the generated C++ class, wrapping the generated function. string class_name; + // Target triple for the architecture we're targeting. + string target_triple; + // Namespaces specifies a list of C++ namespaces to add to the generated // header. If empty, all symbols will be in the global namespace. std::vector namespaces; @@ -42,11 +46,36 @@ struct HeaderOpts { bool gen_program_shape = false; }; +// Describes a generated metadata object file. +struct MetadataResult { + // These are top level "extern C" declarations that are expected to be visible + // wherever program_shape_access_shim is emitted. + std::vector header_variable_decls; + + // program_shape_access_shim is a C++ expression that constructs the + // xla::ProgramShape instance for the CompileResult passed to + // GenerateMetadata. + string program_shape_access_shim; + + // The contents of the object (".o") file. + string object_file_data; +}; + +// Generates a metadata object file according to `opts` and `compile_result`. +// The generated object file is returned via `metadata_result`. +Status GenerateMetadata(const CodegenOpts& opts, + const CompileResult& compile_result, + MetadataResult* metadata_result); + // GenerateHeader uses the meta-information from compile_result to generate a // C++ header giving access to the function in the generated object file. The // header includes API usage documentation. -Status GenerateHeader(const HeaderOpts& opts, const tf2xla::Config& config, - const CompileResult& compile_result, string* header); +// +// metadata_result is an instance of MetadataResult obtained by a previous +// invocation to GenerateMetadata. +Status GenerateHeader(const CodegenOpts& opts, const tf2xla::Config& config, + const CompileResult& compile_result, + const MetadataResult& metadata_result, string* header); // ParseCppClass parses `cpp_class` into its `class_name` and `namespaces` // components. The syntax is [[::],...]. This diff --git a/tensorflow/compiler/aot/codegen_test.cc b/tensorflow/compiler/aot/codegen_test.cc index 75026c57c04a64186a1e5be6c41e4dd7de8520b7..972b7d51ecb3798e61757ac55e973075a23b433a 100644 --- a/tensorflow/compiler/aot/codegen_test.cc +++ b/tensorflow/compiler/aot/codegen_test.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "llvm/Support/TargetSelect.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -123,9 +124,39 @@ TEST_F(ParseCppClassTest, ParseFail) { ExpectFail("good::0bad"); } -TEST(GenerateHeader, Golden) { - HeaderOpts opts; +static void CompareWithGoldenFile( + const string& tensorflow_relative_golden_file_name, + const string& expected_contents) { + // To update the golden file, flip update_golden to true and run the + // following: + // bazel test --test_strategy=local \ + // third_party/tensorflow/compiler/aot:codegen_test + const bool update_golden = false; + const string golden_file_name = io::JoinPath( + testing::TensorFlowSrcRoot(), tensorflow_relative_golden_file_name); + + if (update_golden) { + TF_EXPECT_OK( + WriteStringToFile(Env::Default(), golden_file_name, expected_contents)); + } + + string golden_file_contents; + TF_ASSERT_OK(ReadFileToString(Env::Default(), golden_file_name, + &golden_file_contents)); + EXPECT_EQ(golden_file_contents, expected_contents); +} + +TEST(CodegenTest, Golden) { + // Normally CpuCompiler::CpuCompiler does this, but in this test we've + // bypassed the Cpu compiler so we have to do this manually. + llvm::InitializeNativeTarget(); + llvm::InitializeNativeTargetAsmPrinter(); + LLVMInitializeX86Target(); + LLVMInitializeX86TargetMC(); + + CodegenOpts opts; opts.class_name = "MyClass"; + opts.target_triple = "x86_64-pc-linux"; opts.namespaces = {"foo", "bar"}; opts.gen_name_to_index = true; opts.gen_program_shape = true; @@ -150,25 +181,22 @@ TEST(GenerateHeader, Golden) { {xla::ShapeUtil::MakeShape(xla::U32, {5, 6})})); compile_result.entry_point = "entry_point"; compile_result.pointer_size = 8; + + MetadataResult metadata_result; + TF_ASSERT_OK(GenerateMetadata(opts, compile_result, &metadata_result)); + + // The other fields in metadata_result are tested as part of the generated + // header test. + + CompareWithGoldenFile("compiler/aot/codegen_test_o.golden", + metadata_result.object_file_data); + string header; - TF_EXPECT_OK(GenerateHeader(opts, config, compile_result, &header)); + TF_ASSERT_OK( + GenerateHeader(opts, config, compile_result, metadata_result, &header)); - // Compare against the golden file. - const string golden_name = io::JoinPath(testing::TensorFlowSrcRoot(), - "compiler/aot/codegen_test_h.golden"); - // To update the golden file, flip update_golden to true and run the - // following: - // bazel test --test_strategy=local \ - // third_party/tensorflow/compiler/aot:codegen_test - const bool update_golden = false; - if (update_golden) { - TF_EXPECT_OK(WriteStringToFile(Env::Default(), golden_name, header)); - } - string golden_data; - TF_EXPECT_OK(ReadFileToString(Env::Default(), golden_name, &golden_data)); - EXPECT_EQ(header, golden_data); + CompareWithGoldenFile("compiler/aot/codegen_test_h.golden", header); } - } // namespace } // namespace tfcompile } // namespace tensorflow diff --git a/tensorflow/compiler/aot/codegen_test_h.golden b/tensorflow/compiler/aot/codegen_test_h.golden index 95ab3a7332f51070732bf5d62c7926c84e3b738d..ac3b5873318873b5fdf41bd556a0b2abddc2b30b 100644 --- a/tensorflow/compiler/aot/codegen_test_h.golden +++ b/tensorflow/compiler/aot/codegen_test_h.golden @@ -21,6 +21,8 @@ extern "C" void entry_point( void* result, const xla::ExecutableRunOptions* run_options, const void** args, void** temps, tensorflow::int64* profile_counters); +extern "C" char __tfcompile_foo_bar_MyClass_ProgramShape_protobuf_array_contents[]; + namespace foo { namespace bar { @@ -235,12 +237,10 @@ class MyClass : public tensorflow::XlaCompiledCpuFunction { // Shape of the args and results. static const xla::ProgramShape* StaticProgramShape() { static const xla::ProgramShape* kShape = []() { - static const char kProto[] = {10,14,16,11,26,2,1,2,42,6,10,2,1,0,32,1,10,14,16,5,26,2,3,4,42,6,10,2,1,0,32,1,18,18,16,13,34,14,16,8,26,2,5,6,42,6,10,2,1,0,32,1}; - static constexpr int kProtoSize = 52; - xla::ProgramShape* shape = new xla::ProgramShape; - shape->ParseFromArray(kProto, kProtoSize); - return shape; - }(); + xla::ProgramShape* proto = new xla::ProgramShape; + proto->ParseFromArray(&__tfcompile_foo_bar_MyClass_ProgramShape_protobuf_array_contents[0], 52); + return proto; + }(); return kShape; } }; diff --git a/tensorflow/compiler/aot/codegen_test_o.golden b/tensorflow/compiler/aot/codegen_test_o.golden new file mode 100644 index 0000000000000000000000000000000000000000..eb001c5d45bdfefc76629d7303d89f5480432235 Binary files /dev/null and b/tensorflow/compiler/aot/codegen_test_o.golden differ diff --git a/tensorflow/compiler/aot/embedded_protocol_buffers.cc b/tensorflow/compiler/aot/embedded_protocol_buffers.cc new file mode 100644 index 0000000000000000000000000000000000000000..6489929a576d6469c4ff1358ca5ee9d27fb578bb --- /dev/null +++ b/tensorflow/compiler/aot/embedded_protocol_buffers.cc @@ -0,0 +1,158 @@ +/* Copyright 2017 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/compiler/aot/embedded_protocol_buffers.h" + +#include +#include + +#include "llvm/ADT/Triple.h" +#include "llvm/ExecutionEngine/ObjectMemoryBuffer.h" +#include "llvm/IR/GlobalVariable.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/TargetRegistry.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/Target/TargetOptions.h" +#include "tensorflow/compiler/tf2xla/str_util.h" +#include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" +#include "tensorflow/compiler/xla/util.h" + +namespace tensorflow { +namespace tfcompile { + +using xla::llvm_ir::AsStringRef; + +static std::unique_ptr CreateModuleWithEmbeddedProtocolBuffer( + llvm::LLVMContext* llvm_context, llvm::TargetMachine* target_machine, + const ::tensorflow::protobuf::MessageLite& proto, + StringPiece unique_identifier, string* protobuf_array_symbol_name, + int64* protobuf_array_size) { + string protobuf_array_contents = proto.SerializeAsString(); + *protobuf_array_symbol_name = + strings::StrCat(unique_identifier, "_protobuf_array_contents"); + *protobuf_array_size = protobuf_array_contents.size(); + + std::unique_ptr module = + MakeUnique("embedded_data_module", *llvm_context); + + llvm::Constant* protobuf_array_initializer = + llvm::ConstantDataArray::getString(*llvm_context, + AsStringRef(protobuf_array_contents), + /*AddNull=*/false); + new llvm::GlobalVariable( + *module, protobuf_array_initializer->getType(), + /*isConstant=*/true, llvm::GlobalValue::ExternalLinkage, + protobuf_array_initializer, AsStringRef(*protobuf_array_symbol_name)); + + return module; +} + +static string CreateCPPShimExpression(StringPiece qualified_cpp_protobuf_name, + StringPiece protobuf_array_symbol_name, + int64 protobuf_array_size) { + string code = + "[]() {\n" + " {{PROTOBUF_NAME}}* proto = new {{PROTOBUF_NAME}};\n" + " proto->ParseFromArray(&{{ARRAY_SYMBOL}}[0], {{ARRAY_SIZE}});\n" + " return proto;\n" + " }()"; + + str_util::ReplaceAllPairs( + &code, + { + {"{{ARRAY_SYMBOL}}", strings::StrCat(protobuf_array_symbol_name)}, + {"{{ARRAY_SIZE}}", strings::StrCat(protobuf_array_size)}, + {"{{PROTOBUF_NAME}}", strings::StrCat(qualified_cpp_protobuf_name)}, + }); + return code; +} + +static StatusOr CodegenModule(llvm::TargetMachine* target_machine, + std::unique_ptr module) { + llvm::SmallVector stream_buffer; + llvm::raw_svector_ostream ostream(stream_buffer); + llvm::legacy::PassManager codegen_passes; + + if (target_machine->addPassesToEmitFile( + codegen_passes, ostream, llvm::TargetMachine::CGFT_ObjectFile)) { + return xla::InternalError( + "Could not create pass pipeline to generate object file"); + } + + codegen_passes.run(*module); + + return string(stream_buffer.begin(), stream_buffer.end()); +} + +static StatusOr> +GetTargetMachineFromTriple(StringPiece target_triple) { + std::string error; + std::string normalized_triple = + llvm::Triple::normalize(AsStringRef(target_triple)); + const llvm::Target* target = + llvm::TargetRegistry::lookupTarget(normalized_triple, error); + if (target == nullptr) { + return xla::InternalError("TargetRegistry::lookupTarget failed: %s", + error.c_str()); + } + + return WrapUnique(target->createTargetMachine( + normalized_triple, /*CPU=*/"", + /*Features=*/"", llvm::TargetOptions(), llvm::None)); +} + +StatusOr CreateEmbeddedProtocolBuffer( + StringPiece target_triple, StringPiece symbol_prefix, + StringPiece qualified_cpp_protobuf_name, + const ::tensorflow::protobuf::MessageLite* proto) { + TF_ASSIGN_OR_RETURN(std::unique_ptr target_machine, + GetTargetMachineFromTriple(target_triple)); + + llvm::LLVMContext llvm_context; + string object_file, cpp_shim, cpp_variable_decl; + + if (proto) { + string protobuf_array_symbol_name; + int64 protobuf_array_size; + + std::unique_ptr module_with_serialized_proto = + CreateModuleWithEmbeddedProtocolBuffer( + &llvm_context, target_machine.get(), *proto, symbol_prefix, + &protobuf_array_symbol_name, &protobuf_array_size); + TF_ASSIGN_OR_RETURN(object_file, + CodegenModule(target_machine.get(), + std::move(module_with_serialized_proto))); + cpp_shim = CreateCPPShimExpression(qualified_cpp_protobuf_name, + protobuf_array_symbol_name, + protobuf_array_size); + + cpp_variable_decl = strings::StrCat("extern \"C\" char ", + protobuf_array_symbol_name, "[];"); + } else { + TF_ASSIGN_OR_RETURN( + object_file, + CodegenModule(target_machine.get(), + MakeUnique("empty_module", llvm_context))); + cpp_shim = "nullptr"; + } + + return {{cpp_shim, cpp_variable_decl, object_file}}; +} + +} // namespace tfcompile +} // namespace tensorflow diff --git a/tensorflow/compiler/aot/embedded_protocol_buffers.h b/tensorflow/compiler/aot/embedded_protocol_buffers.h new file mode 100644 index 0000000000000000000000000000000000000000..8436e0ff67f352a24e3d16b46f16c1ad2f3a5957 --- /dev/null +++ b/tensorflow/compiler/aot/embedded_protocol_buffers.h @@ -0,0 +1,73 @@ +/* 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. +==============================================================================*/ + +// This file defines utilities to help "embed" protocol buffers into object +// (".o") files. These C++ binaries and shared objects can link in these .o to +// get access to said protocol buffers at runtime. + +#ifndef TENSORFLOW_COMPILER_AOT_EMBEDDED_PROTOCOL_BUFFERS_H_ +#define TENSORFLOW_COMPILER_AOT_EMBEDDED_PROTOCOL_BUFFERS_H_ + +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/platform/protobuf.h" + +namespace tensorflow { +namespace tfcompile { +using xla::StatusOr; + +// Represents a protocol buffer embedded into an object file and describes a way +// to access it at runtime. +struct EmbeddedProtocolBuffer { + // cpp_shim_expression is a C++ expression that creates an instance of said + // protocol buffer when executed. + string cpp_shim_expression; + + // cpp_variable_decl is an "extern C" array declaration that is used in + // cpp_shim_expression. It must be visible wherever cpp_shim_expression is + // emitted. + string cpp_variable_decl; + + // The contents of the object (".o") file the protocol buffer is embbed in. + // This needs to be linked in to any program that wants to execute + // cpp_variable_decl . + string object_file_data; +}; + +// Creates an object file that contains `proto`. +// +// `proto` is allowed to be nullptr, in which case the generated C++ shim +// expression is just `nullptr`, and the generated object file does not define +// any symbols. +// +// `target_triple` is the target triple for the target architecture for the +// generated object file. +// +// `symbol_prefix` is prefix that is guaranteed to be unique across the binary +// or DSO the generated object file will be linked into. +// +// `qualified_cpp_protobuf_name` is a qualified ("qualified" as in C++ +// namespace qualified) protocol buffer name. This needs is only used in +// EmbeddedProtocolBuffer::cpp_shim_expression so relatively qualified +// names are fine as long as they're valid wherever cpp_shim_expression +// is emitted. +StatusOr CreateEmbeddedProtocolBuffer( + StringPiece target_triple, StringPiece symbol_prefix, + StringPiece qualified_cpp_protobuf_name, + const ::tensorflow::protobuf::MessageLite* proto); + +} // namespace tfcompile +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_AOT_EMBEDDED_PROTOCOL_BUFFERS_H_ diff --git a/tensorflow/compiler/aot/flags.cc b/tensorflow/compiler/aot/flags.cc index 7c2f27e550d44c2487f91acf1029c962ac3f5d01..8c95cb8f90ee031fdbb97fabd9d86f848b42e4c5 100644 --- a/tensorflow/compiler/aot/flags.cc +++ b/tensorflow/compiler/aot/flags.cc @@ -59,8 +59,13 @@ void AppendMainFlags(std::vector* flag_list, MainFlags* flags) { "namespaces may precede the class name, separated by double-colons. " "The class will be generated in the given namespace(s), or if no " "namespaces are given, within the global namespace."}, - {"out_object", &flags->out_object, "Output object file name."}, + {"out_function_object", &flags->out_function_object, + "Output object file containing the generated function for the " + "TensorFlow model."}, {"out_header", &flags->out_header, "Output header file name."}, + {"out_metadata_object", &flags->out_metadata_object, + "Output object file name containing optional metadata for the generated " + "function."}, {"out_session_module", &flags->out_session_module, "Output session module proto."}, {"gen_name_to_index", &flags->gen_name_to_index, diff --git a/tensorflow/compiler/aot/flags.h b/tensorflow/compiler/aot/flags.h index 3519659e3af7cd345f30080a07ce91fb858623fb..d266fbead61f7eb43863d1c67c0f86926ae9452d 100644 --- a/tensorflow/compiler/aot/flags.h +++ b/tensorflow/compiler/aot/flags.h @@ -34,7 +34,8 @@ struct MainFlags { string target_features; string entry_point; string cpp_class; - string out_object; + string out_function_object; + string out_metadata_object; string out_header; string out_session_module; diff --git a/tensorflow/compiler/aot/tests/BUILD b/tensorflow/compiler/aot/tests/BUILD index 7dfd49cc3b92f83fd64ca62bd2230938ce2d0a65..28aab6eb614ca7123d9e00f7f5cc3661b62e23f7 100644 --- a/tensorflow/compiler/aot/tests/BUILD +++ b/tensorflow/compiler/aot/tests/BUILD @@ -74,7 +74,9 @@ tf_library( # compile but the others in this directory succeed, you may need to # expand the "required by all tf_library targets" list in tfcompile.bzl. include_standard_runtime_deps = False, - tags = ["manual"], + tags = [ + "manual", + ], ) tf_library( @@ -84,7 +86,9 @@ tf_library( cpp_class = "AddWithCkptComp", freeze_checkpoint = "test_graph_tfadd_with_ckpt.ckpt", graph = "test_graph_tfadd_with_ckpt.pb", - tags = ["manual"], + tags = [ + "manual", + ], ) tf_library( @@ -95,7 +99,9 @@ tf_library( freeze_checkpoint = "test_graph_tfadd_with_ckpt_saver.ckpt", freeze_saver = "test_graph_tfadd_with_ckpt_saver.saver", graph = "test_graph_tfadd_with_ckpt_saver.pb", - tags = ["manual"], + tags = [ + "manual", + ], ) tf_library( @@ -104,7 +110,9 @@ tf_library( config = "test_graph_tffunction.config.pbtxt", cpp_class = "FunctionComp", graph = "test_graph_tffunction.pb", - tags = ["manual"], + tags = [ + "manual", + ], ) tf_library( @@ -113,7 +121,9 @@ tf_library( config = "test_graph_tfgather.config.pbtxt", cpp_class = "GatherComp", graph = "test_graph_tfgather.pb", - tags = ["manual"], + tags = [ + "manual", + ], ) tf_library( @@ -122,7 +132,9 @@ tf_library( config = "test_graph_tfmatmul.config.pbtxt", cpp_class = "foo::bar::MatMulComp", graph = "test_graph_tfmatmul.pb", - tags = ["manual"], + tags = [ + "manual", + ], ) tf_library( @@ -131,7 +143,9 @@ tf_library( config = "test_graph_tfmatmulandadd.config.pbtxt", cpp_class = "MatMulAndAddComp", graph = "test_graph_tfmatmulandadd.pb", - tags = ["manual"], + tags = [ + "manual", + ], tfcompile_flags = "--gen_name_to_index --gen_program_shape", ) @@ -141,13 +155,17 @@ tf_library( config = "test_graph_tfsplits.config.pbtxt", cpp_class = "SplitsComp", graph = "test_graph_tfsplits.pb", - tags = ["manual"], + tags = [ + "manual", + ], ) tf_cc_test( name = "tfcompile_test", srcs = ["tfcompile_test.cc"], - tags = ["manual"], + tags = [ + "manual", + ], deps = [ ":test_graph_tfadd", ":test_graph_tfadd_with_ckpt", diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index 542451ed2d14fbceca00c6ccb6e28c1c3a0d4321..9dff1be09fede6f65f82c2f36d94be07e781949f 100644 --- a/tensorflow/compiler/aot/tfcompile.bzl +++ b/tensorflow/compiler/aot/tfcompile.bzl @@ -4,7 +4,7 @@ To use from your BUILD file, add the following line to load the macro: -load("@org_tensorflow//tensorflow/compiler/aot:tfcompile.bzl", "tf_library") +load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library") Then call the macro like this: @@ -16,14 +16,15 @@ tf_library( ) """ -load("@org_tensorflow//tensorflow:tensorflow.bzl", "if_android", "tf_copts") +load("//tensorflow:tensorflow.bzl", + "if_android", "tf_cc_test", "tf_copts") def tf_library(name, graph, config, freeze_checkpoint=None, freeze_saver=None, cpp_class=None, gen_test=True, gen_benchmark=True, visibility=None, testonly=None, tfcompile_flags=None, - tfcompile_tool="@org_tensorflow//tensorflow/compiler/aot:tfcompile", + tfcompile_tool="//tensorflow/compiler/aot:tfcompile", include_standard_runtime_deps=True, deps=None, tags=None): """Runs tfcompile to compile a TensorFlow graph into executable code. @@ -102,6 +103,7 @@ def tf_library(name, graph, config, # Now run freeze_graph to convert variables into constants. freeze_args = (" --input_graph=$(location " + graph + ")" + + " --checkpoint_version=1" + " --input_binary=" + str(not graph.endswith(".pbtxt")) + " --input_checkpoint=$(location " + freeze_checkpoint + ")" + " --output_graph=$(location " + freeze_file + ")" + @@ -119,16 +121,17 @@ def tf_library(name, graph, config, out_nodes_file, ] + freeze_saver_srcs, outs=[freeze_file], - cmd=("$(location @org_tensorflow//tensorflow/python/tools:freeze_graph)" + + cmd=("$(location //tensorflow/python/tools:freeze_graph)" + freeze_args), - tools=["@org_tensorflow//tensorflow/python/tools:freeze_graph"], + tools=["//tensorflow/python/tools:freeze_graph"], tags=tags, ) tfcompile_graph = freeze_file # Rule that runs tfcompile to produce the header and object file. header_file = name + ".h" - object_file = name + ".o" + metadata_object_file = name + "_tfcompile_metadata.o" + function_object_file = name + "_tfcompile_function.o" ep = ("__" + PACKAGE_NAME + "__" + name).replace("/", "_") if type(tfcompile_flags) == type(""): flags = tfcompile_flags @@ -142,7 +145,8 @@ def tf_library(name, graph, config, ], outs=[ header_file, - object_file, + metadata_object_file, + function_object_file, ], cmd=("$(location " + tfcompile_tool + ")" + " --graph=$(location " + tfcompile_graph + ")" + @@ -151,7 +155,8 @@ def tf_library(name, graph, config, " --cpp_class=" + cpp_class + " --target_triple=" + target_llvm_triple() + " --out_header=$(@D)/" + header_file + - " --out_object=$(@D)/" + object_file + + " --out_metadata_object=$(@D)/" + metadata_object_file + + " --out_function_object=$(@D)/" + function_object_file + " " + flags), tools=[tfcompile_tool], visibility=visibility, @@ -202,7 +207,7 @@ def tf_library(name, graph, config, need_xla_data_proto = (flags and flags.find("--gen_program_shape") != -1) native.cc_library( name=name, - srcs=[object_file], + srcs=[function_object_file, metadata_object_file], hdrs=[header_file], visibility=visibility, testonly=testonly, @@ -210,22 +215,19 @@ def tf_library(name, graph, config, # These deps are required by all tf_library targets even if # include_standard_runtime_deps is False. Without them, the # generated code will fail to compile. - "@org_tensorflow//tensorflow/compiler/tf2xla:xla_compiled_cpu_function", - "@org_tensorflow//tensorflow/core:framework_lite", + "//tensorflow/compiler/tf2xla:xla_compiled_cpu_function", + "//tensorflow/core:framework_lite", ] + (need_xla_data_proto and [ # If we're generating the program shape, we must depend on the proto. - "@org_tensorflow//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla:xla_data_proto", ] or []) + (include_standard_runtime_deps and [ # TODO(cwhipkey): only depend on kernel code that the model actually needed. - "@org_tensorflow//tensorflow/compiler/tf2xla/kernels:index_ops_kernel_argmax_float_1d", - "@org_tensorflow//tensorflow/compiler/tf2xla/kernels:index_ops_kernel_argmax_float_2d", - "@org_tensorflow//tensorflow/compiler/xla/service/cpu:cpu_runtime_avx", - "@org_tensorflow//tensorflow/compiler/xla/service/cpu:cpu_runtime_neon", - "@org_tensorflow//tensorflow/compiler/xla/service/cpu:cpu_runtime_sse4_1", - "@org_tensorflow//tensorflow/compiler/xla/service/cpu:runtime_conv2d", - "@org_tensorflow//tensorflow/compiler/xla/service/cpu:runtime_matmul", - "@org_tensorflow//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_conv2d", - "@org_tensorflow//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_matmul", + "//tensorflow/compiler/tf2xla/kernels:index_ops_kernel_argmax_float_1d", + "//tensorflow/compiler/tf2xla/kernels:index_ops_kernel_argmax_float_2d", + "//tensorflow/compiler/xla/service/cpu:runtime_conv2d", + "//tensorflow/compiler/xla/service/cpu:runtime_matmul", + "//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_conv2d", + "//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_matmul", "//third_party/eigen3", ] or []) + (deps or []), tags=tags, @@ -251,28 +253,32 @@ def tf_library(name, graph, config, name=("gen_" + test_name), testonly=1, srcs=[ - "@org_tensorflow//tensorflow/compiler/aot:test.cc", + "//tensorflow/compiler/aot:test.cc", header_file, ], outs=[test_file], cmd=("sed " + sed_replace + - " $(location @org_tensorflow//tensorflow/compiler/aot:test.cc) " + + " $(location //tensorflow/compiler/aot:test.cc) " + "> $(OUTS)"), tags=tags, ) - # The cc_test rule for the generated code. - native.cc_test( + # The cc_test rule for the generated code. To ensure that this works + # reliably across build configurations, we must use tf_cc_test instead of + # native.cc_test. This is related to how we build + # //tensorflow/core:lib -- see the note in tensorflow/core/BUILD + # for more details. + tf_cc_test( name=test_name, srcs=[test_file], deps=[ ":" + name, - "@org_tensorflow//tensorflow/compiler/aot:runtime", - "@org_tensorflow//tensorflow/compiler/aot:tf_library_test_main", - "@org_tensorflow//tensorflow/compiler/xla:executable_run_options", + "//tensorflow/compiler/aot:runtime", + "//tensorflow/compiler/aot:tf_library_test_main", + "//tensorflow/compiler/xla:executable_run_options", "//third_party/eigen3", - "@org_tensorflow//tensorflow/core:lib", - "@org_tensorflow//tensorflow/core:test", + "//tensorflow/core:lib", + "//tensorflow/core:test", ], tags=tags, ) @@ -280,7 +286,7 @@ def tf_library(name, graph, config, if gen_benchmark: benchmark_name = name + "_benchmark" benchmark_file = benchmark_name + ".cc" - benchmark_main = ("@org_tensorflow//tensorflow/compiler/aot:" + + benchmark_main = ("//tensorflow/compiler/aot:" + "benchmark_main.template") # Rule to rewrite benchmark.cc to produce the benchmark_file. @@ -298,7 +304,9 @@ def tf_library(name, graph, config, tags=tags, ) - # The cc_benchmark rule for the generated code. + # The cc_benchmark rule for the generated code. This does not need the + # tf_cc_binary since we (by deliberate design) do not depend on + # //tensorflow/core:lib. # # Note: to get smaller size on android for comparison, compile with: # --copt=-fvisibility=hidden @@ -312,12 +320,12 @@ def tf_library(name, graph, config, linkopts = if_android(["-pie", "-s"]), deps=[ ":" + name, - "@org_tensorflow//tensorflow/compiler/aot:benchmark", - "@org_tensorflow//tensorflow/compiler/aot:runtime", - "@org_tensorflow//tensorflow/compiler/xla:executable_run_options", + "//tensorflow/compiler/aot:benchmark", + "//tensorflow/compiler/aot:runtime", + "//tensorflow/compiler/xla:executable_run_options", "//third_party/eigen3", ] + if_android([ - "@org_tensorflow//tensorflow/compiler/aot:benchmark_extra_android", + "//tensorflow/compiler/aot:benchmark_extra_android", ]), tags=tags, ) @@ -327,11 +335,11 @@ def target_llvm_triple(): # TODO(toddw): Add target_triple for other targets. For details see: # http://llvm.org/docs/doxygen/html/Triple_8h_source.html return select({ - "@org_tensorflow//tensorflow:android_armeabi": "armv5-none-android", - "@org_tensorflow//tensorflow:android_arm": "armv7-none-android", - "@org_tensorflow//tensorflow:android_arm64": "aarch64-none-android", - "@org_tensorflow//tensorflow:android_x86": "i686-none-android", - "@org_tensorflow//tensorflow:linux_ppc64le": "ppc64le-ibm-linux-gnu", - "@org_tensorflow//tensorflow:darwin": "x86_64-none-darwin", + "//tensorflow:android_armeabi": "armv5-none-android", + "//tensorflow:android_arm": "armv7-none-android", + "//tensorflow:android_arm64": "aarch64-none-android", + "//tensorflow:android_x86": "i686-none-android", + "//tensorflow:linux_ppc64le": "ppc64le-ibm-linux-gnu", + "//tensorflow:darwin": "x86_64-none-darwin", "//conditions:default": "x86_64-pc-linux", }) diff --git a/tensorflow/compiler/aot/tfcompile_main.cc b/tensorflow/compiler/aot/tfcompile_main.cc index 6ab3d474187c7df2131f94c9f42f0d0f2f9d99d7..e2f01179d4e2e4f6ef72b2761d06e130ffa3a94f 100644 --- a/tensorflow/compiler/aot/tfcompile_main.cc +++ b/tensorflow/compiler/aot/tfcompile_main.cc @@ -91,19 +91,26 @@ Status Main(const MainFlags& flags) { // Write output files. Env* env = Env::Default(); const std::vector& obj = compile_result.aot->object_file_data(); - TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_object, + TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_function_object, StringPiece(obj.data(), obj.size()))); - HeaderOpts header_opts; - header_opts.gen_name_to_index = flags.gen_name_to_index; - header_opts.gen_program_shape = flags.gen_program_shape; + CodegenOpts codegen_opts; + codegen_opts.gen_name_to_index = flags.gen_name_to_index; + codegen_opts.gen_program_shape = flags.gen_program_shape; + codegen_opts.target_triple = flags.target_triple; if (flags.cpp_class.empty()) { return errors::InvalidArgument("Must specify --cpp_class"); } - TF_RETURN_IF_ERROR(ParseCppClass(flags.cpp_class, &header_opts.class_name, - &header_opts.namespaces)); - string header; + TF_RETURN_IF_ERROR(ParseCppClass(flags.cpp_class, &codegen_opts.class_name, + &codegen_opts.namespaces)); + + MetadataResult metadata_result; TF_RETURN_IF_ERROR( - GenerateHeader(header_opts, config, compile_result, &header)); + GenerateMetadata(codegen_opts, compile_result, &metadata_result)); + TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_metadata_object, + metadata_result.object_file_data)); + string header; + TF_RETURN_IF_ERROR(GenerateHeader(codegen_opts, config, compile_result, + metadata_result, &header)); TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_header, header)); return Status::OK(); } @@ -114,7 +121,8 @@ Status Main(const MainFlags& flags) { int main(int argc, char** argv) { tensorflow::tfcompile::MainFlags flags; flags.target_triple = "x86_64-pc-linux"; - flags.out_object = "out.o"; + flags.out_function_object = "out_model.o"; + flags.out_metadata_object = "out_helper.o"; flags.out_header = "out.h"; flags.entry_point = "entry"; diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index 13343967c4f86b610ecacf421166dadd1aa79da0..a711319607f4ff2b83aa0ebe50e215b3d0e2258e 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -110,19 +110,6 @@ cc_library( alwayslink = True, ) -# Internal targets below this point. - -cc_library( - name = "common", - srcs = [ - "defs.cc", - ], - hdrs = [ - "defs.h", - ], - visibility = [":friends"], -) - cc_library( name = "xla_device", srcs = [ @@ -135,6 +122,8 @@ cc_library( "xla_device_context.h", "xla_device_ops.h", ], + # Public visibility is needed for external TF/XLA backends. + visibility = ["//visibility:public"], deps = [ ":common", ":jit_compilation_passes", @@ -164,6 +153,19 @@ cc_library( ], ) +# Internal targets below this point. + +cc_library( + name = "common", + srcs = [ + "defs.cc", + ], + hdrs = [ + "defs.h", + ], + visibility = [":friends"], +) + cc_library( name = "xla_compilation_cache", srcs = ["xla_compilation_cache.cc"], diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc index 0de163d3a8f082eab4d8d802485da1bbc56e8180..9c372a012789fc25ca0a711349c09ca62edc6754 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc @@ -30,12 +30,14 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/optimization_registry.h" +#include "tensorflow/core/common_runtime/shape_refiner.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/graph_def_util.h" #include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/graph_def_builder.h" #include "tensorflow/core/graph/tensor_id.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/lib/gtl/map_util.h" @@ -141,8 +143,7 @@ struct NodeSlot { // everything to use it. static const char* const kArgOp = "_Arg"; static const char* const kRetValOp = "_Retval"; -static const char* const kSendToHostOp = "_XlaSendToHost"; -static const char* const kRecvFromHostOp = "_XlaRecvFromHost"; +static const char* const kHostComputeOp = "_XlaHostCompute"; static const char* const kSendFromHostOp = "_XlaSendFromHost"; static const char* const kRecvAtHostOp = "_XlaRecvAtHost"; @@ -171,7 +172,8 @@ class Encapsulator { // Write a copy of the input graph to 'graph_out', where the subgraphs are // replaced with calls to the new functions. - Status BuildOutputGraph(bool parallel_checking, Graph* graph_out); + Status BuildOutputGraph(bool parallel_checking, Graph* graph_out, + FunctionLibraryDefinition* library); private: // A subgraph of the input, all marked with a common 'group_attribute' @@ -201,21 +203,29 @@ class Encapsulator { // .. . // RAH --> C --> SFH // - // The compiled cluster is as follows. STH is a SendToHost node which is the - // source of a channel to the RAH node above. RFH is a RecvFromHost node which - // is the destination of a channel from the SFH node above. There is a control - // edge that ensures RFH follows STH, which is used in shape inference to - // ensure that the shapes on the STH host channel are known before the RFH - // channel is compiled. + // The compiled cluster is as follows. HC is a HostCompute node which is the + // source of a channel to the RAH node above and the destination of a channel + // from the SFH node above. // - // Arg --> B --> STH ..> RFH --> D --> Retval + // Arg --> B --> HC --> D --> Retval // - // The channels STH/RAH and SFH/RFH each transmit a tuple, so there is at most - // one RAH and SFH in each compiled cluster. This design is preferred over - // adding separate Arg/Retval nodes for each transmitted value because it - // simplifies the host code that would like to limit communication between - // host and device and, e.g., raise only one interrupt per channel rather than - // one per transmitted value. + // The channels HC/RAH and SFH/HC each transmit multiple tensors, so there is + // at most one RAH and SFH in each outside_compilation cluster. This design is + // preferred over adding separate Arg/Retval nodes for each transmitted value + // because it allows optimizations to the host code that would like to limit + // communication between host and device and, e.g., raise only one interrupt + // per channel rather than one per transmitted value. + // + // The shapes of the outputs from the HC node in general cannot be determined + // until the shapes of its inputs are known at compile time, since e.g., + // above, the shape of C's outputs aren't known until the shape of its inputs + // are known. If the shapes of the HC's outputs can be determined during the + // rewrite, they are stored in the node's 'shapes' attr. Otherwise a minimal + // graph is stored in the shape_inference_graph attr. This graph can be used + // when compiling the HC Op to determined the shape of the SFH inputs given + // the shapes of any ancestor RAH outputs. If it can be determined that the + // shape of the SFH inputs will not be inferrable even once the shapes of the + // RAH outputs are known, an error is returned by the rewriter. class Subgraph { public: // Creates a graph to build the subgraph in, if it doesn't already exist, @@ -246,6 +256,10 @@ class Encapsulator { const std::unordered_map& node_images, Graph* graph_out); + // Returns the names of all the outside_compilation subgraphs in this + // Subgraph. + void GetOutsideCompilationSubgraphNames(std::vector* names) const; + // Returns the Node that inputs to the function should be wired up to. Node* GetCallNodeForInputs() const; @@ -305,15 +319,9 @@ class Encapsulator { void RecordOutsideCompilationOutputOrControl( const string& outside_compilation_id, const Edge* edge); - // Adds the SendToHost nodes for each outside_compilation subgraph once the - // edges have all been recorded via RecordOutsideCompilationInputOrControl. - Status AddSendsToOutsideCompilation( - const std::unordered_map& node_images); - - // Adds the RecvFromHost nodes for each outside_compilation subgraph once - // the edges have all been recorded via - // RecordOutsideCompilationOutputOrControl. - Status AddRecvsFromOutsideCompilation( + // Adds the HostCompute nodes for each outside_compilation subgraph. + Status AddHostComputes( + const string& subgraph_name, const std::unordered_map& node_images); // Creates the sequencer node if it doesn't exist, adding it to graph_out. @@ -323,10 +331,16 @@ class Encapsulator { // all the downstream nodes of call_node_outputs. void ConnectSequencerToOutputs(Graph* graph_out); + Status AddShapeInferenceInfo( + const string& outside_compilation_subgraph_name, + const std::vector& shapes, GraphDef* inference_graph); + + Status ReplaceFunctionDef(FunctionLibraryDefinition* library); + private: struct OutsideCompilationSubgraph { // Map from source (producer node/slot) tensors in the original graph to - // input index (slot number in the SendToHost/RecvAtHost nodes that will + // input index (slot number in the HostCompute/RecvAtHost nodes that will // be created) for the outside_compilation subgraph. std::unordered_map inputs; @@ -335,14 +349,14 @@ class Encapsulator { // outside_compilation subgraph. These are recorded by // RecordOutsideCompilationInputOrControl while walking all the subgraph // edges, and lifted control edges within the subgraph are added by - // AddSendsToOutsideCompilation once the _SendToHost node has been + // AddSendsToOutsideCompilation once the _HostCompute node has been // created. The matching control edge from _RecvAtHost to the // destination is added by CopyEdgeToOutputGraph. std::unordered_set control_inputs; // Maps from source (producer node/slot) and destination (consumer // node/slot) tensors in the original graph to output index (slot number - // in the SendFromHost/RecvFromHost nodes that will be created) for the + // in the SendFromHost/HostCompute nodes that will be created) for the // outside_compilation subgraph. std::unordered_map outputs_by_src; std::unordered_map outputs_by_dst; @@ -352,13 +366,13 @@ class Encapsulator { // containing compiled subgraph. These are recorded by // RecordOutsideCompilationOutputOrControl while walking all the subgraph // edges, and lifted control edges within the subgraph are added by - // AddRecvsFromToOutsideCompilation once the _RecvFromHost node has been + // AddRecvsFromToOutsideCompilation once the _HostCompute node has been // created. The matching control edge from the source to _SendFromHost to // the destination is added by CopyEdgeToOutputGraph. std::unordered_set control_outputs; - // _SendToHost node in the subgraph. Not owned. - Node* send_to_host = nullptr; + // Name of the _HostCompute node in the subgraph. + string host_compute_name; // _RecvAtHost node in the output graph. Not owned. Node* recv_at_host = nullptr; @@ -516,6 +530,59 @@ class Encapsulator { const std::unordered_map& node_images, bool parallel_checking, Graph* graph_out); + // Constructs a minimal shape inference graph that can be used to determine + // the shape of send_node at the time that the subgraph is compiled. + // recv_at_host_nodes contains the names of all the recv_at_host nodes that + // send_node might depend on. These recv_at_host nodes have shapes that are + // not known during the rewrite pass, but will be known at compile time. + // + // If the shapes of all the inputs to send_node can be determined during the + // rewrite pass, on exit graphdef_out is empty and the shapes are returned in + // static_shape_out. Otherwise graphdef_out contains a graph that can be used + // for shape inference at compile time, where all the source nodes of the + // graph are either constants with known shapes, or nodes named in + // recv_at_host_nodes. + // + // A non-OK status is returned if neither of the above conditions can be + // satisfied, e.g., because send_node depends on a node that doesn't have a + // registered shape inference function. + Status DoStaticShapeInferenceForOutsideCompilationSend( + const Graph& graph_in, const ShapeRefiner& shape_refiner, + const std::unordered_set& recv_at_host_nodes, Node* send_node, + FunctionLibraryDefinition* library, + std::vector* static_shape_out, + std::unique_ptr* graphdef_out); + + // Makes a copy of graph containing only nodes that are ancestors of at least + // one node in send_from_host_nodes and store it in pruned_graph. On exit + // nodes_images contains a mapping from nodes in graph to nodes in + // pruned_graph. All functions in the copied graph are inlined. + Status MakePrunedGraphCopyAndInline( + const Graph& graph, const std::vector& sink_nodes, + std::unique_ptr* pruned_graph, + std::unordered_map* node_images, + FunctionLibraryDefinition* library); + + // Makes a copy of graph containing only nodes that are ancestors of a + // send_from_host node in an outside_compilation subgraph, and store it in + // pruned_graph. Also perform shape inference on the pruned graph, using + // shape_refiner. On exit node_images contains a mapping from nodes in graph + // to nodes in pruned_graph. + Status MakeGraphForOutsideCompilationSends( + const Graph& graph, std::unique_ptr* pruned_graph, + ShapeRefiner* shape_refiner, + std::unordered_map* node_images, + FunctionLibraryDefinition* library); + + // Performs static shape inference, as far as possible, for the send_from_host + // nodes in each outside_compilation subgraph. Where it is not possible to + // determine the shape statically, stores a serialized GraphDef in the + // HostCompute 'shape_inference_graph' attr, to be used at compile time for + // final inference. If the shapes are known statically they are stored in the + // HostCompute 'shapes' attr. + Status GetShapeInfoForOutsideCompilationSends( + Graph* graph_out, FunctionLibraryDefinition* library); + const string group_attribute_; const string outside_compilation_attribute_; const Graph* graph_in_; @@ -682,16 +749,20 @@ void Encapsulator::Subgraph::RecordOutsideCompilationOutputOrControl( } } -Status Encapsulator::Subgraph::AddSendsToOutsideCompilation( +Status Encapsulator::Subgraph::AddHostComputes( + const string& subgraph_name, const std::unordered_map& node_images) { for (auto& oc_subgraph_iter : outside_compilation_subgraphs_) { const string& oc_subgraph_name = oc_subgraph_iter.first; OutsideCompilationSubgraph& oc_subgraph = oc_subgraph_iter.second; - if (!oc_subgraph.inputs.empty() || !oc_subgraph.control_inputs.empty()) { - // Build a _SendToHost node sending all the args of the appropriate - // types. - std::vector dtypes(oc_subgraph.inputs.size(), DT_INVALID); + if (!oc_subgraph.inputs.empty() || !oc_subgraph.control_inputs.empty() || + !oc_subgraph.outputs_by_src.empty() || + !oc_subgraph.control_outputs.empty()) { + // Build a _HostCompute node. std::vector inputs(oc_subgraph.inputs.size()); + std::vector input_dtypes(oc_subgraph.inputs.size(), DT_INVALID); + std::vector output_dtypes(oc_subgraph.outputs_by_src.size(), + DT_INVALID); for (const auto& input_src : oc_subgraph.inputs) { const Node* src_node = input_src.first.node; @@ -700,94 +771,64 @@ Status Encapsulator::Subgraph::AddSendsToOutsideCompilation( int input_index = input_src.second; DataType dtype = src_node->output_type(src_slot); - dtypes[input_index] = dtype; inputs[input_index].Reset(src_image->name(), src_slot, dtype); + input_dtypes[input_index] = dtype; } - NodeDef send_def; - NodeDefBuilder builder( - strings::StrCat("outside_compilation_", oc_subgraph_name, "_send"), - kSendToHostOp); - builder.Attr("dtypes", dtypes); + for (const auto& output : oc_subgraph.outputs_by_src) { + DataType dtype = output.first.dtype; + int output_index = output.second; + output_dtypes[output_index] = dtype; + } + + NodeDef host_compute_def; + NodeDefBuilder builder(strings::StrCat("outside_compilation_", + oc_subgraph_name, "_host_compute"), + kHostComputeOp); builder.Input(inputs); - Status s = builder.Finalize(&send_def); + builder.Attr("Tinputs", input_dtypes); + builder.Attr("Toutputs", output_dtypes); + builder.Attr("key", + strings::StrCat("host_compute_channel_", subgraph_name, "_", + oc_subgraph_name)); + Status s = builder.Finalize(&host_compute_def); if (!s.ok()) return s; - oc_subgraph.send_to_host = graph_->AddNode(send_def, &s); + Node* host_compute = graph_->AddNode(host_compute_def, &s); if (!s.ok()) return s; + oc_subgraph.host_compute_name = host_compute->name(); - // Connect the _SendToHost node to its producers in the subgraph. + // Connect the _HostCompute node to its producers in the subgraph. for (auto& input_src : oc_subgraph.inputs) { const Node* src_node = input_src.first.node; Node* src_image = node_images.at(src_node); int src_slot = input_src.first.slot; int input_index = input_src.second; - graph_->AddEdge(src_image, src_slot, oc_subgraph.send_to_host, - input_index); + graph_->AddEdge(src_image, src_slot, host_compute, input_index); } - // Connect the _SendToHost node to its control edge producers in the + // Connect the _HostCompute node to its control edge producers in the // subgraph. for (const auto& src_node : oc_subgraph.control_inputs) { Node* src_image = node_images.at(src_node); - graph_->AddControlEdge(src_image, oc_subgraph.send_to_host); - } - } - } - - return Status::OK(); -} - -Status Encapsulator::Subgraph::AddRecvsFromOutsideCompilation( - const std::unordered_map& node_images) { - for (auto& oc_subgraph_iter : outside_compilation_subgraphs_) { - const string& oc_subgraph_name = oc_subgraph_iter.first; - OutsideCompilationSubgraph& oc_subgraph = oc_subgraph_iter.second; - if (!oc_subgraph.outputs_by_src.empty() || - !oc_subgraph.control_outputs.empty()) { - // Build a _RecvFromHost node producing all the outputs of the appropriate - // types. - std::vector dtypes(oc_subgraph.outputs_by_src.size(), - DT_INVALID); - - for (const auto& output : oc_subgraph.outputs_by_src) { - DataType dtype = output.first.dtype; - int output_index = output.second; - dtypes[output_index] = dtype; + graph_->AddControlEdge(src_image, host_compute); } - NodeDef recv_def; - NodeDefBuilder builder( - strings::StrCat("outside_compilation_", oc_subgraph_name, "_recv"), - kRecvFromHostOp); - builder.Attr("dtypes", dtypes); - Status s = builder.Finalize(&recv_def); - if (!s.ok()) return s; - - Node* recv = graph_->AddNode(recv_def, &s); - if (!s.ok()) return s; - - // Connect the consumers in the subgraph to the _RecvFromHost node. + // Connect the consumers in the subgraph to the _HostCompute node. for (const auto& output : oc_subgraph.outputs_by_dst) { const Node* dst_node = output.first.node; Node* dst_image = node_images.at(dst_node); int dst_slot = output.first.slot; int output_index = output.second; - graph_->AddEdge(recv, output_index, dst_image, dst_slot); + graph_->AddEdge(host_compute, output_index, dst_image, dst_slot); } - // Connect the control edge consumers in the subgraph to the _RecvFromHost + // Connect the control edge consumers in the subgraph to the _HostCompute // node. for (const auto& dst_node : oc_subgraph.control_outputs) { Node* dst_image = node_images.at(dst_node); - graph_->AddControlEdge(recv, dst_image); - } - - // Add a control edge in the subgraph so that the _SendToHost node, if - // any, is compiled before the _RecvFromHost node. - if (oc_subgraph.send_to_host != nullptr) { - graph_->AddControlEdge(oc_subgraph.send_to_host, recv); + graph_->AddControlEdge(host_compute, dst_image); } } } @@ -882,6 +923,63 @@ Status Encapsulator::Subgraph::BuildFunctionDef( return Status::OK(); } +Status Encapsulator::Subgraph::AddShapeInferenceInfo( + const string& outside_compilation_subgraph_name, + const std::vector& shapes, GraphDef* inference_graph) { + OutsideCompilationSubgraph& oc_subgraph = + outside_compilation_subgraphs_.at(outside_compilation_subgraph_name); + + Node* host_compute = nullptr; + for (Node* n : graph_->nodes()) { + if (n->name() == oc_subgraph.host_compute_name) { + host_compute = n; + break; + } + } + if (host_compute == nullptr) { + return errors::InvalidArgument( + "After rewriting subgraph ", outside_compilation_subgraph_name, + " there is no HostCompute Op for outside compilation subgraph ", + oc_subgraph.host_compute_name); + } + + if (inference_graph == nullptr) { + host_compute->AddAttr("shape_inference_graph", ""); + host_compute->AddAttr("shapes", shapes); + } else { + string serialized_graph; + if (!inference_graph->SerializeToString(&serialized_graph)) { + return errors::Internal( + "Failed to serialize graph for outside compilation subgraph ", + oc_subgraph.host_compute_name); + } + host_compute->AddAttr("shape_inference_graph", serialized_graph); + host_compute->AddAttr("shapes", std::vector()); + } + return Status::OK(); +} + +Status Encapsulator::Subgraph::ReplaceFunctionDef( + FunctionLibraryDefinition* library) { + const string& name = call_node_def_.name(); + + FunctionDef fdef; + TF_RETURN_IF_ERROR(GraphToFunctionDef(*graph_, name, &fdef)); + + if (VLOG_IS_ON(1)) { + VLOG(2) << "Replace function def " << name; + dump_graph::DumpGraphToFile( + strings::StrCat("replace_encapsulate_fdef_graph_", name), *graph_, + library); + dump_graph::DumpFunctionDefToFile( + strings::StrCat("replace_encapsulate_fdef_", name), fdef); + } + + TF_RETURN_IF_ERROR(library->RemoveFunction(name)); + TF_RETURN_IF_ERROR(library->AddFunctionDef(fdef)); + return Status::OK(); +} + Status Encapsulator::Subgraph::BuildParallelCheckOp( const std::unordered_map& node_images, Graph* graph_out) { @@ -980,7 +1078,9 @@ Status Encapsulator::Subgraph::AddRecvAtHostNode( NodeDefBuilder builder(strings::StrCat("outside_compilation_", subgraph_name, "_", oc_subgraph_name, "_recv"), kRecvAtHostOp); - builder.Attr("dtypes", dtypes); + builder.Attr("Toutputs", dtypes); + builder.Attr("key", strings::StrCat("host_compute_channel_", subgraph_name, + "_", oc_subgraph_name)); Status s = builder.Finalize(&recv_def); if (!s.ok()) return s; @@ -1020,7 +1120,9 @@ Status Encapsulator::Subgraph::AddSendFromHostNode( NodeDefBuilder builder(strings::StrCat("outside_compilation_", subgraph_name, "_", oc_subgraph_name, "_send"), kSendFromHostOp); - builder.Attr("dtypes", dtypes); + builder.Attr("Tinputs", dtypes); + builder.Attr("key", strings::StrCat("host_compute_channel_", subgraph_name, + "_", oc_subgraph_name)); builder.Input(inputs); Status s = builder.Finalize(&send_def); if (!s.ok()) return s; @@ -1062,6 +1164,13 @@ Status Encapsulator::Subgraph::AddOutsideCompilationHostIONodes( return Status::OK(); } +void Encapsulator::Subgraph::GetOutsideCompilationSubgraphNames( + std::vector* names) const { + for (auto& entry : outside_compilation_subgraphs_) { + names->push_back(entry.first); + } +} + Status Encapsulator::GetFunctionNameAttr( Node const* node, string* attr, string* outside_compilation_attr) const { Status s = GetNodeAttr(node->attrs(), group_attribute_, attr); @@ -1220,8 +1329,7 @@ Status Encapsulator::SplitIntoSubgraphs() { // single input and output node for it. for (auto& entry : subgraphs_) { Subgraph& subgraph = entry.second; - TF_RETURN_IF_ERROR(subgraph.AddSendsToOutsideCompilation(node_images)); - TF_RETURN_IF_ERROR(subgraph.AddRecvsFromOutsideCompilation(node_images)); + TF_RETURN_IF_ERROR(subgraph.AddHostComputes(entry.first, node_images)); } MarkGuaranteedConstants(*graph_in_, src_arg_pairs); @@ -1509,8 +1617,346 @@ Status Encapsulator::AddEdgesToOutputGraph( return Status::OK(); } -Status Encapsulator::BuildOutputGraph(bool parallel_checking, - Graph* graph_out) { +namespace { + +// Adds a dummy Const node to graph_out. The "constant" has the type of +// data_type and the shape indicated in 'shape'. The dummy node is not a valid +// Const node because it does not have any value defined, but this doesn't +// matter because it will only be used subsequently for shape inference. (It +// would be possible to add a switch statement over data_type to create a value +// for the constant, but that would entail maintaining the logic as new types +// are added, and is not necessary.) +Node* AddDummyShapedNode(DataType data_type, const TensorShapeProto& shape, + Graph* graph_out) { + TensorProto dummy_proto; + dummy_proto.set_dtype(data_type); + *dummy_proto.mutable_tensor_shape() = shape; + // Don't set any value field in the proto, since it is only going to be used + // for shape inference. + + GraphDefBuilder::Options options(graph_out, /*status=*/nullptr); + NodeBuilder node_builder(options.GetNameForOp("KnownShape"), "Const", + options.op_registry()); + node_builder.Attr("dtype", data_type).Attr("value", dummy_proto); + return options.FinalizeBuilder(&node_builder); +} + +// Adds a copy of node_in to graph_out and adds the mapping to +// copied_node_images. +Status CopyShapeInferenceNodeToGraph( + Node* node_in, const Node* send_node, + const std::unordered_map& dummy_node_images, + FunctionLibraryDefinition* library, + std::unordered_map* copied_node_images, Graph* graph_out) { + // Once all the ancestor nodes have been added to graph_out, add this node + // and connect it to its ancestors. + Node* node_out = graph_out->CopyNode(node_in); + (*copied_node_images)[node_in] = node_out; + // Don't bother to build the shape inference graph if there's a node with no + // shape inference function, since it would just result in an error later at + // compile time. + const OpRegistrationData* op_reg_data; + TF_RETURN_IF_ERROR(library->LookUp(node_in->type_string(), &op_reg_data)); + if (op_reg_data->shape_inference_fn == nullptr) { + return errors::InvalidArgument( + "Shape inference is not possible for outside_compilation " + "SendFromHost node ", + send_node->name(), " because it depends on node ", node_in->name(), + " which does not have a shape inference function registered."); + } + // Add all the edges to the newly copied node. + for (const Edge* in_edge : node_in->in_edges()) { + if (!in_edge->IsControlEdge()) { + Node* src = in_edge->src(); + const auto iter = dummy_node_images.find(src); + if (iter == dummy_node_images.end()) { + // The src is a copied node so use the original output port. + graph_out->AddEdge((*copied_node_images)[in_edge->src()], + in_edge->src_output(), node_out, + in_edge->dst_input()); + } else { + // The src is a dummy node so use output port 0. + graph_out->AddEdge(iter->second, 0, node_out, in_edge->dst_input()); + } + } + } + return Status::OK(); +} + +} // namespace + +Status Encapsulator::DoStaticShapeInferenceForOutsideCompilationSend( + const Graph& graph_in, const ShapeRefiner& shape_refiner, + const std::unordered_set& recv_at_host_nodes, Node* send_node, + FunctionLibraryDefinition* library, + std::vector* static_shape_out, + std::unique_ptr* graphdef_out) { + // Maps from nodes in graph_in to nodes in graph_out. + // + // When an edge has fully defined shape the source node in graph_in is + // replaced in graph_out by a dummy constant node. The mapping from nodes + // in graph_in to dummy nodes is stored in dummy_node_images. + // + // When a node in graph_in has at least one ancestor that doesn't have fully + // defined shape, it is copied into graph_out. The mapping from nodes in + // graph_in to copied nodes is stored in copied_node_images. + // + // The two types of node are treated differently because, when adding edges to + // graph_out, an output from a dummy node always uses port 0, whereas an + // output from a copied node uses the same port that was used in graph_in. + std::unordered_map dummy_node_images; + std::unordered_map copied_node_images; + + std::unique_ptr graph_out(new Graph(graph_in.op_registry())); + graph_out->set_versions(graph_in.versions()); + static_shape_out->resize(send_node->num_inputs()); + + // We don't use the standard ReverseDFS because we want to cut off traversal + // whenever we find an output with fully defined shape. + // TODO(misard) make this work properly in the presence of control flow. + struct Work { + Node* node; + bool leave; // Are we entering or leaving node? + }; + std::vector stack({{send_node, false}}); + std::vector visited(graph_in.num_node_ids(), false); + while (!stack.empty()) { + Work w = stack.back(); + stack.pop_back(); + Node* n = w.node; + + if (w.leave) { + TF_RETURN_IF_ERROR(CopyShapeInferenceNodeToGraph( + n, send_node, dummy_node_images, library, &copied_node_images, + graph_out.get())); + } else { + if (visited[n->id()]) continue; + visited[n->id()] = true; + + // Arrange to revisit when all done with all inputs. + stack.push_back(Work{n, true}); + + bool has_parent_with_unknown_shape = false; + for (const Edge* in_edge : n->in_edges()) { + if (!in_edge->IsControlEdge()) { + Node* src_node = in_edge->src(); + int src_port = in_edge->src_output(); + shape_inference::InferenceContext* context = + shape_refiner.GetContext(src_node); + shape_inference::ShapeHandle shape = context->output(src_port); + if (context->FullyDefined(shape)) { + // This ancestor has known shape, so instead of adding it to the + // stack, add a dummy node with that shape to graph_out and + // continue. + TensorShapeProto proto; + context->ShapeHandleToProto(shape, &proto); + dummy_node_images[src_node] = AddDummyShapedNode( + src_node->output_type(src_port), proto, graph_out.get()); + if (n == send_node) { + (*static_shape_out)[in_edge->dst_input()] = proto; + } + } else { + if (!visited[src_node->id()]) { + has_parent_with_unknown_shape = true; + stack.push_back({src_node, false}); + } + } + } + } + if (!has_parent_with_unknown_shape) { + if (n == send_node) { + // The shapes of all the inputs to send_node are statically known. We + // won't have to do any inference at compile time so return now: the + // shapes were stored in static_shape_out above. + graphdef_out->reset(); + return Status::OK(); + } else { + // Any shape that is being processed is either the original send node + // or has at least one output with statically-unknown shape. If the + // latter and it doesn't have any inputs with statically-unknown + // shape, then check that it is of the recv nodes that we can fill in + // the shape of at run-time later. If it isn't one of those, then we + // won't have any additional knowledge at compile time, so we already + // know we won't be able to do shape inference and we can return an + // error now. + if (recv_at_host_nodes.find(n->name()) == recv_at_host_nodes.end()) { + return errors::InvalidArgument( + "Shape inference is not possible for outside_compilation " + "SendFromHost node ", + send_node->name(), " because shape of node ", n->name(), + " will not be known at compilation time."); + } + } + } + } + } + + graphdef_out->reset(new GraphDef()); + graph_out->ToGraphDef(graphdef_out->get()); + + return Status::OK(); +} + +Status Encapsulator::MakePrunedGraphCopyAndInline( + const Graph& graph, const std::vector& sink_nodes, + std::unique_ptr* pruned_graph, + std::unordered_map* node_images, + FunctionLibraryDefinition* library) { + // First copy all ancestor nodes of sink_nodes into a new graph. + pruned_graph->reset(new Graph(library)); + (*pruned_graph)->set_versions(graph.versions()); + ReverseDFSFrom(graph, sink_nodes, + /*enter=*/nullptr, + /*leave=*/[&](Node* n) { + if (!n->IsSource()) { + Node* copied = (*pruned_graph)->CopyNode(n); + node_images->emplace(n, copied); + } + }); + + // Add all the edges between copied nodes. + for (auto entry : *node_images) { + const Node* orig = entry.first; + Node* image = entry.second; + for (const Edge* out_edge : orig->out_edges()) { + auto iter = node_images->find(out_edge->dst()); + if (iter != node_images->end()) { + // The source and destination are both in the copied graph. + (*pruned_graph) + ->AddEdge(image, out_edge->src_output(), iter->second, + out_edge->dst_input()); + } + } + } + + // Find all the function call nodes, and inline them. + std::vector function_nodes; + for (auto node : (*pruned_graph)->nodes()) { + const OpRegistrationData* op_reg_data; + TF_RETURN_IF_ERROR(library->LookUp(node->type_string(), &op_reg_data)); + if (op_reg_data->is_function_op) { + function_nodes.push_back(node); + } + } + for (auto node : function_nodes) { + VLOG(2) << "Inlining function " << node->name(); + const FunctionDef* fdef = library->Find(node->type_string()); + if (fdef == nullptr) { + return errors::Internal("Failed to find function ", node->type_string(), + " in function library."); + } + FunctionBody* fbody = nullptr; + TF_RETURN_IF_ERROR( + FunctionDefToBodyHelper(*fdef, node->attrs(), library, + [library](const string& op, const OpDef** sig) { + return library->LookUpOpDef(op, sig); + }, + &fbody)); + InlineFunctionBody(*library, pruned_graph->get(), node, fbody); + delete fbody; + } + + return Status::OK(); +} + +Status Encapsulator::MakeGraphForOutsideCompilationSends( + const Graph& graph, std::unique_ptr* pruned_graph, + ShapeRefiner* shape_refiner, + std::unordered_map* node_images, + FunctionLibraryDefinition* library) { + // Find all the send_from_host nodes in all subgraphs, to use as roots for the + // pruning. + std::vector send_from_host_nodes; + for (auto& subgraph_entry : subgraphs_) { + Subgraph& subgraph = subgraph_entry.second; + std::vector outside_compilation_names; + subgraph.GetOutsideCompilationSubgraphNames(&outside_compilation_names); + for (const auto& name : outside_compilation_names) { + Node* send_node = subgraph.GetSendFromHostNode(name); + if (send_node != nullptr) { + send_from_host_nodes.push_back(send_node); + } + } + } + + // Make a copy of all the graph nodes needed to evaluate the send_from_host + // nodes, inlining any functions as needed. + TF_RETURN_IF_ERROR(MakePrunedGraphCopyAndInline( + graph, send_from_host_nodes, pruned_graph, node_images, library)); + + // Perform shape inference on the pruned graph. + shape_refiner->set_require_shape_inference_fns(false); + FixupSourceAndSinkEdges(pruned_graph->get()); + std::vector post_order; + GetReversePostOrder(*(*pruned_graph), &post_order); + for (auto node : post_order) { + // Ignore the status returned by the shape_refiner. At this point we want + // the best effort shapes, even if no shape function is registered for a + // node. + Status status = shape_refiner->AddNode(node); + if (!status.ok()) { + VLOG(1) << "Shape inference failed for node: " << status; + } + } + + return Status::OK(); +} + +Status Encapsulator::GetShapeInfoForOutsideCompilationSends( + Graph* graph_out, FunctionLibraryDefinition* library) { + std::unique_ptr pruned_graph; + ShapeRefiner shape_refiner(graph_out->versions(), graph_out->op_registry()); + std::unordered_map node_images; + TF_RETURN_IF_ERROR(MakeGraphForOutsideCompilationSends( + *graph_out, &pruned_graph, &shape_refiner, &node_images, library)); + + for (auto& subgraph_entry : subgraphs_) { + Subgraph& subgraph = subgraph_entry.second; + // Find all the recv_at_host nodes in this subgraph. + std::vector outside_compilation_names; + subgraph.GetOutsideCompilationSubgraphNames(&outside_compilation_names); + std::unordered_set recv_at_host_names; + for (const auto& name : outside_compilation_names) { + Node* recv_node = subgraph.GetRecvAtHostNode(name); + if (recv_node != nullptr) { + recv_at_host_names.insert(recv_node->name()); + } + } + // For each send_from_host node, do as much shape inference as possible + // without knowing the shape of the recv_at_host nodes, and store the + // result, along with enough information to complete the job at compile time + // once the recv_at_host shapes are known. + for (const auto& name : outside_compilation_names) { + Node* send_node = subgraph.GetSendFromHostNode(name); + std::vector static_shape; + std::unique_ptr graphdef; + if (send_node != nullptr) { + TF_RETURN_IF_ERROR(DoStaticShapeInferenceForOutsideCompilationSend( + *pruned_graph, shape_refiner, recv_at_host_names, + node_images[send_node], library, &static_shape, &graphdef)); + if (graphdef == nullptr) { + VLOG(2) << "Send node " << send_node->name() << " shapes"; + for (int i = 0; i < static_shape.size(); ++i) { + VLOG(2) << static_shape[i].DebugString(); + } + } else { + VLOG(2) << "Send node " << send_node->name() << " graph\n" + << graphdef->DebugString(); + } + } + TF_RETURN_IF_ERROR( + subgraph.AddShapeInferenceInfo(name, static_shape, graphdef.get())); + } + if (!outside_compilation_names.empty()) { + TF_RETURN_IF_ERROR(subgraph.ReplaceFunctionDef(library)); + } + } + + return Status::OK(); +} + +Status Encapsulator::BuildOutputGraph(bool parallel_checking, Graph* graph_out, + FunctionLibraryDefinition* library) { // Map from nodes in the input graph to nodes in the output graph. std::unordered_map node_images; @@ -1522,6 +1968,9 @@ Status Encapsulator::BuildOutputGraph(bool parallel_checking, TF_RETURN_IF_ERROR( AddEdgesToOutputGraph(node_images, parallel_checking, graph_out)); + TF_RETURN_IF_ERROR( + GetShapeInfoForOutsideCompilationSends(graph_out, library)); + return Status::OK(); } @@ -1545,7 +1994,7 @@ Status EncapsulateSubgraphsInFunctions( std::unique_ptr out(new Graph(library)); out->set_versions(graph_in.versions()); TF_RETURN_IF_ERROR( - encapsulator.BuildOutputGraph(parallel_checking, out.get())); + encapsulator.BuildOutputGraph(parallel_checking, out.get(), library)); *graph_out = std::move(out); return Status::OK(); diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc index b100861d5e9c04a8f9d32d486e0ee7252b79c62b..aed9cae0f1799c4524da8ee309344849798755d5 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc @@ -29,17 +29,181 @@ limitations under the License. namespace tensorflow { namespace { +template +bool EqualProtoMap(const ::tensorflow::protobuf::Map& a, + const ::tensorflow::protobuf::Map& b, + const std::function& key_to_string, + const std::function& value_to_string, + const std::function& compare, + const string& map_name, string* diff) { + for (const auto& elt_a : a) { + const auto iter = b.find(elt_a.first); + if (iter == b.end()) { + if (diff) { + *diff = strings::StrCat( + map_name, " expected: contains element with key '", + key_to_string(elt_a.first), "' got: map has no such element"); + } + return false; + } + if (!compare(elt_a.first, elt_a.second, iter->second)) { + if (diff) { + *diff = strings::StrCat(map_name, " expected: element with key '", + key_to_string(elt_a.first), " has value '", + value_to_string(elt_a.second), "' got: '", + value_to_string(iter->second), "'"); + } + return false; + } + } + for (const auto& elt_b : b) { + const auto iter = a.find(elt_b.first); + if (iter == a.end()) { + if (diff) { + *diff = strings::StrCat(map_name, " got: contains element with key '", + key_to_string(elt_b.first), + "' expected: map has no such element"); + } + return false; + } + } + return true; +} + +bool EqualFunctionNodeDef(const NodeDef& a, const NodeDef& b, + const string& diff_preamble, string* diff) { + if (a.op() != b.op()) { + if (diff) { + *diff = strings::StrCat(diff_preamble, " mismatch for node ", a.name(), + ", expected op '", a.op(), "' got '", b.op()); + } + return false; + } + if (a.device() != b.device()) { + if (diff) { + *diff = strings::StrCat(diff_preamble, " mismatch for node ", a.name(), + ", expected device '", a.device(), "' got '", + b.device()); + } + return false; + } + if (a.input_size() != b.input_size()) { + if (diff) { + *diff = strings::StrCat(diff_preamble, " mismatch for node ", a.name(), + ", expected ", a.input_size(), " inputs got ", + b.input_size(), " expected:\n", a.DebugString(), + "\ngot:\n", b.DebugString()); + } + return false; + } + for (int i = 0; i < a.input_size(); ++i) { + if (a.input(i) != b.input(i)) { + if (diff) { + *diff = strings::StrCat(diff_preamble, " mismatch for node ", a.name(), + " input ", i, ", expected ", a.input(i), + " got ", b.input(i), " expected:\n", + a.DebugString(), "\ngot:\n", b.DebugString()); + } + return false; + } + } + return EqualProtoMap( + a.attr(), b.attr(), [](const string& s) { return s; }, + [](const AttrValue& v) { return v.DebugString(); }, + [](const string& key, const AttrValue& av, const AttrValue& bv) { + if (key == "shape_inference_graph") { + // Default serialization of GraphDef is unstable because maps don't + // serialize deterministically. Rather than go through the hoops to + // turn on deterministic serialization of this attr just for this + // test, add logic here to compare determinstically. + GraphDef ga; + if (!ga.ParseFromString(av.s())) { + return false; + } + GraphDef gb; + if (!gb.ParseFromString(bv.s())) { + return false; + } + return EqualGraphDef(ga, gb, nullptr); + } else { + return av.DebugString() == bv.DebugString(); + } + }, + strings::StrCat(diff_preamble, " attr mismatch for node ", a.name()), + diff); +} + bool EqualFunctionDef(const FunctionDef& a, const FunctionDef& b, string* diff) { - // TODO(phawkins) use a more sophisticated equality test. - if (a.DebugString() != b.DebugString()) { + if (a.signature().DebugString() != b.signature().DebugString()) { if (diff) { - *diff = strings::StrCat("Definition mismatch for function ", + *diff = strings::StrCat("Signature mismatch for function ", a.signature().name(), ", expected:\n", - a.DebugString(), "\ngot:\n", b.DebugString()); + a.signature().DebugString(), "\ngot:\n", + b.signature().DebugString()); } return false; } + if (!EqualProtoMap( + a.attr(), b.attr(), [](const string& s) { return s; }, + [](const AttrValue& v) { return v.DebugString(); }, + [](const string& key, const AttrValue& av, const AttrValue& bv) { + return av.DebugString() == bv.DebugString(); + }, + strings::StrCat("attr mismatch for function ", a.signature().name()), + diff)) { + return false; + } + if (!EqualProtoMap( + a.ret(), b.ret(), [](const string& s) { return s; }, + [](const string& s) { return s; }, + [](const string& key, const string& av, const string& bv) { + return av == bv; + }, + strings::StrCat("ret mismatch for function ", a.signature().name()), + diff)) { + return false; + } + for (int i = 0; i < a.node_def_size(); ++i) { + bool found = false; + for (int j = 0; j < b.node_def_size(); ++j) { + if (a.node_def(i).name() == b.node_def(j).name()) { + if (!EqualFunctionNodeDef( + a.node_def(i), b.node_def(j), + strings::StrCat("Function ", a.signature().name()), diff)) { + return false; + } + found = true; + break; + } + } + if (!found) { + if (diff) { + *diff = strings::StrCat("Function ", a.signature().name(), + ", expected: has node '", a.node_def(i).name(), + "' got: no node of that name"); + } + return false; + } + } + for (int i = 0; i < b.node_def_size(); ++i) { + bool found = false; + for (int j = 0; j < a.node_def_size(); ++j) { + if (b.node_def(i).name() == a.node_def(j).name()) { + found = true; + break; + } + } + if (!found) { + if (diff) { + *diff = strings::StrCat("Function ", a.signature().name(), + ", got: has node '", b.node_def(i).name(), + "' expected: no node of that name"); + } + return false; + } + } return true; } @@ -84,29 +248,64 @@ bool EqualFunctionDefLibrary(const FunctionDefLibrary& expected, // TODO(misard): remove these fake registrations once there are real Ops to be // compiled. -REGISTER_OP("_XlaSendToHost") - .Input("input: dtypes") - .Attr("dtypes: list(type) >= 0"); - -REGISTER_OP("_XlaRecvFromHost") - .Output("output: dtypes") - .Attr("dtypes: list(type) >= 0"); +REGISTER_OP("_XlaHostCompute") + .Input("inputs: Tinputs") + .Output("outputs: Toutputs") + .Attr("Tinputs: list(type) >= 0") + .Attr("Toutputs: list(type) >= 0") + .Attr("key: string") + .SetShapeFn(::tensorflow::shape_inference::UnknownShape); REGISTER_OP("_XlaSendFromHost") - .Input("input: dtypes") - .Attr("dtypes: list(type) >= 0"); + .Input("input: Tinputs") + .Attr("Tinputs: list(type) >= 0") + .Attr("key: string") + .SetShapeFn(::tensorflow::shape_inference::UnknownShape); REGISTER_OP("_XlaRecvAtHost") - .Output("output: dtypes") - .Attr("dtypes: list(type) >= 0"); - -REGISTER_OP("InputTest").Output("o: float"); - -REGISTER_OP("UnaryTest").Input("a: float").Output("o: float"); + .Output("output: Toutputs") + .Attr("Toutputs: list(type) >= 0") + .Attr("key: string") + .SetShapeFn(::tensorflow::shape_inference::UnknownShape); + +REGISTER_OP("InputTest") + .Output("o: float") + .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { + c->set_output(0, c->UnknownShape()); + return Status::OK(); + }); + +REGISTER_OP("InputTestShaped") + .Output("o: float") + .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { + c->set_output(0, c->Vector(2)); + return Status::OK(); + }); + +REGISTER_OP("UnaryTest") + .Input("a: float") + .Output("o: float") + .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { + ::tensorflow::shape_inference::ShapeHandle o; + TF_RETURN_IF_ERROR(c->Merge(c->UnknownShape(), c->input(0), &o)); + c->set_output(0, o); + return Status::OK(); + }); REGISTER_OP("BinaryTest") .Input("a: float") .Input("b: float") - .Output("o: float"); + .Output("o: float") + .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { + ::tensorflow::shape_inference::ShapeHandle o; + TF_RETURN_IF_ERROR(c->Merge(c->UnknownShape(), c->input(0), &o)); + c->set_output(0, o); + return Status::OK(); + }); +REGISTER_OP("BinaryTest2") + .Input("a: float") + .Input("b: float") + .Output("o: float") + .SetShapeFn(::tensorflow::shape_inference::UnknownShape); REGISTER_OP("AddNLikeTest") .Input("inputs: N * T") @@ -124,22 +323,48 @@ Node* Input(const GraphDefBuilder::Options& opts) { return ops::SourceOp("InputTest", opts); } -Node* RecvAtHost(const gtl::ArraySlice& dtypes, +Node* InputShaped(const GraphDefBuilder::Options& opts) { + return ops::SourceOp("InputTestShaped", opts); +} + +Node* KnownShape(const gtl::ArraySlice& shape, + const GraphDefBuilder::Options& opts) { + if (opts.HaveError()) return nullptr; + NodeBuilder node_builder(opts.GetNameForOp("Const"), "Const", + opts.op_registry()); + TensorProto value; + value.set_dtype(DT_FLOAT); + for (int dim : shape) { + value.mutable_tensor_shape()->add_dim()->set_size(dim); + } + return opts.WithAttr("value", value) + .WithAttr("dtype", DT_FLOAT) + .FinalizeBuilder(&node_builder); +} + +Node* RecvAtHost(const string& key, const gtl::ArraySlice& dtypes, const GraphDefBuilder::Options& opts) { if (opts.HaveError()) return nullptr; NodeBuilder node_builder(opts.GetNameForOp("_XlaRecvAtHost"), "_XlaRecvAtHost", opts.op_registry()); - return opts.WithAttr("dtypes", dtypes).FinalizeBuilder(&node_builder); + return opts.WithAttr("Toutputs", dtypes) + .WithAttr("key", key) + .FinalizeBuilder(&node_builder); } -Node* SendFromHost(const std::vector& inputs, - const gtl::ArraySlice& dtypes, +Node* SendFromHost(const string& key, const std::vector& inputs, const GraphDefBuilder::Options& opts) { if (opts.HaveError()) return nullptr; NodeBuilder node_builder(opts.GetNameForOp("_XlaSendFromHost"), "_XlaSendFromHost", opts.op_registry()); node_builder.Input(inputs); - return opts.WithAttr("dtypes", dtypes).FinalizeBuilder(&node_builder); + std::vector dtypes; + for (const auto& node : inputs) { + dtypes.push_back(node.dt); + } + return opts.WithAttr("key", key) + .WithAttr("Tinputs", dtypes) + .FinalizeBuilder(&node_builder); } Node* Unary(ops::NodeOut a, const GraphDefBuilder::Options& opts) { @@ -151,6 +376,11 @@ Node* Binary(ops::NodeOut a, ops::NodeOut b, return ops::BinaryOp("BinaryTest", std::move(a), std::move(b), opts); } +Node* BinaryUnknownShape(ops::NodeOut a, ops::NodeOut b, + const GraphDefBuilder::Options& opts) { + return ops::BinaryOp("BinaryTest2", std::move(a), std::move(b), opts); +} + Node* AddNLike(const std::vector& inputs, const GraphDefBuilder::Options& opts) { if (opts.HaveError()) return nullptr; @@ -576,6 +806,21 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { FunctionDefLibrary library_expected; GraphDef graphdef_expected; + string shape_string_expected; + { + GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); + Node* recv = + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT, DT_FLOAT}, + shape.opts().WithName("outside_compilation_F1_O1_recv")); + Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), + shape.opts().WithName("E")); + SendFromHost("host_compute_channel_F1_O1", {e}, + shape.opts().WithName("outside_compilation_F1_O1_send")); + GraphDef shape_graph; + TF_EXPECT_OK(shape.ToGraphDef(&shape_graph)); + EXPECT_TRUE(shape_graph.SerializeToString(&shape_string_expected)); + } + *library_expected.add_function() = test::function::XTimesTwo(); *library_expected.add_function() = FunctionDefHelper::Create( "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, @@ -584,19 +829,18 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { {{"c"}, "BinaryTest", {"b_0_arg", "C:o:0"}, {}, {"C"}}, {{"F"}, "BinaryTest", - {"C:o:0", "outside_compilation_O1_recv:output:0"}, + {"C:o:0", "outside_compilation_O1_host_compute:outputs:0"}, {}, - {"outside_compilation_O1_recv"}}, - {{"outside_compilation_O1_send"}, - "_XlaSendToHost", + {"outside_compilation_O1_host_compute"}}, + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", {"C:o:0", "c:o:0"}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", shape_string_expected}, + {"shapes", gtl::ArraySlice({})}}, {"c"}}, - {{"outside_compilation_O1_recv"}, - "_XlaRecvFromHost", - {}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}, - {"outside_compilation_O1_send"}}, }, {{"f_0_retval", "F:o:0"}}); @@ -612,11 +856,11 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { Node* call = b2.opts().FinalizeBuilder(&node_builder); Node* recv = - RecvAtHost({DT_FLOAT, DT_FLOAT}, + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT, DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_recv")); Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), b2.opts().WithName("E").WithControlInputs({recv, b})); - Node* send = SendFromHost({e}, {DT_FLOAT}, + Node* send = SendFromHost("host_compute_channel_F1_O1", {e}, b2.opts() .WithName("outside_compilation_F1_O1_send") .WithControlInput(e)); @@ -674,37 +918,71 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { FunctionDefLibrary library_expected; GraphDef graphdef_expected; + string shape_string_expected_1; + { + GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); + Node* recv = + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT, DT_FLOAT}, + shape1.opts().WithName("outside_compilation_F1_O1_recv")); + Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), + shape1.opts().WithName("E")); + SendFromHost("host_compute_channel_F1_O1", {e}, + shape1.opts().WithName("outside_compilation_F1_O1_send")); + GraphDef shape1_graph; + TF_EXPECT_OK(shape1.ToGraphDef(&shape1_graph)); + EXPECT_TRUE(shape1_graph.SerializeToString(&shape_string_expected_1)); + } + + string shape_string_expected_2; + { + GraphDefBuilder shape2(GraphDefBuilder::kFailImmediately); + Node* recv1 = + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT, DT_FLOAT}, + shape2.opts().WithName("outside_compilation_F1_O1_recv")); + Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1), + shape2.opts().WithName("E")); + Node* recv2 = + RecvAtHost("host_compute_channel_F1_O2", {DT_FLOAT, DT_FLOAT}, + shape2.opts().WithName("outside_compilation_F1_O2_recv")); + Node* h = Binary(ops::NodeOut(recv2, 0), e, shape2.opts().WithName("H")); + SendFromHost("host_compute_channel_F1_O2", {h}, + shape2.opts().WithName("outside_compilation_F1_O2_send")); + GraphDef shape2_graph; + TF_EXPECT_OK(shape2.ToGraphDef(&shape2_graph)); + EXPECT_TRUE(shape2_graph.SerializeToString(&shape_string_expected_2)); + } + *library_expected.add_function() = FunctionDefHelper::Create( "F1", {"a_0_arg:float", "b_0_arg:float"}, {"i_0_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}, {}}, - {{"I"}, "UnaryTest", {"outside_compilation_O2_recv:output:0"}}, + {{"I"}, + "UnaryTest", + {"outside_compilation_O2_host_compute:outputs:0"}}, {{"F"}, "BinaryTest", - {"C:o:0", "outside_compilation_O1_recv:output:0"}, + {"C:o:0", "outside_compilation_O1_host_compute:outputs:0"}, {}, - {"outside_compilation_O1_recv"}}, - {{"outside_compilation_O2_send"}, - "_XlaSendToHost", + {"outside_compilation_O1_host_compute"}}, + {{"outside_compilation_O2_host_compute"}, + "_XlaHostCompute", {"D:o:0", "F:o:0"}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"key", "host_compute_channel_F1_O2"}, + {"shape_inference_graph", shape_string_expected_2}, + {"shapes", gtl::ArraySlice({})}}, {"F"}}, - {{"outside_compilation_O1_send"}, - "_XlaSendToHost", + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", {"C:o:0", "D:o:0"}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", shape_string_expected_1}, + {"shapes", gtl::ArraySlice({})}}, {"D"}}, - {{"outside_compilation_O2_recv"}, - "_XlaRecvFromHost", - {}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}, - {"outside_compilation_O2_send"}}, - {{"outside_compilation_O1_recv"}, - "_XlaRecvFromHost", - {}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}, - {"outside_compilation_O1_send"}}, }, {{"i_0_retval", "I:o:0"}}); @@ -720,23 +998,24 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { Node* call = b2.opts().FinalizeBuilder(&node_builder); Node* recv1 = - RecvAtHost({DT_FLOAT, DT_FLOAT}, + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT, DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_recv")); Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1), b2.opts().WithName("E").WithControlInputs({recv1, b})); - Node* send1 = SendFromHost({e}, {DT_FLOAT}, + Node* send1 = SendFromHost("host_compute_channel_F1_O1", {e}, b2.opts() .WithName("outside_compilation_F1_O1_send") .WithControlInput(e)); Node* recv2 = - RecvAtHost({DT_FLOAT, DT_FLOAT}, + RecvAtHost("host_compute_channel_F1_O2", {DT_FLOAT, DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O2_recv")); Node* g = Binary(e, ops::NodeOut(recv2, 1), b2.opts().WithName("G").WithControlInputs({recv2, e})); Node* h = Binary(ops::NodeOut(recv2, 0), e, b2.opts().WithName("H")); - Node* send2 = SendFromHost( - {h}, {DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O2_send")); + Node* send2 = + SendFromHost("host_compute_channel_F1_O2", {h}, + b2.opts().WithName("outside_compilation_F1_O2_send")); Node* s = NoOp(b2.opts() .WithName("F1_sequencer") @@ -758,8 +1037,8 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { { GraphDefBuilder b1(GraphDefBuilder::kFailImmediately); - Node* a = Input(b1.opts().WithName("A")); - Node* b = Input(b1.opts().WithName("B")); + Node* a = InputShaped(b1.opts().WithName("A")); + Node* b = InputShaped(b1.opts().WithName("B")); Node* c = Unary(a, b1.opts().WithName("C").WithAttr("_encapsulate", "F1")); Node* d = Binary(b, c, b1.opts().WithName("D").WithAttr("_encapsulate", "F1")); @@ -791,6 +1070,24 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { FunctionDefLibrary library_expected; GraphDef graphdef_expected; + string shape_string_expected; + { + GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); + Node* recv = + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT, DT_FLOAT}, + shape.opts().WithName("outside_compilation_F1_O1_recv")); + Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), + shape.opts().WithName("E")); + SendFromHost("host_compute_channel_F1_O1", {e}, + shape.opts().WithName("outside_compilation_F1_O1_send")); + GraphDef shape_graph; + TF_EXPECT_OK(shape.ToGraphDef(&shape_graph)); + EXPECT_TRUE(shape_graph.SerializeToString(&shape_string_expected)); + } + + TensorShapeProto shape_proto_expected; + shape_proto_expected.add_dim()->set_size(2); + *library_expected.add_function() = FunctionDefHelper::Create( "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float", "d_0_retval:float"}, {}, @@ -799,19 +1096,18 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, {{"F"}, "BinaryTest", - {"C:o:0", "outside_compilation_O1_recv:output:0"}, + {"C:o:0", "outside_compilation_O1_host_compute:outputs:0"}, {}, - {"outside_compilation_O1_recv"}}, - {{"outside_compilation_O1_send"}, - "_XlaSendToHost", + {"outside_compilation_O1_host_compute"}}, + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", {"C:o:0", "D:o:0"}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", shape_string_expected}, + {"shapes", gtl::ArraySlice({})}}, {"D"}}, - {{"outside_compilation_O1_recv"}, - "_XlaRecvFromHost", - {}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}, - {"outside_compilation_O1_send"}}, }, {{"d_0_retval", "D:o:0"}, {"f_0_retval", "F:o:0"}}); @@ -822,16 +1118,16 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { {{"G"}, "BinaryTest", {"e_0_arg", "f_0_arg"}}, {{"I"}, "BinaryTest", - {"f_0_arg", "outside_compilation_O1_recv:output:0"}}, - {{"outside_compilation_O1_send"}, - "_XlaSendToHost", + {"f_0_arg", "outside_compilation_O1_host_compute:outputs:0"}}, + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", {"G:o:0"}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}}, - {{"outside_compilation_O1_recv"}, - "_XlaRecvFromHost", - {}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}, - {"outside_compilation_O1_send"}}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"key", "host_compute_channel_F2_O1"}, + {"shape_inference_graph", ""}, + {"shapes", + gtl::ArraySlice({shape_proto_expected})}}}, }, {{"g_0_retval", "G:o:0"}, {"i_0_retval", "I:o:0"}}); @@ -839,15 +1135,15 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { std::unique_ptr lib_def( new FunctionLibraryDefinition(OpRegistry::Global(), library_expected)); GraphDefBuilder b2(GraphDefBuilder::kFailImmediately, lib_def.get()); - Node* a = Input(b2.opts().WithName("A")); - Node* b = Input(b2.opts().WithName("B")); + Node* a = InputShaped(b2.opts().WithName("A")); + Node* b = InputShaped(b2.opts().WithName("B")); Node* recv1 = - RecvAtHost({DT_FLOAT, DT_FLOAT}, + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT, DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_recv")); Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1), b2.opts().WithName("E").WithControlInputs({recv1, b})); - Node* send1 = SendFromHost({e}, {DT_FLOAT}, + Node* send1 = SendFromHost("host_compute_channel_F1_O1", {e}, b2.opts() .WithName("outside_compilation_F1_O1_send") .WithControlInput(e)); @@ -857,12 +1153,14 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { Node* s1 = NoOp( b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1})); - Node* recv2 = RecvAtHost( - {DT_FLOAT}, b2.opts().WithName("outside_compilation_F2_O1_recv")); + Node* recv2 = + RecvAtHost("host_compute_channel_F2_O1", {DT_FLOAT}, + b2.opts().WithName("outside_compilation_F2_O1_recv")); Node* h = Binary(ops::NodeOut(call1, 1), recv2, b2.opts().WithName("H").WithControlInput(s1)); - Node* send2 = SendFromHost( - {h}, {DT_FLOAT}, b2.opts().WithName("outside_compilation_F2_O1_send")); + Node* send2 = + SendFromHost("host_compute_channel_F2_O1", {h}, + b2.opts().WithName("outside_compilation_F2_O1_send")); NodeBuilder node_builder2("F2", "F2", lib_def.get()); node_builder2.Input(e).Input(call1); @@ -888,7 +1186,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { { GraphDefBuilder b1(GraphDefBuilder::kFailImmediately); - Node* a = Input(b1.opts().WithName("A")); + Node* a = InputShaped(b1.opts().WithName("A")); Node* b = Input(b1.opts().WithName("B")); Node* c = Unary(a, b1.opts().WithName("C").WithAttr("_encapsulate", "F1")); Node* d = @@ -908,6 +1206,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { FunctionDefLibrary library_expected; GraphDef graphdef_expected; + TensorShapeProto shape_proto_expected; + shape_proto_expected.add_dim()->set_size(2); + *library_expected.add_function() = FunctionDefHelper::Create( "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, { @@ -915,11 +1216,16 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, {{"F"}, "BinaryTest", - {"D:o:0", "outside_compilation_O1_recv:output:0"}}, - {{"outside_compilation_O1_recv"}, - "_XlaRecvFromHost", + {"D:o:0", "outside_compilation_O1_host_compute:outputs:0"}}, + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", {}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}}, + {{"Tinputs", gtl::ArraySlice({})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", ""}, + {"shapes", + gtl::ArraySlice({shape_proto_expected})}}}, }, {{"f_0_retval", "F:o:0"}}); @@ -927,12 +1233,13 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { std::unique_ptr lib_def( new FunctionLibraryDefinition(OpRegistry::Global(), library_expected)); GraphDefBuilder b2(GraphDefBuilder::kFailImmediately, lib_def.get()); - Node* a = Input(b2.opts().WithName("A")); + Node* a = InputShaped(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); Node* e = Unary(a, b2.opts().WithName("E")); - Node* send1 = SendFromHost( - {e}, {DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_send")); + Node* send1 = + SendFromHost("host_compute_channel_F1_O1", {e}, + b2.opts().WithName("outside_compilation_F1_O1_send")); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); @@ -954,7 +1261,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { { GraphDefBuilder b1(GraphDefBuilder::kFailImmediately); - Node* a = Input(b1.opts().WithName("A")); + Node* a = InputShaped(b1.opts().WithName("A")); Node* b = Input(b1.opts().WithName("B")); Node* c = Unary(a, b1.opts().WithName("C").WithAttr("_encapsulate", "F1")); Node* d = @@ -975,6 +1282,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { FunctionDefLibrary library_expected; GraphDef graphdef_expected; + TensorShapeProto shape_proto_expected; + shape_proto_expected.add_dim()->set_size(2); + *library_expected.add_function() = FunctionDefHelper::Create( "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, { @@ -982,17 +1292,17 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, {{"F"}, "BinaryTest", - {"D:o:0", "outside_compilation_O1_recv:output:0"}}, - {{"outside_compilation_O1_send"}, - "_XlaSendToHost", + {"D:o:0", "outside_compilation_O1_host_compute:outputs:0"}}, + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", {}, - {{"dtypes", gtl::ArraySlice({})}}, + {{"Tinputs", gtl::ArraySlice({})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", ""}, + {"shapes", + gtl::ArraySlice({shape_proto_expected})}}, {"D"}}, - {{"outside_compilation_O1_recv"}, - "_XlaRecvFromHost", - {}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}, - {"outside_compilation_O1_send"}}, }, {{"f_0_retval", "F:o:0"}}); @@ -1000,14 +1310,16 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { std::unique_ptr lib_def( new FunctionLibraryDefinition(OpRegistry::Global(), library_expected)); GraphDefBuilder b2(GraphDefBuilder::kFailImmediately, lib_def.get()); - Node* a = Input(b2.opts().WithName("A")); + Node* a = InputShaped(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); Node* recv1 = - RecvAtHost({}, b2.opts().WithName("outside_compilation_F1_O1_recv")); + RecvAtHost("host_compute_channel_F1_O1", {}, + b2.opts().WithName("outside_compilation_F1_O1_recv")); Node* e = Unary(a, b2.opts().WithName("E").WithControlInput(recv1)); - Node* send1 = SendFromHost( - {e}, {DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_send")); + Node* send1 = + SendFromHost("host_compute_channel_F1_O1", {e}, + b2.opts().WithName("outside_compilation_F1_O1_send")); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); @@ -1055,10 +1367,14 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoOutputs) { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, {{"F"}, "UnaryTest", {"D:o:0"}}, - {{"outside_compilation_O1_send"}, - "_XlaSendToHost", + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", {"D:o:0"}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", ""}, + {"shapes", gtl::ArraySlice({})}}}, }, {{"f_0_retval", "F:o:0"}}); @@ -1069,8 +1385,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoOutputs) { Node* a = Input(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); - Node* recv1 = RecvAtHost( - {DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_recv")); + Node* recv1 = + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT}, + b2.opts().WithName("outside_compilation_F1_O1_recv")); Node* e = Unary(recv1, b2.opts().WithName("E")); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); @@ -1118,16 +1435,19 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, - {{"F"}, "UnaryTest", {"D:o:0"}, {}, {"outside_compilation_O1_recv"}}, - {{"outside_compilation_O1_send"}, - "_XlaSendToHost", + {{"F"}, + "UnaryTest", {"D:o:0"}, - {{"dtypes", gtl::ArraySlice({DT_FLOAT})}}}, - {{"outside_compilation_O1_recv"}, - "_XlaRecvFromHost", {}, - {{"dtypes", gtl::ArraySlice({})}}, - {"outside_compilation_O1_send"}}, + {"outside_compilation_O1_host_compute"}}, + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", + {"D:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", ""}, + {"shapes", gtl::ArraySlice({})}}}, }, {{"f_0_retval", "F:o:0"}}); @@ -1138,10 +1458,11 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { Node* a = Input(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); - Node* recv1 = RecvAtHost( - {DT_FLOAT}, b2.opts().WithName("outside_compilation_F1_O1_recv")); + Node* recv1 = + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT}, + b2.opts().WithName("outside_compilation_F1_O1_recv")); Node* e = Unary(recv1, b2.opts().WithName("E")); - Node* send1 = SendFromHost({}, {}, + Node* send1 = SendFromHost("host_compute_channel_F1_O1", {}, b2.opts() .WithName("outside_compilation_F1_O1_send") .WithControlInput(e)); @@ -1215,5 +1536,110 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputsOrOutputs) { TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); } +// Test for shape inference of outside compilation. +TEST(EncapsulateSubgraphsTest, OutsideCompilationShapeInference) { + FunctionDefLibrary library; + GraphDef graphdef; + + { + *library.add_function() = test::function::XTimesTwo(); + + GraphDefBuilder b1(GraphDefBuilder::kFailImmediately); + Node* a = InputShaped(b1.opts().WithName("A")); + Node* b = Input(b1.opts().WithName("B")); + // Give nodes 'c' and 'd' names that collide after lowercasing. + Node* c = Unary(a, b1.opts().WithName("C")); + Node* d = Unary(b, b1.opts().WithName("c").WithControlInput(c).WithAttr( + "_encapsulate", "F1")); + Node* e = BinaryUnknownShape(c, d, + b1.opts() + .WithName("E") + .WithControlInputs({b, d}) + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* f = Binary(c, e, + b1.opts().WithName("F").WithControlInput(e).WithAttr( + "_encapsulate", "F1")); + Binary(a, f, b1.opts().WithName("G").WithControlInput(e)); + TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); + } + + TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + + FunctionDefLibrary library_expected; + GraphDef graphdef_expected; + + string shape_string_expected; + { + GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); + Node* known = KnownShape({2}, shape.opts().WithName("KnownShape/_0")); + Node* recv = + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT}, + shape.opts().WithName("outside_compilation_F1_O1_recv")); + Node* e = BinaryUnknownShape(known, recv, shape.opts().WithName("E")); + SendFromHost("host_compute_channel_F1_O1", {e}, + shape.opts().WithName("outside_compilation_F1_O1_send")); + GraphDef shape_graph; + TF_EXPECT_OK(shape.ToGraphDef(&shape_graph)); + EXPECT_TRUE(shape_graph.SerializeToString(&shape_string_expected)); + } + + *library_expected.add_function() = test::function::XTimesTwo(); + *library_expected.add_function() = FunctionDefHelper::Create( + "F1", {"b_0_arg:float", "c_0_arg:float"}, {"f_0_retval:float"}, {}, + { + {{"c"}, "UnaryTest", {"b_0_arg"}, {}, {}}, + {{"F"}, + "BinaryTest", + {"c_0_arg", "outside_compilation_O1_host_compute:outputs:0"}, + {}, + {"outside_compilation_O1_host_compute"}}, + {{"outside_compilation_O1_host_compute"}, + "_XlaHostCompute", + {"c:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", shape_string_expected}, + {"shapes", gtl::ArraySlice({})}}, + {"c"}}, + }, + {{"f_0_retval", "F:o:0"}}); + + { + std::unique_ptr lib_def( + new FunctionLibraryDefinition(OpRegistry::Global(), library_expected)); + GraphDefBuilder b2(GraphDefBuilder::kFailImmediately, lib_def.get()); + Node* a = InputShaped(b2.opts().WithName("A")); + Node* b = Input(b2.opts().WithName("B")); + Node* c = Unary(a, b2.opts().WithName("C")); + + NodeBuilder node_builder("F1", "F1", lib_def.get()); + node_builder.Input(b).Input(c); + Node* call = + b2.opts().WithControlInputs({c}).FinalizeBuilder(&node_builder); + + Node* recv = + RecvAtHost("host_compute_channel_F1_O1", {DT_FLOAT}, + b2.opts().WithName("outside_compilation_F1_O1_recv")); + Node* e = BinaryUnknownShape( + c, ops::NodeOut(recv, 0), + b2.opts().WithName("E").WithControlInputs({recv, b})); + Node* send = SendFromHost("host_compute_channel_F1_O1", {e}, + b2.opts() + .WithName("outside_compilation_F1_O1_send") + .WithControlInput(e)); + + Node* s = NoOp( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send})); + + Binary(a, call, b2.opts().WithName("G").WithControlInputs({s, e})); + TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); + } + + TF_EXPECT_GRAPH_EQ(graphdef_expected, graphdef); + TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.cc b/tensorflow/compiler/jit/kernels/xla_launch_op.cc index 4f3f17df9c680c63546d17dcc5a2775a1014f6c3..6353149e4afdf739fe44dd5c76502ef5d98b8477 100644 --- a/tensorflow/compiler/jit/kernels/xla_launch_op.cc +++ b/tensorflow/compiler/jit/kernels/xla_launch_op.cc @@ -45,7 +45,7 @@ namespace tensorflow { // see comment on `AllowsAsynchronousDeallocation()`. class XlaAllocator : public xla::DeviceMemoryAllocator { public: - XlaAllocator(gpu::Platform* platform, OpKernelContext* op_context); + XlaAllocator(const gpu::Platform* platform, OpKernelContext* op_context); ~XlaAllocator() override; xla::StatusOr Allocate(int device_ordinal, uint64 size, bool retry_on_failure) override; @@ -79,7 +79,8 @@ class XlaAllocator : public xla::DeviceMemoryAllocator { std::unordered_map tensors_; }; -XlaAllocator::XlaAllocator(gpu::Platform* platform, OpKernelContext* op_context) +XlaAllocator::XlaAllocator(const gpu::Platform* platform, + OpKernelContext* op_context) : xla::DeviceMemoryAllocator(platform), op_context_(op_context) {} XlaAllocator::~XlaAllocator() = default; @@ -248,23 +249,26 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { xla::LocalClient* client = static_cast(cache->client()); + // Builds an XLA allocator for the device. + XlaAllocator xla_allocator(client->platform(), ctx); + XlaCompiler::Options options; options.client = client; options.device_type = &cache->device_type(); options.flib_def = ctx->function_library()->GetFunctionLibraryDefinition(); options.graph_def_version = ctx->function_library()->graph_def_version(); options.allow_cpu_custom_calls = (platform_id_ == gpu::host::kHostPlatformId); + options.device_allocator = &xla_allocator; const XlaCompiler::CompilationResult* kernel; xla::LocalExecutable* executable; + OP_REQUIRES_OK(ctx, cache->Compile(options, function_, num_constant_args_, - variables, ctx, &kernel, &executable)); + variables, ctx, &kernel, &executable, + /*compile_options=*/nullptr)); VLOG(1) << "Executing XLA Computation..."; - // Builds an XLA allocator for the device. - XlaAllocator xla_allocator(client->platform(), ctx); - std::unique_ptr output; // Build xla::ShapedBuffers that point directly to the Tensor buffers. std::vector> arg_buffers; @@ -372,8 +376,6 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { OP_REQUIRES(ctx, write.input_index >= 0 && write.input_index < ctx->num_inputs(), errors::Internal("Invalid input index for variable write.")); - TensorShape write_shape; - OP_REQUIRES_OK(ctx, XLAShapeToTensorShape(write.shape, &write_shape)); gpu::DeviceMemoryBase buffer = output->buffer({output_num}); @@ -395,7 +397,7 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { // Looks up the owning Tensor by buffer address. OP_REQUIRES_OK( - ctx, xla_allocator.MakeTensorFromBuffer(buffer, write.type, write_shape, + ctx, xla_allocator.MakeTensorFromBuffer(buffer, write.type, write.shape, variable->tensor())); ++output_num; } diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 79b02baba83cb47f4f2f16544ad711a4b6937d90..a0211acbbe9eec77d30c7d14293650de8826f41c 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -190,6 +190,9 @@ Status FindCompilationCandidates( pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice); for (Node* node : graph.op_nodes()) { + VLOG(2) << "FindCompilationCandidates(): Processing " + << node->DebugString(); + DeviceType device_type(""); TF_RETURN_IF_ERROR( DeviceTypeOfDevice(node->assigned_device_name(), &device_type)); @@ -216,6 +219,13 @@ Status FindCompilationCandidates( !IsCompilableWhile(*node, jit_device_type, 0, lib_runtime)) { continue; } + // _Arg nodes in a top-level function represent feeds. + // Do not compile them. + if (node->type_string() == "_Arg") { + VLOG(2) << "Skipping jit compilation for '_Arg'-typed node " + << node->DebugString(); + continue; + } // _Retval nodes in a top-level function represent fetches. // Do not compile them. if (node->type_string() == "_Retval") { @@ -304,6 +314,7 @@ Status MarkForCompilationPass::Run( static_cast(flags->tf_xla_auto_jit); } bool cpu_global_jit = flags->tf_xla_cpu_global_jit; + VLOG(1) << "flags->tf_xla_cpu_global_jit = " << flags->tf_xla_cpu_global_jit; const FunctionLibraryDefinition* fld = options.flib_def; auto is_compilable = [global_jit_level, cpu_global_jit, fld]( diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc index 454f0aeae98d7afd51f12b2cfb1810de275a57f7..1a8858cccef623185709ab5dc2187a313dd130f7 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/core/framework/op.h" #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/graph/graph_def_builder.h" +#include "tensorflow/core/graph/graph_def_builder_util.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" @@ -80,7 +81,7 @@ TEST(XlaCompilationTest, Chains) { ops::UnaryOp("UncompilableUnary", c, builder.opts().WithName("D")); Node* e = ops::UnaryOp("Relu", d, builder.opts().WithName("E")); ops::UnaryOp("Relu", e, builder.opts().WithName("F")); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -105,7 +106,7 @@ TEST(XlaCompilationTest, UncompilableCycles) { Node* b = ops::UnaryOp("UncompilableUnary", a, builder.opts().WithName("B")); ops::BinaryOp("MatMul", a, b, builder.opts().WithName("C")); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -125,7 +126,7 @@ TEST(XlaCompilationTest, CompilableCycles) { .WithAttr("value", Tensor())); Node* b = ops::UnaryOp("Relu", a, builder.opts().WithName("B")); ops::BinaryOp("MatMul", a, b, builder.opts().WithName("C")); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -148,7 +149,7 @@ TEST(XlaCompilationTest, UnsupportedTypes) { .WithAttr("value", Tensor(DT_COMPLEX128, TensorShape()))); Node* b = ops::UnaryOp("Neg", a, builder.opts().WithName("B")); ops::BinaryOp("MatMul", a, b, builder.opts().WithName("C")); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -177,7 +178,7 @@ TEST(XlaCompilationTest, ConcatWithConstArg) { concat_builder.Input(dim).Input({a, a}).Attr("N", 2); builder.opts().FinalizeBuilder(&concat_builder); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -212,7 +213,7 @@ TEST(XlaCompilationTest, FunctionCalls) { Node* c = ops::UnaryOp("Relu", b, builder.opts().WithName("C")); ops::UnaryOp("UncompilableFn", c, builder.opts().WithName("D")); ops::BinaryOp("NoInlineFn", c, c, builder.opts().WithName("E")); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph, &flib_def)); @@ -244,7 +245,7 @@ TEST(XlaCompilationTest, MetadataOpsDontStartClusters) { Node* c = ops::UnaryOp("Rank", b, builder.opts().WithName("C")); Node* d = ops::UnaryOp("Size", c, builder.opts().WithName("D")); ops::UnaryOp("Shape", d, builder.opts().WithName("E")); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); auto clusters = GetClusters(*graph); @@ -330,7 +331,7 @@ TEST(XlaCompilationTest, SymbolicGradients) { d_builder.Input({c, c}); builder.opts().FinalizeBuilder(&d_builder); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -382,7 +383,7 @@ TEST(XlaCompilationTest, CyclesWithAllDifferentScopes) { ops::BinaryOp( "MatMul", a, b, builder.opts().WithName("C").WithAttr(kXlaScopeAttr, "ScopeC")); - TF_CHECK_OK(builder.ToGraph(graph.get())); + TF_CHECK_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -413,7 +414,7 @@ TEST(XlaCompilationTest, CyclesWithSplittingScopes) { ops::BinaryOp( "Add", b, c, builder.opts().WithName("D").WithAttr(kXlaScopeAttr, "Scope2")); - TF_CHECK_OK(builder.ToGraph(graph.get())); + TF_CHECK_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -443,7 +444,7 @@ TEST(XlaCompilationTest, CyclesWithDifferentScopesAndBridge) { "Relu", a, builder.opts().WithName("B").WithAttr(kXlaScopeAttr, "ScopeB")); ops::BinaryOp("MatMul", a, b, builder.opts().WithName("C")); - TF_CHECK_OK(builder.ToGraph(graph.get())); + TF_CHECK_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); @@ -484,7 +485,7 @@ TEST(XlaCompilationTest, Resources) { Node* c = ops::UnaryOp("ResourceOutput", b, builder.opts().WithName("C")); Node* d = ops::UnaryOp("ResourceInput", c, builder.opts().WithName("D")); ops::UnaryOp("Relu", d, builder.opts().WithName("E")); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); auto clusters = GetClusters(*graph); @@ -541,7 +542,7 @@ TEST(XlaCompilationTest, Retval) { .WithAttr("T", DT_FLOAT) .WithAttr("index", 0)); - TF_EXPECT_OK(builder.ToGraph(graph.get())); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } TF_ASSERT_OK(MarkForCompilation(&graph)); diff --git a/tensorflow/compiler/jit/xla_compilation_cache.cc b/tensorflow/compiler/jit/xla_compilation_cache.cc index 3717c2cc24283e0b218f92ec820d16893cbe0c35..6d854a920eb0b4c01b09024ceaef5035e847d392 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.cc +++ b/tensorflow/compiler/jit/xla_compilation_cache.cc @@ -148,8 +148,7 @@ Status BuildArguments(int num_constant_args, XlaCompiler::Argument& arg = (*args)[input_num]; arg.kind = XlaCompiler::Argument::kConstant; arg.type = input.dtype(); - TF_RETURN_IF_ERROR( - TensorShapeToXLAShape(input.dtype(), input.shape(), &arg.shape)); + arg.shape = input.shape(); arg.constant_value = input; ++input_num; } @@ -170,8 +169,7 @@ Status BuildArguments(int num_constant_args, arg.constant_value = input; } arg.type = input.dtype(); - TF_RETURN_IF_ERROR( - TensorShapeToXLAShape(input.dtype(), input.shape(), &arg.shape)); + arg.shape = input.shape(); ++input_num; } @@ -189,8 +187,7 @@ Status BuildArguments(int num_constant_args, if (variable_args[variable_id].present) { const Tensor& value = variable_args[variable_id].value; arg.type = value.dtype(); - TF_RETURN_IF_ERROR( - TensorShapeToXLAShape(value.dtype(), value.shape(), &arg.shape)); + arg.shape = value.shape(); arg.initialized = true; } else { // The values of uninitialized variables are not passed as inputs, since @@ -199,7 +196,7 @@ Status BuildArguments(int num_constant_args, // uninitialized variables. arg.initialized = false; arg.type = DT_INVALID; - arg.shape = xla::Shape(); + arg.shape = TensorShape(); } ++input_num; } @@ -223,6 +220,7 @@ Status XlaCompilationCache::BuildExecutable( xla::ExecutableBuildOptions build_options; build_options.set_device_ordinal(client_->default_device_ordinal()); build_options.set_result_layout(result.xla_output_shape); + build_options.set_device_allocator(options.device_allocator); auto compile_result = client_->Compile(*result.computation, argument_layouts, build_options); @@ -238,7 +236,8 @@ Status XlaCompilationCache::Compile( int num_constant_args, const std::vector& variable_args, OpKernelContext* ctx, const XlaCompiler::CompilationResult** compilation_result, - xla::LocalExecutable** executable) { + xla::LocalExecutable** executable, + const XlaCompiler::CompileOptions* compile_options) { VLOG(1) << "XlaCompilationCache::Compile " << DebugString(); if (VLOG_IS_ON(2)) { @@ -297,9 +296,9 @@ Status XlaCompilationCache::Compile( XlaCompiler compiler(options); entry->compiled = true; - entry->compilation_status = - compiler.CompileFunction(XlaCompiler::CompileOptions(), function, args, - &entry->compilation_result); + entry->compilation_status = compiler.CompileFunction( + compile_options ? *compile_options : XlaCompiler::CompileOptions(), + function, args, &entry->compilation_result); } *compilation_result = &entry->compilation_result; if (entry->compilation_status.ok() && executable) { diff --git a/tensorflow/compiler/jit/xla_compilation_cache.h b/tensorflow/compiler/jit/xla_compilation_cache.h index c3a8f68a157a2d34d4a6716c9951b2b698aead79..0858020716fcf4763e42dc0699ad22cfda756942 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.h +++ b/tensorflow/compiler/jit/xla_compilation_cache.h @@ -66,7 +66,8 @@ class XlaCompilationCache : public ResourceBase { const std::vector& variable_args, OpKernelContext* ctx, const XlaCompiler::CompilationResult** compilation_result, - xla::LocalExecutable** executable); + xla::LocalExecutable** executable, + const XlaCompiler::CompileOptions* compile_options); xla::LocalClient* client() const { return client_; } const DeviceType& device_type() const { return device_type_; } diff --git a/tensorflow/compiler/plugin/BUILD b/tensorflow/compiler/plugin/BUILD index c1edf2448c54ffddd7b70dcdfb1609080ca81b65..da4bc44c7a75c9f8faf16c537a17a1f2d16d5d61 100644 --- a/tensorflow/compiler/plugin/BUILD +++ b/tensorflow/compiler/plugin/BUILD @@ -41,6 +41,15 @@ cc_library( ], ) +# This target is added purely for the purpose of ensuring that `:xla_device` is +# always publicly visible to external XLA backend/plugin developers. +cc_library( + name = "plugin_device", + deps = [ + "//tensorflow/compiler/jit:xla_device", + ], +) + #----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index cea843146a09c02cfcd65664d3737f23c624f91d..782bf82d4149968d5e5fbfb93bbd4ff1dcd75494 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -144,6 +144,21 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "matrix_triangular_solve_op_test", + size = "small", + srcs = ["matrix_triangular_solve_op_test.py"], + tags = ["optonly"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + "//tensorflow/python:training", + ], +) + tf_xla_py_test( name = "clustering_test", size = "small", @@ -240,6 +255,35 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "extract_image_patches_op_test", + size = "small", + srcs = ["extract_image_patches_op_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + +tf_xla_py_test( + name = "fft_test", + size = "medium", + srcs = ["fft_test.py"], + shard_count = 3, + tags = ["optonly"], + deps = [ + ":xla_test", + "//tensorflow/contrib/signal:signal_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:extra_py_tests_deps", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + "//tensorflow/python:spectral_ops", + ], +) + tf_xla_py_test( name = "slice_ops_test", size = "small", @@ -309,6 +353,19 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "matrix_band_part_test", + size = "medium", + srcs = ["matrix_band_part_test.py"], + tags = ["optonly"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + tf_xla_py_test( name = "momentum_test", size = "small", @@ -385,12 +442,6 @@ tf_xla_py_test( # TODO(b/31361304): enable RNG ops on GPU when parallelized. disabled_backends = [ "gpu", - "cpu", - ], - tags = [ - "manual", - "no_oss", - "notap", ], deps = [ ":xla_test", @@ -426,6 +477,19 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "reverse_sequence_op_test", + size = "medium", + srcs = ["reverse_sequence_op_test.py"], + tags = ["optonly"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + tf_xla_py_test( name = "rmsprop_test", size = "small", @@ -576,6 +640,7 @@ tf_xla_py_test( name = "variable_ops_test", size = "small", srcs = ["variable_ops_test.py"], + tags = ["optonly"], deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -602,6 +667,31 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "gather_nd_op_test", + size = "medium", + srcs = ["gather_nd_op_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + +tf_xla_py_test( + name = "scatter_nd_op_test", + size = "medium", + srcs = ["scatter_nd_op_test.py"], + tags = ["optonly"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + cuda_py_test( name = "xla_device_test", size = "small", @@ -726,6 +816,17 @@ tf_library( tfcompile_flags = ["--xla_cpu_multi_thread_eigen=false"], ) +tf_xla_py_test( + name = "fake_quant_ops_test", + size = "medium", + srcs = ["fake_quant_ops_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/tests/binary_ops_test.py b/tensorflow/compiler/tests/binary_ops_test.py index 65706b35d616eb4dce94f0a7056a1604a97ff4c1..30a6d3a74d64f90ad33062df6d1e16e3a575bd63 100644 --- a/tensorflow/compiler/tests/binary_ops_test.py +++ b/tensorflow/compiler/tests/binary_ops_test.py @@ -43,7 +43,7 @@ class BinaryOpsTest(XLATestCase): output = op(pa, pb) result = session.run(output, {pa: a, pb: b}) if equality_test is None: - equality_test = self.assertAllClose + equality_test = self.assertAllCloseAccordingToType equality_test(result, expected, rtol=1e-3) def _testSymmetricBinary(self, op, a, b, expected, equality_test=None): @@ -54,14 +54,20 @@ class BinaryOpsTest(XLATestCase): """Tests closeness of two lists of floats.""" self.assertEqual(len(result), len(expected)) for i in range(len(result)): - self.assertAllClose(result[i], expected[i], rtol) + self.assertAllCloseAccordingToType(result[i], expected[i], rtol) def testFloatOps(self): for dtype in self.float_types: + if dtype == dtypes.bfloat16.as_numpy_dtype: + a = -1.01 + b = 4.1 + else: + a = -1.001 + b = 4.01 self._testBinary( lambda x, y: math_ops.approximate_equal(x, y, tolerance=0.0001), - np.array([[[[-1, 2.00009999], [-3, 4.01]]]], dtype=dtype), - np.array([[[[-1.001, 2], [-3.00009, 4]]]], dtype=dtype), + np.array([[[[-1, 2.00009999], [-3, b]]]], dtype=dtype), + np.array([[[[a, 2], [-3.00009, 4]]]], dtype=dtype), expected=np.array([[[[False, True], [True, False]]]], dtype=dtype)) self._testBinary( @@ -768,15 +774,15 @@ class BinaryOpsTest(XLATestCase): def DISABLED_testSparseMatMul(self): # Binary wrappers for sparse_matmul with different hints def SparseMatmulWrapperTF(a, b): - return tf.sparse_matmul(a, b, a_is_sparse=True) + return math_ops.sparse_matmul(a, b, a_is_sparse=True) def SparseMatmulWrapperFT(a, b): - return tf.sparse_matmul(a, b, b_is_sparse=True) + return math_ops.sparse_matmul(a, b, b_is_sparse=True) def SparseMatmulWrapperTT(a, b): - return tf.sparse_matmul(a, b, a_is_sparse=True, b_is_sparse=True) + return math_ops.sparse_matmul(a, b, a_is_sparse=True, b_is_sparse=True) - self._testMatMul(tf.sparse_matmul) + self._testMatMul(math_ops.sparse_matmul) self._testMatMul(SparseMatmulWrapperTF) self._testMatMul(SparseMatmulWrapperFT) self._testMatMul(SparseMatmulWrapperTT) @@ -1175,6 +1181,50 @@ class BinaryOpsTest(XLATestCase): np.array([4, 5, 6], dtype=np.int32), expected=None) + def testMatrixSetDiag(self): + for dtype in self.numeric_types: + # Square + self._testBinary( + array_ops.matrix_set_diag, + np.array([[0.0, 1.0, 0.0], [1.0, 0.0, 1.0], [1.0, 1.0, 1.0]], + dtype=dtype), + np.array([1.0, 2.0, 3.0], dtype=dtype), + expected=np.array([[1.0, 1.0, 0.0], [1.0, 2.0, 1.0], [1.0, 1.0, 3.0]], + dtype=dtype)) + + self._testBinary( + array_ops.matrix_set_diag, + np.array([[[1.0, 0.0, 3.0], [0.0, 2.0, 0.0], [1.0, 0.0, 3.0]], + [[4.0, 0.0, 4.0], [0.0, 5.0, 0.0], [2.0, 0.0, 6.0]]], + dtype=dtype), + np.array([[-1.0, 0.0, -3.0], [-4.0, -5.0, -6.0]], dtype=dtype), + expected=np.array( + [[[-1.0, 0.0, 3.0], [0.0, 0.0, 0.0], [1.0, 0.0, -3.0]], + [[-4.0, 0.0, 4.0], [0.0, -5.0, 0.0], [2.0, 0.0, -6.0]]], + dtype=dtype)) + + # Rectangular + self._testBinary( + array_ops.matrix_set_diag, + np.array([[0.0, 1.0, 0.0], [1.0, 0.0, 1.0]], dtype=dtype), + np.array([3.0, 4.0], dtype=dtype), + expected=np.array([[3.0, 1.0, 0.0], [1.0, 4.0, 1.0]], dtype=dtype)) + + self._testBinary( + array_ops.matrix_set_diag, + np.array([[0.0, 1.0], [1.0, 0.0], [1.0, 1.0]], dtype=dtype), + np.array([3.0, 4.0], dtype=dtype), + expected=np.array([[3.0, 1.0], [1.0, 4.0], [1.0, 1.0]], dtype=dtype)) + + self._testBinary( + array_ops.matrix_set_diag, + np.array([[[1.0, 0.0, 3.0], [0.0, 2.0, 0.0]], + [[4.0, 0.0, 4.0], [0.0, 5.0, 0.0]]], dtype=dtype), + np.array([[-1.0, -2.0], [-4.0, -5.0]], + dtype=dtype), + expected=np.array([[[-1.0, 0.0, 3.0], [0.0, -2.0, 0.0]], + [[-4.0, 0.0, 4.0], [0.0, -5.0, 0.0]]], + dtype=dtype)) if __name__ == "__main__": googletest.main() diff --git a/tensorflow/compiler/tests/conv2d_test.py b/tensorflow/compiler/tests/conv2d_test.py index 0d617eb37c5d92c87abb0f996b731112257a2b80..62577b70ce96e220d79978f01614b2d9a3647680 100644 --- a/tensorflow/compiler/tests/conv2d_test.py +++ b/tensorflow/compiler/tests/conv2d_test.py @@ -34,7 +34,13 @@ from tensorflow.python.platform import googletest class Conv2DTest(XLATestCase): - def _VerifyValues(self, input_sizes, filter_sizes, stride, padding, expected): + def _VerifyValues(self, + input_sizes=None, + filter_sizes=None, + strides=None, + dilations=None, + padding=None, + expected=None): """Tests that tf.nn.conv2d produces the expected value. Args: @@ -42,7 +48,8 @@ class Conv2DTest(XLATestCase): [batch, input_rows, input_cols, input_depth]. filter_sizes: Filter tensor dimensions in [kernel_rows, kernel_cols, input_depth, output_depth]. - stride: Stride. + strides: Strides. + dilations: RHS dilations. padding: Padding type. expected: Expected output. """ @@ -50,73 +57,136 @@ class Conv2DTest(XLATestCase): total_size_2 = np.prod(filter_sizes) x1 = np.arange(1, total_size_1 + 1, dtype=np.float32).reshape(input_sizes) x2 = np.arange(1, total_size_2 + 1, dtype=np.float32).reshape(filter_sizes) - strides = [1, stride, stride, 1] + strides = [1] + strides + [1] + if dilations is None: + dilations = [1, 1] + dilations = [1] + dilations + [1] with self.test_session() as sess: + t1 = array_ops.placeholder(dtypes.float32, shape=input_sizes) + t2 = array_ops.placeholder(dtypes.float32, shape=filter_sizes) with self.test_scope(): - t1 = array_ops.placeholder(dtypes.float32, shape=input_sizes) - t2 = array_ops.placeholder(dtypes.float32, shape=filter_sizes) out = nn_ops.conv2d( - t1, t2, strides=strides, padding=padding, data_format="NHWC") + t1, + t2, + strides=strides, + padding=padding, + data_format="NHWC", + dilations=dilations) value = sess.run(out, {t1: x1, t2: x2}) - self.assertArrayNear(expected, np.ravel(value), 1e-3) + self.assertAllClose(expected, value, 1e-3) def testConv2D1x1Filter(self): - expected_output = [ + expected_output = np.reshape([ 30.0, 36.0, 42.0, 66.0, 81.0, 96.0, 102.0, 126.0, 150.0, 138.0, 171.0, 204.0, 174.0, 216.0, 258.0, 210.0, 261.0, 312.0 - ] + ], [1, 2, 3, 3]) self._VerifyValues( input_sizes=[1, 2, 3, 3], filter_sizes=[1, 1, 3, 3], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) def testConv2D2x2Filter(self): - expected_output = [2271.0, 2367.0, 2463.0, 2901.0, 3033.0, 3165.0] + expected_output = np.reshape( + [2271.0, 2367.0, 2463.0, 2901.0, 3033.0, 3165.0], [1, 1, 2, 3]) self._VerifyValues( input_sizes=[1, 2, 3, 3], filter_sizes=[2, 2, 3, 3], - stride=1, + strides=[1, 1], + padding="VALID", + expected=expected_output) + + def testConv2D2x2Filter2x1Dilation(self): + expected_output = np.array([[[[72], [82], [92]], [[112], [122], [132]]]]) + self._VerifyValues( + input_sizes=[1, 4, 4, 1], + filter_sizes=[2, 2, 1, 1], + strides=[1, 1], + dilations=[2, 1], padding="VALID", expected=expected_output) def testConv2D1x2Filter(self): - expected_output = [ + expected_output = np.reshape([ 231.0, 252.0, 273.0, 384.0, 423.0, 462.0, 690.0, 765.0, 840.0, 843.0, 936.0, 1029.0 - ] + ], [1, 2, 2, 3]) self._VerifyValues( input_sizes=[1, 2, 3, 3], filter_sizes=[1, 2, 3, 3], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) def testConv2D2x2FilterStride2(self): - expected_output = [2271.0, 2367.0, 2463.0] + expected_output = np.reshape([2271.0, 2367.0, 2463.0], [1, 1, 1, 3]) self._VerifyValues( input_sizes=[1, 2, 3, 3], filter_sizes=[2, 2, 3, 3], - stride=2, + strides=[2, 2], padding="VALID", expected=expected_output) def testConv2D2x2FilterStride2Same(self): - expected_output = [2271.0, 2367.0, 2463.0, 1230.0, 1305.0, 1380.0] + expected_output = np.reshape( + [2271.0, 2367.0, 2463.0, 1230.0, 1305.0, 1380.0], [1, 1, 2, 3]) self._VerifyValues( input_sizes=[1, 2, 3, 3], filter_sizes=[2, 2, 3, 3], - stride=2, + strides=[2, 2], padding="SAME", expected=expected_output) + def testConv2DEmptyDilation(self): + self._VerifyValues( + input_sizes=[0, 2, 3, 3], + filter_sizes=[1, 1, 3, 3], + strides=[1, 1], + dilations=[2, 1], + padding="VALID", + expected=np.zeros([0, 2, 3, 3])) + + def testConv2D2x2FilterDilation(self): + self._VerifyValues( + input_sizes=[1, 2, 3, 3], + filter_sizes=[2, 2, 3, 3], + strides=[1, 1], + dilations=[1, 2], + padding="VALID", + expected=np.reshape([2667, 2781, 2895], [1, 1, 1, 3])) + + def testConv2D1x2FilterDilation(self): + self._VerifyValues( + input_sizes=[1, 2, 3, 3], + filter_sizes=[1, 2, 3, 3], + strides=[1, 1], + dilations=[2, 1], + padding="VALID", + expected=np.array([[[[231, 252, 273], [384, 423, 462]], + [[690, 765, 840], [843, 936, 1029]]]])) + + def testConv2DKernelSizeMatchesInputSizeDilation(self): + self._VerifyValues( + input_sizes=[1, 3, 3, 1], + filter_sizes=[2, 2, 1, 2], + strides=[1, 1], + dilations=[2, 2], + padding="VALID", + expected=np.reshape([108, 128], [1, 1, 1, 2])) + class Conv2DBackpropInputTest(XLATestCase): - def _VerifyValues(self, input_sizes, filter_sizes, out_backprop_sizes, stride, - padding, expected): + def _VerifyValues(self, + input_sizes=None, + filter_sizes=None, + out_backprop_sizes=None, + strides=None, + dilations=None, + padding=None, + expected=None): """Tests that gen_nn_ops.conv2d_backprop_input produces the expected output. Args: @@ -125,7 +195,8 @@ class Conv2DBackpropInputTest(XLATestCase): filter_sizes: Filter tensor dimensions in [kernel_rows, kernel_cols, input_depth, output_depth]. out_backprop_sizes: Output gradients tensor dimensions. - stride: Stride. + strides: Strides. + dilations: Dilations. padding: Padding type. expected: Expected output. """ @@ -134,21 +205,25 @@ class Conv2DBackpropInputTest(XLATestCase): x1 = np.arange(1, total_size_1 + 1, dtype=np.float32).reshape(filter_sizes) x2 = np.arange( 1, total_size_2 + 1, dtype=np.float32).reshape(out_backprop_sizes) - strides = [1, stride, stride, 1] + strides = [1] + strides + [1] + if dilations is not None: + dilations = [1] + dilations + [1] with self.test_session() as sess: + t1 = array_ops.placeholder(dtypes.float32, shape=filter_sizes) + t2 = array_ops.placeholder(dtypes.float32, shape=out_backprop_sizes) with self.test_scope(): - t1 = array_ops.placeholder(dtypes.float32, shape=filter_sizes) - t2 = array_ops.placeholder(dtypes.float32, shape=out_backprop_sizes) out = gen_nn_ops.conv2d_backprop_input( input_sizes=input_sizes, filter=t1, out_backprop=t2, strides=strides, + dilations=dilations, padding=padding, data_format="NHWC") value = sess.run(out, {t1: x1, t2: x2}) - self.assertArrayNear(expected, np.ravel(value), 1e-3) + self.assertAllEqual(input_sizes, value.shape) + self.assertAllClose(expected, np.ravel(value), 1e-3) def testConv2D1x1Filter(self): expected_output = [ @@ -160,7 +235,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 4, 4, 3], filter_sizes=[1, 1, 3, 2], out_backprop_sizes=[1, 4, 4, 2], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) @@ -170,7 +245,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 1, 5, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=3, + strides=[3, 3], padding="VALID", expected=expected_output) @@ -180,7 +255,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 1, 6, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=3, + strides=[3, 3], padding="VALID", expected=expected_output) @@ -190,7 +265,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 1, 7, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=3, + strides=[3, 3], padding="VALID", expected=expected_output) @@ -200,7 +275,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 2, 3, 1], filter_sizes=[2, 2, 1, 1], out_backprop_sizes=[1, 2, 3, 1], - stride=1, + strides=[1, 1], padding="SAME", expected=expected_output) @@ -213,7 +288,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 2, 3, 3], filter_sizes=[2, 2, 3, 3], out_backprop_sizes=[1, 1, 2, 3], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) @@ -226,7 +301,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 2, 3, 3], filter_sizes=[2, 2, 3, 3], out_backprop_sizes=[1, 2, 3, 3], - stride=1, + strides=[1, 1], padding="SAME", expected=expected_output) @@ -236,7 +311,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 3, 3, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 3, 2, 1], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) @@ -246,7 +321,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 3, 3, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 3, 3, 1], - stride=1, + strides=[1, 1], padding="SAME", expected=expected_output) @@ -256,7 +331,7 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 3, 5, 1], filter_sizes=[1, 3, 1, 1], out_backprop_sizes=[1, 2, 2, 1], - stride=2, + strides=[2, 2], padding="VALID", expected=expected_output) @@ -266,15 +341,76 @@ class Conv2DBackpropInputTest(XLATestCase): input_sizes=[1, 2, 3, 1], filter_sizes=[2, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=2, + strides=[2, 2], padding="SAME", expected=expected_output) + def testConv2D2x2Depth3ValidBackpropInputStride1x1Dilation2x1(self): + self._VerifyValues( + input_sizes=[1, 3, 6, 1], + filter_sizes=[2, 2, 1, 1], + out_backprop_sizes=[1, 1, 5, 1], + strides=[1, 1], + dilations=[2, 1], + padding="VALID", + expected=[1, 4, 7, 10, 13, 10, 0, 0, 0, 0, 0, 0, 3, 10, 17, 24, 31, 20]) + + def testConv2D2x2Depth1ValidBackpropInputDilation1x2(self): + self._VerifyValues( + input_sizes=[1, 2, 3, 1], + filter_sizes=[2, 2, 1, 1], + out_backprop_sizes=[1, 1, 1, 1], + strides=[1, 1], + dilations=[1, 2], + padding="VALID", + expected=[1, 0, 2, 3, 0, 4]) + + def testConv2DEmptyBackpropInputDilation1x2(self): + self._VerifyValues( + input_sizes=[0, 2, 3, 1], + filter_sizes=[2, 2, 1, 1], + out_backprop_sizes=[0, 1, 1, 1], + strides=[1, 1], + dilations=[1, 2], + padding="VALID", + expected=np.zeros([0])) + + def testConv2D2x2Depth3ValidBackpropInputDilation2x1(self): + # The GPU version of this test is not very stable. So adjusting the + # error threshold to 1e-4. + self._VerifyValues( + input_sizes=[1, 3, 2, 3], + filter_sizes=[2, 2, 3, 3], + out_backprop_sizes=[1, 1, 1, 3], + strides=[1, 1], + dilations=[2, 1], + padding="VALID", + expected=[ + 14, 32, 50, 68, 86, 104, 0, 0, 0, 0, 0, 0, 122, 140, 158, 176, 194, + 212 + ]) + + def testConv2DKernelSizeMatchesInputSizeBackpropInputDilation2x2(self): + self._VerifyValues( + input_sizes=[1, 3, 3, 1], + filter_sizes=[2, 2, 1, 2], + out_backprop_sizes=[1, 1, 1, 2], + strides=[1, 1], + dilations=[2, 2], + padding="VALID", + expected=[5, 0, 11, 0, 0, 0, 17, 0, 23]) + class Conv2DBackpropFilterTest(XLATestCase): - def _VerifyValues(self, input_sizes, filter_sizes, out_backprop_sizes, stride, - padding, expected): + def _VerifyValues(self, + input_sizes=None, + filter_sizes=None, + out_backprop_sizes=None, + strides=None, + dilations=None, + padding=None, + expected=None): """Tests that gen_nn_ops.conv2d_backprop_filter produces the right output. Args: @@ -283,7 +419,8 @@ class Conv2DBackpropFilterTest(XLATestCase): filter_sizes: Filter tensor dimensions in [kernel_rows, kernel_cols, input_depth, output_depth]. out_backprop_sizes: Output gradients tensor dimensions. - stride: Stride. + strides: Stride. + dilations: Dilations. padding: Padding type. expected: Expected output. """ @@ -293,22 +430,26 @@ class Conv2DBackpropFilterTest(XLATestCase): x1 = np.arange(1, total_size_1 + 1, dtype=np.float32).reshape(input_sizes) x2 = np.arange( 1, total_size_2 + 1, dtype=np.float32).reshape(out_backprop_sizes) - strides = [1, stride, stride, 1] + strides = [1] + strides + [1] + if dilations is not None: + dilations = [1] + dilations + [1] with self.test_session() as sess: + t1 = array_ops.placeholder(dtypes.float32, shape=input_sizes) + t2 = array_ops.placeholder(dtypes.float32, shape=out_backprop_sizes) with self.test_scope(): - t1 = array_ops.placeholder(dtypes.float32, shape=input_sizes) - t2 = array_ops.placeholder(dtypes.float32, shape=out_backprop_sizes) tensor = gen_nn_ops.conv2d_backprop_filter( input=t1, filter_sizes=filter_sizes, out_backprop=t2, strides=strides, + dilations=dilations, padding=padding, data_format="NHWC") value = sess.run(tensor, {t1: x1, t2: x2}) - self.assertArrayNear(expected, np.ravel(value), 1e-3) + self.assertAllEqual(filter_sizes, value.shape) + self.assertAllClose(expected, np.ravel(value), 1e-3) def testConv2D1x1Filter(self): expected_output = [8056, 8432, 8312, 8704, 8568, 8976] @@ -316,7 +457,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 4, 4, 3], filter_sizes=[1, 1, 3, 2], out_backprop_sizes=[1, 4, 4, 2], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) @@ -326,7 +467,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 3, 3, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 3, 2, 1], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) @@ -336,7 +477,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 2, 3, 1], filter_sizes=[2, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) @@ -350,7 +491,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 2, 3, 3], filter_sizes=[2, 2, 3, 3], out_backprop_sizes=[1, 1, 2, 3], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) @@ -360,7 +501,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 1, 5, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=3, + strides=[3, 3], padding="VALID", expected=expected_output) @@ -370,7 +511,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 1, 6, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=3, + strides=[3, 3], padding="VALID", expected=expected_output) @@ -380,7 +521,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 1, 7, 1], filter_sizes=[1, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=3, + strides=[3, 3], padding="VALID", expected=expected_output) @@ -390,7 +531,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 1, 4, 1], filter_sizes=[1, 3, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=1, + strides=[1, 1], padding="VALID", expected=expected_output) @@ -400,7 +541,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 1, 4, 1], filter_sizes=[1, 3, 1, 1], out_backprop_sizes=[1, 1, 4, 1], - stride=1, + strides=[1, 1], padding="SAME", expected=expected_output) @@ -410,7 +551,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 1, 4, 1], filter_sizes=[1, 3, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=2, + strides=[2, 2], padding="SAME", expected=expected_output) @@ -420,7 +561,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 2, 3, 1], filter_sizes=[2, 2, 1, 1], out_backprop_sizes=[1, 2, 3, 1], - stride=1, + strides=[1, 1], padding="SAME", expected=expected_output) @@ -430,7 +571,7 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 3, 5, 1], filter_sizes=[1, 3, 1, 1], out_backprop_sizes=[1, 2, 2, 1], - stride=2, + strides=[2, 2], padding="VALID", expected=expected_output) @@ -440,10 +581,64 @@ class Conv2DBackpropFilterTest(XLATestCase): input_sizes=[1, 2, 3, 1], filter_sizes=[2, 2, 1, 1], out_backprop_sizes=[1, 1, 2, 1], - stride=2, + strides=[2, 2], padding="SAME", expected=expected_output) + def testConv2D2x2Depth3ValidBackpropFilterStride1x1Dilation2x1(self): + self._VerifyValues( + input_sizes=[1, 3, 6, 1], + filter_sizes=[2, 2, 1, 1], + out_backprop_sizes=[1, 1, 5, 1], + strides=[1, 1], + dilations=[2, 1], + padding="VALID", + expected=[55, 70, 235, 250]) + + def testConv2D2x2Depth1ValidBackpropFilterDilation1x2(self): + self._VerifyValues( + input_sizes=[1, 2, 3, 1], + filter_sizes=[2, 2, 1, 1], + out_backprop_sizes=[1, 1, 1, 1], + strides=[1, 1], + dilations=[1, 2], + padding="VALID", + expected=[1, 3, 4, 6]) + + def testConv2DEmptyBackpropFilterDilation1x2(self): + self._VerifyValues( + input_sizes=[1, 2, 3, 1], + filter_sizes=[2, 2, 1, 0], + out_backprop_sizes=[1, 1, 1, 0], + strides=[1, 1], + dilations=[1, 2], + padding="VALID", + expected=np.zeros([0])) + + def testConv2D2x2Depth3ValidBackpropFilterDilation2x2(self): + self._VerifyValues( + input_sizes=[1, 3, 4, 3], + filter_sizes=[2, 2, 3, 3], + out_backprop_sizes=[1, 1, 2, 3], + strides=[1, 1], + dilations=[2, 2], + padding="VALID", + expected=[ + 17, 22, 27, 22, 29, 36, 27, 36, 45, 47, 64, 81, 52, 71, 90, 57, 78, + 99, 137, 190, 243, 142, 197, 252, 147, 204, 261, 167, 232, 297, 172, + 239, 306, 177, 246, 315 + ]) + + def testConv2DKernelSizeMatchesInputSizeBackpropFilterDilation2x2(self): + self._VerifyValues( + input_sizes=[1, 3, 3, 1], + filter_sizes=[2, 2, 1, 2], + out_backprop_sizes=[1, 1, 1, 2], + strides=[1, 1], + dilations=[2, 2], + padding="VALID", + expected=[1, 2, 3, 6, 7, 14, 9, 18]) + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/compiler/tests/extract_image_patches_op_test.py b/tensorflow/compiler/tests/extract_image_patches_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..0361702e7af778176daed941d64e61198090daf2 --- /dev/null +++ b/tensorflow/compiler/tests/extract_image_patches_op_test.py @@ -0,0 +1,134 @@ +# 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. +# ============================================================================== +"""Functional tests for ExtractImagePatches op.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class ExtractImagePatches(XLATestCase): + """Functional tests for ExtractImagePatches op.""" + + def _VerifyValues(self, image, ksizes, strides, rates, padding, patches): + """Tests input-output pairs for the ExtractImagePatches op. + + Args: + image: Input tensor with shape: [batch, in_rows, in_cols, depth]. + ksizes: Patch size specified as: [ksize_rows, ksize_cols]. + strides: Output strides, specified as [stride_rows, stride_cols]. + rates: Atrous rates, specified as [rate_rows, rate_cols]. + padding: Padding type. + patches: Expected output. + """ + ksizes = [1] + ksizes + [1] + strides = [1] + strides + [1] + rates = [1] + rates + [1] + + with self.test_session(): + image_placeholder = array_ops.placeholder(dtypes.float32) + with self.test_scope(): + out_tensor = array_ops.extract_image_patches( + image_placeholder, + ksizes=ksizes, + strides=strides, + rates=rates, + padding=padding, + name="im2col") + feed_dict = {image_placeholder: image} + self.assertAllClose(patches, out_tensor.eval(feed_dict=feed_dict)) + + def testKsize1x1Stride1x1Rate1x1(self): + """Verifies that for 1x1 kernel the output equals the input.""" + # [2, 3, 4, 5] + image = np.reshape(range(120), [2, 3, 4, 5]) + # [2, 3, 4, 5] + patches = np.reshape(range(120), [2, 3, 4, 5]) + for padding in ["VALID", "SAME"]: + self._VerifyValues( + image, + ksizes=[1, 1], + strides=[1, 1], + rates=[1, 1], + padding=padding, + patches=patches) + + def testKsize1x1Stride2x3Rate1x1(self): + """Test for 1x1 kernel and strides.""" + # [2, 4, 5, 3] + image = np.reshape(range(120), [2, 4, 5, 3]) + # [2, 2, 2, 3] + patches = image[:, ::2, ::3, :] + for padding in ["VALID", "SAME"]: + self._VerifyValues( + image, + ksizes=[1, 1], + strides=[2, 3], + rates=[1, 1], + padding=padding, + patches=patches) + + def testKsize2x2Stride1x1Rate1x1Valid(self): + """Test for 2x2 kernel with VALID padding.""" + # [1, 2, 2, 1] + image = [[[[1], [2]], [[3], [4]]]] + # [1, 1, 1, 4] + patches = [[[[1, 2, 3, 4]]]] + self._VerifyValues( + image, + ksizes=[2, 2], + strides=[1, 1], + rates=[1, 1], + padding="VALID", + patches=patches) + + def testKsize2x2Stride1x1Rate1x1Same(self): + """Test for 2x2 kernel with SAME padding.""" + # [1, 2, 2, 1] + image = [[[[1], [2]], [[3], [4]]]] + # [1, 2, 2, 4] + patches = [[[[1, 2, 3, 4], [2, 0, 4, 0]], [[3, 4, 0, 0], [4, 0, 0, 0]]]] + self._VerifyValues( + image, + ksizes=[2, 2], + strides=[1, 1], + rates=[1, 1], + padding="SAME", + patches=patches) + + def testKsize2x2Stride1x1Rate2x2Valid(self): + """Test for 2x2 kernel with 2x2 dilation.""" + # [1, 2, 2, 1] + image = np.arange(16).reshape(1, 4, 4, 1).astype(np.float32) + # [1, 2, 2, 4] + patches = [[[[0, 2, 8, 10], [1, 3, 9, 11]], + [[4, 6, 12, 14], [5, 7, 13, 15]]]] + self._VerifyValues( + image, + ksizes=[2, 2], + strides=[1, 1], + rates=[2, 2], + padding="VALID", + patches=patches) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/fake_quant_ops_test.py b/tensorflow/compiler/tests/fake_quant_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..dfe9400ef0f55ca011d4e23ba5d735899ca2e054 --- /dev/null +++ b/tensorflow/compiler/tests/fake_quant_ops_test.py @@ -0,0 +1,452 @@ +# 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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_array_ops +from tensorflow.python.platform import googletest + + +class FakeQuantWithMinMaxArgsTest(XLATestCase): + """Test cases for FakeQuantWithMinMaxArgs operation.""" + + # 8 bits, wide range. + def testOp_with8BitsNoScalingNoNudging(self): + self._TestOp(0.0, 255.0, 8, False, 0.0, 255.0, 1.0) + + def testOp_with8BitsScalingAndNudgingDown(self): + self._TestOp(0.5, 128.0, 8, False, 0.0, 127.5, 0.5) + + def testOp_with8BitsScalingAndNudgingUp(self): + self._TestOp(-128.0, -0.5, 8, False, -127.5, 0.0, 0.5) + + def testOp_with8BitsScalingAndNudgingBetween(self): + self._TestOp(-0.1, 127.4, 8, False, 0.0, 127.5, 0.5) + + # 8 bits, narrow range. + def testOp_with8BitsNarrowRangeNoScalingNoNudging(self): + self._TestOp(0.0, 254.0, 8, True, 0.0, 254.0, 1.0) + + def testOp_with8BitsNarrowRangeScalingAndNudgingDown(self): + self._TestOp(0.1, 127.1, 8, True, 0.0, 127.0, 0.5) + + def testOp_with8BitsNarrowRangeScalingAndNudgingUp(self): + self._TestOp(-127.1, -0.1, 8, True, -127.0, 0.0, 0.5) + + def testOp_with8BitsNarrowRangeScalingAndNudgingBetween(self): + self._TestOp(-0.1, 126.9, 8, True, 0.0, 127.0, 0.5) + + # 7 bits, wide range. + def testOp_with7BitsNoScalingNoNudging(self): + self._TestOp(0.0, 127.0, 7, False, 0.0, 127.0, 1.0) + + def testOp_with7BitsScalingAndNudgingDown(self): + self._TestOp(0.5, 64.0, 7, False, 0.0, 63.5, 0.5) + + def testOp_with7BitsScalingAndNudgingUp(self): + self._TestOp(-64.0, -0.5, 7, False, -63.5, 0.0, 0.5) + + def testOp_with7BitsScalingAndNudgingBetween(self): + self._TestOp(-0.1, 63.4, 7, False, 0.0, 63.5, 0.5) + + # 7 bits, narrow range. + def testOp_with7BitsNarrowRangeNoScalingNoNudging(self): + self._TestOp(0.0, 126.0, 7, True, 0.0, 126.0, 1.0) + + def testOp_with7BitsNarrowRangeScalingAndNudgingDown(self): + self._TestOp(0.1, 63.1, 7, True, 0.0, 63.0, 0.5) + + def testOp_with7BitsNarrowRangeScalingAndNudgingUp(self): + self._TestOp(-63.1, -0.1, 7, True, -63.0, 0.0, 0.5) + + def testOp_with7BitsNarrowRangeScalingAndNudgingBetween(self): + self._TestOp(-0.1, 62.9, 7, True, 0.0, 63.0, 0.5) + + def _TestOp(self, input_min, input_max, num_bits, narrow_range, + expected_nudged_input_min, expected_nudged_input_max, + expected_step): + inputs = np.array( + [ + expected_nudged_input_min - expected_step, + expected_nudged_input_min - 0.01, expected_nudged_input_min, + expected_nudged_input_min + 0.01, + expected_nudged_input_min + expected_step - 0.01, + expected_nudged_input_min + expected_step, + expected_nudged_input_min + expected_step + 0.01, + expected_nudged_input_max - 0.01, expected_nudged_input_max, + expected_nudged_input_max + 0.01, + expected_nudged_input_max + expected_step + ], + dtype=np.float32) + expected = np.array( + [ + expected_nudged_input_min, expected_nudged_input_min, + expected_nudged_input_min, expected_nudged_input_min, + expected_nudged_input_min + expected_step, + expected_nudged_input_min + expected_step, + expected_nudged_input_min + expected_step, + expected_nudged_input_max, expected_nudged_input_max, + expected_nudged_input_max, expected_nudged_input_max + ], + dtype=np.float32) + + with self.test_session() as session: + with self.test_scope(): + input_placeholder = array_ops.placeholder( + dtypes.float32, inputs.shape, name="inputs") + outputs = array_ops.fake_quant_with_min_max_args( + input_placeholder, + min=input_min, + max=input_max, + num_bits=num_bits, + narrow_range=narrow_range) + result = session.run(outputs, {input_placeholder: inputs}) + self.assertAllCloseAccordingToType( + result, expected, rtol=1e-3, atol=1e-5, bfloat16_rtol=0.03) + + +class FakeQuantWithMinMaxArgsGradientTest(XLATestCase): + """Test cases for FakeQuantWithMinMaxArgsGradient operation.""" + + # 8 bits, wide range. + def testOp_with8BitsNoScalingNoNudging(self): + self._TestOp(0.0, 255.0, 8, False, 0.0, 255.0, 1.0) + + def testOp_with8BitsScalingAndNudgingDown(self): + self._TestOp(0.5, 128.0, 8, False, 0.0, 127.5, 0.5) + + def testOp_with8BitsScalingAndNudgingUp(self): + self._TestOp(-128.0, -0.5, 8, False, -127.5, 0.0, 0.5) + + def testOp_with8BitsScalingAndNudgingBetween(self): + self._TestOp(-0.1, 127.4, 8, False, 0.0, 127.5, 0.5) + + # 8 bits, narrow range. + def testOp_with8BitsNarrowRangeNoScalingNoNudging(self): + self._TestOp(0.0, 254.0, 8, True, 0.0, 254.0, 1.0) + + def testOp_with8BitsNarrowRangeScalingAndNudgingDown(self): + self._TestOp(0.1, 127.1, 8, True, 0.0, 127.0, 0.5) + + def testOp_with8BitsNarrowRangeScalingAndNudgingUp(self): + self._TestOp(-127.1, -0.1, 8, True, -127.0, 0.0, 0.5) + + def testOp_with8BitsNarrowRangeScalingAndNudgingBetween(self): + self._TestOp(-0.1, 126.9, 8, True, 0.0, 127.0, 0.5) + + # 7 bits, wide range. + def testOp_with7BitsNoScalingNoNudging(self): + self._TestOp(0.0, 127.0, 7, False, 0.0, 127.0, 1.0) + + def testOp_with7BitsScalingAndNudgingDown(self): + self._TestOp(0.5, 64.0, 7, False, 0.0, 63.5, 0.5) + + def testOp_with7BitsScalingAndNudgingUp(self): + self._TestOp(-64.0, -0.5, 7, False, -63.5, 0.0, 0.5) + + def testOp_with7BitsScalingAndNudgingBetween(self): + self._TestOp(-0.1, 63.4, 7, False, 0.0, 63.5, 0.5) + + # 7 bits, narrow range. + def testOp_with7BitsNarrowRangeNoScalingNoNudging(self): + self._TestOp(0.0, 126.0, 7, True, 0.0, 126.0, 1.0) + + def testOp_with7BitsNarrowRangeScalingAndNudgingDown(self): + self._TestOp(0.1, 63.1, 7, True, 0.0, 63.0, 0.5) + + def testOp_with7BitsNarrowRangeScalingAndNudgingUp(self): + self._TestOp(-63.1, -0.1, 7, True, -63.0, 0.0, 0.5) + + def testOp_with7BitsNarrowRangeScalingAndNudgingBetween(self): + self._TestOp(-0.1, 62.9, 7, True, 0.0, 63.0, 0.5) + + def _TestOp(self, input_min, input_max, num_bits, narrow_range, + expected_nudged_input_min, expected_nudged_input_max, + expected_step): + inputs = np.array( + [ + expected_nudged_input_min - expected_step, + expected_nudged_input_min - 0.01, expected_nudged_input_min, + expected_nudged_input_min + 0.01, + expected_nudged_input_min + expected_step - 0.01, + expected_nudged_input_min + expected_step, + expected_nudged_input_min + expected_step + 0.01, + expected_nudged_input_max - 0.01, expected_nudged_input_max, + expected_nudged_input_max + 0.01, + expected_nudged_input_max + expected_step + ], + dtype=np.float32) + gradients = np.arange(1, len(inputs) + 1, dtype=np.float32) + expected_backprops = np.array( + [0.0, 0.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 0.0, 0.0], + dtype=np.float32) + + with self.test_session() as session: + with self.test_scope(): + gradient_placeholder = array_ops.placeholder( + dtypes.float32, gradients.shape, name="gradients") + input_placeholder = array_ops.placeholder( + dtypes.float32, inputs.shape, name="inputs") + outputs = gen_array_ops.fake_quant_with_min_max_args_gradient( + gradient_placeholder, + input_placeholder, + min=input_min, + max=input_max, + num_bits=num_bits, + narrow_range=narrow_range) + backprops = session.run(outputs, { + gradient_placeholder: gradients, + input_placeholder: inputs + }) + self.assertAllCloseAccordingToType( + backprops, + expected_backprops, + rtol=1e-3, + atol=1e-5, + bfloat16_rtol=0.03) + + +class FakeQuantWithMinMaxVarsTest(XLATestCase): + """Test cases for FakeQuantWithMinMaxVars operation.""" + + # 8 bits, wide range. + def testOp_with8BitsNoScalingNoNudging(self): + self._TestOp(0.0, 255.0, 8, False, 0.0, 255.0, 1.0) + + def testOp_with8BitsScalingAndNudgingDown(self): + self._TestOp(0.5, 128.0, 8, False, 0.0, 127.5, 0.5) + + def testOp_with8BitsScalingAndNudgingUp(self): + self._TestOp(-128.0, -0.5, 8, False, -127.5, 0.0, 0.5) + + def testOp_with8BitsScalingAndNudgingBetween(self): + self._TestOp(-0.1, 127.4, 8, False, 0.0, 127.5, 0.5) + + # 8 bits, narrow range. + def testOp_with8BitsNarrowRangeNoScalingNoNudging(self): + self._TestOp(0.0, 254.0, 8, True, 0.0, 254.0, 1.0) + + def testOp_with8BitsNarrowRangeScalingAndNudgingDown(self): + self._TestOp(0.1, 127.1, 8, True, 0.0, 127.0, 0.5) + + def testOp_with8BitsNarrowRangeScalingAndNudgingUp(self): + self._TestOp(-127.1, -0.1, 8, True, -127.0, 0.0, 0.5) + + def testOp_with8BitsNarrowRangeScalingAndNudgingBetween(self): + self._TestOp(-0.1, 126.9, 8, True, 0.0, 127.0, 0.5) + + # 7 bits, wide range. + def testOp_with7BitsNoScalingNoNudging(self): + self._TestOp(0.0, 127.0, 7, False, 0.0, 127.0, 1.0) + + def testOp_with7BitsScalingAndNudgingDown(self): + self._TestOp(0.5, 64.0, 7, False, 0.0, 63.5, 0.5) + + def testOp_with7BitsScalingAndNudgingUp(self): + self._TestOp(-64.0, -0.5, 7, False, -63.5, 0.0, 0.5) + + def testOp_with7BitsScalingAndNudgingBetween(self): + self._TestOp(-0.1, 63.4, 7, False, 0.0, 63.5, 0.5) + + # 7 bits, narrow range. + def testOp_with7BitsNarrowRangeNoScalingNoNudging(self): + self._TestOp(0.0, 126.0, 7, True, 0.0, 126.0, 1.0) + + def testOp_with7BitsNarrowRangeScalingAndNudgingDown(self): + self._TestOp(0.1, 63.1, 7, True, 0.0, 63.0, 0.5) + + def testOp_with7BitsNarrowRangeScalingAndNudgingUp(self): + self._TestOp(-63.1, -0.1, 7, True, -63.0, 0.0, 0.5) + + def testOp_with7BitsNarrowRangeScalingAndNudgingBetween(self): + self._TestOp(-0.1, 62.9, 7, True, 0.0, 63.0, 0.5) + + def _TestOp(self, input_min, input_max, num_bits, narrow_range, + expected_nudged_input_min, expected_nudged_input_max, + expected_step): + inputs = np.array( + [ + expected_nudged_input_min - expected_step, + expected_nudged_input_min - 0.01, expected_nudged_input_min, + expected_nudged_input_min + 0.01, + expected_nudged_input_min + expected_step - 0.01, + expected_nudged_input_min + expected_step, + expected_nudged_input_min + expected_step + 0.01, + expected_nudged_input_max - 0.01, expected_nudged_input_max, + expected_nudged_input_max + 0.01, + expected_nudged_input_max + expected_step + ], + dtype=np.float32) + expected = np.array( + [ + expected_nudged_input_min, expected_nudged_input_min, + expected_nudged_input_min, expected_nudged_input_min, + expected_nudged_input_min + expected_step, + expected_nudged_input_min + expected_step, + expected_nudged_input_min + expected_step, + expected_nudged_input_max, expected_nudged_input_max, + expected_nudged_input_max, expected_nudged_input_max + ], + dtype=np.float32) + + with self.test_session() as session: + with self.test_scope(): + input_placeholder = array_ops.placeholder( + dtypes.float32, inputs.shape, name="inputs") + min_placeholder = array_ops.placeholder(dtypes.float32, (), name="min") + max_placeholder = array_ops.placeholder(dtypes.float32, (), name="max") + outputs = array_ops.fake_quant_with_min_max_vars( + input_placeholder, + min_placeholder, + max_placeholder, + num_bits=num_bits, + narrow_range=narrow_range) + result = session.run( + outputs, { + input_placeholder: inputs, + min_placeholder: input_min, + max_placeholder: input_max + }) + self.assertAllCloseAccordingToType( + result, expected, rtol=1e-3, atol=1e-5, bfloat16_rtol=0.03) + + +class FakeQuantWithMinMaxVarsGradientTest(XLATestCase): + """Test cases for FakeQuantWithMinMaxVarsGradient operation.""" + + # 8 bits, wide range. + def testOp_with8BitsNoScalingNoNudging(self): + self._TestOp(0.0, 255.0, 8, False, 0.0, 255.0, 1.0) + + def testOp_with8BitsScalingAndNudgingDown(self): + self._TestOp(0.5, 128.0, 8, False, 0.0, 127.5, 0.5) + + def testOp_with8BitsScalingAndNudgingUp(self): + self._TestOp(-128.0, -0.5, 8, False, -127.5, 0.0, 0.5) + + def testOp_with8BitsScalingAndNudgingBetween(self): + self._TestOp(-0.1, 127.4, 8, False, 0.0, 127.5, 0.5) + + # 8 bits, narrow range. + def testOp_with8BitsNarrowRangeNoScalingNoNudging(self): + self._TestOp(0.0, 254.0, 8, True, 0.0, 254.0, 1.0) + + def testOp_with8BitsNarrowRangeScalingAndNudgingDown(self): + self._TestOp(0.1, 127.1, 8, True, 0.0, 127.0, 0.5) + + def testOp_with8BitsNarrowRangeScalingAndNudgingUp(self): + self._TestOp(-127.1, -0.1, 8, True, -127.0, 0.0, 0.5) + + def testOp_with8BitsNarrowRangeScalingAndNudgingBetween(self): + self._TestOp(-0.1, 126.9, 8, True, 0.0, 127.0, 0.5) + + # 7 bits, wide range. + def testOp_with7BitsNoScalingNoNudging(self): + self._TestOp(0.0, 127.0, 7, False, 0.0, 127.0, 1.0) + + def testOp_with7BitsScalingAndNudgingDown(self): + self._TestOp(0.5, 64.0, 7, False, 0.0, 63.5, 0.5) + + def testOp_with7BitsScalingAndNudgingUp(self): + self._TestOp(-64.0, -0.5, 7, False, -63.5, 0.0, 0.5) + + def testOp_with7BitsScalingAndNudgingBetween(self): + self._TestOp(-0.1, 63.4, 7, False, 0.0, 63.5, 0.5) + + # 7 bits, narrow range. + def testOp_with7BitsNarrowRangeNoScalingNoNudging(self): + self._TestOp(0.0, 126.0, 7, True, 0.0, 126.0, 1.0) + + def testOp_with7BitsNarrowRangeScalingAndNudgingDown(self): + self._TestOp(0.1, 63.1, 7, True, 0.0, 63.0, 0.5) + + def testOp_with7BitsNarrowRangeScalingAndNudgingUp(self): + self._TestOp(-63.1, -0.1, 7, True, -63.0, 0.0, 0.5) + + def testOp_with7BitsNarrowRangeScalingAndNudgingBetween(self): + self._TestOp(-0.1, 62.9, 7, True, 0.0, 63.0, 0.5) + + def _TestOp(self, input_min, input_max, num_bits, narrow_range, + expected_nudged_input_min, expected_nudged_input_max, + expected_step): + inputs = np.array( + [ + expected_nudged_input_min - expected_step, + expected_nudged_input_min - 0.01, expected_nudged_input_min, + expected_nudged_input_min + 0.01, + expected_nudged_input_min + expected_step - 0.01, + expected_nudged_input_min + expected_step, + expected_nudged_input_min + expected_step + 0.01, + expected_nudged_input_max - 0.01, expected_nudged_input_max, + expected_nudged_input_max + 0.01, + expected_nudged_input_max + expected_step + ], + dtype=np.float32) + gradients = np.arange(1, len(inputs) + 1, dtype=np.float32) + expected_backprops_wrt_input = np.array( + [0.0, 0.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 0.0, 0.0], + dtype=np.float32) + expected_backprops_wrt_min = 1.0 + 2.0 + expected_backprops_wrt_max = 10.0 + 11.0 + + with self.test_session() as session: + with self.test_scope(): + gradient_placeholder = array_ops.placeholder( + dtypes.float32, gradients.shape, name="gradients") + input_placeholder = array_ops.placeholder( + dtypes.float32, inputs.shape, name="inputs") + min_placeholder = array_ops.placeholder(dtypes.float32, (), name="min") + max_placeholder = array_ops.placeholder(dtypes.float32, (), name="max") + outputs = array_ops.fake_quant_with_min_max_vars_gradient( + gradient_placeholder, + input_placeholder, + min_placeholder, + max_placeholder, + num_bits=num_bits, + narrow_range=narrow_range) + backprops_wrt_input, backprops_wrt_min, backprops_wrt_max = session.run( + outputs, { + gradient_placeholder: gradients, + input_placeholder: inputs, + min_placeholder: input_min, + max_placeholder: input_max + }) + self.assertAllCloseAccordingToType( + backprops_wrt_input, + expected_backprops_wrt_input, + rtol=1e-3, + atol=1e-5, + bfloat16_rtol=0.03) + self.assertAllCloseAccordingToType( + backprops_wrt_min, + expected_backprops_wrt_min, + rtol=1e-3, + atol=1e-5, + bfloat16_rtol=0.03) + self.assertAllCloseAccordingToType( + backprops_wrt_max, + expected_backprops_wrt_max, + rtol=1e-3, + atol=1e-5, + bfloat16_rtol=0.03) + + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/compiler/tests/fft_test.py b/tensorflow/compiler/tests/fft_test.py new file mode 100644 index 0000000000000000000000000000000000000000..afb5fa4bb4fefe5bc2ecded826143ffc83c2b559 --- /dev/null +++ b/tensorflow/compiler/tests/fft_test.py @@ -0,0 +1,204 @@ +# Copyright 2017 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 FFT via the XLA JIT.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import itertools + +import numpy as np +import scipy.signal as sps + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.contrib.signal.python.ops import spectral_ops as signal +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import spectral_ops +from tensorflow.python.platform import googletest + +BATCH_DIMS = (3, 5) +RTOL = 0.02 # Eigen/cuFFT differ widely from np, especially for FFT3D +ATOL = 1e-3 + + +def pick_10(x): + x = list(x) + np.random.seed(123) + np.random.shuffle(x) + return x[:10] + + +def to_32bit(x): + if x.dtype == np.complex128: + return x.astype(np.complex64) + if x.dtype == np.float64: + return x.astype(np.float32) + return x + + +POWS_OF_2 = 2**np.arange(3, 12) +INNER_DIMS_1D = list((x,) for x in POWS_OF_2) +POWS_OF_2 = 2**np.arange(3, 8) # To avoid OOM on GPU. +INNER_DIMS_2D = pick_10(itertools.product(POWS_OF_2, POWS_OF_2)) +INNER_DIMS_3D = pick_10(itertools.product(POWS_OF_2, POWS_OF_2, POWS_OF_2)) + + +class FFTTest(XLATestCase): + + def _VerifyFftMethod(self, inner_dims, complex_to_input, input_to_expected, + tf_method): + for indims in inner_dims: + print("nfft =", indims) + shape = BATCH_DIMS + indims + data = np.arange(np.prod(shape) * 2) / np.prod(indims) + np.random.seed(123) + np.random.shuffle(data) + data = np.reshape(data.astype(np.float32).view(np.complex64), shape) + data = to_32bit(complex_to_input(data)) + expected = to_32bit(input_to_expected(data)) + with self.test_session() as sess: + with self.test_scope(): + ph = array_ops.placeholder( + dtypes.as_dtype(data.dtype), shape=data.shape) + out = tf_method(ph) + value = sess.run(out, {ph: data}) + self.assertAllClose(expected, value, rtol=RTOL, atol=ATOL) + + def testContribSignalSTFT(self): + ws = 512 + hs = 128 + dims = (ws * 20,) + shape = BATCH_DIMS + dims + data = np.arange(np.prod(shape)) / np.prod(dims) + np.random.seed(123) + np.random.shuffle(data) + data = np.reshape(data.astype(np.float32), shape) + window = sps.get_window("hann", ws) + expected = sps.stft( + data, nperseg=ws, noverlap=ws - hs, boundary=None, window=window)[2] + expected = np.swapaxes(expected, -1, -2) + expected *= window.sum() # scipy divides by window sum + with self.test_session() as sess: + with self.test_scope(): + ph = array_ops.placeholder( + dtypes.as_dtype(data.dtype), shape=data.shape) + out = signal.stft(ph, ws, hs) + + value = sess.run(out, {ph: data}) + self.assertAllClose(expected, value, rtol=RTOL, atol=ATOL) + + def testFFT(self): + self._VerifyFftMethod(INNER_DIMS_1D, lambda x: x, np.fft.fft, + spectral_ops.fft) + + def testFFT2D(self): + self._VerifyFftMethod(INNER_DIMS_2D, lambda x: x, np.fft.fft2, + spectral_ops.fft2d) + + def testFFT3D(self): + self._VerifyFftMethod(INNER_DIMS_3D, lambda x: x, + lambda x: np.fft.fftn(x, axes=(-3, -2, -1)), + spectral_ops.fft3d) + + def testIFFT(self): + self._VerifyFftMethod(INNER_DIMS_1D, lambda x: x, np.fft.ifft, + spectral_ops.ifft) + + def testIFFT2D(self): + self._VerifyFftMethod(INNER_DIMS_2D, lambda x: x, np.fft.ifft2, + spectral_ops.ifft2d) + + def testIFFT3D(self): + self._VerifyFftMethod(INNER_DIMS_3D, lambda x: x, + lambda x: np.fft.ifftn(x, axes=(-3, -2, -1)), + spectral_ops.ifft3d) + + def testRFFT(self): + self._VerifyFftMethod( + INNER_DIMS_1D, np.real, lambda x: np.fft.rfft(x, n=x.shape[-1]), + lambda x: spectral_ops.rfft(x, fft_length=[x.shape[-1].value])) + + def testRFFT2D(self): + + def _tf_fn(x): + return spectral_ops.rfft2d( + x, fft_length=[x.shape[-2].value, x.shape[-1].value]) + + self._VerifyFftMethod( + INNER_DIMS_2D, np.real, + lambda x: np.fft.rfft2(x, s=[x.shape[-2], x.shape[-1]]), _tf_fn) + + def testRFFT3D(self): + + def _to_expected(x): + return np.fft.rfftn( + x, axes=(-3, -2, -1), s=[x.shape[-3], x.shape[-2], x.shape[-1]]) + + def _tf_fn(x): + return spectral_ops.rfft3d( + x, + fft_length=[x.shape[-3].value, x.shape[-2].value, x.shape[-1].value]) + + self._VerifyFftMethod(INNER_DIMS_3D, np.real, _to_expected, _tf_fn) + + def testIRFFT(self): + + def _tf_fn(x): + return spectral_ops.irfft(x, fft_length=[2 * (x.shape[-1].value - 1)]) + + self._VerifyFftMethod( + INNER_DIMS_1D, lambda x: np.fft.rfft(np.real(x), n=x.shape[-1]), + lambda x: np.fft.irfft(x, n=2 * (x.shape[-1] - 1)), _tf_fn) + + def testIRFFT2D(self): + + def _tf_fn(x): + return spectral_ops.irfft2d( + x, fft_length=[x.shape[-2].value, 2 * (x.shape[-1].value - 1)]) + + self._VerifyFftMethod( + INNER_DIMS_2D, + lambda x: np.fft.rfft2(np.real(x), s=[x.shape[-2], x.shape[-1]]), + lambda x: np.fft.irfft2(x, s=[x.shape[-2], 2 * (x.shape[-1] - 1)]), + _tf_fn) + + def testIRFFT3D(self): + + def _to_input(x): + return np.fft.rfftn( + np.real(x), + axes=(-3, -2, -1), + s=[x.shape[-3], x.shape[-2], x.shape[-1]]) + + def _to_expected(x): + return np.fft.irfftn( + x, + axes=(-3, -2, -1), + s=[x.shape[-3], x.shape[-2], 2 * (x.shape[-1] - 1)]) + + def _tf_fn(x): + return spectral_ops.irfft3d( + x, + fft_length=[ + x.shape[-3].value, x.shape[-2].value, 2 * (x.shape[-1].value - 1) + ]) + + self._VerifyFftMethod(INNER_DIMS_3D, _to_input, _to_expected, _tf_fn) + + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/compiler/tests/gather_nd_op_test.py b/tensorflow/compiler/tests/gather_nd_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9378b1db7245c0da3e8298e7dcd972491616b0cd --- /dev/null +++ b/tensorflow/compiler/tests/gather_nd_op_test.py @@ -0,0 +1,147 @@ +# 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 tensorflow.ops.tf.gather_nd.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import errors +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class GatherNdTest(XLATestCase): + + def _runGather(self, params, indices): + with self.test_session(): + paramsp = array_ops.placeholder(params.dtype) + indicesp = array_ops.placeholder(indices.dtype) + with self.test_scope(): + gather_nd_t = array_ops.gather_nd(paramsp, indicesp) + feed_dict = {paramsp: params, indicesp: indices} + return gather_nd_t.eval(feed_dict=feed_dict) + + def testSimpleDtype(self): + for dtype in self.numeric_types: + self.assertAllEqual( + np.array([7, 7, 8], dtype=dtype), + self._runGather( + np.array([8, 1, 2, 3, 7, 5], dtype=dtype), + np.array([[4], [4], [0]], np.int32))) + + def testEmptyIndicesAndParamsOKButJustEmptyParamsFails(self): + with self.test_session(): + params = np.ones((3, 3), dtype=np.float32) + + indices_empty = np.empty((0, 2), dtype=np.int32) + gather_nd_ok_val = self._runGather(params, indices_empty) + self.assertAllClose(np.empty((0,), dtype=np.float32), gather_nd_ok_val) + + indices_empty = np.empty((0, 1), dtype=np.int32) + gather_nd_ok_val = self._runGather(params, indices_empty) + self.assertAllClose(np.empty((0, 3), dtype=np.float32), gather_nd_ok_val) + + params_empty = np.empty((0, 3), dtype=np.float32) + indices_empty = np.empty((0, 2), dtype=np.int32) + gather_nd_ok_val = self._runGather(params_empty, indices_empty) + self.assertAllClose(np.empty((0,), dtype=np.float32), gather_nd_ok_val) + + params_empty = np.empty((0, 3), dtype=np.float32) + indices_nonempty = np.zeros((1, 2), dtype=np.int32) + with self.assertRaisesWithPredicateMatch( + errors.InvalidArgumentError, r"Gather dimension 0 is of size zero"): + self._runGather(params_empty, indices_nonempty) + + def testIndexScalar(self): + params = np.array( + [[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], dtype=np.float32).T + indices = np.array([4, 1], dtype=np.int32) + gather_nd_val = self._runGather(params, indices) + self.assertAllEqual(np.array(7), gather_nd_val) + + def testParamsRankLargerThanIndexIndexScalarSlices(self): + params = np.array( + [[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], dtype=np.float32).T + indices = np.array( + [ + 4, + ], dtype=np.int32) + gather_nd_val = self._runGather(params, indices) + self.assertAllEqual(np.array([-7, 7]), gather_nd_val) + + def testParamsRankLargerThanIndexSlices(self): + params = np.array( + [[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], dtype=np.float32).T + indices = np.array([[4], [4], [0]], np.int32) + gather_nd_val = self._runGather(params, indices) + self.assertAllEqual(np.array([[-7, 7], [-7, 7], [-8, 8]]), gather_nd_val) + + def testHigherRankParamsLargerThanIndexSlices(self): + params = np.array( + [[[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], + [[-80, -10, -20, -30, -70, -50], [80, 10, 20, 30, 70, 50]]], + dtype=np.float32).T + indices = np.array([[4], [4], [0]], np.int32) + gather_nd_val = self._runGather(params, indices) + self.assertAllEqual(params[[4, 4, 0]], gather_nd_val) + + def testEmptyIndicesLastRankMeansCopyEntireTensor(self): + params = np.array( + [[[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], + [[-80, -10, -20, -30, -70, -50], [80, 10, 20, 30, 70, 50]]], + dtype=np.float32).T + indices = np.array([[], []], dtype=np.int32) # Size (2, 0) + gather_nd_val = self._runGather(params, indices) + self.assertAllEqual( + np.vstack((params[np.newaxis, :], params[np.newaxis, :])), + gather_nd_val) + + def testHigherRankParamsAndIndicesLargerThanIndexSlices(self): + params = np.array( + [[[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], + [[-80, -10, -20, -30, -70, -50], [80, 10, 20, 30, 70, 50]]], + dtype=np.float32).T + indices = np.array([[[3], [2], [1]], [[4], [4], [0]]], np.int32) + gather_nd_val = self._runGather(params, indices) + self.assertAllEqual(params[[3, 2, 1, 4, 4, 0]].reshape(2, 3, 2, 2), + gather_nd_val) + + def testHigherRankParams(self): + shape = (10, 20, 5, 1, 17) + params = np.random.rand(*shape).astype(np.float32) + indices = np.vstack( + [np.random.randint(0, s, size=2000, dtype=np.int32) for s in shape]).T + gather_nd_val = self._runGather(params, indices) + + expected = params[tuple(indices.T)] + self.assertAllEqual(expected, gather_nd_val) + + def testHigherRankParamsAndIndices(self): + shape = (10, 20, 5, 1, 17) + params = np.random.rand(*shape).astype(np.float32) + indices = np.vstack( + [np.random.randint(0, s, size=2000, dtype=np.int32) for s in shape]).T + indices_reshaped = indices.reshape([10, 10, 20, 5]) + gather_nd_val = self._runGather(params, indices_reshaped) + expected = params[tuple(indices.T)] + self.assertAllEqual(expected.reshape([10, 10, 20]), gather_nd_val) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/gather_test.py b/tensorflow/compiler/tests/gather_test.py index 13cbe6f312f5175edaec28fa7a8f28064194b0e9..1a8c4519118f69ce51ca9a5eb95a9d706c7766cc 100644 --- a/tensorflow/compiler/tests/gather_test.py +++ b/tensorflow/compiler/tests/gather_test.py @@ -122,6 +122,20 @@ class GatherTest(xla_test.XLATestCase): gather_np = np.take(params, indices, axis=axis) self.assertAllEqual(gather_np, gather_value) + def testIndicesWithDifferentDimensions(self): + with self.test_session(): + for dtype in self.numeric_tf_types: + params = array_ops.placeholder(dtype=dtype) + indices = array_ops.placeholder(dtype=np.int32) + with self.test_scope(): + gather = array_ops.gather(params, indices) + self.assertAllEqual( + 7, gather.eval(feed_dict={params: [4, 7, 2], indices: 1})) + self.assertAllEqual( + [7], gather.eval(feed_dict={params: [4, 7, 2], indices: [1]})) + self.assertAllEqual( + [[7]], gather.eval(feed_dict={params: [4, 7, 2], indices: [[1]]})) + class GatherBenchmark(test.Benchmark): """Microbenchmarks for the gather op.""" diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py index 9e095aa00b2944185a50113d0effc4ed134b03e4..538fa8e8e570b83ed681ecc0501285520cabdecb 100644 --- a/tensorflow/compiler/tests/image_ops_test.py +++ b/tensorflow/compiler/tests/image_ops_test.py @@ -65,7 +65,7 @@ class RGBToHSVTest(XLATestCase): # Verify that processing batch elements together is the same as separate self.assertAllClose(batch1, join1) self.assertAllClose(batch2, join2) - self.assertAllClose(batch2, inp) + self.assertAllCloseAccordingToType(batch2, inp, bfloat16_atol=0.03) def testRGBToHSVRoundTrip(self): data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] @@ -77,21 +77,25 @@ class RGBToHSVTest(XLATestCase): hsv = image_ops.rgb_to_hsv(placeholder) rgb = image_ops.hsv_to_rgb(hsv) rgb_tf = rgb.eval(feed_dict={placeholder: rgb_np}) - self.assertAllClose(rgb_tf, rgb_np) + self.assertAllCloseAccordingToType(rgb_tf, rgb_np, bfloat16_atol=0.03) def testRGBToHSVNumpy(self): """Tests the RGB to HSV conversion matches a reference implementation.""" for nptype in self.float_types: rgb_flat = np.random.random(64 * 3).reshape((64, 3)).astype(nptype) rgb_np = rgb_flat.reshape(4, 4, 4, 3) - hsv_np = np.array([colorsys.rgb_to_hsv(r, g, b) for r, g, b in rgb_flat]) + hsv_np = np.array([ + colorsys.rgb_to_hsv( + r.astype(np.float64), g.astype(np.float64), b.astype(np.float64)) + for r, g, b in rgb_flat + ]) hsv_np = hsv_np.reshape(4, 4, 4, 3) with self.test_session(): placeholder = array_ops.placeholder(nptype) with self.test_scope(): hsv_op = image_ops.rgb_to_hsv(placeholder) hsv_tf = hsv_op.eval(feed_dict={placeholder: rgb_np}) - self.assertAllClose(hsv_tf, hsv_np) + self.assertAllCloseAccordingToType(hsv_tf, hsv_np) class AdjustContrastTest(XLATestCase): @@ -427,7 +431,8 @@ class ResizeBilinearTest(XLATestCase): np.zeros([1, input_shape[0], input_shape[1], 1], dtype=dtype), align_corners=True) out = sess.run(resized, {grads: grads_np[np.newaxis, :, :, np.newaxis]}) - self.assertAllClose(expected[np.newaxis, :, :, np.newaxis], out) + self.assertAllCloseAccordingToType(expected[np.newaxis, :, :, np.newaxis], + out) def testAlignCorners1x2To3x2(self): for dtype in self.float_types: @@ -506,6 +511,42 @@ class ResizeBilinearTest(XLATestCase): [7, 4, 4, 9]], dtype=np.float32)) + def testAlignCorners3x3To9x9(self): + for dtype in self.float_types: + self._assertForwardOpMatchesExpected( + np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=dtype), [9, 9], + expected=np.array( + [[1.0, 1.25, 1.50, 1.75, 2.00, 2.25, 2.50, 2.75, 3.00], [ + 1.75, 2.00, 2.25, 2.50, 2.75, 3.00, 3.25, 3.50, 3.75 + ], [2.50, 2.75, 3.00, 3.25, 3.50, 3.75, 4.00, 4.25, 4.50], [ + 3.25, 3.50, 3.75, 4.00, 4.25, 4.50, 4.75, 5.00, 5.25 + ], [4.00, 4.25, 4.50, 4.75, 5.00, 5.25, 5.50, 5.75, 6.00], [ + 4.75, 5.00, 5.25, 5.50, 5.75, 6.00, 6.25, 6.50, 6.75 + ], [5.50, 5.75, 6.00, 6.25, 6.50, 6.75, 7.00, 7.25, 7.50], [ + 6.25, 6.50, 6.75, 7.00, 7.25, 7.50, 7.75, 8.00, 8.25 + ], [7.00, 7.25, 7.50, 7.75, 8.00, 8.25, 8.50, 8.75, 9.00]], + dtype=np.float32)) + + def testAlignCorners3x3To9x9Grad(self): + for dtype in self.float_types: + self._assertBackwardOpMatchesExpected( + np.array( + [[1.00, 1.25, 1.50, 1.75, 2.00, 2.25, 2.50, 2.75, 3.00], [ + 1.75, 2.00, 2.25, 2.50, 2.75, 3.00, 3.25, 3.50, 3.75 + ], [2.50, 2.75, 3.00, 3.25, 3.50, 3.75, 4.00, 4.25, 4.50], [ + 3.25, 3.50, 3.75, 4.00, 4.25, 4.50, 4.75, 5.00, 5.25 + ], [4.00, 4.25, 4.50, 4.75, 5.00, 5.25, 5.50, 5.75, 6.00], [ + 4.75, 5.00, 5.25, 5.50, 5.75, 6.00, 6.25, 6.50, 6.75 + ], [5.50, 5.75, 6.00, 6.25, 6.50, 6.75, 7.00, 7.25, 7.50], [ + 6.25, 6.50, 6.75, 7.00, 7.25, 7.50, 7.75, 8.00, 8.25 + ], [7.00, 7.25, 7.50, 7.75, 8.00, 8.25, 8.50, 8.75, 9.00]], + dtype=np.float32), + input_shape=[3, 3], + dtype=dtype, + expected=np.array( + [[12.5, 27.5, 21.875], [42.5, 80.0, 57.5], [40.625, 72.5, 50]], + dtype=np.float32)) + if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tests/matrix_band_part_test.py b/tensorflow/compiler/tests/matrix_band_part_test.py new file mode 100644 index 0000000000000000000000000000000000000000..29394f9ea5139b30f88f53de0469b27e37d79195 --- /dev/null +++ b/tensorflow/compiler/tests/matrix_band_part_test.py @@ -0,0 +1,64 @@ +# 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. +# ============================================================================== + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class MatrixBandPartTest(XLATestCase): + + def _testMatrixBandPart(self, dtype, shape): + with self.test_session(): + batch_shape = shape[:-2] + mat = np.ones(shape).astype(dtype) + batch_mat = np.tile(mat, batch_shape + [1, 1]) + for lower in -1, 0, 1, shape[-2] - 1: + for upper in -1, 0, 1, shape[-1] - 1: + band_np = mat + if lower >= 0: + band_np = np.triu(band_np, -lower) + if upper >= 0: + band_np = np.tril(band_np, upper) + if batch_shape: + band_np = np.tile(band_np, batch_shape + [1, 1]) + + placeholder = array_ops.placeholder(dtype) + with self.test_scope(): + band = array_ops.matrix_band_part( + placeholder, + constant_op.constant(lower, dtype=dtypes.int32), + constant_op.constant(upper, dtype=dtypes.int32)) + feed_dict = {placeholder: batch_mat} + self.assertAllEqual(band_np, band.eval(feed_dict=feed_dict)) + + def testMatrixBandPart(self): + for dtype in self.float_types: + for batch_shape in [[], [2,], [1, 3, 2]]: + for rows in 1, 2, 7: + for cols in 1, 2, 7: + self._testMatrixBandPart(dtype, batch_shape + [rows, cols]) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/matrix_triangular_solve_op_test.py b/tensorflow/compiler/tests/matrix_triangular_solve_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..cccb7f5789dce39ef8c3d4b3a7573aaa983b3fbd --- /dev/null +++ b/tensorflow/compiler/tests/matrix_triangular_solve_op_test.py @@ -0,0 +1,130 @@ +# 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 tensorflow.ops.tf.MatrixTriangularSolve.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import itertools + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import linalg_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test + + +def MakePlaceholder(x): + return array_ops.placeholder(dtypes.as_dtype(x.dtype), shape=x.shape) + + +class MatrixTriangularSolveOpTest(XLATestCase): + + def _VerifyTriangularSolveBase(self, sess, placeholder_a, placeholder_ca, + placeholder_b, a, clean_a, b, verification, + atol): + feed_dict = {placeholder_a: a, placeholder_ca: clean_a, placeholder_b: b} + verification_np = sess.run(verification, feed_dict) + self.assertAllClose(b, verification_np, atol=atol) + + def _VerifyTriangularSolve(self, a, b, lower, adjoint, atol): + clean_a = np.tril(a) if lower else np.triu(a) + with self.test_session() as sess: + placeholder_a = MakePlaceholder(a) + placeholder_ca = MakePlaceholder(clean_a) + placeholder_b = MakePlaceholder(b) + with self.test_scope(): + x = linalg_ops.matrix_triangular_solve( + placeholder_a, placeholder_b, lower=lower, adjoint=adjoint) + verification = math_ops.matmul(placeholder_ca, x, adjoint_a=adjoint) + self._VerifyTriangularSolveBase(sess, placeholder_a, placeholder_ca, + placeholder_b, a, clean_a, b, + verification, atol) + + def _VerifyTriangularSolveCombo(self, a, b, atol=1e-4): + transp = lambda x: np.swapaxes(x, -1, -2) + for lower, adjoint in itertools.product([True, False], repeat=2): + self._VerifyTriangularSolve( + a if lower else transp(a), b, lower, adjoint, atol) + + def testBasic(self): + rng = np.random.RandomState(0) + a = np.tril(rng.randn(5, 5)) + b = rng.randn(5, 7) + for dtype in self.float_types: + self._VerifyTriangularSolveCombo(a.astype(dtype), b.astype(dtype)) + + def testBasicNotActuallyTriangular(self): + rng = np.random.RandomState(0) + a = rng.randn(5, 5) # the `a` matrix is not lower-triangular + b = rng.randn(5, 7) + for dtype in self.float_types: + self._VerifyTriangularSolveCombo(a.astype(dtype), b.astype(dtype)) + + def testBasicComplexDtypes(self): + rng = np.random.RandomState(0) + a = np.tril(rng.randn(5, 5) + rng.randn(5, 5) * 1j) + b = rng.randn(5, 7) + rng.randn(5, 7) * 1j + for dtype in self.complex_types: + self._VerifyTriangularSolveCombo(a.astype(dtype), b.astype(dtype)) + + def testBatch(self): + rng = np.random.RandomState(0) + shapes = [((4, 3, 3), (4, 3, 5)), ((1, 2, 2), (1, 2, 1)), + ((1, 1, 1), (1, 1, 2)), ((2, 3, 4, 4), (2, 3, 4, 1))] + tuples = itertools.product(self.float_types, shapes) + for dtype, (a_shape, b_shape) in tuples: + n = a_shape[-1] + a = np.tril(rng.rand(*a_shape) - 0.5) / (2.0 * n) + np.eye(n) + b = rng.randn(*b_shape) + self._VerifyTriangularSolveCombo( + a.astype(dtype), b.astype(dtype), atol=1e-3) + + def testLarge(self): + n = 1024 + rng = np.random.RandomState(0) + a = np.tril(rng.rand(n, n) - 0.5) / (2.0 * n) + np.eye(n) + b = rng.randn(n, n) + self._VerifyTriangularSolve( + a.astype(np.float32), b.astype(np.float32), True, False, 1e-4) + + def testNonSquareCoefficientMatrix(self): + rng = np.random.RandomState(0) + for dtype in self.float_types: + a = rng.randn(3, 4).astype(dtype) + b = rng.randn(4, 4).astype(dtype) + with self.assertRaises(ValueError): + linalg_ops.matrix_triangular_solve(a, b) + with self.assertRaises(ValueError): + linalg_ops.matrix_triangular_solve(a, b) + + def testWrongDimensions(self): + randn = np.random.RandomState(0).randn + for dtype in self.float_types: + lhs = constant_op.constant(randn(3, 3), dtype=dtype) + rhs = constant_op.constant(randn(4, 3), dtype=dtype) + with self.assertRaises(ValueError): + linalg_ops.matrix_triangular_solve(lhs, rhs) + with self.assertRaises(ValueError): + linalg_ops.matrix_triangular_solve(lhs, rhs) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/randomized_tests.cc b/tensorflow/compiler/tests/randomized_tests.cc index 2ab013b7fa9498d00e070cff061ef990f5ed8d92..e72dd4eea9f127e1df96ab166103c4c16372adb6 100644 --- a/tensorflow/compiler/tests/randomized_tests.cc +++ b/tensorflow/compiler/tests/randomized_tests.cc @@ -93,11 +93,11 @@ class OpTestBuilder { public: explicit OpTestBuilder(const string& op_name); - // Adds an input 'tensor'. + // Adds an input 'tensor' as a Placeholder node. OpTestBuilder& Input(const Tensor& tensor); - // Adds a random input tensor with 'type'. If 'dims' is not provided, - // RandomDims() is used. + // Adds a random input tensor with 'type' as a Placeholder node. + // If 'dims' is not provided, RandomDims() is used. OpTestBuilder& RandomInput(DataType type); OpTestBuilder& RandomInput(DataType type, std::vector dims); @@ -1375,6 +1375,121 @@ TEST_F(OpTest, Conj) { }); } +TEST_F(OpTest, FFT) { + Repeatedly([this]() { + std::vector dims = RandomDims(1, kDefaultMaxRank); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("FFT").RandomInput(DT_COMPLEX64, dims)); + }); +} + +TEST_F(OpTest, FFT2D) { + Repeatedly([this]() { + std::vector dims = RandomDims(2, kDefaultMaxRank); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("FFT2D").RandomInput(DT_COMPLEX64, dims)); + }); +} + +TEST_F(OpTest, FFT3D) { + Repeatedly([this]() { + std::vector dims = RandomDims(3, kDefaultMaxRank); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("FFT3D").RandomInput(DT_COMPLEX64, dims)); + }); +} + +TEST_F(OpTest, IFFT) { + Repeatedly([this]() { + std::vector dims = RandomDims(1, kDefaultMaxRank); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("IFFT").RandomInput(DT_COMPLEX64, dims)); + }); +} + +TEST_F(OpTest, IFFT2D) { + Repeatedly([this]() { + std::vector dims = RandomDims(2, kDefaultMaxRank); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("IFFT2D").RandomInput(DT_COMPLEX64, dims)); + }); +} + +TEST_F(OpTest, IFFT3D) { + Repeatedly([this]() { + std::vector dims = RandomDims(3, kDefaultMaxRank); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("IFFT3D").RandomInput(DT_COMPLEX64, dims)); + }); +} + +TEST_F(OpTest, RFFT) { + Repeatedly([this]() { + std::vector dims = RandomDims(1, kDefaultMaxRank, 3); + Tensor fft_shape = test::AsTensor(AsInt32s({dims[dims.size() - 1]})); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("RFFT").RandomInput(DT_FLOAT, dims).Input(fft_shape)); + }); +} + +TEST_F(OpTest, RFFT2D) { + Repeatedly([this]() { + std::vector dims = RandomDims(2, kDefaultMaxRank, 3); + Tensor fft_shape = test::AsTensor( + AsInt32s({dims[dims.size() - 2], dims[dims.size() - 1]})); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("RFFT2D").RandomInput(DT_FLOAT, dims).Input(fft_shape)); + }); +} + +TEST_F(OpTest, RFFT3D) { + Repeatedly([this]() { + std::vector dims = RandomDims(3, kDefaultMaxRank, 3); + Tensor fft_shape = test::AsTensor(AsInt32s( + {dims[dims.size() - 3], dims[dims.size() - 2], dims[dims.size() - 1]})); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("RFFT3D").RandomInput(DT_FLOAT, dims).Input(fft_shape)); + }); +} + +TEST_F(OpTest, IRFFT) { + Repeatedly([this]() { + std::vector dims = RandomDims(1, kDefaultMaxRank, 3); + int64 orig_size = dims[dims.size() - 1]; + dims[dims.size() - 1] = dims[dims.size() - 1] / 2 + 1; + Tensor fft_shape = test::AsTensor(AsInt32s({orig_size})); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("IRFFT") + .RandomInput(DT_COMPLEX64, dims) + .Input(fft_shape)); + }); +} + +TEST_F(OpTest, IRFFT2D) { + Repeatedly([this]() { + std::vector dims = RandomDims(2, kDefaultMaxRank, 3); + std::vector orig_size = {dims[dims.size() - 2], + dims[dims.size() - 1]}; + dims[dims.size() - 1] = dims[dims.size() - 1] / 2 + 1; + Tensor fft_shape = test::AsTensor(AsInt32s({orig_size})); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("IRFFT2D") + .RandomInput(DT_COMPLEX64, dims) + .Input(fft_shape)); + }); +} + +TEST_F(OpTest, IRFFT3D) { + Repeatedly([this]() { + std::vector dims = RandomDims(3, kDefaultMaxRank, 3); + std::vector orig_size = { + dims[dims.size() - 3], dims[dims.size() - 2], dims[dims.size() - 1]}; + dims[dims.size() - 1] = dims[dims.size() - 1] / 2 + 1; + Tensor fft_shape = test::AsTensor(AsInt32s({orig_size})); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("IRFFT3D") + .RandomInput(DT_COMPLEX64, dims) + .Input(fft_shape)); + }); +} + TEST_F(OpTest, Conv2D) { Repeatedly([this]() { WindowedSpatialDims d = ChooseWindowedSpatialDims(2); diff --git a/tensorflow/compiler/tests/reverse_sequence_op_test.py b/tensorflow/compiler/tests/reverse_sequence_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1a5d05094e53cfecd9476d7d87f023e8a02d7458 --- /dev/null +++ b/tensorflow/compiler/tests/reverse_sequence_op_test.py @@ -0,0 +1,93 @@ +# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for tensorflow.ops.reverse_sequence_op.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class ReverseSequenceTest(XLATestCase): + + def _testReverseSequence(self, + x, + batch_axis, + seq_axis, + seq_lengths, + truth, + expected_err_re=None): + with self.test_session(): + p = array_ops.placeholder(dtypes.as_dtype(x.dtype)) + lengths = array_ops.placeholder(dtypes.as_dtype(seq_lengths.dtype)) + with self.test_scope(): + ans = array_ops.reverse_sequence( + p, batch_axis=batch_axis, seq_axis=seq_axis, seq_lengths=lengths) + if expected_err_re is None: + tf_ans = ans.eval(feed_dict={p: x, lengths: seq_lengths}) + self.assertAllClose(tf_ans, truth, atol=1e-10) + else: + with self.assertRaisesOpError(expected_err_re): + ans.eval(feed_dict={p: x, lengths: seq_lengths}) + + def testSimple(self): + x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.int32) + expected = np.array([[1, 2, 3], [6, 5, 4], [8, 7, 9]], dtype=np.int32) + self._testReverseSequence( + x, + batch_axis=0, + seq_axis=1, + seq_lengths=np.array([1, 3, 2], np.int32), + truth=expected) + + def _testBasic(self, dtype, len_dtype): + x = np.asarray( + [[[1, 2, 3, 4], [5, 6, 7, 8]], [[9, 10, 11, 12], [13, 14, 15, 16]], + [[17, 18, 19, 20], [21, 22, 23, 24]]], + dtype=dtype) + x = x.reshape(3, 2, 4, 1, 1) + x = x.transpose([2, 1, 0, 3, 4]) # permute axes 0 <=> 2 + + # reverse dim 2 up to (0:3, none, 0:4) along dim=0 + seq_lengths = np.asarray([3, 0, 4], dtype=len_dtype) + + truth_orig = np.asarray( + [ + [[3, 2, 1, 4], [7, 6, 5, 8]], # reverse 0:3 + [[9, 10, 11, 12], [13, 14, 15, 16]], # reverse none + [[20, 19, 18, 17], [24, 23, 22, 21]] + ], # reverse 0:4 (all) + dtype=dtype) + truth_orig = truth_orig.reshape(3, 2, 4, 1, 1) + truth = truth_orig.transpose([2, 1, 0, 3, 4]) # permute axes 0 <=> 2 + + seq_axis = 0 # permute seq_axis and batch_axis (originally 2 and 0, resp.) + batch_axis = 2 + self._testReverseSequence(x, batch_axis, seq_axis, seq_lengths, truth) + + def testSeqLength(self): + for dtype in self.all_types: + for seq_dtype in self.int_types: + self._testBasic(dtype, seq_dtype) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/scatter_nd_op_test.py b/tensorflow/compiler/tests/scatter_nd_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..638946e234daf28dc4a34e6c33fc0f78b8e8699b --- /dev/null +++ b/tensorflow/compiler/tests/scatter_nd_op_test.py @@ -0,0 +1,188 @@ +# 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 tensorflow.ops.tf.scatter_nd.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import errors +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +def _AsType(v, vtype): + return v.astype(vtype) if isinstance(v, np.ndarray) else vtype(v) + + +def _FlatInnerDims(tensor, ndims=2): + shape = list(tensor.shape) + return tensor.reshape( + [functools.reduce(lambda x, y: x * y, shape[:-ndims + 1], 1)] + + shape[-ndims + 1:]) + + +def _FlatOuterDims(tensor, ndims=2): + shape = list(tensor.shape) + return tensor.reshape( + shape[:ndims - 1] + + [functools.reduce(lambda x, y: x * y, shape[ndims - 1:], 1)]) + + +def _NumpyScatterNd(ref, indices, updates, op): + ixdim = indices.shape[-1] + num_updates = indices.size // ixdim + total_nd = len(ref.shape) + slice_size = 1 + for i in range(ixdim, total_nd): + slice_size *= ref.shape[i] + flat_indices = _FlatInnerDims(indices) + flat_updates = updates.reshape((num_updates, slice_size)) + output_flat = _FlatOuterDims(ref, ixdim + 1) + for ix_updates, ix_output in enumerate(flat_indices): + ix_output = tuple(ix_output) + output_flat[ix_output] = op(output_flat[ix_output], + flat_updates[ix_updates]) + return output_flat.reshape(ref.shape) + + +def _NumpyUpdate(indices, updates, shape): + ref = np.zeros(shape, dtype=updates.dtype) + return _NumpyScatterNd(ref, indices, updates, lambda p, u: u) + + +class ScatterNdTest(XLATestCase): + + def _VariableRankTest(self, + np_scatter, + tf_scatter, + vtype, + itype, + repeat_indices=False): + np.random.seed(8) + ref_shapes = [(3, 6), (3, 6), (3, 6, 9), (3, 6, 9), (3, 6, 9), (3, 6, 9)] + indices_shapes = [(2,), (2, 2), (2,), (2, 2), (2, 3), (2, 3, 3)] + for ref_shape, indices_shape in zip(ref_shapes, indices_shapes): + num_updates = indices_shape[0] + ixdim = indices_shape[-1] + + indexable_area_shape = () + for i in range(ixdim): + indexable_area_shape += (ref_shape[i],) + all_indices = [ + list(coord) + for coord, _ in np.ndenumerate(np.empty(indexable_area_shape, vtype)) + ] + np.random.shuffle(all_indices) + indices = np.array(all_indices[:num_updates]) + + if num_updates > 1 and repeat_indices: + indices = indices[:num_updates // 2] + for _ in range(num_updates - num_updates // 2): + indices = np.append( + indices, [indices[np.random.randint(num_updates // 2)]], axis=0) + np.random.shuffle(indices) + indices = _AsType(indices[:num_updates], itype) + + updates_shape = (num_updates,) + for i in range(ixdim, len(ref_shape)): + updates_shape += (ref_shape[i],) + updates = _AsType(np.random.randn(*(updates_shape)), vtype) + + # Scatter via numpy + np_out = np_scatter(indices, updates, ref_shape) + # Scatter via tensorflow + tf_out = tf_scatter(indices, updates, ref_shape) + + self.assertAllClose(np_out, tf_out) + + def _VariableRankTests(self, np_scatter, tf_scatter): + for vtype in self.numeric_types: + for itype in set([np.int32, np.int64]).intersection(set(self.int_types)): + self._VariableRankTest(np_scatter, tf_scatter, vtype, itype) + + def _runScatterNd(self, indices, updates, shape): + with self.test_session(): + updates_placeholder = array_ops.placeholder(updates.dtype) + indices_placeholder = array_ops.placeholder(indices.dtype) + with self.test_scope(): + output = array_ops.scatter_nd(indices_placeholder, updates_placeholder, + shape) + feed_dict = {updates_placeholder: updates, indices_placeholder: indices} + return output.eval(feed_dict=feed_dict) + + def testSimple(self): + indices = np.array([[4], [3], [1], [7]], dtype=np.int32) + updates = np.array([9, 10, 11, 12], dtype=np.float32) + expected = np.array([0, 11, 0, 10, 9, 0, 0, 12], dtype=np.int32) + self.assertAllEqual(expected, self._runScatterNd(indices, updates, [8])) + + def testSimple2(self): + indices = np.array([[1, 0], [1, 1]], dtype=np.int32) + updates = np.array([11., 12.], dtype=np.float32) + expected = np.array([[0., 0.], [11., 12.], [0., 0.]], dtype=np.float32) + self.assertAllEqual(expected, self._runScatterNd(indices, updates, [3, 2])) + + def testSimple3(self): + indices = np.array([[1]], dtype=np.int32) + updates = np.array([[11., 12.]], dtype=np.float32) + expected = np.array([[0., 0.], [11., 12.], [0., 0.]]) + self.assertAllEqual(expected, self._runScatterNd(indices, updates, [3, 2])) + + def testVariableRankUpdate(self): + self._VariableRankTests(_NumpyUpdate, self._runScatterNd) + + def testExtraIndicesDimensions(self): + indices = np.zeros([1, 1, 2], np.int32) + updates = np.zeros([1, 1], np.int32) + expected = np.zeros([2, 2], dtype=np.int32) + self.assertAllEqual(expected, self._runScatterNd(indices, updates, [2, 2])) + + def testRank3InvalidShape1(self): + indices = np.zeros([3, 2, 2], np.int32) + updates = np.zeros([2, 2, 2], np.int32) + with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, + "Must have updates.shape"): + self._runScatterNd(indices, updates, [2, 2, 2]) + + def testRank3InvalidShape2(self): + indices = np.zeros([2, 2, 1], np.int32) + updates = np.zeros([2, 2], np.int32) + with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, + "Must have updates.shape"): + self._runScatterNd(indices, updates, [2, 2, 2]) + + def testScatterOutOfRange(self): + updates = np.array([-3, -4, -5]).astype(np.float32) + + # Indices all in range, no problem. + indices = np.array([[2], [0], [5]], dtype=np.int32) + self._runScatterNd(indices, updates, [6]) + + # Indices out of range should not fail. It produces implementation-defined + # output. + indices = np.array([[-1], [0], [5]], dtype=np.int32) + self._runScatterNd(indices, updates, [6]) + indices = np.array([[2], [0], [6]], dtype=np.int32) + self._runScatterNd(indices, updates, [6]) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/segment_reduction_ops_test.py b/tensorflow/compiler/tests/segment_reduction_ops_test.py index 260a04421b62310c109d8f0ea72875a50c234bb0..23bc39cf3f7087424719edfb8b6ee35d87295534 100644 --- a/tensorflow/compiler/tests/segment_reduction_ops_test.py +++ b/tensorflow/compiler/tests/segment_reduction_ops_test.py @@ -60,6 +60,14 @@ class SegmentReductionOpsTest(XLATestCase): np.array([0, 1, 2, 3, 4, 5], dtype=dtype), np.array([3, 0, 2, 1, 3, 3], dtype=np.int32), 4)) + def testUnsortedSegmentSum1DIndices1DDataNegativeIndices(self): + for dtype in self.numeric_types: + self.assertAllClose( + np.array([0, 3, 2, 5], dtype=dtype), + self.UnsortedSegmentSum( + np.array([0, 1, 2, 3, 4, 5], dtype=dtype), + np.array([3, -1, 2, 1, -1, 3], dtype=np.int32), 4)) + def testUnsortedSegmentSum1DIndices2DDataDisjoint(self): for dtype in self.numeric_types: data = np.array( diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py index 0a6fe04d3cdd29f1d40d33be1f4319090e7ba3d1..3d3e112f4821ea8e57ea9589a5b4433647ad294b 100644 --- a/tensorflow/compiler/tests/unary_ops_test.py +++ b/tensorflow/compiler/tests/unary_ops_test.py @@ -67,8 +67,10 @@ class UnaryOpsTest(XLATestCase): output = op(pinp) result = session.run(output, {pinp: inp}) if equality_test is None: - equality_test = self.assertAllCloseAccordingToType - equality_test(result, expected, rtol=rtol, atol=atol) + self.assertAllCloseAccordingToType( + result, expected, rtol=rtol, atol=atol, bfloat16_rtol=0.03) + else: + equality_test(result, expected, rtol=rtol, atol=atol) def ListsAreClose(self, result, expected, rtol, atol): """Tests closeness of two lists of floats.""" @@ -152,6 +154,21 @@ class UnaryOpsTest(XLATestCase): def testFloatOps(self): for dtype in self.float_types: + x = np.arange(-0.90, 0.90, 0.25) + self._assertOpOutputMatchesExpected( + math_ops.acos, + x.astype(dtype), + expected=np.arccos(x).astype(dtype)) + self._assertOpOutputMatchesExpected( + math_ops.asin, + x.astype(dtype), + expected=np.arcsin(x).astype(dtype)) + x = np.arange(-3, 3).reshape(1, 3, 2) + self._assertOpOutputMatchesExpected( + math_ops.atan, + x.astype(dtype), + expected=np.arctan(x).astype(dtype)) + self._assertOpOutputMatchesExpected( math_ops.acosh, np.array([1, 2, 3, 4], dtype=dtype), diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 3c7dfef03dfb5d86dd63fd4aa84ad56081833035..fb82c2601c432cee425a46a3b6dc2c55febeda87 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -312,6 +312,7 @@ tf_cc_test( "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:ops", + "//tensorflow/cc:resource_variable_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", diff --git a/tensorflow/compiler/tf2xla/const_analysis.cc b/tensorflow/compiler/tf2xla/const_analysis.cc index 0249500910c6ae441f038fe9ad6178794f1997ac..82923722c54d235716b9138d95a75a441df924ca 100644 --- a/tensorflow/compiler/tf2xla/const_analysis.cc +++ b/tensorflow/compiler/tf2xla/const_analysis.cc @@ -64,7 +64,7 @@ Status BackwardsConstAnalysis(const Graph& g, // Mark any compile-time constant operator arguments as const. const std::unordered_set* const_inputs = XlaOpRegistry::CompileTimeConstantInputs(node->type_string()); - if (!const_inputs) return; + if (!const_inputs || const_inputs->empty()) return; NameRangeMap input_name_ranges; status = diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc index dd67a1dea9656bac9cf3eaa09295a0d42e283706..f8169795ddfb7fd4e93d3f136c51623385868951 100644 --- a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc @@ -37,6 +37,8 @@ namespace tensorflow { namespace { +using xla::StatusOr; + const char* const kArgOp = "_Arg"; const char* const kRetValOp = "_Retval"; @@ -76,6 +78,20 @@ struct Frame { std::unordered_set nodes; }; +// Comparison function used for sorting nodes consistently. +// a) resource variables are last, and +// b) sort lexicographically by name (for deterministic output). +struct NodeCmp { + bool operator()(const Node* lhs, const Node* rhs) const { + bool lhs_is_resource = + lhs->num_inputs() > 0 ? (lhs->input_type(0) == DT_RESOURCE) : false; + bool rhs_is_resource = + rhs->num_inputs() > 0 ? (rhs->input_type(0) == DT_RESOURCE) : false; + return std::tie(lhs_is_resource, lhs->name()) < + std::tie(rhs_is_resource, rhs->name()); + } +}; + // Returns a textual representation of the names of the nodes in the input. template string NodesToString(const T& nodes) { @@ -141,7 +157,7 @@ Status CopySubgraph(const Graph& graph, const Frame* frame, return Status::OK(); } -xla::StatusOr AddNode(const NodeDef& node_def, Graph* graph) { +StatusOr AddNode(const NodeDef& node_def, Graph* graph) { Status status; Node* inserted_node = graph->AddNode(node_def, &status); if (!status.ok()) { @@ -150,7 +166,7 @@ xla::StatusOr AddNode(const NodeDef& node_def, Graph* graph) { return inserted_node; } -xla::StatusOr BuildArgNode(Graph* graph, DataType type, int index) { +StatusOr BuildArgNode(Graph* graph, DataType type, int index) { NodeDef arg_def; NodeDefBuilder builder(strings::StrCat(kArgOp, index), kArgOp); builder.Attr("T", type); @@ -159,7 +175,7 @@ xla::StatusOr BuildArgNode(Graph* graph, DataType type, int index) { return AddNode(arg_def, graph); } -xla::StatusOr BuildRetvalNode(Graph* graph, DataType type, int index) { +StatusOr BuildRetvalNode(Graph* graph, DataType type, int index) { NodeDef ret_def; ret_def.set_op(kRetValOp); ret_def.set_name(strings::StrCat(kRetValOp, index)); @@ -269,7 +285,8 @@ Status BuildLoopBody(const Graph& graph, Frame* frame, Status FunctionalizeLoop(Graph* graph, Frame* frame, FunctionLibraryDefinition* library) { VLOG(2) << "Frame " << frame->name << " before: " - << dump_graph::DumpGraphToFile("functionalize_before", *graph); + << dump_graph::DumpGraphToFile("functionalize_before", *graph, + library); // Split loop-varying Enter nodes with multiple successors. If the same // Tensor is fed as input to multiple loop arguments, we may end up with a @@ -310,16 +327,9 @@ Status FunctionalizeLoop(Graph* graph, Frame* frame, } frame->args = std::move(args); - // Order the arguments so that: - // a) resource variables are last, and - // b) sort lexicographically by name (for deterministic output). - std::sort(frame->args.begin(), frame->args.end(), - [](const Arg& a, const Arg& b) { - bool a_is_resource = (a.enter->input_type(0) == DT_RESOURCE); - bool b_is_resource = (b.enter->input_type(0) == DT_RESOURCE); - return std::tie(a_is_resource, a.enter->name()) < - std::tie(b_is_resource, b.enter->name()); - }); + std::sort( + frame->args.begin(), frame->args.end(), + [](const Arg& a, const Arg& b) { return NodeCmp()(a.enter, b.enter); }); if (frame->loop_cond == nullptr) { return errors::InvalidArgument("Loop ", frame->name, @@ -418,16 +428,36 @@ Status FunctionalizeLoop(Graph* graph, Frame* frame, // identity nodes are values used by the loop body or condition. // The Identity node may have the wrong device so copy the device from // one of its outputs instead. + std::deque possible_exit; for (const Edge* edge : arg.switch_node->out_edges()) { - if (edge->src_output() == 0 && IsExit(edge->dst())) { + if (edge->src_output() == 0) { + possible_exit.push_back(edge); + } + if (IsIdentity(edge->dst())) { + TF_RETURN_IF_ERROR( + SetNodeShardingFromNeighbors(edge->dst(), /*out_edges=*/true)); + } + } + // TODO(b/67425339): Allow general graph between switch and exit. + while (!possible_exit.empty()) { + const Edge* edge = possible_exit.front(); + possible_exit.pop_front(); + if (IsExit(edge->dst())) { if (arg.exit != nullptr) { return errors::InvalidArgument("Duplicate Exit successors to ", arg.switch_node->name()); } arg.exit = edge->dst(); - } else if (StringPiece(edge->dst()->type_string()) == "Identity") { - TF_RETURN_IF_ERROR( - SetNodeShardingFromNeighbors(edge->dst(), /*out_edges=*/true)); + } else { + if (!IsIdentity(edge->dst())) { + return errors::Unimplemented("General graph between switch (", + arg.switch_node->name(), + ") and exit node of frame ", + frame->name, " not supported yet."); + } + for (const Edge* out : edge->dst()->out_edges()) { + possible_exit.push_back(out); + } } } } @@ -441,7 +471,7 @@ Status FunctionalizeLoop(Graph* graph, Frame* frame, TF_RETURN_IF_ERROR(BuildLoopBody(*graph, frame, &arg_types, &body_graph)); VLOG(2) << "Frame " << frame->name << " condition: " - << dump_graph::DumpGraphToFile("loop_condition", *cond_graph) + << dump_graph::DumpGraphToFile("loop_condition", *cond_graph, library) << " body: " << dump_graph::DumpGraphToFile("loop_body", *body_graph); static std::atomic sequence_num(0LL); @@ -522,7 +552,8 @@ Status FunctionalizeLoop(Graph* graph, Frame* frame, frame->parent->nodes.insert(while_node); VLOG(2) << "Frame " << frame->name << " after: " - << dump_graph::DumpGraphToFile("functionalize_after", *graph); + << dump_graph::DumpGraphToFile("functionalize_after", *graph, + library); return Status::OK(); } @@ -542,18 +573,6 @@ class FunctionalizeCond { // Returns a textual representation of the Branch b. static string Branch_Name(FunctionalizeCond::Branch b); - // Comparison function used for sorting nodes consistently. - struct CondCmp { - bool operator()(const Node* lhs, const Node* rhs) const { - bool lhs_is_resource = - lhs->num_inputs() > 0 ? (lhs->input_type(0) == DT_RESOURCE) : false; - bool rhs_is_resource = - rhs->num_inputs() > 0 ? (rhs->input_type(0) == DT_RESOURCE) : false; - return std::tie(lhs_is_resource, lhs->name()) < - std::tie(rhs_is_resource, rhs->name()); - } - }; - // Functionalize all the switch-merge nodes of a loop-free graph into XlaIf // nodes. That is, attempt to transform every remaining switch and merge nodes // in the graph into XlaIf nodes. @@ -561,6 +580,20 @@ class FunctionalizeCond { static Status Functionalize(Graph* graph, FunctionLibraryDefinition* library); private: + // CondArgNode represents a input to the conditional and its corresponding + // switch nodes. + struct CondArgNode { + explicit CondArgNode(Node* input) : input(input) {} + string ToString() const { + return strings::StrCat("input=", input->name(), + " switches=", NodesToString(switches)); + } + + Node* input; + std::vector switches; + }; + using CondArgNodes = std::vector; + struct ForwardFlowNode { explicit ForwardFlowNode(Branch branch = Branch::kNeither) : branch(branch), count(0) {} @@ -571,57 +604,87 @@ class FunctionalizeCond { int count; }; - FunctionalizeCond(Graph* graph, FunctionLibraryDefinition* library) - : library_(library), graph_(graph) {} + // Group of switch nodes that will be part of the same XlaIf. + struct SwitchCluster { + explicit SwitchCluster(Node* predicate) : predicate(predicate) {} + string ToString() const { + return strings::StrCat(name, " predicate=", predicate->name(), + " switches=", NodesToString(switches)); + } + + string name; + Node* predicate; + std::vector switches; + }; + + FunctionalizeCond(Graph* graph, FunctionLibraryDefinition* library, + bool dump_graphs) + : library_(library), graph_(graph), dump_graphs_(dump_graphs) {} // Perform the actual cond functionalization. Iterate over groups of switch // nodes (linked by common predicate), from innermost to outermost, and // extract into XlaIf nodes. Status FunctionalizeInternal(); - // Converts a Merge node to a XlaIf. This encapsulates the process of - // extracting the bodies needed for the then and else branch, creates a XlaIf - // node, removing the nodes of the branches from the graph and replacing the - // merge node with a XlaIf. - Status ConvertCorrespondingMergeToXlaIf( - const std::vector& switch_nodes, - const std::vector& merge_nodes, Node* predicate); + // Determines the branch_map (mapping from node to branch of cond) and + // frontier (the nodes where the cond ends). + StatusOr, + std::unordered_set>> + DetermineBranchMapAndFrontier(const SwitchCluster& switch_cluster); + + // Returns XlaIf node created from subgraph of merge and switch nodes. This + // encapsulates the process of extracting the bodies needed for the then and + // else branch, creates a XlaIf node, removing the nodes of the branches from + // the graph and replacing the merge node with a XlaIf. + StatusOr ConvertToXlaIf(const CondArgNodes& cond_arg_nodes, + const SwitchCluster& switch_cluster, + const std::vector& switches); // Builds a XlaIfOp to replace the Switch-Graph-Merge cluster with. - xla::StatusOr BuildAndAddXlaIfOp( - const std::vector& switch_nodes, - const std::vector& merge_nodes, Node* predicate); + StatusOr BuildAndAddXlaIfOp(const CondArgNodes& cond_arg_nodes, + const SwitchCluster& switch_cluster, + const std::vector& merge_nodes); // Extracts a function body corresponding to the given input edge of the merge // node. - Status ExtractBody(const std::vector& switch_nodes, + Status ExtractBody(const CondArgNodes& cond_arg_nodes, + const std::vector& switches, const std::vector& merge_nodes, int input_edge, Graph* body); // Adds all the input edges to `if_node` corresponding to the arguments. - Status AddInputEdges(const std::vector& cond_args, Node* predicate, + Status AddInputEdges(const CondArgNodes& cond_arg_nodes, Node* predicate, Node* if_node); // Adds all output edges from the `if_node`. Status AddOutputEdges(const std::vector& outputs, Node* if_node); - // Returns the switches of graph_ in postorder. Dead switch nodes are skipped - // and removed from the graph. - std::vector DetermineSwitchOrder(); + // Returns the switch clusters of graph_ in postorder. Dead switch nodes are + // skipped and removed from the graph. + StatusOr> DeterminePredicateSwitchOrder(); // Update the state for destination based on the state of source and the node // being updated. Status Join(const ForwardFlowNode& src_state, const Node* dst, ForwardFlowNode* dst_state); - // Validates that the branch_map and frontier of nodes for the conditional + // Ensure that all nodes in the branch_map are dominated by the switch + // nodes. Returns nodes that are not dominated by the switches but are a + // control dependency of a node in the cond, and remove such control + // dependencies. + StatusOr> EnsureDominanceAndReturnNonDominatedControlNodes( + const std::unordered_map& branch_map, + const std::vector& switches); + + // Validates that the frontier of nodes for the conditional // section are as expected. - Status ValidBranchMapAndFrontier( + Status ValidateFrontier( const std::unordered_map& branch_map, const std::unordered_set& frontier); FunctionLibraryDefinition* library_; Graph* graph_; + bool dump_graphs_; }; bool IsDeadSwitch(const Node* node) { @@ -645,23 +708,11 @@ string FunctionalizeCond::Branch_Name(FunctionalizeCond::Branch b) { return branch_name[b]; } -Status FunctionalizeCond::ValidBranchMapAndFrontier( +Status FunctionalizeCond::ValidateFrontier( const std::unordered_map& branch_map, const std::unordered_set& frontier) { std::unordered_set pending[kNumBranchTypes]; - for (const auto& kv : branch_map) { - if (kv.second.count != kv.first->in_edges().size()) { - return errors::FailedPrecondition("Value ", kv.first->DebugString(), - " not dominated by switch nodes."); - } - if (VLOG_IS_ON(1)) { - // Append attribute to the graph if running with logging to make the - // changes clearer in the visualization. - kv.first->AddAttr("_XlaFunctionalizeBranch", - Branch_Name(kv.second.branch)); - } - } for (Node* n : frontier) { pending[branch_map.at(n).branch].insert(n); } @@ -681,6 +732,13 @@ Status FunctionalizeCond::ValidBranchMapAndFrontier( ") in both Else and Then branch should be in Both."); } } + // An empty frontier indicates a dead switch. Above we attempt to remove dead + // switch nodes, but not all are removed so don't treat it as an error yet. + // TODO(jpienaar): Find out why dead switch nodes remain. + // if (pending[kBoth].empty() && pending[kThenBranch].empty() && + // pending[kElseBranch].empty()) { + // return errors::Internal("Unexpected empty frontier for switch nodes"); + // } return Status::OK(); } @@ -707,44 +765,191 @@ Status FunctionalizeCond::Join(const ForwardFlowNode& src_state, return Status::OK(); } -std::vector FunctionalizeCond::DetermineSwitchOrder() { +StatusOr> +FunctionalizeCond::DeterminePredicateSwitchOrder() { + struct Cluster { + bool operator==(const Cluster& other) const { + return representative == other.representative; + } + int representative = -1; + }; + + // Perform a DFS over the graph and + // * Determine the reverse topological order of the nodes (there should be no + // cycles at this point so the post-order numbering corresponds to the + // reverse topological sorting); + // * Identify dead switches; + // * Initialize the cluster's representative; + std::vector> clusters(graph_->num_node_ids()); std::vector dead_switches; std::vector switch_order; - DFS(*graph_, nullptr, [this, &dead_switches, &switch_order](Node* n) { + std::vector rev_topo_sorted_nodes; + DFS(*graph_, nullptr, [&](Node* n) { + clusters[n->id()].Get().representative = n->id(); if (IsSwitch(n)) { if (IsDeadSwitch(n)) { dead_switches.push_back(n); } else { + rev_topo_sorted_nodes.push_back(n); switch_order.push_back(n); } + } else if (n->IsOp()) { + // Exclude src and sink nodes from further consideration. + rev_topo_sorted_nodes.push_back(n); } }); + std::vector switch_clusters; + // Return early if there are no switches in the graph. + if (switch_order.empty()) { + return switch_clusters; + } + // Remove all dead switch nodes. for (Node* n : dead_switches) { + VLOG(2) << "Removing dead switch: " << n->DebugString(); graph_->RemoveNode(n); } - return switch_order; -} + // Identify switch nodes that are part of the same control flow context by + // considering the operands of operations: an operation is part of the same + // control context as its operands unless the operation is a switch. Control + // dependencies are considered part of the same control flow context if the + // switch depth is the same (see comment below). + + // entry_cluster records the input cluster to a switch node. This is used when + // merging with a merge node where the dst's cluster is merged with the entry + // cluster of the merge node's cluster (which corresponds to a switch cluster + // and so has an entry cluster). + std::unordered_map*> entry_cluster; + + // Returns the output cluster of a node. Where the output cluster is cluster + // where the output of the node is used. For non-merge nodes this is simply + // the cluster they are part of, while for merge nodes it is the entry cluster + // of the cluster they are part of (this will correspond to the entry node of + // a switch node that dominates the merge). + auto find_output_cluster = [&](Node* n) { + UnionFind* cluster = &clusters[n->id()]; + if (!IsMerge(n)) return cluster; + auto it = entry_cluster.find(clusters[n->id()].Get().representative); + // If the cluster is not found in the entry_cluster map then an + // instruction not dominated by a switch node has been merged into the + // cluster of the merge. This indicates a failure of the clustering. + CHECK(it != entry_cluster.end()) + << "Unable to find entry for n=" << n->id() << " (" + << cluster->Get().representative << ")"; + return it->second; + }; -Status FunctionalizeCond::FunctionalizeInternal() { - std::vector switch_order = DetermineSwitchOrder(); - // If there are no switch nodes, then terminate. - if (switch_order.empty()) { - return Status::OK(); + // TODO(jpienaar): This could be combined with DetermineBranchMapAndFrontier. + std::vector switch_depth(graph_->num_node_ids()); + for (auto it = rev_topo_sorted_nodes.rbegin(); + it != rev_topo_sorted_nodes.rend(); ++it) { + Node* n = *it; + + // Compute switch depth. + int new_switch_depth = 0; + for (const Edge* e : n->in_edges()) { + Node* src = e->src(); + new_switch_depth = std::max( + new_switch_depth, switch_depth[src->id()] - (IsMerge(src) ? 1 : 0)); + } + switch_depth[n->id()] = new_switch_depth + (IsSwitch(n) ? 1 : 0); + + // Only merge the input operands of a switch. The switch's clustering itself + // is determined by the interaction of the switch's outputs. + if (IsSwitch(n)) { + Node* input; + TF_CHECK_OK(n->input_node(0, &input)); + entry_cluster[n->id()] = &clusters[input->id()]; + UnionFind* cluster = find_output_cluster(input); + int cluster_depth = switch_depth[cluster->Get().representative]; + // Merge the inputs of the switch node with one another. This results in + // predicates and control input residing in the same cluster. + for (const Edge* e : n->in_edges()) { + Node* src = e->src(); + UnionFind* src_cluster = find_output_cluster(src); + int src_cluster_depth = switch_depth[src_cluster->Get().representative]; + if (cluster_depth != src_cluster_depth) { + return errors::InvalidArgument( + "Unable to functionalize control flow in graph: Switch ('", + n->name(), "') has operands ('", input->name(), "' and '", + src->name(), "') that have different switch depths (", + cluster_depth, " != ", src_cluster_depth, ")"); + } + cluster->Merge(src_cluster); + } + continue; + } + + for (const Edge* e : n->in_edges()) { + Node* src = e->src(); + if (!src->IsOp()) continue; + UnionFind* cluster = find_output_cluster(src); + // Merge a node with its data operands and with its control operands if + // the src and dst are in the same ControlContext. The ControlContext is + // not explicitly available here, and instead the switch depth is used as + // a proxy here. Due to the invariant that control edges can only be from + // a containing scope to an inner scope or from the inner scope to its + // containing scope (for exit nodes), the switch depth will only match if + // the src and dst are in the same ControlContext. Control edges between + // ControlContexts are handled during the extraction. + int src_id = cluster->Get().representative; + int src_depth = switch_depth[src_id]; + if (!e->IsControlEdge() || new_switch_depth == src_depth) { + if (src_depth != new_switch_depth) { + return errors::InvalidArgument( + "Unable to functionalize control flow in graph: Operand ('", + src->name(), "') and operator ('", n->name(), + "') have different switch depths (", src_depth, + " != ", new_switch_depth, ")"); + } + cluster->Merge(&clusters[n->id()]); + } + } } - struct PredicateSwitches { - explicit PredicateSwitches(Node* predicate) : predicate(predicate) {} + if (dump_graphs_) { + // Mark the switch cluster each node is part of. + for (Node* n : graph_->nodes()) { + n->ClearAttr("_XlaFunctionalizeSwitchGroup"); + n->AddAttr("_XlaFunctionalizeSwitchGroup", + clusters[n->id()].Get().representative); + } + LOG(INFO) << "FunctionalizeControlFlow (with_clusters): " + << dump_graph::DumpGraphToFile("functionalize_clustered", *graph_, + library_); + } - Node* predicate; - std::vector switches; + // Verify all the nodes of a cluster are at the same depth. + std::unordered_map> cluster_to_depth_node; + for (Node* n : graph_->nodes()) { + int depth = switch_depth[n->id()]; + int cluster_rep = clusters[n->id()].Get().representative; + auto it = cluster_to_depth_node.find(cluster_rep); + if (it == cluster_to_depth_node.end()) { + cluster_to_depth_node[cluster_rep] = std::make_pair(depth, n); + } else { + if (it->second.first != depth) { + return errors::Internal( + "Illegal clustering created, mismatch in depths:", "\n\t", + n->DebugString(), "(", clusters[n->id()].Get().representative, + ") at depth=", depth, " vs\n\t", it->second.second->DebugString(), + "(", clusters[n->id()].Get().representative, ") at depth ", + it->second.first); + } + } + } + + struct Hash { + size_t operator()(const std::pair& item) const { + return Hash64Combine(hash()(item.first), + std::hash()(item.second.representative)); + } }; // Merge Switch nodes with common predicate. - std::vector predicate_switch_order; - std::unordered_map predicate_index; + std::unordered_map, int, Hash> predicate_index; // The nodes in switch_order are in reverse topological order, but the // clustered switches need not be (i.e., when considered as a cluster one // element of a cluster may be later in the topological order than another @@ -753,98 +958,198 @@ Status FunctionalizeCond::FunctionalizeInternal() { for (auto it = switch_order.rbegin(); it != switch_order.rend(); ++it) { Node* pred; TF_CHECK_OK((*it)->input_node(1, &pred)); - if (predicate_index.find(pred) == predicate_index.end()) { - predicate_index[pred] = predicate_switch_order.size(); - predicate_switch_order.emplace_back(pred); + auto repr = std::make_pair(pred, clusters[(*it)->id()].Get()); + if (predicate_index.find(repr) == predicate_index.end()) { + predicate_index[repr] = switch_clusters.size(); + switch_clusters.emplace_back(pred); + // Generate a name by concatenating with the cluster representative as + // there could be multiple switch clusters with the same predicate. + switch_clusters[predicate_index[repr]].name = + strings::StrCat(pred->name(), "_", repr.second.representative, "_If"); + } + switch_clusters[predicate_index[repr]].switches.push_back(*it); + } + + return switch_clusters; +} + +StatusOr> +FunctionalizeCond::EnsureDominanceAndReturnNonDominatedControlNodes( + const std::unordered_map& branch_map, + const std::vector& switches) { + std::vector old_control_nodes; + for (const auto& kv : branch_map) { + if (kv.second.count != kv.first->in_edges().size()) { + std::vector delete_edges; + for (const Edge* in : kv.first->in_edges()) { + auto it = branch_map.find(in->src()); + if (it == branch_map.end()) { + if (in->IsControlEdge()) { + old_control_nodes.push_back(in->src()); + delete_edges.push_back(in); + } else { + if (IsSwitch(in->src())) { + if (std::find(switches.begin(), switches.end(), in->src()) == + switches.end()) { + return errors::Internal( + "Unexpected switch node found during flow forward: ", + in->src()->DebugString()); + } + continue; + } + return errors::InvalidArgument( + "Value ", kv.first->name(), "'s input, ", in->src()->name(), + ", is not dominated by switch nodes ", NodesToString(switches)); + } + } + } + // Remove control edges from nodes that are not dominated by the switch + // nodes. New control dependencies will be added between these nodes and + // the XlaIf node inserted. + for (const Edge* e : delete_edges) { + graph_->RemoveEdge(e); + } + } + } + return old_control_nodes; +} + +StatusOr< + std::pair, + std::unordered_set>> +FunctionalizeCond::DetermineBranchMapAndFrontier( + const SwitchCluster& switch_cluster) { + std::unordered_map branch_map; + std::unordered_set frontier; + std::vector stack = switch_cluster.switches; + std::vector visited(graph_->num_node_ids(), false); + while (!stack.empty()) { + Node* n = stack.back(); + stack.pop_back(); + + if (visited[n->id()]) { + continue; + } + visited[n->id()] = true; + + // Propagate branch state along each edge of a switch node. + bool sink_only = true; + for (const Edge* e : n->out_edges()) { + Node* out = e->dst(); + if (!out->IsOp()) { + continue; + } + sink_only = false; + // Propagate branch information. + ForwardFlowNode& ffn = branch_map[out]; + if (IsSwitch(n)) { + int index = e->IsControlEdge() ? Branch::kNeither : e->src_output(); + TF_RETURN_IF_ERROR(Join(ForwardFlowNode(Branch(index)), out, &ffn)); + } else { + TF_RETURN_IF_ERROR(Join(branch_map[n], out, &ffn)); + } + if (IsMerge(out)) { + if (out->in_edges().size() == ffn.count) { + frontier.insert(out); + } + } else if (!visited[out->id()]) { + stack.push_back(out); + } + } + if (sink_only) { + if (!IsIdentity(n)) { + VLOG(1) << "Feeding into sink: " << n->DebugString(); + } + } + } + + if (dump_graphs_) { + for (const auto& kv : branch_map) { + // Append attribute to the graph if running with logging to make the + // changes clearer in the visualization. + kv.first->AddAttr("_XlaFunctionalizeBranch", + Branch_Name(kv.second.branch)); } - predicate_switch_order[predicate_index[pred]].switches.push_back(*it); } + return std::make_pair(std::move(branch_map), std::move(frontier)); +} + +Status FunctionalizeCond::FunctionalizeInternal() { + TF_ASSIGN_OR_RETURN(std::vector predicate_switch_order, + DeterminePredicateSwitchOrder()); // Iterate from innermost set of clustered switches to outermost, replacing // matching switch->merge subgraphs with single XlaIf nodes. for (auto it = predicate_switch_order.rbegin(); it != predicate_switch_order.rend(); ++it) { auto& ps = *it; - VLOG(3) << "Flow down from: " << ps.predicate->name() << " -> " - << NodesToString(ps.switches); + VLOG(3) << "Flow down from: " << NodesToString(ps.switches) << " (" + << ps.predicate->name() << ")"; std::unordered_map branch_map; std::unordered_set frontier; - - std::vector stack = ps.switches; - std::vector visited(graph_->num_node_ids(), false); - while (!stack.empty()) { - Node* n = stack.back(); - stack.pop_back(); - - if (visited[n->id()]) { - continue; + TF_ASSIGN_OR_RETURN(std::tie(branch_map, frontier), + DetermineBranchMapAndFrontier(ps)); + + if (dump_graphs_) + LOG(INFO) << "FunctionalizeControlFlow (before XlaIf conversion): " + << dump_graph::DumpGraphToFile("functionalize_bc", *graph_, + library_); + TF_RETURN_IF_ERROR(ValidateFrontier(branch_map, frontier)); + + // Sort the merge and switch nodes using NodeCmp. The switch-nodes are + // further grouped (post sorting) by input to the switch node as in the + // functionalized form each input will be passed in only once. This grouping + // should retain the sorted order. + CondArgNodes cond_arg_nodes; + std::unordered_map input_index; + std::sort(ps.switches.begin(), ps.switches.end(), NodeCmp()); + for (Node* switch_node : ps.switches) { + Node* in; + TF_RETURN_IF_ERROR(switch_node->input_node(0, &in)); + if (input_index.find(in) == input_index.end()) { + input_index[in] = cond_arg_nodes.size(); + cond_arg_nodes.emplace_back(in); } - visited[n->id()] = true; + cond_arg_nodes.at(input_index.at(in)).switches.push_back(switch_node); + } + std::vector merge_nodes(frontier.begin(), frontier.end()); + std::sort(merge_nodes.begin(), merge_nodes.end(), NodeCmp()); - // Propagate branch state along each edge of a switch node. - bool sink_only = true; - for (const Edge* e : n->out_edges()) { - Node* out = e->dst(); - if (!out->IsOp()) { - continue; - } - sink_only = false; - // Propagate branch information. - ForwardFlowNode& ffn = branch_map[out]; - if (IsSwitch(n)) { - int index = e->IsControlEdge() ? Branch::kNeither : e->src_output(); - TF_RETURN_IF_ERROR(Join(ForwardFlowNode(Branch(index)), out, &ffn)); - } else { - TF_RETURN_IF_ERROR(Join(branch_map[n], out, &ffn)); - } - if (IsMerge(out)) { - if (out->in_edges().size() == ffn.count) { - frontier.insert(out); - } - } else if (!visited[out->id()] && ffn.count == out->in_edges().size()) { - // If all predecessors are dominated by the switch nodes, then add - // the output to the stack. - stack.push_back(out); - } - } - if (sink_only) { - if (!IsIdentity(n)) { - VLOG(1) << "Feeding into sink: " << n->DebugString(); - } - } + TF_ASSIGN_OR_RETURN(std::vector old_control_nodes, + EnsureDominanceAndReturnNonDominatedControlNodes( + branch_map, ps.switches)); + + TF_ASSIGN_OR_RETURN(Node * if_node, + ConvertToXlaIf(cond_arg_nodes, ps, merge_nodes)); + for (Node* old : old_control_nodes) { + graph_->AddControlEdge(old, if_node); } - TF_RETURN_IF_ERROR(ValidBranchMapAndFrontier(branch_map, frontier)); - VLOG(2) << "FunctionalizeControlFlow (before XlaIf conversion): " - << dump_graph::DumpGraphToFile("functionalize_bc", *graph_); - std::vector switch_nodes(ps.switches); - std::sort(switch_nodes.begin(), switch_nodes.end(), CondCmp()); - std::vector merge_nodes(frontier.begin(), frontier.end()); - std::sort(merge_nodes.begin(), merge_nodes.end(), CondCmp()); - TF_RETURN_IF_ERROR(ConvertCorrespondingMergeToXlaIf( - switch_nodes, merge_nodes, ps.predicate)); for (auto& del_kv : branch_map) { graph_->RemoveNode(del_kv.first); } - for (Node* node : switch_nodes) { - graph_->RemoveNode(node); + for (auto& kv : cond_arg_nodes) { + for (Node* node : kv.switches) { + graph_->RemoveNode(node); + } } - VLOG(2) << "FunctionalizeControlFlow (after XlaIf conversion): " - << dump_graph::DumpGraphToFile("functionalize_ac", *graph_); + if (dump_graphs_) + LOG(INFO) << "FunctionalizeControlFlow (after XlaIf conversion): " + << dump_graph::DumpGraphToFile("functionalize_ac", *graph_, + library_); } return Status::OK(); } -xla::StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( - const std::vector& switch_nodes, - const std::vector& merge_nodes, Node* predicate) { - VLOG(2) << "Build if op for " << NodesToString(merge_nodes) << " with input " - << NodesToString(switch_nodes); +StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( + const CondArgNodes& cond_arg_nodes, const SwitchCluster& switch_cluster, + const std::vector& merge_nodes) { + VLOG(2) << "Build if op for " << switch_cluster.name; NodeDef if_def; // Create a new If node using the name of the merge node. - NodeDefBuilder builder(strings::StrCat(predicate->name(), "_If"), "XlaIf"); + NodeDefBuilder builder(switch_cluster.name, "XlaIf"); string branch[] = {"else_branch", "then_branch"}; for (int i = 0; i < 2; ++i) { static std::atomic sequence_num(0LL); @@ -854,7 +1159,8 @@ xla::StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( body_name.set_name( strings::StrCat("_functionalize_if_", branch[i], "_", id)); auto body = xla::MakeUnique(graph_->op_registry()); - TF_RETURN_IF_ERROR(ExtractBody(switch_nodes, merge_nodes, i, body.get())); + TF_RETURN_IF_ERROR(ExtractBody(cond_arg_nodes, switch_cluster.switches, + merge_nodes, i, body.get())); VLOG(3) << "Body " << branch[i] << ": " << DebugString(body.get()); FunctionDef body_fdef; TF_RETURN_IF_ERROR(GraphToFunctionDef(*body, body_name.name(), &body_fdef)); @@ -865,16 +1171,22 @@ xla::StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( // Build input type. std::vector inputs; DataTypeVector in_arg_types; - for (const Node* arg : switch_nodes) { - const Edge* in_edge; - TF_RETURN_IF_ERROR(arg->input_edge(0, &in_edge)); - if (in_edge->IsControlEdge()) { - builder.ControlInput(in_edge->src()->name()); - } else { - DataType dtype = arg->input_type(0); - inputs.emplace_back(NodeDefBuilder::NodeOut( - in_edge->src()->name(), in_edge->src_output(), dtype)); - in_arg_types.push_back(dtype); + for (auto& kv : cond_arg_nodes) { + bool inserted = false; + for (const Node* arg : kv.switches) { + const Edge* in_edge; + TF_RETURN_IF_ERROR(arg->input_edge(0, &in_edge)); + if (in_edge->IsControlEdge()) { + builder.ControlInput(in_edge->src()->name()); + } else { + if (!inserted) { + DataType dtype = arg->input_type(0); + inputs.emplace_back(NodeDefBuilder::NodeOut( + in_edge->src()->name(), in_edge->src_output(), dtype)); + in_arg_types.push_back(dtype); + inserted = true; + } + } } } builder.Attr("Tin", in_arg_types); @@ -888,10 +1200,11 @@ xla::StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( builder.Attr("Tout", out_type); builder.Attr("Tcond", DT_BOOL); - builder.Device(predicate->assigned_device_name()); + builder.Device(switch_cluster.predicate->assigned_device_name()); // Conditional should be the first input ... builder.Input( - NodeDefBuilder::NodeOut(predicate->name(), 0, predicate->output_type(0))); + NodeDefBuilder::NodeOut(switch_cluster.predicate->name(), 0, + switch_cluster.predicate->output_type(0))); // ... followed by the other inputs. builder.Input(inputs); @@ -900,7 +1213,8 @@ xla::StatusOr FunctionalizeCond::BuildAndAddXlaIfOp( return if_node; } -Status FunctionalizeCond::ExtractBody(const std::vector& switch_nodes, +Status FunctionalizeCond::ExtractBody(const CondArgNodes& cond_arg_nodes, + const std::vector& switches, const std::vector& merge_nodes, int input_edge, Graph* body) { VLOG(2) << "ExtractBody for " << NodesToString(merge_nodes) << " along edge " @@ -908,16 +1222,20 @@ Status FunctionalizeCond::ExtractBody(const std::vector& switch_nodes, std::vector squash_src_outputs(graph_->num_node_ids(), false); std::vector node_map(graph_->num_node_ids(), nullptr); int arg_count = 0; - for (const auto* arg : switch_nodes) { - DataType dtype = arg->input_type(0); - TF_ASSIGN_OR_RETURN(Node * arg_node, - BuildArgNode(body, dtype, arg_count++)); - node_map.at(arg->id()) = arg_node; - squash_src_outputs.at(arg->id()) = true; + for (auto& kv : cond_arg_nodes) { + Node* arg_node = nullptr; + for (const auto* arg : kv.switches) { + DataType dtype = arg->input_type(0); + if (arg_node == nullptr) { + TF_ASSIGN_OR_RETURN(arg_node, BuildArgNode(body, dtype, arg_count++)); + } + node_map.at(arg->id()) = arg_node; + squash_src_outputs.at(arg->id()) = true; + } } std::vector stack; - stack.reserve(switch_nodes.size()); + stack.reserve(merge_nodes.size()); for (int j = 0; j < merge_nodes.size(); ++j) { Node* node = merge_nodes[j]; TF_ASSIGN_OR_RETURN(node_map.at(node->id()), @@ -930,8 +1248,7 @@ Status FunctionalizeCond::ExtractBody(const std::vector& switch_nodes, node_map.at(in->id()) = body->CopyNode(in); } - if (std::find(switch_nodes.begin(), switch_nodes.end(), in) == - switch_nodes.end()) { + if (std::find(switches.begin(), switches.end(), in) == switches.end()) { body->AddEdge(node_map.at(in->id()), in_edge->src_output(), node_map.at(node->id()), 0); } else { @@ -946,18 +1263,25 @@ Status FunctionalizeCond::ExtractBody(const std::vector& switch_nodes, body); } -Status FunctionalizeCond::AddInputEdges(const std::vector& cond_args, +Status FunctionalizeCond::AddInputEdges(const CondArgNodes& cond_arg_nodes, Node* predicate, Node* if_node) { VLOG(3) << "AddInputEdges for " << if_node->name(); - int i = 0; - graph_->AddEdge(predicate, 0, if_node, i++); - for (const Node* arg : cond_args) { - const Edge* in_edge; - TF_RETURN_IF_ERROR(arg->input_edge(0, &in_edge)); - if (in_edge->IsControlEdge()) { - graph_->AddControlEdge(in_edge->src(), if_node); - } else { - graph_->AddEdge(in_edge->src(), in_edge->src_output(), if_node, i++); + int index = 0; + graph_->AddEdge(predicate, 0, if_node, index++); + for (auto& kv : cond_arg_nodes) { + bool inserted = false; + for (const Node* arg : kv.switches) { + const Edge* in_edge; + TF_RETURN_IF_ERROR(arg->input_edge(0, &in_edge)); + if (in_edge->IsControlEdge()) { + graph_->AddControlEdge(in_edge->src(), if_node); + } else { + if (!inserted) { + graph_->AddEdge(in_edge->src(), in_edge->src_output(), if_node, + index++); + inserted = true; + } + } } } return Status::OK(); @@ -988,36 +1312,39 @@ Status FunctionalizeCond::AddOutputEdges(const std::vector& outputs, return Status::OK(); } -Status FunctionalizeCond::ConvertCorrespondingMergeToXlaIf( - const std::vector& switch_nodes, - const std::vector& merge_nodes, Node* predicate) { - VLOG(1) << "ConvertMergeToXlaIf for " << NodesToString(switch_nodes) << " -> " +StatusOr FunctionalizeCond::ConvertToXlaIf( + const CondArgNodes& cond_arg_nodes, const SwitchCluster& switch_cluster, + const std::vector& merge_nodes) { + VLOG(1) << "ConvertToXlaIf for " << switch_cluster.ToString() << " -> " << NodesToString(merge_nodes); // Extract bodies and builds a If operator. - TF_ASSIGN_OR_RETURN(Node * if_node, - BuildAndAddXlaIfOp(switch_nodes, merge_nodes, predicate)); - TF_RETURN_IF_ERROR(AddInputEdges(switch_nodes, predicate, if_node)); + TF_ASSIGN_OR_RETURN( + Node * if_node, + BuildAndAddXlaIfOp(cond_arg_nodes, switch_cluster, merge_nodes)); + TF_RETURN_IF_ERROR( + AddInputEdges(cond_arg_nodes, switch_cluster.predicate, if_node)); TF_RETURN_IF_ERROR(AddOutputEdges(merge_nodes, if_node)); - return Status::OK(); + return if_node; } Status FunctionalizeCond::Functionalize(Graph* graph, FunctionLibraryDefinition* library) { VLOG(1) << "FunctionalizeCond::Functionalize"; - FunctionalizeCond fc(graph, library); + FunctionalizeCond fc(graph, library, /*dump_graphs=*/VLOG_IS_ON(2)); return fc.FunctionalizeInternal(); } } // namespace -// Transformation that converts Tensorflow's graph control flow constructs into +// Transformation that converts TensorFlow's graph control flow constructs into // functional equivalents. Status FunctionalizeControlFlow(Graph* graph, FunctionLibraryDefinition* library) { VLOG(2) << "FunctionalizeControlFlow (initial): " - << dump_graph::DumpGraphToFile("functionalize_initial", *graph); + << dump_graph::DumpGraphToFile("functionalize_initial", *graph, + library); // Note: BuildControlFlowInfo() requires that the graph's source node is // connected to all source nodes in the graph. Many graphs violate this // invariant. @@ -1029,7 +1356,8 @@ Status FunctionalizeControlFlow(Graph* graph, for (Node* node : graph->op_nodes()) { const ControlFlowInfo& cf = cf_info[node->id()]; - VLOG(2) << "node: " << node->name() << " frame_name: " << cf.frame_name + VLOG(2) << "node: " << node->name() << " (" << node->id() + << ") frame_name: " << cf.frame_name << " frame: " << (cf.frame ? cf.frame->name() : "---") << " parent_frame: " << (cf.parent_frame ? cf.parent_frame->name() : "---"); @@ -1097,7 +1425,8 @@ Status FunctionalizeControlFlow(Graph* graph, TF_RETURN_IF_ERROR(FunctionalizeCond::Functionalize(graph, library)); VLOG(2) << "FunctionalizeControlFlow (final): " - << dump_graph::DumpGraphToFile("functionalize_final", *graph); + << dump_graph::DumpGraphToFile("functionalize_final", *graph, + library); return Status::OK(); } diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc index 71f12a13339b9b5495631b8f9350579f6a0785a3..bc7276c3afd5060d6faeceb4d479416299ecc5da 100644 --- a/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc @@ -38,10 +38,11 @@ namespace { // Returns the names of the "then" and "else" functions for the XlaIf node in a // graph. -Status FindIfThenAndElse(const GraphDef& graph, NameAttrList* then_fn, - NameAttrList* else_fn) { +Status FindIfThenAndElse(const GraphDef& graph, string* op_name, + NameAttrList* then_fn, NameAttrList* else_fn) { for (const NodeDef& node : graph.node()) { if (node.op() == "XlaIf") { + *op_name = node.name(); const NameAttrList* result; TF_RETURN_IF_ERROR(GetNodeAttr(node, "then_branch", &result)); *then_fn = *result; @@ -96,9 +97,10 @@ TEST(FunctionalizeControlFlow, Conditional) { GraphDef graph_def; graph.ToGraphDef(&graph_def); + string op_name; NameAttrList then_fn; NameAttrList else_fn; - TF_EXPECT_OK(FindIfThenAndElse(graph_def, &then_fn, &else_fn)); + TF_EXPECT_OK(FindIfThenAndElse(graph_def, &op_name, &then_fn, &else_fn)); InstantiationResultForTest else_result; TF_EXPECT_OK( InstantiateFunctionForTest(else_fn.name(), library, &else_result)); @@ -109,7 +111,7 @@ TEST(FunctionalizeControlFlow, Conditional) { auto y = ops::Placeholder(scope.WithOpName("y"), DT_INT32); auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32); auto less = ops::Less(scope.WithOpName("cond/Less"), y, x); - auto if_op = ops::XlaIf(scope.WithOpName("cond/Less_If"), less, + auto if_op = ops::XlaIf(scope.WithOpName(op_name), less, std::initializer_list{less, y, x}, then_fn, else_fn, {DT_INT32}); GraphDef expected; diff --git a/tensorflow/compiler/tf2xla/g3doc/cpu_supported_ops.md b/tensorflow/compiler/tf2xla/g3doc/cpu_supported_ops.md index 82b3b46a2f1e97001d1e0c6b993ec243170bc7d8..91351421bcacd26c41b5c9f98ea833730e4aef30 100644 --- a/tensorflow/compiler/tf2xla/g3doc/cpu_supported_ops.md +++ b/tensorflow/compiler/tf2xla/g3doc/cpu_supported_ops.md @@ -6,6 +6,9 @@ Operator | Type Constraint `Acosh` | `T={complex64,double,float}` `Add` | `T={complex64,double,float,int32,int64}` `AddN` | `T={complex64,double,float,int32,int64,uint32,uint64}` +`AdjustContrastv2` | +`AdjustHue` | +`AdjustSaturation` | `All` | `Tidx={int32,int64}` `Angle` | `Tout={double,float}`
`T={complex64}` `Any` | `Tidx={int32,int64}` @@ -34,7 +37,7 @@ Operator | Type Constraint `BroadcastGradientArgs` | `T={int32,int64}` `Cast` | `DstT={bool,complex64,double,float,int32,int64,uint32,uint64}`
`SrcT={bool,complex64,double,float,int32,int64,uint32,uint64}` `Ceil` | `T={double,float}` -`Cholesky` | `T={complex64,double,float}` +`Cholesky` | `T={double,float}` `Complex` | `Tout={complex64}`
`T={double,float}` `ComplexAbs` | `Tout={double,float}`
`T={complex64}` `Concat` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` @@ -68,7 +71,11 @@ Operator | Type Constraint `Exp` | `T={complex64,double,float}` `ExpandDims` | `Tdim={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Expm1` | `T={complex64,double,float}` -`Fill` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`ExtractImagePatches` | `T={double,float,int32,int64,uint32,uint64}` +`FFT` | +`FFT2D` | +`FFT3D` | +`Fill` | `index_type={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Floor` | `T={double,float}` `FloorDiv` | `T={complex64,double,float,int32,int64}` `FloorMod` | `T={double,float,int32,int64}` @@ -80,6 +87,13 @@ Operator | Type Constraint `GatherV2` | `Taxis={int32,int64}`
`Tindices={int32,int64}`
`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}` `Greater` | `T={double,float,int32,int64,uint32,uint64}` `GreaterEqual` | `T={double,float,int32,int64,uint32,uint64}` +`HSVToRGB` | `T={double,float}` +`IFFT` | +`IFFT2D` | +`IFFT3D` | +`IRFFT` | +`IRFFT2D` | +`IRFFT3D` | `Identity` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `IdentityN` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Imag` | `Tout={double,float}`
`T={complex64}` @@ -105,11 +119,14 @@ Operator | Type Constraint `MatMul` | `T={complex64,double,float}` `MatrixDiag` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `MatrixDiagPart` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`MatrixTriangularSolve` | `T={complex64,double,float}` `Max` | `Tidx={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` `MaxPool` | `T={double,float,int32,int64}` `MaxPool3D` | `T={float}` `MaxPool3DGrad` | `TInput={float}`
`T={float}` `MaxPoolGrad` | `T={double,float,int32,int64,uint32,uint64}` +`MaxPoolGradV2` | `T={double,float,int32,int64,uint32,uint64}` +`MaxPoolV2` | `T={double,float,int32,int64}` `Maximum` | `T={double,float,int32,int64}` `Mean` | `Tidx={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` `Min` | `Tidx={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` @@ -131,6 +148,10 @@ Operator | Type Constraint `PreventGradient` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Prod` | `Tidx={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` `QuantizeAndDequantizeV2` | `T={double,float}` +`RFFT` | +`RFFT2D` | +`RFFT3D` | +`RGBToHSV` | `T={double,float}` `RandomStandardNormal` | `dtype={float}` `RandomUniform` | `T={int32,int64}`
`dtype={double,float}` `RandomUniformInt` | `T={int32,int64}`
`Tout={int32,int64}` @@ -146,6 +167,8 @@ Operator | Type Constraint `Relu6Grad` | `T={double,float,int32,int64,uint32,uint64}` `ReluGrad` | `T={double,float,int32,int64,uint32,uint64}` `Reshape` | `Tshape={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`ResizeBilinear` | `T={double,float,int32,int64}` +`ResizeBilinearGrad` | `T={double,float}` `ResourceApplyAdagrad` | `T={double,float}` `ResourceApplyAdam` | `T={double,float}` `ResourceApplyFtrl` | `T={double,float}` @@ -156,6 +179,7 @@ Operator | Type Constraint `ResourceGather` | `Tindices={int32,int64}`
`dtype={complex64,double,float,int32,int64,uint32,uint64}` `ResourceStridedSliceAssign` | `Index={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Reverse` | `T={bool,complex64,double,float,int32,int64}` +`ReverseSequence` | `Tlen={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `ReverseV2` | `T={bool,complex64,double,float,int32,int64}`
`Tidx={int32,int64}` `RightShift` | `T={int32,int64,uint32,uint64}` `Rint` | `T={double,float}` diff --git a/tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md b/tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md index d4b7621ad2858fe17e93d292dd807e4f7c1c336b..b9bdb829d773825005a8921f48d28b6892d8f0cd 100644 --- a/tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md +++ b/tensorflow/compiler/tf2xla/g3doc/gpu_supported_ops.md @@ -6,6 +6,9 @@ Operator | Type Constraint `Acosh` | `T={complex64,double,float}` `Add` | `T={complex64,double,float,int32,int64}` `AddN` | `T={complex64,double,float,int32,int64,uint32,uint64}` +`AdjustContrastv2` | +`AdjustHue` | +`AdjustSaturation` | `All` | `Tidx={int32,int64}` `Angle` | `Tout={double,float}`
`T={complex64}` `Any` | `Tidx={int32,int64}` @@ -34,7 +37,7 @@ Operator | Type Constraint `BroadcastGradientArgs` | `T={int32,int64}` `Cast` | `DstT={bool,complex64,double,float,int32,int64,uint32,uint64}`
`SrcT={bool,complex64,double,float,int32,int64,uint32,uint64}` `Ceil` | `T={double,float}` -`Cholesky` | `T={complex64,double,float}` +`Cholesky` | `T={double,float}` `Complex` | `Tout={complex64}`
`T={double,float}` `ComplexAbs` | `Tout={double,float}`
`T={complex64}` `Concat` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` @@ -68,7 +71,11 @@ Operator | Type Constraint `Exp` | `T={complex64,double,float}` `ExpandDims` | `Tdim={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Expm1` | `T={complex64,double,float}` -`Fill` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`ExtractImagePatches` | `T={double,float,int32,int64,uint32,uint64}` +`FFT` | +`FFT2D` | +`FFT3D` | +`Fill` | `index_type={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Floor` | `T={double,float}` `FloorDiv` | `T={complex64,double,float,int32,int64}` `FloorMod` | `T={double,float,int32,int64}` @@ -80,6 +87,13 @@ Operator | Type Constraint `GatherV2` | `Taxis={int32,int64}`
`Tindices={int32,int64}`
`Tparams={bool,complex64,double,float,int32,int64,uint32,uint64}` `Greater` | `T={double,float,int32,int64,uint32,uint64}` `GreaterEqual` | `T={double,float,int32,int64,uint32,uint64}` +`HSVToRGB` | `T={double,float}` +`IFFT` | +`IFFT2D` | +`IFFT3D` | +`IRFFT` | +`IRFFT2D` | +`IRFFT3D` | `Identity` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `IdentityN` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Imag` | `Tout={double,float}`
`T={complex64}` @@ -105,11 +119,14 @@ Operator | Type Constraint `MatMul` | `T={complex64,double,float}` `MatrixDiag` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `MatrixDiagPart` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`MatrixTriangularSolve` | `T={complex64,double,float}` `Max` | `Tidx={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` `MaxPool` | `T={double,float,int32,int64}` `MaxPool3D` | `T={float}` `MaxPool3DGrad` | `TInput={float}`
`T={float}` `MaxPoolGrad` | `T={double,float,int32,int64,uint32,uint64}` +`MaxPoolGradV2` | `T={double,float,int32,int64,uint32,uint64}` +`MaxPoolV2` | `T={double,float,int32,int64}` `Maximum` | `T={double,float,int32,int64}` `Mean` | `Tidx={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` `Min` | `Tidx={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` @@ -131,6 +148,10 @@ Operator | Type Constraint `PreventGradient` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Prod` | `Tidx={int32,int64}`
`T={complex64,double,float,int32,int64,uint32,uint64}` `QuantizeAndDequantizeV2` | `T={double,float}` +`RFFT` | +`RFFT2D` | +`RFFT3D` | +`RGBToHSV` | `T={double,float}` `Range` | `Tidx={double,float,int32,int64}` `Rank` | `T={bool,complex64,double,float,int32,int64,uint32,uint64}` `ReadVariableOp` | `dtype={bool,complex64,double,float,int32,int64,uint32,uint64}` @@ -143,6 +164,8 @@ Operator | Type Constraint `Relu6Grad` | `T={double,float,int32,int64,uint32,uint64}` `ReluGrad` | `T={double,float,int32,int64,uint32,uint64}` `Reshape` | `Tshape={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` +`ResizeBilinear` | `T={double,float,int32,int64}` +`ResizeBilinearGrad` | `T={double,float}` `ResourceApplyAdagrad` | `T={double,float}` `ResourceApplyAdam` | `T={double,float}` `ResourceApplyFtrl` | `T={double,float}` @@ -153,6 +176,7 @@ Operator | Type Constraint `ResourceGather` | `Tindices={int32,int64}`
`dtype={complex64,double,float,int32,int64,uint32,uint64}` `ResourceStridedSliceAssign` | `Index={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `Reverse` | `T={bool,complex64,double,float,int32,int64}` +`ReverseSequence` | `Tlen={int32,int64}`
`T={bool,complex64,double,float,int32,int64,uint32,uint64}` `ReverseV2` | `T={bool,complex64,double,float,int32,int64}`
`Tidx={int32,int64}` `RightShift` | `T={int32,int64,uint32,uint64}` `Rint` | `T={double,float}` diff --git a/tensorflow/compiler/tf2xla/graph_compiler.cc b/tensorflow/compiler/tf2xla/graph_compiler.cc index 8062f0c03ca60e88bd5c021092dceb105232219f..058a1f2621c64a735bd9d9c9d0ae007f93aa4dea 100644 --- a/tensorflow/compiler/tf2xla/graph_compiler.cc +++ b/tensorflow/compiler/tf2xla/graph_compiler.cc @@ -34,6 +34,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/graph_optimizer.h" #include "tensorflow/core/framework/attr_value_util.h" +#include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph_constructor.h" @@ -59,9 +60,7 @@ Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph, for (int i = 0; i < args->size(); ++i) { XlaCompiler::Argument& arg = (*args)[i]; arg.type = ctx->input_type(i); - - TF_RETURN_IF_ERROR( - TensorShapeToXLAShape(arg.type, ctx->InputShape(i), &arg.shape)); + arg.shape = ctx->InputShape(i); if (arg.type == DT_RESOURCE) { return errors::InvalidArgument( @@ -135,7 +134,7 @@ Status GraphCompiler::Compile() { TF_RET_CHECK(src->id() < output_registry.size()); const NodeOutputs& src_outputs = output_registry[src->id()]; - tensor_inputs_[e->dst_input()] = src_outputs[e->src_output()]; + tensor_inputs_.at(e->dst_input()) = src_outputs.at(e->src_output()); } OpKernelContext op_context(¶ms, n->num_outputs()); @@ -144,7 +143,9 @@ Status GraphCompiler::Compile() { } else { device_->Compute(CHECK_NOTNULL(params.op_kernel), &op_context); Status s = op_context.status(); - TF_RETURN_IF_ERROR(s); + if (!s.ok()) { + return AttachDef(s, n->def()); + } } // Set up outputs. Also check if outputs from the previous computation is diff --git a/tensorflow/compiler/tf2xla/graph_compiler.h b/tensorflow/compiler/tf2xla/graph_compiler.h index ba00160b6d78c1e55cc2e053cd5285344e0179fb..127562eb23d775f17179cc9ee968ec2255cf3a14 100644 --- a/tensorflow/compiler/tf2xla/graph_compiler.h +++ b/tensorflow/compiler/tf2xla/graph_compiler.h @@ -70,7 +70,7 @@ class GraphCompiler { private: // Partially sets params. This partially set params can be reused - // across multple nodes visit. + // across multiple nodes visit. void PartiallySetupParams(OpKernelContext::Params* params); // Tests if a node is a functional node. A functional node represents a diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index 092d852fe3f11d715b53f9ffe750ea69b27cf839..d2fa933cf9c085f92b2f442827a94d72938e4bb2 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -4,6 +4,7 @@ package( default_visibility = ["//tensorflow/compiler/tf2xla:internal"], ) +load("//tensorflow:tensorflow.bzl", "tf_copts") load("//tensorflow:tensorflow.bzl", "tf_kernel_library") tf_kernel_library( @@ -30,6 +31,9 @@ tf_kernel_library( "diag_op.cc", "dynamic_stitch_op.cc", "elu_op.cc", + "extract_image_patches_op.cc", + "fake_quantize_ops.cc", + "fft_ops.cc", "fill_op.cc", "function_ops.cc", "gather_op.cc", @@ -41,6 +45,9 @@ tf_kernel_library( "l2loss_op.cc", "lrn_ops.cc", "matmul_op.cc", + "matrix_band_part_op.cc", + "matrix_set_diag_op.cc", + "matrix_triangular_solve_op.cc", "mirror_pad_op.cc", "no_op.cc", "one_hot_op.cc", @@ -56,7 +63,9 @@ tf_kernel_library( "reshape_op.cc", "retval_op.cc", "reverse_op.cc", + "reverse_sequence_op.cc", "scan_ops.cc", + "scatter_nd_op.cc", "segment_reduction_ops.cc", "select_op.cc", "sendrecv_ops.cc", @@ -80,7 +89,6 @@ tf_kernel_library( "variable_ops.cc", ], hdrs = [ - "gather_op.h", "index_ops.h", "shape_util.h", ], @@ -90,11 +98,15 @@ tf_kernel_library( "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/lib:batch_dot", "//tensorflow/compiler/tf2xla/lib:cholesky", + "//tensorflow/compiler/tf2xla/lib:scatter", + "//tensorflow/compiler/tf2xla/lib:triangular_solve", "//tensorflow/compiler/tf2xla/lib:util", + "//tensorflow/compiler/tf2xla/lib:while_loop", "//tensorflow/compiler/tf2xla/ops:sendrecv_ops", "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:client_library", @@ -105,6 +117,7 @@ tf_kernel_library( "//tensorflow/core:lib", "//tensorflow/core:linalg_ops_op_lib", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:spectral_ops_op_lib", "//tensorflow/core:stateless_random_ops_op_lib", "//tensorflow/core/kernels:bounds_check", "//tensorflow/core/kernels:concat_lib", @@ -165,6 +178,7 @@ tf_kernel_library( cc_library( name = "index_ops_kernel_argmax_float_1d", srcs = ["index_ops_kernel_argmax_float_1d.cc"], + copts = tf_copts(), visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/xla/service/cpu:custom_call_target_registry", @@ -177,6 +191,7 @@ cc_library( cc_library( name = "index_ops_kernel_argmax_float_2d", srcs = ["index_ops_kernel_argmax_float_2d.cc"], + copts = tf_copts(), visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/xla/service/cpu:custom_call_target_registry", diff --git a/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc index a015b8e0e8949f8aaa03a78b0f88b7ea8d6aaa1c..b0ba25b9983c3a9af26728ce4b1c263c844327db 100644 --- a/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc @@ -28,8 +28,9 @@ class BatchMatMulOp : public XlaOpKernel { } void Compile(XlaOpKernelContext* ctx) override { - auto result = - BatchDot(ctx->builder(), ctx->Input(0), ctx->Input(1), adj_x_, adj_y_); + auto result = BatchDot(ctx->builder(), ctx->Input(0), ctx->Input(1), + /*transpose_x=*/adj_x_, /*transpose_y=*/adj_y_, + /*conjugate_x=*/adj_x_, /*conjugate_y=*/adj_y_); OP_REQUIRES_OK(ctx, result.status()); ctx->SetOutput(0, result.ValueOrDie()); } diff --git a/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc b/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc index 87d858f763560be454c162e0cf40307c68217663..fe6651793dc763d13f4a4b0ac294ec3ecf64af8f 100644 --- a/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc @@ -33,7 +33,7 @@ class CholeskyOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("Cholesky"), CholeskyOp); +REGISTER_XLA_OP(Name("Cholesky").TypeConstraint("T", kFloatTypes), CholeskyOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc index da9be687326596cdee71730240e484a04c53db13..81cea6d376d02c956a5257c5475fe5c10b83deb9 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc @@ -209,15 +209,14 @@ class ConvOp : public XlaOpKernel { num_dims(), " dimensions")); OP_REQUIRES( ctx, dilations_[batch_dim] == 1 && dilations_[feature_dim] == 1, - errors::Unimplemented("Current implementation does not yet support " + errors::Unimplemented("Current implementation does not support " "dilations in the batch and depth dimensions.")); for (int i = 0; i < num_spatial_dims_; ++i) { int input_dim = GetTensorSpatialDimIndex(num_dims(), data_format_, i); - OP_REQUIRES( - ctx, dilations_[input_dim] == 1, - errors::Unimplemented("Current implementation does not yet support " - "dilations in the ", - i, "th spatial dimension.")); + OP_REQUIRES(ctx, dilations_[input_dim] >= 1, + errors::Unimplemented("Dilation values must be positive; ", i, + "th spatial dimension had dilation ", + dilations_[input_dim])); } const TensorShape input_shape = ctx->InputShape(0); @@ -248,32 +247,46 @@ class ConvOp : public XlaOpKernel { xla::ComputationBuilder* b = ctx->builder(); xla::ComputationDataHandle filter = ctx->Input(1); + TensorShape expanded_filter_shape = filter_shape; if (depthwise_) { filter = ExpandFilterForDepthwiseConvolution( filter_shape, ctx->input_type(0), filter, b); + expanded_filter_shape = + ExpandedFilterShapeForDepthwiseConvolution(filter_shape); } xla::ConvolutionDimensionNumbers dims; - std::vector window_strides; + std::vector window_strides(num_spatial_dims_); + std::vector lhs_dilation(num_spatial_dims_, 1); + std::vector rhs_dilation(num_spatial_dims_); + std::vector> padding(num_spatial_dims_); + dims.set_input_batch_dimension(batch_dim); dims.set_output_batch_dimension(batch_dim); dims.set_input_feature_dimension(feature_dim); dims.set_output_feature_dimension(feature_dim); + dims.set_kernel_input_feature_dimension(num_spatial_dims_); + dims.set_kernel_output_feature_dimension(num_spatial_dims_ + 1); + for (int i = 0; i < num_spatial_dims_; ++i) { const int64 dim = GetTensorSpatialDimIndex(num_dims(), data_format_, i); dims.add_input_spatial_dimensions(dim); dims.add_kernel_spatial_dimensions(i); dims.add_output_spatial_dimensions(dim); - window_strides.push_back(strides_.at(dim)); + window_strides[i] = strides_.at(dim); + rhs_dilation[i] = dilations_.at(dim); + + int64 unused_output_size; + OP_REQUIRES_OK( + ctx, GetWindowedOutputSizeVerboseV2( + input_shape.dim_size(dim), expanded_filter_shape.dim_size(i), + rhs_dilation[i], window_strides[i], padding_, + &unused_output_size, &padding[i].first, &padding[i].second)); } - dims.set_kernel_input_feature_dimension(num_spatial_dims_); - dims.set_kernel_output_feature_dimension(num_spatial_dims_ + 1); - - xla::Padding xla_padding = - (padding_ == VALID) ? xla::Padding::kValid : xla::Padding::kSame; - xla::ComputationDataHandle conv = b->ConvWithGeneralDimensions( - ctx->Input(0), filter, window_strides, xla_padding, dims); + xla::ComputationDataHandle conv = + b->ConvGeneralDilated(ctx->Input(0), filter, window_strides, padding, + lhs_dilation, rhs_dilation, dims); ctx->SetOutput(0, conv); } @@ -347,15 +360,14 @@ class ConvBackpropInputOp : public XlaOpKernel { num_dims(), " dimensions")); OP_REQUIRES( ctx, dilations_[batch_dim] == 1 && dilations_[feature_dim] == 1, - errors::Unimplemented("Current implementation does not yet support " + errors::Unimplemented("Current implementation does not support " "dilations in the batch and depth dimensions.")); for (int i = 0; i < num_spatial_dims_; ++i) { int input_dim = GetTensorSpatialDimIndex(num_dims(), data_format_, i); - OP_REQUIRES( - ctx, dilations_[input_dim] == 1, - errors::Unimplemented("Current implementation does not yet support " - "dilations in the ", - i, "th spatial dimension.")); + OP_REQUIRES(ctx, dilations_[input_dim] >= 1, + errors::Unimplemented("Dilation values must be positive; ", i, + "th spatial dimension had dilation ", + dilations_[input_dim])); } TensorShape input_shape; @@ -369,10 +381,11 @@ class ConvBackpropInputOp : public XlaOpKernel { : filter_shape; // Reuse dimension computation logic from conv_grad_ops.cc. ConvBackpropDimensions dims; - OP_REQUIRES_OK(ctx, ConvBackpropComputeDimensions( - type_string(), num_spatial_dims_, input_shape, - expanded_filter_shape, out_backprop_shape, strides_, - padding_, data_format_, &dims)); + OP_REQUIRES_OK(ctx, + ConvBackpropComputeDimensionsV2( + type_string(), num_spatial_dims_, input_shape, + expanded_filter_shape, out_backprop_shape, dilations_, + strides_, padding_, data_format_, &dims)); xla::ComputationBuilder* b = ctx->builder(); auto filter = ctx->Input(1); @@ -396,6 +409,7 @@ class ConvBackpropInputOp : public XlaOpKernel { std::vector kernel_spatial_dims(num_spatial_dims_); std::vector> padding(num_spatial_dims_); std::vector lhs_dilation(num_spatial_dims_); + std::vector rhs_dilation(num_spatial_dims_); std::vector ones(num_spatial_dims_, 1); for (int i = 0; i < num_spatial_dims_; ++i) { int64 dim = GetTensorSpatialDimIndex(num_dims(), data_format_, i); @@ -407,6 +421,7 @@ class ConvBackpropInputOp : public XlaOpKernel { padding[i] = {dims.spatial_dims[i].pad_before, dims.spatial_dims[i].pad_after}; lhs_dilation[i] = dims.spatial_dims[i].stride; + rhs_dilation[i] = dilations_[dim]; } // If this is a depthwise convolution, expand the filter. @@ -423,7 +438,7 @@ class ConvBackpropInputOp : public XlaOpKernel { // = gradients (with padding and dilation) mirrored_weights xla::ComputationDataHandle in_backprop = b->ConvGeneralDilated( out_backprop, mirrored_weights, /*window_strides=*/ones, padding, - lhs_dilation, /*rhs_dilation=*/ones, dnums); + lhs_dilation, rhs_dilation, dnums); ctx->SetOutput(0, in_backprop); } @@ -500,15 +515,14 @@ class ConvBackpropFilterOp : public XlaOpKernel { num_dims(), " dimensions")); OP_REQUIRES( ctx, dilations_[n_dim] == 1 && dilations_[c_dim] == 1, - errors::Unimplemented("Current implementation does not yet support " + errors::Unimplemented("Current implementation does not support " "dilations in the batch and depth dimensions.")); for (int i = 0; i < num_spatial_dims_; ++i) { int input_dim = GetTensorSpatialDimIndex(num_dims(), data_format_, i); - OP_REQUIRES( - ctx, dilations_[input_dim] == 1, - errors::Unimplemented("Current implementation does not yet support " - "dilations in the ", - i, "th spatial dimension.")); + OP_REQUIRES(ctx, dilations_[input_dim] >= 1, + errors::Unimplemented("Dilation values must be positive; ", i, + "th spatial dimension had dilation ", + dilations_[input_dim])); } const TensorShape activations_shape = ctx->InputShape(0); @@ -522,10 +536,11 @@ class ConvBackpropFilterOp : public XlaOpKernel { // Reuse dimension computation logic from conv_grad_ops.cc. ConvBackpropDimensions dims; - OP_REQUIRES_OK(ctx, ConvBackpropComputeDimensions( - type_string(), num_spatial_dims_, activations_shape, - expanded_filter_shape, out_backprop_shape, strides_, - padding_, data_format_, &dims)); + OP_REQUIRES_OK(ctx, + ConvBackpropComputeDimensionsV2( + type_string(), num_spatial_dims_, activations_shape, + expanded_filter_shape, out_backprop_shape, dilations_, + strides_, padding_, data_format_, &dims)); xla::ComputationBuilder* b = ctx->builder(); xla::ComputationDataHandle activations = ctx->Input(0); @@ -555,6 +570,7 @@ class ConvBackpropFilterOp : public XlaOpKernel { std::vector> padding(num_spatial_dims_); std::vector rhs_dilation(num_spatial_dims_); + std::vector window_strides(num_spatial_dims_); std::vector ones(num_spatial_dims_, 1); // Tensorflow filter shape is [ H, W, ..., inC, outC ]. @@ -574,8 +590,9 @@ class ConvBackpropFilterOp : public XlaOpKernel { // The padded_in_rows should be such that when we convolve this with the // expanded_out_rows as a filter, we should get filter_rows back. // - const int padded_in_size = dims.spatial_dims[i].expanded_output_size + - dims.spatial_dims[i].filter_size - 1; + const int64 padded_in_size = + dims.spatial_dims[i].expanded_output_size + + (dims.spatial_dims[i].filter_size - 1) * dilations_[dim]; // However it can be smaller than input_rows: in this // case it means some of the inputs are not used. @@ -591,8 +608,7 @@ class ConvBackpropFilterOp : public XlaOpKernel { // and input "C" is not used at all. // // We apply negative padding in this case. - const int total_pad_in_size = - padded_in_size - dims.spatial_dims[i].input_size; + const int64 pad_total = padded_in_size - dims.spatial_dims[i].input_size; // + For the VALID padding, we don't pad anything on the top/left side // and pad the bottom/right side with the remaining space. @@ -602,13 +618,12 @@ class ConvBackpropFilterOp : public XlaOpKernel { // In addition, if the padded input size is smaller than the input size, // we need to ignore some training elements of the input. We do this by // applying negative padding on the right/bottom. - const int before_pad_in_size = - (total_pad_in_size > 0 && padding_ == Padding::SAME) - ? total_pad_in_size / 2 - : 0; + const int64 pad_before = + padding_ == Padding::SAME ? std::max(pad_total / 2, 0) : 0; - padding[i] = {before_pad_in_size, total_pad_in_size - before_pad_in_size}; + padding[i] = {pad_before, pad_total - pad_before}; rhs_dilation[i] = dims.spatial_dims[i].stride; + window_strides[i] = dilations_[dim]; } // Besides padding the input, we will also expand output_rows to @@ -620,8 +635,7 @@ class ConvBackpropFilterOp : public XlaOpKernel { // This is done by specifying the window dilation factors in the // convolution HLO below. auto filter_backprop = - b->ConvGeneralDilated(activations, gradients, - /*window_strides=*/ones, padding, + b->ConvGeneralDilated(activations, gradients, window_strides, padding, /*lhs_dilation=*/ones, rhs_dilation, dnums); if (depthwise_) { diff --git a/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc b/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..b2970eae20a3fb71f06619f476a49d41b22bca56 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc @@ -0,0 +1,169 @@ +/* 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/compiler/tf2xla/type_util.h" +#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/util/tensor_format.h" + +namespace tensorflow { + +namespace { + +class ExtractImagePatchesOp : public XlaOpKernel { + public: + explicit ExtractImagePatchesOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("ksizes", &ksizes_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &strides_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("rates", &dilations_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + const TensorFormat data_format = FORMAT_NHWC; + const int num_dims = ksizes_.size(); + + OP_REQUIRES( + ctx, num_dims >= 3, + errors::InvalidArgument("Kernel size must have at least 3 dimensions")); + const int num_spatial_dims = num_dims - 2; + + OP_REQUIRES(ctx, strides_.size() == num_dims, + errors::InvalidArgument("Sliding window strides field must " + "specify ", + num_dims, " dimensions")); + OP_REQUIRES(ctx, dilations_.size() == num_dims, + errors::InvalidArgument("Dilations field must " + "specify ", + num_dims, " dimensions")); + + int batch_dim = GetTensorBatchDimIndex(num_dims, data_format); + int feature_dim = GetTensorFeatureDimIndex(num_dims, data_format); + OP_REQUIRES( + ctx, ksizes_[batch_dim] == 1 && ksizes_[feature_dim] == 1, + errors::Unimplemented("Current implementation does not yet support " + "kernel sizes > 1 in the batch and depth " + "dimensions.")); + OP_REQUIRES( + ctx, strides_[batch_dim] == 1 && strides_[feature_dim] == 1, + errors::Unimplemented("Current implementation does not yet support " + "strides in the batch and depth dimensions.")); + OP_REQUIRES( + ctx, dilations_[batch_dim] == 1 && dilations_[feature_dim] == 1, + errors::Unimplemented("Current implementation does not support " + "dilations in the batch and depth dimensions.")); + + for (int i = 0; i < num_spatial_dims; ++i) { + int input_dim = GetTensorSpatialDimIndex(num_dims, data_format, i); + OP_REQUIRES( + ctx, ksizes_[input_dim] >= 0, + errors::Unimplemented("Kernel size values must be non-negative; ", i, + "th spatial dimension had dilation ", + dilations_[input_dim])); + OP_REQUIRES(ctx, strides_[input_dim] >= 1, + errors::Unimplemented("Stride values must be positive; ", i, + "th spatial dimension had dilation ", + dilations_[input_dim])); + OP_REQUIRES(ctx, dilations_[input_dim] >= 1, + errors::Unimplemented("Dilation values must be positive; ", i, + "th spatial dimension had dilation ", + dilations_[input_dim])); + } + + xla::PrimitiveType type; + OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(ctx->input_type(0), &type)); + + const TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES( + ctx, input_shape.dims() == num_dims, + errors::InvalidArgument("input must be ", num_dims, "-dimensional", + input_shape.DebugString())); + const int64 depth = input_shape.dim_size(feature_dim); + + xla::ComputationBuilder* builder = ctx->builder(); + + // The following code is equivalent to: + // eye = np.eye(kH * kW * D).reshape([kH, kW, D, kH * kW * kD]) + int64 kernel_size = 1; + std::vector lhs_shape(num_dims, 1); + for (int i = 0; i < num_spatial_dims; ++i) { + int input_dim = GetTensorSpatialDimIndex(num_dims, data_format, i); + lhs_shape[i] = ksizes_[input_dim]; + kernel_size *= ksizes_[input_dim]; + } + lhs_shape[num_spatial_dims] = depth; + lhs_shape[num_spatial_dims + 1] = 1; + + // Builds an identity matrix as a broadcast equality of iotas. + // iota = np.arange(np.prod(ksize), depth) + // filter = np.equal(np.reshape(iota, [-1, 1]), iota).astype(np.float32) + xla::ComputationDataHandle iota; + TF_CHECK_OK(XlaHelpers::Iota(builder, DataType::DT_INT32, + kernel_size * depth, &iota)); + + auto lhs = builder->Reshape(iota, lhs_shape); + auto filter = builder->ConvertElementType( + builder->Eq(lhs, iota, {num_spatial_dims + 1}), type); + + xla::ConvolutionDimensionNumbers dims; + std::vector window_strides(num_spatial_dims); + std::vector lhs_dilation(num_spatial_dims, 1); + std::vector rhs_dilation(num_spatial_dims); + std::vector> padding(num_spatial_dims); + + dims.set_input_batch_dimension(batch_dim); + dims.set_output_batch_dimension(batch_dim); + dims.set_input_feature_dimension(feature_dim); + dims.set_output_feature_dimension(feature_dim); + dims.set_kernel_input_feature_dimension(num_spatial_dims); + dims.set_kernel_output_feature_dimension(num_spatial_dims + 1); + + for (int i = 0; i < num_spatial_dims; ++i) { + const int64 dim = GetTensorSpatialDimIndex(num_dims, data_format, i); + dims.add_input_spatial_dimensions(dim); + dims.add_kernel_spatial_dimensions(i); + dims.add_output_spatial_dimensions(dim); + window_strides[i] = strides_.at(dim); + rhs_dilation[i] = dilations_.at(dim); + + int64 unused_output_size; + OP_REQUIRES_OK( + ctx, GetWindowedOutputSizeVerboseV2( + input_shape.dim_size(dim), ksizes_[dim], rhs_dilation[i], + window_strides[i], padding_, &unused_output_size, + &padding[i].first, &padding[i].second)); + } + + xla::ComputationDataHandle conv = + builder->ConvGeneralDilated(ctx->Input(0), filter, window_strides, + padding, lhs_dilation, rhs_dilation, dims); + ctx->SetOutput(0, conv); + } + + protected: + std::vector ksizes_; + std::vector dilations_; + std::vector strides_; + Padding padding_; + + private: + TF_DISALLOW_COPY_AND_ASSIGN(ExtractImagePatchesOp); +}; + +REGISTER_XLA_OP(Name("ExtractImagePatches"), ExtractImagePatchesOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..453a32c494b42e9922bc35fc526f3306530054fd --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc @@ -0,0 +1,289 @@ +/* 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/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/lib/arithmetic.h" +#include "tensorflow/core/platform/macros.h" + +namespace tensorflow { +namespace { + +// Gymnastics with nudged zero point is to ensure that the real zero maps to +// an integer, which is required for e.g. zero-padding in convolutional layers. +void CpuNudge(const float min, const float max, const float quant_min, + const float quant_max, float* nudged_min, float* nudged_max, + float* scale) { + *scale = (max - min) / (quant_max - quant_min); + + const float zero_point_from_min = quant_min - min / *scale; + float nudged_zero_point; + if (zero_point_from_min <= quant_min) { + nudged_zero_point = quant_min; + } else if (zero_point_from_min >= quant_max) { + nudged_zero_point = quant_max; + } else { + nudged_zero_point = std::round(zero_point_from_min); + } + + *nudged_min = (quant_min - nudged_zero_point) * (*scale); + *nudged_max = (quant_max - nudged_zero_point) * (*scale); +} + +// An XLA version of CpuNudge(). +void XlaNudge(xla::ComputationBuilder* b, const DataType data_type, + const xla::ComputationDataHandle& min, + const xla::ComputationDataHandle& max, + const float quant_min_value, const float quant_max_value, + xla::ComputationDataHandle* nudged_min, + xla::ComputationDataHandle* nudged_max, + xla::ComputationDataHandle* scale) { + *scale = b->Div(b->Sub(max, min), + XlaHelpers::FloatLiteral(b, data_type, + quant_max_value - quant_min_value)); + xla::ComputationDataHandle quant_min = + XlaHelpers::FloatLiteral(b, data_type, quant_min_value); + xla::ComputationDataHandle zero_point_from_min = + b->Sub(quant_min, b->Div(min, *scale)); + xla::ComputationDataHandle quant_max = + XlaHelpers::FloatLiteral(b, data_type, quant_max_value); + xla::ComputationDataHandle nudged_zero_point = + b->Select(b->Le(zero_point_from_min, quant_min), quant_min, + b->Select(b->Ge(zero_point_from_min, quant_max), quant_max, + b->Round(zero_point_from_min))); + *nudged_min = b->Mul(b->Sub(quant_min, nudged_zero_point), *scale); + *nudged_max = b->Mul(b->Sub(quant_max, nudged_zero_point), *scale); +} + +xla::ComputationDataHandle Quantize( + xla::ComputationBuilder* b, const xla::ComputationDataHandle& input, + const DataType data_type, + const xla::ComputationDataHandle& nudged_input_min, + const xla::ComputationDataHandle& nudged_input_max, + const xla::ComputationDataHandle& input_scale) { + xla::ComputationDataHandle one = XlaHelpers::FloatLiteral(b, data_type, 1.0f); + xla::ComputationDataHandle inv_scale = b->Div(one, input_scale); + xla::ComputationDataHandle half = + XlaHelpers::FloatLiteral(b, data_type, 0.5f); + + xla::ComputationDataHandle clamped = + b->Clamp(nudged_input_min, input, nudged_input_max); + xla::ComputationDataHandle clamped_shifted = + b->Sub(clamped, nudged_input_min); + xla::ComputationDataHandle rounded = + b->Floor(b->Add(b->Mul(clamped_shifted, inv_scale), half)); + return b->Add(b->Mul(rounded, input_scale), nudged_input_min); +} + +class FakeQuantWithMinMaxArgsOp : public XlaOpKernel { + public: + explicit FakeQuantWithMinMaxArgsOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + int num_bits; + OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits)); + OP_REQUIRES(ctx, num_bits >= 2 && num_bits <= 16, + errors::InvalidArgument("num_bits is out of range, expected " + "between 2 and 16, was: ", + num_bits)); + bool narrow_range; + OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range)); + quant_min_ = narrow_range ? 1 : 0; + quant_max_ = (1 << num_bits) - 1; + + float input_min, input_max; + OP_REQUIRES_OK(ctx, ctx->GetAttr("min", &input_min)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("max", &input_max)); + CpuNudge(input_min, input_max, quant_min_, quant_max_, &nudged_input_min_, + &nudged_input_max_, &input_scale_); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationDataHandle input = ctx->Input(0); + const DataType data_type = ctx->input_type(0); + + xla::ComputationBuilder* b = ctx->builder(); + xla::ComputationDataHandle nudged_input_min = + XlaHelpers::FloatLiteral(b, data_type, nudged_input_min_); + xla::ComputationDataHandle nudged_input_max = + XlaHelpers::FloatLiteral(b, data_type, nudged_input_max_); + xla::ComputationDataHandle input_scale = + XlaHelpers::FloatLiteral(b, data_type, input_scale_); + xla::ComputationDataHandle output = Quantize( + b, input, data_type, nudged_input_min, nudged_input_max, input_scale); + ctx->SetOutput(0, output); + } + + private: + float quant_min_; + float quant_max_; + float nudged_input_min_; + float nudged_input_max_; + float input_scale_; +}; + +REGISTER_XLA_OP(Name("FakeQuantWithMinMaxArgs"), FakeQuantWithMinMaxArgsOp); + +class FakeQuantWithMinMaxArgsGradOp : public XlaOpKernel { + public: + explicit FakeQuantWithMinMaxArgsGradOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + int num_bits; + OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits)); + OP_REQUIRES(ctx, num_bits >= 2 && num_bits <= 16, + errors::InvalidArgument("num_bits is out of range, expected " + "between 2 and 16, was: ", + num_bits)); + bool narrow_range; + OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range)); + const float quant_min = narrow_range ? 1 : 0; + const float quant_max = (1 << num_bits) - 1; + + float input_min, input_max, scale; + OP_REQUIRES_OK(ctx, ctx->GetAttr("min", &input_min)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("max", &input_max)); + CpuNudge(input_min, input_max, quant_min, quant_max, &nudged_input_min_, + &nudged_input_max_, &scale); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationDataHandle gradient = ctx->Input(0); + const TensorShape gradient_shape = ctx->InputShape(0); + xla::ComputationDataHandle input = ctx->Input(1); + const DataType data_type = ctx->input_type(1); + + xla::ComputationBuilder* b = ctx->builder(); + xla::ComputationDataHandle nudged_input_min = + XlaHelpers::FloatLiteral(b, data_type, nudged_input_min_); + xla::ComputationDataHandle nudged_input_max = + XlaHelpers::FloatLiteral(b, data_type, nudged_input_max_); + + xla::ComputationDataHandle between_nudged_min_max = + b->And(b->Le(nudged_input_min, input), b->Le(input, nudged_input_max)); + xla::ComputationDataHandle zeroes = b->Broadcast( + XlaHelpers::Zero(b, data_type), gradient_shape.dim_sizes()); + xla::ComputationDataHandle output = + b->Select(between_nudged_min_max, gradient, zeroes); + ctx->SetOutput(0, output); + } + + private: + float nudged_input_min_; + float nudged_input_max_; +}; + +REGISTER_XLA_OP(Name("FakeQuantWithMinMaxArgsGradient"), + FakeQuantWithMinMaxArgsGradOp); + +class FakeQuantWithMinMaxVarsOp : public XlaOpKernel { + public: + explicit FakeQuantWithMinMaxVarsOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + int num_bits; + OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits)); + OP_REQUIRES(ctx, num_bits >= 2 && num_bits <= 16, + errors::InvalidArgument("num_bits is out of range, expected " + "between 2 and 16, was: ", + num_bits)); + bool narrow_range; + OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range)); + quant_min_ = narrow_range ? 1 : 0; + quant_max_ = (1 << num_bits) - 1; + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationDataHandle input = ctx->Input(0); + const DataType data_type = ctx->input_type(0); + xla::ComputationDataHandle input_min = ctx->Input(1); + xla::ComputationDataHandle input_max = ctx->Input(2); + + xla::ComputationBuilder* b = ctx->builder(); + xla::ComputationDataHandle nudged_input_min, nudged_input_max, input_scale; + XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_, + &nudged_input_min, &nudged_input_max, &input_scale); + + xla::ComputationDataHandle output = Quantize( + b, input, data_type, nudged_input_min, nudged_input_max, input_scale); + ctx->SetOutput(0, output); + } + + private: + float quant_min_; + float quant_max_; +}; + +REGISTER_XLA_OP(Name("FakeQuantWithMinMaxVars"), FakeQuantWithMinMaxVarsOp); + +class FakeQuantWithMinMaxVarsGradOp : public XlaOpKernel { + public: + explicit FakeQuantWithMinMaxVarsGradOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + int num_bits; + OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits)); + OP_REQUIRES(ctx, num_bits >= 2 && num_bits <= 16, + errors::InvalidArgument("num_bits is out of range, expected " + "between 2 and 16, was: ", + num_bits)); + bool narrow_range; + OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range)); + quant_min_ = narrow_range ? 1 : 0; + quant_max_ = (1 << num_bits) - 1; + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationDataHandle gradient = ctx->Input(0); + const TensorShape gradient_shape = ctx->InputShape(0); + xla::ComputationDataHandle input = ctx->Input(1); + const DataType data_type = ctx->input_type(1); + xla::ComputationDataHandle input_min = ctx->Input(2); + xla::ComputationDataHandle input_max = ctx->Input(3); + + xla::ComputationBuilder* b = ctx->builder(); + xla::ComputationDataHandle nudged_input_min, nudged_input_max, input_scale; + XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_, + &nudged_input_min, &nudged_input_max, &input_scale); + + xla::ComputationDataHandle between_nudged_min_max = + b->And(b->Le(nudged_input_min, input), b->Le(input, nudged_input_max)); + xla::ComputationDataHandle zero = XlaHelpers::Zero(b, data_type); + xla::ComputationDataHandle zeroes = + b->Broadcast(zero, gradient_shape.dim_sizes()); + xla::ComputationDataHandle output0 = + b->Select(between_nudged_min_max, gradient, zeroes); + ctx->SetOutput(0, output0); + + xla::ComputationDataHandle below_min = b->Lt(input, nudged_input_min); + xla::ComputationDataHandle output1 = + b->ReduceAll(b->Select(below_min, gradient, zeroes), zero, + *ctx->GetOrCreateAdd(data_type)); + ctx->SetOutput(1, output1); + + xla::ComputationDataHandle above_max = b->Gt(input, nudged_input_max); + xla::ComputationDataHandle output2 = + b->ReduceAll(b->Select(above_max, gradient, zeroes), zero, + *ctx->GetOrCreateAdd(data_type)); + ctx->SetOutput(2, output2); + } + + private: + float quant_min_; + float quant_max_; +}; + +REGISTER_XLA_OP(Name("FakeQuantWithMinMaxVarsGradient"), + FakeQuantWithMinMaxVarsGradOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..a4f3c1c3ad9a928e0552c388a25ed9fcb08edabb --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc @@ -0,0 +1,122 @@ +/* Copyright 2017 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. +==============================================================================*/ + +// XLA-specific Ops for FFT. + +#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/core/framework/numeric_op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_slice.h" +#include "tensorflow/core/kernels/bounds_check.h" +#include "tensorflow/core/kernels/conv_grad_ops.h" +#include "tensorflow/core/kernels/ops_util.h" +#include "tensorflow/core/util/padding.h" +#include "tensorflow/core/util/tensor_format.h" + +namespace tensorflow { + +namespace { + +using xla::FftType; + +class GenericFftOp : public XlaOpKernel { + public: + explicit GenericFftOp(OpKernelConstruction* ctx, FftType fft_type, + int fft_rank) + : XlaOpKernel(ctx), fft_type_(fft_type), fft_rank_(fft_rank) {} + + void Compile(XlaOpKernelContext* ctx) override { + const TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES( + ctx, TensorShapeUtils::IsVectorOrHigher(input_shape), + errors::InvalidArgument("input must be at least 1 dimensional")); + + std::vector fft_length; + if (fft_type_ == FftType::RFFT || fft_type_ == FftType::IRFFT) { + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &fft_length)); + OP_REQUIRES(ctx, fft_length.size() == fft_rank_, + errors::InvalidArgument("fft_length must be length ", + fft_rank_, " vector")); + } else { + // Innermost axis provides the FFT length. + for (int i = 0; i < fft_rank_; i++) { + fft_length.push_back( + input_shape.dim_size(input_shape.dims() - fft_rank_ + i)); + } + } + + xla::ComputationBuilder* b = ctx->builder(); + xla::ComputationDataHandle fft = + b->Fft(ctx->Input(0), fft_type_, fft_length); + ctx->SetOutput(0, fft); + } + + protected: + const FftType fft_type_; + const int fft_rank_; + + private: + TF_DISALLOW_COPY_AND_ASSIGN(GenericFftOp); +}; + +template +class FFTOp : public GenericFftOp { + public: + explicit FFTOp(OpKernelConstruction* ctx) + : GenericFftOp(ctx, /*fft_type=*/FftType::FFT, /*fft_rank=*/FFTRank) {} +}; +REGISTER_XLA_OP(Name("FFT"), FFTOp<1>); +REGISTER_XLA_OP(Name("FFT2D"), FFTOp<2>); +REGISTER_XLA_OP(Name("FFT3D"), FFTOp<3>); + +template +class IFFTOp : public GenericFftOp { + public: + explicit IFFTOp(OpKernelConstruction* ctx) + : GenericFftOp(ctx, /*fft_type=*/FftType::IFFT, /*fft_rank=*/FFTRank) {} +}; +REGISTER_XLA_OP(Name("IFFT"), IFFTOp<1>); +REGISTER_XLA_OP(Name("IFFT2D"), IFFTOp<2>); +REGISTER_XLA_OP(Name("IFFT3D"), IFFTOp<3>); + +template +class RFFTOp : public GenericFftOp { + public: + explicit RFFTOp(OpKernelConstruction* ctx) + : GenericFftOp(ctx, /*fft_type=*/FftType::RFFT, /*fft_rank=*/FFTRank) {} +}; +REGISTER_XLA_OP(Name("RFFT").CompileTimeConstInput("fft_length"), RFFTOp<1>); +REGISTER_XLA_OP(Name("RFFT2D").CompileTimeConstInput("fft_length"), RFFTOp<2>); +REGISTER_XLA_OP(Name("RFFT3D").CompileTimeConstInput("fft_length"), RFFTOp<3>); + +template +class IRFFTOp : public GenericFftOp { + public: + explicit IRFFTOp(OpKernelConstruction* ctx) + : GenericFftOp(ctx, /*fft_type=*/FftType::IRFFT, /*fft_rank=*/FFTRank) {} +}; +REGISTER_XLA_OP(Name("IRFFT").CompileTimeConstInput("fft_length"), IRFFTOp<1>); +REGISTER_XLA_OP(Name("IRFFT2D").CompileTimeConstInput("fft_length"), + IRFFTOp<2>); +REGISTER_XLA_OP(Name("IRFFT3D").CompileTimeConstInput("fft_length"), + IRFFTOp<3>); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc index 70192cb32408139808ce788e44136d502ba4b66c..7945c05af40df21a798a2cff51fe7f8e935793f6 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* 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. @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/tf2xla/kernels/gather_op.h" #include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h" +#include "tensorflow/compiler/tf2xla/lib/while_loop.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_context.h" @@ -26,46 +26,48 @@ limitations under the License. namespace tensorflow { -xla::ComputationDataHandle XlaComputeGatherDynamicSlice( - XlaOpKernelContext* context, const xla::ComputationDataHandle& input, - const TensorShape& input_shape, const xla::ComputationDataHandle& indices, - const TensorShape& indices_shape, int64 axis, DataType dtype, - DataType index_type, xla::ComputationBuilder* builder) { +Status XlaGather(const xla::ComputationDataHandle& input, + const TensorShape& input_shape, + const xla::ComputationDataHandle& indices, + TensorShape indices_shape, int64 axis, bool indices_are_nd, + DataType dtype, DataType index_type, + xla::ComputationBuilder* builder, + xla::ComputationDataHandle* gather_output) { + // If the indices are N-dimensional, then the minor dimension of indices + // should be of size N and correspond to the N indices. + int64 num_index_dims = 1; + if (indices_are_nd) { + CHECK_GE(indices_shape.dims(), 1); + num_index_dims = indices_shape.dim_size(indices_shape.dims() - 1); + indices_shape.RemoveLastDims(1); + } + // Although the indices Tensor is flattened into rank 1 during the lookup, // and each scalar entry is used as an index into the first dimension of the // input, the output is returned with shape: // input.shape[:axis] + indices.shape + input.shape[axis+1:] - const int num_indices = indices_shape.num_elements(); + + const int64 num_indices = indices_shape.num_elements(); TensorShape input_shape_pre_axis(input_shape); input_shape_pre_axis.RemoveDimRange(axis, input_shape.dims()); TensorShape input_shape_post_axis(input_shape); - input_shape_post_axis.RemoveDimRange(0, axis + 1); - + input_shape_post_axis.RemoveDimRange(0, axis + num_index_dims); // Each slice of the input tensor has shape: - // [, 1, ] + // [, 1, ..., 1, ] TensorShape slice_shape(input_shape); - slice_shape.set_dim(axis, 1); - - // TODO(b/37575001) The tensor in which we construct the output during - // the loop must have rank >= 3 as a workaround for lowering issues. - int64 extra_dims = 0; - if (input_shape.dims() < 3) extra_dims = 3 - input_shape.dims(); + for (int64 i = 0; i < num_index_dims; ++i) { + slice_shape.set_dim(axis + i, 1); + } TensorShape loop_out_shape; - for (int64 k = 0; k < extra_dims; ++k) loop_out_shape.AddDim(1); loop_out_shape.AppendShape(input_shape_pre_axis); loop_out_shape.AddDim(num_indices); loop_out_shape.AppendShape(input_shape_post_axis); - - // Slices are reshaped into the rank >= 3 shape of the loop carried output. TensorShape loop_out_slice_shape; - for (int64 k = 0; k < extra_dims; ++k) loop_out_slice_shape.AddDim(1); loop_out_slice_shape.AppendShape(input_shape_pre_axis); loop_out_slice_shape.AddDim(1); loop_out_slice_shape.AppendShape(input_shape_post_axis); - // Finally, the loop-carried rank >= 3 output is reshaped to the op's - // specified result shape. TensorShape out_shape; out_shape.AppendShape(input_shape_pre_axis); out_shape.AppendShape(indices_shape); @@ -73,132 +75,176 @@ xla::ComputationDataHandle XlaComputeGatherDynamicSlice( // Degenerate case: empty indices. if (num_indices == 0) { - return builder->Broadcast(XlaHelpers::Zero(builder, dtype), - out_shape.dim_sizes()); + *gather_output = builder->Broadcast(XlaHelpers::Zero(builder, dtype), + out_shape.dim_sizes()); + return Status::OK(); + } + + for (int64 i = 0; i < num_index_dims; ++i) { + if (input_shape.dim_size(axis + i) == 0) { + return errors::InvalidArgument("Gather dimension ", axis + i, + " is of size zero in tensor with shape ", + input_shape.DebugString()); + } + } + + // Flatten the major dimensions of indices into a single dimension for ease of + // iteration. If there is an axis dimension, we must leave it alone. + std::vector flat_indices_shape = {num_indices}; + if (indices_are_nd) { + flat_indices_shape.push_back(num_index_dims); } // Specify the shape of the loop-carried Tensor tuple. - xla::PrimitiveType ptype; - TF_CHECK_OK(DataTypeToPrimitiveType(dtype, &ptype)); - xla::PrimitiveType idxtype; - TF_CHECK_OK(DataTypeToPrimitiveType(index_type, &idxtype)); - std::vector tuple_shapes( - {// The iteration counter i is a scalar, incremented each iteration. - xla::ShapeUtil::MakeShape(idxtype, {}), - // The input array has shape input_shape. Loop invariant. - xla::ShapeUtil::MakeShape(ptype, input_shape.dim_sizes()), - // The gather indices are reshaped to rank 1. Loop invariant. - xla::ShapeUtil::MakeShape(idxtype, {num_indices}), - // The output array is rank >= 3, and is updated on each loop iteration. - xla::ShapeUtil::MakeShape(ptype, loop_out_shape.dim_sizes())}); - xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes); // Construct the initial values of the loop-carried Tensors. - auto init_i = XlaHelpers::Zero(builder, index_type); + auto flat_indices = builder->Reshape(indices, flat_indices_shape); auto init_out = builder->Broadcast(XlaHelpers::Zero(builder, dtype), loop_out_shape.dim_sizes()); - // Flatten the indices into 1-D for ease of iteration. - auto indices_1d = builder->Reshape(indices, {num_indices}); - auto init = builder->Tuple({init_i, input, indices_1d, init_out}); - - // Construct the while loop condition (i < num_indices) - xla::ComputationBuilder condb(context->builder()->client(), - "GatherWhileCond"); - condb.Lt(condb.GetTupleElement( - condb.Parameter(0, tuple_shape, "GatherWhileTuple"), 0), - XlaHelpers::IntegerLiteral(&condb, index_type, num_indices)); - auto cond_status = condb.Build(); - auto cond = cond_status.ConsumeValueOrDie(); + auto init = {input, flat_indices, init_out}; // Construct the while loop body's function. The implementation of gather is: // for i in range(num_indices): // index = dynamic-slice(indices, i) // xi = dynamic-slice(input, index) // output = dynamic-update-slice(output, xi, i) - xla::ComputationBuilder bodyb(context->builder()->client(), - "GatherWhileBody"); - { - // The four loop carried values. - auto loop_tuple = bodyb.Parameter(0, tuple_shape, "GatherWhileTuple"); - auto i = bodyb.GetTupleElement(loop_tuple, 0); - auto input = bodyb.GetTupleElement(loop_tuple, 1); - auto indices = bodyb.GetTupleElement(loop_tuple, 2); - auto output = bodyb.GetTupleElement(loop_tuple, 3); - - // Slice from the input array. - auto index = bodyb.DynamicSlice(indices, bodyb.Reshape(i, {1}), {1}); - auto start_indices = bodyb.Pad( - bodyb.Reshape(index, {1}), XlaHelpers::Zero(&bodyb, index_type), + auto body_fn = [&](xla::ComputationDataHandle i, + gtl::ArraySlice loop_vars, + xla::ComputationBuilder* bodyb) { + auto input = loop_vars[0]; + auto indices = loop_vars[1]; + auto output = loop_vars[2]; + + auto zero_index = XlaHelpers::Zero(bodyb, index_type); + + // Slice the i-th index from the indices array. + xla::ComputationDataHandle index; + auto indices_offset = bodyb->Reshape(i, {1}); + if (indices_are_nd) { + // Slice out the entire nd index, if applicable. + indices_offset = bodyb->Pad(indices_offset, zero_index, + xla::MakeEdgePaddingConfig({{0, 1}})); + index = bodyb->DynamicSlice(indices, indices_offset, {1, num_index_dims}); + index = bodyb->Collapse(index, {0, 1}); + } else { + index = bodyb->DynamicSlice(indices, indices_offset, {1}); + } + + // Slice the corresponding data from the input array. + auto start_indices = bodyb->Pad( + index, zero_index, xla::MakeEdgePaddingConfig( {{input_shape_pre_axis.dims(), input_shape_post_axis.dims()}})); - auto slice_i = bodyb.Reshape( - bodyb.DynamicSlice(input, start_indices, slice_shape.dim_sizes()), + auto slice_i = bodyb->Reshape( + bodyb->DynamicSlice(input, start_indices, slice_shape.dim_sizes()), loop_out_slice_shape.dim_sizes()); - // Construct the index into the R3+ output Tensor 0, ..., , 0, ... + // Construct the index into the output Tensor 0, ..., , 0, ... std::vector out_index_vals( - loop_out_shape.dims(), - bodyb.Reshape(XlaHelpers::Zero(&bodyb, index_type), {1})); - out_index_vals[input_shape_pre_axis.dims() + extra_dims] = - bodyb.Reshape(i, {1}); - auto out_index = bodyb.ConcatInDim(out_index_vals, 0); + loop_out_shape.dims(), bodyb->Reshape(zero_index, {1})); + out_index_vals[input_shape_pre_axis.dims()] = bodyb->Reshape(i, {1}); + auto out_index = bodyb->ConcatInDim(out_index_vals, 0); // Update the output Tensor - auto updated_output = bodyb.DynamicUpdateSlice(output, slice_i, out_index); + auto updated_output = bodyb->DynamicUpdateSlice(output, slice_i, out_index); - bodyb.Tuple({bodyb.Add(i, XlaHelpers::One(&bodyb, index_type)), input, - indices, updated_output}); - } - auto body_status = bodyb.Build(); - auto body = body_status.ConsumeValueOrDie(); + return std::vector{input, indices, + updated_output}; + }; // Construct the While loop, extract and reshape the output. - auto gather_while = builder->While(cond, body, init); - auto gather_output = builder->GetTupleElement(gather_while, 3); - return builder->Reshape(gather_output, out_shape.dim_sizes()); + xla::PrimitiveType ptype; + TF_RETURN_IF_ERROR(DataTypeToPrimitiveType(index_type, &ptype)); + TF_ASSIGN_OR_RETURN(auto outputs, XlaForEachIndex(num_indices, ptype, body_fn, + init, "gather", builder)); + *gather_output = builder->Reshape(outputs[2], out_shape.dim_sizes()); + return Status::OK(); } -GatherOpDynamicSlice::GatherOpDynamicSlice(OpKernelConstruction* context) - : XlaOpKernel(context) {} - -void GatherOpDynamicSlice::Compile(XlaOpKernelContext* context) { - xla::ComputationBuilder* builder = context->builder(); - auto input = context->Input(0); - auto input_shape = context->InputShape(0); - auto indices = context->Input(1); - auto indices_shape = context->InputShape(1); - int64 axis = 0; - if (context->num_inputs() == 3) { - const TensorShape axis_shape = context->InputShape(2); - OP_REQUIRES(context, TensorShapeUtils::IsScalar(axis_shape), - errors::InvalidArgument("axis must be scalar")); - DataType axis_type = input_type(2); - OP_REQUIRES(context, axis_type == DT_INT32 || axis_type == DT_INT64, - errors::InvalidArgument("axis must be int32 or int64")); - - OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(2, &axis)); - const auto params_dims = input_shape.dims(); - if (axis < 0) { - axis += params_dims; +class GatherOp : public XlaOpKernel { + public: + explicit GatherOp(OpKernelConstruction* context) : XlaOpKernel(context) {} + + void Compile(XlaOpKernelContext* context) override { + xla::ComputationBuilder* builder = context->builder(); + auto input = context->Input(0); + auto input_shape = context->InputShape(0); + auto indices = context->Input(1); + auto indices_shape = context->InputShape(1); + int64 axis = 0; + if (context->num_inputs() == 3) { + const TensorShape axis_shape = context->InputShape(2); + OP_REQUIRES(context, TensorShapeUtils::IsScalar(axis_shape), + errors::InvalidArgument("axis must be scalar")); + DataType axis_type = input_type(2); + OP_REQUIRES(context, axis_type == DT_INT32 || axis_type == DT_INT64, + errors::InvalidArgument("axis must be int32 or int64")); + + OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(2, &axis)); + const auto params_dims = input_shape.dims(); + if (axis < 0) { + axis += params_dims; + } + OP_REQUIRES( + context, 0 <= axis && axis < params_dims, + errors::InvalidArgument("Expected axis in the range [", -params_dims, + ", ", params_dims, "), but got ", axis)); } - OP_REQUIRES( - context, 0 <= axis && axis < params_dims, - errors::InvalidArgument("Expected axis in the range [", -params_dims, - ", ", params_dims, "), but got ", axis)); - } - DataType index_type = input_type(1); - OP_REQUIRES(context, index_type == DT_INT32 || index_type == DT_INT64, - errors::InvalidArgument("indices must be int32 or int64")); + DataType index_type = input_type(1); + OP_REQUIRES(context, index_type == DT_INT32 || index_type == DT_INT64, + errors::InvalidArgument("indices must be int32 or int64")); - xla::ComputationDataHandle gather = XlaComputeGatherDynamicSlice( - context, input, input_shape, indices, indices_shape, axis, input_type(0), - index_type, builder); - context->SetOutput(0, gather); -} + xla::ComputationDataHandle gather; + OP_REQUIRES_OK( + context, XlaGather(input, input_shape, indices, indices_shape, axis, + /*indices_are_nd=*/false, input_type(0), index_type, + builder, &gather)); + context->SetOutput(0, gather); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(GatherOp); +}; + +REGISTER_XLA_OP(Name("Gather"), GatherOp); +REGISTER_XLA_OP(Name("GatherV2").CompileTimeConstInput("axis"), GatherOp); + +class GatherNdOp : public XlaOpKernel { + public: + explicit GatherNdOp(OpKernelConstruction* context) : XlaOpKernel(context) {} + + void Compile(XlaOpKernelContext* context) override { + DataType params_type = context->input_type(0); + DataType indices_type = context->input_type(1); + + TensorShape params_shape = context->InputShape(0); + TensorShape indices_shape = context->InputShape(1); + OP_REQUIRES(context, TensorShapeUtils::IsVectorOrHigher(params_shape), + errors::InvalidArgument("params must be at least a vector")); + OP_REQUIRES(context, TensorShapeUtils::IsVectorOrHigher(indices_shape), + errors::InvalidArgument("indices must be at least a vector")); + const int64 num_index_dims = + indices_shape.dim_size(indices_shape.dims() - 1); + OP_REQUIRES( + context, num_index_dims <= params_shape.dims(), + errors::InvalidArgument( + "index innermost dimension length must be <= params rank; saw: ", + indices_shape.dim_size(indices_shape.dims() - 1), " vs. ", + params_shape.dims())); + + xla::ComputationBuilder* builder = context->builder(); + auto params = context->Input(0); + auto indices = context->Input(1); + xla::ComputationDataHandle gather; + OP_REQUIRES_OK(context, XlaGather(params, params_shape, indices, + indices_shape, /*axis=*/0, + /*indices_are_nd=*/true, params_type, + indices_type, builder, &gather)); + context->SetOutput(0, gather); + } +}; -REGISTER_XLA_OP(Name("Gather"), GatherOpDynamicSlice); -REGISTER_XLA_OP(Name("GatherV2").CompileTimeConstInput("axis"), - GatherOpDynamicSlice); +REGISTER_XLA_OP(Name("GatherNd"), GatherNdOp); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.h b/tensorflow/compiler/tf2xla/kernels/gather_op.h deleted file mode 100644 index df86e1fcdd1a4860ed7ee0c5017d25ccf9d227ea..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/kernels/gather_op.h +++ /dev/null @@ -1,41 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ - -// Declaration of the Gather Op using the XLA dynamic slice implementation. - -#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_H_ -#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_H_ - -#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" -#include "tensorflow/compiler/xla/client/client_library.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" -#include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/util/bcast.h" - -namespace tensorflow { - -class GatherOpDynamicSlice : public XlaOpKernel { - public: - explicit GatherOpDynamicSlice(OpKernelConstruction* context); - - void Compile(XlaOpKernelContext* context) override; - - private: - TF_DISALLOW_COPY_AND_ASSIGN(GatherOpDynamicSlice); -}; - -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_H_ diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h b/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h index 2c80395c56d73adad7dc1679ba6423fbe103605a..bd8b92c22d71fe89ab8951ec79f411feef6505e3 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h +++ b/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h @@ -30,11 +30,16 @@ namespace tensorflow { // shape input_shape) keyed on indices (of shape indices_shape). // // index_type must be must be DT_INT32 or DT_INT64. -xla::ComputationDataHandle XlaComputeGatherDynamicSlice( - XlaOpKernelContext* ctx, const xla::ComputationDataHandle& input, - const TensorShape& input_shape, const xla::ComputationDataHandle& indices, - const TensorShape& indices_shape, int64 axis, DataType dtype, - DataType index_type, xla::ComputationBuilder* builder); +// If `indices_are_nd` is true, the last dimension of `indices` are treated as +// a multidimensional index values. Otherwise, `indices` is treated as a tensor +// of scalar indices. +Status XlaGather(const xla::ComputationDataHandle& input, + const TensorShape& input_shape, + const xla::ComputationDataHandle& indices, + TensorShape indices_shape, int64 axis, bool indices_are_nd, + DataType dtype, DataType index_type, + xla::ComputationBuilder* builder, + xla::ComputationDataHandle* gather_output); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/identity_op.cc b/tensorflow/compiler/tf2xla/kernels/identity_op.cc index d2b1f7913ecc9113284827b53de8fb0e5b711322..39af662b638cb9d723118e58fcfc983633fed497 100644 --- a/tensorflow/compiler/tf2xla/kernels/identity_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/identity_op.cc @@ -40,6 +40,7 @@ REGISTER_XLA_OP(Name("Identity").CompilationOnly(), IdentityOp); REGISTER_XLA_OP(Name("IdentityN").CompilationOnly(), IdentityOp); REGISTER_XLA_OP(Name("PreventGradient"), IdentityOp); REGISTER_XLA_OP(Name("StopGradient"), IdentityOp); +REGISTER_XLA_OP(Name("Snapshot"), IdentityOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc index c0b8f9c179ee8760ec3d603e1cb2c1db2c11d2f1..f36b3f594826c27b7866d956c855aa3638db9cb4 100644 --- a/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc @@ -113,26 +113,140 @@ xla::ComputationDataHandle MakeBilinearResizeKernel( auto make_1d_kernel = [](int64 n) { std::vector kernel(n * 2 - 1); for (int64 i = 0; i < n; ++i) { - float v = i + 1; + float v = (i + 1.0f) / n; kernel[i] = v; kernel[n * 2 - 2 - i] = v; } return kernel; }; - // Form a block diagonal kernel where each channel interacts only with itself. - xla::Array4D diag(1, 1, channels, channels, 0.0f); - for (int i = 0; i < channels; ++i) { - diag(0, 0, i, i) = 1.0f / (kernel_size[0] * kernel_size[1]); - } + xla::ComputationDataHandle channels_iota; + // DT_INT32 Iota will always return status::OK(). + TF_CHECK_OK( + XlaHelpers::Iota(builder, DataType::DT_INT32, channels, &channels_iota)); + + auto diag = builder->ConvertElementType( + builder->Eq( + builder->Broadcast(channels_iota, {2 * kernel_size[0] - 1, + 2 * kernel_size[1] - 1, channels}), + channels_iota, /*broadcast_dimensions=*/{2}), + xla::PrimitiveType::F32); return builder->Mul( - builder->ConstantR1(make_1d_kernel(kernel_size[0])), - builder->Mul(builder->ConstantR1(make_1d_kernel(kernel_size[1])), - builder->ConstantR4FromArray4D(diag), + builder->Mul(diag, + builder->ConstantR1(make_1d_kernel(kernel_size[1])), /*broadcast_dimensions=*/{1}), + builder->ConstantR1(make_1d_kernel(kernel_size[0])), /*broadcast_dimensions=*/{0}); } +xla::ComputationDataHandle ResizeUsingDilationAndConvolution( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& input, + const int num_spatial_dims, std::vector in_size, + std::vector out_size, const int64 channels) { + // Picture for a 1x3 to 1x4 resize: + // stride = 2, kernel size = 3 + // Input: + // 3 6 9 + // Input with dilation and padding: + // 0 0 3 0 0 6 0 0 9 0 0 + // Convolution kernel: + // 1/3 * [1 2 3 2 1] + // Output: + // 3 5 7 9 + xla::ConvolutionDimensionNumbers dimension_numbers; + dimension_numbers.set_input_batch_dimension(0); + dimension_numbers.set_output_batch_dimension(0); + dimension_numbers.set_input_feature_dimension(3); + dimension_numbers.set_output_feature_dimension(3); + for (int i = 0; i < num_spatial_dims; ++i) { + dimension_numbers.add_input_spatial_dimensions(1 + i); + dimension_numbers.add_output_spatial_dimensions(1 + i); + dimension_numbers.add_kernel_spatial_dimensions(i); + } + dimension_numbers.set_kernel_input_feature_dimension(num_spatial_dims); + dimension_numbers.set_kernel_output_feature_dimension(num_spatial_dims + 1); + + ResizeConvolutionDims dims = + ComputeResizeConvolutionParameters(in_size, out_size); + xla::ComputationDataHandle kernel = + MakeBilinearResizeKernel(builder, dims.kernel_size, channels); + xla::ComputationDataHandle output = builder->ConvGeneralDilated( + input, kernel, dims.stride, + /*padding=*/ + {{dims.kernel_size[0] - 1, dims.kernel_size[0] - 1}, + {dims.kernel_size[1] - 1, dims.kernel_size[1] - 1}}, + /*lhs_dilation=*/dims.kernel_size, + /*rhs_dilation=*/{1, 1}, dimension_numbers); + + // Add broadcasts to handle expanding from a size == 1 dimension to a + // size > 1 dimension. + for (int i = 0; i < num_spatial_dims; ++i) { + if (in_size[i] == 1 && out_size[i] > 1) { + output = builder->Add(output, builder->ConstantR1(out_size[i], 0), + /*broadcast_dimensions=*/{1 + i}); + } + } + return output; +} + +xla::ComputationDataHandle ResizeUsingDilationAndConvolutionGradOp( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& grad, + const int num_spatial_dims, std::vector in_size, + std::vector grad_size, const int64 channels) { + ResizeConvolutionDims dims = + ComputeResizeConvolutionParameters(in_size, grad_size); + + // To form the backward convolution, we keep the kernel unchanged (it is + // already symmetric) and swap the roles of strides and LHS dilation. + xla::ConvolutionDimensionNumbers dimension_numbers; + dimension_numbers.set_input_batch_dimension(0); + dimension_numbers.set_output_batch_dimension(0); + dimension_numbers.set_input_feature_dimension(3); + dimension_numbers.set_output_feature_dimension(3); + for (int i = 0; i < num_spatial_dims; ++i) { + dimension_numbers.add_input_spatial_dimensions(1 + i); + dimension_numbers.add_output_spatial_dimensions(1 + i); + dimension_numbers.add_kernel_spatial_dimensions(i); + } + dimension_numbers.set_kernel_input_feature_dimension(num_spatial_dims); + dimension_numbers.set_kernel_output_feature_dimension(num_spatial_dims + 1); + xla::ComputationDataHandle kernel = + MakeBilinearResizeKernel(builder, dims.kernel_size, channels); + + // Broadcast the input kernel where the forward op expanded from a size == 1 + // dimension to a size > 1 dimension. This has the effect of summing the + // gradient contributions in that dimension. + for (int i = 0; i < num_spatial_dims; ++i) { + if (in_size[i] == 1 && grad_size[i] > 1) { + kernel = builder->Add(kernel, builder->ConstantR1(grad_size[i], 0), + /*broadcast_dimensions=*/{i}); + } + } + + xla::ComputationDataHandle output = builder->ConvGeneralDilated( + grad, kernel, /*window_strides=*/dims.kernel_size, + /*padding=*/ + {{dims.kernel_size[0] - 1, dims.kernel_size[0] - 1}, + {dims.kernel_size[1] - 1, dims.kernel_size[1] - 1}}, + /*lhs_dilation=*/dims.stride, + /*rhs_dilation=*/{1, 1}, dimension_numbers); + + // If in_size[i] > 1 and grad_size[i] == 1, pad the output in dimension i. + // Opposite of the slice performed by the forward op. + xla::PaddingConfig padding = xla::MakeNoPaddingConfig(4); + bool pad_output = false; + for (int i = 0; i < num_spatial_dims; ++i) { + if (in_size[i] > 1 && grad_size[i] == 1) { + pad_output = true; + padding.mutable_dimensions(1 + i)->set_edge_padding_high(in_size[i] - 1); + } + } + if (pad_output) { + output = builder->Pad(output, builder->ConstantR0(0.0f), padding); + } + return output; +} + class ResizeBilinearOp : public XlaOpKernel { public: explicit ResizeBilinearOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -151,8 +265,8 @@ class ResizeBilinearOp : public XlaOpKernel { errors::InvalidArgument("input must be 4-dimensional", input_shape.DebugString())); const int64 batch = input_shape.dim_size(0); - const std::vector in_size = {input_shape.dim_size(1), - input_shape.dim_size(2)}; + std::vector in_size = {input_shape.dim_size(1), + input_shape.dim_size(2)}; const int64 channels = input_shape.dim_size(3); OP_REQUIRES(ctx, in_size[0] > 0 && in_size[1] > 0, errors::InvalidArgument("input size must be positive, got [", @@ -192,47 +306,41 @@ class ResizeBilinearOp : public XlaOpKernel { // Output is always type float. input = b->ConvertElementType(input, xla::F32); - // Picture for a 1x3 to 1x4 resize: - // stride = 2, kernel size = 3 - // Input: - // 3 6 9 - // Input with dilation and padding: - // 0 0 3 0 0 6 0 0 9 0 0 - // Convolution kernel: - // 1/3 * [1 2 3 2 1] - // Output: - // 3 5 7 9 - xla::ConvolutionDimensionNumbers dnums; - dnums.set_input_batch_dimension(0); - dnums.set_output_batch_dimension(0); - dnums.set_input_feature_dimension(3); - dnums.set_output_feature_dimension(3); - for (int i = 0; i < num_spatial_dims; ++i) { - dnums.add_input_spatial_dimensions(1 + i); - dnums.add_output_spatial_dimensions(1 + i); - dnums.add_kernel_spatial_dimensions(i); - } - dnums.set_kernel_input_feature_dimension(num_spatial_dims); - dnums.set_kernel_output_feature_dimension(num_spatial_dims + 1); - - ResizeConvolutionDims dims = - ComputeResizeConvolutionParameters(in_size, out_size); - xla::ComputationDataHandle kernel = - MakeBilinearResizeKernel(b, dims.kernel_size, channels); - xla::ComputationDataHandle output = b->ConvGeneralDilated( - input, kernel, dims.stride, - /*padding=*/ - {{dims.kernel_size[0] - 1, dims.kernel_size[0] - 1}, - {dims.kernel_size[1] - 1, dims.kernel_size[1] - 1}}, - /*lhs_dilation=*/dims.kernel_size, - /*rhs_dilation=*/{1, 1}, dnums); - - // Add broadcasts to handle expanding from a size == 1 dimension to a - // size > 1 dimension. - for (int i = 0; i < num_spatial_dims; ++i) { - if (in_size[i] == 1 && out_size[i] > 1) { - output = b->Add(output, b->ConstantR1(out_size[i], 0), - /*broadcast_dimensions=*/{1 + i}); + // Special Case: + // Instead of doing a ResizeUsingDilationAndConvolution directly, + // while (out_size[0]-1) = c * 2^x * (in_size[0]-1) for x>1 c>1, resize the + // image to 2*(in_size[0]-1)+1 x-times and then resize by scale c(int here). + // Instead of resizing directly we resize it iteratively. + // + // Since bilinear resize can be broken down as 2 sequential linear + // operations along different dimensions. + // Given sufficient numerical stability and a cxd is same as resizing axb -> exf -> cxd. + // + // This makes the convolutions kernels smaller and the operation faster. + xla::ComputationDataHandle output = input; + while (in_size != out_size) { + if (in_size[0] != 1 && in_size[1] != 1) { + std::vector k = { + (static_cast(out_size[0]) - 1) / ((in_size[0] - 1) * 2), + (static_cast(out_size[1]) - 1) / ((in_size[1] - 1) * 2)}; + if ((k[0] == std::floor(k[0])) && (k[1] == std::floor(k[1])) && + k[0] > 1 && k[1] > 1) { + std::vector next_out_size = {(in_size[0] - 1) * 2 + 1, + (in_size[1] - 1) * 2 + 1}; + output = ResizeUsingDilationAndConvolution( + b, input, num_spatial_dims, in_size, next_out_size, channels); + input = output; + in_size = next_out_size; + } else { + output = ResizeUsingDilationAndConvolution( + b, input, num_spatial_dims, in_size, out_size, channels); + in_size = out_size; + } + } else { + output = ResizeUsingDilationAndConvolution(b, input, num_spatial_dims, + in_size, out_size, channels); + in_size = out_size; } } @@ -268,8 +376,8 @@ class ResizeBilinearGradOp : public XlaOpKernel { errors::InvalidArgument("input must be 4-dimensional", input_shape.DebugString())); const int64 batch = input_shape.dim_size(0); - const std::vector in_size = {input_shape.dim_size(1), - input_shape.dim_size(2)}; + std::vector in_size = {input_shape.dim_size(1), + input_shape.dim_size(2)}; const int64 channels = input_shape.dim_size(3); OP_REQUIRES(ctx, in_size[0] > 0 && in_size[1] > 0, errors::InvalidArgument("input size must be positive, got [", @@ -300,58 +408,31 @@ class ResizeBilinearGradOp : public XlaOpKernel { xla::ComputationDataHandle grad = ctx->Input(0); - ResizeConvolutionDims dims = - ComputeResizeConvolutionParameters(in_size, grad_size); - - // To form the backward convolution, we keep the kernel unchanged (it is - // already symmetric) and swap the roles of strides and LHS dilation. - xla::ConvolutionDimensionNumbers dnums; - dnums.set_input_batch_dimension(0); - dnums.set_output_batch_dimension(0); - dnums.set_input_feature_dimension(3); - dnums.set_output_feature_dimension(3); - for (int i = 0; i < num_spatial_dims; ++i) { - dnums.add_input_spatial_dimensions(1 + i); - dnums.add_output_spatial_dimensions(1 + i); - dnums.add_kernel_spatial_dimensions(i); - } - dnums.set_kernel_input_feature_dimension(num_spatial_dims); - dnums.set_kernel_output_feature_dimension(num_spatial_dims + 1); - xla::ComputationDataHandle kernel = - MakeBilinearResizeKernel(b, dims.kernel_size, channels); - - // Broadcast the input kernel where the forward op expanded from a size == 1 - // dimension to a size > 1 dimension. This has the effect of summing the - // gradient contributions in that dimension. - for (int i = 0; i < num_spatial_dims; ++i) { - if (in_size[i] == 1 && grad_size[i] > 1) { - kernel = b->Add(kernel, b->ConstantR1(grad_size[i], 0), - /*broadcast_dimensions=*/{i}); - } - } - - xla::ComputationDataHandle output = b->ConvGeneralDilated( - grad, kernel, /*window_strides=*/dims.kernel_size, - /*padding=*/ - {{dims.kernel_size[0] - 1, dims.kernel_size[0] - 1}, - {dims.kernel_size[1] - 1, dims.kernel_size[1] - 1}}, - /*lhs_dilation=*/dims.stride, - /*rhs_dilation=*/{1, 1}, dnums); - - // If in_size[i] > 1 and grad_size[i] == 1, pad the output in dimension i. - // Opposite of the slice performed by the forward op. - xla::PaddingConfig padding = xla::MakeNoPaddingConfig(4); - bool pad_output = false; - for (int i = 0; i < num_spatial_dims; ++i) { - if (in_size[i] > 1 && grad_size[i] == 1) { - pad_output = true; - padding.mutable_dimensions(1 + i)->set_edge_padding_high(in_size[i] - - 1); + xla::ComputationDataHandle output = grad; + while (in_size != grad_size) { + if (in_size[0] != 1 && in_size[1] != 1) { + std::vector k = { + (static_cast(grad_size[0]) - 1) / ((in_size[0] - 1) * 2), + (static_cast(grad_size[1]) - 1) / ((in_size[1] - 1) * 2)}; + if ((k[0] == std::floor(k[0])) && (k[1] == std::floor(k[1])) && + k[0] > 1 && k[1] > 1) { + std::vector next_grad_size = {(in_size[0] - 1) * 2 + 1, + (in_size[1] - 1) * 2 + 1}; + output = ResizeUsingDilationAndConvolutionGradOp( + b, grad, num_spatial_dims, in_size, next_grad_size, channels); + grad = output; + in_size = next_grad_size; + } else { + output = ResizeUsingDilationAndConvolutionGradOp( + b, grad, num_spatial_dims, in_size, grad_size, channels); + in_size = grad_size; + } + } else { + output = ResizeUsingDilationAndConvolutionGradOp( + b, grad, num_spatial_dims, in_size, grad_size, channels); + in_size = grad_size; } } - if (pad_output) { - output = b->Pad(output, b->ConstantR0(0.0f), padding); - } output = b->ConvertElementType(output, output_type_); ctx->SetOutput(0, output); diff --git a/tensorflow/compiler/tf2xla/kernels/matmul_op.cc b/tensorflow/compiler/tf2xla/kernels/matmul_op.cc index 644abd5905c6ce5a8f61792a1986560bab891040..886baf8115243a22b7255a3961c914d4cf6c2ed5 100644 --- a/tensorflow/compiler/tf2xla/kernels/matmul_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/matmul_op.cc @@ -29,10 +29,12 @@ constexpr std::array kMatmulTypes = { class MatMulOp : public XlaOpKernel { public: explicit MatMulOp(OpKernelConstruction* ctx, bool is_sparse = false) - : XlaOpKernel(ctx) { + : XlaOpKernel(ctx), is_sparse_(is_sparse) { OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_)); if (is_sparse) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("Ta", &a_type_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("Tb", &b_type_)); // SparseMatMul is actually dense matmul with a hint that one or // both of the inputs may contain a lot of zeroes. On CPU these // inputs are dynamically converted to sparse representation @@ -66,14 +68,25 @@ class MatMulOp : public XlaOpKernel { xla::ComputationDataHandle a = ctx->Input(0); xla::ComputationDataHandle b = ctx->Input(1); + if (is_sparse_) { + if (a_type_ == DT_BFLOAT16) { + a = ctx->builder()->ConvertElementType(a, xla::F32); + } + if (b_type_ == DT_BFLOAT16) { + b = ctx->builder()->ConvertElementType(b, xla::F32); + } + } auto lhs = (transpose_a_) ? ctx->builder()->Transpose(a, {1, 0}) : a; auto rhs = (transpose_b_) ? ctx->builder()->Transpose(b, {1, 0}) : b; ctx->SetOutput(0, ctx->builder()->Dot(lhs, rhs)); } private: + bool is_sparse_; bool transpose_a_; bool transpose_b_; + DataType a_type_; + DataType b_type_; }; REGISTER_XLA_OP(Name("MatMul").TypeConstraint("T", kMatmulTypes), MatMulOp); diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..faa415a97b053b4b11d015fefcd430210b98118a --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc @@ -0,0 +1,98 @@ +/* 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/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/framework/tensor_shape.h" + +namespace tensorflow { +namespace { + +class MatrixBandPartOp : public XlaOpKernel { + public: + explicit MatrixBandPartOp(OpKernelConstruction* context) + : XlaOpKernel(context) {} + + void Compile(XlaOpKernelContext* context) override { + const TensorShape input_shape = context->InputShape(0); + // Preliminary validation of sizes. + OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input_shape), + errors::InvalidArgument( + "input must be at least 2-dim, received shape: ", + input_shape.DebugString())); + + const TensorShape num_lower_in_shape = context->InputShape(1); + OP_REQUIRES(context, TensorShapeUtils::IsScalar(num_lower_in_shape), + errors::InvalidArgument("num_lower must be scalar, got shape ", + num_lower_in_shape.DebugString())); + + const TensorShape num_upper_in_shape = context->InputShape(2); + OP_REQUIRES(context, TensorShapeUtils::IsScalar(num_upper_in_shape), + errors::InvalidArgument("num_upper must be scalar, got shape ", + num_upper_in_shape.DebugString())); + + xla::ComputationBuilder* builder = context->builder(); + xla::ComputationDataHandle input = context->Input(0); + xla::ComputationDataHandle num_lower = context->Input(1); + xla::ComputationDataHandle num_upper = context->Input(2); + DataType input_type = context->input_type(0); + DataType index_type = context->input_type(1); + + TensorShape batch_shape = input_shape; + batch_shape.RemoveLastDims(2); + const int64 m = input_shape.dim_size(input_shape.dims() - 2); + const int64 n = input_shape.dim_size(input_shape.dims() - 1); + + // Compute 'offset', which is how many diagonals we are above/below the + // diagonal. + xla::ComputationDataHandle iota_m; + OP_REQUIRES_OK(context, XlaHelpers::Iota(builder, index_type, m, &iota_m)); + + xla::ComputationDataHandle iota_n; + OP_REQUIRES_OK(context, XlaHelpers::Iota(builder, index_type, n, &iota_n)); + + auto offset = builder->Sub(builder->Broadcast(iota_n, {m}), iota_m, + /*broadcast_dimensions=*/{0}); + + // If num_lower or num_upper are negative, include all lower/upper + // diagonals. + auto zero_index = XlaHelpers::Zero(builder, index_type); + num_lower = builder->Select( + builder->Lt(num_lower, zero_index), + XlaHelpers::IntegerLiteral(builder, index_type, m), num_lower); + num_upper = builder->Select( + builder->Lt(num_upper, zero_index), + XlaHelpers::IntegerLiteral(builder, index_type, n), num_upper); + + auto indicator = builder->And(builder->Le(builder->Neg(num_lower), offset), + builder->Le(offset, num_upper)); + indicator = builder->Broadcast(indicator, batch_shape.dim_sizes()); + + auto zero_input = XlaHelpers::Zero(builder, input_type); + auto output = builder->Select( + indicator, input, + builder->Broadcast(zero_input, input_shape.dim_sizes())); + + context->SetOutput(0, output); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(MatrixBandPartOp); +}; +REGISTER_XLA_OP(Name("MatrixBandPart"), MatrixBandPartOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..b2940bdcff75a087c914fdad0cb2426276e41aff --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc @@ -0,0 +1,93 @@ +/* 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/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" + +namespace tensorflow { + +class MatrixSetDiagOp : public XlaOpKernel { + public: + explicit MatrixSetDiagOp(OpKernelConstruction* context) + : XlaOpKernel(context) {} + + void Compile(XlaOpKernelContext* context) override { + const TensorShape input_shape = context->InputShape(0); + const TensorShape diag_shape = context->InputShape(1); + + const int rank = input_shape.dims(); + + // Preliminary validation of sizes. + OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input_shape), + errors::InvalidArgument( + "input must be at least 2-dim, received shape: ", + input_shape.DebugString())); + + // Check to make sure the last dimension of diag is equal to the smaller of + // the last two dimensions of input. + const int64 m = input_shape.dim_size(rank - 2); + const int64 n = input_shape.dim_size(rank - 1); + const int64 min_dim = std::min(m, n); + + TensorShape batch_shape = input_shape; + batch_shape.RemoveLastDims(2); + + TensorShape expected_diag_shape = batch_shape; + expected_diag_shape.AddDim(min_dim); + OP_REQUIRES(context, expected_diag_shape == diag_shape, + errors::InvalidArgument( + "must have diagonal.shape == input.shape[:-2] + " + "min(input.shape[-2:]), but received input shape: ", + input_shape.DebugString(), + " and diagonal shape: ", diag_shape.DebugString())); + + xla::ComputationBuilder* builder = context->builder(); + xla::ComputationDataHandle input = context->Input(0); + xla::ComputationDataHandle diag = context->Input(1); + + auto zero = XlaHelpers::Zero(builder, context->input_type(0)); + + // Create an indicator tensor that is true only on the diagonal. + xla::ComputationDataHandle iota_m; + OP_REQUIRES_OK(context, XlaHelpers::Iota(builder, DT_INT32, m, &iota_m)); + xla::ComputationDataHandle iota_n; + OP_REQUIRES_OK(context, XlaHelpers::Iota(builder, DT_INT32, n, &iota_n)); + auto indicator = builder->Eq(iota_m, + builder->Broadcast(iota_n, {m}), + /*broadcast_dimensions=*/{0}); + indicator = builder->Broadcast(indicator, batch_shape.dim_sizes()); + + // Broadcast diag up to the input shape. Use an implicit broadcast (Add) + // because we need to broadcast on the right. + std::vector diag_broadcast_dims(rank - 1); + std::iota(diag_broadcast_dims.begin(), diag_broadcast_dims.end(), 0); + if (min_dim != m) { + diag_broadcast_dims.back() = rank - 1; + } + diag = builder->Add(diag, builder->Broadcast(zero, input_shape.dim_sizes()), + /*broadcast_dimensions=*/diag_broadcast_dims); + + auto output = builder->Select(indicator, diag, input); + context->SetOutput(0, output); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(MatrixSetDiagOp); +}; + +REGISTER_XLA_OP(Name("MatrixSetDiag"), MatrixSetDiagOp); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..eaed93146460de5a6e8328432302cc75bf36a534 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc @@ -0,0 +1,50 @@ +/* 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/compiler/tf2xla/lib/triangular_solve.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" + +namespace tensorflow { +namespace { + +class MatrixTriangularSolveOp : public XlaOpKernel { + public: + explicit MatrixTriangularSolveOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("lower", &lower_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("adjoint", &adjoint_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + auto result = TriangularSolve( + ctx->builder(), ctx->Input(0), ctx->Input(1), /*left_side=*/true, + /*lower=*/lower_, /*transpose_a=*/adjoint_, /*conjugate_a=*/adjoint_); + if (!result.ok()) { + ctx->SetStatus(result.status()); + return; + } + ctx->SetOutput(0, result.ValueOrDie()); + } + + private: + bool lower_; + bool adjoint_; +}; + +REGISTER_XLA_OP(Name("MatrixTriangularSolve"), MatrixTriangularSolveOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc index 0b5a38967aeb5b4cd66de5220e2c764371440c2d..d4fb5dd4e06c7c70591262c0d63a91c383a2a6e0 100644 --- a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc @@ -37,21 +37,23 @@ class PoolingOp : public XlaOpKernel { public: PoolingOp(OpKernelConstruction* ctx, int num_spatial_dims) : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) { - std::vector ksize_int; - std::vector stride_int; - OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_int)); - OP_REQUIRES(ctx, ksize_int.size() == num_dims(), - errors::InvalidArgument("Sliding window ksize field must " - "specify ", - num_dims(), " dimensions")); - OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_int)); - OP_REQUIRES(ctx, stride_int.size() == num_dims(), - errors::InvalidArgument("Sliding window stride field must " - "specify ", - num_dims(), " dimensions")); - for (int i = 0; i < num_dims(); ++i) { - ksize_.push_back(ksize_int[i]); - stride_.push_back(stride_int[i]); + if (ctx->num_inputs() == 1) { + std::vector ksize_int; + std::vector stride_int; + OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_int)); + OP_REQUIRES(ctx, ksize_int.size() == num_dims(), + errors::InvalidArgument("Sliding window ksize field must " + "specify ", + num_dims(), " dimensions")); + OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_int)); + OP_REQUIRES(ctx, stride_int.size() == num_dims(), + errors::InvalidArgument("Sliding window stride field must " + "specify ", + num_dims(), " dimensions")); + for (int i = 0; i < num_dims(); ++i) { + ksize_.push_back(ksize_int[i]); + stride_.push_back(stride_int[i]); + } } Padding padding; OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding)); @@ -77,6 +79,33 @@ class PoolingOp : public XlaOpKernel { xla::ComputationDataHandle input = ctx->Input(0); const TensorShape input_shape = ctx->InputShape(0); + std::vector ksize = ksize_; + std::vector stride = stride_; + if (ctx->num_inputs() != 1) { + const TensorShape ksize_shape = ctx->InputShape(1); + // Validate input sizes. + OP_REQUIRES(ctx, TensorShapeUtils::IsVector(ksize_shape), + errors::InvalidArgument("ksize must be a vector, not shape ", + ksize_shape.DebugString())); + OP_REQUIRES(ctx, ksize_shape.num_elements() == num_dims(), + errors::InvalidArgument("Sliding window ksize field must " + "specify ", + num_dims(), " dimensions")); + ksize.clear(); + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &ksize)); + + const TensorShape stride_shape = ctx->InputShape(2); + // Validate input sizes. + OP_REQUIRES(ctx, TensorShapeUtils::IsVector(stride_shape), + errors::InvalidArgument("stride must be a vector, not shape ", + stride_shape.DebugString())); + OP_REQUIRES(ctx, stride_shape.num_elements() == num_dims(), + errors::InvalidArgument("Sliding window stride field must " + "specify ", + num_dims(), " dimensions")); + stride.clear(); + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(2, &stride)); + } OP_REQUIRES(ctx, input_shape.dims() == num_dims(), errors::InvalidArgument("Input to ", type_string(), " operator must have ", num_dims(), @@ -84,8 +113,8 @@ class PoolingOp : public XlaOpKernel { const DataType type = input_type(0); xla::ComputationDataHandle pooled = ctx->builder()->ReduceWindow( - input, InitValue(ctx->builder(), type), *Reduction(ctx, type), ksize_, - stride_, padding_); + input, InitValue(ctx->builder(), type), *Reduction(ctx, type), ksize, + stride, padding_); ctx->SetOutput(0, PostProcessOutput(ctx, pooled, type, input_shape)); } @@ -130,6 +159,10 @@ class MaxPool2DOp : public MaxPoolOp { } }; REGISTER_XLA_OP(Name("MaxPool"), MaxPool2DOp); +REGISTER_XLA_OP(Name("MaxPoolV2") + .CompileTimeConstInput("ksize") + .CompileTimeConstInput("strides"), + MaxPool2DOp); class MaxPool3DOp : public MaxPoolOp { public: @@ -243,22 +276,44 @@ class MaxPoolGradOp : public XlaOpKernel { public: MaxPoolGradOp(OpKernelConstruction* ctx, int num_spatial_dims) : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) { - OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_)); + if (ctx->num_inputs() == 3) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("ksize", &ksize_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_)); + } + OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_)); + } + + int num_dims() const { return num_spatial_dims_ + 2; } + + void Compile(XlaOpKernelContext* ctx) override { + if (ctx->num_inputs() != 3) { + OP_REQUIRES( + ctx, ctx->num_inputs() == 5, + errors::InvalidArgument("Must supply ksize and stride arguments.")); + const TensorShape ksize_shape = ctx->InputShape(3); + // Validate input sizes. + OP_REQUIRES(ctx, TensorShapeUtils::IsVector(ksize_shape), + errors::InvalidArgument("ksize must be a vector, not shape ", + ksize_shape.DebugString())); + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(3, &ksize_)); + + const TensorShape stride_shape = ctx->InputShape(4); + // Validate input sizes. + OP_REQUIRES(ctx, TensorShapeUtils::IsVector(stride_shape), + errors::InvalidArgument("stride must be a vector, not shape ", + stride_shape.DebugString())); + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(4, &stride_)); + } + OP_REQUIRES(ctx, ksize_.size() == num_dims(), errors::InvalidArgument("Sliding window ksize field must " "specify ", num_dims(), " dimensions")); - OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &stride_)); OP_REQUIRES(ctx, stride_.size() == num_dims(), errors::InvalidArgument("Sliding window strides field must " "specify ", num_dims(), " dimensions")); - OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_)); - } - int num_dims() const { return num_spatial_dims_ + 2; } - - void Compile(XlaOpKernelContext* ctx) override { const TensorShape tensor_in_shape = ctx->InputShape(0); const TensorShape tensor_out_shape = ctx->InputShape(1); const TensorShape out_backprop_shape = ctx->InputShape(2); @@ -315,6 +370,10 @@ class MaxPool2DGradOp : public MaxPoolGradOp { } }; REGISTER_XLA_OP(Name("MaxPoolGrad"), MaxPool2DGradOp); +REGISTER_XLA_OP(Name("MaxPoolGradV2") + .CompileTimeConstInput("ksize") + .CompileTimeConstInput("strides"), + MaxPool2DGradOp); class MaxPool3DGradOp : public MaxPoolGradOp { public: diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_op.cc index 17a345fc942f78d161218d9a66e9583f3465a569..e51d386926763ecbb5a943dfb6f872e78901dc69 100644 --- a/tensorflow/compiler/tf2xla/kernels/reverse_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reverse_op.cc @@ -52,7 +52,8 @@ class ReverseOp : public XlaOpKernel { xla::Literal lax; OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped(1, {x_shape.dims()}, &lax)); std::vector revdims(x_shape.dims()); - std::copy(lax.preds().begin(), lax.preds().end(), revdims.begin()); + std::copy(lax.data().begin(), lax.data().end(), + revdims.begin()); std::vector dimensions; for (int d = 0; d < x_shape.dims(); ++d) { diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..6bc5d3adb091cd238974c5b69b7a2f8fe639cc68 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc @@ -0,0 +1,182 @@ +/* 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/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/framework/tensor_shape.h" + +namespace tensorflow { +namespace { + +class ReverseSequenceOp : public XlaOpKernel { + public: + explicit ReverseSequenceOp(OpKernelConstruction* context) + : XlaOpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("batch_dim", &batch_dim_)); + OP_REQUIRES_OK(context, context->GetAttr("seq_dim", &seq_dim_)); + } + + void Compile(XlaOpKernelContext* context) override { + const TensorShape input_shape = context->InputShape(0); + const TensorShape seq_lens_shape = context->InputShape(1); + + OP_REQUIRES(context, TensorShapeUtils::IsVector(seq_lens_shape), + errors::InvalidArgument("seq_lens input must be 1-dim, not ", + seq_lens_shape.dims())); + OP_REQUIRES(context, batch_dim_ != seq_dim_, + errors::InvalidArgument("batch_dim == seq_dim == ", seq_dim_)); + OP_REQUIRES( + context, seq_dim_ < input_shape.dims(), + errors::InvalidArgument("seq_dim must be < input.dims()", "( ", + seq_dim_, " vs. ", input_shape.dims(), ")")); + OP_REQUIRES( + context, batch_dim_ < input_shape.dims(), + errors::InvalidArgument("batch_dim must be < input.dims()", "( ", + batch_dim_, " vs. ", input_shape.dims(), ")")); + OP_REQUIRES( + context, + seq_lens_shape.num_elements() == input_shape.dim_size(batch_dim_), + errors::InvalidArgument("len(seq_lens) != input.dims(", batch_dim_, + "), ", "(", seq_lens_shape.num_elements(), + " vs. ", input_shape.dim_size(batch_dim_))); + + xla::ComputationBuilder* builder = context->builder(); + const auto input = context->Input(0); + const auto seq_lens = context->Input(1); + + const int64 batch_size = input_shape.dim_size(batch_dim_); + + const DataType input_type = context->input_type(0); + const DataType seq_lens_type = context->input_type(1); + const int64 max_seq_len = input_shape.dim_size(seq_dim_); + + xla::Shape input_xla_shape; + OP_REQUIRES_OK(context, TensorShapeToXLAShape(input_type, input_shape, + &input_xla_shape)); + xla::Shape seq_lens_xla_shape; + OP_REQUIRES_OK(context, TensorShapeToXLAShape(seq_lens_type, seq_lens_shape, + &seq_lens_xla_shape)); + + const auto tuple_shape = xla::ShapeUtil::MakeTupleShape({ + xla::ShapeUtil::MakeShape(seq_lens_xla_shape.element_type(), {}), + seq_lens_xla_shape, + input_xla_shape, + }); + + // For each entry in the batch, reverse the sequence. + // TODO(b/65689298): generalize the Map() operator to non-scalar cases and + // use it here, instead of a While loop. + + // Condition: lambda (i, _, _): i < batch_size + auto condition_builder = + builder->CreateSubBuilder("reverse_sequence_condition"); + { + auto param = condition_builder->Parameter(0, tuple_shape, "param"); + auto i = condition_builder->GetTupleElement(param, 0); + condition_builder->Lt( + i, XlaHelpers::IntegerLiteral(condition_builder.get(), seq_lens_type, + batch_size)); + } + auto condition = condition_builder->Build(); + OP_REQUIRES_OK(context, condition.status()); + + auto body_builder = builder->CreateSubBuilder("reverse_sequence_body"); + { + auto param = body_builder->Parameter(0, tuple_shape, "param"); + auto i = body_builder->GetTupleElement(param, 0); + auto seq_lens = body_builder->GetTupleElement(param, 1); + auto output = body_builder->GetTupleElement(param, 2); + + // seq_len is the sequence length of the current batch element (rank 1) + auto seq_len = body_builder->DynamicSlice( + seq_lens, body_builder->Reshape(i, {1}), {1}); + + // Indices is the offset of the batch element in the input. + auto indices = body_builder->Broadcast( + XlaHelpers::Zero(body_builder.get(), seq_lens_type), + {input_shape.dims()}); + indices = body_builder->DynamicUpdateSlice( + indices, body_builder->Reshape(i, {1}), + body_builder->Reshape( + XlaHelpers::IntegerLiteral(body_builder.get(), seq_lens_type, + batch_dim_), + {1})); + + // slice_indices is the offset of the start of the reversed sequence in + // the input. + auto slice_indices = body_builder->DynamicUpdateSlice( + indices, + body_builder->Sub(XlaHelpers::IntegerLiteral( + body_builder.get(), seq_lens_type, max_seq_len), + seq_len), + body_builder->Reshape( + XlaHelpers::IntegerLiteral(body_builder.get(), seq_lens_type, + seq_dim_), + {1})); + + // Slice out the reversed sequence. The slice will overflow the end of the + // sequence, and the contents of the overflow are implementation-defined. + // However, we will mask off these elements and replace them with elements + // from the original input so their values do not matter. + TensorShape slice_shape = input_shape; + slice_shape.set_dim(batch_dim_, 1); + auto slice = body_builder->DynamicSlice(output, slice_indices, + slice_shape.dim_sizes()); + + // Shift the reversed sequence to the left. + output = body_builder->DynamicUpdateSlice(output, slice, indices); + + body_builder->Tuple( + {body_builder->Add( + i, XlaHelpers::One(body_builder.get(), seq_lens_type)), + seq_lens, output}); + } + auto body = body_builder->Build(); + OP_REQUIRES_OK(context, body.status()); + + auto loop_output = builder->While( + condition.ValueOrDie(), body.ValueOrDie(), + builder->Tuple({XlaHelpers::Zero(builder, seq_lens_type), seq_lens, + builder->Rev(input, {seq_dim_})})); + auto output = builder->GetTupleElement(loop_output, 2); + + // Mask out elements after the sequence length. + xla::ComputationDataHandle iota; + OP_REQUIRES_OK( + context, XlaHelpers::Iota(builder, seq_lens_type, max_seq_len, &iota)); + std::vector dims(input_shape.dims(), 1); + dims[batch_dim_] = batch_size; + auto mask = builder->Lt(iota, builder->Reshape(seq_lens, dims), {seq_dim_}); + + // Broadcast the mask up to the input shape. + mask = + builder->Or(mask, builder->Broadcast(builder->ConstantR0(false), + input_shape.dim_sizes())); + + output = builder->Select(mask, output, input); + context->SetOutput(0, output); + } + + private: + int32 batch_dim_; + int32 seq_dim_; +}; + +REGISTER_XLA_OP(Name("ReverseSequence"), ReverseSequenceOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..8433a29c4e203cac726ee6bf7f67a863447326ed --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc @@ -0,0 +1,121 @@ +/* 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/compiler/tf2xla/lib/scatter.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/type_util.h" +#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/framework/kernel_def_builder.h" +#include "tensorflow/core/framework/op_kernel.h" + +namespace tensorflow { +namespace { + +// Check whether updates.shape = indices.shape[:batch_dim] + +// buffer_shape[num_index_dims:] +Status ValidateUpdateShape(const TensorShape& buffer_shape, + const TensorShape& indices_shape, + const TensorShape& updates_shape) { + if (indices_shape.dims() < 1) { + return errors::InvalidArgument( + "indices shape must have >= 1 dimension; got ", + indices_shape.DebugString()); + } + + const int64 num_index_dims = indices_shape.dim_size(indices_shape.dims() - 1); + const int64 batch_dim = indices_shape.dims() - 1; + + auto shape_err = [&]() { + return errors::InvalidArgument( + "Must have updates.shape = indices.shape[:batch_dim] + ", + "buffer_shape[num_index_dims:], got updates.shape: ", + updates_shape.DebugString(), + ", indices.shape: ", indices_shape.DebugString(), + ", buffer_shape: ", buffer_shape.DebugString(), + ", num_index_dims: ", num_index_dims, ", and batch_dim: ", batch_dim); + }; + + if (updates_shape.dims() < batch_dim) return shape_err(); + if (buffer_shape.dims() < + num_index_dims + (updates_shape.dims() - batch_dim)) { + return shape_err(); + } + if (updates_shape.dims() != + batch_dim + buffer_shape.dims() - num_index_dims) { + return shape_err(); + } + for (int d = 0; d < batch_dim; ++d) { + if (updates_shape.dim_size(d) != indices_shape.dim_size(d)) { + return shape_err(); + } + } + for (int d = 0; d < updates_shape.dims() - batch_dim; ++d) { + if (updates_shape.dim_size(d + batch_dim) != + buffer_shape.dim_size(d + num_index_dims)) { + return shape_err(); + } + } + return Status::OK(); +} + +class ScatterNdOp : public XlaOpKernel { + public: + explicit ScatterNdOp(OpKernelConstruction* context) : XlaOpKernel(context) {} + + void Compile(XlaOpKernelContext* context) override { + DataType dtype = context->input_type(1); + + TensorShape indices_shape = context->InputShape(0); + TensorShape updates_shape = context->InputShape(1); + + TensorShape buffer_shape; + OP_REQUIRES_OK(context, context->ConstantInputAsShape(2, &buffer_shape)); + + OP_REQUIRES( + context, TensorShapeUtils::IsVectorOrHigher(buffer_shape), + errors::InvalidArgument("Output must be at least 1-D, ", + "got shape: ", buffer_shape.DebugString())); + + OP_REQUIRES( + context, + buffer_shape.num_elements() > 0 || (indices_shape.num_elements() == 0 && + updates_shape.num_elements() == 0), + errors::InvalidArgument( + "Indices and updates specified for empty output. indices shape: ", + indices_shape.DebugString())); + + OP_REQUIRES_OK(context, ValidateUpdateShape(buffer_shape, indices_shape, + updates_shape)); + + xla::ComputationBuilder* builder = context->builder(); + auto buffer = builder->Broadcast(XlaHelpers::Zero(builder, dtype), + buffer_shape.dim_sizes()); + auto indices = context->Input(0); + auto updates = context->Input(1); + auto result = + XlaScatter(buffer, updates, indices, + /*indices_are_vectors=*/true, /*combiner=*/{}, builder); + OP_REQUIRES_OK(context, result.status()); + context->SetOutput(0, result.ValueOrDie()); + } +}; + +REGISTER_XLA_OP(Name("ScatterNd").CompileTimeConstInput("shape"), ScatterNdOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/scatter_op_helpers.h b/tensorflow/compiler/tf2xla/kernels/scatter_op_helpers.h deleted file mode 100644 index a5ab7de17adb734014fe2dcbd60ae5c219c8e486..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/kernels/scatter_op_helpers.h +++ /dev/null @@ -1,39 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ -// Helper methods for XLA Scatter Ops. -#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_SCATTER_OP_HELPERS_H_ -#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_SCATTER_OP_HELPERS_H_ - -#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" -#include "tensorflow/compiler/xla/client/client_library.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" -#include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/util/bcast.h" - -namespace tensorflow { - -// Adds to builder an XLA computation that performs a scatter-add of input (of -// shape input_shape) keyed on indices (of shape indices_shape). The shape -// of the Tensor returned by this is num_segments input_shape[indices.dims():] -// -static xla::ComputationDataHandle XlaComputeScatterAddDynamicSlice( - XlaOpKernelContext* ctx, const xla::ComputationDataHandle& input, - const TensorShape& input_shape, const xla::ComputationDataHandle& indices, - const TensorShape& indices_shape, int64 num_segments, DataType dtype, - xla::ComputationBuilder* builder); - -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_SCATTER_OP_HELPERS_H_ diff --git a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc index 8a67c0b67fcd95f4841c5e011a4e51638eea5b0f..80d6df6c48b0141734dcee1c2a3c413926931feb 100644 --- a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* 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. @@ -13,143 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include -#include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h" -#include "tensorflow/compiler/tf2xla/shape_util.h" -#include "tensorflow/compiler/tf2xla/type_util.h" +#include "tensorflow/compiler/tf2xla/lib/scatter.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/computation_builder.h" -#include "tensorflow/compiler/xla/literal_util.h" -#include "tensorflow/core/framework/kernel_def_builder.h" -#include "tensorflow/core/framework/types.h" namespace tensorflow { - -xla::ComputationDataHandle XlaComputeScatterAddDynamicSlice( - XlaOpKernelContext* ctx, const xla::ComputationDataHandle& input, - const TensorShape& input_shape, const xla::ComputationDataHandle& indices, - const TensorShape& indices_shape, int64 num_segments, DataType dtype, - xla::ComputationBuilder* builder) { - // Flatten data for dynamic indexing via indices_1d. - TensorShape input_shape_i(input_shape); - for (int64 d = 0; d < indices_shape.dims(); ++d) { - input_shape_i.RemoveDim(0); - } - TensorShape flat_shape({indices_shape.num_elements()}); - flat_shape.AppendShape(input_shape_i); - - // output is same as flattened input shape with dim_size(0) = num_segments. - TensorShape out_shape(flat_shape); - out_shape.set_dim(0, num_segments); - - // TODO(b/37575001) The tensor in which we construct the output during - // the loop must have rank >= 3 as a workaround for lowering issues. - int64 extra_dims = 0; - if (out_shape.dims() < 3) { - extra_dims = 3 - out_shape.dims(); - } - TensorShape loop_out_shape; - for (int64 k = 0; k < extra_dims; ++k) { - loop_out_shape.AddDim(1); - } - loop_out_shape.AppendShape(out_shape); - - // Slices from the input data are same shape as the input data, except dim 0. - TensorShape slice_shape(flat_shape); - slice_shape.set_dim(0, 1); - // slices are reshaped into the rank >= 3 shape of the loop-carried output - TensorShape loop_out_slice_shape(loop_out_shape); - loop_out_slice_shape.set_dim(extra_dims, 1); - - // Construct the initial values of the loop-carried variables - // Flatten the indices into 1-D for ease of iteration. - auto indices_1d = builder->Reshape(indices, {indices_shape.num_elements()}); - // Flatten the data for ease of indexing via values in indices_1d. - auto data_flat = builder->Reshape(input, flat_shape.dim_sizes()); - - auto init_i = builder->ConstantR0(0); - auto init_out = builder->Broadcast(XlaHelpers::Zero(builder, dtype), - loop_out_shape.dim_sizes()); - - xla::PrimitiveType ptype; - TF_CHECK_OK(DataTypeToPrimitiveType(dtype, &ptype)); - - std::vector tuple_shapes( - {// The loop iteration counter is a scalar, incremented each iteration. - xla::ShapeUtil::MakeShape(xla::S32, {}), - // The flattened input data is loop invariant. - xla::ShapeUtil::MakeShape(ptype, flat_shape.dim_sizes()), - // The scatter indices tensor is loop invariant. - xla::ShapeUtil::MakeShape(xla::S32, {indices_shape.num_elements()}), - // The output data array is updated each loop iteration. - xla::ShapeUtil::MakeShape(ptype, loop_out_shape.dim_sizes())}); - xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes); - - auto init = builder->Tuple({init_i, data_flat, indices_1d, init_out}); - - // Construct the while loop condition (i < num_indices) - xla::ComputationBuilder condb(ctx->builder()->client(), - "ScatterAddWhileCond"); - condb.Lt(condb.GetTupleElement( - condb.Parameter(0, tuple_shape, "ScatterAddWhileTuple"), 0), - condb.ConstantR0(indices_shape.num_elements())); - auto cond_status = condb.Build(); - // TF_CHECK_OK(cond_status); - auto cond = cond_status.ConsumeValueOrDie(); - - // Construct the while loop body's function. The implementation of scatter is: - // for i in range(num_indices): - // index = dynamic-slice(indices, i) - // xi = dynamic-slice(input, i) - // output = dynamic-update-slice(output, xi, index) - xla::ComputationBuilder bodyb(ctx->builder()->client(), - "ScatterAddWhileBody"); - { - auto input_tuple = bodyb.Parameter(0, tuple_shape, "ScatterAddWhileTuple"); - auto i = bodyb.GetTupleElement(input_tuple, 0); - auto data = bodyb.GetTupleElement(input_tuple, 1); - auto idcs = bodyb.GetTupleElement(input_tuple, 2); - auto output = bodyb.GetTupleElement(input_tuple, 3); - - // Index into the data array at i. - auto zero = bodyb.ConstantR1({0}); - std::vector index_vals(flat_shape.dims(), zero); - index_vals[0] = bodyb.Reshape(i, {1}); - auto index = bodyb.ConcatInDim(index_vals, 0); - - auto data_slice = - bodyb.Reshape(bodyb.DynamicSlice(data, index, slice_shape.dim_sizes()), - loop_out_slice_shape.dim_sizes()); - - // Index into the output array. - // Construct the index into the R3+ output array 0, ..., , 0, ... - std::vector out_index_vals( - loop_out_shape.dims(), zero); - out_index_vals[extra_dims] = - bodyb.DynamicSlice(idcs, bodyb.Reshape(i, {1}), {1}); - auto out_index = bodyb.ConcatInDim(out_index_vals, 0); - - // Slice the output array, update value, and update the output slice. - auto updated_output = bodyb.DynamicUpdateSlice( - output, - bodyb.Add(data_slice, - bodyb.DynamicSlice(output, out_index, - loop_out_slice_shape.dim_sizes())), - out_index); - - auto ip1 = bodyb.Add(i, bodyb.ConstantR0(1)); - bodyb.Tuple({ip1, data, idcs, updated_output}); - } - auto body_status = bodyb.Build(); - // TF_CHECK_OK(body_status); - auto body = body_status.ConsumeValueOrDie(); - - auto gather_while = builder->While(cond, body, init); - auto updated_output = builder->GetTupleElement(gather_while, 3); - return builder->Reshape(updated_output, out_shape.dim_sizes()); -} - namespace { class UnsortedSegmentSum : public XlaOpKernel { @@ -171,10 +41,10 @@ class UnsortedSegmentSum : public XlaOpKernel { // as data with the first indices.rank dimensions are replaced // by a single dimension with size num_segments. auto data = ctx->Input(0); - auto data_shape = ctx->InputShape(0); + TensorShape data_shape = ctx->InputShape(0); auto indices = ctx->Input(1); - auto indices_shape = ctx->InputShape(1); + TensorShape indices_shape = ctx->InputShape(1); int64 num_segments; OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(2, &num_segments)); @@ -192,10 +62,21 @@ class UnsortedSegmentSum : public XlaOpKernel { d, " differs ", data_shape.dim_size(d), " vs. ", indices_shape.dim_size(d))); } - auto result = XlaComputeScatterAddDynamicSlice( - ctx, data, data_shape, indices, indices_shape, num_segments, dtype_, - ctx->builder()); - ctx->SetOutput(0, result); + xla::ComputationBuilder* builder = ctx->builder(); + TensorShape buffer_shape = data_shape; + buffer_shape.RemoveDimRange(0, indices_shape.dims()); + buffer_shape.InsertDim(0, num_segments); + auto buffer = builder->Broadcast(XlaHelpers::Zero(builder, dtype_), + buffer_shape.dim_sizes()); + + auto combiner = + [](xla::ComputationDataHandle a, xla::ComputationDataHandle b, + xla::ComputationBuilder* builder) { return builder->Add(a, b); }; + + auto result = XlaScatter(buffer, /*updates=*/data, indices, + /*indices_are_vectors=*/false, combiner, builder); + OP_REQUIRES_OK(ctx, result.status()); + ctx->SetOutput(0, result.ValueOrDie()); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/shape_op.cc b/tensorflow/compiler/tf2xla/kernels/shape_op.cc index 8fb7a74310b6760af8e6893ea2e3e3868c85f536..05354bca5bb089703fdcceb6f44648bbb98d004b 100644 --- a/tensorflow/compiler/tf2xla/kernels/shape_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/shape_op.cc @@ -121,7 +121,7 @@ class ExpandDimsOp : public XlaOpKernel { xla::Literal literal; OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped(1, {1}, &literal)); - int dim = literal.s32s(0); + int dim = literal.data()[0]; OP_REQUIRES(ctx, (dim >= -1 - input_shape.dims() && dim <= input_shape.dims()), diff --git a/tensorflow/compiler/tf2xla/kernels/shape_util.h b/tensorflow/compiler/tf2xla/kernels/shape_util.h index 575086e118080f6799a54d3ae6409b2b641c4341..ca57be3d47b95d71b07746e50256070e0a4f4c09 100644 --- a/tensorflow/compiler/tf2xla/kernels/shape_util.h +++ b/tensorflow/compiler/tf2xla/kernels/shape_util.h @@ -31,4 +31,4 @@ Status TensorShapeToConstant(const TensorShape& input_shape, } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_TF2XLA_KERNELS_SHAPE_UTIL_H_ +#endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_SHAPE_UTIL_H_ diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc index c912876e6585758935c2b78450f234b9cba49010..1a78c7ab9be701d3d02285ed21604f0f856b3f1f 100644 --- a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc @@ -77,10 +77,8 @@ Status MaybeInitializeStack(xla::ComputationBuilder* builder, // Stack has not been initialized. xla::ComputationDataHandle zero = XlaHelpers::Zero(builder, resource->type()); - TF_RETURN_IF_ERROR(resource->SetValue( - dtype, - builder->Tuple({builder->Broadcast(zero, stack_shape.dim_sizes()), - builder->ConstantR0(0)}))); + TF_RETURN_IF_ERROR(resource->SetTypeAndShape(dtype, elem_shape)); + TF_RETURN_IF_ERROR(resource->SetZeroValue(builder)); } else { // Checks the expected shape matches the actual shape. TensorShape actual_shape; @@ -107,7 +105,9 @@ class StackOp : public XlaOpKernel { OP_REQUIRES( ctx, size >= 0, errors::InvalidArgument( - "XLA compilation requires a fixed stack size upper bound.")); + "XLA compilation requires a fixed stack size upper bound. If " + "you are using tf.while_loop, set the maximum_iterations parameter " + "to fix this issue.")); // We defer initializing the Stack resource until we see the first push. // Otherwise we do not know the shape of the stack elements. @@ -117,8 +117,8 @@ class StackOp : public XlaOpKernel { string name = strings::StrCat("Stack: ", stack_name_); OP_REQUIRES_OK( ctx, xc.CreateResource(XlaResource::kStack, -1, std::move(name), dtype_, - value, &resource)); - resource->set_tensor_array_size(size); + TensorShape(), value, /*tensor_array_size=*/size, + /*tensor_array_gradients=*/{}, &resource)); ctx->SetResourceOutput(0, resource); } @@ -162,11 +162,9 @@ class StackPushOp : public XlaOpKernel { // TODO(phawkins): We don't check the index is in bounds --- there is no // error mechanism in XLA. - OP_REQUIRES_OK( - ctx, - resource->SetValue( - dtype_, b->Tuple({b->DynamicUpdateSlice(ta, update, start_indices), - b->Add(index, b->ConstantR0(1))}))); + OP_REQUIRES_OK(ctx, resource->SetValue(b->Tuple( + {b->DynamicUpdateSlice(ta, update, start_indices), + b->Add(index, b->ConstantR0(1))}))); ctx->SetOutput(0, value); } @@ -206,7 +204,7 @@ class StackPopOp : public XlaOpKernel { xla::ComputationDataHandle index = b->GetTupleElement(state, 1); index = b->Sub(index, b->ConstantR0(1)); - OP_REQUIRES_OK(ctx, resource->SetValue(dtype_, b->Tuple({ta, index}))); + OP_REQUIRES_OK(ctx, resource->SetValue(b->Tuple({ta, index}))); // start_indices of the DynamicSlice are [index, 0, 0, ..., 0]. auto start_indices = diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc index f0525a5fb86d6d6f0aae954a916186cffc7f3a9f..91c169428c7a88a8d107a97445aeea999946e3e9 100644 --- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc @@ -231,6 +231,7 @@ class StridedSliceAssignOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("new_axis_mask", &new_axis_mask_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("shrink_axis_mask", &shrink_axis_mask_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("Index", &index_type_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); } void Compile(XlaOpKernelContext* ctx) override { @@ -252,9 +253,9 @@ class StridedSliceAssignOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, LiteralToHostTensor(strides_literal, index_type_, &strides_tensor)); - DataType lhs_type; TensorShape lhs_shape; - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &lhs_type, &lhs_shape)); + xla::ComputationDataHandle lhs; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, dtype_, &lhs_shape, &lhs)); const TensorShape rhs_shape = ctx->InputShape(4); @@ -282,9 +283,6 @@ class StridedSliceAssignOp : public XlaOpKernel { " does not match r-value shape ", rhs_shape.DebugString(), ". Automatic broadcasting not yet implemented.")); - xla::ComputationDataHandle lhs; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &lhs)); - xla::ComputationDataHandle rhs = ctx->Input(4); gtl::InlinedVector dimensions_to_reverse; @@ -320,13 +318,14 @@ class StridedSliceAssignOp : public XlaOpKernel { lhs, rhs, ctx->builder()->ConstantR1(slice_begin)); } - OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, lhs_type, lhs)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, dtype_, lhs)); } private: int32 begin_mask_, end_mask_; int32 ellipsis_mask_, new_axis_mask_, shrink_axis_mask_; DataType index_type_; + DataType dtype_; }; REGISTER_XLA_OP(Name("ResourceStridedSliceAssign") diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc index 9224072a3cb92b8ff0e99c79e568ca1a76966ed6..000b50af6bd86b7268c016865fb0856c16053ece 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc @@ -62,15 +62,13 @@ Status MaybeInitializeTensorArray(xla::ComputationBuilder* builder, TF_RET_CHECK(resource->tensor_array_size() >= 0) << resource->name() << " size " << resource->tensor_array_size(); - TensorShape ta_shape; - ta_shape.AddDim(resource->tensor_array_size()); - ta_shape.AppendShape(elem_shape); if (!resource->initialized()) { xla::ComputationDataHandle zero = XlaHelpers::Zero(builder, resource->type()); - TF_RETURN_IF_ERROR(resource->SetValue( - dtype, builder->Broadcast(zero, ta_shape.dim_sizes()))); + + TF_RETURN_IF_ERROR(resource->SetTypeAndShape(dtype, elem_shape)); + TF_RETURN_IF_ERROR(resource->SetZeroValue(builder)); } else { // Checks the elem_shape matches the TensorArray shape. auto shape_or_status = builder->GetShape(resource->value()); @@ -80,6 +78,10 @@ Status MaybeInitializeTensorArray(xla::ComputationBuilder* builder, TensorShape shape; TF_RETURN_IF_ERROR( XLAShapeToTensorShape(*shape_or_status.ValueOrDie(), &shape)); + + TensorShape ta_shape; + ta_shape.AddDim(resource->tensor_array_size()); + ta_shape.AppendShape(elem_shape); if (ta_shape != shape) { return errors::InvalidArgument( "Mismatched TensorArray sizes: ", ta_shape.DebugString(), " vs ", @@ -114,10 +116,8 @@ Status CheckTensorArrayIsInitialized(const string& op_name, Status GetTensorArrayShape(const XlaResource* resource, xla::ComputationBuilder* builder, TensorShape* shape) { - TF_RETURN_IF_ERROR(resource->GetShape(builder, shape)); - if (shape->dims() < 1) { - return errors::InvalidArgument("TensorArray rank must be >= 1"); - } + *shape = resource->shape(); + shape->InsertDim(0, resource->tensor_array_size()); return Status::OK(); } @@ -160,8 +160,8 @@ class TensorArrayOp : public XlaOpKernel { // Initializes the TensorArray value if we know the element shape. // Otherwise, defer initialization to the first write. xla::ComputationDataHandle value; + TensorShape shape; if (element_shape_.IsFullyDefined()) { - TensorShape shape; CHECK(element_shape_.AsTensorShape(&shape)); TensorShape ta_shape; ta_shape.AddDim(size); @@ -175,8 +175,8 @@ class TensorArrayOp : public XlaOpKernel { string name = strings::StrCat("TensorArray: ", tensor_array_name_); OP_REQUIRES_OK( ctx, xc.CreateResource(XlaResource::kTensorArray, -1, std::move(name), - dtype_, value, &var)); - var->set_tensor_array_size(size); + dtype_, shape, value, /*tensor_array_size=*/size, + /*tensor_array_gradients=*/{}, &var)); ctx->SetResourceOutput(0, var); Tensor flow(DT_FLOAT, TensorShape({})); @@ -230,7 +230,7 @@ class TensorArrayWriteOp : public XlaOpKernel { xla::ComputationDataHandle written = DynamicAddSlice(b, ta, update, slice_shape.dim_sizes(), start_indices); - OP_REQUIRES_OK(ctx, resource->SetValue(dtype_, written)); + OP_REQUIRES_OK(ctx, resource->SetValue(written)); ctx->SetOutput(0, flow); } @@ -337,8 +337,11 @@ class TensorArrayGatherOp : public XlaOpKernel { } } - xla::ComputationDataHandle gather = XlaComputeGatherDynamicSlice( - ctx, ta, ta_shape, indices, indices_shape, 0, dtype_, index_type, b); + xla::ComputationDataHandle gather; + OP_REQUIRES_OK( + ctx, + XlaGather(ta, ta_shape, indices, indices_shape, /*axis=*/0, + /*indices_are_nd=*/false, dtype_, index_type, b, &gather)); ctx->SetOutput(0, gather); } @@ -421,7 +424,7 @@ class TensorArrayScatterOp : public XlaOpKernel { } } - OP_REQUIRES_OK(ctx, resource->SetValue(dtype_, ta)); + OP_REQUIRES_OK(ctx, resource->SetValue(ta)); ctx->SetOutput(0, flow); } @@ -525,9 +528,8 @@ class TensorArraySplitOp : public XlaOpKernel { value_shape.DebugString(), " vs. ", ta_shape.DebugString())); - OP_REQUIRES_OK( - ctx, resource->SetValue( - dtype_, b->Add(ta, b->Reshape(value, ta_shape.dim_sizes())))); + OP_REQUIRES_OK(ctx, resource->SetValue(b->Add( + ta, b->Reshape(value, ta_shape.dim_sizes())))); ctx->SetOutput(0, flow); } diff --git a/tensorflow/compiler/tf2xla/kernels/training_ops.cc b/tensorflow/compiler/tf2xla/kernels/training_ops.cc index 5534d1bfa1338c7fe3647cd6aa281c4907dfdf8c..f750f7003be288461f5f10455e58932d1b4e4524 100644 --- a/tensorflow/compiler/tf2xla/kernels/training_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/training_ops.cc @@ -32,9 +32,24 @@ class ResourceApplyGradientDescent : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { xla::ComputationDataHandle handle; xla::ComputationBuilder* b = ctx->builder(); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &handle)); + DataType type = ctx->input_type(1); + TensorShape var_shape; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &handle)); + + TensorShape alpha_shape = ctx->InputShape(1); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(alpha_shape), + errors::InvalidArgument("alpha is not a scalar: ", + alpha_shape.DebugString())); + + TensorShape delta_shape = ctx->InputShape(2); + OP_REQUIRES( + ctx, var_shape.IsSameSize(delta_shape), + errors::InvalidArgument("var and delta do not have the same shape: ", + var_shape.DebugString(), " vs ", + delta_shape.DebugString())); + handle = b->Sub(handle, b->Mul(ctx->Input(1), ctx->Input(2))); - OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, ctx->input_type(1), handle)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, handle)); } }; REGISTER_XLA_OP( @@ -52,18 +67,10 @@ class ResourceApplyMomentum : public XlaOpKernel { DataType type = ctx->input_type(2); - DataType var_type, accum_type; TensorShape var_shape, accum_shape; - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &var_type, &var_shape)); - OP_REQUIRES_OK(ctx, - ctx->GetVariableTypeAndShape(1, &accum_type, &accum_shape)); - - OP_REQUIRES( - ctx, type == var_type && type == accum_type, - errors::InvalidArgument( - "Types of variable arguments to ResourceApplyMomentum must match: ", - DataTypeString(type), " vs. ", DataTypeString(var_type), " and ", - DataTypeString(accum_type))); + xla::ComputationDataHandle var, accum; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &var)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &accum_shape, &accum)); OP_REQUIRES(ctx, var_shape.IsSameSize(accum_shape), errors::InvalidArgument( @@ -86,10 +93,6 @@ class ResourceApplyMomentum : public XlaOpKernel { errors::InvalidArgument("momentum is not a scalar: ", momentum_shape.DebugString())); - xla::ComputationDataHandle var, accum; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &var)); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, &accum)); - xla::ComputationDataHandle lr = ctx->Input(2); xla::ComputationDataHandle grad = ctx->Input(3); xla::ComputationDataHandle momentum = ctx->Input(4); @@ -122,18 +125,10 @@ class ResourceApplyAdagrad : public XlaOpKernel { DataType type = ctx->input_type(2); - DataType var_type, accum_type; TensorShape var_shape, accum_shape; - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &var_type, &var_shape)); - OP_REQUIRES_OK(ctx, - ctx->GetVariableTypeAndShape(1, &accum_type, &accum_shape)); - - OP_REQUIRES( - ctx, type == var_type && type == accum_type, - errors::InvalidArgument( - "Types of variable arguments to ResourceApplyAdagrad must match: ", - DataTypeString(type), " vs. ", DataTypeString(var_type), " and ", - DataTypeString(accum_type))); + xla::ComputationDataHandle var, accum; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &var)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &accum_shape, &accum)); OP_REQUIRES(ctx, var_shape.IsSameSize(accum_shape), errors::InvalidArgument( @@ -151,9 +146,6 @@ class ResourceApplyAdagrad : public XlaOpKernel { "var and grad do not have the same shape", var_shape.DebugString(), " ", grad_shape.DebugString())); - xla::ComputationDataHandle var, accum; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &var)); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, &accum)); xla::ComputationDataHandle lr = ctx->Input(2); xla::ComputationDataHandle grad = ctx->Input(3); @@ -175,18 +167,11 @@ class ResourceApplyAdam : public XlaOpKernel { } void Compile(XlaOpKernelContext* ctx) override { - DataType var_type, m_type, v_type; TensorShape var_shape, m_shape, v_shape; - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &var_type, &var_shape)); - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(1, &m_type, &m_shape)); - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(2, &v_type, &v_shape)); - - OP_REQUIRES( - ctx, dtype_ == var_type && dtype_ == m_type && dtype_ == v_type, - errors::InvalidArgument( - "Types of variable arguments to ResourceApplyRMSProp must match: ", - DataTypeString(dtype_), " vs. ", DataTypeString(var_type), " vs. ", - DataTypeString(m_type), " vs. ", DataTypeString(v_type))); + xla::ComputationDataHandle var, m, v; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, dtype_, &var_shape, &var)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, dtype_, &m_shape, &m)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, dtype_, &v_shape, &v)); TensorShape beta1_power_shape = ctx->InputShape(3); TensorShape beta2_power_shape = ctx->InputShape(4); @@ -228,10 +213,6 @@ class ResourceApplyAdam : public XlaOpKernel { "var and grad do not have the same shape", var_shape.DebugString(), " ", grad_shape.DebugString())); - xla::ComputationDataHandle var, m, v; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &var)); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, &m)); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, &v)); xla::ComputationDataHandle beta1_power = ctx->Input(3); xla::ComputationDataHandle beta2_power = ctx->Input(4); xla::ComputationDataHandle lr = ctx->Input(5); @@ -278,18 +259,11 @@ class ResourceApplyRMSProp : public XlaOpKernel { DataType type = ctx->input_type(3); - DataType var_type, ms_type, mom_type; TensorShape var_shape, ms_shape, mom_shape; - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &var_type, &var_shape)); - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(1, &ms_type, &ms_shape)); - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(2, &mom_type, &mom_shape)); - - OP_REQUIRES( - ctx, type == var_type && type == ms_type && type == mom_type, - errors::InvalidArgument( - "Types of variable arguments to ResourceApplyRMSProp must match: ", - DataTypeString(type), " vs. ", DataTypeString(var_type), " vs. ", - DataTypeString(ms_type), " vs. ", DataTypeString(mom_type))); + xla::ComputationDataHandle var, ms, mom; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &var)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &ms_shape, &ms)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, type, &mom_shape, &mom)); TensorShape lr_shape = ctx->InputShape(3); OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(lr_shape), @@ -323,10 +297,6 @@ class ResourceApplyRMSProp : public XlaOpKernel { "var and grad do not have the same shape", var_shape.DebugString(), " ", grad_shape.DebugString())); - xla::ComputationDataHandle var, ms, mom; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &var)); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, &ms)); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, &mom)); xla::ComputationDataHandle lr = ctx->Input(3); xla::ComputationDataHandle rho = ctx->Input(4); xla::ComputationDataHandle momentum = ctx->Input(5); @@ -373,20 +343,11 @@ void CompileFtrl(XlaOpKernelContext* ctx, DataType dtype, bool has_l2_shrinkage) { xla::ComputationBuilder* b = ctx->builder(); - DataType var_type, accum_type, linear_type; TensorShape var_shape, accum_shape, linear_shape; - OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &var_type, &var_shape)); - OP_REQUIRES_OK(ctx, - ctx->GetVariableTypeAndShape(1, &accum_type, &accum_shape)); - OP_REQUIRES_OK(ctx, - ctx->GetVariableTypeAndShape(2, &linear_type, &linear_shape)); - - OP_REQUIRES( - ctx, dtype == var_type && dtype == accum_type && dtype == linear_type, - errors::InvalidArgument( - "Types of variable arguments to ResourceApplyFtrlV2 must match: ", - DataTypeString(dtype), " vs. ", DataTypeString(var_type), " and ", - DataTypeString(accum_type), " and ", DataTypeString(linear_type))); + xla::ComputationDataHandle var, accum, linear; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, dtype, &var_shape, &var)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, dtype, &accum_shape, &accum)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, dtype, &linear_shape, &linear)); OP_REQUIRES(ctx, var_shape.IsSameSize(accum_shape), errors::InvalidArgument( @@ -438,10 +399,6 @@ void CompileFtrl(XlaOpKernelContext* ctx, DataType dtype, errors::InvalidArgument("lr_power is not a scalar: ", lr_power_shape.DebugString())); - xla::ComputationDataHandle var, accum, linear; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &var)); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, &accum)); - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, &linear)); xla::ComputationDataHandle grad = ctx->Input(3); xla::ComputationDataHandle lr = ctx->Input(4); xla::ComputationDataHandle l1 = ctx->Input(5); diff --git a/tensorflow/compiler/tf2xla/kernels/transpose_op.cc b/tensorflow/compiler/tf2xla/kernels/transpose_op.cc index 5c17b7fbf01990704c47ac564e69ba6d479a902e..c167642174b328a968d7f7ce1f0ad6e0ab8a7a68 100644 --- a/tensorflow/compiler/tf2xla/kernels/transpose_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/transpose_op.cc @@ -54,7 +54,8 @@ class TransposeOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped(1, {dims}, &literal)); std::vector perm(dims); - std::copy(literal.s32s().begin(), literal.s32s().end(), perm.begin()); + std::copy(literal.data().begin(), literal.data().end(), + perm.begin()); std::vector transposed_order; // Check whether permutation is a permutation of integers of [0 .. dims). diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc index a266e9013c41b88788dbc99849f01c09f3d61348..0c5ad9e5255ffc3dfcfb83335060ae833937b3ce 100644 --- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc @@ -50,18 +50,41 @@ XLAJIT_MAKE_UNARY(Conj, b->Conj(x)); // Return x if x>0, otherwise -x. XLAJIT_MAKE_UNARY(Abs, b->Abs(x)); +// acos(x) = 2 * atan(sqrt(1 - x^2) / (1 + x)) +XLAJIT_MAKE_UNARY( + Acos, + b->Mul(XlaHelpers::FloatLiteral(b, input_type(0), 2.0), + b->Atan2(b->Pow(b->Sub(XlaHelpers::One(b, input_type(0)), + b->Mul(x, x)), + XlaHelpers::FloatLiteral(b, input_type(0), 0.5)), + b->Add(XlaHelpers::One(b, input_type(0)), x)))); + // acosh(x) = log(x + sqrt(x^2 - 1)) XLAJIT_MAKE_UNARY( Acosh, b->Log(b->Add(x, b->Pow(b->Sub(b->Mul(x, x), XlaHelpers::One(b, input_type(0))), XlaHelpers::FloatLiteral(b, input_type(0), 0.5))))); + +// asin(x) = 2 * atan(x / (1 + sqrt(1 - x^2))) +XLAJIT_MAKE_UNARY( + Asin, + b->Mul(XlaHelpers::FloatLiteral(b, input_type(0), 2.0), + b->Atan2(x, b->Add(XlaHelpers::One(b, input_type(0)), + b->Pow(b->Sub(XlaHelpers::One(b, input_type(0)), + b->Mul(x, x)), + XlaHelpers::FloatLiteral(b, input_type(0), + 0.5)))))); + // asinh(x) = log(x + sqrt(x^2 + 1)) XLAJIT_MAKE_UNARY( Asinh, b->Log(b->Add(x, b->Pow(b->Add(b->Mul(x, x), XlaHelpers::One(b, input_type(0))), XlaHelpers::FloatLiteral(b, input_type(0), 0.5))))); + +XLAJIT_MAKE_UNARY(Atan, b->Atan2(x, XlaHelpers::One(b, input_type(0)))); + // atanh(x) = 0.5 * log((1 + x) / (1 - x)) XLAJIT_MAKE_UNARY( Atanh, b->Mul(b->Log(b->Div(b->Add(XlaHelpers::One(b, input_type(0)), x), diff --git a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc index 68847ae7a2cb926edd9d29007e24b0db7fb5a75f..71173f5aead47702f0ed9e95b827a6fefd9b7efd 100644 --- a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc @@ -33,21 +33,29 @@ class VarIsInitializedOp : public XlaOpKernel { public: explicit VarIsInitializedOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { - xla::ComputationDataHandle handle; - bool initialized = ctx->ReadVariableInput(0, &handle).ok(); - ctx->SetOutput(0, ctx->builder()->ConstantR0(initialized)); + XlaResource* variable; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &variable)); + ctx->SetOutput(0, + ctx->builder()->ConstantR0(variable->initialized())); } }; REGISTER_XLA_OP(Name("VarIsInitializedOp"), VarIsInitializedOp); class ReadVariableOp : public XlaOpKernel { public: - explicit ReadVariableOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + explicit ReadVariableOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + } + void Compile(XlaOpKernelContext* ctx) override { xla::ComputationDataHandle handle; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &handle)); + OP_REQUIRES_OK( + ctx, ctx->ReadVariableInput(0, dtype_, /*shape=*/nullptr, &handle)); ctx->SetOutput(0, handle); } + + private: + DataType dtype_; }; REGISTER_XLA_OP(Name("ReadVariableOp"), ReadVariableOp); @@ -65,10 +73,12 @@ class AssignAddVariableOp : public XlaOpKernel { public: explicit AssignAddVariableOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { + DataType type = ctx->input_type(1); xla::ComputationDataHandle handle; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &handle)); + OP_REQUIRES_OK(ctx, + ctx->ReadVariableInput(0, type, /*shape=*/nullptr, &handle)); handle = ctx->builder()->Add(handle, ctx->Input(1)); - OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, ctx->input_type(1), handle)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, handle)); } }; REGISTER_XLA_OP( @@ -79,10 +89,12 @@ class AssignSubVariableOp : public XlaOpKernel { public: explicit AssignSubVariableOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { + DataType type = ctx->input_type(1); xla::ComputationDataHandle handle; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &handle)); + OP_REQUIRES_OK(ctx, + ctx->ReadVariableInput(0, type, /*shape=*/nullptr, &handle)); handle = ctx->builder()->Sub(handle, ctx->Input(1)); - OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, ctx->input_type(1), handle)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, handle)); } }; REGISTER_XLA_OP( @@ -95,28 +107,21 @@ class ResourceGatherOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { xla::ComputationBuilder* builder = ctx->builder(); - // Get the shape of the resource tensor. - TensorShape resource_shape; - DataType resource_dtype; - OP_REQUIRES_OK( - ctx, ctx->GetVariableTypeAndShape(0, &resource_dtype, &resource_shape)); - - DataType expected_output_dtype = ctx->expected_output_dtype(0); - OP_REQUIRES(ctx, resource_dtype == expected_output_dtype, - errors::InvalidArgument( - "Variable dtype is ", DataTypeString(resource_dtype), - " but expected output dtype is ", - DataTypeString(expected_output_dtype), ".")); + DataType type = ctx->expected_output_dtype(0); + TensorShape resource_shape; xla::ComputationDataHandle resource_handle; - OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &resource_handle)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &resource_shape, + &resource_handle)); auto indices = ctx->Input(1); auto indices_shape = ctx->InputShape(1); DataType index_type = ctx->input_type(1); - xla::ComputationDataHandle gather = XlaComputeGatherDynamicSlice( - ctx, resource_handle, resource_shape, indices, indices_shape, 0, - resource_dtype, index_type, builder); + xla::ComputationDataHandle gather; + OP_REQUIRES_OK( + ctx, XlaGather(resource_handle, resource_shape, indices, indices_shape, + /*axis=*/0, /*indices_are_nd=*/false, type, index_type, + builder, &gather)); ctx->SetOutput(0, gather); } }; diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.cc b/tensorflow/compiler/tf2xla/kernels/while_op.cc index 07a565e34b8477858e813686fedb65b52a65a673..0ff1b65ae9179d506e453f98097cd88083eb2be7 100644 --- a/tensorflow/compiler/tf2xla/kernels/while_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc @@ -39,7 +39,7 @@ Status MakeXlaCompilerArgumentsFromInputs( *has_uninitialized_vars = false; *has_tensor_arrays = false; for (int i = 0; i < ctx->num_inputs(); ++i) { - VLOG(2) << " Input " << i + VLOG(2) << " Input " << i << " type: " << DataTypeString(ctx->input_type(i)) << " shape: " << ctx->InputShape(i).DebugString(); XlaCompiler::Argument& arg = (*args)[i]; @@ -58,9 +58,8 @@ Status MakeXlaCompilerArgumentsFromInputs( } arg.type = resource->type(); - if (arg.initialized) { - TF_RETURN_IF_ERROR(resource->PackedShape(ctx->builder(), &arg.shape)); - } else { + arg.shape = resource->shape(); + if (!arg.initialized) { *has_uninitialized_vars = true; } arg.tensor_array_size = resource->tensor_array_size(); @@ -70,14 +69,13 @@ Status MakeXlaCompilerArgumentsFromInputs( arg.name = resource->name(); VLOG(2) << " resource " << resource->name() << " type: " << DataTypeString(arg.type) - << " shape: " << xla::ShapeUtil::HumanString(arg.shape) + << " shape: " << arg.shape.DebugString() << " initialized: " << arg.initialized; } else { arg.kind = XlaCompiler::Argument::kParameter; arg.type = ctx->input_type(i); - TF_RETURN_IF_ERROR( - TensorShapeToXLAShape(arg.type, ctx->InputShape(i), &arg.shape)); + arg.shape = ctx->InputShape(i); } } return Status::OK(); @@ -120,6 +118,7 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { body_options.use_tuple_arg = true; body_options.return_updated_values_for_all_resources = true; body_options.resolve_compile_time_constants = false; + body_options.is_entry_computation = false; XlaCompiler::CompilationResult body; OP_REQUIRES_OK(ctx, compiler->CompileFunction(body_options, body_name_attr_, arguments, &body)); @@ -153,17 +152,14 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { XlaCompiler::Argument& arg = arguments[update.input_index]; if (!arg.initialized) { VLOG(2) << "Update shape for argument " << update.input_index << " " - << xla::ShapeUtil::HumanString(update.shape); + << update.shape.DebugString(); arg.initialized = true; - xla::Shape shape = update.shape; - if (!update.tensor_array_gradients_accessed.empty()) { - shape = xla::ShapeUtil::GetTupleElementShape(shape, 0); - } - std::unique_ptr zero = - xla::Literal::CreateFromShape(shape); - OP_REQUIRES_OK(ctx, resource->SetValue( - update.type, builder->ConstantLiteral(*zero))); + arg.shape = update.shape; + OP_REQUIRES_OK(ctx, + resource->SetTypeAndShape(update.type, update.shape)); + + OP_REQUIRES_OK(ctx, resource->SetZeroValue(builder)); } // Add any TensorArray gradients touched by the body to the enclosing @@ -181,9 +177,6 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { for (const auto& gradient : resource->tensor_array_gradients()) { arg.tensor_array_gradients.insert(gradient.first); } - - // Recompute the argument shape. - OP_REQUIRES_OK(ctx, resource->PackedShape(ctx->builder(), &arg.shape)); } // Recompile the body with the "correct" resource shapes. VLOG(1) << "Recompiling body with corrected resource shapes"; @@ -197,14 +190,21 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { XlaCompiler::CompileOptions cond_options; cond_options.use_tuple_arg = true; cond_options.resolve_compile_time_constants = false; + cond_options.is_entry_computation = false; XlaCompiler::CompilationResult cond; OP_REQUIRES_OK(ctx, compiler->CompileFunction(cond_options, cond_name_attr_, arguments, &cond)); - xla::Shape body_input_shape = - xla::ShapeUtil::MakeTupleShape(body.xla_input_shapes); - xla::Shape cond_input_shape = - xla::ShapeUtil::MakeTupleShape(cond.xla_input_shapes); + OP_REQUIRES(ctx, body.xla_input_shapes.size() == 1, + errors::FailedPrecondition("Expected one input shape")); + xla::Shape body_input_shape = body.xla_input_shapes[0]; + OP_REQUIRES(ctx, xla::ShapeUtil::IsTuple(body_input_shape), + errors::FailedPrecondition("Expected tuple shape")); + OP_REQUIRES(ctx, cond.xla_input_shapes.size() == 1, + errors::FailedPrecondition("Expected one input shape")); + xla::Shape cond_input_shape = cond.xla_input_shapes[0]; + OP_REQUIRES(ctx, xla::ShapeUtil::IsTuple(cond_input_shape), + errors::FailedPrecondition("Expected tuple shape")); VLOG(2) << "Body shape: " << xla::ShapeUtil::HumanString(body_input_shape) << " -> " << xla::ShapeUtil::HumanString(body.xla_output_shape); @@ -284,13 +284,12 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { OP_REQUIRES_OK(ctx, resource->SetFromPack( arguments[update.input_index].tensor_array_gradients, - builder->GetTupleElement(while_result, pos), - /*reset_initial_values=*/false, builder)); + builder->GetTupleElement(while_result, pos), builder)); } VLOG(2) << "Loop-carried variable: pos: " << update.input_index << " name: " << resource->name() << " modified: " << update.modified << " type: " << DataTypeString(update.type) - << " shape: " << xla::ShapeUtil::HumanString(update.shape); + << " shape: " << update.shape.DebugString(); // Copies the identity of the resource variable from input to output // unchanged, even if the variable was not modified. ctx->op_kernel_context()->set_output( diff --git a/tensorflow/compiler/tf2xla/lib/BUILD b/tensorflow/compiler/tf2xla/lib/BUILD index 21ad21f73737a289390ed1ea767db1078d05b466..488fda74bf7b5c1d66f8d706a1be3cc1fc29a492 100644 --- a/tensorflow/compiler/tf2xla/lib/BUILD +++ b/tensorflow/compiler/tf2xla/lib/BUILD @@ -49,6 +49,25 @@ cc_library( ], ) +cc_library( + name = "scatter", + srcs = ["scatter.cc"], + hdrs = ["scatter.h"], + deps = [ + ":util", + ":while_loop", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/client:computation", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/core:lib", + ], +) + cc_library( name = "triangular_solve", srcs = ["triangular_solve.cc"], @@ -60,6 +79,8 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/core:lib", @@ -105,6 +126,21 @@ cc_library( ], ) +cc_library( + name = "while_loop", + srcs = ["while_loop.cc"], + hdrs = ["while_loop.h"], + deps = [ + ":util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla/client:computation", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/core:lib", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/tf2xla/lib/batch_dot.cc b/tensorflow/compiler/tf2xla/lib/batch_dot.cc index 9b0e6174475c22e325c090bec5f1d56822e106bc..798f0fa78055e800038e8bf41b4f410b670be7dd 100644 --- a/tensorflow/compiler/tf2xla/lib/batch_dot.cc +++ b/tensorflow/compiler/tf2xla/lib/batch_dot.cc @@ -25,11 +25,10 @@ limitations under the License. namespace tensorflow { -// The current implementation simply unrolls the computation along the batch -// dimension. xla::StatusOr BatchDot( xla::ComputationBuilder* builder, xla::ComputationDataHandle x, - xla::ComputationDataHandle y, bool transpose_x, bool transpose_y) { + xla::ComputationDataHandle y, bool transpose_x, bool transpose_y, + bool conjugate_x, bool conjugate_y) { TF_ASSIGN_OR_RETURN(std::unique_ptr x_shape, builder->GetShape(x)); TF_ASSIGN_OR_RETURN(std::unique_ptr y_shape, @@ -89,10 +88,10 @@ xla::StatusOr BatchDot( dimensions); } - if (x_shape->element_type() == xla::C64 && transpose_x) { + if (x_shape->element_type() == xla::C64 && conjugate_x) { x = builder->Conj(x); } - if (y_shape->element_type() == xla::C64 && transpose_y) { + if (y_shape->element_type() == xla::C64 && conjugate_y) { y = builder->Conj(y); } diff --git a/tensorflow/compiler/tf2xla/lib/batch_dot.h b/tensorflow/compiler/tf2xla/lib/batch_dot.h index b46bc7417d29dc5b7e9649ac28cc78b57d4b619c..b230e885f10f45a78cdd6e455da3ba55ce589b96 100644 --- a/tensorflow/compiler/tf2xla/lib/batch_dot.h +++ b/tensorflow/compiler/tf2xla/lib/batch_dot.h @@ -27,7 +27,10 @@ namespace tensorflow { // viewed as an element of a batch), and arranges the individual results // in a single output tensor of the same batch size. Each of the // individual slices can optionally be transposed before multiplication by -// setting the `transpose_x` or `transpose_y` flag to `true`. +// setting the `transpose_x` or `transpose_y` flag to `true`. Similarly, each +// can be elementwise-complex-conjugated by setting the `conjugate_x` or +// `conjugate_y` flag to `true`. To apply a Hermitian adjoint to `x`, set both +// `transpose_x` and `conjugate_x` to `true`, and analogously for `y`. // // The input tensors `x` and `y` are 2-D or higher with shape `[..., r_x, c_x]` // and `[..., r_y, c_y]`. @@ -40,11 +43,10 @@ namespace tensorflow { // It is computed as: // // output[..., :, :] = matrix(x[..., :, :]) * matrix(y[..., :, :]) -// TODO(phawkins): add an option to take the complex conjugate of the LHS or -// RHS. xla::StatusOr BatchDot( xla::ComputationBuilder* builder, xla::ComputationDataHandle x, - xla::ComputationDataHandle y, bool transpose_x, bool transpose_y); + xla::ComputationDataHandle y, bool transpose_x, bool transpose_y, + bool conjugate_x = false, bool conjugate_y = false); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.cc b/tensorflow/compiler/tf2xla/lib/cholesky.cc index b3cc489adf6042acb3f56b3a0a6c8fbe43bde629..e795701181dd80a2ff544743d513bffd52fd2399 100644 --- a/tensorflow/compiler/tf2xla/lib/cholesky.cc +++ b/tensorflow/compiler/tf2xla/lib/cholesky.cc @@ -71,11 +71,14 @@ xla::StatusOr CholeskyUnblocked( SliceInMinorDims(builder, l, {j + 1, 0}, {n, j})); TF_ASSIGN_OR_RETURN(auto r_squared, BatchDot(builder, r, r, /*transpose_x=*/false, - /*transpose_y=*/true)); + /*transpose_y=*/true, /*conjugate_x=*/false, + /*conjugate_y=*/false)); new_d_squared = builder->Sub(new_d_squared, r_squared); TF_ASSIGN_OR_RETURN(br, BatchDot(builder, b, r, /*transpose_x=*/false, - /*transpose_y=*/true)); + /*transpose_y=*/true, + /*conjugate_x=*/false, + /*conjugate_y=*/false)); } auto new_d_inv = builder->Pow( new_d_squared, FloatLiteral(builder, shape->element_type(), -0.5)); @@ -134,7 +137,8 @@ xla::StatusOr Cholesky( SliceInMinorDims(builder, l, {i, 0}, {i + k, i})); TF_ASSIGN_OR_RETURN(auto delta, BatchDot(builder, lhs, rhs, /*transpose_x=*/false, - /*transpose_y=*/true)); + /*transpose_y=*/true, /*conjugate_x=*/false, + /*conjugate_y=*/false)); TF_ASSIGN_OR_RETURN(auto before, SliceInMinorDims(builder, a, {i, i}, {n, i + k})); TF_ASSIGN_OR_RETURN( @@ -155,6 +159,10 @@ xla::StatusOr Cholesky( SliceInMinorDims(builder, a, {i + k, i}, {n, i + k})); TF_ASSIGN_OR_RETURN(auto update, TriangularSolve(builder, factorized, panel, + /*left_side=*/false, + /*lower=*/true, + /*transpose_a=*/true, + /*conjugate_a=*/false, /*block_size=*/8)); TF_ASSIGN_OR_RETURN( l, UpdateSliceInMinorDims(builder, l, update, {i + k, i})); diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.h b/tensorflow/compiler/tf2xla/lib/cholesky.h index 2bead7359baaf3582c1230adf0cd4a90046859d2..e083a383be4be0d1b556b63214fe5f70323b4149 100644 --- a/tensorflow/compiler/tf2xla/lib/cholesky.h +++ b/tensorflow/compiler/tf2xla/lib/cholesky.h @@ -29,6 +29,7 @@ namespace tensorflow { // the block size to use. // TODO(phawkins): check for negative values on the diagonal and return an // error, instead of silently yielding NaNs. +// TODO(mattjj): handle the complex Hermitian case xla::StatusOr Cholesky( xla::ComputationBuilder* builder, xla::ComputationDataHandle a, int64 block_size = 256); diff --git a/tensorflow/compiler/tf2xla/lib/scatter.cc b/tensorflow/compiler/tf2xla/lib/scatter.cc new file mode 100644 index 0000000000000000000000000000000000000000..6009243f9774eea24e8049e2bd50fe32f291132f --- /dev/null +++ b/tensorflow/compiler/tf2xla/lib/scatter.cc @@ -0,0 +1,189 @@ +/* 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/compiler/tf2xla/lib/scatter.h" + +#include +#include + +#include "tensorflow/compiler/tf2xla/lib/util.h" +#include "tensorflow/compiler/tf2xla/lib/while_loop.h" +#include "tensorflow/compiler/xla/client/lib/arithmetic.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/array_slice.h" + +namespace tensorflow { + +xla::StatusOr XlaScatter( + const xla::ComputationDataHandle& buffer, + const xla::ComputationDataHandle& updates, + const xla::ComputationDataHandle& indices, bool indices_are_vectors, + const std::function& combiner, + xla::ComputationBuilder* builder) { + TF_ASSIGN_OR_RETURN(std::unique_ptr buffer_shape, + builder->GetShape(buffer)); + TF_ASSIGN_OR_RETURN(std::unique_ptr updates_shape, + builder->GetShape(updates)); + TF_ASSIGN_OR_RETURN(std::unique_ptr indices_shape, + builder->GetShape(indices)); + gtl::ArraySlice indices_dims = + xla::AsInt64Slice(indices_shape->dimensions()); + gtl::ArraySlice buffer_dims = + xla::AsInt64Slice(buffer_shape->dimensions()); + + // If the indices are N-dimensional, the minor dimension of indices contains + // the indices to update. Otherwise the indices are all scalars. + int64 num_index_dims = 1; + if (indices_are_vectors) { + TF_RET_CHECK(!indices_dims.empty()); + num_index_dims = indices_dims.back(); + if (num_index_dims > xla::ShapeUtil::Rank(*buffer_shape)) { + return errors::InvalidArgument( + "The size of the minor dimension of the indices (shape: ", + xla::ShapeUtil::HumanString(*indices_shape), + ") must be <= the rank of the buffer (shape: ", + xla::ShapeUtil::HumanString(*buffer_shape), ")"); + } + indices_dims.pop_back(); + } + + int64 num_indices = 1; + for (int64 dim : indices_dims) { + num_indices *= dim; + } + + // Degenerate case: nothing to update. Return the buffer unchanged. + if (num_indices == 0) { + return buffer; + } + + // If any of the indexed dimensions are zero in the buffer, the update cannot + // succeed since it updates a slice of size 1. + for (int64 i = 0; i < num_index_dims; ++i) { + if (xla::ShapeUtil::GetDimension(*buffer_shape, i) == 0) { + return errors::InvalidArgument( + "Scatter dimension ", i, " is of size zero in tensor with shape ", + xla::ShapeUtil::HumanString(*buffer_shape)); + } + } + + // Shape of the non-indexed dimensions of the buffer. + std::vector buffer_shape_post_axes( + buffer_dims.begin() + num_index_dims, buffer_dims.end()); + + // Flatten the major dimensions of indices and updates into a single dimension + // for ease of iteration. + std::vector flat_indices_shape({num_indices}); + if (indices_are_vectors) { + flat_indices_shape.push_back(num_index_dims); + } + + std::vector flat_updates_shape({num_indices}); + flat_updates_shape.insert(flat_updates_shape.end(), + buffer_shape_post_axes.begin(), + buffer_shape_post_axes.end()); + + // Construct the initial values of the loop-carried Tensors. + auto flat_indices = builder->Reshape(indices, flat_indices_shape); + auto flat_updates = builder->Reshape(updates, flat_updates_shape); + auto init = {flat_indices, flat_updates, buffer}; + + // Constructs the loop body. The implementation of scatter is essentially: + // for i in range(num_indices): + // index = dynamic-slice(indices, i) + // update = dynamic-slice(updates, i) + // buffer = dynamic-update-slice(buffer, update, index) + auto body_fn = [&](xla::ComputationDataHandle i, + gtl::ArraySlice loop_vars, + xla::ComputationBuilder* body_builder) { + auto indices = loop_vars[0]; + auto updates = loop_vars[1]; + auto buffer = loop_vars[2]; + + auto zero_index = body_builder->ConstantLiteral( + xla::Literal::Zero(indices_shape->element_type())); + + // Slice the i-th index from the indices array. + xla::ComputationDataHandle index; + auto indices_offset = body_builder->Reshape(i, {1}); + if (indices_are_vectors) { + indices_offset = body_builder->Pad(indices_offset, zero_index, + xla::MakeEdgePaddingConfig({{0, 1}})); + + index = body_builder->DynamicSlice(indices, indices_offset, + {1, num_index_dims}); + index = body_builder->Collapse(index, {0, 1}); + } else { + index = body_builder->DynamicSlice(indices, indices_offset, {1}); + } + + // Discard updates with negative indices, since some users expect this. + auto index_in_range = + body_builder->ReduceAll(body_builder->Le(zero_index, index), + body_builder->ConstantR0(true), + xla::CreateScalarAndComputation(body_builder)); + + index = body_builder->Pad( + index, zero_index, + xla::MakeEdgePaddingConfig({{0, buffer_shape_post_axes.size()}})); + + // Slice the i-th index from the updates array. + auto updates_offset = body_builder->Reshape(i, {1}); + updates_offset = body_builder->Pad( + updates_offset, zero_index, + xla::MakeEdgePaddingConfig({{0, buffer_shape_post_axes.size()}})); + std::vector flat_updates_slice_shape({1}); + flat_updates_slice_shape.insert(flat_updates_slice_shape.end(), + buffer_shape_post_axes.begin(), + buffer_shape_post_axes.end()); + auto update = body_builder->DynamicSlice(updates, updates_offset, + flat_updates_slice_shape); + + // Unflatten the major (iteration) dimensions of the slice to their original + // shape. + std::vector updates_slice_shape(num_index_dims, 1); + updates_slice_shape.insert(updates_slice_shape.end(), + buffer_shape_post_axes.begin(), + buffer_shape_post_axes.end()); + update = body_builder->Reshape(update, updates_slice_shape); + + // Apply the update to the buffer. If there is a combiner, use it to merge + // the current values with the update. + if (combiner) { + auto current_value = + body_builder->DynamicSlice(buffer, index, updates_slice_shape); + update = combiner(current_value, update, body_builder); + } + // Apply the update if it is in range. + buffer = body_builder->Select( + index_in_range, body_builder->DynamicUpdateSlice(buffer, update, index), + buffer); + + return std::vector{indices, updates, buffer}; + }; + + TF_ASSIGN_OR_RETURN( + auto outputs, XlaForEachIndex(num_indices, indices_shape->element_type(), + body_fn, init, "scatter", builder)); + return outputs[2]; +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/scatter.h b/tensorflow/compiler/tf2xla/lib/scatter.h new file mode 100644 index 0000000000000000000000000000000000000000..41e6d3b195ebf90662c7b9b42c53fcb0133ab29e --- /dev/null +++ b/tensorflow/compiler/tf2xla/lib/scatter.h @@ -0,0 +1,53 @@ +/* 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_COMPILER_TF2XLA_LIB_SCATTER_H_ +#define TENSORFLOW_COMPILER_TF2XLA_LIB_SCATTER_H_ + +#include + +#include "tensorflow/compiler/xla/client/computation.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace tensorflow { + +// Builds an XLA computation that performs a scatter operation on `buffer`, +// returning an updated buffer. +// For each i0, i1, ..., sets +// buffer[indices[i0, i1, ...], ...] := updates[i0, i1, ...] +// +// If `indices_are_vectors` is false, then each index in indices is a scalar, +// and the shape of `indices` must be a prefix of the shape of updates. +// Otherwise, `indices_are_vectors`, then indices are multidimensional and the +// minor dimension of `indices` represents a vector of indices. +// +// If any indices are negative, the corresponding update is discarded. +// +// If a `combiner` is provided, updates are combined with the existing values in +// the buffer using the combiner function. Otherwise, the updates replace the +// existing values. The order of updates is implementation-defined. +xla::StatusOr XlaScatter( + const xla::ComputationDataHandle& buffer, + const xla::ComputationDataHandle& updates, + const xla::ComputationDataHandle& indices, bool indices_are_vectors, + const std::function& combiner, + xla::ComputationBuilder* builder); + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_SCATTER_H_ diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc b/tensorflow/compiler/tf2xla/lib/triangular_solve.cc index 579944c3a381e7018b7fee5013d0509158ce21cc..7f72a6073df218b9e2bd4cc0c0b5bb10b5cd4b84 100644 --- a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc +++ b/tensorflow/compiler/tf2xla/lib/triangular_solve.cc @@ -24,13 +24,15 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/errors.h" namespace tensorflow { xla::StatusOr TriangularSolve( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a, - xla::ComputationDataHandle b, int64 block_size) { + xla::ComputationDataHandle b, bool left_side, bool lower, bool transpose_a, + bool conjugate_a, int64 block_size) { TF_ASSIGN_OR_RETURN(std::unique_ptr a_shape, builder->GetShape(a)); TF_ASSIGN_OR_RETURN(std::unique_ptr b_shape, @@ -60,14 +62,15 @@ xla::StatusOr TriangularSolve( batch_dimensions.push_back(a_size); } - const int64 n = xla::ShapeUtil::GetDimension(*a_shape, -1); - const int64 m = xla::ShapeUtil::GetDimension(*b_shape, -2); - if (n != xla::ShapeUtil::GetDimension(*a_shape, -2)) { + if (xla::ShapeUtil::GetDimension(*a_shape, -1) != + xla::ShapeUtil::GetDimension(*a_shape, -2)) { return errors::InvalidArgument( "The 'a' arguments to TriangularSolve must be square matrices: ", xla::ShapeUtil::HumanString(*a_shape)); } - if (n != xla::ShapeUtil::GetDimension(*b_shape, -1)) { + const int64 m = xla::ShapeUtil::GetDimension(*b_shape, -2); + const int64 n = xla::ShapeUtil::GetDimension(*b_shape, -1); + if ((left_side ? m : n) != xla::ShapeUtil::GetDimension(*a_shape, -1)) { return errors::InvalidArgument( "Arguments to TriangularSolve have incompatible matrix shapes: ", xla::ShapeUtil::HumanString(*a_shape), " vs ", @@ -89,6 +92,14 @@ xla::StatusOr TriangularSolve( return output; }; + // Applies a complex conjugation operation if `a` is complex and `conjugate_a` + // is true, otherwise returns its argument. + auto maybe_conj = [&](xla::ComputationBuilder* builder, + xla::ComputationDataHandle x) { + auto perform_conj = a_shape->element_type() == xla::C64 && conjugate_a; + return perform_conj ? builder->Conj(x) : x; + }; + std::map base_computations; auto get_base_triangular_solve = [&](int k) -> xla::StatusOr { @@ -103,19 +114,35 @@ xla::StatusOr TriangularSolve( prepend_batch_dims({k, k})), "a"); + std::array b_lastd; + if (left_side) { + b_lastd = {k, n}; + } else { + b_lastd = {m, k}; + } auto b_param = sub->Parameter(1, xla::ShapeUtil::MakeShape(b_shape->element_type(), - prepend_batch_dims({m, k})), + prepend_batch_dims(b_lastd)), "b"); - // TODO(phawkins): it might make sense to use a while loop here, rather - // than unrolling. - // TODO(phawkins): the left-looking variant of the algorithm might be more - // efficient at block size 1. - TF_RETURN_IF_ERROR(TriangularSolve(sub.get(), a_param, b_param, - /*block_size=*/1) - .status()); + // We use a left-looking subroutine on the block diagonal in some common + // cases, while falling back to a recursive call in unsupported cases. The + // left-looking subroutine is written with a While loop and so yields much + // faster compile times. Moreover, the left-looking variant can give + // higher performance on smaller (sub)problems. + if (left_side && lower) { + TF_RETURN_IF_ERROR(TriangularSolveLeftLooking(sub.get(), a_param, + b_param, transpose_a, + conjugate_a) + .status()); + } else { + TF_RETURN_IF_ERROR(TriangularSolve(sub.get(), a_param, b_param, + left_side, lower, transpose_a, + conjugate_a, + /*block_size=*/1) + .status()); + } TF_ASSIGN_OR_RETURN(computation, sub->Build()); } @@ -129,47 +156,396 @@ xla::StatusOr TriangularSolve( // Goto, Kazushige, and Robert Van De Geijn. "High-performance implementation // of the level-3 BLAS." ACM Transactions on Mathematical Software (TOMS) 35.1 // (2008): 4. - for (int64 i = 0; i < n; i += block_size) { - int64 k = std::min(block_size, n - i); - // if k > 1: - // output[..., :, i:i+k] = triangular_solve( - // a[..., i:i+k, ..., i:i+k], b[..., :, i:i+k], side='Right', - // kind='Lower', transpose=True, block_size=1) - // else: - // output[..., :, i] = b[..., :, i] / a[..., i, i] + // In the code comments below, T = lambda x: np.swapaxes(x, -1, -2) if + // conjugate_a is False, or T = lambda x: np.conj(np.swapaxes(x, -1, -2)) if + // conjugate_a is True. + + if (!left_side && lower == transpose_a) { + // for i in range(0, a.shape[-1], block_size): + for (int64 i = 0; i < n; i += block_size) { + int64 k = std::min(block_size, n - i); + + // output[..., :, i:i+k] = triangular_solve( + // a[..., i:i+k, i:i+k], b[..., :, i:i+k], ..., block_size=1) + TF_ASSIGN_OR_RETURN(auto a_slice, + SliceInMinorDims(builder, a, {i, i}, {i + k, i + k})); + TF_ASSIGN_OR_RETURN(auto b_slice, + SliceInMinorDims(builder, b, {0, i}, {m, i + k})); + xla::ComputationDataHandle update; + if (k > 1) { + TF_ASSIGN_OR_RETURN(xla::Computation * solve, + get_base_triangular_solve(k)); + update = builder->Call(*solve, {a_slice, b_slice}); + } else { + update = builder->Div(b_slice, maybe_conj(builder, a_slice)); + } + TF_ASSIGN_OR_RETURN( + output, UpdateSliceInMinorDims(builder, output, update, {0, i})); + + // if i + k < a.shape[-1]: + // a_slice_2 = a[..., i+k:, i:i+k] if lower else a[..., i:i+k, i+k:] + // a_slice_2 = T(a_slice_2) if transpose_a else a_slice_2 + // b[..., :, i+k:] -= np.matmul(output[..., :, i:i+k], a_slice_2) + if (i + k < n) { + xla::ComputationDataHandle a_slice_2; + if (lower) { + TF_ASSIGN_OR_RETURN( + a_slice_2, SliceInMinorDims(builder, a, {i + k, i}, {n, i + k})); + } else { + TF_ASSIGN_OR_RETURN( + a_slice_2, SliceInMinorDims(builder, a, {i, i + k}, {i + k, n})); + } + + TF_ASSIGN_OR_RETURN(auto b_update, + BatchDot(builder, update, a_slice_2, + /*transpose_x=*/false, + /*transpose_y=*/transpose_a, + /*conjugate_x=*/false, + /*conjugate_y=*/conjugate_a)); + TF_ASSIGN_OR_RETURN(auto b_slice_2, + SliceInMinorDims(builder, b, {0, i + k}, {m, n})); + b_update = builder->Sub(b_slice_2, b_update); + TF_ASSIGN_OR_RETURN( + b, UpdateSliceInMinorDims(builder, b, b_update, {0, i + k})); + } + } + + } else if (left_side && lower != transpose_a) { + // for i in range(0, a.shape[-1], block_size): + for (int64 i = 0; i < m; i += block_size) { + int64 k = std::min(block_size, m - i); + + // output[..., i:i+k, :] = triangular_solve( + // a[..., i:i+k, i:i+k], b[..., i:i+k, :], ..., block_size=1) + TF_ASSIGN_OR_RETURN(auto a_slice, + SliceInMinorDims(builder, a, {i, i}, {i + k, i + k})); + TF_ASSIGN_OR_RETURN(auto b_slice, + SliceInMinorDims(builder, b, {i, 0}, {i + k, n})); + xla::ComputationDataHandle update; + if (k > 1) { + TF_ASSIGN_OR_RETURN(xla::Computation * solve, + get_base_triangular_solve(k)); + update = builder->Call(*solve, {a_slice, b_slice}); + } else { + update = builder->Div(b_slice, maybe_conj(builder, a_slice)); + } + TF_ASSIGN_OR_RETURN( + output, UpdateSliceInMinorDims(builder, output, update, {i, 0})); + + // if i + k < a.shape[-1]: + // a_slice_2 = a[..., i+k:, i:i+k] if lower else a[..., i:i+k, i+k:] + // a_slice_2 = T(a_slice_2) if transpose_a else a_slice_2 + // b[..., i+k:, :] -= np.matmul(a_slice_2, output[..., i:i+k, :]) + if (i + k < m) { + xla::ComputationDataHandle a_slice_2; + if (lower) { + TF_ASSIGN_OR_RETURN( + a_slice_2, SliceInMinorDims(builder, a, {i + k, i}, {m, i + k})); + } else { + TF_ASSIGN_OR_RETURN( + a_slice_2, SliceInMinorDims(builder, a, {i, i + k}, {i + k, m})); + } + + TF_ASSIGN_OR_RETURN(auto b_update, BatchDot(builder, a_slice_2, update, + /*transpose_x=*/transpose_a, + /*transpose_y=*/false, + /*conjugate_x=*/conjugate_a, + /*conjugate_y=*/false)); + TF_ASSIGN_OR_RETURN(auto b_slice_2, + SliceInMinorDims(builder, b, {i + k, 0}, {m, n})); + b_update = builder->Sub(b_slice_2, b_update); + TF_ASSIGN_OR_RETURN( + b, UpdateSliceInMinorDims(builder, b, b_update, {i + k, 0})); + } + } + } else if (!left_side && lower != transpose_a) { + // for i in reversed(range(0, a.shape[-1], block_size)): + const int64 last_blk_ix = xla::RoundUpToNearest(n, block_size) - block_size; + for (int64 i = last_blk_ix; i >= 0; i -= block_size) { + int64 k = std::min(block_size, n - i); + + // output[..., :, i:i+k] triangular_solve( + // a[..., i:i+k, i:i+k], b[..., :, i:i+k], ..., block_size=1) + TF_ASSIGN_OR_RETURN(auto a_slice, + SliceInMinorDims(builder, a, {i, i}, {i + k, i + k})); + TF_ASSIGN_OR_RETURN(auto b_slice, + SliceInMinorDims(builder, b, {0, i}, {m, i + k})); + xla::ComputationDataHandle update; + if (k > 1) { + TF_ASSIGN_OR_RETURN(xla::Computation * solve, + get_base_triangular_solve(k)); + update = builder->Call(*solve, {a_slice, b_slice}); + } else { + update = builder->Div(b_slice, maybe_conj(builder, a_slice)); + } + TF_ASSIGN_OR_RETURN( + output, UpdateSliceInMinorDims(builder, output, update, {0, i})); + + // if i - k >= 0: + // a_slice_2 = a[..., i:i+k, :i] if lower else a[..., :i, i:i+k] + // a_slice_2 = T(a_slice_2) if transpose_a else a_slice_2 + // b[..., :, :i] -= np.matmul(out[..., :, i:i+k], a_slice_2) + if (i - k >= 0) { + xla::ComputationDataHandle a_slice_2; + if (lower) { + TF_ASSIGN_OR_RETURN(a_slice_2, + SliceInMinorDims(builder, a, {i, 0}, {i + k, i})); + } else { + TF_ASSIGN_OR_RETURN(a_slice_2, + SliceInMinorDims(builder, a, {0, i}, {i, i + k})); + } + + TF_ASSIGN_OR_RETURN(auto b_update, + BatchDot(builder, update, a_slice_2, + /*transpose_x=*/false, + /*transpose_y=*/transpose_a, + /*conjugate_x=*/false, + /*conjugate_y=*/conjugate_a)); + TF_ASSIGN_OR_RETURN(auto b_slice_2, + SliceInMinorDims(builder, b, {0, 0}, {m, i})); + b_update = builder->Sub(b_slice_2, b_update); + TF_ASSIGN_OR_RETURN( + b, UpdateSliceInMinorDims(builder, b, b_update, {0, 0})); + } + } + } else { // left_side && lower == transpose_a + // for i in reversed(range(0, a.shape[-1], block_size)): + const int64 last_blk_ix = xla::RoundUpToNearest(m, block_size) - block_size; + for (int64 i = last_blk_ix; i >= 0; i -= block_size) { + int64 k = std::min(block_size, m - i); + + // output[..., i:i+k, :] triangular_solve( + // a[..., i:i+k, i:i+k], b[..., i:i+k, :], ..., block_size=1) + TF_ASSIGN_OR_RETURN(auto a_slice, + SliceInMinorDims(builder, a, {i, i}, {i + k, i + k})); + TF_ASSIGN_OR_RETURN(auto b_slice, + SliceInMinorDims(builder, b, {i, 0}, {i + k, n})); + xla::ComputationDataHandle update; + if (k > 1) { + TF_ASSIGN_OR_RETURN(xla::Computation * solve, + get_base_triangular_solve(k)); + update = builder->Call(*solve, {a_slice, b_slice}); + } else { + update = builder->Div(b_slice, maybe_conj(builder, a_slice)); + } + TF_ASSIGN_OR_RETURN( + output, UpdateSliceInMinorDims(builder, output, update, {i, 0})); + + // if i - k >= 0: + // a_slice_2 = a[..., i:i+k, :i] if lower else a[..., :i, i:i+k] + // a_slice_2 = T(a_slice_2) if transpose_a else a_slice_2 + // b[..., :i, :] -= np.matmul(a_slice_2, out[..., i:i+k, :]) + if (i - k >= 0) { + xla::ComputationDataHandle a_slice_2; + if (lower) { + TF_ASSIGN_OR_RETURN(a_slice_2, + SliceInMinorDims(builder, a, {i, 0}, {i + k, i})); + } else { + TF_ASSIGN_OR_RETURN(a_slice_2, + SliceInMinorDims(builder, a, {0, i}, {i, i + k})); + } + + TF_ASSIGN_OR_RETURN(auto b_update, BatchDot(builder, a_slice_2, update, + /*transpose_x=*/transpose_a, + /*transpose_y=*/false, + /*conjugate_x=*/conjugate_a, + /*conjugate_y=*/false)); + TF_ASSIGN_OR_RETURN(auto b_slice_2, + SliceInMinorDims(builder, b, {0, 0}, {i, n})); + b_update = builder->Sub(b_slice_2, b_update); + TF_ASSIGN_OR_RETURN( + b, UpdateSliceInMinorDims(builder, b, b_update, {0, 0})); + } + } + } + + return output; +} + +xla::StatusOr TriangularSolveLeftLooking( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a, + const xla::ComputationDataHandle& b, bool transpose_a, bool conjugate_a) { + TF_ASSIGN_OR_RETURN(std::unique_ptr a_shape, + builder->GetShape(a)); + TF_ASSIGN_OR_RETURN(std::unique_ptr b_shape, + builder->GetShape(b)); + const int64 m = xla::ShapeUtil::GetDimension(*b_shape, -2); + const int64 n = xla::ShapeUtil::GetDimension(*b_shape, -1); + const int64 ndims = xla::ShapeUtil::Rank(*a_shape); + + std::vector batch_dimensions; + for (int i = 0; i < ndims - 2; ++i) { + int64 a_size = a_shape->dimensions(i); + batch_dimensions.push_back(a_size); + } + + auto prepend_batch_dims = [&](std::array indices) { + std::vector output(ndims); + std::copy(batch_dimensions.begin(), batch_dimensions.end(), output.begin()); + std::copy(indices.begin(), indices.end(), + output.begin() + batch_dimensions.size()); + return output; + }; + + auto maybe_conj = [&](xla::ComputationBuilder* builder, + xla::ComputationDataHandle x) { + auto perform_conj = a_shape->element_type() == xla::C64 && conjugate_a; + return perform_conj ? builder->Conj(x) : x; + }; + + // The main computation is performed in a While loop. + + // Allocate the output and set its first or last row, + // output = np.zeros_like(b) + // if transpose_a: + // output[..., m-1:, :] = b[..., m-1:, :] / a[..., m-1:, m-1:] + // else: + // output[..., :1, :] = b[..., :1, :] / a[..., :1, :1] + xla::ComputationDataHandle output = Zeros(builder, *b_shape); + { + auto i = transpose_a ? m - 1 : 0; TF_ASSIGN_OR_RETURN(auto a_slice, - SliceInMinorDims(builder, a, {i, i}, {i + k, i + k})); + SliceInMinorDims(builder, a, {i, i}, {i + 1, i + 1})); TF_ASSIGN_OR_RETURN(auto b_slice, - SliceInMinorDims(builder, b, {0, i}, {m, i + k})); - xla::ComputationDataHandle update; - if (k > 1) { - TF_ASSIGN_OR_RETURN(xla::Computation * solve, - get_base_triangular_solve(k)); - update = builder->Call(*solve, {a_slice, b_slice}); + SliceInMinorDims(builder, b, {i, 0}, {i + 1, n})); + auto update = builder->Div(b_slice, maybe_conj(builder, a_slice)); + TF_ASSIGN_OR_RETURN( + output, UpdateSliceInMinorDims(builder, output, update, {i, 0})); + } + + // Construct the initial loop carry tuple, + // if transpose_a: + // init = (m-2, output, a, b) + // else: + // init = (1, output, a, b) + std::vector tuple_shapes = { + // The loop iteration counter is a scalar, incremented each iteration. + xla::ShapeUtil::MakeShape(xla::S32, {}), + // The output has the shape of b, with one row updated each iteration. + *b_shape, + // The coefficient matrix a is a loop invariant. + *a_shape, + // The right-hand-side matrix b is a loop invariant. + *b_shape}; + xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes); + auto init_i = builder->ConstantR0(transpose_a ? m - 2 : 1); + auto init = builder->Tuple({init_i, output, a, b}); + + // Construct the loop condition function, + // def cond_fun(loop_carry): + // i, output, a, b = loop_carry + // return i >= 0 if transpose_a else i < m + std::unique_ptr condb = + builder->CreateSubBuilder("TriangularSolveLeftLookingWhileCond"); + { + auto i = condb->GetTupleElement( + condb->Parameter(0, tuple_shape, + "TriangularSolveLeftLookingWhileTuple"), + 0); + if (transpose_a) { + condb->Ge(i, condb->ConstantR0(0)); } else { - update = builder->Div(b_slice, a_slice); + condb->Lt(i, condb->ConstantR0(m)); } + } + TF_ASSIGN_OR_RETURN(auto cond, condb->Build()); - TF_ASSIGN_OR_RETURN( - output, UpdateSliceInMinorDims(builder, output, update, {0, i})); - // b[..., :, i+k:] -= np.dot(output[..., :, i:i+k], - // np.transpose(..., a[i+k:, i:i+k])) - if (i + k < n) { - TF_ASSIGN_OR_RETURN(auto a_slice_2, - SliceInMinorDims(builder, a, {i + k, i}, {n, i + k})); - TF_ASSIGN_OR_RETURN(auto b_update, BatchDot(builder, update, a_slice_2, - /*transpose_x=*/false, - /*transpose_y=*/true)); - - TF_ASSIGN_OR_RETURN(auto b_slice_2, - SliceInMinorDims(builder, b, {0, i + k}, {m, n})); - b_update = builder->Sub(b_slice_2, b_update); - TF_ASSIGN_OR_RETURN( - b, UpdateSliceInMinorDims(builder, b, b_update, {0, i + k})); + // Construct the loop body function, + // def body_fun(loop_carry): + // i, output, a, b = loop_carry + // if transpose_a: + // a_row = np.swapaxes(a[..., i+1:, i:i+1], -1 -2) + // else: + // a_row = a[..., i:i+1, :i] + // result_row = b[..., i:i+1, :] - np.matmul(a_row, output[..., :, :]) + // output[..., i:i+1, :] = result_row / a[..., i:i+1, i:i+1] + // if transpose_a: + // return (i - 1, output, a, b) + // else: + // return (i + 1, output, a, b) + // We have to do some extra FLOPs propagating zeros in the matrix multiply + // because we can't have the size of its arguments depend on the loop counter. + std::unique_ptr bodyb = + builder->CreateSubBuilder("TriangularSolveLeftLookingWhileBody"); + { + auto input_tuple = bodyb->Parameter(0, tuple_shape, + "TriangularSolveLeftLookingWhileTuple"); + + // i, output, a, b = loop_carry + auto i = bodyb->GetTupleElement(input_tuple, 0); + auto body_out = bodyb->GetTupleElement(input_tuple, 1); + auto body_a = bodyb->GetTupleElement(input_tuple, 2); + auto body_b = bodyb->GetTupleElement(input_tuple, 3); + auto zero = bodyb->ConstantR0(0); + + // Set up some helper functions. + auto prepend_zeros = [&](std::array starts) { + auto zero = bodyb->Reshape(bodyb->ConstantR0(0), {1}); + std::vector padded_starts(ndims, zero); + padded_starts[ndims - 2] = bodyb->Reshape(starts[0], {1}); + padded_starts[ndims - 1] = bodyb->Reshape(starts[1], {1}); + return bodyb->ConcatInDim(padded_starts, 0); + }; + + auto dynamic_slice = [&](xla::ComputationDataHandle x, + std::array starts, + std::array sizes) { + auto padded_starts = prepend_zeros(starts); + auto padded_sizes = prepend_batch_dims(sizes); + return bodyb->DynamicSlice(x, padded_starts, padded_sizes); + }; + + auto update = [&](xla::ComputationDataHandle x, + xla::ComputationDataHandle update, + std::array starts) { + auto padded_starts = prepend_zeros(starts); + return bodyb->DynamicUpdateSlice(x, update, padded_starts); + }; + + // We'd like to implement this: + // if transpose_a: + // a_row = T(a[..., i+1:, i:i+1]) + // result_row = (b[..., i:i+1, :] + // - np.matmul(a_row, body_out[..., i+1:, :])) + // else: + // result_row = (b[..., i:i+1, :] + // - np.matmul(a[..., i:i+1, :i], body_out[..., :i, :])) + // But since we can't have intermediate array sizes depend on the loop + // counter, we instead exploit the fact that we initialized the output to + // all zeros and use that as zero-padding (doing unnecessary FLOPs). + xla::ComputationDataHandle a_row; + if (transpose_a) { + a_row = dynamic_slice(body_a, {zero, i}, {m, 1}); + } else { + a_row = dynamic_slice(body_a, {i, zero}, {1, m}); } + TF_ASSIGN_OR_RETURN(auto b_update, BatchDot(bodyb.get(), a_row, body_out, + /*transpose_x=*/transpose_a, + /*transpose_y=*/false, + /*conjugate_x=*/conjugate_a, + /*conjugate_y=*/false)); + auto result_row = + bodyb->Sub(dynamic_slice(body_b, {i, zero}, {1, n}), b_update); + + // body_out[..., i:i+1, :] = result_row / a[..., i:i+1, i:i+1] + auto a_elt = dynamic_slice(body_a, {i, i}, {1, 1}); + auto div_result = bodyb->Div(result_row, maybe_conj(bodyb.get(), a_elt)); + body_out = update(body_out, div_result, {i, zero}); + + // if transpose_a: + // return (i - 1, body_out, a, b) + // else: + // return (i + 1, body_out, a, b) + auto next_i = bodyb->Add(i, bodyb->ConstantR0(transpose_a ? -1 : 1)); + bodyb->Tuple({next_i, body_out, body_a, body_b}); } - return output; + TF_ASSIGN_OR_RETURN(auto body, bodyb->Build()); + + // Construct the While loop and return the result, + // return while_loop(cond_fun, body_fun, init)[1] + auto triangular_solve_left_looking_while = builder->While(cond, body, init); + return builder->GetTupleElement(triangular_solve_left_looking_while, 1); } } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.h b/tensorflow/compiler/tf2xla/lib/triangular_solve.h index 501d026411c80359c7efa406ece5929a2e46ac1f..e32223bfdddda800b1fd4de3e4f0c8061e0f81d8 100644 --- a/tensorflow/compiler/tf2xla/lib/triangular_solve.h +++ b/tensorflow/compiler/tf2xla/lib/triangular_solve.h @@ -21,25 +21,50 @@ limitations under the License. namespace tensorflow { -// Solves systems of linear equations with upper or lower triangular matrices by -// backsubstitution. +// Solves systems of linear equations with lower or upper triangular coefficient +// matrices by forward- or back-substitution. Broadcasting along leading +// dimensions, this routine solves one of the matrix systems +// `op(a) * x = b`, or `x * op(a) = b`, +// for the variable `x` given `a` and `b`, where `op(a)` is either +// `op(a) = a`, or `op(a) = transpose(a)`, or `op(a) = conj(transpose(a))`. +// That is, the innermost matrices in the output satisfy a scalar system +// depending on the value of the value of (left_side, transpose_a, conjugate_a) +// according to: +// (F, F, F) => `output[..., i, k] a[..., k, j] = b[..., i, j]`, +// (F, F, T) => `output[..., i, k] a*[..., k, j] = b[..., i, j]`, +// (F, T, F) => `output[..., i, k] a[..., j, k] = b[..., i, j]`, +// (F, T, T) => `output[..., i, k] a*[..., j, k] = b[..., i, j]`, +// (T, F, F) => ` a[..., i, k] output[..., k, j] = b[..., i, j]`, +// (T, F, T) => `a*[..., i, k] output[..., k, j] = b[..., i, j]`, +// (T, T, F) => ` a[..., i, k] output[..., j, k] = b[..., i, j]`, +// (T, T, T) => `a*[..., i, k] output[..., j, k] = b[..., i, j]`, +// where * denotes complex conjugation and where the index `k` is summed over. // -// `a` is a tensor of shape `[..., M, M]` whose inner-most 2 dimensions form -// square matrices. The strictly upper triangular part of each inner-most matrix -// is assumed to be zero and not accessed. -// `b` is a tensor of shape `[..., M, K]`. -// -// The innermost matrices in the output satisfy matrix equations -// `output[..., i, j] * adjoint(a[..., k, j]) = b[..., i, k]`. +// `a` is a tensor of shape `[..., M, M]` whose innermost 2 dimensions form +// square matrices. If lower is true (false), then the strictly upper (lower) +// triangular part of each innermost matrix in `a` is assumed to be zero and is +// not accessed. +// `b` is a tensor of shape `[..., M, K]` if left_side is true, otherwise a +// tensor of shape `[..., K, M]`. +// `left_side` is a boolean, indicating whether to solve a system of the form +// op(a) * x = b (true) or x * op(a) = b (false). +// `lower` is a boolean, indicating whether the argument `a` is lower-triangular +// (true) or upper-triangular (false). +// `transpose_a` is a boolean indicating whether the matrix `a` is transposed. +// `conjugate_a` is a boolean indicating whether the entries of `a` are complex +// conjugated (independently of whether they are transposed), so that when both +// transpose_a and conjugate_a are true the effect is a Hermitian adjoint. // // Uses a blocked algorithm if `block_size` is > 1; if block_size == 1 then no // blocking is used. -// TODO(phawkins): equivalent to the BLAS TRSM routine with side=right, -// kind=lower, and transposed_a=true. Implement the other possible combinations -// of side, kind and transposed_a. xla::StatusOr TriangularSolve( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a, - xla::ComputationDataHandle b, int64 block_size = 256); + xla::ComputationDataHandle b, bool left_side, bool lower, bool transpose_a, + bool conjugate_a, int64 block_size = 256); + +xla::StatusOr TriangularSolveLeftLooking( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a, + const xla::ComputationDataHandle& b, bool transpose_a, bool conjugate_a); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc b/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc index 671d9aa4fe0c042a3cc44468074653d51c2be75d..661707062916263fd0d5d935ce41698a7655df02 100644 --- a/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc +++ b/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc @@ -27,32 +27,134 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace tensorflow { namespace { using TriangularSolveTest = xla::ClientLibraryTestBase; +using TriangularSolveLeftLookingTest = xla::ClientLibraryTestBase; +using complex64 = xla::complex64; -XLA_TEST_F(TriangularSolveTest, Simple) { +xla::Array2D AValsLower() { + return {{2, 0, 0, 0}, {3, 6, 0, 0}, {4, 7, 9, 0}, {5, 8, 10, 11}}; +} + +xla::Array2D AValsUpper() { + return {{2, 3, 4, 5}, {0, 6, 7, 8}, {0, 0, 9, 10}, {0, 0, 0, 11}}; +} + +xla::Array2D BValsRight() { + return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; +} + +xla::Array2D BValsLeft() { + return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; +} + +xla::Array2D AValsLowerComplex() { + return {{2, 0, 0, 0}, + {complex64(3, 1), 6, 0, 0}, + {4, complex64(7, 2), 9, 0}, + {5, 8, complex64(10, 3), 11}}; +} + +xla::Array2D AValsUpperComplex() { + return {{2, 3, complex64(4, 3), 5}, + {0, 6, complex64(7, 2), 8}, + {0, 0, complex64(9, 1), 10}, + {0, 0, 0, 11}}; +} + +xla::Array2D BValsRightComplex() { + return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; +} + +xla::Array2D BValsLeftComplex() { + return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; +} + +xla::Array2D AValsFull() { + return {{2, 0, 1, 2}, {3, 6, 0, 1}, {4, 7, 9, 0}, {5, 8, 10, 11}}; +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTranspose) { xla::ComputationBuilder builder(client_, TestName()); - xla::Array2D a_vals({ - {2, 0, 0, 0}, - {3, 6, 0, 0}, - {4, 7, 9, 0}, - {5, 8, 10, 11}, + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsRight(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/false, /*lower=*/true, + /*transpose_a=*/true, /*conjugate_a=*/false, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {0.5, 0.08333334, 0.04629629, 0.03367003}, + {2.5, -0.25, -0.1388889, -0.1010101}, + {4.5, -0.58333331, -0.32407406, -0.23569024}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightLowerNotranspose) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsRight(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/false, /*lower=*/true, + /*transpose_a=*/false, /*conjugate_a=*/false, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {-0.16414141, -0.06902357, -0.07070707, 0.36363636}, + {0.64393939, 0.06565657, -0.03030303, 0.72727273}, + {1.4520202, 0.2003367, 0.01010101, 1.09090909}, }); - xla::Array2D b_vals({ - {1, 2, 3, 4}, - {5, 6, 7, 8}, - {9, 10, 11, 12}, + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightUpperTranspose) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsRight(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/false, /*lower=*/false, + /*transpose_a=*/true, /*conjugate_a=*/false, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {-0.16414141, -0.06902357, -0.07070707, 0.36363636}, + {0.64393939, 0.06565657, -0.03030303, 0.72727273}, + {1.4520202, 0.2003367, 0.01010101, 1.09090909}, }); + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightUpperNotranspose) { + xla::ComputationBuilder builder(client_, TestName()); + xla::ComputationDataHandle a, b; - auto a_data = CreateR2Parameter(a_vals, 0, "a", &builder, &a); - auto b_data = CreateR2Parameter(b_vals, 1, "b", &builder, &b); - auto result = TriangularSolve(&builder, a, b, /*block_size=*/2); + auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsRight(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/false, /*lower=*/false, + /*transpose_a=*/false, /*conjugate_a=*/false, + /*block_size=*/2); TF_ASSERT_OK(result.status()); xla::Array2D expected({ @@ -62,7 +164,201 @@ XLA_TEST_F(TriangularSolveTest, Simple) { }); ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(2e-3, 2e-3)); + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerTranspose) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/true, /*lower=*/true, + /*transpose_a=*/true, /*conjugate_a=*/false, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {-0.89646465, -0.69444444, -0.49242424}, + {-0.27441077, -0.24074074, -0.20707071}, + {-0.23232323, -0.22222222, -0.21212121}, + {0.90909091, 1., 1.09090909}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotranspose) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/true, /*lower=*/true, + /*transpose_a=*/false, /*conjugate_a=*/false, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {0.5, 1.0, 1.5}, + {0.41666667, 0.33333333, 0.25}, + {0.23148148, 0.18518519, 0.13888889}, + {0.16835017, 0.13468013, 0.1010101}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTranspose) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/true, /*lower=*/false, + /*transpose_a=*/true, /*conjugate_a=*/false, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {0.5, 1.0, 1.5}, + {0.41666667, 0.33333333, 0.25}, + {0.23148148, 0.18518519, 0.13888889}, + {0.16835017, 0.13468013, 0.1010101}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperNotranspose) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/true, /*lower=*/false, + /*transpose_a=*/false, /*conjugate_a=*/false, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {-0.89646465, -0.69444444, -0.49242424}, + {-0.27441077, -0.24074074, -0.20707071}, + {-0.23232323, -0.22222222, -0.21212121}, + {0.90909091, 1., 1.09090909}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTransposeConjugate) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = + CreateR2Parameter(AValsLowerComplex(), 0, "a", &builder, &a); + auto b_data = + CreateR2Parameter(BValsRightComplex(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/false, /*lower=*/true, + /*transpose_a=*/true, /*conjugate_a=*/true, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {0.5, complex64(0.08333333, 0.08333333), + complex64(0.02777778, -0.0462963), complex64(0.06313131, -0.01094276)}, + {2.5, complex64(-0.25, 0.41666667), complex64(-0.23148148, -0.37962963), + complex64(0.08670034, -0.02104377)}, + {4.5, complex64(-0.58333333, 0.75), complex64(-0.49074074, -0.71296296), + complex64(0.11026936, -0.03114478)}, + }); + + ComputeAndCompareR2(&builder, expected, + {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTransposeNoconjugate) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = + CreateR2Parameter(AValsUpperComplex(), 0, "a", &builder, &a); + auto b_data = + CreateR2Parameter(BValsLeftComplex(), 1, "b", &builder, &b); + auto result = TriangularSolve(&builder, a, b, + /*left_side=*/true, /*lower=*/false, + /*transpose_a=*/true, /*conjugate_a=*/false, + /*block_size=*/2); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {0.5, 1., 1.5}, + {0.41666667, 0.33333333, 0.25}, + {complex64(0.20020325, -2.81504065e-01), + complex64(0.13821138, -4.22764228e-01), + complex64(0.07621951, -5.64024390e-01)}, + {complex64(0.19678492, 2.55912786e-01), + complex64(0.17738359, 3.84331116e-01), + complex64(0.15798226, 5.12749446e-01)}, + }); + + ComputeAndCompareR2(&builder, expected, + {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveLeftLookingTest, Simple) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); + auto result = TriangularSolveLeftLooking(&builder, a, b, + /*transpose_a=*/false, + /*conjugate_a=*/false); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {0.5, 1.0, 1.5}, + {0.41666667, 0.33333333, 0.25}, + {0.23148148, 0.18518519, 0.13888889}, + {0.16835017, 0.13468013, 0.1010101}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveLeftLookingTest, NonzeroUpperTriangle) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b; + auto a_data = CreateR2Parameter(AValsFull(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); + auto result = TriangularSolveLeftLooking(&builder, a, b, + /*transpose_a=*/false, + /*conjugate_a=*/false); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected({ + {0.5, 1.0, 1.5}, + {0.41666667, 0.33333333, 0.25}, + {0.23148148, 0.18518519, 0.13888889}, + {0.16835017, 0.13468013, 0.1010101}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); } } // namespace diff --git a/tensorflow/compiler/tf2xla/lib/util.cc b/tensorflow/compiler/tf2xla/lib/util.cc index ce24b61b5dc7176f3caa05e3eb9257399fef7926..f579669bbd852b514e021ce71d635f8ce5e4fe4d 100644 --- a/tensorflow/compiler/tf2xla/lib/util.cc +++ b/tensorflow/compiler/tf2xla/lib/util.cc @@ -57,6 +57,61 @@ xla::ComputationDataHandle FloatLiteral(xla::ComputationBuilder* builder, } } +xla::ComputationDataHandle IntegerLiteral(xla::ComputationBuilder* builder, + xla::PrimitiveType type, + int64 value) { + xla::Literal literal; + switch (type) { + case xla::U8: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::U32: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::U64: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::S8: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::S32: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::S64: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::F32: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::F64: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::C64: + literal = std::move(*xla::Literal::CreateR0(value)); + break; + case xla::PRED: + LOG(FATAL) << "pred element type is not integral"; + case xla::S16: + case xla::U16: + LOG(FATAL) << "u16/s16 literals not yet implemented"; + case xla::BF16: + literal = std::move( + *xla::Literal::CreateR0(static_cast(value))); + break; + case xla::F16: + literal = std::move( + *xla::Literal::CreateR0(static_cast(value))); + break; + case xla::TUPLE: + LOG(FATAL) << "tuple element type is not integral"; + case xla::OPAQUE: + LOG(FATAL) << "opaque element type is not integral"; + default: + LOG(FATAL) << "unhandled element type " << type; + } + return builder->ConstantLiteral(literal); +} + xla::StatusOr SliceInMinorDims( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, gtl::ArraySlice start, gtl::ArraySlice end) { @@ -107,4 +162,15 @@ xla::StatusOr UpdateSliceInMinorDims( return UpdateSlice(builder, x, update, padded_start); } +xla::StatusOr TransposeInMinorDims( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x) { + TF_ASSIGN_OR_RETURN(std::unique_ptr shape, builder->GetShape(x)); + const int64 n_dims = xla::ShapeUtil::Rank(*shape); + TF_RET_CHECK(n_dims >= 2); + std::vector permutation(n_dims); + std::iota(permutation.begin(), permutation.end(), 0); + std::swap(permutation[n_dims - 1], permutation[n_dims - 2]); + return builder->Transpose(x, permutation); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/util.h b/tensorflow/compiler/tf2xla/lib/util.h index fb138b4f736500aac8184770d97fbf930ced69ea..51f8baaf00bd8fd25baa1a87be8cb0089dfb22b5 100644 --- a/tensorflow/compiler/tf2xla/lib/util.h +++ b/tensorflow/compiler/tf2xla/lib/util.h @@ -32,6 +32,11 @@ xla::ComputationDataHandle Zeros(xla::ComputationBuilder* builder, xla::ComputationDataHandle FloatLiteral(xla::ComputationBuilder* builder, xla::PrimitiveType type, double value); +// Returns a integer scalar constant of 'type' with 'value'. +// If 'type' is complex, returns a real value with zero imaginary component. +xla::ComputationDataHandle IntegerLiteral(xla::ComputationBuilder* builder, + xla::PrimitiveType type, int64 value); + // Performs a slice in the minor dimensions of a Tensor. xla::StatusOr SliceInMinorDims( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, @@ -49,6 +54,10 @@ xla::StatusOr UpdateSliceInMinorDims( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, const xla::ComputationDataHandle& update, gtl::ArraySlice start); +// Transposes a stack of matrices `x` by swapping the last two dimensions. +xla::StatusOr TransposeInMinorDims( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x); + } // namespace tensorflow #endif // TENSORFLOW_COMPILER_TF2XLA_LIB_UTIL_H_ diff --git a/tensorflow/compiler/tf2xla/lib/while_loop.cc b/tensorflow/compiler/tf2xla/lib/while_loop.cc new file mode 100644 index 0000000000000000000000000000000000000000..86c02ac2e65c12d3527c4022df0cc603e522ef7a --- /dev/null +++ b/tensorflow/compiler/tf2xla/lib/while_loop.cc @@ -0,0 +1,125 @@ +/* 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/compiler/tf2xla/lib/while_loop.h" +#include "tensorflow/compiler/tf2xla/lib/util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" + +namespace tensorflow { + +xla::StatusOr> XlaWhileLoop( + const LoopConditionFunction& condition_function, + const LoopBodyFunction& body_function, + gtl::ArraySlice initial_values, + StringPiece name, xla::ComputationBuilder* builder) { + int arity = initial_values.size(); + std::vector var_shapes; + var_shapes.reserve(arity); + for (const xla::ComputationDataHandle& input : initial_values) { + TF_ASSIGN_OR_RETURN(auto shape, builder->GetShape(input)); + var_shapes.push_back(std::move(*shape)); + } + xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(var_shapes); + + // Unpacks a tuple into its component parts. + auto unpack_tuple = [](xla::ComputationDataHandle tuple, int arity, + xla::ComputationBuilder* builder) { + std::vector elements(arity); + for (int i = 0; i < arity; ++i) { + elements[i] = builder->GetTupleElement(tuple, i); + } + return elements; + }; + + // Build the condition. + std::unique_ptr cond_builder = + builder->CreateSubBuilder(strings::StrCat(name, "_condition")); + { + auto parameter = cond_builder->Parameter(0, tuple_shape, "parameter"); + + TF_ASSIGN_OR_RETURN( + auto result, + condition_function(unpack_tuple(parameter, arity, cond_builder.get()), + cond_builder.get())); + TF_RETURN_IF_ERROR(cond_builder->SetReturnValue(result)); + } + TF_ASSIGN_OR_RETURN(auto cond, cond_builder->Build()); + + // Build the body. + std::unique_ptr body_builder = + builder->CreateSubBuilder(strings::StrCat(name, "_body")); + { + auto parameter = body_builder->Parameter(0, tuple_shape, "parameter"); + + TF_ASSIGN_OR_RETURN( + auto result, + body_function(unpack_tuple(parameter, arity, body_builder.get()), + body_builder.get())); + + TF_RET_CHECK(result.size() == initial_values.size()); + body_builder->Tuple(result); + } + TF_ASSIGN_OR_RETURN(auto body, body_builder->Build()); + + auto outputs = builder->While(cond, body, builder->Tuple(initial_values)); + + return unpack_tuple(outputs, arity, builder); +} + +xla::StatusOr> XlaForEachIndex( + int64 num_iterations, xla::PrimitiveType num_iterations_type, + const ForEachIndexBodyFunction& body_function, + gtl::ArraySlice initial_values, + StringPiece name, xla::ComputationBuilder* builder) { + auto while_cond_fn = [&](gtl::ArraySlice values, + xla::ComputationBuilder* cond_builder) + -> xla::StatusOr { + return cond_builder->Lt( + values[0], + IntegerLiteral(cond_builder, num_iterations_type, num_iterations)); + }; + auto while_body_fn = [&](gtl::ArraySlice values, + xla::ComputationBuilder* body_builder) + -> xla::StatusOr> { + xla::ComputationDataHandle iteration = values[0]; + + std::vector updated_values; + updated_values.reserve(values.size()); + updated_values.push_back(body_builder->Add( + iteration, + body_builder->ConstantLiteral(xla::Literal::One(num_iterations_type)))); + + values.remove_prefix(1); + TF_ASSIGN_OR_RETURN(std::vector body_outputs, + body_function(iteration, values, body_builder)); + updated_values.insert(updated_values.end(), body_outputs.begin(), + body_outputs.end()); + return updated_values; + }; + + std::vector values; + values.reserve(initial_values.size() + 1); + values.push_back( + builder->ConstantLiteral(xla::Literal::Zero(num_iterations_type))); + values.insert(values.end(), initial_values.begin(), initial_values.end()); + + TF_ASSIGN_OR_RETURN(values, XlaWhileLoop(while_cond_fn, while_body_fn, values, + name, builder)); + values.erase(values.begin(), values.begin() + 1); + return values; +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/while_loop.h b/tensorflow/compiler/tf2xla/lib/while_loop.h new file mode 100644 index 0000000000000000000000000000000000000000..2e67a0c99b6deb65fa16ab2dec1727f5cb5fcb92 --- /dev/null +++ b/tensorflow/compiler/tf2xla/lib/while_loop.h @@ -0,0 +1,74 @@ +/* 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_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_ +#define TENSORFLOW_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_ + +#include +#include + +#include "tensorflow/compiler/xla/client/computation.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/lib/gtl/array_slice.h" + +namespace tensorflow { + +// Function that builds a loop condition. Takes as input a sequence of input +// values, and returns a boolean value representing if the condition succeeds. +typedef std::function( + gtl::ArraySlice, xla::ComputationBuilder*)> + LoopConditionFunction; + +// Function that builds a loop body. Takes as input a sequence of input values +// and returns a sequence of output values. +typedef std::function>( + gtl::ArraySlice, xla::ComputationBuilder*)> + LoopBodyFunction; + +// Helper function for building an XLA while loop, where the values carried by +// the loop are a tuple of values, e.g., (a, b, c): +// while( +// condition: (a, b, c) -> bool, +// body: (a, b, c) -> (a, b, c) +// init: (a, b, c) +// ) +// 'name' is a descriptive name for the loop. +xla::StatusOr> XlaWhileLoop( + const LoopConditionFunction& condition_function, + const LoopBodyFunction& body_function, + gtl::ArraySlice initial_values, + StringPiece name, xla::ComputationBuilder* builder); + +// Builds an XLA loop that repeats a computation `num_iterations` times. +// +// The body function (ForEachIndexBodyFunction) takes as input a pair of +// (current iteration number, loop-carried values), and returns an updated +// vector of the loop-carried values. +typedef std::function>( + xla::ComputationDataHandle, gtl::ArraySlice, + xla::ComputationBuilder*)> + ForEachIndexBodyFunction; + +xla::StatusOr> XlaForEachIndex( + int64 num_iterations, xla::PrimitiveType num_iterations_type, + const ForEachIndexBodyFunction& body_function, + gtl::ArraySlice initial_values, + StringPiece name, xla::ComputationBuilder* builder); + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_ diff --git a/tensorflow/compiler/tf2xla/literal_util.cc b/tensorflow/compiler/tf2xla/literal_util.cc index 576cd9bf9abb43e29d9eb8f706e0f42ac2d038e9..2c3cd658e0462368ac0b51938979b7a6815a7574 100644 --- a/tensorflow/compiler/tf2xla/literal_util.cc +++ b/tensorflow/compiler/tf2xla/literal_util.cc @@ -23,44 +23,52 @@ limitations under the License. namespace tensorflow { Status HostTensorToLiteral(const Tensor& host_tensor, xla::Literal* literal) { - literal->Clear(); + xla::Shape literal_shape; TF_RETURN_IF_ERROR(TensorShapeToXLAShape( - host_tensor.dtype(), host_tensor.shape(), literal->mutable_shape())); + host_tensor.dtype(), host_tensor.shape(), &literal_shape)); - literal->Reserve(host_tensor.NumElements()); + *literal = xla::Literal(literal_shape); // memcpy over the payload ... // TODO(phawkins): handle string types. size_t total_bytes = host_tensor.TotalBytes(); if (total_bytes > 0) { - void* dst_ptr = literal->MutableInternalData(); + void* dst_ptr = literal->untyped_data(); const void* src_ptr = DMAHelper::base(&host_tensor); memcpy(dst_ptr, src_ptr, total_bytes); } return Status::OK(); } -Status LiteralToHostTensor(const xla::Literal& literal, DataType target_type, - Tensor* host_tensor) { +Status CopyLiteralToHostTensor(const xla::Literal& literal, + Tensor* host_tensor) { + TF_RET_CHECK(xla::ShapeUtil::IsArray(literal.shape()) && + xla::ShapeUtil::ElementsIn(literal.shape()) == + host_tensor->NumElements()); xla::PrimitiveType primitive_type; - TF_RETURN_IF_ERROR(DataTypeToPrimitiveType(target_type, &primitive_type)); + TF_RETURN_IF_ERROR( + DataTypeToPrimitiveType(host_tensor->dtype(), &primitive_type)); if (literal.shape().element_type() != primitive_type) { return errors::InvalidArgument( "Cannot convert literal of type ", xla::PrimitiveType_Name(literal.shape().element_type()), - " to tensor of type ", DataTypeString(target_type)); + " to tensor of type ", DataTypeString(host_tensor->dtype())); } - - TensorShape shape; - TF_RETURN_IF_ERROR(XLAShapeToTensorShape(literal.shape(), &shape)); - *host_tensor = Tensor(target_type, shape); size_t total_bytes = host_tensor->TotalBytes(); if (total_bytes > 0) { - const void* src_ptr = literal.InternalData(); + const void* src_ptr = literal.untyped_data(); void* dst_ptr = DMAHelper::base(host_tensor); memcpy(dst_ptr, src_ptr, total_bytes); } return Status::OK(); } +Status LiteralToHostTensor(const xla::Literal& literal, DataType target_type, + Tensor* host_tensor) { + TensorShape shape; + TF_RETURN_IF_ERROR(XLAShapeToTensorShape(literal.shape(), &shape)); + *host_tensor = Tensor(target_type, shape); + return CopyLiteralToHostTensor(literal, host_tensor); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/literal_util.h b/tensorflow/compiler/tf2xla/literal_util.h index fe08e83c2391a8b24696961cacfd909d46e49e7d..f283b0236811f8d52e8fe2982a74c11c92cd20d8 100644 --- a/tensorflow/compiler/tf2xla/literal_util.h +++ b/tensorflow/compiler/tf2xla/literal_util.h @@ -29,7 +29,8 @@ namespace tensorflow { // unsupported type. Status HostTensorToLiteral(const Tensor& host_tensor, xla::Literal* literal); -// Copies 'literal' to 'host_tensor', which is allocated of type . +// Copies 'literal' to freshly allocated 'host_tensor', which is allocated of +// type . // Fails if the literal's primitive type != // DataTypeToPrimitiveType(target_type). Note that is not // derivable from the type of , because multiple tensorflow types map @@ -38,6 +39,12 @@ Status HostTensorToLiteral(const Tensor& host_tensor, xla::Literal* literal); Status LiteralToHostTensor(const xla::Literal& literal, DataType target_type, Tensor* host_tensor); +// Copies the contents of 'literal' to a previously allocated tensor +// 'host_tensor'. The tensor and the literal must have the same number of +// elements and the same type. +Status CopyLiteralToHostTensor(const xla::Literal& literal, + Tensor* host_tensor); + } // namespace tensorflow #endif // TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ diff --git a/tensorflow/compiler/tf2xla/python/BUILD b/tensorflow/compiler/tf2xla/python/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..f0a2ef0651ff6115bd201a3b1c34b3c061a22a3d --- /dev/null +++ b/tensorflow/compiler/tf2xla/python/BUILD @@ -0,0 +1,24 @@ +licenses(["notice"]) # Apache 2.0 + +package( + default_visibility = [ + "//learning/tfx:__subpackages__", + "//tensorflow:internal", + ], +) + +load( + "//tensorflow/core:platform/default/build_config.bzl", + "tf_py_clif_cc", +) + +tf_py_clif_cc( + name = "xla_op_registry", + srcs = ["xla_op_registry.clif"], + pyclif_deps = [ + "//tensorflow/core:framework/kernel_def_pyclif", + ], + deps = [ + "//tensorflow/compiler/tf2xla:xla_compiler", + ], +) diff --git a/tensorflow/compiler/tf2xla/python/xla_op_registry.clif b/tensorflow/compiler/tf2xla/python/xla_op_registry.clif new file mode 100644 index 0000000000000000000000000000000000000000..e1ee6cc656a314876fc1fabbebe1bee39a6b2831 --- /dev/null +++ b/tensorflow/compiler/tf2xla/python/xla_op_registry.clif @@ -0,0 +1,7 @@ +from "third_party/tensorflow/core/framework/kernel_def_pyclif.h" import * # KernelDef + +from "third_party/tensorflow/compiler/tf2xla/xla_op_registry.h": + namespace `tensorflow`: + def `XlaOpRegistry::DeviceKernels` as + device_kernels(device: str, include_compilation_only_kernels: bool) -> + list diff --git a/tensorflow/compiler/tf2xla/tf2xla.cc b/tensorflow/compiler/tf2xla/tf2xla.cc index 906f2290433face4cce3296b2f815d50d8c496ce..6051d7dffd7493d8cffb07c1b5d10500e7e75522 100644 --- a/tensorflow/compiler/tf2xla/tf2xla.cc +++ b/tensorflow/compiler/tf2xla/tf2xla.cc @@ -241,9 +241,7 @@ Status CreateXlaArgs(const Graph& graph, XlaCompiler::Argument arg; arg.kind = XlaCompiler::Argument::kParameter; TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &arg.type)); - TensorShape shape; - TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kShapeAttr, &shape)); - TF_RETURN_IF_ERROR(TensorShapeToXLAShape(arg.type, shape, &arg.shape)); + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kShapeAttr, &arg.shape)); TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kDebugNameAttr, &arg.name)); xla_args->push_back(arg); } diff --git a/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.cc b/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.cc index 79da701fd244a461a60588153b601d5c1870fa89..672e19bd93449ccc31f4af5ded23257b197a3c39 100644 --- a/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.cc +++ b/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.cc @@ -29,7 +29,7 @@ XlaCompiledCpuFunction::XlaCompiledCpuFunction(const StaticData& static_data, arg_names_(static_data.arg_names), result_names_(static_data.result_names), program_shape_(static_data.program_shape), - hlo_profile_printer_(static_data.hlo_profile_printer) { + hlo_profile_printer_data_(static_data.hlo_profile_printer_data) { // Allocate arg and temp buffers. if (alloc_mode == AllocMode::ARGS_RESULTS_PROFILES_AND_TEMPS) { alloc_args_ = tensorflow::tfcompile::runtime::MallocContiguousBuffers( diff --git a/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h b/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h index e0ae3ed9a811bcc49ce8862037a67d293e879e57..48a8c083cacf2f6ecf9dc1817b6174c01385d035 100644 --- a/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h +++ b/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h @@ -26,7 +26,7 @@ limitations under the License. // never use this functionality. namespace xla { class ProgramShape; -class HloProfilePrinter; +class HloProfilePrinterData; } namespace tensorflow { @@ -77,12 +77,14 @@ class XlaCompiledCpuFunction { // [Optional] Arg and result shapes. const xla::ProgramShape* program_shape = nullptr; - // [Optional] Profile printer. Null if profiling is disabled. - const xla::HloProfilePrinter* hlo_profile_printer = nullptr; + // [Optional] Profile printer data. Null if profiling is disabled. + const xla::HloProfilePrinterData* hlo_profile_printer_data = nullptr; // [Optional] The number of profile counters expected in the profile counter // buffer by the generated code and hlo_profile_printer. 0 if profiling is - // disabled. + // disabled. This information is already present in + // hlo_profile_printer_data but xla::HloProfilePrinterData is forward + // declared so we don't have access to that information here. int64 profile_counters_size = 0; }; @@ -205,10 +207,12 @@ class XlaCompiledCpuFunction { // program shape isn't available. const xla::ProgramShape* ProgramShape() const { return program_shape_; } - bool hlo_profiling_enabled() const { return hlo_profile_printer_ != nullptr; } - const xla::HloProfilePrinter& hlo_profile_printer() const { + bool hlo_profiling_enabled() const { + return hlo_profile_printer_data_ != nullptr; + } + const xla::HloProfilePrinterData& hlo_profile_printer_data() const { assert(hlo_profiling_enabled()); - return *hlo_profile_printer_; + return *hlo_profile_printer_data_; } private: @@ -234,7 +238,7 @@ class XlaCompiledCpuFunction { const char** arg_names_ = nullptr; const char** result_names_ = nullptr; const xla::ProgramShape* program_shape_ = nullptr; - const xla::HloProfilePrinter* hlo_profile_printer_ = nullptr; + const xla::HloProfilePrinterData* hlo_profile_printer_data_ = nullptr; }; } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index c55719be5523763b897c3fff1fc158b075678934..15bba46ac62a97592656942afc767a303c9b97f3 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -66,13 +66,14 @@ Status CheckSignature(const DataTypeVector& types, bool XlaCompiler::Argument::operator==( const XlaCompiler::Argument& other) const { - if (std::tie(kind, resource_kind, type, name, tensor_array_size, + if (std::tie(kind, resource_kind, type, name, initialized, tensor_array_size, tensor_array_gradients) != std::tie(other.kind, other.resource_kind, other.type, other.name, - other.tensor_array_size, other.tensor_array_gradients)) { + other.initialized, other.tensor_array_size, + other.tensor_array_gradients)) { return false; } - if (!xla::ShapeUtil::Equal(shape, other.shape)) { + if (shape != other.shape) { return false; } if (constant_value.shape() != other.constant_value.shape()) { @@ -108,6 +109,12 @@ XlaCompiler::XlaCompiler(XlaCompiler::Options options) local_flib_runtime_ = local_pflr_->GetFLR(device_->name()); flib_runtime_ = pflr_->GetFLR(device_->name()); + + // The default variable representation shape is the identity function. + if (!options_.variable_representation_shape_fn) { + options_.variable_representation_shape_fn = + [](const TensorShape& shape, DataType type) { return shape; }; + } } XlaCompiler::~XlaCompiler() = default; @@ -152,7 +159,8 @@ std::unique_ptr XlaCompiler::GetGraph(const FunctionBody* fbody) { std::unique_ptr graph(new Graph(options_.flib_def)); CopyGraph(*fbody->graph, graph.get()); OptimizerOptions opts; - opts.set_do_common_subexpression_elimination(true); + opts.set_opt_level(OptimizerOptions::L0); + opts.set_do_common_subexpression_elimination(false); opts.set_do_function_inlining(true); opts.set_do_constant_folding(true); GraphOptimizer optimizer(opts); @@ -183,8 +191,7 @@ Status XlaCompiler::CompileFunction(const XlaCompiler::CompileOptions& options, CheckSignature(fbody->arg_types, args), "Signature check failure while compiling: ", function.name()); - std::unique_ptr graph(new Graph(options_.flib_def)); - CopyGraph(*fbody->graph, graph.get()); + std::unique_ptr graph = GetGraph(fbody); // _Arg and _Retval nodes don't exist in the stored subgraph for the function; // they are added by the function body looked up. Therefore, they don't have @@ -212,15 +219,6 @@ Status XlaCompiler::CompileFunction(const XlaCompiler::CompileOptions& options, *graph); } - // Optimize the graph before running the compiler. - OptimizerOptions opts; - opts.set_do_common_subexpression_elimination(true); - opts.set_do_function_inlining(true); - opts.set_do_constant_folding(true); - GraphOptimizer optimizer(opts); - optimizer.Optimize(flib_runtime_, flib_runtime_->env(), - /*device=*/nullptr, &graph, /*shape_map=*/nullptr); - VLOG(1) << "===================================================="; TF_RETURN_IF_ERROR( CompileGraph(options, function_id, std::move(graph), args, result)); @@ -230,6 +228,68 @@ Status XlaCompiler::CompileFunction(const XlaCompiler::CompileOptions& options, return Status::OK(); } +// Computes the XLA shape for argument 'arg'. +Status XlaCompiler::XLAShapeForArgument(const XlaCompiler::Argument& arg, + xla::Shape* xla_shape) { + switch (arg.kind) { + case XlaCompiler::Argument::kConstant: + return TensorShapeToXLAShape(arg.type, arg.constant_value.shape(), + xla_shape); + case XlaCompiler::Argument::kParameter: + return TensorShapeToXLAShape(arg.type, arg.shape, xla_shape); + case XlaCompiler::Argument::kResource: { + TF_RET_CHECK(arg.initialized); + + switch (arg.resource_kind) { + case XlaResource::kVariable: { + TensorShape representation_shape = + options_.variable_representation_shape_fn(arg.shape, arg.type); + return TensorShapeToXLAShape(arg.type, representation_shape, + xla_shape); + } + case XlaResource::kTensorArray: { + if (arg.tensor_array_size < 0) { + return errors::InvalidArgument( + "Negative tensor_array_size in XLAShapeForArgument"); + } + TensorShape shape; + shape.AddDim(arg.tensor_array_size); + shape.AppendShape(arg.shape); + TF_RETURN_IF_ERROR(TensorShapeToXLAShape(arg.type, shape, xla_shape)); + + if (!arg.tensor_array_gradients.empty()) { + std::vector tuple_shape( + arg.tensor_array_gradients.size() + 1, *xla_shape); + *xla_shape = xla::ShapeUtil::MakeTupleShape(tuple_shape); + } + return Status::OK(); + } + case XlaResource::kStack: { + if (arg.tensor_array_size < 0) { + return errors::InvalidArgument( + "Negative tensor_array_size in XLAShapeForArgument"); + } + TensorShape shape; + shape.AddDim(arg.tensor_array_size); + shape.AppendShape(arg.shape); + xla::Shape buffer_shape; + TF_RETURN_IF_ERROR( + TensorShapeToXLAShape(arg.type, shape, &buffer_shape)); + *xla_shape = xla::ShapeUtil::MakeTupleShape( + {buffer_shape, xla::ShapeUtil::MakeShape(xla::S32, {})}); + return Status::OK(); + } + + case XlaResource::kInvalid: + return errors::Internal( + "Invalid resource type in XLAShapeForArgument()"); + } + } + case XlaCompiler::Argument::kInvalid: + return errors::Internal("Invalid argument type in XLAShapeForArgument()"); + } +} + namespace { Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr graph, @@ -260,146 +320,6 @@ Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr graph, return Status::OK(); } -// Builds XLA computations for each of the arguments to the computation. -// `args` are the arguments to the computation. -Status BuildArguments(const Graph& graph, - const std::vector& args, - bool use_tuple_arg, xla::ComputationBuilder* builder, - XlaContext* context, std::vector* arg_cores, - std::vector* arg_expressions, - std::vector* input_mapping, - std::vector* input_shapes) { - arg_expressions->resize(args.size()); - *arg_cores = std::vector(args.size(), -1); - - // Argument numbers of arguments and resources that are to be passed to the - // XLA computation as runtime parameters. - std::vector parameters, resources; - parameters.reserve(args.size()); - resources.reserve(args.size()); - - // Fills in constant arguments, and computes non-constant argument order. - for (std::vector::size_type i = 0; i < args.size(); - ++i) { - const XlaCompiler::Argument& arg = args[i]; - XlaExpression& arg_expression = (*arg_expressions)[i]; - switch (arg.kind) { - case XlaCompiler::Argument::kResource: - TF_RET_CHECK(arg.resource_kind != XlaResource::kInvalid); - // TODO(phawkins): this code assumes that resource arguments do not - // alias. - XlaResource* resource; - TF_RETURN_IF_ERROR( - context->CreateResource(arg.resource_kind, i, arg.name, arg.type, - xla::ComputationDataHandle(), &resource)); - resource->set_tensor_array_size(arg.tensor_array_size); - arg_expression.set_resource(resource); - if (arg.initialized) { - resources.push_back(i); - } - break; - case XlaCompiler::Argument::kParameter: - parameters.push_back(i); - break; - case XlaCompiler::Argument::kConstant: - arg_expression.set_constant_value(arg.constant_value); - break; - case XlaCompiler::Argument::kInvalid: - return errors::Internal("Unreachable case in BuildArguments()"); - } - } - - // Append parameters containing variable values after the other runtime - // parameters. - parameters.insert(parameters.end(), resources.begin(), resources.end()); - if (parameters.empty()) { - return Status::OK(); - } - - input_shapes->resize(parameters.size()); - input_mapping->resize(parameters.size()); - for (std::vector::size_type i = 0; i < parameters.size(); ++i) { - const XlaCompiler::Argument& arg = args[parameters[i]]; - // Computes the shapes of non-constant arguments. - (*input_shapes)[i] = arg.shape; - (*input_mapping)[i] = parameters[i]; - } - - // Use the _Arg nodes in the graph to resolve core assignments. - for (const Node* n : graph.nodes()) { - if (StringPiece(n->type_string()) != "_Arg") continue; - int index; - TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); - TF_RET_CHECK(index >= 0 && index < args.size()) - << "_Arg out of bounds: " << index << " vs " << args.size(); - TF_ASSIGN_OR_RETURN( - auto sharding, - ParseShardingFromDevice(*n, std::numeric_limits::max())); - if (sharding.has_value()) { - TF_RET_CHECK(sharding.value().type() == - xla::OpSharding::Type::OpSharding_Type_MAXIMAL); - const int core = sharding.value().tile_assignment_devices(0); - if ((*arg_cores)[index] == -1 || core < (*arg_cores)[index]) { - (*arg_cores)[index] = core; - } - } - } - - // Build parameter handles for non-constant arguments. - std::vector arg_handles(parameters.size()); - if (use_tuple_arg) { - xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(*input_shapes); - xla::ComputationDataHandle tuple = - builder->Parameter(0, tuple_shape, "arg_tuple"); - for (std::vector::size_type i = 0; i < parameters.size(); ++i) { - const int core = (*arg_cores)[parameters[i]]; - xla::ScopedShardingAssignment assign_sharding( - builder, core == -1 ? tensorflow::gtl::optional() - : xla::sharding_builder::AssignDevice(core)); - arg_handles[i] = builder->GetTupleElement(tuple, i); - } - } else { - for (std::vector::size_type i = 0; i < parameters.size(); ++i) { - const int core = (*arg_cores)[parameters[i]]; - xla::ScopedShardingAssignment assign_sharding( - builder, core == -1 ? tensorflow::gtl::optional() - : xla::sharding_builder::AssignDevice(core)); - arg_handles[i] = - builder->Parameter(i, (*input_shapes)[i], strings::StrCat("arg", i)); - } - } - - // Fill in the handles in non-constant arguments. - VLOG(2) << "XLA computation inputs:"; - for (std::vector::size_type i = 0; i < parameters.size(); ++i) { - const XlaCompiler::Argument& arg = args[parameters[i]]; - VLOG(2) << " XLA arg " << i - << " shape: " << xla::ShapeUtil::HumanString((*input_shapes)[i]) - << " name: " << arg.name << " TF arg " << parameters[i]; - XlaExpression& arg_expression = (*arg_expressions)[parameters[i]]; - switch (arg.kind) { - case XlaCompiler::Argument::kResource: { - TF_RET_CHECK(arg.initialized); - XlaResource* resource = arg_expression.resource(); - TF_RETURN_IF_ERROR( - resource->SetFromPack(arg.tensor_array_gradients, arg_handles[i], - /*reset_initial_values=*/true, builder)); - VLOG(2) << " resource: num_gradients: " - << arg.tensor_array_gradients.size(); - break; - } - case XlaCompiler::Argument::kParameter: - arg_expression.set_handle(arg_handles[i]); - break; - case XlaCompiler::Argument::kConstant: - case XlaCompiler::Argument::kInvalid: - return errors::Internal("Unreachable case in BuildArguments()"); - } - } - - return Status::OK(); -} - // Builds the XLA computation. // // `retvals` is the list of retvals produced by _Retval operators, in index @@ -464,6 +384,7 @@ Status BuildComputation( XlaCompiler::ResourceUpdate& update = resource_updates->back(); update.input_index = resource->arg_num(); update.type = resource->type(); + update.shape = resource->shape(); update.modified = modified; for (const auto& grad : resource->tensor_array_gradients()) { update.tensor_array_gradients_accessed.insert(grad.first); @@ -501,6 +422,166 @@ Status BuildComputation( } // namespace +// Builds XLA computations for each of the arguments to the computation. +// `args` are the arguments to the computation. +Status XlaCompiler::BuildArguments( + const Graph& graph, const std::vector& args, + bool use_tuple_arg, xla::ComputationBuilder* builder, XlaContext* context, + std::vector* arg_cores, std::vector* arg_expressions, + std::vector* input_mapping, std::vector* input_shapes, + bool is_entry_computation) { + arg_expressions->resize(args.size()); + *arg_cores = std::vector(args.size(), -1); + + // Argument numbers of arguments and resources that are to be passed to the + // XLA computation as runtime parameters. + input_mapping->clear(); + input_mapping->reserve(args.size()); + std::vector resources; + resources.reserve(args.size()); + + // Fills in constant arguments, and computes non-constant argument order. + for (std::vector::size_type i = 0; i < args.size(); + ++i) { + const XlaCompiler::Argument& arg = args[i]; + XlaExpression& arg_expression = (*arg_expressions)[i]; + switch (arg.kind) { + case XlaCompiler::Argument::kResource: + TF_RET_CHECK(arg.resource_kind != XlaResource::kInvalid); + // TODO(phawkins): this code assumes that resource arguments do not + // alias. + XlaResource* resource; + TF_RETURN_IF_ERROR(context->CreateResource( + arg.resource_kind, i, arg.name, arg.type, arg.shape, + xla::ComputationDataHandle(), + /*tensor_array_size=*/arg.tensor_array_size, + /*tensor_array_gradients=*/arg.tensor_array_gradients, &resource)); + arg_expression.set_resource(resource); + if (arg.initialized) { + resources.push_back(i); + } + break; + case XlaCompiler::Argument::kParameter: { + input_mapping->push_back(i); + break; + } + case XlaCompiler::Argument::kConstant: + arg_expression.set_constant_value(arg.constant_value); + break; + case XlaCompiler::Argument::kInvalid: + return errors::Internal("Unreachable case in BuildArguments()"); + } + } + + // Append parameters containing variable values after the other runtime + // parameters. + input_mapping->insert(input_mapping->end(), resources.begin(), + resources.end()); + if (input_mapping->empty()) { + return Status::OK(); + } + + std::vector arg_shapes(input_mapping->size()); + for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { + // Computes the shapes of non-constant arguments. + TF_RETURN_IF_ERROR( + XLAShapeForArgument(args[(*input_mapping)[i]], &arg_shapes[i])); + } + + if (use_tuple_arg) { + input_shapes->push_back(xla::ShapeUtil::MakeTupleShape(arg_shapes)); + } else { + *input_shapes = arg_shapes; + } + + // Use the _Arg nodes in the graph to resolve core assignments. + for (const Node* n : graph.nodes()) { + if (StringPiece(n->type_string()) != "_Arg") continue; + int index; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); + TF_RET_CHECK(index >= 0 && index < args.size()) + << "_Arg out of bounds: " << index << " vs " << args.size(); + TF_ASSIGN_OR_RETURN( + auto sharding, + ParseShardingFromDevice(*n, std::numeric_limits::max())); + if (sharding.has_value()) { + TF_RET_CHECK(sharding.value().type() == + xla::OpSharding::Type::OpSharding_Type_MAXIMAL); + const int core = sharding.value().tile_assignment_devices(0); + if ((*arg_cores)[index] == -1 || core < (*arg_cores)[index]) { + (*arg_cores)[index] = core; + } + } + } + + // Build parameter handles for non-constant arguments. + std::vector arg_handles(input_mapping->size()); + if (use_tuple_arg) { + xla::ComputationDataHandle tuple; + if (is_entry_computation) { + xla::OpSharding tuple_sharding; + tuple_sharding.set_type(xla::OpSharding::Type::OpSharding_Type_TUPLE); + for (int64 parameter : *input_mapping) { + const int core = (*arg_cores)[parameter]; + const int root_device = 0; + *tuple_sharding.add_tuple_shardings() = + core == -1 ? xla::sharding_builder::AssignDevice(root_device) + : xla::sharding_builder::AssignDevice(core); + } + xla::ScopedShardingAssignment assign_tuple_sharding(builder, + tuple_sharding); + tuple = builder->Parameter(0, (*input_shapes)[0], "arg_tuple"); + } else { + tuple = builder->Parameter(0, (*input_shapes)[0], "arg_tuple"); + } + for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { + const int core = (*arg_cores)[input_mapping->at(i)]; + xla::ScopedShardingAssignment assign_sharding( + builder, core == -1 ? tensorflow::gtl::optional() + : xla::sharding_builder::AssignDevice(core)); + arg_handles[i] = builder->GetTupleElement(tuple, i); + } + } else { + for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { + const int core = (*arg_cores)[input_mapping->at(i)]; + xla::ScopedShardingAssignment assign_sharding( + builder, core == -1 ? tensorflow::gtl::optional() + : xla::sharding_builder::AssignDevice(core)); + arg_handles[i] = + builder->Parameter(i, (*input_shapes)[i], strings::StrCat("arg", i)); + } + } + + // Fill in the handles in non-constant arguments. + VLOG(2) << "XLA computation inputs:"; + for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { + const XlaCompiler::Argument& arg = args[input_mapping->at(i)]; + VLOG(2) << " XLA arg " << i + << " shape: " << xla::ShapeUtil::HumanString(arg_shapes[i]) + << " name: " << arg.name << " TF arg " << input_mapping->at(i); + XlaExpression& arg_expression = (*arg_expressions)[input_mapping->at(i)]; + switch (arg.kind) { + case XlaCompiler::Argument::kResource: { + TF_RET_CHECK(arg.initialized); + XlaResource* resource = arg_expression.resource(); + TF_RETURN_IF_ERROR(resource->SetFromPack(arg.tensor_array_gradients, + arg_handles[i], builder)); + VLOG(2) << " resource: num_gradients: " + << arg.tensor_array_gradients.size(); + break; + } + case XlaCompiler::Argument::kParameter: + arg_expression.set_handle(arg_handles[i]); + break; + case XlaCompiler::Argument::kConstant: + case XlaCompiler::Argument::kInvalid: + return errors::Internal("Unreachable case in BuildArguments()"); + } + } + + return Status::OK(); +} + Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, string const& name, std::unique_ptr graph, @@ -525,14 +606,16 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, xla::ComputationBuilder builder(client(), name); XlaContext* context = new XlaContext(this, &builder, options_.allow_cpu_custom_calls, - options.resolve_compile_time_constants); + options.resolve_compile_time_constants, + &options_.variable_representation_shape_fn); core::ScopedUnref context_unref(context); std::vector arg_expressions; std::vector arg_cores; - TF_RETURN_IF_ERROR(BuildArguments( - *graph, args, options.use_tuple_arg, &builder, context, &arg_cores, - &arg_expressions, &result->input_mapping, &result->xla_input_shapes)); + TF_RETURN_IF_ERROR( + BuildArguments(*graph, args, options.use_tuple_arg, &builder, context, + &arg_cores, &arg_expressions, &result->input_mapping, + &result->xla_input_shapes, options.is_entry_computation)); context->set_args(std::move(arg_expressions)); TF_RETURN_IF_ERROR(ExecuteGraph(context, std::move(graph), device_, @@ -593,13 +676,6 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, ++computation_output; } } - - for (std::vector::size_type i = 0; - i < result->resource_updates.size(); ++i) { - result->resource_updates[i].shape = xla::ShapeUtil::GetTupleElementShape( - result->xla_output_shape, computation_output); - ++computation_output; - } return Status::OK(); } diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index 380e24e96bc713af4453f92a5359995e9ab4734a..c4449bc4be06daff856eff70c6d89be6ddbcf0ee 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -29,6 +29,9 @@ limitations under the License. #include "tensorflow/core/public/version.h" namespace tensorflow { + +class XlaContext; + // The XlaCompiler class is responsible for compilation of a self-contained // subgraph of a TensorFlow computation using the XLA linear algebra runtime. // It does a symbolic execution of the graph starting from specific input @@ -104,9 +107,17 @@ class XlaCompiler { // is the type of the variable's value, not DT_RESOURCE. DataType type; - // The shape of the argument. If the argument is a resource, this is the - // shape of the resource's value. - xla::Shape shape; + // The shape of the argument. For: + // * a parameter: the shape of the parameter. + // * a constant: ignored; the shape given by constant_value is used + // instead. + // * an uninitialized resource: ignored. We don't yet know the shape of an + // uninitialized resource (otherwise we would have initialized it!) + // * an initialized variable: the shape of the variable's value. + // * an initialized TensorArray or Stack resource: the shape of an entry in + // the TensorArray/Stack. Note this is the size of a single entry, not the + // XLA data structure that represents the complete stack/array. + TensorShape shape; // The value of the argument, if it is a compile-time constant. Must be a // host-memory tensor. @@ -152,6 +163,10 @@ class XlaCompiler { // as Tensors at compile-time, rather than as run-time outputs of the // computation. bool resolve_compile_time_constants = true; + + // True when compiling the entry computation, false for subcomputations + // (while, call, etc.) + bool is_entry_computation = true; }; struct OutputDescription { @@ -171,8 +186,9 @@ class XlaCompiler { int input_index; // Type and shape of the tensor to be written back. + // The `shape` field has the same meaning as the Argument::shape field. DataType type; - xla::Shape shape; + TensorShape shape; // Was the value of the variable modified by the computation? // (Always true, unless `return_updated_values_for_all_resources` is true.) @@ -226,11 +242,30 @@ class XlaCompiler { // for CPU. bool allow_cpu_custom_calls = false; + // If set, the XLA representation of variables represented to XLA as the + // shape given by this shape function. Variables are reshaped to this shape + // on write, and reshaped to their original shape on read. + std::function + variable_representation_shape_fn; + // If not nullptr, populate_resource_manager is called with the // compilation device's resource manager when the compilation // device is created, and can be used to create metadata objects // that can be accessed by XLA op kernels. std::function* populate_resource_manager = nullptr; + + // If not nullptr, this memory allocator can be used by the compiler for + // temporary allocations it might want to make during compilation. + // + // For example, the compiler may want to try out different algorithms and + // choose the fastest one, and it might run those algorithms over buffers + // created using this allocator. + // + // The compiler can function correctly without an explicit allocator given + // here, but on some devices (notably, GPUs), TensorFlow tends to eagerly + // allocate most or all available memory on the device, leaving none for the + // compiler to access, unless it can use TensorFlow's allocator. + xla::DeviceMemoryAllocator* device_allocator = nullptr; }; explicit XlaCompiler(Options options); @@ -249,11 +284,10 @@ class XlaCompiler { const std::vector& args, CompilationResult* result); - Status PrepareArguments(xla::ComputationBuilder* builder, NameAttrList func, - const std::vector& types, - const std::vector& shapes, - const std::vector& expressions, - std::vector* args); + // Returns the shape of the XLA parameter for an argument 'arg'. + // See the class comment for more details about the argument passing + // convention. + Status XLAShapeForArgument(const Argument& arg, xla::Shape* xla_shape); // Retrieves the channel handle associated with `key`. Allocates // a new channel handle if none exists. @@ -274,6 +308,17 @@ class XlaCompiler { // Returns the optimized graph object in this function body. std::unique_ptr GetGraph(const FunctionBody* fbody); + // Builds XLA computations for each of the arguments to the computation. + // `args` are the arguments to the computation. + Status BuildArguments(const Graph& graph, + const std::vector& args, + bool use_tuple_arg, xla::ComputationBuilder* builder, + XlaContext* context, std::vector* arg_cores, + std::vector* arg_expressions, + std::vector* input_mapping, + std::vector* input_shapes, + bool is_entry_computation); + // Graph compiler needs to know how to get an optimized graph from a function // body. friend class GraphCompiler; diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index 93aae8485d157cd4afbf804d695d5c0ab8d7946c..a18eeacd41808884fac9ec5d617cb0d274ea27d8 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/ops/data_flow_ops.h" #include "tensorflow/cc/ops/function_ops.h" +#include "tensorflow/cc/ops/resource_variable_ops.h" #include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" @@ -191,10 +192,10 @@ TEST_F(XlaCompilerTest, Simple) { std::vector args(2); args[0].kind = XlaCompiler::Argument::kParameter; args[0].type = DT_INT32; - args[0].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); + args[0].shape = TensorShape({2}); args[1].kind = XlaCompiler::Argument::kParameter; args[1].type = DT_INT32; - args[1].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); + args[1].shape = TensorShape({2}); // Compiles the graph. XlaCompiler compiler(DefaultOptions()); @@ -227,6 +228,42 @@ TEST_F(XlaCompilerTest, Simple) { xla::LiteralTestUtil::ExpectEqual(*expected_literal, *actual_literal); } +TEST_F(XlaCompilerTest, HasSaneErrorOnNonCompileTimeConstantInputToReshape) { + // Builds a graph that adds reshapes a tensor, but with the shape not + // statically known. + Scope scope = Scope::NewRootScope().ExitOnError(); + auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); + auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1); + auto c = ops::Reshape(scope.WithOpName("C"), a, b); + auto d = ops::_Retval(scope.WithOpName("D"), c, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + // Builds a description of the arguments. + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = TensorShape({2}); + args[1].kind = XlaCompiler::Argument::kParameter; + args[1].type = DT_INT32; + args[1].shape = TensorShape({2}); + + // Compiles the graph. + XlaCompiler compiler(DefaultOptions()); + + XlaCompiler::CompilationResult result; + Status status = + compiler.CompileGraph(XlaCompiler::CompileOptions(), "reshape", + std::move(graph), args, &result); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE( + StringPiece(status.error_message()).contains("depends on a parameter")) + << status.error_message(); + EXPECT_TRUE( + StringPiece(status.error_message()).contains("[[Node: C = Reshape")) + << status.error_message(); +} + // Tests handling of compile-time constant outputs. TEST_F(XlaCompilerTest, ConstantOutputs) { // Builds a graph with one compile-time constant output and one data-dependent @@ -245,7 +282,7 @@ TEST_F(XlaCompilerTest, ConstantOutputs) { std::vector args(1); args[0].kind = XlaCompiler::Argument::kParameter; args[0].type = DT_INT32; - args[0].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); + args[0].shape = TensorShape({2}); XlaCompiler::Options options = DefaultOptions(); XlaCompiler compiler(options); @@ -337,7 +374,7 @@ TEST_F(XlaCompilerTest, ResourceManager) { std::vector args(1); args[0].kind = XlaCompiler::Argument::kParameter; args[0].type = DT_INT32; - args[0].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); + args[0].shape = TensorShape({2}); DummyResourceForTest* resource = new DummyResourceForTest(); @@ -384,7 +421,7 @@ TEST_F(XlaCompilerTest, DeterministicCompilation) { std::vector args(1); args[0].kind = XlaCompiler::Argument::kParameter; args[0].type = DT_INT32; - args[0].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); + args[0].shape = TensorShape({2}); // Compiles the graph. auto options = DefaultOptions(); @@ -436,9 +473,7 @@ TEST_F(XlaCompilerTest, CanPassTensorArraysToAndFromComputation) { args[0].resource_kind = XlaResource::kTensorArray; args[0].initialized = true; args[0].type = DT_INT32; - args[0].shape = xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::S32, {2}), - xla::ShapeUtil::MakeShape(xla::S32, {2})}); + args[0].shape = TensorShape({}); args[0].tensor_array_size = 2; args[0].tensor_array_gradients = {"grad2"}; @@ -504,9 +539,7 @@ TEST_F(XlaCompilerTest, UnwrittenTensorArrayGradientsAreNotComputationOutputs) { args[0].resource_kind = XlaResource::kTensorArray; args[0].initialized = true; args[0].type = DT_INT32; - args[0].shape = xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::S32, {2}), - xla::ShapeUtil::MakeShape(xla::S32, {2})}); + args[0].shape = TensorShape({}); args[0].tensor_array_size = 2; args[0].tensor_array_gradients = {"grad1"}; @@ -538,9 +571,7 @@ TEST_F(XlaCompilerTest, NewTensorArrayGradientsAreComputationOutputs) { args[0].resource_kind = XlaResource::kTensorArray; args[0].initialized = true; args[0].type = DT_INT32; - args[0].shape = xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::S32, {2}), - xla::ShapeUtil::MakeShape(xla::S32, {2})}); + args[0].shape = TensorShape({}); args[0].tensor_array_size = 2; args[0].tensor_array_gradients = {"grad1"}; @@ -653,5 +684,128 @@ TEST_F(XlaCompilerTest, LocalFunctionWithWrongArgumentsFail) { << status.error_message(); } +// Tests a simple graph that reads and writes a variable. +TEST_F(XlaCompilerTest, Variables) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); + auto var = ops::_Arg(scope.WithOpName("V"), DT_RESOURCE, 1); + auto write = ops::AssignAddVariableOp(scope, var, a); + auto read = ops::ReadVariableOp( + scope.WithControlDependencies(std::vector{write}), var, + DT_INT32); + auto read_plus_one = ops::Add(scope, read, ops::Const(scope, 1)); + auto d = ops::_Retval(scope.WithOpName("D"), read_plus_one, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + // Builds a description of the arguments. + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = TensorShape({2}); + args[1].kind = XlaCompiler::Argument::kResource; + args[1].resource_kind = XlaResource::kVariable; + args[1].initialized = true; + args[1].type = DT_INT32; + args[1].shape = TensorShape({2}); + + // Compiles the graph. + XlaCompiler compiler(DefaultOptions()); + + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + std::move(graph), args, &result)); + + // Tests that the generated computation works. + std::unique_ptr param0_literal = + xla::Literal::CreateR1({7, 42}); + std::unique_ptr param1_literal = + xla::Literal::CreateR1({-3, 101}); + std::unique_ptr param0_data = + client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); + std::unique_ptr param1_data = + client_->TransferToServer(*param1_literal).ConsumeValueOrDie(); + + std::unique_ptr actual = + client_ + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) + .ConsumeValueOrDie(); + std::unique_ptr actual_literal = + client_->Transfer(*actual).ConsumeValueOrDie(); + + std::unique_ptr expected0 = + xla::Literal::CreateR1({5, 144}); + std::unique_ptr expected1 = + xla::Literal::CreateR1({4, 143}); + std::unique_ptr expected_literal = + xla::Literal::MakeTuple({expected0.get(), expected1.get()}); + xla::LiteralTestUtil::ExpectEqual(*expected_literal, *actual_literal); +} + +// Tests a simple graph that reads and writes a variable, with a +// variable_representation_shape_fn passed to the compiler that flattens all +// variable tensors to vectors. +TEST_F(XlaCompilerTest, VariableRepresentationShapeFunction) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); + auto var = ops::_Arg(scope.WithOpName("V"), DT_RESOURCE, 1); + auto write = ops::AssignAddVariableOp(scope, var, a); + auto read = ops::ReadVariableOp( + scope.WithControlDependencies(std::vector{write}), var, + DT_INT32); + auto read_plus_one = ops::Add(scope, read, ops::Const(scope, 1)); + auto d = ops::_Retval(scope.WithOpName("D"), read_plus_one, 0); + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + // Builds a description of the arguments. + std::vector args(2); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = TensorShape({2, 2}); + args[1].kind = XlaCompiler::Argument::kResource; + args[1].resource_kind = XlaResource::kVariable; + args[1].initialized = true; + args[1].type = DT_INT32; + args[1].shape = TensorShape({2, 2}); + + // Compiles the graph. + XlaCompiler::Options options = DefaultOptions(); + options.variable_representation_shape_fn = [](const TensorShape& shape, + DataType type) { + return TensorShape({shape.num_elements()}); + }; + XlaCompiler compiler(options); + + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + std::move(graph), args, &result)); + + // Tests that the generated computation works. + std::unique_ptr param0_literal = + xla::Literal::CreateR2({{4, 55}, {1, -3}}); + std::unique_ptr param1_literal = + xla::Literal::CreateR1({22, 11, 33, 404}); + std::unique_ptr param0_data = + client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); + std::unique_ptr param1_data = + client_->TransferToServer(*param1_literal).ConsumeValueOrDie(); + + std::unique_ptr actual = + client_ + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) + .ConsumeValueOrDie(); + std::unique_ptr actual_literal = + client_->Transfer(*actual).ConsumeValueOrDie(); + + std::unique_ptr expected0 = + xla::Literal::CreateR2({{27, 67}, {35, 402}}); + std::unique_ptr expected1 = + xla::Literal::CreateR1({26, 66, 34, 401}); + std::unique_ptr expected_literal = + xla::Literal::MakeTuple({expected0.get(), expected1.get()}); + xla::LiteralTestUtil::ExpectEqual(*expected_literal, *actual_literal); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_context.cc b/tensorflow/compiler/tf2xla/xla_context.cc index e8d17e2e0a1ba01f16d4bbbd2895b112f4dd1989..8423921086fec1cf534cf613102fc3839035cb85 100644 --- a/tensorflow/compiler/tf2xla/xla_context.cc +++ b/tensorflow/compiler/tf2xla/xla_context.cc @@ -62,13 +62,16 @@ void XlaContext::set_args(std::vector args) { args_ = std::move(args); } -XlaContext::XlaContext(XlaCompiler* compiler, xla::ComputationBuilder* builder, - bool allow_cpu_custom_calls, - bool resolve_compile_time_constants) +XlaContext::XlaContext( + XlaCompiler* compiler, xla::ComputationBuilder* builder, + bool allow_cpu_custom_calls, bool resolve_compile_time_constants, + const std::function* + variable_representation_shape_fn) : compiler_(compiler), builder_(builder), allow_cpu_custom_calls_(allow_cpu_custom_calls), - resolve_compile_time_constants_(resolve_compile_time_constants) {} + resolve_compile_time_constants_(resolve_compile_time_constants), + variable_representation_shape_fn_(variable_representation_shape_fn) {} string XlaContext::DebugString() { return "TLA JIT context"; } @@ -103,16 +106,23 @@ Status XlaContext::AddConstRetval(int retval_index, DataType dtype, xla::ComputationBuilder* XlaContext::builder() { return builder_; } -Status XlaContext::CreateResource(XlaResource::Kind kind, int arg_num, - string name, DataType type, - const xla::ComputationDataHandle& handle, - XlaResource** resource) { +Status XlaContext::CreateResource( + XlaResource::Kind kind, int arg_num, string name, DataType type, + TensorShape shape, const xla::ComputationDataHandle& handle, + int64 tensor_array_size, const std::set& tensor_array_gradients, + XlaResource** resource) { resources_.emplace_back( - new XlaResource(kind, arg_num, std::move(name), type, handle)); + new XlaResource(kind, arg_num, std::move(name), type, std::move(shape), + handle, tensor_array_size, tensor_array_gradients)); *resource = resources_.back().get(); return Status::OK(); } +TensorShape XlaContext::VariableRepresentationShape(const TensorShape& shape, + DataType type) const { + return (*variable_representation_shape_fn_)(shape, type); +} + const xla::Computation* XlaContext::GetOrCreateMax(const DataType type) { return LookupOrCreate(type, &max_func_, [this, type] { const string type_string = DataTypeString(type); diff --git a/tensorflow/compiler/tf2xla/xla_context.h b/tensorflow/compiler/tf2xla/xla_context.h index 1a7dafe8cdb56cc9b8fcd3ba6e262c21c2a07d90..00fbaba37c542954f690b310a184cff985a05156 100644 --- a/tensorflow/compiler/tf2xla/xla_context.h +++ b/tensorflow/compiler/tf2xla/xla_context.h @@ -44,7 +44,9 @@ class XlaContext : public ResourceBase { // Creates a new XlaContext. XlaContext(XlaCompiler* compiler, xla::ComputationBuilder* builder, - bool allow_cpu_custom_calls, bool resolve_compile_time_constants); + bool allow_cpu_custom_calls, bool resolve_compile_time_constants, + const std::function* + variable_representation_shape_fn); // Virtual method defined by ResourceBase. string DebugString() override; @@ -71,17 +73,26 @@ class XlaContext : public ResourceBase { Status AddConstRetval(int retval_index, DataType dtype, const xla::Literal& literal); - // Creates a resource with resource `kind` and initial type `type` and - // value `handle`. `name` is a descriptive name for use in error messages. + // Creates a resource with resource `kind` and initial value `handle`. `name` + // is a descriptive name for use in error messages. See the `XlaResource` + // constructor for a description of the remaining arguments. // Fails if the resource already exists. Status CreateResource(XlaResource::Kind kind, int arg_num, string name, - DataType type, const xla::ComputationDataHandle& handle, + DataType type, TensorShape shape, + const xla::ComputationDataHandle& handle, + int64 tensor_array_size, + const std::set& tensor_array_gradients, XlaResource** resource); const std::vector>& resources() { return resources_; } + // Returns the XLA shape to be used to represent a variable of TF `shape` + // and `type`. + TensorShape VariableRepresentationShape(const TensorShape& shape, + DataType type) const; + // Get an XLA lambda to compute Max. This is cached in the // XlaContext since it may be used by multiple Ops. There is a // separate specialization of the computation for each DataType. @@ -129,6 +140,11 @@ class XlaContext : public ResourceBase { // Holds ownership of resources. The resources are not ordered. std::vector> resources_; + // A function that describes how variable shapes should be represented + // in XLA. Variable values will be reshaped to this shape. Must be non-null. + const std::function* + variable_representation_shape_fn_; + // Cache of prebuilt computations indexed by their type. using ComputationMap = std::map; diff --git a/tensorflow/compiler/tf2xla/xla_helpers.cc b/tensorflow/compiler/tf2xla/xla_helpers.cc index ec9e535b707beec6ea26dc81c7ee76b1d4da9225..f048662953e20b2a612271e2daeef6e370c4822a 100644 --- a/tensorflow/compiler/tf2xla/xla_helpers.cc +++ b/tensorflow/compiler/tf2xla/xla_helpers.cc @@ -135,57 +135,9 @@ xla::ComputationDataHandle XlaHelpers::Epsilon(xla::ComputationBuilder* b, xla::ComputationDataHandle XlaHelpers::IntegerLiteral( xla::ComputationBuilder* b, DataType data_type, int64 value) { - xla::Literal literal; xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type)); - switch (type) { - case xla::U8: - literal = *xla::Literal::CreateR0(value); - break; - case xla::U32: - literal = *xla::Literal::CreateR0(value); - break; - case xla::U64: - literal = *xla::Literal::CreateR0(value); - break; - case xla::S8: - literal = *xla::Literal::CreateR0(value); - break; - case xla::S32: - literal = *xla::Literal::CreateR0(value); - break; - case xla::S64: - literal = *xla::Literal::CreateR0(value); - break; - case xla::F32: - literal = *xla::Literal::CreateR0(value); - break; - case xla::F64: - literal = *xla::Literal::CreateR0(value); - break; - case xla::C64: - literal = *xla::Literal::CreateR0(value); - break; - case xla::PRED: - LOG(FATAL) << "pred element type is not integral"; - case xla::S16: - case xla::U16: - LOG(FATAL) << "u16/s16 literals not yet implemented"; - case xla::BF16: - literal = *xla::Literal::CreateR0(static_cast(value)); - break; - case xla::F16: - literal = - *xla::Literal::CreateR0(static_cast(value)); - break; - case xla::TUPLE: - LOG(FATAL) << "tuple element type is not integral"; - case xla::OPAQUE: - LOG(FATAL) << "opaque element type is not integral"; - default: - LOG(FATAL) << "unhandled element type " << type; - } - return b->ConstantLiteral(literal); + return ::tensorflow::IntegerLiteral(b, type, value); } xla::ComputationDataHandle XlaHelpers::FloatLiteral(xla::ComputationBuilder* b, @@ -212,8 +164,8 @@ xla::ComputationDataHandle XlaHelpers::FloatLiteral(xla::ComputationBuilder* b, "elements."); } - *output = input; - output->mutable_shape()->Swap(&shape); + *output = input.Clone(); + output->mutable_shape_do_not_use()->Swap(&shape); return Status::OK(); } diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc index 584417bc72c8f6645c05912e857b031cfb394e54..1fe6e69ff2dc838152032ac3d7b21de41684c6f6 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc @@ -182,10 +182,10 @@ XlaJitCompiledCpuFunction::Compile( jit->static_data_.program_shape = jit->program_shape_.get(); if (cpu_executable->hlo_profiling_enabled()) { - jit->static_data_.hlo_profile_printer = - &cpu_executable->hlo_profile_printer(); + jit->static_data_.hlo_profile_printer_data = + &cpu_executable->hlo_profile_printer_data(); jit->static_data_.profile_counters_size = - cpu_executable->hlo_profile_printer().profile_counters_size(); + cpu_executable->hlo_profile_printer_data().profile_counters_size(); } return std::move(jit_unique_ptr); diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc index 73a91bfade74ed5f792b58575a73df95ea352658..c4bb90d58755f16672ca7c6a6738065be6330485 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc @@ -118,13 +118,36 @@ Status XlaOpKernelContext::ConstantInputReshaped( std::iota(layout_indices.rbegin(), layout_indices.rend(), 0); xla::Layout layout = xla::LayoutUtil::MakeLayout(layout_indices); + xla::StatusOr is_constant = builder()->IsConstant(handle); + if (!is_constant.ok()) { + Status status = is_constant.status(); + errors::AppendToMessage(&status, "while evaluating input ", index, " of ", + context_->op_kernel().type_string(), + " operator as a compile-time constant."); + return status; + } + + if (!is_constant.ValueOrDie()) { + return errors::InvalidArgument( + "Input ", index, " to ", context_->op_kernel().type_string(), + " operator must be a compile-time constant.\n" + "\n" + "XLA compilation requires that operator arguments that represent " + "shapes or dimensions be evaluated to concrete values at compile time. " + "This error means that a shape or dimension argument could not be " + "evaluated at compile time, usually because the value of the argument " + "depends on a parameter to the computation, on a variable, or on a " + "stateful operation such as a random number generator."); + } + // Ask the XLA compiler to evaluate the data handle to a literal. xla::StatusOr> computed = builder()->ComputeConstant(handle, &layout); if (!computed.ok()) { - return errors::InvalidArgument( - "Error evaluating ", context_->op_kernel().name(), " input ", index, - ": ", computed.status().error_message()); + return errors::Internal("Error evaluating ", context_->op_kernel().name(), + " input ", index, + "as a compile-time constant.\nError: ", + computed.status().error_message()); } *constant_literal = std::move(*computed.ValueOrDie()); @@ -206,15 +229,15 @@ Status XlaOpKernelContext::ConstantInputAsInt64Literal(int index, xla::Literal literal; TF_RETURN_IF_ERROR(ConstantInput(index, &literal)); switch (literal.shape().element_type()) { - case xla::S32: - out->Clear(); - *out->mutable_shape() = literal.shape(); - out->mutable_shape()->set_element_type(xla::S64); - for (int32 x : literal.s32s()) { - out->add_s64s(x); + case xla::S32: { + *out = xla::Literal( + xla::ShapeUtil::ChangeElementType(literal.shape(), xla::S64)); + auto src_data = literal.data(); + for (int64 i = 0; i < src_data.size(); ++i) { + out->data()[i] = src_data[i]; } return Status::OK(); - + } case xla::S64: *out = std::move(literal); return Status::OK(); @@ -263,7 +286,8 @@ Status XlaOpKernelContext::ConstantInputList( } Status XlaOpKernelContext::ReadVariableInput( - int index, xla::ComputationDataHandle* value) { + int index, DataType type, TensorShape* shape, + xla::ComputationDataHandle* value) { const Tensor& tensor = context_->input(index); const XlaExpression* expression = CastExpressionFromTensor(tensor); XlaResource* variable = expression->resource(); @@ -273,7 +297,24 @@ Status XlaOpKernelContext::ReadVariableInput( return errors::InvalidArgument("Read of uninitialized variable ", variable->name()); } - *value = variable->value(); + if (variable->type() != type) { + return errors::InvalidArgument( + "Type mismatch for read of variable ", variable->name(), ". Expected ", + DataTypeString(type), "; got ", DataTypeString(variable->type())); + } + if (shape) { + *shape = variable->shape(); + } + + XlaContext& xla_context = XlaContext::Get(context_); + TensorShape representation_shape = xla_context.VariableRepresentationShape( + variable->shape(), variable->type()); + if (representation_shape == variable->shape()) { + *value = variable->value(); + } else { + *value = + builder()->Reshape(variable->value(), variable->shape().dim_sizes()); + } return Status::OK(); } @@ -289,12 +330,7 @@ Status XlaOpKernelContext::GetVariableTypeAndShape(int index, DataType* type, variable->name()); } *type = variable->type(); - auto shape_or_status = builder()->GetShape(variable->value()); - if (!shape_or_status.ok()) { - return shape_or_status.status(); - } - TF_RETURN_IF_ERROR( - XLAShapeToTensorShape(*shape_or_status.ValueOrDie(), shape)); + *shape = variable->shape(); return Status::OK(); } @@ -373,8 +409,8 @@ Status XlaOpKernelContext::GetResourceInput(int index, XlaResource** resource) { return Status::OK(); } -Status XlaOpKernelContext::AssignVariable( - int input_index, DataType type, const xla::ComputationDataHandle& handle) { +Status XlaOpKernelContext::AssignVariable(int input_index, DataType type, + xla::ComputationDataHandle handle) { TF_RET_CHECK(handle.handle() != 0); const XlaExpression* expression = @@ -382,17 +418,44 @@ Status XlaOpKernelContext::AssignVariable( XlaResource* variable = expression->resource(); TF_RET_CHECK(variable != nullptr); TF_RET_CHECK(variable->kind() == XlaResource::kVariable); - return variable->SetValue(type, handle); + + auto shape_or_status = builder()->GetShape(handle); + if (!shape_or_status.ok()) { + return shape_or_status.status(); + } + TensorShape shape; + TF_RETURN_IF_ERROR( + XLAShapeToTensorShape(*shape_or_status.ValueOrDie(), &shape)); + + TF_RETURN_IF_ERROR(variable->SetTypeAndShape(type, shape)); + + XlaContext& xla_context = XlaContext::Get(context_); + TensorShape representation_shape = + xla_context.VariableRepresentationShape(shape, type); + if (shape != representation_shape) { + handle = builder()->Reshape(handle, representation_shape.dim_sizes()); + } + return variable->SetValue(handle); } XlaCompiler* XlaOpKernelContext::compiler() const { return XlaContext::Get(context_).compiler(); } -void XlaOpKernelContext::CtxFailure(Status s) { context_->CtxFailure(s); } -void XlaOpKernelContext::CtxFailureWithWarning(Status s) { +void XlaOpKernelContext::CtxFailure(const Status& s) { + context_->CtxFailure(s); +} +void XlaOpKernelContext::CtxFailureWithWarning(const Status& s) { context_->CtxFailureWithWarning(s); } +void XlaOpKernelContext::CtxFailure(const char* file, int line, + const Status& s) { + context_->CtxFailure(file, line, s); +} +void XlaOpKernelContext::CtxFailureWithWarning(const char* file, int line, + const Status& s) { + context_->CtxFailureWithWarning(file, line, s); +} const xla::Computation* XlaOpKernelContext::GetOrCreateMax( const DataType type) { diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.h b/tensorflow/compiler/tf2xla/xla_op_kernel.h index f1ae81a5aa9d507a3e0dd577568377385b1844e6..4e4b97e0cec8d16b9b5686a779b1285906765dbd 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.h +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.h @@ -164,17 +164,24 @@ class XlaOpKernelContext { TensorShape* shape) const; // Reads the current value of the resouce variable referred to by input - // 'index'. - Status ReadVariableInput(int index, xla::ComputationDataHandle* value); + // 'index'. If `shape` is not nullptr, sets `*shape` to the shape of the + // variable. Returns an error if the variable has not been initialized, or if + // its type does not match `type`. + Status ReadVariableInput(int index, DataType type, TensorShape* shape, + xla::ComputationDataHandle* value); // Assigns the value `handle` to the variable referenced by input - // `input_index`. Marks the operator as having side effects. + // `input_index`. The variable must be of `type`. Returns an error if the + // variable has been initialized with a different type or with a + // different shape. Status AssignVariable(int input_index, DataType type, - const xla::ComputationDataHandle& handle); + xla::ComputationDataHandle handle); // Helper routines for the OP_REQUIRES macros - void CtxFailure(Status s); - void CtxFailureWithWarning(Status s); + void CtxFailure(const Status& s); + void CtxFailureWithWarning(const Status& s); + void CtxFailure(const char* file, int line, const Status& s); + void CtxFailureWithWarning(const char* file, int line, const Status& s); // If this kernel invocation is within a function execution, // call_frame() returns the call frame for the function call. diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.cc b/tensorflow/compiler/tf2xla/xla_op_registry.cc index 97bb100fb1e8fdcc1233d876939e28c67f275b2f..bbe808595d958346bd55bf8419306bf3de4cd1d0 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.cc +++ b/tensorflow/compiler/tf2xla/xla_op_registry.cc @@ -161,7 +161,14 @@ void XlaOpRegistry::RegisterCompilationKernels() { const string& op_name = op.first; const std::unique_ptr& op_registration = op.second; const OpDef* op_def; - TF_CHECK_OK(op_registry->LookUpOpDef(op_name, &op_def)); + Status lookup_status = op_registry->LookUpOpDef(op_name, &op_def); + if (!lookup_status.ok()) { + LOG(ERROR) << lookup_status.error_message(); + XLA_LOG_LINES( + ERROR, "Ops registered: \n" + + dynamic_cast(op_registry)->DebugString(true)); + } + TF_CHECK_OK(lookup_status); std::unordered_set type_attrs; for (const OpDef::AttrDef& attr_def : op_def->attr()) { @@ -248,6 +255,8 @@ void XlaOpRegistry::RegisterCompilationKernels() { std::vector XlaOpRegistry::DeviceKernels( const string& compilation_device_name, bool include_compilation_only_kernels) { + // Ensure compilation kernels registered. + RegisterCompilationKernels(); std::vector kernels; XlaOpRegistry& registry = Instance(); mutex_lock lock(registry.mutex_); diff --git a/tensorflow/compiler/tf2xla/xla_resource.cc b/tensorflow/compiler/tf2xla/xla_resource.cc index 9abac8bdaa77c99a57b2f8ac66fe6ed06fbcd102..c2075b44b82ba279d1246ec6bfcf305d12c418a6 100644 --- a/tensorflow/compiler/tf2xla/xla_resource.cc +++ b/tensorflow/compiler/tf2xla/xla_resource.cc @@ -25,51 +25,99 @@ limitations under the License. namespace tensorflow { -XlaResource::XlaResource(Kind kind, int arg_num, string name, - DataType initial_type, - const xla::ComputationDataHandle& initial_value) +XlaResource::XlaResource(Kind kind, int arg_num, string name, DataType type, + TensorShape shape, + const xla::ComputationDataHandle& initial_value, + int64 tensor_array_size, + const std::set& tensor_array_gradients) : kind_(kind), arg_num_(arg_num), name_(std::move(name)), - type_(initial_type), + type_(type), + shape_(std::move(shape)), value_(initial_value), - initial_value_(initial_value) { + initial_value_(initial_value), + tensor_array_size_(tensor_array_size) { CHECK(kind_ != kInvalid); + + for (const string& gradient : tensor_array_gradients) { + tensor_array_gradients_[gradient].reset( + new XlaResource(/*kind=*/kTensorArray, /*arg_num=*/-1, + /*name=*/strings::StrCat("TensorArrayGrad: ", name_), + type_, shape_, xla::ComputationDataHandle(), + tensor_array_size_, /*tensor_array_gradients=*/{})); + } } -Status XlaResource::SetValue(DataType type, - const xla::ComputationDataHandle& value) { - if (type_ == DT_INVALID && type == DT_INVALID) { - return errors::InvalidArgument("Attempted to initialized resource ", name_, - " to an invalid type"); +Status XlaResource::SetTypeAndShape(DataType type, const TensorShape& shape) { + if (type == DT_INVALID) { + return errors::InvalidArgument("Attempted to set type of resource '", name_, + "'' to an invalid type"); } - if (type_ != DT_INVALID && type_ != type) { + if (initialized() && type_ != type) { return errors::InvalidArgument("Type of resource ", name_, " cannot be changed after initialization: " "old type was ", DataTypeString(type_), ", new type is ", DataTypeString(type)); } + if (initialized() && shape_ != shape) { + return errors::InvalidArgument("Shape of resource ", name_, + " cannot be changed after initialization: " + "old shape was ", + shape_.DebugString(), ", new shape is ", + shape.DebugString()); + } type_ = type; - value_ = value; + shape_ = shape; return Status::OK(); } -Status XlaResource::GetXlaShape(xla::ComputationBuilder* builder, - xla::Shape* shape) const { - auto shape_or_status = builder->GetShape(value_); - if (!shape_or_status.ok()) { - return shape_or_status.status(); +Status XlaResource::SetValue(const xla::ComputationDataHandle& value) { + if (type_ == DT_INVALID) { + return errors::InvalidArgument( + "Resource '", name_, + "' must be initialized with a valid type before use."); } - *shape = *shape_or_status.ValueOrDie(); + value_ = value; return Status::OK(); } -Status XlaResource::GetShape(xla::ComputationBuilder* builder, - TensorShape* shape) const { - xla::Shape xla_shape; - TF_RETURN_IF_ERROR(GetXlaShape(builder, &xla_shape)); - TF_RETURN_IF_ERROR(XLAShapeToTensorShape(xla_shape, shape)); +Status XlaResource::SetZeroValue(xla::ComputationBuilder* builder) { + if (type_ == DT_INVALID) { + return errors::InvalidArgument( + "Resource '", name_, + "' must be initialized with a valid type before use."); + } + switch (kind_) { + case kVariable: { + value_ = builder->Broadcast(XlaHelpers::Zero(builder, type_), + shape_.dim_sizes()); + break; + } + case kTensorArray: { + TensorShape ta_shape; + ta_shape.AddDim(tensor_array_size_); + ta_shape.AppendShape(shape_); + value_ = builder->Broadcast(XlaHelpers::Zero(builder, type_), + ta_shape.dim_sizes()); + break; + } + case kStack: { + TensorShape ta_shape; + ta_shape.AddDim(tensor_array_size_); + ta_shape.AppendShape(shape_); + value_ = + builder->Tuple({builder->Broadcast(XlaHelpers::Zero(builder, type_), + ta_shape.dim_sizes()), + builder->ConstantR0(0)}); + break; + } + + case kInvalid: + default: + LOG(FATAL) << "Invalid resource type"; + } return Status::OK(); } @@ -82,36 +130,20 @@ Status XlaResource::GetOrCreateTensorArrayGradient( std::unique_ptr& gradient = tensor_array_gradients_[source]; if (!gradient) { TensorShape ta_shape; - TF_RETURN_IF_ERROR(GetShape(builder, &ta_shape)); + ta_shape.AddDim(tensor_array_size_); + ta_shape.AppendShape(shape_); xla::ComputationDataHandle gradient_value = builder->Broadcast( XlaHelpers::Zero(builder, type_), ta_shape.dim_sizes()); gradient.reset( new XlaResource(/*kind=*/kTensorArray, /*arg_num=*/-1, /*name=*/strings::StrCat("TensorArrayGrad: ", name_), - type_, gradient_value)); - gradient->tensor_array_size_ = tensor_array_size_; + type_, shape_, gradient_value, tensor_array_size_, + /*tensor_array_gradients=*/{})); } *gradient_out = gradient.get(); return Status::OK(); } -Status XlaResource::PackedShape(xla::ComputationBuilder* builder, - xla::Shape* packed_shape) const { - if (tensor_array_gradients_.empty()) { - return GetXlaShape(builder, packed_shape); - } - TF_RET_CHECK(kind_ == kTensorArray); - std::vector elem_shapes(1 + tensor_array_gradients_.size()); - int pos = 0; - TF_RETURN_IF_ERROR(GetXlaShape(builder, &elem_shapes[pos++])); - for (const auto& gradient : tensor_array_gradients_) { - TF_RETURN_IF_ERROR( - gradient.second->GetXlaShape(builder, &elem_shapes[pos++])); - } - *packed_shape = xla::ShapeUtil::MakeTupleShape(elem_shapes); - return Status::OK(); -} - Status XlaResource::Pack(xla::ComputationDataHandle* pack, xla::ComputationBuilder* builder) const { if (tensor_array_gradients_.empty()) { @@ -130,27 +162,32 @@ Status XlaResource::Pack(xla::ComputationDataHandle* pack, Status XlaResource::SetFromPack(const std::set& gradient_sources, const xla::ComputationDataHandle& pack, - bool reset_initial_values, xla::ComputationBuilder* builder) { if (gradient_sources.empty()) { + if (!initialized()) { + initial_value_ = pack; + } value_ = pack; } else { TF_RET_CHECK(kind_ == kTensorArray); int pos = 0; - value_ = builder->GetTupleElement(pack, pos++); + auto v = builder->GetTupleElement(pack, pos++); + if (!initialized()) { + initial_value_ = v; + } + value_ = v; + for (const auto& source : gradient_sources) { XlaResource* gradient; TF_RETURN_IF_ERROR( GetOrCreateTensorArrayGradient(source, builder, &gradient)); - gradient->value_ = builder->GetTupleElement(pack, pos++); - if (reset_initial_values) { - gradient->initial_value_ = gradient->value_; + auto v = builder->GetTupleElement(pack, pos++); + if (!gradient->initialized()) { + gradient->initial_value_ = v; } + gradient->value_ = v; } } - if (reset_initial_values) { - initial_value_ = value_; - } return Status::OK(); } diff --git a/tensorflow/compiler/tf2xla/xla_resource.h b/tensorflow/compiler/tf2xla/xla_resource.h index 6b46089e4f5e10c195bb59f78c33305c2fa3f84d..1bb2c7274ecdf0954768fd96def51194e52deee8 100644 --- a/tensorflow/compiler/tf2xla/xla_resource.h +++ b/tensorflow/compiler/tf2xla/xla_resource.h @@ -36,8 +36,11 @@ class XlaResource { kStack, }; - XlaResource(Kind kind, int arg_num, string name, DataType initial_type, - const xla::ComputationDataHandle& initial_value); + XlaResource(Kind kind, int arg_num, string name, DataType type, + TensorShape shape, + const xla::ComputationDataHandle& initial_value, + int64 tensor_array_size, + const std::set& tensor_array_gradients); XlaResource(const XlaResource&) = delete; XlaResource(XlaResource&&) = delete; @@ -60,6 +63,12 @@ class XlaResource { // a resource is first initialized we do not yet know its type, so we keep // track of its type dynamically. DataType type() const { return type_; } + + // Shape of the resource. For an uninitialized resource, this is ignored. + // For a Variable, this is the shape of the value. For a TensorArray or Stack + // this is the shape of each entry in the TensorArray/Stack. + const TensorShape& shape() const { return shape_; } + const xla::ComputationDataHandle& value() const { return value_; } // Value of the resource at computation entry. Used to detect which @@ -68,17 +77,19 @@ class XlaResource { return initial_value_; } + // A variable is initialized if it has a value. bool initialized() const { return value_.handle() > 0; } - // Sets the current type/value of the resource. - Status SetValue(DataType type, const xla::ComputationDataHandle& value); + // Sets the type and shape of the resource. The type and shape of a resource + // must not change once the variable has been initialized. + Status SetTypeAndShape(DataType type, const TensorShape& shape); - // Returns the shape of the resource as an xla::Shape. - Status GetXlaShape(xla::ComputationBuilder* builder, xla::Shape* shape) const; + // Sets the current value of the resource. Returns an error if the type is not + // set to a valid value. + Status SetValue(const xla::ComputationDataHandle& value); - // Returns the shape of the resource as an TensorShape. Fails if the shape is - // not representable as a TensorShape. - Status GetShape(xla::ComputationBuilder* builder, TensorShape* shape) const; + // Sets the current value of the resource to an all-zero value. + Status SetZeroValue(xla::ComputationBuilder* builder); // Looks up the gradient for `source`, or creates it if it does not already // exist. The call target must be an initialized TensorArray resource. A @@ -96,10 +107,6 @@ class XlaResource { Status Pack(xla::ComputationDataHandle* pack, xla::ComputationBuilder* builder) const; - // Returns the shape of the `pack` value computed by `Pack()`. - Status PackedShape(xla::ComputationBuilder* builder, - xla::Shape* packed_shape) const; - // Updates the resource with values from `pack`. If `gradient_sources` is // non-empty, treats `pack` as a tuple that represents a TensorArray and // its gradients, and unpacks and updates the gradient resources. @@ -108,14 +115,14 @@ class XlaResource { // Opposite of Pack(). Status SetFromPack(const std::set& gradient_sources, const xla::ComputationDataHandle& pack, - bool reset_initial_values, xla::ComputationBuilder* builder); - // TensorArray-specific fields + // TensorArray and Stack specific fields // 'tensor_array_size' stores the expected size of the TensorArray or Stack. // We need to store this since sometimes TensorArrays must be initialized // lazily since we do not know the element shape at construction time. + // Used by both TensorArrays and Stacks. int64 tensor_array_size() const { return tensor_array_size_; } void set_tensor_array_size(int64 size) { tensor_array_size_ = size; } @@ -136,6 +143,7 @@ class XlaResource { const string name_; DataType type_; + TensorShape shape_; xla::ComputationDataHandle value_; xla::ComputationDataHandle initial_value_; diff --git a/tensorflow/compiler/xla/BUILD b/tensorflow/compiler/xla/BUILD index cd69c69889b2487ad12abea275e79fee4f5c51e6..34e733bc8d80b364cec1783006eba0a5468b55ea 100644 --- a/tensorflow/compiler/xla/BUILD +++ b/tensorflow/compiler/xla/BUILD @@ -88,7 +88,6 @@ cc_library( visibility = [":friends"], deps = [ "//tensorflow/core:framework_lite", - "//tensorflow/core:lib", "//third_party/eigen3", ], ) @@ -182,6 +181,7 @@ cc_library( deps = [ ":status", ":status_macros", + ":statusor", ":types", ":xla_data_proto", "//tensorflow/core:lib", @@ -260,6 +260,7 @@ tf_cc_test( srcs = ["shape_util_test.cc"], deps = [ ":shape_util", + ":status_macros", ":test", ":test_helpers", ":types", @@ -300,7 +301,9 @@ cc_library( ":array2d", ":array3d", ":array4d", + ":shape_tree", ":shape_util", + ":sparse_index_array", ":status_macros", ":types", ":util", @@ -627,6 +630,28 @@ tf_cc_test( ], ) +cc_library( + name = "sparse_index_array", + srcs = ["sparse_index_array.cc"], + hdrs = ["sparse_index_array.h"], + deps = [ + ":array2d", + ":shape_util", + ":xla_data_proto", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "sparse_index_array_test", + srcs = ["sparse_index_array_test.cc"], + deps = [ + ":sparse_index_array", + ":test", + "//tensorflow/core:test_main", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/xla/array.h b/tensorflow/compiler/xla/array.h index 71aa057cd3a1c273c0e851497a78f94ba37c778e..46ee4e64c9ae7ca111d9d04bedcb74ff02a42386 100644 --- a/tensorflow/compiler/xla/array.h +++ b/tensorflow/compiler/xla/array.h @@ -121,6 +121,23 @@ class Array { CHECK(idx == num_elements()); } + // Creates a 2D array of Eigen::half from the given nested initializer list of + // float values. + template ::value && + std::is_same::value>::type> + Array(std::initializer_list> values) + : Array(ToInt64Vector({values.size(), values.begin()->size()})) { + int64 idx = 0; + for (const auto& it1 : values) { + for (const auto& it2 : it1) { + values_[idx] = static_cast(it2); + ++idx; + } + } + CHECK(idx == num_elements()); + } + // Creates a 3D array from the given nested initializer list. The outer // initializer list is the first dimension, and so on. Array(InitializerList3D values) @@ -138,6 +155,27 @@ class Array { CHECK(idx == num_elements()); } + // Creates a 3D array of Eigen::half from the given nested initializer list of + // float values. + template ::value && + std::is_same::value>::type> + Array(std::initializer_list>> + values) + : Array(ToInt64Vector({values.size(), values.begin()->size(), + values.begin()->begin()->size()})) { + int64 idx = 0; + for (const auto& it1 : values) { + for (const auto& it2 : it1) { + for (const auto& it3 : it2) { + values_[idx] = static_cast(it3); + ++idx; + } + } + } + CHECK(idx == num_elements()); + } + // Creates a 4D array from the given nested initializer list. The outer // initializer list is the first dimension, and so on. Array(InitializerList4D values) @@ -158,6 +196,31 @@ class Array { CHECK(idx == num_elements()); } + // Creates a 4D array of Eigen::half from the given nested initializer list of + // float values. + template ::value && + std::is_same::value>::type> + Array(std::initializer_list< + std::initializer_list>>> + values) + : Array(ToInt64Vector({values.size(), values.begin()->size(), + values.begin()->begin()->size(), + values.begin()->begin()->begin()->size()})) { + int64 idx = 0; + for (const auto& it1 : values) { + for (const auto& it2 : it1) { + for (const auto& it3 : it2) { + for (const auto& it4 : it3) { + values_[idx] = static_cast(it4); + ++idx; + } + } + } + } + CHECK(idx == num_elements()); + } + Array(const Array& other) : sizes_(other.sizes_), values_(new T[num_elements()]) { std::copy(&other.values_[0], &other.values_[0] + num_elements(), @@ -185,7 +248,7 @@ class Array { // Fills the array with the sequence i*multiplier for i=0,1,... void FillWithMultiples(const T& multiplier) { for (int64 i = 0; i < num_elements(); ++i) { - values_[i] = i * multiplier; + values_[i] = static_cast(i) * multiplier; } } diff --git a/tensorflow/compiler/xla/array2d.h b/tensorflow/compiler/xla/array2d.h index bb85fbee9b97fd6b9b0bf7223a9b820989dcbfa7..41f563486d21e42e88dcf6c751ce4a64da5e3213 100644 --- a/tensorflow/compiler/xla/array2d.h +++ b/tensorflow/compiler/xla/array2d.h @@ -52,6 +52,14 @@ class Array2D : public Array { Array2D(std::initializer_list> values) : Array(values) {} + // Creates an array of Eigen::half from the given nested initializer list of + // float values. + template ::value && + std::is_same::value>::type> + Array2D(std::initializer_list> values) + : Array(values) {} + Array2D(const Array2D& other) : Array(other) {} int64 n1() const { return this->dim(0); } diff --git a/tensorflow/compiler/xla/array2d_test.cc b/tensorflow/compiler/xla/array2d_test.cc index c08e42c20ee684dfad8268aa8223440fbfad8a33..93034a719bfbd6724c007059715754677f3f1e62 100644 --- a/tensorflow/compiler/xla/array2d_test.cc +++ b/tensorflow/compiler/xla/array2d_test.cc @@ -63,6 +63,20 @@ TEST(Array2dTest, InitializerListCtor) { EXPECT_EQ(arr(1, 2), 6); } +TEST(Array2dTest, InitializerListCtorHalf) { + Array2D arr = {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}}; + + EXPECT_EQ(arr.n1(), 2); + EXPECT_EQ(arr.n2(), 3); + + EXPECT_EQ(arr(0, 0), static_cast(1)); + EXPECT_EQ(arr(0, 1), static_cast(2)); + EXPECT_EQ(arr(0, 2), static_cast(3)); + EXPECT_EQ(arr(1, 0), static_cast(4)); + EXPECT_EQ(arr(1, 1), static_cast(5)); + EXPECT_EQ(arr(1, 2), static_cast(6)); +} + TEST(Array2dTest, Accessors) { Array2D arr = {{1, 2, 3}, {4, 5, 6}}; diff --git a/tensorflow/compiler/xla/array3d.h b/tensorflow/compiler/xla/array3d.h index a1c5840a5f3874e27043c821ed4684da2fa6c542..e5eb235d45d160d486d1499db665ed14a8509043 100644 --- a/tensorflow/compiler/xla/array3d.h +++ b/tensorflow/compiler/xla/array3d.h @@ -57,6 +57,16 @@ class Array3D : public Array { values) : Array(values) {} + // Creates an array of Eigen::half from the given nested initializer list of + // float values. + template ::value && + std::is_same::value>::type> + Array3D( + std::initializer_list>> + values) + : Array(values) {} + int64 n1() const { return this->dim(0); } int64 n2() const { return this->dim(1); } int64 n3() const { return this->dim(2); } diff --git a/tensorflow/compiler/xla/array3d_test.cc b/tensorflow/compiler/xla/array3d_test.cc index 6b5f4b343b2113652758bbd5ce0fc803239c1266..691ff6c03594a98a12e0fdd2151c4c2a2c9c128a 100644 --- a/tensorflow/compiler/xla/array3d_test.cc +++ b/tensorflow/compiler/xla/array3d_test.cc @@ -69,6 +69,29 @@ TEST(Array3dTest, InitializerListCtor) { EXPECT_EQ(arr(2, 3, 1), 24); } +TEST(Array3dTest, InitializerListCtorHalf) { + Array3D arr = { + {{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}, {7.0f, 8.0f}}, + {{9.0f, 10.0f}, {11.0f, 12.0f}, {13.0f, 14.0f}, {15.0f, 16.0f}}, + {{17.0f, 18.0f}, {19.0f, 20.0f}, {21.0f, 22.0f}, {23.0f, 24.0f}}}; + + EXPECT_EQ(arr.n1(), 3); + EXPECT_EQ(arr.n2(), 4); + EXPECT_EQ(arr.n3(), 2); + EXPECT_EQ(arr.num_elements(), 24); + + EXPECT_EQ(arr(0, 0, 0), static_cast(1)); + EXPECT_EQ(arr(0, 0, 1), static_cast(2)); + EXPECT_EQ(arr(0, 1, 0), static_cast(3)); + EXPECT_EQ(arr(0, 3, 1), static_cast(8)); + EXPECT_EQ(arr(1, 0, 0), static_cast(9)); + EXPECT_EQ(arr(1, 1, 1), static_cast(12)); + EXPECT_EQ(arr(2, 0, 0), static_cast(17)); + EXPECT_EQ(arr(2, 1, 1), static_cast(20)); + EXPECT_EQ(arr(2, 2, 0), static_cast(21)); + EXPECT_EQ(arr(2, 3, 1), static_cast(24)); +} + TEST(Array3dTest, Fill) { Array3D fullof7(2, 3, 4, 7); for (int64 n1 = 0; n1 < fullof7.n1(); ++n1) { diff --git a/tensorflow/compiler/xla/array4d.h b/tensorflow/compiler/xla/array4d.h index f8b2b2afe5fed9c465c2a1f39308b7f44311b16a..cff70e54bad0116bdd08674b626b3bf99dc89e1f 100644 --- a/tensorflow/compiler/xla/array4d.h +++ b/tensorflow/compiler/xla/array4d.h @@ -82,6 +82,16 @@ class Array4D : public Array { values) : Array(values) {} + // Creates an array of Eigen::half from the given nested initializer list of + // float values. + template ::value && + std::is_same::value>::type> + Array4D(std::initializer_list>>> + values) + : Array(values) {} + // Numerically-named aliases for the various dimensions. This matches the // dimension names used in array3d. int64 n4() const { return this->dim(3); } diff --git a/tensorflow/compiler/xla/array4d_test.cc b/tensorflow/compiler/xla/array4d_test.cc index 3bc8148c911df0aeade364e4ac2e2ee828bacb53..927733ea1eab43feff643c35535cc6d9ea59ba5a 100644 --- a/tensorflow/compiler/xla/array4d_test.cc +++ b/tensorflow/compiler/xla/array4d_test.cc @@ -97,6 +97,36 @@ TEST(Array3dTest, InitializerListCtor) { EXPECT_EQ(arr(2, 3, 1, 0), 24); } +TEST(Array3dTest, InitializerListCtorHalf) { + Array4D arr = { + {{{1.0f}, {2.0f}}, {{3.0f}, {4.0f}}, {{5.0f}, {6.0f}}, {{7.0f}, {8.0f}}}, + {{{9.0f}, {10.0f}}, + {{11.0f}, {12.0f}}, + {{13.0f}, {14.0f}}, + {{15.0f}, {16.0f}}}, + {{{17.0f}, {18.0f}}, + {{19.0f}, {20.0f}}, + {{21.0f}, {22.0f}}, + {{23.0f}, {24.0f}}}}; + + EXPECT_EQ(arr.n1(), 3); + EXPECT_EQ(arr.n2(), 4); + EXPECT_EQ(arr.n3(), 2); + EXPECT_EQ(arr.n4(), 1); + EXPECT_EQ(arr.num_elements(), 24); + + EXPECT_EQ(arr(0, 0, 0, 0), static_cast(1)); + EXPECT_EQ(arr(0, 0, 1, 0), static_cast(2)); + EXPECT_EQ(arr(0, 1, 0, 0), static_cast(3)); + EXPECT_EQ(arr(0, 3, 1, 0), static_cast(8)); + EXPECT_EQ(arr(1, 0, 0, 0), static_cast(9)); + EXPECT_EQ(arr(1, 1, 1, 0), static_cast(12)); + EXPECT_EQ(arr(2, 0, 0, 0), static_cast(17)); + EXPECT_EQ(arr(2, 1, 1, 0), static_cast(20)); + EXPECT_EQ(arr(2, 2, 0, 0), static_cast(21)); + EXPECT_EQ(arr(2, 3, 1, 0), static_cast(24)); +} + TEST(Array4dTest, Fill) { Array4D fullof7(2, 3, 4, 5, 7); fullof7.Each([](tensorflow::gtl::ArraySlice idx, int* cell) { diff --git a/tensorflow/compiler/xla/array_test.cc b/tensorflow/compiler/xla/array_test.cc index 8b9419477479d952126fd831eb44899e7649ca71..e8356c9832d34135f5ffb1a5c7a9d6db6db3a051 100644 --- a/tensorflow/compiler/xla/array_test.cc +++ b/tensorflow/compiler/xla/array_test.cc @@ -60,6 +60,25 @@ TEST(ArrayTest, InitializerListCtor) { EXPECT_EQ(arr(1, 2), 6); } +TEST(ArrayTest, InitializerListCtorHalf) { + Array d2({{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}}); + EXPECT_EQ(d2.dim(0), 2); + EXPECT_EQ(d2.dim(1), 3); + + Array d3({{{1.0f}, {4.0f}}, {{1.0f}, {4.0f}}, {{1.0f}, {4.0f}}}); + EXPECT_EQ(d3.dim(0), 3); + EXPECT_EQ(d3.dim(1), 2); + EXPECT_EQ(d3.dim(2), 1); + + Array d4( + {{{{1.0f}, {4.0f}}, {{1.0f}, {4.0f}}, {{1.0f}, {4.0f}}}, + {{{1.0f}, {4.0f}}, {{1.0f}, {4.0f}}, {{1.0f}, {4.0f}}}}); + EXPECT_EQ(d4.dim(0), 2); + EXPECT_EQ(d4.dim(1), 3); + EXPECT_EQ(d4.dim(2), 2); + EXPECT_EQ(d4.dim(3), 1); +} + TEST(ArrayTest, IndexingReadWrite) { Array arr({2, 3}); diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD index d6b4ebfc39ae039ff27fe9fb8a3487c870832f3e..02356699a25e47be50eb15872df4c9c302fc289b 100644 --- a/tensorflow/compiler/xla/client/BUILD +++ b/tensorflow/compiler/xla/client/BUILD @@ -80,6 +80,18 @@ cc_library( ], ) +cc_library( + name = "executable_build_options", + srcs = ["executable_build_options.cc"], + hdrs = ["executable_build_options.h"], + deps = [ + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:device_memory_allocator", + "//tensorflow/core:lib", + ], +) + cc_library( name = "local_client", srcs = ["local_client.cc"], @@ -87,6 +99,7 @@ cc_library( deps = [ ":client", ":computation", + ":executable_build_options", "//tensorflow/compiler/xla:executable_run_options", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", @@ -98,6 +111,7 @@ cc_library( "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:local_service", "//tensorflow/compiler/xla/service:shaped_buffer", + "//tensorflow/compiler/xla/service:source_map_util", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "@llvm//:support", diff --git a/tensorflow/compiler/xla/client/client.cc b/tensorflow/compiler/xla/client/client.cc index 66937d64aff18817bbd5310e0c24e19556e9d727..d15ccb0c28522c647617153aaa8e738d029dfaba 100644 --- a/tensorflow/compiler/xla/client/client.cc +++ b/tensorflow/compiler/xla/client/client.cc @@ -60,7 +60,7 @@ StatusOr> Client::Transfer( "server provided response without a literal in " "TransferToClient request"); } - return MakeUnique(response.literal()); + return Literal::CreateFromProto(*response.mutable_literal()); } StatusOr> Client::TransferToServer( @@ -142,7 +142,7 @@ StatusOr> Client::TransferFromOutfeed( "TransferToClient request"); } - return MakeUnique(response.literal()); + return Literal::CreateFromProto(response.literal()); } Status Client::ResetDevice() { diff --git a/tensorflow/compiler/xla/client/computation.cc b/tensorflow/compiler/xla/client/computation.cc index 4baea8df6e3331200ee52f500fb7b961428e56be..e6c57bda0f0c4cb969939883efebcf3a6d6be381 100644 --- a/tensorflow/compiler/xla/client/computation.cc +++ b/tensorflow/compiler/xla/client/computation.cc @@ -64,4 +64,14 @@ void Computation::ResetWithoutFreeing() { parent_ = nullptr; } +StatusOr Computation::GetProgramShape() const { + GetComputationShapeRequest request; + *request.mutable_computation() = handle_; + GetComputationShapeResponse response; + + TF_RETURN_IF_ERROR(parent_->GetComputationShape(&request, &response)); + + return std::move(*response.mutable_program_shape()); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/computation.h b/tensorflow/compiler/xla/client/computation.h index b595172486950bf08b057625d7b2dd97ac9b2278..a53fc9e9cf34704bd08ddb5bf062c1ec1107f5fb 100644 --- a/tensorflow/compiler/xla/client/computation.h +++ b/tensorflow/compiler/xla/client/computation.h @@ -60,6 +60,10 @@ class Computation { // Returns true if this object is a null Computation. bool IsNull() const { return parent_ == nullptr; } + // Returns the "program shape" (parameter and return shapes) for this + // computation. + StatusOr GetProgramShape() const; + private: void ResetWithoutFreeing(); diff --git a/tensorflow/compiler/xla/client/computation_builder.cc b/tensorflow/compiler/xla/client/computation_builder.cc index 317dcb4e41723b93e7e50d911f16e48bc3505a09..b1dcad6a49a270935b07e26de2d3945b912359d1 100644 --- a/tensorflow/compiler/xla/client/computation_builder.cc +++ b/tensorflow/compiler/xla/client/computation_builder.cc @@ -34,25 +34,9 @@ limitations under the License. namespace xla { -ComputationDataHandle ComputationBuilder::ParseOpResponse( - const Status& status, OpResponse* response) { - VLOG(2) << "done with op request"; - - if (!status.ok()) { - NoteError(status); - return ComputationDataHandle(); - } - - if (response->output().handle() == 0) { - NoteError(InternalError("No output handle")); - return ComputationDataHandle(); - } - return response->output(); -} - ComputationBuilder::ComputationBuilder(Client* client, const string& computation_name) - : name_(computation_name), first_error_(Status::OK()), client_(client) {} + : name_(computation_name), client_(client) {} ComputationBuilder::~ComputationBuilder() {} @@ -76,9 +60,8 @@ std::unique_ptr ComputationBuilder::CreateSubBuilder( } Status ComputationBuilder::PrepareComputation() { - if (!first_error_.ok()) { - return first_error_; - } + TF_RETURN_IF_ERROR(first_error_); + if (!computation_.IsNull()) { return Status::OK(); } @@ -100,6 +83,49 @@ Status ComputationBuilder::PrepareComputation() { return Status::OK(); } +Status ComputationBuilder::RunOp(OpRequest* op_request, + OpResponse* op_response) { + TF_RETURN_IF_ERROR(first_error_); + TF_RETURN_IF_ERROR(PrepareComputation()); + + // Fill in fields that are set on every OpRequest. + *op_request->mutable_computation() = computation_.handle(); + *op_request->mutable_metadata() = metadata_; + if (sharding_) { + *op_request->mutable_sharding() = *sharding_; + } + + const string& op_name = + OpRequest::descriptor()->FindFieldByNumber(op_request->op_case())->name(); + VLOG(2) << "running op request: " << op_name; + Status status = client_->stub()->Op(op_request, op_response); + VLOG(2) << "done with op request: " << op_name; + return status; +} + +void ComputationBuilder::RunOpAndNoteError(OpRequest* op_request) { + OpResponse op_response; + Status status = RunOp(op_request, &op_response); + if (!status.ok()) { + NoteError(status); + } +} + +ComputationDataHandle ComputationBuilder::RunOpAndParseResponse( + OpRequest* op_request) { + OpResponse op_response; + Status status = RunOp(op_request, &op_response); + if (!status.ok()) { + NoteError(status); + return ComputationDataHandle(); + } + if (op_response.output().handle() == 0) { + NoteError(InternalError("No output handle")); + return ComputationDataHandle(); + } + return op_response.output(); +} + bool ComputationBuilder::MakeWindow( tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, @@ -158,81 +184,75 @@ bool ComputationBuilder::MakeWindow( return true; } -ComputationDataHandle ComputationBuilder::ConstantOp( - const PopulateLiteral& populate) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ConstantRequest request; - Literal literal; - populate(&literal); - *request.mutable_literal() = literal.ToProto(); - VLOG(3) << "created constant: " << request.literal().ShortDebugString(); - OpRequest op_request; - *op_request.mutable_constant_request() = request; - *op_request.mutable_computation() = computation_.handle(); - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making constant request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); -} - ComputationDataHandle ComputationBuilder::ConstantLiteral( const Literal& literal) { - return ConstantOp( - [literal](Literal* mutable_literal) { *mutable_literal = literal; }); + OpRequest op_request; + ConstantRequest* request = op_request.mutable_constant_request(); + *request->mutable_literal() = literal.ToProto(); + VLOG(3) << "created constant: " << request->literal().ShortDebugString(); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Parameter(int64 parameter_number, const Shape& shape, const string& name) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ParameterRequest request; - *request.mutable_shape() = shape; - request.set_parameter(parameter_number); - request.set_name(name); OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_parameter_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making parameter request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + ParameterRequest* request = op_request.mutable_parameter_request(); + *request->mutable_shape() = shape; + request->set_parameter(parameter_number); + request->set_name(name); + return RunOpAndParseResponse(&op_request); } -StatusOr> ComputationBuilder::GetShape( +StatusOr> ComputationBuilder::GetShapeWithoutNoteError( const ComputationDataHandle& operand) { - if (!first_error_.ok()) { - return first_error_; - } - GetLocalShapeRequest request; *request.mutable_computation() = computation_.handle(); *request.mutable_operand() = operand; GetLocalShapeResponse response; VLOG(2) << "making get-shape request"; - Status s = client_->stub()->GetLocalShape(&request, &response); + TF_RETURN_IF_ERROR(client_->stub()->GetLocalShape(&request, &response)); VLOG(2) << "done with request"; - if (!s.ok()) { - NoteError(s); - return first_error_; - } TF_RET_CHECK(response.has_shape()); std::unique_ptr shape = WrapUnique(response.release_shape()); TF_RET_CHECK(shape != nullptr); return std::move(shape); } +StatusOr> ComputationBuilder::GetShape( + const ComputationDataHandle& operand) { + TF_RETURN_IF_ERROR(first_error_); + + auto status_or_shape = GetShapeWithoutNoteError(operand); + if (!status_or_shape.ok()) { + NoteError(status_or_shape.status()); + return first_error_; + } + return status_or_shape; +} + +StatusOr ComputationBuilder::GetProgramShape() { + TF_RETURN_IF_ERROR(first_error_); + + GetComputationShapeRequest request; + *request.mutable_computation() = computation_.handle(); + GetComputationShapeResponse response; + + VLOG(2) << "making get-program-shape-request"; + Status status = client_->stub()->GetComputationShape(&request, &response); + VLOG(2) << "done with get-program-shape-request"; + + if (!status.ok()) { + first_error_ = status; + return status; + } + + TF_RET_CHECK(response.has_program_shape()); + return std::move(*response.mutable_program_shape()); +} + ComputationDataHandle ComputationBuilder::CheckShape( const ComputationDataHandle& operand, const Shape& expected_shape) { std::unique_ptr actual_shape = GetShape(operand).ConsumeValueOrDie(); @@ -258,30 +278,19 @@ ComputationDataHandle ComputationBuilder::Slice( tensorflow::gtl::ArraySlice start_indices, tensorflow::gtl::ArraySlice limit_indices, tensorflow::gtl::ArraySlice strides) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - SliceRequest request; - *request.mutable_operand() = operand; + OpRequest op_request; + SliceRequest* request = op_request.mutable_slice_request(); + *request->mutable_operand() = operand; for (int64 index : start_indices) { - request.add_start_indices(index); + request->add_start_indices(index); } for (int64 index : limit_indices) { - request.add_limit_indices(index); + request->add_limit_indices(index); } for (int64 index : strides) { - request.add_strides(index); + request->add_strides(index); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_slice_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making slice request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::SliceInDim( @@ -307,143 +316,78 @@ ComputationDataHandle ComputationBuilder::DynamicSlice( const ComputationDataHandle& operand, const ComputationDataHandle& start_indices, tensorflow::gtl::ArraySlice slice_sizes) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - DynamicSliceRequest request; - *request.mutable_operand() = operand; - *request.mutable_start_indices() = start_indices; + OpRequest op_request; + DynamicSliceRequest* request = op_request.mutable_dynamic_slice_request(); + *request->mutable_operand() = operand; + *request->mutable_start_indices() = start_indices; for (int64 index : slice_sizes) { - request.add_slice_sizes(index); + request->add_slice_sizes(index); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_dynamic_slice_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making dynamic slice request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::DynamicUpdateSlice( const ComputationDataHandle& operand, const ComputationDataHandle& update, const ComputationDataHandle& start_indices) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - DynamicUpdateSliceRequest request; - *request.mutable_operand() = operand; - *request.mutable_update() = update; - *request.mutable_start_indices() = start_indices; OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_dynamic_update_slice_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making dynamic update slice request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + DynamicUpdateSliceRequest* request = + op_request.mutable_dynamic_update_slice_request(); + *request->mutable_operand() = operand; + *request->mutable_update() = update; + *request->mutable_start_indices() = start_indices; + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::ConcatInDim( tensorflow::gtl::ArraySlice operands, int64 dimension) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ConcatenateRequest request; + OpRequest op_request; + ConcatenateRequest* request = op_request.mutable_concatenate_request(); for (const ComputationDataHandle& operand : operands) { - *request.add_operands() = operand; + *request->add_operands() = operand; } - request.set_dimension(dimension); - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_concatenate_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making concatenate request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + request->set_dimension(dimension); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Broadcast( const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice broadcast_sizes) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - BroadcastRequest request; - *request.mutable_operand() = operand; + OpRequest op_request; + BroadcastRequest* request = op_request.mutable_broadcast_request(); + *request->mutable_operand() = operand; for (int64 size : broadcast_sizes) { - request.add_broadcast_sizes(size); + request->add_broadcast_sizes(size); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_broadcast_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making broadcast request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Pad( const ComputationDataHandle& operand, const ComputationDataHandle& padding_value, const PaddingConfig& padding_config) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - PadRequest request; - *request.mutable_operand() = operand; - *request.mutable_padding_value() = padding_value; - *request.mutable_padding_config() = padding_config; OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_pad_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making pad request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + PadRequest* request = op_request.mutable_pad_request(); + *request->mutable_operand() = operand; + *request->mutable_padding_value() = padding_value; + *request->mutable_padding_config() = padding_config; + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Reshape( const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice dimensions, tensorflow::gtl::ArraySlice new_sizes) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ReshapeRequest request; - *request.mutable_operand() = operand; + OpRequest op_request; + ReshapeRequest* request = op_request.mutable_reshape_request(); + *request->mutable_operand() = operand; for (int64 dimension : dimensions) { - request.add_dimensions(dimension); + request->add_dimensions(dimension); } for (int64 new_size : new_sizes) { - request.add_new_sizes(new_size); + request->add_new_sizes(new_size); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_reshape_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making reshape request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Reshape( @@ -455,7 +399,6 @@ ComputationDataHandle ComputationBuilder::Reshape( StatusOr> shape = GetShape(operand); if (!shape.ok()) { - first_error_ = shape.status(); return ComputationDataHandle(); } std::vector dimensions(shape.ValueOrDie()->dimensions().size()); @@ -485,7 +428,6 @@ ComputationDataHandle ComputationBuilder::Collapse( // dimensions by the product of their sizes. StatusOr> shape_or_status = GetShape(operand); if (!shape_or_status.ok()) { - first_error_ = shape_or_status.status(); return ComputationDataHandle(); } std::unique_ptr original_shape = shape_or_status.ConsumeValueOrDie(); @@ -517,26 +459,11 @@ ComputationDataHandle ComputationBuilder::Collapse( void ComputationBuilder::Trace(const string& tag, const ComputationDataHandle& operand) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return; - } - - TraceRequest request; - request.set_tag(tag); - *request.mutable_operand() = operand; OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_trace_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making trace request"; - Status s = client_->stub()->Op(&op_request, &response); - VLOG(2) << "done with request"; - - if (!s.ok()) { - NoteError(s); - } + TraceRequest* request = op_request.mutable_trace_request(); + request->set_tag(tag); + *request->mutable_operand() = operand; + RunOpAndNoteError(&op_request); } ComputationDataHandle ComputationBuilder::Select( @@ -547,44 +474,23 @@ ComputationDataHandle ComputationBuilder::Select( ComputationDataHandle ComputationBuilder::Tuple( tensorflow::gtl::ArraySlice elements) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - VariadicOpRequest request; - request.set_varop(VAROP_TUPLE); + OpRequest op_request; + VariadicOpRequest* request = op_request.mutable_variadic_op_request(); + request->set_varop(VAROP_TUPLE); for (const ComputationDataHandle& operand : elements) { - *request.add_operands() = operand; + *request->add_operands() = operand; } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_variadic_op_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making variadic op request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::GetTupleElement( const ComputationDataHandle& tuple_data, int64 index) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - GetTupleElementRequest request; - *request.mutable_operand() = tuple_data; - request.set_index(index); OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_get_tuple_element_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making get tuple element op request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + GetTupleElementRequest* request = + op_request.mutable_get_tuple_element_request(); + *request->mutable_operand() = tuple_data; + request->set_index(index); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Eq( @@ -627,7 +533,6 @@ ComputationDataHandle ComputationBuilder::Dot( const ComputationDataHandle& lhs, const ComputationDataHandle& rhs) { StatusOr> lhs_shape_or_status = GetShape(lhs); if (!lhs_shape_or_status.ok()) { - NoteError(lhs_shape_or_status.status()); return ComputationDataHandle(); } std::unique_ptr lhs_shape = lhs_shape_or_status.ConsumeValueOrDie(); @@ -642,33 +547,17 @@ ComputationDataHandle ComputationBuilder::Dot( ComputationDataHandle ComputationBuilder::DotGeneral( const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, const DotDimensionNumbers& dimension_numbers) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - DotRequest request; - *request.mutable_lhs() = lhs; - *request.mutable_rhs() = rhs; - *request.mutable_dimension_numbers() = dimension_numbers; - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_dot_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making Dot request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + DotRequest* request = op_request.mutable_dot_request(); + *request->mutable_lhs() = lhs; + *request->mutable_rhs() = rhs; + *request->mutable_dimension_numbers() = dimension_numbers; + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Conv( const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, tensorflow::gtl::ArraySlice window_strides, Padding padding) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - return ConvWithGeneralDimensions( lhs, rhs, window_strides, padding, CreateDefaultConvDimensionNumbers(window_strides.size())); @@ -678,10 +567,6 @@ ComputationDataHandle ComputationBuilder::ConvWithGeneralPadding( const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, tensorflow::gtl::ArraySlice window_strides, tensorflow::gtl::ArraySlice> padding) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - return ConvGeneral(lhs, rhs, window_strides, padding, CreateDefaultConvDimensionNumbers(window_strides.size())); } @@ -749,13 +634,11 @@ ComputationDataHandle ComputationBuilder::ConvWithGeneralDimensions( StatusOr> lhs_shape_or_status = GetShape(lhs); if (!lhs_shape_or_status.ok()) { - first_error_ = lhs_shape_or_status.status(); return ComputationDataHandle(); } StatusOr> rhs_shape_or_status = GetShape(rhs); if (!rhs_shape_or_status.ok()) { - first_error_ = rhs_shape_or_status.status(); return ComputationDataHandle(); } @@ -810,13 +693,11 @@ ComputationDataHandle ComputationBuilder::ConvGeneralDilated( StatusOr> lhs_shape_or_status = GetShape(lhs); if (!lhs_shape_or_status.ok()) { - first_error_ = lhs_shape_or_status.status(); return ComputationDataHandle(); } StatusOr> rhs_shape_or_status = GetShape(rhs); if (!rhs_shape_or_status.ok()) { - first_error_ = rhs_shape_or_status.status(); return ComputationDataHandle(); } @@ -834,122 +715,78 @@ ComputationDataHandle ComputationBuilder::ConvGeneralDilated( rhs_shape->dimensions(dimension_numbers.kernel_spatial_dimensions(i)); } - ConvolveRequest request; - *request.mutable_lhs() = lhs; - *request.mutable_rhs() = rhs; - *request.mutable_dimension_numbers() = dimension_numbers; + OpRequest op_request; + ConvolveRequest* request = op_request.mutable_convolve_request(); + *request->mutable_lhs() = lhs; + *request->mutable_rhs() = rhs; + *request->mutable_dimension_numbers() = dimension_numbers; if (!MakeWindow(window_dimensions, window_strides, padding, lhs_dilation, - rhs_dilation, request.mutable_window())) { + rhs_dilation, request->mutable_window())) { // Error is recorded in MakeWindow. return ComputationDataHandle(); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_convolve_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - VLOG(2) << "making convolve request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } -ComputationDataHandle ComputationBuilder::Infeed(const Shape& shape, - const string& config) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); +ComputationDataHandle ComputationBuilder::Fft( + const ComputationDataHandle& operand, const FftType fft_type, + const tensorflow::gtl::ArraySlice fft_length) { + OpRequest op_request; + FftRequest* request = op_request.mutable_fft_request(); + *request->mutable_operand() = operand; + request->set_fft_type(fft_type); + for (int64 dim_len : fft_length) { + request->add_fft_length(dim_len); } + return RunOpAndParseResponse(&op_request); +} - InfeedRequest request; - *request.mutable_shape() = shape; - *request.mutable_config() = config; +ComputationDataHandle ComputationBuilder::Infeed(const Shape& shape, + const string& config) { OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_infeed_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making infeed op request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + InfeedRequest* request = op_request.mutable_infeed_request(); + *request->mutable_shape() = shape; + *request->mutable_config() = config; + return RunOpAndParseResponse(&op_request); } void ComputationBuilder::Outfeed(const ComputationDataHandle& operand, const Shape& shape, const string& outfeed_config) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return; - } - - OutfeedRequest request; - request.set_outfeed_config(outfeed_config); - *request.mutable_operand() = operand; - *request.mutable_shape() = shape; OpRequest op_request; - *op_request.mutable_outfeed_request() = request; - *op_request.mutable_computation() = computation_.handle(); - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making outfeed op request"; - tensorflow::Status s = client_->stub()->Op(&op_request, &response); - - if (!s.ok()) { - NoteError(s); - return; - } + OutfeedRequest* request = op_request.mutable_outfeed_request(); + request->set_outfeed_config(outfeed_config); + *request->mutable_operand() = operand; + *request->mutable_shape() = shape; + RunOpAndNoteError(&op_request); } ComputationDataHandle ComputationBuilder::Call( const Computation& computation, tensorflow::gtl::ArraySlice operands) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - CallRequest request; - *request.mutable_to_apply() = computation.handle(); + OpRequest op_request; + CallRequest* request = op_request.mutable_call_request(); + *request->mutable_to_apply() = computation.handle(); for (const ComputationDataHandle& operand : operands) { - *request.add_operands() = operand; + *request->add_operands() = operand; } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_call_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making call op request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::CustomCall( const string& call_target_name, tensorflow::gtl::ArraySlice operands, const Shape& shape) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - CustomCallRequest request; - request.set_call_target_name(call_target_name); + OpRequest op_request; + CustomCallRequest* request = op_request.mutable_custom_call_request(); + request->set_call_target_name(call_target_name); for (const ComputationDataHandle& operand : operands) { - *request.add_operands() = operand; + *request->add_operands() = operand; } - *request.mutable_shape() = shape; - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_custom_call_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making custom call op request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + *request->mutable_shape() = shape; + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Complex( @@ -1114,47 +951,25 @@ ComputationDataHandle ComputationBuilder::IsFinite( ComputationDataHandle ComputationBuilder::Transpose( const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice permutation) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); TransposeRequest* request = op_request.mutable_transpose_request(); *request->mutable_operand() = operand; for (int64 dimension : permutation) { request->add_dimensions(dimension); } - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making transpose request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Rev( const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice dimensions) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ReverseRequest request; - *request.mutable_operand() = operand; + OpRequest op_request; + ReverseRequest* request = op_request.mutable_reverse_request(); + *request->mutable_operand() = operand; for (int64 dimension : dimensions) { - request.add_dimensions(dimension); + request->add_dimensions(dimension); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_reverse_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making reverse op request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Sort( @@ -1182,24 +997,15 @@ ComputationDataHandle ComputationBuilder::ConvertElementType( StatusOr> shape_status = GetShape(operand); if (!shape_status.ok()) { - first_error_ = shape_status.status(); return ComputationDataHandle(); } std::unique_ptr original = shape_status.ConsumeValueOrDie(); - ConvertRequest request; - *request.mutable_operand() = operand; - request.set_new_element_type(new_element_type); OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_convert_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making convert request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + ConvertRequest* request = op_request.mutable_convert_request(); + *request->mutable_operand() = operand; + request->set_new_element_type(new_element_type); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::BitcastConvertType( @@ -1210,24 +1016,15 @@ ComputationDataHandle ComputationBuilder::BitcastConvertType( StatusOr> shape_status = GetShape(operand); if (!shape_status.ok()) { - first_error_ = shape_status.status(); return ComputationDataHandle(); } std::unique_ptr original = shape_status.ConsumeValueOrDie(); - ConvertRequest request; - *request.mutable_operand() = operand; - request.set_new_element_type(new_element_type); OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_bitcast_convert_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making bitcast convert request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + ConvertRequest* request = op_request.mutable_bitcast_convert_request(); + *request->mutable_operand() = operand; + request->set_new_element_type(new_element_type); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::SquareF32( @@ -1255,107 +1052,57 @@ ComputationDataHandle ComputationBuilder::Clamp( ComputationDataHandle ComputationBuilder::UnaryOp( UnaryOperation unop, const ComputationDataHandle& operand) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - UnaryOpRequest request; - request.set_unop(unop); - *request.mutable_operand() = operand; OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_unary_op_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making unop request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + UnaryOpRequest* request = op_request.mutable_unary_op_request(); + request->set_unop(unop); + *request->mutable_operand() = operand; + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::BinaryOp( BinaryOperation binop, const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, tensorflow::gtl::ArraySlice broadcast_dimensions) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - BinaryOpRequest request; - request.set_binop(binop); - *request.mutable_lhs() = lhs; - *request.mutable_rhs() = rhs; + OpRequest op_request; + BinaryOpRequest* request = op_request.mutable_binary_op_request(); + request->set_binop(binop); + *request->mutable_lhs() = lhs; + *request->mutable_rhs() = rhs; for (int64 dimension : broadcast_dimensions) { - request.add_broadcast_dimensions(dimension); + request->add_broadcast_dimensions(dimension); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_binary_op_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making binop request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::RngOp( RandomDistribution distribution, tensorflow::gtl::ArraySlice parameters, const Shape& shape) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - RngRequest request; - request.set_distribution(distribution); + OpRequest op_request; + RngRequest* request = op_request.mutable_rng_request(); + request->set_distribution(distribution); for (const ComputationDataHandle& param : parameters) { - *request.add_parameter() = param; + *request->add_parameter() = param; } - *request.mutable_shape() = shape; - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_rng_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making rngop request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + *request->mutable_shape() = shape; + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::TernaryOp( TernaryOperation triop, const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, const ComputationDataHandle& ehs) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - TernaryOpRequest request; - request.set_triop(triop); - *request.mutable_lhs() = lhs; - *request.mutable_rhs() = rhs; - *request.mutable_ehs() = ehs; OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_ternary_op_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making triop request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + TernaryOpRequest* request = op_request.mutable_ternary_op_request(); + request->set_triop(triop); + *request->mutable_lhs() = lhs; + *request->mutable_rhs() = rhs; + *request->mutable_ehs() = ehs; + return RunOpAndParseResponse(&op_request); } Status ComputationBuilder::SetReturnValue( const ComputationDataHandle& operand) { - if (!first_error_.ok()) { - return first_error_; - } + TF_RETURN_IF_ERROR(first_error_); SetReturnValueRequest request; *request.mutable_computation() = computation_.handle(); @@ -1377,9 +1124,7 @@ Status ComputationBuilder::SetReturnValue( StatusOr ComputationBuilder::IsConstant( const ComputationDataHandle& operand, int64 num_parameters) { - if (!first_error_.ok()) { - return first_error_; - } + TF_RETURN_IF_ERROR(first_error_); IsConstantRequest request; *request.mutable_computation() = computation_.handle(); @@ -1400,9 +1145,7 @@ StatusOr ComputationBuilder::IsConstant( StatusOr> ComputationBuilder::ComputeConstant( const ComputationDataHandle& operand, const Layout* output_layout, tensorflow::gtl::ArraySlice parameters) { - if (!first_error_.ok()) { - return first_error_; - } + TF_RETURN_IF_ERROR(first_error_); ComputeConstantRequest request; *request.mutable_computation() = computation_.handle(); @@ -1431,7 +1174,7 @@ StatusOr> ComputationBuilder::ComputeConstant( "no computed literal in the provided response in ComputeConstant " "request"); } - return MakeUnique(response.literal()); + return Literal::CreateFromProto(response.literal()); } ComputationDataHandle ComputationBuilder::Map( @@ -1439,30 +1182,19 @@ ComputationDataHandle ComputationBuilder::Map( const Computation& computation, tensorflow::gtl::ArraySlice dimensions, tensorflow::gtl::ArraySlice static_operands) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - MapRequest request; + OpRequest op_request; + MapRequest* request = op_request.mutable_map_request(); for (const ComputationDataHandle& operand : operands) { - *request.add_operands() = operand; + *request->add_operands() = operand; } - *request.mutable_to_apply() = computation.handle(); + *request->mutable_to_apply() = computation.handle(); for (int64 dimension : dimensions) { - request.add_dimensions(dimension); + request->add_dimensions(dimension); } for (const ComputationDataHandle& sop : static_operands) { - *request.add_static_operands() = sop; + *request->add_static_operands() = sop; } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_map_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making Map request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::RngNormal( @@ -1477,31 +1209,15 @@ ComputationDataHandle ComputationBuilder::RngUniform( return RngOp(RandomDistribution::RNG_UNIFORM, {a, b}, shape); } -ComputationDataHandle ComputationBuilder::RngBernoulli( - const ComputationDataHandle& mean, const Shape& shape) { - return RngOp(RandomDistribution::RNG_BERNOULLI, {mean}, shape); -} - ComputationDataHandle ComputationBuilder::While( const Computation& condition, const Computation& body, const ComputationDataHandle& init) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - WhileRequest request; - *request.mutable_condition() = condition.handle(); - *request.mutable_body() = body.handle(); - *request.mutable_init() = init; OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_while_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making while request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + WhileRequest* request = op_request.mutable_while_request(); + *request->mutable_condition() = condition.handle(); + *request->mutable_body() = body.handle(); + *request->mutable_init() = init; + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Conditional( @@ -1510,52 +1226,29 @@ ComputationDataHandle ComputationBuilder::Conditional( const Computation& true_computation, const ComputationDataHandle& false_operand, const Computation& false_computation) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ConditionalRequest request; - *request.mutable_predicate() = predicate; - *request.mutable_true_operand() = true_operand; - *request.mutable_true_computation() = true_computation.handle(); - *request.mutable_false_operand() = false_operand; - *request.mutable_false_computation() = false_computation.handle(); OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_conditional_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making conditional op request"; - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + ConditionalRequest* request = op_request.mutable_conditional_request(); + *request->mutable_predicate() = predicate; + *request->mutable_true_operand() = true_operand; + *request->mutable_true_computation() = true_computation.handle(); + *request->mutable_false_operand() = false_operand; + *request->mutable_false_computation() = false_computation.handle(); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::Reduce( const ComputationDataHandle& operand, const ComputationDataHandle& init_value, const Computation& computation, tensorflow::gtl::ArraySlice dimensions_to_reduce) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ReduceRequest request; - *request.mutable_operand() = operand; - *request.mutable_init_value() = init_value; + OpRequest op_request; + ReduceRequest* request = op_request.mutable_reduce_request(); + *request->mutable_operand() = operand; + *request->mutable_init_value() = init_value; for (int64 dimension : dimensions_to_reduce) { - request.add_dimensions(dimension); + request->add_dimensions(dimension); } - *request.mutable_to_apply() = computation.handle(); - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_reduce_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making reduce request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + *request->mutable_to_apply() = computation.handle(); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::ReduceAll( @@ -1567,7 +1260,6 @@ ComputationDataHandle ComputationBuilder::ReduceAll( StatusOr> shape = GetShape(operand); if (!shape.ok()) { - first_error_ = shape.status(); return ComputationDataHandle(); } @@ -1587,7 +1279,6 @@ ComputationDataHandle ComputationBuilder::ReduceWindow( StatusOr> shape = GetShape(operand); if (!shape.ok()) { - first_error_ = shape.status(); return ComputationDataHandle(); } @@ -1613,84 +1304,50 @@ ComputationDataHandle ComputationBuilder::ReduceWindowWithGeneralPadding( tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, tensorflow::gtl::ArraySlice> padding) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ReduceWindowRequest request; - *request.mutable_operand() = operand; - *request.mutable_to_apply() = computation.handle(); - *request.mutable_init_value() = init_value; + OpRequest op_request; + ReduceWindowRequest* request = op_request.mutable_reduce_window_request(); + *request->mutable_operand() = operand; + *request->mutable_to_apply() = computation.handle(); + *request->mutable_init_value() = init_value; if (!MakeWindow(window_dimensions, window_strides, padding, {}, {}, - request.mutable_window())) { + request->mutable_window())) { NoteError(InternalError("failed to make window")); return ComputationDataHandle(); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_reduce_window_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - VLOG(2) << "making reduce-window request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::BatchNormTraining( const ComputationDataHandle& operand, const ComputationDataHandle& scale, const ComputationDataHandle& offset, float epsilon, int64 feature_index) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - BatchNormTrainingRequest request; - *request.mutable_operand() = operand; - *request.mutable_scale() = scale; - *request.mutable_offset() = offset; - request.set_epsilon(epsilon); - request.set_feature_index(feature_index); - OpRequest op_request; - *op_request.mutable_batch_norm_training_request() = request; - *op_request.mutable_computation() = computation_.handle(); - AddCommonFieldsToOpRequest(&op_request); - - OpResponse response; - - VLOG(2) << "making BatchNormTraining request"; - - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + BatchNormTrainingRequest* request = + op_request.mutable_batch_norm_training_request(); + *request->mutable_operand() = operand; + *request->mutable_scale() = scale; + *request->mutable_offset() = offset; + request->set_epsilon(epsilon); + request->set_feature_index(feature_index); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::BatchNormInference( const ComputationDataHandle& operand, const ComputationDataHandle& scale, const ComputationDataHandle& offset, const ComputationDataHandle& mean, const ComputationDataHandle& variance, float epsilon, int64 feature_index) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - BatchNormInferenceRequest request; - *request.mutable_operand() = operand; - *request.mutable_scale() = scale; - *request.mutable_offset() = offset; - *request.mutable_mean() = mean; - *request.mutable_variance() = variance; - request.set_epsilon(epsilon); - request.set_feature_index(feature_index); - OpRequest op_request; - *op_request.mutable_batch_norm_inference_request() = request; - *op_request.mutable_computation() = computation_.handle(); - AddCommonFieldsToOpRequest(&op_request); - - OpResponse response; - - VLOG(2) << "making BatchNormInference request"; - - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + BatchNormInferenceRequest* request = + op_request.mutable_batch_norm_inference_request(); + *request->mutable_operand() = operand; + *request->mutable_scale() = scale; + *request->mutable_offset() = offset; + *request->mutable_mean() = mean; + *request->mutable_variance() = variance; + request->set_epsilon(epsilon); + request->set_feature_index(feature_index); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::BatchNormGrad( @@ -1698,49 +1355,25 @@ ComputationDataHandle ComputationBuilder::BatchNormGrad( const ComputationDataHandle& mean, const ComputationDataHandle& var, const ComputationDataHandle& grad_output, float epsilon, int64 feature_index) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - BatchNormGradRequest request; - *request.mutable_operand() = operand; - *request.mutable_scale() = scale; - *request.mutable_mean() = mean; - *request.mutable_variance() = var; - *request.mutable_grad_output() = grad_output; - request.set_epsilon(epsilon); - request.set_feature_index(feature_index); - OpRequest op_request; - *op_request.mutable_batch_norm_grad_request() = request; - *op_request.mutable_computation() = computation_.handle(); - AddCommonFieldsToOpRequest(&op_request); - - OpResponse response; - - VLOG(2) << "making BatchNormGrad request"; - - Status s = client_->stub()->Op(&op_request, &response); - - return ParseOpResponse(s, &response); + BatchNormGradRequest* request = op_request.mutable_batch_norm_grad_request(); + *request->mutable_operand() = operand; + *request->mutable_scale() = scale; + *request->mutable_mean() = mean; + *request->mutable_variance() = var; + *request->mutable_grad_output() = grad_output; + request->set_epsilon(epsilon); + request->set_feature_index(feature_index); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::CrossReplicaSum( const ComputationDataHandle& operand) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - CrossReplicaSumRequest request; - *request.mutable_operand() = operand; OpRequest op_request; - *op_request.mutable_cross_replica_sum_request() = request; - *op_request.mutable_computation() = computation_.handle(); - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making cross-replica-sum request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + CrossReplicaSumRequest* request = + op_request.mutable_cross_replica_sum_request(); + *request->mutable_operand() = operand; + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::SelectAndScatter( @@ -1755,7 +1388,6 @@ ComputationDataHandle ComputationBuilder::SelectAndScatter( StatusOr> shape = GetShape(operand); if (!shape.ok()) { - first_error_ = shape.status(); return ComputationDataHandle(); } return SelectAndScatterWithGeneralPadding( @@ -1772,98 +1404,53 @@ ComputationDataHandle ComputationBuilder::SelectAndScatterWithGeneralPadding( tensorflow::gtl::ArraySlice> padding, const ComputationDataHandle& source, const ComputationDataHandle& init_value, const Computation& scatter) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - SelectAndScatterRequest request; - *request.mutable_operand() = operand; - *request.mutable_select() = select.handle(); - *request.mutable_source() = source; - *request.mutable_init_value() = init_value; - *request.mutable_scatter() = scatter.handle(); + OpRequest op_request; + SelectAndScatterRequest* request = + op_request.mutable_select_and_scatter_request(); + *request->mutable_operand() = operand; + *request->mutable_select() = select.handle(); + *request->mutable_source() = source; + *request->mutable_init_value() = init_value; + *request->mutable_scatter() = scatter.handle(); if (!MakeWindow(window_dimensions, window_strides, padding, {}, {}, - request.mutable_window())) { + request->mutable_window())) { NoteError(InternalError("failed to make window")); return ComputationDataHandle(); } - OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_select_and_scatter_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - VLOG(2) << "making select-and-scatter request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + return RunOpAndParseResponse(&op_request); } ComputationDataHandle ComputationBuilder::ReducePrecision( const ComputationDataHandle& operand, const int exponent_bits, const int mantissa_bits) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - ReducePrecisionRequest request; - *request.mutable_operand() = operand; - request.set_exponent_bits(exponent_bits); - request.set_mantissa_bits(mantissa_bits); OpRequest op_request; - *op_request.mutable_computation() = computation_.handle(); - *op_request.mutable_reduce_precision_request() = request; - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making reduce-precision request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + ReducePrecisionRequest* request = + op_request.mutable_reduce_precision_request(); + *request->mutable_operand() = operand; + request->set_exponent_bits(exponent_bits); + request->set_mantissa_bits(mantissa_bits); + return RunOpAndParseResponse(&op_request); } void ComputationBuilder::Send(const ComputationDataHandle& operand, const ChannelHandle& handle) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return; - } - - SendRequest request; - *request.mutable_operand() = operand; - *request.mutable_channel_handle() = handle; OpRequest op_request; - *op_request.mutable_send_request() = request; + SendRequest* request = op_request.mutable_send_request(); + *request->mutable_operand() = operand; + *request->mutable_channel_handle() = handle; *op_request.mutable_computation() = computation_.handle(); - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making send request"; - Status s = client_->stub()->Op(&op_request, &response); - VLOG(2) << "done with op request"; - - if (!s.ok()) { - NoteError(s); - return; - } + RunOpAndNoteError(&op_request); } ComputationDataHandle ComputationBuilder::Recv(const Shape& shape, const ChannelHandle& handle) { - if (!first_error_.ok() || !PrepareComputation().ok()) { - return ComputationDataHandle(); - } - - RecvRequest request; - *request.mutable_shape() = shape; - *request.mutable_channel_handle() = handle; OpRequest op_request; - *op_request.mutable_recv_request() = request; - *op_request.mutable_computation() = computation_.handle(); - AddCommonFieldsToOpRequest(&op_request); - OpResponse response; - - VLOG(2) << "making recv request"; - Status s = client_->stub()->Op(&op_request, &response); - return ParseOpResponse(s, &response); + RecvRequest* request = op_request.mutable_recv_request(); + *request->mutable_shape() = shape; + *request->mutable_channel_handle() = handle; + return RunOpAndParseResponse(&op_request); } Computation ComputationBuilder::BuildAndNoteError() { @@ -1892,13 +1479,6 @@ StatusOr ComputationBuilder::Build() { return {std::move(computation_)}; } -void ComputationBuilder::AddCommonFieldsToOpRequest(OpRequest* request) const { - *request->mutable_metadata() = metadata_; - if (sharding_) { - *request->mutable_sharding() = *sharding_; - } -} - /* static */ ConvolutionDimensionNumbers ComputationBuilder::CreateDefaultConvDimensionNumbers(int num_spatial_dims) { ConvolutionDimensionNumbers dimension_numbers; diff --git a/tensorflow/compiler/xla/client/computation_builder.h b/tensorflow/compiler/xla/client/computation_builder.h index 7293b35c0f31f8dad19777bb310ca6df3a9468d4..7cae91e9e04bba8f28f2348c552a941e4f7a36b4 100644 --- a/tensorflow/compiler/xla/client/computation_builder.h +++ b/tensorflow/compiler/xla/client/computation_builder.h @@ -67,7 +67,7 @@ class ComputationBuilder { // OpMetadata is often applied to a series of XLA HLO instructions. As a // result, OpMetadata is set on the Computation Builder. All subsequent // instructions generated via this Computation Builder will have the same - // OpMetadata attached until a call to ClearOpMetdata. + // OpMetadata attached until a call to ClearOpMetadata. void SetOpMetadata(const OpMetadata& metadata) { metadata_ = metadata; } // Clears the HloMetadata state. @@ -101,6 +101,9 @@ class ComputationBuilder { StatusOr> GetShape( const ComputationDataHandle& operand); + // Retrieves the (inferred) result for the current computation's shape. + StatusOr GetProgramShape(); + // Checks that the operand has the given expected shape. Returns the operand // if yes, fails with a CHECK error if no. ComputationDataHandle CheckShape(const ComputationDataHandle& operand, @@ -410,14 +413,24 @@ class ComputationBuilder { tensorflow::gtl::ArraySlice rhs_dilation, const ConvolutionDimensionNumbers& dimension_numbers); + // Enqueues an FFT instruction onto the computation, of the given type and + // with the given FFT length. + ComputationDataHandle Fft(const ComputationDataHandle& operand, + FftType fft_type, + tensorflow::gtl::ArraySlice fft_length); + // Enqueues an infeed instruction onto the computation, which writes data of // the given shape to the infeed buffer of the device. ComputationDataHandle Infeed(const Shape& shape, const string& config = ""); // Enqueues an outfeed instruction onto the computation. This instruction // generates outgoing data transfers for the given data. - void Outfeed(const ComputationDataHandle& operand, const Shape& shape, - const string& outfeed_config); + // + // shape_with_layout communicates the laid out shape that we want to outfeed + // -- if !ShapeUtil::Compatible(GetShape(operand), shape_with_layout) an error + // will occur. + void Outfeed(const ComputationDataHandle& operand, + const Shape& shape_with_layout, const string& outfeed_config); // Enqueues a call instruction onto the computation. ComputationDataHandle Call( @@ -678,11 +691,6 @@ class ComputationBuilder { const ComputationDataHandle& b, const Shape& shape); - // Enqueues a B(1, p) random number generation instruction onto the - // computation. - ComputationDataHandle RngBernoulli(const ComputationDataHandle& mean, - const Shape& shape); - // Enqueues a while node onto the computation. ComputationDataHandle While(const Computation& condition, const Computation& body, @@ -710,7 +718,7 @@ class ComputationBuilder { ComputationDataHandle Recv(const Shape& shape, const ChannelHandle& handle); // Returns true if 'operand' is a compile-time constant. A compile-time - // constant does not depend on parameters with higher index then + // constant does not depend on parameters with index greater than or equal to // `num_parameters`, or on stateful operators such as `RngNormal` or `Infeed`. // Unlike `ComputeConstant`, `IsConstant` tests whether a computation is a // compile-time constant without evaluating the computation. @@ -829,8 +837,6 @@ class ComputationBuilder { Status first_error() const { return first_error_; } private: - using PopulateLiteral = std::function; - // Limited checking of convolution parameters. Returns false on // error. bool VerifyConvolution(const Shape& lhs_shape, const Shape& rhs_shape, @@ -849,11 +855,6 @@ class ComputationBuilder { tensorflow::gtl::ArraySlice rhs_dilation, Window* window); - // Internal helper method that makes a request for a constant operation -- the - // provided function is used to populate the literal before sending the - // request. - ComputationDataHandle ConstantOp(const PopulateLiteral& populate); - // Internal helper method that does the building for an arbitrary unary op. ComputationDataHandle UnaryOp(UnaryOperation binop, const ComputationDataHandle& operand); @@ -883,19 +884,28 @@ class ComputationBuilder { // This is used before any given operation is enqueued. Status PrepareComputation(); - // Helper function for parsing a method response and either returning the - // output computation data handle (on success) or a vacuous computation data - // handle (on failure). - ComputationDataHandle ParseOpResponse(const Status& status, - OpResponse* response); - // Notes that the error occurred by: // * storing it internally and capturing a backtrace if it's the first error // (this deferred value will be produced on the call to Build()) // * dying if die_immediately_on_error_ is true void NoteError(const Status& error); - void AddCommonFieldsToOpRequest(OpRequest* request) const; + // Helper function that runs the given op_request, filling in op_response. + // Before the op is run, PrepareComputation is called, and common fields in + // the op_request are filled in. + Status RunOp(OpRequest* op_request, OpResponse* op_response); + + // Helper function that calls RunOp and calls NoteError on failures. + void RunOpAndNoteError(OpRequest* op_request); + + // Helper function that calls RunOp and either returns the output computation + // data handle (on success) or a vacuous computation data handle (on failure). + ComputationDataHandle RunOpAndParseResponse(OpRequest* op_request); + + // Helper function that implements GetShape without noting errors. This makes + // it easier to ensure the real GetShape will note errors on every error path. + StatusOr> GetShapeWithoutNoteError( + const ComputationDataHandle& operand); string name_; // Name to use for the built computation. @@ -929,68 +939,66 @@ class ComputationBuilder { template ComputationDataHandle ComputationBuilder::ConstantR0(NativeT value) { - return ConstantOp([value](Literal* literal) { literal->PopulateR0(value); }); + return ConstantLiteral(*Literal::CreateR0(value)); } template ComputationDataHandle ComputationBuilder::ConstantR1( tensorflow::gtl::ArraySlice values) { - return ConstantOp( - [&values](Literal* literal) { literal->PopulateR1(values); }); + return ConstantLiteral(*Literal::CreateR1(values)); } template ComputationDataHandle ComputationBuilder::ConstantR1(int64 length, NativeT value) { - return ConstantOp([length, value](Literal* literal) { - literal->PopulateWithValue(value, {length}); - }); + Literal literal(ShapeUtil::MakeShape( + primitive_util::NativeToPrimitiveType(), {length})); + literal.PopulateWithValue(value); + return ConstantLiteral(literal); } inline ComputationDataHandle ComputationBuilder::ConstantR1( const tensorflow::core::Bitmap& values) { - return ConstantOp( - [&values](Literal* literal) { literal->PopulateR1(values); }); + return ConstantLiteral(*Literal::CreateR1(values)); } template ComputationDataHandle ComputationBuilder::ConstantR2( std::initializer_list> values) { - return ConstantOp( - [&values](Literal* literal) { literal->PopulateR2(values); }); + return ConstantLiteral(*Literal::CreateR2(values)); } template ComputationDataHandle ComputationBuilder::ConstantFromArrayWithLayout( const Array& values, const Layout& layout) { - return ConstantOp([&values, &layout](Literal* literal) { - literal->PopulateFromArrayWithLayout(values, layout); - }); + return ConstantLiteral( + *Literal::CreateFromArrayWithLayout(values, layout)); } template ComputationDataHandle ComputationBuilder::ConstantFromArray( const Array& values) { - return ConstantOp( - [&values](Literal* literal) { literal->PopulateFromArray(values); }); + return ConstantLiteral(*Literal::CreateFromArray(values)); } template ComputationDataHandle ComputationBuilder::ConstantR2FromArray2DWithLayout( const Array2D& values, const Layout& layout) { - return ConstantFromArrayWithLayout(values, layout); + return ConstantLiteral( + *Literal::CreateFromArrayWithLayout(values, layout)); } template ComputationDataHandle ComputationBuilder::ConstantR2FromArray2D( const Array2D& values) { - return ConstantFromArray(values); + return ConstantLiteral(*Literal::CreateR2FromArray2D(values)); } template ComputationDataHandle ComputationBuilder::ConstantR3FromArray3DWithLayout( const Array3D& values, const Layout& layout) { - return ConstantFromArrayWithLayout(values, layout); + return ConstantLiteral( + *Literal::CreateR3FromArray3DWithLayout(values, layout)); } template diff --git a/tensorflow/compiler/xla/client/executable_build_options.cc b/tensorflow/compiler/xla/client/executable_build_options.cc new file mode 100644 index 0000000000000000000000000000000000000000..804e34f5e75ce2d153ac7627b94a543fda88e810 --- /dev/null +++ b/tensorflow/compiler/xla/client/executable_build_options.cc @@ -0,0 +1,79 @@ +/* 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/compiler/xla/client/executable_build_options.h" + +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/lib/strings/stringprintf.h" + +namespace xla { + +ExecutableBuildOptions& ExecutableBuildOptions::set_device_allocator( + DeviceMemoryAllocator* allocator) { + device_allocator_ = allocator; + return *this; +} + +DeviceMemoryAllocator* ExecutableBuildOptions::device_allocator() const { + return device_allocator_; +} + +ExecutableBuildOptions& ExecutableBuildOptions::set_device_ordinal( + int device_ordinal) { + CHECK_GE(device_ordinal, 0); + device_ordinal_ = device_ordinal; + return *this; +} + +int ExecutableBuildOptions::device_ordinal() const { return device_ordinal_; } + +ExecutableBuildOptions& ExecutableBuildOptions::set_result_layout( + const Shape& shape_with_layout) { + result_layout_set_ = true; + result_layout_ = shape_with_layout; + return *this; +} + +const Shape* ExecutableBuildOptions::result_layout() const { + return result_layout_set_ ? &result_layout_ : nullptr; +} + +string ExecutableBuildOptions::ToString() const { + string result_layout = "nullopt"; + if (result_layout_set_) { + result_layout = ShapeUtil::HumanStringWithLayout(result_layout_); + } + string generate_hlo_graph = "nullopt"; + if (generate_hlo_graph_.has_value()) { + generate_hlo_graph = generate_hlo_graph_.value(); + } + return tensorflow::strings::Printf( + "ExecutableBuildOptions{device_ordinal=%d, result_layout=%s, " + "generate_hlo_graph=%s}", + device_ordinal_, result_layout.c_str(), generate_hlo_graph.c_str()); +} + +ExecutableBuildOptions& ExecutableBuildOptions::set_generate_hlo_graph( + string regex) { + generate_hlo_graph_ = std::move(regex); + return *this; +} + +const tensorflow::gtl::optional& +ExecutableBuildOptions::generate_hlo_graph() const { + return generate_hlo_graph_; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/client/executable_build_options.h b/tensorflow/compiler/xla/client/executable_build_options.h new file mode 100644 index 0000000000000000000000000000000000000000..3a52dbac9adb155ad9a7d91a8102707f70fe2fbf --- /dev/null +++ b/tensorflow/compiler/xla/client/executable_build_options.h @@ -0,0 +1,74 @@ +/* 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_COMPILER_XLA_CLIENT_EXECUTABLE_BUILD_OPTIONS_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_EXECUTABLE_BUILD_OPTIONS_H_ + +#include "tensorflow/compiler/xla/service/device_memory_allocator.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/optional.h" + +namespace xla { + +// Class containing options for building an LocalExecutable with +// LocalClient::Compile. +class ExecutableBuildOptions { + public: + // If set, this is the device to build the computation for. Valid + // device_ordinal values are: 0 to # of devices - 1. These values are + // identical to the device ordinal values used by StreamExecutor. The built + // executable will be executable on any device equivalent to the specified + // device as determined by Backend::devices_equivalent(). A value of -1 + // indicates this option has not been set. + ExecutableBuildOptions& set_device_ordinal(int device_ordinal); + int device_ordinal() const; + + // If set, this specifies the layout of the result of the computation. If not + // set, the service will chose the layout of the result. A Shape is used to + // store the layout to accommodate tuple result shapes. A value of nullptr + // indicates the option has not been set. + ExecutableBuildOptions& set_result_layout(const Shape& shape_with_layout); + const Shape* result_layout() const; + + // If set, this specifies an allocator that can be used to allocate temporary + // space on the device during compilation. For example, the compiler might + // want to run various algorithms on the device and pick the fastest one -- it + // might allocate buffers for use by these algorithms using this allocator. + // + // This does not need to be the same as the DeviceMemoryAllocator passed when + // running the executable. + ExecutableBuildOptions& set_device_allocator( + DeviceMemoryAllocator* allocator); + DeviceMemoryAllocator* device_allocator() const; + + // If set, specifies a regexp of HLO graphs to dump (as in DebugOptions). + ExecutableBuildOptions& set_generate_hlo_graph(string regex); + const tensorflow::gtl::optional& generate_hlo_graph() const; + + // Returns a string representation of the build options, suitable for + // debugging. + string ToString() const; + + private: + int device_ordinal_ = -1; + Shape result_layout_; + bool result_layout_set_ = false; + tensorflow::gtl::optional generate_hlo_graph_; + DeviceMemoryAllocator* device_allocator_ = nullptr; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_EXECUTABLE_BUILD_OPTIONS_H_ diff --git a/tensorflow/compiler/xla/client/lib/testing.cc b/tensorflow/compiler/xla/client/lib/testing.cc index 5f2b55713e342aa3d0251386d57cb52481fe748d..b63a1465ea755b906853860d47768ecbeaa0dcdd 100644 --- a/tensorflow/compiler/xla/client/lib/testing.cc +++ b/tensorflow/compiler/xla/client/lib/testing.cc @@ -31,14 +31,43 @@ limitations under the License. namespace xla { namespace { +// Calculates the number of bytes required to store the data within the +// specified shape. In case of a (nested) tuple shape this is the total byte +// size of all sub-shapes within the tuple. +int64 DataSizeOfShape(const Shape& shape) { + if (ShapeUtil::IsArray(shape)) { + return ShapeUtil::ByteSizeOf(shape); + } + + int64 total_size = 0; + for (const Shape& s : shape.tuple_shapes()) { + total_size += DataSizeOfShape(s); + } + return total_size; +} + +// Create a ComputationDataHandle for an op what generates fake data with the +// given shape. +ComputationDataHandle BuildFakeDataOpOnDevice(const Shape& shape, + ComputationBuilder* builder) { + if (ShapeUtil::IsArray(shape)) { + return builder->Broadcast( + builder->ConstantLiteral(Literal::One(shape.element_type())), + AsInt64Slice(shape.dimensions())); + } + std::vector parts; + for (const Shape& s : shape.tuple_shapes()) { + parts.push_back(BuildFakeDataOpOnDevice(s, builder)); + } + return builder->Tuple(parts); +} + std::unique_ptr MakeFakeDataViaDeviceOrDie(const Shape& shape, Client* client) { ComputationBuilder b( client, tensorflow::strings::StrCat("make_fake_", ShapeUtil::HumanString(shape))); - // TODO(b/26811613): Replace this when RNG is supported on all backends. - b.Broadcast(b.ConstantLiteral(Literal::One(shape.element_type())), - AsInt64Slice(shape.dimensions())); + BuildFakeDataOpOnDevice(shape, &b); Computation computation = b.Build().ConsumeValueOrDie(); auto execution_options = CreateDefaultExecutionOptions(); @@ -51,7 +80,7 @@ std::unique_ptr MakeFakeDataViaDeviceOrDie(const Shape& shape, std::unique_ptr MakeFakeDataOrDie(const Shape& shape, Client* client) { - if (ShapeUtil::ByteSizeOf(shape) < (1LL << 20)) { + if (DataSizeOfShape(shape) < (1LL << 20)) { StatusOr> literal_status = MakeFakeLiteral(shape); if (!literal_status.ok()) { // If we got an Unimplemented error, fall back to making the fake data via diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc index f37368459064ad0e1e997024eec881598b6d59fd..91396f055fe4a3ecbd436139be9470e2a35e1c63 100644 --- a/tensorflow/compiler/xla/client/local_client.cc +++ b/tensorflow/compiler/xla/client/local_client.cc @@ -21,30 +21,14 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/service/service_executable_run_options.h" +#include "tensorflow/compiler/xla/service/source_map_util.h" #include "tensorflow/compiler/xla/status_macros.h" namespace se = ::perftools::gputools; -namespace xla { - -ExecutableBuildOptions& ExecutableBuildOptions::set_device_ordinal( - int device_ordinal) { - device_ordinal_ = device_ordinal; - return *this; -} +using xla::source_map_util::InvalidParameterArgument; -int ExecutableBuildOptions::device_ordinal() const { return device_ordinal_; } - -ExecutableBuildOptions& ExecutableBuildOptions::set_result_layout( - const Shape& shape_with_layout) { - result_layout_set_ = true; - result_layout_ = shape_with_layout; - return *this; -} - -const Shape* ExecutableBuildOptions::result_layout() const { - return result_layout_set_ ? &result_layout_ : nullptr; -} +namespace xla { namespace { StatusOr BorrowStreamForDevice(int device_ordinal, @@ -57,16 +41,18 @@ StatusOr BorrowStreamForDevice(int device_ordinal, } // namespace LocalExecutable::LocalExecutable(std::unique_ptr executable, - Backend* backend, int device_ordinal, - const ExecutableBuildOptions& build_options) + Backend* backend, + ExecutableBuildOptions build_options) : executable_(std::move(executable)), backend_(backend), - build_device_ordinal_(device_ordinal), - build_options_(build_options) {} + build_options_(std::move(build_options)) { + CHECK_GE(build_options_.device_ordinal(), 0) + << "Must have a valid device ordinal that the executable was built for."; +} tensorflow::Status LocalExecutable::ValidateExecutionOptions( const tensorflow::gtl::ArraySlice arguments, - const ExecutableRunOptions& options, const Backend& backend) { + const ExecutableRunOptions& run_options, const Backend& backend) { const ComputationLayout& computation_layout = executable_->module_config().entry_computation_layout(); @@ -79,9 +65,10 @@ tensorflow::Status LocalExecutable::ValidateExecutionOptions( for (int i = 0; i < arguments.size(); ++i) { if (!computation_layout.parameter_layout(i).MatchesLayoutInShape( arguments[i]->on_host_shape())) { - return InvalidArgument( - "argument does not match shape or layout of computation parameter " - "%d: expected %s, got %s", + return InvalidParameterArgument( + executable_.get(), i, + "Argument does not match shape or layout of computation parameter " + "%d: want %s, got %s", i, ShapeUtil::HumanString(computation_layout.parameter_layout(i).shape()) .c_str(), @@ -89,14 +76,14 @@ tensorflow::Status LocalExecutable::ValidateExecutionOptions( } } - if (options.stream() != nullptr) { - if (!options.stream()->ok()) { + if (run_options.stream() != nullptr) { + if (!run_options.stream()->ok()) { return InvalidArgument("stream is uninitialized or in an error state"); } // Check stream matches service platform. const se::Platform* stream_platform = - options.stream()->parent()->platform(); + run_options.stream()->parent()->platform(); if (stream_platform != backend_->platform()) { return InvalidArgument( "stream is for platform %s, but service targets platform %s", @@ -106,7 +93,7 @@ tensorflow::Status LocalExecutable::ValidateExecutionOptions( // Cannot specify device_ordinal with a stream. The stream determines these // values. - if (options.device_ordinal() != -1) { + if (run_options.device_ordinal() != -1) { return InvalidArgument( "cannot set both device ordinal and stream options in " "ExecutableRunOptions; the stream determines the device ordinal"); @@ -115,34 +102,34 @@ tensorflow::Status LocalExecutable::ValidateExecutionOptions( // Verify that the device the executable was built for is equivalent to the // device it will run on. - int run_device_ordinal = options.device_ordinal() == -1 + int run_device_ordinal = run_options.device_ordinal() == -1 ? backend_->default_device_ordinal() - : options.device_ordinal(); - TF_ASSIGN_OR_RETURN( - bool devices_equivalent, - backend_->devices_equivalent(run_device_ordinal, build_device_ordinal_)); + : run_options.device_ordinal(); + TF_ASSIGN_OR_RETURN(bool devices_equivalent, + backend_->devices_equivalent( + run_device_ordinal, build_options_.device_ordinal())); if (!devices_equivalent) { TF_ASSIGN_OR_RETURN(se::StreamExecutor * run_executor, backend_->stream_executor(run_device_ordinal)); TF_ASSIGN_OR_RETURN(se::StreamExecutor * build_executor, - backend_->stream_executor(build_device_ordinal_)); + backend_->stream_executor(build_device_ordinal())); return InvalidArgument( "executable is built for device %s of type \"%s\"; cannot run it on " "device %s of type \"%s\"", - backend_->device_name(build_device_ordinal_).c_str(), + backend_->device_name(build_device_ordinal()).c_str(), build_executor->GetDeviceDescription().name().c_str(), backend_->device_name(run_device_ordinal).c_str(), run_executor->GetDeviceDescription().name().c_str()); } - if (!options.allocator()) { + if (!run_options.allocator()) { return InvalidArgument("an allocator must be provided to ExecuteLocally"); } - if (options.allocator()->platform() != backend.platform()) { + if (run_options.allocator()->platform() != backend.platform()) { return InvalidArgument( "allocator platform (%s) does not match service platform (%s)", - options.allocator()->platform()->Name().c_str(), + run_options.allocator()->platform()->Name().c_str(), backend.platform()->Name().c_str()); } @@ -151,23 +138,22 @@ tensorflow::Status LocalExecutable::ValidateExecutionOptions( StatusOr> LocalExecutable::Run( const tensorflow::gtl::ArraySlice arguments, - const ExecutableRunOptions& options) { - TF_RETURN_IF_ERROR(ValidateExecutionOptions(arguments, options, *backend_)); - - ExecutableRunOptions actual_options = options; + ExecutableRunOptions run_options) { + TF_RETURN_IF_ERROR( + ValidateExecutionOptions(arguments, run_options, *backend_)); Backend::StreamPtr stream; - if (options.stream() == nullptr) { + if (run_options.stream() == nullptr) { // NB! The lifetime of `stream` needs to match the lifetime of // `actual_options` (otherwise we will end up using a returned stream in // ExecuteOnStreamWrapper), which is why it isn't declared in the inner "if" // scope. TF_ASSIGN_OR_RETURN( - stream, BorrowStreamForDevice(options.device_ordinal(), backend_)); - actual_options.set_stream(stream.get()); + stream, BorrowStreamForDevice(run_options.device_ordinal(), backend_)); + run_options.set_stream(stream.get()); } - if (options.allocator() == nullptr) { - actual_options.set_allocator(backend_->memory_allocator()); + if (run_options.allocator() == nullptr) { + run_options.set_allocator(backend_->memory_allocator()); } // For local client execution on CPU backends: @@ -176,7 +162,7 @@ StatusOr> LocalExecutable::Run( // *) The thread pool used for XLA CPU ops is from // backend_->eigen_intra_op_thread_pool(). ServiceExecutableRunOptions service_options( - actual_options, backend_->StreamBorrower(), + run_options, backend_->StreamBorrower(), backend_->eigen_intra_op_thread_pool()); if (executable_->dumping()) { @@ -185,9 +171,10 @@ StatusOr> LocalExecutable::Run( TF_ASSIGN_OR_RETURN( std::unique_ptr result, executable_->ExecuteOnStreamWrapper( - &service_options, options.execution_profile(), arguments)); - return ScopedShapedBuffer::MakeScoped(result.get(), - actual_options.allocator()); + &service_options, run_options.execution_profile(), arguments)); + + return MakeUnique(std::move(*result), + run_options.allocator()); } StatusOr> LocalExecutable::ExecuteAndDump( @@ -263,16 +250,19 @@ StatusOr> LocalClient::Compile( const Computation& computation, const tensorflow::gtl::ArraySlice argument_layouts, const ExecutableBuildOptions& options) { - int device_ordinal = options.device_ordinal() == -1 - ? default_device_ordinal() - : options.device_ordinal(); - TF_ASSIGN_OR_RETURN(std::unique_ptr executable, - local_service_->CompileExecutable( - computation.handle(), argument_layouts, - options.result_layout(), device_ordinal)); + ExecutableBuildOptions updated_options = options; + if (options.device_ordinal() == -1) { + updated_options.set_device_ordinal(default_device_ordinal()); + VLOG(3) << "Set device ordinal to default value of: " + << updated_options.device_ordinal(); + } + TF_ASSIGN_OR_RETURN( + std::unique_ptr executable, + local_service_->CompileExecutable(computation.handle(), argument_layouts, + updated_options)); return WrapUnique(new LocalExecutable(std::move(executable), local_service_->mutable_backend(), - device_ordinal, options)); + updated_options)); } StatusOr> @@ -318,4 +308,8 @@ StatusOr> LocalClient::TransferFromOutfeedLocal( return std::move(literal); } +StatusOr LocalClient::ReplicaNumberToDeviceOrdinal(int replica_number) { + return local_service_->ReplicaNumberToDeviceOrdinal(replica_number); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h index 3ca0d2ef5513cfb6b0dbfbc63b311f81a318356e..b52a30f5a0b92e0094e6b0de3241c10a5a909cad 100644 --- a/tensorflow/compiler/xla/client/local_client.h +++ b/tensorflow/compiler/xla/client/local_client.h @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client.h" #include "tensorflow/compiler/xla/client/computation.h" +#include "tensorflow/compiler/xla/client/executable_build_options.h" #include "tensorflow/compiler/xla/executable_run_options.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" @@ -33,39 +34,13 @@ limitations under the License. namespace xla { -// Class containing options for building an LocalExecutable with -// LocalClient::Compile. -class ExecutableBuildOptions { - public: - // If set, this is the device to build the computation for. Valid - // device_ordinal values are: 0 to # of devices - 1. These values are - // identical to the device ordinal values used by StreamExecutor. The built - // executable will be executable on any device equivalent to the specified - // device as determined by Backend::devices_equivalent(). A value of -1 - // indicates this option has not been set. - ExecutableBuildOptions& set_device_ordinal(int device_ordinal); - int device_ordinal() const; - - // If set, this specifies the layout of the result of the computation. If not - // set, the service will chose the layout of the result. A Shape is used to - // store the layout to accommodate tuple result shapes. A value of nullptr - // indicates the option has not been set. - ExecutableBuildOptions& set_result_layout(const Shape& shape_with_layout); - const Shape* result_layout() const; - - private: - int device_ordinal_ = -1; - Shape result_layout_; - bool result_layout_set_ = false; -}; - class LocalExecutable { public: // Run the compiled computation with the given arguments and options and // return the result. StatusOr> Run( const tensorflow::gtl::ArraySlice arguments, - const ExecutableRunOptions& options); + ExecutableRunOptions run_options); // Return the layout (contained in a shape) of the result produced by the // computation. @@ -88,8 +63,7 @@ class LocalExecutable { // Constructor invoked by LocalClient. LocalExecutable(std::unique_ptr executable, Backend* backend, - int device_ordinal, - const ExecutableBuildOptions& build_options); + ExecutableBuildOptions build_options); // Validates that the given arguments and options satisfy various constraints // of the computation. @@ -117,19 +91,19 @@ class LocalExecutable { StatusOr> LiteralFromShapedBuffer( const ShapedBuffer& shaped_buffer); + // The ordinal of the device which this executable was compiled for. The + // executable can run on all equivalent devices (as determined by + // Backend::devices_equivalent). + int build_device_ordinal() const { return build_options_.device_ordinal(); } + // Compiled computation. std::unique_ptr executable_; // Execution backend. - Backend* backend_; - - // The ordinal of the device which this executable was compiled for. The - // executable can run on all equivalent devices (as determined by - // Backend::devices_equivalent). - int build_device_ordinal_; + Backend* backend_ = nullptr; // Options used to build the executable. - const ExecutableBuildOptions& build_options_; + const ExecutableBuildOptions build_options_; }; // An XLA Client specialization for use when the client and service run in @@ -176,6 +150,13 @@ class LocalClient : public Client { StatusOr> TransferFromOutfeedLocal( const Shape& shape, int device_ordinal); + // Returns the device ordinal that corresponds to the given replica number. + // + // This returns an error if there is not a one-to-one correspondence of + // replicas to device ordinals, but is useful as a short term mechanism for + // the "easy" case where a single replica is a single device. + StatusOr ReplicaNumberToDeviceOrdinal(int replica_number); + // Returns the platform that the underlying service targets. perftools::gputools::Platform* platform() const; diff --git a/tensorflow/compiler/xla/executable_run_options.cc b/tensorflow/compiler/xla/executable_run_options.cc index 33d5b6f1d4d15d5143a3421c87eab9b7a7d11345..392ad9010ab81923a089c7b00a79ddc281af92bb 100644 --- a/tensorflow/compiler/xla/executable_run_options.cc +++ b/tensorflow/compiler/xla/executable_run_options.cc @@ -83,7 +83,7 @@ ExecutableRunOptions& ExecutableRunOptions::set_device_assignment( return *this; } -DeviceAssignment* ExecutableRunOptions::device_assignment() const { +const DeviceAssignment* ExecutableRunOptions::device_assignment() const { return device_assignment_; } diff --git a/tensorflow/compiler/xla/executable_run_options.h b/tensorflow/compiler/xla/executable_run_options.h index deb3ddb203d263d25bef0499a8a53a6098d0de0c..d4fcbf0493c936ebcd0639a432e56b62ee15672c 100644 --- a/tensorflow/compiler/xla/executable_run_options.h +++ b/tensorflow/compiler/xla/executable_run_options.h @@ -82,7 +82,7 @@ class ExecutableRunOptions { ExecutableRunOptions& set_device_assignment( DeviceAssignment* device_assignment); - DeviceAssignment* device_assignment() const; + const DeviceAssignment* device_assignment() const; private: DeviceMemoryAllocator* allocator_ = nullptr; diff --git a/tensorflow/compiler/xla/execution_options_util.h b/tensorflow/compiler/xla/execution_options_util.h index 562da78e837ea6c4a01f0d1170797340fd421ad8..a8ca27ec8dfdc01267ccc9efa6c39093c43d4e2d 100644 --- a/tensorflow/compiler/xla/execution_options_util.h +++ b/tensorflow/compiler/xla/execution_options_util.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ +#ifndef TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ #include "tensorflow/compiler/xla/xla.pb.h" @@ -26,4 +26,4 @@ ExecutionOptions CreateDefaultExecutionOptions(); } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ +#endif // TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ diff --git a/tensorflow/compiler/xla/index_util.cc b/tensorflow/compiler/xla/index_util.cc index 2ee23927d86612a59470dd3d3a219d00055ec65b..ffd1fb79e986f82e1c2721f0eefbf3b4c0838e41 100644 --- a/tensorflow/compiler/xla/index_util.cc +++ b/tensorflow/compiler/xla/index_util.cc @@ -149,4 +149,33 @@ namespace xla { return stride; } +/* static */ bool IndexUtil::IndexInBounds( + const Shape& shape, tensorflow::gtl::ArraySlice index) { + int64 rank = ShapeUtil::Rank(shape); + if (rank != index.size()) { + return false; + } + for (int64 d = 0; d < rank; ++d) { + if (index[d] >= shape.dimensions(d)) { + return false; + } + } + return true; +} + +/* static */ int IndexUtil::CompareIndices( + tensorflow::gtl::ArraySlice lhs, + tensorflow::gtl::ArraySlice rhs) { + int64 rank = lhs.size(); + CHECK_EQ(rhs.size(), rank); + for (int64 dim = 0; dim < rank; ++dim) { + if (lhs[dim] < rhs[dim]) { + return -1; + } else if (lhs[dim] > rhs[dim]) { + return 1; + } + } + return 0; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/index_util.h b/tensorflow/compiler/xla/index_util.h index c9838966a5b67397eb5fc4afe3ab9d98e82eb2b1..142006f2626e83d3254f2de65fc28fd5d6694e53 100644 --- a/tensorflow/compiler/xla/index_util.h +++ b/tensorflow/compiler/xla/index_util.h @@ -37,7 +37,7 @@ class IndexUtil { static int64 MultidimensionalIndexToLinearIndex( const Shape& shape, tensorflow::gtl::ArraySlice multi_index); - // Coverts a linear index into multidimensional index (eg {x, y, z}) based on + // Converts a linear index into multidimensional index (eg {x, y, z}) based on // the shape and its layout. The first index in the returned multidimensional // index is dimension 0. static std::vector LinearIndexToMultidimensionalIndex( @@ -69,6 +69,18 @@ class IndexUtil { // sizeof(dimension(3)) * sizeof(dimension(2)) == 4 * 10 static int64 GetDimensionStride(const Shape& shape, int64 dimension); + // Returns true iff the given multi-index is contained in the bounds for the + // shape. + static bool IndexInBounds(const Shape& shape, + tensorflow::gtl::ArraySlice index); + + // Compares the given indices in lexicographic order. lhs[0] and rhs[0] are + // compared first, and lhs[rank-1] and rhs[rank-1] last. If lhs is larger, + // then -1 is returned. If rhs is larger, then 1 is returned. Otherwise, 0 is + // returned. + static int CompareIndices(tensorflow::gtl::ArraySlice lhs, + tensorflow::gtl::ArraySlice rhs); + private: TF_DISALLOW_COPY_AND_ASSIGN(IndexUtil); }; diff --git a/tensorflow/compiler/xla/iterator_util.h b/tensorflow/compiler/xla/iterator_util.h index a39999705eddc5728dce028dab64b7358395757e..a8bb8c7a7e6784e555f4e9dad73ecc78c668ac42 100644 --- a/tensorflow/compiler/xla/iterator_util.h +++ b/tensorflow/compiler/xla/iterator_util.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_ITERATOR_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_ITERATOR_UTIL_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_ITERATOR_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_ITERATOR_UTIL_H_ #include #include @@ -95,4 +95,4 @@ UnwrappingIterator MakeUnwrappingIterator(NestedIter iter) { } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_ITERATOR_UTIL_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_ITERATOR_UTIL_H_ diff --git a/tensorflow/compiler/xla/layout_util.cc b/tensorflow/compiler/xla/layout_util.cc index 6435226fbe6a76ae391c241cbf61aedfb3b96573..fdc4bbdd8b162b7115788e267c2a53e73c186123 100644 --- a/tensorflow/compiler/xla/layout_util.cc +++ b/tensorflow/compiler/xla/layout_util.cc @@ -64,6 +64,13 @@ void SetDefaultLayoutToContainer( return layout; } +/* static */ Layout LayoutUtil::MakeSparseLayout(int64 max_sparse_elements) { + Layout layout; + layout.set_format(SPARSE); + layout.set_max_sparse_elements(max_sparse_elements); + return layout; +} + namespace { // Internal helper that creates a default layout for an array of the given rank. @@ -234,7 +241,7 @@ Layout CreateDefaultLayoutForRank(int64 rank) { LayoutUtil::ClearLayout(program_shape->mutable_result()); } -/* static */ bool LayoutUtil::IsDense(const Shape& shape) { +/* static */ bool LayoutUtil::IsDenseArray(const Shape& shape) { return ShapeUtil::IsArray(shape) && shape.has_layout() && IsDense(shape.layout()); } @@ -260,7 +267,7 @@ Layout CreateDefaultLayoutForRank(int64 rank) { shape.layout().padded_dimensions_size() == 0) { return false; } - CHECK(IsDense(shape)); + CHECK(IsDenseArray(shape)); CHECK_EQ(shape.dimensions_size(), shape.layout().padded_dimensions_size()); for (int64 i = 0; i < shape.dimensions_size(); ++i) { if (shape.layout().padded_dimensions(i) > shape.dimensions(i)) { @@ -272,21 +279,35 @@ Layout CreateDefaultLayoutForRank(int64 rank) { /* static */ tensorflow::gtl::ArraySlice LayoutUtil::PaddedDimensions( const Shape& shape) { - CHECK(IsDense(shape)); + CHECK(IsDenseArray(shape)); return AsInt64Slice(shape.layout().padded_dimensions()); } /* static */ int64 LayoutUtil::PaddedDimension(const Shape& shape, int64 index) { - CHECK(IsDense(shape)); + CHECK(IsDenseArray(shape)); return shape.layout().padded_dimensions(index); } /* static */ PaddingValue LayoutUtil::GetPaddingValue(const Shape& shape) { - CHECK(IsDense(shape)); + CHECK(IsDenseArray(shape)); return shape.layout().padding_value(); } +/* static */ bool LayoutUtil::IsSparseArray(const Shape& shape) { + return ShapeUtil::IsArray(shape) && shape.has_layout() && + IsSparse(shape.layout()); +} + +/* static */ bool LayoutUtil::IsSparse(const Layout& layout) { + return layout.format() == SPARSE; +} + +/* static */ int64 LayoutUtil::MaxSparseElements(const Layout& layout) { + CHECK(IsSparse(layout)); + return layout.max_sparse_elements(); +} + /* static */ bool LayoutUtil::HasLayout(const Shape& shape) { if (ShapeUtil::IsTuple(shape)) { // Tuple shape: all subshapes must have a layout. @@ -313,7 +334,7 @@ Layout CreateDefaultLayoutForRank(int64 rank) { /* static */ tensorflow::gtl::ArraySlice LayoutUtil::MinorToMajor( const Shape& shape) { - CHECK(IsDense(shape)); + CHECK(IsDenseArray(shape)); return AsInt64Slice(shape.layout().minor_to_major()); } @@ -350,6 +371,11 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } /* static */ string LayoutUtil::HumanString(const Layout& layout) { + if (IsSparse(layout)) { + return tensorflow::strings::StrCat("sparse{", layout.max_sparse_elements(), + "}"); + } + CHECK(IsDense(layout)); return tensorflow::strings::StrCat( "{", tensorflow::str_util::Join(layout.minor_to_major(), ","), "}"); } @@ -419,6 +445,7 @@ tensorflow::Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src, /* static */ bool LayoutUtil::AreDimensionsConsecutive( const Layout& layout, tensorflow::gtl::ArraySlice dims) { + CHECK(IsDense(layout)); std::vector positions_in_layout; for (int64 dim : dims) { positions_in_layout.push_back( @@ -433,4 +460,9 @@ tensorflow::Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src, return true; } +std::ostream& operator<<(std::ostream& out, const Layout& layout) { + out << LayoutUtil::HumanString(layout); + return out; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/layout_util.h b/tensorflow/compiler/xla/layout_util.h index f73cc957649e7832287ef75da5104e8967e519b9..6c54eb2201b66a4a0c5695bceb14bb2367133935 100644 --- a/tensorflow/compiler/xla/layout_util.h +++ b/tensorflow/compiler/xla/layout_util.h @@ -36,6 +36,10 @@ class LayoutUtil { // convenience function for protobuf construction.) static Layout MakeLayout(tensorflow::gtl::ArraySlice minor_to_major); + // Creates a sparse layout with the given maximum number of elements. (This is + // a convenience function for protobuf construction.) + static Layout MakeSparseLayout(int64 max_sparse_elements); + // Returns default layout for the given shape. static Layout GetDefaultLayoutForShape(const Shape& shape); @@ -72,7 +76,7 @@ class LayoutUtil { static void ClearLayout(ProgramShape* program_shape); // Returns whether the given Shape is an array and has a dense format layout. - static bool IsDense(const Shape& shape); + static bool IsDenseArray(const Shape& shape); // Returns whether the given Layout has a dense format. static bool IsDense(const Layout& layout); @@ -107,6 +111,17 @@ class LayoutUtil { // an array and has a dense layout. static PaddingValue GetPaddingValue(const Shape& shape); + // Returns whether the given Shape is an array (i.e. not a tuple) and has a + // sparse format layout. + static bool IsSparseArray(const Shape& shape); + + // Returns whether the given Layout has a sparse format. + static bool IsSparse(const Layout& layout); + + // Returns the maximum number of elements that can be stored in a sparse + // layout. + static int64 MaxSparseElements(const Layout& layout); + // Returns whether the given shape has a layout. For tuple shapes, true is // returned only if all elements have layouts. static bool HasLayout(const Shape& shape); @@ -122,7 +137,7 @@ class LayoutUtil { static tensorflow::gtl::ArraySlice MinorToMajor(const Shape& shape); static tensorflow::gtl::ArraySlice MinorToMajor(const Layout& layout); - // Major(0) is the most major logical dimension number, major(1) is the + // Major(0) is the most major logical dimension number, Major(1) is the // second-most-major logical dimension number and so on. // // This can be used to translate physical dimension numbers to logical @@ -184,6 +199,8 @@ class LayoutUtil { TF_DISALLOW_COPY_AND_ASSIGN(LayoutUtil); }; +std::ostream& operator<<(std::ostream& out, const Layout& layout); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_LAYOUT_UTIL_H_ diff --git a/tensorflow/compiler/xla/layout_util_test.cc b/tensorflow/compiler/xla/layout_util_test.cc index 331bb9afa94e9e7c97d9c880dbac31c60ac0da18..4fd1d818e3e3b417eee9f6b14bb598bfb9480c6e 100644 --- a/tensorflow/compiler/xla/layout_util_test.cc +++ b/tensorflow/compiler/xla/layout_util_test.cc @@ -14,6 +14,9 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/layout_util.h" + +#include + #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" @@ -30,6 +33,14 @@ class LayoutUtilTest : public ::testing::Test { *shape.mutable_layout() = LayoutUtil::MakeLayout(minor_to_major); return shape; } + + Shape MakeShapeWithSparseLayout(PrimitiveType element_type, + tensorflow::gtl::ArraySlice dimensions, + int64 max_sparse_elements) { + Shape shape = ShapeUtil::MakeShape(element_type, dimensions); + *shape.mutable_layout() = LayoutUtil::MakeSparseLayout(max_sparse_elements); + return shape; + } }; TEST_F(LayoutUtilTest, TupleLayoutComparison) { @@ -81,6 +92,29 @@ TEST_F(LayoutUtilTest, CopyLayoutArray) { EXPECT_FALSE(dst.has_layout()); } +TEST_F(LayoutUtilTest, CopyLayoutSparse) { + Shape src = MakeShapeWithSparseLayout(F32, {2, 3}, 2); + Shape dst = MakeShapeWithLayout(F32, {2, 3}, {1, 0}); + + EXPECT_FALSE(LayoutUtil::LayoutsInShapesEqual(src, dst)); + EXPECT_IS_OK(LayoutUtil::CopyLayoutBetweenShapes(src, &dst)); + EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(src, dst)); + + // Should work if destination has no layout. + dst.clear_layout(); + EXPECT_FALSE(LayoutUtil::LayoutsInShapesEqual(src, dst)); + EXPECT_IS_OK(LayoutUtil::CopyLayoutBetweenShapes(src, &dst)); + EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(src, dst)); + + // If source is cleared, then destination should be cleared. + src.clear_layout(); + EXPECT_FALSE(LayoutUtil::LayoutsInShapesEqual(src, dst)); + EXPECT_TRUE(dst.has_layout()); + EXPECT_IS_OK(LayoutUtil::CopyLayoutBetweenShapes(src, &dst)); + EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(src, dst)); + EXPECT_FALSE(dst.has_layout()); +} + TEST_F(LayoutUtilTest, CopyLayoutTuple) { Shape src = ShapeUtil::MakeTupleShape( {MakeShapeWithLayout(F32, {2, 3}, {0, 1}), @@ -100,6 +134,25 @@ TEST_F(LayoutUtilTest, CopyLayoutTuple) { EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(src, dst)); } +TEST_F(LayoutUtilTest, CopyLayoutTupleSparse) { + Shape src = ShapeUtil::MakeTupleShape( + {MakeShapeWithSparseLayout(F32, {2, 3}, 4), + MakeShapeWithSparseLayout(F32, {42, 123}, 4), + ShapeUtil::MakeTupleShape( + {MakeShapeWithLayout(F32, {}, {}), + MakeShapeWithSparseLayout(F32, {1, 2, 3}, 6)})}); + Shape dst = ShapeUtil::MakeTupleShape( + {MakeShapeWithLayout(F32, {2, 3}, {1, 0}), + MakeShapeWithLayout(F32, {42, 123}, {1, 0}), + ShapeUtil::MakeTupleShape( + {MakeShapeWithLayout(F32, {}, {}), + MakeShapeWithLayout(F32, {1, 2, 3}, {1, 2, 0})})}); + + EXPECT_FALSE(LayoutUtil::LayoutsInShapesEqual(src, dst)); + EXPECT_IS_OK(LayoutUtil::CopyLayoutBetweenShapes(src, &dst)); + EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(src, dst)); +} + TEST_F(LayoutUtilTest, CopyLayoutNotCompatibleSameRank) { Shape src = MakeShapeWithLayout(F32, {123, 42, 7}, {2, 0, 1}); Shape dst = MakeShapeWithLayout(F32, {2, 3, 5}, {1, 0}); @@ -107,6 +160,13 @@ TEST_F(LayoutUtilTest, CopyLayoutNotCompatibleSameRank) { EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(src, dst)); } +TEST_F(LayoutUtilTest, CopyLayoutSparseNotCompatibleSameRank) { + Shape src = MakeShapeWithSparseLayout(F32, {123, 42, 7}, 6); + Shape dst = MakeShapeWithLayout(F32, {2, 3, 5}, {1, 0}); + ASSERT_IS_OK(LayoutUtil::CopyLayoutBetweenShapes(src, &dst)); + EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(src, dst)); +} + TEST_F(LayoutUtilTest, CopyLayoutNotCompatibleDifferentRank) { Shape src = MakeShapeWithLayout(F32, {123, 42, 7}, {2, 0, 1}); Shape dst = MakeShapeWithLayout(F32, {2, 3}, {1, 0}); @@ -116,6 +176,15 @@ TEST_F(LayoutUtilTest, CopyLayoutNotCompatibleDifferentRank) { ::testing::ContainsRegex("cannot copy layout from shape")); } +TEST_F(LayoutUtilTest, CopyLayoutSparseNotCompatibleDifferentRank) { + Shape src = MakeShapeWithLayout(F32, {123, 42, 7}, {2, 0, 1}); + Shape dst = MakeShapeWithSparseLayout(F32, {2, 3}, 4); + auto status = LayoutUtil::CopyLayoutBetweenShapes(src, &dst); + EXPECT_FALSE(status.ok()); + EXPECT_THAT(status.error_message(), + ::testing::ContainsRegex("cannot copy layout from shape")); +} + TEST_F(LayoutUtilTest, CopyLayoutNotCompatibleTuple) { Shape src = ShapeUtil::MakeTupleShape({MakeShapeWithLayout(F32, {2, 3}, {0, 1}), @@ -221,5 +290,16 @@ TEST_F(LayoutUtilTest, DefaultLayoutGettersMajorToMinor) { ShapeUtil::MakeShape(F32, {10, 20, 30, 15, 25})))); } +TEST_F(LayoutUtilTest, SparseLayoutMaxElements) { + EXPECT_EQ(LayoutUtil::MaxSparseElements(LayoutUtil::MakeSparseLayout(101)), + 101); +} + +TEST_F(LayoutUtilTest, StreamOut) { + std::ostringstream oss; + oss << LayoutUtil::MakeLayout({0, 1, 2}); + EXPECT_EQ(oss.str(), "{0,1,2}"); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc index e88bffd0ba2dacb837c568023f5da1338fea40f3..c8ed3e3a2b009ddffdfb79a9a6ced8d5e736bee6 100644 --- a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc +++ b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc @@ -221,8 +221,19 @@ void AllocateFlags() { flag_values->xla_gpu_disable_multi_streaming(), "If true, multi-streaming in the GPU backend is disabled."), tensorflow::Flag( - "xla_dump_hlo_proto_to", flag_values->mutable_xla_dump_hlo_proto_to(), - "Dump compilation artifacts as proto binary into this directory."), + "xla_dump_optimized_hlo_proto_to", + flag_values->mutable_xla_dump_optimized_hlo_proto_to(), + "Dump Hlo after all hlo passes are executed as proto binary into " + "this directory."), + tensorflow::Flag( + "xla_dump_unoptimized_hlo_proto_to", + flag_values->mutable_xla_dump_unoptimized_hlo_proto_to(), + "Dump HLO before any hlo passes are executed as proto binary into " + "this directory."), + tensorflow::Flag("xla_dump_per_pass_hlo_proto_to", + flag_values->mutable_xla_dump_per_pass_hlo_proto_to(), + "Dump HLO after each pass as an HloProto in binary file " + "format into this directory."), tensorflow::Flag( "xla_test_all_output_layouts", bool_setter_for(&DebugOptions::set_xla_test_all_output_layouts), diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h index d0ef8e66ab0bcbf88035ae31fe32eb161e32e998..b53157f59c61cf4e0850e006ad3656f4be63a936 100644 --- a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h +++ b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ +#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ +#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ #include @@ -35,4 +35,4 @@ xla::DebugOptions GetDebugOptionsFromFlags(); } // namespace legacy_flags } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ +#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h b/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h index 0c238e6a5decffb0339f428e4ea676944479cf1b..e9cf435d83d8345e974d83f8e5340dafeba8e3b2 100644 --- a/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h +++ b/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ +#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ +#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ #include #include "tensorflow/compiler/xla/service/hlo_opcode.h" @@ -148,4 +148,4 @@ inline bool parse_xla_reduce_precision_option( } // namespace legacy_flags } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ +#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ diff --git a/tensorflow/compiler/xla/literal_util.cc b/tensorflow/compiler/xla/literal_util.cc index f493460e795deaf66fa57b9fa42918fdd05bdac6..e0a9b148b443e90a0c4f3e19660b6234d49eef84 100644 --- a/tensorflow/compiler/xla/literal_util.cc +++ b/tensorflow/compiler/xla/literal_util.cc @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* 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. @@ -27,14 +27,20 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" + +using tensorflow::strings::Printf; +using tensorflow::strings::StrCat; + +namespace xla { + namespace { -using tensorflow::int64; constexpr bool kLittleEndian = __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__; @@ -46,9 +52,8 @@ void ConvertEndianShort(char* bytes, int64 size) { std::swap(bytes[i], bytes[i + 1]); } } -} // namespace -namespace xla { +} // namespace std::ostream& operator<<(std::ostream& out, const Literal& literal) { out << literal.ToString(); @@ -78,52 +83,134 @@ Literal::StrideConfig::StrideConfig( } } +Literal::Literal(const Shape& shape) + : Literal(shape, /*allocate_arrays=*/true) {} + +Literal::Literal(const Shape& shape, bool allocate_arrays) + : shape_(shape), pieces_(shape), owns_buffers_(true) { + CHECK(LayoutUtil::HasLayout(shape)); + for (auto& pair : pieces_) { + const ShapeIndex& index = pair.first; + Piece& piece = pair.second; + + piece.set_subshape(&ShapeUtil::GetSubshape(shape_, index)); + const Shape& subshape = piece.subshape(); + if (ShapeUtil::IsArray(subshape)) { + if (allocate_arrays) { + piece.set_buffer(new char[piece.size_bytes()]); + if (LayoutUtil::IsSparseArray(subshape)) { + piece.set_sparse_indices(new SparseIndexArray( + LayoutUtil::MaxSparseElements(subshape.layout()), + ShapeUtil::Rank(subshape))); + } + } else { + piece.set_buffer(nullptr); + } + } + } +} + +Literal::~Literal() { DeallocateBuffers(); } + +void Literal::DeallocateBuffers() { + if (owns_buffers_) { + for (auto& pair : pieces_) { + Piece& piece = pair.second; + if (piece.buffer() != nullptr) { + delete[] piece.buffer(); + delete piece.sparse_indices(); + } + } + } +} + +Literal::Literal(Literal&& other) { + shape_ = std::move(other.shape_); + pieces_ = std::move(other.pieces_); + // We need to iterate through the pieces to set the subshape pointer + // properly. It must refer to subshapes within shape_. + for (auto& pair : pieces_) { + const ShapeIndex& index = pair.first; + Piece& piece = pair.second; + piece.set_subshape(&ShapeUtil::GetSubshape(shape_, index)); + } + owns_buffers_ = other.owns_buffers_; + + other.shape_ = ShapeUtil::MakeNil(); + other.pieces_ = ShapeTree(other.shape_); + other.piece({}).set_subshape(&other.shape_); +} + +Literal& Literal::operator=(Literal&& other) { + DeallocateBuffers(); + shape_ = std::move(other.shape_); + pieces_ = std::move(other.pieces_); + // We need to iterate through the pieces to set the subshape pointer + // properly. It must refer to subshapes within shape_. + for (auto& pair : pieces_) { + const ShapeIndex& index = pair.first; + Piece& piece = pair.second; + piece.set_subshape(&ShapeUtil::GetSubshape(shape_, index)); + } + owns_buffers_ = other.owns_buffers_; + + other.shape_ = ShapeUtil::MakeNil(); + other.pieces_ = ShapeTree(other.shape_); + other.piece({}).set_subshape(&other.shape_); + return *this; +} + std::unique_ptr Literal::CreateFromShape(const Shape& shape) { - auto literal = MakeUnique(); - *literal->mutable_shape() = shape; - if (ShapeUtil::IsTuple(shape)) { - int64 num_elements = ShapeUtil::TupleElementCount(shape); - literal->tuple_literals_.resize(num_elements); - for (int i = 0; i < num_elements; ++i) { - std::unique_ptr elem = - CreateFromShape(ShapeUtil::GetTupleElementShape(shape, i)); - literal->tuple_literals_[i] = std::move(*elem); + auto literal = MakeUnique(shape); + for (auto& pair : literal->pieces_) { + Piece& piece = pair.second; + if (ShapeUtil::IsArray(piece.subshape())) { + memset(piece.untyped_data(), 0, piece.size_bytes()); } - } else { - literal->Reserve(ShapeUtil::ElementsIn(literal->shape())); } return literal; } +const SparseIndexArray* Literal::sparse_indices( + const ShapeIndex& shape_index) const { + return piece(shape_index).sparse_indices(); +} + +SparseIndexArray* Literal::sparse_indices(const ShapeIndex& shape_index) { + return piece(shape_index).sparse_indices(); +} + /* static */ std::unique_ptr Literal::CreateFromDimensions( PrimitiveType primitive_type, tensorflow::gtl::ArraySlice dimensions) { return CreateFromShape(ShapeUtil::MakeShape(primitive_type, dimensions)); } -template -Status Literal::CopyRange(const Literal& src_literal, - tensorflow::gtl::ArraySlice src_base, - tensorflow::gtl::ArraySlice dest_base, - tensorflow::gtl::ArraySlice copy_size) { - const Shape& src_shape = src_literal.shape(); - const Shape& dest_shape = shape(); - tensorflow::gtl::ArraySlice src_data = src_literal.GetArraySlice(); - tensorflow::gtl::MutableArraySlice dest_data = GetMutableArraySlice(); - - TF_RET_CHECK(ShapeUtil::Rank(src_shape) == src_base.size()); - TF_RET_CHECK(ShapeUtil::Rank(dest_shape) == dest_base.size()); +template +Status Literal::CopySliceFromInternal( + const Literal& src_literal, tensorflow::gtl::ArraySlice src_base, + tensorflow::gtl::ArraySlice dest_base, + tensorflow::gtl::ArraySlice copy_size) { + TF_RET_CHECK(ShapeUtil::Rank(src_literal.shape()) == src_base.size()); + TF_RET_CHECK(ShapeUtil::Rank(shape()) == dest_base.size()); + + auto linear_index = [](const Shape& shape, + tensorflow::gtl::ArraySlice multi_index) { + return IndexUtil::MultidimensionalIndexToLinearIndex(shape, multi_index); + }; - if (ShapeUtil::Rank(src_shape) == 0 || ShapeUtil::Rank(dest_shape) == 0) { + if (ShapeUtil::Rank(src_literal.shape()) == 0 || + ShapeUtil::Rank(shape()) == 0) { // If any of the two shapes are scalars, we can just call the StridedCopy() // directly, and we know we will be copying only one value. TF_RET_CHECK(copy_size.empty()); - StridedCopy(dest_data, LinearIndex(dest_base), 0, src_data, - src_literal.LinearIndex(src_base), 0, 1); - } else if (!ShapeUtil::HasZeroElements(dest_shape) && - !ShapeUtil::HasZeroElements(src_shape)) { - // Perform copy if neither src literal nor dest literal has dimensions with - // zero element, otherwise it's a no-op. + StridedCopy(data(), linear_index(shape(), dest_base), 0, + src_literal.data(), + linear_index(src_literal.shape(), src_base), 0, 1); + } else if (!ShapeUtil::HasZeroElements(shape()) && + !ShapeUtil::HasZeroElements(src_literal.shape())) { + // Perform copy if neither src nor dest has dimensions with zero element, + // otherwise it's a no-op. TF_RET_CHECK(src_base.size() == dest_base.size()); TF_RET_CHECK(src_base.size() == copy_size.size()); @@ -133,7 +220,8 @@ Status Literal::CopyRange(const Literal& src_literal, // proper stride size at the matching dimension. DimensionVector src_indexes(src_base.size(), 0); DimensionVector dest_indexes(dest_base.size(), 0); - StrideConfig stride_config(src_shape, dest_shape, copy_size); + Literal::StrideConfig stride_config(src_literal.shape(), shape(), + copy_size); auto copy_proc = [&](const std::vector& indexes) { // Map from multi-dimensional index, to source index. @@ -143,89 +231,296 @@ Status Literal::CopyRange(const Literal& src_literal, std::transform(indexes.begin(), indexes.end(), dest_base.begin(), dest_indexes.begin(), std::plus()); - int64 src_index = src_literal.LinearIndex(src_indexes); - int64 dest_index = LinearIndex(dest_indexes); + int64 src_index = linear_index(src_literal.shape(), src_indexes); + int64 dest_index = linear_index(shape(), dest_indexes); - StridedCopy(dest_data, dest_index, stride_config.dest_stride, src_data, - src_index, stride_config.source_stride, - stride_config.minor_loop_size); + // `this->` is needed to workaround MSVC bug: #16882 + StridedCopy(this->data(), dest_index, stride_config.dest_stride, + src_literal.data(), src_index, + stride_config.source_stride, stride_config.minor_loop_size); return true; }; - ShapeUtil::ForEachIndex(src_shape, stride_config.base, + ShapeUtil::ForEachIndex(src_literal.shape(), stride_config.base, stride_config.dimensions, stride_config.step, copy_proc); } return Status::OK(); } -Status Literal::Copy(const Literal& src_literal, - tensorflow::gtl::ArraySlice src_base, - tensorflow::gtl::ArraySlice dest_base, - tensorflow::gtl::ArraySlice copy_size) { +std::vector Literal::DecomposeTuple() { + CHECK(ShapeUtil::IsTuple(shape())); + std::vector elements; + for (int i = 0; i < ShapeUtil::TupleElementCount(shape()); ++i) { + elements.push_back(Literal(ShapeUtil::GetSubshape(shape(), {i}), + /*allocate_arrays=*/false)); + Literal& element = elements.back(); + for (auto& pair : element.pieces_) { + const ShapeIndex& index = pair.first; + Piece& dest_piece = pair.second; + ShapeIndex src_index = {i}; + for (int64 j : index) { + src_index.push_back(j); + } + Piece& src_piece = piece(src_index); + + // Move the respective buffer and sparse indices over to the element + // Literal. + dest_piece.set_buffer(src_piece.buffer()); + src_piece.set_buffer(nullptr); + dest_piece.set_sparse_indices(src_piece.sparse_indices()); + src_piece.set_sparse_indices(nullptr); + } + } + // Set this literal to be nil-shaped. + *this = Literal(); + return elements; +} + +/* static */ Literal Literal::MoveIntoTuple( + tensorflow::gtl::MutableArraySlice elements) { + std::vector element_shapes; + for (const Literal& element : elements) { + element_shapes.push_back(element.shape()); + } + Literal literal(ShapeUtil::MakeTupleShape(element_shapes), + /*allocate_arrays=*/false); + for (int i = 0; i < elements.size(); ++i) { + TF_CHECK_OK( + literal.MoveFrom(std::move(elements[i]), /*dest_shape_index=*/{i})); + } + return literal; +} + +namespace { + +// Copies the elements in 'src' to 'dest'. The shape and layout of the data in +// the array slices are indicated by dest_shape and src_shape respectively. +template +void CopyElementsBetween(tensorflow::gtl::MutableArraySlice dest, + tensorflow::gtl::ArraySlice src, + const Shape& dest_shape, const Shape& src_shape) { + CHECK(ShapeUtil::Compatible(dest_shape, src_shape)); + if (ShapeUtil::HasZeroElements(dest_shape)) { + return; + } + std::vector index(ShapeUtil::Rank(dest_shape)); + do { + dest[IndexUtil::MultidimensionalIndexToLinearIndex(dest_shape, index)] = + src[IndexUtil::MultidimensionalIndexToLinearIndex(src_shape, index)]; + } while (IndexUtil::BumpIndices(dest_shape, &index)); +} + +} // namespace + +Status Literal::Piece::CopyFrom(const Literal::Piece& src) { + if (ShapeUtil::Equal(subshape(), src.subshape())) { + // If the layouts are equal it's faster just to memcpy. + memcpy(buffer(), src.buffer(), src.size_bytes()); + } else { + TF_RET_CHECK(ShapeUtil::Compatible(src.subshape(), subshape())); + std::vector origin(ShapeUtil::Rank(subshape()), 0); + switch (subshape().element_type()) { +#define COPY_ELEMENTS(XLA_T, NATIVE_T) \ + case (XLA_T): \ + CopyElementsBetween(data(), src.data(), \ + subshape(), src.subshape()); \ + break; + COPY_ELEMENTS(U8, uint8); + COPY_ELEMENTS(U16, uint16); + COPY_ELEMENTS(U32, uint32); + COPY_ELEMENTS(U64, uint64); + COPY_ELEMENTS(S8, int8); + COPY_ELEMENTS(S16, int16); + COPY_ELEMENTS(S32, int32); + COPY_ELEMENTS(S64, int64); + COPY_ELEMENTS(F16, half); + COPY_ELEMENTS(BF16, bfloat16); + COPY_ELEMENTS(F32, float); + COPY_ELEMENTS(F64, double); + COPY_ELEMENTS(C64, complex64); + COPY_ELEMENTS(PRED, bool); +#undef COPY_ELEMENTS + default: + return Unimplemented( + "Unhandled primitive type %s", + PrimitiveType_Name(subshape().element_type()).c_str()); + } + } + return Status::OK(); +} + +Status Literal::CopyFrom(const Literal& src_literal, + const ShapeIndex& dest_shape_index, + const ShapeIndex& src_shape_index) { + const Shape& dest_subshape = + ShapeUtil::GetSubshape(shape(), dest_shape_index); + const Shape& src_subshape = + ShapeUtil::GetSubshape(src_literal.shape(), src_shape_index); + if (!ShapeUtil::Compatible(dest_subshape, src_subshape)) { + return InvalidArgument( + "Destination subshape incompatible with source subshape: %s vs %s", + ShapeUtil::HumanString(dest_subshape).c_str(), + ShapeUtil::HumanString(src_subshape).c_str()); + } + + for (auto& pair : pieces_) { + const ShapeIndex& index = pair.first; + Piece& piece = pair.second; + if (!ShapeUtil::IsArray(piece.subshape())) { + continue; + } + + // Determine if this index is in the part of this literal that we want to + // copy over from src_literal. + bool in_subtree_to_copy = true; + for (int i = 0; i < dest_shape_index.size(); ++i) { + if (index[i] != dest_shape_index[i]) { + in_subtree_to_copy = false; + break; + } + } + if (!in_subtree_to_copy) { + continue; + } + + // Construct the index of the corresponding piece in the source literal. + ShapeIndex src_piece_index = src_shape_index; + for (int64 i = dest_shape_index.size(); i < index.size(); ++i) { + src_piece_index.push_back(index[i]); + } + + TF_RETURN_IF_ERROR(piece.CopyFrom(src_literal.piece(src_piece_index))); + } + return Status::OK(); +} + +Status Literal::MoveFrom(Literal&& src_literal, + const ShapeIndex& dest_shape_index) { + const Shape& dest_subshape = + ShapeUtil::GetSubshape(shape(), dest_shape_index); + if (!ShapeUtil::Equal(dest_subshape, src_literal.shape())) { + return InvalidArgument( + "Destination subshape not equal to source shape: %s vs %s", + ShapeUtil::HumanString(dest_subshape).c_str(), + ShapeUtil::HumanString(src_literal.shape()).c_str()); + } + + if (!(owns_buffers_ && src_literal.owns_buffers_)) { + return InvalidArgument( + "Source and destination literals must both own their buffers (ie, not " + "be views)"); + } + + for (auto& pair : src_literal.pieces_) { + const ShapeIndex& src_index = pair.first; + Piece& src_piece = pair.second; + if (!ShapeUtil::IsArray(src_piece.subshape())) { + continue; + } + + ShapeIndex dest_index = dest_shape_index; + for (int64 i : src_index) { + dest_index.push_back(i); + } + Piece& dest_piece = piece(dest_index); + delete[] dest_piece.buffer(); + dest_piece.set_buffer(src_piece.buffer()); + delete dest_piece.sparse_indices(); + dest_piece.set_sparse_indices(src_piece.sparse_indices()); + } + + src_literal.shape_ = ShapeUtil::MakeNil(); + src_literal.pieces_ = ShapeTree(src_literal.shape_); + src_literal.piece({}).set_subshape(&src_literal.shape_); + return Status::OK(); +} + +Status Literal::CopySliceFrom(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_base, + tensorflow::gtl::ArraySlice dest_base, + tensorflow::gtl::ArraySlice copy_size) { + TF_RET_CHECK(ShapeUtil::IsArray(shape())) << ShapeUtil::HumanString(shape()); + TF_RET_CHECK(ShapeUtil::IsArray(src_literal.shape())) + << ShapeUtil::HumanString(src_literal.shape()); TF_RET_CHECK(ShapeUtil::SameElementType(src_literal.shape(), shape())); - switch (src_literal.shape().element_type()) { + + switch (shape().element_type()) { case U8: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case U16: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case U32: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case U64: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case S8: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case S16: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case S32: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case S64: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case F16: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case BF16: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case F32: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case F64: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case C64: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case PRED: - return CopyRange(src_literal, src_base, dest_base, copy_size); + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); default: break; } - return Unimplemented("Unhandled primitive type %d", - src_literal.shape().element_type()); + return Unimplemented("Unhandled primitive type %d", shape().element_type()); } /* static */ Literal Literal::Zero(PrimitiveType primitive_type) { switch (primitive_type) { case U8: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case U32: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case U64: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case S8: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case S32: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case S64: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case F16: - return *Literal::CreateR0(static_cast(0.0f)); + return std::move(*Literal::CreateR0(static_cast(0.0f))); case BF16: - return *Literal::CreateR0(static_cast(0.0f)); + return std::move( + *Literal::CreateR0(static_cast(0.0f))); case F32: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case F64: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case C64: - return *Literal::CreateR0(0); + return std::move(*Literal::CreateR0(0)); case PRED: - return *Literal::CreateR0(false); + return std::move(*Literal::CreateR0(false)); case S16: case U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; @@ -241,29 +536,30 @@ Status Literal::Copy(const Literal& src_literal, /* static */ Literal Literal::One(PrimitiveType primitive_type) { switch (primitive_type) { case U8: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case U32: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case U64: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case S8: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case S32: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case S64: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case F16: - return *Literal::CreateR0(static_cast(1.0f)); + return std::move(*Literal::CreateR0(static_cast(1.0f))); case BF16: - return *Literal::CreateR0(static_cast(1.0f)); + return std::move( + *Literal::CreateR0(static_cast(1.0f))); case F32: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case F64: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case C64: - return *Literal::CreateR0(1); + return std::move(*Literal::CreateR0(1)); case PRED: - return *Literal::CreateR0(true); + return std::move(*Literal::CreateR0(true)); case S16: case U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; @@ -279,35 +575,42 @@ Status Literal::Copy(const Literal& src_literal, /* static */ Literal Literal::MinValue(PrimitiveType primitive_type) { switch (primitive_type) { case U8: - return *Literal::CreateR0(std::numeric_limits::min()); + return std::move( + *Literal::CreateR0(std::numeric_limits::min())); case U32: - return *Literal::CreateR0(std::numeric_limits::min()); + return std::move( + *Literal::CreateR0(std::numeric_limits::min())); case U64: - return *Literal::CreateR0(std::numeric_limits::min()); + return std::move( + *Literal::CreateR0(std::numeric_limits::min())); case S8: - return *Literal::CreateR0(std::numeric_limits::min()); + return std::move( + *Literal::CreateR0(std::numeric_limits::min())); case S32: - return *Literal::CreateR0(std::numeric_limits::min()); + return std::move( + *Literal::CreateR0(std::numeric_limits::min())); case S64: - return *Literal::CreateR0(std::numeric_limits::min()); + return std::move( + *Literal::CreateR0(std::numeric_limits::min())); case F32: - return *Literal::CreateR0(-std::numeric_limits::infinity()); + return std::move( + *Literal::CreateR0(-std::numeric_limits::infinity())); case F64: - return *Literal::CreateR0( - -std::numeric_limits::infinity()); + return std::move( + *Literal::CreateR0(-std::numeric_limits::infinity())); case C64: LOG(FATAL) << "C64 element type has no minimum value"; case PRED: - return *Literal::CreateR0(false); + return std::move(*Literal::CreateR0(false)); case S16: case U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; case F16: - return *Literal::CreateR0( - static_cast(-std::numeric_limits::infinity())); + return std::move(*Literal::CreateR0( + static_cast(-std::numeric_limits::infinity()))); case BF16: - return *Literal::CreateR0( - static_cast(-std::numeric_limits::infinity())); + return std::move(*Literal::CreateR0( + static_cast(-std::numeric_limits::infinity()))); case TUPLE: LOG(FATAL) << "tuple element type has no minimum value"; case OPAQUE: @@ -320,33 +623,40 @@ Status Literal::Copy(const Literal& src_literal, /* static */ Literal Literal::MaxValue(PrimitiveType primitive_type) { switch (primitive_type) { case U8: - return *Literal::CreateR0(std::numeric_limits::max()); + return std::move( + *Literal::CreateR0(std::numeric_limits::max())); case U32: - return *Literal::CreateR0(std::numeric_limits::max()); + return std::move( + *Literal::CreateR0(std::numeric_limits::max())); case U64: - return *Literal::CreateR0(std::numeric_limits::max()); + return std::move( + *Literal::CreateR0(std::numeric_limits::max())); case S8: - return *Literal::CreateR0(std::numeric_limits::max()); + return std::move( + *Literal::CreateR0(std::numeric_limits::max())); case S32: - return *Literal::CreateR0(std::numeric_limits::max()); + return std::move( + *Literal::CreateR0(std::numeric_limits::max())); case S64: - return *Literal::CreateR0(std::numeric_limits::max()); + return std::move( + *Literal::CreateR0(std::numeric_limits::max())); case F32: - return *Literal::CreateR0(std::numeric_limits::infinity()); + return std::move( + *Literal::CreateR0(std::numeric_limits::infinity())); case F64: - return *Literal::CreateR0( - std::numeric_limits::infinity()); + return std::move( + *Literal::CreateR0(std::numeric_limits::infinity())); case PRED: - return *Literal::CreateR0(true); + return std::move(*Literal::CreateR0(true)); case S16: case U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; case F16: - return *Literal::CreateR0( - static_cast(std::numeric_limits::infinity())); + return std::move(*Literal::CreateR0( + static_cast(std::numeric_limits::infinity()))); case BF16: - return *Literal::CreateR0( - static_cast(std::numeric_limits::infinity())); + return std::move(*Literal::CreateR0( + static_cast(std::numeric_limits::infinity()))); case TUPLE: LOG(FATAL) << "tuple element type has no maximum value"; case OPAQUE: @@ -358,17 +668,29 @@ Status Literal::Copy(const Literal& src_literal, /* static */ std::unique_ptr Literal::CreateR1( const tensorflow::core::Bitmap& values) { - auto literal = MakeUnique(); + auto literal = MakeUnique( + ShapeUtil::MakeShape(PRED, {static_cast(values.bits())})); literal->PopulateR1(values); return literal; } +void Literal::PopulateR1(const tensorflow::core::Bitmap& values) { + CHECK(ShapeUtil::IsArray(shape())); + CHECK_EQ(ShapeUtil::Rank(shape()), 1); + CHECK_EQ(element_count(), values.bits()); + CHECK_EQ(shape().element_type(), PRED); + for (int64 i = 0; i < static_cast(values.bits()); ++i) { + Set({i}, values.get(i)); + } +} + /* static */ std::unique_ptr Literal::CreateR1U8( tensorflow::StringPiece value) { - auto literal = MakeUnique(); - *literal->mutable_shape() = - ShapeUtil::MakeShape(U8, {static_cast(value.size())}); - literal->set_u8s(tensorflow::StringPiece(value.ToString())); + auto literal = MakeUnique( + ShapeUtil::MakeShape(U8, {static_cast(value.size())})); + for (int i = 0; i < value.size(); ++i) { + literal->Set({i}, value[i]); + } return literal; } @@ -382,26 +704,14 @@ Status Literal::Copy(const Literal& src_literal, std::unique_ptr Literal::Relayout( const Layout& new_layout, const ShapeIndex& shape_index) const { - std::unique_ptr outer_result = CloneToUnique(); - - const Literal* copy_from = this; - Literal* copy_to = outer_result.get(); - for (int64 i = 0; i < shape_index.size(); i++) { - *ShapeUtil::GetMutableSubshape(copy_to->mutable_shape(), {shape_index, i}) - ->mutable_layout() = new_layout; - copy_from = ©_from->tuple_literals_[shape_index[i]]; - copy_to = ©_to->tuple_literals_[shape_index[i]]; - } - - DimensionVector base(ShapeUtil::Rank(copy_from->shape()), 0); - DimensionVector copy_size(copy_from->shape().dimensions().begin(), - copy_from->shape().dimensions().end()); - - CHECK(ShapeUtil::IsArray(copy_from->shape())); - CHECK(ShapeUtil::IsArray(copy_to->shape())); - *copy_to->mutable_shape()->mutable_layout() = new_layout; - TF_CHECK_OK(copy_to->Copy(*copy_from, base, base, copy_size)); - return outer_result; + // Create new shape with 'new_layout' set at the given shape index. + Shape new_shape = shape(); + Shape* subshape = ShapeUtil::GetMutableSubshape(&new_shape, shape_index); + TF_CHECK_OK(LayoutUtil::ValidateLayoutForShape(new_layout, *subshape)); + *subshape->mutable_layout() = new_layout; + auto result = MakeUnique(new_shape); + TF_CHECK_OK(result->CopyFrom(*this)); + return result; } std::unique_ptr Literal::Relayout( @@ -415,11 +725,9 @@ std::unique_ptr Literal::Relayout( result->shape(), [this, &result](const Shape& subshape, const ShapeIndex& index) { if (ShapeUtil::IsArray(subshape)) { - DimensionVector base(ShapeUtil::Rank(subshape), 0); - DimensionVector copy_size(subshape.dimensions().begin(), - subshape.dimensions().end()); - TF_CHECK_OK(result->GetSubliteral(index).Copy(GetSubliteral(index), - base, base, copy_size)); + TF_CHECK_OK(result->CopyFrom(*this, + /*dest_shape_index=*/index, + /*src_shape_index=*/index)); } }); return result; @@ -427,7 +735,7 @@ std::unique_ptr Literal::Relayout( StatusOr> Literal::Reshape( tensorflow::gtl::ArraySlice dimensions) const { - if (ShapeUtil::IsTuple(shape())) { + if (!ShapeUtil::IsArray(shape())) { return InvalidArgument("Reshape does not support tuples."); } std::unique_ptr output; @@ -439,8 +747,7 @@ StatusOr> Literal::Reshape( } // Because the layout is monotonic, we can simply reuse the same sequence of // values without changing their order. - *output->mutable_shape() = - ShapeUtil::MakeShape(shape().element_type(), dimensions); + output->shape_ = ShapeUtil::MakeShape(shape().element_type(), dimensions); int64 elements_before = ShapeUtil::ElementsIn(shape()); int64 elements_after = ShapeUtil::ElementsIn(output->shape()); @@ -456,7 +763,7 @@ StatusOr> Literal::Reshape( std::unique_ptr Literal::Transpose( tensorflow::gtl::ArraySlice permutation) const { - CHECK(!ShapeUtil::IsTuple(shape())) << "Tuple is not supported for transpose"; + CHECK(ShapeUtil::IsArray(shape())) << "Tuple is not supported for transpose"; CHECK(IsPermutation(permutation, ShapeUtil::Rank(shape()))) << "Given permutation is not a permutation of dimension numbers"; // To transpose the array, we just permute the dimensions and layout, and @@ -479,7 +786,7 @@ std::unique_ptr Literal::Transpose( // dimension has within the transposed array, a layout is affine if // MinMaj(Di) == TMinMaj(T(Di)), with TMinMaj() being the minor to major // vector of the affine layout. - CHECK(LayoutUtil::IsDense(permuted_shape)); + CHECK(LayoutUtil::IsDenseArray(permuted_shape)); Layout* layout = permuted_shape.mutable_layout(); layout->clear_minor_to_major(); for (auto index : LayoutUtil::MinorToMajor(shape())) { @@ -488,15 +795,15 @@ std::unique_ptr Literal::Transpose( std::unique_ptr new_literal = CreateFromShape(permuted_shape); DCHECK_GE(ShapeUtil::ByteSizeOf(new_literal->shape()), ShapeUtil::ByteSizeOf(shape())); - std::memcpy(new_literal->MutableInternalData(), InternalData(), - ShapeUtil::ByteSizeOf(shape())); + std::memcpy(new_literal->root_piece().buffer(), root_piece().buffer(), + root_piece().size_bytes()); return new_literal; } std::unique_ptr Literal::Slice( tensorflow::gtl::ArraySlice start_indices, tensorflow::gtl::ArraySlice limit_indices) const { - CHECK(!ShapeUtil::IsTuple(shape())) << "tuple is not supported for reshape"; + CHECK(ShapeUtil::IsArray(shape())) << "tuple is not supported for slice"; DimensionVector result_dimensions; for (int64 dnum = 0; dnum < ShapeUtil::Rank(shape()); ++dnum) { @@ -510,9 +817,7 @@ std::unique_ptr Literal::Slice( ShapeUtil::MakeShapeWithLayout(shape().element_type(), result_dimensions, LayoutUtil::MinorToMajor(shape())); - auto result_literal = MakeUnique(); - *result_literal->mutable_shape() = result_shape; - result_literal->Reserve(ShapeUtil::ElementsIn(result_shape)); + auto result_literal = MakeUnique(result_shape); DimensionVector new_indices(ShapeUtil::Rank(result_shape)); switch (result_shape.element_type()) { @@ -526,6 +831,16 @@ std::unique_ptr Literal::Slice( result_literal->Set(indices, value); }); return result_literal; + case C64: + result_literal->EachCell( + [&](tensorflow::gtl::ArraySlice indices, complex64 /*value*/) { + for (int64 i = 0; i < ShapeUtil::Rank(result_shape); ++i) { + new_indices[i] = indices[i] + start_indices[i]; + } + complex64 value = Get(new_indices); + result_literal->Set(indices, value); + }); + return result_literal; case S32: result_literal->EachCell( [&](tensorflow::gtl::ArraySlice indices, int32 /*value*/) { @@ -552,48 +867,116 @@ std::unique_ptr Literal::Slice( } } +Literal Literal::Clone() const { + Literal result(shape()); + TF_CHECK_OK(result.CopyFrom(*this)); + return result; +} + std::unique_ptr Literal::CloneToUnique() const { - auto unique = MakeUnique(); - *unique = *this; - return unique; + auto result = MakeUnique(shape()); + TF_CHECK_OK(result->CopyFrom(*this)); + return result; } -string Literal::GetAsString( - tensorflow::gtl::ArraySlice multi_index) const { - switch (shape().element_type()) { +string Literal::GetAsString(tensorflow::gtl::ArraySlice multi_index, + const ShapeIndex& shape_index) const { + const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); + CHECK(LayoutUtil::IsDenseArray(subshape)); + switch (subshape.element_type()) { case PRED: - return Get(multi_index) ? "true" : "false"; - case U8: - return tensorflow::strings::StrCat(Get(multi_index)); + return Get(multi_index, shape_index) ? "true" : "false"; + case S8: + return StrCat(Get(multi_index, shape_index)); + case S16: + return StrCat(Get(multi_index, shape_index)); case S32: - return tensorflow::strings::StrCat(Get(multi_index)); + return StrCat(Get(multi_index, shape_index)); case S64: - return tensorflow::strings::StrCat(Get(multi_index)); + return StrCat(Get(multi_index, shape_index)); + case U8: + return StrCat(Get(multi_index, shape_index)); + case U16: + return StrCat(Get(multi_index, shape_index)); case U32: - return tensorflow::strings::StrCat(Get(multi_index)); + return StrCat(Get(multi_index, shape_index)); case U64: - return tensorflow::strings::StrCat(Get(multi_index)); + return StrCat(Get(multi_index, shape_index)); + case F16: + return StrCat(Get(multi_index, shape_index)); case F32: - return tensorflow::strings::StrCat(Get(multi_index)); + return StrCat(Get(multi_index, shape_index)); + case BF16: + return StrCat( + static_cast(Get(multi_index, shape_index))); case F64: - return tensorflow::strings::StrCat(Get(multi_index)); + return StrCat(Get(multi_index, shape_index)); case C64: { - complex64 c = Get(multi_index); - return tensorflow::strings::StrCat("(", c.real(), ", ", c.imag(), ")"); + complex64 c = Get(multi_index, shape_index); + return StrCat("(", c.real(), ", ", c.imag(), ")"); } + default: + LOG(FATAL) << PrimitiveType_Name(subshape.element_type()); + } +} + +string Literal::GetSparseElementAsString(int64 sparse_element_number, + const ShapeIndex& shape_index) const { + const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); + CHECK(LayoutUtil::IsSparseArray(subshape)); + switch (subshape.element_type()) { + case PRED: + return GetSparseElement(sparse_element_number, shape_index) + ? "true" + : "false"; + case S8: + return StrCat(GetSparseElement(sparse_element_number, shape_index)); + case S16: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); + case S32: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); + case S64: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); + case U8: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); + case U16: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); + case U32: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); + case U64: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); case F16: - return tensorflow::strings::StrCat(Get(multi_index)); + return StrCat(GetSparseElement(sparse_element_number, shape_index)); + case F32: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); case BF16: - return tensorflow::strings::StrCat( - static_cast(Get(multi_index))); + return StrCat(static_cast( + GetSparseElement(sparse_element_number, shape_index))); + case F64: + return StrCat( + GetSparseElement(sparse_element_number, shape_index)); + case C64: { + complex64 c = + GetSparseElement(sparse_element_number, shape_index); + return StrCat("(", c.real(), ", ", c.imag(), ")"); + } default: - return tensorflow::strings::StrCat( - "[", PrimitiveType_Name(shape().element_type()), "]"); + LOG(FATAL) << "Invalid element type for sparse arrays: " + << PrimitiveType_Name(subshape.element_type()); } } StatusOr Literal::GetIntegralAsS64( tensorflow::gtl::ArraySlice multi_index) const { + CHECK(LayoutUtil::IsDenseArray(shape())); switch (shape().element_type()) { case PRED: return Get(multi_index); @@ -614,13 +997,83 @@ StatusOr Literal::GetIntegralAsS64( } } -int64 Literal::LinearIndex( - tensorflow::gtl::ArraySlice multi_index) const { - return IndexUtil::MultidimensionalIndexToLinearIndex(shape(), multi_index); +tensorflow::gtl::ArraySlice Literal::GetSparseIndex( + int64 sparse_element_number, const ShapeIndex& shape_index) const { + const Piece& p = piece(shape_index); + CHECK_GE(sparse_element_number, 0); + CHECK_LT(sparse_element_number, p.sparse_indices()->index_count()); + return p.sparse_indices()->At(sparse_element_number); } -string Literal::ToString(bool print_layout) const { - std::vector pieces; +void Literal::SortSparseElements(const ShapeIndex& shape_index) { + piece(shape_index).SortSparseElements(); +} + +void Literal::Piece::SortSparseElements() { + switch (subshape().element_type()) { + case PRED: + SortSparseElementsInternal(); + break; + case S8: + SortSparseElementsInternal(); + break; + case U8: + SortSparseElementsInternal(); + break; + case S16: + SortSparseElementsInternal(); + break; + case U16: + SortSparseElementsInternal(); + break; + case S32: + SortSparseElementsInternal(); + break; + case U32: + SortSparseElementsInternal(); + break; + case S64: + SortSparseElementsInternal(); + break; + case U64: + SortSparseElementsInternal(); + break; + case F32: + SortSparseElementsInternal(); + break; + case F64: + SortSparseElementsInternal(); + break; + case C64: + SortSparseElementsInternal(); + break; + case F16: + SortSparseElementsInternal(); + break; + case BF16: + SortSparseElementsInternal(); + break; + default: + LOG(FATAL) << "Element type not valid for sparse array: " + << PrimitiveType_Name(subshape().element_type()); + } +} + +template +void Literal::Piece::SortSparseElementsInternal() { + CHECK(LayoutUtil::IsSparseArray(subshape())); + int64 num_elements = sparse_indices()->index_count(); + auto values = data(); + CHECK_LE(num_elements, values.size()); + sparse_indices()->SortWithValues( + tensorflow::gtl::MutableArraySlice(values.data(), num_elements)); +} + +namespace { + +void ToStringHelper(const Literal& literal, const ShapeIndex& shape_index, + bool print_layout, std::vector* pieces) { + const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); auto shape_to_string = [print_layout](const Shape& shape) { if (print_layout) { @@ -630,328 +1083,236 @@ string Literal::ToString(bool print_layout) const { } }; + // TODO(b/32894291): refactor this code to reduce code duplication. + if (ShapeUtil::IsTuple(subshape)) { + pieces->push_back(shape_to_string(subshape)); + pieces->push_back(" (\n"); + std::vector tuple_pieces; + for (int i = 0; i < ShapeUtil::TupleElementCount(subshape); ++i) { + ShapeIndex element_index = shape_index; + element_index.push_back(i); + std::vector element_pieces; + ToStringHelper(literal, element_index, print_layout, &element_pieces); + tuple_pieces.push_back(tensorflow::str_util::Join(element_pieces, "")); + } + pieces->push_back(tensorflow::str_util::Join(tuple_pieces, ",\n")); + pieces->push_back("\n)"); + return; + } + + if (LayoutUtil::IsSparseArray(subshape)) { + pieces->push_back(shape_to_string(subshape)); + pieces->push_back("{"); + int64 rank = ShapeUtil::Rank(subshape); + int64 num_elements = literal.sparse_element_count(); + for (int64 i = 0; i < num_elements; ++i) { + if (i > 0) { + pieces->push_back(", "); + } + if (rank == 1) { + pieces->push_back(StrCat(literal.GetSparseIndex(i)[0])); + pieces->push_back(": "); + } else { + pieces->push_back("["); + pieces->push_back( + tensorflow::str_util::Join(literal.GetSparseIndex(i), ", ")); + pieces->push_back("]: "); + } + pieces->push_back(literal.GetSparseElementAsString(i)); + } + pieces->push_back("}"); + return; + } + + CHECK(LayoutUtil::IsDenseArray(subshape)); + auto element_to_string = - [this](tensorflow::gtl::ArraySlice indices) -> string { - PrimitiveType element_type = shape().element_type(); + [&](tensorflow::gtl::ArraySlice indices) -> string { + PrimitiveType element_type = subshape.element_type(); if (element_type == PRED) { // We display predicates in a densely packed form. - return Get(indices) ? "1" : "0"; + return literal.Get(indices, shape_index) ? "1" : "0"; } return ((!indices.empty() && indices.back() > 0) ? ", " : "") + - GetAsString(indices); + literal.GetAsString(indices, shape_index); }; - // TODO(b/32894291): refactor this code to reduce code duplication. - if (ShapeUtil::IsTuple(shape())) { - pieces.push_back(shape_to_string(shape())); - pieces.push_back(" (\n"); - pieces.push_back(tensorflow::str_util::Join( - tuple_literals(), ",\n", [](string* out, const Literal& element) { - tensorflow::strings::StrAppend(out, element.ToString()); - })); - pieces.push_back("\n)"); - } else if (ShapeUtil::Rank(shape()) == 0) { - pieces.push_back(GetAsString({})); - } else if (ShapeUtil::Rank(shape()) == 1) { - pieces.push_back("{"); - for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { - pieces.push_back(element_to_string({i0})); + if (ShapeUtil::Rank(subshape) == 0) { + pieces->push_back(literal.GetAsString({}, shape_index)); + } else if (ShapeUtil::Rank(subshape) == 1) { + pieces->push_back("{"); + for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { + pieces->push_back(element_to_string({i0})); } - pieces.push_back("}"); - } else if (ShapeUtil::Rank(shape()) == 2) { - pieces.push_back(shape_to_string(shape())); - pieces.push_back(" {\n"); - for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { - pieces.push_back(" { "); - for (int64 i1 = 0; i1 < shape().dimensions(1); ++i1) { - pieces.push_back(element_to_string({i0, i1})); + pieces->push_back("}"); + } else if (ShapeUtil::Rank(subshape) == 2) { + pieces->push_back(shape_to_string(subshape)); + pieces->push_back(" {\n"); + for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { + pieces->push_back(" { "); + for (int64 i1 = 0; i1 < subshape.dimensions(1); ++i1) { + pieces->push_back(element_to_string({i0, i1})); } - pieces.push_back(" "); - pieces.push_back(i0 == shape().dimensions(0) - 1 ? "}\n" : "},\n"); + pieces->push_back(" "); + pieces->push_back(i0 == subshape.dimensions(0) - 1 ? "}\n" : "},\n"); } - pieces.push_back("}"); - } else if (ShapeUtil::Rank(shape()) == 3) { - pieces.push_back(shape_to_string(shape())); - pieces.push_back(" {\n"); - for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { - pieces.push_back(i0 > 0 ? ",\n{" : "{"); - for (int64 i1 = 0; i1 < shape().dimensions(1); ++i1) { - pieces.push_back(i1 > 0 ? ",\n { " : " { "); - for (int64 i2 = 0; i2 < shape().dimensions(2); ++i2) { - pieces.push_back(element_to_string({i0, i1, i2})); + pieces->push_back("}"); + } else if (ShapeUtil::Rank(subshape) == 3) { + pieces->push_back(shape_to_string(subshape)); + pieces->push_back(" {\n"); + for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { + pieces->push_back(i0 > 0 ? ",\n{" : "{"); + for (int64 i1 = 0; i1 < subshape.dimensions(1); ++i1) { + pieces->push_back(i1 > 0 ? ",\n { " : " { "); + for (int64 i2 = 0; i2 < subshape.dimensions(2); ++i2) { + pieces->push_back(element_to_string({i0, i1, i2})); } - pieces.push_back(" }"); + pieces->push_back(" }"); } - pieces.push_back(" }"); + pieces->push_back(" }"); } - pieces.push_back("\n}"); - } else if (ShapeUtil::Rank(shape()) == 4) { - pieces.push_back(shape_to_string(shape())); - pieces.push_back(" {\n"); - for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { - pieces.push_back(tensorflow::strings::Printf(" { /*i0=%lld*/\n", i0)); - for (int64 i1 = 0; i1 < shape().dimensions(1); ++i1) { - pieces.push_back( - tensorflow::strings::Printf(" { /*i1=%lld*/\n", i1)); - for (int64 i2 = 0; i2 < shape().dimensions(2); ++i2) { - pieces.push_back(" {"); - for (int64 i3 = 0; i3 < shape().dimensions(3); ++i3) { - pieces.push_back(element_to_string({i0, i1, i2, i3})); + pieces->push_back("\n}"); + } else if (ShapeUtil::Rank(subshape) == 4) { + pieces->push_back(shape_to_string(subshape)); + pieces->push_back(" {\n"); + for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { + pieces->push_back(Printf(" { /*i0=%lld*/\n", i0)); + for (int64 i1 = 0; i1 < subshape.dimensions(1); ++i1) { + pieces->push_back(Printf(" { /*i1=%lld*/\n", i1)); + for (int64 i2 = 0; i2 < subshape.dimensions(2); ++i2) { + pieces->push_back(" {"); + for (int64 i3 = 0; i3 < subshape.dimensions(3); ++i3) { + pieces->push_back(element_to_string({i0, i1, i2, i3})); } - pieces.push_back(i2 == shape().dimensions(2) - 1 ? "}\n" : "},\n"); + pieces->push_back(i2 == subshape.dimensions(2) - 1 ? "}\n" : "},\n"); } - pieces.push_back(i1 == shape().dimensions(1) - 1 ? " }\n" - : " },\n"); + pieces->push_back(i1 == subshape.dimensions(1) - 1 ? " }\n" + : " },\n"); } - pieces.push_back(i0 == shape().dimensions(0) - 1 ? " }\n" : " },\n"); + pieces->push_back(i0 == subshape.dimensions(0) - 1 ? " }\n" : " },\n"); } - pieces.push_back("}"); - } else if (ShapeUtil::Rank(shape()) == 5) { - pieces.push_back(shape_to_string(shape())); - pieces.push_back(" {\n"); - for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { - pieces.push_back(tensorflow::strings::Printf(" { /*i0=%lld*/\n", i0)); - for (int64 i1 = 0; i1 < shape().dimensions(1); ++i1) { - pieces.push_back( - tensorflow::strings::Printf(" { /*i1=%lld*/\n", i1)); - for (int64 i2 = 0; i2 < shape().dimensions(2); ++i2) { - pieces.push_back( - tensorflow::strings::Printf(" { /*i2=%lld*/\n", i2)); - for (int64 i3 = 0; i3 < shape().dimensions(3); ++i3) { - pieces.push_back(" {"); - for (int64 i4 = 0; i4 < shape().dimensions(4); ++i4) { - pieces.push_back(element_to_string({i0, i1, i2, i3, i4})); + pieces->push_back("}"); + } else if (ShapeUtil::Rank(subshape) == 5) { + pieces->push_back(shape_to_string(subshape)); + pieces->push_back(" {\n"); + for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { + pieces->push_back(Printf(" { /*i0=%lld*/\n", i0)); + for (int64 i1 = 0; i1 < subshape.dimensions(1); ++i1) { + pieces->push_back(Printf(" { /*i1=%lld*/\n", i1)); + for (int64 i2 = 0; i2 < subshape.dimensions(2); ++i2) { + pieces->push_back(Printf(" { /*i2=%lld*/\n", i2)); + for (int64 i3 = 0; i3 < subshape.dimensions(3); ++i3) { + pieces->push_back(" {"); + for (int64 i4 = 0; i4 < subshape.dimensions(4); ++i4) { + pieces->push_back(element_to_string({i0, i1, i2, i3, i4})); } - pieces.push_back(i3 == shape().dimensions(3) - 1 ? "}\n" : "},\n"); + pieces->push_back(i3 == subshape.dimensions(3) - 1 ? "}\n" + : "},\n"); } - pieces.push_back(i2 == shape().dimensions(2) - 1 ? " }\n" - : " },\n"); + pieces->push_back(i2 == subshape.dimensions(2) - 1 ? " }\n" + : " },\n"); } - pieces.push_back(i1 == shape().dimensions(1) - 1 ? " }\n" - : " },\n"); + pieces->push_back(i1 == subshape.dimensions(1) - 1 ? " }\n" + : " },\n"); } - pieces.push_back(i0 == shape().dimensions(0) - 1 ? " }\n" : " },\n"); + pieces->push_back(i0 == subshape.dimensions(0) - 1 ? " }\n" : " },\n"); } - pieces.push_back("}"); + pieces->push_back("}"); } else { - pieces.push_back(shape_to_string(shape())); - pieces.push_back(" {"); - EachCellAsString( + pieces->push_back(shape_to_string(subshape)); + pieces->push_back(" {"); + literal.EachCellAsString( [&](tensorflow::gtl::ArraySlice indices, const string& value) { - pieces.push_back(" "); - pieces.push_back(value); + pieces->push_back(" "); + pieces->push_back(value); }); - pieces.push_back("}"); + pieces->push_back("}"); } +} +} // namespace + +int64 Literal::sparse_element_count() const { + CHECK(LayoutUtil::IsSparseArray(shape())); + return sparse_indices()->index_count(); +} + +string Literal::ToString(bool print_layout) const { + std::vector pieces; + ToStringHelper(*this, {}, print_layout, &pieces); return tensorflow::str_util::Join(pieces, ""); } /* static */ std::unique_ptr Literal::MakeTuple( tensorflow::gtl::ArraySlice elements) { - auto literal = MakeUnique(); - std::vector shape; - for (const Literal* tuple_element : elements) { - *literal->add_tuple_literals() = *tuple_element; - shape.push_back(tuple_element->shape()); + std::vector element_shapes; + for (const Literal* element : elements) { + element_shapes.push_back(element->shape()); + } + auto literal = MakeUnique(ShapeUtil::MakeTupleShape(element_shapes)); + for (int i = 0; i < elements.size(); ++i) { + TF_CHECK_OK(literal->CopyFrom(*elements[i], /*dest_shape_index=*/{i})); } - *literal->mutable_shape() = ShapeUtil::MakeTupleShape(shape); return literal; } /* static */ std::unique_ptr Literal::MakeTupleOwned( std::vector> elements) { - auto literal = MakeUnique(); - std::vector shape; - for (auto& tuple_element : elements) { - shape.push_back(tuple_element->shape()); - *literal->add_tuple_literals() = std::move(*tuple_element); + std::vector element_shapes; + element_shapes.reserve(elements.size()); + for (const auto& element : elements) { + element_shapes.push_back(element->shape()); + } + auto literal = MakeUnique(ShapeUtil::MakeTupleShape(element_shapes)); + for (int64 i = 0; i < elements.size(); ++i) { + TF_CHECK_OK( + literal->MoveFrom(std::move(*elements[i]), /*dest_shape_index=*/{i})); } - *literal->mutable_shape() = ShapeUtil::MakeTupleShape(shape); return literal; } -const void* Literal::InternalData() const { - return const_cast( - const_cast(this)->MutableInternalData()); +void Literal::EachCellAsString( + const std::function indices, + const string& value)>& per_cell) const { + if (ShapeUtil::HasZeroElements(shape())) { + return; + } + std::vector indices = IndexUtil::LinearIndexToMultidimensionalIndex( + shape(), /*linear_index=*/0); + do { + per_cell(indices, GetAsString(indices)); + } while (IndexUtil::BumpIndices(shape(), &indices)); } -void* Literal::MutableInternalData() { - // NOTE: We access the vectors directly to avoid the const reference - // created by the accessor functions. - switch (shape().element_type()) { - case PRED: - case U8: - return reinterpret_cast(u8s_.data()); - case S32: - return reinterpret_cast(s32s_.data()); - case S64: - return reinterpret_cast(s64s_.data()); - case U32: - return reinterpret_cast(u32s_.data()); - case U64: - return reinterpret_cast(u64s_.data()); - case F32: - return reinterpret_cast(f32s_.data()); - case F64: - return reinterpret_cast(f64s_.data()); - case C64: - return reinterpret_cast(c64s_.data()); - case F16: - return reinterpret_cast(f16s_.data()); - case BF16: - return reinterpret_cast(bf16s_.data()); - default: - LOG(FATAL) << "primitive type not supported in literals: " - << PrimitiveType_Name(shape().element_type()); +namespace { +template +std::unique_ptr ConvertBetweenNativeTypes(const Literal& src_literal) { + CHECK(ShapeUtil::IsArray(src_literal.shape())); + auto result_literal = MakeUnique(ShapeUtil::ChangeElementType( + src_literal.shape(), + primitive_util::NativeToPrimitiveType())); + auto src_data = src_literal.data(); + auto dest_data = result_literal->template data(); + int64 num_elements = src_literal.element_count(); + + for (int64 i = 0; i < num_elements; ++i) { + dest_data[i] = static_cast(src_data[i]); } -} - -void Literal::Reserve(int64 num_elements) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - switch (shape().element_type()) { - case PRED: - Resize(num_elements, false); - break; - case S8: - Resize(num_elements, 0); - break; - case U8: - Resize(num_elements, 0); - break; - case S32: - Resize(num_elements, 0); - break; - case S64: - Resize(num_elements, 0); - break; - case U32: - Resize(num_elements, 0); - break; - case U64: - Resize(num_elements, 0); - break; - case F32: - Resize(num_elements, 0); - break; - case F64: - Resize(num_elements, 0); - break; - case C64: - Resize(num_elements, 0); - break; - case F16: - Resize(num_elements, static_cast(0.0f)); - break; - case BF16: - Resize(num_elements, static_cast(0.0f)); - break; - default: - LOG(FATAL) << "primitive type not supported in literals: " - << PrimitiveType_Name(shape().element_type()); - } -} - -tensorflow::Status Literal::ValidateLiteral() const { - TF_CHECK_OK(ShapeUtil::ValidateShape(shape())); - int64 expected = ShapeUtil::ElementsIn(shape()); - int64 actual = -1; - switch (shape().element_type()) { - case PRED: - case U8: - actual = u8s_size(); - break; - case S32: - actual = s32s_size(); - break; - case U32: - actual = u32s_size(); - break; - case S64: - actual = s64s_size(); - break; - case U64: - actual = u64s_size(); - break; - case F32: - actual = f32s_size(); - break; - case F64: - actual = f64s_size(); - break; - case C64: - actual = c64s_size(); - break; - case F16: - actual = f16s().size() / sizeof(half); - break; - case BF16: - actual = bf16s().size(); - break; - default: - return tensorflow::errors::Unimplemented( - "unhandled element type for literal validation: " + - PrimitiveType_Name(shape().element_type())); - } - - if (expected != actual) { - return tensorflow::errors::InvalidArgument(tensorflow::strings::Printf( - "literal has bad number of elements for its shape %s: want %lld " - "got %lld", - ShapeUtil::HumanString(shape()).c_str(), expected, actual)); - } - - return tensorflow::Status::OK(); -} - -void Literal::EachCellAsString( - const std::function indices, - const string& value)>& per_cell) const { - if (ShapeUtil::HasZeroElements(shape())) { - return; - } - std::vector indices = IndexUtil::LinearIndexToMultidimensionalIndex( - shape(), /*linear_index=*/0); - do { - per_cell(indices, GetAsString(indices)); - } while (IndexUtil::BumpIndices(shape(), &indices)); -} - -namespace { -template -std::unique_ptr ConvertBetweenNativeTypes(const Literal& src_literal) { - auto result_literal = MakeUnique(); - Shape* result_shape = result_literal->mutable_shape(); - *result_shape = src_literal.shape(); - result_shape->set_element_type( - primitive_util::NativeToPrimitiveType()); - result_literal->Reserve(ShapeUtil::ElementsIn(*result_shape)); - tensorflow::gtl::ArraySlice src_data = - src_literal.GetArraySlice(); - tensorflow::gtl::MutableArraySlice dest_data = - result_literal->GetMutableArraySlice(); - int64 num_elements = ShapeUtil::ElementsIn(src_literal.shape()); - - for (int64 i = 0; i < num_elements; ++i) { - dest_data[i] = static_cast(src_data[i]); - } - return result_literal; + return result_literal; } template std::unique_ptr ConvertToC64(const Literal& src_literal) { - auto result_literal = MakeUnique(); - Shape* result_shape = result_literal->mutable_shape(); - *result_shape = src_literal.shape(); - result_shape->set_element_type(C64); - result_literal->Reserve(ShapeUtil::ElementsIn(*result_shape)); + CHECK(ShapeUtil::IsArray(src_literal.shape())); + auto result_literal = MakeUnique( + ShapeUtil::ChangeElementType(src_literal.shape(), C64)); using NativeSrcT = typename primitive_util::PrimitiveTypeToNative::type; tensorflow::gtl::ArraySlice src_data = - src_literal.GetArraySlice(); + src_literal.data(); tensorflow::gtl::MutableArraySlice dest_data = - result_literal->GetMutableArraySlice(); - int64 num_elements = ShapeUtil::ElementsIn(src_literal.shape()); + result_literal->data(); + int64 num_elements = src_literal.element_count(); for (int64 i = 0; i < num_elements; ++i) { dest_data[i] = complex64(static_cast(src_data[i]), 0); } @@ -996,10 +1357,12 @@ StatusOr> ConvertIfDestTypeMatches( PrimitiveType_Name(primitive_dest_type).c_str()); } } + } // namespace StatusOr> Literal::Convert( PrimitiveType primitive_dest_type) const { + TF_RET_CHECK(ShapeUtil::IsArray(shape())); switch (shape().element_type()) { #define CONVERT_IF_DEST_TYPE_MATCHES(type) \ case (type): \ @@ -1024,356 +1387,192 @@ StatusOr> Literal::Convert( } } -namespace { - -// Helper function which compares whether the elements of literal1 are equal to -// the elements of literal2. Recursively iterates through the entire -// multidimensional index space and compares the literal elements -// one-by-one. literal1 and literal2 must be compatible (same dimensions and -// type). template -bool EqualElements(const Literal& literal1, const Literal& literal2, - int dimension, std::vector* multi_index) { - if (dimension == ShapeUtil::Rank(literal1.shape())) { - return (literal1.Get(*multi_index) == - literal2.Get(*multi_index)); - } - for (int64 i = 0; i < literal1.shape().dimensions(dimension); ++i) { - (*multi_index)[dimension] = i; - if (!EqualElements(literal1, literal2, dimension + 1, - multi_index)) { +bool Literal::Piece::EqualElementsInternal( + const Literal::Piece& other, std::vector* multi_index) const { + if (multi_index->size() == ShapeUtil::Rank(subshape())) { + return (Get(*multi_index) == other.Get(*multi_index)); + } + for (int64 i = 0; i < subshape().dimensions(multi_index->size()); ++i) { + multi_index->push_back(i); + if (!EqualElementsInternal(other, multi_index)) { return false; } + multi_index->pop_back(); } return true; } -} // namespace +bool Literal::Piece::EqualElements(const Literal::Piece& other) const { + DCHECK(ShapeUtil::Compatible(subshape(), other.subshape())); + + std::vector multi_index; + switch (subshape().element_type()) { + case PRED: + return EqualElementsInternal(other, &multi_index); + case U8: + return EqualElementsInternal(other, &multi_index); + case S32: + return EqualElementsInternal(other, &multi_index); + case S64: + return EqualElementsInternal(other, &multi_index); + case U32: + return EqualElementsInternal(other, &multi_index); + case U64: + return EqualElementsInternal(other, &multi_index); + case F32: + return EqualElementsInternal(other, &multi_index); + case F64: + return EqualElementsInternal(other, &multi_index); + case F16: + return EqualElementsInternal(other, &multi_index); + case BF16: + return EqualElementsInternal(other, &multi_index); + case C64: + return EqualElementsInternal(other, &multi_index); + default: + LOG(FATAL) << "Unimplemented: Literal::Piece::EqualElements for type " + << PrimitiveType_Name(subshape().element_type()); + } +} bool Literal::operator==(const Literal& other) const { if (!ShapeUtil::Compatible(shape(), other.shape())) { return false; } - if (ShapeUtil::IsTuple(shape())) { - // Because the shapes are compatible, they must have the same number of - // tuple elements. - CHECK_EQ(tuple_literals_size(), other.tuple_literals_size()); - for (int i = 0; i < tuple_literals_size(); ++i) { - if (tuple_literals(i) != other.tuple_literals(i)) { - return false; - } + for (const auto& pair : pieces_) { + const ShapeIndex& index = pair.first; + const Piece& piece = pair.second; + if (!ShapeUtil::IsArray(piece.subshape())) { + continue; } - return true; - } else { - std::vector multi_index(ShapeUtil::Rank(shape()), 0); - switch (shape().element_type()) { - case PRED: - return EqualElements(*this, other, 0, &multi_index); - case U8: - return EqualElements(*this, other, 0, &multi_index); - case S32: - return EqualElements(*this, other, 0, &multi_index); - case S64: - return EqualElements(*this, other, 0, &multi_index); - case U32: - return EqualElements(*this, other, 0, &multi_index); - case U64: - return EqualElements(*this, other, 0, &multi_index); - case F32: - return EqualElements(*this, other, 0, &multi_index); - case F64: - return EqualElements(*this, other, 0, &multi_index); - case F16: - return EqualElements(*this, other, 0, &multi_index); - case BF16: - return EqualElements(*this, other, 0, &multi_index); - case C64: - return EqualElements(*this, other, 0, &multi_index); - default: - LOG(FATAL) << "Unimplemented: Literal::Equal for type " - << PrimitiveType_Name(shape().element_type()); + + const Piece& other_piece = other.piece(index); + if (!piece.EqualElements(other_piece)) { + return false; } } + return true; } -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_preds(); - return tensorflow::gtl::MutableArraySlice( - reinterpret_cast(values->data()), values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_u8s(); - return tensorflow::gtl::MutableArraySlice( - reinterpret_cast(values->data()), values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_u8s(); - return tensorflow::gtl::MutableArraySlice(values->data(), - values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_s16s(); - return tensorflow::gtl::MutableArraySlice(values->data(), - values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_u16s(); - return tensorflow::gtl::MutableArraySlice(values->data(), - values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_s32s(); - return tensorflow::gtl::MutableArraySlice(values->data(), - values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_u32s(); - return tensorflow::gtl::MutableArraySlice(values->data(), - values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - static_assert(sizeof(int64) == sizeof(tensorflow::protobuf_int64) && - alignof(int64) == alignof(tensorflow::protobuf_int64), - "The int64 and tensorflow::protobuf_int64 types are not " - "compatible"); - auto values = mutable_s64s(); - // Because of the fact that tensorflow::protobuf_int64 is defined as int64_t - // while tensorflow::int64 is defined as long long, a reinterpret_cast<> is - // necessary from the raw data pointer returned by the mutable_data() API. - return tensorflow::gtl::MutableArraySlice( - reinterpret_cast(values->data()), values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - static_assert(sizeof(uint64) == sizeof(tensorflow::protobuf_uint64) && - alignof(uint64) == alignof(tensorflow::protobuf_uint64), - "The uint64 and tensorflow::protobuf_uint64 types are not " - "compatible"); - auto values = mutable_u64s(); - // Because of the fact that tensorflow::protobuf_uint64 is defined as uint64_t - // while tensorflow::uint64 is defined as unsigned long long, a - // reinterpret_cast<> is necessary from the raw data pointer returned by the - // mutable_data() API. - return tensorflow::gtl::MutableArraySlice( - reinterpret_cast(values->data()), values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_f32s(); - return tensorflow::gtl::MutableArraySlice(values->data(), - values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_f64s(); - return tensorflow::gtl::MutableArraySlice(values->data(), - values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_c64s(); - return {values->data(), values->size()}; -} - -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { - auto values = mutable_f16s(); - return tensorflow::gtl::MutableArraySlice(values->data(), - values->size()); -} - -template <> -tensorflow::gtl::MutableArraySlice -Literal::GetMutableArraySlice() { - auto values = mutable_bf16s(); - return {values->data(), values->size()}; -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), PRED); - return tensorflow::gtl::ArraySlice( - reinterpret_cast(preds().data()), preds().size()); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), U8); - return tensorflow::gtl::ArraySlice( - reinterpret_cast(u8s().data()), u8s().size()); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), S8); - return tensorflow::gtl::ArraySlice( - reinterpret_cast(u8s().data()), u8s().size()); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), U16); - return tensorflow::gtl::ArraySlice(u16s().data(), u16s().size()); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), S16); - return tensorflow::gtl::ArraySlice(s16s().data(), s16s().size()); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), U32); - return u32s(); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), U64); - return u64s(); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), S32); - return s32s(); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), S64); - return s64s(); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), F64); - return f64s(); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), F16); - return tensorflow::gtl::ArraySlice(f16s().data(), - f16s().size() / sizeof(half)); -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { - CHECK_EQ(shape().element_type(), BF16); - return {bf16s().data(), bf16s().size()}; -} - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() - const { - CHECK_EQ(shape().element_type(), C64); - return c64s(); -} +namespace { template -static bool AllElementsEqualValue(const Literal& literal, NativeT value) { - for (int64 i = 0; i < ShapeUtil::ElementsIn(literal.shape()); ++i) { - auto multi_index = - IndexUtil::LinearIndexToMultidimensionalIndex(literal.shape(), i); - if (literal.Get(multi_index) != value) { +static bool AllElementsEqualValue(tensorflow::gtl::ArraySlice data, + NativeT value) { + for (int64 i = 0; i < data.size(); ++i) { + if (data[i] != value) { return false; } } return true; } +} // namespace + bool Literal::IsAll(int8 value) const { - switch (shape().element_type()) { - case U8: - if (value >= 0) { - return AllElementsEqualValue(*this, value); - } - return false; - case U32: - if (value >= 0) { - return AllElementsEqualValue(*this, value); - } - return false; - case U64: - if (value >= 0) { - return AllElementsEqualValue(*this, value); - } - return false; - case S8: - return AllElementsEqualValue(*this, value); - case S32: - return AllElementsEqualValue(*this, value); - case S64: - return AllElementsEqualValue(*this, value); - case F32: - return AllElementsEqualValue(*this, value); - case F64: - return AllElementsEqualValue(*this, value); - case F16: - return AllElementsEqualValue(*this, static_cast(value)); - case BF16: - return AllElementsEqualValue(*this, - static_cast(value)); - case PRED: - if (value == 0) { - return AllElementsEqualValue(*this, false); - } - if (value == 1) { - return AllElementsEqualValue(*this, true); + for (const auto& pair : pieces_) { + const Piece& piece = pair.second; + if (!ShapeUtil::IsArray(piece.subshape())) { + continue; + } + + auto piece_is_all = [&]() { + switch (shape().element_type()) { + case U8: + if (value >= 0) { + return AllElementsEqualValue(piece.data(), value); + } + return false; + case U32: + if (value >= 0) { + return AllElementsEqualValue(piece.data(), value); + } + return false; + case U64: + if (value >= 0) { + return AllElementsEqualValue(piece.data(), value); + } + return false; + case S8: + return AllElementsEqualValue(piece.data(), value); + case S32: + return AllElementsEqualValue(piece.data(), value); + case S64: + return AllElementsEqualValue(piece.data(), value); + case F32: + return AllElementsEqualValue(piece.data(), value); + case F64: + return AllElementsEqualValue(piece.data(), value); + case F16: + return AllElementsEqualValue(piece.data(), + static_cast(value)); + case BF16: + return AllElementsEqualValue(piece.data(), + static_cast(value)); + case PRED: + if (value == 0) { + return AllElementsEqualValue(piece.data(), false); + } + if (value == 1) { + return AllElementsEqualValue(piece.data(), true); + } + return false; + default: + return false; } return false; - default: + }; + + if (!piece_is_all()) { return false; + } } + return true; } bool Literal::IsAllFloat(float value) const { - switch (shape().element_type()) { - case F32: - return AllElementsEqualValue(*this, value); - case F64: - return AllElementsEqualValue(*this, value); - case F16: - return AllElementsEqualValue(*this, static_cast(value)); - case BF16: - return AllElementsEqualValue(*this, - static_cast(value)); - default: + for (const auto& pair : pieces_) { + const Piece& piece = pair.second; + if (!ShapeUtil::IsArray(piece.subshape())) { + continue; + } + + auto piece_is_all = [&]() { + switch (shape().element_type()) { + case F32: + return AllElementsEqualValue(piece.data(), value); + case F64: + return AllElementsEqualValue(piece.data(), value); + case F16: + return AllElementsEqualValue(piece.data(), + static_cast(value)); + case BF16: + return AllElementsEqualValue(piece.data(), + static_cast(value)); + default: + return false; + } + }; + if (!piece_is_all()) { return false; + } } + return true; } bool Literal::IsAllComplex(complex64 value) const { switch (shape().element_type()) { case C64: - return AllElementsEqualValue(*this, value); + return AllElementsEqualValue(root_piece().data(), + value); default: return false; } } bool Literal::IsZero(tensorflow::gtl::ArraySlice indices) const { + CHECK(ShapeUtil::IsArray(shape())); switch (shape().element_type()) { case U8: return Get(indices) == 0; @@ -1404,247 +1603,294 @@ bool Literal::IsZero(tensorflow::gtl::ArraySlice indices) const { } } -template <> -/* static */ void Literal::Resize(int64 num_elements, bool value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_preds()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, int8 value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_u8s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, uint8 value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_u8s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, int32 value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_s32s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, uint32 value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_u32s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, int64 value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_s64s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, uint64 value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_u64s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, float value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_f32s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, double value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_f64s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, half value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_f16s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, bfloat16 value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_bf16s()->resize(num_elements, value); -} - -template <> -void Literal::Resize(int64 num_elements, complex64 value) { - CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); - mutable_c64s()->resize(num_elements, value); -} +namespace { template void CopyToRepeatedField(RepeatedFieldT* dest, - const std::vector& src) { + const tensorflow::gtl::ArraySlice src) { *dest = RepeatedFieldT(src.begin(), src.end()); } -template <> -void CopyToRepeatedField, complex64>( - tensorflow::protobuf::RepeatedField* dest, - const std::vector& src) { - *dest = tensorflow::protobuf::RepeatedField( - reinterpret_cast(src.data()), - reinterpret_cast(src.data()) + src.size() * 2); -} +} // namespace -LiteralProto Literal::ToProto() const { - LiteralProto proto; - proto.Clear(); - *proto.mutable_shape() = shape(); - switch (shape().element_type()) { +void Literal::Piece::WriteToProto(LiteralProto* proto) const { + *proto->mutable_shape() = subshape(); + switch (subshape().element_type()) { case PRED: - CopyToRepeatedField(proto.mutable_preds(), preds()); + CopyToRepeatedField(proto->mutable_preds(), data()); break; case U8: - *proto.mutable_u8s() = u8s_string(); - break; - case S32: - CopyToRepeatedField(proto.mutable_s32s(), s32s()); - break; - case S64: - CopyToRepeatedField(proto.mutable_s64s(), s64s()); + proto->set_u8s(static_cast(data().data()), + element_count()); break; case U32: - CopyToRepeatedField(proto.mutable_u32s(), u32s()); + CopyToRepeatedField(proto->mutable_u32s(), data()); break; case U64: - CopyToRepeatedField(proto.mutable_u64s(), u64s()); + CopyToRepeatedField(proto->mutable_u64s(), data()); + break; + case S32: + CopyToRepeatedField(proto->mutable_s32s(), data()); + break; + case S64: + CopyToRepeatedField(proto->mutable_s64s(), data()); break; case F16: - *proto.mutable_f16s() = - string(reinterpret_cast(f16s_.data()), - f16s_.size() * sizeof(half)); + *proto->mutable_f16s() = string( + reinterpret_cast(data().data()), size_bytes()); if (!kLittleEndian) { - ConvertEndianShort(const_cast(proto.mutable_f16s()->data()), - proto.f16s().size()); + ConvertEndianShort(const_cast(proto->mutable_f16s()->data()), + proto->f16s().size()); } break; case BF16: - *proto.mutable_bf16s() = - string(reinterpret_cast(bf16s_.data()), - bf16s_.size() * sizeof(bfloat16)); + *proto->mutable_bf16s() = string( + reinterpret_cast(data().data()), size_bytes()); if (!kLittleEndian) { - ConvertEndianShort(const_cast(proto.mutable_bf16s()->data()), - proto.bf16s().size()); + ConvertEndianShort(const_cast(proto->mutable_bf16s()->data()), + proto->bf16s().size()); } break; case F32: - CopyToRepeatedField(proto.mutable_f32s(), f32s()); + CopyToRepeatedField(proto->mutable_f32s(), data()); break; case F64: - CopyToRepeatedField(proto.mutable_f64s(), f64s()); + CopyToRepeatedField(proto->mutable_f64s(), data()); break; case C64: - CopyToRepeatedField(proto.mutable_c64s(), c64s()); - break; - case TUPLE: - for (const auto& tuple : tuple_literals()) { - *proto.add_tuple_literals() = tuple.ToProto(); + for (complex64 value : data()) { + proto->add_c64s(value.real()); + proto->add_c64s(value.imag()); } break; + case TUPLE: + // Nothing to do but assign the shape which is done above. + return; default: - LOG(FATAL) << "Unhandled primitive type " << shape().element_type(); + LOG(FATAL) << "Unhandled primitive type " << subshape().element_type(); } - - return proto; } -template -void CopyFromRepeatedField(std::vector* dest, - const RepeatedFieldT& src) { - *dest = std::vector(src.begin(), src.end()); +const void* Literal::Piece::untyped_data() const { + CHECK(ShapeUtil::IsArray(subshape())) << ShapeUtil::HumanString(subshape()); + return buffer(); } -template <> -void CopyFromRepeatedField, - complex64>( - std::vector* dest, - const tensorflow::protobuf::RepeatedField& src) { - *dest = std::vector( - reinterpret_cast(src.data()), - reinterpret_cast(src.data()) + src.size() / 2); +void* Literal::Piece::untyped_data() { + CHECK(ShapeUtil::IsArray(subshape())) << ShapeUtil::HumanString(subshape()); + return buffer(); } -void Literal::CopyFromProto(const LiteralProto& literal_proto) { - if (!literal_proto.has_shape()) { - return; +namespace { + +template +Status CopyFromRepeatedField(tensorflow::gtl::MutableArraySlice dest, + const RepeatedFieldT& src) { + if (dest.size() != src.size()) { + return InvalidArgument( + "Expected %lu elements in LiteralProto repeated field, has %d", + dest.size(), src.size()); } + std::copy(src.begin(), src.end(), dest.begin()); + return Status::OK(); +} - *mutable_shape() = literal_proto.shape(); - switch (shape().element_type()) { +} // namespace + +Status Literal::Piece::CopyFromProto(const LiteralProto& proto) { + // These conditions should have been checked in Literal::CreateFromProto. + TF_RET_CHECK(proto.has_shape()); + TF_RET_CHECK(LayoutUtil::HasLayout(proto.shape())); + TF_RET_CHECK(ShapeUtil::Equal(proto.shape(), subshape())); + + switch (subshape().element_type()) { case PRED: - CopyFromRepeatedField(mutable_preds(), literal_proto.preds()); - break; - case U8: - set_u8s(literal_proto.u8s()); + TF_RETURN_IF_ERROR(CopyFromRepeatedField(data(), proto.preds())); break; + case U8: { + auto u8_data = data(); + TF_RET_CHECK(proto.u8s().size() == u8_data.size()); + std::copy(proto.u8s().begin(), proto.u8s().end(), u8_data.begin()); + } break; case S32: - CopyFromRepeatedField(mutable_s32s(), literal_proto.s32s()); + TF_RETURN_IF_ERROR(CopyFromRepeatedField(data(), proto.s32s())); break; case S64: - CopyFromRepeatedField(mutable_s64s(), literal_proto.s64s()); + TF_RETURN_IF_ERROR(CopyFromRepeatedField(data(), proto.s64s())); break; case U32: - CopyFromRepeatedField(mutable_u32s(), literal_proto.u32s()); + TF_RETURN_IF_ERROR(CopyFromRepeatedField(data(), proto.u32s())); break; case U64: - CopyFromRepeatedField(mutable_u64s(), literal_proto.u64s()); + TF_RETURN_IF_ERROR(CopyFromRepeatedField(data(), proto.u64s())); break; case F16: { - const string& s(literal_proto.f16s()); - CHECK_EQ(0, s.size() % sizeof(half)); - f16s_ = std::vector(s.size() / sizeof(half)); - memcpy(f16s_.data(), s.data(), s.size()); - + const string& s(proto.f16s()); + TF_RET_CHECK(data().size() * sizeof(half) == s.size()); + memcpy(untyped_data(), s.data(), s.size()); if (!kLittleEndian) { - ConvertEndianShort(reinterpret_cast(f16s_.data()), s.size()); + ConvertEndianShort(reinterpret_cast(untyped_data()), s.size()); } - break; - } - case BF16: { - const string& s(literal_proto.bf16s()); - CHECK_EQ(0, s.size() % sizeof(bfloat16)); - bf16s_ = std::vector(s.size() / sizeof(bfloat16)); - memcpy(bf16s_.data(), s.data(), s.size()); + } break; + case BF16: { + const string& s(proto.bf16s()); + TF_RET_CHECK(data().size() * sizeof(bfloat16) == s.size()); + memcpy(untyped_data(), s.data(), s.size()); if (!kLittleEndian) { - ConvertEndianShort(reinterpret_cast(bf16s_.data()), s.size()); + ConvertEndianShort(reinterpret_cast(untyped_data()), s.size()); } - break; - } + } break; case F32: - CopyFromRepeatedField(mutable_f32s(), literal_proto.f32s()); + TF_RETURN_IF_ERROR(CopyFromRepeatedField(data(), proto.f32s())); break; case F64: - CopyFromRepeatedField(mutable_f64s(), literal_proto.f64s()); - break; - case C64: - CopyFromRepeatedField(mutable_c64s(), literal_proto.c64s()); + TF_RETURN_IF_ERROR(CopyFromRepeatedField(data(), proto.f64s())); break; - case TUPLE: - for (const auto& proto : literal_proto.tuple_literals()) { - mutable_tuple_literals()->push_back(Literal(proto)); + case C64: { + auto complex_data = data(); + TF_RET_CHECK(proto.c64s_size() == complex_data.size() * 2); + for (int64 i = 0; i < complex_data.size(); ++i) { + complex_data[i] = complex64{proto.c64s(i * 2), proto.c64s(i * 2 + 1)}; } + } break; + case TUPLE: + LOG(FATAL) << "Should not be called on tuple shapes: " + << ShapeUtil::HumanString(subshape()); break; default: - LOG(FATAL) << "Unhandled primitive type " << shape().element_type(); + LOG(FATAL) << "Unhandled primitive type " << subshape().element_type(); } + return Status::OK(); } -const Literal& Literal::GetSubliteral(const ShapeIndex& index) const { - return const_cast(this)->GetSubliteral(index); +LiteralProto Literal::ToProto() const { + LiteralProto proto; + for (const auto& pair : pieces_) { + const ShapeIndex& index = pair.first; + const Piece& piece = pair.second; + + LiteralProto* proto_piece = &proto; + for (int64 i : index) { + while (proto_piece->tuple_literals_size() <= i) { + proto_piece->add_tuple_literals(); + } + proto_piece = proto_piece->mutable_tuple_literals(i); + } + piece.WriteToProto(proto_piece); + } + + if (LayoutUtil::IsSparseArray(shape())) { + CopyToRepeatedField(proto.mutable_sparse_indices(), + sparse_indices()->data()); + } + + return proto; +} + +/* static */ +StatusOr> Literal::CreateFromProto( + const LiteralProto& proto) { + if (!proto.has_shape()) { + return InvalidArgument("LiteralProto has no shape"); + } + if (!LayoutUtil::HasLayout(proto.shape())) { + return InvalidArgument("LiteralProto has no layout"); + } + + auto literal = MakeUnique(proto.shape()); + + for (auto& pair : literal->pieces_) { + const ShapeIndex& index = pair.first; + Piece& piece = pair.second; + const LiteralProto* proto_element = &proto; + for (int64 i : index) { + TF_RET_CHECK(i < proto_element->tuple_literals_size()); + proto_element = &proto_element->tuple_literals(i); + } + + if (ShapeUtil::IsTuple(piece.subshape())) { + if (proto_element->tuple_literals_size() != + ShapeUtil::TupleElementCount(piece.subshape())) { + return InvalidArgument( + "Expected %lld tuple elements in LiteralProto, has %d", + ShapeUtil::TupleElementCount(piece.subshape()), + proto_element->tuple_literals_size()); + } + continue; + } + + TF_RET_CHECK(ShapeUtil::IsArray(piece.subshape())); + TF_RETURN_IF_ERROR(piece.CopyFromProto(*proto_element)); + } + return std::move(literal); +} + +const void* Literal::untyped_data(const ShapeIndex& shape_index) const { + return piece(shape_index).untyped_data(); +} + +void* Literal::untyped_data(const ShapeIndex& shape_index) { + return piece(shape_index).untyped_data(); +} + +int64 Literal::size_bytes(const ShapeIndex& shape_index) const { + return piece(shape_index).size_bytes(); +} + +string Literal::GetR1U8AsString() const { + CHECK(ShapeUtil::IsArray(shape())); + CHECK_EQ(ShapeUtil::Rank(shape()), 1); + CHECK_EQ(shape().element_type(), U8); + return string(tensorflow::bit_cast(data().data()), + ShapeUtil::ElementsIn(shape())); +} + +/* static */ const LiteralView LiteralView::Create( + const Literal& literal, const ShapeIndex& view_root) { + return LiteralView(literal, view_root); +} + +LiteralView::LiteralView(const Literal& literal, const ShapeIndex& view_root) { + shape_ = ShapeUtil::GetSubshape(literal.shape(), view_root); + pieces_ = ShapeTree(shape_); + owns_buffers_ = false; + for (auto& pair : pieces_) { + const ShapeIndex& index = pair.first; + Piece& piece = pair.second; + + ShapeIndex src_index = view_root; + for (int64 i : index) { + src_index.push_back(i); + } + const Piece& src_piece = literal.piece(src_index); + piece.set_buffer(src_piece.buffer()); + piece.set_sparse_indices(src_piece.sparse_indices()); + piece.set_subshape(&ShapeUtil::GetSubshape(shape_, index)); + } +} + +LiteralView::~LiteralView() {} + +LiteralView::LiteralView(const LiteralView& other) { CopyFrom(other); } + +LiteralView& LiteralView::operator=(const LiteralView& other) { + CopyFrom(other); + return *this; } -Literal& Literal::GetSubliteral(const ShapeIndex& index) { - Literal* subliteral = this; - for (int64 i : index) { - subliteral = &subliteral->tuple_literals_.at(i); +void LiteralView::CopyFrom(const LiteralView& other) { + // We can't use the default copy-constructor/copy-assignment because + // Piece::subshape_ points to subshapes within the Shape of the owning + // Literal/LiteralView. + shape_ = other.shape(); + pieces_ = other.pieces_; + for (auto& pair : pieces_) { + const ShapeIndex& index = pair.first; + Piece& piece = pair.second; + piece.set_subshape(&ShapeUtil::GetSubshape(shape_, index)); } - return *subliteral; + owns_buffers_ = false; } } // namespace xla diff --git a/tensorflow/compiler/xla/literal_util.h b/tensorflow/compiler/xla/literal_util.h index 6254fafaf3ee2867b50ff85c64f456e9f0ed7654..d996004888ab521790b4c5a10da2a93f0d98d12f 100644 --- a/tensorflow/compiler/xla/literal_util.h +++ b/tensorflow/compiler/xla/literal_util.h @@ -34,7 +34,9 @@ limitations under the License. #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/shape_tree.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/sparse_index_array.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" @@ -50,153 +52,70 @@ limitations under the License. namespace xla { -// Utility class for dealing with XLA literal values. Most methods are -// templated by native (host) type which corresponds to a unique XLA -// PrimitiveType. See ComputationBuilder for details. Not all primitive types -// defined in xla_data.proto have a corresponding native type or even have a -// storage location in the Literal proto yet (for example, primitive type F16). +// Class representing literal values in XLA. +// +// TODO(b/67651157): The methods in this class should be reduced to a minimal +// set of methods which construct Literals and accessors methods. Other methods +// which perform computation on Literals (Reshape, Slice, etc) should be moved +// elsewhere, and perhaps combined with evaluator code which operates on +// Literals. class Literal { public: - Literal() {} + Literal() : Literal(ShapeUtil::MakeNil()) {} - Literal(const Literal& other) = default; - Literal(Literal&&) = default; + // Create a literal of the given shape. The literal is allocated sufficient + // memory to hold the shape. Memory is uninitialized. + explicit Literal(const Shape& shape); + virtual ~Literal(); - explicit Literal(const LiteralProto& other) { CopyFromProto(other); } - - Literal& operator=(const Literal& other) = default; - Literal& operator=(Literal&&) = default; + // Literals are moveable, but not copyable. To copy a literal use + // Literal::Clone or Literal::CloneToUnique. This prevents inadvertent copies + // of literals which can be expensive. + Literal(const Literal& other) = delete; + Literal& operator=(const Literal& other) = delete; + Literal(Literal&& other); + Literal& operator=(Literal&& other); // Literals are equal if they have compatible shapes and the same data - // values. Layout is not checked. + // values. Layout is not compared. bool operator==(const Literal& other) const; bool operator!=(const Literal& other) const { return !(*this == other); } + // Serialize to and from a proto. + static StatusOr> CreateFromProto( + const LiteralProto& proto); LiteralProto ToProto() const; - bool has_shape() const { - return shape_.element_type() != PRIMITIVE_TYPE_INVALID; - } - - // Basic accessor functions. Names mirror the original protobuf - // functions for convenience. - string DebugString() const { return ToProto().DebugString(); } - string ShortDebugString() const { return ToProto().ShortDebugString(); } - - // Return the nested literal at the given shape index. - const Literal& GetSubliteral(const ShapeIndex& index) const; - Literal& GetSubliteral(const ShapeIndex& index); - - void Clear() { - shape_.Clear(); - u8s_.clear(); - s16s_.clear(); - s32s_.clear(); - s64s_.clear(); - u16s_.clear(); - u32s_.clear(); - u64s_.clear(); - f16s_.clear(); - f32s_.clear(); - f64s_.clear(); - c64s_.clear(); - tuple_literals_.clear(); - } - - int preds_size() const { return u8s().size(); } - const std::vector& preds() const { - static_assert(sizeof(uint8) == sizeof(bool), - "The uint8 and bool types should be the same size"); - return u8s_; - } - std::vector* mutable_preds() { - static_assert(sizeof(uint8) == sizeof(bool), - "The uint8 and bool types should be the same size"); - return &u8s_; - } - - int s16s_size() const { return s16s().size(); } - int32 s16s(int i) const { return s16s_[i]; } - const std::vector& s16s() const { return s16s_; } - std::vector* mutable_s16s() { return &s16s_; } - - int s32s_size() const { return s32s().size(); } - int32 s32s(int i) const { return s32s_[i]; } - const std::vector& s32s() const { return s32s_; } - std::vector* mutable_s32s() { return &s32s_; } - - int s64s_size() const { return s64s().size(); } - void add_s64s(int64 value) { s64s_.push_back(value); } - const std::vector& s64s() const { return s64s_; } - std::vector* mutable_s64s() { return &s64s_; } - - int u16s_size() const { return u16s().size(); } - uint32 u16s(int i) const { return u16s_[i]; } - const std::vector& u16s() const { return u16s_; } - std::vector* mutable_u16s() { return &u16s_; } - - int u32s_size() const { return u32s().size(); } - uint32 u32s(int i) const { return u32s_[i]; } - const std::vector& u32s() const { return u32s_; } - std::vector* mutable_u32s() { return &u32s_; } - - int u64s_size() const { return u64s().size(); } - const std::vector& u64s() const { return u64s_; } - std::vector* mutable_u64s() { return &u64s_; } - - int f16s_size() const { return f16s().size(); } - half f16s(int i) const { return f16s_[i]; } - const std::vector& f16s() const { return f16s_; } - std::vector* mutable_f16s() { return &f16s_; } - - int f32s_size() const { return f32s().size(); } - float f32s(int i) const { return f32s_[i]; } - void add_f32s(float value) { f32s_.push_back(value); } - const std::vector& f32s() const { return f32s_; } - std::vector& f32s() { return f32s_; } - std::vector* mutable_f32s() { return &f32s_; } - - int f64s_size() const { return f64s().size(); } - const std::vector& f64s() const { return f64s_; } - std::vector* mutable_f64s() { return &f64s_; } - - int c64s_size() const { return c64s().size(); } - const std::vector& c64s() const { return c64s_; } - std::vector* mutable_c64s() { return &c64s_; } - - int bf16s_size() const { return bf16s().size(); } - bfloat16 bf16s(int i) const { return bf16s_[i]; } - const std::vector& bf16s() const { return bf16s_; } - std::vector* mutable_bf16s() { return &bf16s_; } - - int tuple_literals_size() const { return tuple_literals().size(); } - const Literal& tuple_literals(int i) const { return tuple_literals_[i]; } - Literal* add_tuple_literals() { - tuple_literals_.push_back(Literal()); - return &tuple_literals_.back(); - } - std::vector* mutable_tuple_literals() { return &tuple_literals_; } - const std::vector& tuple_literals() const { return tuple_literals_; } - - int u8s_size() const { return u8s().size(); } - const std::vector& u8s() const { return u8s_; } - void set_u8s(const std::vector& value) { u8s_ = value; } - void set_u8s(tensorflow::StringPiece value) { - u8s_ = std::vector(value.size()); - u8s_.clear(); - append_u8s(value); - } - - void append_u8s(tensorflow::StringPiece value) { - u8s_.insert(u8s_.end(), value.begin(), value.end()); - } - - string u8s_string() const { return string(u8s().begin(), u8s().end()); } + // Return the shape of the literal. + const Shape& shape() const { return shape_; } - std::vector* mutable_u8s() { return &u8s_; } + // TODO(b/67651157): Remove this accessor. Literal users should not be able to + // mutate the shape as this can produce malformed Literals. + Shape* mutable_shape_do_not_use() { return &shape_; } - const Shape& shape() const { return shape_; } - Shape* mutable_shape() { return &shape_; } + // Returns a (Mutable)ArraySlice view of the array for this literal for the + // given NativeT (e.g., float). CHECKs if the subshape of the literal at the + // given ShapeIndex is not array. See primitive_util.h for the mapping from + // XLA type to native type. + template + tensorflow::gtl::ArraySlice data( + const ShapeIndex& shape_index = {}) const; + template + tensorflow::gtl::MutableArraySlice data( + const ShapeIndex& shape_index = {}); + + // Returns a pointer to the sparse index array. Returns nullptr if the literal + // is not a sparse array. + const SparseIndexArray* sparse_indices( + const ShapeIndex& shape_index = {}) const; + SparseIndexArray* sparse_indices(const ShapeIndex& shape_index = {}); + + // Returns a pointer to (or size of) the underlying buffer holding the array + // at the given shape index. CHECKs if the subshape of the literal at the + // given ShapeIndex is not array. + const void* untyped_data(const ShapeIndex& shape_index = {}) const; + void* untyped_data(const ShapeIndex& shape_index = {}); + int64 size_bytes(const ShapeIndex& shape_index = {}) const; // Creates a new literal of a given rank. To minimize ambiguity (for users // and the compiler) these CreateR[0-2] methods should explicitly specify the @@ -244,6 +163,60 @@ class Literal { values, const Layout& layout); + // Returns this literal's data as a string. This literal must be a rank-1 U8 + // array. + string GetR1U8AsString() const; + + // Creates a literal with a sparse layout and the given indices and values. + // The shape is initialized from the given dimensions. The minor dimension of + // the indices array must equal the rank of the shape (i.e. size of the + // dimensions array). The major dimension of the indices array must equal the + // number of elements in the values array. The maximum number of elements in + // the array is taken from the max_indices() value of the index array. + // + // XLA assumes that sparse literals are in sorted order for all operations. If + // the `sort` argument is true, then the indices and values will be sorted + // while copying them into the literal. If you have ensured that the indices + // and values are already sorted, then you may set the `sort` argument to + // false to skip the sorting step. + // + // For example: + // + // CreateSparse( + // {12, 12, 12}, + // SparseIndexArray(10, 3, + // Array2D{ + // {0, 1, 2}, + // {3, 4, 5}, + // {6, 7, 8}, + // {9, 10, 11}, + // }), + // {1.0, 2.0 3.0, 4.0}) + // + // This creates an array with shape F64[12,12,12]sparse{10}, that has the + // following non-zero values: + // + // [0, 1, 2]: 1.0 + // [3, 4, 5]: 2.0 + // [6, 7, 8]: 3.0 + // [9, 10, 11]: 4.0 + // + template + static std::unique_ptr CreateSparse( + tensorflow::gtl::ArraySlice dimensions, SparseIndexArray indices, + tensorflow::gtl::ArraySlice values, bool sort = true); + + // Populates a literal with a sparse layout with the given indices and values. + // Each index in the indices array is CHECKed against the dimensions in the + // literal's shape. If sort is true, then the indices and values will be + // sorted. If sort is false, then the indices and values are assumed to + // already be in sorted order. See CreateSparse for an example of how data + // are populated. + template + void PopulateSparse(SparseIndexArray indices, + tensorflow::gtl::ArraySlice values, + bool sort = true); + // Creates a new Literal object with the shape specified as parameter. // The content of the literal values is the default value of the primitive // type of literal itself (0 for numeric types, and false for predicates). @@ -257,6 +230,23 @@ class Literal { PrimitiveType primitive_type, tensorflow::gtl::ArraySlice dimensions); + // Copy values from 'src_literal' rooted at 'src_shape_index' into this + // literal rooted at 'dest_shape_index'. The subshape of this literal rooted + // at 'dest_shape_index' must be compatible with the subshape of 'src_literal' + // rooted at 'src_shape_index', but need not be arrays. + Status CopyFrom(const Literal& src_literal, + const ShapeIndex& dest_shape_index = {}, + const ShapeIndex& src_shape_index = {}); + + // Similar to CopyFrom, but with move semantincs. The subshape of this literal + // rooted at 'dest_shape_index' must be *equal* to the shape 'src_literal' + // (layouts and shapes must match), but need not be arrays. The memory + // allocated in this literal for the subshape at dest_shape_index is + // deallocated, and the respective buffers are replaced with those in + // src_literal. Upon return, src_literal is set to a nil shape (empty tuple). + Status MoveFrom(Literal&& src_literal, + const ShapeIndex& dest_shape_index = {}); + // Copies the values from src_literal, starting at src_base shape indexes, // to this literal, starting at dest_base, where the copy size in each // dimension is specified by copy_size. @@ -266,10 +256,24 @@ class Literal { // Note: if either src_literal or this literal contains dimensions with zero // element, then copy_size must be 0 in these dimensions while the // corresponding base indices being 0. - Status Copy(const Literal& src_literal, - tensorflow::gtl::ArraySlice src_base, - tensorflow::gtl::ArraySlice dest_base, - tensorflow::gtl::ArraySlice copy_size); + // This literal and 'src_literal' must be arrays. + Status CopySliceFrom(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_base, + tensorflow::gtl::ArraySlice dest_base, + tensorflow::gtl::ArraySlice copy_size); + + // Returns a vector containing the tuple elements of this Literal as separate + // Literals. This Literal must be tuple-shaped and can be a nested tuple. The + // elements are moved into the new Literals; no data is copied. Upon return + // this Literal is set to a nil shape (empty tuple) + std::vector DecomposeTuple(); + + // This operation is the inverse of DecomposeTuple. The given elements are + // moved into the tuple elements of a new tuple-shaped Literal which is + // returned. Upon return, each of the Literals in 'elements' is set to a nil + // shape (empty tuple). + static Literal MoveIntoTuple( + tensorflow::gtl::MutableArraySlice elements); // Creates a new value that has the equivalent value as this literal, but // conforms to new_layout; e.g. a literal matrix that was in {0, 1} @@ -293,6 +297,7 @@ class Literal { // Creates a new literal by reshaping this literal to have the given // dimensions. The total number of elements must not change; The // implementation currently only supports monotonic dim0-major layouts. + // This literal must be an array. StatusOr> Reshape( tensorflow::gtl::ArraySlice dimensions) const; @@ -302,6 +307,7 @@ class Literal { // in the result literal (i.e., new_order[i] = old_order[permutation[i]]). // For example, a transpose call on a literal of shape [3 x 8 x 4] and // `permutation` = {2, 0, 1} returns a new literal of shape [4 x 3 x 8]. + // This literal must be an array. std::unique_ptr Transpose( tensorflow::gtl::ArraySlice permutation) const; @@ -310,6 +316,7 @@ class Literal { // same rank and layout as for the given literal. The number of indices in // start_indices and limit_indices must be the rank of the literal, and the // indices follow the order of the dimensions. + // This literal must be an array. std::unique_ptr Slice( tensorflow::gtl::ArraySlice start_indices, tensorflow::gtl::ArraySlice limit_indices) const; @@ -317,25 +324,26 @@ class Literal { // Creates a literal with a prepended dimension with bound "times"; e.g. a // f32[3x2] with times=4 will produce a f32[4x3x2] with the 3x2 from this // literal replicated four times. + // This literal must be an array. template std::unique_ptr Replicate(int64 times) const; // Converts this literal to another primitive type. Returns an error if the - // conversion is not possible. + // conversion is not possible. This literal must be array-shaped. StatusOr> Convert( PrimitiveType primitive_dest_type) const; - // Creates a literal value zero of the given primitive type. + // Creates a scalar literal value zero of the given primitive type. static Literal Zero(PrimitiveType primitive_type); - // Creates a literal value one of the given primitive type. + // Creates a scalar literal value one of the given primitive type. static Literal One(PrimitiveType primitive_type); - // Creates a literal value containing the minimum value of the given + // Creates a scalar literal value containing the minimum value of the given // primitive type. For floating-point types, returns -inf. static Literal MinValue(PrimitiveType primitive_type); - // Creates a literal value containing the maximum value of the given + // Creates a scalar literal value containing the maximum value of the given // primitive type. For floating-point types, returns inf. static Literal MaxValue(PrimitiveType primitive_type); @@ -344,7 +352,7 @@ class Literal { static std::unique_ptr CreateFullWithDescendingLayout( tensorflow::gtl::ArraySlice dimensions, NativeT value); - // Creates a new literal from an array. The variants not ending with + // Creates a new literal from an Array type. The variants not ending with // WithLayout use the default XLA layout for the literal's linear // representation in memory. template @@ -393,35 +401,50 @@ class Literal { std::initializer_list> values, int64 projection_p, int64 projection_z); - // Clones this literal into an owned unique_ptr version. + // Clones this literal into a new Literal, or new std::unique_ptr. + Literal Clone() const; std::unique_ptr CloneToUnique() const; - // Returns the linear index of the given index within this literal's - // element_type repeated field. - int64 LinearIndex(tensorflow::gtl::ArraySlice multi_index) const; + // Gets or sets an element in the literal at the given index. The multi_index + // is CHECKed against the dimension sizes. + template + NativeT Get(tensorflow::gtl::ArraySlice multi_index, + const ShapeIndex& shape_index) const; + template + void Set(tensorflow::gtl::ArraySlice multi_index, + const ShapeIndex& shape_index, NativeT value); - // Gets or sets an element in the literal at the given index. The index is - // CHECKed against the dimension sizes. + // Overloads of Get and Set for array literals. CHECKs if the literal is not + // array-shaped and dense. template NativeT Get(tensorflow::gtl::ArraySlice multi_index) const; template void Set(tensorflow::gtl::ArraySlice multi_index, NativeT value); - // Returns a (Mutable)ArraySlice view of the array for this literal for the - // given NativeT (e.g., float). These functions map native type to XLA - // PrimitiveType via template specialization. The unspecialized forms below - // aborts to handle the error case where the given native type does not map to - // an XLA primitive type. + // Returns the multi-index of the element in a sparse literal at the given + // sparse element number. The sparse element number is the position with in + // the sparse array's list of (index, value) pairs, and is checked against the + // total number of (index, value) pairs in the sparse array. + tensorflow::gtl::ArraySlice GetSparseIndex( + int64 sparse_element_number, const ShapeIndex& shape_index = {}) const; + + // Returns the value of the element in a sparse literal at the given sparse + // element number. The sparse element number is the position with in the + // sparse array's list of (index, value) pairs, and is checked against the + // total number of (index, value) pairs in the sparse array. template - tensorflow::gtl::ArraySlice GetArraySlice() const { - static_assert(!std::is_same::value, - "Cannot map native type to primitive type."); - } + NativeT GetSparseElement(int64 sparse_element_number, + const ShapeIndex& shape_index = {}) const; + + // Appends the given element to the literal. If the elements are not appended + // in sorted order, then SortSparseElements should be called before calling + // other methods. This literal must have a sparse layout. template - tensorflow::gtl::MutableArraySlice GetMutableArraySlice() { - static_assert(!std::is_same::value, - "Cannot map native type to primitive type."); - } + void AppendSparseElement(tensorflow::gtl::ArraySlice multi_index, + NativeT value, const ShapeIndex& shape_index = {}); + + // Sorts the elements in a sparse array. + void SortSparseElements(const ShapeIndex& shape_index = {}); // Returns the element value at index (0, ..., 0), however many zeroes are // required for that index. @@ -430,10 +453,16 @@ class Literal { // As Get(), but determines the correct type and converts the value // into text. - string GetAsString(tensorflow::gtl::ArraySlice multi_index) const; + string GetAsString(tensorflow::gtl::ArraySlice multi_index, + const ShapeIndex& shape_index = {}) const; + + // As GetSparseElement(), but determines the correct type and converts the + // value into text. + string GetSparseElementAsString(int64 sparse_element_number, + const ShapeIndex& shape_index = {}) const; // As Get(), but determines the correct type and converts the value into - // int64. + // int64. This literal must be an array. StatusOr GetIntegralAsS64( tensorflow::gtl::ArraySlice multi_index) const; @@ -441,7 +470,8 @@ class Literal { template static std::unique_ptr MakeIdentityR2(int64 size); - // Returns a tuple literal composed of given literals. + // Returns a tuple literal composed of given literals. Data is copied from the + // given elements into the returned literal. static std::unique_ptr MakeTuple( tensorflow::gtl::ArraySlice elements); @@ -455,11 +485,29 @@ class Literal { static std::unique_ptr MakeTupleOwned( std::vector> elements); - // Validates that the data payload of the literal matches the literal shape; - // if it does not, an appropriate status is returned. - tensorflow::Status ValidateLiteral() const; + // This overload lets you pass a braced list of unique_ptrs to + // MakeTupleOwned: + // + // Literal::MakeTupleOwned(Literal::CreateR1(...), ...). + // + // Simply relying on the MakeTupleOwned(std::vector>) + // overload doesn't work because std::initializer_list's elements are always + // const. + // + // The arguments to this function must all be unique_ptr. + template + static std::unique_ptr MakeTupleOwned( + std::unique_ptr... elements) { + std::array, sizeof...(Ts)> arr{ + std::move(elements)...}; + std::vector> v; + v.insert(v.begin(), std::make_move_iterator(arr.begin()), + std::make_move_iterator(arr.end())); + return MakeTupleOwned(std::move(v)); + } // Returns a string representation of the literal value. + // Warning: this function can take minutes for multi-million element Literals. string ToString(bool print_layout = false) const; // Invokes the "per cell" callback for each element in the provided @@ -469,6 +517,8 @@ class Literal { // This function is useful if you want a polymorphic representation // of the tensor's elements (turning it to a string for something // like representation in a protobuf). + // + // This literal must have a dense layout. void EachCellAsString( const std::function indices, const string& value)>& per_cell) const; @@ -477,80 +527,45 @@ class Literal { NativeT value)> per_cell) const; - // Templated methods which populate the given repeated field in this literal - // with the given value(s). The Shape field of this literal is set - // to match the array dimensions and type. Examples: + // Populate this literal with the given values. Examples: // // // Populate with floats. // Array2D float_values = ... // literal.PopulateR2FromArray2D(values); // // // Populate with int32s. - // literal.PopulateR2({{1, 2}, {3, 4}}); + // literal.PopulateR2({{1, 2}, {3, 4}}); // - template - void PopulateR0(NativeT values); + // The shape and element type of this literal must match given values. For + // example, in the call above to literal.PopulateR2(), 'literal' must be a 2x2 + // array of S32. template void PopulateR1(tensorflow::gtl::ArraySlice values); void PopulateR1(const tensorflow::core::Bitmap& values); template void PopulateR2(std::initializer_list> values); template - void PopulateR2WithLayout( - std::initializer_list> values, - const Layout& layout); - template void PopulateFromArray(const Array& values); template - void PopulateFromArrayWithLayout(const Array& values, - const Layout& layout); - template void PopulateR2FromArray2D(const Array2D& values); template - void PopulateR2FromArray2DWithLayout(const Array2D& values, - const Layout& layout); - template void PopulateR3FromArray3D(const Array3D& values); template - void PopulateR3FromArray3DWithLayout(const Array3D& values, - const Layout& layout); - template void PopulateR4FromArray4D(const Array4D& values); - template - void PopulateR4FromArray4DWithLayout(const Array4D& values, - const Layout& layout); // Populates literal values by calling the generator function for every cell // in this literal object. // // generator must be a callable of the type // NativeT(tensorflow::gtl::ArraySlice indexes) or compatible. + // + // This literal must have a dense layout. template Status Populate(const FnType& generator); - // Creates a Literal of the given dimensions with all elements set to the - // given value. - template - void PopulateWithValue(NativeT value, - tensorflow::gtl::ArraySlice dimensions); - - // Returns a pointer to the underlying vector corresponding to the Literal's - // shape. - const void* InternalData() const; - void* MutableInternalData(); - - // Allocates space in the underlying vector of this literal sufficient to hold - // num_elements of this literal's primitive type. Values in the vector are set - // to zero. num_elements must equal the number of elements in the literal's - // shape. - void Reserve(int64 num_elements); - - // Allocates space in the underlying vector of this literal sufficient to hold - // num_elements of this literal's primitive type and sets each element in this - // literal to the given value. num_elements must equal the number of elements - // in this literal's shape. + // Fills this literal with the given value. template - void Resize(int64 num_elements, NativeT value); + void PopulateWithValue(NativeT value); // Returns whether every element in this literal is equal to value. // @@ -560,7 +575,7 @@ class Literal { // // If value doesn't fit in this literal's type, returns false. Values of 1/0 // are considered equal to true/false; other values are not considered equal - // to true. + // to true. Also if this literal is not array-shaped false is returned. bool IsAll(int8 value) const; // Like IsAll(const Literal&, int8), except we check whether the literal is @@ -571,7 +586,7 @@ class Literal { // This casts value to the type of literal, then compares using ==. The usual // admonishments about floating-point equality checks apply. We expect you to // use this to check for values that can be expressed precisely as a float, - // e.g. -0.5. + // e.g. -0.5. Also if this literal is not array-shaped false is returned. bool IsAllFloat(float value) const; // Like IsAll(const Literal&, int8), except we check whether the literal is @@ -583,23 +598,38 @@ class Literal { // admonishments about floating-point equality checks apply. We expect you to // use this to check for complex values that can be expressed precisely as // float pairs e.g. (-0.5, 1.0). + // + // This literal must have a dense layout. bool IsAllComplex(complex64 value) const; // Returns whether this literal is zero at the specified index. This literal - // must be an array. + // must be an array with a dense layout. bool IsZero(tensorflow::gtl::ArraySlice indices) const; - private: - // Copy from a LiteralProto instance. - void CopyFromProto(const LiteralProto& literal_proto); + // Return the count of the elements in the array at the given shape index in + // this literal. + int64 element_count(const ShapeIndex& index = {}) const { + return ShapeUtil::ElementsIn(ShapeUtil::GetSubshape(shape(), index)); + } - // Internal template helper for the Copy() API, matching its arguments one by - // one. - template - Status CopyRange(const Literal& src_literal, - tensorflow::gtl::ArraySlice src_base, - tensorflow::gtl::ArraySlice dest_base, - tensorflow::gtl::ArraySlice copy_size); + // Return the count of the elements in the sparse array at the given shape + // index in this literal, which will be no larger than + // LayoutUtil::MaxSparseElements(SetSubshape(shape(), index).layout()). + int64 sparse_element_count() const; + + protected: + // 'allocate_arrays' indicates whether to allocate memory for the arrays in + // the shape. If false, buffer pointers inside of the Literal::Pieces are set + // to nullptr. + Literal(const Shape& shape, bool allocate_arrays); + + // Internal template helper for the Literal::CopySliceFrom(), matching its + // arguments one by one. + template + Status CopySliceFromInternal(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_base, + tensorflow::gtl::ArraySlice dest_base, + tensorflow::gtl::ArraySlice copy_size); // Utility structure which is used to create the optimal configuration for // a ShapeUtil::ForEachIndex() scan across two literals. @@ -624,163 +654,243 @@ class Literal { int64 minor_loop_size = 1; }; - Shape shape_; - std::vector u8s_; - std::vector s16s_; - std::vector s32s_; - std::vector s64s_; - std::vector u16s_; - std::vector u32s_; - std::vector u64s_; - std::vector bf16s_; - std::vector f16s_; - std::vector f32s_; - std::vector f64s_; - std::vector c64s_; - std::vector tuple_literals_; -}; - -std::ostream& operator<<(std::ostream& out, const Literal& literal); - -// Declarations of template specializations for GetArraySlice and -// GetMutableArraySlice. The specializations map native type to XLA primitive -// type. -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; - -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; + // A data structure representing a subshape at a particular ShapeIndex within + // the literal. For array-shaped ShapeIndexes, this data structure holds the + // pointer to the memory allocated for the array data. + class Piece { + public: + // Return the buffer holding the array data for this piece as an array + // slice. This piece must be array-shaped. + template + tensorflow::gtl::ArraySlice data() const; + template + tensorflow::gtl::MutableArraySlice data(); + + // Return the buffer holding the array data for this piece as a void*. This + // piece must be array-shaped. + void* untyped_data(); + const void* untyped_data() const; + + // Gets or sets an element in the array at the given index. The multi_index + // is CHECKed against the dimension sizes of the array. This piece must be + // array-shaped. + template + NativeT Get(tensorflow::gtl::ArraySlice index) const; + template + void Set(tensorflow::gtl::ArraySlice index, NativeT value); + + // Gets/sets the buffer holding the array data. + char* buffer() const { return buffer_; } + void set_buffer(char* buffer) { buffer_ = buffer; } + + // The array of multi-indices that provide the locations of non-zero + // elements in a sparse array. Only used if + // LayoutUtil::IsSparseArray(shape()) is true. + SparseIndexArray* sparse_indices() const { return sparse_indices_; } + void set_sparse_indices(SparseIndexArray* sparse_indices) { + sparse_indices_ = sparse_indices; + } -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; + // Gets or sets the subshape of this piece. This reference points to a + // subshape within the shape in the containing Literal (Literal::shape_). + const Shape& subshape() const { return *subshape_; } + void set_subshape(const Shape* subshape) { subshape_ = subshape; } -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; + // Returns the size in bytes of the buffer holding the array data. + int64 size_bytes() const { return ShapeUtil::ByteSizeOf(subshape()); } -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; + // Returns the number of elements in this piece's array. + int64 element_count() const { return ShapeUtil::ElementsIn(subshape()); } -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; + // Copy the data from 'src' into this piece's buffer. Shapes of this piece + // and src must be compatible. + Status CopyFrom(const Piece& src); -template <> -inline tensorflow::gtl::ArraySlice Literal::GetArraySlice() - const { - DCHECK(shape().element_type() == F32); - return f32s(); -} + // Returns true if this piece and 'other' contain the same data. This piece + // and 'other' must be array-shaped and compatible. + bool EqualElements(const Piece& other) const; -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; + // Writes the shape and data (if array-shaped) into the given proto. + void WriteToProto(LiteralProto* proto) const; -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; + // Copies the data from the given proto into this piece. The shape of this + // piece must be equal (not just compatible) to the shape of the proto. + Status CopyFromProto(const LiteralProto& proto); -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; + // Sorts the elements in a sparse array. + void SortSparseElements(); -template <> -tensorflow::gtl::ArraySlice Literal::GetArraySlice() - const; + private: + // Recursive helper for EqualElements. + template + bool EqualElementsInternal(const Piece& other, + std::vector* multi_index) const; -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // Helper for SortSparseElements that has the element type as a template + // parameter. + template + void SortSparseElementsInternal(); -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // For array-shaped pieces, this is the buffer holding the literal data. + char* buffer_ = nullptr; -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // For sparse arrays, this is the array of indices. + SparseIndexArray* sparse_indices_ = nullptr; -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // The shape of piece. This points into the shape of the containing Literal + // (Literal::shape_). + const Shape* subshape_ = nullptr; + }; -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // Returns the piece at the given ShapeIndex. + Piece& piece(const ShapeIndex& shape_index) { + return *pieces_.mutable_element(shape_index); + } + const Piece& piece(const ShapeIndex& shape_index) const { + return pieces_.element(shape_index); + } -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // Returns the piece at the root of the shape (empty ShapeIndex). + Piece& root_piece() { return piece({}); } + const Piece& root_piece() const { return piece({}); } -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // Deallocate the buffers held by this literal (if the literal owns the + // buffer). + void DeallocateBuffers(); -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + Shape shape_; + ShapeTree pieces_; -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // Whether the buffers held in pieces_ are owned by this Literal. + bool owns_buffers_; -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + // LiteralView must access and manipulate Pieces of other Literals. + friend class LiteralView; +}; // namespace xla -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); +std::ostream& operator<<(std::ostream& out, const Literal& literal); -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); +// A read-only view of a Literal. A LiteralView contains pointers to buffers +// owned by the viewed Literal. +// +// TODO(b/71550060): Replace LiteralView with Literal slice classes (immutable +// and mutable) similar to (Mutable)ArraySlice. +class LiteralView : public Literal { + public: + // Create and return a view of the given literal rooted at the given shape + // index within the given literal. A factory is used rather than a public + // constructor because only const LiteralViews are supported. It's still + // possible to create non-const LiteralViews via the copy constructors, but + // the factory method makes it a bit less likely. Implementing literal slices + // will fix this undesirable situation (b/71550060). + static const LiteralView Create(const Literal& literal, + const ShapeIndex& view_root = {}); -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + LiteralView(const LiteralView& other); + LiteralView& operator=(const LiteralView& other); -template <> -tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + virtual ~LiteralView(); -template <> -void Literal::Resize(int64 num_elements, bool value); + private: + LiteralView(const Literal& literal, const ShapeIndex& view_root); -template <> -void Literal::Resize(int64 num_elements, int8 value); + // Helper for the copy constructor and copy assignment operator. + void CopyFrom(const LiteralView& other); +}; -template <> -void Literal::Resize(int64 num_elements, uint8 value); +template +tensorflow::gtl::ArraySlice Literal::Piece::data() const { + CHECK(ShapeUtil::IsArray(subshape())) << ShapeUtil::HumanString(subshape()); + CHECK_EQ(subshape().element_type(), + primitive_util::NativeToPrimitiveType()) + << "Attempting to access " + << PrimitiveType_Name(primitive_util::NativeToPrimitiveType()) + << " type, but literal element type is " + << PrimitiveType_Name(subshape().element_type()); + return tensorflow::gtl::ArraySlice( + reinterpret_cast(buffer()), + ShapeUtil::ElementsIn(subshape())); +} -template <> -void Literal::Resize(int64 num_elements, int32 value); +template +tensorflow::gtl::MutableArraySlice Literal::Piece::data() { + CHECK(ShapeUtil::IsArray(subshape())) << ShapeUtil::HumanString(subshape()); + CHECK_EQ(subshape().element_type(), + primitive_util::NativeToPrimitiveType()) + << "Attempting to access " + << PrimitiveType_Name(primitive_util::NativeToPrimitiveType()) + << " type, but literal element type is " + << PrimitiveType_Name(subshape().element_type()); + return tensorflow::gtl::MutableArraySlice( + reinterpret_cast(buffer()), ShapeUtil::ElementsIn(subshape())); +} -template <> -void Literal::Resize(int64 num_elements, uint32 value); +template +NativeT Literal::Piece::Get( + tensorflow::gtl::ArraySlice multi_index) const { + CHECK(LayoutUtil::IsDenseArray(subshape())); + return data()[IndexUtil::MultidimensionalIndexToLinearIndex( + subshape(), multi_index)]; +} -template <> -void Literal::Resize(int64 num_elements, int64 value); +template +void Literal::Piece::Set(tensorflow::gtl::ArraySlice multi_index, + NativeT value) { + CHECK(LayoutUtil::IsDenseArray(subshape())); + data()[IndexUtil::MultidimensionalIndexToLinearIndex( + subshape(), multi_index)] = value; +} -template <> -void Literal::Resize(int64 num_elements, uint64 value); +template +tensorflow::gtl::ArraySlice Literal::data( + const ShapeIndex& shape_index) const { + return piece(shape_index).data(); +} -template <> -void Literal::Resize(int64 num_elements, float value); +template +tensorflow::gtl::MutableArraySlice Literal::data( + const ShapeIndex& shape_index) { + return piece(shape_index).data(); +} -template <> -void Literal::Resize(int64 num_elements, double value); +template +inline NativeT Literal::Get(tensorflow::gtl::ArraySlice multi_index, + const ShapeIndex& shape_index) const { + return piece(shape_index).Get(multi_index); +} -template <> -void Literal::Resize(int64 num_elements, half value); +template +inline NativeT Literal::Get( + tensorflow::gtl::ArraySlice multi_index) const { + return root_piece().Get(multi_index); +} -template <> -void Literal::Resize(int64 num_elements, bfloat16 value); +template +inline void Literal::Set(tensorflow::gtl::ArraySlice multi_index, + const ShapeIndex& shape_index, NativeT value) { + return piece(shape_index).Set(multi_index, value); +} -template <> -void Literal::Resize(int64 num_elements, complex64 value); +template +inline void Literal::Set(tensorflow::gtl::ArraySlice multi_index, + NativeT value) { + return root_piece().Set(multi_index, value); +} template /* static */ std::unique_ptr Literal::CreateR0(NativeT value) { - auto literal = MakeUnique(); - literal->PopulateR0(value); + auto literal = MakeUnique(ShapeUtil::MakeShape( + primitive_util::NativeToPrimitiveType(), {})); + literal->Set({}, value); return literal; } template /* static */ std::unique_ptr Literal::CreateR1( tensorflow::gtl::ArraySlice values) { - auto literal = MakeUnique(); + auto literal = MakeUnique( + ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), + {static_cast(values.size())})); literal->PopulateR1(values); return literal; } @@ -789,8 +899,12 @@ template /* static */ std::unique_ptr Literal::CreateR2WithLayout( std::initializer_list> values, const Layout& layout) { - auto literal = MakeUnique(); - literal->PopulateR2WithLayout(values, layout); + auto literal = MakeUnique(ShapeUtil::MakeShapeWithLayout( + primitive_util::NativeToPrimitiveType(), + {static_cast(values.size()), + static_cast(values.begin()->size())}, + AsInt64Slice(layout.minor_to_major()))); + literal->PopulateR2(values); return literal; } @@ -863,6 +977,21 @@ template return CreateR4FromArray4DWithLayout(tmp, layout); } +template +/* static */ std::unique_ptr Literal::CreateSparse( + tensorflow::gtl::ArraySlice dimensions, SparseIndexArray indices, + tensorflow::gtl::ArraySlice values, bool sort) { + int64 num_elements = values.size(); + int64 rank = dimensions.size(); + CHECK_EQ(num_elements, indices.index_count()); + CHECK_EQ(rank, indices.rank()); + auto literal = MakeUnique(ShapeUtil::MakeShapeWithSparseLayout( + primitive_util::NativeToPrimitiveType(), dimensions, + indices.max_indices())); + literal->PopulateSparse(indices, values, sort); + return literal; +} + template /* static */ std::unique_ptr Literal::CreateR4( std::initializer_list template /* static */ std::unique_ptr Literal::CreateFromArrayWithLayout( const Array& values, const Layout& layout) { - auto literal = MakeUnique(); - literal->PopulateFromArrayWithLayout(values, layout); + auto literal = MakeUnique(ShapeUtil::MakeShapeWithLayout( + primitive_util::NativeToPrimitiveType(), values.dimensions(), + AsInt64Slice(layout.minor_to_major()))); + literal->PopulateFromArray(values); return literal; } @@ -975,81 +1106,33 @@ template return CreateFromArrayWithLayout(values, layout); } -template -NativeT Literal::Get(tensorflow::gtl::ArraySlice multi_index) const { - int64 linear_index = LinearIndex(multi_index); - return GetArraySlice().at(linear_index); -} - template NativeT Literal::GetFirstElement() const { - return GetArraySlice().at(0); -} - -template <> -inline uint8 Literal::Get( - tensorflow::gtl::ArraySlice multi_index) const { - CHECK(shape().element_type() == U8); - int64 linear_index = LinearIndex(multi_index); - return u8s()[linear_index]; -} - -template <> -inline int8 Literal::Get( - tensorflow::gtl::ArraySlice multi_index) const { - CHECK(shape().element_type() == S8); - int64 linear_index = LinearIndex(multi_index); - return u8s()[linear_index]; -} - -template <> -inline half Literal::Get( - tensorflow::gtl::ArraySlice multi_index) const { - CHECK(shape().element_type() == F16); - int64 linear_index = LinearIndex(multi_index); - return GetArraySlice()[linear_index]; -} - -template <> -inline bfloat16 Literal::Get( - tensorflow::gtl::ArraySlice multi_index) const { - CHECK(shape().element_type() == BF16); - int64 linear_index = LinearIndex(multi_index); - return GetArraySlice()[linear_index]; + return data().at(0); } template -void Literal::Set(tensorflow::gtl::ArraySlice multi_index, - NativeT value) { - int64 linear_index = LinearIndex(multi_index); - GetMutableArraySlice().at(linear_index) = value; -} - -template <> -inline void Literal::Set(tensorflow::gtl::ArraySlice multi_index, - uint8 value) { - int64 linear_index = LinearIndex(multi_index); - (*mutable_u8s())[linear_index] = value; -} - -template <> -inline void Literal::Set(tensorflow::gtl::ArraySlice multi_index, - int8 value) { - return Set(multi_index, value); -} - -template <> -inline void Literal::Set(tensorflow::gtl::ArraySlice multi_index, - int64 value) { - int64 linear_index = LinearIndex(multi_index); - (*mutable_s64s())[linear_index] = value; +NativeT Literal::GetSparseElement(int64 sparse_element_number, + const ShapeIndex& shape_index) const { + CHECK( + LayoutUtil::IsSparseArray(ShapeUtil::GetSubshape(shape(), shape_index))); + return data(shape_index)[sparse_element_number]; } -template <> -/* static */ inline void Literal::Set( - tensorflow::gtl::ArraySlice multi_index, uint64 value) { - int64 linear_index = LinearIndex(multi_index); - (*mutable_u64s())[linear_index] = value; +template +void Literal::AppendSparseElement( + tensorflow::gtl::ArraySlice multi_index, NativeT value, + const ShapeIndex& shape_index) { + Piece& p = piece(shape_index); + const Shape& subshape = p.subshape(); + CHECK(LayoutUtil::IsSparseArray(subshape)); + int64 rank = ShapeUtil::Rank(subshape); + CHECK_EQ(multi_index.size(), rank); + int64 last_element = p.sparse_indices()->index_count(); + CHECK_LT(last_element, LayoutUtil::MaxSparseElements(subshape.layout())); + p.sparse_indices()->Append(multi_index); + CHECK_LT(last_element, p.data().size()); + p.data()[last_element] = value; } // Returns an identity matrix (rank 2) with the given row and column count. @@ -1076,51 +1159,31 @@ void Literal::EachCell( } while (IndexUtil::BumpIndices(shape(), &indices)); } -template -inline void Literal::PopulateR0(NativeT value) { - *mutable_shape() = ShapeUtil::MakeShape( - primitive_util::NativeToPrimitiveType(), {}); - Resize(1, value); -} - template inline void Literal::PopulateR1(tensorflow::gtl::ArraySlice values) { - *mutable_shape() = - ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), - {static_cast(values.size())}); - Reserve(values.size()); + CHECK(ShapeUtil::IsArray(shape())); + CHECK_EQ(ShapeUtil::Rank(shape()), 1); + CHECK_EQ(ShapeUtil::ElementsIn(shape()), values.size()); + CHECK_EQ(shape().element_type(), + primitive_util::NativeToPrimitiveType()); for (int64 i = 0; i < values.size(); ++i) { Set({i}, values[i]); } } -inline void Literal::PopulateR1(const tensorflow::core::Bitmap& values) { - *mutable_shape() = - ShapeUtil::MakeShape(PRED, {static_cast(values.bits())}); - Reserve(values.bits()); - for (int64 i = 0; i < static_cast(values.bits()); ++i) { - Set({i}, values.get(i)); - } -} - template -void Literal::PopulateR2WithLayout( - std::initializer_list> values, - const Layout& layout) { - *mutable_shape() = ShapeUtil::MakeShapeWithLayout( - primitive_util::NativeToPrimitiveType(), - {static_cast(values.size()), - static_cast(values.begin()->size())}, - LayoutUtil::MinorToMajor(layout)); +void Literal::PopulateR2( + std::initializer_list> values) { + CHECK(ShapeUtil::IsArray(shape())); + CHECK_EQ(ShapeUtil::Rank(shape()), 2); + CHECK_EQ(shape().element_type(), + primitive_util::NativeToPrimitiveType()); const int64 dim0_size = values.size(); const int64 dim1_size = values.begin()->size(); CHECK_EQ(dim0_size, shape().dimensions(0)); CHECK_EQ(dim1_size, shape().dimensions(1)); - const int64 num_elements = dim1_size * dim0_size; - Reserve(num_elements); - int64 dim0 = 0; for (auto inner_list : values) { int64 dim1 = 0; @@ -1134,70 +1197,65 @@ void Literal::PopulateR2WithLayout( } template -void Literal::PopulateR2( - std::initializer_list> values) { - PopulateR2WithLayout(values, LayoutUtil::GetDefaultLayoutForR2()); -} - -template -void Literal::PopulateFromArrayWithLayout(const Array& values, - const Layout& layout) { - CHECK_EQ(layout.format(), DENSE); - *mutable_shape() = ShapeUtil::MakeShapeWithLayout( - primitive_util::NativeToPrimitiveType(), values.dimensions(), - LayoutUtil::MinorToMajor(layout)); - Reserve(values.num_elements()); +void Literal::PopulateFromArray(const Array& values) { + CHECK(ShapeUtil::IsArray(shape())); + CHECK_EQ(shape().element_type(), + primitive_util::NativeToPrimitiveType()); + CHECK_EQ(ShapeUtil::Rank(shape()), values.num_dimensions()); + for (int dim = 0; dim < values.num_dimensions(); ++dim) { + CHECK_EQ(values.dim(dim), shape().dimensions(dim)); + } values.Each([this](tensorflow::gtl::ArraySlice indices, NativeT value) { this->Set(indices, value); }); } -template -void Literal::PopulateFromArray(const Array& values) { - PopulateFromArrayWithLayout( - values, LayoutUtil::GetDefaultLayoutForRank(values.num_dimensions())); -} - -template -void Literal::PopulateR2FromArray2DWithLayout(const Array2D& values, - const Layout& layout) { - PopulateFromArrayWithLayout(values, layout); -} - template void Literal::PopulateR2FromArray2D(const Array2D& values) { PopulateFromArray(values); } -template -void Literal::PopulateR3FromArray3DWithLayout(const Array3D& values, - const Layout& layout) { - PopulateFromArrayWithLayout(values, layout); -} - template void Literal::PopulateR3FromArray3D(const Array3D& values) { PopulateFromArray(values); } template -void Literal::PopulateR4FromArray4DWithLayout(const Array4D& values, - const Layout& layout) { - PopulateFromArrayWithLayout(values, layout); +void Literal::PopulateR4FromArray4D(const Array4D& values) { + PopulateFromArray(values); } template -void Literal::PopulateR4FromArray4D(const Array4D& values) { - PopulateFromArray(values); +void Literal::PopulateSparse(SparseIndexArray indices, + tensorflow::gtl::ArraySlice values, + bool sort) { + CHECK(LayoutUtil::IsSparseArray(shape())); + int rank = ShapeUtil::Rank(shape()); + CHECK_EQ(indices.rank(), rank); + int64 max_elements = LayoutUtil::MaxSparseElements(shape().layout()); + CHECK_LE(indices.max_indices(), max_elements); + int64 num_elements = values.size(); + CHECK_LE(num_elements, max_elements); + CHECK_EQ(num_elements, indices.index_count()); + auto root_data = root_piece().data(); + root_data.remove_suffix(max_elements - values.size()); + std::copy(values.begin(), values.end(), root_data.begin()); + *this->root_piece().sparse_indices() = std::move(indices); + if (sort) { + auto root_data = this->root_piece().data(); + root_data.remove_suffix(root_data.size() - num_elements); + this->root_piece().sparse_indices()->SortWithValues(root_data); + } + DCHECK(this->root_piece().sparse_indices()->Validate(shape())); } template Status Literal::Populate(const FnType& generator) { const Shape& this_shape = shape(); const int64 rank = ShapeUtil::Rank(this_shape); + TF_RET_CHECK(LayoutUtil::IsDenseArray(this_shape)); TF_RET_CHECK(this_shape.element_type() == primitive_util::NativeToPrimitiveType()); - tensorflow::gtl::MutableArraySlice data = - GetMutableArraySlice(); + tensorflow::gtl::MutableArraySlice literal_data = data(); if (rank > 0) { StrideConfig stride_config(this_shape, this_shape, AsInt64Slice(this_shape.dimensions())); @@ -1206,11 +1264,12 @@ Status Literal::Populate(const FnType& generator) { ShapeUtil::GetDimension(this_shape, stride_config.minor_dimension); auto init_function = [&](const std::vector& indexes) { - const int64 index = LinearIndex(indexes); + const int64 index = + IndexUtil::MultidimensionalIndexToLinearIndex(shape(), indexes); std::copy(indexes.begin(), indexes.end(), minor_scan_indexes.begin()); for (int64 i = 0; i < minor_dimension_size; ++i) { minor_scan_indexes[stride_config.minor_dimension] = i; - data.at(index + i) = generator(minor_scan_indexes); + literal_data.at(index + i) = generator(minor_scan_indexes); } return true; }; @@ -1219,31 +1278,27 @@ Status Literal::Populate(const FnType& generator) { init_function); } else { // For scalars. - data.at(0) = generator({}); + literal_data.at(0) = generator({}); } return Status::OK(); } template -void Literal::PopulateWithValue(NativeT value, - tensorflow::gtl::ArraySlice dimensions) { - *mutable_shape() = ShapeUtil::MakeShape( - primitive_util::NativeToPrimitiveType(), dimensions); - Resize(ShapeUtil::ElementsIn(shape()), value); +void Literal::PopulateWithValue(NativeT value) { + CHECK(ShapeUtil::IsArray(shape())); + CHECK_EQ(shape().element_type(), + primitive_util::NativeToPrimitiveType()); + for (NativeT& element : data()) { + element = value; + } } template /* static */ std::unique_ptr Literal::CreateFullWithDescendingLayout( tensorflow::gtl::ArraySlice dimensions, NativeT value) { - Shape this_shape = ShapeUtil::MakeShapeWithDescendingLayout( - primitive_util::NativeToPrimitiveType(), dimensions); - auto literal = MakeUnique(); - *literal->mutable_shape() = this_shape; - literal->Reserve(ShapeUtil::ElementsIn(this_shape)); - std::vector index(dimensions.size(), 0); - do { - literal->Set(index, value); - } while (IndexUtil::BumpIndices(this_shape, &index)); + auto literal = MakeUnique(ShapeUtil::MakeShapeWithDescendingLayout( + primitive_util::NativeToPrimitiveType(), dimensions)); + literal->PopulateWithValue(value); return literal; } @@ -1254,14 +1309,12 @@ std::unique_ptr Literal::Replicate(int64 times) const { for (int64 bound : shape().dimensions()) { bounds.push_back(bound); } - auto literal = MakeUnique(); - *literal->mutable_shape() = - ShapeUtil::MakeShape(shape().element_type(), bounds); + auto literal = + MakeUnique(ShapeUtil::MakeShape(shape().element_type(), bounds)); int64 elements = ShapeUtil::ElementsIn(literal->shape()); if (elements == 0) { return literal; } - literal->Reserve(elements); DimensionVector output_indices(bounds.size(), 0); tensorflow::gtl::ArraySlice input_indices = output_indices; diff --git a/tensorflow/compiler/xla/literal_util_test.cc b/tensorflow/compiler/xla/literal_util_test.cc index 7ff64c4134155e7fe22ab99584970a7d6d6e8803..b3583c2eb75de8297d5e7507430491f119bd4462 100644 --- a/tensorflow/compiler/xla/literal_util_test.cc +++ b/tensorflow/compiler/xla/literal_util_test.cc @@ -31,6 +31,7 @@ namespace xla { namespace { using ::testing::ElementsAre; +using ::testing::HasSubstr; class LiteralUtilTest : public ::testing::Test { protected: @@ -192,6 +193,34 @@ TEST_F(LiteralUtilTest, CreateR3FromArray3d) { ASSERT_EQ(expected, result); } +TEST_F(LiteralUtilTest, CreateSparse) { + std::vector dimensions = {8, 8, 8}; + Array2D indices = { + {3, 4, 5}, + {1, 2, 3}, + {2, 3, 4}, + {3, 5, 6}, + }; + std::vector values = {7, 8, 9, 10}; + auto literal = Literal::CreateSparse( + dimensions, SparseIndexArray(indices.n1() + 3, indices), values); + + Array2D expected_indices = { + {1, 2, 3}, + {2, 3, 4}, + {3, 4, 5}, + {3, 5, 6}, + }; + std::vector expected_values = {8, 9, 7, 10}; + + EXPECT_EQ(literal->sparse_indices()->data(), + tensorflow::gtl::ArraySlice( + expected_indices.data(), expected_indices.num_elements())); + EXPECT_EQ(tensorflow::gtl::ArraySlice(literal->data().data(), + expected_values.size()), + tensorflow::gtl::ArraySlice(expected_values)); +} + TEST_F(LiteralUtilTest, LiteralR4F32ProjectedStringifies) { // clang-format off auto literal = Literal::CreateR4Projected({ @@ -293,29 +322,28 @@ TEST_F(LiteralUtilTest, NonScalarEquality) { auto matrix_different = Literal::CreateR2({{4.0, 3.0}, {1.0, 2.0}}); auto vector_literal = Literal::CreateR1({1.0, 2.0, 3.0, 4.0}); auto scalar = Literal::CreateR0(1.0); + Literal nil(ShapeUtil::MakeNil()); EXPECT_EQ(*matrix, *matrix); EXPECT_EQ(*matrix, *matrix_clone); EXPECT_NE(*matrix, *matrix_different); EXPECT_NE(*matrix, *vector_literal); EXPECT_NE(*matrix, *scalar); + EXPECT_NE(*matrix, nil); + EXPECT_EQ(nil, nil); } TEST_F(LiteralUtilTest, DifferentLayoutEquality) { // Test equality with literals which have different layouts. - auto colmajor = MakeUnique(); - *colmajor->mutable_shape() = ShapeUtil::MakeShape(F32, {2, 2}); - *colmajor->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); - colmajor->Reserve(4); + auto colmajor = + MakeUnique(ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {0, 1})); colmajor->Set({0, 0}, 1.0); colmajor->Set({0, 1}, 2.0); colmajor->Set({1, 0}, 3.0); colmajor->Set({1, 1}, 4.0); - auto rowmajor = MakeUnique(); - *rowmajor->mutable_shape() = ShapeUtil::MakeShape(F32, {2, 2}); - *rowmajor->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({1, 0}); - rowmajor->Reserve(4); + auto rowmajor = + MakeUnique(ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {1, 0})); rowmajor->Set({0, 0}, 1.0); rowmajor->Set({0, 1}, 2.0); rowmajor->Set({1, 0}, 3.0); @@ -597,24 +625,26 @@ TEST_F(LiteralUtilTest, TestR4RelayoutEquivalence) { TEST_F(LiteralUtilTest, TestR2LinearLayout) { // Test expected memory layout of R2 dim0-minor (column-major) literal. - auto mat_dim0minor = Literal::CreateR2WithLayout({{1, 2, 3}, {4, 5, 6}}, - layout_r2_dim0minor_); - EXPECT_EQ(mat_dim0minor->s32s_size(), 6); - EXPECT_THAT(mat_dim0minor->s32s(), ElementsAre(1, 4, 2, 5, 3, 6)); + auto mat_dim0minor = Literal::CreateR2WithLayout( + {{1, 2, 3}, {4, 5, 6}}, layout_r2_dim0minor_); + EXPECT_EQ(mat_dim0minor->element_count(), 6); + EXPECT_THAT(mat_dim0minor->data(), ElementsAre(1, 4, 2, 5, 3, 6)); // Test expected memory layout when using Relayout to row major. auto relaid_mat_to_dim0major = mat_dim0minor->Relayout(layout_r2_dim0major_); - EXPECT_THAT(relaid_mat_to_dim0major->s32s(), ElementsAre(1, 2, 3, 4, 5, 6)); + EXPECT_THAT(relaid_mat_to_dim0major->data(), + ElementsAre(1, 2, 3, 4, 5, 6)); // Test expected memory layout of R2 created with dim0-major (row-major). - auto mat_dim0major = Literal::CreateR2WithLayout({{1, 2, 3}, {4, 5, 6}}, - layout_r2_dim0major_); - EXPECT_EQ(mat_dim0major->s32s_size(), 6); - EXPECT_THAT(mat_dim0major->s32s(), ElementsAre(1, 2, 3, 4, 5, 6)); + auto mat_dim0major = Literal::CreateR2WithLayout( + {{1, 2, 3}, {4, 5, 6}}, layout_r2_dim0major_); + EXPECT_EQ(mat_dim0major->element_count(), 6); + EXPECT_THAT(mat_dim0major->data(), ElementsAre(1, 2, 3, 4, 5, 6)); // Test expected memory layout when using Relayout to column major. auto relaid_mat_to_dim0minor = mat_dim0major->Relayout(layout_r2_dim0minor_); - EXPECT_THAT(relaid_mat_to_dim0minor->s32s(), ElementsAre(1, 4, 2, 5, 3, 6)); + EXPECT_THAT(relaid_mat_to_dim0minor->data(), + ElementsAre(1, 4, 2, 5, 3, 6)); } TEST_F(LiteralUtilTest, TestR3LinearLayout) { @@ -634,27 +664,27 @@ TEST_F(LiteralUtilTest, TestR3LinearLayout) { auto lit_dim0minor = Literal::CreateR3FromArray3DWithLayout(arr3d, layout_r3_dim0minor_); - EXPECT_EQ(lit_dim0minor->s32s_size(), 12); + EXPECT_EQ(lit_dim0minor->element_count(), 12); std::vector expected_dim0minor{1, 7, 4, 10, 2, 8, 5, 11, 3, 9, 6, 12}; - EXPECT_THAT(lit_dim0minor->s32s(), + EXPECT_THAT(lit_dim0minor->data(), testing::ElementsAreArray(expected_dim0minor)); // Test expected memory layout when using Relayout to row major. auto relaid_lit_to_dim0major = lit_dim0minor->Relayout(layout_r3_dim0major_); std::vector expected_dim0major{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; - EXPECT_THAT(relaid_lit_to_dim0major->s32s(), + EXPECT_THAT(relaid_lit_to_dim0major->data(), testing::ElementsAreArray(expected_dim0major)); // Test expected memory layout of R3 created with dim0-major (row-major). auto lit_dim0major = Literal::CreateR3FromArray3DWithLayout(arr3d, layout_r3_dim0major_); - EXPECT_EQ(lit_dim0major->s32s_size(), 12); - EXPECT_THAT(lit_dim0major->s32s(), + EXPECT_EQ(lit_dim0major->element_count(), 12); + EXPECT_THAT(lit_dim0major->data(), testing::ElementsAreArray(expected_dim0major)); // Test expected memory layout when using Relayout to column major. auto relaid_lit_to_dim0minor = lit_dim0major->Relayout(layout_r3_dim0minor_); - EXPECT_THAT(relaid_lit_to_dim0minor->s32s(), + EXPECT_THAT(relaid_lit_to_dim0minor->data(), testing::ElementsAreArray(expected_dim0minor)); } @@ -687,28 +717,28 @@ TEST_F(LiteralUtilTest, SliceR3U32Full) { } TEST_F(LiteralUtilTest, PopulateR1S64) { - Literal output; + Literal output(ShapeUtil::MakeShape(S64, {1})); output.PopulateR1({77}); auto expected = Literal::CreateR1({77}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateR1U64) { - Literal output; + Literal output(ShapeUtil::MakeShape(U64, {2})); output.PopulateR1({{77, 88}}); auto expected = Literal::CreateR1({{77, 88}}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateR1C64) { - Literal output; + Literal output(ShapeUtil::MakeShape(C64, {1})); output.PopulateR1({{77, 88}}); auto expected = Literal::CreateR1({{77, 88}}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateR2C64) { - Literal output; + Literal output(ShapeUtil::MakeShape(C64, {2, 2})); output.PopulateR2({{{7, 8}, {9, 10}}, {{1, 2}, {3, 4}}}); auto expected = Literal::CreateR2({{{7, 8}, {9, 10}}, {{1, 2}, {3, 4}}}); @@ -716,78 +746,78 @@ TEST_F(LiteralUtilTest, PopulateR2C64) { } TEST_F(LiteralUtilTest, PopulateWithValueR0BF16) { - Literal output; + Literal output(ShapeUtil::MakeShape(BF16, {})); bfloat16 h(0.25f); - output.PopulateWithValue(h, {}); + output.PopulateWithValue(h); auto expected = Literal::CreateR0(h); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR1BF16) { - Literal output; + Literal output(ShapeUtil::MakeShape(BF16, {3})); bfloat16 h(0.5f); - output.PopulateWithValue(h, {3}); + output.PopulateWithValue(h); auto expected = Literal::CreateR1({h, h, h}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR2BF16) { - Literal output; + Literal output(ShapeUtil::MakeShape(BF16, {2, 2})); bfloat16 h(2.0f); - output.PopulateWithValue(h, {2, 2}); + output.PopulateWithValue(h); auto expected = Literal::CreateR2({{h, h}, {h, h}}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR0F32) { - Literal output; - output.PopulateWithValue(2.5f, {}); + Literal output(ShapeUtil::MakeShape(F32, {})); + output.PopulateWithValue(2.5f); auto expected = Literal::CreateR0(2.5f); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR1S64) { - Literal output; - output.PopulateWithValue(-7, {3}); + Literal output(ShapeUtil::MakeShape(S64, {3})); + output.PopulateWithValue(-7); auto expected = Literal::CreateR1({-7, -7, -7}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR2U64) { - Literal output; - output.PopulateWithValue(42, {2, 2}); + Literal output(ShapeUtil::MakeShape(U64, {2, 2})); + output.PopulateWithValue(42); auto expected = Literal::CreateR2({{42, 42}, {42, 42}}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR2C64) { - Literal output; - output.PopulateWithValue({4, 2}, {2, 2}); + Literal output(ShapeUtil::MakeShape(C64, {2, 2})); + output.PopulateWithValue({4, 2}); auto expected = Literal::CreateR2({{{4, 2}, {4, 2}}, {{4, 2}, {4, 2}}}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR0F16) { - Literal output; + Literal output(ShapeUtil::MakeShape(F16, {})); half h(0.25f); - output.PopulateWithValue(h, {}); + output.PopulateWithValue(h); auto expected = Literal::CreateR0(h); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR1F16) { - Literal output; + Literal output(ShapeUtil::MakeShape(F16, {3})); half h(0.5f); - output.PopulateWithValue(h, {3}); + output.PopulateWithValue(h); auto expected = Literal::CreateR1({h, h, h}); EXPECT_EQ(output, *expected); } TEST_F(LiteralUtilTest, PopulateWithValueR2F16) { - Literal output; + Literal output(ShapeUtil::MakeShape(F16, {2, 2})); half h(2.0f); - output.PopulateWithValue(h, {2, 2}); + output.PopulateWithValue(h); auto expected = Literal::CreateR2({{h, h}, {h, h}}); EXPECT_EQ(output, *expected); } @@ -803,7 +833,7 @@ TEST_F(LiteralUtilTest, ReplicateR2U32) { EXPECT_EQ(*output, *expected); } -TEST_F(LiteralUtilTest, Copy) { +TEST_F(LiteralUtilTest, CopySliceFrom) { const int64 dimensions[] = {17, 15, 34, 21}; const int64 layouts[][4] = { {3, 2, 1, 0}, {0, 2, 1, 3}, {0, 1, 2, 3}, {2, 0, 3, 1}, {1, 3, 0, 2}}; @@ -826,7 +856,7 @@ TEST_F(LiteralUtilTest, Copy) { const int64 src_base[] = {3, 1, 5, 7}; const int64 dest_base[] = {6, 4, 12, 2}; const int64 copy_size[] = {7, 8, 11, 9}; - TF_EXPECT_OK(blank->Copy(*source, src_base, dest_base, copy_size)); + TF_EXPECT_OK(blank->CopySliceFrom(*source, src_base, dest_base, copy_size)); std::vector source_indexes(TF_ARRAYSIZE(dimensions), 0); std::vector blank_indexes(TF_ARRAYSIZE(dimensions), 0); @@ -849,16 +879,16 @@ TEST_F(LiteralUtilTest, Copy) { } } -TEST_F(LiteralUtilTest, CopyScalars) { +TEST_F(LiteralUtilTest, CopyFromScalars) { auto zero = Literal::CreateR0(0); auto nine = Literal::CreateR0(9); - TF_EXPECT_OK(zero->Copy(*nine, {}, {}, {})); + TF_EXPECT_OK(zero->CopyFrom(*nine)); EXPECT_EQ(*zero, *nine); auto vect = Literal::CreateR1({3, 4, 9, 12, 5, 17, 21}); - TF_EXPECT_OK(zero->Copy(*vect, {5}, {}, {})); + TF_EXPECT_OK(zero->CopySliceFrom(*vect, {5}, {}, {})); EXPECT_EQ(zero->Get({}), 17); - TF_EXPECT_OK(vect->Copy(*zero, {}, {4}, {})); + TF_EXPECT_OK(vect->CopySliceFrom(*zero, {}, {4}, {})); EXPECT_EQ(vect->Get({4}), 17); } @@ -872,7 +902,7 @@ TEST_F(LiteralUtilTest, CopyFromAndToZeroElement) { const auto empty = Literal::CreateFromShape(empty_r1_shape); auto nine = Literal::CreateR1({9}); - TF_EXPECT_OK(nine->Copy(*empty, {0}, {0}, {0})); + TF_EXPECT_OK(nine->CopySliceFrom(*empty, {0}, {0}, {0})); EXPECT_EQ(*nine, *const_nine); } @@ -881,18 +911,101 @@ TEST_F(LiteralUtilTest, CopyFromAndToZeroElement) { const auto empty = Literal::CreateFromShape(empty_r1_shape); auto nine = Literal::CreateR1({9}); - TF_EXPECT_OK(empty->Copy(*nine, {0}, {0}, {0})); + TF_EXPECT_OK(empty->CopySliceFrom(*nine, {0}, {0}, {0})); EXPECT_EQ(*empty, *const_empty); } } +TEST_F(LiteralUtilTest, CopyFromNilShape) { + Literal nil_literal0(ShapeUtil::MakeNil()); + Literal nil_literal1(ShapeUtil::MakeNil()); + // This doesn't actually do any copying, but it should succeed. + TF_ASSERT_OK(nil_literal0.CopyFrom(nil_literal1)); +} + +TEST_F(LiteralUtilTest, CopyFromArrays) { + auto scalar_42 = Literal::CreateR0(42.0); + auto scalar_123 = Literal::CreateR0(123.0); + EXPECT_NE(*scalar_42, *scalar_123); + TF_ASSERT_OK(scalar_42->CopyFrom(*scalar_123, /*dest_shape_index=*/{}, + /*src_shape_index=*/{})); + EXPECT_EQ(*scalar_42, *scalar_123); + EXPECT_EQ(scalar_42->Get({}), 123.0f); + + auto matrix_1234 = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto matrix_5678 = Literal::CreateR2({{5.0, 6.0}, {7.0, 8.0}}); + EXPECT_NE(*matrix_1234, *matrix_5678); + EXPECT_EQ(matrix_1234->Get({0, 0}), 1.0f); + TF_ASSERT_OK(matrix_1234->CopyFrom(*matrix_5678, /*dest_shape_index=*/{}, + /*src_shape_index=*/{})); + EXPECT_EQ(*matrix_1234, *matrix_5678); + EXPECT_EQ(matrix_1234->Get({0, 0}), 5.0f); +} + +TEST_F(LiteralUtilTest, CopyFromTuples) { + auto matrix = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + Literal nil_literal(ShapeUtil::MakeNil()); + auto nested_tuple = Literal::MakeTuple( + {matrix.get(), + Literal::MakeTuple({Literal::CreateR0(42).get(), + Literal::CreateR1({23.0, 44.0}).get(), + &nil_literal}) + .get()}); + // Create a tuple the same shape as the inner tuple of nested_tuple but with + // different values.. + auto tuple = Literal::MakeTuple({Literal::CreateR0(-5).get(), + Literal::CreateR1({2.0, 4.0}).get(), + &nil_literal}); + + EXPECT_EQ(*matrix, LiteralView::Create(*nested_tuple, {0})); + EXPECT_EQ(nested_tuple->Get({}, {1, 0}), 42); + EXPECT_EQ(nested_tuple->Get({0}, {1, 1}), 23.0); + EXPECT_EQ(nested_tuple->Get({1}, {1, 1}), 44.0); + + // Overwrite the inner tuple element of nested_tuple with the contents of + // 'tuple'. + TF_ASSERT_OK(nested_tuple->CopyFrom(*tuple, /*dest_shape_index=*/{1}, + /*src_shape_index=*/{})); + + // The matrix element should be unchanged. + EXPECT_EQ(*matrix, LiteralView::Create(*nested_tuple, {0})); + + // The tuple element should have been copied from 'tuple'. + EXPECT_EQ(nested_tuple->Get({}, {1, 0}), -5); + EXPECT_EQ(nested_tuple->Get({0}, {1, 1}), 2.0); + EXPECT_EQ(nested_tuple->Get({1}, {1, 1}), 4.0); +} +TEST_F(LiteralUtilTest, CopyBetweenSameTuple) { + auto tuple = Literal::MakeTuple( + {Literal::CreateR0(-2).get(), Literal::CreateR0(4).get()}); + + EXPECT_EQ(tuple->Get({}, {0}), -2); + EXPECT_EQ(tuple->Get({}, {1}), 4); + + // Copy from one element to the other. + TF_ASSERT_OK(tuple->CopyFrom(*tuple, /*dest_shape_index=*/{1}, + /*src_shape_index=*/{0})); + + EXPECT_EQ(tuple->Get({}, {0}), -2); + EXPECT_EQ(tuple->Get({}, {1}), -2); +} + +TEST_F(LiteralUtilTest, CopyFromDifferentShapes) { + auto matrix = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto vector = Literal::CreateR1({5.0, 7.0}); + Status status = matrix->CopyFrom(*vector); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), + HasSubstr("Destination subshape incompatible")); +} + TEST_F(LiteralUtilTest, F16) { // Verify that the internal data views are consistent and that they // are in little endian format // TODO - modify if we make the data format machine endianess dependent auto m1 = Literal::CreateFromShape(ShapeUtil::MakeShape(F16, {2, 2})); Literal* l1 = m1.get(); - const char* d1 = static_cast(l1->InternalData()); + const char* d1 = reinterpret_cast(l1->data().data()); EXPECT_EQ(d1[0], 0); EXPECT_EQ(d1[1], 0); EXPECT_EQ(d1[2], 0); @@ -901,13 +1014,12 @@ TEST_F(LiteralUtilTest, F16) { EXPECT_EQ(d1[5], 0); EXPECT_EQ(d1[6], 0); EXPECT_EQ(d1[7], 0); - EXPECT_EQ(l1->InternalData(), l1->MutableInternalData()); half h1(1.0f); half h2(2.0f); auto m2 = Literal::CreateR2({{h1, h2}, {h2, h1}}); Literal* l2 = m2.get(); - const char* d2 = static_cast(l2->InternalData()); + const char* d2 = reinterpret_cast(l2->data().data()); EXPECT_EQ(d2[0], 0); EXPECT_EQ(d2[1], 0x3C); EXPECT_EQ(d2[2], 0); @@ -916,7 +1028,6 @@ TEST_F(LiteralUtilTest, F16) { EXPECT_EQ(d2[5], 0x40); EXPECT_EQ(d2[6], 0); EXPECT_EQ(d2[7], 0x3C); - EXPECT_EQ(l2->InternalData(), l2->MutableInternalData()); } TEST_F(LiteralUtilTest, Populate) { @@ -941,7 +1052,9 @@ TEST_F(LiteralUtilTest, Populate) { auto generator = [&](tensorflow::gtl::ArraySlice indexes) -> uint32 { // Offsets from linear index just to avoid R0 literals to be initialized // with zero. - return literal->LinearIndex(indexes) + 17; + return IndexUtil::MultidimensionalIndexToLinearIndex(literal->shape(), + indexes) + + 17; }; TF_EXPECT_OK(literal->Populate(generator)); @@ -1118,16 +1231,18 @@ TEST_F(LiteralUtilTest, CopyFromProto_Bool) { for (int len = 0; len < 25; ++len) { p.mutable_shape()->clear_dimensions(); p.mutable_shape()->add_dimensions(len); + LayoutUtil::SetToDefaultLayout(p.mutable_shape()); p.clear_preds(); for (int i = 0; i < len; ++i) { p.add_preds((i % 2) == (len % 2)); } - Literal literal(p); - ASSERT_EQ(len, literal.preds_size()); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr literal, + Literal::CreateFromProto(p)); + ASSERT_EQ(len, literal->data().size()); int i = 0; - for (auto it = literal.preds().begin(); it < literal.preds().end(); ++it) { - EXPECT_EQ((i % 2) == (len % 2), *it); + for (bool value : literal->data()) { + EXPECT_EQ((i % 2) == (len % 2), value); ++i; } } @@ -1141,8 +1256,7 @@ TEST_F(LiteralUtilTest, ToProto_f16) { auto m = Literal::CreateR2({{h1, h2}, {h2, h1}}); Literal* l = m.get(); EXPECT_EQ(4, ShapeUtil::ElementsIn(l->shape())); - EXPECT_EQ(4, l->f16s().size()); - EXPECT_EQ(4, l->f16s_size()); + EXPECT_EQ(4, l->data().size()); LiteralProto p = l->ToProto(); EXPECT_EQ(4, ShapeUtil::ElementsIn(p.shape())); @@ -1168,17 +1282,12 @@ TEST_F(LiteralUtilTest, CopyFromProto_f16) { p.mutable_shape()->set_element_type(F16); p.mutable_shape()->clear_dimensions(); p.mutable_shape()->add_dimensions(4); + LayoutUtil::SetToDefaultLayout(p.mutable_shape()); p.clear_f16s(); p.set_f16s(half_vals, 8); - - Literal literal(p); - ASSERT_EQ(4, literal.f16s_size()); - ASSERT_EQ(h1, literal.f16s(0)); - ASSERT_EQ(h2, literal.f16s(1)); - ASSERT_EQ(h2, literal.f16s(2)); - ASSERT_EQ(h1, literal.f16s(3)); - - const std::vector& r = literal.f16s(); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr literal, + Literal::CreateFromProto(p)); + auto r = literal->data(); ASSERT_EQ(4, r.size()); ASSERT_EQ(h1, r[0]); ASSERT_EQ(h2, r[1]); @@ -1186,24 +1295,402 @@ TEST_F(LiteralUtilTest, CopyFromProto_f16) { ASSERT_EQ(h1, r[3]); } -TEST_F(LiteralUtilTest, Subliterals) { +TEST_F(LiteralUtilTest, LiteralViewTest) { + auto scalar = Literal::CreateR0(1.0); + auto matrix = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto tuple = Literal::MakeTuple({scalar.get(), matrix.get()}); + auto nested_tuple = Literal::MakeTuple({tuple.get(), scalar.get()}); + Literal nil(ShapeUtil::MakeNil()); + + EXPECT_EQ(LiteralView::Create(*scalar, {}), *scalar); + EXPECT_EQ(LiteralView::Create(*matrix, {}), *matrix); + EXPECT_EQ(LiteralView::Create(*tuple, {}), *tuple); + EXPECT_EQ(LiteralView::Create(*nested_tuple, {}), *nested_tuple); + EXPECT_EQ(LiteralView::Create(nil, {}), nil); + + EXPECT_EQ(LiteralView::Create(*tuple, {0}), *scalar); + EXPECT_EQ(LiteralView::Create(*tuple, {1}), *matrix); + + EXPECT_EQ(LiteralView::Create(*nested_tuple, {0}), *tuple); + EXPECT_EQ(LiteralView::Create(*nested_tuple, {0, 0}), *scalar); + EXPECT_EQ(LiteralView::Create(*nested_tuple, {0, 1}), *matrix); + EXPECT_EQ(LiteralView::Create(*nested_tuple, {1}), *scalar); +} + +TEST_F(LiteralUtilTest, MutatingLiteralView) { + auto scalar = Literal::CreateR0(1.0); + auto matrix = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto tuple = Literal::MakeTuple({scalar.get(), matrix.get()}); + auto nested_tuple = Literal::MakeTuple({tuple.get(), scalar.get()}); + // Verify that changing the underlying data beneath the view changes the + // data of the view itself. + const auto nested_tuple_view = LiteralView::Create(*nested_tuple); + EXPECT_EQ( + nested_tuple->Get(/*multi_index=*/{}, /*shape_index=*/{0, 0}), + 1.0f); + EXPECT_EQ(nested_tuple_view.Get(/*multi_index=*/{}, + /*shape_index=*/{0, 0}), + 1.0f); + nested_tuple->Set(/*multi_index=*/{}, /*shape_index=*/{0, 0}, 555.0f); + EXPECT_EQ( + nested_tuple->Get(/*multi_index=*/{}, /*shape_index=*/{0, 0}), + 555.0f); + EXPECT_EQ(nested_tuple_view.Get(/*multi_index=*/{}, + /*shape_index=*/{0, 0}), + 555.0f); +} + +TEST_F(LiteralUtilTest, LiteralViewOfALiteralView) { auto scalar = Literal::CreateR0(1.0); auto matrix = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); auto tuple = Literal::MakeTuple({scalar.get(), matrix.get()}); auto nested_tuple = Literal::MakeTuple({tuple.get(), scalar.get()}); - EXPECT_EQ(&scalar->GetSubliteral(/*index=*/{}), scalar.get()); - EXPECT_EQ(&matrix->GetSubliteral(/*index=*/{}), matrix.get()); - EXPECT_EQ(&tuple->GetSubliteral(/*index=*/{}), tuple.get()); - EXPECT_EQ(&nested_tuple->GetSubliteral(/*index=*/{}), nested_tuple.get()); + const auto nested_tuple_view = LiteralView::Create(*nested_tuple); + const auto tuple_view = + LiteralView::Create(nested_tuple_view, /*view_root=*/{0}); + const auto matrix_view = LiteralView::Create(tuple_view, /*view_root=*/{1}); + EXPECT_EQ(matrix_view, *Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}})); +} + +TEST_F(LiteralUtilTest, LiteralMove) { + std::unique_ptr matrix = + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + Literal literal(std::move(*matrix)); + + EXPECT_TRUE( + ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {2, 2}), literal.shape())); + EXPECT_EQ(literal.Get({0, 0}), 1.0); + EXPECT_EQ(literal.Get({0, 1}), 2.0); + EXPECT_EQ(literal.Get({1, 0}), 3.0); + EXPECT_EQ(literal.Get({1, 1}), 4.0); +} - EXPECT_EQ(tuple->GetSubliteral(/*index=*/{0}), *scalar); - EXPECT_EQ(tuple->GetSubliteral(/*index=*/{1}), *matrix); +TEST_F(LiteralUtilTest, DecomposeTuple) { + Literal nil_literal(ShapeUtil::MakeNil()); + auto nested_tuple = Literal::MakeTuple( + {Literal::CreateR2({{1, 2}, {3, 4}}).get(), + Literal::MakeTuple({Literal::CreateR0(42).get(), + Literal::CreateR1({23.0, 44.0}).get(), + &nil_literal}) + .get(), + &nil_literal}); + + EXPECT_FALSE(ShapeUtil::IsNil(nested_tuple->shape())); + std::vector elements = nested_tuple->DecomposeTuple(); + EXPECT_TRUE(ShapeUtil::IsNil(nested_tuple->shape())); + + ASSERT_EQ(elements.size(), 3); + + EXPECT_TRUE(ShapeUtil::Compatible(elements[0].shape(), + ShapeUtil::MakeShape(S32, {2, 2}))); + EXPECT_EQ(elements[0].Get({0, 0}), 1); + EXPECT_EQ(elements[0].Get({0, 1}), 2); + EXPECT_EQ(elements[0].Get({1, 0}), 3); + EXPECT_EQ(elements[0].Get({1, 1}), 4); + + EXPECT_TRUE(ShapeUtil::Compatible( + elements[1].shape(), + ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(S32, {}), + ShapeUtil::MakeShape(F64, {2}), + ShapeUtil::MakeNil()}))); + EXPECT_EQ(elements[1].Get({}, /*shape_index=*/{0}), 42); + EXPECT_EQ(elements[1].Get({0}, /*shape_index=*/{1}), 23.0); + EXPECT_EQ(elements[1].Get({1}, /*shape_index=*/{1}), 44.0); + + EXPECT_TRUE(ShapeUtil::Compatible(elements[2].shape(), ShapeUtil::MakeNil())); +} + +TEST_F(LiteralUtilTest, DecomposeEmptyTuple) { + Literal nil_literal(ShapeUtil::MakeNil()); + std::vector elements = nil_literal.DecomposeTuple(); + EXPECT_EQ(elements.size(), 0); +} + +TEST_F(LiteralUtilTest, MoveIntoTuple) { + std::vector elements; + elements.push_back(std::move(*Literal::CreateR0(1.0))); + elements.push_back(std::move(*Literal::CreateR1({4, 8}))); + elements.push_back(std::move( + *Literal::MakeTuple({Literal::CreateR0(42).get(), + Literal::CreateR1({23.0, 44.0}).get()}) + + )); + + Literal literal = Literal::MoveIntoTuple(&elements); + ASSERT_TRUE(ShapeUtil::IsTuple(literal.shape())); + ASSERT_EQ(ShapeUtil::TupleElementCount(literal.shape()), 3); + + EXPECT_EQ(literal.Get({}, /*shape_index=*/{0}), 1.0); + EXPECT_EQ(literal.Get({0}, /*shape_index=*/{1}), 4); + EXPECT_EQ(literal.Get({1}, /*shape_index=*/{1}), 8); + EXPECT_EQ(literal.Get({}, /*shape_index=*/{2, 0}), 42); + EXPECT_EQ(literal.Get({0}, /*shape_index=*/{2, 1}), 23.0); + EXPECT_EQ(literal.Get({1}, /*shape_index=*/{2, 1}), 44.0); + + for (const Literal& element : elements) { + EXPECT_TRUE(ShapeUtil::IsNil(element.shape())); + } +} + +TEST_F(LiteralUtilTest, MoveIntoEmptyTuple) { + Literal literal = Literal::MoveIntoTuple({}); + ASSERT_TRUE(ShapeUtil::IsTuple(literal.shape())); + ASSERT_EQ(ShapeUtil::TupleElementCount(literal.shape()), 0); +} + +TEST_F(LiteralUtilTest, LiteralMoveAssignment) { + Literal literal; + EXPECT_TRUE(ShapeUtil::Equal(ShapeUtil::MakeNil(), literal.shape())); + + std::unique_ptr matrix = + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + literal = std::move(*matrix); + + EXPECT_TRUE( + ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {2, 2}), literal.shape())); + EXPECT_EQ(literal.Get({0, 0}), 1.0); + EXPECT_EQ(literal.Get({0, 1}), 2.0); + EXPECT_EQ(literal.Get({1, 0}), 3.0); + EXPECT_EQ(literal.Get({1, 1}), 4.0); +} + +TEST_F(LiteralUtilTest, LiteralViewCopy) { + std::unique_ptr matrix = + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + const auto matrix_view = LiteralView::Create(*matrix); + LiteralView matrix_view_copy(matrix_view); + + EXPECT_EQ(matrix_view_copy.Get({0, 0}), 1.0); + EXPECT_EQ(matrix_view_copy.Get({0, 1}), 2.0); + EXPECT_EQ(matrix_view_copy.Get({1, 0}), 3.0); + EXPECT_EQ(matrix_view_copy.Get({1, 1}), 4.0); +} + +TEST_F(LiteralUtilTest, GetSetTuple) { + auto tuple = Literal::MakeTuple( + {Literal::CreateR0(42.0).get(), + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}).get()}); + EXPECT_EQ(tuple->Get(/*multi_index=*/{}, /*shape_index=*/{0}), 42.0); + tuple->Set(/*multi_index=*/{}, /*shape_index=*/{0}, -5.0); + EXPECT_EQ(tuple->Get(/*multi_index=*/{}, /*shape_index=*/{0}), -5.0); + + EXPECT_EQ(tuple->Get(/*multi_index=*/{1, 0}, /*shape_index=*/{1}), + 3.0); + tuple->Set(/*multi_index=*/{1, 0}, /*shape_index=*/{1}, -4.0); + EXPECT_EQ(tuple->Get(/*multi_index=*/{1, 0}, /*shape_index=*/{1}), + -4.0); +} + +TEST_F(LiteralUtilTest, CreateFromShapeZeroInitialized) { + // Literals constructed using CreateFromShape should be zero initialized. + std::unique_ptr scalar_f32 = + Literal::CreateFromShape(ShapeUtil::MakeShape(F32, {})); + EXPECT_EQ(scalar_f32->Get({}), 0.0); + EXPECT_TRUE(scalar_f32->IsAll(0)); + + std::unique_ptr vector_s32 = + Literal::CreateFromShape(ShapeUtil::MakeShape(S32, {3})); + EXPECT_EQ(vector_s32->Get({0}), 0); + EXPECT_EQ(vector_s32->Get({1}), 0); + EXPECT_EQ(vector_s32->Get({2}), 0); + EXPECT_TRUE(vector_s32->IsAll(0)); + + std::unique_ptr tuple = + Literal::CreateFromShape(ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F64, {}), ShapeUtil::MakeShape(PRED, {2}), + ShapeUtil::MakeShape(U64, {2, 1}), ShapeUtil::MakeShape(C64, {})})); + + EXPECT_EQ(tuple->Get({}, {0}), 0.0); + EXPECT_EQ(tuple->Get({0}, {1}), false); + EXPECT_EQ(tuple->Get({1}, {1}), false); + EXPECT_EQ(tuple->Get({0, 0}, {2}), 0); + EXPECT_EQ(tuple->Get({1, 0}, {2}), 0); + EXPECT_EQ(tuple->Get({}, {3}), complex64(0.0f, 0.0f)); +} + +TEST_F(LiteralUtilTest, ProtoRoundTrip) { + // Test serializing then deserializing a Literal through a proto. + auto one_f32 = Literal::CreateR0(1.0); + auto two_f32 = Literal::CreateR0(2.0); + auto vector_int8 = Literal::CreateR1({-128, 0, 2, 4, 7, 56, 127}); + auto vector_c64 = Literal::CreateR1({{1.0, 2.0}, {3.0, 4.0}}); + auto vector_bfloat16 = Literal::CreateR1( + {bfloat16{-1.0}, bfloat16{2.0}, bfloat16{-3.0}}); + auto vector_half = + Literal::CreateR1({half{10.0}, half{20.0}, half{-30.0}}); + auto matrix_pred = + Literal::CreateR2({{true, false, true}, {false, false, true}}); + auto tuple = Literal::MakeTuple( + {one_f32.get(), vector_half.get(), matrix_pred.get(), matrix_pred.get()}); + Literal nil_literal(ShapeUtil::MakeNil()); + auto nested_tuple = Literal::MakeTuple( + {tuple.get(), vector_bfloat16.get(), tuple.get(), &nil_literal}); + + auto to_from_proto = [](const Literal& literal) -> Literal { + return std::move(*Literal::CreateFromProto(literal.ToProto()).ValueOrDie()); + }; + + EXPECT_EQ(*one_f32, to_from_proto(*one_f32)); + EXPECT_EQ(*vector_c64, to_from_proto(*vector_c64)); + EXPECT_EQ(*vector_bfloat16, to_from_proto(*vector_bfloat16)); + EXPECT_EQ(*matrix_pred, to_from_proto(*matrix_pred)); + EXPECT_EQ(*tuple, to_from_proto(*tuple)); + EXPECT_EQ(*nested_tuple, to_from_proto(*nested_tuple)); + EXPECT_EQ(nil_literal, to_from_proto(nil_literal)); + + EXPECT_NE(*one_f32, *two_f32); + EXPECT_NE(*one_f32, to_from_proto(*two_f32)); +} + +TEST_F(LiteralUtilTest, InvalidProtoNoValues) { + // Proto contains a shape, but no values. + LiteralProto proto; + *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {3}); + Status status = Literal::CreateFromProto(proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), + HasSubstr("Expected 3 elements in LiteralProto")); +} + +TEST_F(LiteralUtilTest, InvalidProtoNoShape) { + // Proto contains values, but no shape. + LiteralProto proto; + proto.add_preds(false); + proto.add_preds(true); + proto.add_preds(false); + Status status = Literal::CreateFromProto(proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), HasSubstr("LiteralProto has no shape")); +} + +TEST_F(LiteralUtilTest, InvalidProtoWrongContainer) { + // Proto contains values in wrong container. + LiteralProto proto; + *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {3}); + proto.add_preds(false); + proto.add_preds(true); + proto.add_preds(false); + Status status = Literal::CreateFromProto(proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), + HasSubstr("Expected 3 elements in LiteralProto")); +} + +TEST_F(LiteralUtilTest, InvalidProtoTooFewValues) { + // Proto contains too few values. + LiteralProto proto; + *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {42, 2}); + proto.add_f32s(1.0); + proto.add_f32s(2.0); + proto.add_f32s(3.0); + Status status = Literal::CreateFromProto(proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), + HasSubstr("Expected 84 elements in LiteralProto")); +} + +TEST_F(LiteralUtilTest, InvalidProtoTooManyValues) { + // Proto contains too many values. + LiteralProto proto; + *proto.mutable_shape() = ShapeUtil::MakeShape(S32, {2}); + proto.add_s32s(42); + proto.add_s32s(-10); + proto.add_s32s(100); + Status status = Literal::CreateFromProto(proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), + HasSubstr("Expected 2 elements in LiteralProto")); +} + +TEST_F(LiteralUtilTest, InvalidProtoMissingLayout) { + // Proto shape missing layout. + LiteralProto proto; + *proto.mutable_shape() = ShapeUtil::MakeShape(PRED, {2, 2}); + LayoutUtil::ClearLayout(proto.mutable_shape()); + proto.add_preds(true); + proto.add_preds(false); + proto.add_preds(true); + proto.add_preds(false); + Status status = Literal::CreateFromProto(proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), HasSubstr("LiteralProto has no layout")); +} + +TEST_F(LiteralUtilTest, InvalidProtoTooFewTupleElements) { + // Proto has the too few tuple elements. + LiteralProto proto; + *proto.mutable_shape() = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(PRED, {2}), ShapeUtil::MakeShape(F32, {})}); + LiteralProto* element0 = proto.add_tuple_literals(); + *element0->mutable_shape() = + ShapeUtil::GetTupleElementShape(proto.shape(), 0); + element0->add_preds(false); + element0->add_preds(true); + + Status status = Literal::CreateFromProto(proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), HasSubstr("Expected 2 tuple elements")); +} + +TEST_F(LiteralUtilTest, InvalidProtoTooManyTupleElements) { + // Proto has the too many tuple elements. + LiteralProto proto; + *proto.mutable_shape() = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(PRED, {2}), ShapeUtil::MakeShape(F32, {})}); + LiteralProto* element0 = proto.add_tuple_literals(); + *element0->mutable_shape() = + ShapeUtil::GetTupleElementShape(proto.shape(), 0); + element0->add_preds(false); + element0->add_preds(true); + LiteralProto* element1 = proto.add_tuple_literals(); + *element1->mutable_shape() = + ShapeUtil::GetTupleElementShape(proto.shape(), 1); + element1->add_f32s(42.0); + LiteralProto* element2 = proto.add_tuple_literals(); + *element2->mutable_shape() = ShapeUtil::MakeShape(F32, {}); + element2->add_f32s(123.0); + + Status status = Literal::CreateFromProto(proto).status(); + ASSERT_FALSE(status.ok()); + ASSERT_THAT(status.error_message(), HasSubstr("Expected 2 tuple elements")); +} + +TEST_F(LiteralUtilTest, SortSparseElements) { + auto literal = + Literal::CreateSparse({10, 10, 10}, SparseIndexArray(10, 3), {}); + literal->AppendSparseElement({2, 3, 4}, 2.0); + literal->AppendSparseElement({3, 4, 5}, 3.0); + literal->AppendSparseElement({1, 2, 3}, 1.0); + literal->SortSparseElements(); + ASSERT_EQ(literal->ToString(false), + "f32[10,10,10]{[1, 2, 3]: 1, [2, 3, 4]: 2, [3, 4, 5]: 3}"); +} - EXPECT_EQ(nested_tuple->GetSubliteral(/*index=*/{0}), *tuple); - EXPECT_EQ(nested_tuple->GetSubliteral(/*index=*/{0, 0}), *scalar); - EXPECT_EQ(nested_tuple->GetSubliteral(/*index=*/{0, 1}), *matrix); - EXPECT_EQ(nested_tuple->GetSubliteral(/*index=*/{1}), *scalar); +TEST_F(LiteralUtilTest, GetSparseElementAsString) { + std::vector dimensions = {10, 10, 10}; + SparseIndexArray indices(10, {{1, 2, 3}, {2, 3, 4}, {3, 4, 5}}); + + ASSERT_EQ( + Literal::CreateSparse(dimensions, indices, {true, false, true}) + ->GetSparseElementAsString(1), + "false"); + ASSERT_EQ(Literal::CreateSparse(dimensions, indices, {1, 2, 3}) + ->GetSparseElementAsString(1), + tensorflow::strings::StrCat(int64{2})); + ASSERT_EQ(Literal::CreateSparse(dimensions, indices, {1.0, 2.0, 3.0}) + ->GetSparseElementAsString(1), + tensorflow::strings::StrCat(double{2.0})); + ASSERT_EQ(Literal::CreateSparse(dimensions, indices, + {half{1.0}, half{2.0}, half{3.0}}) + ->GetSparseElementAsString(1), + tensorflow::strings::StrCat(half{2.0})); + ASSERT_EQ( + Literal::CreateSparse( + dimensions, indices, + std::vector{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}}) + ->GetSparseElementAsString(1), + tensorflow::strings::StrCat("(", float{3.0}, ", ", float{4.0}, ")")); } } // namespace diff --git a/tensorflow/compiler/xla/map_util.h b/tensorflow/compiler/xla/map_util.h index 51d0d5f86f00c539951e8e2baa6296337a5a21e9..8db8c6f3de84a6c46625eadbb6b0f83d2262e5f7 100644 --- a/tensorflow/compiler/xla/map_util.h +++ b/tensorflow/compiler/xla/map_util.h @@ -16,6 +16,11 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ #define TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ +#include +#include + +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/platform/logging.h" namespace xla { @@ -44,6 +49,41 @@ typename Collection::value_type::second_type& FindOrDie( return it->second; } +// Like FindOrDie but returns an error instead of dying if `key` is not in +// `container`. +template +StatusOr< + std::reference_wrapper> +MaybeFind(const Collection& collection, + const typename Collection::value_type::first_type& key) { + typename Collection::const_iterator it = collection.find(key); + if (it == collection.end()) { + std::ostringstream os; + os << key; + return NotFound("key not found: %s", os.str().c_str()); + } + return {it->second}; +} + +// Returns a const reference to the value associated with the given key if it +// exists, otherwise returns a const reference to the provided default value. +// +// WARNING: If a temporary object is passed as the default "value," +// this function will return a reference to that temporary object, +// which will be destroyed at the end of the statement. A common +// example: if you have a map with string values, and you pass a char* +// as the default "value," either use the returned value immediately +// or store it in a string (not string&). +template +const typename Collection::value_type::second_type& FindOrDefault( + const Collection& collection, + const typename Collection::value_type::first_type& key, + const typename Collection::value_type::second_type& value) { + auto it = collection.find(key); + if (it != collection.end()) return it->second; + return value; +} + // Inserts the key-value pair into the collection. Dies if key was already // present. template @@ -60,6 +100,12 @@ bool ContainsKey(const Collection& collection, const Key& key) { return collection.find(key) != collection.end(); } +// Inserts `value` into `set`. Dies if it was already present. +template +void InsertOrDie(Set* const set, const typename Set::value_type& value) { + CHECK(set->insert(value).second) << "duplicate value: " << value; +} + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_MAP_UTIL_H_ diff --git a/tensorflow/compiler/xla/packed_literal_reader.cc b/tensorflow/compiler/xla/packed_literal_reader.cc index 70e0f5a74711c8ceef1b6d4225141aa1cc9c6219..857aae0a7982a57bb3057a6f267f5f033a0fdde4 100644 --- a/tensorflow/compiler/xla/packed_literal_reader.cc +++ b/tensorflow/compiler/xla/packed_literal_reader.cc @@ -44,11 +44,11 @@ StatusOr> PackedLiteralReader::Read( VLOG(3) << "reading shape from file: " << ShapeUtil::HumanString(shape) << " layout: " << (layout == nullptr ? "" : layout->ShortDebugString()); - auto result = MakeUnique(); - *result->mutable_shape() = shape; + Shape literal_shape = shape; if (layout != nullptr) { - TF_RETURN_IF_ERROR(LayoutUtil::ValidateLayoutForShape(*layout, shape)); - *result->mutable_shape()->mutable_layout() = *layout; + TF_RETURN_IF_ERROR( + LayoutUtil::ValidateLayoutForShape(*layout, literal_shape)); + *literal_shape.mutable_layout() = *layout; } if (shape.element_type() != F32) { @@ -57,10 +57,12 @@ StatusOr> PackedLiteralReader::Read( PrimitiveType_Name(shape.element_type()).c_str()); } + auto result = MakeUnique(literal_shape); + result->PopulateWithValue(std::numeric_limits::quiet_NaN()); + int64 elements = ShapeUtil::ElementsIn(shape); - result->Resize(elements, std::numeric_limits::quiet_NaN()); - std::vector* field = result->mutable_f32s(); - char* data = tensorflow::bit_cast(field->data()); + tensorflow::gtl::ArraySlice field = result->data(); + char* data = tensorflow::bit_cast(field.data()); uint64 bytes = elements * sizeof(float); tensorflow::StringPiece sp; auto s = file_->Read(offset_, bytes, &sp, data); diff --git a/tensorflow/compiler/xla/primitive_util.cc b/tensorflow/compiler/xla/primitive_util.cc index 2bce56b7bd2f91f20ea670d0e7ccaa432c2b5f9f..143c9a2366be5786b7ef2148580caeb97d67d2d8 100644 --- a/tensorflow/compiler/xla/primitive_util.cc +++ b/tensorflow/compiler/xla/primitive_util.cc @@ -20,79 +20,6 @@ limitations under the License. namespace xla { namespace primitive_util { -template <> -PrimitiveType NativeToPrimitiveType() { - return PRED; -} - -// Unsigned integer -template <> -PrimitiveType NativeToPrimitiveType() { - return U8; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return U16; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return U32; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return U64; -} - -// Signed integer -template <> -PrimitiveType NativeToPrimitiveType() { - return S8; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return S16; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return S32; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return S64; -} - -// Floating point -template <> -PrimitiveType NativeToPrimitiveType() { - return F32; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return F64; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return BF16; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return F16; -} - -template <> -PrimitiveType NativeToPrimitiveType() { - return C64; -} - bool IsFloatingPointType(PrimitiveType type) { return type == F16 || type == F32 || type == F64 || type == BF16; } diff --git a/tensorflow/compiler/xla/primitive_util.h b/tensorflow/compiler/xla/primitive_util.h index cb4583d198b454be1432134a9f6a77dbbbe5bdd8..b26a10ade63a5dad3bf8f9f3a2a33c3c5e67bdb2 100644 --- a/tensorflow/compiler/xla/primitive_util.h +++ b/tensorflow/compiler/xla/primitive_util.h @@ -47,49 +47,81 @@ PrimitiveType NativeToPrimitiveType() { } // Declarations of specializations for each native type which correspond to a -// XLA primitive type. +// XLA primitive type. As an optimization, these are declared inline in the +// header. template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return PRED; +} // Unsigned integer template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return U8; +} template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return U16; +} template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return U32; +} template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return U64; +} // Signed integer template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return S8; +} template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return S16; +} template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return S32; +} template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return S64; +} // Floating point template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return F32; +} + template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return F64; +} + template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return F16; +} + template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return BF16; +} // Complex template <> -PrimitiveType NativeToPrimitiveType(); +inline PrimitiveType NativeToPrimitiveType() { + return C64; +} bool IsFloatingPointType(PrimitiveType type); diff --git a/tensorflow/compiler/xla/python/BUILD b/tensorflow/compiler/xla/python/BUILD index 3b674a42722912711901f95aadd8df7f320d9263..e2972f06016ab3555c4fc0cc4616993fe6764b1e 100644 --- a/tensorflow/compiler/xla/python/BUILD +++ b/tensorflow/compiler/xla/python/BUILD @@ -49,7 +49,10 @@ cc_library( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:executable_build_options", "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/service:shaped_buffer", + "//tensorflow/core:framework_lite", "//tensorflow/core:lib", ], ) diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc index 522e5ec158d90607029dab09d2ecb679bd01ab27..cb7bb21e092c80d7360c23f3d6b00409a75dce23 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.cc +++ b/tensorflow/compiler/xla/python/local_computation_builder.cc @@ -16,37 +16,230 @@ limitations under the License. #include "tensorflow/compiler/xla/python/local_computation_builder.h" #include "tensorflow/compiler/xla/executable_run_options.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/platform/default/thread_annotations.h" namespace xla { namespace swig { +// TODO(b/34473877) Ideally XLA would support AllReduce among arbitrary sets of +// device handles instead of needing to set the number of replicas at XLA +// service initialization time. +tensorflow::mutex g_local_client_mutex(tensorflow::LINKER_INITIALIZED); +int g_replica_count GUARDED_BY(g_local_client_mutex) = 1; +LocalClient* g_local_client GUARDED_BY(g_local_client_mutex) = nullptr; + +Status InitializeReplicaCount(int replica_count) { + if (replica_count < 1) { + return InvalidArgument("Replica count must be >= 1; got %d.", + replica_count); + } + tensorflow::mutex_lock lock(g_local_client_mutex); + if (g_local_client != nullptr) { + return FailedPrecondition( + "Attempted to set the replica count to %d, but a local XLA service was " + "previously created with a replica count of %d.", + replica_count, g_replica_count); + } + g_replica_count = replica_count; + return Status::OK(); +} + +int GetReplicaCount() { + tensorflow::mutex_lock lock(g_local_client_mutex); + return g_replica_count; +} + +LocalClient* GetOrCreateLocalClient() { + tensorflow::mutex_lock lock(g_local_client_mutex); + if (g_local_client != nullptr) { + return g_local_client; + } + LocalClientOptions options; + options.set_number_of_replicas(g_replica_count); + g_local_client = ClientLibrary::GetOrCreateLocalClient(options).ValueOrDie(); + CHECK(g_local_client != nullptr); + return g_local_client; +} + +Status TransferToInfeedLocal(const Literal& literal) { + VLOG(1) << "Infeeding literal without replica number; shape: " + << literal.shape(); + LocalClient* client = GetOrCreateLocalClient(); + return client->TransferToInfeedLocal(literal, /*device_ordinal=*/0); +} + +Status TransferToInfeedLocalReplica(const Literal& literal, + int replica_number) { + VLOG(1) << "Infeeding shape " << literal.shape() + << " to replica number: " << replica_number; + LocalClient* client = GetOrCreateLocalClient(); + TF_ASSIGN_OR_RETURN(int device_ordinal, + client->ReplicaNumberToDeviceOrdinal(replica_number)); + return client->TransferToInfeedLocal(literal, device_ordinal); +} + +StatusOr> TransferFromOutfeedLocalReplica( + const Shape& shape, int replica_number) { + VLOG(1) << "Outfeeding literal from replica number: " << replica_number + << " shape: " << shape; + LocalClient* client = GetOrCreateLocalClient(); + TF_ASSIGN_OR_RETURN(int device_ordinal, + client->ReplicaNumberToDeviceOrdinal(replica_number)); + return client->TransferFromOutfeedLocal(shape, device_ordinal); +} + +LocalShapedBuffer::LocalShapedBuffer( + std::unique_ptr shaped_buffer) + : shaped_buffer_(std::move(shaped_buffer)) {} + +const std::unique_ptr& LocalShapedBuffer::shaped_buffer() + const { + return shaped_buffer_; +} + +static StatusOr> ToBuffer( + LocalClient* client, int device_ordinal, const Literal& arg) { + return client->LiteralToShapedBuffer(arg, device_ordinal, + client->backend().memory_allocator()); +} + +/* static */ +LocalShapedBuffer* LocalShapedBuffer::FromLiteral( + const Literal& argument, + const tensorflow::gtl::optional& shape_with_layout) { + LocalClient* client = GetOrCreateLocalClient(); + std::unique_ptr buf; + if (shape_with_layout) { + std::unique_ptr relaid = + argument.Relayout(shape_with_layout.value()); + buf = ToBuffer(client, /*device_ordinal=*/0, *relaid).ConsumeValueOrDie(); + } else { + buf = ToBuffer(client, /*device_ordinal=*/0, argument).ConsumeValueOrDie(); + } + return new LocalShapedBuffer(std::move(buf)); +} + +std::unique_ptr LocalShapedBuffer::ToLiteral() const { + LocalClient* client = GetOrCreateLocalClient(); + return client->ShapedBufferToLiteral(*shaped_buffer()).ConsumeValueOrDie(); +} + CompiledLocalComputation::CompiledLocalComputation( std::unique_ptr executable) : executable_(std::move(executable)) {} -std::unique_ptr CompiledLocalComputation::Execute( - const std::vector& arguments) { - LocalClient* client = ClientLibrary::LocalClientOrDie(); - - // Transfer arguments in - std::vector> scoped_buffers; - scoped_buffers.reserve(arguments.size()); - for (const Literal& argument : arguments) { - scoped_buffers.push_back( - client - ->LiteralToShapedBuffer(argument, - /*device_ordinal=*/0, - client->backend().memory_allocator()) - .ConsumeValueOrDie()); +StatusOr> CompiledLocalComputation::Execute( + const std::vector& arguments, + const std::vector>& shapes_with_layout) { + LocalClient* client = GetOrCreateLocalClient(); + + VLOG(1) << "Execution requested with " << GetReplicaCount() << " replicas."; + + // Each replica populates a StatusOr result, but only replica zero actually + // retrieves its literal value. + std::vector>> results(GetReplicaCount()); + { + tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(), "xlarun", + GetReplicaCount()); + + for (int replica = 0; replica < GetReplicaCount(); ++replica) { + pool.Schedule([this, client, replica, &arguments, &shapes_with_layout, + &results] { + StatusOr device_ordinal_status = + client->ReplicaNumberToDeviceOrdinal(replica); + if (!device_ordinal_status.ok()) { + results[replica] = device_ordinal_status.status(); + return; + } + const int device_ordinal = device_ordinal_status.ValueOrDie(); + VLOG(3) << "Replica " << replica + << " mapped to device ordinal for execution: " + << device_ordinal; + + // Transfer arguments in + std::vector> scoped_buffers; + scoped_buffers.reserve(arguments.size()); + for (int i = 0; i < arguments.size(); ++i) { + const Literal& argument = arguments[i]; + const tensorflow::gtl::optional& shape_with_layout = + shapes_with_layout[i]; + + StatusOr> pushed; + if (shape_with_layout) { + std::unique_ptr relaid = + argument.Relayout(shape_with_layout.value()); + pushed = ToBuffer(client, device_ordinal, *relaid); + } else { + pushed = ToBuffer(client, device_ordinal, argument); + } + if (!pushed.ok()) { + results[replica] = pushed.status(); + return; + } + + scoped_buffers.push_back(std::move(pushed).ValueOrDie()); + } + + // Execute + std::vector argument_buffers; + argument_buffers.reserve(scoped_buffers.size()); + for (auto& buffer : scoped_buffers) { + argument_buffers.push_back(buffer.get()); + } + + DeviceAssignment device_assignment = + client->backend() + .computation_placer() + ->AssignDevices(GetReplicaCount(), /*computation_count=*/1) + .ConsumeValueOrDie(); + + ExecutableRunOptions options; + options.set_device_ordinal(device_ordinal); + options.set_allocator(client->backend().memory_allocator()); + options.set_inter_op_thread_pool( + client->backend().inter_op_thread_pool()); + options.set_intra_op_thread_pool( + client->backend().eigen_intra_op_thread_pool_device()); + options.set_device_assignment(&device_assignment); + StatusOr> result_buffer_status = + executable_->Run(argument_buffers, options); + if (!result_buffer_status.ok()) { + results[replica] = result_buffer_status.status(); + return; + } + + // Transfer result out + results[replica] = + client->ShapedBufferToLiteral(*result_buffer_status.ValueOrDie()); + }); + } } - // Execute + for (int replica = 0; replica < GetReplicaCount(); ++replica) { + const auto& statusor = results[replica]; + if (!statusor.ok()) { + return InternalError( + "Failed running replica %d (other replicas may have failed as well): " + "%s.", + replica, statusor.status().ToString().c_str()); + } + } + + return std::move(results[0]); +} + +LocalShapedBuffer* CompiledLocalComputation::ExecuteWithShapedBuffers( + tensorflow::gtl::ArraySlice argument_handles) { + LocalClient* client = GetOrCreateLocalClient(); + std::vector argument_buffers; - argument_buffers.reserve(scoped_buffers.size()); - for (auto& buffer : scoped_buffers) { - argument_buffers.push_back(buffer.get()); + argument_buffers.reserve(argument_handles.size()); + for (auto& handle : argument_handles) { + argument_buffers.push_back(handle->shaped_buffer().get()); } + + // Execute ExecutableRunOptions options; options.set_allocator(client->backend().memory_allocator()); options.set_inter_op_thread_pool(client->backend().inter_op_thread_pool()); @@ -55,38 +248,54 @@ std::unique_ptr CompiledLocalComputation::Execute( std::unique_ptr result_buffer = executable_->Run(argument_buffers, options).ConsumeValueOrDie(); - // Transfer result out - return client->ShapedBufferToLiteral(*result_buffer).ConsumeValueOrDie(); + return new LocalShapedBuffer(std::move(result_buffer)); } -LocalComputation::LocalComputation(std::unique_ptr computation) +LocalComputation::LocalComputation(Computation computation) : computation_(std::move(computation)) {} -CompiledLocalComputation* LocalComputation::Compile( - const std::vector& argument_shapes) { +StatusOr LocalComputation::Compile( + const std::vector& argument_shapes, + const ExecutableBuildOptions* build_options) { std::vector argument_shape_pointers; argument_shape_pointers.reserve(argument_shapes.size()); for (auto& argument_shape : argument_shapes) { argument_shape_pointers.push_back(&argument_shape); } - LocalClient* client = ClientLibrary::LocalClientOrDie(); + LocalClient* client = GetOrCreateLocalClient(); ExecutableBuildOptions options; - return new CompiledLocalComputation( - client->Compile(*computation_, argument_shape_pointers, options) - .ValueOrDie()); + if (build_options != nullptr) { + options = *build_options; + } + TF_ASSIGN_OR_RETURN( + auto local_executable, + client->Compile(computation_, argument_shape_pointers, options)); + return new CompiledLocalComputation(std::move(local_executable)); } const Computation& LocalComputation::computation() const { - return *computation_; + return computation_; +} + +StatusOr LocalComputation::GetReturnValueShape() const { + TF_ASSIGN_OR_RETURN(ProgramShape program_shape, + computation_.GetProgramShape()); + return std::move(*program_shape.mutable_result()); } LocalComputationBuilder::LocalComputationBuilder(const string& computation_name) - : builder_(ClientLibrary::LocalClientOrDie(), computation_name) {} + : builder_(GetOrCreateLocalClient(), computation_name) {} -LocalComputation* LocalComputationBuilder::Build() { - return new LocalComputation(std::unique_ptr( - new Computation(builder_.Build().ConsumeValueOrDie()))); +void LocalComputationBuilder::SetOpMetadata(const OpMetadata& metadata) { + builder_.SetOpMetadata(metadata); +} + +void LocalComputationBuilder::ClearOpMetadata() { builder_.ClearOpMetadata(); } + +StatusOr LocalComputationBuilder::Build() { + TF_ASSIGN_OR_RETURN(Computation computation, builder_.Build()); + return new LocalComputation(std::move(computation)); } ComputationDataHandle LocalComputationBuilder::Parameter(int64 parameter_number, @@ -100,6 +309,21 @@ std::unique_ptr LocalComputationBuilder::GetShape( return builder_.GetShape(operand).ConsumeValueOrDie(); } +StatusOr LocalComputationBuilder::GetReturnValueShape() { + TF_ASSIGN_OR_RETURN(ProgramShape program_shape, builder_.GetProgramShape()); + return program_shape.result(); +} + +ComputationDataHandle LocalComputationBuilder::Infeed(const Shape& shape) { + return builder_.Infeed(shape); +} + +void LocalComputationBuilder::Outfeed(const ComputationDataHandle& operand, + const Shape& shape, + const string& outfeed_config) { + builder_.Outfeed(operand, shape, outfeed_config); +} + ComputationDataHandle LocalComputationBuilder::ConstantLiteral( const Literal& literal) { return builder_.ConstantLiteral(literal); @@ -111,6 +335,13 @@ ComputationDataHandle LocalComputationBuilder::Broadcast( return builder_.Broadcast(operand, broadcast_sizes); } +ComputationDataHandle LocalComputationBuilder::Pad( + const ComputationDataHandle& operand, + const ComputationDataHandle& padding_value, + const PaddingConfig& padding_config) { + return builder_.Pad(operand, padding_value, padding_config); +} + ComputationDataHandle LocalComputationBuilder::Reshape( const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice dimensions, @@ -156,10 +387,17 @@ ComputationDataHandle LocalComputationBuilder::ConcatInDim( return builder_.ConcatInDim(operands, dimension); } -ComputationDataHandle LocalComputationBuilder::Select( - const ComputationDataHandle& pred, const ComputationDataHandle& on_true, - const ComputationDataHandle& on_false) { - return builder_.Select(pred, on_true, on_false); +ComputationDataHandle +LocalComputationBuilder::SelectAndScatterWithGeneralPadding( + const ComputationDataHandle& operand, const LocalComputation& select, + tensorflow::gtl::ArraySlice window_dimensions, + tensorflow::gtl::ArraySlice window_strides, + tensorflow::gtl::ArraySlice> padding, + const ComputationDataHandle& source, + const ComputationDataHandle& init_value, const LocalComputation& scatter) { + return builder_.SelectAndScatterWithGeneralPadding( + operand, select.computation(), window_dimensions, window_strides, padding, + source, init_value, scatter.computation()); } ComputationDataHandle LocalComputationBuilder::Tuple( @@ -177,6 +415,12 @@ ComputationDataHandle LocalComputationBuilder::Dot( return builder_.Dot(lhs, rhs); } +ComputationDataHandle LocalComputationBuilder::DotGeneral( + const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, + const DotDimensionNumbers& dimension_numbers) { + return builder_.DotGeneral(lhs, rhs, dimension_numbers); +} + ComputationDataHandle LocalComputationBuilder::ConvGeneralDilated( const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, tensorflow::gtl::ArraySlice window_strides, @@ -230,12 +474,47 @@ ComputationDataHandle LocalComputationBuilder::Reduce( dimensions_to_reduce); } +ComputationDataHandle LocalComputationBuilder::ReduceWindowWithGeneralPadding( + const ComputationDataHandle& operand, + const ComputationDataHandle& init_value, + const LocalComputation& local_computation, + tensorflow::gtl::ArraySlice window_dimensions, + tensorflow::gtl::ArraySlice window_strides, + tensorflow::gtl::ArraySlice> padding) { + return builder_.ReduceWindowWithGeneralPadding( + operand, init_value, local_computation.computation(), window_dimensions, + window_strides, padding); +} + +ComputationDataHandle LocalComputationBuilder::RngNormal( + const ComputationDataHandle& mu, const ComputationDataHandle& sigma, + const Shape& shape) { + return builder_.RngNormal(mu, sigma, shape); +} + +ComputationDataHandle LocalComputationBuilder::RngUniform( + const ComputationDataHandle& a, const ComputationDataHandle& b, + const Shape& shape) { + return builder_.RngUniform(a, b, shape); +} + ComputationDataHandle LocalComputationBuilder::While( const LocalComputation& condition, const LocalComputation& body, const ComputationDataHandle& init) { return builder_.While(condition.computation(), body.computation(), init); } +ComputationDataHandle LocalComputationBuilder::Conditional( + const ComputationDataHandle& predicate, + const ComputationDataHandle& true_operand, + const LocalComputation& true_computation, + const ComputationDataHandle& false_operand, + const LocalComputation& false_computation) { + return builder_.Conditional(predicate, true_operand, + true_computation.computation(), false_operand, + false_computation.computation()); +} + #define _FORWARD(method_name, return_sig, args_sig, args) \ return_sig LocalComputationBuilder::method_name args_sig { \ return builder_.method_name args; \ @@ -252,6 +531,15 @@ ComputationDataHandle LocalComputationBuilder::While( tensorflow::gtl::ArraySlice broadcast_dimensions), \ (lhs, rhs, broadcast_dimensions)) +#define _FORWARD_TRIOP(method_name) \ + _FORWARD( \ + method_name, ComputationDataHandle, \ + (const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, \ + const ComputationDataHandle& ehs), \ + (lhs, rhs, ehs)) + +_FORWARD_TRIOP(Select) +_FORWARD_TRIOP(Clamp) _FORWARD_BINOP(Eq) _FORWARD_BINOP(Ne) _FORWARD_BINOP(Ge) @@ -272,6 +560,7 @@ _FORWARD_UNOP(Abs) _FORWARD_UNOP(Exp) _FORWARD_UNOP(Floor) _FORWARD_UNOP(Ceil) +_FORWARD_UNOP(Round) _FORWARD_UNOP(Log) _FORWARD_UNOP(Sign) _FORWARD_UNOP(Cos) @@ -288,6 +577,19 @@ _FORWARD_UNOP(Sort) #undef _FORWARD #undef _FORWARD_UNOP #undef _FORWARD_BINOP +#undef _FORWARD_TRIOP + +void DeleteLocalShapedBuffer(LocalShapedBuffer* local_shaped_buffer) { + delete local_shaped_buffer; +} + +void DeleteCompiledLocalComputation(CompiledLocalComputation* computation) { + delete computation; +} + +void DeleteLocalComputation(LocalComputation* computation) { + delete computation; +} } // namespace swig diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h index 5638d2e1943d0b7cbfd96df75f22fb630e54523a..d3e9503ea10b011520ec5148a756ef4d421f244c 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.h +++ b/tensorflow/compiler/xla/python/local_computation_builder.h @@ -18,7 +18,9 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/client/executable_build_options.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/service/shaped_buffer.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/gtl/array_slice.h" @@ -26,6 +28,48 @@ namespace xla { namespace swig { +// Initializes the number of replicas that XLA will be initialized with (when +// first obtaining a handle to the local XLA service). If this is called after +// the handle to the local XLA service has been established, then an error is +// returned. +Status InitializeReplicaCount(int replica_count); + +// Returns the replica count that is currently set, regardless of whether the +// local XLA service has been instantiated yet or not. +int GetReplicaCount(); + +// Wraps the local client's infeed-transfer function. +// +// The default device ordinal (0) is used. +Status TransferToInfeedLocal(const Literal& literal); + +// Transfers the given literal to the infeed of the given replica. +// +// The replica number is resolved to an appropriate device ordinal. +Status TransferToInfeedLocalReplica(const Literal& literal, int replica_number); + +// Transfers a literal of the given shape from the outfeed of the given replica. +// +// The replica number is resolved to an appropriate device ordinal. +StatusOr > TransferFromOutfeedLocalReplica( + const Shape& shape, int replica_number); + +// Wraps a ScopedShapedBuffer produced by copying a literal "to +// device," i.e. copying a literal to a scoped buffer via the local +// client. +class LocalShapedBuffer { + public: + static LocalShapedBuffer* FromLiteral( + const Literal& argument, + const tensorflow::gtl::optional& shape_with_layout); + LocalShapedBuffer(std::unique_ptr shaped_buffer); + const std::unique_ptr& shaped_buffer() const; + std::unique_ptr ToLiteral() const; + + private: + std::unique_ptr shaped_buffer_; +}; + // Wraps a LocalExecutable produced by compiling a // LocalComputation. The Execute method forwards to that of the // underlying LocalExecutable, and additionally handles tranferring @@ -35,7 +79,17 @@ namespace swig { class CompiledLocalComputation { public: CompiledLocalComputation(std::unique_ptr executable); - std::unique_ptr Execute(const std::vector& arguments); + + // Execute the computation with the given argument literals, and + // with optionally-specified argument layouts. The literals will be + // re-laid out according to the corresponding elements of + // shapes_with_layout. + StatusOr > Execute( + const std::vector& arguments, + const std::vector >& shapes_with_layout); + + LocalShapedBuffer* ExecuteWithShapedBuffers( + tensorflow::gtl::ArraySlice argument_handles); private: std::unique_ptr executable_; @@ -47,12 +101,19 @@ class CompiledLocalComputation { // made available to Python via SWIG. class LocalComputation { public: - LocalComputation(std::unique_ptr computation); - CompiledLocalComputation* Compile(const std::vector& argument_shapes); + LocalComputation(Computation computation); + + StatusOr Compile( + const std::vector& argument_shapes, + const ExecutableBuildOptions* build_options); + const Computation& computation() const; + // Returns the return-value shape for this computation. + StatusOr GetReturnValueShape() const; + private: - std::unique_ptr computation_; + Computation computation_; }; // Wraps the ComputationBuilder API in order to: @@ -66,19 +127,35 @@ class LocalComputationBuilder { public: LocalComputationBuilder(const string& computation_name); - LocalComputation* Build(); + void SetOpMetadata(const OpMetadata& metadata); + void ClearOpMetadata(); + + // Returns an owned LocalComputation to the caller on success. + StatusOr Build(); ComputationDataHandle Parameter(int64 parameter_number, const Shape& shape, const string& name); std::unique_ptr GetShape(const ComputationDataHandle& operand); + // Returns the shape of the current return value for the computation. + StatusOr GetReturnValueShape(); + + ComputationDataHandle Infeed(const Shape& shape); + + void Outfeed(const ComputationDataHandle& operand, const Shape& shape, + const string& outfeed_config); + ComputationDataHandle ConstantLiteral(const Literal& literal); ComputationDataHandle Broadcast( const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice broadcast_sizes); + ComputationDataHandle Pad(const ComputationDataHandle& operand, + const ComputationDataHandle& padding_value, + const PaddingConfig& padding_config); + ComputationDataHandle Reshape(const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice dimensions, tensorflow::gtl::ArraySlice new_sizes); @@ -106,9 +183,13 @@ class LocalComputationBuilder { tensorflow::gtl::ArraySlice operands, int64 dimension); - ComputationDataHandle Select(const ComputationDataHandle& pred, - const ComputationDataHandle& on_true, - const ComputationDataHandle& on_false); + ComputationDataHandle SelectAndScatterWithGeneralPadding( + const ComputationDataHandle& operand, const LocalComputation& select, + tensorflow::gtl::ArraySlice window_dimensions, + tensorflow::gtl::ArraySlice window_strides, + tensorflow::gtl::ArraySlice > padding, + const ComputationDataHandle& source, + const ComputationDataHandle& init_value, const LocalComputation& scatter); ComputationDataHandle Tuple( tensorflow::gtl::ArraySlice elements); @@ -119,6 +200,10 @@ class LocalComputationBuilder { ComputationDataHandle Dot(const ComputationDataHandle& lhs, const ComputationDataHandle& rhs); + ComputationDataHandle DotGeneral( + const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, + const DotDimensionNumbers& dimension_numbers); + ComputationDataHandle ConvGeneralDilated( const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, tensorflow::gtl::ArraySlice window_strides, @@ -153,10 +238,32 @@ class LocalComputationBuilder { const LocalComputation& local_computation, tensorflow::gtl::ArraySlice dimensions_to_reduce); + ComputationDataHandle ReduceWindowWithGeneralPadding( + const ComputationDataHandle& operand, + const ComputationDataHandle& init_value, + const LocalComputation& local_computation, + tensorflow::gtl::ArraySlice window_dimensions, + tensorflow::gtl::ArraySlice window_strides, + tensorflow::gtl::ArraySlice > padding); + + ComputationDataHandle RngNormal(const ComputationDataHandle& mu, + const ComputationDataHandle& sigma, + const Shape& shape); + + ComputationDataHandle RngUniform(const ComputationDataHandle& a, + const ComputationDataHandle& b, + const Shape& shape); + ComputationDataHandle While(const LocalComputation& condition, const LocalComputation& body, const ComputationDataHandle& init); + ComputationDataHandle Conditional(const ComputationDataHandle& predicate, + const ComputationDataHandle& true_operand, + const LocalComputation& true_computation, + const ComputationDataHandle& false_operand, + const LocalComputation& false_computation); + #define _FORWARD(method_name, return_sig, args_sig) \ return_sig method_name args_sig; @@ -170,6 +277,14 @@ class LocalComputationBuilder { (const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, \ tensorflow::gtl::ArraySlice broadcast_dimensions)) +#define _FORWARD_TRIOP(method_name) \ + _FORWARD( \ + method_name, ComputationDataHandle, \ + (const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, \ + const ComputationDataHandle& ehs)) + + _FORWARD_TRIOP(Select) + _FORWARD_TRIOP(Clamp) _FORWARD_BINOP(Eq) _FORWARD_BINOP(Ne) _FORWARD_BINOP(Ge) @@ -190,6 +305,7 @@ class LocalComputationBuilder { _FORWARD_UNOP(Exp) _FORWARD_UNOP(Floor) _FORWARD_UNOP(Ceil) + _FORWARD_UNOP(Round) _FORWARD_UNOP(Log) _FORWARD_UNOP(Sign) _FORWARD_UNOP(Cos) @@ -206,19 +322,16 @@ class LocalComputationBuilder { #undef _FORWARD #undef _FORWARD_UNOP #undef _FORWARD_BINOP +#undef _FORWARD_TRIOP private: ComputationBuilder builder_; }; -static void DeleteLocalComputation(LocalComputation* computation) { - delete computation; -} - -static void DeleteCompiledLocalComputation( - CompiledLocalComputation* computation) { - delete computation; -} +// Functions for freeing resources from the Python side. +void DeleteLocalShapedBuffer(LocalShapedBuffer* local_shaped_buffer); +void DeleteCompiledLocalComputation(CompiledLocalComputation* computation); +void DeleteLocalComputation(LocalComputation* computation); } // namespace swig diff --git a/tensorflow/compiler/xla/python/local_computation_builder.i b/tensorflow/compiler/xla/python/local_computation_builder.i index 8b2f84e045f9e1860253dcda290179485defe6b6..456e341f877e529f7fc5ebc81d85862bfa291943 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.i +++ b/tensorflow/compiler/xla/python/local_computation_builder.i @@ -27,11 +27,14 @@ limitations under the License. // ArraySlice <- sequence of int // Literal <-> (nested tuple of) numpy ndarray // std::vector <- sequence of (nested tuple of) ndarray -// Shape <-> pair holding (dtype, dimensions) -// std::vector <- sequence of shape information pairs +// Shape -> pair holding (dtype, dimensions) +// <- object duck-typed as xla_client.Shape +// std::vector <- sequence of xla_client.Shape objects // PrimitiveType <- int // ArraySlice> <- sequence of int pairs +// PaddingConfig proto <- corresponding Python proto // ConvolutionDimensionNumbers proto <- corresponding Python proto +// DotDimensionNumbers proto <- corresponding Python proto // // Arrows indicate whether a conversion only ever occurs in one // direction, or whether it is maintained bidirectionally. @@ -54,7 +57,7 @@ limitations under the License. // translates to a tuple-shaped XLA Literal, whose component subshapes // are a 2x3 F32-shaped literal followed by two tuple-shaped literals. // -// The Python objects corresponding to C++ Shapes have the type: +// Shapes output by C++ become Python objects with the type: // // T = (dtype, S) // S = DIMENSIONS | TUPLE_SHAPES @@ -99,6 +102,11 @@ limitations under the License. // wrapped by xla_client in order to set up a custom destructor that // triggers memory deallocation on the C++ side. +%module(threads="1") local_computation_builder + +// Keep the GIL except where explicitly specified. +%nothread; + %include "tensorflow/python/platform/base.i" %{ @@ -157,6 +165,60 @@ tensorflow::ImportNumpy(); $result = numpy::LongToPyIntOrPyLong($1.handle()); } +%typemap(out) StatusOr { + if ($1.ok()) { + auto* value = $1.ValueOrDie(); + { + auto* $1 = value; + $typemap(out, xla::swig::CompiledLocalComputation*) + } + } else { + PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); + return NULL; + } +} + +%typemap(out) StatusOr< std::unique_ptr > { + if ($1.ok()) { + std::unique_ptr value = $1.ConsumeValueOrDie(); + $result = numpy::PyObjectFromXlaLiteral(*value); + } else { + PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); + return NULL; + } +} + +%typemap(out) StatusOr { + if ($1.ok()) { + auto* value = $1.ValueOrDie(); + { + auto* $1 = value; + $typemap(out, xla::swig::LocalComputation*) + } + } else { + PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); + return NULL; + } +} + +%typemap(out) StatusOr { + if ($1.ok()) { + $result = numpy::PyShapeInfoFromXlaShape($1.ConsumeValueOrDie()); + } else { + PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); + return NULL; + } +} + +%typemap(out) Status { + if (!$1.ok()) { + PyErr_SetString( + PyExc_RuntimeError, $1.ToString().c_str()); + return NULL; + } + $result = Py_None; +} + // ArraySlice %typemap(in) tensorflow::gtl::ArraySlice @@ -221,17 +283,52 @@ tensorflow::ImportNumpy(); $1 = temps; } +// LocalShapedBuffer* + +%typemap(in) tensorflow::gtl::ArraySlice + (std::vector temps) { + if (!PySequence_Check($input)) { + PyErr_SetString(PyExc_TypeError, "Argument is not a sequence"); + return NULL; + } + const int size = PySequence_Size($input); + temps.reserve(size); + for (int i = 0; i < size; ++i) { + PyObject* o = PySequence_GetItem($input, i); + LocalShapedBuffer* lsbp; + if ((SWIG_ConvertPtr(o, (void**) &lsbp, $descriptor(xla::swig::LocalShapedBuffer*), + SWIG_POINTER_EXCEPTION)) == -1) { + return NULL; + } + temps.push_back(lsbp); + Py_DECREF(o); + } + $1 = temps; +} + // Literal -%typemap(in) const Literal& (std::unique_ptr temp) { - temp = numpy::XlaLiteralFromPyObject($input); - $1 = &*temp; +%typemap(in) const Literal& (StatusOr< std::unique_ptr > literal_status) { + literal_status = numpy::XlaLiteralFromPyObject($input); + if (!literal_status.ok()) { + PyErr_SetString(PyExc_RuntimeError, literal_status.status().ToString().c_str()); + return NULL; + } + $1 = literal_status.ValueOrDie().get(); } %typemap(out) std::unique_ptr { $result = numpy::PyObjectFromXlaLiteral(*$1); } +%typemap(out) StatusOr< std::unique_ptr > { + if (!$1.ok()) { + PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); + return NULL; + } + $result = numpy::PyObjectFromXlaLiteral(*$1.ValueOrDie()); +} + %typemap(in) const std::vector& (std::vector temps) { if (!PySequence_Check($input)) { PyErr_SetString(PyExc_TypeError, "Argument is not a sequence"); @@ -240,22 +337,58 @@ tensorflow::ImportNumpy(); const int size = PySequence_Size($input); for (int i = 0; i < size; ++i) { PyObject* o = PySequence_GetItem($input, i); - temps.push_back(*numpy::XlaLiteralFromPyObject(o)); + StatusOr< std::unique_ptr > literal_status = numpy::XlaLiteralFromPyObject(o); + if (!literal_status.ok()) { + PyErr_SetString(PyExc_RuntimeError, literal_status.status().ToString().c_str()); + Py_DECREF(o); + return NULL; + } + temps.push_back(std::move(*literal_status.ConsumeValueOrDie())); Py_DECREF(o); } $1 = &temps; } +// OpMetadata + +%typemap(in) const OpMetadata& (OpMetadata temp) { + StatusOr statusor = numpy::OpMetadataFromPyObject($input); + if (!statusor.ok()) { + PyErr_SetString(PyExc_RuntimeError, statusor.status().ToString().c_str()); + return NULL; + } + temp = std::move(statusor).ValueOrDie(); + $1 = &temp; +} + // Shape %typemap(in) const Shape& (Shape temp) { - if (!numpy::CheckPyShapeInfo($input)) { + StatusOr statusor = numpy::XlaShapeFromPyShape($input); + if (!statusor.ok()) { + PyErr_SetString(PyExc_RuntimeError, statusor.status().ToString().c_str()); return NULL; } - temp = numpy::XlaShapeFromPyShapeInfo($input); + temp = std::move(statusor).ValueOrDie(); $1 = &temp; } +%typemap(in) const tensorflow::gtl::optional& ( + tensorflow::gtl::optional temp) { + if ($input == Py_None) { + temp = tensorflow::gtl::nullopt; + $1 = &temp; + } else { + StatusOr statusor = numpy::XlaShapeFromPyShape($input); + if (!statusor.ok()) { + PyErr_SetString(PyExc_RuntimeError, statusor.status().ToString().c_str()); + return NULL; + } + temp = std::move(statusor).ValueOrDie(); + $1 = &temp; + } +} + %typemap(out) std::unique_ptr { $result = numpy::PyShapeInfoFromXlaShape(*$1); } @@ -268,12 +401,37 @@ tensorflow::ImportNumpy(); const int size = PySequence_Size($input); for (int i = 0; i < size; ++i) { PyObject* o = PySequence_GetItem($input, i); - if (!numpy::CheckPyShapeInfo(o)) { - Py_DECREF(o); + StatusOr statusor = numpy::XlaShapeFromPyShape(o); + Py_DECREF(o); + if (!statusor.ok()) { + PyErr_SetString(PyExc_RuntimeError, statusor.status().ToString().c_str()); return NULL; } - temps.push_back(numpy::XlaShapeFromPyShapeInfo(o)); - Py_DECREF(o); + temps.push_back(statusor.ConsumeValueOrDie()); + } + $1 = &temps; +} + +%typemap(in) const std::vector >& ( + std::vector > temps) { + if (!PySequence_Check($input)) { + PyErr_SetString(PyExc_TypeError, "Argument is not a sequence"); + return NULL; + } + const int size = PySequence_Size($input); + for (int i = 0; i < size; ++i) { + PyObject* o = PySequence_GetItem($input, i); + if (o == Py_None) { + temps.push_back(tensorflow::gtl::nullopt); + } else { + StatusOr statusor = numpy::XlaShapeFromPyShape(o); + Py_DECREF(o); + if (!statusor.ok()) { + PyErr_SetString(PyExc_RuntimeError, statusor.status().ToString().c_str()); + return NULL; + } + temps.push_back(statusor.ConsumeValueOrDie()); + } } $1 = &temps; } @@ -363,6 +521,177 @@ tensorflow::ImportNumpy(); $1 = temps; } +// DotDimensionNumbers + +%typemap(in) const DotDimensionNumbers& + (DotDimensionNumbers dimension_numbers) { + int length; + + /* lhs_contracting_dimensions */ + PyObject* lhs_contracting_dimensions = PyObject_GetAttrString( + $input, "lhs_contracting_dimensions"); + if (!lhs_contracting_dimensions) { + return NULL; + } + + length = PySequence_Size(lhs_contracting_dimensions); + if (length == -1) { + Py_DECREF(lhs_contracting_dimensions); + return NULL; + } + + for (int i = 0; i < length; ++i) { + PyObject* item = PySequence_GetItem(lhs_contracting_dimensions, i); + if (!item) { + Py_DECREF(lhs_contracting_dimensions); + return NULL; + } + const int64 dimension = numpy::PyIntOrPyLongToLong(item); + if (dimension == -1 && PyErr_Occurred()) { + Py_DECREF(item); + Py_DECREF(lhs_contracting_dimensions); + return NULL; + } + dimension_numbers.add_lhs_contracting_dimensions(dimension); + Py_DECREF(item); + } + Py_DECREF(lhs_contracting_dimensions); + + /* rhs_contracting_dimensions */ + PyObject* rhs_contracting_dimensions = PyObject_GetAttrString( + $input, "rhs_contracting_dimensions"); + if (!lhs_contracting_dimensions) { + return NULL; + } + + length = PySequence_Size(rhs_contracting_dimensions); + if (length == -1) { + Py_DECREF(rhs_contracting_dimensions); + return NULL; + } + + for (int i = 0; i < length; ++i) { + PyObject* item = PySequence_GetItem(rhs_contracting_dimensions, i); + if (!item) { + Py_DECREF(rhs_contracting_dimensions); + return NULL; + } + const int64 dimension = numpy::PyIntOrPyLongToLong(item); + if (dimension == -1 && PyErr_Occurred()) { + Py_DECREF(item); + Py_DECREF(rhs_contracting_dimensions); + return NULL; + } + dimension_numbers.add_rhs_contracting_dimensions(dimension); + Py_DECREF(item); + } + Py_DECREF(rhs_contracting_dimensions); + + /* lhs_batch_dimensions */ + PyObject* lhs_batch_dimensions = PyObject_GetAttrString( + $input, "lhs_batch_dimensions"); + if (!lhs_batch_dimensions) { + return NULL; + } + + length = PySequence_Size(lhs_batch_dimensions); + if (length == -1) { + Py_DECREF(lhs_batch_dimensions); + return NULL; + } + + for (int i = 0; i < length; ++i) { + PyObject* item = PySequence_GetItem(lhs_batch_dimensions, i); + if (!item) { + Py_DECREF(lhs_batch_dimensions); + return NULL; + } + const int64 dimension = numpy::PyIntOrPyLongToLong(item); + if (dimension == -1 && PyErr_Occurred()) { + Py_DECREF(item); + Py_DECREF(lhs_batch_dimensions); + return NULL; + } + dimension_numbers.add_lhs_batch_dimensions(dimension); + Py_DECREF(item); + } + Py_DECREF(lhs_batch_dimensions); + + /* rhs_batch_dimensions */ + PyObject* rhs_batch_dimensions = PyObject_GetAttrString( + $input, "rhs_batch_dimensions"); + if (!rhs_batch_dimensions) { + return NULL; + } + + length = PySequence_Size(rhs_batch_dimensions); + if (length == -1) { + Py_DECREF(rhs_batch_dimensions); + return NULL; + } + + for (int i = 0; i < length; ++i) { + PyObject* item = PySequence_GetItem(rhs_batch_dimensions, i); + if (!item) { + Py_DECREF(rhs_batch_dimensions); + return NULL; + } + const int64 dimension = numpy::PyIntOrPyLongToLong(item); + if (dimension == -1 && PyErr_Occurred()) { + Py_DECREF(item); + Py_DECREF(rhs_batch_dimensions); + return NULL; + } + dimension_numbers.add_rhs_batch_dimensions(dimension); + Py_DECREF(item); + } + Py_DECREF(rhs_batch_dimensions); + + $1 = &dimension_numbers; +} + +// PaddingConfig + +%typemap(in) const PaddingConfig& + (PaddingConfig padding_config) { + PyObject* dimensions = PyObject_GetAttrString($input, "dimensions"); + if (!dimensions) { + return NULL; + } + + int length = PySequence_Size(dimensions); + if (length == -1) { + Py_DECREF(dimensions); + return NULL; + } + + for (int i = 0; i < length; ++i) { + PyObject* item = PySequence_GetItem(dimensions, i); + if (!item) { + Py_DECREF(dimensions); + return NULL; + } + int64 edge_padding_low, edge_padding_high, interior_padding; + if (!GetIntAttr(item, "edge_padding_low", &edge_padding_low) + || !GetIntAttr(item, "edge_padding_high", &edge_padding_high) + || !GetIntAttr(item, "interior_padding", &interior_padding)) { + Py_DECREF(item); + Py_DECREF(dimensions); + return NULL; + } + Py_DECREF(item); + + PaddingConfig::PaddingConfigDimension* dimension = + padding_config.add_dimensions(); + dimension->set_edge_padding_low(edge_padding_low); + dimension->set_edge_padding_high(edge_padding_high); + dimension->set_interior_padding(interior_padding); + } + Py_DECREF(dimensions); + + $1 = &padding_config; +} + // ConvolutionDimensionNumbers %typemap(in) const ConvolutionDimensionNumbers& @@ -483,21 +812,76 @@ tensorflow::ImportNumpy(); $1 = &dimension_numbers; } +// ExecutableBuildOptions + +%typemap(in) const ExecutableBuildOptions* + (ExecutableBuildOptions build_options) { + if ($input == Py_None) { + $1 = NULL; + } else { + PyObject* o = PyObject_GetAttrString($input, "generate_hlo_graph"); + if (!o) { + return NULL; + } + if (o != Py_None) { + if (!PyString_Check(o)) { + PyErr_SetString(PyExc_TypeError, "ExecutableBuildOptions.generate_hlo_graph must be a string or None."); + return NULL; + } + build_options.set_generate_hlo_graph(PyString_AsString(o)); + } + Py_DECREF(o); + + o = PyObject_GetAttrString($input, "result_shape"); + if (o == nullptr) { + return nullptr; + } + if (o != Py_None) { + StatusOr statusor = numpy::XlaShapeFromPyShape(o); + if (!statusor.ok()) { + PyErr_SetString(PyExc_TypeError, tensorflow::strings::StrCat("ExecutableBuildOptions.result_shape could not be created from Python shape value: ", statusor.status().ToString()).c_str()); + Py_DECREF(o); + return NULL; + } + build_options.set_result_layout(statusor.ValueOrDie()); + } + Py_DECREF(o); + + $1 = &build_options; + } +} + %ignoreall %unignore xla; %unignore xla::swig; +%unignore xla::swig::InitializeReplicaCount; +%unignore xla::swig::GetReplicaCount; +%unignore xla::swig::TransferToInfeedLocal; +%unignore xla::swig::TransferToInfeedLocalReplica; +%unignore xla::swig::TransferFromOutfeedLocalReplica; +%unignore xla::swig::LocalShapedBuffer; +%unignore xla::swig::LocalShapedBuffer::FromLiteral; +%unignore xla::swig::LocalShapedBuffer::ToLiteral; %unignore xla::swig::CompiledLocalComputation; %unignore xla::swig::CompiledLocalComputation::Execute; +%unignore xla::swig::CompiledLocalComputation::ExecuteWithShapedBuffers; %unignore xla::swig::LocalComputation; %unignore xla::swig::LocalComputation::Compile; +%unignore xla::swig::LocalComputation::GetReturnValueShape; %unignore xla::swig::LocalComputationBuilder; %unignore xla::swig::LocalComputationBuilder::LocalComputationBuilder; %unignore xla::swig::LocalComputationBuilder::Build; +%unignore xla::swig::LocalComputationBuilder::SetOpMetadata; +%unignore xla::swig::LocalComputationBuilder::ClearOpMetadata; %unignore xla::swig::LocalComputationBuilder::Parameter; %unignore xla::swig::LocalComputationBuilder::GetShape; +%unignore xla::swig::LocalComputationBuilder::GetReturnValueShape; +%unignore xla::swig::LocalComputationBuilder::Infeed; +%unignore xla::swig::LocalComputationBuilder::Outfeed; %unignore xla::swig::LocalComputationBuilder::ConstantLiteral; %unignore xla::swig::LocalComputationBuilder::ConstantR0; %unignore xla::swig::LocalComputationBuilder::Broadcast; +%unignore xla::swig::LocalComputationBuilder::Pad; %unignore xla::swig::LocalComputationBuilder::Reshape; %unignore xla::swig::LocalComputationBuilder::Collapse; %unignore xla::swig::LocalComputationBuilder::CrossReplicaSum; @@ -505,6 +889,7 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::DynamicSlice; %unignore xla::swig::LocalComputationBuilder::DynamicUpdateSlice; %unignore xla::swig::LocalComputationBuilder::ConcatInDim; +%unignore xla::swig::LocalComputationBuilder::SelectAndScatterWithGeneralPadding; %unignore xla::swig::LocalComputationBuilder::Select; %unignore xla::swig::LocalComputationBuilder::Tuple; %unignore xla::swig::LocalComputationBuilder::GetTupleElement; @@ -512,9 +897,15 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::Call; %unignore xla::swig::LocalComputationBuilder::Transpose; %unignore xla::swig::LocalComputationBuilder::Rev; +%unignore xla::swig::LocalComputationBuilder::Clamp; %unignore xla::swig::LocalComputationBuilder::Map; %unignore xla::swig::LocalComputationBuilder::Reduce; +%unignore xla::swig::LocalComputationBuilder::ReduceWindowWithGeneralPadding; +%unignore xla::swig::LocalComputationBuilder::RngNormal; +%unignore xla::swig::LocalComputationBuilder::RngUniform; +%unignore xla::swig::LocalComputationBuilder::RngBernoulli; %unignore xla::swig::LocalComputationBuilder::While; +%unignore xla::swig::LocalComputationBuilder::Conditional; %unignore xla::swig::LocalComputationBuilder::Eq; %unignore xla::swig::LocalComputationBuilder::Ne; %unignore xla::swig::LocalComputationBuilder::Ge; @@ -522,6 +913,7 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::Lt; %unignore xla::swig::LocalComputationBuilder::Le; %unignore xla::swig::LocalComputationBuilder::Dot; +%unignore xla::swig::LocalComputationBuilder::DotGeneral; %unignore xla::swig::LocalComputationBuilder::ConvGeneralDilated; %unignore xla::swig::LocalComputationBuilder::Add; %unignore xla::swig::LocalComputationBuilder::Sub; @@ -537,6 +929,7 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::Exp; %unignore xla::swig::LocalComputationBuilder::Floor; %unignore xla::swig::LocalComputationBuilder::Ceil; +%unignore xla::swig::LocalComputationBuilder::Round; %unignore xla::swig::LocalComputationBuilder::Log; %unignore xla::swig::LocalComputationBuilder::Sign; %unignore xla::swig::LocalComputationBuilder::Cos; @@ -549,9 +942,12 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::ReciprocalF32; %unignore xla::swig::LocalComputationBuilder::Neg; %unignore xla::swig::LocalComputationBuilder::Sort; +%unignore xla::swig::DeleteLocalShapedBuffer; %unignore xla::swig::DeleteLocalComputation; %unignore xla::swig::DeleteCompiledLocalComputation; +%thread; %include "tensorflow/compiler/xla/python/local_computation_builder.h" +%nothread; %unignoreall diff --git a/tensorflow/compiler/xla/python/numpy_bridge.cc b/tensorflow/compiler/xla/python/numpy_bridge.cc index b30bdc3669de3992a08ab70ef49b0aa17cc855f3..3d87480728aab1d4ebbc71c6c7504d37cae5edaf 100644 --- a/tensorflow/compiler/xla/python/numpy_bridge.cc +++ b/tensorflow/compiler/xla/python/numpy_bridge.cc @@ -139,97 +139,221 @@ static int NumpyTypenum(PyObject* o) { return reinterpret_cast(o)->type_num; } -bool CheckPyShapeInfo(PyObject* o) { - // The object is a tuple (a pair) - if (!PyTuple_Check(o)) { - PyErr_SetString(PyExc_TypeError, "Shape record must be a tuple"); - return false; +// Extracts the string held inside r and returns it as a C++ string. +// +// NOTE: this is an internal helper for conversion to a C++, and so decrefs r. +static string ExtractStringAndDecref(PyObject* r) { + auto error = [r] { + return tensorflow::strings::Printf("", r); + }; + if (r == nullptr) { + return error(); } - if (PyTuple_Size(o) != 2) { - PyErr_SetString(PyExc_ValueError, "Shape record tuple must be of length 2"); - return false; +#if PY_MAJOR_VERSION < 3 + string result = PyString_AsString(r); +#else + PyObject* bytes = PyUnicode_AsEncodedString(r, 0, 0); + if (bytes == nullptr) { + return error(); } + CHECK(PyBytes_Check(bytes)); + string result = PyBytes_AsString(bytes); + Py_DECREF(bytes); +#endif + Py_DECREF(r); + return result; +} - // It has a first element, which is a numpy dtype object - PyObject* first = PyTuple_GetItem(o, 0); - if (!first) { - return false; - } - if (first->ob_type != &PyArrayDescr_Type) { - PyErr_SetString( - PyExc_TypeError, - "Shape record does not have a numpy dtype as its first element"); - return false; - } - const int np_type = NumpyTypenum(first); - if (!NumpyTypeIsValid(np_type)) { - PyErr_SetString(PyExc_ValueError, - "Shape record has an invalid integer dtype"); - return false; - } +// Safely returns a str of the given Python object o as a C++ string. +static string PyObjectCppStr(PyObject* o) { + PyObject* s = PyObject_Str(o); + return ExtractStringAndDecref(s); +} - // It has a second element, which is a tuple, either of shape - // records or of Python ints - PyObject* second = PyTuple_GetItem(o, 1); - if (!second) { - return false; - } - if (!PyTuple_Check(second)) { - PyErr_SetString(PyExc_TypeError, - "Shape record does not have a tuple as its second element"); - return false; - } - const int length = PyTuple_Size(second); - const PrimitiveType element_type = NumpyTypeToPrimitiveType(np_type); - for (int i = 0; i < length; i++) { - PyObject* dimension = PyTuple_GetItem(second, i); - if (element_type == TUPLE) { - if (!CheckPyShapeInfo(dimension)) { - return false; - } - } else if (!CheckPyIntOrLong(dimension)) { - PyErr_SetString(PyExc_TypeError, - "Non-tuple shape record has a non-integer dimension"); - return false; +// Safely returns a repr of the given Python object o as a C++ string. +static string PyObjectCppRepr(PyObject* o) { + PyObject* r = PyObject_Repr(o); + return ExtractStringAndDecref(r); +} + +StatusOr XlaShapeFromPyShape(PyObject* o) { + auto error = [o](const string& prefix) { + return InvalidArgument("%s; got %s", prefix.c_str(), + PyObjectCppRepr(o).c_str()); + }; + + auto get_attr = [o, &error](const string& field) -> StatusOr { + PyObject* result = + PyObject_GetAttrString(o, const_cast(field.c_str())); + if (result == nullptr) { + return error(tensorflow::strings::StrCat( + "Failed to get attribute of Shape object:", field)); } - } + return result; + }; - return true; -} + auto call_method = [o, &error](const string& method) -> StatusOr { + PyObject* result = + PyObject_CallMethod(o, const_cast(method.c_str()), nullptr); + if (result == nullptr) { + return error(tensorflow::strings::StrCat( + "Failed to call method of shape object:", method)); + } + return result; + }; + + PyObject* np_type; + TF_ASSIGN_OR_RETURN(np_type, get_attr("np_dtype")); + if (np_type->ob_type != &PyArrayDescr_Type) { + return error("Shape attribute np_dtype is not an integer numpy dtype"); + } + if (!NumpyTypeIsValid(NumpyTypenum(np_type))) { + return error("Shape attribute np_dtype is not a valid integer numpy dtype"); + } + const PrimitiveType element_type = + NumpyTypeToPrimitiveType(NumpyTypenum(np_type)); + Py_DECREF(np_type); -// Precondition: CheckPyShapeInfo(o) -Shape XlaShapeFromPyShapeInfo(PyObject* o) { - const int np_type = NumpyTypenum(PyTuple_GetItem(o, 0)); - const PrimitiveType element_type = NumpyTypeToPrimitiveType(np_type); - PyObject* py_dimensions = PyTuple_GetItem(o, 1); - const int length = PyTuple_Size(py_dimensions); if (element_type == TUPLE) { + PyObject* py_subshapes; + TF_ASSIGN_OR_RETURN(py_subshapes, call_method("tuple_shapes")); + if (!PyTuple_Check(py_subshapes)) { + return error( + "Return value of Shape method tuple_shapes() is not a tuple"); + } + const int length = PyTuple_Size(py_subshapes); std::vector subshapes; subshapes.reserve(length); for (int i = 0; i < length; i++) { - subshapes.push_back( - XlaShapeFromPyShapeInfo(PyTuple_GetItem(py_dimensions, i))); + TF_ASSIGN_OR_RETURN( + const Shape& subshape, + XlaShapeFromPyShape(PyTuple_GetItem(py_subshapes, i))); + subshapes.push_back(subshape); } + Py_DECREF(py_subshapes); return ShapeUtil::MakeTupleShape(subshapes); } else { + PyObject* py_dimensions; + PyObject* py_minor_to_major; + TF_ASSIGN_OR_RETURN(py_dimensions, call_method("dimensions")); + TF_ASSIGN_OR_RETURN(py_minor_to_major, call_method("minor_to_major")); + if (!PyTuple_Check(py_dimensions)) { + return error("Return value of Shape method dimensions() is not a tuple"); + } + if (py_minor_to_major != Py_None && !PyTuple_Check(py_minor_to_major)) { + return error( + "Return value of Shape method minor_to_major() is neither a tuple " + "nor None"); + } + const int length = PyTuple_Size(py_dimensions); + if (py_minor_to_major != Py_None && + length != PyTuple_Size(py_minor_to_major)) { + return error( + "Shape methods dimensions() and minor_to_major() return " + "different-length tuples"); + } std::vector dimensions(length); + std::vector minor_to_major(length); for (int i = 0; i < length; i++) { dimensions[i] = PyIntOrPyLongToLong(PyTuple_GetItem(py_dimensions, i)); - if (dimensions[i] == -1) { - CHECK(!PyErr_Occurred()); + if (dimensions[i] == -1 && PyErr_Occurred()) { + return error("Dimension is not an int"); + } + + if (py_minor_to_major != Py_None) { + minor_to_major[i] = + PyIntOrPyLongToLong(PyTuple_GetItem(py_minor_to_major, i)); + if (minor_to_major[i] == -1 && PyErr_Occurred()) { + return error("Minor-to-major value is not an int"); + } } } - return ShapeUtil::MakeShape(element_type, dimensions); + bool with_layout = py_minor_to_major != Py_None; + Py_DECREF(py_dimensions); + Py_DECREF(py_minor_to_major); + if (with_layout) { + return ShapeUtil::MakeShapeWithLayout(element_type, dimensions, + minor_to_major); + } else { + return ShapeUtil::MakeShape(element_type, dimensions); + } } } +// Helper that retrieves the member with attr_name, stringifies it if is not +// None, and returns it as a C++ string. +static tensorflow::gtl::optional GetAttrAsString( + PyObject* o, const string& attr_name) { + if (!PyObject_HasAttrString(o, attr_name.c_str())) { + return tensorflow::gtl::nullopt; + } + PyObject* attr = PyObject_GetAttrString(o, attr_name.c_str()); + if (attr == Py_None) { + Py_DECREF(attr); + return tensorflow::gtl::nullopt; + } + string result = PyObjectCppStr(attr); + Py_DECREF(attr); + return result; +} + +// Helper that retrieves the member with attr_name, checks that it is an integer +// if it is not None, and returns it as an int32 value. +static tensorflow::gtl::optional GetAttrAsInt32( + PyObject* o, const string& attr_name) { + if (!PyObject_HasAttrString(o, attr_name.c_str())) { + return tensorflow::gtl::nullopt; + } + PyObject* attr = PyObject_GetAttrString(o, attr_name.c_str()); + if (attr == Py_None) { + Py_DECREF(attr); + return tensorflow::gtl::nullopt; + } + if (!CheckPyIntOrLong(attr)) { + Py_DECREF(attr); + return tensorflow::gtl::nullopt; + } + long value = PyIntOrPyLongToLong(attr); // NOLINT + Py_DECREF(attr); + if (value == -1 && PyErr_Occurred() != nullptr) { + return tensorflow::gtl::nullopt; + } + if (static_cast(value) != value) { + return tensorflow::gtl::nullopt; + } + return value; +} + +StatusOr OpMetadataFromPyObject(PyObject* o) { + OpMetadata result; + tensorflow::gtl::optional op_type = GetAttrAsString(o, "op_type"); + if (op_type.has_value()) { + result.set_op_type(op_type.value()); + } + tensorflow::gtl::optional op_name = GetAttrAsString(o, "op_name"); + if (op_name.has_value()) { + result.set_op_name(op_name.value()); + } + tensorflow::gtl::optional source_file = + GetAttrAsString(o, "source_file"); + if (source_file.has_value()) { + result.set_source_file(source_file.value()); + } + tensorflow::gtl::optional source_line = + GetAttrAsInt32(o, "source_line"); + if (source_line.has_value()) { + result.set_source_line(source_line.value()); + } + return result; +} + PyObject* PyObjectFromXlaLiteral(const Literal& literal) { if (ShapeUtil::IsTuple(literal.shape())) { - const std::vector& tuple_literals = literal.tuple_literals(); int num_elements = ShapeUtil::TupleElementCount(literal.shape()); PyObject* tuple = PyTuple_New(num_elements); for (int i = 0; i < num_elements; i++) { - PyTuple_SET_ITEM(tuple, i, PyObjectFromXlaLiteral(tuple_literals[i])); + PyTuple_SET_ITEM( + tuple, i, PyObjectFromXlaLiteral(LiteralView::Create(literal, {i}))); } return tuple; } else { @@ -247,14 +371,15 @@ PyObject* PyObjectFromXlaLiteral(const Literal& literal) { } } -std::unique_ptr XlaLiteralFromPyObject(PyObject* o) { +StatusOr> XlaLiteralFromPyObject(PyObject* o) { if (PyTuple_Check(o)) { int num_elements = PyTuple_Size(o); std::vector> elements; elements.reserve(num_elements); for (int i = 0; i < num_elements; i++) { PyObject* element = PyTuple_GetItem(o, i); - elements.push_back(XlaLiteralFromPyObject(element)); + TF_ASSIGN_OR_RETURN(auto literal, XlaLiteralFromPyObject(element)); + elements.push_back(std::move(literal)); } return Literal::MakeTupleOwned(std::move(elements)); } else if (PyArray_Check(o)) { @@ -267,16 +392,17 @@ std::unique_ptr XlaLiteralFromPyObject(PyObject* o) { int np_type = PyArray_TYPE(py_array); auto literal = Literal::CreateFromDimensions( NumpyTypeToPrimitiveType(np_type), dimensions); - CopyNumpyArrayToLiteral(np_type, py_array, literal.get()); - return literal; + TF_RETURN_IF_ERROR( + CopyNumpyArrayToLiteral(np_type, py_array, literal.get())); + return std::move(literal); } else { - LOG(FATAL) - << "Non-tuple or Numpy array encountered in conversion to XLA literal"; + return InvalidArgument( + "Non-tuple or Numpy array encountered in conversion to XLA literal."); } } -void CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, - Literal* literal) { +Status CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, + Literal* literal) { switch (np_type) { case NPY_BOOL: CopyNumpyArrayToLiteral(py_array, literal); @@ -306,8 +432,10 @@ void CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, CopyNumpyArrayToLiteral(py_array, literal); break; default: - LOG(FATAL) << "No XLA literal container for Numpy type" << np_type; + return InvalidArgument( + "No XLA literal container for Numpy type number: %d", np_type); } + return Status::OK(); } void CopyLiteralToNumpyArray(int np_type, const Literal& literal, diff --git a/tensorflow/compiler/xla/python/numpy_bridge.h b/tensorflow/compiler/xla/python/numpy_bridge.h index 4e6ecbb0e8b58979ec1f1484e722725c391106fb..adfcc3b8588dce01718bb19dea936bace483be4d 100644 --- a/tensorflow/compiler/xla/python/numpy_bridge.h +++ b/tensorflow/compiler/xla/python/numpy_bridge.h @@ -56,15 +56,15 @@ bool NumpyTypeIsValid(int np_type); // The return value is a new reference. PyObject* PyShapeInfoFromXlaShape(const Shape& shape); -// Returns the outcome of a best-effort check that the Python object -// is a pair of the form (numpy dtype, dimensions), as produced by -// PyShapeInfoFromXlaShape. -bool CheckPyShapeInfo(PyObject* o); - -// Performs the inverse conversion to that of PyShapeInfoFromXlaShape. +// Converts a Python object with a method interface mathing that of +// xla_client.Shape into an XLA Shape object. // // The return value is a new reference. -Shape XlaShapeFromPyShapeInfo(PyObject* o); +StatusOr XlaShapeFromPyShape(PyObject* o); + +// Converts a PyObject that represents operation metadata into protocol buffer +// form. +StatusOr OpMetadataFromPyObject(PyObject* o); // Converts an XLA literal to a Python object, either a Numpy ndarray // or a nested Python tuple thereof. @@ -82,13 +82,13 @@ PyObject* PyObjectFromXlaLiteral(const Literal& literal); // To avoid transferring ownership of the data buffers that underlie // PyArrays and XLA literals, this function makes deep copies of all // array data. -std::unique_ptr XlaLiteralFromPyObject(PyObject* o); +StatusOr > XlaLiteralFromPyObject(PyObject* o); // The following functions copy array data from the buffers underlying Numpy // ndarrays into those underlying XLA literals, and vice versa. -void CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, - Literal* literal); +Status CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, + Literal* literal); void CopyLiteralToNumpyArray(int np_type, const Literal& literal, PyArrayObject* py_array); @@ -96,14 +96,14 @@ void CopyLiteralToNumpyArray(int np_type, const Literal& literal, template void CopyNumpyArrayToLiteral(PyArrayObject* py_array, Literal* literal) { NativeT* source = static_cast(PyArray_DATA(py_array)); - auto dest = literal->GetMutableArraySlice(); + auto dest = literal->data(); std::copy(source, source + PyArray_SIZE(py_array), dest.data()); } template void CopyLiteralToNumpyArray(const Literal& literal, PyArrayObject* py_array) { NativeT* dest = static_cast(PyArray_DATA(py_array)); - auto source = literal.GetArraySlice(); + auto source = literal.data(); std::copy(source.begin(), source.end(), dest); } diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py index e2ddb15559459e3e7d09f98284bbc6af7f14df13..9bda9d09294bc75acaa35d8e4a512820046e8920 100644 --- a/tensorflow/compiler/xla/python/xla_client.py +++ b/tensorflow/compiler/xla/python/xla_client.py @@ -18,8 +18,11 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections import enum # pylint: disable=g-bad-import-order +import inspect import itertools +import os import numpy as np @@ -27,16 +30,66 @@ from tensorflow.compiler.xla import xla_data_pb2 from tensorflow.compiler.xla.python import pywrap_xla as c_api +# Most functions are snake_case for consistency with other modules, +# whereas method names of ComputationBuilder and LocalComputation are +# CamelCase for consistency with XLA. +# pylint: disable=invalid-name + + +_OP_METADATA_FIELDS = [ + 'op_type', + 'op_name', + 'source_file', + 'source_line', +] +OpMetadata = collections.namedtuple('OpMetadata', _OP_METADATA_FIELDS) + + +def OpMetadataToProto(pyobj): + proto = xla_data_pb2.OpMetadata() + for field in _OP_METADATA_FIELDS: + attr = getattr(pyobj, field) + if attr is not None: + setattr(proto, field, attr) + return proto + + +def CurrentSourceInfoMetadata(op_type=None, op_name=None, skip_frames=1): + """Helper for use in source mapping that returns an OpMetadata object.""" + full_filename, lineno = inspect.stack()[skip_frames][1:3] + filename = os.path.basename(full_filename) + return OpMetadata( + op_type=op_type, + op_name=op_name, + source_file=filename, + source_line=lineno) + + class PaddingType(enum.Enum): VALID = 1 SAME = 2 +def _convert_padding_type_to_pad_values(padding_type, lhs_dims, rhs_dims, + window_strides): + """Maps PaddingType (VALID or SAME) to pad values (list of pairs of ints).""" + if padding_type == PaddingType.VALID: + return [(0, 0)] * len(window_strides) + + out_shape = np.ceil(np.true_divide(lhs_dims, window_strides)).astype(int) + pad_sizes = [max((out_size - 1) * stride + filter_size - in_size, 0) + for out_size, stride, filter_size, in_size + in zip(out_shape, window_strides, rhs_dims, lhs_dims)] + return [(pad_size // 2, pad_size - pad_size // 2) + for pad_size in pad_sizes] + + _UNARY_OPS = [ 'Not', 'Abs', 'Exp', 'Floor', + 'Round', 'Ceil', 'Log', 'Sign', @@ -70,26 +123,63 @@ _BINARY_OPS = [ 'Pow', ] -# Most functions are snake_case for consistency with other modules, -# whereas method names of ComputationBuilder and LocalComputation are -# CamelCase for consistency with XLA. -# pylint: disable=invalid-name - XLA_ELEMENT_TYPE_TO_DTYPE = { xla_data_pb2.F32: np.dtype(np.float32), xla_data_pb2.F64: np.dtype(np.float64), xla_data_pb2.S32: np.dtype(np.int32), xla_data_pb2.S64: np.dtype(np.int64), + xla_data_pb2.U32: np.dtype(np.uint32), + xla_data_pb2.U64: np.dtype(np.uint64), xla_data_pb2.PRED: np.dtype(np.bool), xla_data_pb2.TUPLE: np.dtype(np.object), } +# Note the conversion on the key. Numpy has a known issue wherein dtype hashing +# doesn't work as expected (https://github.com/numpy/numpy/issues/7242). Thus, +# when keying by dtype in this dict, we use the string form of dtypes. DTYPE_TO_XLA_ELEMENT_TYPE = { str(v): k for k, v in XLA_ELEMENT_TYPE_TO_DTYPE.items() } +class LocalBuffer(object): + """Represents a handle to data owned by XLA. + + The referent is ready for use in executing a local, compiled + Computation. On XLA platforms involving a device (e.g. GPU), this + means the referent is in device memory. + """ + + def __init__(self, c_local_shaped_buffer): + self.c_local_shaped_buffer = c_local_shaped_buffer + self._delete = c_api.DeleteLocalShapedBuffer + + @staticmethod + def from_py(npval, layout_fn=None): + npval = require_numpy_array_layout(npval) + if layout_fn: + shape = Shape.from_numpy(npval) + shape = shape.map_leaves(layout_fn) + else: + shape = None + return LocalBuffer(c_api.LocalShapedBuffer.FromLiteral(npval, shape)) + + def to_py(self): + return self.c_local_shaped_buffer.ToLiteral() + + def delete(self): + if self.c_local_shaped_buffer is not None: + self._delete(self.c_local_shaped_buffer) + self.c_local_shaped_buffer = None + + def is_deleted(self): + return self.c_local_shaped_buffer is None + + def __del__(self): + self.delete() + + class Shape(object): """XLA shape. @@ -98,9 +188,23 @@ class Shape(object): represents an XLA tuple. """ - def __init__(self, np_dtype, dimensions): + def __init__(self, np_dtype, dimensions, minor_to_major=None): + assert isinstance(dimensions, tuple) self.np_dtype = np_dtype self._dimensions = dimensions + self._minor_to_major = minor_to_major + self._check_minor_to_major() + + def __eq__(self, other): + # pylint: disable=protected-access + return (self.np_dtype == other.np_dtype and + self._dimensions == other._dimensions and + self._minor_to_major == other._minor_to_major) + + def __repr__(self): + return ('xla_client.Shape(np_dtype={!r}, dimensions={!r}, ' + 'minor_to_major={!r})').format(self.np_dtype, self._dimensions, + self._minor_to_major) def element_type(self): return DTYPE_TO_XLA_ELEMENT_TYPE[str(self.np_dtype)] @@ -113,11 +217,49 @@ class Shape(object): raise ValueError('Tuple shape has no dimensions') return self._dimensions + def minor_to_major(self): + return self._minor_to_major + def tuple_shapes(self): if not self.is_tuple(): raise ValueError('Shape is not a tuple shape') return self._dimensions + def rank(self): + return len(self.dimensions()) + + def map_leaves(self, f): + """Map f over each leaf-level array subshape. + + Args: + f: The function to apply. Whenever f returns None, the identity is + applied instead. + + Returns: + A new Shape with the mapped leaves. + """ + if self.is_tuple(): + children = tuple(child.map_leaves(f) for child in self.tuple_shapes()) + return Shape(np.dtype('O'), children) + else: + mapped = f(self) + return self if mapped is None else mapped + + def _check_minor_to_major(self): + mtm = self._minor_to_major + if self.is_tuple(): + assert mtm is None, self + if mtm is not None: + assert self.rank() == len(mtm), self + assert sorted(mtm) == range(len(mtm)), self + + def update_minor_to_major(self, minor_to_major): + if not isinstance(minor_to_major, tuple): + raise TypeError('minor_to_major must be a tuple') + updated = Shape(self.np_dtype, tuple(self.dimensions()), minor_to_major) + updated._check_minor_to_major() # pylint: disable=protected-access + return updated + @staticmethod def from_numpy(npval): @@ -134,23 +276,10 @@ def _wrap_shape(shape_info): dtype, dims = shape_info element_type = DTYPE_TO_XLA_ELEMENT_TYPE[str(dtype)] if element_type == xla_data_pb2.TUPLE: - dims = [_wrap_shape(subshape_info) for subshape_info in dims] + dims = tuple(_wrap_shape(subshape_info) for subshape_info in dims) return Shape(dtype, dims) -def _unwrap_shape(shape): - if shape.is_tuple(): - components = tuple( - _unwrap_shape(subshape) for subshape in shape.tuple_shapes()) - else: - components = shape.dimensions() - return (shape.np_dtype, components) - - -def _unwrap_shapes(shapes): - return [_unwrap_shape(shape) for shape in shapes] - - def _wrap_data_handle(handle): cdh = xla_data_pb2.ComputationDataHandle() cdh.handle = handle @@ -172,6 +301,51 @@ def require_numpy_array_layout(value): return np.require(value, requirements=['C', 'A']) +class CompileOptions(object): + """Python object for XLA compile options. + + These options can be passed to the 'compile' step when using a local XLA + client. + """ + + def __init__(self): + self.generate_hlo_graph = None + + +def transfer_to_infeed(value, replica_number=None): + """Transfers the given value into the XLA infeed queue. + + XLA's infeed queue is a single queue that feeds the "XLA virtual machine" with + a totally ordered stream of values. This is dequeued from XLA computations via + the Infeed() operation. + + Args: + value: the value that the caller would like to enqueue into the XLA infeed + queue + replica_number: the replica number to infeed the value to -- if not + provided, then the default replica (trivially replica 0) is used. + """ + if replica_number is None: + c_api.TransferToInfeedLocal(require_numpy_array_layout(value)) + else: + c_api.TransferToInfeedLocalReplica( + require_numpy_array_layout(value), replica_number) + + +def transfer_from_outfeed(shape, replica_number=None): + """Transfers a literal of the given shape from replica_number's outfeed. + + Args: + shape: The shape of the value to transfer from outfeed. + replica_number: The replica number ordinal to transfer the outfeed value + from. (Each replica has a distinct outfeed queue.) + + Returns: + The literal value that is produced from the outfeed queue. + """ + return c_api.TransferFromOutfeedLocalReplica(shape, replica_number or 0) + + class LocalComputation(object): """Python wrapper for a local XLA Computation. @@ -186,26 +360,81 @@ class LocalComputation(object): # Ensure a reference to C-based destructor for use in __del__. if is_compiled: + assert isinstance(c_local_computation, c_api.CompiledLocalComputation) self._delete = c_api.DeleteCompiledLocalComputation else: + assert isinstance(c_local_computation, c_api.LocalComputation) self._delete = c_api.DeleteLocalComputation - def Compile(self, argument_shapes=()): + def Compile(self, argument_shapes=(), compile_options=None, layout_fn=None): + """Compiles an un-compiled local computation. + + Local computations are the result of a "LocalComputationBuild'ing" process + -- they start in uncompiled form, and via a call to Compile() turn into a + compiled local computation. + + Raises: + ValueError: if this is already a compiled local computation. + + Arguments: + argument_shapes: parameter shapes -- they are first laid out by layout_fn + if layout_fn is provided. Otherwise, the default layout for those shapes + will be used. + compile_options: options to use for compilation, includes an optional + laid out result shape for the computation. + layout_fn: lambda that is used to lay out the argument/result shapes. + + Returns: + A newly *compiled* local computation instance. + """ if self.is_compiled: raise ValueError('Attempt to compile a compiled local XLA computation.') + + if layout_fn: + argument_shapes = [ + shape.map_leaves(layout_fn) for shape in argument_shapes + ] + result_shape = _wrap_shape(self.c_local_computation.GetReturnValueShape()) + result_shape = result_shape.map_leaves(layout_fn) + compile_options = compile_options or CompileOptions() + compile_options.result_shape = result_shape return LocalComputation( - self.c_local_computation.Compile(_unwrap_shapes(argument_shapes)), + self.c_local_computation.Compile(argument_shapes, compile_options), is_compiled=True) - def CompileWithExampleArguments(self, arguments=()): + def CompileWithExampleArguments(self, + arguments=(), + compile_options=None, + layout_fn=None): return self.Compile( - argument_shapes=[Shape.from_numpy(arg) for arg in arguments]) + argument_shapes=[Shape.from_numpy(arg) for arg in arguments], + compile_options=compile_options, + layout_fn=layout_fn) - def Execute(self, arguments=()): + def Execute(self, arguments=(), layout_fn=None): + """Execute with Python values as arguments and return value.""" if not self.is_compiled: raise ValueError('Cannot execute an uncompiled local XLA computation.') + argument_shapes = [Shape.from_numpy(arg) for arg in arguments] + if layout_fn: + argument_shapes = [ + shape.map_leaves(layout_fn) for shape in argument_shapes + ] + else: + argument_shapes = [None for shape in argument_shapes] arguments = tuple(map(require_numpy_array_layout, arguments)) - return self.c_local_computation.Execute(arguments) + return self.c_local_computation.Execute(arguments, argument_shapes) + + def ExecuteWithLocalBuffers(self, arguments=()): + """Execute with LocalBuffer arguments and return value.""" + if not self.is_compiled: + raise ValueError('Cannot execute an uncompiled local XLA computation.') + arguments = tuple(arguments) + if any(arg.is_deleted() for arg in arguments): + raise ValueError('Executing with deleted local buffer argument') + return LocalBuffer( + self.c_local_computation.ExecuteWithShapedBuffers( + [arg.c_local_shaped_buffer for arg in arguments])) def __del__(self): self._delete(self.c_local_computation) @@ -233,6 +462,35 @@ class ComputationBuilder(object): def Build(self): return LocalComputation(self._client.Build(), is_compiled=False) + def SetOpMetadata(self, op_metadata): + """Set metadata for operations that are about to be enqueued.""" + self._client.SetOpMetadata(op_metadata) + + def ClearOpMetadata(self): + """Clear metadata for operations that are about to be enqueued.""" + self._client.ClearOpMetadata() + + def Infeed(self, shape): + """Enqueues an infeed op onto the computation. + + Infeed operations dequeue data of the given shape from the device's infeed + queue for subsequent use in the computation. + + Returns: + A ComputationDataHandle message. + """ + return _wrap_data_handle(self._client.Infeed(shape)) + + def Outfeed(self, operand): + """Enqueues an outfeed op onto the computation. + + Outfeed operations enqueue data, using the given operand, onto the XLA + outfeed queue for subsequent dequeue via the client API. + """ + self._client.Outfeed( + _unwrap_data_handle(operand), self.GetShape(operand), + ''.encode('utf-8')) + def Constant(self, value): """Enqueues a constant op onto the computation. @@ -321,8 +579,7 @@ class ComputationBuilder(object): parameter_num = next(self._parameter_numbering) return _wrap_data_handle( - self._client.Parameter( - parameter_num, _unwrap_shape(shape), name.encode('utf8'))) + self._client.Parameter(parameter_num, shape, name.encode('utf8'))) def ParameterFromNumpy(self, value, name=None, parameter_num=None): """Enqueues a Parameter op onto the computation. @@ -382,9 +639,32 @@ class ComputationBuilder(object): def GetShape(self, operand): return _wrap_shape(self._client.GetShape(_unwrap_data_handle(operand))) + def GetReturnValueShape(self): + return _wrap_shape(self._client.GetReturnValueShape()) + def GetComputationStats(self): raise NotImplementedError() + def Pad(self, operand, padding_value, padding_config): + """Enqueues a Pad operation onto the computation. + + Args: + operand: ComputationDataHandle representing the array to pad. + padding_value: ComputationDataHandle representing the scalar pad value. + padding_config: either an xla_data_pb2.PaddingConfig or a list of integer + triples (edge_padding_low, edge_padding_high, interior_padding) + representing the configuration of the padding operation. + + Returns: + A ComputationDataHandle representing the added pad op. + """ + if not isinstance(padding_config, xla_data_pb2.PaddingConfig): + padding_config = GetPaddingConfigFromTriples(padding_config) + return _wrap_data_handle( + self._client.Pad(_unwrap_data_handle(operand), + _unwrap_data_handle(padding_value), + padding_config)) + def Reshape(self, operand, dimensions, new_sizes): """Reshape op.""" return _wrap_data_handle( @@ -423,6 +703,43 @@ class ComputationBuilder(object): return _wrap_data_handle( self._client.Rev(_unwrap_data_handle(operand), dimensions)) + def Clamp(self, min, operand, max): # pylint: disable=redefined-builtin + """Clamp op.""" + return _wrap_data_handle( + self._client.Clamp(_unwrap_data_handle(min), + _unwrap_data_handle(operand), + _unwrap_data_handle(max))) + + def SelectAndScatter(self, operand, select, window_dimensions, window_strides, + padding, source, init_value, scatter): + """Select and scatter op, used by the gradient of ReduceWindow. + + Args: + operand: ComputationDataHandle for array of dimension N and type T over + which the windows slide. + select: Computation of type (T, T) -> Pred to apply to the elements of + each window to indicate which element is selected. + window_dimensions: sequence of N integers for dimensions of the window. + window_strides: sequence of N integers for the strides of the window. + padding: PaddingType representing either 'SAME' or 'VALID ' padding. + source: ComputationDataHandle for array of type T with values to scatter. + init_value: ComputationDataHandle of scalar type T for initial out value. + scatter: Computation of type (T, T) -> T to apply to each scatter source + element with its destination element. + + Returns: + A ComputationDataHandle representing the added SelectAndScatter op. + """ + pads = _convert_padding_type_to_pad_values( + padding, self.GetShape(operand).dimensions(), + window_dimensions, window_strides) + return _wrap_data_handle( + self._client.SelectAndScatterWithGeneralPadding( + _unwrap_data_handle(operand), select.c_local_computation, + window_dimensions, window_strides, pads, + _unwrap_data_handle(source), _unwrap_data_handle(init_value), + scatter.c_local_computation)) + def Select(self, pred, on_true, on_false): """Element-wise selection op. @@ -573,13 +890,74 @@ class ComputationBuilder(object): computation_to_apply.c_local_computation, dimensions)) + def ReduceWindow(self, operand, init_value, computation_to_apply, + window_dimensions, window_strides, padding): + """Enqueues a windowed reduction operation onto the computation. + + Args: + operand: reduction operand (ComputationDataHandle). + init_value: reduction initial value (ComputationDataHandle). + computation_to_apply: a binary reduction function (Computation). + window_dimensions: dimensions of window (sequence of integers). + window_strides: strides for window (sequence of integers). + padding: PaddingType representing either 'SAME' or 'VALID' padding. + + Returns: + A ComputationDataHandle representing the added ReduceWindow op. + """ + pads = _convert_padding_type_to_pad_values( + padding, self.GetShape(operand).dimensions(), window_dimensions, + window_strides) + return _wrap_data_handle( + self._client.ReduceWindowWithGeneralPadding( + _unwrap_data_handle(operand), + _unwrap_data_handle(init_value), + computation_to_apply.c_local_computation, + window_dimensions, window_strides, pads)) + + def RngNormal(self, mu, sigma, dims): + """Enqueues an RngNormal operation onto the computation. + + Args: + mu: A ComputationDataHandle to an F32 scalar specifying the mean. + sigma: A ComputationDataHandle to an F32 scalar specifying the standard + deviation. + dims: A 1D array-like of nonnegative integers specifying the dimensions. + + Returns: a ComputationDataHandle to the generated array of F32 values. + """ + shape = Shape(self.GetShape(mu).np_dtype, dims) + return _wrap_data_handle( + self._client.RngNormal( + _unwrap_data_handle(mu), _unwrap_data_handle(sigma), shape)) + + def RngUniform(self, a, b, dims): + """Enqueues an RngUniform operation onto the computation. + + Args: + a: a ComputationDataHandle to an F32, S32, or U32 scalar (consistent with + the type of b) specifying the low end of the interval [a, b) over which + values are generated. + b: a ComputationDataHandle to an F32, S32, or U32 scalar (consistent with + the type of a) specifying the high end of the interval [a, b) over which + values are generated. + dims: A 1D array-like of nonnegative integers specifying the dimensions. + + Returns: a ComputationDataHandle to the generated array of values with the + same numeric type (F32, S32, or U32) as the arguments a and b. + """ + shape = Shape(self.GetShape(a).np_dtype, dims) + return _wrap_data_handle( + self._client.RngUniform( + _unwrap_data_handle(a), _unwrap_data_handle(b), shape)) + def While(self, cond, body, init): """Enqueues a While operation onto the computation. Args: cond: a Computation for the loop condition, which has type T -> PRED body: a Computation for the loop body, which has type T -> T - init: an ComputationDataHandle for the initial parameter, which has type T + init: a ComputationDataHandle for the initial parameter, which has type T Returns: a ComputationDataHandle representing the While operation. """ @@ -588,11 +966,58 @@ class ComputationBuilder(object): body.c_local_computation, _unwrap_data_handle(init))) + def Conditional(self, pred, true_operand, true_computation, false_operand, + false_computation): + """Enqueues a Conditional operation onto the computation. + + Args: + predicate: a ComputationDataHandle to test, which has scalar type PRED + true_operand: a ComputationDataHandle of type T_0 + true_computation: a Computation to apply to true_operand, type T_0 -> S + false_operand: a ComputationDatahandle of type T_1 + false_computation: a Computation to apply to false_operand, type T_1 -> S + + Returns: a ComputationDataHandle representing the Conditional operation. + """ + return _wrap_data_handle( + self._client.Conditional( + _unwrap_data_handle(pred), _unwrap_data_handle(true_operand), + true_computation.c_local_computation, + _unwrap_data_handle(false_operand), + false_computation.c_local_computation)) + def Dot(self, lhs, rhs): - """Matrix multiplication between lhs and rhs.""" + """Enqueues a dot operation onto the computation. + + Args: + lhs: ComputationDataHandle for the rank 1 or rank 2 left-hand-side array. + rhs: ComputationDataHandle for the rank 1 or rank 2 right-hand-side array. + + Returns: a ComputationDataHandle representing the Dot operation. + """ return _wrap_data_handle( self._client.Dot(_unwrap_data_handle(lhs), _unwrap_data_handle(rhs))) + def DotGeneral(self, lhs, rhs, dimension_numbers): + """Enqueues a general dot operation onto the computation. + + Args: + lhs: ComputationDataHandle for the left-hand-side array. + rhs: ComputationDataHandle for the right-hand-side array. + dimension_numbers: either an xla_data_pb2.DotDimensionNumbers or a nested + tuple ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)) of lists of + integers representing the dimensions to treat as contracting dimensions + and batch dimensions on each input operand. + + Returns: a ComputationDataHandle representing the DotGeneral operation. + """ + if not isinstance(dimension_numbers, xla_data_pb2.DotDimensionNumbers): + dimension_numbers = GetDotDimensionsFromLists(dimension_numbers) + return _wrap_data_handle( + self._client.DotGeneral( + _unwrap_data_handle(lhs), _unwrap_data_handle(rhs), + dimension_numbers)) + def Conv(self, lhs, rhs, window_strides, padding): """Enqueues a Conv operation onto the computation. @@ -604,18 +1029,9 @@ class ComputationBuilder(object): Returns: a ComputationDataHandle representing the Conv operation. """ - if padding == PaddingType.SAME: - lhs_dims = self.GetShape(lhs).dimensions() - rhs_dims = self.GetShape(rhs).dimensions() - in_shape, filter_shape = lhs_dims[2:], rhs_dims[2:] - out_shape = np.ceil(np.true_divide(in_shape, window_strides)).astype(int) - pad_sizes = [max((out_size - 1) * stride + filter_size - in_size, 0) - for out_size, stride, filter_size, in_size - in zip(out_shape, window_strides, filter_shape, in_shape)] - pads = [(pad_size // 2, pad_size - pad_size // 2) - for pad_size in pad_sizes] - else: - pads = [(0, 0)] * len(window_strides) + pads = _convert_padding_type_to_pad_values( + padding, self.GetShape(lhs).dimensions()[2:], + self.GetShape(rhs).dimensions()[2:], window_strides) dimension_numbers = self._GetConvDimensionNumbers(len(window_strides)) return _wrap_data_handle( self._client.ConvGeneralDilated(_unwrap_data_handle(lhs), @@ -705,3 +1121,46 @@ def _forward_methods_to_local_builder(): _forward_methods_to_local_builder() + + +def initialize_replica_count(replica_count): + """Initializes the desired replica count to use on XLA service init. + + Args: + replica_count: number of replicas that are desired for set up during XLA + initialization. + + Raises: + A runtime exception if the XLA service has already been initialized. + """ + c_api.InitializeReplicaCount(replica_count) + + +def get_replica_count(): + """Returns the current replica count used for the XLA service. + + Note: this will return a value whether the XLA service has been initialized + yet or not. + """ + return c_api.GetReplicaCount() + + +def GetPaddingConfigFromTriples(triples): + """Create PaddingConfig proto from list of triples of integers.""" + padding_config = xla_data_pb2.PaddingConfig() + for lo, hi, interior in triples: + dimension = padding_config.dimensions.add() + dimension.edge_padding_low = lo + dimension.edge_padding_high = hi + dimension.interior_padding = interior + return padding_config + + +def GetDotDimensionsFromLists(dimension_numbers): + (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers + dot_dims_proto = xla_data_pb2.DotDimensionNumbers() + dot_dims_proto.lhs_contracting_dimensions.extend(lhs_contract) + dot_dims_proto.rhs_contracting_dimensions.extend(rhs_contract) + dot_dims_proto.lhs_batch_dimensions.extend(lhs_batch) + dot_dims_proto.rhs_batch_dimensions.extend(rhs_batch) + return dot_dims_proto diff --git a/tensorflow/compiler/xla/python/xla_client_test.py b/tensorflow/compiler/xla/python/xla_client_test.py index 4a2fe68c2d9f1c96e1cd1ce1b5d67e4dadee32c8..c9d09cd5d57e001fd48d2dba9f2b0ee18374231b 100644 --- a/tensorflow/compiler/xla/python/xla_client_test.py +++ b/tensorflow/compiler/xla/python/xla_client_test.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function import itertools +import threading import numpy as np @@ -34,10 +35,13 @@ class LocalComputationTest(unittest.TestCase): name = self.id() return xla_client.ComputationBuilder(name) + def _Execute(self, c, arguments): + compiled_c = c.Build().CompileWithExampleArguments(arguments) + return compiled_c.Execute(arguments) + def _ExecuteAndAssertWith(self, assert_func, c, arguments, expected): assert expected is not None - compiled_c = c.Build().CompileWithExampleArguments(arguments) - result = compiled_c.Execute(arguments) + result = self._Execute(c, arguments) # Numpy's comparison methods are a bit too lenient by treating inputs as # "array-like", meaning that scalar 4 will be happily compared equal to # [[4]]. We'd like to be more strict so assert shapes as well. @@ -82,7 +86,8 @@ class ComputationsWithConstantsTest(LocalComputationTest): def testConstantScalarSumF32(self): c = self._NewComputation() - c.Add(c.ConstantF32Scalar(1.11), c.ConstantF32Scalar(3.14)) + root = c.Add(c.ConstantF32Scalar(1.11), c.ConstantF32Scalar(3.14)) + self.assertEqual(c.GetShape(root), c.GetReturnValueShape()) self._ExecuteAndCompareClose(c, expected=4.25) def testConstantScalarSumF64(self): @@ -309,6 +314,48 @@ class ParametersTest(LocalComputationTest): expected=[-4.3, 1.3, -6.3, 3.3]) +class LocalBufferTest(LocalComputationTest): + """Tests focusing on execution with LocalBuffers.""" + + def _Execute(self, c, arguments): + compiled_c = c.Build().CompileWithExampleArguments(arguments) + arg_buffers = [xla_client.LocalBuffer.from_py(arg) for arg in arguments] + result_buffer = compiled_c.ExecuteWithLocalBuffers(arg_buffers) + return result_buffer.to_py() + + def testConstantSum(self): + c = self._NewComputation() + c.Add(c.ConstantF32Scalar(1.11), c.ConstantF32Scalar(3.14)) + self._ExecuteAndCompareClose(c, expected=4.25) + + def testOneParameterSum(self): + c = self._NewComputation() + c.Add(c.ParameterFromNumpy(NumpyArrayF32(0.)), c.ConstantF32Scalar(3.14)) + self._ExecuteAndCompareClose( + c, + arguments=[NumpyArrayF32(1.11)], + expected=4.25) + + def testTwoParameterSum(self): + c = self._NewComputation() + c.Add(c.ParameterFromNumpy(NumpyArrayF32(0.)), + c.ParameterFromNumpy(NumpyArrayF32(0.))) + self._ExecuteAndCompareClose( + c, + arguments=[NumpyArrayF32(1.11), NumpyArrayF32(3.14)], + expected=4.25) + + def testCannotCallWithDeletedBuffers(self): + c = self._NewComputation() + c.Add(c.ParameterFromNumpy(NumpyArrayF32(0.)), c.ConstantF32Scalar(3.14)) + arg = NumpyArrayF32(1.11) + compiled_c = c.Build().CompileWithExampleArguments([arg]) + arg_buffer = xla_client.LocalBuffer.from_py(arg) + arg_buffer.delete() + with self.assertRaises(ValueError): + compiled_c.ExecuteWithLocalBuffers([arg_buffer]) + + class SingleOpTest(LocalComputationTest): """Tests for single ops. @@ -398,6 +445,30 @@ class SingleOpTest(LocalComputationTest): c.Dot(c.Constant(lhs), c.Constant(rhs)) self._ExecuteAndCompareClose(c, expected=np.dot(lhs, rhs)) + def testDotGeneral(self): + c = self._NewComputation() + rng = np.random.RandomState(0) + lhs = NumpyArrayF32(rng.randn(10, 3, 4)) + rhs = NumpyArrayF32(rng.randn(10, 4, 5)) + dimension_numbers = (([2], [1]), ([0], [0])) + c.DotGeneral(c.Constant(lhs), c.Constant(rhs), dimension_numbers) + self._ExecuteAndCompareClose(c, expected=np.matmul(lhs, rhs)) + + def testDotGeneralWithDotDimensionNumbersProto(self): + c = self._NewComputation() + rng = np.random.RandomState(0) + lhs = NumpyArrayF32(rng.randn(10, 3, 4)) + rhs = NumpyArrayF32(rng.randn(10, 4, 5)) + + dimension_numbers = xla_client.xla_data_pb2.DotDimensionNumbers() + dimension_numbers.lhs_contracting_dimensions.append(2) + dimension_numbers.rhs_contracting_dimensions.append(1) + dimension_numbers.lhs_batch_dimensions.append(0) + dimension_numbers.rhs_batch_dimensions.append(0) + + c.DotGeneral(c.Constant(lhs), c.Constant(rhs), dimension_numbers) + self._ExecuteAndCompareClose(c, expected=np.matmul(lhs, rhs)) + def testConvF32Same(self): c = self._NewComputation() a = lambda *dims: np.arange(np.prod(dims)).reshape(dims).astype("float32") @@ -450,6 +521,12 @@ class SingleOpTest(LocalComputationTest): c.Exp(c.Constant(arr)) self._ExecuteAndCompareClose(c, expected=np.exp(arr)) + def testRound(self): + c = self._NewComputation() + arr = NumpyArrayF32([3.3, 12.1]) + c.Round(c.Constant(arr)) + self._ExecuteAndCompareClose(c, expected=np.round(arr)) + def testLog(self): c = self._NewComputation() arr = NumpyArrayF32([3.3, 12.1]) @@ -598,6 +675,38 @@ class SingleOpTest(LocalComputationTest): c.Constant(NumpyArrayF32([1.0, 0.0, 2.0, 7.0, 12.0]))) self._ExecuteAndCompareExact(c, expected=[1.0, 0.0, 2.0, 4.0, 9.0]) + def testPad(self): + c = self._NewComputation() + c.Pad( + c.Constant(NumpyArrayF32([[1.0, 2.0], [3.0, 4.0]])), + c.Constant(NumpyArrayF32(0.0)), + [(1, 2, 1), (0, 1, 0)]) + self._ExecuteAndCompareClose(c, expected=[[0.0, 0.0, 0.0], + [1.0, 2.0, 0.0], + [0.0, 0.0, 0.0], + [3.0, 4.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0]]) + + def testPadWithPaddingConfig(self): + c = self._NewComputation() + padding_config = xla_client.xla_data_pb2.PaddingConfig() + for lo, hi, interior in [(1, 2, 1), (0, 1, 0)]: + dimension = padding_config.dimensions.add() + dimension.edge_padding_low = lo + dimension.edge_padding_high = hi + dimension.interior_padding = interior + c.Pad( + c.Constant(NumpyArrayF32([[1.0, 2.0], [3.0, 4.0]])), + c.Constant(NumpyArrayF32(0.0)), + padding_config) + self._ExecuteAndCompareClose(c, expected=[[0.0, 0.0, 0.0], + [1.0, 2.0, 0.0], + [0.0, 0.0, 0.0], + [3.0, 4.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0]]) + def testReshape(self): c = self._NewComputation() c.Reshape( @@ -621,6 +730,23 @@ class SingleOpTest(LocalComputationTest): self._ExecuteAndCompareExact( c, expected=[[[6, 5], [8, 7]], [[2, 1], [4, 3]]]) + def testClampF32(self): + c = self._NewComputation() + c.Clamp( + c.Constant(NumpyArrayF32(-1)), + c.Constant(NumpyArrayF32([-2, -1, 0, 1, 2, 3])), + c.Constant(NumpyArrayF32(2))) + self._ExecuteAndCompareExact(c, expected=[-1, -1, 0, 1, 2, 2]) + + # TODO(b/72689392): re-enable when bug S32 resolved + def DISABLED_testClampS32(self): + c = self._NewComputation() + c.Clamp( + c.Constant(NumpyArrayS32(-1)), + c.Constant(NumpyArrayS32([-2, -1, 0, 1, 2, 3])), + c.Constant(NumpyArrayS32(2))) + self._ExecuteAndCompareExact(c, expected=[-1, 0, 1, 2, 2]) + def testSelect(self): c = self._NewComputation() c.Select( @@ -676,6 +802,42 @@ class SingleOpTest(LocalComputationTest): self._ExecuteAndCompareExact( c, expected=[[10, 20, 30, 40], [10, 20, 30, 40], [10, 20, 30, 40]]) + def testRngNormal(self): + shape = (2, 3) + c = self._NewComputation() + c.RngNormal(c.Constant(NumpyArrayF32(0.)), c.Constant(NumpyArrayF32(1.)), + dims=shape) + result = c.Build().Compile().Execute() + # since the result is random, we just check shape and uniqueness + self.assertEqual(result.shape, shape) + self.assertEqual(len(np.unique(result)), np.prod(shape)) + + def testRngUniformF32(self): + lo, hi = 2., 4. + shape = (2, 3) + c = self._NewComputation() + c.RngUniform(c.Constant(NumpyArrayF32(lo)), c.Constant(NumpyArrayF32(hi)), + dims=shape) + result = c.Build().Compile().Execute() + # since the result is random, we just check shape, uniqueness, and range + self.assertEqual(result.shape, shape) + self.assertEqual(len(np.unique(result)), np.prod(shape)) + self.assertTrue(np.all(lo <= result)) + self.assertTrue(np.all(result < hi)) + + def testRngUniformS32(self): + lo, hi = 2, 4 + shape = (2, 3) + c = self._NewComputation() + c.RngUniform(c.Constant(NumpyArrayS32(lo)), c.Constant(NumpyArrayS32(hi)), + dims=shape) + result = c.Build().Compile().Execute() + # since the result is random, we just check shape, integrality, and range + self.assertEqual(result.shape, shape) + self.assertEqual(result.dtype, np.int32) + self.assertTrue(np.all(lo <= result)) + self.assertTrue(np.all(result < hi)) + class EmbeddedComputationsTest(LocalComputationTest): """Tests for XLA graphs with embedded computations (such as maps).""" @@ -720,6 +882,13 @@ class EmbeddedComputationsTest(LocalComputationTest): c.Mul(c.ParameterFromNumpy(NumpyArrayF32(0)), c.ConstantF32Scalar(2.0)) return c.Build() + def _CreateMulF32ByParamComputation(self): + """Computation (f32) -> f32 that multiplies one parameter by the other.""" + c = self._NewComputation("mul_f32_by_param") + c.Mul(c.ParameterFromNumpy(NumpyArrayF32(0)), + c.ParameterFromNumpy(NumpyArrayF32(0))) + return c.Build() + def _CreateMulF64By2Computation(self): """Computation (f64) -> f64 that multiplies its parameter by 2.""" c = self._NewComputation("mul_f64_by2") @@ -770,6 +939,20 @@ class EmbeddedComputationsTest(LocalComputationTest): c.Lt(c.ParameterFromNumpy(NumpyArrayF64(0)), c.ConstantF64Scalar(10.)) return c.Build() + def _CreateBinaryGeF32Computation(self): + """Computation (f32, f32) -> bool that tests first_param >= second_param.""" + c = self._NewComputation("param0_lt_param1") + c.Ge(c.ParameterFromNumpy(NumpyArrayF32(0)), + c.ParameterFromNumpy(NumpyArrayF32(0))) + return c.Build() + + def _CreateBinaryGeF64Computation(self): + """Computation (f64, f64) -> bool that tests first_param >= second_param.""" + c = self._NewComputation("param0_lt_param1") + c.Ge(c.ParameterFromNumpy(NumpyArrayF64(0)), + c.ParameterFromNumpy(NumpyArrayF64(0))) + return c.Build() + def _MakeSample3DArrayF32(self): return NumpyArrayF32([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]]) @@ -846,6 +1029,38 @@ class EmbeddedComputationsTest(LocalComputationTest): self._CreateBinaryDivF64Computation(), [0]) self._ExecuteAndCompareClose(c, expected=[0.2, 0.4, 0.75, 1.0]) + def DISABLED_testMapWithStaticOperands(self): + c = self._NewComputation() + factor = c.ConstantF32Scalar(3.0) + c.Map([c.Constant(NumpyArrayF32([1.0, 2.0, 3.0, 4.0]))], + self._CreateMulF32ByParamComputation(), [0], + static_operands=[factor]) + self._ExecuteAndCompareClose(c, expected=[3.0, 6.0, 9.0, 12.0]) + + def testSelectAndScatterF32(self): + c = self._NewComputation() + c.SelectAndScatter(c.Constant(NumpyArrayF32([[1., 2., 6.], [4., 5., 3.]])), + select=self._CreateBinaryGeF32Computation(), + window_dimensions=(2, 1), + window_strides=(1, 2), + padding=xla_client.PaddingType.VALID, + source=c.Constant(NumpyArrayF32([[0.1, 0.2]])), + init_value=c.Constant(NumpyArrayF32(1)), + scatter=self._CreateBinaryAddF32Computation()) + self._ExecuteAndCompareClose(c, expected=[[1., 1., 1.2], [1.1, 1., 1.]]) + + def testSelectAndScatterF64(self): + c = self._NewComputation() + c.SelectAndScatter(c.Constant(NumpyArrayF64([[1., 2., 6.], [4., 5., 3.]])), + select=self._CreateBinaryGeF64Computation(), + window_dimensions=(2, 1), + window_strides=(1, 2), + padding=xla_client.PaddingType.VALID, + source=c.Constant(NumpyArrayF64([[0.1, 0.2]])), + init_value=c.Constant(NumpyArrayF64(1)), + scatter=self._CreateBinaryAddF64Computation()) + self._ExecuteAndCompareClose(c, expected=[[1., 1., 1.2], [1.1, 1., 1.]]) + def testReduce1DtoScalarF32(self): c = self._NewComputation() c.Reduce( @@ -944,6 +1159,66 @@ class EmbeddedComputationsTest(LocalComputationTest): _ReduceAndTest(1, 2) _ReduceAndTest(0, 1, 2) + def testReduceWindowValidUnitStridesF32(self): + input_array = NumpyArrayF32([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + c = self._NewComputation() + c.ReduceWindow(operand=c.Constant(input_array), + init_value=c.ConstantF32Scalar(0), + computation_to_apply=self._CreateBinaryAddF32Computation(), + window_dimensions=(2, 1), window_strides=(1, 1), + padding=xla_client.PaddingType.VALID) + self._ExecuteAndCompareClose(c, expected=[[5., 7., 9.]]) + + def testReduceWindowSameUnitStridesF32(self): + input_array = NumpyArrayF32([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + c = self._NewComputation() + c.ReduceWindow(operand=c.Constant(input_array), + init_value=c.ConstantF32Scalar(0), + computation_to_apply=self._CreateBinaryAddF32Computation(), + window_dimensions=(2, 1), window_strides=(1, 1), + padding=xla_client.PaddingType.SAME) + self._ExecuteAndCompareClose(c, expected=[[5., 7., 9.], [4., 5., 6.]]) + + def testReduceWindowValidGeneralStridesF32(self): + input_array = NumpyArrayF32([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + c = self._NewComputation() + c.ReduceWindow(operand=c.Constant(input_array), + init_value=c.ConstantF32Scalar(0), + computation_to_apply=self._CreateBinaryAddF32Computation(), + window_dimensions=(2, 1), window_strides=(1, 2), + padding=xla_client.PaddingType.VALID) + self._ExecuteAndCompareClose(c, expected=[[5., 9.]]) + + def testReduceWindowValidUnitStridesF64(self): + input_array = NumpyArrayF64([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + c = self._NewComputation() + c.ReduceWindow(operand=c.Constant(input_array), + init_value=c.ConstantF64Scalar(0), + computation_to_apply=self._CreateBinaryAddF64Computation(), + window_dimensions=(2, 1), window_strides=(1, 1), + padding=xla_client.PaddingType.VALID) + self._ExecuteAndCompareClose(c, expected=[[5., 7., 9.]]) + + def testReduceWindowSameUnitStridesF64(self): + input_array = NumpyArrayF64([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + c = self._NewComputation() + c.ReduceWindow(operand=c.Constant(input_array), + init_value=c.ConstantF64Scalar(0), + computation_to_apply=self._CreateBinaryAddF64Computation(), + window_dimensions=(2, 1), window_strides=(1, 1), + padding=xla_client.PaddingType.SAME) + self._ExecuteAndCompareClose(c, expected=[[5., 7., 9.], [4., 5., 6.]]) + + def testReduceWindowValidGeneralStridesF64(self): + input_array = NumpyArrayF64([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + c = self._NewComputation() + c.ReduceWindow(operand=c.Constant(input_array), + init_value=c.ConstantF64Scalar(0), + computation_to_apply=self._CreateBinaryAddF64Computation(), + window_dimensions=(2, 1), window_strides=(1, 2), + padding=xla_client.PaddingType.VALID) + self._ExecuteAndCompareClose(c, expected=[[5., 9.]]) + def testWhileF32(self): cond = self._CreateTestF32Lt10Computation() body = self._CreateMulF32By2Computation() @@ -960,6 +1235,74 @@ class EmbeddedComputationsTest(LocalComputationTest): c.While(cond, body, init) self._ExecuteAndCompareClose(c, expected=16.) + def testConditionalTrue(self): + c = self._NewComputation() + pred = c.ConstantPredScalar(True) + true_operand = c.ConstantF32Scalar(3.) + true_computation = self._CreateMulF32By2Computation() + false_operand = c.ConstantF32Scalar(2.) + false_computation = self._CreateConstantF32Computation() + c.Conditional(pred, true_operand, true_computation, false_operand, + false_computation) + self._ExecuteAndCompareClose(c, expected=6.) + + def testConditionalFalse(self): + c = self._NewComputation() + pred = c.ConstantPredScalar(False) + true_operand = c.ConstantF32Scalar(3.) + true_computation = self._CreateMulF32By2Computation() + false_operand = c.ConstantF32Scalar(2.) + false_computation = self._CreateConstantF32Computation() + c.Conditional(pred, true_operand, true_computation, false_operand, + false_computation) + self._ExecuteAndCompareClose(c, expected=1.) + + def testInfeedS32Values(self): + to_infeed = NumpyArrayS32([1, 2, 3, 4]) + c = self._NewComputation() + c.Infeed(xla_client.Shape.from_numpy(to_infeed[0])) + compiled_c = c.Build().CompileWithExampleArguments() + for item in to_infeed: + xla_client.transfer_to_infeed(item) + + for item in to_infeed: + result = compiled_c.Execute() + self.assertEqual(result, item) + + def testInfeedThenOutfeedS32(self): + to_round_trip = NumpyArrayS32([1, 2, 3, 4]) + c = self._NewComputation() + x = c.Infeed(xla_client.Shape.from_numpy(to_round_trip[0])) + c.Outfeed(x) + + compiled_c = c.Build().CompileWithExampleArguments() + + for want in to_round_trip: + execution = threading.Thread(target=compiled_c.Execute) + execution.start() + xla_client.transfer_to_infeed(want) + got = xla_client.transfer_from_outfeed( + xla_client.Shape.from_numpy(to_round_trip[0])) + execution.join() + self.assertEqual(want, got) + + +class ErrorTest(LocalComputationTest): + + def setUp(self): + self.f32_scalar_2 = NumpyArrayF32(2.0) + self.s32_scalar_2 = NumpyArrayS32(2) + + def testInvokeWithWrongElementType(self): + c = self._NewComputation() + c.SetOpMetadata(xla_client.CurrentSourceInfoMetadata()) + c.ParameterFromNumpy(self.s32_scalar_2) + c.ClearOpMetadata() + self.assertRaisesRegexp( + RuntimeError, r"Invalid argument shape.*xla_client_test.py.*" + r"expected s32\[\], got f32\[\]", + lambda: c.Build().CompileWithExampleArguments([self.f32_scalar_2])) + if __name__ == "__main__": unittest.main() diff --git a/tensorflow/compiler/xla/reference_util.cc b/tensorflow/compiler/xla/reference_util.cc index 3be2567e14085b3501d74433f783d9e15e1574e9..a9acdae380af5b7f9efb3d08302fc717108f5e40 100644 --- a/tensorflow/compiler/xla/reference_util.cc +++ b/tensorflow/compiler/xla/reference_util.cc @@ -281,6 +281,51 @@ ReferenceUtil::ReduceWindow1DAdd( return result; } +/* static */ std::unique_ptr> ReferenceUtil::ReduceWindow3DAdd( + const Array3D& operand, float init, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, Padding padding) { + std::vector dim_lengths{operand.n1(), operand.n2(), operand.n3()}; + auto padding_both = xla::MakePadding(dim_lengths, window, stride, padding); + + std::vector window_counts(window.size(), 0); + std::vector pad_low(window.size(), 0); + for (int64 i = 0; i < window.size(); ++i) { + window_counts[i] = + WindowCount(dim_lengths[i], window[i], stride[i], padding); + pad_low[i] = padding_both[i].first; + } + auto result = MakeUnique>(window_counts[0], window_counts[1], + window_counts[2]); + + for (int64 i0 = 0; i0 < window_counts[0]; ++i0) { + for (int64 i1 = 0; i1 < window_counts[1]; ++i1) { + for (int64 i2 = 0; i2 < window_counts[2]; ++i2) { + int64 i0_base = i0 * stride[0] - pad_low[0]; + int64 i1_base = i1 * stride[1] - pad_low[1]; + int64 i2_base = i2 * stride[2] - pad_low[2]; + + float val = init; + for (int64 i0_win = 0; i0_win < window[0]; ++i0_win) { + for (int64 i1_win = 0; i1_win < window[1]; ++i1_win) { + for (int64 i2_win = 0; i2_win < window[2]; ++i2_win) { + if (i0_base + i0_win >= 0 && i1_base + i1_win >= 0 && + i2_base + i2_win >= 0 && i0_base + i0_win < operand.n1() && + i1_base + i1_win < operand.n2() && + i2_base + i2_win < operand.n3()) { + val += operand(i0_base + i0_win, i1_base + i1_win, + i2_base + i2_win); + } + } + } + } + (*result)(i0, i1, i2) = val; + } + } + } + return result; +} + /* static */ std::unique_ptr> ReferenceUtil::ReduceWindow4DGeneric( const Array4D& operand, float init, diff --git a/tensorflow/compiler/xla/reference_util.h b/tensorflow/compiler/xla/reference_util.h index 58e1a844610678f64677838e93f0379b63f65d39..3ec96f2f38b8f91e1549419b60481327fa9bbd5f 100644 --- a/tensorflow/compiler/xla/reference_util.h +++ b/tensorflow/compiler/xla/reference_util.h @@ -173,6 +173,10 @@ class ReferenceUtil { const Array2D& operand, float init, const tensorflow::gtl::ArraySlice& window, const tensorflow::gtl::ArraySlice& stride, Padding padding); + static std::unique_ptr> ReduceWindow3DAdd( + const Array3D& operand, float init, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, Padding padding); static std::unique_ptr> ReduceWindow4DAdd( const Array4D& operand, float init, const tensorflow::gtl::ArraySlice& window, @@ -195,6 +199,7 @@ class ReferenceUtil { const std::function& reduce_func, const tensorflow::gtl::ArraySlice& window, const tensorflow::gtl::ArraySlice& stride, Padding padding); + // With arbitrary padding. static std::unique_ptr> ReduceWindow4DGeneric( const Array4D& operand, float init, const std::function& reduce_func, diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index ae7d2d93bde72adb4ba4230b87b6b02dd1428baa..83c67ed9368bc617a90c528f200b566ee8754edd 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -29,6 +29,11 @@ xla_proto_library( deps = ["//tensorflow/compiler/xla:xla_data_proto"], ) +xla_proto_library( + name = "hlo_profile_printer_data", + srcs = ["hlo_profile_printer_data.proto"], +) + # Filegroup used to collect source files for dependency checking. filegroup( name = "c_srcs", @@ -38,6 +43,81 @@ filegroup( ]), ) +cc_library( + name = "bfloat16_support", + srcs = ["bfloat16_support.cc"], + hdrs = ["bfloat16_support.h"], + deps = [ + ":hlo", + ], +) + +cc_library( + name = "bfloat16_conversion_folding", + srcs = ["bfloat16_conversion_folding.cc"], + hdrs = ["bfloat16_conversion_folding.h"], + deps = [ + ":bfloat16_support", + ":hlo", + ":hlo_pass", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "bfloat16_conversion_folding_test", + srcs = ["bfloat16_conversion_folding_test.cc"], + deps = [ + ":bfloat16_conversion_folding", + ":bfloat16_support", + ":hlo", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", + ], +) + +cc_library( + name = "bfloat16_normalization", + srcs = ["bfloat16_normalization.cc"], + hdrs = ["bfloat16_normalization.h"], + deps = [ + ":bfloat16_support", + ":hlo", + ":hlo_pass", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "bfloat16_normalization_test", + srcs = ["bfloat16_normalization_test.cc"], + deps = [ + ":bfloat16_normalization", + ":bfloat16_support", + ":hlo", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", + ], +) + cc_library( name = "shape_inference", srcs = ["shape_inference.cc"], @@ -452,8 +532,10 @@ cc_library( ":hlo_evaluator", ":hlo_execution_profile", ":hlo_module_config", + ":hlo_proto_util", ":platform_util", ":session_proto", + ":source_map_util", ":transfer_manager", ":user_computation", ":versioned_computation_handle", @@ -502,6 +584,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:executable_build_options", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", ], @@ -905,6 +988,7 @@ tf_cc_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", ], ) @@ -1084,6 +1168,7 @@ tf_cc_test( "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep @@ -1145,6 +1230,34 @@ tf_cc_test( ], ) +cc_library( + name = "implicit_broadcast_remover", + srcs = ["implicit_broadcast_remover.cc"], + hdrs = ["implicit_broadcast_remover.h"], + deps = [ + ":hlo", + ":hlo_dce", + ":hlo_pass", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "implicit_broadcast_remover_test", + srcs = ["implicit_broadcast_remover_test.cc"], + deps = [ + ":hlo_matchers", + ":implicit_broadcast_remover", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + ], +) + cc_library( name = "dot_decomposer", srcs = ["dot_decomposer.cc"], @@ -1680,6 +1793,7 @@ tf_cc_test( ":hlo", ":hlo_graph_dumper", ":hlo_matchers", + ":hlo_runner", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test", @@ -1687,7 +1801,6 @@ tf_cc_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], ) @@ -1814,7 +1927,9 @@ tf_cc_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", "//tensorflow/core:lib", + "//tensorflow/core:test", ], ) @@ -1845,6 +1960,7 @@ cc_library( ":hlo", ":hlo_graph_dumper", ":hlo_pass", + ":hlo_proto_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", @@ -1938,6 +2054,16 @@ cc_library( ], ) +tf_cc_test( + name = "hlo_element_type_converter_test", + srcs = ["hlo_element_type_converter_test.cc"], + deps = [ + ":hlo_element_type_converter", + ":hlo_matchers", + "//tensorflow/compiler/xla/tests:hlo_test_base", + ], +) + cc_library( name = "device_memory_allocator", srcs = ["device_memory_allocator.cc"], @@ -2070,6 +2196,7 @@ cc_library( "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla:xla_proto", "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", "//tensorflow/core:regexp_internal", ], alwayslink = 1, @@ -2249,11 +2376,96 @@ cc_library( srcs = ["hlo_profile_printer.cc"], hdrs = ["hlo_profile_printer.h"], deps = [ + ":hlo_profile_printer_data", ":human_readable_profile_builder", "//tensorflow/compiler/xla:types", ], ) +cc_library( + name = "tuple_util", + srcs = ["tuple_util.cc"], + hdrs = ["tuple_util.h"], + deps = [ + ":hlo", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "tuple_util_test", + srcs = ["tuple_util_test.cc"], + deps = [ + ":tuple_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", + ], +) + +cc_library( + name = "while_util", + srcs = ["while_util.cc"], + hdrs = ["while_util.h"], + deps = [ + ":call_inliner", + ":hlo", + ":tuple_util", + ], +) + +tf_cc_test( + name = "while_util_test", + srcs = ["while_util_test.cc"], + deps = [ + ":while_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", + ], +) + +cc_library( + name = "while_loop_invariant_code_motion", + srcs = ["while_loop_invariant_code_motion.cc"], + hdrs = ["while_loop_invariant_code_motion.h"], + deps = [ + ":hlo", + ":hlo_pass", + ":tuple_util", + ":while_util", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "while_loop_invariant_code_motion_test", + srcs = ["while_loop_invariant_code_motion_test.cc"], + deps = [ + ":hlo_matchers", + ":while_loop_invariant_code_motion", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/core:test", + ], +) + +cc_library( + name = "source_map_util", + srcs = ["source_map_util.cc"], + hdrs = ["source_map_util.h"], + deps = [ + ":executable", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc index 8d476d771499dd21ded24d2f073b228b193b7382..fb857559f972a220a19b108baa4c441e09b90e1f 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc @@ -498,13 +498,14 @@ static HloInstruction* BuildTupleConstant(HloComputation* computation, if (ShapeUtil::IsTuple(literal.shape())) { std::vector elems; elems.reserve(ShapeUtil::TupleElementCount(literal.shape())); - for (const Literal& child : literal.tuple_literals()) { - elems.push_back(BuildTupleConstant(computation, child)); + for (int i = 0; i < ShapeUtil::TupleElementCount(literal.shape()); ++i) { + elems.push_back( + BuildTupleConstant(computation, LiteralView::Create(literal, {i}))); } return computation->AddInstruction(HloInstruction::CreateTuple(elems)); } else { return computation->AddInstruction( - HloInstruction::CreateConstant(MakeUnique(literal))); + HloInstruction::CreateConstant(literal.CloneToUnique())); } } @@ -1617,9 +1618,12 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { reduce, HloInstruction::CreateBroadcast(reduce->shape(), init_value, {})); } + // A Transpose feeding a reduce can simply permute the reduction dimensions - // field. - if (arg->opcode() == HloOpcode::kTranspose) { + // field if the output of the reduce is a vector or scalar. Higher ranked + // result may require a transpose of the output. + if (ShapeUtil::Rank(reduce->shape()) <= 1 && + arg->opcode() == HloOpcode::kTranspose) { auto transpose_dimensions = arg->dimensions(); std::vector new_reduce_dimensions; for (auto dim : dimensions) { @@ -1740,6 +1744,63 @@ Status AlgebraicSimplifierVisitor::HandleReduceWindow( } } + // If the pad puts a single non-identity value in each window that we're + // reducing, then this is a broadcast. + HloInstruction* pad_operand = operand->mutable_operand(0); + auto is_effective_broadcast = [&] { + if (window_util::HasStride(window)) { + VLOG(10) << "Window has stride."; + return false; + } + if (!window_util::HasSymmetricPadding(pad_config)) { + VLOG(10) << "Window has uneven padding."; + return false; + } + for (int64 i = 0; i < pad_config.dimensions_size(); ++i) { + const auto& pad_dimension = pad_config.dimensions(i); + if ((pad_dimension.edge_padding_low() != 0 || + pad_dimension.edge_padding_high() != 0) && + pad_operand->shape().dimensions(i) != 1) { + VLOG(10) << "Found non-trivial dimension being padded: " << i; + return false; + } + } + VLOG(10) << "Found to be padding trivial dimensions only."; + + for (int64 i = 0; i < window.dimensions_size(); ++i) { + const auto& pad_dimension = pad_config.dimensions(i); + const WindowDimension& window_dimension = window.dimensions(i); + bool dimension_has_padding = (pad_dimension.edge_padding_low() != 0 || + pad_dimension.edge_padding_high() != 0); + if (dimension_has_padding && + window_dimension.size() < pad_dimension.edge_padding_low() + 1) { + VLOG(10) << "Found window did not cover single unpadded element in " + "dimension: " + << i; + return false; + } + if (pad_operand->shape().dimensions(i) != 1 && + window_dimension.size() != 1) { + VLOG(10) << "Found window covers more than one element in non-trivial " + "dimension: " + << i; + return false; + } + } + VLOG(10) << "Found window covers a single unpadded element."; + return true; + }; + if (is_effective_broadcast()) { + VLOG(10) << "Replacing pad/reduce-window with (implicit) broadcast."; + auto fadd = [this](std::unique_ptr x) { + return computation_->AddInstruction(std::move(x)); + }; + return ReplaceWithNewInstruction( + reduce_window, HloInstruction::CreateBroadcastSequence( + /*output_shape=*/reduce_window->shape(), + /*operand=*/pad_operand, fadd)); + } + // Carry out the folding of the pad into reduce_window. VLOG(10) << "Folding pad into reduce-window."; Window new_window = window; @@ -1757,7 +1818,7 @@ Status AlgebraicSimplifierVisitor::HandleReduceWindow( return ReplaceWithNewInstruction( reduce_window, HloInstruction::CreateReduceWindow( /*shape=*/reduce_window->shape(), - /*operand=*/operand->mutable_operand(0), + /*operand=*/pad_operand, /*init_value=*/reduce_window->mutable_operand(1), /*window=*/new_window, /*reduce_computation=*/function)); diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc index e7c4dfb0a1690683bbdb7e61067392b48fdba8a5..0f08eb3a3267c4b7b04958270a5788fc48d3fa04 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc @@ -30,6 +30,7 @@ limitations under the License. #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" #include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/lib/strings/str_util.h" @@ -60,13 +61,12 @@ TEST_F(AlgebraicSimplifierTest, AddZero) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, param0, zero)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } @@ -82,13 +82,12 @@ TEST_F(AlgebraicSimplifierTest, AddConstOnLHS) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, constant, param0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_THAT(root, op::Add(param0, op::Constant())); } @@ -109,13 +108,12 @@ TEST_F(AlgebraicSimplifierTest, AddReassociateMergeConstants) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, add1, constant2)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_THAT(root, op::Add(param0, op::Add(constant1, constant2))); } @@ -132,13 +130,12 @@ TEST_F(AlgebraicSimplifierTest, AddBroadcastZeroR0Operand) { builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kAdd, bcast, param0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } @@ -155,13 +152,12 @@ TEST_F(AlgebraicSimplifierTest, AddBroadcastZeroR1Operand) { builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kAdd, bcast, param0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } @@ -177,13 +173,12 @@ TEST_F(AlgebraicSimplifierTest, SubZero) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kSubtract, param0, zero)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kSubtract); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } @@ -199,13 +194,12 @@ TEST_F(AlgebraicSimplifierTest, SubConstCanonicalization) { builder.AddInstruction(HloInstruction::CreateBinary( r0f32, HloOpcode::kSubtract, param0, constant)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kSubtract); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_THAT(root, op::Add(param0, op::Negate(constant))); } @@ -225,15 +219,14 @@ TEST_F(AlgebraicSimplifierTest, LhsDivOfDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, div, param2)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Divide(op::Divide(param0, param1), param2)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Divide(param0, op::Multiply(param1, param2))); @@ -254,15 +247,14 @@ TEST_F(AlgebraicSimplifierTest, RhsDivOfDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, div)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Divide(param0, op::Divide(param1, param2))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Divide(op::Multiply(param0, param2), param1)); @@ -288,8 +280,7 @@ TEST_F(AlgebraicSimplifierTest, DivOfDivAndDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kDivide, div0, div1)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT( computation->root_instruction(), @@ -297,7 +288,7 @@ TEST_F(AlgebraicSimplifierTest, DivOfDivAndDiv) { AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT( computation->root_instruction(), @@ -319,15 +310,14 @@ TEST_F(AlgebraicSimplifierTest, DivOfExp) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, exp)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Divide(param0, op::Exp(param1))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Multiply(param0, op::Exp(op::Negate(param1)))); @@ -348,15 +338,14 @@ TEST_F(AlgebraicSimplifierTest, DivOfPower) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, power)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Divide(param0, op::Power(param1, param2))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Multiply(param0, op::Power(param1, op::Negate(param2)))); @@ -379,15 +368,14 @@ TEST_F(AlgebraicSimplifierTest, DivOfBroadcastingPower) { builder.AddInstruction( HloInstruction::CreateBinary(r1f32, HloOpcode::kDivide, param0, power)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Divide(param0, op::Power(param1, param2))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); ASSERT_THAT(computation->root_instruction(), op::Multiply(param0, op::Power(param1, op::Negate(param2)))); @@ -410,12 +398,11 @@ TEST_F(AlgebraicSimplifierTest, DivideByConstant) { builder.AddInstruction(HloInstruction::CreateBinary(r1f32, HloOpcode::kDivide, param0, constant)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Multiply(param0, op::Divide(op::Constant(), constant))); @@ -437,11 +424,10 @@ TEST_F(AlgebraicSimplifierTest, PowerOfPower) { builder.AddInstruction(HloInstruction::CreateBinary(r1f32, HloOpcode::kPower, inner_power, exp2)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Power(base, op::Multiply(exp1, exp2))); } @@ -450,24 +436,23 @@ TEST_F(AlgebraicSimplifierTest, PowerOfPower) { // numbers. TEST_F(AlgebraicSimplifierTest, PowerOfPowerComplex) { Shape r0c64 = ShapeUtil::MakeShape(C64, {}); - Shape r1f32 = ShapeUtil::MakeShape(F32, {7}); + Shape r1c64 = ShapeUtil::MakeShape(C64, {7}); HloComputation::Builder builder(TestName()); HloInstruction* base = builder.AddInstruction( - HloInstruction::CreateParameter(0, r1f32, "param0")); + HloInstruction::CreateParameter(0, r1c64, "param0")); HloInstruction* exp1 = builder.AddInstruction( HloInstruction::CreateParameter(1, r0c64, "param1")); HloInstruction* exp2 = builder.AddInstruction( HloInstruction::CreateParameter(2, r0c64, "param2")); HloInstruction* inner_power = builder.AddInstruction( - HloInstruction::CreateBinary(r1f32, HloOpcode::kPower, base, exp1)); - builder.AddInstruction(HloInstruction::CreateBinary(r1f32, HloOpcode::kPower, + HloInstruction::CreateBinary(r1c64, HloOpcode::kPower, base, exp1)); + builder.AddInstruction(HloInstruction::CreateBinary(r1c64, HloOpcode::kPower, inner_power, exp2)); - auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); + module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_FALSE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_FALSE(simplifier.Run(&module()).ValueOrDie()); } // Test that A/1 is simplified to A for a scalar. @@ -481,13 +466,12 @@ TEST_F(AlgebraicSimplifierTest, DivOneScalar) { HloInstruction* div = builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, one)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, div); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } @@ -503,13 +487,12 @@ TEST_F(AlgebraicSimplifierTest, DivOneArray) { HloInstruction* div = builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kDivide, param0, one)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, div); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } @@ -528,13 +511,12 @@ TEST_F(AlgebraicSimplifierTest, ComplexOfRealImagC) { HloInstruction* cplx = builder.AddInstruction( HloInstruction::CreateBinary(r2c64, HloOpcode::kComplex, real, imag)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, cplx); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } @@ -553,13 +535,12 @@ TEST_F(AlgebraicSimplifierTest, RealOfComplex) { HloInstruction* real = builder.AddInstruction( HloInstruction::CreateUnary(r2f32, HloOpcode::kReal, cplx)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, real); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } @@ -578,13 +559,12 @@ TEST_F(AlgebraicSimplifierTest, ImagOfComplex) { HloInstruction* imag = builder.AddInstruction( HloInstruction::CreateUnary(r2f32, HloOpcode::kImag, cplx)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, imag); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param1); } @@ -606,13 +586,12 @@ TEST_F(AlgebraicSimplifierTest, SelectMakeTuple) { HloInstruction* add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, get, param2)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, add); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); root = computation->root_instruction(); EXPECT_THAT(root, op::Add(param1, param2)); } @@ -632,15 +611,14 @@ TEST_F(AlgebraicSimplifierTest, ExpDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, exp0, exp1)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Divide(op::Exp(param0), op::Exp(param1))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Exp(op::Subtract(param0, param1))); @@ -661,15 +639,14 @@ TEST_F(AlgebraicSimplifierTest, ExpMul) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kMultiply, exp0, exp1)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Multiply(op::Exp(param0), op::Exp(param1))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Exp(op::Add(param0, param1))); @@ -688,15 +665,14 @@ TEST_F(AlgebraicSimplifierTest, PowExp) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, exp0, param1)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Power(op::Exp(param0), param1)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Exp(op::Multiply(param0, param1))); @@ -715,15 +691,14 @@ TEST_F(AlgebraicSimplifierTest, LnPow) { builder.AddInstruction( HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, pow)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Log(op::Power(param0, param1))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Multiply(op::Log(param0), param1)); @@ -740,14 +715,13 @@ TEST_F(AlgebraicSimplifierTest, LnExp) { builder.AddInstruction( HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, exp0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Log(op::Exp(param0))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), param0); } @@ -769,15 +743,14 @@ TEST_F(AlgebraicSimplifierTest, LnExpDiv) { builder.AddInstruction( HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, div)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Log(op::Divide(op::Exp(param0), op::Exp(param1)))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Subtract(param0, param1)); } @@ -794,14 +767,13 @@ TEST_F(AlgebraicSimplifierTest, Pow0Scalar) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, zero)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Power(param0, zero)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); EXPECT_THAT(root, op::Constant()); @@ -819,14 +791,13 @@ TEST_F(AlgebraicSimplifierTest, Pow0Vector) { builder.AddInstruction( HloInstruction::CreateBinary(r1f32, HloOpcode::kPower, param0, zero)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Power(param0, zero)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); EXPECT_THAT(root, op::Broadcast()); @@ -848,14 +819,13 @@ TEST_F(AlgebraicSimplifierTest, Pow1) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, one)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Power(param0, one)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), param0); } @@ -871,14 +841,13 @@ TEST_F(AlgebraicSimplifierTest, Pow2) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, two)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Power(param0, two)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Multiply(param0, param0)); } @@ -894,14 +863,13 @@ TEST_F(AlgebraicSimplifierTest, PowNegative1) { builder.AddInstruction(HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, negative_one)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Power(param0, negative_one)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); EXPECT_THAT(root, op::Divide(op::Broadcast(), param0)); @@ -940,16 +908,15 @@ TEST_F(AlgebraicSimplifierTest, ZeroSizedConvolution) { dim->set_base_dilation(1); dim->set_window_reversal(false); // Create add computation. - std::unique_ptr module = CreateNewModule(); builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {3, 3, 3}), lhs, rhs, window, dnums)); - module->AddEntryComputation(builder.Build()); + module().AddEntryComputation(builder.Build()); HloPassFix simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - EXPECT_THAT(module->entry_computation()->root_instruction(), + EXPECT_THAT(module().entry_computation()->root_instruction(), op::Convolution(lhs, rhs)); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_THAT(module->entry_computation()->root_instruction(), + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + EXPECT_THAT(module().entry_computation()->root_instruction(), op::Broadcast(op::Constant())); } @@ -968,7 +935,6 @@ TEST_F(AlgebraicSimplifierTest, ZeroSizedReduceWindow) { dim->set_base_dilation(1); } // Create add computation. - std::unique_ptr module = CreateNewModule(); HloComputation* add_computation = nullptr; { HloComputation::Builder builder(TestName() + ".add"); @@ -979,20 +945,20 @@ TEST_F(AlgebraicSimplifierTest, ZeroSizedReduceWindow) { HloInstruction::CreateParameter(1, scalar_shape, "p1")); builder.AddInstruction( HloInstruction::CreateBinary(scalar_shape, HloOpcode::kAdd, p0, p1)); - add_computation = module->AddEmbeddedComputation(builder.Build()); + add_computation = module().AddEmbeddedComputation(builder.Build()); } builder.AddInstruction(HloInstruction::CreateReduceWindow( ShapeUtil::MakeShape(F32, {5, 2}), param, builder.AddInstruction( HloInstruction::CreateConstant(Literal::CreateR0(0.0f))), window, add_computation)); - module->AddEntryComputation(builder.Build()); + module().AddEntryComputation(builder.Build()); HloPassFix simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - EXPECT_THAT(module->entry_computation()->root_instruction(), + EXPECT_THAT(module().entry_computation()->root_instruction(), op::ReduceWindow(param, op::Constant())); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_THAT(module->entry_computation()->root_instruction(), + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + EXPECT_THAT(module().entry_computation()->root_instruction(), op::Broadcast(op::Constant())); } @@ -1013,14 +979,13 @@ TEST_F(AlgebraicSimplifierTest, ZeroSizedPad) { builder.AddInstruction( HloInstruction::CreateConstant(Literal::CreateR0(0.0f))), padding)); - std::unique_ptr module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); - EXPECT_THAT(module->entry_computation()->root_instruction(), + module().AddEntryComputation(builder.Build()); + EXPECT_THAT(module().entry_computation()->root_instruction(), op::Pad(param, op::Constant())); HloPassFix simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_THAT(module->entry_computation()->root_instruction(), + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + EXPECT_THAT(module().entry_computation()->root_instruction(), op::Broadcast(op::Constant())); } @@ -1038,17 +1003,16 @@ TEST_F(AlgebraicSimplifierTest, ReshapeBroadcast) { ShapeUtil::MakeShape(F32, {3, 2}), broadcast)); auto computation = builder.Build(); - auto module = CreateNewModule(); - module->AddEntryComputation(std::move(computation)); + module().AddEntryComputation(std::move(computation)); - EXPECT_THAT(module->entry_computation()->root_instruction(), + EXPECT_THAT(module().entry_computation()->root_instruction(), op::Reshape(op::Broadcast(op::Reshape(op)))); HloPassFix simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); - EXPECT_THAT(module->entry_computation()->root_instruction(), op); + EXPECT_THAT(module().entry_computation()->root_instruction(), op); } // Test that convert(A, $TYPE) is simplified to A if A is of type $TYPE. @@ -1059,14 +1023,13 @@ TEST_F(AlgebraicSimplifierTest, ConvertBetweenSameType) { builder.AddInstruction( HloInstruction::CreateConvert(ShapeUtil::MakeShape(F32, {}), input)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Convert(input)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), input); } @@ -1080,14 +1043,13 @@ TEST_F(AlgebraicSimplifierTest, RemoveCopy) { builder.AddInstruction( HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), param0); } @@ -1101,14 +1063,13 @@ TEST_F(AlgebraicSimplifierTest, RemoveUnaryConcatenate) { builder.AddInstruction( HloInstruction::CreateConcatenate(param0->shape(), {param0}, 0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Concatenate(param0)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), param0); } @@ -1131,8 +1092,7 @@ TEST_F(AlgebraicSimplifierTest, RemoveEmptyConcatenateOperands) { builder.AddInstruction(HloInstruction::CreateConcatenate( result_shape, {empty_literal, param0, param0, empty_slice, param1}, 0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT( computation->root_instruction(), @@ -1140,7 +1100,7 @@ TEST_F(AlgebraicSimplifierTest, RemoveEmptyConcatenateOperands) { AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Concatenate(param0, param0, param1)); @@ -1162,15 +1122,14 @@ TEST_F(AlgebraicSimplifierTest, OnlyEmptyConcatenateOperands) { builder.AddInstruction(HloInstruction::CreateConcatenate( result_shape, {empty_literal, empty_slice}, 0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Concatenate(empty_literal, empty_slice)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), empty_literal); } @@ -1187,14 +1146,13 @@ TEST_F(AlgebraicSimplifierTest, ConcatenateOfBroadcastBecomesPad) { HloInstruction* broadcast = builder.AddInstruction( HloInstruction::CreateBroadcast(r1f32, param1, {})); builder.AddInstruction(HloInstruction::CreateConcatenate( - param0->shape(), {broadcast, param0}, 0)); + ShapeUtil::MakeShape(F32, {200}), {broadcast, param0}, 0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Pad(param0, param1)); } @@ -1208,8 +1166,7 @@ TEST_F(AlgebraicSimplifierTest, CopyWithDifferentLayout) { HloInstruction* copy = builder.AddInstruction( HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); // Set to different layouts. *param0->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); @@ -1219,7 +1176,7 @@ TEST_F(AlgebraicSimplifierTest, CopyWithDifferentLayout) { AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); + EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); // Copy has not been removed. EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); @@ -1235,8 +1192,7 @@ TEST_F(AlgebraicSimplifierTest, CopyWithSameLayout) { HloInstruction* copy = builder.AddInstruction( HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); // Set to same layouts. *param0->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); @@ -1246,7 +1202,7 @@ TEST_F(AlgebraicSimplifierTest, CopyWithSameLayout) { AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); // Copy has been removed. EXPECT_THAT(computation->root_instruction(), param0); @@ -1267,14 +1223,13 @@ TEST_F(AlgebraicSimplifierTest, NoBitcastAdded) { *reshape->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2, 3, 4, 5}); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); + EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); // Reshape is not replaced with a bitcast. EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); @@ -1313,8 +1268,7 @@ TEST_F(AlgebraicSimplifierTest, ReshapeReplacedWithBitcast) { builder.AddInstruction(HloInstruction::CreateTuple( {transformable_reshape, dimensions_wrong_reshape, layout_wrong_reshape})); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Tuple(transformable_reshape, dimensions_wrong_reshape, @@ -1322,7 +1276,7 @@ TEST_F(AlgebraicSimplifierTest, ReshapeReplacedWithBitcast) { AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, bitcasting_callback()); - simplifier.Run(module.get()).ValueOrDie(); + simplifier.Run(&module()).ValueOrDie(); // Verify that only the first reshape is replaced. EXPECT_THAT( @@ -1343,8 +1297,7 @@ TEST_F(AlgebraicSimplifierTest, ReshapeAfterEffectiveUnary) { builder.AddInstruction( HloInstruction::CreateBinary(ShapeUtil::MakeShape(F32, {1, 2, 3, 4, 5}), HloOpcode::kMaximum, movable_reshape, zero)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Maximum(op::Reshape(param), zero)); @@ -1352,7 +1305,7 @@ TEST_F(AlgebraicSimplifierTest, ReshapeAfterEffectiveUnary) { AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, bitcasting_callback()); - simplifier.Run(module.get()).ValueOrDie(); + simplifier.Run(&module()).ValueOrDie(); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Maximum(param, zero))); } @@ -1370,8 +1323,7 @@ TEST_F(AlgebraicSimplifierTest, ReshapeToScalarNotHoistedAfterEffectiveUnary) { HloInstruction::CreateConstant(Literal::CreateR1({1., 2., 3.}))); builder.AddInstruction(HloInstruction::CreateBinary( ShapeUtil::MakeShape(F32, {3}), HloOpcode::kMaximum, reshape, zero)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Maximum(op::Reshape(param), zero)); @@ -1379,7 +1331,7 @@ TEST_F(AlgebraicSimplifierTest, ReshapeToScalarNotHoistedAfterEffectiveUnary) { AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, bitcasting_callback()); - simplifier.Run(module.get()).ValueOrDie(); + simplifier.Run(&module()).ValueOrDie(); EXPECT_THAT(computation->root_instruction(), op::Maximum(op::Reshape(param), zero)); @@ -1404,9 +1356,8 @@ TEST_F(AlgebraicSimplifierTest, FailureToSinkReshapeDoesntAffectChangedBit) { AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, bitcasting_callback()); - auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + module().AddEntryComputation(builder.Build()); + EXPECT_TRUE(simplifier.Run(&module()).ValueOrDie()); } // Regression test for a bug where if we failed to sink a reshape, we'd set the @@ -1423,14 +1374,14 @@ TEST_F(AlgebraicSimplifierTest, FailureToSinkBroadcastDoesntAffectChangedBit) { builder.AddInstruction(HloInstruction::CreateConstant( Literal::CreateR2({{0, 0}, {0, 0}}))))); - builder.AddInstruction(HloInstruction::CreateBroadcast( - ShapeUtil::MakeShape(F32, {2, 2, 2}), add, /*broadcast_dimensions=*/{0})); + builder.AddInstruction( + HloInstruction::CreateBroadcast(ShapeUtil::MakeShape(F32, {2, 2, 2}), add, + /*broadcast_dimensions=*/{0, 1})); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, bitcasting_callback()); - auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + module().AddEntryComputation(builder.Build()); + EXPECT_TRUE(simplifier.Run(&module()).ValueOrDie()); } TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast1) { @@ -1447,14 +1398,13 @@ TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast1) { *transpose->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2, 3}); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Transpose(param)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); // Verify that the reshape is replaced. EXPECT_THAT(computation->root_instruction(), op::Bitcast(param)); @@ -1474,14 +1424,13 @@ TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast2) { *transpose->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({3, 1, 2, 0}); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Transpose(param)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); // Verify that the reshape is replaced. EXPECT_THAT(computation->root_instruction(), op::Bitcast(param)); @@ -1500,15 +1449,14 @@ TEST_F(AlgebraicSimplifierTest, ReshapesMerged) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {1, 2, 1, 1, 2, 1}), reshape1)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Reshape(param0))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); } @@ -1528,14 +1476,13 @@ TEST_F(AlgebraicSimplifierTest, CopiesMerged) { ShapeUtil::MakeShapeWithLayout(F32, {2, 2, 2}, {0, 2, 1}), HloOpcode::kCopy, copy1)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Copy(op::Copy(param0))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); } @@ -1553,14 +1500,13 @@ TEST_F(AlgebraicSimplifierTest, TransposesMerged) { builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {4, 3, 2}), transpose1, {1, 0, 2})); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Transpose(transpose1)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Transpose(param0)); EXPECT_EQ(std::vector({2, 1, 0}), @@ -1575,17 +1521,16 @@ TEST_F(AlgebraicSimplifierTest, ReshapeAndBroadcastMerged) { auto reshape1 = builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {1, 5, 1}), param0)); builder.AddInstruction(HloInstruction::CreateBroadcast( - ShapeUtil::MakeShape(F32, {1, 2, 3, 5, 1}), reshape1, {0, 2, 3})); + ShapeUtil::MakeShape(F32, {1, 2, 3, 5, 1}), reshape1, {0, 3, 2})); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Broadcast(op::Reshape(param0))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Broadcast(param0)); } @@ -1600,15 +1545,14 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshapeMerged) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {2, 3, 7, 2, 1, 3, 2}), broadcast1)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Broadcast(param0))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Broadcast(param0)); } @@ -1622,15 +1566,14 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x1_3) { builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {3}), broadcast)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Broadcast(param))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); + EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Broadcast(param))); @@ -1645,15 +1588,14 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4_6x1x1x4) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 1, 1, 4}), broadcast)); - auto module = CreateNewModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); + HloComputation* computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Broadcast(param))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Broadcast(param)); EXPECT_THAT(computation->root_instruction()->dimensions(), @@ -1669,15 +1611,14 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x2x1_6x1x1x1) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 1, 1, 1}), broadcast)); - auto module = CreateNewModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); + HloComputation* computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Broadcast(param))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Broadcast(param)); const std::vector broadcast_dims = @@ -1695,15 +1636,14 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4x2_6x8) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 8}), broadcast)); - auto module = CreateNewModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); + HloComputation* computation = module().AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Broadcast(param))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); + EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Broadcast(param))); @@ -2409,12 +2349,11 @@ TEST_F(AlgebraicSimplifierTest, IteratorInvalidation) { call_builder.AddInstruction( HloInstruction::CreateCall(r1f32, {zero, one}, dot_computation.get())); - auto module = CreateNewModule(); - module->AddEmbeddedComputation(std::move(dot_computation)); - module->AddEntryComputation(call_builder.Build()); + module().AddEmbeddedComputation(std::move(dot_computation)); + module().AddEntryComputation(call_builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); } // Test that a constant with tuple shape becomes a tuple of constants. @@ -2427,12 +2366,11 @@ TEST_F(AlgebraicSimplifierTest, ConstantTupleBecomesTupleOfConstants) { Literal::CreateR1(constant_vector).get()}); builder.AddInstruction(HloInstruction::CreateConstant(std::move(value))); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Tuple(op::Constant(), op::Constant())); } @@ -2452,11 +2390,10 @@ TEST_F(AlgebraicSimplifierTest, TrivialDynamicSlice) { HloInstruction::CreateConstant(Literal::CreateR1({0, 0, 0}))), /*slice_sizes=*/{10, 100, 1000})); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Parameter()); } @@ -2486,15 +2423,153 @@ TEST_F(AlgebraicSimplifierTest, TrivialDynamicUpdateSlice) { builder.AddInstruction( HloInstruction::CreateConstant(Literal::CreateR1({0, 0, 0}))))); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::DynamicSlice(op::Parameter(), op::Parameter())); } +struct PadReduceWindowEffectiveBroadcastCase { + std::vector input_spatials; + std::vector symmetric_pad_spatials; + std::vector reduce_window_spatials; + // Whether to use `B F S0 S1` form vs `B S0 S1 F` form. + // + // This doesn't test any different functionality but is useful for making sure + // kBroadcast nodes are well formed. + bool prepend_a; + bool should_become_broadcast; + + string ToTestCaseName() const { + return tensorflow::strings::StrCat( + tensorflow::str_util::Join(input_spatials, ","), ";", + tensorflow::str_util::Join(symmetric_pad_spatials, ","), ";", + tensorflow::str_util::Join(reduce_window_spatials, ","), ";", prepend_a, + ";", should_become_broadcast); + } +}; + +void PrintTo(const PadReduceWindowEffectiveBroadcastCase& c, std::ostream* os) { + *os << c.ToTestCaseName(); +} + +class PadReduceWindowEffectiveBroadcastTest + : public AlgebraicSimplifierTest, + public ::testing::WithParamInterface< + PadReduceWindowEffectiveBroadcastCase> {}; + +TEST_P(PadReduceWindowEffectiveBroadcastTest, DoIt) { + const auto& param = GetParam(); + + // a and b are parallel bounds we can either turn into a B F S0 S1 or + // `B S0 S1 F` kind of pattern. + auto decorate_spatials = [¶m](tensorflow::gtl::ArraySlice spatials, + int64 a, int64 b) { + std::vector result; + if (param.prepend_a) { + result.push_back(a); + } + for (int64 s : spatials) { + result.push_back(s); + } + if (!param.prepend_a) { + result.push_back(a); + } + result.push_back(b); + return result; + }; + + HloComputation::Builder builder(TestName()); + const Shape input_shape = ShapeUtil::MakeShape( + F32, decorate_spatials(param.input_spatials, 128, 2048)); + HloInstruction* input = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape, "input")); + + PaddingConfig padding = window_util::MakeSymmetricPadding( + decorate_spatials(param.symmetric_pad_spatials, 0, 0)); + TF_ASSERT_OK_AND_ASSIGN( + const Shape pad_shape, + ShapeInference::InferPadShape(input->shape(), + ShapeUtil::MakeShape(F32, {}), padding)); + HloInstruction* pad = builder.AddInstruction(HloInstruction::CreatePad( + pad_shape, input, + builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))), + padding)); + + HloComputation* add_computation = nullptr; + { + HloComputation::Builder builder(TestName() + ".add"); + const Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + HloInstruction* p0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "p0")); + HloInstruction* p1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "p1")); + builder.AddInstruction( + HloInstruction::CreateBinary(scalar_shape, HloOpcode::kAdd, p0, p1)); + add_computation = module().AddEmbeddedComputation(builder.Build()); + } + + Window window = window_util::MakeWindow( + decorate_spatials(param.reduce_window_spatials, 1, 1)); + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + TF_ASSERT_OK_AND_ASSIGN(const Shape output_shape, + ShapeInference::InferReduceWindowShape( + pad->shape(), zero->shape(), window, + add_computation->ComputeProgramShape())); + builder.AddInstruction(HloInstruction::CreateReduceWindow( + output_shape, pad, zero, window, add_computation)); + + auto computation = module().AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module())); + ASSERT_TRUE(run_successful); + + EXPECT_TRUE( + ShapeUtil::Equal(computation->root_instruction()->shape(), output_shape)); + + if (param.should_become_broadcast) { + EXPECT_THAT(computation->root_instruction(), op::Broadcast(::testing::_)); + } else { + EXPECT_THAT(computation->root_instruction(), + op::ReduceWindow(::testing::_, zero)); + } +} + +const std::vector& +PadReduceWindowEffectiveBroadcastCases() { + static auto* cases = new std::vector{ + {/*input_spatials=*/{1, 1}, /*symmetric_pad_amount=*/{6, 6}, + /*reduce_window_spatials=*/{7, 7}, /*prepend_a=*/true, + /*should_become_broadcast=*/true}, // + {/*input_spatials=*/{1, 1}, /*symmetric_pad_amount=*/{6, 6}, + /*reduce_window_spatials=*/{7, 7}, /*prepend_a=*/false, + /*should_become_broadcast=*/true}, // + {/*input_spatials=*/{2, 2}, /*symmetric_pad_amount=*/{6, 6}, + /*reduce_window_spatials=*/{7, 7}, /*prepend_a=*/true, + /*should_become_broadcast=*/false}, // + {/*input_spatials=*/{1, 1}, /*symmetric_pad_amount=*/{2, 2}, + /*reduce_window_spatials=*/{5, 5}, /*prepend_a=*/true, + /*should_become_broadcast=*/true}, // + {/*input_spatials=*/{1, 1}, /*symmetric_pad_amount=*/{2, 2}, + /*reduce_window_spatials=*/{1, 1}, /*prepend_a=*/true, + /*should_become_broadcast=*/false}, // + {/*input_spatials=*/{5, 1}, /*symmetric_pad_amount=*/{0, 2}, + /*reduce_window_spatials=*/{2, 5}, /*prepend_a=*/true, + /*should_become_broadcast=*/false}, // + }; + return *cases; +} + +INSTANTIATE_TEST_CASE_P( + PadReduceWindowEffectiveBroadcastInstantiation, + PadReduceWindowEffectiveBroadcastTest, + ::testing::ValuesIn(PadReduceWindowEffectiveBroadcastCases())); + class DotStrengthReductionTest : public AlgebraicSimplifierTest, public ::testing::WithParamInterface< @@ -2528,11 +2603,10 @@ TEST_P(DotStrengthReductionTest, DotStrengthReduction) { dot_dnums.add_rhs_contracting_dimensions(0); builder.AddInstruction( HloInstruction::CreateDot(dot_shape, lhs, rhs, dot_dnums)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(module.get())); + TF_ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(&module())); const bool dot_should_be_transformed = m == 1 || k == 1 || n == 1; const bool computation_should_be_modified = dot_should_be_transformed || (transpose_lhs && transpose_rhs); @@ -2560,7 +2634,7 @@ struct DotOfConcatTestSpec { }; class DotOfConcatSimplificationTest - : public HloTestBase, + : public HloVerifiedTestBase, public ::testing::WithParamInterface {}; // Test that we transform @@ -2606,11 +2680,10 @@ TEST_P(DotOfConcatSimplificationTest, ConstantLHS) { builder.AddInstruction( HloInstruction::CreateDot(dot_shape, lhs, rhs, dot_dnums)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(module.get())); + TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module())); ASSERT_TRUE(run_successful); EXPECT_TRUE( @@ -2651,17 +2724,17 @@ TEST_P(DotOfConcatSimplificationTest, ConstantRHS) { HloInstruction* lhs2 = builder.AddInstruction( HloInstruction::CreateParameter(2, lhs2_shape, "lhs2")); HloInstruction* lhs3 = builder.AddInstruction( - HloInstruction::CreateParameter(3, lhs2_shape, "lhs3")); + HloInstruction::CreateParameter(3, lhs3_shape, "lhs3")); Shape lhs_shape = ShapeUtil::MakeShape(F32, {spec.m, spec.k}); HloInstruction* lhs = builder.AddInstruction(HloInstruction::CreateConcatenate( lhs_shape, {lhs0, lhs1, lhs2, lhs3}, 1)); - Shape rhs_shape = ShapeUtil::MakeShape(F32, {spec.k, spec.m}); + Shape rhs_shape = ShapeUtil::MakeShape(F32, {spec.k, spec.n}); auto* rhs = builder.AddInstruction( HloInstruction::CreateConstant(Literal::CreateR2F32Linspace( - /*from=*/10.0, /*to=*/10000.0, /*rows=*/spec.k, /*cols=*/spec.m))); + /*from=*/10.0, /*to=*/10000.0, /*rows=*/spec.k, /*cols=*/spec.n))); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(1); @@ -2671,11 +2744,10 @@ TEST_P(DotOfConcatSimplificationTest, ConstantRHS) { builder.AddInstruction( HloInstruction::CreateDot(dot_shape, lhs, rhs, dot_dnums)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build()); + auto computation = module().AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(module.get())); + TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module())); ASSERT_TRUE(run_successful); EXPECT_TRUE( ShapeUtil::Equal(computation->root_instruction()->shape(), dot_shape)); diff --git a/tensorflow/compiler/xla/service/allocation_tracker.h b/tensorflow/compiler/xla/service/allocation_tracker.h index 8b25cbb482720f7debe95bb5ff74afe696bd8b73..807af8694972083d097604a67ee46d2f73d9545a 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.h +++ b/tensorflow/compiler/xla/service/allocation_tracker.h @@ -67,7 +67,7 @@ class AllocationTracker { // The device that the memory is allocated on. int device_ordinal; - // This is the number of times this memory allocation is refered to by + // This is the number of times this memory allocation is referred to by // registered data handles. int ref_count; }; diff --git a/tensorflow/compiler/xla/service/batchnorm_expander.cc b/tensorflow/compiler/xla/service/batchnorm_expander.cc index b806d61663e2ca371d90dfe39d7fe66becfe4bc2..27ddfd47aa3096afd3e245af1ac3cedd9b48ce4a 100644 --- a/tensorflow/compiler/xla/service/batchnorm_expander.cc +++ b/tensorflow/compiler/xla/service/batchnorm_expander.cc @@ -286,9 +286,11 @@ Status BatchNormExpanderVisitor::HandleBatchNormTraining( int64 instruction_count_after = computation_->instruction_count(); CHECK_EQ(instruction_count_after, instruction_count_before + added_instructions.size()); + HloSharding operand_sharding = + batch_norm->sharding().GetAsShapeTree(batch_norm->shape()).element({0}); for (HloInstruction* inst : added_instructions) { if (ShapeUtil::Equal(inst->shape(), operand_shape)) { - inst->set_sharding(batch_norm->sharding()); + inst->set_sharding(operand_sharding); } else { inst->set_sharding(HloSharding::Replicate()); } @@ -578,9 +580,11 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( int64 instruction_count_after = computation_->instruction_count(); CHECK_EQ(instruction_count_after, instruction_count_before + added_instructions.size()); + HloSharding activation_sharding = + batch_norm->sharding().GetAsShapeTree(batch_norm->shape()).element({0}); for (HloInstruction* inst : added_instructions) { if (ShapeUtil::Equal(inst->shape(), activation_shape)) { - inst->set_sharding(batch_norm->sharding()); + inst->set_sharding(activation_sharding); } else { inst->set_sharding(HloSharding::Replicate()); } diff --git a/tensorflow/compiler/xla/service/bfloat16_conversion_folding.cc b/tensorflow/compiler/xla/service/bfloat16_conversion_folding.cc new file mode 100644 index 0000000000000000000000000000000000000000..cde990e176ddb57a8e93ecc3c60260b2dbae32a8 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_conversion_folding.cc @@ -0,0 +1,184 @@ +/* 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/compiler/xla/service/bfloat16_conversion_folding.h" + +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +class BFloat16ConversionFoldingVisitor : public DfsHloVisitorWithDefault { + public: + explicit BFloat16ConversionFoldingVisitor( + HloComputation* computation, const BFloat16Support* bfloat16_support) + : computation_(computation), bfloat16_support_(bfloat16_support) {} + + Status DefaultAction(HloInstruction* hlo) override; + + static bool Run(HloComputation* computation, + const BFloat16Support* bfloat16_support) { + BFloat16ConversionFoldingVisitor visitor(computation, bfloat16_support); + TF_CHECK_OK(computation->Accept(&visitor)); + return visitor.changed_; + } + + private: + // Checks if the HLO has a BF16 -> F32 conversion as input, or a F32 -> BF16 + // conversion as output, and folds them to the HLO itself if feasible. + Status TryFoldBF16Conversions(HloInstruction* hlo); + + // Folds the F32 -> BF16 conversions from the HLO's output. + // + // Precondition: all of the HLO's users are F32 -> BF16 conversions. + Status FoldOutputConversions(HloInstruction* hlo); + + // Folds the BF16 -> F32 conversion operand to the HLO. + // + // Precondition: the operand is a F32 -> BF16 conversion. + Status FoldOperandConversion(HloInstruction* hlo, int64 operand_index); + + HloComputation* computation_; + const BFloat16Support* bfloat16_support_; + bool changed_ = false; +}; + +Status BFloat16ConversionFoldingVisitor::FoldOutputConversions( + HloInstruction* hlo) { + std::vector materialized_users = hlo->users(); + hlo->mutable_shape()->set_element_type(BF16); + for (auto user : materialized_users) { + CHECK_EQ(user->opcode(), HloOpcode::kConvert); + TF_RETURN_IF_ERROR(user->ReplaceAllUsesWith(hlo)); + changed_ = true; + } + return Status::OK(); +} + +Status BFloat16ConversionFoldingVisitor::FoldOperandConversion( + HloInstruction* hlo, int64 operand_index) { + // The operand is a convert from BF16 to F32. + auto operand = hlo->mutable_operand(operand_index); + CHECK_EQ(operand->opcode(), HloOpcode::kConvert); + TF_RETURN_IF_ERROR( + hlo->ReplaceOperandWith(operand_index, operand->mutable_operand(0))); + changed_ = true; + return Status::OK(); +} + +Status BFloat16ConversionFoldingVisitor::TryFoldBF16Conversions( + HloInstruction* hlo) { + std::vector bf16_to_f32_operands; + bool has_other_f32_operands = false; + for (int64 i = 0; i < hlo->operands().size(); ++i) { + auto operand = hlo->operand(i); + if (operand->shape().element_type() == F32) { + if (operand->opcode() == HloOpcode::kConvert && + operand->operand(0)->shape().element_type() == BF16 && + bfloat16_support_->SupportsBF16Operand(*hlo, i)) { + // Operand is a convert from BF16 to F32 and we support BF16 input + // directly in the current HLO at the operand index. + bf16_to_f32_operands.push_back(i); + } else { + has_other_f32_operands = true; + } + continue; + } + } + + bool fold_output_conversion = hlo->user_count() > 0 && + hlo->shape().element_type() == F32 && + bfloat16_support_->SupportsBF16Output(*hlo) && + hlo != computation_->root_instruction(); + if (fold_output_conversion) { + for (auto user : hlo->users()) { + if (user->opcode() == HloOpcode::kConvert && + user->shape().element_type() == BF16) { + continue; + } + // We should not change the output type if any user is not a conversion + // from F32 to BF16. + fold_output_conversion = false; + break; + } + } + + if (!bfloat16_support_->SupportsMixedPrecisions(*hlo)) { + if (has_other_f32_operands || + (!fold_output_conversion && hlo->shape().element_type() == F32)) { + // Some of the operands/output will remain F32, but we cannot use mixed + // precisions, so we cannot do anything here. + return Status::OK(); + } + } + + if (fold_output_conversion) { + TF_RETURN_IF_ERROR(FoldOutputConversions(hlo)); + } + + for (int64 i : bf16_to_f32_operands) { + TF_RETURN_IF_ERROR(FoldOperandConversion(hlo, i)); + } + return Status::OK(); +} + +Status BFloat16ConversionFoldingVisitor::DefaultAction(HloInstruction* hlo) { + // Do not fold BF16 conversions for instructions related to tuples, entry and + // exit of a computation, fusion, convert, and control flow. + if (hlo->opcode() == HloOpcode::kTuple || // + hlo->opcode() == HloOpcode::kGetTupleElement || // + hlo->opcode() == HloOpcode::kInfeed || // + hlo->opcode() == HloOpcode::kOutfeed || // + hlo->opcode() == HloOpcode::kConstant || // + hlo->opcode() == HloOpcode::kParameter || // + hlo->opcode() == HloOpcode::kFusion || // + hlo->opcode() == HloOpcode::kConvert || // + hlo->opcode() == HloOpcode::kCall || // + hlo->opcode() == HloOpcode::kCustomCall || // + hlo->opcode() == HloOpcode::kWhile || // + hlo->opcode() == HloOpcode::kConditional) { + return Status::OK(); + } + if (hlo == computation_->root_instruction() && + !bfloat16_support_->SupportsMixedPrecisions(*hlo)) { + // If hlo is the root instruction, we cannot change its output, so folding + // can only happen when it supports mixed precision so that we can change + // its operands. + return Status::OK(); + } + return TryFoldBF16Conversions(hlo); +} + +StatusOr BFloat16ConversionFolding::Run(HloModule* module) { + XLA_VLOG_LINES( + 2, "BFloat16ConversionFolding::Run(), before:\n" + module->ToString()); + bool changed = false; + for (auto* comp : module->MakeNonfusionComputations()) { + if (BFloat16ConversionFoldingVisitor::Run(comp, bfloat16_support_)) { + changed = true; + } + } + XLA_VLOG_LINES( + 2, "BFloat16ConversionFolding::Run(), after:\n" + module->ToString()); + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/bfloat16_conversion_folding.h b/tensorflow/compiler/xla/service/bfloat16_conversion_folding.h new file mode 100644 index 0000000000000000000000000000000000000000..c9398387098fad84ba28735c30e426fedd9b0cb0 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_conversion_folding.h @@ -0,0 +1,52 @@ +/* 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_COMPILER_XLA_SERVICE_BFLOAT16_CONVERSION_FOLDING_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_CONVERSION_FOLDING_H_ + +#include "tensorflow/compiler/xla/service/bfloat16_support.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// A pass which folds F32 <-> BF16 conversions to their operands or users, when +// it is supported by the backend. +// +// This pass follows the passed-in backend-specific BF16 support rules, but can +// introduce mixed precision in individual HLOs which breaks the assumption of +// some other HLO passes. So it should be used at the end of the HLO +// optimization pipeline followed by a DCE pass. If other passes are needed +// after this pass, run BFloat16MixedPrecisionRemoval first to undo some of the +// changed made by this pass. +class BFloat16ConversionFolding : public HloPassInterface { + public: + explicit BFloat16ConversionFolding(const BFloat16Support* bfloat16_support) + : bfloat16_support_(bfloat16_support) {} + + ~BFloat16ConversionFolding() override = default; + tensorflow::StringPiece name() const override { return "bfloat16-fold"; } + + // Run BF16 conversion folding on the given computation. Returns whether the + // computation was changed. + StatusOr Run(HloModule* module) override; + + private: + const BFloat16Support* bfloat16_support_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_CONVERSION_FOLDING_H_ diff --git a/tensorflow/compiler/xla/service/bfloat16_conversion_folding_test.cc b/tensorflow/compiler/xla/service/bfloat16_conversion_folding_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..cb37759439debf41a305ec7dccaa548e1bf234cd --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_conversion_folding_test.cc @@ -0,0 +1,209 @@ +/* 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/compiler/xla/service/bfloat16_conversion_folding.h" +#include "tensorflow/compiler/xla/service/bfloat16_support.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { + +class TestBFloat16Support : public BFloat16Support { + public: + TestBFloat16Support() {} + ~TestBFloat16Support() override {} + + bool SupportsBF16Operand(const HloInstruction& hlo, + int64 operand_index) const override { + if (hlo.opcode() == HloOpcode::kAdd || + hlo.opcode() == HloOpcode::kSubtract || + hlo.opcode() == HloOpcode::kTuple || + hlo.opcode() == HloOpcode::kGetTupleElement) { + return true; + } + return false; + } + + bool SupportsBF16Output(const HloInstruction& hlo) const override { + if (hlo.opcode() == HloOpcode::kAdd || + hlo.opcode() == HloOpcode::kSubtract || + hlo.opcode() == HloOpcode::kTuple || + hlo.opcode() == HloOpcode::kGetTupleElement) { + return true; + } + return false; + } + + bool SupportsMixedPrecisions(const HloInstruction& hlo) const override { + if (hlo.opcode() == HloOpcode::kAdd || hlo.opcode() == HloOpcode::kTuple || + hlo.opcode() == HloOpcode::kGetTupleElement) { + return true; + } + return false; + } +}; + +class BFloat16ConversionFoldingTest : public HloTestBase { + protected: + bool FoldConversions(HloModule* module) { + TestBFloat16Support bfloat16_support_; + BFloat16ConversionFolding fold(&bfloat16_support_); + StatusOr result = fold.Run(module); + EXPECT_IS_OK(result.status()); + return result.ValueOrDie(); + } +}; + +TEST_F(BFloat16ConversionFoldingTest, FoldIfSupported) { + auto builder = HloComputation::Builder(TestName()); + Shape f32_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, f32_shape, "b")); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32_shape, "c")); + + HloInstruction* add0 = builder.AddInstruction( + HloInstruction::CreateBinary(f32_shape, HloOpcode::kAdd, a, b)); + HloInstruction* convert0 = + builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, add0)); + HloInstruction* convert1 = builder.AddInstruction( + HloInstruction::CreateConvert(f32_shape, convert0)); + + HloInstruction* add1 = builder.AddInstruction( + HloInstruction::CreateBinary(f32_shape, HloOpcode::kAdd, convert1, c)); + builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, add1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(FoldConversions(module.get())); + + EXPECT_EQ(computation->root_instruction(), add1); + EXPECT_EQ(add0->shape().element_type(), BF16); + EXPECT_EQ(add1->shape().element_type(), BF16); + EXPECT_EQ(add1->operand(0), add0); +} + +TEST_F(BFloat16ConversionFoldingTest, DoNotFoldIfUnsupported) { + auto builder = HloComputation::Builder(TestName()); + Shape f32_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, f32_shape, "b")); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32_shape, "c")); + + HloInstruction* mul0 = builder.AddInstruction( + HloInstruction::CreateBinary(f32_shape, HloOpcode::kMultiply, a, b)); + HloInstruction* convert0 = + builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, mul0)); + HloInstruction* convert1 = builder.AddInstruction( + HloInstruction::CreateConvert(f32_shape, convert0)); + + HloInstruction* mul1 = builder.AddInstruction(HloInstruction::CreateBinary( + f32_shape, HloOpcode::kMultiply, convert1, c)); + HloInstruction* convert2 = + builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, mul1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_FALSE(FoldConversions(module.get())); + + EXPECT_EQ(computation->root_instruction(), convert2); + EXPECT_EQ(mul0->shape().element_type(), F32); + EXPECT_EQ(mul1->shape().element_type(), F32); + EXPECT_EQ(mul1->operand(0), convert1); +} + +TEST_F(BFloat16ConversionFoldingTest, DoNotFoldUnsupportedMixedPrecision) { + auto builder = HloComputation::Builder(TestName()); + Shape f32_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, f32_shape, "b")); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32_shape, "c")); + + HloInstruction* sub0 = builder.AddInstruction( + HloInstruction::CreateBinary(f32_shape, HloOpcode::kSubtract, a, b)); + HloInstruction* convert0 = + builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, sub0)); + HloInstruction* convert1 = builder.AddInstruction( + HloInstruction::CreateConvert(f32_shape, convert0)); + + HloInstruction* sub1 = builder.AddInstruction(HloInstruction::CreateBinary( + f32_shape, HloOpcode::kSubtract, convert1, c)); + HloInstruction* convert2 = + builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, sub1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_FALSE(FoldConversions(module.get())); + + EXPECT_EQ(computation->root_instruction(), convert2); + EXPECT_EQ(sub0->shape().element_type(), F32); + EXPECT_EQ(sub1->shape().element_type(), F32); + EXPECT_EQ(sub1->operand(0), convert1); +} + +TEST_F(BFloat16ConversionFoldingTest, DoNotFoldTuple) { + auto builder = HloComputation::Builder(TestName()); + Shape f32_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, bf16_shape, "b")); + HloInstruction* convert0 = + builder.AddInstruction(HloInstruction::CreateConvert(f32_shape, b)); + + HloInstruction* tuple = + builder.AddInstruction(HloInstruction::CreateTuple({a, convert0})); + HloInstruction* gte = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(f32_shape, tuple, 0)); + HloInstruction* convert1 = + builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, gte)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_FALSE(FoldConversions(module.get())); + + EXPECT_EQ(computation->root_instruction(), convert1); + EXPECT_EQ(gte->shape().element_type(), F32); + EXPECT_EQ(tuple->operand(1), convert0); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/bfloat16_normalization.cc b/tensorflow/compiler/xla/service/bfloat16_normalization.cc new file mode 100644 index 0000000000000000000000000000000000000000..b032c040e8aff49f9e0fc1ff9a1c1e79ea4bb77f --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_normalization.cc @@ -0,0 +1,351 @@ +/* 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/compiler/xla/service/bfloat16_normalization.h" + +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +class BFloat16NormalizationVisitor : public DfsHloVisitorWithDefault { + public: + explicit BFloat16NormalizationVisitor(HloComputation* computation, + const BFloat16Support* bfloat16_support) + : computation_(computation), bfloat16_support_(bfloat16_support) {} + + Status DefaultAction(HloInstruction* hlo) override; + + // Special handling for cross-replica-sum which can have a tuple output. + Status HandleCrossReplicaSum(HloInstruction* crs) override; + + static bool Run(HloComputation* computation, + const BFloat16Support* bfloat16_support) { + BFloat16NormalizationVisitor visitor(computation, bfloat16_support); + TF_CHECK_OK(computation->Accept(&visitor)); + return visitor.changed_; + } + + private: + // Checks if the HLO uses BF16 in an unsupported way, and if so, inserts + // conversions between F32 and BF16 to make it supported. + Status HandleInstruction(HloInstruction* hlo); + + // Inserts a conversion HLO that changes the given HLO's output type. + Status InsertConvertAfterOutput(HloInstruction* hlo, PrimitiveType to, + HloComputation* computation); + + // Changes the output type to the specified type, then inserts a conversion + // to the original type. + Status ChangeOutputTypeThenInsertConvertBack(HloInstruction* hlo, + PrimitiveType to, + HloComputation* computation); + + // Inserts a conversion HLO that changes the given HLO's operand type. + Status InsertConvertBeforeOperand(HloInstruction* hlo, int64 operand_idx, + PrimitiveType to, + HloComputation* computation); + + // Inserts conversion HLOs to replace the called computations' BF16 + // operands/outputs to F32. + Status ConvertCalledComputations( + HloInstruction* hlo, + tensorflow::gtl::ArraySlice bf16_called_comps); + + HloComputation* computation_; + const BFloat16Support* bfloat16_support_; + bool changed_ = false; +}; + +Status BFloat16NormalizationVisitor::InsertConvertAfterOutput( + HloInstruction* hlo, PrimitiveType to, HloComputation* computation) { + bool is_root = computation->root_instruction() == hlo; + std::vector materialized_users = hlo->users(); + // Use inst's shape temporarily, in order to pass checks in ReplaceUseWith. + auto convert = computation->AddInstruction( + HloInstruction::CreateConvert(hlo->shape(), hlo)); + for (auto* user : materialized_users) { + TF_RETURN_IF_ERROR(hlo->ReplaceUseWith(user, convert)); + } + if (is_root) { + computation->set_root_instruction(convert); + } + convert->mutable_shape()->set_element_type(to); + changed_ = true; + return Status::OK(); +} + +Status BFloat16NormalizationVisitor::ChangeOutputTypeThenInsertConvertBack( + HloInstruction* hlo, PrimitiveType to, HloComputation* computation) { + auto original_type = hlo->shape().element_type(); + hlo->mutable_shape()->set_element_type(to); + return InsertConvertAfterOutput(hlo, original_type, computation); +} + +Status BFloat16NormalizationVisitor::InsertConvertBeforeOperand( + HloInstruction* hlo, int64 operand_idx, PrimitiveType to, + HloComputation* computation) { + auto operand = hlo->mutable_operand(operand_idx); + auto convert = computation->AddInstruction(HloInstruction::CreateConvert( + ShapeUtil::ChangeElementType(operand->shape(), to), operand)); + TF_RETURN_IF_ERROR(hlo->ReplaceOperandWith(operand_idx, convert)); + changed_ = true; + return Status::OK(); +} + +Status BFloat16NormalizationVisitor::ConvertCalledComputations( + HloInstruction* hlo, + tensorflow::gtl::ArraySlice bf16_called_comps) { + std::map cloned_computations; + for (auto& comp : bf16_called_comps) { + auto cloned = comp->parent()->AddEmbeddedComputation(comp->Clone()); + cloned_computations[comp] = cloned; + changed_ = true; + } + hlo->ReplaceCalledComputations([&](HloComputation* comp) { + auto it = cloned_computations.find(comp); + if (it != cloned_computations.end()) { + return it->second; + } + return comp; + }); + for (auto& comp_pair : cloned_computations) { + auto comp = comp_pair.second; + if (comp->root_instruction()->shape().element_type() == BF16) { + TF_RETURN_IF_ERROR( + InsertConvertAfterOutput(comp->root_instruction(), F32, comp)); + } + for (auto* param : comp->parameter_instructions()) { + if (param->shape().element_type() == BF16) { + // This changes the parameter to F32 then inserts a convert after it. + TF_RETURN_IF_ERROR( + ChangeOutputTypeThenInsertConvertBack(param, F32, comp)); + } + } + } + return Status::OK(); +} + +Status BFloat16NormalizationVisitor::HandleCrossReplicaSum( + HloInstruction* crs) { + if (!ShapeUtil::IsTuple(crs->shape())) { + return HandleInstruction(crs); + } + + std::vector operand_types(crs->operand_count()); + std::vector output_types(crs->operand_count()); + bool has_f32 = false; + bool has_bf16 = false; + bool has_bf16_output = false; + for (int64 i = 0; i < crs->operand_count(); ++i) { + operand_types[i] = crs->operand(i)->shape().element_type(); + output_types[i] = ShapeUtil::GetSubshape(crs->shape(), {i}).element_type(); + if (operand_types[i] == F32 || output_types[i] == F32) { + has_f32 = true; + } else if (operand_types[i] == BF16) { + has_bf16 = true; + } + if (output_types[i] == BF16) { + has_bf16 = true; + has_bf16_output = true; + } + } + + for (int64 i = 0; i < crs->operand_count(); ++i) { + if (operand_types[i] != BF16) { + continue; + } + if (bfloat16_support_->SupportsBF16Operand(*crs, i) && + (bfloat16_support_->SupportsMixedPrecisions(*crs) || !has_f32)) { + continue; + } + TF_RETURN_IF_ERROR(InsertConvertBeforeOperand(crs, i, F32, computation_)); + has_f32 = true; + } + + if (!has_bf16_output) { + return Status::OK(); + } + + if (bfloat16_support_->SupportsBF16Output(*crs) && + (bfloat16_support_->SupportsMixedPrecisions(*crs) || !has_f32)) { + return Status::OK(); + } + + std::vector output_elements(crs->operand_count()); + auto original_shape = crs->shape(); + for (int64 i = 0; i < crs->operand_count(); ++i) { + auto subshape = ShapeUtil::GetMutableSubshape(crs->mutable_shape(), {i}); + if (output_types[i] != BF16) { + output_elements[i] = computation_->AddInstruction( + HloInstruction::CreateGetTupleElement(*subshape, crs, i)); + continue; + } + subshape->set_element_type(F32); + auto gte = computation_->AddInstruction( + HloInstruction::CreateGetTupleElement(*subshape, crs, i)); + output_elements[i] = + computation_->AddInstruction(HloInstruction::CreateConvert( + ShapeUtil::ChangeElementType(*subshape, BF16), gte)); + } + auto tuple = computation_->AddInstruction( + HloInstruction::CreateTuple(output_elements)); + + std::vector materialized_users = crs->users(); + // Use the crs' shape temporarily, in order to pass checks in + // ReplaceUseWith. + *tuple->mutable_shape() = crs->shape(); + for (auto* user : materialized_users) { + TF_RETURN_IF_ERROR(crs->ReplaceUseWith(user, tuple)); + } + *tuple->mutable_shape() = original_shape; + return Status::OK(); +} + +Status BFloat16NormalizationVisitor::HandleInstruction(HloInstruction* hlo) { + std::vector bf16_operands; + std::vector f32_operands; + bool has_f32 = false; + bool has_bf16 = false; + + for (int64 i = 0; i < hlo->operand_count(); ++i) { + if (hlo->operand(i)->shape().element_type() == F32) { + f32_operands.push_back(i); + has_f32 = true; + } else if (hlo->operand(i)->shape().element_type() == BF16) { + bf16_operands.push_back(i); + has_bf16 = true; + } + } + + if (hlo->shape().element_type() == F32) { + has_f32 = true; + } else if (hlo->shape().element_type() == BF16) { + has_bf16 = true; + } + + std::vector bf16_called_comps; + for (auto* comp : hlo->called_computations()) { + bool comp_has_bf16 = false; + if (comp->root_instruction()->shape().element_type() == F32) { + has_f32 = true; + } else if (comp->root_instruction()->shape().element_type() == BF16) { + has_bf16 = true; + comp_has_bf16 = true; + } + for (auto* param : comp->parameter_instructions()) { + if (param->shape().element_type() == F32) { + has_f32 = true; + } else if (param->shape().element_type() == BF16) { + has_bf16 = true; + comp_has_bf16 = true; + } + } + if (comp_has_bf16) { + bf16_called_comps.push_back(comp); + } + } + + if (!bfloat16_support_->SupportsMixedPrecisions(*hlo) && has_bf16 && + has_f32) { + // Resolve unsupported mixed precision. + // + // See if we can change everything to BF16. + if (hlo->called_computations().empty() && + hlo->shape().element_type() == BF16) { + bool can_use_bf16 = true; + for (int i : f32_operands) { + if (bfloat16_support_->EffectiveOperandPrecisionIsOutputPrecision(*hlo, + i) && + bfloat16_support_->SupportsBF16Operand(*hlo, i)) { + continue; + } + can_use_bf16 = false; + break; + } + if (can_use_bf16) { + for (int i : f32_operands) { + TF_RETURN_IF_ERROR( + InsertConvertBeforeOperand(hlo, i, BF16, computation_)); + } + return Status::OK(); + } + } + if (hlo->shape().element_type() == BF16) { + TF_RETURN_IF_ERROR( + ChangeOutputTypeThenInsertConvertBack(hlo, F32, computation_)); + } + for (int i : bf16_operands) { + TF_RETURN_IF_ERROR(InsertConvertBeforeOperand(hlo, i, F32, computation_)); + } + return ConvertCalledComputations(hlo, bf16_called_comps); + } + + for (int i : bf16_operands) { + if (!bfloat16_support_->SupportsBF16Operand(*hlo, i)) { + TF_RETURN_IF_ERROR(InsertConvertBeforeOperand(hlo, i, F32, computation_)); + } + } + + if (hlo->shape().element_type() == BF16 && + !bfloat16_support_->SupportsBF16Output(*hlo)) { + TF_RETURN_IF_ERROR( + ChangeOutputTypeThenInsertConvertBack(hlo, F32, computation_)); + } + + return Status::OK(); +} + +Status BFloat16NormalizationVisitor::DefaultAction(HloInstruction* hlo) { + // Do not change instructions related to entry and exit of a computation, + // tuples, fusion, convert, and control flow. + if (hlo->opcode() == HloOpcode::kTuple || // + hlo->opcode() == HloOpcode::kGetTupleElement || // + hlo->opcode() == HloOpcode::kInfeed || // + hlo->opcode() == HloOpcode::kOutfeed || // + hlo->opcode() == HloOpcode::kConstant || // + hlo->opcode() == HloOpcode::kParameter || // + hlo->opcode() == HloOpcode::kFusion || // + hlo->opcode() == HloOpcode::kConvert || // + hlo->opcode() == HloOpcode::kCall || // + hlo->opcode() == HloOpcode::kCustomCall || // + hlo->opcode() == HloOpcode::kWhile || // + hlo->opcode() == HloOpcode::kConditional) { + return Status::OK(); + } + return HandleInstruction(hlo); +} + +StatusOr BFloat16Normalization::Run(HloModule* module) { + XLA_VLOG_LINES( + 2, "BFloat16Normalization::Run(), before:\n" + module->ToString()); + bool changed = false; + for (auto* comp : module->MakeComputationPostOrder()) { + if (BFloat16NormalizationVisitor::Run(comp, bfloat16_support_)) { + changed = true; + } + } + XLA_VLOG_LINES(2, + "BFloat16Normalization::Run(), after:\n" + module->ToString()); + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/bfloat16_normalization.h b/tensorflow/compiler/xla/service/bfloat16_normalization.h new file mode 100644 index 0000000000000000000000000000000000000000..2a60fe0af3218484acb95e6c69815d551350764c --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_normalization.h @@ -0,0 +1,92 @@ +/* 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_COMPILER_XLA_SERVICE_BFLOAT16_NORMALIZATION_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_NORMALIZATION_H_ + +#include "tensorflow/compiler/xla/service/bfloat16_support.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// A pass which adds F32 <-> BF16 conversions for HLO instructions that do not +// support BF16 input/output or mixed precision, according to the passed-in +// backend-specific BF16 support rules. +class BFloat16Normalization : public HloPassInterface { + public: + explicit BFloat16Normalization(const BFloat16Support* bfloat16_support) + : bfloat16_support_(bfloat16_support) {} + + ~BFloat16Normalization() override = default; + tensorflow::StringPiece name() const override { return "bf16-normalization"; } + + // Run BF16 normalization on the given computation. Returns whether the + // computation was changed. + StatusOr Run(HloModule* module) override; + + private: + const BFloat16Support* bfloat16_support_; +}; + +// A pass that unconditionally removes the mixed F32/BF16 uses in HLO +// instructions (excluding convert) by adding F32 <-> BF16 conversions. Unlike +// BFloat16Normalization, this pass does not use a backend-specific +// BFloat16Support, and does not change HLOs that have BF16 data if they do not +// use mixed precision; it removes mixed precision even if the backend supports +// it. This pass is used to make the HLO module valid for other HLO passes which +// do not support mixed precision. +class BFloat16MixedPrecisionRemoval : public HloPassInterface { + public: + BFloat16MixedPrecisionRemoval() {} + + ~BFloat16MixedPrecisionRemoval() override = default; + + tensorflow::StringPiece name() const override { + return "bf16-mixed-precision-removal"; + } + + // Run mixed precision removal on the given computation. Returns whether the + // computation was changed. + StatusOr Run(HloModule* module) override { + BFloat16Normalization normalization(&no_mixed_precision_support_); + return normalization.Run(module); + } + + private: + class BFloat16SupportForMixedPrecisionRemoval : public BFloat16Support { + public: + BFloat16SupportForMixedPrecisionRemoval() {} + + ~BFloat16SupportForMixedPrecisionRemoval() override = default; + + bool SupportsBF16Operand(const HloInstruction& hlo, + int64 operand_index) const override { + return true; + } + + bool SupportsBF16Output(const HloInstruction& hlo) const override { + return true; + } + + bool SupportsMixedPrecisions(const HloInstruction& hlo) const override { + return false; + } + } no_mixed_precision_support_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_NORMALIZATION_H_ diff --git a/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..66c3085842c4afe7ffc4d5891883e4cce9389d45 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc @@ -0,0 +1,248 @@ +/* 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/compiler/xla/service/bfloat16_normalization.h" +#include "tensorflow/compiler/xla/service/bfloat16_support.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { + +class TestBFloat16Support : public BFloat16Support { + public: + TestBFloat16Support() {} + ~TestBFloat16Support() override {} + + bool SupportsBF16Operand(const HloInstruction& hlo, + int64 operand_index) const override { + if (hlo.opcode() == HloOpcode::kAdd || + hlo.opcode() == HloOpcode::kSubtract || + hlo.opcode() == HloOpcode::kReduce || + hlo.opcode() == HloOpcode::kTuple || + hlo.opcode() == HloOpcode::kGetTupleElement) { + return true; + } + return false; + } + + bool SupportsBF16Output(const HloInstruction& hlo) const override { + if (hlo.opcode() == HloOpcode::kAdd || hlo.opcode() == HloOpcode::kReduce || + hlo.opcode() == HloOpcode::kSubtract || + hlo.opcode() == HloOpcode::kTuple || + hlo.opcode() == HloOpcode::kGetTupleElement) { + return true; + } + return false; + } + + bool SupportsMixedPrecisions(const HloInstruction& hlo) const override { + if (hlo.opcode() == HloOpcode::kAdd || hlo.opcode() == HloOpcode::kTuple || + hlo.opcode() == HloOpcode::kGetTupleElement) { + return true; + } + return false; + } +}; + +class BFloat16NormalizationTest : public HloTestBase { + protected: + bool Normalize(HloModule* module) { + TestBFloat16Support bfloat16_support_; + BFloat16Normalization normalization(&bfloat16_support_); + StatusOr result = normalization.Run(module); + EXPECT_IS_OK(result.status()); + return result.ValueOrDie(); + } +}; + +TEST_F(BFloat16NormalizationTest, NoopIfSupported) { + auto builder = HloComputation::Builder(TestName()); + Shape f32_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, bf16_shape, "b")); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32_shape, "c")); + + HloInstruction* add0 = builder.AddInstruction( + HloInstruction::CreateBinary(bf16_shape, HloOpcode::kAdd, a, b)); + + HloInstruction* add1 = builder.AddInstruction( + HloInstruction::CreateBinary(f32_shape, HloOpcode::kAdd, add0, c)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_FALSE(Normalize(module.get())); + + EXPECT_EQ(computation->root_instruction(), add1); + EXPECT_EQ(add0->shape().element_type(), BF16); + EXPECT_EQ(add1->shape().element_type(), F32); +} + +TEST_F(BFloat16NormalizationTest, ResolveIfUnsupportedBF16) { + auto builder = HloComputation::Builder(TestName()); + Shape f32_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, bf16_shape, "b")); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32_shape, "c")); + + HloInstruction* mul0 = builder.AddInstruction( + HloInstruction::CreateBinary(bf16_shape, HloOpcode::kMultiply, a, b)); + + HloInstruction* mul1 = builder.AddInstruction( + HloInstruction::CreateBinary(bf16_shape, HloOpcode::kMultiply, mul0, c)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(Normalize(module.get())); + + EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kConvert); + EXPECT_EQ(computation->root_instruction()->operand(0), mul1); + EXPECT_EQ(mul0->shape().element_type(), F32); + EXPECT_EQ(mul1->shape().element_type(), F32); + EXPECT_EQ(mul1->operand(0)->opcode(), HloOpcode::kConvert); +} + +TEST_F(BFloat16NormalizationTest, ResolveUnsupportedMixedPrecisionSubtraction) { + auto builder = HloComputation::Builder(TestName()); + Shape f32_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, bf16_shape, "b")); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32_shape, "c")); + + HloInstruction* sub0 = builder.AddInstruction( + HloInstruction::CreateBinary(bf16_shape, HloOpcode::kSubtract, a, b)); + + HloInstruction* sub1 = builder.AddInstruction( + HloInstruction::CreateBinary(bf16_shape, HloOpcode::kSubtract, sub0, c)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(Normalize(module.get())); + + EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kConvert); + EXPECT_EQ(computation->root_instruction()->operand(0), sub1); + EXPECT_EQ(sub0->shape().element_type(), F32); + EXPECT_EQ(sub1->shape().element_type(), F32); + EXPECT_EQ(sub1->operand(0)->opcode(), HloOpcode::kConvert); +} + +TEST_F(BFloat16NormalizationTest, ResolveUnsupportedMixedPrecisionReduce) { + Shape f32_input_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape f32_output_shape = ShapeUtil::MakeShape(F32, {4}); + + Shape bf16_scalar_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + auto reduce_comp_builder = HloComputation::Builder("reduce_comp"); + auto reduce_comp_param0 = reduce_comp_builder.AddInstruction( + HloInstruction::CreateParameter(0, bf16_scalar_shape, "param0")); + auto reduce_comp_param1 = reduce_comp_builder.AddInstruction( + HloInstruction::CreateParameter(1, bf16_scalar_shape, "param1")); + reduce_comp_builder.AddInstruction( + HloInstruction::CreateBinary(bf16_scalar_shape, HloOpcode::kAdd, + reduce_comp_param0, reduce_comp_param1)); + + auto module = CreateNewModule(); + auto reduce_computation = + module->AddEmbeddedComputation(reduce_comp_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + HloInstruction* input = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_input_shape, "a")); + HloInstruction* init = builder.AddInstruction( + HloInstruction::CreateParameter(1, bf16_scalar_shape, "init")); + HloInstruction* reduce = builder.AddInstruction(HloInstruction::CreateReduce( + f32_output_shape, input, init, {0}, reduce_computation)); + + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(Normalize(module.get())); + + EXPECT_EQ(computation->root_instruction(), reduce); + EXPECT_EQ(reduce->called_computations().size(), 1); + EXPECT_EQ(reduce->called_computations()[0]->num_parameters(), 2); + EXPECT_EQ(reduce->called_computations()[0] + ->parameter_instruction(0) + ->shape() + .element_type(), + F32); + EXPECT_EQ(reduce->called_computations()[0] + ->parameter_instruction(1) + ->shape() + .element_type(), + F32); + EXPECT_EQ(reduce->called_computations()[0] + ->root_instruction() + ->shape() + .element_type(), + F32); + EXPECT_EQ(reduce->shape().element_type(), F32); + EXPECT_EQ(reduce->operand(0), input); + EXPECT_EQ(input->shape().element_type(), F32); + EXPECT_EQ(reduce->operand(1)->opcode(), HloOpcode::kConvert); + EXPECT_EQ(reduce->operand(1)->shape().element_type(), F32); +} + +TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleCrossReplicaSum) { + auto builder = HloComputation::Builder(TestName()); + Shape f32_shape = ShapeUtil::MakeShape(F32, {2, 4}); + Shape bf16_shape = ShapeUtil::MakeShape(BF16, {2, 4}); + + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32_shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateParameter(1, bf16_shape, "b")); + + HloInstruction* crs = + builder.AddInstruction(HloInstruction::CreateCrossReplicaSum( + ShapeUtil::MakeTupleShape({f32_shape, bf16_shape}), {a, b})); + HloInstruction* gte = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(bf16_shape, crs, 1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(Normalize(module.get())); + + EXPECT_EQ(computation->root_instruction(), gte); + EXPECT_EQ(gte->shape().element_type(), BF16); + EXPECT_EQ(crs->operand(1)->shape().element_type(), F32); + EXPECT_EQ(ShapeUtil::GetSubshape(crs->shape(), {1}).element_type(), F32); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/bfloat16_support.cc b/tensorflow/compiler/xla/service/bfloat16_support.cc new file mode 100644 index 0000000000000000000000000000000000000000..3fd9e24601f27633c8063e4574c7c4f91f30dcff --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_support.cc @@ -0,0 +1,111 @@ +/* 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/compiler/xla/service/bfloat16_support.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" + +namespace xla { + +bool BFloat16Support::SupportsBF16Operand(const HloInstruction& hlo, + int64 operand_index) const { + switch (hlo.opcode()) { + case HloOpcode::kCall: + case HloOpcode::kConditional: + case HloOpcode::kCustomCall: + case HloOpcode::kGetTupleElement: + case HloOpcode::kTuple: + case HloOpcode::kWhile: + return true; + case HloOpcode::kConvert: + CHECK_EQ(operand_index, 0); + return hlo.operand(0)->shape().element_type() == BF16; + default: + break; + } + return false; +} + +bool BFloat16Support::SupportsBF16Output(const HloInstruction& hlo) const { + switch (hlo.opcode()) { + case HloOpcode::kCall: + case HloOpcode::kConditional: + case HloOpcode::kCustomCall: + case HloOpcode::kGetTupleElement: + case HloOpcode::kTuple: + case HloOpcode::kWhile: + return true; + case HloOpcode::kConvert: + return hlo.shape().element_type() == BF16; + default: + break; + } + return false; +} + +bool BFloat16Support::SupportsMixedPrecisions(const HloInstruction& hlo) const { + switch (hlo.opcode()) { + case HloOpcode::kCall: + case HloOpcode::kConditional: + case HloOpcode::kConvert: + case HloOpcode::kCustomCall: + case HloOpcode::kGetTupleElement: + case HloOpcode::kTuple: + case HloOpcode::kWhile: + return true; + default: + break; + } + return false; +} + +/* static */ +bool BFloat16Support::EffectiveOperandPrecisionIsOutputPrecision( + const HloInstruction& hlo, int64 operand_index) { + switch (hlo.opcode()) { + case HloOpcode::kAbs: + case HloOpcode::kBroadcast: + case HloOpcode::kClamp: + case HloOpcode::kConcatenate: + case HloOpcode::kCopy: + case HloOpcode::kGetTupleElement: + case HloOpcode::kMaximum: + case HloOpcode::kMinimum: + case HloOpcode::kPad: + case HloOpcode::kReshape: + case HloOpcode::kReverse: + case HloOpcode::kSlice: + case HloOpcode::kSort: + case HloOpcode::kTranspose: + case HloOpcode::kTuple: + return true; + case HloOpcode::kDynamicSlice: + return operand_index == 0; + case HloOpcode::kDynamicUpdateSlice: + return operand_index == 0 || operand_index == 1; + case HloOpcode::kSelect: + return operand_index == 1 || operand_index == 2; + default: + break; + } + return false; +} + +bool BFloat16Support::EffectiveOperandPrecisionIsBF16( + const HloInstruction& hlo, int64 operand_index) const { + return false; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/bfloat16_support.h b/tensorflow/compiler/xla/service/bfloat16_support.h new file mode 100644 index 0000000000000000000000000000000000000000..29f662d22b4e5486662a1387407d41e0fd2ed1b3 --- /dev/null +++ b/tensorflow/compiler/xla/service/bfloat16_support.h @@ -0,0 +1,60 @@ +/* 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_COMPILER_XLA_SERVICE_BFLOAT16_SUPPORT_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_SUPPORT_H_ + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" + +namespace xla { + +class BFloat16Support { + public: + BFloat16Support() {} + virtual ~BFloat16Support() {} + + // Returns whether the backend supports BF16 operand for the HLO instruction + // at the given index. + virtual bool SupportsBF16Operand(const HloInstruction& hlo, + int64 operand_index) const; + + // Returns whether the backend supports BF16 output for the HLO instruction. + virtual bool SupportsBF16Output(const HloInstruction& hlo) const; + + // Returns whether the backend support mixed precision: the operands, output, + // and parameters/output of the called computations can have different + // precisions (BF16 and F32). + virtual bool SupportsMixedPrecisions(const HloInstruction& hlo) const; + + // Returns whether the given HLO inherits its BF16 operand precision at the + // given index, so even if the output is F32, elements in the output that + // depend on the BF16 operand will still have BF16 effective precision even if + // they have F32 format. Similarly, this also means if the output is BF16 then + // increasing the operand precision from BF16 to F32 will not change the + // output. This typically includes HLOs that pass elements from the operand to + // the output without arithmetic operations. + static bool EffectiveOperandPrecisionIsOutputPrecision( + const HloInstruction& hlo, int64 operand_index); + + // Returns if the backend only uses BF16 precision for the operand at the + // specified index, even if the operand is F32. + virtual bool EffectiveOperandPrecisionIsBF16(const HloInstruction& hlo, + int64 operand_index) const; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_SUPPORT_H_ diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc index 7ece79d781acfaffc21d6a29e8a12e68622a1617..b1e693da9d5af4babe619b8796007f2da318f6a8 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment.cc @@ -45,6 +45,8 @@ using ::tensorflow::gtl::FlatMap; using ::tensorflow::gtl::FlatSet; using ::tensorflow::strings::Appendf; using ::tensorflow::strings::HumanReadableNumBytes; +using ::tensorflow::strings::Printf; +using ::tensorflow::strings::StrAppend; size_t BufferAllocation::Slice::Hasher::operator()(Slice s) const { uint64 h = std::hash()(s.index()); @@ -73,9 +75,10 @@ void BufferAllocation::AddAssignment(const LogicalBuffer& buffer, int64 offset, CHECK_LE(offset, size_) << "LogicalBuffer " << buffer << " offset out of range"; CHECK_LE(offset + size, size_) - << "LogicalBuffer " << buffer << " size out of range"; + << "LogicalBuffer " << buffer + << " size out of range at offset: " << offset << " with size: " << size; CHECK_EQ(buffer.color(), color()) - << "Buffer color " << buffer.color() + << "Buffer color " << buffer.color() << " for buffer " << buffer << " does not match allocation color " << color() << "."; OffsetSize offset_size; offset_size.offset = offset; @@ -92,6 +95,9 @@ BufferAllocationProto BufferAllocation::ToProto() const { proto.set_color(color_.value()); if (is_entry_computation_parameter_) { proto.set_is_entry_computation_parameter(true); + for (int64 idx : param_shape_index()) { + proto.add_parameter_shape_index(idx); + } proto.set_parameter_number(parameter_number_); } proto.set_maybe_live_out(maybe_live_out_); @@ -111,25 +117,24 @@ BufferAllocationProto BufferAllocation::ToProto() const { string BufferAllocation::ToString() const { string output; - tensorflow::strings::StrAppend( - &output, tensorflow::strings::Printf("allocation %lld: %p, size %lld", - index_, this, size())); + Appendf(&output, "allocation %lld: %p, size %lld", index_, this, size()); if (color().value() != 0) { - tensorflow::strings::StrAppend(&output, ", color ", color().value()); + StrAppend(&output, ", color ", color().value()); } if (is_entry_computation_parameter()) { - tensorflow::strings::StrAppend(&output, ", parameter ", parameter_number()); + StrAppend(&output, ", parameter ", parameter_number(), " at ShapeIndex ", + param_shape_index().ToString()); } if (is_thread_local()) { - tensorflow::strings::StrAppend(&output, ", thread-local"); + StrAppend(&output, ", thread-local"); } if (maybe_live_out()) { - tensorflow::strings::StrAppend(&output, ", maybe-live-out"); + StrAppend(&output, ", maybe-live-out"); } if (IsPreallocatedTempBuffer()) { - tensorflow::strings::StrAppend(&output, ", preallocated-temp"); + StrAppend(&output, ", preallocated-temp"); } - tensorflow::strings::StrAppend(&output, ":\n"); + StrAppend(&output, ":\n"); // Dump the assigned buffers ordered by id. std::vector sorted_buffers; for (const auto& buffer_offset_size : assigned_buffers_) { @@ -141,12 +146,11 @@ string BufferAllocation::ToString() const { }); for (const LogicalBuffer* buffer : sorted_buffers) { const OffsetSize& offset_size = FindOrDie(assigned_buffers_, buffer); - tensorflow::strings::StrAppend( - &output, - tensorflow::strings::Printf( - " %s [%lld,%lld]: %s\n", buffer->ToString().c_str(), - offset_size.offset, offset_size.size, - ShapeUtil::HumanStringWithLayout(buffer->shape()).c_str())); + StrAppend(&output, + tensorflow::strings::Printf( + " %s [%lld,%lld]: %s\n", buffer->ToString().c_str(), + offset_size.offset, offset_size.size, + ShapeUtil::HumanStringWithLayout(buffer->shape()).c_str())); } return output; } @@ -839,20 +843,19 @@ Status BufferAssigner::AssignBuffersForComputation( /*is_thread_local=*/false, /*is_reusable=*/false); allocation->set_entry_computation_parameter( - instruction->parameter_number()); + instruction->parameter_number(), buffer->index()); VLOG(3) << "New allocation #" << allocation->index() << " for entry computation parameter: " << *buffer; continue; } - if (is_thread_local || instruction->opcode() == HloOpcode::kCustomCall) { - // Custom call operations never have reusable buffers. Also we do not - // reuse thread-local buffers for now, because they are dynamically - // allocated and their lifetimes are hard to compute. + if (is_thread_local) { + // We do not reuse thread-local buffers for now, because they are + // dynamically allocated and their lifetimes are hard to compute. BufferAllocation* allocation = assignment->NewAllocation( *buffer, buffer_size, is_thread_local, /*is_reusable=*/false); VLOG(3) << "New allocation #" << allocation->index() - << " for thread-local/CustomCall: " << *buffer; + << " for thread-local: " << *buffer; continue; } @@ -997,14 +1000,15 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering( auto color = single_colored_set.first; VLOG(2) << "Simulating heap for color " << color; int64 alignment = assignment->color_alignment_(color); + HeapSimulator::Options options; + options.buffers_to_assign = &single_colored_set.second; TF_ASSIGN_OR_RETURN( const HeapSimulator::Result result, HeapSimulator::Run(MakeUnique( MakeUnique(alignment)), assignment->module(), module_sequence, assignment->points_to_analysis(), - assignment->buffer_size_, - &single_colored_set.second)); + assignment->buffer_size_, options)); AssignBuffersFromHeapSimulator(result, assignment, single_colored_set.first); } @@ -1024,14 +1028,15 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering( auto color = single_colored_set.first; VLOG(2) << "Simulating heap for color " << color; int64 alignment = assignment->color_alignment_(color); + HeapSimulator::Options options; + options.buffers_to_assign = &single_colored_set.second; TF_ASSIGN_OR_RETURN( const HeapSimulator::Result result, HeapSimulator::Run(MakeUnique( MakeUnique(alignment)), *computation, *instruction_sequence, assignment->points_to_analysis(), - assignment->buffer_size_, - &single_colored_set.second)); + assignment->buffer_size_, options)); AssignBuffersFromHeapSimulator(result, assignment, single_colored_set.first); } @@ -1120,140 +1125,6 @@ void BufferAssigner::AddSetToColocatedBufferSets( } } -// Conceptually the same as AddSetToColocatedBufferSets, but specific to the -// colocated buffers for while instructions. 'colocated_set' contains the -// buffers for a single while instruction that must be colocated. The idea here -// is to apply a memory-saving heuristic for separate while instructions whose -// buffers are disjoint in liveness, by using the colocation mechanism to force -// buffer sharing. This often reduces memory for multi-layer RNNs. -// -// TODO(b/32491382): We should be able to remove this heuristic after we -// implement module-level liveness analysis, which would let us directly detect -// buffer sharing opportunities between the while instruction buffer and the -// buffers from the predicate and body computation, as well as sharing across -// different while instructions. -void BufferAssigner::AddWhileSetToColocatedBufferSets( - const std::vector& colocated_set, - const LogicalBuffer* while_init_buffer, - const LogicalBuffer* while_result_buffer, const HloInstruction* while_hlo, - const HloComputation& computation, const BufferLiveness& buffer_liveness, - const LogicalBuffer::SizeFunction& buffer_size, - std::vector* colocated_buffer_sets) { - CHECK(!colocated_set.empty()); - const TuplePointsToAnalysis& points_to_analysis = - buffer_liveness.points_to_analysis(); - - // Parallel while loops cannot safely share colocated buffer sets. - if (buffer_liveness.hlo_ordering().SequentialOrder(computation) == nullptr) { - AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); - return; - } - - // Scan 'colocated_buffer_sets' in reverse order for locality; colocated sets - // are added in postorder over computations and instructions. - const int64 init_buffer_size = buffer_size(*while_init_buffer); - const bool is_live_out = buffer_liveness.MaybeLiveOut(*while_result_buffer); - for (int i = colocated_buffer_sets->size() - 1; i >= 0; --i) { - const ColocatedBufferSet& predecessor_set = (*colocated_buffer_sets)[i]; - - // Skip predecessor sets not associated with while loops. - if (std::all_of(predecessor_set.begin(), predecessor_set.end(), - [](const LogicalBuffer* buffer) { - return buffer->instruction()->opcode() != - HloOpcode::kWhile; - })) { - continue; - } - - // Skip predecessor sets already associated with 'while_hlo'. - if (std::any_of(predecessor_set.begin(), predecessor_set.end(), - [&while_hlo](const LogicalBuffer* buffer) { - return buffer->instruction() == while_hlo; - })) { - continue; - } - - // Skip predecessor sets with entry parameter if the while result is live - // out. - if (is_live_out && - std::any_of(predecessor_set.begin(), predecessor_set.end(), - [](const LogicalBuffer* buffer) { - auto* instruction = buffer->instruction(); - auto* computation = instruction->parent(); - auto* module = computation->parent(); - return instruction->opcode() == HloOpcode::kParameter && - computation == module->entry_computation(); - })) { - continue; - } - - // Build vector of predecessor while result and init buffers, which are - // checked for liveness interference below. We must check both the result - // and init buffers because they're aliased together, but - // TuplePointsToAnalysis is unaware of this aliasing. - std::vector predecessor_while_buffers; - for (const LogicalBuffer* buffer : predecessor_set) { - const HloInstruction* instruction = buffer->instruction(); - if (instruction->opcode() == HloOpcode::kWhile && - buffer_size(*buffer) == init_buffer_size && - instruction->parent() == &computation) { - predecessor_while_buffers.push_back(buffer); - // Add the init buffer at the same index, which must also exist in the - // predecessor set, and must be unambiguous. - const PointsToSet& init_points_to = - points_to_analysis.GetPointsToSet(instruction->operand(0)); - const auto& init_buffers = init_points_to.element(buffer->index()); - CHECK_EQ(init_buffers.size(), 1); - CHECK_GT(predecessor_set.count(init_buffers[0]), 0); - predecessor_while_buffers.push_back(init_buffers[0]); - } - } - if (predecessor_while_buffers.empty()) { - continue; - } - - // Skip predecessor set if the live range of any predecessor - // buffers overlaps with 'while_init_buffer' or - // 'while_result_buffer' (we need to check both since they're - // aliased together, but the points-to analysis is unaware of this - // aliasing). Note that tuple element buffer forwarding can cause - // the same buffer to appear on both sides of the interference - // comparison below. - auto may_interfere_with_init_or_result = [&](const LogicalBuffer* buffer) { - if (while_init_buffer->id() != buffer->id() && - buffer_liveness.MayInterfere(*while_init_buffer, *buffer)) { - return true; - } - - if (while_result_buffer->id() != buffer->id() && - buffer_liveness.MayInterfere(*while_result_buffer, *buffer)) { - return true; - } - - return false; - }; - - if (std::any_of(predecessor_while_buffers.begin(), - predecessor_while_buffers.end(), - may_interfere_with_init_or_result)) { - continue; - } - - // All our checks have passed; merge 'predecessor_set' with 'colocated_set', - // and add the merged set to 'colocated_buffer_sets'. This forces the - // colocation of buffers across different while instructions. - FlatSet unique; - unique.insert(predecessor_set.begin(), predecessor_set.end()); - unique.insert(colocated_set.begin(), colocated_set.end()); - std::vector merged_set(unique.begin(), unique.end()); - AddSetToColocatedBufferSets(merged_set, colocated_buffer_sets); - return; - } - - // Failed to merge into predecessor set; add 'colocated_set' as-is. - AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); -} - namespace { // Checks that points-to set of 'instruction' is unambiguous and distinct @@ -1270,8 +1141,130 @@ const LogicalBuffer* AddBufferToColocatedSet( return colocated_set->back(); } +// Given the interference map of a graph (the list of interfering node indices +// for each node), perform graph coloring such that interfering nodes are +// assigned to different colors. Returns the assigned color of the nodes, where +// the colors are represented as integer values [0, color_count). +std::vector ColorInterferenceGraph( + const std::vector>& interference_map) { + const int64 node_count = interference_map.size(); + + // Sort the nodes such that we assign nodes with more interference first. This + // relies on the common heuristic of assigning the most constrained node + // first, but it would be good to investigate other ordering heuristics too. + std::vector nodes(node_count); + std::iota(nodes.begin(), nodes.end(), 0); + std::sort(nodes.begin(), nodes.end(), + [&interference_map](const int64 i, const int64 j) { + return interference_map[i].size() > interference_map[j].size(); + }); + + const int64 kColorUnassigned = -1; + std::vector assigned_colors(node_count, kColorUnassigned); + for (int64 node : nodes) { + // Mark the colors that are already assigned to the neighbors. + std::vector available_colors(node_count, true); + for (int64 neighbor : interference_map[node]) { + int64 color = assigned_colors[neighbor]; + if (color != kColorUnassigned) { + available_colors[color] = false; + } + } + + // Find the color that is not yet assigned to the neighbors. + int64 color = kColorUnassigned; + for (color = 0; color < available_colors.size(); ++color) { + if (available_colors[color]) { + break; + } + } + CHECK_NE(color, kColorUnassigned); + assigned_colors[node] = color; + } + return assigned_colors; +} + } // namespace +std::vector +BufferAssigner::MergeColocatedBufferSets( + const std::vector& colocated_buffer_sets, + const BufferLiveness& buffer_liveness, + const LogicalBuffer::SizeFunction& buffer_size) { + VLOG(1) << "colocation sets count before coalescing:" + << colocated_buffer_sets.size(); + + // Returns true if the given buffer is for the entry parameter. + auto is_entry_parameter = [](const LogicalBuffer& buffer) { + auto* instruction = buffer.instruction(); + auto* computation = instruction->parent(); + auto* module = computation->parent(); + return instruction->opcode() == HloOpcode::kParameter && + computation == module->entry_computation(); + }; + + // Returns true if the two colocated buffer sets (specified by their indices + // into the colocated_buffer_sets) can be merged into a single set. + auto cannot_merge_buffer_sets = [&colocated_buffer_sets, &buffer_liveness, + &buffer_size, + &is_entry_parameter](int64 i, int64 j) { + for (auto& buffer_a : colocated_buffer_sets[i]) { + for (auto& buffer_b : colocated_buffer_sets[j]) { + // Do not merge if the set includes live outs or entry parameters. + if ((buffer_liveness.MaybeLiveOut(*buffer_a) && + is_entry_parameter(*buffer_b)) || + (buffer_liveness.MaybeLiveOut(*buffer_b) && + is_entry_parameter(*buffer_a))) { + return true; + } + // Do not merge if the buffers interfere with each other. + if (buffer_a->id() != buffer_b->id() && + buffer_liveness.MayInterfere(*buffer_a, *buffer_b)) { + return true; + } + // Do not merge if the buffer sizes are different. + if (buffer_size(*buffer_a) != buffer_size(*buffer_b)) { + return true; + } + } + } + return false; + }; + + // Build the interference map among the colocated buffer sets (nodes), by + // adding an edge between any two nodes that cannot be merged into a single + // colocated buffer set. + std::vector> interference_map( + colocated_buffer_sets.size()); + for (int64 i = 0; i < colocated_buffer_sets.size(); ++i) { + for (int64 j = i + 1; j < colocated_buffer_sets.size(); ++j) { + if (cannot_merge_buffer_sets(i, j)) { + interference_map[i].push_back(j); + interference_map[j].push_back(i); + } + } + } + + // Assign a color to each colocation set in colocated_buffer_sets, such that + // the sets that can be merged are assigned with the same color. + auto assigned_colors = ColorInterferenceGraph(interference_map); + + // Merge the buffer sets with the same color. + CHECK(!assigned_colors.empty()); + int64 num_sets = + *std::max_element(assigned_colors.begin(), assigned_colors.end()) + 1; + std::vector new_colocated_buffer_sets(num_sets); + for (int64 i = 0; i < colocated_buffer_sets.size(); ++i) { + const auto& buffer_set = colocated_buffer_sets[i]; + new_colocated_buffer_sets[assigned_colors[i]].insert(buffer_set.begin(), + buffer_set.end()); + } + + VLOG(1) << "colocation sets count after coalescing:" + << colocated_buffer_sets.size(); + return new_colocated_buffer_sets; +} + // Builds sets of buffers in 'colocated_buffer_sets' which should be colocated // in the same allocation (currently just supports kWhile, kCall, and // kConditional). @@ -1297,12 +1290,11 @@ void BufferAssigner::BuildColocatedBufferSets( const Shape& /*subshape*/, const ShapeIndex& index) { std::vector colocated_set; // Add while.init. - auto* init_buffer = - AddBufferToColocatedSet(while_hlo->operand(0), index, - points_to_analysis, &colocated_set); + AddBufferToColocatedSet(while_hlo->operand(0), index, + points_to_analysis, &colocated_set); // Add while.result. - auto* result_buffer = AddBufferToColocatedSet( - while_hlo, index, points_to_analysis, &colocated_set); + AddBufferToColocatedSet(while_hlo, index, points_to_analysis, + &colocated_set); // Add while.cond.parameter. AddBufferToColocatedSet( while_hlo->while_condition()->parameter_instruction(0), index, @@ -1315,10 +1307,7 @@ void BufferAssigner::BuildColocatedBufferSets( AddBufferToColocatedSet( while_hlo->while_body()->root_instruction(), index, points_to_analysis, &colocated_set); - AddWhileSetToColocatedBufferSets( - colocated_set, init_buffer, result_buffer, while_hlo, - *computation, buffer_liveness, buffer_size, - colocated_buffer_sets); + AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); }); } else if (opcode == HloOpcode::kCall) { const HloInstruction* call_hlo = instruction; @@ -1358,9 +1347,62 @@ void BufferAssigner::BuildColocatedBufferSets( index, points_to_analysis, &colocated_set); AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); }); + + // Add true_operand and conditional.true_computation.parameter(0) as a + // colocated buffer set. Note that this has to be done for each subshape + // in the true_operand of the conditional. + ShapeUtil::ForEachSubshape( + conditional_hlo->operand(1)->shape(), + [this, conditional_hlo, &points_to_analysis, colocated_buffer_sets]( + const Shape& /*subshape*/, const ShapeIndex& index) { + std::vector true_set; + // Add conditional.true_operand. + AddBufferToColocatedSet(conditional_hlo->operand(1), index, + points_to_analysis, &true_set); + // Add conditional.true_computation.parameter_instruction(0). + AddBufferToColocatedSet( + conditional_hlo->true_computation()->parameter_instruction(0), + index, points_to_analysis, &true_set); + AddSetToColocatedBufferSets(true_set, colocated_buffer_sets); + }); + + // Add false_operand and conditional.false_computation.parameter(0) as a + // colocated buffer set. Note that this has to be done for each subshape + // in the false_operand of the conditional. + ShapeUtil::ForEachSubshape( + conditional_hlo->operand(2)->shape(), + [this, conditional_hlo, &points_to_analysis, colocated_buffer_sets]( + const Shape& /*subshape*/, const ShapeIndex& index) { + std::vector false_set; + // Add conditional.false_operand. + AddBufferToColocatedSet(conditional_hlo->operand(2), index, + points_to_analysis, &false_set); + // Add conditional.false_computation.parameter_instruction(0). + AddBufferToColocatedSet( + conditional_hlo->false_computation()->parameter_instruction( + 0), + index, points_to_analysis, &false_set); + AddSetToColocatedBufferSets(false_set, colocated_buffer_sets); + }); } } } + + if (colocated_buffer_sets->empty()) { + return; + } + + // Try to find more coalescing opportunities among the colocated buffer sets. + // + // TODO(b/32491382): We should be able to remove this by using the + // module-level liveness analysis, which would let us directly detect buffer + // sharing opportunities between the while instruction buffer and the buffers + // from the predicate and body computation, as well as sharing across + // different while instructions. + std::vector new_colocated_buffer_sets = + MergeColocatedBufferSets(*colocated_buffer_sets, buffer_liveness, + buffer_size); + std::swap(*colocated_buffer_sets, new_colocated_buffer_sets); } // Assigns all colocated buffer sets in 'colocated_buffer_sets' to the same @@ -1372,39 +1414,47 @@ void BufferAssigner::AssignColocatedBufferSets( FlatSet* colocated_allocations) { for (const ColocatedBufferSet& colocated_buffer_set : colocated_buffer_sets) { BufferAllocation* allocation = nullptr; - // Set 'entry_parameter_number' if entry param in 'colocated_buffer_set'. + // Set 'entry_parameter_number' and 'entry_parameter_shape_idx' if entry + // param in 'colocated_buffer_set'. int64 entry_parameter_number = -1; + const ShapeIndex* entry_parameter_shape_idx = nullptr; for (const LogicalBuffer* buffer : colocated_buffer_set) { const HloInstruction* instruction = buffer->instruction(); const HloComputation* computation = instruction->parent(); if (instruction->opcode() == HloOpcode::kParameter && computation == computation->parent()->entry_computation()) { entry_parameter_number = instruction->parameter_number(); + entry_parameter_shape_idx = &buffer->index(); break; } } for (const LogicalBuffer* buffer : colocated_buffer_set) { + const int64 buffer_size = assignment->buffer_size_(*buffer); if (allocation == nullptr) { // TODO(b/32491382) Avoid current trivial solution of using new // allocations for each colocated buffer set. When liveness has // module-level scope, we can allow buffers to be shared across // computations (in some cases). - allocation = assignment->NewAllocation( - *buffer, assignment->buffer_size_(*buffer), - /*is_thread_local=*/false, /*is_reusable=*/true); + allocation = assignment->NewAllocation(*buffer, buffer_size, + /*is_thread_local=*/false, + /*is_reusable=*/true); if (entry_parameter_number >= 0) { // This colocated buffer set contains an entry parameter and other // logical buffers which use the parameter as read-only in a while // body computation (which updates in place). // Set 'entry_computation_parameter' to indicate that it contains // an entry parameter, and to prevent reuse in MaybeAssignBuffer. - allocation->set_entry_computation_parameter(entry_parameter_number); + allocation->set_entry_computation_parameter( + entry_parameter_number, *entry_parameter_shape_idx); } colocated_allocations->insert(allocation->index()); } else { + CHECK_EQ(buffer_size, allocation->size()) + << "Buffer: " << *buffer << " size mismatch in colocated buffer " + << "allocation: " << *allocation; assignment->AddAssignment(allocation, *buffer, /*offset=*/0, - assignment->buffer_size_(*buffer)); + buffer_size); } colocated_buffers->insert(buffer); } diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h index 08a40bfeb2a2a78c25805308e73154c6cc667f21..6b7fd0014d103ef0617afcc5cb3f663554a01aa4 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.h +++ b/tensorflow/compiler/xla/service/buffer_assignment.h @@ -91,6 +91,13 @@ class BufferAllocation { return parameter_number_; } + // If this allocation is for a parameter of the entry computation, this + // function returns which subshape of the parameter the allocation is for. + const ShapeIndex& param_shape_index() const { + CHECK(is_entry_computation_parameter_); + return param_shape_index_; + } + // Returns whether this allocation is assigned a LogicalBuffer which may // be live out of the entry computation. bool maybe_live_out() const { return maybe_live_out_; } @@ -203,9 +210,11 @@ class BufferAllocation { // Adds a LogicalBuffer to the set assigned to this buffer. void AddAssignment(const LogicalBuffer& buffer, int64 offset, int64 size); - void set_entry_computation_parameter(int64 parameter_number) { + void set_entry_computation_parameter(int64 parameter_number, + ShapeIndex param_shape_index) { is_entry_computation_parameter_ = true; parameter_number_ = parameter_number; + param_shape_index_ = std::move(param_shape_index); } void set_maybe_live_out(bool value) { maybe_live_out_ = value; } void set_index(Index index) { index_ = index; } @@ -235,6 +244,10 @@ class BufferAllocation { // indicates the index (starting from 0) of the parameter. int64 parameter_number_ = 0; + // If this buffer is for an entry computation parameter, which subshape of the + // parameter is it for? + ShapeIndex param_shape_index_; + // Whether the allocation contains a LogicalBuffer which may be live-out of // the entry computation. Note that this flag is conservatively computed by // TuplePointsToAnalysis. That is, an allocation marked `maybe_live_out_` @@ -528,15 +541,13 @@ class BufferAssigner { const std::vector& colocated_set, std::vector* colocated_buffer_sets); - // Conceptually the same as AddSetToColocatedBufferSets, but specific to the - // colocated buffers for while instructions. - void AddWhileSetToColocatedBufferSets( - const std::vector& colocated_set, - const LogicalBuffer* while_init_buffer, - const LogicalBuffer* while_result_buffer, const HloInstruction* while_hlo, - const HloComputation& computation, const BufferLiveness& buffer_liveness, - const LogicalBuffer::SizeFunction& buffer_size, - std::vector* colocated_buffer_sets); + // Given a list of colocated buffer sets (each colocated buffer set represents + // the logical buffers that would be assigned to the same physical buffer), + // try to merge the sets if the buffers can be shared. Returns the merged set. + std::vector MergeColocatedBufferSets( + const std::vector& colocated_buffer_sets, + const BufferLiveness& buffer_liveness, + const LogicalBuffer::SizeFunction& buffer_size); // Split a set of buffers into several sets, each of which contains buffers // colored with the same color. diff --git a/tensorflow/compiler/xla/service/buffer_assignment_test.cc b/tensorflow/compiler/xla/service/buffer_assignment_test.cc index 6fc9d783f1b34de8c0f93c6aa342591891d08eaf..cd73654b8f666c4b96c000235cc3ad2cd0a46c17 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment_test.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment_test.cc @@ -614,7 +614,7 @@ TEST_F(BufferAssignmentTest, TrivialMap) { BufferAllocation map_buffer = GetAssignedOutputAllocation(*buffers, map); EXPECT_NE(param0_buffer.index(), map_buffer.index()); - // The final computation node of the map is an add of an f32 parm and a + // The final computation node of the map is an add of an f32 param and a // constant. EXPECT_EQ(HloOpcode::kAdd, inner_last->opcode()); const BufferAllocation& inner_add_buffer = @@ -1587,6 +1587,117 @@ TEST_F(WhileBufferAssignmentTest, TwoForwardWhileLoops) { assignment->GetUniqueSlice(while1, {1}).ConsumeValueOrDie()); } +// Tests that the colocated buffers for while instructions are properly assigned +// during buffer assignment such that the result tuple elements are not assigned +// to the same buffer. +// +// %infeed --> %while.0 --> %while.1 --+ +// +-- %tuple +// %zero --> %add --> %while.2 --+ +// +// Execution Order: +// %infeed -> %while.0 -> %while.1 -> %zero -> %add -> %while.2 -> %tuple +// +// The HLO computation used in this test requires specific ordering to expose +// the bug (b/72496031). During buffer assignment, the visitation order of +// colocated buffers is %while.2 -> while.0 -> while.1, and the buffer +// assignment was coalescing the colocated buffers for all 3 while instructions, +// therefore assigning the same buffer to the two result tuple elements. +TEST_F(WhileBufferAssignmentTest, ColocatedBuffers) { + const Shape r0s32 = ShapeUtil::MakeShape(S32, {}); + + // Builds a condition computation: x -> x < 4 + auto build_cond = [&]() { + auto builder = HloComputation::Builder("cond"); + auto const4 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(4))); + auto param = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0s32, "x")); + builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kLt, param, const4)); + return builder.Build(); + }; + + // Builds a body computation: x -> x + 9 + auto build_body = [&]() { + auto builder = HloComputation::Builder("body"); + auto const9 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(9))); + auto param = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0s32, "x")); + builder.AddInstruction( + HloInstruction::CreateBinary(r0s32, HloOpcode::kAdd, param, const9)); + return builder.Build(); + }; + + // Build the entry computation as described in the comment above. + auto module = xla::MakeUnique(TestName()); + auto builder = HloComputation::Builder("entry"); + + auto infeed = builder.AddInstruction(HloInstruction::CreateInfeed(r0s32, "")); + auto cond0 = module->AddEmbeddedComputation(build_cond()); + auto body0 = module->AddEmbeddedComputation(build_body()); + auto while0 = builder.AddInstruction( + HloInstruction::CreateWhile(r0s32, cond0, body0, infeed)); + + auto cond1 = module->AddEmbeddedComputation(build_cond()); + auto body1 = module->AddEmbeddedComputation(build_body()); + auto while1 = builder.AddInstruction( + HloInstruction::CreateWhile(r0s32, cond1, body1, while0)); + + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0))); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r0s32, HloOpcode::kAdd, zero, zero)); + auto cond2 = module->AddEmbeddedComputation(build_cond()); + auto body2 = module->AddEmbeddedComputation(build_body()); + auto while2 = builder.AddInstruction( + HloInstruction::CreateWhile(r0s32, cond2, body2, add)); + + auto tuple = + builder.AddInstruction(HloInstruction::CreateTuple({while2, while1})); + module->AddEntryComputation(builder.Build()); + + // Run CopyInsertion and check if the graph constructed above doesn't need + // any copies inserted for BufferAssignment to run. + int64 instruction_count = module->instruction_count(); + CopyInsertion copy_insertion; + ASSERT_IS_OK(copy_insertion.Run(module.get()).status()); + ASSERT_EQ(instruction_count, module->instruction_count()); + + // Create a sequential order among all the instructions in the entry + // computation, since the issue this test stresses depends on the order the + // nodes are traversed during BufferAssignment. + SequentialHloOrdering::HloModuleSequence sequence; + sequence[module->entry_computation()] = {infeed, while0, while1, zero, + add, while2, tuple}; + TF_ASSERT_OK_AND_ASSIGN( + auto assignment, + BufferAssigner::Run( + module.get(), + xla::MakeUnique(module.get(), sequence), + backend().compiler()->BufferSizeBytesFunction(), + [](LogicalBuffer::Color) { return 1; })); + + // The result tuple elements must be assigned with different buffers. + TF_ASSERT_OK_AND_ASSIGN(auto slice0, assignment->GetUniqueSlice(tuple, {0})); + TF_ASSERT_OK_AND_ASSIGN(auto slice1, assignment->GetUniqueSlice(tuple, {1})); + EXPECT_NE(slice0, slice1); + + // while0 and while1 result buffers must be equal to slice1. + TF_ASSERT_OK_AND_ASSIGN(auto slice_while0, + assignment->GetUniqueSlice(while0, {})); + TF_ASSERT_OK_AND_ASSIGN(auto slice_while1, + assignment->GetUniqueSlice(while1, {})); + EXPECT_EQ(slice1, slice_while0); + EXPECT_EQ(slice1, slice_while1); + + // while2 result buffer must be equal to slice0. + TF_ASSERT_OK_AND_ASSIGN(auto slice_while2, + assignment->GetUniqueSlice(while2, {})); + EXPECT_EQ(slice0, slice_while2); +} + TEST_F(WhileBufferAssignmentTest, OneForwardBackwardWhileLoopSet) { auto module = xla::MakeUnique(TestName()); auto builder = HloComputation::Builder("entry"); diff --git a/tensorflow/compiler/xla/service/buffer_liveness.cc b/tensorflow/compiler/xla/service/buffer_liveness.cc index 513bfa3b7f7b45696093d03c1dd8250c548d260a..37982aaef9eddd64ef6b57ad5a9cf8dd6a565097 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness.cc +++ b/tensorflow/compiler/xla/service/buffer_liveness.cc @@ -102,8 +102,8 @@ bool BufferLiveness::live_range_strictly_before(const LogicalBuffer& a, return false; } - // Every user of 'a' must be a predecessor of 'b' or 'b' itself. for (const BufferAlias& alias : points_to_analysis_->GetBufferAliases(a)) { + // Every user of 'a' must be a predecessor of 'b' or 'b' itself. for (auto user : alias.instruction()->users()) { if (DoesNotUseOperandBuffer(alias.instruction(), alias.index(), user, points_to_analysis())) { @@ -114,6 +114,17 @@ bool BufferLiveness::live_range_strictly_before(const LogicalBuffer& a, return false; } } + + // If the root instruction aliases the buffer 'a', the live range of 'a' is + // until the end of the computation and can never be strictly before another + // buffer defined in the same computation. This is needed to prevent the + // root instruction's buffers from being reused by later instructions even + // when the root is not the last instruction in the schedule. + if (alias.instruction()->parent()->root_instruction() == + alias.instruction() && + alias.instruction()->parent() == b.instruction()->parent()) { + return false; + } } // If 'b' is a user of 'a' then the buffers interfere unless 'a.instruction' diff --git a/tensorflow/compiler/xla/service/buffer_liveness_test.cc b/tensorflow/compiler/xla/service/buffer_liveness_test.cc index 13825fe05bb1b98045f1a3dac3d7272a2d1151fb..f623aef67a4f98b447a9a15634a78deb60cfe6f1 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness_test.cc +++ b/tensorflow/compiler/xla/service/buffer_liveness_test.cc @@ -311,6 +311,48 @@ TEST_F(BufferLivenessTest, OverlappedBuffersSequentialOrder) { EXPECT_FALSE(InstructionsMayInterfere(*liveness, add, exp)); } +TEST_F(BufferLivenessTest, RootInstructionIsNotLastInSequentialOrder) { + // Tests that when the root instruction is not the last instruction in the + // schedule, the live range of its buffers interfere with the buffers of the + // later instructions. + // + // Two sets of independent instructions are executed in the computation. + // param --> add (root) + // recv --> recv-done --> send --> send-done + // + // Sequential order: + // param, add (root), recv, recv-done, send, send-done + auto builder = HloComputation::Builder(TestName()); + auto param = + builder.AddInstruction(HloInstruction::CreateParameter(0, vec_, "param")); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(vec_, HloOpcode::kAdd, param, param)); + auto recv = builder.AddInstruction( + HloInstruction::CreateRecv(vec_, /*channel_id=*/0)); + auto recv_done = builder.AddInstruction(HloInstruction::CreateRecvDone(recv)); + auto send = builder.AddInstruction( + HloInstruction::CreateSend(recv_done, /*channel_id=*/1)); + auto send_done = builder.AddInstruction(HloInstruction::CreateSendDone(send)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build(add)); + + SequentialHloOrdering::HloModuleSequence module_sequence; + std::vector order = {param, add, recv, + recv_done, send, send_done}; + module_sequence.emplace(computation, order); + auto liveness = + BufferLiveness::Run(module.get(), xla::MakeUnique( + module.get(), module_sequence)) + .ConsumeValueOrDie(); + + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, add)); + // Check the root instruction (add) buffer interferes with the recv buffer. + EXPECT_TRUE( + liveness->MayInterfere(GetBuffer(*liveness, add, /*index=*/{}), + GetBuffer(*liveness, recv, /*index=*/{0}))); +} + TEST_F(BufferLivenessTest, TupleLiveOut) { // Verify MaybeLiveOut with nested tuples. Result of computation looks like: // diff --git a/tensorflow/compiler/xla/service/call_inliner.cc b/tensorflow/compiler/xla/service/call_inliner.cc index 3aa7f5c4d5829ccc0e8df697c1363754128ff436..482ccc5b67109258f544e5657ecfa0e8f62192c0 100644 --- a/tensorflow/compiler/xla/service/call_inliner.cc +++ b/tensorflow/compiler/xla/service/call_inliner.cc @@ -82,6 +82,10 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault { return outer_->ReplaceInstruction(call_, new_root); } + CallInliner::InlinedInstructionMap ConsumeInstructionMap() { + return std::move(subcomputation_hlo_to_new_hlo_); + } + private: // Resolves the callee subcomputation_hlo to the new (inline) HLO in the // caller computation, or returns a NotFound error if that subcomputation HLO @@ -112,13 +116,13 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault { HloInstruction* call_; HloComputation* outer_; - std::unordered_map - subcomputation_hlo_to_new_hlo_; + CallInliner::InlinedInstructionMap subcomputation_hlo_to_new_hlo_; }; } // namespace -/* static */ Status CallInliner::Inline(HloInstruction* call) { +/* static */ StatusOr CallInliner::Inline( + HloInstruction* call) { TF_RET_CHECK(call->opcode() == HloOpcode::kCall) << "Instruction was not a call op: " << call->opcode(); const auto& callees = call->called_computations(); @@ -126,7 +130,8 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault { HloComputation* callee = callees[0]; // We visit the callee, cloning its body into its caller. SubcomputationInsertionVisitor visitor(call); - return callee->Accept(&visitor); + TF_RETURN_IF_ERROR(callee->Accept(&visitor)); + return visitor.ConsumeInstructionMap(); } StatusOr CallInliner::Run(HloModule* module) { @@ -140,7 +145,7 @@ StatusOr CallInliner::Run(HloModule* module) { VLOG(1) << "Visiting callsite: " << callsite.ToString(); if (callsite.instruction()->opcode() == HloOpcode::kCall) { HloInstruction* call = callsite.instruction(); - TF_RETURN_IF_ERROR(Inline(call)); + TF_RETURN_IF_ERROR(Inline(call).status()); did_mutate = true; } } diff --git a/tensorflow/compiler/xla/service/call_inliner.h b/tensorflow/compiler/xla/service/call_inliner.h index 2dbd38bf1ac90d3efa1453e6af6f791668d5e72a..a8345a394d46c90a48305313dac0bcd9b06938ac 100644 --- a/tensorflow/compiler/xla/service/call_inliner.h +++ b/tensorflow/compiler/xla/service/call_inliner.h @@ -27,8 +27,12 @@ namespace xla { // called function, and proceed recursively. class CallInliner : public HloPassInterface { public: - // Inlines one call instruction. - static Status Inline(HloInstruction* call); + using InlinedInstructionMap = + std::unordered_map; + + // Inlines one call instruction. Returns a mapping from the original + // instructions to their inlined versions. + static StatusOr Inline(HloInstruction* call); ~CallInliner() override = default; tensorflow::StringPiece name() const override { return "CallInliner"; } diff --git a/tensorflow/compiler/xla/service/call_inliner_test.cc b/tensorflow/compiler/xla/service/call_inliner_test.cc index 865ed993da121d26ceb61123f1822d93814cbb9b..738d00881dd057fc13c115006c15e8f5b6d14a1d 100644 --- a/tensorflow/compiler/xla/service/call_inliner_test.cc +++ b/tensorflow/compiler/xla/service/call_inliner_test.cc @@ -135,7 +135,7 @@ TEST_F(CallInlinerTest, InlineWithoutRunningPass) { HloInstruction::CreateCall(pred, {}, false_computation)); auto computation = module->AddEntryComputation(call_false_builder.Build()); - TF_ASSERT_OK(CallInliner::Inline(call)); + TF_ASSERT_OK(CallInliner::Inline(call).status()); EXPECT_THAT(computation->root_instruction(), op::Constant()); EXPECT_THAT(computation->root_instruction()->control_successors(), ElementsAre(op::Constant())); diff --git a/tensorflow/compiler/xla/service/compile_only_service.cc b/tensorflow/compiler/xla/service/compile_only_service.cc index 9e96898d9b4215e67c8686d372e4b4e6edd1d88b..dab73596e1639eed62151197048ee8d29570b20a 100644 --- a/tensorflow/compiler/xla/service/compile_only_service.cc +++ b/tensorflow/compiler/xla/service/compile_only_service.cc @@ -101,12 +101,13 @@ CompileOnlyService::CompileAheadOfTime( TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, CreateModuleConfig(*program_shape, instance.argument_layouts, - &execution_options)); + &execution_options, *user_computation)); TF_ASSIGN_OR_RETURN(std::unique_ptr hlo_module, computation_tracker_.BuildHloModule( versioned_handle, *module_config, /*include_unreachable_instructions=*/true)); + TF_RETURN_IF_ERROR(MaybeDumpHloModule(*hlo_module)); hlo_modules.push_back(std::move(hlo_module)); } diff --git a/tensorflow/compiler/xla/service/compiler.h b/tensorflow/compiler/xla/service/compiler.h index fc67330f5cbdbcb0d1a259d284599916a908d1fe..74fd24edf88d44b2dfdc87556b0af43987e69e08 100644 --- a/tensorflow/compiler/xla/service/compiler.h +++ b/tensorflow/compiler/xla/service/compiler.h @@ -72,8 +72,18 @@ class AotCompilationOptions { // Returns the ID of the platform to which these options apply. virtual perftools::gputools::Platform::Id PlatformId() const = 0; + // Optional allocator that may be used for allocating temp space on the device + // during compilation. + DeviceMemoryAllocator* device_allocator() const { return device_allocator_; } + void set_device_allocator(DeviceMemoryAllocator* device_allocator) { + device_allocator_ = device_allocator; + } + protected: AotCompilationOptions() = default; + + private: + DeviceMemoryAllocator* device_allocator_ = nullptr; }; // Abstract compiler interface that is subclassed for compilation on a @@ -99,9 +109,16 @@ class Compiler { // Runs Hlo passes to optimize the given Hlo module, returns the optimized // module. + // + // If device_allocator is not null, the compiler may use it to allocate temp + // space on the device for use during compilation. For example, the compiler + // may allocate buffers on the device and then run variants of a given + // algorithm over those buffers, to see which variant is fastest. Any space + // allocated should be deallocated before this function returns. virtual StatusOr> RunHloPasses( std::unique_ptr module, - perftools::gputools::StreamExecutor* executor) = 0; + perftools::gputools::StreamExecutor* executor, + DeviceMemoryAllocator* device_allocator) = 0; // Compiles the HLO module for execution on a device given by the executor, // and returns an executable object or an error status. No HLO passes are @@ -112,21 +129,27 @@ class Compiler { // The compiler may optionally specialize to the individual device // (not just type of device) indicated by the executor. // + // device_allocator is optional; see RunHloPasses. + // // Use the overload below to compile computations that run in parallel. virtual StatusOr> RunBackend( std::unique_ptr module, - perftools::gputools::StreamExecutor* executor) = 0; + perftools::gputools::StreamExecutor* executor, + DeviceMemoryAllocator* device_allocator) = 0; // Compiles a set of HLO modules that can run in parallel, potentially // communicating data between the modules, and returns a corresponding // sequence of executable objects. // + // device_allocator is optional; see RunHloPasses. + // // TODO(b/68666782): Remove this method after adding support for multiple // modules to RunHloPasses and RunBackends. virtual StatusOr>> Compile( std::vector> modules, std::vector> - stream_exec) = 0; + stream_exec, + DeviceMemoryAllocator* device_allocator) = 0; // Compiles the HLO module for ahead-of-time execution. This is intended for // use in static compilation. diff --git a/tensorflow/compiler/xla/service/copy_insertion_test.cc b/tensorflow/compiler/xla/service/copy_insertion_test.cc index 8388574716ad1b78eb8868a8cd732005050b3310..153f062d015e49db11c4c9ae0a2a61e76c020f02 100644 --- a/tensorflow/compiler/xla/service/copy_insertion_test.cc +++ b/tensorflow/compiler/xla/service/copy_insertion_test.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_runner.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" @@ -1723,8 +1724,242 @@ void BM_ParallelWhiles(int num_iters, int num_whiles) { } } +std::unique_ptr MakeBenchmarkWhileBody( + const int num_tuple_inputs) { + auto builder = HloComputation::Builder("benchmark_loop_body"); + const Shape element_shape = ShapeUtil::MakeShape(F32, {}); + std::vector input_shape(num_tuple_inputs, element_shape); + const Shape loop_state_shape = ShapeUtil::MakeTupleShape(input_shape); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, loop_state_shape, "loop_state")); + std::vector gte_nodes(num_tuple_inputs); + for (int i = 0; i < num_tuple_inputs; ++i) { + gte_nodes[i] = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(element_shape, param, i)); + } + builder.AddInstruction(HloInstruction::CreateTuple(gte_nodes)); + return builder.Build(); +} + +void BM_ManyElementTuple(int num_iters, const int num_tuple_inputs) { + tensorflow::testing::StopTiming(); + HloModuleConfig config; + config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); + CopyInsertion copy_insertion; + const Shape element_shape = ShapeUtil::MakeShape(F32, {}); + std::vector tuple_params(num_tuple_inputs); + for (int i = 0; i < num_iters; ++i) { + auto builder = HloComputation::Builder("BM_ParallelWhiles"); + HloModule module("BM_ManyElementTuple", VersionedComputationHandle(), + config); + for (int j = 0; j < num_tuple_inputs; ++j) { + tuple_params[j] = builder.AddInstruction( + HloInstruction::CreateParameter(j, element_shape, "")); + } + HloInstruction* init = + builder.AddInstruction(HloInstruction::CreateTuple(tuple_params)); + HloComputation* condition = + module.AddEmbeddedComputation(MakeTrivialCondition(init->shape())); + HloComputation* body = + module.AddEmbeddedComputation(MakeBenchmarkWhileBody(num_tuple_inputs)); + HloInstruction* xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(init->shape(), condition, body, init)); + builder.AddInstruction(HloInstruction::CreateGetTupleElement( + ShapeUtil::MakeShape(F32, {}), xla_while, 0)); + module.AddEntryComputation(builder.Build()); + tensorflow::testing::StartTiming(); + ASSERT_IS_OK(copy_insertion.Run(&module).status()); + tensorflow::testing::StopTiming(); + } +} + BENCHMARK(BM_SequentialWhiles)->Arg(512)->Arg(1024)->Arg(2048)->Arg(4096); BENCHMARK(BM_ParallelWhiles)->Arg(512)->Arg(1024)->Arg(2048)->Arg(4096); +BENCHMARK(BM_ManyElementTuple)->Arg(1024)->Arg(12288); + +TEST_F(CopyInsertionTest, SimpleControlFlowTest) { + const string& hlo_string = R"( +HloModule TestModule + +if-body.v5 { + constant.3 = s32[] constant(-1) + p.1 = (s32[], (s32[], s32[], s32[]), (s32[])) parameter(0) + get-tuple-element.18 = (s32[], s32[], s32[]) get-tuple-element(p.1), index=1 + get-tuple-element.65 = s32[] get-tuple-element(get-tuple-element.18), index=0 + get-tuple-element.66 = s32[] get-tuple-element(get-tuple-element.18), index=1 + add.3 = s32[] add(get-tuple-element.65, get-tuple-element.66) + tuple.33 = (s32[]) tuple(add.3) + ROOT tuple.34 = (s32[], (s32[], s32[], s32[]), (s32[])) tuple(constant.3, get-tuple-element.18, tuple.33) +} + +if-condition.v4 { + p.2 = (s32[], (s32[], s32[], s32[]), (s32[])) parameter(0) + get-tuple-element.67 = s32[] get-tuple-element(p.2), index=0 + constant.4 = s32[] constant(0) + ROOT equal-to = pred[] equal-to(get-tuple-element.67, constant.4) +} + +_functionalize_body_1__.v28 { + arg_tuple.4 = (s32[], s32[], s32[], s32[]) parameter(0) + get-tuple-element.68 = s32[] get-tuple-element(arg_tuple.4), index=0 + constant.7 = s32[] constant(1) + add.4 = s32[] add(get-tuple-element.68, constant.7) + get-tuple-element.69 = s32[] get-tuple-element(arg_tuple.4), index=1 + get-tuple-element.70 = s32[] get-tuple-element(arg_tuple.4), index=2 + less-than-or-equal-to = pred[] less-than-or-equal-to(get-tuple-element.69, get-tuple-element.70) + constant.8 = s32[] constant(0) + select = s32[] select(less-than-or-equal-to, constant.8, constant.7) + get-tuple-element.71 = s32[] get-tuple-element(arg_tuple.4), index=3 + tuple.35 = (s32[], s32[], s32[]) tuple(get-tuple-element.69, get-tuple-element.71, get-tuple-element.70) + tuple.36 = (s32[]) tuple(constant.8) + tuple.37 = (s32[], (s32[], s32[], s32[]), (s32[])) tuple(select, tuple.35, tuple.36) + while = (s32[], (s32[], s32[], s32[]), (s32[])) while(tuple.37), condition=if-condition.v4, body=if-body.v5 + get-tuple-element.72 = (s32[]) get-tuple-element(while), index=2 + get-tuple-element.73 = s32[] get-tuple-element(get-tuple-element.72), index=0 + ROOT tuple.38 = (s32[], s32[], s32[], s32[]) tuple(add.4, get-tuple-element.69, get-tuple-element.70, get-tuple-element.73) +} + +cond_wrapper.v3.1 { + inputs.1 = (s32[], s32[], s32[], s32[]) parameter(0) + get-tuple-element.75 = s32[] get-tuple-element(inputs.1), index=0 + constant.11 = s32[] constant(7) + ROOT less-than.2 = pred[] less-than(get-tuple-element.75, constant.11) +} + +_functionalize_body_2__.v25 { + arg_tuple.5 = (s32[], s32[], s32[], s32[], s32[]) parameter(0) + get-tuple-element.76 = s32[] get-tuple-element(arg_tuple.5), index=0 + get-tuple-element.77 = s32[] get-tuple-element(arg_tuple.5), index=2 + get-tuple-element.78 = s32[] get-tuple-element(arg_tuple.5), index=3 + get-tuple-element.79 = s32[] get-tuple-element(arg_tuple.5), index=4 + tuple.39 = (s32[], s32[], s32[], s32[]) tuple(get-tuple-element.76, get-tuple-element.77, get-tuple-element.78, get-tuple-element.79) + while.2 = (s32[], s32[], s32[], s32[]) while(tuple.39), condition=cond_wrapper.v3.1, body=_functionalize_body_1__.v28 + get-tuple-element.80 = s32[] get-tuple-element(while.2), index=0 + get-tuple-element.81 = s32[] get-tuple-element(arg_tuple.5), index=1 + constant.12 = s32[] constant(1) + add.5 = s32[] add(get-tuple-element.81, constant.12) + get-tuple-element.82 = s32[] get-tuple-element(while.2), index=3 + ROOT tuple.40 = (s32[], s32[], s32[], s32[], s32[]) tuple(get-tuple-element.80, add.5, get-tuple-element.77, get-tuple-element.78, get-tuple-element.82) +} + +cond_wrapper.v3.2 { + inputs.2 = (s32[], s32[], s32[], s32[], s32[]) parameter(0) + get-tuple-element.83 = s32[] get-tuple-element(inputs.2), index=1 + constant.13 = s32[] constant(5) + ROOT less-than.3 = pred[] less-than(get-tuple-element.83, constant.13) +} + +ENTRY TestComputation { + arg_tuple.6 = (s32[], s32[], s32[], s32[], s32[]) parameter(0) + ROOT while.3 = (s32[], s32[], s32[], s32[], s32[]) while(arg_tuple.6), condition=cond_wrapper.v3.2, body=_functionalize_body_2__.v25 +} +)"; + auto module_or_status = + HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest()); + auto module = module_or_status.ConsumeValueOrDie(); + InsertCopies(module.get()); +} + +TEST_F(CopyInsertionTest, ControlFlowTest) { + const string& hlo_string = R"( +HloModule TestModule + +if-body.v5 { + constant.3 = s32[] constant(-1) + p.1 = (s32[], (s32[], s32[], s32[]), (s32[])) parameter(0) + get-tuple-element.18 = (s32[], s32[], s32[]) get-tuple-element(p.1), index=1 + get-tuple-element.65 = s32[] get-tuple-element(get-tuple-element.18), index=0 + get-tuple-element.66 = s32[] get-tuple-element(get-tuple-element.18), index=1 + add.3 = s32[] add(get-tuple-element.65, get-tuple-element.66) + tuple.33 = (s32[]) tuple(add.3) + ROOT tuple.34 = (s32[], (s32[], s32[], s32[]), (s32[])) tuple(constant.3, get-tuple-element.18, tuple.33) +} + +if-condition.v4 { + p.2 = (s32[], (s32[], s32[], s32[]), (s32[])) parameter(0) + get-tuple-element.67 = s32[] get-tuple-element(p.2), index=0 + constant.4 = s32[] constant(0) + ROOT equal-to = pred[] equal-to(get-tuple-element.67, constant.4) +} + +if-body.v5.1 { + constant.5 = s32[] constant(-1) + p.3 = (s32[], (s32[], s32[], s32[]), (s32[])) parameter(0) + get-tuple-element.68 = (s32[], s32[], s32[]) get-tuple-element(p.3), index=1 + get-tuple-element.70 = s32[] get-tuple-element(get-tuple-element.68), index=2 + multiply.1 = s32[] multiply(get-tuple-element.70, get-tuple-element.70) + tuple.35 = (s32[]) tuple(multiply.1) + ROOT tuple.36 = (s32[], (s32[], s32[], s32[]), (s32[])) tuple(constant.5, get-tuple-element.68, tuple.35) +} + +if-condition.v4.1 { + p.4 = (s32[], (s32[], s32[], s32[]), (s32[])) parameter(0) + get-tuple-element.71 = s32[] get-tuple-element(p.4), index=0 + constant.6 = s32[] constant(1) + ROOT equal-to.1 = pred[] equal-to(get-tuple-element.71, constant.6) +} + +_functionalize_body_1__.v28 { + arg_tuple.4 = (s32[], s32[], s32[], s32[]) parameter(0) + get-tuple-element.72 = s32[] get-tuple-element(arg_tuple.4), index=0 + constant.7 = s32[] constant(1) + add.4 = s32[] add(get-tuple-element.72, constant.7) + get-tuple-element.73 = s32[] get-tuple-element(arg_tuple.4), index=1 + get-tuple-element.74 = s32[] get-tuple-element(arg_tuple.4), index=2 + less-than-or-equal-to = pred[] less-than-or-equal-to(get-tuple-element.73, get-tuple-element.74) + constant.8 = s32[] constant(0) + select = s32[] select(less-than-or-equal-to, constant.8, constant.7) + get-tuple-element.75 = s32[] get-tuple-element(arg_tuple.4), index=3 + tuple.37 = (s32[], s32[], s32[]) tuple(get-tuple-element.73, get-tuple-element.75, get-tuple-element.74) + tuple.38 = (s32[]) tuple(constant.8) + tuple.39 = (s32[], (s32[], s32[], s32[]), (s32[])) tuple(select, tuple.37, tuple.38) + while = (s32[], (s32[], s32[], s32[]), (s32[])) while(tuple.39), condition=if-condition.v4, body=if-body.v5 + while.1 = (s32[], (s32[], s32[], s32[]), (s32[])) while(while), condition=if-condition.v4.1, body=if-body.v5.1 + get-tuple-element.76 = (s32[]) get-tuple-element(while.1), index=2 + get-tuple-element.77 = s32[] get-tuple-element(get-tuple-element.76), index=0 + ROOT tuple.40 = (s32[], s32[], s32[], s32[]) tuple(add.4, get-tuple-element.73, get-tuple-element.74, get-tuple-element.77) +} + +cond_wrapper.v3.1 { + inputs.1 = (s32[], s32[], s32[], s32[]) parameter(0) + get-tuple-element.78 = s32[] get-tuple-element(inputs.1), index=0 + constant.11 = s32[] constant(7) + ROOT less-than.2 = pred[] less-than(get-tuple-element.78, constant.11) +} + +_functionalize_body_2__.v25 { + arg_tuple.5 = (s32[], s32[], s32[], s32[], s32[]) parameter(0) + get-tuple-element.79 = s32[] get-tuple-element(arg_tuple.5), index=0 + get-tuple-element.80 = s32[] get-tuple-element(arg_tuple.5), index=2 + get-tuple-element.81 = s32[] get-tuple-element(arg_tuple.5), index=3 + get-tuple-element.82 = s32[] get-tuple-element(arg_tuple.5), index=4 + tuple.41 = (s32[], s32[], s32[], s32[]) tuple(get-tuple-element.79, get-tuple-element.80, get-tuple-element.81, get-tuple-element.82) + while.2 = (s32[], s32[], s32[], s32[]) while(tuple.41), condition=cond_wrapper.v3.1, body=_functionalize_body_1__.v28 + get-tuple-element.83 = s32[] get-tuple-element(while.2), index=0 + get-tuple-element.84 = s32[] get-tuple-element(arg_tuple.5), index=1 + constant.12 = s32[] constant(1) + add.5 = s32[] add(get-tuple-element.84, constant.12) + get-tuple-element.85 = s32[] get-tuple-element(while.2), index=3 + ROOT tuple.42 = (s32[], s32[], s32[], s32[], s32[]) tuple(get-tuple-element.83, add.5, get-tuple-element.80, get-tuple-element.81, get-tuple-element.85) +} + +cond_wrapper.v3.2 { + inputs.2 = (s32[], s32[], s32[], s32[], s32[]) parameter(0) + get-tuple-element.86 = s32[] get-tuple-element(inputs.2), index=1 + constant.13 = s32[] constant(5) + ROOT less-than.3 = pred[] less-than(get-tuple-element.86, constant.13) +} + +ENTRY TestComputation { + arg_tuple.6 = (s32[], s32[], s32[], s32[], s32[]) parameter(0) + ROOT while.3 = (s32[], s32[], s32[], s32[], s32[]) while(arg_tuple.6), condition=cond_wrapper.v3.2, body=_functionalize_body_2__.v25 +} +)"; + auto module_or_status = + HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest()); + auto module = module_or_status.ConsumeValueOrDie(); + InsertCopies(module.get()); +} } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD index c9fbeae77ca17d466c4d9bb1ab4a27e3d94382bc..c13a0b1cdf0b5be0b69db98b2b9587f30ca4c304 100644 --- a/tensorflow/compiler/xla/service/cpu/BUILD +++ b/tensorflow/compiler/xla/service/cpu/BUILD @@ -81,6 +81,7 @@ cc_library( ":conv_canonicalization", ":cpu_copy_insertion", ":cpu_executable", + ":cpu_hlo_support_checker", ":cpu_instruction_fusion", ":cpu_layout_assignment", ":cpu_options", @@ -126,6 +127,7 @@ cc_library( "//tensorflow/compiler/xla/service:reshape_mover", "//tensorflow/compiler/xla/service:transpose_folding", "//tensorflow/compiler/xla/service:tuple_simplifier", + "//tensorflow/compiler/xla/service:while_loop_invariant_code_motion", "//tensorflow/compiler/xla/service:while_loop_simplifier", "//tensorflow/compiler/xla/service:zero_sized_hlo_elimination", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", # fixdeps: keep @@ -148,19 +150,21 @@ cc_library( cc_library( name = "simple_orc_jit", - srcs = ["simple_orc_jit.cc"], + srcs = [ + "simple_orc_jit.cc", + "windows_compatibility.cc", + "windows_compatibility.h", + ], hdrs = ["simple_orc_jit.h"], deps = [ ":compiler_functor", ":cpu_runtime", - ":cpu_runtime_avx", - ":cpu_runtime_neon", - ":cpu_runtime_sse4_1", ":custom_call_target_registry", ":disassembler", ":external_constant_pool", ":orc_jit_memory_mapper", ":runtime_conv2d", + ":runtime_fft", ":runtime_fork_join", ":runtime_matmul", ":runtime_single_threaded_conv2d", @@ -255,6 +259,7 @@ cc_library( ":parallel_loop_emitter", ":shape_partition", ":simple_orc_jit", + ":target_machine_features", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", @@ -276,7 +281,6 @@ cc_library( "//tensorflow/compiler/xla/service/llvm_ir:ops", "//tensorflow/compiler/xla/service/llvm_ir:tuple_ops", "//tensorflow/core:lib", - "@llvm//:analysis", "@llvm//:code_gen", "@llvm//:core", "@llvm//:support", @@ -284,6 +288,20 @@ cc_library( ], ) +cc_library( + name = "target_machine_features", + srcs = [ + "target_machine_features.cc", + ], + hdrs = ["target_machine_features.h"], + deps = [ + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/core:lib", + "@llvm//:analysis", + "@llvm//:target", + ], +) + cc_library( name = "ir_function", srcs = ["ir_function.cc"], @@ -325,6 +343,8 @@ cc_library( deps = [ ":cpu_options", ":cpu_runtime", + ":target_machine_features", + ":vector_support_library", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:types", @@ -336,7 +356,6 @@ cc_library( "//tensorflow/compiler/xla/service/llvm_ir:kernel_support_library", "//tensorflow/compiler/xla/service/llvm_ir:llvm_loop", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", - "//tensorflow/compiler/xla/service/llvm_ir:vector_support_library", "//tensorflow/core:lib", "@llvm//:core", ], @@ -386,9 +405,6 @@ cc_library( hdrs = ["compiler_functor.h"], deps = [ ":cpu_runtime", - ":cpu_runtime_avx", - ":cpu_runtime_neon", - ":cpu_runtime_sse4_1", ":disassembler", ":llvm_ir_runtime", "//tensorflow/compiler/xla:statusor", @@ -408,43 +424,6 @@ cc_library( ], ) -cc_library( - name = "cpu_runtime_sse4_1", - srcs = ["cpu_runtime_sse4_1.cc"], - hdrs = ["cpu_runtime_sse4_1.h"], - copts = ["-DEIGEN_AVOID_STL_ARRAY"], - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/core:framework_lite", - "//third_party/eigen3", - ], -) - -cc_library( - name = "cpu_runtime_avx", - srcs = ["cpu_runtime_avx.cc"], - hdrs = ["cpu_runtime_avx.h"], - copts = ["-DEIGEN_AVOID_STL_ARRAY"], - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/core:framework_lite", - "//third_party/eigen3", - ], -) - -cc_library( - name = "cpu_runtime_neon", - srcs = ["cpu_runtime_neon.cc"], - hdrs = ["cpu_runtime_neon.h"], - # runtime_copts() enables -mfpu=neon - copts = ["-DEIGEN_AVOID_STL_ARRAY"] + runtime_copts(), - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/core:framework_lite", - "//third_party/eigen3", - ], -) - cc_library( name = "cpu_runtime", srcs = [ @@ -475,6 +454,7 @@ cc_library( "llvm_ir_runtime.h", ], deps = [ + ":vector_support_library", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/core:lib", "@llvm//:core", @@ -499,6 +479,24 @@ cc_library( ], ) +cc_library( + name = "runtime_fft", + srcs = [ + "runtime_fft.cc", + "runtime_fft_impl.h", + ], + hdrs = ["runtime_fft.h"], + copts = runtime_copts(), + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/compiler/xla:executable_run_options", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:framework", + "//tensorflow/core:framework_lite", + "//third_party/eigen3", + ], +) + cc_library( name = "runtime_matvec", srcs = ["runtime_matvec.cc"], @@ -801,6 +799,21 @@ cc_library( ], ) +cc_library( + name = "vector_support_library", + srcs = ["vector_support_library.cc"], + hdrs = ["vector_support_library.h"], + deps = [ + ":target_machine_features", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", + "@llvm//:core", + "@llvm//:support", + ], +) + tf_cc_test( name = "cpu_copy_insertion_test", srcs = ["cpu_copy_insertion_test.cc"], @@ -821,6 +834,32 @@ tf_cc_test( ], ) +cc_library( + name = "cpu_hlo_support_checker", + srcs = ["cpu_hlo_support_checker.cc"], + hdrs = ["cpu_hlo_support_checker.h"], + deps = [ + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo_pass", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "cpu_hlo_support_checker_test", + srcs = ["cpu_hlo_support_checker_test.cc"], + deps = [ + ":cpu_hlo_support_checker", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc index 04b4a8c5c80eeefdbe10001ba5c462affbc9b21d..ed290fcdf8bb69f1bbad57fa5a0926376bc9405a 100644 --- a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc +++ b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc @@ -37,9 +37,6 @@ limitations under the License. #include "llvm/Transforms/IPO/PassManagerBuilder.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" -#include "tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h" -#include "tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h" -#include "tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h" #include "tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/statusor.h" @@ -50,15 +47,6 @@ limitations under the License. namespace xla { namespace cpu { -/* static */ CompilerFunctor::VectorIntrinsics -CompilerFunctor::AllIntrinsics() { - VectorIntrinsics intrinsics; - intrinsics.sse_intrinsics = true; - intrinsics.avx_intrinsics = true; - intrinsics.neon_intrinsics = true; - return intrinsics; -} - /* Create filtered versions of the LLVM Pass Managers to filter out some of the expensive passes. Profiling: @@ -192,89 +180,28 @@ operator()(llvm::Module& module) const { std::move(object_file), std::move(memory_buffer)); } -namespace { -// Returns the set of vectorized library functions supported for the target. -std::vector VectorFunctionsForTargetLibraryInfoImpl( - llvm::Triple::ArchType arch, llvm::StringRef feature_string, - CompilerFunctor::VectorIntrinsics const& available_intrinsics) { - std::vector vector_functions; - - const llvm::VecDesc four_wide_vector_functions_neon[] = { - {"expf", runtime::kExpV4F32NEONSymbolName, 4}, - {"llvm.exp.f32", runtime::kExpV4F32NEONSymbolName, 4}, - - {"logf", runtime::kLogV4F32NEONSymbolName, 4}, - {"llvm.log.f32", runtime::kLogV4F32NEONSymbolName, 4}, - }; - - const llvm::VecDesc four_wide_vector_functions_sse[] = { - {"expf", runtime::kExpV4F32SSESymbolName, 4}, - {"llvm.exp.f32", runtime::kExpV4F32SSESymbolName, 4}, - - {"logf", runtime::kLogV4F32SSESymbolName, 4}, - {"llvm.log.f32", runtime::kLogV4F32SSESymbolName, 4}, - }; - - const llvm::VecDesc eight_wide_vector_functions_avx[] = { - {"expf", runtime::kExpV8F32AVXSymbolName, 8}, - {"llvm.exp.f32", runtime::kExpV8F32AVXSymbolName, 8}, - - {"logf", runtime::kLogV8F32AVXSymbolName, 8}, - {"llvm.log.f32", runtime::kLogV8F32AVXSymbolName, 8}, - }; - - // These functions are generated by XLA as LLVM IR, so they're always - // available. - const llvm::VecDesc ir_vector_functions[] = { +static std::vector VectorFunctionsForTargetLibraryInfoImpl() { + std::vector result = { {"tanhf", runtime::kTanhV4F32SymbolName, 4}, {"llvm.tanh.f32", runtime::kTanhV4F32SymbolName, 4}, {"tanhf", runtime::kTanhV8F32SymbolName, 8}, {"llvm.tanh.f32", runtime::kTanhV8F32SymbolName, 8}, - }; - llvm::SmallVector features; - feature_string.split(features, ',', -1, /*KeepEmpty=*/false); - auto has_feature = [&features](const llvm::StringRef feature) { - return std::find(features.begin(), features.end(), feature) != - features.end(); - }; + {"expf", runtime::kExpV4F32SymbolName, 4}, + {"llvm.exp.f32", runtime::kExpV4F32SymbolName, 4}, - switch (arch) { - case llvm::Triple::x86: - case llvm::Triple::x86_64: { - if (has_feature("+sse4.1") && available_intrinsics.sse_intrinsics) { - vector_functions.insert(vector_functions.end(), - std::begin(four_wide_vector_functions_sse), - std::end(four_wide_vector_functions_sse)); - } - if (has_feature("+avx") && available_intrinsics.avx_intrinsics) { - vector_functions.insert(vector_functions.end(), - std::begin(eight_wide_vector_functions_avx), - std::end(eight_wide_vector_functions_avx)); - } - break; - } - case llvm::Triple::arm: - case llvm::Triple::aarch64: { - if (has_feature("+neon") && available_intrinsics.neon_intrinsics) { - vector_functions.insert(vector_functions.end(), - std::begin(four_wide_vector_functions_neon), - std::end(four_wide_vector_functions_neon)); - } - break; - } - default: - break; - } + {"expf", runtime::kExpV8F32SymbolName, 8}, + {"llvm.exp.f32", runtime::kExpV8F32SymbolName, 8}, - vector_functions.insert(vector_functions.end(), - std::begin(ir_vector_functions), - std::end(ir_vector_functions)); + {"logf", runtime::kLogV4F32SymbolName, 4}, + {"llvm.log.f32", runtime::kLogV4F32SymbolName, 4}, - return vector_functions; + {"logf", runtime::kLogV8F32SymbolName, 8}, + {"llvm.log.f32", runtime::kLogV8F32SymbolName, 8}, + }; + return result; } -} // namespace void CompilerFunctor::AddTargetInfoPasses( llvm::legacy::PassManagerBase* passes) const { @@ -282,9 +209,7 @@ void CompilerFunctor::AddTargetInfoPasses( auto target_library_info_impl = MakeUnique(target_triple); target_library_info_impl->addVectorizableFunctions( - VectorFunctionsForTargetLibraryInfoImpl( - target_triple.getArch(), target_machine_->getTargetFeatureString(), - available_intrinsics_)); + VectorFunctionsForTargetLibraryInfoImpl()); passes->add( new llvm::TargetLibraryInfoWrapperPass(*target_library_info_impl)); passes->add(createTargetTransformInfoWrapperPass( diff --git a/tensorflow/compiler/xla/service/cpu/compiler_functor.h b/tensorflow/compiler/xla/service/cpu/compiler_functor.h index 8cdd049e7b773bdc455db627ff1749997d621ee4..1a8283a702223a7414c1ffcd99c1ac42c04ac068 100644 --- a/tensorflow/compiler/xla/service/cpu/compiler_functor.h +++ b/tensorflow/compiler/xla/service/cpu/compiler_functor.h @@ -31,21 +31,10 @@ namespace cpu { // Orc JIT compile layer. class CompilerFunctor { public: - // Describes the set of vector intrinsics available to the generated code. - struct VectorIntrinsics { - bool sse_intrinsics; - bool avx_intrinsics; - bool neon_intrinsics; - }; - - // Returns a VectorIntrinsics where all intrinsics are available. - static VectorIntrinsics AllIntrinsics(); - explicit CompilerFunctor( llvm::TargetMachine* target_machine, const Disassembler* disassembler, int opt_level, bool optimize_for_size, bool enable_fast_math, bool disable_expensive_passes, - const VectorIntrinsics& available_intrinsics, LLVMCompiler::ModuleHook pre_optimization_hook = nullptr, LLVMCompiler::ModuleHook post_optimization_hook = nullptr) : target_machine_(target_machine), @@ -54,7 +43,6 @@ class CompilerFunctor { optimize_for_size_(optimize_for_size), enable_fast_math_(enable_fast_math), disable_expensive_passes_(disable_expensive_passes), - available_intrinsics_(available_intrinsics), pre_optimization_hook_(pre_optimization_hook), post_optimization_hook_(post_optimization_hook) {} @@ -78,7 +66,6 @@ class CompilerFunctor { const bool optimize_for_size_; const bool enable_fast_math_; const bool disable_expensive_passes_; - const VectorIntrinsics available_intrinsics_; LLVMCompiler::ModuleHook pre_optimization_hook_; LLVMCompiler::ModuleHook post_optimization_hook_; }; diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc index 508721312252810043234fb825797cde718260be..f9cc9651846cca7bd6ab7e9e61590cec4e2400da 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc @@ -31,6 +31,7 @@ limitations under the License. #include "llvm/IR/Function.h" #include "llvm/IR/LLVMContext.h" #include "llvm/IR/Module.h" +#include "llvm/IR/Verifier.h" #include "llvm/Object/ObjectFile.h" #include "llvm/Support/CommandLine.h" #include "llvm/Support/TargetRegistry.h" @@ -50,6 +51,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/conv_canonicalization.h" #include "tensorflow/compiler/xla/service/cpu/cpu_copy_insertion.h" #include "tensorflow/compiler/xla/service/cpu/cpu_executable.h" +#include "tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.h" #include "tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.h" #include "tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.h" #include "tensorflow/compiler/xla/service/cpu/cpu_options.h" @@ -85,6 +87,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/reshape_mover.h" #include "tensorflow/compiler/xla/service/transpose_folding.h" #include "tensorflow/compiler/xla/service/tuple_simplifier.h" +#include "tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h" #include "tensorflow/compiler/xla/service/while_loop_simplifier.h" #include "tensorflow/compiler/xla/service/zero_sized_hlo_elimination.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -164,42 +167,16 @@ namespace { // first module is compiled. std::once_flag llvm_command_line_options_initialized; -void InitializeLLVMCommandLineOptions(const HloModuleConfig& config) { - auto options = config.debug_options().xla_backend_extra_options(); - if (!options.empty()) { - std::vector fake_argv_storage; - fake_argv_storage.push_back(""); - for (const auto& it : options) { - // Skip options the XLA backend itself consumes. - if (!tensorflow::StringPiece(it.first).starts_with("xla_")) { - if (it.second.empty()) { - fake_argv_storage.push_back(it.first); - } else { - fake_argv_storage.push_back(it.first + "=" + it.second); - } - } - } - - VLOG(2) << "Passing argv to LLVM:"; - std::vector fake_argv; - for (const auto& s : fake_argv_storage) { - fake_argv.push_back(s.c_str()); - VLOG(2) << s; - } - llvm::cl::ParseCommandLineOptions(fake_argv.size(), &fake_argv[0]); - } -} - // This visitor records which HLO instructions should have profiling information // recorded. class CollectProfileCandidates : public DfsHloVisitorWithDefault { public: - static StatusOr> + static StatusOr> GetCandidatesForComputation( HloComputation* computation, const std::unordered_map& assigned_indices) { - std::unordered_map hlo_to_profile_idx; + std::unordered_map hlo_to_profile_idx; CollectProfileCandidates profile_candidates_for_computation( &hlo_to_profile_idx, assigned_indices); TF_RETURN_IF_ERROR( @@ -209,7 +186,7 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault { private: CollectProfileCandidates( - std::unordered_map* hlo_to_profile_idx, + std::unordered_map* hlo_to_profile_idx, const std::unordered_map& assigned_indices) : hlo_to_profile_idx_(hlo_to_profile_idx), assigned_indices_(assigned_indices) {} @@ -249,7 +226,7 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault { return Status::OK(); } - std::unordered_map* hlo_to_profile_idx_; + std::unordered_map* hlo_to_profile_idx_; const std::unordered_map& assigned_indices_; }; } // namespace @@ -257,7 +234,8 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault { Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile) { // Optimization pipeline. HloPassPipeline pipeline("CPU"); - pipeline.AddInvariantChecker(ShapeSizeBytesFunction()); + pipeline.AddInvariantChecker(); + pipeline.AddPass(); ReducePrecisionInsertion::AddPasses( &pipeline, module->config().debug_options(), @@ -275,7 +253,7 @@ Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile) { { auto& pass = pipeline.AddPass>("simplification"); - pass.AddInvariantChecker(ShapeSizeBytesFunction()); + pass.AddInvariantChecker(); pass.AddPass( /*rewrite_training_op=*/true, @@ -291,6 +269,7 @@ Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile) { // elimination has to come after that pass. pipeline.AddPass(); + pass.AddPass(); pass.AddPass(); pass.AddPass(); pass.AddPass(); @@ -439,11 +418,27 @@ Status InitializeModuleHooks( return Status::OK(); } +Status VerifyLlvmModule(const llvm::Module& llvm_module) { + XLA_SCOPED_LOGGING_TIMER("CpuCompiler - Running LLVM verifier"); + + std::string err; + llvm::raw_string_ostream err_stream(err); + + // verifyModule() returns true if the module is broken. + TF_RET_CHECK(!llvm::verifyModule(llvm_module, &err_stream)) + << "Invalid LLVM IR before optimizations:\n" + << err_stream.str() + << "\nThis probably indicates a bug in the HLO -> LLVM IR lowering. " + "Rerun with --xla_dump_ir_to to get the IR. "; + return Status::OK(); +} + } // namespace StatusOr> CpuCompiler::RunHloPasses( std::unique_ptr module, - perftools::gputools::StreamExecutor* /*stream_exec*/) { + perftools::gputools::StreamExecutor* /*stream_exec*/, + DeviceMemoryAllocator* /*device_allocator*/) { VLOG(2) << "Before optimization:"; XLA_VLOG_LINES(2, module->ToString()); @@ -456,7 +451,8 @@ StatusOr> CpuCompiler::RunHloPasses( StatusOr> CpuCompiler::RunBackend( std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec) { + perftools::gputools::StreamExecutor* stream_exec, + DeviceMemoryAllocator* /*device_allocator*/) { const string timer_message = "Compiling [" + module->name() + "] for CPU using JIT"; XLA_SCOPED_LOGGING_TIMER(timer_message); @@ -464,7 +460,7 @@ StatusOr> CpuCompiler::RunBackend( VLOG(1) << "Compiling: " << module->name(); TF_RET_CHECK(stream_exec != nullptr); std::call_once(llvm_command_line_options_initialized, - &InitializeLLVMCommandLineOptions, module->config()); + &llvm_ir::InitializeLLVMCommandLineOptions, module->config()); ModuleHook pre_optimization_ir_hook; ModuleHook post_optimization_ir_hook; @@ -487,17 +483,19 @@ StatusOr> CpuCompiler::RunBackend( llvm_module->setDataLayout(jit->data_layout()); llvm_module->setTargetTriple(jit->target_triple().getTriple()); - HloComputation* computation = module->entry_computation(); - std::unordered_map hlo_to_profile_idx; + HloComputation* entry_computation = module->entry_computation(); + std::unordered_map instruction_to_profile_idx; + std::unordered_map computation_to_profile_idx; std::unique_ptr hlo_profile_index_map; - std::unique_ptr hlo_profile_printer; + std::unique_ptr hlo_profile_printer_data; if (module->config().hlo_profiling_enabled()) { hlo_profile_index_map = MakeUnique(*module); TF_ASSIGN_OR_RETURN( - hlo_to_profile_idx, + instruction_to_profile_idx, CollectProfileCandidates::GetCandidatesForComputation( - computation, hlo_profile_index_map->instruction_to_profile_idx())); + entry_computation, + hlo_profile_index_map->instruction_to_profile_idx())); auto shape_size_bytes = [](const Shape& shape) { // On the cpu, opaques are pointers. @@ -508,8 +506,11 @@ StatusOr> CpuCompiler::RunBackend( }; HloCostAnalysis cost_analysis(shape_size_bytes); - hlo_profile_printer = - CreateHloProfilePrinter(*hlo_profile_index_map, cost_analysis); + TF_RETURN_IF_ERROR(entry_computation->Accept(&cost_analysis)); + hlo_profile_printer_data = + CreateHloProfilePrinterData(*hlo_profile_index_map, cost_analysis); + computation_to_profile_idx = + hlo_profile_index_map->computation_to_profile_idx(); } std::unique_ptr cpu_executable; @@ -518,8 +519,8 @@ StatusOr> CpuCompiler::RunBackend( // ownership is std::moved. const bool embed_ir_in_executable = module->config().debug_options().xla_embed_ir_in_executable(); - const string xla_dump_hlo_proto_to = - module->config().debug_options().xla_dump_hlo_proto_to(); + const string xla_dump_optimized_hlo_proto_to = + module->config().debug_options().xla_dump_optimized_hlo_proto_to(); if (options::CpuParallelBackendRequested(module->config())) { VLOG(1) << "Using parallel cpu backend"; @@ -539,10 +540,10 @@ StatusOr> CpuCompiler::RunBackend( // print one ourselves. XLA_VLOG_LINES(2, assignment->ToString()); - if (!xla_dump_hlo_proto_to.empty()) { + if (!xla_dump_optimized_hlo_proto_to.empty()) { HloProto proto = MakeHloProto(*module, *assignment); TF_RETURN_IF_ERROR(protobuf_util::DumpProtoToDirectory( - proto, xla_dump_hlo_proto_to, module->name())); + proto, xla_dump_optimized_hlo_proto_to, module->name())); } // If we are using the parallel CPU backend, we need to create map from @@ -550,7 +551,7 @@ StatusOr> CpuCompiler::RunBackend( std::map parallel_computations; std::unordered_map> aligned_constants; - for (auto instruction : computation->MakeInstructionPostOrder()) { + for (auto instruction : entry_computation->MakeInstructionPostOrder()) { // Parameters and constants don't get their own computation. if (instruction->opcode() == HloOpcode::kParameter) { continue; @@ -558,7 +559,7 @@ StatusOr> CpuCompiler::RunBackend( if (instruction->opcode() == HloOpcode::kConstant) { // Copy the constant out of the ProtocolBuffer so that we can give it a // higher alignment. - const void* data = instruction->literal().InternalData(); + const void* data = instruction->literal().untyped_data(); int64 size = CpuExecutable::ShapeSizeBytes(instruction->shape()); auto iter = aligned_constants.emplace( instruction, xla::MakeUnique(size)); @@ -575,22 +576,15 @@ StatusOr> CpuCompiler::RunBackend( parallel_computations.emplace(to_apply, instruction); } - // We always profile the entire computation as a whole, even if hlo - // profiling is disabled. When hlo profiling is diabled, we pass in a - // profile counter array of just one element, which corresponds to the whole - // computation. - size_t entry_computation_profile_idx = - hlo_profile_index_map ? hlo_profile_index_map->GetProfileIndexFor( - *module->entry_computation()) - : 0; IrEmitter ir_emitter(*module, *assignment, llvm_module.get(), - hlo_to_profile_idx, entry_computation_profile_idx, + std::move(instruction_to_profile_idx), + std::move(computation_to_profile_idx), jit->target_machine(), jit->external_constant_pool()); std::unique_ptr> function_names( new HloInstructionMap()); for (auto embedded_computation : - computation->MakeEmbeddedComputationsList()) { + entry_computation->MakeEmbeddedComputationsList()) { if (embedded_computation->IsFusionComputation()) { continue; } @@ -620,13 +614,14 @@ StatusOr> CpuCompiler::RunBackend( if (embed_ir_in_executable) { ir_module_string = llvm_ir::DumpModuleToString(*llvm_module); } + TF_RETURN_IF_ERROR(VerifyLlvmModule(*llvm_module)); // JIT compile the LLVM IR module to in-memory machine code. jit->AddModule(std::move(llvm_module)); cpu_executable.reset(new ParallelCpuExecutable( std::move(jit), std::move(assignment), std::move(module), std::move(function_names), std::move(aligned_constants), - std::move(hlo_profile_printer), std::move(hlo_profile_index_map))); + std::move(hlo_profile_printer_data), std::move(hlo_profile_index_map))); if (embed_ir_in_executable) { static_cast(*cpu_executable) @@ -654,19 +649,11 @@ StatusOr> CpuCompiler::RunBackend( // print one ourselves. XLA_VLOG_LINES(2, assignment->ToString()); - if (!xla_dump_hlo_proto_to.empty()) { + if (!xla_dump_optimized_hlo_proto_to.empty()) { HloProto proto = MakeHloProto(*module, *assignment); TF_RETURN_IF_ERROR(protobuf_util::DumpProtoToDirectory( - proto, xla_dump_hlo_proto_to, module->name())); + proto, xla_dump_optimized_hlo_proto_to, module->name())); } - // We always profile the entire computation as a whole, even if hlo - // profiling is disabled. When hlo profiling is diabled, we pass in a - // profile counter array of just one element, which corresponds to the whole - // computation. - size_t entry_computation_profile_idx = - hlo_profile_index_map ? hlo_profile_index_map->GetProfileIndexFor( - *module->entry_computation()) - : 0; // Each computation is a single function. Emit all embedded computations // before the entry computation. The order of computations returned from @@ -674,11 +661,12 @@ StatusOr> CpuCompiler::RunBackend( // before a caller computation. IrEmitter ir_emitter(*module, *assignment, llvm_module.get(), - hlo_to_profile_idx, entry_computation_profile_idx, + std::move(instruction_to_profile_idx), + std::move(computation_to_profile_idx), jit->target_machine(), jit->external_constant_pool()); for (auto embedded_computation : - computation->MakeEmbeddedComputationsList()) { + entry_computation->MakeEmbeddedComputationsList()) { if (embedded_computation->IsFusionComputation()) { continue; } @@ -690,25 +678,29 @@ StatusOr> CpuCompiler::RunBackend( &module_sequence.at(embedded_computation)) .status()); } - string function_name_prefix = - computation->name().empty() ? "__compute" : computation->name(); + string function_name_prefix = entry_computation->name().empty() + ? "__compute" + : entry_computation->name(); TF_ASSIGN_OR_RETURN( llvm::Function * entry_function, - ir_emitter.EmitComputation(computation, function_name_prefix, + ir_emitter.EmitComputation(entry_computation, function_name_prefix, /*is_top_level_computation=*/true, - &module_sequence.at(computation))); + &module_sequence.at(entry_computation))); string function_name = llvm_ir::AsString(entry_function->getName()); string ir_module_string; if (embed_ir_in_executable) { ir_module_string = llvm_ir::DumpModuleToString(*llvm_module); } + TF_RETURN_IF_ERROR(VerifyLlvmModule(*llvm_module)); + + XLA_VLOG_LINES(2, "LLVM IR:\n" + llvm_ir::DumpModuleToString(*llvm_module)); // JIT compile the LLVM IR module to in-memory machine code. jit->AddModule(std::move(llvm_module)); cpu_executable.reset(new CpuExecutable( std::move(jit), std::move(assignment), std::move(module), function_name, - std::move(hlo_profile_printer), std::move(hlo_profile_index_map))); + std::move(hlo_profile_printer_data), std::move(hlo_profile_index_map))); if (embed_ir_in_executable) { static_cast(*cpu_executable) @@ -725,7 +717,8 @@ CpuCompiler::CompileAheadOfTime(std::vector> modules, const AotCompilationOptions& aot_options) { TF_RET_CHECK(!modules.empty()); std::call_once(llvm_command_line_options_initialized, - &InitializeLLVMCommandLineOptions, modules[0]->config()); + &llvm_ir::InitializeLLVMCommandLineOptions, + modules[0]->config()); // We can pass just one llvm::TargetOptions when we compile the LLVM module, // so we bail if the configs have conflicting flags. At the moment, the only @@ -835,21 +828,21 @@ CpuCompiler::CompileAheadOfTime(std::vector> modules, // print one ourselves. XLA_VLOG_LINES(2, assignment->ToString()); - const string xla_dump_hlo_proto_to = - module->config().debug_options().xla_dump_hlo_proto_to(); - if (!xla_dump_hlo_proto_to.empty()) { + const string xla_dump_optimized_hlo_proto_to = + module->config().debug_options().xla_dump_optimized_hlo_proto_to(); + if (!xla_dump_optimized_hlo_proto_to.empty()) { HloProto proto = MakeHloProto(*module, *assignment); TF_RETURN_IF_ERROR(protobuf_util::DumpProtoToDirectory( - proto, xla_dump_hlo_proto_to, module->name())); + proto, xla_dump_optimized_hlo_proto_to, module->name())); } - IrEmitter ir_emitter( - *module, *assignment, &llvm_module, - /*hlo_to_profile_idx=*/ - std::unordered_map{}, - /*entry_computation_profile_idx=*/tensorflow::gtl::nullopt, - target_machine.get(), - /*external_constant_pool=*/nullptr); + IrEmitter ir_emitter(*module, *assignment, &llvm_module, + /*instruction_to_profile_idx=*/ + std::unordered_map{}, + /*computation_to_profile_idx=*/ + std::unordered_map{}, + target_machine.get(), + /*external_constant_pool=*/nullptr); HloComputation* computation = module->entry_computation(); for (auto embedded_computation : computation->MakeEmbeddedComputationsList()) { @@ -879,14 +872,23 @@ CpuCompiler::CompileAheadOfTime(std::vector> modules, *module, user_pre_optimization_hook_, user_post_optimization_hook_, &pre_optimization_ir_dump_hook, &post_optimization_ir_dump_hook)); + // Run the LLVM verifier over the unoptimized LLVM IR. If it fails, run the + // pre-optimization IR dump hook before returning. + { + Status verify_status = VerifyLlvmModule(llvm_module); + if (!verify_status.ok() && pre_optimization_ir_dump_hook) { + pre_optimization_ir_dump_hook(llvm_module).IgnoreError(); + } + TF_RETURN_IF_ERROR(verify_status); + } + Disassembler disassembler(*target_machine); CompilerFunctor compiler_functor( target_machine.get(), &disassembler, opt_level, options::OptimizeForSizeRequested(module->config()), module->config().debug_options().xla_enable_fast_math(), module->config().debug_options().xla_llvm_disable_expensive_passes(), - CompilerFunctor::AllIntrinsics(), pre_optimization_ir_dump_hook, - post_optimization_ir_dump_hook); + pre_optimization_ir_dump_hook, post_optimization_ir_dump_hook); llvm::object::OwningBinary object_file = compiler_functor(llvm_module); llvm::StringRef object_file_data_ref = object_file.getBinary()->getData(); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.h b/tensorflow/compiler/xla/service/cpu/cpu_compiler.h index ebed7058d8f7968c6e03ef90d0da6b2325037eb0..3498139ab95d21383c6dc008ae5614b7bfe91148 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.h @@ -118,11 +118,13 @@ class CpuCompiler : public LLVMCompiler { StatusOr> RunHloPasses( std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec) override; + perftools::gputools::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) override; StatusOr> RunBackend( std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec) override; + perftools::gputools::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) override; StatusOr>> CompileAheadOfTime(std::vector> modules, diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc index 028f827337979de14ec557a8f0d7a47f095bf55e..802d0a6fb46890b31d14b1fbf3b2e7d6520caccc 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc @@ -55,9 +55,9 @@ CpuExecutable::CpuExecutable( std::unique_ptr assignment, std::unique_ptr hlo_module, const string& entry_function_name, - std::unique_ptr hlo_profile_printer, + std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map) - : Executable(std::move(hlo_module), std::move(hlo_profile_printer), + : Executable(std::move(hlo_module), std::move(hlo_profile_printer_data), std::move(hlo_profile_index_map)), jit_(std::move(jit)), assignment_(std::move(assignment)) { @@ -147,17 +147,13 @@ Status CpuExecutable::ExecuteComputeFunction( uint64 start_micros = tensorflow::Env::Default()->NowMicros(); - // Allocate profiling counters for each hlo instruction that we would like to - // profile. Even when not Hlo profiling, we allocate a counter for the entire - // computation, which we use to update ExecutionProfile below. - std::vector* profile_counters = nullptr; - std::vector profile_counter_for_entry_computation; - if (hlo_execution_profile) { - profile_counters = hlo_execution_profile->mutable_profile_counters(); - } else { - profile_counters = &profile_counter_for_entry_computation; - profile_counter_for_entry_computation.push_back(0); - } + size_t profile_counters_size = + hlo_execution_profile ? hlo_execution_profile->profile_counters().size() + : 0; + int64* profile_counters = + hlo_execution_profile + ? hlo_execution_profile->mutable_profile_counters()->data() + : nullptr; // Call the computation function following the calling convention. std::vector buffer_pointers; @@ -172,7 +168,7 @@ Status CpuExecutable::ExecuteComputeFunction( VLOG(3) << tensorflow::strings::Printf( " func(void* result, void* params[%zu], void* temps[%zu], " "uint64 profile_counters[%zu])", - args_array.size(), buffer_pointers.size(), profile_counters->size()); + args_array.size(), buffer_pointers.size(), profile_counters_size); VLOG(3) << tensorflow::strings::Printf(" result = %p", result_buffer); auto ptr_printer = [](string* out, const void* p) { tensorflow::strings::StrAppend(out, tensorflow::strings::Printf("%p", p)); @@ -184,11 +180,11 @@ Status CpuExecutable::ExecuteComputeFunction( " temps = [%s]", tensorflow::str_util::Join(buffer_pointers, ", ", ptr_printer).c_str()); VLOG(3) << tensorflow::strings::Printf(" profile_counters = %p", - profile_counters->data()); + profile_counters); } compute_function_(result_buffer, run_options, args_array.data(), - buffer_pointers.data(), profile_counters->data()); + buffer_pointers.data(), profile_counters); uint64 end_micros = tensorflow::Env::Default()->NowMicros(); @@ -196,13 +192,11 @@ Status CpuExecutable::ExecuteComputeFunction( tensorflow::mutex_lock lock(mutex_); const double nanoseconds = (end_micros - start_micros) * 1000.0; execution_profile_.set_compute_time_ns(std::max(nanoseconds, 1.0)); - + // If hlo profiling was disabled then the cycle count is left empty. if (hlo_execution_profile) { execution_profile_.set_compute_cycle_count( hlo_execution_profile->total_cycles_executed( *module().entry_computation())); - } else { - execution_profile_.set_compute_cycle_count(profile_counters->back()); } } diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.h b/tensorflow/compiler/xla/service/cpu/cpu_executable.h index 50443a59954e222f65fc935e83effdaf6d6c8bf0..267b89a10b3c038dc2048f0ad5b5b343c88ef0f9 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.h @@ -51,7 +51,7 @@ class CpuExecutable : public Executable { std::unique_ptr assignment, std::unique_ptr hlo_module, const string& entry_function_name, - std::unique_ptr hlo_profile_printer, + std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map); ~CpuExecutable() override {} diff --git a/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.cc b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.cc new file mode 100644 index 0000000000000000000000000000000000000000..7bd4741a04b1135d9780e0cf765b7b33378526e1 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.cc @@ -0,0 +1,48 @@ +/* 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/compiler/xla/service/cpu/cpu_hlo_support_checker.h" + +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/errors.h" + +namespace xla { + +StatusOr CpuHloSupportChecker::Run(HloModule* module) { + for (auto* computation : module->computations()) { + for (const auto& instruction : computation->instructions()) { + TF_RETURN_IF_ERROR( + ShapeUtil::ValidateShapeWithOptionalLayout(instruction->shape())); + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + instruction->shape(), + [&instruction](const Shape& subshape, const ShapeIndex&) { + if (LayoutUtil::IsSparseArray(subshape)) { + return xla::Unimplemented( + "CPU backend does not support HLO instruction %s with shape " + "containing a sparse layout: %s", + instruction->ToString().c_str(), + ShapeUtil::HumanStringWithLayout(instruction->shape()) + .c_str()); + } + return Status::OK(); + })); + } + } + return false; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.h b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.h new file mode 100644 index 0000000000000000000000000000000000000000..2924b6365943f0a3ec998d7a77767a76cbb576ae --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.h @@ -0,0 +1,42 @@ +/* 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_COMPILER_XLA_SERVICE_CPU_CPU_HLO_SUPPORT_CHECKER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_HLO_SUPPORT_CHECKER_H_ + +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// This pass should run early in the HLO pipeline and checks for HLO constructs +// which are not supported by the CPU backend and cannot be removed via HLO +// transformations (eg, sparse layouts). +class CpuHloSupportChecker : public HloPassInterface { + public: + CpuHloSupportChecker() = default; + ~CpuHloSupportChecker() override = default; + + tensorflow::StringPiece name() const override { + return "cpu_hlo_support_checker"; + } + + // Note: always returns false (no instructions are ever modified by this + // pass). + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_HLO_SUPPORT_CHECKER_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0f463e6de623fc6ab43d685ff2a5d6882ba7b8a2 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker_test.cc @@ -0,0 +1,72 @@ +/* Copyright 2017 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/compiler/xla/service/cpu/cpu_hlo_support_checker.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/core/lib/core/error_codes.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +using ::testing::HasSubstr; + +class CpuHloSupportCheckerTest : public HloTestBase { + protected: + CpuHloSupportChecker& checker() { return checker_; } + + private: + CpuHloSupportChecker checker_; +}; + +TEST_F(CpuHloSupportCheckerTest, Add) { + HloComputation::Builder builder(TestName()); + const Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "param1")); + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape, HloOpcode::kAdd, param0, param1)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + TF_ASSERT_OK(checker().Run(module.get()).status()); +} + +TEST_F(CpuHloSupportCheckerTest, SparseUnimplemented) { + HloComputation::Builder builder(TestName()); + const Shape sparse_shape = ShapeUtil::MakeShapeWithSparseLayout(F32, {10}, 2); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, sparse_shape, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, sparse_shape, "param1")); + builder.AddInstruction(HloInstruction::CreateBinary( + sparse_shape, HloOpcode::kAdd, param0, param1)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + Status status = checker().Run(module.get()).status(); + ASSERT_EQ(status.code(), tensorflow::error::UNIMPLEMENTED); + EXPECT_THAT(status.error_message(), + HasSubstr("CPU backend does not support")); + EXPECT_THAT(status.error_message(), + HasSubstr(ShapeUtil::HumanStringWithLayout(sparse_shape))); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc index 7908dc173d79a4a9dcb6127ac344267e27d2b5f2..40ace963270e8cead47cc731cc326351178dff7d 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc @@ -35,12 +35,17 @@ extern const char* const kEigenMatMulF32SymbolName = "__xla_cpu_runtime_EigenMatMulF32"; extern const char* const kEigenMatMulF64SymbolName = "__xla_cpu_runtime_EigenMatMulF64"; +extern const char* const kEigenConvF16SymbolName = + "__xla_cpu_runtime_EigenConvF16"; extern const char* const kEigenConvF32SymbolName = "__xla_cpu_runtime_EigenConvF32"; +extern const char* const kEigenFftSymbolName = "__xla_cpu_runtime_EigenFft"; extern const char* const kEigenSingleThreadedMatMulF32SymbolName = "__xla_cpu_runtime_EigenSingleThreadedMatMulF32"; extern const char* const kEigenSingleThreadedMatMulF64SymbolName = "__xla_cpu_runtime_EigenSingleThreadedMatMulF64"; +extern const char* const kEigenSingleThreadedConvF16SymbolName = + "__xla_cpu_runtime_EigenSingleThreadedConvF16"; extern const char* const kEigenSingleThreadedConvF32SymbolName = "__xla_cpu_runtime_EigenSingleThreadedConvF32"; extern const char* const kAcquireInfeedBufferForDequeueSymbolName = diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime.h index 2ade455b8a0a43dda8c93bbb79891439da2e4f75..2141dfe1cedd6f9674acc348152574b4fd30895b 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.h @@ -43,9 +43,12 @@ namespace runtime { // because it is a symbol in the cpu_runtime library. extern const char* const kEigenMatMulF32SymbolName; extern const char* const kEigenMatMulF64SymbolName; +extern const char* const kEigenConvF16SymbolName; extern const char* const kEigenConvF32SymbolName; +extern const char* const kEigenFftSymbolName; extern const char* const kEigenSingleThreadedMatMulF32SymbolName; extern const char* const kEigenSingleThreadedMatMulF64SymbolName; +extern const char* const kEigenSingleThreadedConvF16SymbolName; extern const char* const kEigenSingleThreadedConvF32SymbolName; extern const char* const kAcquireInfeedBufferForDequeueSymbolName; extern const char* const kReleaseInfeedBufferAfterDequeueSymbolName; diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h deleted file mode 100644 index 74ae6d00c91be07c0d181ea324e570c73c6b2e77..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h +++ /dev/null @@ -1,53 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ - -// This header declares functions which may be called by the generated code on -// the CPU. Calls to these functions must be resolved explicitly in the JIT in -// xla::cpu::SimpleResolver. It also defines a per-CpuExecutable context -// which is used to cache expensive state and resources utilized by the -// aforementioned functions. - -#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_AVX_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_AVX_H_ - -#include "tensorflow/core/platform/macros.h" - -namespace xla { -namespace cpu { -namespace runtime { - -extern const char *const kExpV8F32AVXSymbolName; -extern const char *const kLogV8F32AVXSymbolName; - -typedef float V8F32AVX __attribute__((__vector_size__(32))); -} // namespace runtime -} // namespace cpu -} // namespace xla - -extern "C" { - -#ifdef __AVX__ -// The following functions are vectorized versions of a selection of libm -// library functions. -// References to these functions are created by the LLVM vectorizer. -xla::cpu::runtime::V8F32AVX __xla_cpu_runtime_ExpV8F32AVX( - xla::cpu::runtime::V8F32AVX x); - -xla::cpu::runtime::V8F32AVX __xla_cpu_runtime_LogV8F32AVX( - xla::cpu::runtime::V8F32AVX x); -#endif -} - -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_AVX_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.cc deleted file mode 100644 index abe792b2787ce8baf56ee62585a0ab886d922a23..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.cc +++ /dev/null @@ -1,46 +0,0 @@ -/* Copyright 2017 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/compiler/xla/service/cpu/cpu_runtime_neon.h" - -#define EIGEN_USE_THREADS - -#include "third_party/eigen3/Eigen/Core" - -#ifdef __ARM_NEON__ - -xla::cpu::runtime::V4F32NEON __xla_cpu_runtime_ExpV4F32NEON( - xla::cpu::runtime::V4F32NEON x) { - return Eigen::internal::pexp(x); -} - -xla::cpu::runtime::V4F32NEON __xla_cpu_runtime_LogV4F32NEON( - xla::cpu::runtime::V4F32NEON x) { - Eigen::internal::Packet4f p = x; - return Eigen::internal::plog(p); -} - -#endif // __ARM_NEON__ - -namespace xla { -namespace cpu { -namespace runtime { - -const char *const kExpV4F32NEONSymbolName = "__xla_cpu_runtime_ExpV4F32NEON"; -const char *const kLogV4F32NEONSymbolName = "__xla_cpu_runtime_LogV4F32NEON"; - -} // namespace runtime -} // namespace cpu -} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h deleted file mode 100644 index 645a43858fb8c3d8e7e94709333c88503b6cc52d..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h +++ /dev/null @@ -1,64 +0,0 @@ -/* Copyright 2017 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_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_NEON_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_NEON_H_ - -// This header declares functions which may be called by the generated code on -// the CPU. Calls to these functions must be resolved explicitly in the JIT in -// xla::cpu::SimpleResolver. - -#include "tensorflow/core/platform/macros.h" - -#ifdef __ARM_NEON__ -// For the other runtimes (AVX, SSE4.1) we define the vector type directly using -// __attribute__((__vector_size__(*))). Unfortunately, the typedef for the ARM -// NEON SIMD types is not portable, so the type has to come from -#include -#endif // __ARM_NEON__ - -namespace xla { -namespace cpu { -namespace runtime { - -extern const char *const kExpV4F32NEONSymbolName; -extern const char *const kLogV4F32NEONSymbolName; - -#ifdef __ARM_NEON__ -typedef float32x4_t V4F32NEON; -#else -// On non-ARM platforms ensure the declaration is present -struct V4F32NEON; -#endif // __ARM_NEON__ - -} // namespace runtime -} // namespace cpu -} // namespace xla - -extern "C" { - -#ifdef __ARM_NEON__ -// The following functions are vectorized versions of a selection of libm -// library functions. -// References to these functions are created by the LLVM vectorizer. -xla::cpu::runtime::V4F32NEON __xla_cpu_runtime_ExpV4F32NEON( - xla::cpu::runtime::V4F32NEON x); - -xla::cpu::runtime::V4F32NEON __xla_cpu_runtime_LogV4F32NEON( - xla::cpu::runtime::V4F32NEON x); -#endif // __ARM_NEON__ -} - -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_NEON_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.cc deleted file mode 100644 index a9a45db5a424d2faecbd437542c41fbd7fdf0bb8..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.cc +++ /dev/null @@ -1,47 +0,0 @@ -/* Copyright 2017 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/compiler/xla/service/cpu/cpu_runtime_sse4_1.h" - -#define EIGEN_USE_THREADS - -#include "third_party/eigen3/Eigen/Core" - -#ifdef __SSE4_1__ - -xla::cpu::runtime::V4F32SSE __xla_cpu_runtime_ExpV4F32SSE( - xla::cpu::runtime::V4F32SSE x) { - Eigen::internal::Packet4f p = x; - return Eigen::internal::pexp(p); -} - -xla::cpu::runtime::V4F32SSE __xla_cpu_runtime_LogV4F32SSE( - xla::cpu::runtime::V4F32SSE x) { - Eigen::internal::Packet4f p = x; - return Eigen::internal::plog(p); -} - -#endif // __SSE4_1__ - -namespace xla { -namespace cpu { -namespace runtime { - -const char *const kExpV4F32SSESymbolName = "__xla_cpu_runtime_ExpV4F32SSE"; -const char *const kLogV4F32SSESymbolName = "__xla_cpu_runtime_LogV4F32SSE"; - -} // namespace runtime -} // namespace cpu -} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h deleted file mode 100644 index 80ca4243a2c284452421a07f92233b705fe48b9d..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h +++ /dev/null @@ -1,54 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ - -// This header declares functions which may be called by the generated code on -// the CPU. Calls to these functions must be resolved explicitly in the JIT in -// xla::cpu::SimpleResolver. It also defines a per-CpuExecutable context -// which is used to cache expensive state and resources utilized by the -// aforementioned functions. - -#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_SSE4_1_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_SSE4_1_H_ - -#include "tensorflow/core/platform/macros.h" - -namespace xla { -namespace cpu { -namespace runtime { - -extern const char *const kExpV4F32SSESymbolName; -extern const char *const kLogV4F32SSESymbolName; - -typedef float V4F32SSE __attribute__((__vector_size__(16))); - -} // namespace runtime -} // namespace cpu -} // namespace xla - -extern "C" { - -#ifdef __SSE4_1__ -// The following functions are vectorized versions of a selection of libm -// library functions. -// References to these functions are created by the LLVM vectorizer. -xla::cpu::runtime::V4F32SSE __xla_cpu_runtime_ExpV4F32SSE( - xla::cpu::runtime::V4F32SSE x); - -xla::cpu::runtime::V4F32SSE __xla_cpu_runtime_LogV4F32SSE( - xla::cpu::runtime::V4F32SSE x); -#endif -} - -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_SSE4_1_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc index b53719fcc260d706eab3d7460c42af4a1b5e775f..f5e61aef534da57ce13d3ee9bbeaeaec31f53d2e 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc @@ -98,7 +98,7 @@ Status CpuTransferManager::TransferLiteralToInfeed(se::StreamExecutor* executor, if (!ShapeUtil::IsTuple(shape)) { int64 size = GetByteSizeRequirement(shape); - return TransferBufferToInfeed(executor, size, literal.InternalData()); + return TransferBufferToInfeed(executor, size, literal.untyped_data()); } if (ShapeUtil::IsNestedTuple(shape)) { @@ -111,20 +111,20 @@ Status CpuTransferManager::TransferLiteralToInfeed(se::StreamExecutor* executor, // enqueue the resulting destination device addresses with the // infeed manager. std::vector buffers; - buffers.reserve(literal.tuple_literals_size()); + buffers.reserve(ShapeUtil::TupleElementCount(shape)); auto cleanup = tensorflow::gtl::MakeCleanup([&buffers]() { for (cpu::runtime::XfeedBuffer* b : buffers) { b->Done(Cancelled("Failed to infeed buffer to device.")); } }); - for (const auto& tuple_element : literal.tuple_literals()) { - const Shape& tuple_element_shape = tuple_element.shape(); + for (int64 i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { + const Shape& tuple_element_shape = ShapeUtil::GetSubshape(shape, {i}); int64 tuple_element_size = GetByteSizeRequirement(tuple_element_shape); TF_ASSIGN_OR_RETURN( cpu::runtime::XfeedBuffer * buffer, TransferBufferToInfeedInternal(executor, tuple_element_size, - tuple_element.InternalData())); + literal.untyped_data({i}))); buffers.push_back(buffer); } @@ -187,14 +187,14 @@ Status CpuTransferManager::TransferLiteralFromOutfeed( literal_shape.element_type(), dimensions)); TF_ASSIGN_OR_RETURN(Shape received_shape, TransferArrayBufferFromOutfeed( - executor, literal->MutableInternalData(), size)); + executor, literal->untyped_data(), size)); TF_RET_CHECK(ShapeUtil::Compatible(received_shape, literal->shape())) << "Shape received from outfeed " << ShapeUtil::HumanString(received_shape) << " did not match the shape that was requested for outfeed: " << ShapeUtil::HumanString(literal_shape); TF_RET_CHECK(size == GetByteSizeRequirement(received_shape)); - *literal->mutable_shape() = received_shape; + *literal->mutable_shape_do_not_use() = received_shape; return Status::OK(); } @@ -217,7 +217,7 @@ Status CpuTransferManager::TransferLiteralFromOutfeed( auto empty = Literal::CreateFromDimensions( tuple_element_shape.element_type(), dimensions); int64 size = GetByteSizeRequirement(tuple_element_shape); - buffer_data.push_back({empty->MutableInternalData(), size}); + buffer_data.push_back({empty->untyped_data(), size}); elements.push_back(std::move(empty)); } @@ -233,7 +233,7 @@ Status CpuTransferManager::TransferLiteralFromOutfeed( GetByteSizeRequirement(received_shape)); for (int64 i = 0; i < literal_shape.tuple_shapes_size(); ++i) { - *elements[i]->mutable_shape() = received_shape.tuple_shapes(i); + *elements[i]->mutable_shape_do_not_use() = received_shape.tuple_shapes(i); } *literal = std::move(*Literal::MakeTupleOwned(std::move(elements))); TF_RET_CHECK(ShapeUtil::Equal(literal->shape(), literal_shape)); diff --git a/tensorflow/compiler/xla/service/cpu/custom_call_target_registry.h b/tensorflow/compiler/xla/service/cpu/custom_call_target_registry.h index 2994642356d55df26c31553ef28dc653503d05be..664125ecc95ca5ac10be4201b9120ddbdb9b9821 100644 --- a/tensorflow/compiler/xla/service/cpu/custom_call_target_registry.h +++ b/tensorflow/compiler/xla/service/cpu/custom_call_target_registry.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ // This file is depended on by kernels that have to build for mobile devices. // For this reason, we avoid relying on TensorFlow and instead only use the @@ -71,4 +71,4 @@ class RegisterCustomCallTarget { } // namespace cpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc index 74f71e5ad575134d78f834a9e63723c22ae49111..cfe7c9c3af0be109ac8a86753e880e2bcbceba41 100644 --- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc @@ -23,10 +23,11 @@ limitations under the License. #include "llvm/IR/Module.h" #include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" +#include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" +#include "tensorflow/compiler/xla/service/cpu/vector_support_library.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" -#include "tensorflow/compiler/xla/service/llvm_ir/vector_support_library.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/util.h" @@ -175,7 +176,7 @@ class ColumnMajorMatrixVectorProductEmitter { } // Load a tile of values from the RHS. For the RHS a "tile" is a contiguous - // sequnce of `count` values, each one broadcasted to the vector width. + // sequence of `count` values, each one broadcasted to the vector width. std::vector LoadRhsTile(llvm::Value* offset, int64 count) { llvm::Value* base_pointer = vsl_.ComputeOffsetPointer(rhs_, offset); std::vector result; @@ -526,7 +527,8 @@ DotOpEmitter::DotOpEmitter( const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, const llvm_ir::IrArray* addend_array, llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder, - const HloModuleConfig& hlo_module_config) + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features) : dot_(dot), transpose_lhs_(transpose_lhs), transpose_rhs_(transpose_rhs), @@ -536,20 +538,22 @@ DotOpEmitter::DotOpEmitter( addend_array_(addend_array), executable_run_options_value_(executable_run_options_value), ir_builder_(ir_builder), - hlo_module_config_(hlo_module_config) {} + hlo_module_config_(hlo_module_config), + target_machine_features_(target_machine_features) {} /* static */ tensorflow::Status DotOpEmitter::EmitDotOperation( const HloInstruction& dot, bool transpose_lhs, bool transpose_rhs, const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, const llvm_ir::IrArray* addend_array, llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder, - const HloModuleConfig& hlo_module_config) { + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features) { PrimitiveType type = target_array.GetShape().element_type(); - TF_RET_CHECK(F32 == type || F64 == type || C64 == type); + TF_RET_CHECK(F16 == type || F32 == type || F64 == type || C64 == type); DotOpEmitter dot_emitter(dot, transpose_lhs, transpose_rhs, target_array, lhs_array, rhs_array, addend_array, executable_run_options_value, ir_builder, - hlo_module_config); + hlo_module_config, target_machine_features); return dot_emitter.Emit(); } @@ -624,10 +628,23 @@ bool DotOpEmitter::EmitLlvmIrDotIfProfitable() { const bool optimize_for_size = options::OptimizeForSizeRequested(hlo_module_config_); + const int target_vector_register_element_size = + target_machine_features_.vector_register_num_elements( + *ir_builder_->GetInsertBlock()->getParent(), primitive_type); + + // We may not always know the vector register size for the target we're + // compiling against, in which case target_vector_register_element_size is 0. + // In these cases we choose a default LLVM IR register size. + const int kUnknownTargetVectorRegisterSize = 4; + const int vector_register_element_size = + target_vector_register_element_size == 0 + ? kUnknownTargetVectorRegisterSize + : target_vector_register_element_size; + if (is_column_major_matrix_vector) { VLOG(2) << "Emitting column major matrix-vector multiply with m = " << m << " and k = " << k; - int64 tile_rows = 8; + int64 tile_rows = vector_register_element_size; int64 tile_cols = tiling_factor; string kernel_name = tensorflow::strings::StrCat( @@ -651,7 +668,7 @@ bool DotOpEmitter::EmitLlvmIrDotIfProfitable() { VLOG(2) << "Emitting row major matrix-vector multiply with m = " << m << " and k = " << k; int64 tile_rows = tiling_factor; - int64 tile_cols = 8; + int64 tile_cols = vector_register_element_size; string kernel_name = tensorflow::strings::StrCat( "row_major_gemv_", PrimitiveType_Name(primitive_type), "_", tile_rows, diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h index 2118965a70872846204974e25555340baca718cf..9d748eb81f7850f3ccdb10f076eecfdc8326c05f 100644 --- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h @@ -18,6 +18,7 @@ limitations under the License. #include "llvm/IR/IRBuilder.h" #include "tensorflow/compiler/xla/service/cpu/cpu_options.h" +#include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" @@ -59,7 +60,8 @@ class DotOpEmitter { const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, const llvm_ir::IrArray* addend_array, llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder, - const HloModuleConfig& hlo_module_config); + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features); private: DotOpEmitter(const HloInstruction& dot, bool transpose_lhs, @@ -69,7 +71,8 @@ class DotOpEmitter { const llvm_ir::IrArray* addend_array, llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder, - const HloModuleConfig& hlo_module_config); + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features); // Emits the IR to perform the dot operation. tensorflow::Status Emit(); @@ -142,6 +145,7 @@ class DotOpEmitter { llvm::Value* executable_run_options_value_; llvm::IRBuilder<>* ir_builder_; const HloModuleConfig& hlo_module_config_; + const TargetMachineFeatures& target_machine_features_; }; } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc index ebd96c4c42759b71b79408c73814605301af03c1..99c5e16db70c6a203b4751c1ed8a106c0dc260e6 100644 --- a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc @@ -33,8 +33,14 @@ StatusOr CpuElementalIrEmitter::EmitFloatUnaryOp( switch (op->opcode()) { case HloOpcode::kTanh: { PrimitiveType element_type = op->shape().element_type(); + bool cast_result_to_fp16 = false; string function_name; switch (element_type) { + case F16: + cast_result_to_fp16 = true; + operand_value = ir_builder_->CreateFPCast(operand_value, + ir_builder_->getFloatTy()); + TF_FALLTHROUGH_INTENDED; case F32: function_name = "tanhf"; break; @@ -44,7 +50,7 @@ StatusOr CpuElementalIrEmitter::EmitFloatUnaryOp( default: return Unimplemented("tanh"); } - // Create function declaration for 'tanhf'. + // Create a function declaration. llvm::Function* function = llvm::cast(module_->getOrInsertFunction( llvm_ir::AsStringRef(function_name), operand_value->getType(), @@ -52,8 +58,12 @@ StatusOr CpuElementalIrEmitter::EmitFloatUnaryOp( function->setCallingConv(llvm::CallingConv::C); function->setDoesNotThrow(); function->setDoesNotAccessMemory(); - // Create instruction to call 'tanhf'. - return ir_builder_->CreateCall(function, operand_value); + // Create an instruction to call the function. + llvm::Value* result = ir_builder_->CreateCall(function, operand_value); + if (cast_result_to_fp16) { + result = ir_builder_->CreateFPCast(result, ir_builder_->getHalfTy()); + } + return result; } default: return ElementalIrEmitter::EmitFloatUnaryOp(op, operand_value); @@ -63,7 +73,13 @@ StatusOr CpuElementalIrEmitter::EmitFloatUnaryOp( StatusOr CpuElementalIrEmitter::EmitAtan2( PrimitiveType prim_type, llvm::Value* lhs, llvm::Value* rhs) const { string function_name; + bool cast_result_to_fp16 = false; switch (prim_type) { + case F16: + cast_result_to_fp16 = true; + lhs = ir_builder_->CreateFPCast(lhs, ir_builder_->getFloatTy()); + rhs = ir_builder_->CreateFPCast(rhs, ir_builder_->getFloatTy()); + TF_FALLTHROUGH_INTENDED; case F32: function_name = "atan2f"; break; @@ -73,7 +89,7 @@ StatusOr CpuElementalIrEmitter::EmitAtan2( default: return Unimplemented("atan2"); } - // Create function declaration for 'atan2'. + // Create a function declaration. llvm::Function* function = llvm::cast(module_->getOrInsertFunction( llvm_ir::AsStringRef(function_name), lhs->getType(), lhs->getType(), @@ -81,8 +97,12 @@ StatusOr CpuElementalIrEmitter::EmitAtan2( function->setCallingConv(llvm::CallingConv::C); function->setDoesNotThrow(); function->setDoesNotAccessMemory(); - // Create instruction to call 'atan2'. - return ir_builder_->CreateCall(function, {lhs, rhs}); + // Create an instruction to call the function. + llvm::Value* result = ir_builder_->CreateCall(function, {lhs, rhs}); + if (cast_result_to_fp16) { + result = ir_builder_->CreateFPCast(result, ir_builder_->getHalfTy()); + } + return result; } llvm_ir::ElementGenerator CpuElementalIrEmitter::MakeElementGenerator( diff --git a/tensorflow/compiler/xla/service/cpu/external_constant_pool.cc b/tensorflow/compiler/xla/service/cpu/external_constant_pool.cc index 7a97021dda82942e2ddc35d25e620b73b1e378bb..7dcc4ca7fa08b478f24065275ffa69725dc51682 100644 --- a/tensorflow/compiler/xla/service/cpu/external_constant_pool.cc +++ b/tensorflow/compiler/xla/service/cpu/external_constant_pool.cc @@ -38,7 +38,7 @@ void ExternalConstantPool::Insert(string name, const Literal& literal, CHECK(raw_pointer != nullptr) << "failed to allocate " << literal_size << " bytes with alignment of " << alignment; - std::memcpy(raw_pointer, literal.InternalData(), literal_size); + std::memcpy(raw_pointer, literal.untyped_data(), literal_size); entries_.emplace(std::move(name), static_cast(raw_pointer)); } diff --git a/tensorflow/compiler/xla/service/cpu/external_constant_pool.h b/tensorflow/compiler/xla/service/cpu/external_constant_pool.h index 9c00d476b1fca6c3174af4ebb62dbbde324fd0ea..8008a56df4dbf16e7b57aee8a344058bb0d5883d 100644 --- a/tensorflow/compiler/xla/service/cpu/external_constant_pool.h +++ b/tensorflow/compiler/xla/service/cpu/external_constant_pool.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_EXTERNAL_CONSTANT_POOL_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_EXTERNAL_CONSTANT_POOL_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_EXTERNAL_CONSTANT_POOL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_EXTERNAL_CONSTANT_POOL_H_ #include @@ -62,4 +62,4 @@ class ExternalConstantPool { } // namespace cpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_EXTERNAL_CONSTANT_POOL_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_EXTERNAL_CONSTANT_POOL_H_ diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index e3671d3243e2efec570cb35113c9c4907c23e2b0..4dffaee87f6b33933b58c8c58478eec918569197 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc @@ -62,6 +62,7 @@ limitations under the License. #include "tensorflow/core/lib/core/bits.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/lib/strings/stringprintf.h" @@ -72,6 +73,7 @@ namespace { using llvm_ir::AsStringRef; using llvm_ir::IrName; using llvm_ir::SetToFirstInsertPoint; +namespace gtl = tensorflow::gtl; } // namespace namespace cpu { @@ -79,16 +81,16 @@ namespace cpu { IrEmitter::IrEmitter( const HloModule& hlo_module, const BufferAssignment& assignment, llvm::Module* llvm_module, - std::unordered_map hlo_to_profile_idx, - tensorflow::gtl::optional entry_computation_profile_idx, + std::unordered_map instruction_to_profile_idx, + std::unordered_map computation_to_profile_idx, llvm::TargetMachine* target_machine, ExternalConstantPool* external_constant_pool) : assignment_(assignment), module_(llvm_module), arch_type_(llvm::Triple(llvm_module->getTargetTriple()).getArch()), ir_builder_(llvm_module->getContext()), - hlo_to_profile_idx_(std::move(hlo_to_profile_idx)), - entry_computation_profile_idx_(std::move(entry_computation_profile_idx)), + instruction_to_profile_idx_(std::move(instruction_to_profile_idx)), + computation_to_profile_idx_(std::move(computation_to_profile_idx)), alias_analysis_(hlo_module, assignment, &llvm_module->getContext()), hlo_module_config_(hlo_module.config()), parallel_cpu_backend_( @@ -120,8 +122,7 @@ StatusOr IrEmitter::EmitComputation( // readcyclecounter if it is unavailable. bool use_rdtscp = arch_type_ == llvm::Triple::ArchType::x86 || arch_type_ == llvm::Triple::ArchType::x86_64; - profiling_state_ = ProfilingState(is_top_level_computation_, use_rdtscp, - GetProfileCountersArgument()); + profiling_state_ = ProfilingState(use_rdtscp, GetProfileCountersArgument()); if (instruction_order == nullptr) { TF_RETURN_IF_ERROR(computation->Accept(this)); } else { @@ -478,7 +479,7 @@ Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) { Status IrEmitter::HandleSort(HloInstruction* sort) { // TODO(b/26783907): Implement sort on CPU. - return Unimplemented("Sort is not supported on CPU (b/26783907)."); + return Unimplemented("Sort is not implemented on CPU."); } Status IrEmitter::HandleTuple(HloInstruction* tuple) { @@ -492,7 +493,7 @@ Status IrEmitter::HandleTuple(HloInstruction* tuple) { } Status IrEmitter::HandleMap(HloInstruction* map) { - tensorflow::gtl::ArraySlice operands(map->operands()); + gtl::ArraySlice operands(map->operands()); HloComputation* function = map->to_apply(); // The called computation should have been emitted previously. llvm::Function* mapped_ir_function = FindOrDie(emitted_functions_, function); @@ -521,7 +522,7 @@ Status IrEmitter::HandleReduceWindow(HloInstruction* reduce_window) { // TODO(b/31410564): Implement dilation for reduce-window. if (window_util::HasDilation(window)) { return Unimplemented( - "Dilation for reduce-window not implemented on CPU. See b/31410564."); + "Dilation for ReduceWindow is not implemented on CPU."); } // The called computation should have been emitted previously. @@ -624,8 +625,7 @@ Status IrEmitter::HandleSelectAndScatter(HloInstruction* select_and_scatter) { // TODO(b/31410564): Implement dilation for select-and-scatter. if (window_util::HasDilation(window)) { return Unimplemented( - "Dilation for select-and-scatter not implemented on CPU. " - "See b/31410564."); + "Dilation for SelectAndScatter is not implemented on CPU. "); } // The select and scatter computations should have been emitted previously. @@ -801,7 +801,7 @@ Status IrEmitter::HandleDot(HloInstruction* dot) { auto rhs = dot->operand(1); TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( /*instruction=*/*dot, /*operands=*/{lhs, rhs}, - /*supported_types=*/{F32, F64, C64})); + /*supported_types=*/{F16, F32, F64, C64})); const DotDimensionNumbers& dnums = dot->dot_dimension_numbers(); if (dnums.lhs_batch_dimensions_size() > 0 || dnums.rhs_batch_dimensions_size() > 0) { @@ -839,7 +839,8 @@ Status IrEmitter::HandleDot(HloInstruction* dot) { return DotOpEmitter::EmitDotOperation( *dot, /*transpose_lhs=*/false, /*transpose_rhs=*/false, target_array, lhs_array, rhs_array, /*addend_array=*/nullptr, - GetExecutableRunOptionsArgument(), &ir_builder_, hlo_module_config_); + GetExecutableRunOptionsArgument(), &ir_builder_, hlo_module_config_, + target_machine_features_); } Status IrEmitter::HandleConvolution(HloInstruction* convolution) { @@ -848,7 +849,7 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) { const auto& window = convolution->window(); TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( /*instruction=*/*convolution, /*operands=*/{lhs, rhs}, - /*supported_types=*/{F32, C64})); + /*supported_types=*/{F16, F32, C64})); const ConvolutionDimensionNumbers& dnums = convolution->convolution_dimension_numbers(); @@ -927,25 +928,30 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) { int64 rhs_col_dilation = one_dim_convolution ? 1 : window.dimensions(1).window_dilation(); - // Args have been computed, make the call. - llvm::Type* float_ptr_type = ir_builder_.getFloatTy()->getPointerTo(); + PrimitiveType primitive_type = lhs->shape().element_type(); + llvm::Type* ir_ptr_type = primitive_type == F16 + ? ir_builder_.getHalfTy()->getPointerTo() + : ir_builder_.getFloatTy()->getPointerTo(); llvm::Type* int64_type = ir_builder_.getInt64Ty(); llvm::Type* int8_ptr_type = ir_builder_.getInt8Ty()->getPointerTo(); llvm::FunctionType* conv_type = llvm::FunctionType::get( ir_builder_.getVoidTy(), - {int8_ptr_type, float_ptr_type, float_ptr_type, float_ptr_type, - int64_type, int64_type, int64_type, int64_type, - int64_type, int64_type, int64_type, int64_type, - int64_type, int64_type, int64_type, int64_type, - int64_type, int64_type, int64_type, int64_type, - int64_type, int64_type, int64_type, int64_type}, + {int8_ptr_type, ir_ptr_type, ir_ptr_type, ir_ptr_type, int64_type, + int64_type, int64_type, int64_type, int64_type, int64_type, + int64_type, int64_type, int64_type, int64_type, int64_type, + int64_type, int64_type, int64_type, int64_type, int64_type, + int64_type, int64_type, int64_type, int64_type}, /*isVarArg=*/false); bool multi_threaded_eigen = hlo_module_config_.debug_options().xla_cpu_multi_thread_eigen(); const char* fn_name = - (multi_threaded_eigen - ? runtime::kEigenConvF32SymbolName - : runtime::kEigenSingleThreadedConvF32SymbolName); + primitive_type == F16 + ? (multi_threaded_eigen + ? runtime::kEigenConvF16SymbolName + : runtime::kEigenSingleThreadedConvF16SymbolName) + : (multi_threaded_eigen + ? runtime::kEigenConvF32SymbolName + : runtime::kEigenSingleThreadedConvF32SymbolName); llvm::Function* conv_func = llvm::cast( module_->getOrInsertFunction(fn_name, conv_type)); conv_func->setCallingConv(llvm::CallingConv::C); @@ -955,9 +961,9 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) { conv_func, { GetExecutableRunOptionsArgument(), ir_builder_.CreateBitCast( - GetEmittedValueFor(convolution), float_ptr_type), - ir_builder_.CreateBitCast(lhs_address, float_ptr_type), - ir_builder_.CreateBitCast(rhs_address, float_ptr_type), + GetEmittedValueFor(convolution), ir_ptr_type), + ir_builder_.CreateBitCast(lhs_address, ir_ptr_type), + ir_builder_.CreateBitCast(rhs_address, ir_ptr_type), ir_builder_.getInt64(input_batch), ir_builder_.getInt64(input_rows), ir_builder_.getInt64(input_cols), @@ -1135,6 +1141,55 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) { }); } +Status IrEmitter::HandleFft(HloInstruction* fft) { + auto operand = fft->operand(0); + TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( + /*instruction=*/*fft, /*operands=*/{operand}, + /*supported_types=*/{F32, C64})); + TF_RET_CHECK(LayoutUtil::IsMonotonicWithDim0Major(operand->shape().layout())); + TF_RET_CHECK(LayoutUtil::IsMonotonicWithDim0Major(fft->shape().layout())); + VLOG(3) << "operand=" << ShapeUtil::HumanStringWithLayout(operand->shape()); + VLOG(3) << "fft=" << ShapeUtil::HumanStringWithLayout(fft->shape()); + + llvm::Value* operand_address = GetEmittedValueFor(operand); + TF_RETURN_IF_ERROR(EmitTargetAddressForOp(fft)); + + const std::vector& fft_length = fft->fft_length(); + int64 input_batch = 1; + for (int i = 0; i < fft->shape().dimensions_size() - fft_length.size(); i++) { + input_batch *= fft->shape().dimensions(i); + } + + // Args have been computed, make the call. + llvm::Type* int8_ptr_type = ir_builder_.getInt8Ty()->getPointerTo(); + llvm::Type* int32_type = ir_builder_.getInt32Ty(); + llvm::Type* int64_type = ir_builder_.getInt64Ty(); + llvm::FunctionType* fft_type = llvm::FunctionType::get( + ir_builder_.getVoidTy(), + {int8_ptr_type, int8_ptr_type, int8_ptr_type, int32_type, int32_type, + int64_type, int64_type, int64_type, int64_type}, + /*isVarArg=*/false); + const char* fn_name = runtime::kEigenFftSymbolName; + llvm::Function* fft_func = llvm::cast( + module_->getOrInsertFunction(fn_name, fft_type)); + fft_func->setCallingConv(llvm::CallingConv::C); + fft_func->setDoesNotThrow(); + fft_func->setOnlyAccessesInaccessibleMemOrArgMem(); + const int fft_rank = fft_length.size(); + ir_builder_.CreateCall( + fft_func, + {GetExecutableRunOptionsArgument(), + ir_builder_.CreateBitCast(GetEmittedValueFor(fft), int8_ptr_type), + ir_builder_.CreateBitCast(operand_address, int8_ptr_type), + ir_builder_.getInt32(fft->fft_type()), ir_builder_.getInt32(fft_rank), + ir_builder_.getInt64(input_batch), + ir_builder_.getInt64(fft_rank > 0 ? fft_length[0] : 0), + ir_builder_.getInt64(fft_rank > 1 ? fft_length[1] : 0), + ir_builder_.getInt64(fft_rank > 2 ? fft_length[2] : 0)}); + + return Status::OK(); +} + Status IrEmitter::HandleCrossReplicaSum(HloInstruction* crs) { if (hlo_module_config_.replica_count() == 1) { // When there is a single replica, a cross replica sum is the identity @@ -1145,8 +1200,7 @@ Status IrEmitter::HandleCrossReplicaSum(HloInstruction* crs) { } // TODO(b/33011107): Support cross replica sum on CPU. - return Unimplemented( - "Cross replica sum is not implemented on CPU. See b/33011107."); + return Unimplemented("CrossReplicaSum is not implemented on CPU."); } // Fills up the free variables in 'index_with_free_var' with values from @@ -1176,205 +1230,6 @@ static llvm_ir::IrArray::Index FillReducedDimensionIndex( return index_with_free_var; } -Status IrEmitter::HandleBatchNormTraining(HloInstruction* batch_norm_training) { - // The output of BatchNormTraining is a tuple of three element: - // - An N-dimensional array containing normalized values. - // - A 1 dimensional array containing the mean value for each feature. - // - A 1 dimensional array containing the variance value for each feature. - HloInstruction* operand = batch_norm_training->operands()[0]; - HloInstruction* scale = batch_norm_training->operands()[1]; - HloInstruction* offset = batch_norm_training->operands()[2]; - float epsilon = batch_norm_training->epsilon(); - int64 feature_index = batch_norm_training->feature_index(); - TF_RET_CHECK(ShapeUtil::IsTuple(batch_norm_training->shape()) && - ShapeUtil::TupleElementCount(batch_norm_training->shape()) == 3); - - const Shape& output_shape = - ShapeUtil::GetTupleElementShape(batch_norm_training->shape(), 0); - const Shape& feature_shape = - ShapeUtil::GetTupleElementShape(batch_norm_training->shape(), 1); - - // Reduce vector of the non-feature dimensions. - std::vector dimensions_to_reduce; - - for (int64 i = 0; i < operand->shape().dimensions_size(); ++i) { - if (i != feature_index) { - dimensions_to_reduce.push_back(i); - } - } - - // Get the second and third allocations in the output tuple, which should be - // used to store the result of mean and variance value calculation. - TF_ASSIGN_OR_RETURN( - const BufferAllocation::Slice slice_mean, - assignment_.GetUniqueSlice(batch_norm_training, /*index=*/{1})); - TF_ASSIGN_OR_RETURN( - const BufferAllocation::Slice slice_var, - assignment_.GetUniqueSlice(batch_norm_training, /*index=*/{2})); - const int feature_count = output_shape.dimensions(feature_index); - const int size_in_elements = ShapeUtil::ElementsIn(output_shape); - TF_RET_CHECK(ShapeUtil::ElementsIn(operand->shape()) == size_in_elements); - const int elements_per_feature = size_in_elements / feature_count; - - llvm::Value* mean = EmitTempBufferPointer(slice_mean, feature_shape); - llvm_ir::IrArray mean_array(mean, feature_shape); - - llvm::Value* var = EmitTempBufferPointer(slice_var, feature_shape); - llvm_ir::IrArray var_array(var, feature_shape); - - // This loop calculates mean and variance for each feature. - // - // In theory this could be swapped by multi-output fusion. We will evaluate - // this when it's ready. - // - // For variance calculation, we use a simplified formula so we can fuse the - // computation into the same loop to calculate mean: Var=E(X^2) - E(X)^2. - TF_RETURN_IF_ERROR( - llvm_ir::LoopEmitter( - [&](const llvm_ir::IrArray::Index& index) { - PrimitiveType element_type = operand->shape().element_type(); - // Used to calculate E(X). - llvm::Value* sum_address = llvm_ir::EmitAllocaAtFunctionEntry( - llvm_ir::PrimitiveTypeToIrType(element_type, module_), - "sum_address", &ir_builder_, - MinimumAlignmentForPrimitiveType(element_type)); - - // Used to calculate E(X^2). - llvm::Value* sum_square_address = - llvm_ir::EmitAllocaAtFunctionEntry( - llvm_ir::PrimitiveTypeToIrType(element_type, module_), - "sum_square_address", &ir_builder_, - MinimumAlignmentForPrimitiveType(element_type)); - - ir_builder_.CreateStore( - llvm::ConstantFP::get(ir_builder_.getFloatTy(), 0.0), - sum_address); - - ir_builder_.CreateStore( - llvm::ConstantFP::get(ir_builder_.getFloatTy(), 0.0), - sum_square_address); - - llvm_ir::ForLoopNest loops(IrName(batch_norm_training, "inner"), - &ir_builder_); - - const llvm_ir::IrArray::Index reduced_dims_index = - loops.AddLoopsForShapeOnDimensions( - operand->shape(), dimensions_to_reduce, "reduction_dim"); - - SetToFirstInsertPoint(loops.GetInnerLoopBodyBasicBlock(), - &ir_builder_); - - llvm_ir::IrArray operand_array(GetIrArrayFor(operand)); - llvm_ir::IrArray::Index input_index = - FillReducedDimensionIndex(reduced_dims_index, index); - llvm::Value* new_value = - operand_array.EmitReadArrayElement(input_index, &ir_builder_); - - llvm::Value* new_value_square = - ir_builder_.CreateFMul(new_value, new_value); - - llvm::Value* current_sum = ir_builder_.CreateLoad(sum_address); - llvm::Value* current_sum_square = - ir_builder_.CreateLoad(sum_square_address); - // Update sum. - ir_builder_.CreateStore( - ir_builder_.CreateFAdd(current_sum, new_value), sum_address); - - // Update sum square. - ir_builder_.CreateStore( - ir_builder_.CreateFAdd(current_sum_square, new_value_square), - sum_square_address); - - SetToFirstInsertPoint(loops.GetOuterLoopExitBasicBlock(), - &ir_builder_); - - llvm::Value* sum = ir_builder_.CreateLoad(sum_address); - llvm::Value* elements_per_feature_value = llvm::ConstantFP::get( - ir_builder_.getFloatTy(), elements_per_feature); - llvm::Value* mean = - ir_builder_.CreateFDiv(sum, elements_per_feature_value); - llvm::Value* mean_square = ir_builder_.CreateFMul(mean, mean); - llvm::Value* sum_square = - ir_builder_.CreateLoad(sum_square_address); - - // Var=E(X^2) - E(X)^2. - llvm::Value* var = ir_builder_.CreateFSub( - ir_builder_.CreateFDiv(sum_square, elements_per_feature_value), - mean_square); - - var_array.EmitWriteArrayElement(index, var, &ir_builder_); - return mean; - }, - mean_array, &ir_builder_) - .EmitLoop(IrName(batch_norm_training, "mean_var"))); - - TF_RETURN_IF_ERROR(EmitTargetAddressForOp(batch_norm_training)); - TF_ASSIGN_OR_RETURN( - const BufferAllocation::Slice slice, - assignment_.GetUniqueSlice(batch_norm_training, /*index=*/{0})); - - llvm::Value* normalized = EmitTempBufferPointer(slice, output_shape); - - llvm_ir::IrArray target_array(normalized, output_shape); - - AddAliasingInformationToIrArray(*batch_norm_training, &target_array); - - TF_RETURN_IF_ERROR( - llvm_ir::LoopEmitter( - [this, mean_array, var_array, epsilon, operand, dimensions_to_reduce, - feature_index, offset, scale](const llvm_ir::IrArray::Index& index) { - // The following logic normalizes the input value, scales and shifts - // it: - // - // normalized = (input - mean) / sqrt(variance + epsilon) - // result = normalized * scale + offset - - // Current index in the feature dimension. - llvm_ir::IrArray::Index feature_index_value(1, - index[feature_index]); - - llvm::Value* mean = mean_array.EmitReadArrayElement( - feature_index_value, &ir_builder_); - llvm::Value* var = var_array.EmitReadArrayElement( - feature_index_value, &ir_builder_); - - llvm_ir::IrArray operand_array(GetIrArrayFor(operand)); - llvm::Value* input = - operand_array.EmitReadArrayElement(index, &ir_builder_); - - llvm::Value* variance_with_epsilon = ir_builder_.CreateFAdd( - var, llvm::ConstantFP::get(ir_builder_.getFloatTy(), epsilon)); - llvm::Function* func_llvm_sqrt = llvm::Intrinsic::getDeclaration( - module_, llvm::Intrinsic::sqrt, {ir_builder_.getFloatTy()}); - llvm::Value* variance_sqrt = - ir_builder_.CreateCall(func_llvm_sqrt, {variance_with_epsilon}); - llvm::Value* normalized = ir_builder_.CreateFDiv( - ir_builder_.CreateFSub(input, mean), variance_sqrt); - llvm_ir::IrArray offset_array(GetIrArrayFor(offset)); - llvm::Value* offset = offset_array.EmitReadArrayElement( - feature_index_value, &ir_builder_); - llvm_ir::IrArray scale_array(GetIrArrayFor(scale)); - llvm::Value* scale = scale_array.EmitReadArrayElement( - feature_index_value, &ir_builder_); - llvm::Value* result = ir_builder_.CreateFAdd( - ir_builder_.CreateFMul(normalized, scale), offset); - - return result; - }, - target_array, &ir_builder_) - .EmitLoop(IrName(batch_norm_training, "normalize"))); - - llvm_ir::EmitTuple(GetIrArrayFor(batch_norm_training), - {normalized, mean, var}, &ir_builder_, module_); - return Status::OK(); -} - -Status IrEmitter::HandleBatchNormGrad(HloInstruction* batch_norm_grad) { - // TODO(b/62843645) Implement BatchNormGrad on CPU backend. - return Unimplemented( - "BatchNormGrad is not implemented on CPU. See b/62843645."); -} - Status IrEmitter::HandleParameter(HloInstruction* parameter) { VLOG(2) << "HandleParameter: " << parameter->ToString(); auto param_number = parameter->parameter_number(); @@ -1420,6 +1275,52 @@ Status IrEmitter::HandleParameter(HloInstruction* parameter) { return Status::OK(); } +// Returns true if the relative order of the unreduced dimensions stays the same +// through the reduce operation. +static bool ReductionPreservesLayout(const HloInstruction& reduce) { + DCHECK_EQ(reduce.opcode(), HloOpcode::kReduce); + + // Maps dimensions that were not reduced from their dimension numbers in the + // source shape to their dimensions numbers in the destination shape. + // + // So if we reduce f32[A,B,C,D] on dimensions 1 and 2, this map contains + // [0->0, 3->1]. + gtl::FlatMap unreduced_dim_map; + + gtl::FlatSet reduced_dims(reduce.dimensions().begin(), + reduce.dimensions().end()); + + const Shape& operand_shape = reduce.operand(0)->shape(); + const Shape& result_shape = reduce.shape(); + + int64 delta = 0; + for (int64 i = 0; i < operand_shape.dimensions_size(); i++) { + if (reduced_dims.count(i)) { + delta++; + } else { + InsertOrDie(&unreduced_dim_map, i, i - delta); + } + } + + // Iterate dimensions minor to major and check that the corresponding + // dimensions in the source and target shapes are equivalent. + int64 result_dim_idx = 0; + for (int64 operand_dim_idx = 0; + operand_dim_idx < operand_shape.dimensions_size(); operand_dim_idx++) { + int64 operand_dim = operand_shape.layout().minor_to_major(operand_dim_idx); + if (!reduced_dims.count(operand_dim)) { + if (FindOrDie(unreduced_dim_map, operand_dim) != + result_shape.layout().minor_to_major(result_dim_idx++)) { + return false; + } + } + } + + CHECK_EQ(result_dim_idx, result_shape.dimensions_size()); + + return true; +} + IrEmitter::ReductionGenerator IrEmitter::MatchReductionGenerator( HloComputation* function, string* failure_reason) const { CHECK_EQ(function->num_parameters(), 2); @@ -1436,7 +1337,7 @@ IrEmitter::ReductionGenerator IrEmitter::MatchReductionGenerator( if (ShapeUtil::ElementIsComplex(root_shape)) { // TODO(b/65408531): Complex add could by done via bitcast to // Complex multiply would be more challenging. We could perhaps use a - // strided load to get all reals in a vector, all imags in a vector, or use + // strided load to get all reals in a vector, all images in a vector, or use // CreateShuffleVector on a bitcast to float x [2N]. *failure_reason = "complex values not supported"; return nullptr; @@ -1583,7 +1484,7 @@ IrEmitter::EmitInnerLoopForVectorizedReduction( const ReductionGenerator& reduction_generator, const llvm_ir::IrArray::Index& output_index, const ShardedVectorType& accumulator_type, HloInstruction* init_value, - HloInstruction* arg, tensorflow::gtl::ArraySlice dimensions, + HloInstruction* arg, gtl::ArraySlice dimensions, unsigned element_alignment) { ShardedVector accumulator; accumulator.reserve(accumulator_type.size()); @@ -1687,8 +1588,12 @@ void IrEmitter::EmitShardedVectorStore( StatusOr IrEmitter::EmitVectorizedReduce( HloInstruction* reduce, HloInstruction* arg, HloInstruction* init_value, - tensorflow::gtl::ArraySlice dimensions, HloComputation* function, + gtl::ArraySlice dimensions, HloComputation* function, string* failure_reason) { + if (!ReductionPreservesLayout(*reduce)) { + return false; + } + ReductionGenerator reduction_generator = MatchReductionGenerator(function, failure_reason); if (!reduction_generator) { @@ -1832,7 +1737,7 @@ StatusOr IrEmitter::EmitVectorizedReduce( Status IrEmitter::HandleReduce(HloInstruction* reduce) { auto arg = reduce->mutable_operand(0); auto init_value = reduce->mutable_operand(1); - tensorflow::gtl::ArraySlice dimensions(reduce->dimensions()); + gtl::ArraySlice dimensions(reduce->dimensions()); HloComputation* function = reduce->to_apply(); if (!options::VectorizedReduceDisabled(hlo_module_config_)) { string vectorization_failure_reason; @@ -1909,12 +1814,12 @@ Status IrEmitter::HandleReduce(HloInstruction* reduce) { Status IrEmitter::HandleSend(HloInstruction* send) { // TODO(b/33942983): Support Send/Recv on CPU. - return Unimplemented("Send is not implemented on CPU. See b/33942983."); + return Unimplemented("Send is not implemented on CPU."); } Status IrEmitter::HandleSendDone(HloInstruction* send_done) { // TODO(b/33942983): Support Send/Recv on CPU. - return Unimplemented("Send-done is not implemented on CPU. See b/33942983."); + return Unimplemented("Send-done is not implemented on CPU."); } Status IrEmitter::HandleSlice(HloInstruction* slice) { @@ -1952,7 +1857,7 @@ Status IrEmitter::HandleSlice(HloInstruction* slice) { // // * Implement the memcpy within the innermost loop. - tensorflow::gtl::FlatSet inner_dims; + gtl::FlatSet inner_dims; for (int64 dim : LayoutUtil::MinorToMajor(layout)) { if (operand->shape().dimensions(dim) != slice->shape().dimensions(dim)) { break; @@ -2079,12 +1984,12 @@ Status IrEmitter::HandleDynamicUpdateSlice( Status IrEmitter::HandleRecv(HloInstruction* recv) { // TODO(b/33942983): Support Send/Recv on CPU. - return Unimplemented("Recv is not implemented on CPU. See b/33942983."); + return Unimplemented("Recv is not implemented on CPU."); } Status IrEmitter::HandleRecvDone(HloInstruction* recv_done) { // TODO(b/33942983): Support Send/Recv on CPU. - return Unimplemented("Recv-done is not implemented on CPU. See b/33942983."); + return Unimplemented("Recv-done is not implemented on CPU."); } Status IrEmitter::HandlePad(HloInstruction* pad) { @@ -2093,10 +1998,10 @@ Status IrEmitter::HandlePad(HloInstruction* pad) { for (auto& padding_dimension : pad->padding_config().dimensions()) { if (padding_dimension.edge_padding_low() < 0 || padding_dimension.edge_padding_high() < 0) { - return Unimplemented( - "Negative padding not supported in the CPU backend (b/34628603); " - "this should have been eliminated at the HLO level: %s", - pad->padding_config().ShortDebugString().c_str()); + return InternalErrorStrCat( + "Encountered negative padding in IrEmitter on CPU. " + "This should have been eliminated at the HLO level. ", + pad->ToString()); } } @@ -2190,7 +2095,7 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) { *root, root->operand(0)->IsRank2Transpose(), root->operand(1)->IsRank2Transpose(), target_array, lhs_array, rhs_array, /*addend_array=*/nullptr, GetExecutableRunOptionsArgument(), - &ir_builder_, hlo_module_config_)); + &ir_builder_, hlo_module_config_, target_machine_features_)); return Status::OK(); } else if (llvm_ir::CanEmitFusedDynamicUpdateSliceInPlace(fusion, assignment_)) { @@ -2238,7 +2143,7 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) { TF_RETURN_IF_ERROR(DotOpEmitter::EmitDotOperation( *dot, /*transpose_lhs=*/false, /*transpose_rhs=*/false, target_array, lhs_array, rhs_array, &addend_array, GetExecutableRunOptionsArgument(), - &ir_builder_, hlo_module_config_)); + &ir_builder_, hlo_module_config_, target_machine_features_)); return Status::OK(); } else { return Unimplemented("Fusion kind not implemented on CPU"); @@ -2280,8 +2185,7 @@ Status IrEmitter::HandleCall(HloInstruction* call) { } Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) { - tensorflow::gtl::ArraySlice operands( - custom_call->operands()); + gtl::ArraySlice operands(custom_call->operands()); tensorflow::StringPiece custom_call_target(custom_call->custom_call_target()); llvm::Type* i8_ptr_type = ir_builder_.getInt8PtrTy(); llvm::AllocaInst* operands_alloca = @@ -2412,8 +2316,7 @@ Status IrEmitter::HandleWhile(HloInstruction* xla_while) { } StatusOr IrEmitter::EmitFastConcatenate( - HloInstruction* concatenate, - tensorflow::gtl::ArraySlice operands, + HloInstruction* concatenate, gtl::ArraySlice operands, string* failure_reason) { if (ShouldEmitParallelLoopFor(*concatenate)) { *failure_reason = @@ -2552,8 +2455,7 @@ void IrEmitter::EmitTransferElements(llvm::Value* target, llvm::Value* source, } Status IrEmitter::HandleConcatenate(HloInstruction* concatenate) { - tensorflow::gtl::ArraySlice operands( - concatenate->operands()); + gtl::ArraySlice operands(concatenate->operands()); string failure_reason; TF_ASSIGN_OR_RETURN( bool successful, @@ -2639,56 +2541,51 @@ Status IrEmitter::FinishVisit(HloInstruction* root) { llvm::Value* root_value = GetEmittedValueFor(root); VLOG(2) << " value: " << llvm_ir::DumpToString(*root_value); - llvm::Value* prof_counter = [&]() { - // For the parallel cpu backend, we record the total for each embedded - // computation callee with its caller kCall HLO. - if (parallel_cpu_backend_ && is_top_level_computation_) { - auto* computation = root->parent(); - auto* entry_computation = computation->parent()->entry_computation(); - if (computation != entry_computation) { - for (HloInstruction* instruction : entry_computation->instructions()) { - if (instruction->opcode() == HloOpcode::kCall && - instruction->to_apply()->root_instruction() == root) { - return GetProfileCounterFor(*instruction); - } + auto record_complete_computation = [&](llvm::Value* prof_counter) { + if (prof_counter) { + profiling_state_.RecordCompleteComputation(&ir_builder_, prof_counter); + } + }; + + // For the parallel cpu backend, we record the total for each embedded + // computation callee with its caller kCall HLO. + if (parallel_cpu_backend_ && is_top_level_computation_) { + auto* computation = root->parent(); + auto* entry_computation = computation->parent()->entry_computation(); + if (computation != entry_computation) { + for (HloInstruction* instruction : entry_computation->instructions()) { + if (instruction->opcode() == HloOpcode::kCall && + instruction->to_apply()->root_instruction() == root) { + record_complete_computation(GetProfileCounterFor(*instruction)); + return Status::OK(); } } } - - // Otherwise we record the total computation cycles in a dedicated slot for - // the entry computation. - return GetProfileCounterForEntryComputation(); - }(); - - if (prof_counter) { - profiling_state_.RecordCompleteComputation(&ir_builder_, prof_counter); } + + // For the entry computation this increment is cumulative of embedded + // computations since it includes cycles spent in computations invoked by + // While, Call etc. + record_complete_computation(GetProfileCounterFor(*root->parent())); return Status::OK(); } -llvm::Value* IrEmitter::GetProfileCounterFor(const HloInstruction& hlo) { - auto it = hlo_to_profile_idx_.find(&hlo); - if (it == hlo_to_profile_idx_.end()) { +template +llvm::Value* IrEmitter::GetProfileCounterCommon( + const T& hlo, + const std::unordered_map& profile_index_map) { + auto it = profile_index_map.find(&hlo); + if (it == profile_index_map.end()) { return nullptr; } - size_t prof_counter_idx = it->second; + int64 prof_counter_idx = it->second; string counter_name = IrName("prof_counter", hlo.name()); return ir_builder_.CreateGEP(GetProfileCountersArgument(), ir_builder_.getInt64(prof_counter_idx), AsStringRef(counter_name)); } -llvm::Value* IrEmitter::GetProfileCounterForEntryComputation() { - if (entry_computation_profile_idx_) { - return ir_builder_.CreateGEP( - GetProfileCountersArgument(), - ir_builder_.getInt64(*entry_computation_profile_idx_), - "prof_counter.computation"); - } - return nullptr; -} - void IrEmitter::ProfilingState::UpdateProfileCounter( llvm::IRBuilder<>* ir_builder, llvm::Value* prof_counter, llvm::Value* cycle_end, llvm::Value* cycle_start) { @@ -2751,8 +2648,7 @@ void IrEmitter::ProfilingState::RecordCycleDelta(llvm::IRBuilder<>* ir_builder, void IrEmitter::ProfilingState::RecordCompleteComputation( llvm::IRBuilder<>* ir_builder, llvm::Value* prof_counter) { - if (is_top_level_computation_ && last_read_cycle_end_ && - first_read_cycle_start_) { + if (last_read_cycle_end_ && first_read_cycle_start_) { UpdateProfileCounter(ir_builder, prof_counter, last_read_cycle_end_, first_read_cycle_start_); } @@ -2760,7 +2656,7 @@ void IrEmitter::ProfilingState::RecordCompleteComputation( Status IrEmitter::Preprocess(HloInstruction* hlo) { VLOG(3) << "Visiting: " << hlo->ToString(); - if (hlo_to_profile_idx_.count(hlo)) { + if (instruction_to_profile_idx_.count(hlo)) { profiling_state_.RecordCycleStart(&ir_builder_, hlo); } return Status::OK(); @@ -2872,7 +2768,7 @@ llvm::Value* IrEmitter::EmitTempBufferPointer( // for a single element_type value, and loads it after call. llvm::Value* IrEmitter::EmitElementFunctionCall( llvm::Function* function, const Shape& return_shape, - tensorflow::gtl::ArraySlice parameter_addresses, + gtl::ArraySlice parameter_addresses, tensorflow::StringPiece name) { llvm::Value* return_value_buffer = EmitArrayFunctionCall( function, return_shape, 1, parameter_addresses, name); @@ -2892,8 +2788,7 @@ llvm::Value* IrEmitter::EmitElementFunctionCall( // temps) // return return_value_buffer -- address of the return value. void IrEmitter::EmitArrayFunctionCallInto( - llvm::Function* function, - tensorflow::gtl::ArraySlice parameter_addresses, + llvm::Function* function, gtl::ArraySlice parameter_addresses, llvm::Value* return_value_buffer, tensorflow::StringPiece name) { ir_builder_.CreateCall( function, GetArrayFunctionCallArguments( @@ -2906,7 +2801,7 @@ void IrEmitter::EmitArrayFunctionCallInto( llvm::Value* IrEmitter::EmitArrayFunctionCall( llvm::Function* function, const Shape& return_shape, int64 element_count, - tensorflow::gtl::ArraySlice parameter_addresses, + gtl::ArraySlice parameter_addresses, tensorflow::StringPiece name) { llvm::Value* elements = llvm::ConstantInt::get(ir_builder_.getInt64Ty(), element_count); @@ -3016,8 +2911,8 @@ Status IrEmitter::EmitMemcpy(const HloInstruction& source, Status IrEmitter::ElementTypesSameAndSupported( const HloInstruction& instruction, - tensorflow::gtl::ArraySlice operands, - tensorflow::gtl::ArraySlice supported_types) { + gtl::ArraySlice operands, + gtl::ArraySlice supported_types) { for (auto operand : operands) { TF_RET_CHECK( ShapeUtil::SameElementType(operands[0]->shape(), operand->shape())); @@ -3061,27 +2956,5 @@ StatusOr IrEmitter::EmitScalarCall( ShapeUtil::MakeShape(return_type, {}), argument_addrs, name); } - -llvm::TargetTransformInfo* TargetMachineFeatures::GetTargetTransformInfoFor( - const llvm::Function& function) { - auto it = target_transform_infos_.find(&function); - if (it == target_transform_infos_.end()) { - // Using a dummy function analysis manager is kind of hacky, but LLVM's - // TargetTransformInfoWrapperPass::getTTI does the same thing. - // - // TODO(sanjoy): Fix this within LLVM by directly exposing - // TargetTransformInfo factories from TargetMachine. - llvm::FunctionAnalysisManager DummyFAM; - llvm::TargetTransformInfo target_transform_info = - target_machine_->getTargetIRAnalysis().run(function, DummyFAM); - auto emplace_result = target_transform_infos_.emplace( - &function, std::move(target_transform_info)); - CHECK(emplace_result.second); - it = emplace_result.first; - } - - return &it->second; -} - } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h index 2341e3ea72ff312f2ca54b9495aff4065b34cd81..509440251497cd7337284c39dae05c5f6c28e7c2 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h @@ -24,7 +24,6 @@ limitations under the License. #include #include "llvm/ADT/Triple.h" -#include "llvm/Analysis/TargetTransformInfo.h" #include "llvm/IR/Function.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Module.h" @@ -33,6 +32,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/cpu/external_constant_pool.h" #include "tensorflow/compiler/xla/service/cpu/ir_function.h" +#include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -52,41 +52,6 @@ limitations under the License. namespace xla { namespace cpu { - -// Wraps an llvm::TargetMachine and parses out some information that feeds into -// code LLVM IR generation decisions. -class TargetMachineFeatures { - public: - TargetMachineFeatures(llvm::TargetMachine* target_machine) - : target_machine_(target_machine) {} - - // Return the vectorization factor, which is the number of bytes of data - // explicitly vectorized routines will try to process at once. - int vectorization_factor_in_bytes() const { - // Ideally this should be a function of the cache line size (which we can - // get from llvm::TargetTransformInfo::getCacheLineSize) of the target - // machine. Guess a value of 128 bytes for now. - return 128; - } - - // Return the size of the largest vector size in bytes. We need to pass in - // "function" since llvm functions can contain annotations for specializing - // them to specific micro-architectures (though currently XLA does not use - // this functionality). - int vector_register_byte_size(const llvm::Function& function) { - llvm::TargetTransformInfo* tti = GetTargetTransformInfoFor(function); - return tti->getRegisterBitWidth(/*Vector=*/true) / 8; - } - - private: - llvm::TargetTransformInfo* GetTargetTransformInfoFor( - const llvm::Function& function); - - tensorflow::gtl::FlatMap - target_transform_infos_; - llvm::TargetMachine* target_machine_; -}; - // This class is the top-level API for the XLA HLO --> LLVM IR compiler. It // implements the DfsHloVisitor interface and emits HLO computations as LLVM IR // functions. @@ -98,20 +63,21 @@ class IrEmitter : public DfsHloVisitorWithDefault { // assignment: a BufferAssignment from which we know which temporary buffers // are used by the HLO nodes. // llvm_module: the LLVM module to emit IR into. - // hlo_to_profile_idx: the mapping from HLO to its index in the profiling - // array. - // entry_computation_profile_idx: the index in the profiling array - // for the entry computation. + // instruction_to_profile_idx: the mapping from HLO instructions to their + // index in the profiling array. + // computation_to_profile_idx: the mapping from HLO computations to their + // index in the profiling array. // external_constant_pool: if non-null, points to an ExternalConstantPool // instance into which the Ir emitter can spill // constants. - IrEmitter( - const HloModule& hlo_module, const BufferAssignment& assignment, - llvm::Module* llvm_module, - std::unordered_map hlo_to_profile_idx, - tensorflow::gtl::optional entry_computation_profile_idx, - llvm::TargetMachine* target_machine, - ExternalConstantPool* external_constant_pool); + IrEmitter(const HloModule& hlo_module, const BufferAssignment& assignment, + llvm::Module* llvm_module, + std::unordered_map + instruction_to_profile_idx, + std::unordered_map + computation_to_profile_idx, + llvm::TargetMachine* target_machine, + ExternalConstantPool* external_constant_pool); ~IrEmitter() override; // Emit and return the given HLO computation as an LLVM IR @@ -158,8 +124,7 @@ class IrEmitter : public DfsHloVisitorWithDefault { Status HandleSelect(HloInstruction* select) override; Status HandleDot(HloInstruction* dot) override; Status HandleConvolution(HloInstruction* convolution) override; - Status HandleBatchNormTraining(HloInstruction* batch_norm_training) override; - Status HandleBatchNormGrad(HloInstruction* batch_norm_grad) override; + Status HandleFft(HloInstruction* fft) override; Status HandleCrossReplicaSum(HloInstruction* crs) override; Status HandleInfeed(HloInstruction* infeed) override; Status HandleOutfeed(HloInstruction* outfeed) override; @@ -194,14 +159,23 @@ class IrEmitter : public DfsHloVisitorWithDefault { // Private helper to initialize an IR function for the computation. void InitializeIrFunction(const string& function_name); - // Convenience function to generate a GEP into the profile counter parameter - // which would correspond to the index for a given HLO. - llvm::Value* GetProfileCounterFor(const HloInstruction& hlo); + template + llvm::Value* GetProfileCounterCommon( + const T& hlo, + const std::unordered_map& profile_index_map); + + // Convenience functions to generate a GEP into the profile counter parameter + // which would correspond to the index for a given HLO instruction or + // computation. + llvm::Value* GetProfileCounterFor(const HloInstruction& instruction) { + return GetProfileCounterCommon(instruction, + instruction_to_profile_idx_); + } - // Convenience function to generate a GEP into the profile counter parameter - // corresponding to the index for the entry computation. Returns nullptr if - // profiling the entry computation is disabled. - llvm::Value* GetProfileCounterForEntryComputation(); + llvm::Value* GetProfileCounterFor(const HloComputation& computation) { + return GetProfileCounterCommon(computation, + computation_to_profile_idx_); + } // Gets the IR Value emitted previously for the given hlo. // @@ -445,9 +419,13 @@ class IrEmitter : public DfsHloVisitorWithDefault { std::unique_ptr compute_function_; llvm::IRBuilder<> ir_builder_; - // Maps HLOs to their index into the profile counter array. - std::unordered_map hlo_to_profile_idx_; - const tensorflow::gtl::optional entry_computation_profile_idx_; + // Maps HLO instructions to their index into the profile counter array. + const std::unordered_map + instruction_to_profile_idx_; + + // Maps HLO computations to their index into the profile counter array. + const std::unordered_map + computation_to_profile_idx_; // Maps HLOs to Values emitted for them. std::unordered_map emitted_value_; @@ -470,15 +448,9 @@ class IrEmitter : public DfsHloVisitorWithDefault { // profiling a computation. class ProfilingState { public: - ProfilingState() - : is_top_level_computation_(false), - use_rdtscp_(false), - prof_counters_(nullptr) {} - ProfilingState(bool is_top_level_computation, bool use_rdtscp, - llvm::Value* prof_counters) - : is_top_level_computation_(is_top_level_computation), - use_rdtscp_(use_rdtscp), - prof_counters_(prof_counters) {} + ProfilingState() : use_rdtscp_(false), prof_counters_(nullptr) {} + ProfilingState(bool use_rdtscp, llvm::Value* prof_counters) + : use_rdtscp_(use_rdtscp), prof_counters_(prof_counters) {} // Record the cycle counter before an HLO executes. void RecordCycleStart(llvm::IRBuilder<>* ir_builder, HloInstruction* hlo); @@ -500,9 +472,6 @@ class IrEmitter : public DfsHloVisitorWithDefault { llvm::Value* cycle_start); private: - // Is this IrEmitter for a top-level computation? - bool is_top_level_computation_; - // Should we use the x86-specific rdtscp or the generic readcyclecounter // intrinsic? bool use_rdtscp_; diff --git a/tensorflow/compiler/xla/service/cpu/ir_function.cc b/tensorflow/compiler/xla/service/cpu/ir_function.cc index ca8c290dd1c4959e42026c3917d37f8fc95a1011..2d6f2f3818a7bd4424aaa7d918ca86abef15c0e9 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_function.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_function.cc @@ -209,9 +209,9 @@ std::vector GetArrayFunctionCallArguments( parameter_addresses[i], ir_builder->getInt8PtrTy(), AsStringRef(tensorflow::strings::StrCat(name, "_parameter_", i, "_address_as_i8ptr"))); - llvm::Value* slot_in_param_adresses = ir_builder->CreateInBoundsGEP( + llvm::Value* slot_in_param_addresses = ir_builder->CreateInBoundsGEP( parameter_addresses_buffer, {ir_builder->getInt64(i)}); - ir_builder->CreateStore(parameter_as_i8ptr, slot_in_param_adresses); + ir_builder->CreateStore(parameter_as_i8ptr, slot_in_param_addresses); } const auto to_int8_ptr = [=](llvm::Value* ptr) { diff --git a/tensorflow/compiler/xla/service/cpu/ir_function.h b/tensorflow/compiler/xla/service/cpu/ir_function.h index 1fd2da4dce23982ed030f3aa8ec604182d0ebab8..557aa4a6bfc2ef70cafca4b226f8d8f15ea01e2b 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_function.h +++ b/tensorflow/compiler/xla/service/cpu/ir_function.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ #include "llvm/IR/Function.h" #include "llvm/IR/IRBuilder.h" @@ -131,4 +131,4 @@ Status EmitCallToParallelForkJoin( } // namespace cpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ diff --git a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc index 0f71258ff074f6e9e3bbbda06681fc8475ccad59..2e5cc96098241415b82f225afc81981f3e1069e0 100644 --- a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc +++ b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc @@ -20,6 +20,8 @@ limitations under the License. #include "llvm/IR/Intrinsics.h" #include "llvm/IR/Verifier.h" #include "llvm/Transforms/Utils/Cloning.h" +#include "tensorflow/compiler/xla/service/cpu/vector_support_library.h" +#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/platform/logging.h" namespace xla { @@ -28,6 +30,10 @@ namespace runtime { const char* const kTanhV4F32SymbolName = "__xla_cpu_runtime_TanhV4F32"; const char* const kTanhV8F32SymbolName = "__xla_cpu_runtime_TanhV8F32"; +const char* const kExpV4F32SymbolName = "__xla_cpu_runtime_ExpV4F32"; +const char* const kExpV8F32SymbolName = "__xla_cpu_runtime_ExpV8F32"; +const char* const kLogV4F32SymbolName = "__xla_cpu_runtime_LogV4F32AVX"; +const char* const kLogV8F32SymbolName = "__xla_cpu_runtime_LogV8F32AVX"; namespace { llvm::Function* EmitVectorF32TanhIfNeeded(llvm::Module* module, @@ -42,62 +48,257 @@ llvm::Function* EmitVectorF32TanhIfNeeded(llvm::Module* module, } llvm::LLVMContext* context = &module->getContext(); - llvm::Type* float_type = llvm::Type::getFloatTy(*context); - llvm::VectorType* vector_type = - llvm::VectorType::get(float_type, vector_width); llvm::BasicBlock* vector_tanh_body = llvm::BasicBlock::Create(*context, "body", vector_tanh_function); llvm::IRBuilder<> ir_builder(vector_tanh_body); - llvm::FastMathFlags fast_math_flags; fast_math_flags.setFast(); ir_builder.setFastMathFlags(fast_math_flags); + VectorSupportLibrary vsl(F32, vector_width, &ir_builder, "tanh_f32"); + llvm::Value* input = &*vector_tanh_function->arg_begin(); - CHECK_EQ(input->getType(), vector_type); + CHECK_EQ(input->getType(), vsl.vector_type()); // This implements the same rational interpolant as implemented in Eigen3. - llvm::Value* input_clamped = llvm_ir::EmitFloatMin( - llvm_ir::EmitFloatMax(input, llvm::ConstantFP::get(vector_type, -9.0), - &ir_builder), - llvm::ConstantFP::get(vector_type, 9.0), &ir_builder); - - std::array numerator_coeffs( - {-2.76076847742355e-16f, 2.00018790482477e-13f, -8.60467152213735e-11f, - 5.12229709037114e-08f, 1.48572235717979e-05f, 6.37261928875436e-04f, - 4.89352455891786e-03f}); - - std::array denominator_coeffs( - {1.19825839466702e-06f, 1.18534705686654e-04f, 2.26843463243900e-03f, - 4.89352518554385e-03f}); - - llvm::Value* input_squared = - ir_builder.CreateFMul(input_clamped, input_clamped); - llvm::Value* numerator = - llvm::ConstantFP::get(vector_type, numerator_coeffs[0]); + llvm::Value* input_clamped = + vsl.Clamp(input, /*low=*/GetIeeeF32(-9.0), /*high=*/GetIeeeF32(9.0)); + + std::array numerator_coeffs{ + -2.76076847742355e-16f, 2.00018790482477e-13f, -8.60467152213735e-11f, + 5.12229709037114e-08f, 1.48572235717979e-05f, 6.37261928875436e-04f, + 4.89352455891786e-03f}; + + std::array denominator_coeffs{ + 1.19825839466702e-06f, 1.18534705686654e-04f, 2.26843463243900e-03f, + 4.89352518554385e-03f}; + + llvm::Value* input_squared = vsl.Mul(input_clamped, input_clamped); + llvm::Value* numerator = vsl.SplatFloat(GetIeeeF32(numerator_coeffs[0])); for (int i = 1; i < numerator_coeffs.size(); i++) { - numerator = ir_builder.CreateFAdd( - ir_builder.CreateFMul(input_squared, numerator), - llvm::ConstantFP::get(vector_type, numerator_coeffs[i])); + numerator = + vsl.MulAdd(input_squared, numerator, GetIeeeF32(numerator_coeffs[i])); } - numerator = ir_builder.CreateFMul(input_clamped, numerator); - llvm::Value* denominator = - llvm::ConstantFP::get(vector_type, denominator_coeffs[0]); + numerator = vsl.Mul(input_clamped, numerator); + + llvm::Value* denominator = vsl.SplatFloat(GetIeeeF32(denominator_coeffs[0])); for (int i = 1; i < denominator_coeffs.size(); i++) { - denominator = ir_builder.CreateFAdd( - ir_builder.CreateFMul(input_squared, denominator), - llvm::ConstantFP::get(vector_type, denominator_coeffs[i])); + denominator = vsl.MulAdd(input_squared, denominator, + GetIeeeF32(denominator_coeffs[i])); } - llvm::Value* result = ir_builder.CreateFDiv(numerator, denominator); + llvm::Value* result = vsl.Div(numerator, denominator); ir_builder.CreateRet(result); DCHECK(!llvm::verifyFunction(*vector_tanh_function)); return vector_tanh_function; } + +llvm::Function* EmitVectorF32ExpIfNeeded(llvm::Module* module, + llvm::StringRef function_name, + int vector_width, + bool enable_fast_math) { + llvm::Function* vector_exp_function = module->getFunction(function_name); + if (vector_exp_function == nullptr) { + // If the function declaration is not present in the module, there can't be + // any calls to resolve. Don't emit the function in this case. + return nullptr; + } + + llvm::LLVMContext* context = &module->getContext(); + + llvm::BasicBlock* vector_exp_body = + llvm::BasicBlock::Create(*context, "body", vector_exp_function); + + llvm::IRBuilder<> ir_builder(vector_exp_body); + llvm::FastMathFlags fast_math_flags; + fast_math_flags.setFast(); + ir_builder.setFastMathFlags(fast_math_flags); + + VectorSupportLibrary vsl(F32, vector_width, &ir_builder, "exp_f32"); + + // This implements the same polynomial approximation as implemented in Eigen3. + + const llvm::APFloat half = GetIeeeF32(0.5); + const llvm::APFloat one = GetIeeeF32(1.0); + + const llvm::APFloat exp_hi = GetIeeeF32(88.3762626647950); + const llvm::APFloat exp_lo = GetIeeeF32(-88.3762626647949); + + const llvm::APFloat cephes_LOG2EF = GetIeeeF32(1.44269504088896341); + const llvm::APFloat cephes_exp_C1 = GetIeeeF32(0.693359375); + const llvm::APFloat cephes_exp_C2 = GetIeeeF32(-2.12194440e-4); + + const llvm::APFloat cephes_exp_p0 = GetIeeeF32(1.9875691500E-4); + const llvm::APFloat cephes_exp_p1 = GetIeeeF32(1.3981999507E-3); + const llvm::APFloat cephes_exp_p2 = GetIeeeF32(8.3334519073E-3); + const llvm::APFloat cephes_exp_p3 = GetIeeeF32(4.1665795894E-2); + const llvm::APFloat cephes_exp_p4 = GetIeeeF32(1.6666665459E-1); + const llvm::APFloat cephes_exp_p5 = GetIeeeF32(5.0000001201E-1); + + llvm::Value* input = &*vector_exp_function->arg_begin(); + llvm::Value* input_clamped = + vsl.Clamp(input, /*low=*/exp_lo, /*high=*/exp_hi); + llvm::Value* fx = vsl.Floor(vsl.MulAdd(input_clamped, cephes_LOG2EF, half)); + llvm::Value* tmp = vsl.Mul(cephes_exp_C1, fx); + llvm::Value* z = vsl.Mul(cephes_exp_C2, fx); + llvm::Value* x = vsl.Sub(input_clamped, tmp); + x = vsl.Sub(x, z); + z = vsl.Mul(x, x); + + llvm::Value* y = vsl.MulAdd(x, cephes_exp_p0, cephes_exp_p1); + y = vsl.MulAdd(y, x, cephes_exp_p2); + y = vsl.MulAdd(y, x, cephes_exp_p3); + y = vsl.MulAdd(y, x, cephes_exp_p4); + y = vsl.MulAdd(y, x, cephes_exp_p5); + y = vsl.MulAdd(y, z, x); + y = vsl.Add(one, y); + + // VectorSupportLibrary (intentionally) can't juggle more than one type at a + // time so drop down to IRBuilder for this bit. + llvm::Value* vector_constant_0x7f = + ir_builder.CreateVectorSplat(vector_width, ir_builder.getInt32(0x7f)); + llvm::Value* vector_constant_23 = + ir_builder.CreateVectorSplat(vector_width, ir_builder.getInt32(23)); + llvm::Type* i32_vector_type = + llvm::VectorType::get(ir_builder.getInt32Ty(), vector_width); + // fx is clamped so we don't have to worry about it being out of range for + // i32. + llvm::Value* emm0 = ir_builder.CreateFPToSI(fx, i32_vector_type); + emm0 = ir_builder.CreateAdd(emm0, vector_constant_0x7f); + emm0 = ir_builder.CreateShl(emm0, vector_constant_23); + llvm::Value* emm0_f32 = ir_builder.CreateBitCast(emm0, vsl.vector_type()); + + llvm::Value* result = vsl.Max(vsl.Mul(y, emm0_f32), input); + + ir_builder.CreateRet(result); + + DCHECK(!llvm::verifyFunction(*vector_exp_function)); + return vector_exp_function; +} + +llvm::Function* EmitVectorF32LogIfNeeded(llvm::Module* module, + llvm::StringRef function_name, + int vector_width, + bool enable_fast_math) { + llvm::Function* vector_log_function = module->getFunction(function_name); + if (vector_log_function == nullptr) { + // If the function declaration is not present in the module, there can't be + // any calls to resolve. Don't emit the function in this case. + return nullptr; + } + + llvm::LLVMContext* context = &module->getContext(); + + llvm::BasicBlock* vector_log_body = + llvm::BasicBlock::Create(*context, "body", vector_log_function); + + llvm::IRBuilder<> ir_builder(vector_log_body); + llvm::FastMathFlags fast_math_flags; + fast_math_flags.setFast(); + ir_builder.setFastMathFlags(fast_math_flags); + + llvm::Value* input = &*vector_log_function->arg_begin(); + VectorSupportLibrary vsl(F32, vector_width, &ir_builder, "log_f32"); + + const llvm::APFloat half = GetIeeeF32(0.5); + const llvm::APFloat one = GetIeeeF32(1.0); + + // This implements the same polynomial approximation as implemented in Eigen3. + // Returns NaN for x < 0, -INF for x = 0 + const llvm::APFloat cephes_SQRTHF = GetIeeeF32(0.707106781186547524); + const llvm::APFloat cephes_log_p0 = GetIeeeF32(7.0376836292E-2); + const llvm::APFloat cephes_log_p1 = GetIeeeF32(-1.1514610310E-1); + const llvm::APFloat cephes_log_p2 = GetIeeeF32(1.1676998740E-1); + const llvm::APFloat cephes_log_p3 = GetIeeeF32(-1.2420140846E-1); + const llvm::APFloat cephes_log_p4 = GetIeeeF32(+1.4249322787E-1); + const llvm::APFloat cephes_log_p5 = GetIeeeF32(-1.6668057665E-1); + const llvm::APFloat cephes_log_p6 = GetIeeeF32(+2.0000714765E-1); + const llvm::APFloat cephes_log_p7 = GetIeeeF32(-2.4999993993E-1); + const llvm::APFloat cephes_log_p8 = GetIeeeF32(+3.3333331174E-1); + const llvm::APFloat cephes_log_q1 = GetIeeeF32(-2.12194440e-4); + const llvm::APFloat cephes_log_q2 = GetIeeeF32(0.693359375); + + // The smallest non denormalized float number. + const llvm::APFloat min_norm_pos = GetIeeeF32FromBitwiseRep(0x00800000); + const llvm::APFloat minus_inf = GetIeeeF32FromBitwiseRep(0xff800000); + const llvm::APFloat inv_mant_mask = GetIeeeF32FromBitwiseRep(~0x7f800000); + + // invalid_mask is set if x is negative or NaN (and therefore output + // must be NaN). + llvm::Value* invalid_mask = vsl.FCmpULEMask(input, vsl.GetZeroVector()); + llvm::Value* iszero_mask = vsl.FCmpEQMask(input, vsl.GetZeroVector()); + + // Cut off denormalized stuff. + input = vsl.Max(min_norm_pos, input); + + // VectorSupportLibrary (intentionally) can't juggle more than one type at a + // time so drop down to IRBuilder for this bit. + llvm::Value* vector_constant_0x7f = + ir_builder.CreateVectorSplat(vector_width, ir_builder.getInt32(0x7f)); + llvm::Value* vector_constant_23 = + ir_builder.CreateVectorSplat(vector_width, ir_builder.getInt32(23)); + llvm::Type* i32_vector_type = + llvm::VectorType::get(ir_builder.getInt32Ty(), vector_width); + + llvm::Value* emm0 = ir_builder.CreateLShr( + ir_builder.CreateBitCast(input, i32_vector_type), vector_constant_23); + + // Keep only the fractional part. + input = vsl.FloatAnd(input, inv_mant_mask); + input = vsl.FloatOr(input, half); + + emm0 = ir_builder.CreateSub(emm0, vector_constant_0x7f); + llvm::Value* e = + vsl.Add(one, ir_builder.CreateSIToFP(emm0, vsl.vector_type())); + + // part2: + // if( x < SQRTHF ) { + // e -= 1; + // x = x + x - 1.0; + // } else { x = x - 1.0; } + llvm::Value* mask = vsl.FCmpOLTMask(input, cephes_SQRTHF); + llvm::Value* tmp = vsl.FloatAnd(input, mask); + input = vsl.Sub(input, one); + e = vsl.Sub(e, vsl.FloatAnd(mask, one)); + input = vsl.Add(input, tmp); + + llvm::Value* x2 = vsl.Mul(input, input); + llvm::Value* x3 = vsl.Mul(x2, input); + + llvm::Value *y, *y1, *y2; + y = vsl.MulAdd(input, cephes_log_p0, cephes_log_p1); + y1 = vsl.MulAdd(input, cephes_log_p3, cephes_log_p4); + y2 = vsl.MulAdd(input, cephes_log_p6, cephes_log_p7); + y = vsl.MulAdd(y, input, cephes_log_p2); + y1 = vsl.MulAdd(y1, input, cephes_log_p5); + y2 = vsl.MulAdd(y2, input, cephes_log_p8); + y = vsl.MulAdd(y, x3, y1); + y = vsl.MulAdd(y, x3, y2); + y = vsl.Mul(y, x3); + + y1 = vsl.Mul(cephes_log_q1, e); + tmp = vsl.Mul(half, x2); + y = vsl.Add(y, y1); + input = vsl.Sub(input, tmp); + y2 = vsl.Mul(cephes_log_q2, e); + input = vsl.Add(input, y); + input = vsl.Add(input, y2); + + // Negative arg will be NAN, 0 will be -INF. + llvm::Value* or_lhs = + vsl.FloatAndNot(iszero_mask, vsl.FloatOr(input, invalid_mask)); + llvm::Value* or_rhs = vsl.FloatAnd(iszero_mask, minus_inf); + llvm::Value* result = vsl.FloatOr(or_lhs, or_rhs); + + ir_builder.CreateRet(result); + + DCHECK(!llvm::verifyFunction(*vector_log_function)); + return vector_log_function; +} } // namespace void RewriteIRRuntimeFunctions(llvm::Module* module, bool enable_fast_math) { @@ -108,11 +309,28 @@ void RewriteIRRuntimeFunctions(llvm::Module* module, bool enable_fast_math) { EmitVectorF32TanhIfNeeded(module, kTanhV8F32SymbolName, /*vector_width=*/8, enable_fast_math); + auto* exp_v4f32 = + EmitVectorF32ExpIfNeeded(module, kExpV4F32SymbolName, + /*vector_width=*/4, enable_fast_math); + auto* exp_v8f32 = + EmitVectorF32ExpIfNeeded(module, kExpV8F32SymbolName, + /*vector_width=*/8, enable_fast_math); + + auto* log_v4f32 = + EmitVectorF32LogIfNeeded(module, kLogV4F32SymbolName, + /*vector_width=*/4, enable_fast_math); + auto* log_v8f32 = + EmitVectorF32LogIfNeeded(module, kLogV8F32SymbolName, + /*vector_width=*/8, enable_fast_math); + // Gather all the call sites, force inline them and then delete the vector // function bodies. + // + // TODO(b/73081976): Should we avoid inlining these intrinsics in some cases? std::vector calls_to_inline; - for (auto* function : {tanh_v4f32, tanh_v8f32}) { + for (auto* function : + {tanh_v4f32, tanh_v8f32, exp_v4f32, exp_v8f32, log_v4f32, log_v8f32}) { if (function != nullptr) { for (auto* user : function->users()) { calls_to_inline.push_back(llvm::cast(user)); @@ -125,7 +343,8 @@ void RewriteIRRuntimeFunctions(llvm::Module* module, bool enable_fast_math) { CHECK(llvm::InlineFunction(call_to_inline, inline_function_info)); } - for (auto* function : {tanh_v4f32, tanh_v8f32}) { + for (auto* function : + {tanh_v4f32, tanh_v8f32, exp_v4f32, exp_v8f32, log_v4f32, log_v8f32}) { if (function != nullptr) { function->eraseFromParent(); } diff --git a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h index 7f31fb98b0d03c16ef40bff9822227e01f6be46b..5553972677512617ccb6ac4f57a4d33400b664e3 100644 --- a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h +++ b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_LLVM_IR_RUNTINE_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_LLVM_IR_RUNTINE_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_LLVM_IR_RUNTIME_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_LLVM_IR_RUNTIME_H_ #include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" @@ -25,6 +25,10 @@ namespace runtime { extern const char* const kTanhV4F32SymbolName; extern const char* const kTanhV8F32SymbolName; +extern const char* const kExpV4F32SymbolName; +extern const char* const kExpV8F32SymbolName; +extern const char* const kLogV4F32SymbolName; +extern const char* const kLogV8F32SymbolName; // The following CPU runtime functions have LLVM-IR only implementations: // @@ -40,4 +44,4 @@ void RewriteIRRuntimeFunctions(llvm::Module* module, bool enable_fast_math); } // namespace cpu } // namespace xla -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_LLVM_IR_RUNTINE_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_LLVM_IR_RUNTIME_H_ diff --git a/tensorflow/compiler/xla/service/cpu/orc_jit_memory_mapper.h b/tensorflow/compiler/xla/service/cpu/orc_jit_memory_mapper.h index 2d29550fd5bd659770cc6300e56b57bf1763e671..f8963841158b71a30aa926e3b2b153c42bf78eb1 100644 --- a/tensorflow/compiler/xla/service/cpu/orc_jit_memory_mapper.h +++ b/tensorflow/compiler/xla/service/cpu/orc_jit_memory_mapper.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_ORC_JIT_MEMORY_MAPPER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_ORC_JIT_MEMORY_MAPPER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_ORC_JIT_MEMORY_MAPPER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_ORC_JIT_MEMORY_MAPPER_H_ #include @@ -53,4 +53,4 @@ class Registrar { } // namespace cpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_ORC_JIT_MEMORY_MAPPER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_ORC_JIT_MEMORY_MAPPER_H_ diff --git a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc index d1b88b27f068962fb86477fcad3e4390b1636c2b..cd997f07890cdc1d9a546ede58cc1d992b6416ae 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc @@ -61,9 +61,9 @@ ParallelCpuExecutable::ParallelCpuExecutable( std::unique_ptr> function_names, std::unordered_map> aligned_constants, - std::unique_ptr hlo_profile_printer, + std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map) - : Executable(std::move(hlo_module), std::move(hlo_profile_printer), + : Executable(std::move(hlo_module), std::move(hlo_profile_printer_data), std::move(hlo_profile_index_map)), jit_(std::move(jit)), assignment_(std::move(assignment)), diff --git a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h index 90ac94ef9288b2e860cb30c47ed44a7b96e4825d..c393e9b8ea39bfb4c605ebba8e2cd29726bc4af9 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h +++ b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h @@ -55,7 +55,7 @@ class ParallelCpuExecutable : public Executable { std::unordered_map> aligned_constants, - std::unique_ptr hlo_profile_printer, + std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map); ~ParallelCpuExecutable() override {} diff --git a/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.h b/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.h index 9335d2818e99eb3588537d80dabddda08c1c020e..ce92e36a944de33b991d97460f0b2e859ad56081 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Value.h" @@ -70,4 +70,4 @@ class ParallelLoopEmitter : public llvm_ir::LoopEmitter { } // namespace cpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ diff --git a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc index 4b44ac8941e222d5954121bbb9654062e41f55d6..deb21bf4ef5895cfdbec5c2449b6ce7b306a7008 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc @@ -126,7 +126,7 @@ int64 ParallelTaskAssignment::GetTargetParallelTaskCount( HloInstruction* instruction) { // Currently, we do not assign parallel tasks to instructions with at least // one of the following properties: - // *) Internal threading (library calls to kConv, kDot, and kCustomCall). + // *) Internal threading (library calls to kConv, kDot, kFft, kCustomCall). // *) Emit custom loops (kSelectAndScatter, FusionKind::kTransposeDot). // *) Tuple-shaped. // TODO(b/27458679) Parallelize instructions which are skipped here. @@ -137,6 +137,7 @@ int64 ParallelTaskAssignment::GetTargetParallelTaskCount( instruction->opcode() == HloOpcode::kSelectAndScatter || instruction->opcode() == HloOpcode::kGetTupleElement || instruction->opcode() == HloOpcode::kBitcast || + instruction->opcode() == HloOpcode::kFft || (instruction->opcode() == HloOpcode::kConvolution && PotentiallyImplementedAsEigenConvolution(*instruction)) || PotentiallyImplementedAsEigenDot(*instruction) || diff --git a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.h b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.h index 5801ec8d270cdaed7f2f65c24987a9ea643edb02..7140dabe516cd7ea9260456e994e8b63b68c60d6 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.h +++ b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_TASK_ASSIGNMENT_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_TASK_ASSIGNMENT_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_TASK_ASSIGNMENT_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_TASK_ASSIGNMENT_H_ #include "tensorflow/compiler/xla/service/hlo_cost_analysis.h" #include "tensorflow/compiler/xla/service/hlo_module.h" @@ -99,4 +99,4 @@ class ParallelTaskAssigner : public HloPassInterface { } // namespace cpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_TASK_ASSIGNMENT_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_PARALLEL_TASK_ASSIGNMENT_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_conv2d.cc b/tensorflow/compiler/xla/service/cpu/runtime_conv2d.cc index c2f64eb27a554d17ebe2a94dba334fe378bd7254..3905e7ff2a14d25813e345399e692f9e0f4bd0af 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_conv2d.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_conv2d.cc @@ -34,7 +34,26 @@ TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenConvF32( int64 lhs_col_dilation, int64 rhs_row_dilation, int64 rhs_col_dilation) { const xla::ExecutableRunOptions* run_options = static_cast(run_options_ptr); - tensorflow::xla::EigenConvF32Impl( + tensorflow::xla::EigenConvImpl( + *run_options->intra_op_thread_pool(), out, lhs, rhs, input_batch, + input_rows, input_cols, input_channels, kernel_rows, kernel_cols, + kernel_channels, kernel_filters, output_rows, output_cols, row_stride, + col_stride, padding_top, padding_bottom, padding_left, padding_right, + lhs_row_dilation, lhs_col_dilation, rhs_row_dilation, rhs_col_dilation); +} + +TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenConvF16( + const void* run_options_ptr, Eigen::half* out, Eigen::half* lhs, + Eigen::half* rhs, int64 input_batch, int64 input_rows, int64 input_cols, + int64 input_channels, int64 kernel_rows, int64 kernel_cols, + int64 kernel_channels, int64 kernel_filters, int64 output_rows, + int64 output_cols, int64 row_stride, int64 col_stride, int64 padding_top, + int64 padding_bottom, int64 padding_left, int64 padding_right, + int64 lhs_row_dilation, int64 lhs_col_dilation, int64 rhs_row_dilation, + int64 rhs_col_dilation) { + const xla::ExecutableRunOptions* run_options = + static_cast(run_options_ptr); + tensorflow::xla::EigenConvImpl( *run_options->intra_op_thread_pool(), out, lhs, rhs, input_batch, input_rows, input_cols, input_channels, kernel_rows, kernel_cols, kernel_channels, kernel_filters, output_rows, output_cols, row_stride, diff --git a/tensorflow/compiler/xla/service/cpu/runtime_conv2d.h b/tensorflow/compiler/xla/service/cpu/runtime_conv2d.h index 05ae094691fd9a7ca83b902145c0750fafdc529a..39e20ed45639040110b99ddb52eb6f6dab26dfaa 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_conv2d.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_conv2d.h @@ -34,6 +34,20 @@ extern void __xla_cpu_runtime_EigenConvF32( tensorflow::int64 lhs_col_dilation, tensorflow::int64 rhs_row_dilation, tensorflow::int64 rhs_col_dilation); +extern void __xla_cpu_runtime_EigenConvF16( + const void* /* xla::ExecutableRunOptions* */ run_options_ptr, + Eigen::half* out, Eigen::half* lhs, Eigen::half* rhs, + tensorflow::int64 input_batch, tensorflow::int64 input_rows, + tensorflow::int64 input_cols, tensorflow::int64 input_channels, + tensorflow::int64 kernel_rows, tensorflow::int64 kernel_cols, + tensorflow::int64 kernel_channels, tensorflow::int64 kernel_filters, + tensorflow::int64 output_rows, tensorflow::int64 output_cols, + tensorflow::int64 row_stride, tensorflow::int64 col_stride, + tensorflow::int64 padding_top, tensorflow::int64 padding_bottom, + tensorflow::int64 padding_left, tensorflow::int64 padding_right, + tensorflow::int64 lhs_row_dilation, tensorflow::int64 lhs_col_dilation, + tensorflow::int64 rhs_row_dilation, tensorflow::int64 rhs_col_dilation); + } // extern "C" #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_CONV2D_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_conv2d_impl.h b/tensorflow/compiler/xla/service/cpu/runtime_conv2d_impl.h index 02f45fee0f1b8cd1125ec6a97f01e0028137bb69..85af63bb032ce33bdd188d6e5bcd78a726d5d9fa 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_conv2d_impl.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_conv2d_impl.h @@ -24,26 +24,27 @@ limitations under the License. namespace tensorflow { namespace xla { -template -void EigenConvF32Impl(const EigenDevice& device, float* out, float* lhs, - float* rhs, int64 input_batch, int64 input_rows, - int64 input_cols, int64 input_channels, int64 kernel_rows, - int64 kernel_cols, int64 kernel_channels, - int64 kernel_filters, int64 output_rows, - int64 output_cols, int64 row_stride, int64 col_stride, - int64 padding_top, int64 padding_bottom, - int64 padding_left, int64 padding_right, - int64 lhs_row_dilation, int64 lhs_col_dilation, - int64 rhs_row_dilation, int64 rhs_col_dilation) { - const Eigen::TensorMap, +template +void EigenConvImpl(const EigenDevice& device, ScalarType* out, ScalarType* lhs, + ScalarType* rhs, int64 input_batch, int64 input_rows, + int64 input_cols, int64 input_channels, int64 kernel_rows, + int64 kernel_cols, int64 kernel_channels, + int64 kernel_filters, int64 output_rows, int64 output_cols, + int64 row_stride, int64 col_stride, int64 padding_top, + int64 padding_bottom, int64 padding_left, + int64 padding_right, int64 lhs_row_dilation, + int64 lhs_col_dilation, int64 rhs_row_dilation, + int64 rhs_col_dilation) { + const Eigen::TensorMap, Eigen::Aligned> input(lhs, input_batch, input_rows, input_cols, input_channels); - const Eigen::TensorMap, + const Eigen::TensorMap, Eigen::Aligned> kernel(rhs, kernel_rows, kernel_cols, kernel_channels, kernel_filters); - Eigen::TensorMap, Eigen::Aligned> + Eigen::TensorMap, + Eigen::Aligned> output(out, input_batch, output_rows, output_cols, kernel_filters); Eigen::array, 1> contract_dims; @@ -75,7 +76,7 @@ void EigenConvF32Impl(const EigenDevice& device, float* out, float* lhs, row_stride, rhs_col_dilation, rhs_row_dilation, lhs_col_dilation, lhs_row_dilation, padding_left, padding_right, padding_top, - padding_bottom, 0.0f) + padding_bottom, static_cast(0.0f)) .reshape(pre_contract_dims) .contract(kernel.reshape(kernel_dims), contract_dims) .reshape(post_contract_dims); diff --git a/tensorflow/compiler/xla/service/cpu/runtime_fft.cc b/tensorflow/compiler/xla/service/cpu/runtime_fft.cc new file mode 100644 index 0000000000000000000000000000000000000000..848d2d22414e8fc9bca82de90f7676011d8992fd --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/runtime_fft.cc @@ -0,0 +1,37 @@ +/* Copyright 2017 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/compiler/xla/service/cpu/runtime_fft.h" + +#define EIGEN_USE_THREADS + +#include "tensorflow/compiler/xla/executable_run_options.h" +#include "tensorflow/compiler/xla/service/cpu/runtime_fft_impl.h" +#include "tensorflow/core/platform/dynamic_annotations.h" +#include "tensorflow/core/platform/types.h" + +using tensorflow::int32; +using tensorflow::int64; + +TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenFft( + const void* run_options_ptr, void* out, void* operand, int32 fft_type, + int32 fft_rank, int64 input_batch, int64 fft_length0, int64 fft_length1, + int64 fft_length2) { + const xla::ExecutableRunOptions* run_options = + static_cast(run_options_ptr); + tensorflow::xla::EigenFftImpl(*run_options->intra_op_thread_pool(), out, + operand, fft_type, fft_rank, input_batch, + fft_length0, fft_length1, fft_length2); +} diff --git a/tensorflow/compiler/xla/service/cpu/runtime_fft.h b/tensorflow/compiler/xla/service/cpu/runtime_fft.h new file mode 100644 index 0000000000000000000000000000000000000000..f20c5aa0aa2dcbc700f47c718e75baae18650d1a --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/runtime_fft.h @@ -0,0 +1,31 @@ +/* Copyright 2017 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_COMPILER_XLA_SERVICE_CPU_RUNTIME_FFT_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FFT_H_ + +#include "tensorflow/core/platform/types.h" + +extern "C" { + +extern void __xla_cpu_runtime_EigenFft( + const void* /* xla::ExecutableRunOptions* */ run_options_ptr, void* out, + void* operand, tensorflow::int32 fft_type, tensorflow::int32 fft_rank, + tensorflow::int64 input_batch, tensorflow::int64 fft_length0, + tensorflow::int64 fft_length1, tensorflow::int64 fft_length2); + +} // extern "C" + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FFT_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_fft_impl.h b/tensorflow/compiler/xla/service/cpu/runtime_fft_impl.h new file mode 100644 index 0000000000000000000000000000000000000000..984cb0616e02475babad7160d0f43bb23de0b50e --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/runtime_fft_impl.h @@ -0,0 +1,240 @@ +/* Copyright 2017 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_COMPILER_XLA_SERVICE_CPU_RUNTIME_FFT_IMPL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FFT_IMPL_H_ + +#include + +#include "third_party/eigen3/Eigen/Core" +#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/framework/numeric_types.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/platform/types.h" + +// 'tensorflow' namespace is used so that int64 and other types don't require +// qualification. +namespace tensorflow { +namespace xla { + +namespace internal { + +// Computes either a forward or reverse complex-to-complex FFT. +template +void EigenFftC2C(const EigenDevice& device, complex64* out, complex64* operand, + int64 input_batch, int64 fft_length0, int64 fft_length1, + int64 fft_length2) { + // Create the axes (which are always trailing). + const auto axes = Eigen::ArrayXi::LinSpaced(FFTRank, 1, FFTRank); + constexpr auto direction = Forward ? Eigen::FFT_FORWARD : Eigen::FFT_REVERSE; + + const std::array fft_shape = { + {fft_length0, fft_length1, fft_length2}}; + + Eigen::DSizes dims; + dims[0] = input_batch; + for (int i = 0; i < FFTRank; i++) { + dims[i + 1] = fft_shape[i]; + } + const Eigen::TensorMap, + Eigen::Aligned> + input(operand, dims); + Eigen::TensorMap, + Eigen::Aligned> + output(out, dims); + output.device(device) = input.template fft(axes); +} + +// Computes a forward real->complex FFT, slicing out redundant negative +// frequencies from the innermost dimension. +template +void EigenFftR2C(const EigenDevice& device, complex64* out, float* operand, + int64 input_batch, int64 fft_length0, int64 fft_length1, + int64 fft_length2) { + const std::array fft_shape = { + {fft_length0, fft_length1, fft_length2}}; + + Eigen::DSizes in_dims; + in_dims[0] = input_batch; + Eigen::DSizes out_dims; + out_dims[0] = input_batch; + TensorShape temp_shape{input_batch}; + for (int i = 0; i < FFTRank; i++) { + in_dims[i + 1] = fft_shape[i]; + out_dims[i + 1] = i == FFTRank - 1 ? fft_shape[i] / 2 + 1 : fft_shape[i]; + temp_shape.AddDim(fft_shape[i]); + } + const Eigen::TensorMap, + Eigen::Aligned> + input(operand, in_dims); + Eigen::TensorMap, + Eigen::Aligned> + output(out, out_dims); + + // Create the axes (which are always trailing). + const auto axes = Eigen::ArrayXi::LinSpaced(FFTRank, 1, FFTRank); + + // Compute the full FFT using a temporary tensor. + Tensor temp(DataTypeToEnum::v(), temp_shape); + auto full_fft = temp.flat_inner_dims(); + const Eigen::DSizes zero_start_indices; + full_fft.device(device) = + input.template fft(axes); + + // Slice away the negative frequency components. + output.device(device) = full_fft.slice(zero_start_indices, out_dims); +} + +// Computes a reverse complex->real FFT, reconstructing redundant negative +// frequencies using reverse conjugate on innermost dimension after doing IFFT +// on outer dimensions. +template +void EigenFftC2R(const EigenDevice& device, float* out, complex64* operand, + int64 input_batch, int64 fft_length0, int64 fft_length1, + int64 fft_length2) { + const std::array fft_shape = { + {fft_length0, fft_length1, fft_length2}}; + + Eigen::DSizes in_dims; + in_dims[0] = input_batch; + Eigen::DSizes out_dims; + out_dims[0] = input_batch; + TensorShape temp_shape{input_batch}; + for (int i = 0; i < FFTRank; i++) { + in_dims[i + 1] = i == FFTRank - 1 ? fft_shape[i] / 2 + 1 : fft_shape[i]; + out_dims[i + 1] = fft_shape[i]; + temp_shape.AddDim(fft_shape[i]); + } + const Eigen::TensorMap, + Eigen::Aligned> + input(operand, in_dims); + Eigen::TensorMap, + Eigen::Aligned> + output(out, out_dims); + + // Calculate the shape of the temporary tensor for the full FFT and the + // region we will slice from input given fft_shape. We slice input to + // fft_shape on its inner-most dimensions, except the last (which we + // slice to fft_shape[-1] / 2 + 1). + Tensor temp(DataTypeToEnum::v(), temp_shape); + auto full_fft = temp.flat_inner_dims(); + + // Calculate the starting point and range of the source of + // negative frequency part. + auto neg_sizes = in_dims; + neg_sizes[FFTRank] = fft_shape[FFTRank - 1] - in_dims[FFTRank]; + Eigen::DSizes neg_target_indices; + neg_target_indices[FFTRank] = in_dims[FFTRank]; + + const Eigen::DSizes zero_start_indices; + Eigen::DSizes neg_start_indices; + neg_start_indices[FFTRank] = 1; + + full_fft.slice(zero_start_indices, in_dims).device(device) = input; + + // First, conduct IFFTs on outer dimensions. We save computation (and + // avoid touching uninitialized memory) by slicing full_fft to the + // subregion we wrote input to. + if (FFTRank > 1) { + const auto outer_axes = + Eigen::ArrayXi::LinSpaced(FFTRank - 1, 1, FFTRank - 1); + full_fft.slice(zero_start_indices, in_dims).device(device) = + full_fft.slice(zero_start_indices, in_dims) + .template fft(outer_axes); + } + + // Reconstruct the full FFT by appending reversed and conjugated + // spectrum as the negative frequency part. + Eigen::array reverse_last_axis; + for (auto i = 0; i <= FFTRank; i++) { + reverse_last_axis[i] = i == FFTRank; + } + + if (neg_sizes[FFTRank] != 0) { + full_fft.slice(neg_target_indices, neg_sizes).device(device) = + full_fft.slice(neg_start_indices, neg_sizes) + .reverse(reverse_last_axis) + .conjugate(); + } + + auto inner_axis = Eigen::array{FFTRank}; + output.device(device) = + full_fft.template fft(inner_axis); +} + +template +void EigenFftWithRank(const EigenDevice& device, void* out, void* operand, + int32 fft_type, int64 input_batch, int64 fft_length0, + int64 fft_length1, int64 fft_length2) { + CHECK(::xla::FftType_IsValid(fft_type)) << fft_type; + switch (fft_type) { + case ::xla::FftType::FFT: + EigenFftC2C( + device, static_cast(out), + static_cast(operand), input_batch, fft_length0, + fft_length1, fft_length2); + break; + case ::xla::FftType::IFFT: + EigenFftC2C( + device, static_cast(out), + static_cast(operand), input_batch, fft_length0, + fft_length1, fft_length2); + break; + case ::xla::FftType::RFFT: + EigenFftR2C( + device, static_cast(out), static_cast(operand), + input_batch, fft_length0, fft_length1, fft_length2); + break; + case ::xla::FftType::IRFFT: + EigenFftC2R( + device, static_cast(out), static_cast(operand), + input_batch, fft_length0, fft_length1, fft_length2); + break; + default: + LOG(FATAL) << "Unsupported FFT type: " << fft_type; + } +} + +} // namespace internal + +template +void EigenFftImpl(const EigenDevice& device, void* out, void* operand, + int32 fft_type, int32 fft_rank, int64 input_batch, + int64 fft_length0, int64 fft_length1, int64 fft_length2) { + switch (fft_rank) { + case 1: + internal::EigenFftWithRank<1, EigenDevice>( + device, out, operand, fft_type, input_batch, fft_length0, 0, 0); + break; + case 2: + internal::EigenFftWithRank<2, EigenDevice>(device, out, operand, fft_type, + input_batch, fft_length0, + fft_length1, 0); + break; + case 3: + internal::EigenFftWithRank<3, EigenDevice>(device, out, operand, fft_type, + input_batch, fft_length0, + fft_length1, fft_length2); + break; + default: + LOG(FATAL) << "Unsupported FFT rank " << fft_rank; + } +} + +} // namespace xla +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FFT_IMPL_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_fork_join.h b/tensorflow/compiler/xla/service/cpu/runtime_fork_join.h index fcf1cc62078d3847435a2e75e3ca9d109cf8b200..1cf0ec6e3df400e35fa4e755a0b25b4ce7966e8f 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_fork_join.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_fork_join.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FORK_JOIN_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FORK_JOIN_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FORK_JOIN_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FORK_JOIN_H_ #include "tensorflow/core/platform/types.h" @@ -30,4 +30,4 @@ extern void __xla_cpu_runtime_ParallelForkJoin( } // extern "C" -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FORK_JOIN_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_FORK_JOIN_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matvec.h b/tensorflow/compiler/xla/service/cpu/runtime_matvec.h index cb7e0a81f09e2702de565012e1fcac8b7cd841ab..1bd8dfb377acc1f7cfbe9a92773f87f0ef25de3a 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_matvec.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_matvec.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ #include "tensorflow/core/platform/types.h" @@ -42,4 +42,4 @@ void EigenMatVecF64(double* out, double* lhs, double* rhs, tensorflow::int64 m, } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.cc b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.cc index d0b0e11ac0f9fd06e384c2bb5e6296edd0825f5c..5afccc6a86e2df468e3e3e874cf0f4d4e1342a88 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.cc @@ -21,6 +21,24 @@ limitations under the License. using tensorflow::int64; +TF_ATTRIBUTE_NO_SANITIZE_MEMORY void +__xla_cpu_runtime_EigenSingleThreadedConvF16( + const void* run_options_ptr, Eigen::half* out, Eigen::half* lhs, + Eigen::half* rhs, int64 input_batch, int64 input_rows, int64 input_cols, + int64 input_channels, int64 kernel_rows, int64 kernel_cols, + int64 kernel_channels, int64 kernel_filters, int64 output_rows, + int64 output_cols, int64 row_stride, int64 col_stride, int64 padding_top, + int64 padding_bottom, int64 padding_left, int64 padding_right, + int64 lhs_row_dilation, int64 lhs_col_dilation, int64 rhs_row_dilation, + int64 rhs_col_dilation) { + tensorflow::xla::EigenConvImpl( + Eigen::DefaultDevice(), out, lhs, rhs, input_batch, input_rows, + input_cols, input_channels, kernel_rows, kernel_cols, kernel_channels, + kernel_filters, output_rows, output_cols, row_stride, col_stride, + padding_top, padding_bottom, padding_left, padding_right, + lhs_row_dilation, lhs_col_dilation, rhs_row_dilation, rhs_col_dilation); +} + TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenSingleThreadedConvF32( const void* run_options_ptr, float* out, float* lhs, float* rhs, @@ -30,7 +48,7 @@ __xla_cpu_runtime_EigenSingleThreadedConvF32( int64 row_stride, int64 col_stride, int64 padding_top, int64 padding_bottom, int64 padding_left, int64 padding_right, int64 lhs_row_dilation, int64 lhs_col_dilation, int64 rhs_row_dilation, int64 rhs_col_dilation) { - tensorflow::xla::EigenConvF32Impl( + tensorflow::xla::EigenConvImpl( Eigen::DefaultDevice(), out, lhs, rhs, input_batch, input_rows, input_cols, input_channels, kernel_rows, kernel_cols, kernel_channels, kernel_filters, output_rows, output_cols, row_stride, col_stride, diff --git a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.h b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.h index 8ae1a42149bde26ca2f510ad47e76ae47f34a977..f216bd0152aa93b8753d881938c63a9cabea899b 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.h @@ -20,6 +20,20 @@ limitations under the License. extern "C" { +extern void __xla_cpu_runtime_EigenSingleThreadedConvF16( + const void* /* xla::ExecutableRunOptions* */ run_options_ptr, + Eigen::half* out, Eigen::half* lhs, Eigen::half* rhs, + tensorflow::int64 input_batch, tensorflow::int64 input_rows, + tensorflow::int64 input_cols, tensorflow::int64 input_channels, + tensorflow::int64 kernel_rows, tensorflow::int64 kernel_cols, + tensorflow::int64 kernel_channels, tensorflow::int64 kernel_filters, + tensorflow::int64 output_rows, tensorflow::int64 output_cols, + tensorflow::int64 row_stride, tensorflow::int64 col_stride, + tensorflow::int64 padding_top, tensorflow::int64 padding_bottom, + tensorflow::int64 padding_left, tensorflow::int64 padding_right, + tensorflow::int64 lhs_row_dilation, tensorflow::int64 lhs_col_dilation, + tensorflow::int64 rhs_row_dilation, tensorflow::int64 rhs_col_dilation); + extern void __xla_cpu_runtime_EigenSingleThreadedConvF32( const void* /* xla::ExecutableRunOptions* */ run_options_ptr, float* out, float* lhs, float* rhs, tensorflow::int64 input_batch, diff --git a/tensorflow/compiler/xla/service/cpu/shape_partition.h b/tensorflow/compiler/xla/service/cpu/shape_partition.h index 7a2d00421cfdc8e41ec48698a16665621de16bda..33d02b70e61e3311c9af934e80874939fbe3adae 100644 --- a/tensorflow/compiler/xla/service/cpu/shape_partition.h +++ b/tensorflow/compiler/xla/service/cpu/shape_partition.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ #include @@ -102,4 +102,4 @@ class ShapePartitionIterator { } // namespace cpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc index c942cd6bf12c58873d5195f7454249763e639f91..cfed551eedf9c429752c52bb4411ce1200845d11 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc @@ -15,29 +15,28 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/simple_orc_jit.h" -#include #include #include #include #include #include "llvm/ExecutionEngine/ExecutionEngine.h" +#include "llvm/ExecutionEngine/JITSymbol.h" #include "llvm/ExecutionEngine/SectionMemoryManager.h" #include "llvm/IR/Mangler.h" #include "llvm/Support/CodeGen.h" #include "llvm/Support/Host.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" -#include "tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h" -#include "tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h" -#include "tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h" #include "tensorflow/compiler/xla/service/cpu/custom_call_target_registry.h" #include "tensorflow/compiler/xla/service/cpu/orc_jit_memory_mapper.h" #include "tensorflow/compiler/xla/service/cpu/runtime_conv2d.h" +#include "tensorflow/compiler/xla/service/cpu/runtime_fft.h" #include "tensorflow/compiler/xla/service/cpu/runtime_fork_join.h" #include "tensorflow/compiler/xla/service/cpu/runtime_matmul.h" #include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_conv2d.h" #include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h" +#include "tensorflow/compiler/xla/service/cpu/windows_compatibility.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/logging.h" @@ -46,7 +45,7 @@ namespace cpu { namespace { // A simple SymbolResolver that delegates to the host dynamic linker. -class SimpleResolver : public llvm::JITSymbolResolver { +class SimpleResolver : public llvm::LegacyJITSymbolResolver { public: explicit SimpleResolver(ExternalConstantPool* external_constant_pool) : external_constant_pool_(external_constant_pool) {} @@ -99,27 +98,6 @@ llvm::StringRef GetHostCpuName() { cpu_name.consume_back("-avx512"); return cpu_name; } - -CompilerFunctor::VectorIntrinsics GetAvailableIntrinsics() { - CompilerFunctor::VectorIntrinsics intrinsics; -#ifdef __SSE4_1__ - intrinsics.sse_intrinsics = true; -#else - intrinsics.sse_intrinsics = false; -#endif -#ifdef __AVX__ - intrinsics.avx_intrinsics = true; -#else - intrinsics.avx_intrinsics = false; -#endif -#ifdef __ARM_NEON__ - intrinsics.neon_intrinsics = true; -#else - intrinsics.neon_intrinsics = false; -#endif - return intrinsics; -} - } // namespace SimpleOrcJIT::SimpleOrcJIT(const llvm::TargetOptions& target_options, @@ -138,34 +116,59 @@ SimpleOrcJIT::SimpleOrcJIT(const llvm::TargetOptions& target_options, /*MAttrs=*/DetectMachineAttributes()))), disassembler_(*target_machine_), data_layout_(target_machine_->createDataLayout()), - object_layer_([] { - return std::make_shared( - orc_jit_memory_mapper::GetInstance()); - }), - compile_layer_( - object_layer_, - CompilerFunctor(target_machine_.get(), &disassembler_, opt_level, - optimize_for_size, enable_fast_math, - disable_expensive_passes, GetAvailableIntrinsics(), - std::move(pre_optimization_hook), - std::move(post_optimization_hook))) { + execution_session_(string_pool_), + symbol_resolver_(llvm::orc::createLegacyLookupResolver( + [this](const std::string& name) -> llvm::JITSymbol { + if (const uint8* from_constant_pool = + external_constant_pool_.Find(string(name))) { + return llvm::JITEvaluatedSymbol( + reinterpret_cast(from_constant_pool), + llvm::JITSymbolFlags::None); + } + + void* func_addr = CustomCallTargetRegistry::Global()->Lookup(name); + if (func_addr == nullptr) { + return nullptr; + } + llvm::JITEvaluatedSymbol symbol_info( + reinterpret_cast(func_addr), + llvm::JITSymbolFlags::None); + return symbol_info; + }, + [](llvm::Error Err) { + cantFail(std::move(Err), "lookupFlags failed"); + })), + object_layer_(execution_session_, + [this](llvm::orc::VModuleKey) { + llvm::orc::RTDyldObjectLinkingLayer::Resources result; + result.MemMgr = + std::make_shared( + orc_jit_memory_mapper::GetInstance()); + result.Resolver = symbol_resolver_; + return result; + }), + compile_layer_(object_layer_, + CompilerFunctor(target_machine_.get(), &disassembler_, + opt_level, optimize_for_size, + enable_fast_math, disable_expensive_passes, + std::move(pre_optimization_hook), + std::move(post_optimization_hook))) { VLOG(1) << "CPU target: " << target_machine_->getTargetCPU().str() << " features: " << target_machine_->getTargetFeatureString().str(); } -SimpleOrcJIT::ModuleHandleT SimpleOrcJIT::AddModule( +SimpleOrcJIT::VModuleKeyT SimpleOrcJIT::AddModule( std::unique_ptr module) { - auto handle = cantFail(compile_layer_.addModule( - std::move(module), MakeUnique(external_constant_pool()))); - module_handles_.push_back(handle); - return handle; + auto key = execution_session_.allocateVModule(); + cantFail(compile_layer_.addModule(key, std::move(module))); + module_keys_.push_back(key); + return key; } -void SimpleOrcJIT::RemoveModule(SimpleOrcJIT::ModuleHandleT handle) { - module_handles_.erase( - std::remove(module_handles_.begin(), module_handles_.end(), handle), - module_handles_.end()); - cantFail(compile_layer_.removeModule(handle)); +void SimpleOrcJIT::RemoveModule(SimpleOrcJIT::VModuleKeyT key) { + module_keys_.erase(std::remove(module_keys_.begin(), module_keys_.end(), key), + module_keys_.end()); + cantFail(compile_layer_.removeModule(key)); } llvm::JITSymbol SimpleOrcJIT::FindSymbol(const std::string& name) { @@ -177,10 +180,10 @@ llvm::JITSymbol SimpleOrcJIT::FindSymbol(const std::string& name) { // Resolve symbol from last module to first, allowing later redefinitions of // symbols shadow earlier ones. - for (auto& handle : - llvm::make_range(module_handles_.rbegin(), module_handles_.rend())) { + for (auto& key : + llvm::make_range(module_keys_.rbegin(), module_keys_.rend())) { if (auto symbol = - compile_layer_.findSymbolIn(handle, mangled_name, + compile_layer_.findSymbolIn(key, mangled_name, /*ExportedSymbolsOnly=*/true)) { return symbol; } @@ -207,24 +210,15 @@ bool RegisterKnownJITSymbols() { REGISTER_CPU_RUNTIME_SYMBOL(AcquireInfeedBufferForDequeue); REGISTER_CPU_RUNTIME_SYMBOL(AcquireOutfeedBufferForPopulation); + REGISTER_CPU_RUNTIME_SYMBOL(EigenConvF16); REGISTER_CPU_RUNTIME_SYMBOL(EigenConvF32); + REGISTER_CPU_RUNTIME_SYMBOL(EigenFft); REGISTER_CPU_RUNTIME_SYMBOL(EigenMatMulF32); REGISTER_CPU_RUNTIME_SYMBOL(EigenMatMulF64); + REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedConvF16); REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedConvF32); REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedMatMulF32); REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedMatMulF64); -#ifdef __ARM_NEON__ - REGISTER_CPU_RUNTIME_SYMBOL(ExpV4F32NEON); - REGISTER_CPU_RUNTIME_SYMBOL(LogV4F32NEON); -#endif -#ifdef __SSE4_1__ - REGISTER_CPU_RUNTIME_SYMBOL(ExpV4F32SSE); - REGISTER_CPU_RUNTIME_SYMBOL(LogV4F32SSE); -#endif -#ifdef __AVX__ - REGISTER_CPU_RUNTIME_SYMBOL(ExpV8F32AVX); - REGISTER_CPU_RUNTIME_SYMBOL(LogV8F32AVX); -#endif REGISTER_CPU_RUNTIME_SYMBOL(ParallelForkJoin); REGISTER_CPU_RUNTIME_SYMBOL(ReleaseInfeedBufferAfterDequeue); REGISTER_CPU_RUNTIME_SYMBOL(ReleaseOutfeedBufferAfterPopulation); @@ -271,15 +265,15 @@ bool RegisterKnownJITSymbols() { REGISTER_LIBM_SYMBOL(ilogb, int (*)(double)); REGISTER_LIBM_SYMBOL(ldexp, double (*)(double, int)); REGISTER_LIBM_SYMBOL(lgamma, double (*)(double)); - REGISTER_LIBM_SYMBOL(llrint, long long (*)(double)); - REGISTER_LIBM_SYMBOL(llround, long long (*)(double)); + REGISTER_LIBM_SYMBOL(llrint, long long (*)(double)); // NOLINT(runtime/int) + REGISTER_LIBM_SYMBOL(llround, long long (*)(double)); // NOLINT(runtime/int) REGISTER_LIBM_SYMBOL(log, double (*)(double)); REGISTER_LIBM_SYMBOL(log10, double (*)(double)); REGISTER_LIBM_SYMBOL(log1p, double (*)(double)); REGISTER_LIBM_SYMBOL(log2, double (*)(double)); REGISTER_LIBM_SYMBOL(logb, double (*)(double)); - REGISTER_LIBM_SYMBOL(lrint, long (*)(double)); - REGISTER_LIBM_SYMBOL(lround, long (*)(double)); + REGISTER_LIBM_SYMBOL(lrint, long (*)(double)); // NOLINT(runtime/int) + REGISTER_LIBM_SYMBOL(lround, long (*)(double)); // NOLINT(runtime/int) REGISTER_LIBM_SYMBOL(modf, double (*)(double, double*)); REGISTER_LIBM_SYMBOL(nan, double (*)(const char*)); REGISTER_LIBM_SYMBOL(nearbyint, double (*)(double)); @@ -290,7 +284,8 @@ bool RegisterKnownJITSymbols() { REGISTER_LIBM_SYMBOL(remquo, double (*)(double, double, int*)); REGISTER_LIBM_SYMBOL(rint, double (*)(double)); REGISTER_LIBM_SYMBOL(round, double (*)(double)); - REGISTER_LIBM_SYMBOL(scalbln, double (*)(double, long)); + REGISTER_LIBM_SYMBOL(scalbln, + double (*)(double, long)); // NOLINT(runtime/int) REGISTER_LIBM_SYMBOL(scalbn, double (*)(double, int)); REGISTER_LIBM_SYMBOL(sin, double (*)(double)); #ifdef __APPLE__ diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h index ded01e9e4d7442296f7406dd035e6ab385458238..50993afc8f73617a2c65310ae73b3ab00519f550 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h @@ -21,8 +21,10 @@ limitations under the License. #include #include "llvm/ADT/Triple.h" +#include "llvm/ExecutionEngine/Orc/Core.h" #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" +#include "llvm/ExecutionEngine/Orc/SymbolStringPool.h" #include "llvm/IR/Module.h" #include "llvm/Target/TargetMachine.h" #include "tensorflow/compiler/xla/service/cpu/compiler_functor.h" @@ -48,7 +50,7 @@ class SimpleOrcJIT { std::function( llvm::Module&)>; using CompileLayerT = llvm::orc::IRCompileLayer; - using ModuleHandleT = CompileLayerT::ModuleHandleT; + using VModuleKeyT = llvm::orc::VModuleKey; // Create a new JIT, targeting the host architecture. // The |target_options| parameter allows customization of certain code @@ -78,12 +80,12 @@ class SimpleOrcJIT { return target_machine_->getTargetTriple(); } - // Add a module to the JIT. Returns an opaque handle that can be used to later + // Add a module to the JIT. Returns an opaque key that can be used to later // remove this module. - ModuleHandleT AddModule(std::unique_ptr module); + VModuleKeyT AddModule(std::unique_ptr module); // Remove a module from the JIT and free the memory associated with it. - void RemoveModule(ModuleHandleT handle); + void RemoveModule(VModuleKeyT key); // Get the runtime address of the compiled symbol whose name is given. Returns // nullptr if the symbol cannot be found. @@ -96,10 +98,13 @@ class SimpleOrcJIT { } private: - std::vector module_handles_; + std::vector module_keys_; std::unique_ptr target_machine_; const Disassembler disassembler_; const llvm::DataLayout data_layout_; + llvm::orc::SymbolStringPool string_pool_; + llvm::orc::ExecutionSession execution_session_; + std::shared_ptr symbol_resolver_; ObjLayerT object_layer_; CompileLayerT compile_layer_; ExternalConstantPool external_constant_pool_; diff --git a/tensorflow/compiler/xla/service/cpu/target_machine_features.cc b/tensorflow/compiler/xla/service/cpu/target_machine_features.cc new file mode 100644 index 0000000000000000000000000000000000000000..eeb049737dddd11ef2ce229df772baec3ac03dd8 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/target_machine_features.cc @@ -0,0 +1,35 @@ +/* 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/compiler/xla/service/cpu/target_machine_features.h" + +namespace xla { +namespace cpu { + +llvm::TargetTransformInfo* TargetMachineFeatures::GetTargetTransformInfoFor( + const llvm::Function& function) const { + auto it = target_transform_info_cache_.find(&function); + if (it == target_transform_info_cache_.end()) { + auto emplace_result = target_transform_info_cache_.emplace( + &function, target_machine_->getTargetTransformInfo(function)); + CHECK(emplace_result.second); + it = emplace_result.first; + } + + return &it->second; +} + +} // namespace cpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/target_machine_features.h b/tensorflow/compiler/xla/service/cpu/target_machine_features.h new file mode 100644 index 0000000000000000000000000000000000000000..703942615e552dccde7ddec8c8b90e8a486652af --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/target_machine_features.h @@ -0,0 +1,84 @@ +/* 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_COMPILER_XLA_SERVICE_CPU_TARGET_MACHINE_FEATURES_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TARGET_MACHINE_FEATURES_H_ + +#include "llvm/Analysis/TargetTransformInfo.h" +#include "llvm/Target/TargetMachine.h" +#include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/core/lib/gtl/flatmap.h" + +namespace xla { +namespace cpu { + +// Wraps an llvm::TargetMachine and parses out some information that feeds into +// LLVM IR code generation decisions. +class TargetMachineFeatures { + public: + static constexpr int kX86AvxVectorByteSize = 32; + + TargetMachineFeatures(llvm::TargetMachine* target_machine) + : target_machine_(target_machine) {} + + // Return the vectorization factor, which is the number of bytes of data + // explicitly vectorized routines will try to process at once. + int vectorization_factor_in_bytes() const { + // Ideally this should be a function of the cache line size (which we can + // get from llvm::TargetTransformInfo::getCacheLineSize) of the target + // machine. Guess a value of 128 bytes for now. + return 128; + } + + // Return the size of the largest vector size in bytes. We need to pass in + // "function" since llvm functions can contain annotations for specializing + // them to specific micro-architectures (though currently XLA does not use + // this functionality). + int vector_register_byte_size(const llvm::Function& function) const { + llvm::TargetTransformInfo* tti = GetTargetTransformInfoFor(function); + return tti->getRegisterBitWidth(/*Vector=*/true) / 8; + } + + // Return the number of elements of type `type` that can fit into the largest + // vector register available. We need to pass in "function" since llvm + // functions can contain annotations for specializing them to specific + // micro-architectures (though currently XLA does not use this functionality). + int vector_register_num_elements(const llvm::Function& function, + PrimitiveType type) const { + return vector_register_byte_size(function) / + (primitive_util::BitWidth(type) / 8); + } + + private: + llvm::TargetTransformInfo* GetTargetTransformInfoFor( + const llvm::Function& function) const; + + // This cache saves us from having to create a llvm::TargetTransformInfo for + // every call to GetTargetTransformInfoFor (creating a TargetTransformInfo + // costs one heap allocation on X86). + // + // This is mutated from within `GetTargetTransformInfoFor` which is + // semantically a getter (and thus `const`); and is therefore declared + // mutable. Making this mutable is okay because it has cache semantics. + mutable tensorflow::gtl::FlatMap + target_transform_info_cache_; + llvm::TargetMachine* target_machine_; +}; + +} // namespace cpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TARGET_MACHINE_FEATURES_H_ diff --git a/tensorflow/compiler/xla/service/llvm_ir/vector_support_library.cc b/tensorflow/compiler/xla/service/cpu/vector_support_library.cc similarity index 64% rename from tensorflow/compiler/xla/service/llvm_ir/vector_support_library.cc rename to tensorflow/compiler/xla/service/cpu/vector_support_library.cc index 0f6d8483da88ba4bf3f26961c0cbc8d855faa82c..150db1cb6edec1af6724a8bca6a5f6272f1a7416 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/vector_support_library.cc +++ b/tensorflow/compiler/xla/service/cpu/vector_support_library.cc @@ -13,11 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/service/llvm_ir/vector_support_library.h" +#include "tensorflow/compiler/xla/service/cpu/vector_support_library.h" +#include "llvm/Support/raw_ostream.h" +#include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" namespace xla { +namespace cpu { VectorSupportLibrary::VectorSupportLibrary(PrimitiveType primitive_type, int64 vector_size, llvm::IRBuilder<>* ir_builder, @@ -33,8 +36,27 @@ VectorSupportLibrary::VectorSupportLibrary(PrimitiveType primitive_type, vector_pointer_type_ = llvm::PointerType::getUnqual(vector_type_); } +static string TypeToString(llvm::Type* type) { + std::string o; + llvm::raw_string_ostream ostream(o); + type->print(ostream); + return ostream.str(); +} + +void VectorSupportLibrary::AssertCorrectTypes( + std::initializer_list values) { + for (llvm::Value* v : values) { + llvm::Type* type = v->getType(); + if (type != scalar_type() && type != vector_type()) { + LOG(FATAL) << "Expected either " << TypeToString(scalar_type()) << " or " + << TypeToString(vector_type()) << " but got " + << TypeToString(type); + } + } +} + llvm::Value* VectorSupportLibrary::Mul(llvm::Value* lhs, llvm::Value* rhs) { - CHECK(lhs->getType() == scalar_type() || lhs->getType() == vector_type()); + AssertCorrectTypes({lhs, rhs}); return MulInternal(lhs, rhs); } @@ -48,10 +70,128 @@ llvm::Value* VectorSupportLibrary::MulInternal(llvm::Value* lhs, } llvm::Value* VectorSupportLibrary::Add(llvm::Value* lhs, llvm::Value* rhs) { - CHECK(lhs->getType() == scalar_type() || lhs->getType() == vector_type()); + AssertCorrectTypes({lhs, rhs}); return AddInternal(lhs, rhs); } +llvm::Value* VectorSupportLibrary::Sub(llvm::Value* lhs, llvm::Value* rhs) { + AssertCorrectTypes({lhs, rhs}); + return ir_builder()->CreateFSub(lhs, rhs); +} + +llvm::Value* VectorSupportLibrary::Max(llvm::Value* lhs, llvm::Value* rhs) { + AssertCorrectTypes({lhs, rhs}); + if (scalar_type_->isFloatingPointTy()) { + return llvm_ir::EmitFloatMax(lhs, rhs, ir_builder_); + } else { + LOG(FATAL) << "Max for integers is unimplemented"; + } +} + +llvm::Value* VectorSupportLibrary::Floor(llvm::Value* a) { + AssertCorrectTypes({a}); + return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::floor, {a}, + {a->getType()}, ir_builder()); +} + +llvm::Value* VectorSupportLibrary::Div(llvm::Value* lhs, llvm::Value* rhs) { + AssertCorrectTypes({lhs, rhs}); + if (scalar_type_->isFloatingPointTy()) { + return ir_builder()->CreateFDiv(lhs, rhs, name()); + } else { + LOG(FATAL) << "Division for integers is unimplemented"; + } +} + +llvm::Value* VectorSupportLibrary::Clamp(llvm::Value* a, + const llvm::APFloat& low, + const llvm::APFloat& high) { + AssertCorrectTypes({a}); + llvm::Type* type = a->getType(); + CHECK(low.compare(high) == llvm::APFloat::cmpLessThan); + CHECK(scalar_type_->isFloatingPointTy()); + return llvm_ir::EmitFloatMin( + llvm_ir::EmitFloatMax(a, GetConstantFloat(type, low), ir_builder_), + GetConstantFloat(type, high), ir_builder_); +} + +llvm::Value* VectorSupportLibrary::FCmpEQMask(llvm::Value* lhs, + llvm::Value* rhs) { + AssertCorrectTypes({lhs, rhs}); + return I1ToFloat(ir_builder()->CreateFCmpOEQ(lhs, rhs, name())); +} + +llvm::Value* VectorSupportLibrary::FCmpOLTMask(llvm::Value* lhs, + llvm::Value* rhs) { + AssertCorrectTypes({lhs, rhs}); + return I1ToFloat(ir_builder()->CreateFCmpOLT(lhs, rhs, name())); +} + +llvm::Value* VectorSupportLibrary::FCmpULEMask(llvm::Value* lhs, + llvm::Value* rhs) { + AssertCorrectTypes({lhs, rhs}); + return I1ToFloat(ir_builder()->CreateFCmpULE(lhs, rhs, name())); +} + +llvm::Value* VectorSupportLibrary::I1ToFloat(llvm::Value* i1) { + bool is_vector = llvm::isa(i1->getType()); + llvm::Type* integer_type = IntegerTypeForFloatSize(is_vector); + return ir_builder()->CreateBitCast( + ir_builder()->CreateSExt(i1, integer_type, name()), + is_vector ? vector_type() : scalar_type(), name()); +} + +llvm::Type* VectorSupportLibrary::IntegerTypeForFloatSize(bool vector) { + CHECK(scalar_type()->isFloatingPointTy()); + const llvm::DataLayout& data_layout = + ir_builder()->GetInsertBlock()->getModule()->getDataLayout(); + int64 float_size_bits = data_layout.getTypeSizeInBits(scalar_type()); + llvm::Type* scalar_int_type = ir_builder()->getIntNTy(float_size_bits); + if (vector) { + return llvm::VectorType::get(scalar_int_type, vector_size()); + } else { + return scalar_int_type; + } +} + +llvm::Value* VectorSupportLibrary::BroadcastScalar(llvm::Value* x) { + CHECK_EQ(x->getType(), scalar_type()); + return ir_builder()->CreateVectorSplat(vector_size(), x, name()); +} + +llvm::Value* VectorSupportLibrary::FloatAnd(llvm::Value* lhs, + llvm::Value* rhs) { + AssertCorrectTypes({lhs, rhs}); + llvm::Type* int_type = + IntegerTypeForFloatSize(lhs->getType() == vector_type()); + return ir_builder()->CreateBitCast( + ir_builder()->CreateAnd( + ir_builder()->CreateBitCast(lhs, int_type, name()), + ir_builder()->CreateBitCast(rhs, int_type, name()), name()), + vector_type()); +} + +llvm::Value* VectorSupportLibrary::FloatNot(llvm::Value* lhs) { + AssertCorrectTypes({lhs}); + llvm::Type* int_type = + IntegerTypeForFloatSize(lhs->getType() == vector_type()); + return ir_builder()->CreateBitCast( + ir_builder()->CreateNot( + ir_builder()->CreateBitCast(lhs, int_type, name()), name()), + vector_type()); +} + +llvm::Value* VectorSupportLibrary::FloatOr(llvm::Value* lhs, llvm::Value* rhs) { + AssertCorrectTypes({lhs, rhs}); + llvm::Type* int_type = + IntegerTypeForFloatSize(lhs->getType() == vector_type()); + return ir_builder()->CreateBitCast( + ir_builder()->CreateOr(ir_builder()->CreateBitCast(lhs, int_type, name()), + ir_builder()->CreateBitCast(rhs, int_type, name()), + name()), + vector_type(), name()); +} + llvm::Value* VectorSupportLibrary::AddInternal(llvm::Value* lhs, llvm::Value* rhs) { if (scalar_type_->isFloatingPointTy()) { @@ -91,6 +231,7 @@ llvm::Value* VectorSupportLibrary::LoadScalar(llvm::Value* pointer) { void VectorSupportLibrary::StoreVector(llvm::Value* value, llvm::Value* pointer) { + AssertCorrectTypes({value}); if (pointer->getType() != vector_pointer_type()) { pointer = ir_builder()->CreateBitCast(pointer, vector_pointer_type()); } @@ -100,6 +241,7 @@ void VectorSupportLibrary::StoreVector(llvm::Value* value, void VectorSupportLibrary::StoreScalar(llvm::Value* value, llvm::Value* pointer) { + AssertCorrectTypes({value}); if (pointer->getType() != scalar_pointer_type()) { pointer = ir_builder()->CreateBitCast(pointer, scalar_pointer_type(), name()); @@ -206,9 +348,10 @@ llvm::Value* VectorSupportLibrary::ExtractHighHalf(llvm::Value* vector) { std::vector VectorSupportLibrary::ComputeHorizontalSums( std::vector vectors, llvm::Value* init_values) { - // TODO(sanjoy): Move this magic constant to TargetMachineFeatures. - const int kAvxVectorWidth = 8; - if (vector_size() == kAvxVectorWidth && vectors.size() == kAvxVectorWidth) { + const int x86_avx_vector_elements = + TargetMachineFeatures::kX86AvxVectorByteSize / scalar_byte_size(); + if (vector_size() == x86_avx_vector_elements && + vectors.size() == x86_avx_vector_elements) { return ComputeAvxOptimizedHorizontalSums(std::move(vectors), init_values); } @@ -277,4 +420,5 @@ llvm::Value* LlvmVariable::Get() const { void LlvmVariable::Set(llvm::Value* new_value) { ir_builder_->CreateStore(new_value, alloca_); } +} // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/llvm_ir/vector_support_library.h b/tensorflow/compiler/xla/service/cpu/vector_support_library.h similarity index 60% rename from tensorflow/compiler/xla/service/llvm_ir/vector_support_library.h rename to tensorflow/compiler/xla/service/cpu/vector_support_library.h index f404687ab6864bd0702d142ff691a394b78278a5..6479bf76aab581ae3ec2923d98dab53720cab203 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/vector_support_library.h +++ b/tensorflow/compiler/xla/service/cpu/vector_support_library.h @@ -13,17 +13,29 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_VECTOR_SUPPORT_LIBRARY_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_VECTOR_SUPPORT_LIBRARY_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_VECTOR_SUPPORT_LIBRARY_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_VECTOR_SUPPORT_LIBRARY_H_ #include #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Value.h" +#include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { +namespace cpu { + +// Simple wrappers around llvm::APFloat::APFloat to make the calling code more +// obvious. + +inline llvm::APFloat GetIeeeF32(float f) { return llvm::APFloat(f); } +inline llvm::APFloat GetIeeeF32FromBitwiseRep(int32 bitwise_value) { + return llvm::APFloat(llvm::APFloat::IEEEsingle(), + llvm::APInt(/*numBits=*/32, /*val=*/bitwise_value)); +} + // A thin wrapper around llvm_util.h to make code generating vector math flow // more readable. class VectorSupportLibrary { @@ -39,16 +51,96 @@ class VectorSupportLibrary { llvm::Value* Mul(int64 lhs, llvm::Value* rhs) { return Mul(ir_builder()->getInt64(lhs), rhs); } + llvm::Value* Mul(const llvm::APFloat& lhs, llvm::Value* rhs) { + return Mul(GetConstantFloat(rhs->getType(), lhs), rhs); + } + + // If your call resolved to these then you probably wanted the versions taking + // APFloat. + llvm::Value* Mul(double lhs, llvm::Value* rhs) = delete; + llvm::Value* Mul(float lhs, llvm::Value* rhs) = delete; llvm::Value* Add(llvm::Value* lhs, llvm::Value* rhs); llvm::Value* Add(int64 lhs, llvm::Value* rhs) { return Add(ir_builder()->getInt64(lhs), rhs); } + llvm::Value* Add(const llvm::APFloat& lhs, llvm::Value* rhs) { + return Add(GetConstantFloat(rhs->getType(), lhs), rhs); + } + + // If your call resolved to these then you probably wanted the versions taking + // APFloat. + llvm::Value* Add(double lhs, llvm::Value* rhs) = delete; + llvm::Value* Add(float lhs, llvm::Value* rhs) = delete; + + llvm::Value* Sub(llvm::Value* lhs, llvm::Value* rhs); + llvm::Value* Sub(llvm::Value* lhs, const llvm::APFloat& rhs) { + return Sub(lhs, GetConstantFloat(lhs->getType(), rhs)); + } + llvm::Value* Max(llvm::Value* lhs, llvm::Value* rhs); + llvm::Value* Max(const llvm::APFloat& lhs, llvm::Value* rhs) { + return Max(GetConstantFloat(rhs->getType(), lhs), rhs); + } + llvm::Value* Div(llvm::Value* lhs, llvm::Value* rhs); llvm::Value* MulAdd(llvm::Value* a, llvm::Value* b, llvm::Value* c) { return Add(c, Mul(a, b)); } + llvm::Value* MulAdd(llvm::Value* a, llvm::Value* b, const llvm::APFloat& c) { + return Add(GetConstantFloat(vector_type(), c), Mul(a, b)); + } + + llvm::Value* MulAdd(llvm::Value* a, const llvm::APFloat& b, + const llvm::APFloat& c) { + return Add(GetConstantFloat(a->getType(), c), + Mul(a, GetConstantFloat(a->getType(), b))); + } + + llvm::Value* Floor(llvm::Value* a); + + llvm::Value* Clamp(llvm::Value* a, const llvm::APFloat& low, + const llvm::APFloat& high); + llvm::Value* SplatFloat(const llvm::APFloat& d) { + return GetConstantFloat(vector_type(), d); + } + + // These compare instructions return a floating point typed mask instead of an + // i1. For instance, on a vector typed input, lanes where the predicate is + // true get a float with all ones and other lanes get a float with all zeros. + // This is slightly odd from the perspective of LLVM's type system, but it + // makes kernel IR generation code written using VectorSupportLibrary (its + // raison d'etre) less cluttered. + + llvm::Value* FCmpEQMask(llvm::Value* lhs, llvm::Value* rhs); + llvm::Value* FCmpULEMask(llvm::Value* lhs, llvm::Value* rhs); + llvm::Value* FCmpOLTMask(llvm::Value* lhs, llvm::Value* rhs); + llvm::Value* FCmpOLTMask(llvm::Value* lhs, const llvm::APFloat& rhs) { + return FCmpOLTMask(lhs, GetConstantFloat(lhs->getType(), rhs)); + } + + // These boolean operations operate on the bitwise values of the floating + // point inputs. They return a (vector of) float(s) but like in the mask + // generating predicates above this type system oddity makes the kernel IR + // generation code less cluttered. + llvm::Value* FloatAnd(llvm::Value* lhs, llvm::Value* rhs); + llvm::Value* FloatAnd(llvm::Value* lhs, const llvm::APFloat& rhs) { + return FloatAnd(lhs, GetConstantFloat(lhs->getType(), rhs)); + } + llvm::Value* FloatOr(llvm::Value* lhs, llvm::Value* rhs); + llvm::Value* FloatOr(llvm::Value* lhs, const llvm::APFloat& rhs) { + return FloatOr(lhs, GetConstantFloat(lhs->getType(), rhs)); + } + llvm::Value* FloatNot(llvm::Value* lhs); + llvm::Value* FloatAndNot(llvm::Value* lhs, llvm::Value* rhs) { + return FloatAnd(FloatNot(lhs), rhs); + } + + llvm::Value* BroadcastScalar(llvm::Value* x); + llvm::Value* BroadcastScalar(const llvm::APFloat& d) { + return BroadcastScalar(GetConstantFloat(scalar_type(), d)); + } + llvm::Value* ComputeOffsetPointer(llvm::Value* base_pointer, llvm::Value* offset_elements); llvm::Value* ComputeOffsetPointer(llvm::Value* base_pointer, @@ -127,6 +219,9 @@ class VectorSupportLibrary { llvm::Type* vector_pointer_type() const { return vector_pointer_type_; } llvm::Type* scalar_type() const { return scalar_type_; } llvm::Type* scalar_pointer_type() const { return scalar_pointer_type_; } + int64 scalar_byte_size() const { + return primitive_util::BitWidth(primitive_type_) / 8; + } const std::string& name() const { return name_; } @@ -139,6 +234,11 @@ class VectorSupportLibrary { llvm::Value* AddReduce(llvm::Value* vector); + // Checks that each value in `values` is either of type scalar_type() or + // vector_type(). This LOG(FATAL)'s so it should only be called in cases + // where a mismatching type is a programmer bug. + void AssertCorrectTypes(std::initializer_list values); + // Perform an X86 AVX style horizontal add between `lhs` and `rhs`. The // resulting IR for an 8-float wide vector is expected to lower to a single // vhaddps instruction on a CPU that supports vhaddps, and not be too bad in @@ -158,6 +258,16 @@ class VectorSupportLibrary { std::vector ComputeAvxOptimizedHorizontalSums( std::vector vectors, llvm::Value* init_values); + llvm::Type* IntegerTypeForFloatSize(bool vector); + llvm::Value* I1ToFloat(llvm::Value* i1); + llvm::Value* GetConstantFloat(llvm::Type* type, const llvm::APFloat& f) { + llvm::Constant* scalar_value = llvm::ConstantFP::get(type->getContext(), f); + if (llvm::isa(type)) { + return llvm::ConstantVector::getSplat(vector_size(), scalar_value); + } + return scalar_value; + } + int64 vector_size_; PrimitiveType primitive_type_; llvm::IRBuilder<>* ir_builder_; @@ -201,6 +311,7 @@ class ScalarVariable : public LlvmVariable { Set(initial_value); } }; +} // namespace cpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_VECTOR_SUPPORT_LIBRARY_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_VECTOR_SUPPORT_LIBRARY_H_ diff --git a/tensorflow/compiler/xla/service/cpu/windows_compatibility.cc b/tensorflow/compiler/xla/service/cpu/windows_compatibility.cc new file mode 100644 index 0000000000000000000000000000000000000000..ab308ee6cb16ba95e24694b59a4b5737765bbb8b --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/windows_compatibility.cc @@ -0,0 +1,32 @@ +/* Copyright 2017 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/compiler/xla/service/cpu/windows_compatibility.h" + +#ifdef _MSC_VER + +#include + +void sincos(double x, double *sinv, double *cosv) { + *sinv = sin(x); + *cosv = cos(x); +} + +void sincosf(float x, float *sinv, float *cosv) { + *sinv = sinf(x); + *cosv = cosf(x); +} + +#endif // _MSC_VER diff --git a/tensorflow/compiler/xla/service/cpu/windows_compatibility.h b/tensorflow/compiler/xla/service/cpu/windows_compatibility.h new file mode 100644 index 0000000000000000000000000000000000000000..262f379d8b6017f4a7e0156b724bfee7e8ec5b9a --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/windows_compatibility.h @@ -0,0 +1,31 @@ +/* Copyright 2017 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_COMPILER_XLA_SERVICE_CPU_WINDOWS_COMPATIBILITY_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_WINDOWS_COMPATIBILITY_H_ + +#ifdef _MSC_VER + +extern "C" { + +// MSVC does not have sincos[f]. +void sincos(double x, double *sinv, double *cosv); +void sincosf(float x, float *sinv, float *cosv); + +} + +#endif // _MSC_VER + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_WINDOWS_COMPATIBILITY_H_ diff --git a/tensorflow/compiler/xla/service/cpu/xfeed_manager.cc b/tensorflow/compiler/xla/service/cpu/xfeed_manager.cc index d0f214202908266371639af8f431ad8269ad0e35..47543b2082f55cf7b8cf60f1c5bbb16a0a609912 100644 --- a/tensorflow/compiler/xla/service/cpu/xfeed_manager.cc +++ b/tensorflow/compiler/xla/service/cpu/xfeed_manager.cc @@ -41,6 +41,8 @@ void XfeedQueueManager::EnqueueBuffersAtomically( tensorflow::mutex_lock l(mu_); bool was_empty = enqueued_buffers_.empty(); for (XfeedBuffer* b : buffers) { + VLOG(3) << "Enqueueing " << queue_name_ << " buffer (of " << buffers.size() + << " buffers) with length: " << b->length(); enqueued_buffers_.push_back(b); } if (was_empty && !buffers.empty()) { @@ -54,9 +56,11 @@ void XfeedQueueManager::EnqueueBuffersAtomically( XfeedBuffer* XfeedQueueManager::BlockingDequeueBuffer() { tensorflow::mutex_lock l(mu_); + VLOG(3) << "Waiting for an available buffer."; while (enqueued_buffers_.empty()) { cv_.wait(l); } + VLOG(3) << "A buffer is available!"; CHECK(current_buffer_ == nullptr); current_buffer_ = enqueued_buffers_.front(); enqueued_buffers_.pop_front(); @@ -65,6 +69,9 @@ XfeedBuffer* XfeedQueueManager::BlockingDequeueBuffer() { void XfeedQueueManager::ReleaseCurrentBuffer(int32 length, void* data, StatusOr shape) { + VLOG(3) << "Releasing buffer with shape: " + << (shape.ok() ? ShapeUtil::HumanString(shape.ValueOrDie()) + : ""); tensorflow::mutex_lock l(mu_); CHECK(current_buffer_ != nullptr); CHECK_EQ(length, current_buffer_->length()); diff --git a/tensorflow/compiler/xla/service/cpu/xfeed_manager.h b/tensorflow/compiler/xla/service/cpu/xfeed_manager.h index 6af55700052007a2ca419d52b63dddea2052bd0b..b4ace232607e14fbfec01d48946f0031d96cd027 100644 --- a/tensorflow/compiler/xla/service/cpu/xfeed_manager.h +++ b/tensorflow/compiler/xla/service/cpu/xfeed_manager.h @@ -50,7 +50,7 @@ class XfeedBuffer { // Reusable component for managing the infeed and outfeed queue state. class XfeedQueueManager { public: - XfeedQueueManager() = default; + XfeedQueueManager(string queue_name) : queue_name_(queue_name) {} // Calls the completion callback for any enqueued buffers that have // not been dequeued by the runtime, and empties the @@ -86,6 +86,8 @@ class XfeedQueueManager { void ReleaseCurrentBuffer(int32 length, void* data, StatusOr shape); private: + const string queue_name_; + tensorflow::mutex mu_; // Condition variable that is signaled every time a buffer is @@ -112,8 +114,8 @@ class XfeedManager { XfeedQueueManager* outfeed() { return &outfeed_; } private: - XfeedQueueManager infeed_; - XfeedQueueManager outfeed_; + XfeedQueueManager infeed_ = {"infeed"}; + XfeedQueueManager outfeed_ = {"outfeed"}; }; } // namespace runtime diff --git a/tensorflow/compiler/xla/service/device_memory_allocator.cc b/tensorflow/compiler/xla/service/device_memory_allocator.cc index 2e4b0a5230516b5308aeed892de9a49565a09f2e..78e7aa48accdbb51a8477455f5f9c004828c068f 100644 --- a/tensorflow/compiler/xla/service/device_memory_allocator.cc +++ b/tensorflow/compiler/xla/service/device_memory_allocator.cc @@ -24,7 +24,7 @@ limitations under the License. namespace xla { StreamExecutorMemoryAllocator::StreamExecutorMemoryAllocator( - perftools::gputools::Platform* platform, + const perftools::gputools::Platform* platform, tensorflow::gtl::ArraySlice stream_executors) : DeviceMemoryAllocator(platform), diff --git a/tensorflow/compiler/xla/service/device_memory_allocator.h b/tensorflow/compiler/xla/service/device_memory_allocator.h index 00caefab667cba6abfef200050ca18f229fc0320..39dfad84c1c1c1c461c24de555ecd919cea47d83 100644 --- a/tensorflow/compiler/xla/service/device_memory_allocator.h +++ b/tensorflow/compiler/xla/service/device_memory_allocator.h @@ -33,7 +33,7 @@ class DeviceMemoryAllocator { public: // Parameter platform indicates which platform the allocator allocates memory // on. Must be non-null. - explicit DeviceMemoryAllocator(perftools::gputools::Platform* platform) + explicit DeviceMemoryAllocator(const perftools::gputools::Platform* platform) : platform_(platform) {} virtual ~DeviceMemoryAllocator() {} @@ -49,14 +49,14 @@ class DeviceMemoryAllocator { int device_ordinal, perftools::gputools::DeviceMemoryBase* mem) = 0; // Return the platform that the allocator allocates memory on. - perftools::gputools::Platform* platform() const { return platform_; } + const perftools::gputools::Platform* platform() const { return platform_; } // Can we call Deallocate() as soon as a computation has been scheduled on // a stream, or do we have to wait for the computation to complete first? virtual bool AllowsAsynchronousDeallocation() const = 0; protected: - perftools::gputools::Platform* platform_; + const perftools::gputools::Platform* platform_; }; // Default memory allocator for a platform which uses @@ -64,7 +64,7 @@ class DeviceMemoryAllocator { class StreamExecutorMemoryAllocator : public DeviceMemoryAllocator { public: StreamExecutorMemoryAllocator( - perftools::gputools::Platform* platform, + const perftools::gputools::Platform* platform, tensorflow::gtl::ArraySlice stream_executors); diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h index 0d54e325e618a7b1aae38407958ddf7b41ef1cda..a803b3171f9afa6297553c5507c4f9aa45e420ab 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h @@ -103,6 +103,7 @@ class DfsHloVisitorBase { return HandleElementwiseBinary(hlo); } virtual Status HandleConvolution(HloInstructionPtr hlo) = 0; + virtual Status HandleFft(HloInstructionPtr fft) = 0; virtual Status HandleCrossReplicaSum(HloInstructionPtr hlo) = 0; virtual Status HandleCompare(HloInstructionPtr hlo) { return HandleElementwiseBinary(hlo); diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h index 133aa2509405738de8388708b0c61a82023e2738..170adb3d241b3648bc53f96dde9866f0b794f80a 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h @@ -85,6 +85,9 @@ class DfsHloVisitorWithDefaultBase Status HandleConvolution(HloInstructionPtr convolution) override { return DefaultAction(convolution); } + Status HandleFft(HloInstructionPtr fft) override { + return DefaultAction(fft); + } Status HandleCrossReplicaSum(HloInstructionPtr crs) override { return DefaultAction(crs); } diff --git a/tensorflow/compiler/xla/service/dot_decomposer.h b/tensorflow/compiler/xla/service/dot_decomposer.h index 5ff0ab34eac0cd0fbc264b408c57653c944402a6..1959b687f16d6909a3283021c8635b3e65e6e412 100644 --- a/tensorflow/compiler/xla/service/dot_decomposer.h +++ b/tensorflow/compiler/xla/service/dot_decomposer.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_DOT_DECOMPOSER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_DOT_DECOMPOSER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_DOT_DECOMPOSER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_DOT_DECOMPOSER_H_ #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" @@ -41,4 +41,4 @@ class DotDecomposer : public HloPassInterface { } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_DOT_DECOMPOSER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_DOT_DECOMPOSER_H_ diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc index 37929294327d2a57bb0ab1c48e90b6843cba6ae4..4468adbadbf823f1420a8b665a26f66cb7d36b43 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc @@ -428,7 +428,7 @@ StatusOr ElementalIrEmitter::EmitFloatUnaryOp( llvm::Intrinsic::round, {operand_value}, {operand_value->getType()}, ir_builder_); case HloOpcode::kSign: { - // TODO(b/32151903): Ensure consistent sign behavior for -0.0 + // TODO(b/32151903): Ensure consistent sign behavior for -0.0. auto type = operand_value->getType(); auto zero = llvm::ConstantFP::get(type, 0.0); auto oeq = ir_builder_->CreateFCmpOEQ(operand_value, zero); @@ -870,7 +870,10 @@ llvm::Value* ElementalIrEmitter::EmitFloatMin(llvm::Value* lhs_value, StatusOr ElementalIrEmitter::EmitErfInv(PrimitiveType prim_type, llvm::Value* x) const { if (prim_type != F32) { - return Unimplemented("inverse erf only implemented for F32 (b/34339814)"); + // TODO(b/34339814): Implement inverse erf for F64. + return Unimplemented( + "Inverse erf is only implemented for element " + "type F32."); } auto getFloat = [&](const float f) { return llvm::ConstantFP::get(ir_builder_->getFloatTy(), f); @@ -1040,17 +1043,9 @@ StatusOr ElementalIrEmitter::EmitIntegerBinaryOp( is_signed ? llvm::CmpInst::ICMP_SGE : llvm::CmpInst::ICMP_UGE, lhs_value, rhs_value, ir_builder_); case HloOpcode::kMinimum: - return ir_builder_->CreateSelect( - ir_builder_->CreateICmp( - is_signed ? llvm::ICmpInst::ICMP_SLE : llvm::ICmpInst::ICMP_ULE, - lhs_value, rhs_value), - lhs_value, rhs_value); + return EmitIntegralMin(lhs_value, rhs_value, is_signed); case HloOpcode::kMaximum: - return ir_builder_->CreateSelect( - ir_builder_->CreateICmp( - is_signed ? llvm::ICmpInst::ICMP_SGE : llvm::ICmpInst::ICMP_UGE, - lhs_value, rhs_value), - lhs_value, rhs_value); + return EmitIntegralMax(lhs_value, rhs_value, is_signed); case HloOpcode::kAnd: return ir_builder_->CreateAnd(lhs_value, rhs_value); case HloOpcode::kOr: @@ -1067,6 +1062,26 @@ StatusOr ElementalIrEmitter::EmitIntegerBinaryOp( } } +llvm::Value* ElementalIrEmitter::EmitIntegralMax(llvm::Value* lhs_value, + llvm::Value* rhs_value, + bool is_signed) const { + return ir_builder_->CreateSelect( + ir_builder_->CreateICmp( + is_signed ? llvm::ICmpInst::ICMP_SGE : llvm::ICmpInst::ICMP_UGE, + lhs_value, rhs_value), + lhs_value, rhs_value); +} + +llvm::Value* ElementalIrEmitter::EmitIntegralMin(llvm::Value* lhs_value, + llvm::Value* rhs_value, + bool is_signed) const { + return ir_builder_->CreateSelect( + ir_builder_->CreateICmp( + is_signed ? llvm::ICmpInst::ICMP_SLE : llvm::ICmpInst::ICMP_ULE, + lhs_value, rhs_value), + lhs_value, rhs_value); +} + llvm_ir::IrArray::Index ElementalIrEmitter::ElementwiseSourceIndex( const llvm_ir::IrArray::Index& target_index, const HloInstruction& hlo, int64 operand_no) const { @@ -1264,14 +1279,6 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeRngElementGenerator( get_next_uniform_float()))); return ir_builder_->CreateFAdd(ir_builder_->CreateFMul(r, s), m); } - case RNG_BERNOULLI: { - TF_ASSIGN_OR_RETURN(llvm::Value * p, - operand_to_generator.at(hlo->operand(0))(index)); - return ir_builder_->CreateZExt( - ir_builder_->CreateFCmpOLT(get_next_uniform_float(), p), - llvm_ir::PrimitiveTypeToIrType(hlo->shape().element_type(), - module_)); - } default: return InvalidArgument( "unhandled distribution %s", @@ -1371,7 +1378,18 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( TF_ASSIGN_OR_RETURN(llvm::Value * max_value, operand_to_generator.at(hlo->operand(2))( ElementwiseSourceIndex(index, *hlo, 2))); - return EmitFloatMin(max_value, EmitFloatMax(min_value, arg_value)); + PrimitiveType prim_type = hlo->shape().element_type(); + if (primitive_util::IsFloatingPointType(prim_type)) { + return EmitFloatMin(max_value, EmitFloatMax(min_value, arg_value)); + } else if (primitive_util::IsIntegralType(prim_type)) { + bool is_signed = primitive_util::IsSignedIntegralType(prim_type); + return EmitIntegralMin( + max_value, EmitIntegralMax(min_value, arg_value, is_signed), + is_signed); + } else { + return Unimplemented("Clamp unimplemented for %s", + PrimitiveType_Name(prim_type).c_str()); + } }; case HloOpcode::kReducePrecision: return [this, hlo, &operand_to_generator]( diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/elemental_ir_emitter.h index 1a48eb5fcb960b60d524ea56a43e15269576db76..c516a826d9e382bc738e54635426db639d17108c 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.h +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.h @@ -86,6 +86,12 @@ class ElementalIrEmitter { virtual llvm::Value* EmitFloatMin(llvm::Value* lhs_value, llvm::Value* rhs_value) const; + llvm::Value* EmitIntegralMax(llvm::Value* lhs_value, llvm::Value* rhs_value, + bool is_signed) const; + + llvm::Value* EmitIntegralMin(llvm::Value* lhs_value, llvm::Value* rhs_value, + bool is_signed) const; + virtual StatusOr EmitErfInv(PrimitiveType prim_type, llvm::Value* value) const; diff --git a/tensorflow/compiler/xla/service/executable.cc b/tensorflow/compiler/xla/service/executable.cc index c9a6ad5edbd35041f180e06c7ae0229b2b00b95a..90481c7a88f90edea5399ee44aee2d2c77fc115f 100644 --- a/tensorflow/compiler/xla/service/executable.cc +++ b/tensorflow/compiler/xla/service/executable.cc @@ -73,7 +73,7 @@ StatusOr> Executable::ExecuteOnStreamWrapper( std::unique_ptr profile_ptr = module_config().debug_options().xla_hlo_profile() && hlo_profiling_enabled() - ? MakeUnique(&hlo_profile_printer(), + ? MakeUnique(&hlo_profile_printer_data(), &hlo_profile_index_map()) : nullptr; @@ -87,6 +87,10 @@ StatusOr> Executable::ExecuteOnStreamWrapper( VLOG(1) << "done with block-host-until-done"; // Merge in run-time profile information from execution_profile. + // + // TODO(b/71713097): This is buggy -- even though the mutex takes care of + // C++ level races, some other concurrent ExecuteOnStreamWrapper call could + // have rewritten the execution_profile before we get to it. profile->MergeFrom(execution_profile()); // Overall execution time (in nanoseconds) from the executor timer. diff --git a/tensorflow/compiler/xla/service/executable.h b/tensorflow/compiler/xla/service/executable.h index 5ecfdffe211c571b1bb2bc30ff2acd3021c735ae..0aee535ee780ef000bc5e9963ff48786b3a61eb2 100644 --- a/tensorflow/compiler/xla/service/executable.h +++ b/tensorflow/compiler/xla/service/executable.h @@ -44,13 +44,14 @@ namespace xla { // interface that is used for launching compiled programs across platforms. class Executable { public: - explicit Executable(std::unique_ptr hlo_module, - std::unique_ptr hlo_profile_printer, - std::unique_ptr hlo_profile_index_map) + explicit Executable( + std::unique_ptr hlo_module, + std::unique_ptr hlo_profile_printer_data, + std::unique_ptr hlo_profile_index_map) : hlo_module_(std::move(hlo_module)), - hlo_profile_printer_(std::move(hlo_profile_printer)), + hlo_profile_printer_data_(std::move(hlo_profile_printer_data)), hlo_profile_index_map_(std::move(hlo_profile_index_map)) { - CHECK_EQ(hlo_profile_printer_.get() == nullptr, + CHECK_EQ(hlo_profile_printer_data_.get() == nullptr, hlo_profile_index_map_.get() == nullptr); } virtual ~Executable() {} @@ -116,9 +117,9 @@ class Executable { "Equality test on this executable is not implemented."); } - const HloProfilePrinter& hlo_profile_printer() const { + const HloProfilePrinterData& hlo_profile_printer_data() const { CHECK(hlo_profiling_enabled()); - return *hlo_profile_printer_; + return *hlo_profile_printer_data_; } const HloProfileIndexMap& hlo_profile_index_map() const { @@ -129,7 +130,9 @@ class Executable { // Returns whether this executable was compiled with HLO profilings support // enabled. If not, the caller should not expect an hlo_execution_profile // passed to ExecuteOnStream above to be populated during execution. - bool hlo_profiling_enabled() const { return hlo_profile_printer_ != nullptr; } + bool hlo_profiling_enabled() const { + return hlo_profile_printer_data_ != nullptr; + } const HloModule& module() const { return *hlo_module_; } @@ -179,7 +182,7 @@ class Executable { // execution. int64 execution_count_ = 0; - std::unique_ptr hlo_profile_printer_; + std::unique_ptr hlo_profile_printer_data_; std::unique_ptr hlo_profile_index_map_; }; diff --git a/tensorflow/compiler/xla/service/generic_transfer_manager.cc b/tensorflow/compiler/xla/service/generic_transfer_manager.cc index 271a856efd66f9f977ac4e201161ba4b505f31e1..78dc0ad4fcd167c93f19d0c2b18ea72d666897ef 100644 --- a/tensorflow/compiler/xla/service/generic_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/generic_transfer_manager.cc @@ -89,7 +89,7 @@ GenericTransferManager::TransferLiteralFromDevice( /*source=*/device_buffer.buffer(index), /*size=*/GetByteSizeRequirement(subshape), /*destination=*/ - literal->GetSubliteral(index).MutableInternalData())); + literal->untyped_data(index))); } return Status::OK(); @@ -124,17 +124,17 @@ Status GenericTransferManager::TransferLiteralToDevice( TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == device_memory.size()); // Element is array-shaped: transfer array data to device buffer. - const Literal& subliteral = literal.GetSubliteral(index); + const auto subliteral = LiteralView::Create(literal, index); std::unique_ptr relayed_out_literal; const void* source; if (LayoutUtil::Equal(device_subshape.layout(), subliteral.shape().layout())) { - source = subliteral.InternalData(); + source = subliteral.untyped_data(); } else { // Relayout data before transferring. relayed_out_literal = subliteral.Relayout(device_subshape.layout(), /*shape_index=*/{}); - source = relayed_out_literal->InternalData(); + source = relayed_out_literal->untyped_data(); } return TransferBufferToDevice( executor, diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD index d5d89de6d4dbd2a206ed561ca00993119c9d2bc0..9da4fb97fa27a238fead74985cb481a9be1f4a65 100644 --- a/tensorflow/compiler/xla/service/gpu/BUILD +++ b/tensorflow/compiler/xla/service/gpu/BUILD @@ -23,6 +23,15 @@ filegroup( load("//tensorflow:tensorflow.bzl", "tf_cc_test") +cc_library( + name = "gpu_constants", + srcs = ["gpu_constants.cc"], + hdrs = ["gpu_constants.h"], + deps = [ + "//tensorflow/compiler/xla:types", + ], +) + cc_library( name = "partition_assignment", srcs = [ @@ -120,9 +129,13 @@ cc_library( hdrs = [ "ir_emitter.h", "ir_emitter_context.h", + "ir_emitter_nested.h", + "ir_emitter_unnested.h", ], deps = [ + ":cudnn_convolution_runner", ":elemental_ir_emitter", + ":gpu_constants", ":gpu_executable", ":hlo_to_ir_bindings", ":ir_emission_utils", @@ -203,6 +216,7 @@ cc_library( srcs = ["buffer_allocations.cc"], hdrs = ["buffer_allocations.h"], deps = [ + ":gpu_constants", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", @@ -217,9 +231,11 @@ cc_library( cc_library( name = "gpu_executable", srcs = [ + "conditional_thunk.cc", "convolution_thunk.cc", "copy_thunk.cc", "cudnn_batchnorm_thunk.cc", + "fft_thunk.cc", "for_thunk.cc", "gemm_thunk.cc", "gpu_executable.cc", @@ -231,9 +247,11 @@ cc_library( "while_thunk.cc", ], hdrs = [ + "conditional_thunk.h", "convolution_thunk.h", "copy_thunk.h", "cudnn_batchnorm_thunk.h", + "fft_thunk.h", "for_thunk.h", "gemm_thunk.h", "gpu_executable.h", @@ -247,6 +265,7 @@ cc_library( ], deps = [ ":buffer_allocations", + ":cudnn_convolution_runner", ":infeed_manager", ":ir_emission_utils", ":partition_assignment", @@ -272,6 +291,7 @@ cc_library( "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core/platform/default/build_config:cublas_plugin", "//tensorflow/core/platform/default/build_config:cudnn_plugin", + "//tensorflow/core/platform/default/build_config:cufft_plugin", "//tensorflow/core/platform/default/build_config:stream_executor_cuda", # build_cleaner: keep ], ) @@ -293,9 +313,41 @@ cc_library( ) cc_library( - name = "convolution_folding", - srcs = ["convolution_folding.cc"], - hdrs = ["convolution_folding.h"], + name = "cudnn_convolution_algorithm_picker", + srcs = ["cudnn_convolution_algorithm_picker.cc"], + hdrs = ["cudnn_convolution_algorithm_picker.h"], + deps = [ + ":cudnn_convolution_runner", + ":gpu_executable", + ":ir_emission_utils", + "//tensorflow/compiler/xla/service:device_memory_allocator", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_pass", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + ], +) + +cc_library( + name = "cudnn_convolution_runner", + srcs = ["cudnn_convolution_runner.cc"], + hdrs = ["cudnn_convolution_runner.h"], + deps = [ + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:stream_executor_no_cuda", + ], +) + +cc_library( + name = "cudnn_convolution_rewriter", + srcs = ["cudnn_convolution_rewriter.cc"], + hdrs = ["cudnn_convolution_rewriter.h"], deps = [ ":ir_emission_utils", "//tensorflow/compiler/xla:literal_util", @@ -309,15 +361,18 @@ cc_library( ) tf_cc_test( - name = "convolution_folding_test", - srcs = ["convolution_folding_test.cc"], + name = "cudnn_convolution_rewriter_test", + srcs = ["cudnn_convolution_rewriter_test.cc"], deps = [ - ":convolution_folding", + ":cudnn_convolution_rewriter", + ":ir_emission_utils", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_matchers", "//tensorflow/compiler/xla/service:shape_inference", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep "//tensorflow/core:test", ], ) @@ -430,10 +485,13 @@ cc_library( srcs = ["gpu_compiler.cc"], hdrs = ["gpu_compiler.h"], deps = [ - ":convolution_folding", + ":cudnn_convolution_algorithm_picker", + ":cudnn_convolution_rewriter", ":fusion_merger", + ":gpu_constants", ":gpu_copy_insertion", ":gpu_executable", + ":gpu_hlo_support_checker", ":gpu_layout_assignment", ":hlo_schedule", ":instruction_fusion", @@ -459,6 +517,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo_constant_folding", "//tensorflow/compiler/xla/service:hlo_cse", "//tensorflow/compiler/xla/service:hlo_dce", + "//tensorflow/compiler/xla/service:hlo_element_type_converter", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/compiler/xla/service:hlo_pass_pipeline", "//tensorflow/compiler/xla/service:hlo_proto", @@ -478,6 +537,7 @@ cc_library( "//tensorflow/core:cuda_libdevice_path", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "//tensorflow/core:regexp_internal", "//tensorflow/core:stream_executor_no_cuda", "@llvm//:core", "@llvm//:support", @@ -494,7 +554,6 @@ cc_library( "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_pass", - "@llvm//:core", ], ) @@ -606,6 +665,32 @@ tf_cc_test( ], ) +cc_library( + name = "gpu_hlo_support_checker", + srcs = ["gpu_hlo_support_checker.cc"], + hdrs = ["gpu_hlo_support_checker.h"], + deps = [ + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo_pass", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "gpu_hlo_support_checker_test", + srcs = ["gpu_hlo_support_checker_test.cc"], + deps = [ + ":gpu_hlo_support_checker", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc b/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc index 9fdf717b5d463010e2709b6209c070f25555de72..2029c303d47e9a62135b003c3bd9be6f8b3438d4 100644 --- a/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc +++ b/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc @@ -19,6 +19,7 @@ limitations under the License. #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_constants.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" @@ -48,6 +49,15 @@ StatusOr> BufferAllocations::Builder::Build( // If buffer #i's address is already registered (e.g. external arguments or // result buffers), use that registered buffer. if (registered_buffers_.count(i)) { + se::DeviceMemoryBase address = FindOrDie(registered_buffers_, i); + if (reinterpret_cast(address.opaque()) % + kCudaMallocAlignBytes != + 0) { + return InternalError( + "Address of registered buffer %lld must be a multiple of %llx, but " + "was %p", + i, kCudaMallocAlignBytes, address.opaque()); + } buffer_allocations->SetBuffer(i, FindOrDie(registered_buffers_, i)); continue; } @@ -67,6 +77,14 @@ StatusOr> BufferAllocations::Builder::Build( tensorflow::strings::HumanReadableNumBytes(buffer_size).c_str(), i); } + if (reinterpret_cast(buffer_address.opaque()) % + kCudaMallocAlignBytes != + 0) { + return InternalError( + "Address returned by memory_allocator->Allocate must be a " + "multiple of %llx, but was %p", + kCudaMallocAlignBytes, buffer_address.opaque()); + } } buffer_allocations->SetBuffer(i, buffer_address); if (allocation.IsPreallocatedTempBuffer()) { @@ -80,6 +98,14 @@ StatusOr> BufferAllocations::Builder::Build( } } + if (VLOG_IS_ON(2)) { + for (BufferAllocation::Index i = 0; i < num_buffers; ++i) { + const auto& buf = buffer_allocations->buffers_[i]; + VLOG(2) << "Buffer " << i << " -> " << buf.opaque() << " (" << buf.size() + << "B)"; + } + } + return std::move(buffer_allocations); } diff --git a/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc b/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc new file mode 100644 index 0000000000000000000000000000000000000000..790ca535b11ee47724ef6227de40726d940d6153 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc @@ -0,0 +1,72 @@ +/* Copyright 2017 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/compiler/xla/service/gpu/conditional_thunk.h" + +#include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" + +namespace xla { +namespace gpu { + +ConditionalThunk::ConditionalThunk( + const BufferAllocation::Slice& predicate_buffer_index, + const BufferAllocation::Slice& true_operand_buffer_index, + const BufferAllocation::Slice& false_operand_buffer_index, + ThunkSequence true_thunk_sequence, ThunkSequence false_thunk_sequence, + const HloInstruction* hlo) + : Thunk(Kind::kConditional, hlo), + predicate_buffer_index_(predicate_buffer_index), + true_operand_buffer_index_(true_operand_buffer_index), + false_operand_buffer_index_(false_operand_buffer_index), + true_thunk_(std::move(true_thunk_sequence), hlo), + false_thunk_(std::move(false_thunk_sequence), hlo) {} + +Status ConditionalThunk::Initialize(const GpuExecutable& executable) { + TF_RETURN_IF_ERROR(true_thunk_.Initialize(executable)); + TF_RETURN_IF_ERROR(false_thunk_.Initialize(executable)); + return Status::OK(); +} + +Status ConditionalThunk::ExecuteOnStream( + const BufferAllocations& buffer_allocations, + perftools::gputools::Stream* stream) { + // Copy the predicate value from device. + bool predicate; + perftools::gputools::DeviceMemoryBase predicate_address = + buffer_allocations.GetDeviceAddress(predicate_buffer_index_); + stream->ThenMemcpy(&predicate, predicate_address, sizeof(bool)); + + Status block_status = stream->BlockHostUntilDone(); + if (!block_status.ok()) { + return InternalError("Failed to retrieve predicate value on stream %p: %s.", + stream, block_status.error_message().c_str()); + } + + // Execute the true or the false computation depending on the value of the + // predicate. + if (predicate) { + TF_RETURN_IF_ERROR(true_thunk_.ExecuteOnStream(buffer_allocations, stream)); + } else { + TF_RETURN_IF_ERROR( + false_thunk_.ExecuteOnStream(buffer_allocations, stream)); + } + + return Status::OK(); +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/conditional_thunk.h b/tensorflow/compiler/xla/service/gpu/conditional_thunk.h new file mode 100644 index 0000000000000000000000000000000000000000..7725c46a3b4b51af34a4dd977885353ff32c21f6 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/conditional_thunk.h @@ -0,0 +1,65 @@ +/* Copyright 2017 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_COMPILER_XLA_SERVICE_GPU_CONDITIONAL_THUNK_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CONDITIONAL_THUNK_H_ + +#include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" +#include "tensorflow/compiler/xla/service/gpu/sequential_thunk.h" +#include "tensorflow/compiler/xla/service/gpu/thunk.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { +namespace gpu { + +// ConditionalThunk implements the conditional instruction on GPU by reading the +// predicate of the conditional and executing the true or the false computation +// depending on the value of the predicate. +// +// ConditionalThunk assumes that the buffers of the conditional result and the +// result of the true and false computations share the same allocation. Also, +// the buffers of the true operand of the conditional and that of the parameter +// instruction of the true computation share the same allocation. Similarly, the +// buffers of the false operand and that of the parameter instruction of the +// false computation share the same allocation. +class ConditionalThunk : public Thunk { + public: + ConditionalThunk(const BufferAllocation::Slice& predicate_buffer_index, + const BufferAllocation::Slice& true_operand_buffer_index, + const BufferAllocation::Slice& false_operand_buffer_index, + ThunkSequence true_thunk_sequence, + ThunkSequence false_thunk_sequence, + const HloInstruction* hlo); + + ConditionalThunk(const ConditionalThunk&) = delete; + ConditionalThunk& operator=(const ConditionalThunk&) = delete; + + Status Initialize(const GpuExecutable& executable) override; + Status ExecuteOnStream(const BufferAllocations& buffer_allocations, + perftools::gputools::Stream* stream) override; + + private: + BufferAllocation::Slice predicate_buffer_index_; + BufferAllocation::Slice true_operand_buffer_index_; + BufferAllocation::Slice false_operand_buffer_index_; + SequentialThunk true_thunk_; + SequentialThunk false_thunk_; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CONDITIONAL_THUNK_H_ diff --git a/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc b/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc index 899cc5c83b99f1bb6154f883ca17871863e1f457..461747b699b542ae0c8735aea34cc9e57c1fb387 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/strings/strcat.h" @@ -36,366 +37,70 @@ using se::dnn::DataLayout; using se::dnn::FilterDescriptor; using se::dnn::FilterLayout; -ConvolveScratchAllocator::ConvolveScratchAllocator( - int device_ordinal, DeviceMemoryAllocator* memory_allocator) - : device_ordinal_(device_ordinal), memory_allocator_(memory_allocator) {} - -ConvolveScratchAllocator::~ConvolveScratchAllocator() { - for (auto& allocated_buffer : allocated_buffers_) { - if (!memory_allocator_->Deallocate(device_ordinal_, &allocated_buffer) - .ok()) { - // The program can still continue with failed deallocation. - LOG(ERROR) << "Failed to deallocate the allocated buffer: " - << allocated_buffer.opaque(); - } - } -} - -int64 ConvolveScratchAllocator::GetMemoryLimitInBytes(se::Stream* stream) { - constexpr int64 kConvolveScratchSize = 1LL << 32; // 4GB by default. - return kConvolveScratchSize; -} - -se::port::StatusOr> -ConvolveScratchAllocator::AllocateBytes(se::Stream* stream, int64 byte_size) { - CHECK_GE(byte_size, 0) << "byte_size must be positive."; - if (byte_size > GetMemoryLimitInBytes(stream)) { - return se::port::Status( - se::port::error::RESOURCE_EXHAUSTED, - tensorflow::strings::Printf( - "Allocating %lld bytes exceeds the memory limit of %lld bytes.", - byte_size, GetMemoryLimitInBytes(stream))); - } - - auto status_or_memory = - memory_allocator_->Allocate(device_ordinal_, byte_size, - /*retry_on_failure=*/false); - if (!status_or_memory.ok()) { - return se::port::Status(se::port::error::RESOURCE_EXHAUSTED, - tensorflow::strings::Printf( - "Failed to allocate %lld bytes on device %d.", - byte_size, device_ordinal_)); - } - se::DeviceMemoryBase allocated_buffer = status_or_memory.ValueOrDie(); - allocated_buffers_.push_back(allocated_buffer); - total_allocated_bytes_ += byte_size; - return se::DeviceMemory(allocated_buffer); -} - -string ConvolutionKindToString( - ConvolutionThunk::ConvolutionKind convolution_kind) { - switch (convolution_kind) { - case ConvolutionThunk::ConvolutionKind::kForward: - return "forward"; - case ConvolutionThunk::ConvolutionKind::kBackwardFilter: - return "backward_filter"; - case ConvolutionThunk::ConvolutionKind::kBackwardInput: - return "backward_input"; - } - return "unknown convolution kind"; -} - ConvolutionThunk::ConvolutionThunk( - ConvolutionKind convolution_kind, - const BufferAllocation::Slice& input_buffer, + CudnnConvKind convolution_kind, const BufferAllocation::Slice& input_buffer, const BufferAllocation::Slice& filter_buffer, - const BufferAllocation::Slice& output_buffer, const Shape& input_shape, + const BufferAllocation::Slice& output_buffer, + const BufferAllocation::Slice& tuple_result_buffer, + const BufferAllocation::Slice& scratch_buffer, const Shape& input_shape, const Shape& filter_shape, const Shape& output_shape, const Window& window, - const ConvolutionDimensionNumbers& dim_nums, const HloInstruction* hlo) + const ConvolutionDimensionNumbers& dim_nums, int64 algorithm, + bool tensor_ops_enabled, const HloInstruction* hlo) : Thunk(Kind::kConvolution, hlo), convolution_kind_(convolution_kind), input_buffer_(input_buffer), filter_buffer_(filter_buffer), output_buffer_(output_buffer), + tuple_result_buffer_(tuple_result_buffer), + scratch_buffer_(scratch_buffer), input_shape_(input_shape), filter_shape_(filter_shape), output_shape_(output_shape), window_(window), - dim_nums_(dim_nums) {} - -tensorflow::Status ConvolutionThunk::ExecuteOnStream( - const BufferAllocations& buffer_allocations, se::Stream* stream) { - VLOG(3) << "Convolution kind: " << ConvolutionKindToString(convolution_kind_); - VLOG(3) << "input shape: { " << input_shape_.ShortDebugString() << " }"; - VLOG(3) << "filter shape: { " << filter_shape_.ShortDebugString() << " }"; - VLOG(3) << "Output shape: { " << output_shape_.ShortDebugString() << " }"; - VLOG(3) << "Dim nums: { " << dim_nums_.ShortDebugString() << " }"; - VLOG(3) << "Window: { " << window_.ShortDebugString() << " }"; - - const int num_dimensions = window_.dimensions_size(); - CHECK_LE(num_dimensions, 3); - // cuDNN does not support 1D convolutions. We therefore express 1D - // convolutions as 2D convolutions where the first spatial dimension is 1. - // This matches the behavior of TF (see definition of conv1d in - // tensorflow/python/ops/nn_ops.py). - const int effective_num_dimensions = std::max(2, num_dimensions); - - CHECK_EQ(F32, output_shape_.element_type()); - CHECK_EQ(num_dimensions, dim_nums_.input_spatial_dimensions_size()); - CHECK_EQ(num_dimensions, dim_nums_.kernel_spatial_dimensions_size()); - CHECK_EQ(num_dimensions, dim_nums_.output_spatial_dimensions_size()); - for (const WindowDimension& dim : window_.dimensions()) { - CHECK_EQ(dim.padding_low(), dim.padding_high()); - } - - // cuDNN's convolution APIs support the BDYX layout for activations/output and - // the OIYX layout for weights. - BatchDescriptor input_descriptor(effective_num_dimensions); - input_descriptor.set_layout(DataLayout::kBatchDepthYX) - .set_feature_map_count( - input_shape_.dimensions(dim_nums_.input_feature_dimension())) - .set_count(input_shape_.dimensions(dim_nums_.input_batch_dimension())); - for (int dim = 0; dim < num_dimensions; ++dim) { - // Note that the dimensions are reversed. The same holds below. - input_descriptor.set_spatial_dim( - static_cast(effective_num_dimensions - dim - 1), - input_shape_.dimensions(dim_nums_.input_spatial_dimensions(dim))); - } - - FilterDescriptor filter_descriptor(effective_num_dimensions); - filter_descriptor.set_layout(FilterLayout::kOutputInputYX) - .set_input_feature_map_count( - filter_shape_.dimensions(dim_nums_.kernel_input_feature_dimension())) - .set_output_feature_map_count(filter_shape_.dimensions( - dim_nums_.kernel_output_feature_dimension())); - for (int dim = 0; dim < num_dimensions; ++dim) { - filter_descriptor.set_spatial_dim( - static_cast(effective_num_dimensions - dim - 1), - filter_shape_.dimensions(dim_nums_.kernel_spatial_dimensions(dim))); - } - - ConvolutionDescriptor convolution_descriptor(effective_num_dimensions); - for (int dim = 0; dim < num_dimensions; ++dim) { - convolution_descriptor - .set_zero_padding( - static_cast(effective_num_dimensions - dim - 1), - window_.dimensions(dim).padding_low()) - .set_filter_stride( - static_cast(effective_num_dimensions - dim - 1), - window_.dimensions(dim).stride()); - } - - BatchDescriptor output_descriptor(effective_num_dimensions); - output_descriptor.set_layout(DataLayout::kBatchDepthYX) - .set_feature_map_count( - output_shape_.dimensions(dim_nums_.output_feature_dimension())) - .set_count(output_shape_.dimensions(dim_nums_.output_batch_dimension())); - for (int dim = 0; dim < num_dimensions; ++dim) { - output_descriptor.set_spatial_dim( - static_cast(effective_num_dimensions - dim - 1), - output_shape_.dimensions(dim_nums_.output_spatial_dimensions(dim))); - } - - // Add a singleton dimension in the 1D convolution case. - if (num_dimensions == 1) { - input_descriptor.set_spatial_dim(static_cast(0), 1); - output_descriptor.set_spatial_dim(static_cast(0), 1); - filter_descriptor.set_spatial_dim(static_cast(0), 1); - convolution_descriptor - .set_zero_padding(static_cast(0), 0) - .set_filter_stride(static_cast(0), 1); - } - - se::DeviceMemory input_data( - buffer_allocations.GetDeviceAddress(input_buffer_)); - se::DeviceMemory filter_data( - buffer_allocations.GetDeviceAddress(filter_buffer_)); - se::DeviceMemory output_data( - buffer_allocations.GetDeviceAddress(output_buffer_)); - return ConvolveWithTune(input_descriptor, input_data, filter_descriptor, - filter_data, output_descriptor, output_data, - convolution_descriptor, buffer_allocations, stream); -} - -tensorflow::Status ConvolutionThunk::Convolve( - const BatchDescriptor& input_descriptor, se::DeviceMemory input_data, - const FilterDescriptor& filter_descriptor, - se::DeviceMemory filter_data, - const BatchDescriptor& output_descriptor, - se::DeviceMemory output_data, - const ConvolutionDescriptor& convolution_descriptor, - const se::dnn::AlgorithmConfig& algorithm_config, se::Stream* stream, - ConvolveScratchAllocator* scratch_allocator, - se::dnn::ProfileResult* profile_result) { - bool launch_ok; - switch (convolution_kind_) { - case ConvolutionKind::kBackwardFilter: - launch_ok = - stream - ->ThenConvolveBackwardFilterWithAlgorithm( - input_descriptor, input_data, output_descriptor, output_data, - convolution_descriptor, filter_descriptor, &filter_data, - scratch_allocator, algorithm_config, profile_result) - .ok(); - break; - case ConvolutionKind::kBackwardInput: - launch_ok = stream - ->ThenConvolveBackwardDataWithAlgorithm( - filter_descriptor, filter_data, output_descriptor, - output_data, convolution_descriptor, input_descriptor, - &input_data, scratch_allocator, algorithm_config, - profile_result) - .ok(); - break; - case ConvolutionKind::kForward: - launch_ok = - stream - ->ThenConvolveWithAlgorithm( - input_descriptor, input_data, filter_descriptor, filter_data, - convolution_descriptor, output_descriptor, &output_data, - scratch_allocator, algorithm_config, profile_result) - .ok(); - break; - } - if (launch_ok) { - return tensorflow::Status::OK(); - } - return InternalError( - "Unable to launch convolution for thunk %p with type %s and algorithm " - "(%lld, %lld)", - this, ConvolutionKindToString(convolution_kind_).c_str(), - algorithm_config.algorithm().algo_id(), - algorithm_config.algorithm_no_scratch().algo_id()); -} + dim_nums_(dim_nums), + algorithm_(algorithm), + tensor_ops_enabled_(tensor_ops_enabled) {} -std::vector ConvolutionThunk::GetAlgorithms( - bool with_winograd_nonfused, se::StreamExecutor* stream_exec) const { - std::vector algorithms; - switch (convolution_kind_) { - case ConvolutionKind::kBackwardFilter: - CHECK(stream_exec->GetConvolveBackwardFilterAlgorithms( - with_winograd_nonfused, &algorithms)); - break; - case ConvolutionKind::kBackwardInput: - CHECK(stream_exec->GetConvolveBackwardDataAlgorithms( - with_winograd_nonfused, &algorithms)); - break; - case ConvolutionKind::kForward: - CHECK(stream_exec->GetConvolveAlgorithms(with_winograd_nonfused, - &algorithms)); - break; - } - return algorithms; -} - -static string AlgorithmToString(const se::dnn::AlgorithmDesc& algo) { - if (algo.tensor_ops_enabled()) { - return tensorflow::strings::StrCat(algo.algo_id(), "+TC"); - } - return tensorflow::strings::StrCat(algo.algo_id()); -} - -// Determines whether we can safely perform a winograd non-fused convolution for -// the given input and output descriptors. This works around b/68264959, an -// integer overflow in cuDNNv5 and cuDNNv6. -static bool ShouldIncludeWinogradNonfusedAlgo( - const BatchDescriptor& input_descriptor, - const BatchDescriptor& output_descriptor) { - int64 batch = input_descriptor.count(); - int64 in_depths = input_descriptor.feature_map_count(); - int64 in_rows = input_descriptor.height(); - int64 in_cols = input_descriptor.width(); - int64 out_depths = output_descriptor.feature_map_count(); - - int64 total_size = 16 * std::ceil(batch / 16.0) * - std::max(in_depths, out_depths) * in_cols * in_rows * - sizeof(float); - int64 threshold = 1L << 31; - - return total_size < threshold; -} - -tensorflow::Status ConvolutionThunk::ConvolveWithTune( - const BatchDescriptor& input_descriptor, se::DeviceMemory input_data, - const FilterDescriptor& filter_descriptor, - se::DeviceMemory filter_data, - const BatchDescriptor& output_descriptor, - se::DeviceMemory output_data, - const ConvolutionDescriptor& convolution_descriptor, +Status ConvolutionThunk::ExecuteOnStream( const BufferAllocations& buffer_allocations, se::Stream* stream) { - // TODO(b/29126320): Try cudnn v5's new auto-tuner when it's rolled out. - if (!best_algorithm_.has_value()) { - best_algorithm_.emplace(); - - // Auto-tuning either is disabled or only happens in the first run of this - // function. - VLOG(2) << "Profiling for best convolution algorithm used for " - "ConvolutionThunk: " - << this; - - bool with_winograd_nonfused = - ShouldIncludeWinogradNonfusedAlgo(input_descriptor, output_descriptor); - - se::dnn::ProfileResult best_result; - se::dnn::ProfileResult best_result_without_scratch; - std::vector algorithms = - GetAlgorithms(with_winograd_nonfused, stream->parent()); - for (auto algorithm : algorithms) { - ConvolveScratchAllocator scratch_allocator( - buffer_allocations.device_ordinal(), - buffer_allocations.memory_allocator()); - se::dnn::ProfileResult profile_result; - VLOG(3) << "Trying algorithm " << AlgorithmToString(algorithm) - << " for ConvolutionThunk: " << this; - bool launch_ok = - Convolve(input_descriptor, input_data, filter_descriptor, filter_data, - output_descriptor, output_data, convolution_descriptor, - se::dnn::AlgorithmConfig(algorithm, algorithm), stream, - &scratch_allocator, &profile_result) - .ok(); - if (launch_ok && profile_result.is_valid()) { - VLOG(3) << "Run of algorithm " << AlgorithmToString(algorithm) - << " for ConvolutionThunk " << this << " succeeded, taking " - << profile_result.elapsed_time_in_ms() - << "ms. (Best result: " << best_result.elapsed_time_in_ms() - << "ms)"; - if (profile_result.elapsed_time_in_ms() < - best_result.elapsed_time_in_ms()) { - best_result = profile_result; - } - if (scratch_allocator.TotalAllocatedBytes() == 0 && - profile_result.elapsed_time_in_ms() < - best_result_without_scratch.elapsed_time_in_ms()) { - best_result_without_scratch = profile_result; - } - } else { - VLOG(3) << "Run of algorithm " << AlgorithmToString(algorithm) - << " for ConvolutionThunk " << this << " failed."; - } - } - - if (best_result.is_valid()) { - best_algorithm_->set_algorithm(best_result.algorithm()); - } else { - LOG(ERROR) << "No convolution algorithm works with profiling. Fall back " - "to the default algorithm."; - best_algorithm_->set_algorithm(AlgorithmDesc()); + se::DeviceMemoryBase input_data = + buffer_allocations.GetDeviceAddress(input_buffer_); + se::DeviceMemoryBase filter_data = + buffer_allocations.GetDeviceAddress(filter_buffer_); + se::DeviceMemoryBase output_data = + buffer_allocations.GetDeviceAddress(output_buffer_); + se::DeviceMemoryBase scratch = + buffer_allocations.GetDeviceAddress(scratch_buffer_); + + se::dnn::AlgorithmConfig algorithm_config( + se::dnn::AlgorithmDesc(algorithm_, tensor_ops_enabled_)); + + TF_RETURN_IF_ERROR(RunCudnnConvolution( + convolution_kind_, input_shape_, filter_shape_, output_shape_, input_data, + filter_data, output_data, scratch, window_, dim_nums_, algorithm_config, + stream)); + + // Figure out which of output/input/filter is the result produced by + // this op, and write the result tuple. + void* result_ptr = [&] { + switch (convolution_kind_) { + case CudnnConvKind::kForward: + return output_data.opaque(); + case CudnnConvKind::kBackwardInput: + return input_data.opaque(); + case CudnnConvKind::kBackwardFilter: + return filter_data.opaque(); } + }(); + void* ptrs[] = {result_ptr, scratch.opaque()}; + se::DeviceMemory tuple_addr( + buffer_allocations.GetDeviceAddress(tuple_result_buffer_)); + stream->ThenMemcpyH2D(ptrs, &tuple_addr); - if (best_result_without_scratch.is_valid()) { - best_algorithm_->set_algorithm_no_scratch( - best_result_without_scratch.algorithm()); - } else { - LOG(ERROR) << "No convolution algorithm without scratch works with " - "profiling. Fall back " - "to the default algorithm."; - best_algorithm_->set_algorithm_no_scratch(AlgorithmDesc()); - } - } - - { - VLOG(2) << "Using convolution algorithm (" - << AlgorithmToString(best_algorithm_->algorithm()) << ", " - << AlgorithmToString(best_algorithm_->algorithm_no_scratch()) - << ") for ConvolutionThunk: " << this; - ConvolveScratchAllocator scratch_allocator( - buffer_allocations.device_ordinal(), - buffer_allocations.memory_allocator()); - return Convolve(input_descriptor, input_data, filter_descriptor, - filter_data, output_descriptor, output_data, - convolution_descriptor, *best_algorithm_, stream, - &scratch_allocator, nullptr); + if (!stream->ok()) { + return InternalError("ConvolutionThunk::ExecuteOnStream failed."); } + return Status::OK(); } } // namespace gpu diff --git a/tensorflow/compiler/xla/service/gpu/convolution_thunk.h b/tensorflow/compiler/xla/service/gpu/convolution_thunk.h index 46c94d0bf1e486fb91e63109efb8e4ba778c4120..900d9cb6243088b56a1825fb3ab8c06cf8d74726 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/convolution_thunk.h @@ -18,6 +18,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" +#include "tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h" #include "tensorflow/compiler/xla/service/gpu/gpu_executable.h" #include "tensorflow/compiler/xla/service/gpu/thunk.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -30,106 +31,47 @@ limitations under the License. namespace xla { namespace gpu { -// A one-time scratch allocator for forward and backward convolution. The -// scratch buffers allocated are released on destruction. -// -// Not thread-safe. -class ConvolveScratchAllocator : public perftools::gputools::ScratchAllocator { - public: - ConvolveScratchAllocator(int device_ordinal, - DeviceMemoryAllocator* memory_allocator); - - ~ConvolveScratchAllocator() override; - - int64 GetMemoryLimitInBytes(perftools::gputools::Stream* stream) override; - - int64 TotalAllocatedBytes() { return total_allocated_bytes_; } - - perftools::gputools::port::StatusOr> - AllocateBytes(perftools::gputools::Stream* stream, int64 byte_size) override; - - private: - const int device_ordinal_; - DeviceMemoryAllocator* memory_allocator_; - std::vector allocated_buffers_; - int64 total_allocated_bytes_ = 0; -}; - // This class stores everything that StreamExecutor needs to launch a BNN // convolution. It is generated by IrEmitter. // // This is thread-compatible. class ConvolutionThunk : public Thunk { public: - // ConvolutionThunk performs one of the following types of convolution. - enum class ConvolutionKind { - kBackwardFilter, // Backward convolution for filter. - kBackwardInput, // Backward convolution for input. - kForward, // Forward convolution. - }; - - // Constructs a thunk for launching a DNN convolution. + // Constructs a thunk for launching a DNN convolution. When run, it will + // write a tuple (result, scratch_memory) into `tuple_result_buffer`. + // + // `algorithm` is a cudnn algorithm number. `algorithm == -1` indicates that + // we should use the default (i.e. baseline) cudnn algorithm. + // + // Note that "output" here doesn't refer to the output from running this + // thunk, but rather to the "output" of a hypothetical forward convolution + // that corresponds to this input+filter+output triple. That is, the result + // generated by this thunk is "output" for forward convs, "input" for + // backward-input convs, and "filter" for backward-filter convs. + // // Semantics of null hlo_instruction argument are as in Thunk. - ConvolutionThunk(ConvolutionKind convolution_kind, + ConvolutionThunk(CudnnConvKind convolution_kind, const BufferAllocation::Slice& input_buffer, const BufferAllocation::Slice& filter_buffer, const BufferAllocation::Slice& output_buffer, + const BufferAllocation::Slice& tuple_result_buffer, + const BufferAllocation::Slice& scratch_buffer, const Shape& input_shape, const Shape& filter_shape, const Shape& output_shape, const Window& window, - const ConvolutionDimensionNumbers& dnums, - const HloInstruction* hlo); + const ConvolutionDimensionNumbers& dim_nums, int64 algorithm, + bool tensor_ops_enabled, const HloInstruction* hlo); ConvolutionThunk(const ConvolutionThunk&) = delete; ConvolutionThunk& operator=(const ConvolutionThunk&) = delete; - // Does the convolution for the thunk on "stream". Auto-tuning happens on the - // first run of this function. - tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; - - // Returns true if the next run of ExecuteOnStream will do autotuning. If so, - // we want the GPU to be quiescent during autotuning, so as not to introduce - // noise in our results. - bool ShouldHaltAllActivityBeforeRunning( - perftools::gputools::Stream*) override { - return !best_algorithm_.has_value(); - } - - // Return true if scratch memory is needed to execute the thunk, that is - // either the best algorithm hasn't been chosen or the best algorithm is not - // the same as the no-scratch algorithm. This is because that the execution - // of the thunk is asynchronous, and the scratch allocator goes out of - // scope before the thunk finishes execution. Returning true tells the stream - // executor to make future thunks wait for this thunk to avoid reusing the - // deallocated scratch memory until this thunk is done with it. - bool ShouldBlockFutureThunks() { - if (!best_algorithm_.has_value()) { - return true; - } - - const perftools::gputools::dnn::AlgorithmDesc& best_alg = - best_algorithm_->algorithm(); - const perftools::gputools::dnn::AlgorithmDesc& no_scratch_best_alg = - best_algorithm_->algorithm_no_scratch(); - return (!best_alg.is_default() || !no_scratch_best_alg.is_default() || - !(best_alg == no_scratch_best_alg)); - } + // Does the convolution for the thunk on "stream". + Status ExecuteOnStream(const BufferAllocations& buffer_allocations, + perftools::gputools::Stream* stream) override; private: - tensorflow::Status ConvolveWithTune( - const perftools::gputools::dnn::BatchDescriptor& input_descriptor, - perftools::gputools::DeviceMemory input_data, - const perftools::gputools::dnn::FilterDescriptor& filter_descriptor, - perftools::gputools::DeviceMemory filter_data, - const perftools::gputools::dnn::BatchDescriptor& output_descriptor, - perftools::gputools::DeviceMemory output_data, - const perftools::gputools::dnn::ConvolutionDescriptor& - convolution_descriptor, - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream); + class ScratchAllocator; - tensorflow::Status Convolve( + Status Convolve( const perftools::gputools::dnn::BatchDescriptor& input_descriptor, perftools::gputools::DeviceMemory input_data, const perftools::gputools::dnn::FilterDescriptor& filter_descriptor, @@ -139,40 +81,27 @@ class ConvolutionThunk : public Thunk { const perftools::gputools::dnn::ConvolutionDescriptor& convolution_descriptor, const perftools::gputools::dnn::AlgorithmConfig& algorithm_config, - perftools::gputools::Stream* stream, - ConvolveScratchAllocator* scratch_allocator, + perftools::gputools::Stream* stream, ScratchAllocator* scratch_allocator, perftools::gputools::dnn::ProfileResult* profile_result); - // Returns the convolve algorithms that can be used for this ConvolutionThunk. - std::vector GetAlgorithms( - bool with_winograd_nonfused, - perftools::gputools::StreamExecutor* stream_exec) const; - - // Fastest cuDNN convolution algorithm for this thunk learned from - // auto-tuning. If auto-tuning is disabled or failed, best_algorithm_ is set - // to the default value, indicating cuDNN's convolution will choose the best - // algorithm from some heuristics based on its parameters. - tensorflow::gtl::optional - best_algorithm_; - - const ConvolutionKind convolution_kind_; + const CudnnConvKind convolution_kind_; const BufferAllocation::Slice input_buffer_; const BufferAllocation::Slice filter_buffer_; const BufferAllocation::Slice output_buffer_; + const BufferAllocation::Slice tuple_result_buffer_; + const BufferAllocation::Slice scratch_buffer_; const Shape input_shape_; const Shape filter_shape_; const Shape output_shape_; const Window window_; - const ConvolutionDimensionNumbers dim_nums_; + int64 algorithm_; + bool tensor_ops_enabled_; }; -string ConvolutionKindToString( - ConvolutionThunk::ConvolutionKind convolution_kind); - } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.cc b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.cc new file mode 100644 index 0000000000000000000000000000000000000000..1792893ae401bf16d2dd9e861607e8f3821a505e --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.cc @@ -0,0 +1,369 @@ +/* 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/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h" +#include "tensorflow/compiler/xla/service/gpu/convolution_thunk.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" +#include "tensorflow/core/lib/gtl/optional.h" +#include "tensorflow/core/lib/strings/numbers.h" +#include "tensorflow/core/lib/strings/strcat.h" + +namespace xla { +namespace gpu { +namespace { + +namespace se = perftools::gputools; + +using se::DeviceMemoryBase; +using se::dnn::AlgorithmConfig; +using se::dnn::AlgorithmDesc; +using tensorflow::gtl::nullopt; +using tensorflow::gtl::optional; + +class ScratchAllocator : public se::ScratchAllocator { + public: + ScratchAllocator(int device_ordinal, DeviceMemoryAllocator* memory_allocator) + : device_ordinal_(device_ordinal), memory_allocator_(memory_allocator) {} + + ~ScratchAllocator() override; + + int64 GetMemoryLimitInBytes(se::Stream* stream) override { + return 1LL << 32; // 4GB. TODO(jlebar): Tune this? + } + int64 TotalAllocatedBytes() { return total_allocated_bytes_; } + + se::port::StatusOr> AllocateBytes( + se::Stream* stream, int64 byte_size) override; + + private: + const int device_ordinal_; + DeviceMemoryAllocator* memory_allocator_; + std::vector allocated_buffers_; + int64 total_allocated_bytes_ = 0; +}; + +ScratchAllocator::~ScratchAllocator() { + for (auto& allocated_buffer : allocated_buffers_) { + if (!memory_allocator_->Deallocate(device_ordinal_, &allocated_buffer) + .ok()) { + // The program can still continue with failed deallocation. + LOG(ERROR) << "Failed to deallocate the allocated buffer: " + << allocated_buffer.opaque(); + } + } +} + +se::port::StatusOr> ScratchAllocator::AllocateBytes( + se::Stream* stream, int64 byte_size) { + CHECK_GE(byte_size, 0) << "byte_size must be positive."; + if (byte_size > GetMemoryLimitInBytes(stream)) { + return se::port::Status( + se::port::error::RESOURCE_EXHAUSTED, + tensorflow::strings::Printf( + "Allocating %lld bytes exceeds the memory limit of %lld bytes.", + byte_size, GetMemoryLimitInBytes(stream))); + } + + auto status_or_memory = + memory_allocator_->Allocate(device_ordinal_, byte_size, + /*retry_on_failure=*/false); + if (!status_or_memory.ok()) { + return se::port::Status(se::port::error::RESOURCE_EXHAUSTED, + tensorflow::strings::Printf( + "Failed to allocate %lld bytes on device %d.", + byte_size, device_ordinal_)); + } + se::DeviceMemoryBase allocated_buffer = status_or_memory.ValueOrDie(); + allocated_buffers_.push_back(allocated_buffer); + total_allocated_bytes_ += byte_size; + return se::DeviceMemory(allocated_buffer); +} + +// Determines whether we can safely perform a winograd non-fused convolution for +// the given input and output shapes. This works around b/68264959, an integer +// overflow in cuDNNv5 and cuDNNv6. +// +// TODO(jlebar): We shouldn't need this check for cuDNNv7. +bool ShouldIncludeWinogradNonfusedAlgo( + const Shape& input_shape, const Shape& output_shape, + const ConvolutionDimensionNumbers& dnums) { + int64 batch = input_shape.dimensions(dnums.input_batch_dimension()); + int64 in_depths = input_shape.dimensions(dnums.input_feature_dimension()); + int64 in_rows = input_shape.dimensions(dnums.input_spatial_dimensions(0)); + int64 in_cols = + dnums.input_spatial_dimensions_size() == 1 + ? 1 + : input_shape.dimensions(dnums.input_spatial_dimensions(1)); + int64 out_depths = output_shape.dimensions(dnums.output_feature_dimension()); + + int64 total_size = CeilOfRatio(batch, int64{16}) * + std::max(in_depths, out_depths) * in_cols * in_rows * + sizeof(float); + + const int64 threshold = 1L << 31; + return total_size < threshold; +} + +std::vector GetAlgorithms(CudnnConvKind kind, + bool with_winograd_nonfused, + se::StreamExecutor* stream_exec_) { + std::vector algorithms; + switch (kind) { + case CudnnConvKind::kBackwardFilter: + CHECK(stream_exec_->GetConvolveBackwardFilterAlgorithms( + with_winograd_nonfused, &algorithms)); + break; + case CudnnConvKind::kBackwardInput: + CHECK(stream_exec_->GetConvolveBackwardDataAlgorithms( + with_winograd_nonfused, &algorithms)); + break; + case CudnnConvKind::kForward: + CHECK(stream_exec_->GetConvolveAlgorithms(with_winograd_nonfused, + &algorithms)); + break; + } + + return algorithms; +} + +string AlgorithmToString(const AlgorithmDesc& algo) { + if (algo.tensor_ops_enabled()) { + return tensorflow::strings::StrCat(algo.algo_id(), "+TC"); + } + return tensorflow::strings::StrCat(algo.algo_id()); +} + +string NumBytesToString(int64 bytes) { + return tensorflow::strings::StrCat( + tensorflow::strings::HumanReadableNumBytes(bytes), " (", bytes, "B)"); +} + +} // anonymous namespace + +// We could have caching here so that we don't redo this work for two identical +// convolutions. Unfortunately our cache key would have to be a tuple +// containing the protos passed to this function, and we have no utility for +// hashing protos. We could write our own hash functions, but they'd silently +// break if we ever added a field to one of the protos. Perhaps we could hack +// using the binary-encoded proto as the hash key, on the assumption that two +// protos being binary-equal is a sufficient, if not necessary, condition for +// proper equality. But that would still leave us open to having unnecessary +// cache misses and doing extra work. Overall, caching doesn't seem worth the +// trouble, but we may want to revisit this if we ever find a model where +// caching would speed up compilation a lot. +optional> +CudnnConvolutionAlgorithmPicker::PickBestAlgorithm( + CudnnConvKind kind, const Shape& input_shape, const Shape& filter_shape, + const Shape& output_shape, const Window& window, + const ConvolutionDimensionNumbers& dnums, HloInstruction* instr) { + // Create a stream for us to do our work on. + se::Stream stream{stream_exec_}; + stream.Init(); + const auto device_ordinal = stream_exec_->device_ordinal(); + + // allocator either points to this->allocator_ or, if that's null, to a + // StreamExecutorMemoryAllocator for stream_exec_. + DeviceMemoryAllocator* allocator; + optional se_allocator; + if (allocator_ != nullptr) { + allocator = allocator_; + } else { + se_allocator.emplace( + stream_exec_->platform(), + tensorflow::gtl::ArraySlice({stream_exec_})); + allocator = &*se_allocator; + } + + // Allocate space for the input, filter, and output of the convolution. We + // use a ScratchAllocator for this instead of calling allocator_ directly so + // that our allocations don't leak. + // + // We don't put any data in these buffers, because (in theory, anyway) the + // speed of a conv isn't affected by the data being convolved. + ScratchAllocator input_output_allocator(device_ordinal, allocator); + se::port::StatusOr input_buf = + input_output_allocator.AllocateBytes(&stream, + ShapeUtil::ByteSizeOf(input_shape)); + se::port::StatusOr filter_buf = + input_output_allocator.AllocateBytes(&stream, + ShapeUtil::ByteSizeOf(filter_shape)); + se::port::StatusOr output_buf = + input_output_allocator.AllocateBytes(&stream, + ShapeUtil::ByteSizeOf(output_shape)); + if (!input_buf.ok() || !filter_buf.ok() || !output_buf.ok()) { + LOG(WARNING) + << "Couldn't allocate space for input/filter/output of convolution " + << instr->ToString() << ". Falling back to default algorithm."; + return nullopt; + } + + const bool use_winograd_nonfused = + ShouldIncludeWinogradNonfusedAlgo(input_shape, output_shape, dnums); + se::dnn::ProfileResult best_result; + int64 best_result_bytes_used = 0; + + for (const AlgorithmDesc& alg : + GetAlgorithms(kind, use_winograd_nonfused, stream_exec_)) { + ScratchAllocator scratch_allocator(device_ordinal, allocator); + se::dnn::ProfileResult profile_result; + VLOG(3) << "Trying algorithm " << AlgorithmToString(alg) << " for " + << instr->ToString(); + + bool launch_ok = RunCudnnConvolution( + kind, input_shape, filter_shape, output_shape, + input_buf.ValueOrDie(), filter_buf.ValueOrDie(), + output_buf.ValueOrDie(), &scratch_allocator, window, + dnums, AlgorithmConfig(alg), &stream, &profile_result) + .ok(); + + if (launch_ok && profile_result.is_valid()) { + int64 scratch_bytes_used = scratch_allocator.TotalAllocatedBytes(); + VLOG(3) << "Run of algorithm " << AlgorithmToString(alg) + << " succeeded, taking " << profile_result.elapsed_time_in_ms() + << "ms and using " << NumBytesToString(scratch_bytes_used) + << " of scratch (Best result: " + << best_result.elapsed_time_in_ms() << "ms, " + << NumBytesToString(best_result_bytes_used) << " of scratch)"; + if (profile_result.elapsed_time_in_ms() < + best_result.elapsed_time_in_ms()) { + best_result = profile_result; + best_result_bytes_used = scratch_bytes_used; + } + } else { + VLOG(3) << "Run of algorithm " << AlgorithmToString(alg) << " failed."; + } + } + if (best_result.is_valid()) { + VLOG(2) << "Best algorithm for " << instr->ToString() << ": " + << AlgorithmToString(best_result.algorithm()) << ", takes " + << best_result.elapsed_time_in_ms() << "ms, and uses " + << best_result_bytes_used << "B of scratch memory."; + return std::make_tuple(best_result.algorithm().algo_id(), + best_result.algorithm().tensor_ops_enabled(), + best_result_bytes_used); + } + + LOG(WARNING) << "All algorithms tried for convolution " << instr->ToString() + << " failed. Falling back to default algorithm."; + return nullopt; +} + +StatusOr CudnnConvolutionAlgorithmPicker::RunOnInstruction( + HloInstruction* instr) { + CHECK(IsCustomCallToDnnConvolution(*instr)); + + const auto& call_target = instr->custom_call_target(); + const auto& lhs_shape = instr->operand(0)->shape(); + const auto& rhs_shape = instr->operand(1)->shape(); + const auto& conv_result_shape = instr->shape().tuple_shapes(0); + optional> alg_scratch_and_tc; + if (call_target == kCudnnConvForwardCallTarget) { + alg_scratch_and_tc = PickBestAlgorithm( + CudnnConvKind::kForward, /*input_shape=*/lhs_shape, + /*filter_shape=*/rhs_shape, /*output_shape=*/conv_result_shape, + instr->window(), instr->convolution_dimension_numbers(), instr); + } else if (call_target == kCudnnConvBackwardInputCallTarget) { + alg_scratch_and_tc = PickBestAlgorithm( + CudnnConvKind::kBackwardInput, /*input_shape=*/conv_result_shape, + /*filter_shape=*/rhs_shape, /*output_shape=*/lhs_shape, instr->window(), + instr->convolution_dimension_numbers(), instr); + } else if (call_target == kCudnnConvBackwardFilterCallTarget) { + alg_scratch_and_tc = PickBestAlgorithm( + CudnnConvKind::kBackwardFilter, /*input_shape=*/lhs_shape, + /*filter_shape=*/conv_result_shape, /*output_shape=*/rhs_shape, + instr->window(), instr->convolution_dimension_numbers(), instr); + } else { + LOG(FATAL) << "Unknown custom call target for cudnn conv: " + << instr->ToString(); + } + + if (!alg_scratch_and_tc.has_value()) { + return false; + } + + int64 algorithm; + bool tensor_ops_enabled; + int64 scratch_bytes; + + std::tie(algorithm, tensor_ops_enabled, scratch_bytes) = *alg_scratch_and_tc; + + VLOG(1) << "Setting cudnn conv to use algorithm " << algorithm << " and " + << NumBytesToString(scratch_bytes) + << " of scratch memory: " << instr->ToString() + << " tensor_ops_enabled: " << tensor_ops_enabled; + + // Replace instr with a new CustomCall which has the correct algorithm, and + // whose output shape has the appropriate amount of scratch memory. + HloComputation* computation = instr->parent(); + Shape new_call_shape = + ShapeUtil::MakeTupleShape({instr->shape().tuple_shapes(0), + ShapeUtil::MakeShape(U8, {scratch_bytes})}); + HloInstruction* algorithm_hlo = computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(algorithm))); + HloInstruction* tensor_ops_enabled_hlo = + computation->AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR0(tensor_ops_enabled))); + + HloInstruction* new_call = + computation->AddInstruction(HloInstruction::CreateCustomCall( + new_call_shape, + {instr->mutable_operand(0), instr->mutable_operand(1), algorithm_hlo, + tensor_ops_enabled_hlo}, + instr->custom_call_target())); + new_call->set_window(instr->window()); + new_call->set_convolution_dimension_numbers( + instr->convolution_dimension_numbers()); + + // Repackage new_call so it has the same shape as the original call, namely + // (conv_result, u8[0]). + HloInstruction* new_tuple = + computation->AddInstruction(HloInstruction::CreateTuple( + {computation->AddInstruction(HloInstruction::CreateGetTupleElement( + new_call_shape.tuple_shapes(0), new_call, 0)), + computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({})))})); + + TF_RETURN_IF_ERROR(instr->parent()->ReplaceInstruction(instr, new_tuple)); + return true; +} + +StatusOr CudnnConvolutionAlgorithmPicker::RunOnComputation( + HloComputation* computation) { + std::vector convs; + for (auto* instr : computation->instructions()) { + if (IsCustomCallToDnnConvolution(*instr)) { + convs.push_back(instr); + } + } + + bool changed = false; + for (auto* instr : convs) { + TF_ASSIGN_OR_RETURN(bool result, RunOnInstruction(instr)); + changed |= result; + } + return changed; +} + +StatusOr CudnnConvolutionAlgorithmPicker::Run(HloModule* module) { + bool changed = false; + for (HloComputation* computation : module->MakeNonfusionComputations()) { + TF_ASSIGN_OR_RETURN(bool result, RunOnComputation(computation)); + changed |= result; + } + return changed; +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h new file mode 100644 index 0000000000000000000000000000000000000000..516210ec2e500cf03774d27408300ac3346e7b4f --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h @@ -0,0 +1,62 @@ +/* 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_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_ALGORITHM_PICKER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_ALGORITHM_PICKER_H_ + +#include "tensorflow/compiler/xla/service/device_memory_allocator.h" +#include "tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/core/lib/gtl/optional.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { +namespace gpu { + +// Modifies CustomCalls to cudnn convolutions, choosing the best algorithm for +// each and adding explicit scratch space to the CustomCalls. +class CudnnConvolutionAlgorithmPicker : public HloPassInterface { + public: + // If the `allocator` parameter is not null, we will use it to allocate temp + // memory while timing the various convolution algorithms. If it's null, + // we'll use the default allocator on the StreamExecutor. + CudnnConvolutionAlgorithmPicker( + perftools::gputools::StreamExecutor* stream_exec, + DeviceMemoryAllocator* allocator) + : stream_exec_(stream_exec), allocator_(allocator) {} + + tensorflow::StringPiece name() const override { + return "cudnn-convolution-algorithm-picker"; + } + + StatusOr Run(HloModule* module) override; + + private: + StatusOr RunOnComputation(HloComputation* computation); + StatusOr RunOnInstruction(HloInstruction* instr); + tensorflow::gtl::optional> PickBestAlgorithm( + CudnnConvKind kind, const Shape& input_shape, const Shape& filter_shape, + const Shape& output_shape, const Window& window, + const ConvolutionDimensionNumbers& dnums, HloInstruction* instr); + + perftools::gputools::StreamExecutor* stream_exec_; // never null + DeviceMemoryAllocator* allocator_; // may be null +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_ALGORITHM_PICKER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/convolution_folding.cc b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.cc similarity index 83% rename from tensorflow/compiler/xla/service/gpu/convolution_folding.cc rename to tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.cc index b0626ca3bc9f843e513d4727932f0e2d5fa37748..e0c73aa73acb7f3313eb54fb07390cb76590433e 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_folding.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.cc @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* 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. @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/service/gpu/convolution_folding.h" +#include "tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.h" #include #include @@ -33,14 +33,32 @@ namespace xla { namespace gpu { namespace { + +bool CanImplementAsCudnnForwardConv(HloInstruction* conv) { + const ConvolutionDimensionNumbers& dnums = + conv->convolution_dimension_numbers(); + if (dnums.input_spatial_dimensions_size() > 3) { + return false; + } + + // CuDNN does not accept zero-element arguments + if (ShapeUtil::HasZeroElements(conv->operand(0)->shape()) || + ShapeUtil::HasZeroElements(conv->operand(1)->shape())) { + return false; + } + + if (window_util::HasWindowReversal(conv->window())) { + return false; + } + return true; +} + // Try to match a backward filter pattern that contains "conv". // Precondition: "conv" is a kConvolution. -std::tuple, Window, - ConvolutionDimensionNumbers> -MatchBackwardFilter(HloInstruction* conv) { +std::tuple MatchBackwardFilter( + HloInstruction* conv) { const auto no_match_result = - std::make_tuple(false, std::vector(), Window(), - ConvolutionDimensionNumbers()); + std::make_tuple(false, Window(), ConvolutionDimensionNumbers()); // Step 1: match the instruction pattern without considering the paddings and // dimension numbers just yet. We may need some generic pattern matcher // similar to third_party/llvm/llvm/include/llvm/IR/PatternMatch.h @@ -190,18 +208,15 @@ MatchBackwardFilter(HloInstruction* conv) { backward_conv_dnums.add_kernel_spatial_dimensions(output_spatial_dims[i]); } - return std::make_tuple(true, std::vector({conv}), - backward_conv_window, backward_conv_dnums); + return std::make_tuple(true, backward_conv_window, backward_conv_dnums); } // Try to match a backward input pattern that contains "conv". // Precondition: "conv" is a kConvolution. -std::tuple, Window, - ConvolutionDimensionNumbers> -MatchBackwardInput(HloInstruction* conv) { +std::tuple MatchBackwardInput( + HloInstruction* conv) { const auto no_match_result = - std::make_tuple(false, std::vector(), Window(), - ConvolutionDimensionNumbers()); + std::make_tuple(false, Window(), ConvolutionDimensionNumbers()); // Match instruction pattern. CHECK_EQ(HloOpcode::kConvolution, conv->opcode()); @@ -374,58 +389,82 @@ MatchBackwardInput(HloInstruction* conv) { dnums.set_kernel_output_feature_dimension( conv->convolution_dimension_numbers().kernel_input_feature_dimension()); - return std::make_tuple(true, - std::vector({conv, reverse_filter}), - new_window, dnums); + return std::make_tuple(true, new_window, dnums); } -} // namespace -StatusOr ConvolutionFolding::Run(HloModule* module) { - HloComputation* entry_computation = module->entry_computation(); - std::vector convs; - for (auto* hlo : entry_computation->instructions()) { - if (hlo->opcode() == HloOpcode::kConvolution) { - convs.push_back(hlo); - } - } +// Tries to rewrite a single convolution into a call to cudnn. +StatusOr RunOnInstruction(HloInstruction* conv) { + CHECK_EQ(conv->opcode(), HloOpcode::kConvolution); - bool changed = false; - for (HloInstruction* conv : convs) { + HloInstruction* custom_call = [&]() -> HloInstruction* { bool match; - std::vector hlos_to_fuse; Window window; ConvolutionDimensionNumbers dnums; - std::tie(match, hlos_to_fuse, window, dnums) = MatchBackwardFilter(conv); + + std::tie(match, window, dnums) = MatchBackwardFilter(conv); if (match) { - VLOG(2) << "Fuse instructions"; - for (HloInstruction* hlo_to_fuse : hlos_to_fuse) { - VLOG(2) << " " << hlo_to_fuse->ToString(); - } - HloInstruction* backward_convolution = - entry_computation->CreateFusionInstructionForBackwardConvolution( - hlos_to_fuse, HloInstruction::FusionKind::kConvBackwardFilter, - window, dnums); - VLOG(2) << "to backward filter convolution"; - VLOG(2) << " " << backward_convolution->ToString(); - changed = true; - continue; + return CreateCudnnConvBackwardFilter( + conv->shape(), conv->mutable_operand(0), conv->mutable_operand(1), + window, dnums); } - std::tie(match, hlos_to_fuse, window, dnums) = MatchBackwardInput(conv); + std::tie(match, window, dnums) = MatchBackwardInput(conv); if (match) { - VLOG(2) << "Fuse instructions"; - for (HloInstruction* hlo_to_fuse : hlos_to_fuse) { - VLOG(2) << " " << hlo_to_fuse->ToString(); - } - HloInstruction* backward_convolution = - entry_computation->CreateFusionInstructionForBackwardConvolution( - hlos_to_fuse, HloInstruction::FusionKind::kConvBackwardInput, - window, dnums); - VLOG(2) << "to backward input convolution"; - VLOG(2) << " " << backward_convolution->ToString(); - changed = true; - continue; + // Backward input conv subsumes the conv plus the reverse in operand 1. + HloInstruction* reverse = conv->mutable_operand(1); + CHECK_EQ(reverse->opcode(), HloOpcode::kReverse); + HloInstruction* rhs = reverse->mutable_operand(0); + + return CreateCudnnConvBackwardInput( + conv->shape(), conv->mutable_operand(0), rhs, window, dnums); } + + // If all else fails, try a forward convolution. + if (CanImplementAsCudnnForwardConv(conv)) { + return CreateCudnnConvForward(conv->shape(), conv->mutable_operand(0), + conv->mutable_operand(1), conv->window(), + conv->convolution_dimension_numbers()); + } + + return nullptr; + }(); + + if (custom_call == nullptr) { + return false; + } + + // The CustomCall returns a tuple (conv_result, scratch_memory). Extract out + // the conv result and replace `conv` with it. + TF_RETURN_IF_ERROR(conv->parent()->ReplaceWithNewInstruction( + conv, + HloInstruction::CreateGetTupleElement(conv->shape(), custom_call, 0))); + return true; +} + +// Rewrites the convolutions in the given computation into calls to cudnn. +// Returns true if it made any changes. +StatusOr RunOnComputation(HloComputation* computation) { + std::vector convs; + for (auto* hlo : computation->instructions()) { + if (hlo->opcode() == HloOpcode::kConvolution) { + convs.push_back(hlo); + } + } + + bool changed = false; + for (HloInstruction* conv : convs) { + TF_ASSIGN_OR_RETURN(bool result, RunOnInstruction(conv)); + changed |= result; + } + return changed; +} +} // namespace + +StatusOr CudnnConvolutionRewriter::Run(HloModule* module) { + bool changed = false; + for (HloComputation* computation : module->MakeNonfusionComputations()) { + TF_ASSIGN_OR_RETURN(bool result, RunOnComputation(computation)); + changed |= result; } return changed; } diff --git a/tensorflow/compiler/xla/service/gpu/convolution_folding.h b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.h similarity index 63% rename from tensorflow/compiler/xla/service/gpu/convolution_folding.h rename to tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.h index f9c898721f8dd6b8b7e74c82bb2085cc437eaad5..0c0578d88840fed1d77f7456c9acef27dec380f5 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_folding.h +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.h @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* 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. @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CONVOLUTION_FOLDING_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CONVOLUTION_FOLDING_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_REWRITER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_REWRITER_H_ #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" @@ -22,10 +22,12 @@ limitations under the License. namespace xla { namespace gpu { -class ConvolutionFolding : public HloPassInterface { +// Rewrites plain convolutions, backwards-filter convolutions, and +// backwards-input convolutions into CustomCall HLOs that call into cuDNN. +class CudnnConvolutionRewriter : public HloPassInterface { public: tensorflow::StringPiece name() const override { - return "convolution-folding"; + return "cudnn-convolution-rewriter"; } StatusOr Run(HloModule* module) override; @@ -34,4 +36,4 @@ class ConvolutionFolding : public HloPassInterface { } // namespace gpu } // namespace xla -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CONVOLUTION_FOLDING_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_REWRITER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/convolution_folding_test.cc b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter_test.cc similarity index 82% rename from tensorflow/compiler/xla/service/gpu/convolution_folding_test.cc rename to tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter_test.cc index 34e6bdb117d47a3d7e1eb3bae5806e130e94ea79..65588b6aaf24da628ea586eb52c462b78b8daaa7 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_folding_test.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter_test.cc @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* 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. @@ -13,23 +13,29 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/service/gpu/convolution_folding.h" +#include "tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/shape_inference.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/platform/test.h" namespace xla { namespace gpu { +namespace { -class ConvolutionFoldingTest : public HloTestBase { +namespace op = xla::testing::opcode_matchers; + +class CudnnConvolutionRewriterTest : public HloTestBase { public: - ConvolutionFoldingTest() { + CudnnConvolutionRewriterTest() { for (int i = 0; i < 2; ++i) { WindowDimension* window_dim = default_conv_window_.add_dimensions(); window_dim->set_size(1); @@ -44,7 +50,8 @@ class ConvolutionFoldingTest : public HloTestBase { // the batch and feature dimension in the activations, and treat the batch // dimension in gradients as the input feature dimension in the filter. // - // TODO(jingyue): Add more tests on NCHW input order which TF also supports. + // TODO(jingyue): Add more tests on NCHW input order, which TF also + // supports. tf_default_dnums_for_backward_filter_.set_input_batch_dimension(3); tf_default_dnums_for_backward_filter_.set_input_feature_dimension(0); tf_default_dnums_for_backward_filter_.add_input_spatial_dimensions(1); @@ -74,9 +81,8 @@ class ConvolutionFoldingTest : public HloTestBase { } protected: - bool FoldConvolution(HloModule* module) { - ConvolutionFolding convolution_folding; - return convolution_folding.Run(module).ValueOrDie(); + bool RunPass(HloModule* module) { + return CudnnConvolutionRewriter().Run(module).ValueOrDie(); } // A convolution window with stride 1 and zero padding. The size fields are @@ -86,7 +92,7 @@ class ConvolutionFoldingTest : public HloTestBase { ConvolutionDimensionNumbers tf_default_dnums_for_backward_input_; }; -TEST_F(ConvolutionFoldingTest, BackwardFilterConvolve) { +TEST_F(CudnnConvolutionRewriterTest, BackwardFilterConvolve) { HloComputation::Builder builder(TestName()); HloInstruction* activations = builder.AddInstruction(HloInstruction::CreateParameter( @@ -108,14 +114,13 @@ TEST_F(ConvolutionFoldingTest, BackwardFilterConvolve) { auto module = CreateNewModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); - EXPECT_EQ(HloOpcode::kFusion, - entry_computation->root_instruction()->opcode()); - EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardFilter == - entry_computation->root_instruction()->fusion_kind()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); } -TEST_F(ConvolutionFoldingTest, +TEST_F(CudnnConvolutionRewriterTest, BackwardFilterConvolveEquivalentToForwardConvolution) { HloComputation::Builder builder(TestName()); HloInstruction* activations = @@ -135,12 +140,17 @@ TEST_F(ConvolutionFoldingTest, tf_default_dnums_for_backward_filter_)); auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); + HloComputation* entry_computation = + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); } // Extracted from block35 training. -TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithPaddedActivations) { +TEST_F(CudnnConvolutionRewriterTest, + BackwardFilterConvolveWithPaddedActivations) { auto builder = HloComputation::Builder(TestName()); HloInstruction* activations = builder.AddInstruction(HloInstruction::CreateParameter( @@ -162,15 +172,15 @@ TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithPaddedActivations) { auto module = CreateNewModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); - EXPECT_EQ(HloOpcode::kFusion, - entry_computation->root_instruction()->opcode()); - EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardFilter == - entry_computation->root_instruction()->fusion_kind()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); } // Extracted from inception v3 training. -TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithPaddedGradients) { +TEST_F(CudnnConvolutionRewriterTest, + BackwardFilterConvolveWithPaddedGradients) { auto builder = HloComputation::Builder(TestName()); HloInstruction* activations = builder.AddInstruction(HloInstruction::CreateParameter( @@ -192,14 +202,13 @@ TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithPaddedGradients) { auto module = CreateNewModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); - EXPECT_EQ(HloOpcode::kFusion, - entry_computation->root_instruction()->opcode()); - EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardFilter == - entry_computation->root_instruction()->fusion_kind()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); } -TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithUnevenPadding) { +TEST_F(CudnnConvolutionRewriterTest, BackwardFilterConvolveWithUnevenPadding) { auto builder = HloComputation::Builder(TestName()); HloInstruction* activations = builder.AddInstruction(HloInstruction::CreateParameter( @@ -221,14 +230,13 @@ TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithUnevenPadding) { auto module = CreateNewModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); - EXPECT_EQ(HloOpcode::kFusion, - entry_computation->root_instruction()->opcode()); - EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardFilter == - entry_computation->root_instruction()->fusion_kind()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); } -TEST_F(ConvolutionFoldingTest, BackwardInputConvolveEvenPadding) { +TEST_F(CudnnConvolutionRewriterTest, BackwardInputConvolveEvenPadding) { auto builder = HloComputation::Builder(TestName()); HloInstruction* output = builder.AddInstruction(HloInstruction::CreateParameter( @@ -272,14 +280,15 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveEvenPadding) { auto module = CreateNewModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); - EXPECT_EQ(HloOpcode::kFusion, - entry_computation->root_instruction()->opcode()); - EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardInput == - entry_computation->root_instruction()->fusion_kind()); + EXPECT_TRUE(RunPass(module.get())); + + ASSERT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardInputCallTarget), 0)); + const HloInstruction* custom_call = + entry_computation->root_instruction()->operand(0); for (int i = 0; i < 2; ++i) { - const WindowDimension& window_dim = - entry_computation->root_instruction()->window().dimensions(i); + const WindowDimension& window_dim = custom_call->window().dimensions(i); // Low padding of the backward input convolution // = kernel_size - 1 - low padding on gradients. EXPECT_EQ(3, window_dim.padding_low()); @@ -291,7 +300,7 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveEvenPadding) { // Convolve([abc], [x], base_dilation=2) // = Convolve([abc], Reverse([x]), base_dilation=2) // = BackwardInputConvolve([abc], [x], stride=2) -TEST_F(ConvolutionFoldingTest, BackwardInputConvolve1x1Filter) { +TEST_F(CudnnConvolutionRewriterTest, BackwardInputConvolve1x1Filter) { auto builder = HloComputation::Builder(TestName()); // NHWC dimension order. HloInstruction* output = @@ -316,17 +325,16 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolve1x1Filter) { auto module = CreateNewModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); - EXPECT_EQ(HloOpcode::kFusion, - entry_computation->root_instruction()->opcode()); - EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardInput == - entry_computation->root_instruction()->fusion_kind()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardInputCallTarget), 0)); } // BackwardInputConvolve([abc], [x], stride=1) is equivalent to // ForwardConvolve([abc], [x], stride=1). No need to fold it into backward input // convolution. -TEST_F(ConvolutionFoldingTest, +TEST_F(CudnnConvolutionRewriterTest, BackwardInputConvolve1x1FilterEquivalentToForwardConvolve) { auto builder = HloComputation::Builder(TestName()); // NHWC dimension order. @@ -347,8 +355,12 @@ TEST_F(ConvolutionFoldingTest, tf_default_dnums_for_backward_input_)); auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConvolution(module.get())); + HloComputation* entry_computation = + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT( + entry_computation->root_instruction(), + op::GetTupleElement(op::CustomCall(kCudnnConvForwardCallTarget), 0)); } // Extracted from Inception V3 training. @@ -365,7 +377,8 @@ TEST_F(ConvolutionFoldingTest, // 20x10x10x192 // // Gradients are padded unevenly. -TEST_F(ConvolutionFoldingTest, BackwardInputConvolveUnevenPaddingOnGradients) { +TEST_F(CudnnConvolutionRewriterTest, + BackwardInputConvolveUnevenPaddingOnGradients) { auto builder = HloComputation::Builder(TestName()); HloInstruction* output = builder.AddInstruction(HloInstruction::CreateParameter( @@ -397,14 +410,14 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveUnevenPaddingOnGradients) { auto module = CreateNewModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); - EXPECT_EQ(HloOpcode::kFusion, - entry_computation->root_instruction()->opcode()); - EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardInput == - entry_computation->root_instruction()->fusion_kind()); + EXPECT_TRUE(RunPass(module.get())); + ASSERT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardInputCallTarget), 0)); + const HloInstruction* custom_call = + entry_computation->root_instruction()->operand(0); for (int i = 0; i < 2; ++i) { - const WindowDimension& window_dim = - entry_computation->root_instruction()->window().dimensions(i); + const WindowDimension& window_dim = custom_call->window().dimensions(i); EXPECT_EQ(0, window_dim.padding_low()); EXPECT_EQ(0, window_dim.padding_high()); EXPECT_EQ(2, window_dim.stride()); @@ -413,7 +426,7 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveUnevenPaddingOnGradients) { // Similar to BackwardInputConvolveUnevenPadding, but the low padding of the // gradients exceeds kernel_size - 1. Therefore, this pattern cannot be fused. -TEST_F(ConvolutionFoldingTest, BackwardInputConvolveLowPaddingTooLarge) { +TEST_F(CudnnConvolutionRewriterTest, BackwardInputConvolveLowPaddingTooLarge) { auto builder = HloComputation::Builder(TestName()); HloInstruction* output = builder.AddInstruction(HloInstruction::CreateParameter( @@ -442,8 +455,12 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveLowPaddingTooLarge) { .ValueOrDie())); auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConvolution(module.get())); + HloComputation* entry_computation = + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT( + entry_computation->root_instruction(), + op::GetTupleElement(op::CustomCall(kCudnnConvForwardCallTarget), 0)); } // Extracted from //learning/brain/google/xla/benchmarks/resnet.py @@ -460,7 +477,7 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveLowPaddingTooLarge) { // // We should fuse BC even though padding on activations is uneven, because // PadInsertion will canonicalize the fusion HLO. -TEST_F(ConvolutionFoldingTest, +TEST_F(CudnnConvolutionRewriterTest, BackwardInputConvolveUnevenPaddingOnActivations) { auto builder = HloComputation::Builder(TestName()); // The gradients are in NCHW layout. @@ -493,13 +510,12 @@ TEST_F(ConvolutionFoldingTest, auto module = CreateNewModule(); const HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(module.get())); - const HloInstruction* backward_conv = entry_computation->root_instruction(); - EXPECT_EQ(HloOpcode::kFusion, backward_conv->opcode()); - EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardInput == - backward_conv->fusion_kind()); + EXPECT_TRUE(RunPass(module.get())); + ASSERT_THAT(entry_computation->root_instruction(), + op::GetTupleElement( + op::CustomCall(kCudnnConvBackwardInputCallTarget), 0)); const WindowDimension& backward_conv_col_dim = - backward_conv->window().dimensions(1); + entry_computation->root_instruction()->operand(0)->window().dimensions(1); EXPECT_EQ(0, backward_conv_col_dim.padding_low()); EXPECT_EQ(1, backward_conv_col_dim.padding_high()); } @@ -515,7 +531,7 @@ TEST_F(ConvolutionFoldingTest, // // We currently don't fuse BC because PadInsertion doesn't support negative // padding on the gradients of backward convolution (b/32744257). -TEST_F(ConvolutionFoldingTest, +TEST_F(CudnnConvolutionRewriterTest, BackwardInputConvolveNegativePaddingHighOnActivations) { auto builder = HloComputation::Builder(TestName()); // The gradients are in NCHW layout. @@ -544,9 +560,14 @@ TEST_F(ConvolutionFoldingTest, .ValueOrDie())); auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConvolution(module.get())); + HloComputation* entry_computation = + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(RunPass(module.get())); + EXPECT_THAT( + entry_computation->root_instruction(), + op::GetTupleElement(op::CustomCall(kCudnnConvForwardCallTarget), 0)); } +} // anonymous namespace } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.cc b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.cc new file mode 100644 index 0000000000000000000000000000000000000000..e4ae839e1dd4cb3a744a3f6a3329cabdaeb3f38d --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.cc @@ -0,0 +1,262 @@ +/* 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/compiler/xla/service/gpu/cudnn_convolution_runner.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/util.h" + +namespace xla { +namespace gpu { +namespace { + +namespace se = ::perftools::gputools; + +using se::DeviceMemory; +using se::DeviceMemoryBase; +using se::Stream; +using se::dnn::AlgorithmConfig; +using se::dnn::BatchDescriptor; +using se::dnn::ConvolutionDescriptor; +using se::dnn::DataLayout; +using se::dnn::DimIndex; +using se::dnn::FilterDescriptor; +using se::dnn::FilterLayout; +using se::dnn::ProfileResult; + +// A StreamExecutor ScratchAllocator that wraps a single XLA allocation, +// returning it (in its entirety) the first time Allocate() is called. +class ScratchBufAllocator : public se::ScratchAllocator { + public: + explicit ScratchBufAllocator(se::DeviceMemoryBase scratch) + : scratch_(scratch) {} + + ~ScratchBufAllocator() override = default; + + int64 GetMemoryLimitInBytes(se::Stream* /*stream*/) override { + return scratch_.size(); + } + + se::port::StatusOr> AllocateBytes( + se::Stream* stream, int64 byte_size) override { + if (allocated_) { + return se::port::InternalError( + "Can't allocate twice from a ScratchBufAllocator."); + } + if (byte_size > scratch_.size()) { + return se::port::InternalError(tensorflow::strings::StrCat( + "Can't allocate ", byte_size, + " bytes from a ScratchBufAllocator of size ", scratch_.size())); + } + + allocated_ = true; + return se::DeviceMemory(scratch_); + } + + private: + se::DeviceMemoryBase scratch_; + bool allocated_ = false; +}; + +template +Status RunCudnnConvolution( + CudnnConvKind kind, const Shape& input_shape, const Shape& filter_shape, + const Shape& output_shape, DeviceMemory input_buf, + DeviceMemory filter_buf, DeviceMemory output_buf, + se::ScratchAllocator* scratch_allocator, const Window& window, + const ConvolutionDimensionNumbers& dnums, AlgorithmConfig algorithm, + Stream* stream, ProfileResult* profile_result /*= nullptr*/) { + VLOG(3) << "Convolution Algorithm: " << algorithm.algorithm().algo_id(); + VLOG(3) << "tensor_ops_enabled: " + << algorithm.algorithm().tensor_ops_enabled(); + VLOG(3) << "Convolution kind: " << CudnnConvKindToString(kind); + VLOG(3) << "input shape: { " << ShapeUtil::HumanString(input_shape) << " }"; + VLOG(3) << "filter shape: { " << ShapeUtil::HumanString(filter_shape) << " }"; + VLOG(3) << "Output shape: { " << ShapeUtil::HumanString(output_shape) << " }"; + VLOG(3) << "Window: { " << window.ShortDebugString() << " }"; + VLOG(3) << "Dim nums: { " << dnums.ShortDebugString() << " }"; + + const int num_dimensions = window.dimensions_size(); + CHECK_LE(num_dimensions, 3); + // cuDNN does not support 1D convolutions. We therefore express 1D + // convolutions as 2D convolutions where the first spatial dimension is 1. + // This matches the behavior of TF (see definition of conv1d in + // tensorflow/python/ops/nn_ops.py). + const int effective_num_dimensions = std::max(2, num_dimensions); + + if (std::is_same::value) { + CHECK_EQ(F32, output_shape.element_type()) + << ShapeUtil::HumanString(output_shape); + } else if (std::is_same::value) { + CHECK_EQ(F16, output_shape.element_type()) + << ShapeUtil::HumanString(output_shape); + } else { + LOG(FATAL) << ShapeUtil::HumanString(output_shape); + } + + CHECK_EQ(num_dimensions, dnums.input_spatial_dimensions_size()); + CHECK_EQ(num_dimensions, dnums.kernel_spatial_dimensions_size()); + CHECK_EQ(num_dimensions, dnums.output_spatial_dimensions_size()); + for (const WindowDimension& dim : window.dimensions()) { + CHECK_EQ(dim.padding_low(), dim.padding_high()); + } + + // cuDNN's convolution APIs support the BDYX layout for activations/output and + // the OIYX layout for weights. + BatchDescriptor input_descriptor(effective_num_dimensions); + input_descriptor.set_layout(DataLayout::kBatchDepthYX) + .set_feature_map_count( + input_shape.dimensions(dnums.input_feature_dimension())) + .set_count(input_shape.dimensions(dnums.input_batch_dimension())); + for (int dim = 0; dim < num_dimensions; ++dim) { + // Note that the dimensions are reversed. The same holds below. + input_descriptor.set_spatial_dim( + static_cast(effective_num_dimensions - dim - 1), + input_shape.dimensions(dnums.input_spatial_dimensions(dim))); + } + + FilterDescriptor filter_descriptor(effective_num_dimensions); + filter_descriptor.set_layout(FilterLayout::kOutputInputYX) + .set_input_feature_map_count( + filter_shape.dimensions(dnums.kernel_input_feature_dimension())) + .set_output_feature_map_count( + filter_shape.dimensions(dnums.kernel_output_feature_dimension())); + for (int dim = 0; dim < num_dimensions; ++dim) { + filter_descriptor.set_spatial_dim( + static_cast(effective_num_dimensions - dim - 1), + filter_shape.dimensions(dnums.kernel_spatial_dimensions(dim))); + } + + ConvolutionDescriptor convolution_descriptor(effective_num_dimensions); + for (int dim = 0; dim < num_dimensions; ++dim) { + convolution_descriptor + .set_zero_padding( + static_cast(effective_num_dimensions - dim - 1), + window.dimensions(dim).padding_low()) + .set_filter_stride( + static_cast(effective_num_dimensions - dim - 1), + window.dimensions(dim).stride()); + } + + BatchDescriptor output_descriptor(effective_num_dimensions); + output_descriptor.set_layout(DataLayout::kBatchDepthYX) + .set_feature_map_count( + output_shape.dimensions(dnums.output_feature_dimension())) + .set_count(output_shape.dimensions(dnums.output_batch_dimension())); + for (int dim = 0; dim < num_dimensions; ++dim) { + output_descriptor.set_spatial_dim( + static_cast(effective_num_dimensions - dim - 1), + output_shape.dimensions(dnums.output_spatial_dimensions(dim))); + } + + // Add a singleton dimension in the 1D convolution case. + if (num_dimensions == 1) { + input_descriptor.set_spatial_dim(static_cast(0), 1); + output_descriptor.set_spatial_dim(static_cast(0), 1); + filter_descriptor.set_spatial_dim(static_cast(0), 1); + convolution_descriptor.set_zero_padding(static_cast(0), 0) + .set_filter_stride(static_cast(0), 1); + } + + switch (kind) { + case CudnnConvKind::kForward: + stream->ThenConvolveWithAlgorithm( + input_descriptor, input_buf, filter_descriptor, filter_buf, + convolution_descriptor, output_descriptor, &output_buf, + scratch_allocator, algorithm, profile_result); + break; + case CudnnConvKind::kBackwardInput: + stream->ThenConvolveBackwardDataWithAlgorithm( + filter_descriptor, filter_buf, output_descriptor, output_buf, + convolution_descriptor, input_descriptor, &input_buf, + scratch_allocator, algorithm, profile_result); + break; + case CudnnConvKind::kBackwardFilter: + stream->ThenConvolveBackwardFilterWithAlgorithm( + input_descriptor, input_buf, output_descriptor, output_buf, + convolution_descriptor, filter_descriptor, &filter_buf, + scratch_allocator, algorithm, profile_result); + break; + } + + if (!stream->ok()) { + return InternalError( + "Unable to launch convolution with type %s and algorithm (%lld, %lld)", + CudnnConvKindToString(kind).c_str(), algorithm.algorithm().algo_id(), + algorithm.algorithm_no_scratch().algo_id()); + } + return Status::OK(); +} + +} // anonymous namespace + +string CudnnConvKindToString(CudnnConvKind kind) { + switch (kind) { + case CudnnConvKind::kForward: + return "forward"; + case CudnnConvKind::kBackwardFilter: + return "backward_filter"; + case CudnnConvKind::kBackwardInput: + return "backward_input"; + } +} + +Status RunCudnnConvolution( + CudnnConvKind kind, const Shape& input_shape, const Shape& filter_shape, + const Shape& output_shape, perftools::gputools::DeviceMemoryBase input_buf, + perftools::gputools::DeviceMemoryBase filter_buf, + perftools::gputools::DeviceMemoryBase output_buf, + perftools::gputools::DeviceMemoryBase scratch_buf, const Window& window, + const ConvolutionDimensionNumbers& dnums, + perftools::gputools::dnn::AlgorithmConfig algorithm, + perftools::gputools::Stream* stream, + perftools::gputools::dnn::ProfileResult* profile_result) { + ScratchBufAllocator scratch_allocator(scratch_buf); + return RunCudnnConvolution(kind, input_shape, filter_shape, output_shape, + input_buf, filter_buf, output_buf, + &scratch_allocator, window, dnums, algorithm, + stream, profile_result); +} + +Status RunCudnnConvolution( + CudnnConvKind kind, const Shape& input_shape, const Shape& filter_shape, + const Shape& output_shape, perftools::gputools::DeviceMemoryBase input_buf, + perftools::gputools::DeviceMemoryBase filter_buf, + perftools::gputools::DeviceMemoryBase output_buf, + perftools::gputools::ScratchAllocator* scratch_allocator, + const Window& window, const ConvolutionDimensionNumbers& dnums, + perftools::gputools::dnn::AlgorithmConfig algorithm, + perftools::gputools::Stream* stream, + perftools::gputools::dnn::ProfileResult* profile_result) { + PrimitiveType output_primitive_type = output_shape.element_type(); + CHECK(output_primitive_type == F32 || output_primitive_type == F16) + << ShapeUtil::HumanString(output_shape); + if (output_primitive_type == F32) { + return RunCudnnConvolution( + kind, input_shape, filter_shape, output_shape, + se::DeviceMemory(input_buf), se::DeviceMemory(filter_buf), + se::DeviceMemory(output_buf), scratch_allocator, window, dnums, + algorithm, stream, profile_result); + } + return RunCudnnConvolution(kind, input_shape, filter_shape, output_shape, + se::DeviceMemory(input_buf), + se::DeviceMemory(filter_buf), + se::DeviceMemory(output_buf), + scratch_allocator, window, dnums, algorithm, + stream, profile_result); +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h new file mode 100644 index 0000000000000000000000000000000000000000..3dbfa2730da359d3c7937140508017c4a7b02d6c --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h @@ -0,0 +1,98 @@ +/* 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_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_RUNNER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_RUNNER_H_ + +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { +namespace gpu { + +// This file contains low-level routines for running cudnn convolutions. + +// Different types of convolutions supported by cudnn. +// +// A way to think about these is that a convolution is defined by three arrays +// -- the "input", the "filter", and the "output" -- and given any two of these, +// we can compute the third. For example, a backward-input convolution takes as +// input a filter and an "output" and produces an "input" such that if one were +// to do a forward convolution of "input" using filter, the result would be +// something with the same shape as "output". +// +// This way of thinking is not correct if you look at the values produced. For +// example, a backward-input convolution is not actually the mathematical +// inverse of a forward convolution. But it's right as far as the shapes and +// "connectivity" (i.e. which elements of the input affect which elements of +// the output) are concerned. +enum class CudnnConvKind { + kForward, // input + filter => output + kBackwardInput, // filter + output => input + kBackwardFilter, // input + output => filter +}; + +// Converts a CudnnConvKind value to a string. +string CudnnConvKindToString(CudnnConvKind kind); + +// Calls into cudnn to run the specified convolution. +// +// Note that depending on the value of CudnnConvKind, the result of this call +// may be written into input_buf, filter_buf, or output_buf! +// +// At the moment we only support cudnn convolutions over float and half, and +// convolution with half data type is implemented with cudnn PSEUDO_HALF +// configuration, that is, the input values are half and the internal +// computation type is float. +// +// We provide one overload which takes a scratch buffer, and another which takes +// an allocator which is responsible for allocating the scratch space. In +// theory the second one shouldn't be necessary -- users of this function could +// just ask cudnn how much scratch space it needs for a particular convolution. +// But in practice, StreamExecutor does not expose such an API, and in the name +// of parsimony, perhaps it's better not to add it. Instead, the first time you +// call a convolution, you should call the version that takes a scratch +// allocator and take note of how much memory is used. The next time you call +// the same conv, you can provide an explicitly preallocated scratch buffer of +// that size, if you like. +Status RunCudnnConvolution( + CudnnConvKind kind, const Shape& input_shape, const Shape& filter_shape, + const Shape& output_shape, perftools::gputools::DeviceMemoryBase input_buf, + perftools::gputools::DeviceMemoryBase filter_buf, + perftools::gputools::DeviceMemoryBase output_buf, + perftools::gputools::DeviceMemoryBase scratch_buf, const Window& window, + const ConvolutionDimensionNumbers& dnums, + perftools::gputools::dnn::AlgorithmConfig algorithm, + perftools::gputools::Stream* stream, + perftools::gputools::dnn::ProfileResult* profile_result = nullptr); + +Status RunCudnnConvolution( + CudnnConvKind kind, const Shape& input_shape, const Shape& filter_shape, + const Shape& output_shape, perftools::gputools::DeviceMemoryBase input_buf, + perftools::gputools::DeviceMemoryBase filter_buf, + perftools::gputools::DeviceMemoryBase output_buf, + perftools::gputools::ScratchAllocator* scratch_allocator, + const Window& window, const ConvolutionDimensionNumbers& dnums, + perftools::gputools::dnn::AlgorithmConfig algorithm, + perftools::gputools::Stream* stream, + perftools::gputools::dnn::ProfileResult* profile_result = nullptr); + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONVOLUTION_RUNNER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc index 4b511cb4bb94addfae53d6b2e6d6f86d5b9afd84..5af7a77ea858563fbea05af8efd54f96a74aee93 100644 --- a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc @@ -72,9 +72,27 @@ StatusOr GpuElementalIrEmitter::EmitLibdeviceMathCall( tensorflow::gtl::ArraySlice input_types, PrimitiveType output_type) const { // The libdevice math functions differentiate between "double" and "float" by - // appending an 'f' to the function's name. + // appending an 'f' to the function's name. libdevice doesn't have f16 math + // functions, so we convert the operands to f32 before calling the function + // and then convert the result back to f16. string munged_callee = callee_name; + bool cast_result_to_fp16 = false; + std::vector converted_operands(operands.begin(), + operands.end()); + std::vector converted_input_types(input_types.begin(), + input_types.end()); switch (output_type) { + case F16: + cast_result_to_fp16 = true; + for (int64 i = 0; i < operands.size(); ++i) { + if (input_types[i] == F16) { + converted_operands[i] = ir_builder_->CreateFPCast( + converted_operands[i], ir_builder_->getFloatTy()); + converted_input_types[i] = F32; + } + } + output_type = F32; + TF_FALLTHROUGH_INTENDED; case F32: StrAppend(&munged_callee, "f"); break; @@ -84,7 +102,13 @@ StatusOr GpuElementalIrEmitter::EmitLibdeviceMathCall( return Unimplemented("Bad type for libdevice math call: %s", PrimitiveType_Name(output_type).c_str()); } - return EmitMathCall(munged_callee, operands, input_types, output_type); + llvm::Value* result = EmitMathCall(munged_callee, converted_operands, + converted_input_types, output_type) + .ValueOrDie(); + if (cast_result_to_fp16) { + result = ir_builder_->CreateFPCast(result, ir_builder_->getHalfTy()); + } + return result; } StatusOr GpuElementalIrEmitter::EmitLlvmIntrinsicMathCall( @@ -92,10 +116,13 @@ StatusOr GpuElementalIrEmitter::EmitLlvmIntrinsicMathCall( tensorflow::gtl::ArraySlice operands, tensorflow::gtl::ArraySlice input_types, PrimitiveType output_type) const { - // llvm intrinsics differentiate between float/double functions via the ".f32" - // and ".f64" suffixes. + // llvm intrinsics differentiate between half/float/double functions via + // the suffixes ".f16", ".f32" and ".f64". string munged_callee = callee_name; switch (output_type) { + case F16: + StrAppend(&munged_callee, ".f16"); + break; case F32: StrAppend(&munged_callee, ".f32"); break; @@ -233,12 +260,6 @@ StatusOr GpuElementalIrEmitter::EmitFloatUnaryOp( PrimitiveType input_type = op->operand(0)->shape().element_type(); PrimitiveType output_type = op->shape().element_type(); switch (op->opcode()) { - case HloOpcode::kFloor: - return EmitLibdeviceMathCall("__nv_floor", {operand_value}, {input_type}, - output_type); - case HloOpcode::kCeil: - return EmitLibdeviceMathCall("__nv_ceil", {operand_value}, {input_type}, - output_type); case HloOpcode::kTanh: return EmitLibdeviceMathCall("__nv_tanh", {operand_value}, {input_type}, output_type); diff --git a/tensorflow/compiler/xla/service/gpu/fft_thunk.cc b/tensorflow/compiler/xla/service/gpu/fft_thunk.cc new file mode 100644 index 0000000000000000000000000000000000000000..66931bdc8b1030b2b2e7731ce6327c1e908d4ee6 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/fft_thunk.cc @@ -0,0 +1,234 @@ +/* Copyright 2017 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/compiler/xla/service/gpu/fft_thunk.h" + +#include + +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace se = ::perftools::gputools; + +namespace xla { +namespace gpu { + +FftScratchAllocator::FftScratchAllocator( + int device_ordinal, DeviceMemoryAllocator* memory_allocator) + : device_ordinal_(device_ordinal), memory_allocator_(memory_allocator) {} + +FftScratchAllocator::~FftScratchAllocator() { + for (auto& allocated_buffer : allocated_buffers_) { + if (!memory_allocator_->Deallocate(device_ordinal_, &allocated_buffer) + .ok()) { + // The program can still continue with failed deallocation. + LOG(ERROR) << "Failed to deallocate the allocated buffer: " + << allocated_buffer.opaque(); + } + } +} + +int64 FftScratchAllocator::GetMemoryLimitInBytes(se::Stream* stream) { + constexpr int64 kFftScratchSize = 1LL << 32; // 4GB by default. + return kFftScratchSize; +} + +se::port::StatusOr> FftScratchAllocator::AllocateBytes( + se::Stream* stream, int64 byte_size) { + CHECK_GE(byte_size, 0) << "byte_size must be positive."; + if (byte_size > GetMemoryLimitInBytes(stream)) { + return se::port::Status( + se::port::error::RESOURCE_EXHAUSTED, + tensorflow::strings::Printf( + "Allocating %lld bytes exceeds the memory limit of %lld bytes.", + byte_size, GetMemoryLimitInBytes(stream))); + } + + auto status_or_memory = + memory_allocator_->Allocate(device_ordinal_, byte_size, + /*retry_on_failure=*/false); + if (!status_or_memory.ok()) { + return tensorflow::errors::ResourceExhausted( + "Failed to allocate %lld bytes on device %d.", byte_size, + device_ordinal_); + } + se::DeviceMemoryBase allocated_buffer = status_or_memory.ValueOrDie(); + allocated_buffers_.push_back(allocated_buffer); + total_allocated_bytes_ += byte_size; + return se::DeviceMemory(allocated_buffer); +} + +namespace { + +se::fft::Type FftTypeToSeType(FftType type) { + switch (type) { + case FftType::FFT: + return se::fft::Type::kC2CForward; + case FftType::IFFT: + return se::fft::Type::kC2CInverse; + case FftType::IRFFT: + return se::fft::Type::kC2R; + case FftType::RFFT: + return se::fft::Type::kR2C; + default: + LOG(FATAL) << "unsupported fft type"; + } +} + +string FftTypeToString(se::fft::Type type) { + switch (type) { + case se::fft::Type::kC2CForward: + return "FFT"; + case se::fft::Type::kC2CInverse: + return "IFFT"; + case se::fft::Type::kC2R: + return "IRFFT"; + case se::fft::Type::kR2C: + return "RFFT"; + default: + LOG(FATAL) << "unknown fft type"; + } +} + +} // namespace + +FftThunk::FftThunk(FftType fft_type, + tensorflow::gtl::ArraySlice fft_length, + const BufferAllocation::Slice& input_buffer, + const BufferAllocation::Slice& output_buffer, + const Shape& input_shape, const Shape& output_shape, + const HloInstruction* hlo) + : Thunk(Kind::kFft, hlo), + fft_type_(FftTypeToSeType(fft_type)), + fft_length_(fft_length.begin(), fft_length.end()), + scale_factor_(1.0f), + input_buffer_(input_buffer), + output_buffer_(output_buffer), + input_shape_(input_shape), + output_shape_(output_shape) {} + +tensorflow::Status FftThunk::ExecuteOnStream( + const BufferAllocations& buffer_allocations, se::Stream* stream) { + VLOG(3) << "FFT type: " << FftTypeToString(fft_type_); + VLOG(3) << "Input shape: " << ShapeUtil::HumanStringWithLayout(input_shape_); + VLOG(3) << "Output shape: " + << ShapeUtil::HumanStringWithLayout(output_shape_); + + FftScratchAllocator scratch_allocator(buffer_allocations.device_ordinal(), + buffer_allocations.memory_allocator()); + + if (fft_plan_ == nullptr) { + const int64 fft_rank = fft_length_.size(); + CHECK_LE(fft_rank, 3); + int batch_size = 1; + for (int i = 0; i < input_shape_.dimensions_size() - fft_rank; ++i) { + batch_size *= input_shape_.dimensions(i); + } + uint64 fft_length[3]; + uint64 input_embed[3]; + const uint64 input_stride = 1; + uint64 input_distance = 1; + uint64 output_embed[3]; + const uint64 output_stride = 1; + uint64 output_distance = 1; + + for (int i = 0; i < fft_rank; ++i) { + auto dim_offset = input_shape_.dimensions_size() - fft_rank + i; + fft_length[i] = static_cast(fft_length_[i]); + input_embed[i] = input_shape_.dimensions(dim_offset); + input_distance *= input_shape_.dimensions(dim_offset); + output_embed[i] = output_shape_.dimensions(dim_offset); + output_distance *= output_shape_.dimensions(dim_offset); + } + + constexpr bool kInPlaceFft = false; + fft_plan_ = + stream->parent()->AsFft()->CreateBatchedPlanWithScratchAllocator( + stream, fft_rank, fft_length, input_embed, input_stride, + input_distance, output_embed, output_stride, output_distance, + fft_type_, kInPlaceFft, batch_size, &scratch_allocator); + scale_factor_ = 1.0f / output_distance; + } else { + stream->parent()->AsFft()->UpdatePlanWithScratchAllocator( + stream, fft_plan_.get(), &scratch_allocator); + } + + bool launch_ok; + switch (fft_type_) { + case se::fft::Type::kC2CForward: { + se::DeviceMemory input_data( + buffer_allocations.GetDeviceAddress(input_buffer_)); + se::DeviceMemory output_data( + buffer_allocations.GetDeviceAddress(output_buffer_)); + launch_ok = + stream->ThenFft(fft_plan_.get(), input_data, &output_data).ok(); + break; + } + case se::fft::Type::kC2CInverse: { + se::DeviceMemory input_data( + buffer_allocations.GetDeviceAddress(input_buffer_)); + se::DeviceMemory output_data( + buffer_allocations.GetDeviceAddress(output_buffer_)); + launch_ok = + stream->ThenFft(fft_plan_.get(), input_data, &output_data).ok(); + if (launch_ok) { + launch_ok = + stream + ->ThenBlasScal(ShapeUtil::ElementsIn(output_shape_), + complex64(scale_factor_), &output_data, 1) + .ok(); + } + break; + } + case se::fft::Type::kR2C: { + se::DeviceMemory input_data( + buffer_allocations.GetDeviceAddress(input_buffer_)); + se::DeviceMemory output_data( + buffer_allocations.GetDeviceAddress(output_buffer_)); + launch_ok = + stream->ThenFft(fft_plan_.get(), input_data, &output_data).ok(); + break; + } + case se::fft::Type::kC2R: { + se::DeviceMemory input_data( + buffer_allocations.GetDeviceAddress(input_buffer_)); + se::DeviceMemory output_data( + buffer_allocations.GetDeviceAddress(output_buffer_)); + launch_ok = + stream->ThenFft(fft_plan_.get(), input_data, &output_data).ok(); + if (launch_ok) { + launch_ok = stream + ->ThenBlasScal(ShapeUtil::ElementsIn(output_shape_), + scale_factor_, &output_data, 1) + .ok(); + } + break; + } + default: + LOG(FATAL) << "unsupported fft type"; + } + if (launch_ok) { + return tensorflow::Status::OK(); + } + return InternalError("Unable to launch fft for thunk %p with type %s", this, + FftTypeToString(fft_type_).c_str()); +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/fft_thunk.h b/tensorflow/compiler/xla/service/gpu/fft_thunk.h new file mode 100644 index 0000000000000000000000000000000000000000..52fb8c376d7acea0f15aaa865c23fa2382717338 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/fft_thunk.h @@ -0,0 +1,98 @@ +/* Copyright 2017 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_COMPILER_XLA_SERVICE_GPU_FFT_THUNK_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FFT_THUNK_H_ + +#include "tensorflow/compiler/xla/service/buffer_assignment.h" +#include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_executable.h" +#include "tensorflow/compiler/xla/service/gpu/thunk.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/optional.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { +namespace gpu { + +// A one-time scratch allocator for FFT. The scratch buffers allocated are +// released on destruction. +// +// Not thread-safe in that AllocateBytes, destructor are not locked. +class FftScratchAllocator : public perftools::gputools::ScratchAllocator { + public: + FftScratchAllocator(int device_ordinal, + DeviceMemoryAllocator* memory_allocator); + + ~FftScratchAllocator() override; + + int64 GetMemoryLimitInBytes(perftools::gputools::Stream* stream) override; + + int64 TotalAllocatedBytes() { return total_allocated_bytes_; } + + perftools::gputools::port::StatusOr> + AllocateBytes(perftools::gputools::Stream* stream, int64 byte_size) override; + + private: + const int device_ordinal_; + DeviceMemoryAllocator* memory_allocator_; + std::vector allocated_buffers_; + int64 total_allocated_bytes_ = 0; +}; + +// This class stores everything that StreamExecutor needs to launch an FFT. +// It is generated by IrEmitter. +// +// This is thread-compatible. +class FftThunk : public Thunk { + public: + // Constructs a thunk for launching an FFT on a stream. + // Semantics of null hlo_instruction argument are as in Thunk. + FftThunk(FftType fft_type, tensorflow::gtl::ArraySlice fft_length, + const BufferAllocation::Slice& input_buffer, + const BufferAllocation::Slice& output_buffer, + const Shape& input_shape, const Shape& output_shape, + const HloInstruction* hlo); + + FftThunk(const FftThunk&) = delete; // Cannot share fft_plan_ + FftThunk& operator=(const FftThunk&) = delete; // Cannot share fft_plan_ + + // Does the FFT for the thunk on "stream". + tensorflow::Status ExecuteOnStream( + const BufferAllocations& buffer_allocations, + perftools::gputools::Stream* stream) override; + + private: + const perftools::gputools::fft::Type fft_type_; + const std::vector fft_length_; + + float scale_factor_; + + std::unique_ptr fft_plan_; + + const BufferAllocation::Slice input_buffer_; + const BufferAllocation::Slice output_buffer_; + + const Shape input_shape_; + const Shape output_shape_; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FFT_THUNK_H_ diff --git a/tensorflow/compiler/xla/service/gpu/for_thunk.h b/tensorflow/compiler/xla/service/gpu/for_thunk.h index 525a2af941e77a27c0e01543e00e8a4c3e4b9f62..832494d17e9c4e1d9e92e18ef331df1cf3689024 100644 --- a/tensorflow/compiler/xla/service/gpu/for_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/for_thunk.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FOR_THUNK_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FOR_THUNK_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FOR_THUNK_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FOR_THUNK_H_ #include @@ -49,4 +49,4 @@ class ForThunk : public Thunk { } // namespace gpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FOR_THUNK_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FOR_THUNK_H_ diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger.h b/tensorflow/compiler/xla/service/gpu/fusion_merger.h index bd720f8584f6254c43a3e2a1a5399aa919eebbc0..4c523a66de977cd32423b25f0d165c4f4ba51c4a 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger.h +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSION_MERGER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSION_MERGER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSION_MERGER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSION_MERGER_H_ #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" @@ -44,4 +44,4 @@ class FusionMerger : public HloPassInterface { } // namespace gpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSION_MERGER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSION_MERGER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc index 68d7b2d5c2f4724a29e6059680c1d9e7b6116271..28ebd034ee0c89137f4e6eb417d8a37f4a00af7a 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc @@ -18,12 +18,14 @@ limitations under the License. #include #include #include +#include // NOLINT(build/c++11): only using std::call_once, not mutex. #include #include "llvm/IR/DiagnosticInfo.h" #include "llvm/IR/DiagnosticPrinter.h" #include "llvm/IR/LLVMContext.h" #include "llvm/IR/Module.h" +#include "llvm/IR/Verifier.h" #include "tensorflow/compiler/xla/protobuf_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/algebraic_simplifier.h" @@ -33,17 +35,20 @@ limitations under the License. #include "tensorflow/compiler/xla/service/call_inliner.h" #include "tensorflow/compiler/xla/service/dot_decomposer.h" #include "tensorflow/compiler/xla/service/flatten_call_graph.h" -#include "tensorflow/compiler/xla/service/gpu/convolution_folding.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_rewriter.h" +#include "tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h" +#include "tensorflow/compiler/xla/service/gpu/cudnn_convolution_rewriter.h" #include "tensorflow/compiler/xla/service/gpu/fusion_merger.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_constants.h" #include "tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.h" #include "tensorflow/compiler/xla/service/gpu/gpu_executable.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.h" #include "tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h" #include "tensorflow/compiler/xla/service/gpu/hlo_schedule.h" #include "tensorflow/compiler/xla/service/gpu/instruction_fusion.h" #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" -#include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter_context.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h" #include "tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h" #include "tensorflow/compiler/xla/service/gpu/pad_insertion.h" #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" @@ -54,6 +59,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_constant_folding.h" #include "tensorflow/compiler/xla/service/hlo_cse.h" #include "tensorflow/compiler/xla/service/hlo_dce.h" +#include "tensorflow/compiler/xla/service/hlo_element_type_converter.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_pass_fix.h" #include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h" @@ -77,9 +83,11 @@ limitations under the License. #include "tensorflow/core/platform/cuda_libdevice_path.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/regexp.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/platform/subprocess.h" #include "tensorflow/core/platform/tracing.h" +#include "tensorflow/stream_executor/cuda/cuda_diagnostics.h" namespace se = ::perftools::gputools; @@ -93,14 +101,6 @@ namespace gpu { namespace { using tensorflow::port::Tracing; -using tensorflow::strings::StrCat; - -// Any address of a variable residing in global memory or returned by one of the -// memory allocation routines from the driver or runtime API is always aligned -// to at least 256 bytes. -// -// http://docs.nvidia.com/cuda/cuda-c-programming-guide/#device-memory-accesses -constexpr int64 kMemoryAlignment = 256; // Returns the directory containing nvvm libdevice files. config_cuda_data_dir // should be equal to config().debug_options().xla_gpu_cuda_data_dir() of the @@ -128,23 +128,29 @@ string GetLibdeviceDir(const string& config_cuda_data_dir) { } // Runs optimization passes on the given HLO module. -tensorflow::Status OptimizeHloModule( - HloModule* hlo_module, - const HloCostAnalysis::ShapeSizeFunction& shape_size_function) { +tensorflow::Status OptimizeHloModule(HloModule* hlo_module, + se::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) { { HloPassPipeline pipeline("optimization"); - pipeline.AddInvariantChecker(shape_size_function); + pipeline.AddInvariantChecker(); + pipeline.AddPass(); ReducePrecisionInsertion::AddPasses( &pipeline, hlo_module->config().debug_options(), ReducePrecisionInsertion::PassTiming::BEFORE_OPTIMIZATION); // TODO(b/64094172): make Call work on GPU instead of inlining. pipeline.AddPass(); + // Convert BF16 operations to F32 operations so that the GPU backend can + // support BF16 operations without directly implementing a BF16 lowering for + // most ops. + pipeline.AddPass(BF16, F32); pipeline.AddPass(); + { auto& pass = pipeline.AddPass>("simplification"); - pass.AddInvariantChecker(shape_size_function); + pass.AddInvariantChecker(); // If cudnn batchnorms are enabled, rewrite batchnorm HLOs to cudnn calls // where possible. Not every batchnorm op can be implemented as a call to @@ -171,7 +177,7 @@ tensorflow::Status OptimizeHloModule( pass.AddPass(); pass.AddPass(); } - pipeline.AddPass(); + pipeline.AddPass( [](const HloInstruction& dot, const TransposeFolding::OperandIndices& candidate_operands) { @@ -183,16 +189,68 @@ tensorflow::Status OptimizeHloModule( pipeline.AddPass(); TF_RETURN_IF_ERROR(pipeline.Run(hlo_module).status()); } + + { + // Convert convolutions into CustomCalls to cudnn, then canonicalize them + // (PadInsertion). + HloPassPipeline pipeline("conv_canonicalization"); + pipeline.AddInvariantChecker(); + pipeline.AddPass(); + pipeline.AddPass(); + + // Choose the fastest algorithm for each conv. + // + // In theory doing this here is way too early: It needs to happen after + // layout assignment, because the layout of the inputs/outputs affects the + // speed of the conv. But currently we only allow only one input/output + // layout when calling cudnn, so there's no ambiguity. + // + // We pick the algorithm at this early stage so we can generate better HLO. + // After CudnnConvolutionRewriter, our convolutions are CustomCalls which + // return a tuple (conv_result, scratch_memory), and the each conv uses 0 + // bytes of scratch: + // + // customcall = (f32[...], f32[0]) + // return gte(customcall, 0) + // + // The algorithm picker then chooses the best algorithm, and potentially + // increases the scratch space. It replaces customcall with new_tuple, + // giving us the following: + // + // new_customcall = (f32[...], f32[N]) + // new_tuple = tuple(gte(new_customcall, 0), constant f32[0]) + // return gte(new_tuple, 0) + // + // The new tuple and gte instructions then be simplified away, because + // nobody is expected to use the scratch value. + // + // However, if we were to run CudnnConvolutionAlgorithmPicker after layout + // assignment, fusion would already have run, and the gte(customcall, 0) + // would probably already be into a fusion node. We can't simplify across + // HloComputation boundaries, so in this case we wouldn't be able to + // simplify away the new_tuple bits. + // + // We'll need to revisit this if we ever allow multiple layouts for the + // inputs/outputs of a cudnn convolution. + pipeline.AddPass(stream_exec, + device_allocator); + // Clean up new_tuple described above. + pipeline.AddPass(); + pipeline.AddPass(); + + TF_RETURN_IF_ERROR(pipeline.Run(hlo_module).status()); + } + { HloPassFix fusion("fusion"); - fusion.AddInvariantChecker(shape_size_function); + fusion.AddInvariantChecker(); fusion.AddPass(/*may_duplicate=*/false); fusion.AddPass(/*may_duplicate=*/true); fusion.AddPass(); TF_RETURN_IF_ERROR(fusion.Run(hlo_module).status()); HloPassPipeline reduce_pipeline("reduce-precision"); - reduce_pipeline.AddInvariantChecker(shape_size_function); + reduce_pipeline.AddInvariantChecker(); ReducePrecisionInsertion::AddPasses( &reduce_pipeline, hlo_module->config().debug_options(), ReducePrecisionInsertion::PassTiming::AFTER_FUSION); @@ -210,19 +268,18 @@ tensorflow::Status OptimizeHloModule( // Modifies the given HLO module so that it will be accepted by IrEmitter. // Unlike optimization passes, the passes are necessary for correctness. -tensorflow::Status PrepareHloModuleForIrEmitting( - HloModule* hlo_module, - const HloCostAnalysis::ShapeSizeFunction& shape_size_function) { +tensorflow::Status PrepareHloModuleForIrEmitting(HloModule* hlo_module) { // In some cases, we have to place the result of an instruction in a temporary // buffer. For instance, the buffer that holds an external parameter is // assumed immutable at this point, and should not be reused for output // (b/27180329). Therefore, in that case, we set the output to be a copy of // the parameter. HloPassPipeline pipeline("GPU-ir-emit-prepare"); - pipeline.AddInvariantChecker(shape_size_function); - pipeline.AddPass(); + pipeline.AddInvariantChecker(); + pipeline.AddPass( hlo_module->mutable_entry_computation_layout()); + // The LayoutAssignment pass may leave behind kCopy instructions which are // duplicate or NOPs, so remove them with algebraic simplification and CSE. pipeline.AddPass>( @@ -241,6 +298,103 @@ tensorflow::Status PrepareHloModuleForIrEmitting( return pipeline.Run(hlo_module).status(); } +// Prints a warning if the ptxas at ptxas_path has known bugs. +// +// Only prints a warning the first time it's called for a particular value of +// ptxas_path. +void WarnIfBadPtxasVersion(const string& ptxas_path) { + static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED); + static std::unordered_set* seen_ptxas_paths GUARDED_BY(mu) = + new std::unordered_set(); + + tensorflow::mutex_lock lock(mu); + if (!seen_ptxas_paths->insert(ptxas_path).second) { + // Already checked this ptx binary, nothing to do. + return; + } + + tensorflow::SubProcess ptxas; + ptxas.SetProgram(ptxas_path, {ptxas_path, "--version"}); + ptxas.SetChannelAction(tensorflow::CHAN_STDOUT, tensorflow::ACTION_PIPE); + if (!ptxas.Start()) { + LOG(WARNING) << "Couldn't invoke " << ptxas_path << " --version"; + return; + } + + string out; + int exit_code = ptxas.Communicate(/*stdin_input=*/nullptr, &out, + /*stderr_output=*/nullptr); + if (exit_code != 0) { + LOG(WARNING) << "Running " << ptxas_path << " --version returned " + << exit_code; + return; + } + + int64 vmaj, vmin, vdot; + string vmaj_str, vmin_str, vdot_str; + if (!RE2::PartialMatch(out, R"(\bV(\d+)\.(\d+)\.(\d+)\b)", &vmaj_str, + &vmin_str, &vdot_str) || + !tensorflow::strings::safe_strto64(vmaj_str, &vmaj) || + !tensorflow::strings::safe_strto64(vmin_str, &vmin) || + !tensorflow::strings::safe_strto64(vdot_str, &vdot)) { + LOG(WARNING) << "Couldn't parse ptxas version in output of " << ptxas_path + << " --version:\n" + << out; + return; + } + + // ptxas 9.0 before 9.0.276 and ptxas 9.1 before 9.1.121 miscompile some + // address calculations with large offsets (e.g. "load ptr + large_constant"), + // b/70245379. + if ((vmaj == 9 && vmin == 0 && vdot < 276) || + (vmaj == 9 && vmin == 1 && vdot < 121)) { + LOG(WARNING) << "*** WARNING *** You are using ptxas " << vmaj << "." + << vmin << "." << vdot + << ", which is in range [9.0.0, 9.0.276) + [9.1.0, 9.1.121). " + "These versions are known to miscompile XLA code, leading " + "to incorrect results or invalid-address errors."; + } +} + +// Prints a warning if the ptx->sass JIT in the driver has known bugs. +// +// Using such a driver only a problem if we fail to use ptxas to compile our ptx +// and have to use the driver instead, so you should only call this function if +// we're going to use the driver JIT. +// +// Only prints a warning the first time it's called. +void WarnIfBadDriverJITVersion() { + static std::once_flag run_once; + std::call_once(run_once, [] { + auto version_or_status = se::cuda::Diagnostician::FindKernelDriverVersion(); + if (!version_or_status.ok()) { + LOG(WARNING) << "Couldn't read CUDA driver version."; + return; + } + se::cuda::DriverVersion version = version_or_status.ValueOrDie(); + + // The following versions of the driver JIT miscompile some address + // calculations with large offsets (e.g. "load ptr + large_constant"), + // b/70245379: + // + // - 384.x before 384.108 + // - 387.x before 387.40 + // - 390.x before 390.10. + auto vmaj = std::get<0>(version); + auto vmin = std::get<1>(version); + if ((vmaj == 384 && vmin < 108) || // + (vmaj == 387 && vmin < 40) || // + (vmaj == 390 && vmin < 10)) { + LOG(WARNING) + << "*** WARNING *** Invoking the PTX->SASS JIT from driver version " + << se::cuda::DriverVersionToString(version) + << ", which is in range [384.0.0, 384.108.0) + [387.0.0, 387.40.0) + " + "[390.0.0, 390.10.0). These versions are known to miscompile XLA " + "code, leading to incorrect results or invalid-address errors."; + } + }); +} + // Compiles the given PTX string using ptxas and returns the resulting machine // code (i.e. a cubin) as a byte array. StatusOr> CompilePtx(const string& ptx, int cc_major, @@ -252,6 +406,8 @@ StatusOr> CompilePtx(const string& ptx, int cc_major, auto env = tensorflow::Env::Default(); TF_RETURN_IF_ERROR(env->FileExists(ptxas_path)); + WarnIfBadPtxasVersion(ptxas_path); + // Write ptx into a temporary file. string ptx_path; if (!env->LocalTempFilename(&ptx_path)) { @@ -275,8 +431,9 @@ StatusOr> CompilePtx(const string& ptx, int cc_major, tensorflow::Env::Default()->DeleteFile(cubin_path).IgnoreError(); }); tensorflow::SubProcess ptxas_info_dumper; - std::vector ptxas_args = {ptxas_path, ptx_path, "-o", cubin_path, - StrCat("-arch=sm_", cc_major, cc_minor)}; + std::vector ptxas_args = { + ptxas_path, ptx_path, "-o", cubin_path, + tensorflow::strings::StrCat("-arch=sm_", cc_major, cc_minor)}; if (VLOG_IS_ON(2)) { ptxas_args.push_back("-v"); } @@ -310,22 +467,24 @@ GpuCompiler::GpuCompiler() .getPointerSize(0 /* default address space */)) {} StatusOr> GpuCompiler::RunHloPasses( - std::unique_ptr module, se::StreamExecutor* /*stream_exec*/) { + std::unique_ptr module, se::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) { XLA_SCOPED_LOGGING_TIMER("GpuCompiler::RunHloPasses"); Tracing::TraceMe annotation("HLO Transforms", module->name(), /*is_expensive=*/true); - TF_RETURN_IF_ERROR(OptimizeHloModule(module.get(), ShapeSizeBytesFunction())); + TF_RETURN_IF_ERROR( + OptimizeHloModule(module.get(), stream_exec, device_allocator)); return std::move(module); } StatusOr> GpuCompiler::RunBackend( - std::unique_ptr module, se::StreamExecutor* stream_exec) { + std::unique_ptr module, se::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) { XLA_SCOPED_LOGGING_TIMER("GpuCompiler::RunBackend"); TF_RET_CHECK(stream_exec != nullptr); - TF_RETURN_IF_ERROR( - PrepareHloModuleForIrEmitting(module.get(), ShapeSizeBytesFunction())); + TF_RETURN_IF_ERROR(PrepareHloModuleForIrEmitting(module.get())); llvm::LLVMContext llvm_context; std::string buffer; @@ -356,19 +515,21 @@ StatusOr> GpuCompiler::RunBackend( TF_ASSIGN_OR_RETURN( std::unique_ptr buffer_assignment, BufferAssigner::Run(module.get(), hlo_schedule->ConsumeHloOrdering(), - BufferSizeBytesFunction(), [](LogicalBuffer::Color) { - return kMemoryAlignment; + BufferSizeBytesFunction(), + /*color_alignment=*/[](LogicalBuffer::Color) { + return kCudaMallocAlignBytes; })); - // BufferAssignment::ToString() includes a header, so no need for us to - // print one ourselves. + // BufferAssignment::Stats::ToString() and BufferAssignment::ToString() + // include headers, so no need for us to print them ourselves. + XLA_VLOG_LINES(1, buffer_assignment->GetStats().ToString()); XLA_VLOG_LINES(2, buffer_assignment->ToString()); XLA_VLOG_LINES(2, module->ToString()); - const string xla_dump_hlo_proto_to = - module->config().debug_options().xla_dump_hlo_proto_to(); - if (!xla_dump_hlo_proto_to.empty()) { + const string xla_dump_optimized_hlo_proto_to = + module->config().debug_options().xla_dump_optimized_hlo_proto_to(); + if (!xla_dump_optimized_hlo_proto_to.empty()) { HloProto proto = MakeHloProto(*module, *buffer_assignment); TF_RETURN_IF_ERROR(protobuf_util::DumpProtoToDirectory( - proto, xla_dump_hlo_proto_to, module->name())); + proto, xla_dump_optimized_hlo_proto_to, module->name())); } IrEmitterContext ir_emitter_context(module.get(), buffer_assignment.get(), @@ -406,6 +567,20 @@ StatusOr> GpuCompiler::RunBackend( /*optimized=*/false)); } + { + XLA_SCOPED_LOGGING_TIMER("GpuCompiler::RunBackend - Running LLVM verifier"); + + std::string err; + llvm::raw_string_ostream err_stream(err); + + // verifyModule() returns true if the module is broken. + TF_RET_CHECK(!llvm::verifyModule(llvm_module, &err_stream)) + << "Invalid LLVM IR before optimizations:\n" + << err_stream.str() + << "\nThis probably indicates a bug in the HLO -> LLVM IR lowering. " + "Rerun with --xla_dump_ir_to to get the IR. "; + } + string libdevice_dir; { tensorflow::mutex_lock lock(mutex_); @@ -456,7 +631,7 @@ StatusOr> GpuCompiler::RunBackend( // Write PTX to IR dump directory, if IR dumping was requested. if (!ir_dump_directory.empty()) { const string ptx_outfile = tensorflow::io::JoinPath( - ir_dump_directory, StrCat(module->name(), ".ptx")); + ir_dump_directory, tensorflow::strings::StrCat(module->name(), ".ptx")); auto status = [&] { auto* env = tensorflow::Env::Default(); TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(ir_dump_directory)); @@ -479,13 +654,14 @@ StatusOr> GpuCompiler::RunBackend( XLA_VLOG_LINES(2, thunk_schedule->ToString()); std::unique_ptr profile_index_map; - std::unique_ptr profile_printer; + std::unique_ptr profile_printer; if (module->config().hlo_profiling_enabled()) { HloCostAnalysis cost_analysis(ShapeSizeBytesFunction()); + TF_RETURN_IF_ERROR(module->entry_computation()->Accept(&cost_analysis)); profile_index_map = MakeUnique(*module); profile_printer = - CreateHloProfilePrinter(*profile_index_map, cost_analysis); + CreateHloProfilePrinterData(*profile_index_map, cost_analysis); } auto* gpu_executable = new GpuExecutable( @@ -554,6 +730,10 @@ std::vector GpuCompiler::CompilePtxOrGetCachedResult(const string& ptx, "GPU driver compile the ptx. " << maybe_cubin.status(); } + + // We're going to use the driver to JIT our PTX->SASS, so warn if + // the JIT in the driver has known bugs. + WarnIfBadDriverJITVersion(); } } cache_value->compilation_done = true; diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.h b/tensorflow/compiler/xla/service/gpu/gpu_compiler.h index 18e34340205b6f51497e26c45520799d21c55a46..c352d4d8462fadb266c55ad437de998e86a6528e 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.h @@ -51,11 +51,13 @@ class GpuCompiler : public LLVMCompiler { StatusOr> RunHloPasses( std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec) override; + perftools::gputools::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) override; StatusOr> RunBackend( std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec) override; + perftools::gputools::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) override; StatusOr>> CompileAheadOfTime(std::vector> module, diff --git a/tensorflow/compiler/xla/service/gpu/gpu_constants.cc b/tensorflow/compiler/xla/service/gpu/gpu_constants.cc new file mode 100644 index 0000000000000000000000000000000000000000..aa360c7f73de2f0f9cf59c22b552b8e60ddb3a87 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/gpu_constants.cc @@ -0,0 +1,25 @@ +/* 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/compiler/xla/service/gpu/gpu_constants.h" + +namespace xla { +namespace gpu { + +// http://docs.nvidia.com/cuda/cuda-c-programming-guide/#device-memory-accesses +const int64 kCudaMallocAlignBytes = 256; + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_constants.h b/tensorflow/compiler/xla/service/gpu/gpu_constants.h new file mode 100644 index 0000000000000000000000000000000000000000..eb1ca4c6c95a23d2a08f5f9c3cbc85e7d47d4f89 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/gpu_constants.h @@ -0,0 +1,31 @@ +/* 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_COMPILER_XLA_SERVICE_GPU_GPU_CONSTANTS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_CONSTANTS_H_ + +#include "tensorflow/compiler/xla/types.h" + +namespace xla { +namespace gpu { + +// Minimum alignment of cudaMalloc. We require that buffers created by our +// DeviceMemoryAllocator, and all input/output buffers, have this alignment. +extern const int64 kCudaMallocAlignBytes; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_CONSTANTS_H_ diff --git a/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.cc b/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.cc index e67087d822e2f3367c48b08be66f5f60791be638..916b556fd43a453a4da2c96217e74c367f8c7653 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.cc @@ -36,7 +36,7 @@ namespace gpu { StatusOr GpuCopyInsertion::FindOrInsertCopy( HloInstruction* hlo) { - HloInstruction*& copy = inserted_copies_[hlo]; + HloInstruction*& copy = hlo_to_copy_map_[hlo]; if (copy == nullptr) { TF_ASSIGN_OR_RETURN(copy, hlo->parent()->DeepCopyInstruction(hlo)); } @@ -78,6 +78,12 @@ StatusOr GpuCopyInsertion::Run(HloModule* module) { for (int64 i = 0; i < hlo->operand_count() - 2; ++i) { TF_RETURN_IF_ERROR(copy_operand_if_constant(i)); } + } else if (IsCustomCallToDnnConvolution(*hlo)) { + // The last two arguments to a CUDNN convolution are two HLO constants for + // cudnn algorithm and tensor_ops_enabled flag, which shouldn't be copied. + for (int64 i = 0; i < hlo->operand_count() - 2; ++i) { + TF_RETURN_IF_ERROR(copy_operand_if_constant(i)); + } } else if (ImplementedAsLibraryCall(*hlo)) { // For all other library calls, materialize all the operands into memory. for (int64 i = 0; i < hlo->operand_count(); ++i) { @@ -86,27 +92,34 @@ StatusOr GpuCopyInsertion::Run(HloModule* module) { } } - // Init values of a while node cannot be constants. Insert copies for any - // constants found at the operand of a while. - tensorflow::gtl::FlatSet copied_constants; + // Init values of while and conditional nodes cannot be constants. Insert + // copies for any constants found at the operands of these nodes. + tensorflow::gtl::FlatSet inserted_copies; for (HloComputation* computation : module->computations()) { for (HloInstruction* instruction : computation->instructions()) { - if (instruction->opcode() != HloOpcode::kWhile) { + if (instruction->opcode() != HloOpcode::kWhile && + instruction->opcode() != HloOpcode::kConditional) { continue; } - for (auto& pair : - dataflow->GetInstructionValueSet(instruction->operand(0))) { - const HloValueSet& value_set = pair.second; - for (const HloValue* value : value_set.values()) { - if (value->defining_instruction()->opcode() == - HloOpcode::kConstant && - !ContainsKey(copied_constants, value->defining_instruction())) { - HloInstruction* constant = value->defining_instruction(); - TF_ASSIGN_OR_RETURN(HloInstruction * copy, - FindOrInsertCopy(constant)); - TF_RETURN_IF_ERROR(constant->ReplaceAllUsesWith(copy)); - copied_constants.insert(constant); - changed = true; + for (auto operand : instruction->operands()) { + // Skip the operands that have already been replaced with a copy in a + // previous iteration (which is possible when a constant is used as an + // operand in multiple places). + if (ContainsKey(inserted_copies, operand)) { + continue; + } + for (auto& pair : dataflow->GetInstructionValueSet(operand)) { + const HloValueSet& value_set = pair.second; + for (const HloValue* value : value_set.values()) { + if (value->defining_instruction()->IsConstant() && + !ContainsKey(hlo_to_copy_map_, value->defining_instruction())) { + HloInstruction* constant = value->defining_instruction(); + TF_ASSIGN_OR_RETURN(HloInstruction * copy, + FindOrInsertCopy(constant)); + TF_RETURN_IF_ERROR(constant->ReplaceAllUsesWith(copy)); + inserted_copies.insert(copy); + changed = true; + } } } } diff --git a/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.h b/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.h index 4d77f337e6eb20f7d79acc0829fde26bbe443f25..0c6f9b511f3aac5f62182273b827adcd068cd633 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.h @@ -32,13 +32,13 @@ class GpuCopyInsertion : public HloPassInterface { StatusOr Run(HloModule* module) override; protected: - // Returns a copy of `hlo`. Looks in inserted_copies_ first to avoid making + // Returns a copy of `hlo`. Looks in hlo_to_copy_map_ first to avoid making // duplicate copies. StatusOr FindOrInsertCopy(HloInstruction* hlo); // A map containing all copies inserted to materialize operands of library // calls. The key is the copied instruction and the value is the copy. - tensorflow::gtl::FlatMap inserted_copies_; + tensorflow::gtl::FlatMap hlo_to_copy_map_; }; } // namespace gpu diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc index 5b019e5289b8cb86ece1c728b506ca81e2fb3555..623d6714de501000e38b7698620925f66425f157 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc @@ -66,7 +66,9 @@ class HloExecutionProfiler { // If profiling is enabled, sets the total cycle count on the profile from the // execution timer. - ~HloExecutionProfiler() { + void FinishExecution() { + CHECK(!finished_execution_) << "Call FinishExecution only once!"; + finished_execution_ = true; if (do_profile_) { stream_->ThenStopTimer(execution_timer_.get()); stream_->BlockHostUntilDone().IgnoreError(); @@ -101,6 +103,7 @@ class HloExecutionProfiler { const HloComputation* computation_; std::unique_ptr execution_timer_; std::unique_ptr per_op_timer_; + bool finished_execution_ = false; }; } // namespace @@ -113,9 +116,9 @@ GpuExecutable::GpuExecutable( std::unique_ptr thunk_schedule, std::unique_ptr hlo_module, std::unique_ptr assignment, - std::unique_ptr hlo_profile_printer, + std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map) - : Executable(std::move(hlo_module), std::move(hlo_profile_printer), + : Executable(std::move(hlo_module), std::move(hlo_profile_printer_data), std::move(hlo_profile_index_map)), ptx_(ptx), cubin_(cubin), @@ -143,9 +146,12 @@ Status GpuExecutable::ExecuteThunks( if (do_profile) { LOG(WARNING) << "PROFILING: profiling is enabled"; } + HloExecutionProfiler profiler(do_profile, hlo_execution_profile, main_stream, hlo_module_->entry_computation()); + uint64 start_micros = tensorflow::Env::Default()->NowMicros(); + // Stream 0 indicates `main_stream` and substreams start from stream 1. std::vector::SmartPtr> sub_streams; while (sub_streams.size() + 1 < thunk_schedule_->StreamCount()) { @@ -222,6 +228,22 @@ Status GpuExecutable::ExecuteThunks( } } + profiler.FinishExecution(); + uint64 end_micros = tensorflow::Env::Default()->NowMicros(); + + { + tensorflow::mutex_lock lock(mutex_); + const double nanoseconds = (end_micros - start_micros) * 1000.0; + execution_profile_.set_compute_time_ns(std::max(nanoseconds, 1.0)); + + // If hlo profiling was disabled then the cycle count is left empty. + if (do_profile) { + execution_profile_.set_compute_cycle_count( + hlo_execution_profile->total_cycles_executed( + *module().entry_computation())); + } + } + return Status::OK(); } @@ -240,9 +262,16 @@ StatusOr> GpuExecutable::ExecuteOnStream( ++i) { const BufferAllocation& allocation = assignment_->GetAllocation(i); if (allocation.is_entry_computation_parameter()) { - auto param_no = allocation.parameter_number(); - buffer_allocations_builder.RegisterBuffer( - i, arguments[param_no]->root_buffer()); + // The caller must give us a buffer for ShapeIndex {} of every parameter. + // It can optionally give us a buffer for other ShapeIndices, but we + // ignore them: Because we can't rely on these sub-buffers' addresses + // being available, our generated code can't use them. Instead, it must + // chase pointers starting at the tuple root. + if (allocation.param_shape_index().empty()) { + auto param_no = allocation.parameter_number(); + buffer_allocations_builder.RegisterBuffer( + i, arguments[param_no]->root_buffer()); + } } } se::StreamExecutor* executor = run_options->stream()->parent(); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.h b/tensorflow/compiler/xla/service/gpu/gpu_executable.h index 00da64dfade8ddb0694c0ee7ac158c9f2e15a508..b19cfd43debd0a5490495d176fa2f1fcd625da07 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.h @@ -54,7 +54,7 @@ class GpuExecutable : public Executable { std::unique_ptr thunk_schedule, std::unique_ptr hlo_module, std::unique_ptr assignment, - std::unique_ptr hlo_profile_printer, + std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map); // This should be called after set_ir_module_string. diff --git a/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.cc b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.cc new file mode 100644 index 0000000000000000000000000000000000000000..4944c41f7d8dc7a78a3cd094aee4d7087c74857e --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.cc @@ -0,0 +1,48 @@ +/* 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/compiler/xla/service/gpu/gpu_hlo_support_checker.h" + +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/errors.h" + +namespace xla { + +StatusOr GpuHloSupportChecker::Run(HloModule* module) { + for (auto* computation : module->computations()) { + for (const auto& instruction : computation->instructions()) { + TF_RETURN_IF_ERROR( + ShapeUtil::ValidateShapeWithOptionalLayout(instruction->shape())); + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + instruction->shape(), + [&instruction](const Shape& subshape, const ShapeIndex&) { + if (LayoutUtil::IsSparseArray(subshape)) { + return xla::Unimplemented( + "GPU backend does not support HLO instruction %s with shape " + "containing a sparse layout: %s", + instruction->ToString().c_str(), + ShapeUtil::HumanStringWithLayout(instruction->shape()) + .c_str()); + } + return Status::OK(); + })); + } + } + return false; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.h b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.h new file mode 100644 index 0000000000000000000000000000000000000000..d63e213d2b1efab4bcff75541cc5ab33d7a07976 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.h @@ -0,0 +1,42 @@ +/* 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_COMPILER_XLA_SERVICE_GPU_GPU_HLO_SUPPORT_CHECKER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_HLO_SUPPORT_CHECKER_H_ + +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// his pass should run early in the HLO pipeline and checks for HLO constructs +// which are not supported by the GPU backend and cannot be removed via HLO +// transformations (eg, sparse layouts). +class GpuHloSupportChecker : public HloPassInterface { + public: + GpuHloSupportChecker() = default; + ~GpuHloSupportChecker() override = default; + + tensorflow::StringPiece name() const override { + return "gpu_hlo_support_checker"; + } + + // Note: always returns false (no instructions are ever modified by this + // pass). + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_HLO_SUPPORT_CHECKER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker_test.cc b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0a4089df4c954cafcbe241189ee79a0995683513 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker_test.cc @@ -0,0 +1,72 @@ +/* Copyright 2017 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/compiler/xla/service/gpu/gpu_hlo_support_checker.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/core/lib/core/error_codes.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +using ::testing::HasSubstr; + +class GpuHloSupportCheckerTest : public HloTestBase { + protected: + GpuHloSupportChecker& checker() { return checker_; } + + private: + GpuHloSupportChecker checker_; +}; + +TEST_F(GpuHloSupportCheckerTest, Add) { + HloComputation::Builder builder(TestName()); + const Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "param1")); + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape, HloOpcode::kAdd, param0, param1)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + TF_ASSERT_OK(checker().Run(module.get()).status()); +} + +TEST_F(GpuHloSupportCheckerTest, SparseUnimplemented) { + HloComputation::Builder builder(TestName()); + const Shape sparse_shape = ShapeUtil::MakeShapeWithSparseLayout(F32, {10}, 2); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, sparse_shape, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, sparse_shape, "param1")); + builder.AddInstruction(HloInstruction::CreateBinary( + sparse_shape, HloOpcode::kAdd, param0, param1)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + Status status = checker().Run(module.get()).status(); + ASSERT_EQ(status.code(), tensorflow::error::UNIMPLEMENTED); + EXPECT_THAT(status.error_message(), + HasSubstr("GPU backend does not support")); + EXPECT_THAT(status.error_message(), + HasSubstr(ShapeUtil::HumanStringWithLayout(sparse_shape))); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc index 58915f1f62f0c0f320443058a798333c498ffe47..89f1e625884568bf7370b3801d851ef4846c2a98 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc @@ -28,122 +28,114 @@ limitations under the License. namespace xla { namespace gpu { +// cuDNN convolutions are called with specific layouts on the input, output, +// and filter: +// +// input: DataLayout::kBatchDepthYX +// output: DataLayout::kBatchDepthYX +// filter: FilterLayout::kOutputInputYX +// +// The order dimensions in the constant name is major-to-minor (eg, the +// most-major dimension of the input is batch, most-minor is X). The +// specific dimension numbers these named dimensions correspond to is +// determined by the ConvolutionDimensionNumbers argument. Y is spatial +// dimension 0, and X is spatial dimension 1. +// +// TODO(b/29399649): Be more flexible about handling layouts of cuDNN calls. +static Status AddBackendConstraintsToDnnConvCustomCall( + HloInstruction* instr, LayoutConstraints* constraints) { + CHECK(IsCustomCallToDnnConvolution(*instr)) << instr->ToString(); + Shape input_shape; + Shape filter_shape; + Shape output_shape; + const auto& target = instr->custom_call_target(); + if (target == kCudnnConvForwardCallTarget) { + input_shape = instr->operand(0)->shape(); + filter_shape = instr->operand(1)->shape(); + output_shape = instr->shape().tuple_shapes(0); + } else if (target == kCudnnConvBackwardInputCallTarget) { + input_shape = instr->shape().tuple_shapes(0); + filter_shape = instr->operand(1)->shape(); + output_shape = instr->operand(0)->shape(); + } else if (target == kCudnnConvBackwardFilterCallTarget) { + input_shape = instr->operand(0)->shape(); + filter_shape = instr->shape().tuple_shapes(0); + output_shape = instr->operand(1)->shape(); + } else { + LOG(FATAL) << "Unexpected custom call target: " + << instr->custom_call_target(); + } + + // Construct minor-to-major dimension orders for operands and result. + // cuDNN's convolution APIs support the BDYX layout for activations/output + // and the OIYX layout for weights. + // TODO(b/29399649): Be more flexible about handling layouts of cuDNN + // calls after we switch to cuDNN v5. + const ConvolutionDimensionNumbers& dimension_numbers = + instr->convolution_dimension_numbers(); + std::vector input_layout; + for (int i = dimension_numbers.input_spatial_dimensions_size() - 1; i >= 0; + --i) { + input_layout.push_back(dimension_numbers.input_spatial_dimensions(i)); + } + input_layout.push_back(dimension_numbers.input_feature_dimension()); + input_layout.push_back(dimension_numbers.input_batch_dimension()); + *input_shape.mutable_layout() = LayoutUtil::MakeLayout(input_layout); + + std::vector filter_layout; + for (int i = dimension_numbers.kernel_spatial_dimensions_size() - 1; i >= 0; + --i) { + filter_layout.push_back(dimension_numbers.kernel_spatial_dimensions(i)); + } + filter_layout.push_back(dimension_numbers.kernel_input_feature_dimension()); + filter_layout.push_back(dimension_numbers.kernel_output_feature_dimension()); + *filter_shape.mutable_layout() = LayoutUtil::MakeLayout(filter_layout); + + std::vector output_layout; + for (int i = dimension_numbers.output_spatial_dimensions_size() - 1; i >= 0; + --i) { + output_layout.push_back(dimension_numbers.output_spatial_dimensions(i)); + } + output_layout.push_back(dimension_numbers.output_feature_dimension()); + output_layout.push_back(dimension_numbers.output_batch_dimension()); + *output_shape.mutable_layout() = LayoutUtil::MakeLayout(output_layout); + + // The custom call returns a tuple of (actual_result, scratch_buffer); + // call_result_buf is the logical buffer for actual_result, the thing that + // contains the result of the conv call. + TF_ASSIGN_OR_RETURN(const LogicalBuffer* call_result_buf, + constraints->points_to_analysis().GetBufferDefinedAt( + instr, /*index=*/{0})); + + // Set layouts of the instructions' shapes. + if (target == kCudnnConvForwardCallTarget) { + TF_RETURN_IF_ERROR(constraints->SetOperandLayout(input_shape, instr, 0)); + TF_RETURN_IF_ERROR(constraints->SetOperandLayout(filter_shape, instr, 1)); + TF_RETURN_IF_ERROR( + constraints->SetBufferLayout(output_shape.layout(), *call_result_buf)); + } else if (target == kCudnnConvBackwardInputCallTarget) { + TF_RETURN_IF_ERROR(constraints->SetOperandLayout(output_shape, instr, 0)); + TF_RETURN_IF_ERROR(constraints->SetOperandLayout(filter_shape, instr, 1)); + TF_RETURN_IF_ERROR( + constraints->SetBufferLayout(input_shape.layout(), *call_result_buf)); + } else if (target == kCudnnConvBackwardFilterCallTarget) { + TF_RETURN_IF_ERROR(constraints->SetOperandLayout(input_shape, instr, 0)); + TF_RETURN_IF_ERROR(constraints->SetOperandLayout(output_shape, instr, 1)); + TF_RETURN_IF_ERROR( + constraints->SetBufferLayout(filter_shape.layout(), *call_result_buf)); + } else { + LOG(FATAL) << "Unexpected custom call target: " + << instr->custom_call_target(); + } + return Status::OK(); +} + Status GpuLayoutAssignment::AddBackendConstraints( LayoutConstraints* constraints) { for (auto* instruction : constraints->computation()->instructions()) { - // cuDNN is called with specific layouts on the input, output, and filter: - // - // input: DataLayout::kBatchDepthYX - // output: DataLayout::kBatchDepthYX - // filter: FilterLayout::kOutputInputYX - // - // The order dimensions in the constant name is major-to-minor (eg, the - // most-major dimension of the input is batch, most-minor is X). The - // specific dimension numbers these named dimensions correspond to is - // determined by the ConvolutionDimensionNumbers argument. Y is spatial - // dimension 0, and X is spatial dimension 1. - // - // TODO(b/29399649): Be more flexible about handling layouts of cuDNN calls. - if (ImplementedAsDnnConvolution(*instruction)) { - HloInstruction* input = nullptr; - HloInstruction* filter = nullptr; - HloInstruction* output = nullptr; - if (instruction->opcode() == HloOpcode::kConvolution) { - input = instruction->mutable_operand(0); - filter = instruction->mutable_operand(1); - output = instruction; - } else { - CHECK_EQ(HloOpcode::kFusion, instruction->opcode()); - switch (instruction->fusion_kind()) { - case HloInstruction::FusionKind::kConvBackwardFilter: - // filter = BackwardFilterConvolve(input, output) - input = instruction->mutable_operand(0); - filter = instruction; - output = instruction->mutable_operand(1); - break; - case HloInstruction::FusionKind::kConvBackwardInput: - // input = BackwardInputConvolve(output, filter) - input = instruction; - filter = instruction->mutable_operand(1); - output = instruction->mutable_operand(0); - break; - default: - LOG(FATAL) << "Not a convolution-fusion"; - } - } - - // Construct minor-to-major dimension orders for operands and result. - // cuDNN's convolution APIs support the BDYX layout for activations/output - // and the OIYX layout for weights. - // TODO(b/29399649): Be more flexible about handling layouts of cuDNN - // calls after we switch to cuDNN v5. - const ConvolutionDimensionNumbers& dimension_numbers = - instruction->convolution_dimension_numbers(); - std::vector input_layout; - for (int i = dimension_numbers.input_spatial_dimensions_size() - 1; - i >= 0; --i) { - input_layout.push_back(dimension_numbers.input_spatial_dimensions(i)); - } - input_layout.push_back(dimension_numbers.input_feature_dimension()); - input_layout.push_back(dimension_numbers.input_batch_dimension()); - Shape input_shape(input->shape()); - *input_shape.mutable_layout() = LayoutUtil::MakeLayout(input_layout); - - std::vector filter_layout; - for (int i = dimension_numbers.kernel_spatial_dimensions_size() - 1; - i >= 0; --i) { - filter_layout.push_back(dimension_numbers.kernel_spatial_dimensions(i)); - } - filter_layout.push_back( - dimension_numbers.kernel_input_feature_dimension()); - filter_layout.push_back( - dimension_numbers.kernel_output_feature_dimension()); - Shape filter_shape(filter->shape()); - *filter_shape.mutable_layout() = LayoutUtil::MakeLayout(filter_layout); - - std::vector output_layout; - for (int i = dimension_numbers.output_spatial_dimensions_size() - 1; - i >= 0; --i) { - output_layout.push_back(dimension_numbers.output_spatial_dimensions(i)); - } - output_layout.push_back(dimension_numbers.output_feature_dimension()); - output_layout.push_back(dimension_numbers.output_batch_dimension()); - Shape output_shape(output->shape()); - *output_shape.mutable_layout() = LayoutUtil::MakeLayout(output_layout); - - // Set layouts of the instructions' shapes. - if (instruction->opcode() == HloOpcode::kConvolution) { - TF_RETURN_IF_ERROR( - constraints->SetOperandLayout(input_shape, output, 0)); - TF_RETURN_IF_ERROR( - constraints->SetOperandLayout(filter_shape, output, 1)); - TF_RETURN_IF_ERROR( - constraints->SetInstructionLayout(output_shape, output)); - } else { - CHECK_EQ(HloOpcode::kFusion, instruction->opcode()); - switch (instruction->fusion_kind()) { - case HloInstruction::FusionKind::kConvBackwardFilter: - // filter = BackwardFilterConvolve(input, output) - TF_RETURN_IF_ERROR( - constraints->SetOperandLayout(input_shape, filter, 0)); - TF_RETURN_IF_ERROR( - constraints->SetInstructionLayout(filter_shape, filter)); - TF_RETURN_IF_ERROR( - constraints->SetOperandLayout(output_shape, filter, 1)); - break; - case HloInstruction::FusionKind::kConvBackwardInput: - // input = BackwardInputConvolve(output, filter) - TF_RETURN_IF_ERROR( - constraints->SetInstructionLayout(input_shape, input)); - TF_RETURN_IF_ERROR( - constraints->SetOperandLayout(output_shape, input, 0)); - TF_RETURN_IF_ERROR( - constraints->SetOperandLayout(filter_shape, input, 1)); - break; - default: - LOG(FATAL) << "Not a convolution-fusion"; - } - } + if (IsCustomCallToDnnConvolution(*instruction)) { + TF_RETURN_IF_ERROR( + AddBackendConstraintsToDnnConvCustomCall(instruction, constraints)); } } return Status::OK(); @@ -151,9 +143,12 @@ Status GpuLayoutAssignment::AddBackendConstraints( bool GpuLayoutAssignment::CustomCallRequiresMajorFirstLayout( const HloInstruction* instruction) { - // Inputs to cudnn batchnorm custom calls don't need the major-first layout - // (i.e. {n, n-1, ...0}) -- we can handle any layout. - return !IsCustomCallToDnnBatchNorm(*instruction); + // - Inputs to cudnn batchnorm custom calls don't need the major-first layout + // (i.e. {n, n-1, ...0}) -- we can handle any layout. + // - Inputs to cudnn convolution require custom layouts handled in + // AddBackendConstraints. + return !IsCustomCallToDnnBatchNorm(*instruction) && + !IsCustomCallToDnnConvolution(*instruction); } Status GpuLayoutAssignment::PropagateOperandConstraint( diff --git a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc index 98f1f6cc5408b2b86c71f62706ef326ec80e0948..4c45d2e94aebce5496da94841f6a1ae9015615c1 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc @@ -31,8 +31,6 @@ namespace xla { namespace gpu { namespace { -using tensorflow::strings::StrCat; - using LayoutAssignmentTest = HloTestBase; TEST_F(LayoutAssignmentTest, Elementwise) { @@ -116,7 +114,7 @@ TEST_F(LayoutAssignmentTest, BatchNormInference) { for (const Shape& input_shape : AllLayoutsOf(shape)) { for (const Shape& result_shape : AllLayoutsOf(shape)) { - SCOPED_TRACE(StrCat( + SCOPED_TRACE(tensorflow::strings::StrCat( "input_shape=", ShapeUtil::HumanStringWithLayout(input_shape), ", result_shape=", ShapeUtil::HumanStringWithLayout(result_shape))); @@ -188,7 +186,7 @@ TEST_F(LayoutAssignmentTest, BatchNormTraining) { // Enumerate all combinations of shapes. for (const Shape& input_shape : AllLayoutsOf(shape)) { for (const Shape& result_shape : AllLayoutsOf(shape)) { - SCOPED_TRACE(StrCat( + SCOPED_TRACE(tensorflow::strings::StrCat( "input_shape=", ShapeUtil::HumanStringWithLayout(input_shape), ", result_shape=", ShapeUtil::HumanStringWithLayout(result_shape))); @@ -260,7 +258,7 @@ TEST_F(LayoutAssignmentTest, BatchNormGrad) { for (const Shape& input_shape : AllLayoutsOf(shape)) { for (const Shape& result_shape : AllLayoutsOf(shape)) { for (int constrained_param_no : {0, 4}) { - SCOPED_TRACE(StrCat( + SCOPED_TRACE(tensorflow::strings::StrCat( "input_shape=", ShapeUtil::HumanStringWithLayout(input_shape), ", result_shape=", ShapeUtil::HumanStringWithLayout(result_shape))); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc index ae92daef8882de2e7d64b69f68452061cb5507f2..af9897769fda371e47af06c19abce9a06015e094 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc @@ -54,7 +54,7 @@ Status GpuTransferManager::TransferLiteralToInfeed(se::StreamExecutor* executor, if (!ShapeUtil::IsTuple(shape)) { int64 size = GetByteSizeRequirement(shape); - return TransferBufferToInfeed(executor, size, literal.InternalData()); + return TransferBufferToInfeed(executor, size, literal.untyped_data()); } if (ShapeUtil::IsNestedTuple(shape)) { @@ -67,20 +67,21 @@ Status GpuTransferManager::TransferLiteralToInfeed(se::StreamExecutor* executor, // enqueue the resulting destination device addresses with the // infeed manager. std::vector buffers; - buffers.reserve(literal.tuple_literals_size()); + buffers.reserve(ShapeUtil::TupleElementCount(shape)); auto cleanup = tensorflow::gtl::MakeCleanup([buffers]() { for (gpu::InfeedBuffer* b : buffers) { b->Done(); } }); - for (const auto& tuple_element : literal.tuple_literals()) { - const Shape& tuple_element_shape = tuple_element.shape(); + for (int64 i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { + const Shape& tuple_element_shape = + ShapeUtil::GetTupleElementShape(shape, i); int64 tuple_element_size = GetByteSizeRequirement(tuple_element_shape); TF_ASSIGN_OR_RETURN( gpu::InfeedBuffer * buffer, TransferBufferToInfeedInternal(executor, tuple_element_size, - tuple_element.InternalData())); + literal.untyped_data({i}))); buffers.push_back(buffer); } diff --git a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc index c2115c49993ef71c4b6dd584e7e0498807666613..061210352cf12e6802d066d311fd2cb481673f15 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc +++ b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc @@ -22,12 +22,17 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/service/llvm_ir/tuple_ops.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" namespace xla { namespace gpu { +using tensorflow::strings::StrAppend; +using tensorflow::strings::StrCat; + void HloToIrBindings::EmitBasePointersForHlos( tensorflow::gtl::ArraySlice io_hlos, tensorflow::gtl::ArraySlice non_io_hlos) { @@ -191,7 +196,11 @@ static bool BuffersInvariantWithinConsumer( llvm_ir::IrArray HloToIrBindings::GetIrArray(const HloInstruction& hlo, const HloInstruction& consumer, const ShapeIndex& shape_index) { - llvm_ir::IrArray ir_array(GetBasePointer(hlo, shape_index), + llvm::Value* base_ptr = GetBasePointer(hlo, shape_index); + CHECK_NE(base_ptr, nullptr) + << "Buffer not assigned for shape_index " << shape_index.ToString() + << " of " << hlo.ToString(); + llvm_ir::IrArray ir_array(base_ptr, ShapeUtil::GetSubshape(hlo.shape(), shape_index)); alias_analysis_.AddAliasingInformationToIrArray(hlo, &ir_array); @@ -223,5 +232,54 @@ void HloToIrBindings::UnbindAllLocalIrValues() { } } +string HloToIrBindings::ToString() const { + string s = StrCat("** HloToIrBindings **\n"); + StrAppend(&s, " is_nested_=", is_nested_, "\n"); + StrAppend(&s, + " temp_buffer_base_=", llvm_ir::DumpToString(*temp_buffer_base_), + "\n"); + + if (base_ptrs_.empty()) { + return s; + } + + // Iterate over all computations in the module in topological order, and print + // out the base pointers we have in each computation in topological order. + for (const HloComputation* computation : + base_ptrs_.begin()->first->GetModule()->MakeComputationPostOrder()) { + bool is_first = true; + for (const HloInstruction* instr : + computation->MakeInstructionPostOrder()) { + auto it = base_ptrs_.find(instr); + if (it == base_ptrs_.end()) { + continue; + } + if (is_first) { + StrAppend(&s, " Base pointers for computation ", computation->name(), + ":\n"); + is_first = false; + } + StrAppend(&s, " ", instr->ToString()); + + const ShapeTree& shape_tree = it->second; + if (!ShapeUtil::IsTuple(instr->shape())) { + const llvm::Value* val = shape_tree.begin()->second; + StrAppend(&s, " -> ", llvm_ir::DumpToString(*val), "\n"); + continue; + } + + StrAppend(&s, "\n"); + for (auto shape_it = shape_tree.begin(); shape_it != shape_tree.end(); + ++shape_it) { + llvm::Value* val = shape_it->second; + StrAppend(&s, " ", shape_it->first.ToString(), " -> ", + (val != nullptr ? llvm_ir::DumpToString(*val) : "null"), + "\n"); + } + } + } + return s; +} + } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h index 62ae1769a1f2fb3b9acaf35bdf18a793232500b0..3d34311b4368d17cb074aaf33c71fc865e96387e 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h +++ b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h @@ -66,13 +66,14 @@ class HloToIrBindings { } llvm::Value* GetTempBufferBase() const { return temp_buffer_base_; } + void SetTempBufferBase(llvm::Value* v) { temp_buffer_base_ = v; } // A helper method that returns the base pointer of the IrArray containing the // output of "inst".at the given ShapeIndex. llvm::Value* GetBasePointer(const HloInstruction& hlo, const ShapeIndex& shape_index = {}) const { auto it = base_ptrs_.find(&hlo); - CHECK(it != base_ptrs_.end()); + CHECK(it != base_ptrs_.end()) << hlo.ToString(); return it->second.element(shape_index); } @@ -87,6 +88,8 @@ class HloToIrBindings { const HloInstruction& consumer, const ShapeIndex& shape_index = {}); + string ToString() const; + private: // Emits IR to resolve (possibly) recursive GetTupleElement instructions. llvm::Value* EmitGetTupleElement(const HloInstruction* gte, @@ -111,7 +114,7 @@ class HloToIrBindings { std::unordered_map> base_ptrs_; // The address of the memory block that contains all temporary buffers. - llvm::Value* temp_buffer_base_; + llvm::Value* temp_buffer_base_ = nullptr; llvm_ir::AliasAnalysis alias_analysis_; }; diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc index 1d47ffde4331868cbc8a8afb2d01b11e77a7fab0..2d6dad27a59978da6e4719afc50ebee5e641dde0 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc @@ -137,49 +137,6 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfDotUnfused) { .ValueOrDie()); } -TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfConvolutionUnfused) { - HloComputation::Builder builder(TestName()); - auto input = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeShape(F32, {1, 1, 1, 3}), "input")); - auto filter = builder.AddInstruction(HloInstruction::CreateParameter( - 1, ShapeUtil::MakeShape(F32, {1, 1, 1, 2}), "filter")); - - Window conv_window; - WindowDimension* conv_window_row = conv_window.add_dimensions(); - conv_window_row->set_size(1); - WindowDimension* conv_window_col = conv_window.add_dimensions(); - conv_window_col->set_size(2); - conv_window_col->set_padding_high(1); - - ConvolutionDimensionNumbers conv_dnums; - conv_dnums.set_input_batch_dimension(0); - conv_dnums.set_output_batch_dimension(0); - conv_dnums.set_input_feature_dimension(1); - conv_dnums.set_output_feature_dimension(1); - conv_dnums.add_input_spatial_dimensions(2); - conv_dnums.add_output_spatial_dimensions(2); - conv_dnums.add_input_spatial_dimensions(3); - conv_dnums.add_output_spatial_dimensions(3); - conv_dnums.set_kernel_output_feature_dimension(0); - conv_dnums.set_kernel_input_feature_dimension(1); - conv_dnums.add_kernel_spatial_dimensions(2); - conv_dnums.add_kernel_spatial_dimensions(3); - - auto conv = builder.AddInstruction( - HloInstruction::CreateConvolve(ShapeUtil::MakeShape(F32, {1, 1, 1, 3}), - input, filter, conv_window, conv_dnums)); - auto transpose = builder.AddInstruction(HloInstruction::CreateTranspose( - ShapeUtil::MakeShape(F32, {3, 1, 1, 1}), conv, {3, 2, 1, 0})); - builder.AddInstruction( - HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {3}), transpose)); - - auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module.get()) - .ValueOrDie()); -} - TEST_F(InstructionFusionTest, GetTupleElementFused) { HloComputation::Builder builder(TestName()); Shape data_shape = ShapeUtil::MakeShape(F32, {8}); diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc index 76566a9e3dbbc936ff90fe3f440ede14bf4e5233..2f65edffea81db7dba1f8545f92b27ea622044e7 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc @@ -90,43 +90,6 @@ bool ImplementedAsGemm(const HloInstruction& hlo) { return false; } -bool ImplementedAsDnnConvolution(const HloInstruction& hlo) { - // We can only do this if the HLO is unnested. - if (hlo.parent() != hlo.GetModule()->entry_computation()) { - return false; - } - - // Forward convolution. - if (hlo.opcode() == HloOpcode::kConvolution) { - const ConvolutionDimensionNumbers& dnums = - hlo.convolution_dimension_numbers(); - if (dnums.input_spatial_dimensions_size() > 3) { - return false; - } - - // CuDNN does not accept zero-element arguments - if (ShapeUtil::HasZeroElements(hlo.operand(0)->shape()) || - ShapeUtil::HasZeroElements(hlo.operand(1)->shape())) { - return false; - } - - if (window_util::HasWindowReversal(hlo.window())) { - return false; - } - - return true; - } - - // Backward convolution. - if (hlo.opcode() == HloOpcode::kFusion && - (hlo.fusion_kind() == HloInstruction::FusionKind::kConvBackwardFilter || - hlo.fusion_kind() == HloInstruction::FusionKind::kConvBackwardInput)) { - return true; - } - - return false; -} - const char* const kCudnnBatchNormForwardInferenceCallTarget = "__cudnn$batchNormalizationForwardInference"; const char* const kCudnnBatchNormForwardTrainingCallTarget = @@ -144,9 +107,76 @@ bool IsCustomCallToDnnBatchNorm(const HloInstruction& hlo) { target == kCudnnBatchNormBackwardCallTarget; } +const char* const kCudnnConvForwardCallTarget = "__cudnn$convForward"; +const char* const kCudnnConvBackwardInputCallTarget = + "__cudnn$convBackwardInput"; +const char* const kCudnnConvBackwardFilterCallTarget = + "__cudnn$convBackwardFilter"; + +bool IsCustomCallToDnnConvolution(const HloInstruction& hlo) { + if (hlo.opcode() != HloOpcode::kCustomCall) { + return false; + } + const auto& target = hlo.custom_call_target(); + return target == kCudnnConvForwardCallTarget || + target == kCudnnConvBackwardInputCallTarget || + target == kCudnnConvBackwardFilterCallTarget; +} + bool ImplementedAsLibraryCall(const HloInstruction& hlo) { - return ImplementedAsGemm(hlo) || ImplementedAsDnnConvolution(hlo) || - IsCustomCallToDnnBatchNorm(hlo); + return ImplementedAsGemm(hlo) || IsCustomCallToDnnBatchNorm(hlo) || + IsCustomCallToDnnConvolution(hlo); +} + +static HloInstruction* CreateCudnnConv( + const char* call_target, const Shape& shape, HloInstruction* lhs, + HloInstruction* rhs, const Window& window, + const ConvolutionDimensionNumbers& dnums) { + HloComputation* computation = lhs->parent(); + + // This call returns a tuple of (conv_result, scratch_memory), where + // conv_result is the actual result of the convolution, and scratch_memory is + // temporary memory used by cudnn. + // + // At the moment, we don't know how much scratch memory this conv is going to + // use, so we put u8[0] in this place. Later on another pass will choose + // which conv algorithm to use, and at that point we'll modify the shape of + // this second tuple element. + Shape call_shape = + ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeShape(U8, {0})}); + + // Our CustomCall takes three arguments: The conv lhs and rhs, and the cudnn + // algorithm to use. It's up to a later pass to choose the algorithm, so to + // indicate that we haven't yet made a choice, we speicfy -1 for that arg. + HloInstruction* negative_one = computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(-1))); + HloInstruction* custom_call = + computation->AddInstruction(HloInstruction::CreateCustomCall( + call_shape, {lhs, rhs, negative_one}, call_target)); + custom_call->set_window(window); + custom_call->set_convolution_dimension_numbers(dnums); + return custom_call; +} + +HloInstruction* CreateCudnnConvForward( + const Shape& shape, HloInstruction* input, HloInstruction* kernel, + const Window& window, const ConvolutionDimensionNumbers& dnums) { + return CreateCudnnConv(kCudnnConvForwardCallTarget, shape, input, kernel, + window, dnums); +} + +HloInstruction* CreateCudnnConvBackwardInput( + const Shape& shape, HloInstruction* output, HloInstruction* reverse_filter, + const Window& window, const ConvolutionDimensionNumbers& dnums) { + return CreateCudnnConv(kCudnnConvBackwardInputCallTarget, shape, output, + reverse_filter, window, dnums); +} + +HloInstruction* CreateCudnnConvBackwardFilter( + const Shape& shape, HloInstruction* input, HloInstruction* output, + const Window& window, const ConvolutionDimensionNumbers& dnums) { + return CreateCudnnConv(kCudnnConvBackwardFilterCallTarget, shape, input, + output, window, dnums); } bool IsReductionToVector(const HloInstruction& reduce) { diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h index d24ed9879d084e96862885efaae2f79a256cd71d..59455f389e733fee2d6cace7486f919a0c5e834e 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h @@ -22,6 +22,9 @@ limitations under the License. #include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" +// TODO(jlebar): Move functions related to cublas/cudnn to a separate file; they +// don't belong in "ir_emission_utils". + namespace xla { namespace gpu { @@ -30,9 +33,6 @@ constexpr int64 kWarpSize = 32; // Returns true if `hlo` will be implemented as a call to BLAS gemm. bool ImplementedAsGemm(const HloInstruction& hlo); -// Returns true if `hlo` will be implemented as a call to cuDNN convolution. -bool ImplementedAsDnnConvolution(const HloInstruction& hlo); - // A call to cuDNN for batch normalization is represented as CustomCall HLO with // a call target equal to one of these strings. // @@ -58,6 +58,61 @@ extern const char* const kCudnnBatchNormBackwardCallTarget; // sequence of generic HLOs or to a cuDNN CustomCall. bool IsCustomCallToDnnBatchNorm(const HloInstruction& hlo); +// A call to cuDNN for convolution (forward, backward filter, or backward input) +// is represented as a CustomCall HLO with a call target equal to one of these +// strings. +// +// These CustomCalls have window() and convolution_dimension_numbers() set like +// regular convolution ops. They have the same LHS and RHS operands, plus two +// additional constant operands: an int64 operand for the cudnn algorithm and +// a bool operand for whether tensor_ops is enabled. A value of -1 for the cudnn +// algorithm means that the implementation is free to choose the best algorithm +// it can. +// +// These calls output a tuple (conv_result, scratch_memory), where conv_result +// is the actual result of the convolution, and scratch_memory is temporary +// memory used by cudnn. Callers shouldn't inspect scratch_memory, as its value +// is not well-defined. +// +// CudnnConvolutionRewriter lowers kConvolution HLOs to these custom calls. +// When it does so, it chooses algorithm -1 and 0 bytes of scratch space. Later +// on in the pipeline, CudnnConvolutionAlgorithmChooser chooses an explicit +// algorithm for each conv and sets the amount of scratch space needed. +// +// (Representing the scratch memory as an output may seem strange at first, but +// it's quite sensible, from a certain point of view. The scratch buffer is a +// location in memory that the conv can write into, but which it can't legally +// read from, at least until it's written something first. But that's exactly +// the definition of an output buffer.) +extern const char* const kCudnnConvForwardCallTarget; +extern const char* const kCudnnConvBackwardInputCallTarget; +extern const char* const kCudnnConvBackwardFilterCallTarget; + +// Returns true if `hlo` will be implemented as a call to a cuDNN convolution +// routine. +// +// This returns true if `hlo` is a CustomCall HLO with a call target equal to +// one of the kCudnnConvFoo constants above, but returns *false* for HLOs with a +// kConvolution opcode. +bool IsCustomCallToDnnConvolution(const HloInstruction& hlo); + +// Creates a CustomCall for a cudnn forward/backward-input/backward-filter conv. +// Note that these CustomCalls return a tuple (conv_result, scratch_memory). If +// you want just the conv result, you'll need to get-tuple-element the value +// returned by this function. +// +// The created cudnn call will use the default cudnn algorithm and no scratch +// space. +HloInstruction* CreateCudnnConvForward( + const Shape& shape, HloInstruction* input, HloInstruction* kernel, + const Window& window, const ConvolutionDimensionNumbers& dnums); +HloInstruction* CreateCudnnConvBackwardInput( + const Shape& shape, HloInstruction* output, HloInstruction* reverse_filter, + const Window& window, const ConvolutionDimensionNumbers& dnums); +HloInstruction* CreateCudnnConvBackwardFilter( + const Shape& shape, HloInstruction* input, HloInstruction* output, + const Window& window, const ConvolutionDimensionNumbers& dnums); + // Returns true if `hlo` will be implemented as a library call, e.g. cuBLAS gemm // or cuDNN convolution. bool ImplementedAsLibraryCall(const HloInstruction& hlo); diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc index d85d6aae12ce54e005d8ffb79a7914564885b6dd..a3df67a87344d6ece2ea9047321ad9542c13f8cf 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc @@ -27,6 +27,8 @@ limitations under the License. #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/service/elemental_ir_emitter.h" #include "tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emitter_nested.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h" #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h" @@ -269,7 +271,7 @@ bool IrEmitter::MaybeEmitDirectAtomicOperation( // cas_new_output_address = alloca(atomic_size); // cas_old_output_address = alloca(atomic_size); // if (atomic_size != element_size) { -// atomic_address = output_address & ((int64)(-2)); +// atomic_address = output_address & ((int64)(-4)); // new_output_address = cas_new_output_address + (output_address & 3); // } else { // atomic_address = output_address; @@ -326,7 +328,7 @@ Status IrEmitter::EmitAtomicOperationUsingCAS(const HloComputation& computation, ir_builder_.CreatePtrToInt(output_address, address_int_type); llvm::Value* mask = llvm::ConstantInt::get(address_int_type, 3); llvm::Value* offset = ir_builder_.CreateAnd(atomic_memory_address, mask); - mask = llvm::ConstantInt::get(address_int_type, -2); + mask = llvm::ConstantInt::get(address_int_type, -4); atomic_memory_address = ir_builder_.CreateAnd(atomic_memory_address, mask); atomic_memory_address = ir_builder_.CreateIntToPtr(atomic_memory_address, atomic_address_type); @@ -605,10 +607,17 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) { "Hit a case for convolution that is not implemented on GPU."); } +Status IrEmitter::HandleFft(HloInstruction* fft) { + if (ShapeUtil::HasZeroElements(fft->shape())) { + // Emit no code for an empty output. + return Status::OK(); + } + return Unimplemented("Hit a case for fft that is not implemented on GPU."); +} + Status IrEmitter::HandleCrossReplicaSum(HloInstruction* crs) { // TODO(b/33011107): Support cross replica sum on GPU. - return Unimplemented( - "Cross replica sum not implemented on GPU. See b/33011107."); + return Unimplemented("CrossReplicaSum is not implemented on GPU."); } Status IrEmitter::HandleParameter(HloInstruction* parameter) { @@ -702,11 +711,13 @@ Status IrEmitter::HandleCustomCall(HloInstruction*) { } Status IrEmitter::HandleInfeed(HloInstruction*) { - return Unimplemented("Infeed is not supported on GPU (b/30467474)."); + // TODO(b/30467474): Implement infeed on GPU. + return Unimplemented("Infeed is not supported on GPU."); } Status IrEmitter::HandleOutfeed(HloInstruction*) { - return Unimplemented("Outfeed is not supported on GPU (b/34359662)."); + // TODO(b/34359662): Implement outfeed on GPU. + return Unimplemented("Outfeed is not supported on GPU."); } Status IrEmitter::HandleRng(HloInstruction* random) { @@ -750,37 +761,6 @@ Status IrEmitter::HandleBatchNormGrad(HloInstruction*) { "to a cudnn CustomCall using CudnnBatchNormRewriter."); } -Status IrEmitter::HandleConditional(HloInstruction* conditional) { - auto pred = conditional->operand(0); - auto true_arg = conditional->operand(1); - auto false_arg = conditional->operand(2); - - llvm::Value* conditional_result = GetBasePointer(*conditional); - - llvm::LoadInst* pred_value = ir_builder_.CreateLoad( - GetBasePointer(*pred), - llvm_ir::AsStringRef(IrName(conditional, "load_predicate_value"))); - llvm::Value* pred_cond = ir_builder_.CreateICmpNE( - pred_value, - llvm::ConstantInt::get(llvm_ir::PrimitiveTypeToIrType(PRED, module_), 0), - llvm_ir::AsStringRef(IrName(conditional, "boolean_predicate"))); - llvm_ir::LlvmIfData if_data = llvm_ir::EmitIfThenElse( - pred_cond, IrName(conditional, "if_then_else"), &ir_builder_); - - SetToFirstInsertPoint(if_data.true_block, &ir_builder_); - TF_RETURN_IF_ERROR(EmitCallToNestedComputation( - *conditional->true_computation(), {GetBasePointer(*true_arg)}, - conditional_result)); - - SetToFirstInsertPoint(if_data.false_block, &ir_builder_); - TF_RETURN_IF_ERROR(EmitCallToNestedComputation( - *conditional->false_computation(), {GetBasePointer(*false_arg)}, - conditional_result)); - - SetToFirstInsertPoint(if_data.after_block, &ir_builder_); - return Status::OK(); -} - llvm_ir::IrArray::Index IrEmitter::EmitOperandArrayLoopNest( const llvm_ir::IrArray& operand_array, int64 reduction_dimension, tensorflow::StringPiece name_suffix, llvm_ir::ForLoopNest* loop_nest) { diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter.h b/tensorflow/compiler/xla/service/gpu/ir_emitter.h index 41d013c13dc732f56a9ff35e1f12eea436f68144..b0accc08d479258d65a18202122e4c9e90ff78d0 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter.h @@ -13,19 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// An XLA HLO graph may contain multiple computations. These computations -// fall into two types, nested and unnested. We translate each nested -// computation (e.g. the computation operand of a Map operator) to a device -// function. For each unnested computation composed of top-level -// HloInstructions, we generate a CUDA kernel for each HloInstruction. -// -// This file declares classes that translate an XLA HLO graph to LLVM IR for -// GPUs. IrEmitterNested emits LLVM IR for nested computations, and -// IrEmitterUnnested for unnested computations. The logic of emitting LLVM IR -// for each individual HloInstruction is largely the same between these two -// classes. Therefore, we implement the common logic in the Handle* functions in -// the superclass IrEmitter. - #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_H_ @@ -60,25 +47,35 @@ limitations under the License. namespace xla { namespace gpu { -// This class is the top-level API for the XLA HLO --> LLVM IR compiler. -// It implements the DfsHloVisitor interface and emits an LLVM IR program that -// implements the input HLO graph. +// Abstract base class for translating HLO graphs to LLVM IR for a GPU. +// +// There are two concrete subclasses of IrEmitter: IrEmitterNested and +// IrEmitterUnnested. In the unnested variety, each HLO gets its own kernel +// function, whereas in the nested version the whole computation is emitted as +// one *non-kernel* function. +// +// In XLA, kernel functions never call other kernel functions. This means that +// if we have a kernel -- e.g. implementing a kReduce HLO -- that wants to use +// an HLO computation as a "subroutine" -- e.g. the HLO computation that +// specifies how to reduce two elements -- then the subroutine computation must +// be emitted using IrEmitterNested. // -// Note: if `T` is a subclass of `IrEmitter` and a handler is not overridden in -// either `IrEmitter` or `T`, the handler in `DfsHloVisitorWithDefault` -// calls `T::DefaultAction`. +// Fusion nodes are a special case. A fusion node is emitted using +// IrEmitterUnnested, but the code is generated using FusedIrEmitter, which is +// not a subclass of gpu::IrEmitter, and in fact is better understood as an IR +// generator generator. See comments on that class. class IrEmitter : public DfsHloVisitorWithDefault { public: IrEmitter(const IrEmitter&) = delete; IrEmitter& operator=(const IrEmitter&) = delete; - // The following methods implement the DfsHloVisitorWithDefault interface. Status DefaultAction(HloInstruction* hlo) override; Status HandleConstant(HloInstruction* constant) override; Status HandleBitcast(HloInstruction* bitcast) override; Status HandleGetTupleElement(HloInstruction* get_tuple_element) override; Status HandleDot(HloInstruction* dot) override; Status HandleConvolution(HloInstruction* convolution) override; + Status HandleFft(HloInstruction* fft) override; Status HandleCrossReplicaSum(HloInstruction* crs) override; Status HandleInfeed(HloInstruction* infeed) override; Status HandleOutfeed(HloInstruction* outfeed) override; @@ -95,7 +92,6 @@ class IrEmitter : public DfsHloVisitorWithDefault { Status HandleCall(HloInstruction* call) override; Status HandleCustomCall(HloInstruction* custom_call) override; Status HandleRng(HloInstruction* random) override; - Status HandleConditional(HloInstruction* conditional) override; Status HandleBatchNormInference(HloInstruction* batch_norm) override; Status HandleBatchNormTraining(HloInstruction* batch_norm) override; Status HandleBatchNormGrad(HloInstruction* batch_norm) override; @@ -217,187 +213,6 @@ class IrEmitter : public DfsHloVisitorWithDefault { std::map computation_to_ir_function_; }; -// Emits LLVM IR for unnested computations. Each HloInstruction is translated to -// a separate CUDA kernel. These kernels are inserted into the resultant module -// sorted in reverse postorder of the XLA HLO graph. -class IrEmitterUnnested : public IrEmitter { - public: - IrEmitterUnnested(const HloModuleConfig& hlo_module_config, - const HloComputation* hlo_computation, - IrEmitterContext* ir_emitter_context); - IrEmitterUnnested(const IrEmitterUnnested&) = delete; - IrEmitterUnnested& operator=(const IrEmitterUnnested&) = delete; - - // Transfers the ownship of thunk_sequence_ out. - std::unique_ptr ConsumeThunkSequence() { - return std::move(thunk_sequence_); - } - - Status DefaultAction(HloInstruction* hlo) override; - - // IrEmitterUnnested handles the following instructions differently from - // IrEmitter. - Status HandleCopy(HloInstruction* copy) override; - Status HandleConditional(HloInstruction* conditional) override; - Status HandleConvolution(HloInstruction* convolution) override; - Status HandleCustomCall(HloInstruction* custom_call) override; - Status HandleDot(HloInstruction* dot) override; - Status HandleFusion(HloInstruction* fusion) override; - Status HandleGetTupleElement(HloInstruction* get_tuple_element) override; - Status HandleReduce(HloInstruction* reduce) override; - Status HandleSelectAndScatter(HloInstruction* instruction) override; - Status HandleTuple(HloInstruction* tuple) override; - Status HandleWhile(HloInstruction* xla_while) override; - Status HandleInfeed(HloInstruction* xla_infeed) override; - Status HandleRng(HloInstruction* random) override; - Status HandleSelect(HloInstruction* select) override; - - Status EmitTargetElementLoop( - const HloInstruction& hlo, - const llvm_ir::ElementGenerator& body_emitter) override; - - // Same as `EmitTargetElementLoop`, but in given `thunk` rather than - // `LastThunk()`. - Status EmitTargetElementLoopInThunk( - const HloInstruction& hlo, const llvm_ir::ElementGenerator& body_emitter, - KernelThunk* thunk); - - private: - // Builds the appropriate thunk for the instruction hlo and returns the owning - // pointer to it. The caller needs to make sure `inst` outlives the lifetime - // of the returned Thunk object. - std::unique_ptr BuildThunk(const HloInstruction* hlo); - - // Builds the prototype of the IR kernel for `inst` and adds it to the module. - llvm::Function* BuildKernelPrototype( - const HloInstruction& inst, - tensorflow::gtl::ArraySlice escaped_hlos); - - // Emits the base pointers for `hlo` and its operands. `io_hlos` will store - // all input/output HLOs among `hlo` and its operands. - llvm::Function* EmitBasePointersForHloAndItsOperands( - const HloInstruction& hlo, std::vector* io_hlos); - - // EmitColumnReduction and EmitRowReduction emit code for column and row - // reduction of a matrix and/or 3D tensor. Row and column reduction have - // different memory access pattern, so for performance their implementations - // are significantly different. - // - // Emits code that reduces a matrix of shape [height x width] to a vector of - // [width]. Other parameters have the same meaning as those of - // `EmitReductionToVector`. Note that input shape might not be - // [height x width], but can be bitcast to [height x weight] with "height" - // being the major dimension. - Status EmitColumnReduction(int64 height, int64 width, HloInstruction* reduce, - const Shape& input_shape, - const llvm_ir::ElementGenerator& input_gen, - const llvm_ir::ElementGenerator& init_value_gen, - HloComputation* reducer); - - // Emits code that reduces a 3D tensor of shape [depth x height x width] to a - // vector of shape [height]. Other parameters have the same meaning as those - // of `EmitReductionToVector`. Note that input shape might not be - // [depth x height x width], but can be bitcast to [depth x height x weight] - // with "depth" being the most major dimension. - Status EmitRowReduction(int64 depth, int64 height, int64 width, - HloInstruction* reduce, const Shape& input_shape, - const llvm_ir::ElementGenerator& input_gen, - const llvm_ir::ElementGenerator& init_value_gen, - HloComputation* reducer); - - // Figures out whether `reduce` is a row or column reduction, and which - // dimensions to reduce, and calls either `EmitRowReduction` or - // `EmitColumnReduction` as appropriate. `input_shape` is the shape of the - // input array, which is the operand of the Reduce instruction if unfused or - // of the Fusion instruction if fused. `input_gen` and `init_value_gen` - // generate elements of the input and the initial value. Other parameters mean - // the same as for `HandleReduce`. - // - // Prerequisite: `IsReductionToVector(*reduce)` - Status EmitReductionToVector( - HloInstruction* reduce, const Shape& input_shape, - const llvm_ir::ElementGenerator& input_gen, - const llvm_ir::ElementGenerator& init_value_gen, - tensorflow::gtl::ArraySlice dimensions_to_reduce, - HloComputation* reducer); - - // Emits code to initialize buffer of `inst` in given `thunk`. - Status EmitInitializer(const HloInstruction* inst, KernelThunk* thunk); - - // Returns a KernelThunk that invokes the kernel emitted for `inst`. The - // caller needs to make sure `inst` outlives the lifetime of the returned - // Thunk object. - std::unique_ptr BuildKernelThunk(const HloInstruction* inst); - - // Returns a ConvolutionThunk that calls DNN to implement `inst`. - std::unique_ptr BuildConvolutionThunk(const HloInstruction* inst); - - // Returns a GemmThunk that calls gemm to implement `inst`. The caller needs - // to make sure `inst` outlives the lifetime of the returned Thunk object. - std::unique_ptr BuildGemmThunk(const HloInstruction* inst); - - // Returns a thunk that calls host-to-device cuMemcpy to implement `inst`. - std::unique_ptr BuildHostToDeviceCopyThunk(const HloInstruction* inst); - - // Returns a thunk that calls device-to-device cuMemcpy to implement `inst`. - std::unique_ptr BuildDeviceToDeviceCopyThunk( - const HloInstruction* inst); - - // Returns an InfeedThunk that performs device-to-device memcpy to implement - // `inst`. - std::unique_ptr BuildInfeedThunk(const HloInstruction* inst); - - // Returns a WhileThunk that invokes thunk sequences for 'condition' and - // 'body' sub-computations of while instruction 'hlo'. - std::unique_ptr BuildWhileThunk(const HloInstruction* hlo); - - // Returns a ForThunk which executes 'loop_limit' invocations of a thunk - // sequence from the 'body' sub-computation of the while instruction 'hlo'. - std::unique_ptr BuildForThunk(const HloInstruction* hlo, - const int64 loop_limit); - - Status Postprocess(HloInstruction* hlo) override; - - // Returns the last generated thunk. - Thunk* LastThunk() const { return thunk_sequence_->back().get(); } - - // The thunk sequence this IrEmitter generates for the input computation. - std::unique_ptr thunk_sequence_; - - // The HloComputation that this IrEmitter emits code for. - const HloComputation* hlo_computation_; -}; - -// Emits LLVM IR for a nested computation to the resultant function. -class IrEmitterNested : public IrEmitter { - public: - // Constructs an LLVM IR emitter for a nested HLO computation. `function` is - // the containing IR function this emitter produces IR to. See - // IrEmitter::IrEmitter for the meanings of other arguments. - IrEmitterNested(const HloModuleConfig& hlo_module_config, - const HloComputation& nested_computation, - IrEmitterContext* ir_emitter_context); - IrEmitterNested(const IrEmitterNested&) = delete; - IrEmitterNested& operator=(const IrEmitterNested&) = delete; - - // Overrides the default empty implementation. Binds the given instruction - // "parameter" with the parameter of the IR function. - Status HandleParameter(HloInstruction* parameter) override; - - llvm::Function* GetEmittedFunction() const { return emitted_function_; } - - Status EmitTargetElementLoop( - const HloInstruction& hlo, - const llvm_ir::ElementGenerator& body_emitter) override; - - private: - llvm::Function* EmitBasePointersForNestedComputation( - const HloComputation& nested_computation, - std::vector* io_hlos); - - llvm::Function* emitted_function_; -}; - } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.cc index 5225ff36ff3a8a1b049479c34aa301de8724f73e..71aada080ae8df70bffce3e1854b5fbd833efd23 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.cc @@ -16,12 +16,13 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/service/gpu/ir_emitter_nested.h" + #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Function.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instructions.h" #include "tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h" -#include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter_context.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.h new file mode 100644 index 0000000000000000000000000000000000000000..ca11cf2c182b0600b931b19d2d7fb3983e36441a --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.h @@ -0,0 +1,72 @@ +/* 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_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_NESTED_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_NESTED_H_ + +#include "llvm/IR/Function.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" + +namespace xla { +namespace gpu { + +// Emits LLVM IR for a "nested computation" into a non-kernel device function. +// +// This is used to emit code for HloComputations that don't require a separate +// kernel call. For example, IrEmitterNested is used to emit code for a kReduce +// HLO's elementwise reduction computation. Notably, IrEmitterNested is *not* +// used to emit code for fusion nodes -- fusion nodes use FusedIrEmitter, which +// is a different beast altogether. +// +// IrEmitterNested generates a non-kernel function with the following +// parameters: +// +// - N pointers to the buffers of each of the N parameters to the computation, +// - a pointer to the output buffer of the computation, and +// - a pointer to the top-level temp buffer. +// +class IrEmitterNested : public IrEmitter { + public: + // Constructs an LLVM IR emitter for a nested HLO computation. `function` is + // the containing IR function this emitter produces IR to. See + // IrEmitter::IrEmitter for the meanings of other arguments. + IrEmitterNested(const HloModuleConfig& hlo_module_config, + const HloComputation& nested_computation, + IrEmitterContext* ir_emitter_context); + IrEmitterNested(const IrEmitterNested&) = delete; + IrEmitterNested& operator=(const IrEmitterNested&) = delete; + + // Overrides the default empty implementation. Binds the given instruction + // "parameter" with the parameter of the IR function. + Status HandleParameter(HloInstruction* parameter) override; + + llvm::Function* GetEmittedFunction() const { return emitted_function_; } + + Status EmitTargetElementLoop( + const HloInstruction& hlo, + const llvm_ir::ElementGenerator& body_emitter) override; + + private: + llvm::Function* EmitBasePointersForNestedComputation( + const HloComputation& nested_computation, + std::vector* io_hlos); + + llvm::Function* emitted_function_; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_NESTED_H_ diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc index 1be3473f52b837c30700a14330c672072651ebf8..aa2a0a9800bab142481e1def785c9052526fcd8c 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc @@ -17,6 +17,8 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h" + #include "llvm/ADT/StringRef.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Function.h" @@ -28,15 +30,18 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor.h" +#include "tensorflow/compiler/xla/service/gpu/conditional_thunk.h" #include "tensorflow/compiler/xla/service/gpu/convolution_thunk.h" #include "tensorflow/compiler/xla/service/gpu/copy_thunk.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_thunk.h" +#include "tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h" +#include "tensorflow/compiler/xla/service/gpu/fft_thunk.h" #include "tensorflow/compiler/xla/service/gpu/for_thunk.h" #include "tensorflow/compiler/xla/service/gpu/gemm_thunk.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_constants.h" #include "tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h" #include "tensorflow/compiler/xla/service/gpu/infeed_thunk.h" #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" -#include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter_context.h" #include "tensorflow/compiler/xla/service/gpu/kernel_thunk.h" #include "tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h" @@ -70,6 +75,10 @@ namespace gpu { namespace { using llvm_ir::IrName; +using tensorflow::gtl::ArraySlice; +using tensorflow::gtl::nullopt; +using tensorflow::gtl::optional; +using tensorflow::strings::StrCat; // If a dimensions is smaller than this, untiled transposition may be more // efficient. @@ -132,6 +141,38 @@ void UpdateLaunchDimensions(const LaunchDimensions& launch_dims, Thunk* thunk, llvm::MDString::get(llvm_context, "reqntidx"), llvm::ConstantAsMetadata::get(threads_per_block_ir_value)})); } + +// Tries to get a Slice for the given instruction at the given index, but +// returns nullopt if we might not know the slice's address at runtime without +// dereferencing a containing tuple. +// +// In particular, when XLA accepts a parameter of tuple type, the caller has the +// option of telling XLA what are the values inside of the tuple, or just giving +// XLA a pointer to the top-level tuple and letting us chase the pointers on the +// GPU. We therefore cannot rely having these pointers to parameter sub-buffers +// being present when we run the program. +optional GetKnownAtRuntimeSlice( + const HloInstruction* instr, const ShapeIndex& index, + const BufferAssignment& buffer_assn) { + auto maybe_slice = buffer_assn.GetUniqueSlice(instr, index); + if (!maybe_slice.ok()) { + return nullopt; + } + // BufferAllocation gives a slice and alloc to every buffer accessed by XLA, + // but we don't necessarily know the runtime address of sub-buffers of input + // parameters. + const BufferAllocation::Slice& slice = maybe_slice.ValueOrDie(); + const BufferAllocation* alloc = slice.allocation(); + if (alloc->IsInputOrOutput() && !alloc->maybe_live_out() && + !alloc->param_shape_index().empty()) { + return nullopt; + } + + // Otherwise, we will know the address of this slice at runtime without having + // to dereference a tuple. + return slice; +} + } // namespace IrEmitterUnnested::IrEmitterUnnested(const HloModuleConfig& hlo_module_config, @@ -149,16 +190,20 @@ Status IrEmitterUnnested::Postprocess(HloInstruction* hlo) { } namespace { -bool ImplementedAsHostToDeviceMemcpy(const HloInstruction& hlo) { - // `hlo` needs to satisfy three conditions to be implemented as a +bool ImplementedAsHostToDeviceMemcpy(const BufferAssignment& buffer_assignment, + const HloInstruction& hlo) { + // `hlo` needs to satisfy the following conditions to be implemented as a // host-to-device cuMemcpy. // // 1. `hlo` is a kCopy instruction. // 2. `hlo`'s only operand is a kConstant instruction. // 3. `hlo` and its operand have the same shape (thus the same layout too). + // 4. The address of `hlo`'s buffer is known at runtime (without dereferencing + // pointers in a tuple). return hlo.opcode() == HloOpcode::kCopy && hlo.operand(0)->opcode() == HloOpcode::kConstant && - ShapeUtil::Equal(hlo.operand(0)->shape(), hlo.shape()); + ShapeUtil::Equal(hlo.operand(0)->shape(), hlo.shape()) && + GetKnownAtRuntimeSlice(&hlo, {}, buffer_assignment).has_value(); } bool ImplementedAsDeviceToDeviceMemcpy( @@ -172,52 +217,50 @@ bool ImplementedAsDeviceToDeviceMemcpy( // instance) which means the source buffer also resides on the device. return hlo.opcode() == HloOpcode::kCopy && ShapeUtil::Equal(hlo.operand(0)->shape(), hlo.shape()) && - buffer_assignment.HasTopLevelAllocation(hlo.operand(0)); + GetKnownAtRuntimeSlice(&hlo, {}, buffer_assignment).has_value() && + GetKnownAtRuntimeSlice(hlo.operand(0), {}, buffer_assignment) + .has_value(); } } // namespace llvm::Function* IrEmitterUnnested::BuildKernelPrototype( const HloInstruction& inst, - tensorflow::gtl::ArraySlice escaped_hlos) { + tensorflow::gtl::ArraySlice args) { // Compute the kernel name. The opcode string may contain "-" which cannot be // in a PTX function name, so sanitize the name before uniquifying it. string kernel_name = ir_emitter_context_->name_uniquer()->GetUniqueName( llvm_ir::SanitizeFunctionName(inst.name())); - // Create the kernel and adds it to the module. + // Create the kernel and add it to the module. llvm::Module* module = ir_emitter_context_->llvm_module(); llvm::LLVMContext& context = module->getContext(); - int num_escaped_hlos = escaped_hlos.size(); llvm::FunctionType* kernel_type = llvm::FunctionType::get( - llvm::Type::getVoidTy(context), // The type of function result. - std::vector(num_escaped_hlos + 1, - ir_builder_.getInt8PtrTy()), - false); // Not a variadic argument function. + /*Result=*/llvm::Type::getVoidTy(context), + std::vector(args.size(), ir_builder_.getInt8PtrTy()), + /*isVarArg=*/false); llvm::Function* kernel = llvm::Function::Create(kernel_type, llvm::GlobalValue::ExternalLinkage, kernel_name.c_str(), module); - // Add dereferenceable information to each of the escaped HLO parameters. - for (size_t arg_no = 0; arg_no < escaped_hlos.size(); ++arg_no) { - const HloInstruction* escaped_hlo = escaped_hlos[arg_no]; - const Shape& escaped_hlo_shape = escaped_hlo->shape(); - int64 escaped_hlo_size = llvm_ir::ByteSizeOf( - escaped_hlo_shape, ir_emitter_context_->llvm_module()->getDataLayout()); - kernel->addDereferenceableAttr(arg_no + 1, escaped_hlo_size); - } - - // The last argument is a pointer to the temporary buffer memory block. - // We know that it doesn't alias any of the escaped arguments (the inputs + - // the result). We also know how many bytes can be dereferenced in it. - const llvm::Argument& temp_buffer = *std::prev(kernel->arg_end()); - int64 temp_buffer_arg_no = temp_buffer.getArgNo(); - int64 temp_allocation_total_size = - ir_emitter_context_->buffer_assignment().temp_allocation_total_size(); - if (temp_allocation_total_size != 0) { - kernel->addDereferenceableAttr(temp_buffer_arg_no + 1, - temp_allocation_total_size); + // Add dereferenceable and alignment information to each of the kernel's + // parameters. + auto arg_it = kernel->arg_begin(); + for (size_t arg_no = 0; arg_no < args.size(); ++arg_no) { + const BufferAllocation* alloc = args[arg_no]; + llvm::Argument* fn_arg = &*arg_it; + ++arg_it; + + kernel->addDereferenceableAttr(arg_no + 1, alloc->size()); + kernel->addParamAttr( + arg_no, llvm::Attribute::get(context, llvm::Attribute::Alignment, + kCudaMallocAlignBytes)); + + if (alloc->IsPreallocatedTempBuffer()) { + fn_arg->setName("temp_buf"); + } else { + fn_arg->setName(llvm_ir::AsStringRef(StrCat("alloc", alloc->index()))); + } } - kernel->addAttribute(temp_buffer_arg_no + 1, llvm::Attribute::NoAlias); // TODO(b/65380986): Investigate if adding fast math flags for generated // kernels makes sense. @@ -233,10 +276,9 @@ llvm::Function* IrEmitterUnnested::BuildKernelPrototype( // Update the insert point to the entry basic block. llvm::BasicBlock* entry_bb = - llvm::BasicBlock::Create(context, - "entry", // The name of the basic block. - kernel); // The parent/owner of "entry_bb". - // Emit a "return void" at entry_bb's end, and sets the insert point before + llvm::BasicBlock::Create(context, /*Name=*/"entry", /*Parent=*/kernel); + + // Emit a "return void" at entry_bb's end, and set the insert point before // that return instruction. ir_builder_.SetInsertPoint(llvm::ReturnInst::Create(context, entry_bb)); @@ -263,15 +305,11 @@ Status IrEmitterUnnested::HandleDot(HloInstruction* dot) { } Status IrEmitterUnnested::HandleConditional(HloInstruction* conditional) { - thunk_sequence_->push_back(BuildKernelThunk(conditional)); - return IrEmitter::HandleConditional(conditional); + thunk_sequence_->emplace_back(BuildConditionalThunk(conditional)); + return Status::OK(); } Status IrEmitterUnnested::HandleConvolution(HloInstruction* convolution) { - if (ImplementedAsDnnConvolution(*convolution)) { - thunk_sequence_->emplace_back(BuildConvolutionThunk(convolution)); - return Status::OK(); - } thunk_sequence_->emplace_back(BuildKernelThunk(convolution)); return IrEmitter::HandleConvolution(convolution); } @@ -370,9 +408,87 @@ Status IrEmitterUnnested::HandleCustomCall(HloInstruction* custom_call) { return Status::OK(); } + if (IsCustomCallToDnnConvolution(*custom_call)) { + const auto& assn = ir_emitter_context_->buffer_assignment(); + const auto& lhs_shape = custom_call->operand(0)->shape(); + const auto& rhs_shape = custom_call->operand(1)->shape(); + const auto& conv_result_shape = custom_call->shape().tuple_shapes(0); + auto lhs_slice = GetAllocationSlice(*custom_call->operand(0)); + auto rhs_slice = GetAllocationSlice(*custom_call->operand(1)); + auto tuple_result_slice = GetAllocationSlice(*custom_call); + auto conv_result_slice = assn.GetUniqueSlice(custom_call, {0}).ValueOrDie(); + auto scratch_slice = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); + + const HloInstruction* algorithm_inst = custom_call->operand(2); + CHECK(algorithm_inst->IsConstant()) << algorithm_inst->ToString(); + int64 algorithm = algorithm_inst->literal().Get({}); + + const HloInstruction* tensor_ops_enabled_inst = custom_call->operand(3); + CHECK(tensor_ops_enabled_inst->IsConstant()) + << tensor_ops_enabled_inst->ToString(); + bool tensor_ops_enabled = tensor_ops_enabled_inst->literal().Get({}); + + const auto& target = custom_call->custom_call_target(); + std::unique_ptr thunk; + if (target == kCudnnConvForwardCallTarget) { + thunk = MakeUnique( + CudnnConvKind::kForward, + /*input_buffer=*/lhs_slice, + /*filter_buffer=*/rhs_slice, + /*output_buffer=*/conv_result_slice, + /*tuple_result_buffer=*/tuple_result_slice, + /*scratch_buffer=*/scratch_slice, + /*input_shape=*/lhs_shape, + /*filter_shape=*/rhs_shape, + /*output_shape=*/conv_result_shape, // + custom_call->window(), custom_call->convolution_dimension_numbers(), + algorithm, tensor_ops_enabled, custom_call); + } else if (target == kCudnnConvBackwardInputCallTarget) { + thunk = MakeUnique( + CudnnConvKind::kBackwardInput, + /*input_buffer=*/conv_result_slice, + /*filter_buffer=*/rhs_slice, + /*output_buffer=*/lhs_slice, + /*tuple_result_buffer=*/tuple_result_slice, + /*scratch_buffer=*/scratch_slice, + /*input_shape=*/conv_result_shape, + /*filter_shape=*/rhs_shape, + /*output_shape=*/lhs_shape, // + custom_call->window(), custom_call->convolution_dimension_numbers(), + algorithm, tensor_ops_enabled, custom_call); + } else if (target == kCudnnConvBackwardFilterCallTarget) { + thunk = MakeUnique( + CudnnConvKind::kBackwardFilter, + /*input_buffer=*/lhs_slice, + /*filter_buffer=*/conv_result_slice, + /*output_buffer=*/rhs_slice, + /*tuple_result_buffer=*/tuple_result_slice, + /*scratch_buffer=*/scratch_slice, + /*input_shape=*/lhs_shape, + /*filter_shape=*/conv_result_shape, + /*output_shape=*/rhs_shape, // + custom_call->window(), custom_call->convolution_dimension_numbers(), + algorithm, tensor_ops_enabled, custom_call); + } else { + LOG(FATAL) << "Unexpected custom call target: " + << custom_call->custom_call_target(); + } + + thunk_sequence_->emplace_back(std::move(thunk)); + return Status::OK(); + } + return IrEmitter::HandleCustomCall(custom_call); } +Status IrEmitterUnnested::HandleFft(HloInstruction* fft) { + TF_RET_CHECK( + LayoutUtil::IsMonotonicWithDim0Major(fft->operand(0)->shape().layout())); + TF_RET_CHECK(LayoutUtil::IsMonotonicWithDim0Major(fft->shape().layout())); + thunk_sequence_->emplace_back(BuildFftThunk(fft)); + return Status::OK(); +} + Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { HloInstruction* root = fusion->fused_expression_root(); // HandleFusion specializes reduction from a multi-dimensional array to a 1D @@ -482,10 +598,6 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { thunk_sequence_->emplace_back(BuildGemmThunk(fusion)); return Status::OK(); } - if (ImplementedAsDnnConvolution(*fusion)) { - thunk_sequence_->emplace_back(BuildConvolutionThunk(fusion)); - return Status::OK(); - } thunk_sequence_->emplace_back(BuildKernelThunk(fusion)); return IrEmitter::HandleFusion(fusion); } @@ -787,7 +899,8 @@ int64 EmitTranspose021Tiled(llvm_ir::IrArray input, llvm_ir::IrArray output, } // namespace Status IrEmitterUnnested::HandleCopy(HloInstruction* copy) { - if (ImplementedAsHostToDeviceMemcpy(*copy)) { + if (ImplementedAsHostToDeviceMemcpy(ir_emitter_context_->buffer_assignment(), + *copy)) { thunk_sequence_->emplace_back(BuildHostToDeviceCopyThunk(copy)); return Status::OK(); } @@ -821,6 +934,194 @@ Status IrEmitterUnnested::HandleCopy(HloInstruction* copy) { return IrEmitter::HandleCopy(copy); } +Status IrEmitterUnnested::EmitReductionToScalar( + HloInstruction* reduce, const Shape& input_shape, + const llvm_ir::ElementGenerator& input_gen, + const llvm_ir::ElementGenerator& init_value_gen, HloComputation* reducer) { + // Number of elements processed by a single thread. + constexpr int64 kTileSize = 16; + int64 num_elems = ShapeUtil::ElementsIn(input_shape); + + // Round up the number of tiles to a multiple of the warp size. This is + // necessary for correctness. We launch one thread per tile, and if the + // number of threads isn't a multiple of the number of the warp size, our + // shuffles will read from inactive threads, producing undefined values. + int64 num_tiles = + RoundUpToNearest(CeilOfRatio(num_elems, kTileSize), kWarpSize); + + // Check whether every thread will process a full tile's worth of elements + // without reading outside the bounds of the input. If this is true, we can + // skip some bounds checks in the final algorithm. + bool all_threads_in_bounds = num_tiles * kTileSize == num_elems; + + // __global__ void full_reduce_kernel() { + // x_in_tiles = threadIdx.x + blockIdx.x * blockDim.x; + // x = x_in_tiles * kTileSize; + // + // partial_result = init_value; + // if (all_threads_in_bounds || x + kTileSize <= num_elems) { + // for (i = 0; i < kTileSize; ++i) { + // partial_result = Reducer(partial_result, input[x + i]); + // } + // } else { + // for (i = 0; i < kTileSize; ++i) { + // if (x + i < num_elems) { + // partial_result = Reducer(partial_result, input[x + i]); + // } + // } + // } + // for (i = warpSize / 2; i > 0; i /= 2) { + // partial_result = Reducer(partial_result, + // __shfl_down(partial_result, i)); + // } + // if (lane_id == 0) { + // AtomicReducer(&output[y], partial_result); + // } + // } + // + // // Choose num_blocks and threads_per_block such that: + // // + // // num_blocks * threads_per_block = + // // RoundUpToNextMultipleOf(Ceil(num_elems / kTileSize), warpSize), + // // + // // and threads_per_block is a multiple of warpSize. + // reduce_kernel<<>>(); + // + auto loop_body_emitter = + [=](const llvm_ir::IrArray::Index& tile_index) -> Status { + llvm::Type* element_ir_type = + llvm_ir::PrimitiveTypeToIrType(input_shape.element_type(), module_); + llvm::Value* partial_reduction_result_address = ir_builder_.CreateAlloca( + element_ir_type, /*ArraySize=*/nullptr, "partial_reduction_result"); + { + TF_ASSIGN_OR_RETURN(llvm::Value * init_ir_value, + init_value_gen(llvm_ir::IrArray::Index({}))); + ir_builder_.CreateStore(init_ir_value, partial_reduction_result_address); + } + + llvm::Value* x_in_tiles = tile_index[0]; + + // Emit an inner for-loop that reduces the elements in the tile. + auto emit_tile_element_loop = [=](bool tile_in_bounds) -> Status { + std::unique_ptr tile_element_loop = + llvm_ir::ForLoop::EmitForLoop("element_id_in_tile", + ir_builder_.getInt64(0), + ir_builder_.getInt64(kTileSize), + ir_builder_.getInt64(1), &ir_builder_); + + // Emit the body of the partial reduction loop. + llvm_ir::SetToFirstInsertPoint(tile_element_loop->GetBodyBasicBlock(), + &ir_builder_); + llvm::Value* x = ir_builder_.CreateNSWAdd( + ir_builder_.CreateNSWMul(x_in_tiles, ir_builder_.getInt64(kTileSize)), + tile_element_loop->GetIndVarValue()); + // Unless we know the tile is entirely in bounds, we have to emit a + // x-in-bounds check before reading from the input. + if (!tile_in_bounds) { + llvm_ir::LlvmIfData if_data = llvm_ir::EmitIfThenElse( + ir_builder_.CreateICmpULT(x, ir_builder_.getInt64(num_elems)), + "x_in_bounds", &ir_builder_); + + // Emit code that reads the input element and accumulates it to + // the partial reduction result. + llvm_ir::SetToFirstInsertPoint(if_data.true_block, &ir_builder_); + } + llvm_ir::IrArray::Index input_index( + /*linear=*/x, input_shape, &ir_builder_); + llvm::Value* input_address = ir_builder_.CreateAlloca(element_ir_type); + TF_ASSIGN_OR_RETURN(llvm::Value * input_ir_value, input_gen(input_index)); + ir_builder_.CreateStore(input_ir_value, input_address); + return (EmitCallToNestedComputation( + *reducer, {partial_reduction_result_address, input_address}, + partial_reduction_result_address)); + }; + + // x_end = kTileSize + x_in_tiles * kTileSize, i.e., the location that's + // immediately beyond the tile. + llvm::Value* x_end = ir_builder_.CreateNSWAdd( + ir_builder_.getInt64(kTileSize), + ir_builder_.CreateNSWMul(x_in_tiles, ir_builder_.getInt64(kTileSize))); + // The tile is entirely in bound if all_threads_in_bounds or + // x_end <= num_elems. + llvm::Value* tile_in_bounds = ir_builder_.CreateOr( + ir_builder_.CreateICmpULE(x_end, ir_builder_.getInt64(num_elems)), + ir_builder_.getInt1(all_threads_in_bounds)); + llvm_ir::LlvmIfData if_tile_in_bounds_data = + llvm_ir::EmitIfThenElse(tile_in_bounds, "tile_in_bounds", &ir_builder_); + llvm_ir::SetToFirstInsertPoint(if_tile_in_bounds_data.true_block, + &ir_builder_); + TF_RETURN_IF_ERROR(emit_tile_element_loop(/*tile_in_bounds=*/true)); + llvm_ir::SetToFirstInsertPoint(if_tile_in_bounds_data.false_block, + &ir_builder_); + TF_RETURN_IF_ERROR(emit_tile_element_loop(/*tile_in_bounds=*/false)); + + // After the if-then-else statement on tile_in_bounds, emit calls to + // shfl_down that accumulate the partial reduction results of all threads + // from the warp. + llvm_ir::SetToFirstInsertPoint(if_tile_in_bounds_data.after_block, + &ir_builder_); + int bit_width = llvm_ir::GetSizeInBits(element_ir_type); + // bitcast cannot be applied to aggregate types (even packed ones), so we + // instead bitcast addresses of load/store to intN* of the same bit-width. + llvm::Type* shuffle_ir_type = element_ir_type->isStructTy() + ? ir_builder_.getIntNTy(bit_width) + : element_ir_type; + for (int shuffle_distance = kWarpSize / 2; shuffle_distance >= 1; + shuffle_distance /= 2) { + llvm::Value* partial_reduction_result = ir_builder_.CreateLoad( + ir_builder_.CreateBitCast(partial_reduction_result_address, + shuffle_ir_type->getPointerTo()), + "partial_reduction_result"); + llvm::Value* result_from_other_lane = ir_builder_.CreateAlloca( + element_ir_type, nullptr, "result_from_other_lane"); + ir_builder_.CreateStore( + EmitShuffleDown(partial_reduction_result, + ir_builder_.getInt32(shuffle_distance), &ir_builder_), + ir_builder_.CreateBitCast(result_from_other_lane, + shuffle_ir_type->getPointerTo())); + TF_RETURN_IF_ERROR(EmitCallToNestedComputation( + *reducer, {partial_reduction_result_address, result_from_other_lane}, + partial_reduction_result_address)); + } + + const HloInstruction* output = + reduce->IsFused() ? reduce->parent()->FusionInstruction() : reduce; + + // Emit an atomic operation that accumulates the partial reduction result of + // lane 0 (which holds the partially accumulated result for its warp) to the + // output element. + llvm::Value* lane_id = ir_builder_.CreateURem( + x_in_tiles, ir_builder_.getInt64(kWarpSize), "lane_id"); + llvm_ir::LlvmIfData if_lane_id_is_zero_data = llvm_ir::EmitIfThenElse( + ir_builder_.CreateICmpEQ(lane_id, ir_builder_.getInt64(0)), + "lane_id_is_zero", &ir_builder_); + llvm_ir::SetToFirstInsertPoint(if_lane_id_is_zero_data.true_block, + &ir_builder_); + llvm::Value* output_address = + GetIrArray(*output, *output) + .EmitArrayElementAddress( + llvm_ir::IrArray::Index(/*linear=*/ir_builder_.getInt64(0), + output->shape(), &ir_builder_), + &ir_builder_, "output_element_address"); + return EmitAtomicOperationForNestedComputation( + *reducer, output_address, partial_reduction_result_address); + }; + + // Emit a parallel loop that iterates through all input tiles, one per thread. + Shape tiled_input_shape = ShapeUtil::MakeShapeWithLayout( + reduce->shape().element_type(), {num_tiles}, {0}); + LaunchDimensions launch_dimensions = CalculateLaunchDimensions( + tiled_input_shape, ir_emitter_context_->device_description()); + CHECK(LastThunk()->kind() == Thunk::Kind::kSequential); + UpdateLaunchDimensions( + launch_dimensions, + static_cast(LastThunk())->thunks().back().get(), + ir_emitter_context_->llvm_module()); + return ParallelLoopEmitter(loop_body_emitter, tiled_input_shape, + launch_dimensions, &ir_builder_) + .EmitLoop(IrName(reduce)); +} + Status IrEmitterUnnested::EmitColumnReduction( int64 height, int64 width, HloInstruction* reduce, const Shape& input_shape, const llvm_ir::ElementGenerator& input_gen, @@ -1017,7 +1318,7 @@ Status IrEmitterUnnested::EmitRowReduction( // // Three optimizations are performed. // - // 1. To coalesc global memory accesses, dilate the tile with a factor of 32 + // 1. To coalesce global memory accesses, dilate the tile with a factor of 32 // (i.e. the warp size). For example, suppose the width is 8x32=256. Instead // of making each tile consecutive, we let make tile 0 column // [0,32,64,...,224], tile 1 column [1,33,65,...,225], and so on. This ensures @@ -1306,14 +1607,11 @@ Status IrEmitterUnnested::EmitReductionToVector( // the dimensions to keep are contiguous, by prerequisite of // `EmitReductionToVector`, we only need to check whether the minormost // dimension of the input is to keep. - // - // If the output is scalar, we could emit either a row or a column reduction. - // Some tests have shown scalar reduction is no more efficient as row - // reduction, and is simpler to emit as column reduction, so we emit a column - // reduction in this case. - if (input_dims_to_keep.empty() || - input_dims_to_keep.front() == - LayoutUtil::Minor(input_shape.layout(), 0)) { + if (input_dims_to_keep.empty()) { + return EmitReductionToScalar(reduce, input_shape, input_gen, init_value_gen, + reducer); + } else if (input_dims_to_keep.front() == + LayoutUtil::Minor(input_shape.layout(), 0)) { // Column reduction. Treat the result of "input" as a matrix whose width // is the most minor dimension and height the product of other dimensions, // and treat "reduce" as a column reduction of the input matrix. @@ -1396,24 +1694,24 @@ Status IrEmitterUnnested::HandleReduce(HloInstruction* reduce) { } Status IrEmitterUnnested::HandleTuple(HloInstruction* tuple) { - tensorflow::gtl::ArraySlice operands(tuple->operands()); - bool all_tuple_elements_have_buffer = std::all_of( - operands.begin(), operands.end(), [this](HloInstruction* tuple_element) { + bool all_tuple_elements_have_buffer = + c_all_of(tuple->operands(), [&](HloInstruction* tuple_element) { return ir_emitter_context_->buffer_assignment().HasTopLevelAllocation( tuple_element); }); - // Tuples (especially output tuples) can take too many tuple elements, - // causing the kernel emitted exceeds the parameter space limit - // (b/31336476). As an optimization, if all tuple elements have a buffer, we - // collect their buffer addresses in a host array, and then copy that array - // to the tuple's buffer. + // Tuples (especially tuples that are the final result of a computation) can + // be so huge that if we were to emit a kernel that took each tuple element as + // a parameter, we would exceed the max allowable number of parameters to a + // GPU kernel, b/31336476. As an optimization, if all tuple elements have a + // buffer, we collect their buffer addresses in a host array, and then copy + // that array to the tuple's buffer. // // Some tuple elements (e.g. const or bitcast of const) might not have a - // buffer -- their contents are stored in code. In that case, we fall back - // to emitting kernels which have access to their buffer addresses in code. + // buffer -- their contents are stored in code. In that case, we fall back to + // emitting kernels which have access to their buffer addresses in code. if (all_tuple_elements_have_buffer) { std::vector tuple_element_buffers; - for (const HloInstruction* tuple_element : operands) { + for (const HloInstruction* tuple_element : tuple->operands()) { tuple_element_buffers.push_back(GetAllocationSlice(*tuple_element)); } thunk_sequence_->emplace_back(MakeUnique( @@ -1455,8 +1753,7 @@ Status IrEmitterUnnested::HandleSelectAndScatter( // TODO(b/31410564): Implement dilation rate for select-and-scatter. if (window_util::HasDilation(window)) { return Unimplemented( - "Dilation for select-and-scatter not implemented on GPU. " - "See b/31410564."); + "Dilation for SelectAndScatter not implemented on GPU."); } // kSelectAndScatter is implemented as two kernel launches: the first launch @@ -1665,62 +1962,202 @@ Status IrEmitterUnnested::HandleInfeed(HloInstruction* infeed) { return Status::OK(); } -llvm::Function* IrEmitterUnnested::EmitBasePointersForHloAndItsOperands( - const HloInstruction& hlo, std::vector* io_hlos) { - const BufferAssignment& buffer_assignment = - ir_emitter_context_->buffer_assignment(); - // GetTupleElement instructions are implemented by emitting IR that indexes - // and loads the target tuple element pointer from its operand (possibly - // recursively). For this reason, GetTupleElement instructions are associated - // with their operand buffer in 'io_hlos' and 'non_io_hlos' below. - std::vector non_io_hlos; - for (const HloInstruction* operand : hlo.operands()) { - const HloInstruction* to_lookup = operand->LatestNonGteAncestor(); - if (buffer_assignment.HasTopLevelAllocation(to_lookup) && - buffer_assignment.GetUniqueTopLevelSlice(to_lookup) - .ConsumeValueOrDie() - .allocation() - ->IsInputOrOutput()) { - io_hlos->push_back(operand); - } else { - non_io_hlos.push_back(operand); +// Figures out how to access the buffers for all subshapes of hlo's operands and +// for hlo itself (i.e. all the buffers produced by HLO). +// +// Returns a map keyed on the pair {HloInstruction, ShapeIndex}. The value for +// this key is a pair {Slice, ShapeIndex}, where the slice tells you the root +// buffer to look in, and the ShapeIndex describes how to dereference starting +// at that buffer to get to the buffer in question. +// +// For example, if {hlo, {1}} is mapped to {slice, {3, 4}}, then the buffer for +// hlo at ShapeIndex {1} (i.e. the buffer for the second tuple element of hlo) +// is found at slice[3][4]. That is, slice is a void***, which we dereference +// twice -- first at index 3, and then at index 4 -- to get the address of our +// buffer. +// +// This function conservatively assumes that we'll touch all sub-buffers of +// every operand and of the output. +static std::map, + std::pair> +GetHloBufferSlices(const HloInstruction* hlo, + const BufferAssignment& buffer_assn) { + std::map, + std::pair> + slices; + + // Tries to find a slice plus an array of indices i1, ..., iN such that the + // sub-buffer for instr at index can be found at slice[i1]...[iN]. + auto find_slice_for = [&](const HloInstruction* instr, + const ShapeIndex& index) + -> optional> { + // Simple, common case: Is the buffer for instr known at runtime? If so, + // we're done. + auto slice = GetKnownAtRuntimeSlice(instr, index, buffer_assn); + if (slice.has_value()) { + return {{*slice, ShapeIndex()}}; } - } - CHECK_NE(HloOpcode::kGetTupleElement, hlo.opcode()); - if (buffer_assignment.HasTopLevelAllocation(&hlo) && - buffer_assignment.GetUniqueTopLevelSlice(&hlo) - .ConsumeValueOrDie() - .allocation() - ->IsInputOrOutput()) { - io_hlos->push_back(&hlo); - } else { - non_io_hlos.push_back(&hlo); + // If we don't know the buffer for instr at index, see if we know the buffer + // for instr at index without its last element. If so, we can dynamically + // find the buffer for instr by dereferencing a pointer in that buffer. + // Continue looking this way until we run out of elements in 'index'. + ShapeIndex new_index = index; + ShapeIndex gte_indices; + while (!new_index.empty()) { + gte_indices.push_front(new_index.back()); + new_index.pop_back(); + auto slice = GetKnownAtRuntimeSlice(instr, new_index, buffer_assn); + if (slice.has_value()) { + return {{*slice, gte_indices}}; + } + } + + // If *that* didn't work, check whether instr is a GTE instruction. If it + // is, see if we can get a buffer for its parent, and continue walking up + // parents until we find a defined buffer or we hit something that's not a + // GTE. + const HloInstruction* parent = instr; + while (parent->opcode() == HloOpcode::kGetTupleElement) { + gte_indices.push_front(parent->tuple_index()); + parent = parent->operand(0); + + auto slice = GetKnownAtRuntimeSlice(parent, {}, buffer_assn); + if (slice.has_value()) { + return {{*slice, gte_indices}}; + } + } + + return nullopt; + }; + + // Adds entries for all subshapes of instr to `slices`. + auto add_slices_for = [&](const HloInstruction* instr) { + // GPU constants don't have buffers; don't bother looking for one. + if (instr->IsConstant()) { + return; + } + + ShapeUtil::ForEachSubshape( + instr->shape(), [&](const Shape& /*shape*/, const ShapeIndex& index) { + if (slices.count({instr, index})) { + // HLOs can have duplicate operands; don't bother redoing work. + return; + } + auto maybe_slice = find_slice_for(instr, index); + if (maybe_slice.has_value()) { + slices[{instr, index}] = *maybe_slice; + } else { + VLOG(1) << "Couldn't find buffer for " << instr->ToString() + << " at index " << index.ToString(); + } + }); + }; + + add_slices_for(hlo); + for (const HloInstruction* operand : hlo->operands()) { + // Conservatively assume we'll need the buffers for all subshapes of the + // operand. + add_slices_for(operand); } - llvm::Function* kernel = BuildKernelPrototype(hlo, *io_hlos); - // bindings_ is reused because the bindings of kConstant to their underlying - // llvm::Constant can be shared for all HLOs in this computation. - bindings_.EmitBasePointersForHlos(*io_hlos, non_io_hlos); - return kernel; + return slices; } std::unique_ptr IrEmitterUnnested::BuildKernelThunk( const HloInstruction* inst) { - std::vector io_hlos; - llvm::Function* kernel = - EmitBasePointersForHloAndItsOperands(*inst, &io_hlos); + const BufferAssignment& buffer_assn = + ir_emitter_context_->buffer_assignment(); + + std::map, + std::pair> + hlo_slices = GetHloBufferSlices(inst, buffer_assn); + + // Figure out which buffer allocations need to be passed as arguments to our + // kernel. This is simply all of the allocations referenced in hlo_slices, + // plus the XLA temp buffer (if we have it). We always include the temp + // buffer because even if the kernel itself doesn't use it, a nested + // subcomputation within the kernel (e.g. a kMap's computation) might. + std::unordered_set buffers_needed; + for (const auto& kv : hlo_slices) { + buffers_needed.insert(kv.second.first.allocation()); + } + tensorflow::gtl::optional temp_buffer; + for (const BufferAllocation& alloc : buffer_assn.Allocations()) { + if (alloc.IsPreallocatedTempBuffer()) { + if (!temp_buffer.has_value()) { + temp_buffer = &alloc; + } else { + LOG(FATAL) << "Multiple temp buffers found, but only one is allowed!"; + } + } + } + if (temp_buffer.has_value()) { + buffers_needed.insert(*temp_buffer); + } + + // We'll pass a pointer to each of the elements of `buffers` to our kernel, in + // this order. + std::vector buffers(buffers_needed.begin(), + buffers_needed.end()); + std::sort(buffers.begin(), buffers.end(), + [](const BufferAllocation* a, const BufferAllocation* b) { + return a->index() < b->index(); + }); + + llvm::Function* kernel = BuildKernelPrototype(*inst, buffers); + + // Build a map from a BufferAllocation to the corresponding argument in our + // kernel. + std::unordered_map kernel_args; + { + auto arg_it = kernel->arg_begin(); + auto buffers_it = buffers.begin(); + for (; arg_it != kernel->arg_end(); ++arg_it, ++buffers_it) { + kernel_args[*buffers_it] = arg_it; + } + } + + // For each buffer our kernel might want to touch, bind it to a value derived + // from our kernel args. + for (const auto& kv : hlo_slices) { + const HloInstruction* instr = kv.first.first; + const ShapeIndex& index = kv.first.second; + const BufferAllocation::Slice& slice = kv.second.first; + const ShapeIndex& gte_index = kv.second.second; + + VLOG(3) << "Buffer for " << instr->ToString() << " at " << index.ToString() + << " is found in slice " << slice.ToString() << " at GTE index " + << gte_index.ToString(); + + llvm::Value* loc = + ir_builder_.CreateInBoundsGEP(kernel_args.at(slice.allocation()), + {ir_builder_.getInt64(slice.offset())}); + + // If gte_index is nonempty, we have to dereference `loc` to get to the + // value we're ultimately interested in. + llvm::Type* int8_double_pointer = + llvm::PointerType::get(ir_builder_.getInt8PtrTy(), /*AddressSpace=*/0); + for (int64 idx : gte_index) { + loc = ir_builder_.CreateBitCast(loc, int8_double_pointer); + loc = ir_builder_.CreateLoad( + ir_builder_.CreateInBoundsGEP(loc, {ir_builder_.getInt64(idx)})); + } + + bindings_.BindHloToIrValue(*instr, loc, index); + } - // Compute the input buffer indices. - std::vector io_buffers; - io_buffers.reserve(io_hlos.size()); - for (const HloInstruction* io_hlo : io_hlos) { - io_buffers.push_back(GetAllocationSlice(*io_hlo->LatestNonGteAncestor())); + // Bind the temp buffer so that nested subcomputations can find it if they + // need. + if (temp_buffer.has_value()) { + bindings_.SetTempBufferBase(kernel_args.at(*temp_buffer)); + } else { + bindings_.SetTempBufferBase( + llvm::ConstantPointerNull::get(ir_builder_.getInt8PtrTy())); } - // Create a KernelThunk that launches the kernel that implements "inst". - return MakeUnique(io_buffers, - llvm_ir::AsString(kernel->getName()), inst); + return MakeUnique(buffers, llvm_ir::AsString(kernel->getName()), + inst); } std::unique_ptr IrEmitterUnnested::BuildHostToDeviceCopyThunk( @@ -1728,7 +2165,7 @@ std::unique_ptr IrEmitterUnnested::BuildHostToDeviceCopyThunk( const HloInstruction* operand = inst->operand(0); CHECK_EQ(HloOpcode::kConstant, operand->opcode()); return MakeUnique( - /*source_address=*/operand->literal().InternalData(), + /*source_address=*/operand->literal().untyped_data(), /*destination_buffer=*/GetAllocationSlice(*inst), /*mem_size=*/ llvm_ir::ByteSizeOf(operand->shape(), @@ -1809,50 +2246,14 @@ std::unique_ptr IrEmitterUnnested::BuildGemmThunk( LOG(FATAL) << "Cannot build a GemmThunk for " << inst->ToString(); } -std::unique_ptr IrEmitterUnnested::BuildConvolutionThunk( +std::unique_ptr IrEmitterUnnested::BuildFftThunk( const HloInstruction* inst) { - const HloInstruction* lhs = inst->operand(0); - const HloInstruction* rhs = inst->operand(1); - if (inst->opcode() == HloOpcode::kConvolution) { - // Forward covolution. - return MakeUnique( - ConvolutionThunk::ConvolutionKind::kForward, - /*input_buffer=*/GetAllocationSlice(*lhs), - /*filter_buffer=*/GetAllocationSlice(*rhs), - /*output_buffer=*/GetAllocationSlice(*inst), - /*input_shape=*/lhs->shape(), - /*filter_shape=*/rhs->shape(), - /*output_shape=*/inst->shape(), inst->window(), - inst->convolution_dimension_numbers(), inst); - } - - // Backward filter convolution, which takes the input (activations) and the - // gradients, and computes the filter. - CHECK_EQ(HloOpcode::kFusion, inst->opcode()); - switch (inst->fusion_kind()) { - case HloInstruction::FusionKind::kConvBackwardFilter: - return MakeUnique( - ConvolutionThunk::ConvolutionKind::kBackwardFilter, - /*input_buffer=*/GetAllocationSlice(*lhs), - /*filter_buffer=*/GetAllocationSlice(*inst), - /*output_buffer=*/GetAllocationSlice(*rhs), - /*input_shape=*/lhs->shape(), - /*filter_shape=*/inst->shape(), - /*output_shape=*/rhs->shape(), inst->window(), - inst->convolution_dimension_numbers(), inst); - case HloInstruction::FusionKind::kConvBackwardInput: - return MakeUnique( - ConvolutionThunk::ConvolutionKind::kBackwardInput, - /*input_buffer=*/GetAllocationSlice(*inst), - /*filter_buffer=*/GetAllocationSlice(*rhs), - /*output_buffer=*/GetAllocationSlice(*lhs), - /*input_shape=*/inst->shape(), - /*filter_shape=*/rhs->shape(), - /*output_shape=*/lhs->shape(), inst->window(), - inst->convolution_dimension_numbers(), inst); - default: - LOG(FATAL) << "Not a convolution-fusion"; - } + const HloInstruction* operand = inst->operand(0); + return MakeUnique(inst->fft_type(), inst->fft_length(), + /*input_buffer=*/GetAllocationSlice(*operand), + /*output_buffer=*/GetAllocationSlice(*inst), + /*input_shape=*/operand->shape(), + /*output_shape=*/inst->shape(), inst); } Status IrEmitterUnnested::EmitInitializer(const HloInstruction* hlo, @@ -1890,6 +2291,24 @@ Status IrEmitterUnnested::EmitInitializer(const HloInstruction* hlo, namespace { +// Checks that the buffers corresponding to the given two HLOs share the same +// allocation. +Status CheckHloBuffersShareAllocation( + const HloInstruction* a, const HloInstruction* b, const ShapeIndex& index, + const BufferAssignment& buffer_assignment) { + const BufferAllocation::Slice slice_a = + buffer_assignment.GetUniqueSlice(a, index).ConsumeValueOrDie(); + const BufferAllocation::Slice slice_b = + buffer_assignment.GetUniqueSlice(b, index).ConsumeValueOrDie(); + if (slice_a != slice_b) { + return InternalError( + "instruction %s %s does not share allocation with instruction %s %s", + a->ToString().c_str(), slice_a.ToString().c_str(), + b->ToString().c_str(), slice_b.ToString().c_str()); + } + return Status::OK(); +} + // Checks that all buffers used during while loop iteration share the same // buffer allocation. This includes buffers for while result, while init // operand, condition parameter, body parameter and body result. @@ -1899,37 +2318,65 @@ Status CheckWhileBuffersShareAllocation( const BufferAssignment& buffer_assignment) { return ShapeUtil::ForEachSubshapeWithStatus( xla_while->shape(), - [&buffer_assignment, &xla_while](const Shape& /*subshape*/, - const ShapeIndex& index) -> Status { - auto check = [&buffer_assignment](const HloInstruction* a, - const HloInstruction* b, - const ShapeIndex& index) -> Status { - const BufferAllocation::Slice slice_a = - buffer_assignment.GetUniqueSlice(a, index).ConsumeValueOrDie(); - const BufferAllocation::Slice slice_b = - buffer_assignment.GetUniqueSlice(b, index).ConsumeValueOrDie(); - if (slice_a != slice_b) { - return InternalError( - "instruction %s %s does not share allocation with " - "instruction %s %s", - a->ToString().c_str(), slice_a.ToString().c_str(), - b->ToString().c_str(), slice_b.ToString().c_str()); - } - return Status::OK(); - }; + [&](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { const HloInstruction* condition_parameter = xla_while->while_condition()->parameter_instruction(0); const HloComputation* body = xla_while->while_body(); const HloInstruction* body_parameter = body->parameter_instruction(0); const HloInstruction* body_result = body->root_instruction(); - TF_RETURN_IF_ERROR(check(xla_while, xla_while->operand(0), index)); - TF_RETURN_IF_ERROR(check(xla_while, condition_parameter, index)); - TF_RETURN_IF_ERROR(check(xla_while, body_parameter, index)); - TF_RETURN_IF_ERROR(check(xla_while, body_result, index)); + TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( + xla_while, xla_while->operand(0), index, buffer_assignment)); + TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( + xla_while, condition_parameter, index, buffer_assignment)); + TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( + xla_while, body_parameter, index, buffer_assignment)); + TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( + xla_while, body_result, index, buffer_assignment)); return Status::OK(); }); } +// Checks that the buffers used in a conditional instruction are shared with the +// operands and result as follows: +// * The result buffer of the conditional should share the allocation with the +// result buffers of the true and false computations. +// * The buffer of operand 1 should share the allocation with the buffer of +// the parameter 0 instruction of the true computation. +// * The buffer of operand 2 should share the allocation with the buffer of +// the parameter 0 instruction of the false computation. +Status CheckConditionalBuffersShareAllocation( + const HloInstruction* conditional, + const BufferAssignment& buffer_assignment) { + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + conditional->shape(), + [&](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { + TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( + conditional, conditional->true_computation()->root_instruction(), + index, buffer_assignment)); + TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( + conditional, conditional->false_computation()->root_instruction(), + index, buffer_assignment)); + return Status::OK(); + })); + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + conditional->operand(1)->shape(), + [&](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { + return CheckHloBuffersShareAllocation( + conditional->operand(1), + conditional->true_computation()->parameter_instruction(0), index, + buffer_assignment); + })); + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + conditional->operand(2)->shape(), + [&](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { + return CheckHloBuffersShareAllocation( + conditional->operand(2), + conditional->false_computation()->parameter_instruction(0), index, + buffer_assignment); + })); + return Status::OK(); +} + } // namespace std::unique_ptr IrEmitterUnnested::BuildWhileThunk( @@ -1972,9 +2419,36 @@ std::unique_ptr IrEmitterUnnested::BuildForThunk( ir_emitter_body.ConsumeThunkSequence(), hlo); } +std::unique_ptr IrEmitterUnnested::BuildConditionalThunk( + const HloInstruction* hlo) { + // Check that the buffers used in conditional are shared with the operands and + // result appropriately. + TF_CHECK_OK(CheckConditionalBuffersShareAllocation( + hlo, ir_emitter_context_->buffer_assignment())); + + HloComputation* true_computation = hlo->true_computation(); + IrEmitterUnnested ir_emitter_true(hlo_module_config_, true_computation, + ir_emitter_context_); + TF_CHECK_OK(true_computation->root_instruction()->Accept(&ir_emitter_true)); + + HloComputation* false_computation = hlo->false_computation(); + IrEmitterUnnested ir_emitter_false(hlo_module_config_, false_computation, + ir_emitter_context_); + TF_CHECK_OK(false_computation->root_instruction()->Accept(&ir_emitter_false)); + + return MakeUnique( + GetAllocationSlice(*hlo->operand(0)), + GetAllocationSlice(*hlo->operand(1)), + GetAllocationSlice(*hlo->operand(2)), + std::move(*ir_emitter_true.ConsumeThunkSequence()), + std::move(*ir_emitter_false.ConsumeThunkSequence()), hlo); +} + Status IrEmitterUnnested::EmitTargetElementLoopInThunk( const HloInstruction& hlo, const llvm_ir::ElementGenerator& element_generator, KernelThunk* thunk) { + VLOG(3) << bindings_.ToString(); + const Shape& element_shape = hlo.IsMultiOutputFusion() ? ShapeUtil::GetSubshape(hlo.shape(), {0}) : hlo.shape(); diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h new file mode 100644 index 0000000000000000000000000000000000000000..688760efbd2c725a4bf48e45eb6f2734b63d25e1 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h @@ -0,0 +1,205 @@ +/* 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_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_UNNESTED_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_UNNESTED_H_ + +#include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" +#include "tensorflow/compiler/xla/service/gpu/thunk.h" + +namespace xla { +namespace gpu { + +// Emits LLVM IR for an "unnested computation". +// +// An unnested computation is an HloComputation which you run by executing one +// or more kernels for each HloInstruction it contains. Examples of unnested +// computations: +// +// - An HloModule's root computation, +// - The body of an HLO while loop, +// - The true/false computation of an HLO conditional. +// +// Note the opportunity for confusion -- the while loop's computation is nested +// within the root computation, but it's emitted using IrEmitterUnnested! Don't +// think about it too hard. +// +// Examples of things that are not unnested computations: +// +// - The reducer of a kReduce HLO. This is emited using IrEmitterNested. +// - The body of a fusion node. IrEmitterUnenested emits the relevant code +// within a kernel function using FusedIrEmitter. (FusedIrEmitter is not +// really an IrEmitter, but is more an "IR generator generator".) +// +class IrEmitterUnnested : public IrEmitter { + public: + IrEmitterUnnested(const HloModuleConfig& hlo_module_config, + const HloComputation* hlo_computation, + IrEmitterContext* ir_emitter_context); + IrEmitterUnnested(const IrEmitterUnnested&) = delete; + IrEmitterUnnested& operator=(const IrEmitterUnnested&) = delete; + + // Transfers the ownship of thunk_sequence_ out. + std::unique_ptr ConsumeThunkSequence() { + return std::move(thunk_sequence_); + } + + Status DefaultAction(HloInstruction* hlo) override; + + // IrEmitterUnnested handles the following instructions differently from + // IrEmitter. + Status HandleCopy(HloInstruction* copy) override; + Status HandleConditional(HloInstruction* conditional) override; + Status HandleConvolution(HloInstruction* convolution) override; + Status HandleCustomCall(HloInstruction* custom_call) override; + Status HandleDot(HloInstruction* dot) override; + Status HandleFft(HloInstruction* fft) override; + Status HandleFusion(HloInstruction* fusion) override; + Status HandleGetTupleElement(HloInstruction* get_tuple_element) override; + Status HandleReduce(HloInstruction* reduce) override; + Status HandleSelectAndScatter(HloInstruction* instruction) override; + Status HandleTuple(HloInstruction* tuple) override; + Status HandleWhile(HloInstruction* xla_while) override; + Status HandleInfeed(HloInstruction* xla_infeed) override; + Status HandleRng(HloInstruction* random) override; + Status HandleSelect(HloInstruction* select) override; + + Status EmitTargetElementLoop( + const HloInstruction& hlo, + const llvm_ir::ElementGenerator& body_emitter) override; + + // Same as `EmitTargetElementLoop`, but in given `thunk` rather than + // `LastThunk()`. + Status EmitTargetElementLoopInThunk( + const HloInstruction& hlo, const llvm_ir::ElementGenerator& body_emitter, + KernelThunk* thunk); + + private: + // Builds the appropriate thunk for the instruction hlo and returns the owning + // pointer to it. The caller needs to make sure `inst` outlives the lifetime + // of the returned Thunk object. + std::unique_ptr BuildThunk(const HloInstruction* hlo); + + // Builds the prototype of the IR kernel for `inst` and adds it to the module. + // This kernel takes as arguments pointers to the given buffer allocations. + llvm::Function* BuildKernelPrototype( + const HloInstruction& inst, + tensorflow::gtl::ArraySlice args); + + // EmitColumnReduction and EmitRowReduction emit code for column and row + // reduction of a matrix and/or 3D tensor. Row and column reduction have + // different memory access pattern, so for performance their implementations + // are significantly different. + // + // Emits code that reduces a matrix of shape [height x width] to a vector of + // [width]. Other parameters have the same meaning as those of + // `EmitReductionToVector`. Note that input shape might not be + // [height x width], but can be bitcast to [height x weight] with "height" + // being the major dimension. + Status EmitColumnReduction(int64 height, int64 width, HloInstruction* reduce, + const Shape& input_shape, + const llvm_ir::ElementGenerator& input_gen, + const llvm_ir::ElementGenerator& init_value_gen, + HloComputation* reducer); + + // Emits code that reduces a 3D tensor of shape [depth x height x width] to a + // vector of shape [height]. Other parameters have the same meaning as those + // of `EmitReductionToVector`. Note that input shape might not be + // [depth x height x width], but can be bitcast to [depth x height x weight] + // with "depth" being the most major dimension. + Status EmitRowReduction(int64 depth, int64 height, int64 width, + HloInstruction* reduce, const Shape& input_shape, + const llvm_ir::ElementGenerator& input_gen, + const llvm_ir::ElementGenerator& init_value_gen, + HloComputation* reducer); + + // Emits code that reduces a tensor of arbitrary rank to a scalar. + Status EmitReductionToScalar(HloInstruction* reduce, const Shape& input_shape, + const llvm_ir::ElementGenerator& input_gen, + const llvm_ir::ElementGenerator& init_value_gen, + HloComputation* reducer); + + // Figures out whether `reduce` is a row or column reduction, and which + // dimensions to reduce, and calls either `EmitRowReduction` or + // `EmitColumnReduction` as appropriate. `input_shape` is the shape of the + // input array, which is the operand of the Reduce instruction if unfused or + // of the Fusion instruction if fused. `input_gen` and `init_value_gen` + // generate elements of the input and the initial value. Other parameters mean + // the same as for `HandleReduce`. + // + // Prerequisite: `IsReductionToVector(*reduce)` + Status EmitReductionToVector( + HloInstruction* reduce, const Shape& input_shape, + const llvm_ir::ElementGenerator& input_gen, + const llvm_ir::ElementGenerator& init_value_gen, + tensorflow::gtl::ArraySlice dimensions_to_reduce, + HloComputation* reducer); + + // Emits code to initialize buffer of `inst` in given `thunk`. + Status EmitInitializer(const HloInstruction* inst, KernelThunk* thunk); + + // Returns a KernelThunk that invokes the kernel emitted for `inst`. The + // caller needs to make sure `inst` outlives the lifetime of the returned + // Thunk object. + std::unique_ptr BuildKernelThunk(const HloInstruction* inst); + + // Returns a FftThunk that calls cuFFT to implement `inst`. + std::unique_ptr BuildFftThunk(const HloInstruction* inst); + + // Returns a GemmThunk that calls gemm to implement `inst`. The caller needs + // to make sure `inst` outlives the lifetime of the returned Thunk object. + std::unique_ptr BuildGemmThunk(const HloInstruction* inst); + + // Returns a thunk that calls host-to-device cuMemcpy to implement `inst`. + std::unique_ptr BuildHostToDeviceCopyThunk(const HloInstruction* inst); + + // Returns a thunk that calls device-to-device cuMemcpy to implement `inst`. + std::unique_ptr BuildDeviceToDeviceCopyThunk( + const HloInstruction* inst); + + // Returns an InfeedThunk that performs device-to-device memcpy to implement + // `inst`. + std::unique_ptr BuildInfeedThunk(const HloInstruction* inst); + + // Returns a WhileThunk that invokes thunk sequences for 'condition' and + // 'body' sub-computations of while instruction 'hlo'. + std::unique_ptr BuildWhileThunk(const HloInstruction* hlo); + + // Returns a ForThunk which executes 'loop_limit' invocations of a thunk + // sequence from the 'body' sub-computation of the while instruction 'hlo'. + std::unique_ptr BuildForThunk(const HloInstruction* hlo, + const int64 loop_limit); + + // Returns a ConditionalThunk that executes the thunk sequence for + // 'true_computation' or 'false_computation' depending on the value of the + // predicate in the given conditional instruction. + std::unique_ptr BuildConditionalThunk(const HloInstruction* hlo); + + Status Postprocess(HloInstruction* hlo) override; + + // Returns the last generated thunk. + Thunk* LastThunk() const { return thunk_sequence_->back().get(); } + + // The thunk sequence this IrEmitter generates for the input computation. + std::unique_ptr thunk_sequence_; + + // The HloComputation that this IrEmitter emits code for. + const HloComputation* hlo_computation_; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_UNNESTED_H_ diff --git a/tensorflow/compiler/xla/service/gpu/kernel_thunk.cc b/tensorflow/compiler/xla/service/gpu/kernel_thunk.cc index 96606993696354f36e143b3b994bbe6afb902df3..c20a781a33fe89af4740ed31dd5bfb1a64473057 100644 --- a/tensorflow/compiler/xla/service/gpu/kernel_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/kernel_thunk.cc @@ -29,10 +29,10 @@ namespace xla { namespace gpu { KernelThunk::KernelThunk( - tensorflow::gtl::ArraySlice io_buffers, + tensorflow::gtl::ArraySlice args, const string& kernel_name, const HloInstruction* hlo_instruction) : Thunk(Kind::kKernel, hlo_instruction), - io_buffers_(io_buffers.begin(), io_buffers.end()), + args_(args.begin(), args.end()), kernel_name_(kernel_name) {} tensorflow::Status KernelThunk::Initialize(const GpuExecutable& executable) { @@ -42,7 +42,7 @@ tensorflow::Status KernelThunk::Initialize(const GpuExecutable& executable) { return tensorflow::Status::OK(); } - loader_spec_.reset(new se::MultiKernelLoaderSpec(io_buffers_.size() + 1)); + loader_spec_.reset(new se::MultiKernelLoaderSpec(args_.size())); tensorflow::StringPiece ptx = executable.ptx(); // Convert tensorflow::StringPiece to se::port::StringPiece because // StreamExecutor uses the latter. @@ -81,15 +81,16 @@ tensorflow::Status KernelThunk::ExecuteOnStream( kernel = &it->second; } + VLOG(3) << "Launching " << kernel->name(); // Launch the kernel with potentially multiple blocks and threads. static constexpr int kKernelArgsLimit = 1024; auto kernel_args = MakeUnique>(); - for (const BufferAllocation::Slice io_buffer : io_buffers_) { - kernel_args->add_device_memory_argument( - buffer_allocations.GetDeviceAddress(io_buffer)); + for (const BufferAllocation* arg : args_) { + const auto& buf = buffer_allocations.GetDeviceAddress(arg->index()); + kernel_args->add_device_memory_argument(buf); + VLOG(3) << " Arg: alloc #" << arg->index() << ": " << buf.opaque() << " (" + << buf.size() << "B)"; } - kernel_args->add_device_memory_argument( - buffer_allocations.GetTempBufferBase()); if (!stream->parent()->Launch( stream, se::ThreadDim(launch_dimensions.threads_per_block()), se::BlockDim(launch_dimensions.block_count()), *kernel, diff --git a/tensorflow/compiler/xla/service/gpu/kernel_thunk.h b/tensorflow/compiler/xla/service/gpu/kernel_thunk.h index 350b5aaf360b0dad7f7b04d73f4c32bad55d3ce9..9ae455e2fcc253a7a08ff95764721048a16b0bf7 100644 --- a/tensorflow/compiler/xla/service/gpu/kernel_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/kernel_thunk.h @@ -46,7 +46,7 @@ class KernelThunk : public Thunk { // Constructs a thunk for the given kernel. // // `hlo_instruction` is as in Thunk. Other arguments are as the class members. - KernelThunk(tensorflow::gtl::ArraySlice io_buffers, + KernelThunk(tensorflow::gtl::ArraySlice args, const string& kernel_name, const HloInstruction* hlo_instruction); KernelThunk(const KernelThunk&) = delete; KernelThunk& operator=(const KernelThunk&) = delete; @@ -63,8 +63,8 @@ class KernelThunk : public Thunk { perftools::gputools::Stream* stream) override; private: - // The indices of the input/output buffers. - const std::vector io_buffers_; + // Buffers passed to the kernel as arguments. + const std::vector args_; // Entry kernel name for the computation. const string kernel_name_; diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc index 059943d48cd34b0ac487b91c3f3079ee3f761229..defd281d74bd38f7da3f268e0f55970fc1af8263 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc @@ -252,7 +252,7 @@ void EmitBitcodeToFile(const Module& module, tensorflow::StringPiece filename) { LOG(FATAL) << "opening bitcode file for writing: " << error_code.message(); } - llvm::WriteBitcodeToFile(&module, outfile.os()); + llvm::WriteBitcodeToFile(module, outfile.os()); outfile.keep(); } @@ -440,7 +440,7 @@ StatusOr CompileModuleToPtx(llvm::Module* module, // One-time module initializer. // Must be called only once -- DO NOT CALL DIRECTLY. -void GPUBackendInit() { +void GPUBackendInit(const HloModuleConfig& hlo_module_config) { // Feed all customized flags here, so we can override them with llvm_cl_opts // without redeploy the compiler for development purpose. @@ -466,6 +466,8 @@ void GPUBackendInit() { // between those loads. FeedLLVMWithFlags({"-memdep-block-scan-limit=500"}); + llvm_ir::InitializeLLVMCommandLineOptions(hlo_module_config); + // Initialize the NVPTX target; it's the only target we link with, so call its // specific initialization functions instead of the catch-all InitializeAll*. LLVMInitializeNVPTXTarget(); @@ -485,7 +487,7 @@ StatusOr CompileToPtx(llvm::Module* module, const HloModuleConfig& hlo_module_config, const string& libdevice_dir_path) { static std::once_flag backend_init_flag; - std::call_once(backend_init_flag, GPUBackendInit); + std::call_once(backend_init_flag, GPUBackendInit, hlo_module_config); string ptx; { diff --git a/tensorflow/compiler/xla/service/gpu/pad_insertion.cc b/tensorflow/compiler/xla/service/gpu/pad_insertion.cc index c29fee0879c02021fdc23ac0e02ab398cf40f99e..25846dc6cd4633c7becb6e62d6bc9585348a6eac 100644 --- a/tensorflow/compiler/xla/service/gpu/pad_insertion.cc +++ b/tensorflow/compiler/xla/service/gpu/pad_insertion.cc @@ -27,8 +27,8 @@ namespace gpu { namespace { bool IsForwardConvolutionCanonical(const HloInstruction& conv) { - CHECK_EQ(HloOpcode::kConvolution, conv.opcode()); - return window_util::HasEvenPadding(conv.window()) && + CHECK_EQ(conv.custom_call_target(), kCudnnConvForwardCallTarget); + return window_util::HasSymmetricPadding(conv.window()) && !window_util::HasNegativePadding(conv.window()) && !window_util::HasDilation(conv.window()); } @@ -43,10 +43,16 @@ HloInstruction* MaybePaddedAndSlicedInput( const Window& conv_window, const ConvolutionDimensionNumbers& conv_dnums, HloInstruction* input) { HloComputation* computation = input->parent(); - if (!window_util::HasEvenPadding(conv_window) || + if (!window_util::HasSymmetricPadding(conv_window) || window_util::HasBaseDilation(conv_window)) { // If padding is uneven or has dilation, we insert a kPad instruction that // applies positive padding and dilation. + // + // TODO(phawkins): If conv_window has asymmetric padding, perhaps instead of + // moving all the padding into an explicit pad op, we should keep as much + // padding inside of cudnn as possible, on the assumption that padding + // within cudnn is basically free, whereas a kPad's cost increases as the + // amount of padding increases. PaddingConfig padding_config = MakeNoPaddingConfig(input->shape().dimensions_size()); for (size_t i = 0; i < conv_dnums.input_spatial_dimensions().size(); ++i) { @@ -167,14 +173,17 @@ bool PadInsertion::CanonicalizeForwardConvolution(HloInstruction* conv) { dim->set_window_dilation(1); } + // The conv CustomCall returns a tuple (conv_result, scratch_buffer). Extract + // out the shape of conv_result. + Shape old_conv_shape = conv->shape().tuple_shapes(0); + VLOG(1) << "Canonicalizing forward conv"; - auto new_conv = HloInstruction::CreateConvolve( - conv->shape(), new_input, new_kernel, new_conv_window, - conv->convolution_dimension_numbers()); + auto new_conv = CreateCudnnConvForward(old_conv_shape, new_input, new_kernel, + new_conv_window, + conv->convolution_dimension_numbers()); VLOG(1) << "Replacing:\n " << conv->ToString() << "\nwith:\n " << new_conv->ToString(); - TF_CHECK_OK( - conv->parent()->ReplaceWithNewInstruction(conv, std::move(new_conv))); + TF_CHECK_OK(conv->parent()->ReplaceInstruction(conv, new_conv)); return true; } @@ -190,7 +199,9 @@ void IncreasePaddingHighBy(int64 delta, WindowDimension* window_dim) { bool PadInsertion::CanonicalizeBackwardFilterConvolution( HloInstruction* backward_conv) { - if (window_util::HasEvenPadding(backward_conv->window())) { + CHECK_EQ(backward_conv->custom_call_target(), + kCudnnConvBackwardFilterCallTarget); + if (window_util::HasSymmetricPadding(backward_conv->window())) { return false; } @@ -202,15 +213,11 @@ bool PadInsertion::CanonicalizeBackwardFilterConvolution( // ABCD0 = Pad(ABCD, padding_high=1) // BackwardFilterConv(ABCD0, xyz, padding_low=pading_high=1) // We choose the lesser of padding_low and padding_high as the new padding. - HloInstruction* forward_conv = backward_conv->fused_expression_root(); HloInstruction* input = backward_conv->mutable_operand(0); - Window new_forward_conv_window = forward_conv->window(); Window new_backward_conv_window = backward_conv->window(); // input_padding_config is the config of the kPad to be inserted. PaddingConfig input_padding_config = MakeNoPaddingConfig(ShapeUtil::Rank(input->shape())); - ConvolutionDimensionNumbers forward_conv_dnums = - forward_conv->convolution_dimension_numbers(); ConvolutionDimensionNumbers backward_conv_dnums = backward_conv->convolution_dimension_numbers(); for (size_t i = 0; i < backward_conv->window().dimensions_size(); ++i) { @@ -222,11 +229,7 @@ bool PadInsertion::CanonicalizeBackwardFilterConvolution( // cuDNN convolution (which doesn't support negative padding) to fail. return false; } - // If the backward convolution has uneven padding on the activations, we - // move some padding on the larger end to "internal" padding, so that the - // backward convolution produces larger weight gradients which get sliced - // later. Therefore, the amount of new padding (low or high) is the minimum - // of the amount of old padding low and old padding high. + // Compute the new, even padding for the backward conv operation. int64 new_conv_padding = std::min(padding_low, padding_high); int64 dim = backward_conv_dnums.input_spatial_dimensions(i); input_padding_config.mutable_dimensions(dim)->set_edge_padding_low( @@ -237,14 +240,9 @@ bool PadInsertion::CanonicalizeBackwardFilterConvolution( // Since we move some padding from the backward convolution to the kPad, we // need to accordingly reduce the padding amount of the backward convolution // and its inner forward convolution. - IncreasePaddingLowBy(-(padding_low - new_conv_padding), - new_backward_conv_window.mutable_dimensions(i)); - IncreasePaddingHighBy(-(padding_high - new_conv_padding), - new_backward_conv_window.mutable_dimensions(i)); - IncreasePaddingLowBy(-(padding_low - new_conv_padding), - new_forward_conv_window.mutable_dimensions(i)); - IncreasePaddingHighBy(-(padding_high - new_conv_padding), - new_forward_conv_window.mutable_dimensions(i)); + auto* new_dim = new_backward_conv_window.mutable_dimensions(i); + new_dim->set_padding_low(new_conv_padding); + new_dim->set_padding_high(new_conv_padding); } // Create a new backward convolution replacing the old one. @@ -260,19 +258,12 @@ bool PadInsertion::CanonicalizeBackwardFilterConvolution( .ConsumeValueOrDie(), input, padding, input_padding_config)); - HloInstruction* new_forward_conv = - computation->AddInstruction(HloInstruction::CreateConvolve( - ShapeInference::InferConvolveShape( - padded_input->shape(), output->shape(), new_forward_conv_window, - forward_conv_dnums) - .ConsumeValueOrDie(), - padded_input, output, new_forward_conv_window, forward_conv_dnums)); - - // Fuse the new forward convolution to the new backward convolution. - HloInstruction* new_backward_conv = - computation->CreateFusionInstructionForBackwardConvolution( - {new_forward_conv}, HloInstruction::FusionKind::kConvBackwardFilter, - new_backward_conv_window, backward_conv_dnums); + // The shape of the backward_conv CustomCall is a tuple (conv_result, + // scratch_buffer). Extract out the shape of conv_result. + Shape backward_conv_shape = backward_conv->shape().tuple_shapes(0); + HloInstruction* new_backward_conv = CreateCudnnConvBackwardFilter( + backward_conv_shape, padded_input, output, new_backward_conv_window, + backward_conv_dnums); VLOG(1) << "Canonicalizing backward filter conv"; VLOG(1) << "Replacing:\n " << backward_conv->ToString() << "\nwith:\n " @@ -285,18 +276,19 @@ bool PadInsertion::CanonicalizeBackwardFilterConvolution( bool PadInsertion::CanonicalizeBackwardInputConvolution( HloInstruction* backward_conv) { - if (window_util::HasEvenPadding(backward_conv->window())) { + if (window_util::HasSymmetricPadding(backward_conv->window())) { return false; } - HloInstruction* forward_conv = backward_conv->fused_expression_root(); - HloInstruction* reverse_filter = forward_conv->mutable_operand(1); - Window new_forward_conv_window = forward_conv->window(); Window new_backward_conv_window = backward_conv->window(); - ConvolutionDimensionNumbers forward_conv_dnums = - forward_conv->convolution_dimension_numbers(); ConvolutionDimensionNumbers backward_conv_dnums = backward_conv->convolution_dimension_numbers(); + + // The backward_conv CustomCall returns a tuple (conv_result, scratch_memory). + // Get the shape of conv_result. + Shape backward_conv_shape = backward_conv->shape().tuple_shapes(0); + + Shape new_backward_conv_shape = backward_conv_shape; for (size_t i = 0; i < backward_conv->window().dimensions_size(); ++i) { int64 padding_low = backward_conv->window().dimensions(i).padding_low(); int64 padding_high = backward_conv->window().dimensions(i).padding_high(); @@ -315,41 +307,38 @@ bool PadInsertion::CanonicalizeBackwardInputConvolution( // where the amount of padding low is larger, we can canonicalize it to // [B A] = BackwardInputConvolve([a b], [x y z], padding=(low=1,high=1)) // [A] = Slice([B A]) - // For consistency, we need to increase the low padding of the inner - // convolution by 1 as well because the input is larger now. if (padding_low > padding_high) { IncreasePaddingLowBy(padding_high - padding_low, new_backward_conv_window.mutable_dimensions(i)); - IncreasePaddingLowBy(padding_low - padding_high, - new_forward_conv_window.mutable_dimensions(i)); } else if (padding_low < padding_high) { IncreasePaddingHighBy(padding_low - padding_high, new_backward_conv_window.mutable_dimensions(i)); - IncreasePaddingHighBy(padding_high - padding_low, - new_forward_conv_window.mutable_dimensions(i)); } + // Decreasing the padding by X *increases* the size of our output by X. + int64 dim = backward_conv_dnums.output_spatial_dimensions(i); + new_backward_conv_shape.set_dimensions( + dim, new_backward_conv_shape.dimensions(dim) + + std::abs(padding_low - padding_high)); } // Create a new backward convolution replacing the old one. HloComputation* computation = backward_conv->parent(); HloInstruction* output = backward_conv->mutable_operand(0); HloInstruction* filter = backward_conv->mutable_operand(1); - HloInstruction* new_reverse_filter = - computation->AddInstruction(HloInstruction::CreateReverse( - filter->shape(), filter, reverse_filter->dimensions())); - HloInstruction* new_forward_conv = - computation->AddInstruction(HloInstruction::CreateConvolve( - ShapeInference::InferConvolveShape( - output->shape(), new_reverse_filter->shape(), - new_forward_conv_window, forward_conv_dnums) - .ConsumeValueOrDie(), - output, new_reverse_filter, new_forward_conv_window, - forward_conv_dnums)); + + HloInstruction* new_backward_conv_call = CreateCudnnConvBackwardInput( + new_backward_conv_shape, output, filter, new_backward_conv_window, + backward_conv_dnums); + + // The CustomCall created above returns a tuple (conv_result, scratch_memory). + // Extract out the two elements. HloInstruction* new_backward_conv = - computation->CreateFusionInstructionForBackwardConvolution( - {new_forward_conv, new_reverse_filter}, - HloInstruction::FusionKind::kConvBackwardInput, - new_backward_conv_window, backward_conv_dnums); + computation->AddInstruction(HloInstruction::CreateGetTupleElement( + new_backward_conv_shape, new_backward_conv_call, 0)); + HloInstruction* new_backward_conv_scratch = + computation->AddInstruction(HloInstruction::CreateGetTupleElement( + new_backward_conv_call->shape().tuple_shapes(1), + new_backward_conv_call, 1)); // Slice the new backward convolution. // @@ -377,22 +366,25 @@ bool PadInsertion::CanonicalizeBackwardInputConvolution( } // Replace the old backward convolution with the slice. - CHECK(ShapeUtil::Compatible( + Shape slice_shape = ShapeInference::InferSliceShape(new_backward_conv->shape(), start_indices, limit_indices, strides) - .ConsumeValueOrDie(), - backward_conv->shape())); + .ConsumeValueOrDie(); + CHECK(ShapeUtil::Compatible(slice_shape, backward_conv_shape)) + << ShapeUtil::HumanString(slice_shape) << " vs " + << ShapeUtil::HumanString(backward_conv_shape); - auto slice = - HloInstruction::CreateSlice(backward_conv->shape(), new_backward_conv, - start_indices, limit_indices, strides); + HloInstruction* slice = computation->AddInstruction( + HloInstruction::CreateSlice(backward_conv_shape, new_backward_conv, + start_indices, limit_indices, strides)); + HloInstruction* new_tuple = computation->AddInstruction( + HloInstruction::CreateTuple({slice, new_backward_conv_scratch})); VLOG(1) << "Canonicalizing backward input conv"; VLOG(1) << "Replacing:\n " << backward_conv->ToString() << "\nwith:\n " - << slice->ToString(); + << new_tuple->ToString(); - TF_CHECK_OK( - computation->ReplaceWithNewInstruction(backward_conv, std::move(slice))); + TF_CHECK_OK(computation->ReplaceInstruction(backward_conv, new_tuple)); return true; } @@ -400,18 +392,17 @@ StatusOr PadInsertion::Run(HloModule* module) { bool changed = false; for (HloInstruction* instruction : module->entry_computation()->MakeInstructionPostOrder()) { - if (instruction->opcode() == HloOpcode::kConvolution) { - changed |= CanonicalizeForwardConvolution(instruction); - } else if (instruction->opcode() == HloOpcode::kFusion) { - switch (instruction->fusion_kind()) { - case HloInstruction::FusionKind::kConvBackwardFilter: - changed |= CanonicalizeBackwardFilterConvolution(instruction); - break; - case HloInstruction::FusionKind::kConvBackwardInput: - changed |= CanonicalizeBackwardInputConvolution(instruction); - break; - default: - break; + if (IsCustomCallToDnnConvolution(*instruction)) { + const auto& target = instruction->custom_call_target(); + if (target == kCudnnConvForwardCallTarget) { + changed |= CanonicalizeForwardConvolution(instruction); + } else if (target == kCudnnConvBackwardFilterCallTarget) { + changed |= CanonicalizeBackwardFilterConvolution(instruction); + } else if (target == kCudnnConvBackwardInputCallTarget) { + changed |= CanonicalizeBackwardInputConvolution(instruction); + } else { + LOG(FATAL) << "Unknown custom call target for cudnn conv: " + << instruction->ToString(); } } } diff --git a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc index 457e6094d90413440658452937bff2ccfe6cbe5c..388dcc008b07a76ff9ed07df04181e49a8734f51 100644 --- a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc @@ -88,6 +88,23 @@ llvm_ir::IrArray::Index ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( /*HasNUW=*/true, /*HasNSW=*/true), thread_id, "linear_index", /*HasNUW=*/true, /*HasNSW=*/true); + // Add an @llvm.assume(linear_index < threads_per_block * num_blocks). + // + // This might seem obvious from the computation above, but LLVM does not + // currently determine the range of linear_index precisely. InstCombine uses + // known-bits, which, when applied to the task of determining a value's range, + // is imprecise for everything other than powers of 2. And + // CorrelatedValuePropagation is, as a cost-saving measure, disabled for + // conditions in the same basic block as their operands. + llvm_ir::EmitCallToIntrinsic( + llvm::Intrinsic::assume, + {ir_builder_->CreateICmpULT( + linear_index, + ir_builder_->getInt64(launch_dimensions_.threads_per_block() * + launch_dimensions_.block_count()), + "linear_index_in_range")}, + {}, ir_builder_); + auto if_in_bounds = llvm_ir::EmitIfThenElse( ir_builder_->CreateICmpULT( linear_index, ir_builder_->getInt64(ShapeUtil::ElementsIn(shape_))), diff --git a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h index 934e7e1919f08a16daf09ec634e2f9dc0c7cc723..8ed63a854a74fc06c3c389f40fe1f5970885deac 100644 --- a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h +++ b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h @@ -42,6 +42,11 @@ class ParallelLoopEmitter : public llvm_ir::LoopEmitter { const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder); + // Constructs a loop emitter for a loop that generates on element of each of N + // arrays on each iteration. + // + // This is used in multi-output fusion. target_element_generator should + // produce a struct with N elements, one for each of target_arrays. ParallelLoopEmitter( const llvm_ir::ElementGenerator& target_element_generator, tensorflow::gtl::ArraySlice target_arrays, diff --git a/tensorflow/compiler/xla/service/gpu/thunk.h b/tensorflow/compiler/xla/service/gpu/thunk.h index 191c7675c6d14274bfd1f5fdf9e27cf37fbd231f..2c3032d79be221e8cacb178ffb1817459b603cc0 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk.h +++ b/tensorflow/compiler/xla/service/gpu/thunk.h @@ -41,11 +41,13 @@ class GpuExecutable; class Thunk { public: enum class Kind { + kConditional, kConvolution, kCopy, kCudnnBatchNormBackward, kCudnnBatchNormForwardInference, kCudnnBatchNormForwardTraining, + kFft, kGemm, kInfeed, kKernel, diff --git a/tensorflow/compiler/xla/service/gpu/while_transformer.cc b/tensorflow/compiler/xla/service/gpu/while_transformer.cc index ab94d7d5436e8edd12f68f7e0c395c53f303e6eb..e6caec8625f0d622dbb92bcc20802d254fe23f94 100644 --- a/tensorflow/compiler/xla/service/gpu/while_transformer.cc +++ b/tensorflow/compiler/xla/service/gpu/while_transformer.cc @@ -62,7 +62,7 @@ namespace { // &tagged_instructions)); // // Instructions that are "tagged" with a context-specific string will -// be returned in 'tagged_instructions' for further procesing (i.e. parsing +// be returned in 'tagged_instructions' for further processing (i.e. parsing // constants or recording the tuple_index). // class ExprTree { diff --git a/tensorflow/compiler/xla/service/gpu/while_transformer.h b/tensorflow/compiler/xla/service/gpu/while_transformer.h index a4f527fce0e4e280e24efc1f33ea68a0b71555b9..fe3a954e1828ee4a323872eea81f64c7e780ad24 100644 --- a/tensorflow/compiler/xla/service/gpu/while_transformer.h +++ b/tensorflow/compiler/xla/service/gpu/while_transformer.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_WHILE_TRANSFORMER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_WHILE_TRANSFORMER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_WHILE_TRANSFORMER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_WHILE_TRANSFORMER_H_ #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/statusor.h" @@ -40,4 +40,4 @@ StatusOr> CanTransformWhileToFor( } // namespace gpu } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_GPU_WHILE_TRANSFORMER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_WHILE_TRANSFORMER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc b/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc index f16daa0b5481474e754c880ead1945297ca50168..2f290f61bd527e9827472a78256f015e066e44be 100644 --- a/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc +++ b/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc @@ -117,9 +117,7 @@ class WhileTransformerTest : public HloTestBase { } void RunCopyInsertionPass() { - HloVerifier verifier([](const Shape& shape) { - return ShapeUtil::ByteSizeOf(shape, /*pointer_size=*/sizeof(void*)); - }); + HloVerifier verifier; TF_ASSERT_OK(verifier.Run(module_.get()).status()); CopyInsertion copy_insertion; TF_ASSERT_OK(copy_insertion.Run(module_.get()).status()); diff --git a/tensorflow/compiler/xla/service/heap_simulator.cc b/tensorflow/compiler/xla/service/heap_simulator.cc index 34e2f7ee206c6a74073d8f4e867e862feb4aff49..cde5877e29f36abc61c5417ce960e2c7699e2749 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.cc +++ b/tensorflow/compiler/xla/service/heap_simulator.cc @@ -64,10 +64,8 @@ StatusOr HeapSimulator::Run( std::unique_ptr algorithm, const HloModule& module, const SequentialHloOrdering::HloModuleSequence& module_sequence, const TuplePointsToAnalysis& points_to_analysis, - const LogicalBuffer::SizeFunction& size_fn, - const FlatSet* buffers_to_assign) { - HeapSimulator heap(std::move(algorithm), size_fn, buffers_to_assign, - &module_sequence); + const LogicalBuffer::SizeFunction& size_fn, const Options& options) { + HeapSimulator heap(std::move(algorithm), size_fn, options, &module_sequence); const HloComputation* entry_computation = module.entry_computation(); const std::vector& instruction_sequence = FindOrDie(module_sequence, entry_computation); @@ -81,9 +79,8 @@ StatusOr HeapSimulator::Run( std::unique_ptr algorithm, const HloComputation& computation, const std::vector& instruction_sequence, const TuplePointsToAnalysis& points_to_analysis, - const LogicalBuffer::SizeFunction& size_fn, - const FlatSet* buffers_to_assign) { - HeapSimulator heap(std::move(algorithm), size_fn, buffers_to_assign, + const LogicalBuffer::SizeFunction& size_fn, const Options& options) { + HeapSimulator heap(std::move(algorithm), size_fn, options, /*module_sequence=*/nullptr); TF_RETURN_IF_ERROR(heap.RunComputation(computation, instruction_sequence, points_to_analysis)); @@ -199,15 +196,17 @@ Status HeapSimulator::RunComputation( // We can only share with the operand buffer if it is about to be freed; // we must be the last user of the buffer. bool shared = false; - for (const LogicalBuffer* operand_buffer : operand_buffers_to_free) { - if (buffer->instruction()->IsUserOf(operand_buffer->instruction()) && - buffer->instruction()->opcode() != HloOpcode::kCopy && - CanShareOperandBufferWithUser( - operand_buffer->instruction(), operand_buffer->index(), - buffer->instruction(), buffer->index(), points_to_analysis)) { - ShareBuffer(buffer, operand_buffer, instruction); - shared = true; - break; + if (options_.may_reuse_operand_buffers) { + for (const LogicalBuffer* operand_buffer : operand_buffers_to_free) { + if (buffer->instruction()->IsUserOf(operand_buffer->instruction()) && + buffer->instruction()->opcode() != HloOpcode::kCopy && + CanShareOperandBufferWithUser( + operand_buffer->instruction(), operand_buffer->index(), + buffer->instruction(), buffer->index(), points_to_analysis)) { + ShareBuffer(buffer, operand_buffer, instruction); + shared = true; + break; + } } } @@ -266,13 +265,12 @@ Status HeapSimulator::RunComputation( HeapSimulator::HeapSimulator( std::unique_ptr algorithm, - const LogicalBuffer::SizeFunction& size_fn, - const FlatSet* buffers_to_assign, + const LogicalBuffer::SizeFunction& size_fn, const Options& options, const SequentialHloOrdering::HloModuleSequence* module_sequence) : no_fragmentation_stats_(MakeUnique()), algorithm_(std::move(algorithm)), size_fn_(size_fn), - buffers_to_assign_(buffers_to_assign), + options_(options), module_sequence_(module_sequence) { debug_trace_.set_whole_module_simulation(module_sequence_ != nullptr); } @@ -280,13 +278,16 @@ HeapSimulator::HeapSimulator( HeapSimulator::~HeapSimulator() {} bool HeapSimulator::IgnoreBuffer(const LogicalBuffer* buffer) const { - // Buffers for constants are ignored, as with BufferAssigner. Also ignore - // buffers that we're not meant to assign. + // Buffers for constants are ignored unless the alloc_constants option is + // set. Also ignore buffers that we're not meant to assign. // // TODO(b/32248867): For consistency, constants should get allocations. - return buffer->instruction()->opcode() == HloOpcode::kConstant || - (buffers_to_assign_ != nullptr && - buffers_to_assign_->count(buffer) == 0); + if (!options_.alloc_constants && + buffer->instruction()->opcode() == HloOpcode::kConstant) { + return true; + } + return options_.buffers_to_assign != nullptr && + options_.buffers_to_assign->count(buffer) == 0; } // Alloc always calls the underlying heap algorithm. @@ -400,8 +401,8 @@ HeapSimulator::Result HeapSimulator::Finish() { } // If we were told to assign specific buffers, make sure we've assigned // exactly that many buffers. - if (buffers_to_assign_ != nullptr) { - CHECK_EQ(buffers_to_assign_->size(), result.chunk_map.size()); + if (options_.buffers_to_assign != nullptr) { + CHECK_EQ(options_.buffers_to_assign->size(), result.chunk_map.size()); } } diff --git a/tensorflow/compiler/xla/service/heap_simulator.h b/tensorflow/compiler/xla/service/heap_simulator.h index 88a8698d16132372fc8f4e87eba3b99125aab876..636f19dd39f09721bd82fc4b44785f196f281ad7 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.h +++ b/tensorflow/compiler/xla/service/heap_simulator.h @@ -67,6 +67,23 @@ class HeapSimulator { HeapSimulatorTrace debug_trace; }; + // The different options to be passed to the Run() APIs. + struct Options { + Options() + : may_reuse_operand_buffers(true), + alloc_constants(false), + buffers_to_assign(nullptr) {} + + // Whether a buffer about to be Free()-ed, can be recycled for a new born + // one, hence collapsing Free()+Alloc() calls (default true). + bool may_reuse_operand_buffers; + // Whether to issue Alloc() and Free() calls for constants (default false). + bool alloc_constants; + // If 'buffers_to_assign' is provided, only those buffers are assigned + // offsets, otherwise all buffers defined by the instructions are assigned. + const tensorflow::gtl::FlatSet* buffers_to_assign; + }; + // Run the heap simulation with the given algorithm, assuming the given // module_sequence, which must contain a topologically-consistent total // ordering of all instructions within each computation. The result is invalid @@ -76,15 +93,12 @@ class HeapSimulator { // to running on a per-computation basis, since we can re-use buffer space for // called sub-computations. // - // If 'buffers_to_assign' is provided, only those buffers are assigned - // offsets, otherwise all buffers defined by the instructions are assigned. static StatusOr Run( std::unique_ptr algorithm, const HloModule& module, const SequentialHloOrdering::HloModuleSequence& module_sequence, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_fn, - const tensorflow::gtl::FlatSet* buffers_to_assign = - nullptr); + const Options& options = Options()); // Same as above, but runs on a single computation. The 'instruction_sequence' // must contain a topologically-consistent total ordering of all instructions @@ -96,8 +110,7 @@ class HeapSimulator { const std::vector& instruction_sequence, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_fn, - const tensorflow::gtl::FlatSet* buffers_to_assign = - nullptr); + const Options& options = Options()); private: // If 'module_sequence' is non-null, it is used to find kCall and kWhile @@ -105,8 +118,7 @@ class HeapSimulator { // be run recursively. I.e. the simulation is run over the whole module. HeapSimulator( std::unique_ptr algorithm, - const LogicalBuffer::SizeFunction& size_fn, - const tensorflow::gtl::FlatSet* buffers_to_assign, + const LogicalBuffer::SizeFunction& size_fn, const Options& options, const SequentialHloOrdering::HloModuleSequence* module_sequence); ~HeapSimulator(); @@ -130,7 +142,7 @@ class HeapSimulator { const std::unique_ptr no_fragmentation_stats_; const std::unique_ptr algorithm_; const LogicalBuffer::SizeFunction size_fn_; - const tensorflow::gtl::FlatSet* buffers_to_assign_; + const Options options_; const SequentialHloOrdering::HloModuleSequence* module_sequence_; // In addition to Alloc and Free, the heap simulator exposes a concept of diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto index e4aed7593c51a2d1bb156493666b3c1b03dcc626..36db711c6c3570efdf678261ad38bbdb08cf94aa 100644 --- a/tensorflow/compiler/xla/service/hlo.proto +++ b/tensorflow/compiler/xla/service/hlo.proto @@ -123,6 +123,12 @@ message HloInstructionProto { // Describes the dimension numbers used for a dot operation xla.DotDimensionNumbers dot_dimension_numbers = 30; + + // FFT type (FFT, IFFT, etc). + xla.FftType fft_type = 31; + + // FFT length. + repeated int64 fft_length = 32; } // Serialization of HloComputation. @@ -194,6 +200,7 @@ message BufferAllocationProto { bool is_reusable = 4; bool is_entry_computation_parameter = 5; int64 parameter_number = 6; + repeated int64 parameter_shape_index = 10; bool maybe_live_out = 7; int64 color = 8; repeated Assigned assigned = 9; diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index a63affa06caf75f1ccab084bd114e39ba7c91a38..5432419e4a2dd2916da32ac6566851bf52fd68ca 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -461,20 +461,6 @@ HloInstruction* HloComputation::CreateFusionInstruction( return fusion_instruction; } -HloInstruction* HloComputation::CreateFusionInstructionForBackwardConvolution( - tensorflow::gtl::ArraySlice instructions_to_fuse, - HloInstruction::FusionKind fusion_kind, const Window& window, - const ConvolutionDimensionNumbers& conv_dnums) { - CHECK(HloInstruction::FusionKind::kConvBackwardFilter == fusion_kind || - HloInstruction::FusionKind::kConvBackwardInput == fusion_kind); - HloInstruction* root = instructions_to_fuse.front(); - HloInstruction* fusion_instruction = - AddInstruction(HloInstruction::CreateFusionForBackwardConvolution( - root->shape(), fusion_kind, window, conv_dnums, root)); - FuseInstructionsInto(instructions_to_fuse, fusion_instruction); - return fusion_instruction; -} - StatusOr HloComputation::DeepCopyHelper( HloInstruction* instruction, const ShapeTree* indices_to_copy, ShapeTree* copies_added, ShapeIndex* index) { @@ -577,8 +563,11 @@ Status HloComputation::ReplaceWithNewInstruction( Status HloComputation::ReplaceInstruction(HloInstruction* old_instruction, HloInstruction* new_instruction) { - TF_RET_CHECK(ShapeUtil::Compatible(old_instruction->shape(), - new_instruction->shape())); + TF_RET_CHECK( + ShapeUtil::Compatible(old_instruction->shape(), new_instruction->shape())) + << ShapeUtil::HumanString(old_instruction->shape()) << " vs " + << ShapeUtil::HumanString(new_instruction->shape()); + VLOG(10) << "transformed " << old_instruction->ToString() << " to " << new_instruction->ToString(); // Try to add metadata for HLO instructions that are created to replace diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h index 6436815f910405477ec21a33dec75ef71df08602..061c59abe5e315917161ed737f89de53d71bb1b6 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.h +++ b/tensorflow/compiler/xla/service/hlo_computation.h @@ -224,15 +224,6 @@ class HloComputation { tensorflow::gtl::ArraySlice instructions_to_fuse, HloInstruction::FusionKind fusion_kind); - // Creates a fusion instruction that represents a backward convolution. This - // is similar to CreateFusionInstruction but takes window and conv_dnums which - // indicate the window and convolution dimension numbers of the backward - // convolution. - HloInstruction* CreateFusionInstructionForBackwardConvolution( - tensorflow::gtl::ArraySlice instructions_to_fuse, - HloInstruction::FusionKind fusion_kind, const Window& window, - const ConvolutionDimensionNumbers& conv_dnums); - // Create a deep copy of the given instruction and return the instruction // producing the copied result. All instructions performing the copy are added // to the computation. For array-shaped values, this method trivially returns diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc index b933695b823871c6c0174da6d6f99e618219442a..9cd5a1e2b71a7aa768e478289e8e4cc13030fcc3 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc @@ -392,6 +392,21 @@ Status HloCostAnalysis::HandleConvolution(const HloInstruction* convolution) { return Status::OK(); } +Status HloCostAnalysis::HandleFft(const HloInstruction* fft) { + auto real_shape = + ShapeUtil::IsTuple(fft->operand(0)->shape()) + ? ShapeUtil::GetTupleElementShape(fft->operand(0)->shape(), 0) + : fft->operand(0)->shape(); + constexpr int kFmaPerComplexMul = 4; + int64 log_factors = 1; + for (int64 dim : fft->fft_length()) { + log_factors *= tensorflow::Log2Floor(dim); + } + current_properties_[kFlopsKey] = kFmaFlops * kFmaPerComplexMul * log_factors * + ShapeUtil::ElementsIn(real_shape); + return Status::OK(); +} + Status HloCostAnalysis::HandleCrossReplicaSum(const HloInstruction* crs) { // We assume 2 replicas, so that each output element is the sum of two input // elements. @@ -454,7 +469,13 @@ Status HloCostAnalysis::HandleCall(const HloInstruction* call) { } Status HloCostAnalysis::HandleCustomCall(const HloInstruction*) { - return Unimplemented("Custom-call is not implemented for HLO cost analysis."); + // We can't do anything sane with CustomCalls, since we don't know what they + // do, and returning an error status will stop iteration over this + // computation, which is probably also not what we want. So just punt and + // return OK. This will cause all of the properties to be reported as 0, + // which is fine. + current_should_compute_bottleneck_time_ = false; + return Status::OK(); } Status HloCostAnalysis::HandleSort(const HloInstruction* sort) { diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.h b/tensorflow/compiler/xla/service/hlo_cost_analysis.h index fade19522cf0c30eab037aa355de1f9203f80014..e5783539e5436f09fa58bf7889118380ee90fea0 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.h @@ -67,6 +67,7 @@ class HloCostAnalysis : public ConstDfsHloVisitor { Status HandleCopy(const HloInstruction* copy) override; Status HandleDot(const HloInstruction* dot) override; Status HandleConvolution(const HloInstruction* convolution) override; + Status HandleFft(const HloInstruction* fft) override; Status HandleCrossReplicaSum(const HloInstruction* crs) override; Status HandleInfeed(const HloInstruction* infeed) override; Status HandleOutfeed(const HloInstruction* outfeed) override; diff --git a/tensorflow/compiler/xla/service/hlo_cse.cc b/tensorflow/compiler/xla/service/hlo_cse.cc index d35ba19a730555433099072c51ca5cf3774d4b99..279edd4ba8772a9c576f76f554de8ec68631b953 100644 --- a/tensorflow/compiler/xla/service/hlo_cse.cc +++ b/tensorflow/compiler/xla/service/hlo_cse.cc @@ -32,6 +32,7 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/inlined_vector.h" namespace xla { @@ -91,6 +92,10 @@ bool CombineConstants(HloComputation* computation, bool is_layout_sensitive) { StatusOr HloCSE::Run(HloModule* module) { bool changed = false; + const std::function + eq_instructions = std::equal_to(); + const std::function + eq_computations = std::equal_to(); for (auto* computation : module->computations()) { changed |= CombineConstants(computation, is_layout_sensitive_); @@ -110,11 +115,12 @@ StatusOr HloCSE::Run(HloModule* module) { // of this instruction. const HloInstruction* operand = instruction->operand(0); - std::vector equivalent_instructions; + tensorflow::gtl::InlinedVector + equivalent_instructions; for (HloInstruction* user : operand->users()) { - if (user != instruction && user->Identical(*instruction) && - (!is_layout_sensitive_ || - ShapeUtil::Equal(user->shape(), instruction->shape()))) { + if (user != instruction && + user->Identical(*instruction, eq_instructions, eq_computations, + is_layout_sensitive_)) { equivalent_instructions.push_back(user); } } diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc index 80d89d851e239a78f43bb2078e724d76a188ced4..ccbbe8f1966d59b4ab2904dcc6ea724aaf4a7603 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc @@ -154,7 +154,11 @@ bool HloDataflowAnalysis::Phi( tensorflow::gtl::ArraySlice inputs) { CHECK(ssa_form_); VLOG(4) << "Phi(" << instruction->name() << ")"; - + VLOG(5) << "instruction value set = " + << GetInstructionValueSet(instruction).ToString(); + for (const InstructionValueSet* input : inputs) { + VLOG(5) << "input value set = " << input->ToString(); + } for (const InstructionValueSet* input : inputs) { DCHECK(ShapeUtil::Compatible(instruction->shape(), input->shape())); } @@ -171,9 +175,14 @@ bool HloDataflowAnalysis::Phi( value_set.values().size() == 1 ? value_set.values()[0] : nullptr; // Construct a vector of unique value IDs of the inputs. + // Don't add value ids where the input is equal to the definition. std::vector input_value_ids; for (const InstructionValueSet* input : inputs) { for (const HloValue* value : input->element(index).values()) { + if (value->defining_instruction() == instruction && + value->defining_index() == index) { + continue; + } input_value_ids.push_back(value->id()); } } @@ -190,6 +199,7 @@ bool HloDataflowAnalysis::Phi( current_value->defining_instruction() == instruction && current_value->defining_index() == index); if (current_value_defined_here) { + VLOG(5) << "current_value_defined_here: " << current_value->ToString(); CHECK(current_value->is_phi()); auto it = std::find(input_value_ids.begin(), input_value_ids.end(), current_value->id()); @@ -197,7 +207,7 @@ bool HloDataflowAnalysis::Phi( input_value_ids.erase(it); } } - + VLOG(5) << "after input_value_ids.size = " << input_value_ids.size(); if (input_value_ids.empty()) { // A value set which has at least one element should never have its value // set reduced to zero elements. During dataflow value sets only can go @@ -575,16 +585,23 @@ bool HloDataflowAnalysis::UpdateInstructionValueSet( void HloDataflowAnalysis::Propagate() { std::queue worklist; + tensorflow::gtl::FlatSet workset; + auto add_to_worklist = [&worklist, &workset](HloInstruction* instruction) { + if (workset.insert(instruction).second) { + worklist.push(instruction); + } + }; for (HloComputation* computation : module_->computations()) { for (HloInstruction* instruction : computation->instructions()) { - worklist.push(instruction); + add_to_worklist(instruction); } } while (!worklist.empty()) { HloInstruction* instruction = worklist.front(); worklist.pop(); + workset.erase(workset.find(instruction)); VLOG(3) << "Worklist top: " << instruction->name(); VLOG(3) << ToString(); @@ -598,9 +615,10 @@ void HloDataflowAnalysis::Propagate() { VLOG(4) << "New value set for " << instruction->name() << ": " << GetInstructionValueSet(instruction); - // Instruction value was updated. Add users to work list. + // Instruction value was updated. Add users to work list if we haven't + // already. for (HloInstruction* user : instruction->users()) { - worklist.push(user); + add_to_worklist(user); // If user sequentially calls a computation, then the respective // parameter(s) of the computation need to be updated. @@ -615,10 +633,10 @@ void HloDataflowAnalysis::Propagate() { // Note that the same instruction can be used in both operand 1 and // operand 2. if (user->operand(1) == instruction) { - worklist.push(user->true_computation()->parameter_instruction(0)); + add_to_worklist(user->true_computation()->parameter_instruction(0)); } if (user->operand(2) == instruction) { - worklist.push(user->false_computation()->parameter_instruction(0)); + add_to_worklist(user->false_computation()->parameter_instruction(0)); } } else { for (HloComputation* called_computation : user->called_computations()) { @@ -626,7 +644,7 @@ void HloDataflowAnalysis::Propagate() { call_graph_->GetNode(called_computation); if (call_graph_node.context() == CallContext::kSequential) { for (int64 operand_number : user->OperandIndices(instruction)) { - worklist.push( + add_to_worklist( called_computation->parameter_instruction(operand_number)); } } @@ -642,13 +660,13 @@ void HloDataflowAnalysis::Propagate() { for (const CallSite& callsite : call_graph_node.caller_callsites()) { if ((callsite.instruction()->opcode() == HloOpcode::kCall) || (callsite.instruction()->opcode() == HloOpcode::kConditional)) { - worklist.push(callsite.instruction()); + add_to_worklist(callsite.instruction()); } else if (callsite.instruction()->opcode() == HloOpcode::kWhile) { // Add the while itself, and the body and condition parameters. - worklist.push(callsite.instruction()); - worklist.push( + add_to_worklist(callsite.instruction()); + add_to_worklist( callsite.instruction()->while_body()->parameter_instruction(0)); - worklist.push( + add_to_worklist( callsite.instruction()->while_condition()->parameter_instruction( 0)); } diff --git a/tensorflow/compiler/xla/service/hlo_element_type_converter.cc b/tensorflow/compiler/xla/service/hlo_element_type_converter.cc index 1773bb401d380031f6c860d295e76d2f62c9e5ff..c782d1b0add17c70e0f54826917df251d5a613e2 100644 --- a/tensorflow/compiler/xla/service/hlo_element_type_converter.cc +++ b/tensorflow/compiler/xla/service/hlo_element_type_converter.cc @@ -54,45 +54,96 @@ bool HasOperandType(HloInstruction* hlo, PrimitiveType type) { return false; } +// Finds out the Tuple Shape of the new instruction after converting the element +// type of the operands of the original instruction from `from_type` to +// `to_type`. +// +// This routine assumes the resulting `shape` of the original instruction is a +// non-nested tuple. This assumption is currently safe as only kTuple, kInfeed, +// kOutfeed, kCall, kCustomCall and kBatchNorm* HLO instructions can produce +// results with tuple shapes, and this routine is only called to convert the +// result shapes of kBatchNorm* HLO instructions, which are non-nested tuples. +Shape GetConvertedTupleShape(const Shape& shape, PrimitiveType from_type, + PrimitiveType to_type) { + std::vector new_tuple_subshapes; + for (int64 i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { + Shape subshape = ShapeUtil::GetTupleElementShape(shape, i); + CHECK(!ShapeUtil::IsTuple(subshape)); + if (subshape.element_type() == from_type) { + subshape = ShapeUtil::ChangeElementType(subshape, to_type); + } + new_tuple_subshapes.push_back(subshape); + } + return ShapeUtil::MakeTupleShape(new_tuple_subshapes); +} + +// Converts the elements of the result of `hlo` to produce a new tuple with +// shape `to_shape`. +// +// This routine assumes `hlo` is an instruction that produces a non-nested Tuple +// as a result. +HloInstruction* ConvertTupleElements(HloInstruction* hlo, + const Shape& to_shape) { + const Shape& shape = hlo->shape(); + HloComputation* computation = hlo->parent(); + std::vector tuple_elements; + for (int64 i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { + const Shape& ele_shape = ShapeUtil::GetTupleElementShape(shape, i); + HloInstruction* element = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(ele_shape, hlo, i)); + const Shape& to_ele_shape = ShapeUtil::GetTupleElementShape(to_shape, i); + CHECK(!ShapeUtil::IsTuple(ele_shape)); + if (ele_shape.element_type() != to_ele_shape.element_type()) { + element = computation->AddInstruction( + HloInstruction::CreateConvert(to_ele_shape, element)); + } + tuple_elements.push_back(element); + } + return computation->AddInstruction( + HloInstruction::CreateTuple(tuple_elements)); +} + } // namespace HloElementTypeConverter::HloElementTypeConverter( PrimitiveType eliminate_type, PrimitiveType replace_with_type) : eliminate_type_(eliminate_type), replace_with_type_(replace_with_type) {} +// This routine converts the arithmetic operations in the given module that use +// eliminate_type_ to operations that use replace_with_type_. StatusOr HloElementTypeConverter::Run(HloModule* module) { XLA_VLOG_LINES( 3, "HloElementTypeConverter::Run(), before:\n" + module->ToString()); + + if (eliminate_type_ == replace_with_type_) { + return false; + } + bool changed = false; for (auto* computation : module->computations()) { for (auto* hlo : computation->MakeInstructionPostOrder()) { + const auto opcode = hlo->opcode(); // These are ops where it does not make sense to convert them. - if (hlo->opcode() == HloOpcode::kParameter || - hlo->opcode() == HloOpcode::kConstant || - hlo->opcode() == HloOpcode::kTuple || - hlo->opcode() == HloOpcode::kConvert || - hlo->opcode() == HloOpcode::kGetTupleElement || - hlo->opcode() == HloOpcode::kInfeed || - hlo->opcode() == HloOpcode::kOutfeed) { + if (opcode == HloOpcode::kParameter || opcode == HloOpcode::kConstant || + opcode == HloOpcode::kTuple || opcode == HloOpcode::kConvert || + opcode == HloOpcode::kGetTupleElement || + opcode == HloOpcode::kInfeed || opcode == HloOpcode::kOutfeed) { continue; } // We cannot change a CustomCall since we have no way of adjusting the // called binary to expect the updated type. - if (hlo->opcode() == HloOpcode::kCustomCall) { + if (opcode == HloOpcode::kCustomCall) { continue; } // These are ops with embedded computations where it suffices to convert // the embedded computations instead of converting the ops themselves. - if (hlo->opcode() == HloOpcode::kWhile || - hlo->opcode() == HloOpcode::kCall || - hlo->opcode() == HloOpcode::kFusion || - hlo->opcode() == HloOpcode::kMap || - hlo->opcode() == HloOpcode::kReduce || - hlo->opcode() == HloOpcode::kReduceWindow || - hlo->opcode() == HloOpcode::kSelectAndScatter || - hlo->opcode() == HloOpcode::kConditional) { + if (opcode == HloOpcode::kWhile || opcode == HloOpcode::kCall || + opcode == HloOpcode::kFusion || opcode == HloOpcode::kMap || + opcode == HloOpcode::kReduce || opcode == HloOpcode::kReduceWindow || + opcode == HloOpcode::kSelectAndScatter || + opcode == HloOpcode::kConditional) { continue; } TF_RET_CHECK(hlo->called_computations().empty()) << hlo->ToString(); @@ -106,6 +157,11 @@ StatusOr HloElementTypeConverter::Run(HloModule* module) { continue; } + // Handle instructions that perform arithmetic operations and contain + // operands with eliminate_type_. + // + // First, convert the operands with eliminate_type_ to operands with + // replace_with_type_. std::vector new_operands; for (HloInstruction* operand : hlo->operands()) { if (operand->shape().element_type() == eliminate_type_) { @@ -114,6 +170,10 @@ StatusOr HloElementTypeConverter::Run(HloModule* module) { new_operands.push_back(operand); } + // Then find out the result type of the new instruction with the same + // opcode but using the converted operands, create the new instruction, + // and convert the result of the new instruction back to match the result + // type of the original instruction. HloInstruction* new_hlo; if (hlo->shape().element_type() == eliminate_type_) { Shape shape = @@ -121,10 +181,20 @@ StatusOr HloElementTypeConverter::Run(HloModule* module) { new_hlo = computation->AddInstruction( hlo->CloneWithNewOperands(shape, new_operands, hlo->GetModule())); new_hlo = ToElementType(new_hlo, eliminate_type_); + } else if (ShapeUtil::IsTuple(hlo->shape())) { + Shape old_shape = hlo->shape(); + Shape new_shape = GetConvertedTupleShape(hlo->shape(), eliminate_type_, + replace_with_type_); + new_hlo = computation->AddInstruction(hlo->CloneWithNewOperands( + new_shape, new_operands, hlo->GetModule())); + // Convert the elements of the result of `new_hlo` to produce a new + // tuple with shape `old_shape`. + new_hlo = ConvertTupleElements(new_hlo, old_shape); } else { new_hlo = computation->AddInstruction(hlo->CloneWithNewOperands( hlo->shape(), new_operands, hlo->GetModule())); } + TF_RETURN_IF_ERROR(computation->ReplaceInstruction(hlo, new_hlo)); changed = true; } diff --git a/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc b/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..cb94d9f19b825d1321263a4737b66a6bf198a772 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc @@ -0,0 +1,121 @@ +/* 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/compiler/xla/service/hlo_element_type_converter.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" + +namespace xla { +namespace { + +namespace op = xla::testing::opcode_matchers; + +class HloElementTypeConverterTest : public HloTestBase { + public: + std::unique_ptr CreateModuleFromHloString( + const string& hlo_string) { + return HloRunner::CreateModuleFromString(hlo_string, + GetDebugOptionsForTest()) + .ValueOrDie(); + } +}; + +TEST_F(HloElementTypeConverterTest, CustomCallsNotConverted) { + const string& hlo_string = R"( + HloModule custom_call + ENTRY CustomCall { + constant = bf16[1]{0} constant({12345}) + ROOT custom-call = bf16[1,2,3]{0,2,1} custom-call(constant), + custom_call_target="foo" + } + )"; + auto module = CreateModuleFromHloString(hlo_string); + HloElementTypeConverter type_converter(BF16, F32); + TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); + EXPECT_FALSE(converted); +} + +TEST_F(HloElementTypeConverterTest, InfeedsOutfeedsNotConverted) { + const string& hlo_string = R"( + HloModule InfeedOutfeed + ENTRY RoundTrip16MiBR1.v2 { + ROOT infeed = bf16[4]{0} infeed() + outfeed = () outfeed(infeed) + } + )"; + auto module = CreateModuleFromHloString(hlo_string); + HloElementTypeConverter type_converter(BF16, F32); + TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); + EXPECT_FALSE(converted); +} + +TEST_F(HloElementTypeConverterTest, OperationsInNestedTuplesConverted) { + const string& hlo_string = R"( + HloModule NestedTuples + ENTRY NestedTuples.v5 { + constant.4 = bf16[] constant(42) + constant.2 = f32[2]{0} constant({1, 2}) + constant.3 = bf16[] constant(42) + add = bf16[] add(constant.2, constant.3) + tuple = (f32[2]{0}, bf16[]) tuple(constant.2, add) + constant.5 = bf16[2]{0} constant({22, 44}) + ROOT tuple.1 = ((f32[2]{0}, bf16[]), bf16[2]{0}) tuple(tuple, constant.5) + } + )"; + + auto module = CreateModuleFromHloString(hlo_string); + HloElementTypeConverter type_converter(BF16, F32); + TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); + EXPECT_TRUE(converted); + const HloInstruction* bf16_op = + module->entry_computation()->root_instruction()->operand(0)->operand(1); + EXPECT_THAT(bf16_op, op::Convert(op::Add(op::Constant(), op::Convert()))); +} + +TEST_F(HloElementTypeConverterTest, BatchNormGradBF16Converted) { + const string& hlo_string = R"( + HloModule BatchNormGrad + ENTRY BatchNormGrad.v6 { + constant.4 = bf16[2,2,2,1]{3,2,1,0} constant(bf16[2,2,2,1] { { /*i0=0*/ + { /*i1=0*/ {0}, {0} }, { /*i1=1*/ {0}, {0} } }, { /*i0=1*/ { /*i1=0*/ {0}, + {0} }, { /*i1=1*/ {0}, {0} } } }) + constant.5 = bf16[2]{0} constant({1, 1}) + constant.6 = bf16[2]{0} constant({0, 0}) + constant.7 = bf16[2]{0} constant({1, 1}) + constant.8 = bf16[2,2,2,1]{3,2,1,0} constant(bf16[2,2,2,1] { { /*i0=0*/ + { /*i1=0*/ {1}, {2} }, { /*i1=1*/ {3}, {4} } }, { /*i0=1*/ { /*i1=0*/ + {5}, {6} }, { /*i1=1*/ {7}, {8} } } }) + ROOT batch-norm-grad = (bf16[2,2,2,1]{3,2,1,0}, bf16[2]{0}, bf16[2]{0}) + batch-norm-grad(constant.4, constant.5, constant.6, constant.7, + constant.8), epsilon=0, feature_index=2 + } + )"; + + auto module = CreateModuleFromHloString(hlo_string); + HloElementTypeConverter type_converter(BF16, F32); + TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); + EXPECT_TRUE(converted); + const HloInstruction* tuple_instr = + module->entry_computation()->root_instruction(); + ::testing::Matcher batch_norm = + op::BatchNormGrad(); + EXPECT_THAT(tuple_instr, + op::Tuple(op::Convert(op::GetTupleElement(batch_norm, 0)), + op::Convert(op::GetTupleElement(batch_norm, 1)), + op::Convert(op::GetTupleElement(batch_norm, 2)))); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc index 173f0e2c42bed2ea461eef27d811e0a626c4fee3..8016b38d15330842c0e11f192587b6035a0dae01 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.cc +++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc @@ -40,9 +40,11 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/core/lib/core/bitmap.h" +#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/lib/gtl/optional.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/protobuf.h" #include "tensorflow/core/platform/types.h" @@ -166,6 +168,34 @@ StatusOr> ElementWiseUnaryOpImpl( return std::move(result); } +// For one particular placement of a window in a base shape (the placement is +// represented as `window_count_index`), iterates inside the window. Translates +// the window index into base index. If the base index is within bound, call `f` +// with the base index. +void IterateThroughWindow( + const Shape& window_shape, const Window& window, const Shape& base_shape, + const tensorflow::gtl::ArraySlice& window_count_index, + const std::function&)>& f) { + const int64 rank = ShapeUtil::Rank(base_shape); + DimensionVector window_index(rank); + std::fill(window_index.begin(), window_index.end(), 0); + do { + std::vector base_index(rank); + bool out_of_bound = false; + for (int64 i = 0; i < rank; ++i) { + base_index[i] = window_count_index[i] * window.dimensions(i).stride() + + window_index[i] - window.dimensions(i).padding_low(); + if (base_index[i] < 0 || base_index[i] >= base_shape.dimensions(i)) { + out_of_bound = true; + break; + } + } + if (!out_of_bound) { + f(base_index); + } + } while (IndexUtil::BumpIndices(window_shape, &window_index)); +} + } // namespace template @@ -261,7 +291,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { parent_->evaluated_[broadcast] = Literal::CreateFromShape(broadcast->shape()); auto output = parent_->evaluated_[broadcast].get(); - auto operand_to_broadcast = + const Literal& operand_to_broadcast = parent_->GetEvaluatedLiteralFor(broadcast->operand(0)); std::vector broadcast_indices( ShapeUtil::Rank(broadcast->operand(0)->shape()), 0); @@ -467,6 +497,27 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { return HandleSign(sign); } + template ::value>::type* = nullptr> + Status HandleAtan2(HloInstruction* atan2) { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[atan2], + ElementWiseBinaryOp(atan2, [](ElementwiseT lhs_elem, + ElementwiseT rhs_elem) { + return std::atan2(lhs_elem, rhs_elem); + })); + return Status::OK(); + } + + template ::value>::type* = nullptr> + Status HandleAtan2(HloInstruction* atan2) { + return InvalidArgument("Unsupported type for Atan2"); + } + + Status HandleAtan2(HloInstruction* atan2) override { + return HandleAtan2(atan2); + } + Status HandleTanh(HloInstruction* tanh) override { TF_ASSIGN_OR_RETURN(parent_->evaluated_[tanh], ElementWiseUnaryOp(tanh, [](ElementwiseT elem_operand) { @@ -743,6 +794,10 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { TF_ASSIGN_OR_RETURN( parent_->evaluated_[shr], ElementWiseBinaryOp(shr, [](NativeT lhs_elem, NativeT rhs_elem) { + // If shift amount is greater than the number of bits, then return 0. + if (rhs_elem >= sizeof(UnsignedT) * CHAR_BIT) { + return static_cast(0); + } return static_cast(static_cast(lhs_elem) >> rhs_elem); })); @@ -815,7 +870,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { << " but is inferred to be: " << ShapeUtil::HumanString(inferred_return_shape); - auto operand_literal = parent_->GetEvaluatedLiteralFor(operand); + const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); auto result = Literal::CreateFromShape(result_shape); TF_RETURN_IF_ERROR(result->Populate( @@ -924,14 +979,21 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { out_index[output_spatial_dim] * window_dim.stride() - window_dim.padding_low() + rhs_spatial_index[ki] * window_dim.window_dilation(); - // Skip if the lhs (input) index is to be dilated. - if (undilated_index % window_dim.base_dilation() != 0) { + // Skip if the lhs (input) index is to be dilated. As an + // optimization, skip this mod if there's no dilation. + if (window_dim.base_dilation() > 1 && + undilated_index % window_dim.base_dilation() != 0) { goto cnt; } - // Calculate the actual lhs (input) index after dilation. - lhs_index[input_spatial_dim] = - undilated_index / window_dim.base_dilation(); + // Calculate the actual lhs (input) index after dilation. As an + // optimization, skip this integer divide if there's no dilation. + if (window_dim.base_dilation() > 1) { + lhs_index[input_spatial_dim] = + undilated_index / window_dim.base_dilation(); + } else { + lhs_index[input_spatial_dim] = undilated_index; + } // Skip if input index is not in bound. if (!(lhs_index[input_spatial_dim] >= 0 && @@ -970,55 +1032,119 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { CHECK(ShapeUtil::IsArray(lhs->shape())); CHECK(ShapeUtil::IsArray(rhs->shape())); - // Dot only supports operands of rank 1 and 2. - const auto dot_rank = ShapeUtil::Rank(dot->shape()); + const auto& dnums = dot->dot_dimension_numbers(); + const auto lhs_rank = ShapeUtil::Rank(lhs->shape()); const auto rhs_rank = ShapeUtil::Rank(rhs->shape()); - CHECK(lhs_rank > 0 && lhs_rank <= 2); - CHECK(rhs_rank > 0 && rhs_rank <= 2); - CHECK_EQ(dot_rank, lhs_rank + rhs_rank - 2); CHECK(ShapeUtil::SameElementType(lhs->shape(), rhs->shape())); CHECK(ShapeUtil::SameElementType(lhs->shape(), dot->shape())); - // Check contracted dimensions are the same. - // - // Determine the index of the contracted dimensions for input tensors. - // dimensions -1 of lhs and dimension 0 of rhs are contracted. - const int64 lhs_contracted_dimension = - ShapeUtil::GetDimensionNumber(lhs->shape(), -1); - const int64 rhs_contracted_dimension = 0; - CHECK_EQ(lhs->shape().dimensions(lhs_contracted_dimension), - rhs->shape().dimensions(rhs_contracted_dimension)) + // There must be 1 and only 1 Contracting dimension for lhs and rhs. + CHECK_EQ(dnums.lhs_contracting_dimensions_size(), 1); + CHECK_EQ(dnums.rhs_contracting_dimensions_size(), 1); + const int64 lhs_contracting_dimension = dnums.lhs_contracting_dimensions(0); + const int64 rhs_contracting_dimension = dnums.rhs_contracting_dimensions(0); + // Contracted dimension sizes must be the same. + CHECK_EQ(lhs->shape().dimensions(lhs_contracting_dimension), + rhs->shape().dimensions(rhs_contracting_dimension)) << "lhs contracted dimension: " - << lhs->shape().dimensions(lhs_contracted_dimension) + << lhs->shape().dimensions(lhs_contracting_dimension) << " rhs contracted dimension: " - << rhs->shape().dimensions(rhs_contracted_dimension); + << rhs->shape().dimensions(rhs_contracting_dimension); const int64 contracted_dimension_size = - lhs->shape().dimensions(lhs_contracted_dimension); + lhs->shape().dimensions(lhs_contracting_dimension); const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); auto result = Literal::CreateFromShape(dot->shape()); + + CHECK_EQ(dnums.lhs_batch_dimensions_size(), + dnums.rhs_batch_dimensions_size()); + + std::vector lhs_non_contracting_dims; + for (int64 i = 0; i < lhs_rank; i++) { + if (i != lhs_contracting_dimension) { + lhs_non_contracting_dims.push_back(i); + } + } + + std::vector rhs_non_batch_non_contracting_dims; + tensorflow::gtl::FlatSet batch_dims_set( + dnums.rhs_batch_dimensions().begin(), + dnums.rhs_batch_dimensions().end()); + for (int64 i = 0; i < rhs_rank; i++) { + if (i != rhs_contracting_dimension && batch_dims_set.count(i) == 0) { + rhs_non_batch_non_contracting_dims.push_back(i); + } + } + + const int64 batch_dim_size = dnums.lhs_batch_dimensions_size(); + const int64 lhs_non_contracting_size = lhs_non_contracting_dims.size(); + + DimensionVector lhs_index(lhs_rank); + DimensionVector rhs_index(rhs_rank); TF_RETURN_IF_ERROR(result->Populate( - [&](tensorflow::gtl::ArraySlice multi_index) { + [&](tensorflow::gtl::ArraySlice result_index) { ElementwiseT result_val = static_cast(0); - std::vector lhs_index(lhs_rank, 0); - std::vector rhs_index(rhs_rank, 0); - // Set index for non-contracted dimension for lhs and rhs. - if (lhs_rank > 1) { - lhs_index[0] = multi_index[0]; + // Find the corresponding non-contracting indices for lhs and rhs. + // + // For `result_index`, its batch dimension, if exists, will be at the + // same dimension as the batch dimension of lhs and rhs. More + // specifically: + // - For lhs, the non-contracting dimensions, including the batch + // dimension have the same index as the `result_index`. + // - For rhs, the batch dimension is set seperately from other + // non-contracting dimensions, since these other non-contracting + // dimensions in rhs follow the non-contracting dimensions of lhs in + // the resulting index. + // + // As an example, for a resulting index: + // result_index [result_batch, result_x, result_y] + // the effecting lhs and rhs indices are: + // lhs [result_batch, lhs_non_contracting_dim, contracting_dim + // rhs [result_batch, contracting_dim, rhs_non_contracting_dim] + // `result_x` is only affected by the lhs_non_contracting_dim and + // likewise `result_y` only depends on rhs_non_contracting_dim. + // + // so we can look up the lhs and rhs indices by: + // + // lhs: + // batch index is the same as `result_batch`. + // non-contracting dimension is the same as + // result_index[lhs_non_contracting_dim] + // rhs: + // batch index: the same as `result_batch`. + // non-contracting dimension index: *not* the same as + // result_index[rhs_non_contractng_dim], since the + // non-contracting dimensions of lhs are included in the + // result_index first. Instead, the non_contracting_dim of rhs must + // be calculated as following: + // lhs_non_contracting_dimensions_size + + // (rhs_non_batch_non_contracting_dim - batch_dim_size) - 1 + // + // Note that (rhs_non_batch_contracting_dim - batch_dim_size) is + // the index offset to the result_index that only depends on + // the non_batch and non-contracting dimensions of rhs. -1 at the + // end translates size to index. + for (auto i : lhs_non_contracting_dims) { + lhs_index[i] = result_index[i]; + } + for (auto i : dnums.rhs_batch_dimensions()) { + rhs_index[i] = result_index[i]; } - if (rhs_rank > 1) { - rhs_index[1] = multi_index[multi_index.size() - 1]; + for (auto i : rhs_non_batch_non_contracting_dims) { + const int64 rhs_non_batch_non_contracting_dim = + lhs_non_contracting_size + (i - batch_dim_size) - 1; + rhs_index[i] = result_index[rhs_non_batch_non_contracting_dim]; } // Accumulates resulting product along the contracted dimension. for (int64 i = 0; i < contracted_dimension_size; ++i) { - lhs_index[lhs_contracted_dimension] = i; - rhs_index[rhs_contracted_dimension] = i; + lhs_index[lhs_contracting_dimension] = i; + rhs_index[rhs_contracting_dimension] = i; result_val += static_cast(lhs_literal.Get(lhs_index)) * @@ -1058,7 +1184,8 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { return scalar; })); - auto evaluated_operand = parent_->GetEvaluatedLiteralFor(pad->operand(0)); + const Literal& evaluated_operand = + parent_->GetEvaluatedLiteralFor(pad->operand(0)); std::vector input_index(ShapeUtil::Rank(evaluated_operand.shape()), 0); @@ -1280,6 +1407,11 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl(map)); break; } + case F16: { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], + MapImpl(map)); + break; + } case F32: { TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl(map)); break; @@ -1391,6 +1523,111 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } + Status HandleSelectAndScatter(HloInstruction* select_and_scatter) override { + auto operand = select_and_scatter->operand(0); + auto source = select_and_scatter->operand(1); + const Window& window = select_and_scatter->window(); + + const Literal& init_literal = + parent_->GetEvaluatedLiteralFor(select_and_scatter->operand(2)); + TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape())); + auto init_scalar = init_literal.Get({}); + + auto result = Literal::CreateFromShape(select_and_scatter->shape()); + + // Initialize result array with the init value. + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice output_index) { + return init_scalar; + })); + + std::vector window_dimension_sizes; + for (const auto& window_dimension : window.dimensions()) { + window_dimension_sizes.push_back(window_dimension.size()); + } + const Shape window_shape = ShapeUtil::MakeShape( + operand->shape().element_type(), window_dimension_sizes); + + HloComputation* select = select_and_scatter->select(); + HloComputation* scatter = select_and_scatter->scatter(); + + const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); + const Literal& source_literal = parent_->GetEvaluatedLiteralFor(source); + + int64 rank = ShapeUtil::Rank(operand_literal.shape()); + + HloEvaluator embedded_evaluator; + DimensionVector source_index(rank); + + std::fill(source_index.begin(), source_index.end(), 0); + do { + // For each element in `source`, we place a window in `operand`. For each + // window placement, we iterate inside the window twice: + // + // 1. Find the selected index by applying `select` function to all + // elements. E.g., If the `select` function is GreaterEqual, the first + // iteration through the window finds the biggest value and returns its + // index. + // + // 2. Using the selected index, scatter value from `source` to result. We + // do this by iterating through the window, and compare each index with + // the selected index. + tensorflow::gtl::optional selected_val; + tensorflow::gtl::optional> selected_index; + + IterateThroughWindow( + window_shape, window, operand_literal.shape(), source_index, + [&](const std::vector& operand_index) { + auto curr_val = operand_literal.Get(operand_index); + if (!selected_val) { + selected_val = curr_val; + selected_index = operand_index; + } + const auto curr_val_literal = Literal::CreateR0(curr_val); + const auto selected_val_literal = + Literal::CreateR0(*selected_val); + + const std::vector args = { + curr_val_literal.get(), selected_val_literal.get()}; + std::unique_ptr computed_result = + embedded_evaluator.Evaluate(*select, args) + .ConsumeValueOrDie(); + bool selected = computed_result->Get({}); + if (selected) { + selected_val = curr_val; + selected_index = operand_index; + } + embedded_evaluator.ResetVisitStates(); + }); + + IterateThroughWindow( + window_shape, window, operand_literal.shape(), source_index, + [&](const std::vector& operand_index) { + if (std::equal(operand_index.begin(), operand_index.end(), + selected_index->begin())) { + auto source = source_literal.Get(source_index); + auto scattered = result->Get(operand_index); + const auto source_literal = Literal::CreateR0(source); + const auto scattered_literal = + Literal::CreateR0(scattered); + + const std::vector args = { + source_literal.get(), scattered_literal.get()}; + std::unique_ptr computed_result = + embedded_evaluator.Evaluate(*scatter, args) + .ConsumeValueOrDie(); + result->Set(operand_index, computed_result->Get({})); + // Clear visit states so that the we can use the evaluator again + // on the same computation. + embedded_evaluator.ResetVisitStates(); + } + }); + } while (IndexUtil::BumpIndices(source->shape(), &source_index)); + + parent_->evaluated_[select_and_scatter] = std::move(result); + return Status::OK(); + } + Status HandleReduceWindow(HloInstruction* reduce_window) override { auto operand = reduce_window->operand(0); const Window& window = reduce_window->window(); @@ -1439,39 +1676,28 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { std::fill(window_index.begin(), window_index.end(), 0); std::fill(operand_index.begin(), operand_index.end(), 0); - do { - // Set curr_val to 0 if out of bound (padded). - ReturnT curr_val = static_cast(0); - bool out_of_bound = false; - for (int i = 0; i < operand_index.size(); ++i) { - operand_index[i] = - output_index[i] * window.dimensions(i).stride() + - window_index[i] - window.dimensions(i).padding_low(); - if (operand_index[i] < 0 || - operand_index[i] >= operand_literal.shape().dimensions(i)) { - out_of_bound = true; - break; - } - } - if (!out_of_bound) { - curr_val = operand_literal.Get(operand_index); - } - // Evaluate computation with specified literal operands. - const auto curr_val_literal = Literal::CreateR0(curr_val); - const auto result_val_literal = - Literal::CreateR0(result_val); - const std::vector args = {curr_val_literal.get(), - result_val_literal.get()}; - std::unique_ptr computed_result = - embedded_evaluator.Evaluate(*function, args) - .ConsumeValueOrDie(); - - // Clear visit states so that the we can use the evaluate again on - // the same computation. - embedded_evaluator.ResetVisitStates(); - - result_val = computed_result->Get({}); - } while (IndexUtil::BumpIndices(window_shape, &window_index)); + IterateThroughWindow( + window_shape, window, operand_literal.shape(), output_index, + [&](const std::vector& operand_index) { + auto curr_val = operand_literal.Get(operand_index); + + // Evaluate computation with specified literal operands. + const auto curr_val_literal = + Literal::CreateR0(curr_val); + const auto result_val_literal = + Literal::CreateR0(result_val); + const std::vector args = { + curr_val_literal.get(), result_val_literal.get()}; + std::unique_ptr computed_result = + embedded_evaluator.Evaluate(*function, args) + .ConsumeValueOrDie(); + + // Clear visit states so that the we can use the evaluate again + // on the same computation. + embedded_evaluator.ResetVisitStates(); + + result_val = computed_result->Get({}); + }); return result_val; })); @@ -1493,7 +1719,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { << ShapeUtil::HumanString(inferred_return_shape); const int64 rank = ShapeUtil::Rank(operand->shape()); - auto operand_literal = parent_->GetEvaluatedLiteralFor(operand); + const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); auto func = [&](tensorflow::gtl::ArraySlice out_index) { DimensionVector operand_index(rank); for (int64 i = 0; i < rank; ++i) { @@ -1554,13 +1780,121 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { return HandleCos(cos); } + template ::value>::type* = nullptr> + Status HandleReducePrecision(HloInstruction* reduce_precision) { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[reduce_precision], + ElementWiseUnaryOp(reduce_precision, [reduce_precision]( + ElementwiseT elem) { + uint32_t value_as_int = tensorflow::bit_cast(elem); + const uint32_t mantissa_bits = reduce_precision->mantissa_bits(); + const uint32_t exponent_bits = reduce_precision->exponent_bits(); + + // Code is based on the CPU/GPU implementation in LLVM-emitting code. + // + // Bits in float type: + // mantissa : bits [0:22] + // exponent : bits [23:30] + // sign : bits [31] + if (mantissa_bits < 23) { + const uint32_t last_mantissa_bit_mask = 1u << (23 - mantissa_bits); + + // Compute rounding bias for round-to-nearest with ties to even. + // This is equal to a base value of 0111... plus one bit if the last + // remaining mantissa bit is 1. + const uint32_t base_rounding_bias = + (last_mantissa_bit_mask >> 1) - 1; + const uint32_t x_last_mantissa_bit = + (value_as_int & last_mantissa_bit_mask) >> (23 - mantissa_bits); + const uint32_t x_rounding_bias = + x_last_mantissa_bit + base_rounding_bias; + + // Add rounding bias, and mask out truncated bits. Note that the + // case where adding the rounding bias overflows into the exponent + // bits is correct; the non-masked mantissa bits will all be zero, + // and the exponent will be incremented by one. + const uint32_t truncation_mask = ~(last_mantissa_bit_mask - 1); + value_as_int = value_as_int + x_rounding_bias; + value_as_int = value_as_int & truncation_mask; + } + if (exponent_bits < 8) { + // Masks for f32 values. + const uint32_t f32_sign_bit_mask = 1u << 31; + const uint32_t f32_exp_bits_mask = 0xffu << 23; + + // An exponent of 2^(n-1)-1 -- that is, 0111... with the zero in the + // most- significant bit -- is equal to 1.0f for all exponent sizes. + // Adding 2^(n-1)-1 to this gives us the highest non-infinite + // exponent for a bit- size of n, and subtracting 2^(n-1)-1 from + // this gives us the lowest' exponent (corresponding to 0.0f). + // + // Thus, the f32 exponent corresponding to the highest non-infinite + // exponent for a bit size of n is (2^7-1) + 2^(n-1)-1, and the f32 + // exponent corresponding to the lowest exponent for a bit size of n + // is (2^7-1) - 2^(n-1)-1. + // + // Note that we have already checked that exponents_bits >= 1. + const uint32_t f32_exponent_bias = (1 << 7) - 1; + const uint32_t reduced_exponent_bias = + (1 << (exponent_bits - 1)) - 1; + const uint32_t reduced_max_exponent = + f32_exponent_bias + reduced_exponent_bias; + const uint32_t reduced_min_exponent = + f32_exponent_bias - reduced_exponent_bias; + + // Do we overflow or underflow? + const uint32_t x_exponent = value_as_int & f32_exp_bits_mask; + const bool x_overflows = x_exponent > (reduced_max_exponent << 23); + const bool x_underflows = + x_exponent <= (reduced_min_exponent << 23); + + // Compute appropriately-signed values of zero and infinity. + const uint32_t x_signed_zero = value_as_int & f32_sign_bit_mask; + const uint32_t x_signed_inf = x_signed_zero | f32_exp_bits_mask; + + // Force to zero or infinity if overflow or underflow. (Note that + // this truncates all denormal values to zero, rather than rounding + // them.) + value_as_int = x_overflows ? x_signed_inf : value_as_int; + value_as_int = x_underflows ? x_signed_zero : value_as_int; + } + + float reduced_result = tensorflow::bit_cast(value_as_int); + if (std::isnan(elem)) { + reduced_result = mantissa_bits > 0 + ? elem + : std::numeric_limits::infinity(); + } + return reduced_result; + })); + return Status::OK(); + } + + template ::value>::type* = nullptr> + Status HandleReducePrecision(HloInstruction* reduce_precision) { + return InvalidArgument("Double not supported for reduce precision"); + } + + template < + typename NativeT, + typename std::enable_if::value || + is_complex_t::value>::type* = nullptr> + Status HandleReducePrecision(HloInstruction* reduce_precision) { + return InvalidArgument("Unsupported type for reduce precision"); + } + + Status HandleReducePrecision(HloInstruction* reduce_precision) override { + return HandleReducePrecision(reduce_precision); + } + private: template StatusOr> DynamicSlice( const Literal& operand_literal, const Literal& start_indices_literal, const Shape& result_shape) { - const auto& start_indices_typed = - start_indices_literal.GetArraySlice(); + auto start_indices_typed = start_indices_literal.data(); std::vector start(start_indices_typed.begin(), start_indices_typed.end()); @@ -1588,12 +1922,11 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { StatusOr> DynamicUpdateSlice( const Literal& operand_literal, const Literal& update_literal, const Literal& start_indices_literal) { - const auto& start_indices_typed = - start_indices_literal.GetArraySlice(); + auto start_indices_typed = start_indices_literal.data(); const std::vector start(start_indices_typed.begin(), start_indices_typed.end()); - auto result = MakeUnique(operand_literal); + auto result = operand_literal.CloneToUnique(); std::vector result_index(ShapeUtil::Rank(result->shape()), 0); auto func = [&](const std::vector& update_index) { @@ -1717,9 +2050,7 @@ HloEvaluator::HloEvaluator() { }); typed_visitors_[S32] = MakeUnique>(this); typed_visitors_[S64] = MakeUnique>(this); - typed_visitors_[F16] = MakeUnique([](HloInstruction*) { - return Unimplemented("HloEvaluator: unhandled primitive type: F16."); - }); + typed_visitors_[F16] = MakeUnique>(this); typed_visitors_[F32] = MakeUnique>(this); typed_visitors_[F64] = MakeUnique>(this); typed_visitors_[C64] = MakeUnique>(this); @@ -1751,8 +2082,8 @@ StatusOr> HloEvaluator::Evaluate( TF_RETURN_IF_ERROR(module.entry_computation()->Accept(this)); - return MakeUnique( - GetEvaluatedLiteralFor(module.entry_computation()->root_instruction())); + return GetEvaluatedLiteralFor(module.entry_computation()->root_instruction()) + .CloneToUnique(); } template @@ -1769,8 +2100,7 @@ StatusOr> HloEvaluator::Evaluate( } TF_RETURN_IF_ERROR(computation.Accept(this)); - return MakeUnique( - GetEvaluatedLiteralFor(computation.root_instruction())); + return GetEvaluatedLiteralFor(computation.root_instruction()).CloneToUnique(); } template @@ -1795,14 +2125,14 @@ StatusOr> HloEvaluator::Evaluate( << input_literal->ToString(); TF_RET_CHECK(ShapeUtil::Equal(operand->shape(), input_literal->shape())); - evaluated_[operand] = MakeUnique(*input_literal); + evaluated_[operand] = input_literal->CloneToUnique(); } } TF_RETURN_IF_ERROR(Preprocess(instruction)); TF_RETURN_IF_ERROR(instruction->Visit(this)); TF_RETURN_IF_ERROR(Postprocess(instruction)); - return MakeUnique(GetEvaluatedLiteralFor(instruction)); + return GetEvaluatedLiteralFor(instruction).CloneToUnique(); } StatusOr> HloEvaluator::Evaluate( @@ -1823,7 +2153,7 @@ StatusOr> HloEvaluator::Evaluate( TF_RETURN_IF_ERROR(Preprocess(instruction)); TF_RETURN_IF_ERROR(instruction->Visit(this)); TF_RETURN_IF_ERROR(Postprocess(instruction)); - return MakeUnique(GetEvaluatedLiteralFor(instruction)); + return GetEvaluatedLiteralFor(instruction).CloneToUnique(); } std::unique_ptr HloEvaluator::TryEvaluate( @@ -1880,7 +2210,7 @@ Status HloEvaluator::HandleParameter(HloInstruction* parameter) { << ", but input literal shape is: " << ShapeUtil::HumanString(input_literal->shape()); - evaluated_[parameter] = MakeUnique(*input_literal); + evaluated_[parameter] = input_literal->CloneToUnique(); return Status::OK(); } @@ -1931,7 +2261,7 @@ Status HloEvaluator::HandleConcatenate(HloInstruction* concatenate) { for (auto operand : operands) { const Shape& operand_shape = operand->shape(); - TF_RETURN_IF_ERROR(result_literal->Copy( + TF_RETURN_IF_ERROR(result_literal->CopySliceFrom( GetEvaluatedLiteralFor(operand), source_indices, dest_indices, AsInt64Slice(operand_shape.dimensions()))); dest_indices[concat_dim] += @@ -2089,16 +2419,17 @@ Status HloEvaluator::HandleGetTupleElement(HloInstruction* get_tuple_element) { const Literal& operand_tuple_literal = GetEvaluatedLiteralFor(operand); - evaluated_[get_tuple_element] = - MakeUnique(operand_tuple_literal.tuple_literals(index)); - - return Status::OK(); + evaluated_[get_tuple_element] = MakeUnique( + ShapeUtil::GetTupleElementShape(operand->shape(), index)); + return evaluated_[get_tuple_element]->CopyFrom(operand_tuple_literal, + /*dest_shape_index=*/{}, + /*src_shape_index=*/{index}); } Status HloEvaluator::HandleCopy(HloInstruction* copy) { TF_RET_CHECK(ShapeUtil::Compatible(copy->shape(), copy->operand(0)->shape())); - auto result = MakeUnique(GetEvaluatedLiteralFor(copy->operand(0))); + auto result = GetEvaluatedLiteralFor(copy->operand(0)).CloneToUnique(); evaluated_[copy] = std::move(result); return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.h b/tensorflow/compiler/xla/service/hlo_evaluator.h index 02bb8b0a47065c359603a113f49626bf3ad344d8..3b2b697e492a78a06a4e5ae6bf056ff8676f2ff5 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.h +++ b/tensorflow/compiler/xla/service/hlo_evaluator.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ #include @@ -195,4 +195,4 @@ class HloEvaluator : public DfsHloVisitorWithDefault { } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ diff --git a/tensorflow/compiler/xla/service/hlo_execution_profile.cc b/tensorflow/compiler/xla/service/hlo_execution_profile.cc index 0111cfd5a3d7889f80370f9e3e744457bc4091e4..f0df93b61d29c1535d8a89fbd65e669de5b43729 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile.cc +++ b/tensorflow/compiler/xla/service/hlo_execution_profile.cc @@ -32,7 +32,7 @@ HloProfileIndexMap::HloProfileIndexMap(const HloModule& module) { InsertOrDie(&computation_to_profile_idx_, computation, current_profile_index++); for (const HloInstruction* instruction : computation->instructions()) { - // For simplicity we track all instrutions here, but we could skip + // For simplicity we track all instructions here, but we could skip // non-executing instructions like constants and parameters. InsertOrDie(&instruction_to_profile_idx_, instruction, current_profile_index++); @@ -40,83 +40,75 @@ HloProfileIndexMap::HloProfileIndexMap(const HloModule& module) { } } -std::unique_ptr CreateHloProfilePrinter( +std::unique_ptr CreateHloProfilePrinterData( const HloProfileIndexMap& hlo_profile_index_map, const HloCostAnalysis& cost_analysis) { - using HloComputationInfo = HloProfilePrinter::HloComputationInfo; - using HloInstructionInfo = HloProfilePrinter::HloInstructionInfo; - - HloComputationInfo* computation_infos = - new HloComputationInfo[hlo_profile_index_map.computation_count()]; - - // There are two "indices" in play here. The first one is the index of the - // HloComputationInfo or HloInstructionInfo in the array that contains said - // HloComputationInfo or HloInstructionInfo. The second index is the index of - // the HloComputationInfo or HloInstructionInfo in the profile counters array, - // as decided by hlo_profile_index_map. The latter index is always referred - // to as "profile_index". - - size_t computation_index_in_static_data = 0; - size_t max_profile_index = hlo_profile_index_map.total_count(); - for (const auto& pair : hlo_profile_index_map.computation_to_profile_idx()) { - CHECK_LT(pair.second, max_profile_index); + using HloComputationInfo = HloProfilePrinterData::HloComputationInfo; + using HloInstructionInfo = HloProfilePrinterData::HloInstructionInfo; + + size_t profile_counters_size = hlo_profile_index_map.total_count(); + + std::unique_ptr profile_printer_data = + MakeUnique(); + profile_printer_data->set_profile_counters_size(profile_counters_size); + profile_printer_data->mutable_computation_infos()->Reserve( + hlo_profile_index_map.computation_count()); + + const auto& computation_to_profile_idx_map = + hlo_profile_index_map.computation_to_profile_idx(); + + // computation_to_profile_idx_map's order is not deterministic so create a + // deterministic computation_and_profile_idx_list so that we end up with a + // deterministic HloProfilePrinterData protobuf. + + std::vector> + computation_and_profile_idx_list(computation_to_profile_idx_map.begin(), + computation_to_profile_idx_map.end()); + + // The profile indices were computed deterministically in + // HloProfileIndexMap::HloProfileIndexMap. + c_sort(computation_and_profile_idx_list, + [](const std::pair& left, + const std::pair& right) { + return left.second < right.second; + }); + + for (const auto& pair : computation_and_profile_idx_list) { + CHECK_LT(pair.second, profile_counters_size); const HloComputation* computation = pair.first; - size_t current_computation_index = computation_index_in_static_data++; HloComputationInfo* computation_info = - &computation_infos[current_computation_index]; + profile_printer_data->add_computation_infos(); - computation_info->name = strdup(computation->name().c_str()); - computation_info->profile_index = pair.second; - computation_info->instructions = - new HloInstructionInfo[computation->instruction_count()]; - computation_info->instructions_size = computation->instruction_count(); + computation_info->set_name(computation->name()); + computation_info->set_profile_index(pair.second); + computation_info->mutable_instruction_infos()->Reserve( + computation->instruction_count()); - size_t instruction_index_in_static_data = 0; for (const HloInstruction* hlo : computation->instructions()) { - HloProfilePrinter::HloInstructionInfo* instruction_info = - &computation_info->instructions[instruction_index_in_static_data++]; - instruction_info->long_name = strdup(hlo->ToString().c_str()); - instruction_info->short_name = strdup( - hlo->ToString(HloPrintOptions().set_compact_operands(true)).c_str()); - instruction_info->category = strdup(hlo->ToCategory().c_str()); - instruction_info->flop_count = cost_analysis.flop_count(*hlo); - instruction_info->transcendental_count = - cost_analysis.transcendental_count(*hlo); - instruction_info->bytes_accessed = cost_analysis.bytes_accessed(*hlo); - instruction_info->optimal_seconds = cost_analysis.optimal_seconds(*hlo); - instruction_info->profile_index = - hlo_profile_index_map.GetProfileIndexFor(*hlo); - CHECK_LT(instruction_info->profile_index, max_profile_index); + HloInstructionInfo* instruction_info = + computation_info->add_instruction_infos(); + instruction_info->set_long_name(hlo->ToString()); + instruction_info->set_short_name( + hlo->ToString(HloPrintOptions().set_compact_operands(true))); + instruction_info->set_category(hlo->ToCategory()); + instruction_info->set_flop_count(cost_analysis.flop_count(*hlo)); + instruction_info->set_transcendental_count( + cost_analysis.transcendental_count(*hlo)); + instruction_info->set_bytes_accessed(cost_analysis.bytes_accessed(*hlo)); + instruction_info->set_optimal_seconds( + cost_analysis.optimal_seconds(*hlo)); + instruction_info->set_profile_index( + hlo_profile_index_map.GetProfileIndexFor(*hlo)); } } - auto deleter = [](HloProfilePrinter::HloComputationInfo* computation_infos, - int64 computation_infos_size) { - for (int64 i = 0; i < computation_infos_size; i++) { - HloInstructionInfo* instruction_infos = computation_infos[i].instructions; - for (int64 j = 0; j < computation_infos[i].instructions_size; j++) { - // We can't make instruction_infos[j].long_name etc. non-const pointers - // since they may point into static storage, so we have a const_cast - // here. - free(const_cast(instruction_infos[j].long_name)); - free(const_cast(instruction_infos[j].short_name)); - free(const_cast(instruction_infos[j].category)); - } - delete[] instruction_infos; - free(const_cast(computation_infos[i].name)); - } - delete[] computation_infos; - }; - - return MakeUnique( - computation_infos, hlo_profile_index_map.computation_count(), - /*profile_counters_size=*/max_profile_index, deleter); + return profile_printer_data; } HloExecutionProfile::HloExecutionProfile( - const HloProfilePrinter* hlo_profile_printer, + const HloProfilePrinterData* hlo_profile_printer_data, const HloProfileIndexMap* hlo_profile_index_map) - : hlo_profile_printer_(*hlo_profile_printer), + : hlo_profile_printer_data_(*hlo_profile_printer_data), hlo_profile_index_map_(*hlo_profile_index_map), profile_counters_( /*count*/ hlo_profile_index_map_.total_count(), diff --git a/tensorflow/compiler/xla/service/hlo_execution_profile.h b/tensorflow/compiler/xla/service/hlo_execution_profile.h index 470fd4ce3c205d84152238f4b18daad77e403f68..6fb91b9bef9d1df82b8806ce79cc147823edeb3d 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile.h +++ b/tensorflow/compiler/xla/service/hlo_execution_profile.h @@ -77,8 +77,8 @@ class HloProfileIndexMap { std::unordered_map computation_to_profile_idx_; }; -// Create an instance of `HloProfilePrinter` that owns its memory. -std::unique_ptr CreateHloProfilePrinter( +// Create an instance of `HloProfilePrinterData`. +std::unique_ptr CreateHloProfilePrinterData( const HloProfileIndexMap& hlo_profile_index_map, const HloCostAnalysis& cost_analysis); @@ -90,7 +90,7 @@ class HloExecutionProfile { public: using DeviceDescription = perftools::gputools::DeviceDescription; - HloExecutionProfile(const HloProfilePrinter* hlo_profile_printer, + HloExecutionProfile(const HloProfilePrinterData* hlo_profile_printer_data, const HloProfileIndexMap* hlo_profile_index_map); // Record how many cycles this HLO took to execute. @@ -117,17 +117,19 @@ class HloExecutionProfile { // debugging; e.g. emits cycle counts, execution time at the nominal device // frequency, and the effective throughput given the provided cost_analysis // for the operations in a given computation. Returns an empty string if it - // wasn't possible to generate a printable version. cost_analysis should be a - // clean analysis that can be used to visit the computation. + // wasn't possible to generate a printable version. string ToString(const DeviceDescription& device_description) const { - return hlo_profile_printer_.ToString(profile_counters_.data(), - device_description.clock_rate_ghz()); + return PrintHloProfile(hlo_profile_printer_data_, profile_counters_.data(), + device_description.clock_rate_ghz()); } std::vector* mutable_profile_counters() { return &profile_counters_; } + const std::vector& profile_counters() const { + return profile_counters_; + } private: - const HloProfilePrinter& hlo_profile_printer_; + const HloProfilePrinterData& hlo_profile_printer_data_; const HloProfileIndexMap& hlo_profile_index_map_; // Stores per-Hlo profile counters. This is the only thing that changes when diff --git a/tensorflow/compiler/xla/service/hlo_execution_profile_test.cc b/tensorflow/compiler/xla/service/hlo_execution_profile_test.cc index b1e6729e2bccad4bdbe075a635d8a9b1ede6fecb..a0cb28246d3be541e798e85552436f64a3521f22 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile_test.cc +++ b/tensorflow/compiler/xla/service/hlo_execution_profile_test.cc @@ -73,8 +73,8 @@ TEST_F(HloExecutionProfileTest, Basic) { HloCostAnalysis cost_analysis(shape_size_function); HloProfileIndexMap profile_index_map(*hlo_module); - std::unique_ptr profile_printer = - CreateHloProfilePrinter(profile_index_map, cost_analysis); + std::unique_ptr profile_printer = + CreateHloProfilePrinterData(profile_index_map, cost_analysis); HloExecutionProfile execution_profile(profile_printer.get(), &profile_index_map); diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc index b9f76531f36bbe757f023219023775b125c1e7a5..44fcd36370dcd0cf77601aa1cd2b92810947bd5f 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc @@ -34,6 +34,7 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/map_util.h" #include "tensorflow/core/lib/gtl/optional.h" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/lib/strings/numbers.h" @@ -508,8 +509,17 @@ stylesheet=" // The "to_node" value may be a NULL, indicating that this points to the // "root" tag rather than a normal node. - int64 from_node_id = node_ids_.at(from_node); - int64 to_node_id = to_node ? node_ids_.at(to_node) : root_node_id_; + int64 from_node_id = + tensorflow::gtl::FindWithDefault(node_ids_, from_node, -1); + if (from_node_id == -1) { + LOG(FATAL) << from_node->name() << " was added to edges but not to nodes"; + } + int64 to_node_id = + to_node ? tensorflow::gtl::FindWithDefault(node_ids_, to_node, -1) + : root_node_id_; + if (to_node != nullptr && to_node_id == -1) { + LOG(FATAL) << to_node->name() << " was added to edges but not to nodes"; + } add_hover_css_rule("node", from_node_id, kBlue); add_hover_css_rule("node", to_node_id, kRed); @@ -653,12 +663,15 @@ string HloDotDumper::DumpComputation(const HloComputation* comp) { string HloDotDumper::DumpRootTag() { const HloInstruction* from = GetNodeForEdge(computation_->root_instruction()); - auto from_id = InstructionId(from); - if (!filter_.Show(from)) { + // We didn't display constants as separate nodes; so if the root is a + // constant, we don't add root tag or edge for it. + if (!filter_.Show(from) || from->opcode() == HloOpcode::kConstant) { return ""; } + auto from_id = InstructionId(from); + // The ID of the root computation is otherwise unused, so it makes a good ID // to use for the root-tag node. However, the edge_ids_ map requires a // HloInstruction* pointer for the 'to' value, so we use a NULL value there @@ -948,6 +961,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { return kGreen; case HloOpcode::kConvolution: case HloOpcode::kDot: + case HloOpcode::kFft: return kDarkBlue; case HloOpcode::kReducePrecision: return kRed; @@ -1049,14 +1063,19 @@ string HloDotDumper::GetInstructionNodeExtraInfo(const HloInstruction* instr) { // node -- there the shape and layout is present in the output node. if (instr->opcode() != HloOpcode::kFusion || !ShouldShowFusionSubcomputation(instr)) { - string instr_shape = ShapeUtil::HumanString(instr->shape()); - - // Show layout of non-tuple shapes with more than one dimension. - if (LayoutUtil::HasLayout(instr->shape()) && - instr->shape().dimensions_size() > 1 && - !ShapeUtil::IsTuple(instr->shape())) { - StrAppend(&instr_shape, "{", - Join(LayoutUtil::MinorToMajor(instr->shape()), ","), "}"); + // Show layout of instructions with more than one dimension. Don't show + // layout on tuples or tensors with just one dimension (which only have one + // possible layout) to avoid visual noise. + bool shape_is_multidim = false; + ShapeUtil::ForEachSubshape(instr->shape(), + [&](const Shape& s, const ShapeIndex&) { + shape_is_multidim |= s.dimensions_size() > 1; + }); + string instr_shape; + if (instr->opcode() != HloOpcode::kTuple && shape_is_multidim) { + instr_shape = ShapeUtil::HumanStringWithLayout(instr->shape()); + } else { + instr_shape = ShapeUtil::HumanString(instr->shape()); } // Some instructions have giant tuples as their shapes, so truncate the @@ -1407,9 +1426,11 @@ void DumpText(const HloModule& module, const string& label, string MaybeDumpHloModule(const HloModule& module, const string& label, const HloExecutionProfile* profile) { - VLOG(2) << "MaybeDumpHloModule called on module " << module.name(); - string graph_url; const DebugOptions& debug_options = module.config().debug_options(); + VLOG(2) << "MaybeDumpHloModule called on module " << module.name() + << " with generate_hlo_graph regex \"" + << debug_options.xla_generate_hlo_graph() << "\""; + string graph_url; if (!debug_options.xla_generate_hlo_graph().empty() && RE2::PartialMatch(module.name(), debug_options.xla_generate_hlo_graph())) { diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc index 8e1531c87f9c6e133e2d6763b046b1d5dcbcd09f..1f00aa41dc783f9e5657f5fa654884a31fae0fe7 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc @@ -117,5 +117,18 @@ TEST(HloGraphDumperTest, NestedFusion) { HasSubstr(inner_sum->name())); } +TEST(HloGraphDumperTest, Constant) { + HloComputation::Builder b("b"); + auto instruction = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(-42))); + instruction->set_name("i_am_a_constant_root_instruction"); + HloModule m(TestName()); + HloComputation* root_computation = m.AddEntryComputation(b.Build()); + string graph = hlo_graph_dumper::DumpGraph( + *root_computation, /*label=*/"an_empty_graph", DebugOptions()); + EXPECT_THAT(graph, HasSubstr("an_empty_graph")); + EXPECT_THAT(graph, Not(HasSubstr("i_am_a_constant_root_instruction"))); +} + } // anonymous namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc index 9805818e4cfc4c212818f70acb44da6856b44eb6..0981f1f4fe57751d5b7059b4b08099385369e4b9 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction.cc @@ -101,7 +101,8 @@ StatusOr> HloInstruction::CreateFromProto( instruction->metadata_ = proto.metadata(); if (proto.has_literal()) { - instruction->literal_ = MakeUnique(proto.literal()); + TF_ASSIGN_OR_RETURN(instruction->literal_, + Literal::CreateFromProto(proto.literal())); } instruction->parameter_number_ = proto.parameter_number(); @@ -144,6 +145,10 @@ StatusOr> HloInstruction::CreateFromProto( instruction->infeed_config_ = proto.infeed_config(); instruction->custom_call_target_ = proto.custom_call_target(); instruction->outfeed_shape_ = proto.outfeed_shape(); + instruction->fft_type_ = proto.fft_type(); + for (int64 fft_len : proto.fft_length()) { + instruction->fft_length_.push_back(fft_len); + } return std::move(instruction); } @@ -162,8 +167,7 @@ StatusOr> HloInstruction::CreateFromProto( auto instruction = WrapUnique(new HloInstruction(HloOpcode::kTrace, ShapeUtil::MakeNil())); instruction->operands_.push_back(operand); - instruction->literal_.reset(new Literal); - instruction->literal_->append_u8s(tag); + instruction->literal_ = Literal::CreateR1U8(tag); return instruction; } @@ -334,6 +338,16 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, return instruction; } +/* static */ std::unique_ptr HloInstruction::CreateFft( + const Shape& shape, HloInstruction* operand, FftType fft_type, + tensorflow::gtl::ArraySlice fft_length) { + auto instruction = WrapUnique(new HloInstruction(HloOpcode::kFft, shape)); + instruction->AppendOperand(operand); + instruction->fft_type_ = fft_type; + instruction->fft_length_.assign(fft_length.begin(), fft_length.end()); + return instruction; +} + /* static */ std::unique_ptr HloInstruction::CreateDot( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, const DotDimensionNumbers& dimension_numbers) { @@ -390,6 +404,9 @@ HloInstruction::CreateCrossReplicaSum( tensorflow::StringPiece outfeed_config) { std::unique_ptr instruction = WrapUnique(new HloInstruction(HloOpcode::kOutfeed, ShapeUtil::MakeNil())); + CHECK(ShapeUtil::Compatible(operand->shape(), shape)) + << "Outfeed shape " << shape << " must be compatible with operand shape " + << operand->shape(); instruction->AppendOperand(operand); instruction->outfeed_config_ = outfeed_config.ToString(); instruction->outfeed_shape_ = shape; @@ -655,6 +672,58 @@ HloInstruction::CreateSelectAndScatter( return instruction; } +/* static */ std::unique_ptr +HloInstruction::CreateBroadcastSequence( + const Shape& output_shape, HloInstruction* operand, + const std::function)>& + adder) { + CHECK(ShapeUtil::IsScalar(operand->shape()) || + ShapeUtil::Rank(operand->shape()) == ShapeUtil::Rank(output_shape)); + Shape broadcast_shape = ShapeUtil::ChangeElementType( + output_shape, operand->shape().element_type()); + // Do explicit broadcast for scalar. + if (ShapeUtil::IsScalar(operand->shape())) { + auto broadcast = + HloInstruction::CreateBroadcast(broadcast_shape, operand, {}); + broadcast->set_metadata(operand->metadata()); + if (operand->has_sharding()) { + broadcast->set_sharding(operand->sharding()); + } + return broadcast; + } + // Do explicit broadcast for degenerate broadcast. + std::vector broadcast_dimensions; + std::vector reshaped_dimensions; + for (int i = 0; i < ShapeUtil::Rank(operand->shape()); i++) { + if (operand->shape().dimensions(i) == output_shape.dimensions(i)) { + broadcast_dimensions.push_back(i); + reshaped_dimensions.push_back(operand->shape().dimensions(i)); + } else { + CHECK_EQ(operand->shape().dimensions(i), 1) + << "An explicit broadcast sequence requires the broadcasted " + "dimensions to be trivial; operand: " + << operand->ToString() << "; output_shape: " << output_shape; + } + } + // Eliminate the size one dimensions. + HloInstruction* reshaped_operand = adder(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(operand->shape().element_type(), + reshaped_dimensions), + operand)); + reshaped_operand->set_metadata(operand->metadata()); + if (operand->has_sharding()) { + reshaped_operand->set_sharding(operand->sharding()); + } + // Broadcast 'reshape' up to the larger size. + auto broadcast = HloInstruction::CreateBroadcast( + broadcast_shape, reshaped_operand, broadcast_dimensions); + broadcast->set_metadata(operand->metadata()); + if (operand->has_sharding()) { + broadcast->set_sharding(operand->sharding()); + } + return broadcast; +} + /* static */ std::unique_ptr HloInstruction::CreatePad( const Shape& shape, HloInstruction* operand, HloInstruction* padding_value, const PaddingConfig& padding_config) { @@ -694,10 +763,23 @@ HloInstruction::CreateSelectAndScatter( return instruction; } +// We put the fusion kind into the instruction's name for transpose-dot fusions, +// since those fusions are really just describing a type of dot rather than +// generating a novel computation. +static string FusionNodeName(HloInstruction::FusionKind fusion_kind) { + switch (fusion_kind) { + case HloInstruction::FusionKind::kTransposeDot: + return "dot_fusion"; + default: + return "fusion"; + } +} + /* static */ std::unique_ptr HloInstruction::CreateFusion( const Shape& shape, FusionKind fusion_kind, HloInstruction* fused_root) { auto instruction = WrapUnique(new HloInstruction(HloOpcode::kFusion, shape)); instruction->fusion_kind_ = fusion_kind; + instruction->name_ = FusionNodeName(fusion_kind); instruction->set_parent(fused_root->parent()); instruction->set_metadata(fused_root->metadata()); instruction->CloneAndFuseInternal(fused_root); @@ -713,23 +795,12 @@ HloInstruction::CreateSelectAndScatter( instruction->AppendOperand(operand); } instruction->fusion_kind_ = fusion_kind; + instruction->name_ = FusionNodeName(fusion_kind); instruction->called_computations_.push_back(fusion_computation); fusion_computation->SetFusionInstruction(instruction.get()); return instruction; } -/* static */ std::unique_ptr -HloInstruction::CreateFusionForBackwardConvolution( - const Shape& shape, FusionKind fusion_kind, const Window& window, - const ConvolutionDimensionNumbers& conv_dnums, HloInstruction* fused_root) { - std::unique_ptr fusion = - CreateFusion(shape, fusion_kind, fused_root); - fusion->window_ = MakeUnique(window); - fusion->convolution_dimension_numbers_ = - MakeUnique(conv_dnums); - return fusion; -} - void HloInstruction::MergeFusionInstruction( HloInstruction* instruction_to_merge) { CHECK_EQ(opcode_, HloOpcode::kFusion); @@ -1167,6 +1238,9 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( clone = CreateDot(shape, new_operands[0], new_operands[1], *dot_dimension_numbers_); break; + case HloOpcode::kFft: + CHECK_EQ(new_operands.size(), 1); + return CreateFft(shape, new_operands[0], fft_type_, fft_length_); case HloOpcode::kCrossReplicaSum: clone = CreateCrossReplicaSum(shape, new_operands); break; @@ -1537,8 +1611,9 @@ bool HloInstruction::HasConstantOperand() const { bool HloInstruction::IdenticalSlowPath( const HloInstruction& other, - std::function - eq_computations) const { + const std::function& + eq_computations, + const std::function& eq_shapes) const { // Perform opcode specific checks. switch (opcode()) { // The result of these instructions only depend upon their opcode and @@ -1586,8 +1661,12 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kTuple: return true; - // These opcodes have complex or special behavior so just return false. case HloOpcode::kFusion: + return fusion_kind() == other.fusion_kind() && + eq_computations(fused_instructions_computation(), + other.fused_instructions_computation()); + + // These opcodes have complex or special behavior so just return false. case HloOpcode::kRng: case HloOpcode::kTrace: case HloOpcode::kWhile: @@ -1597,7 +1676,7 @@ bool HloInstruction::IdenticalSlowPath( return parameter_number() == other.parameter_number() && // Check the shape too because `this` and `other` may be in // different HloComputations. - ShapeUtil::Compatible(shape(), other.shape()); + eq_shapes(shape(), other.shape()); case HloOpcode::kBatchNormTraining: case HloOpcode::kBatchNormInference: @@ -1631,6 +1710,11 @@ bool HloInstruction::IdenticalSlowPath( return protobuf_util::ProtobufEquals(dot_dimension_numbers(), other.dot_dimension_numbers()); + // FFT has various types & lengths. + case HloOpcode::kFft: + return fft_type() == other.fft_type() && + fft_length() == other.fft_length(); + // Reduction results are determined by the reduction dimension and the // reduction computation. case HloOpcode::kReduce: @@ -1648,18 +1732,18 @@ bool HloInstruction::IdenticalSlowPath( protobuf_util::ProtobufEquals(window(), other.window()); case HloOpcode::kReshape: - return ShapeUtil::Compatible(shape(), other.shape()); + return eq_shapes(shape(), other.shape()); // Transpose result is determined by the final shape and the permutation. case HloOpcode::kTranspose: - return ShapeUtil::Compatible(shape(), other.shape()) && + return eq_shapes(shape(), other.shape()) && dimensions() == other.dimensions(); // Remaining instructions with special values. case HloOpcode::kBitcast: - return ShapeUtil::Equal(shape(), other.shape()); + return eq_shapes(shape(), other.shape()); case HloOpcode::kBroadcast: - return ShapeUtil::Compatible(shape(), other.shape()) && + return eq_shapes(shape(), other.shape()) && dimensions() == other.dimensions(); case HloOpcode::kConcatenate: return dimensions() == other.dimensions(); @@ -1673,10 +1757,10 @@ bool HloInstruction::IdenticalSlowPath( slice_limits_ == other.slice_limits_ && slice_strides_ == other.slice_strides_; case HloOpcode::kDynamicSlice: - return ShapeUtil::Compatible(shape(), other.shape()) && + return eq_shapes(shape(), other.shape()) && dynamic_slice_sizes_ == other.dynamic_slice_sizes_; case HloOpcode::kDynamicUpdateSlice: - return ShapeUtil::Compatible(shape(), other.shape()); + return eq_shapes(shape(), other.shape()); case HloOpcode::kCall: case HloOpcode::kMap: return eq_computations(to_apply(), other.to_apply()); @@ -1721,7 +1805,8 @@ void HloInstruction::RemoveUser(HloInstruction* user) { Status HloInstruction::ReplaceUseWith(HloInstruction* user, HloInstruction* new_producer) { - TF_RET_CHECK(ShapeUtil::Compatible(shape(), new_producer->shape())) + TF_RET_CHECK( + ShapeUtil::CompatibleIgnoringFpPrecision(shape(), new_producer->shape())) << "this shape: " << ShapeUtil::HumanString(shape()) << ", replacement shape: " << ShapeUtil::HumanString(new_producer->shape()); @@ -1744,8 +1829,8 @@ Status HloInstruction::ReplaceOperandWith(int64 operand_num, TF_RET_CHECK(operand_num >= 0); TF_RET_CHECK(operand_num < operand_count()); HloInstruction* old_operand = mutable_operand(operand_num); - TF_RET_CHECK( - ShapeUtil::Compatible(old_operand->shape(), new_operand->shape())) + TF_RET_CHECK(ShapeUtil::CompatibleIgnoringFpPrecision(old_operand->shape(), + new_operand->shape())) << old_operand->shape().ShortDebugString() << " is not compatible with " << new_operand->shape().ShortDebugString(); operands_[operand_num] = new_operand; @@ -2055,6 +2140,10 @@ std::vector HloInstruction::ExtraAttributesToString( if (dot_dimension_numbers_ != nullptr) { extra.push_back(DotDimensionNumbersToString()); } + if (opcode() == HloOpcode::kFft) { + extra.push_back(StrCat("fft_type=", FftType_Name(fft_type()))); + extra.push_back(StrCat("fft_length={", Join(fft_length(), ","), "}")); + } if (options.print_subcomputation_references()) { if (opcode() == HloOpcode::kWhile) { @@ -2206,6 +2295,10 @@ HloInstructionProto HloInstruction::ToProto() const { proto.set_infeed_config(infeed_config_); proto.set_custom_call_target(custom_call_target_); *proto.mutable_outfeed_shape() = outfeed_shape_; + proto.set_fft_type(fft_type_); + for (int64 fft_len : fft_length_) { + proto.add_fft_length(fft_len); + } return proto; } @@ -2227,42 +2320,27 @@ string HloInstruction::ToCategory() const { return category; } + // Give transpose-dot and backwards-conv fusions the categories "dot" and + // "convolution" so they match the categories of proper kDot and kConvolution + // ops. These fusion categories are really just a way of expressing a + // particular kind of dot or conv, so they should have the same category as a + // vanilla dot/conv. if (opcode() == HloOpcode::kFusion) { - if (operands().size() == 2) { - bool saw_rank_1 = false; - bool saw_higher_rank = false; - for (const auto* operand : operands()) { - if (!ShapeUtil::IsTuple(operand->shape())) { - saw_rank_1 |= ShapeUtil::Rank(operand->shape()) == 1; - saw_higher_rank |= ShapeUtil::Rank(operand->shape()) > 1; - } - } - if (saw_rank_1 && saw_higher_rank) { - return "rank-1-broadcast binary fusion"; - } - } switch (fusion_kind()) { case FusionKind::kLoop: - if (IsElementwise()) { - return "elementwise fusion"; - } else { - return "non-elementwise fusion"; - } + return "loop fusion"; case FusionKind::kInput: return "input fusion"; case FusionKind::kOutput: return "output fusion"; case FusionKind::kTransposeDot: - return "dot fusion"; - case FusionKind::kConvBackwardFilter: - case FusionKind::kConvBackwardInput: - return "convolution fusion"; + return "dot"; case FusionKind::kCustom: return "custom fusion"; } } - if (IsElementwise() && opcode() != HloOpcode::kFusion) { + if (IsElementwise()) { return "non-fusion elementwise"; } @@ -2278,7 +2356,7 @@ void HloInstruction::set_tracing(HloInstruction* trace_instruction) { string HloInstruction::TracingTag() const { CHECK_EQ(HloOpcode::kTrace, opcode()); CHECK(literal_ != nullptr); - return literal_->u8s_string(); + return literal_->GetR1U8AsString(); } bool HloInstruction::IsFused() const { return parent_->IsFusionComputation(); } @@ -2421,6 +2499,8 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleSelect(this); case HloOpcode::kConvolution: return visitor->HandleConvolution(this); + case HloOpcode::kFft: + return visitor->HandleFft(this); case HloOpcode::kCrossReplicaSum: return visitor->HandleCrossReplicaSum(this); case HloOpcode::kTuple: @@ -3029,10 +3109,6 @@ string ToString(HloInstruction::FusionKind kind) { return "kOutput"; case HloInstruction::FusionKind::kTransposeDot: return "kTransposeDot"; - case HloInstruction::FusionKind::kConvBackwardFilter: - return "kConvBackwardFilter"; - case HloInstruction::FusionKind::kConvBackwardInput: - return "kConvBackwardInput"; case HloInstruction::FusionKind::kCustom: return "kCustom"; } @@ -3052,12 +3128,6 @@ StatusOr StringToFusionKind( if (kind_name == "kTransposeDot") { return HloInstruction::FusionKind::kTransposeDot; } - if (kind_name == "kConvBackwardFilter") { - return HloInstruction::FusionKind::kConvBackwardFilter; - } - if (kind_name == "kConvBackwardInput") { - return HloInstruction::FusionKind::kConvBackwardInput; - } if (kind_name == "kCustom") { return HloInstruction::FusionKind::kCustom; } @@ -3165,7 +3235,13 @@ string HloInstruction::ConvolutionDimensionNumbersToString() const { result += "_"; append_dims(rhs_dims, operand(1)->shape()); result += "->"; - append_dims(output_dims, shape()); + + // A convolution can be represented as a kConvolution HLO or as a CustomCall + // that returns a tuple, the first element of which is the result of the + // convolution. + Shape this_shape = + ShapeUtil::IsTuple(shape()) ? shape().tuple_shapes(0) : shape(); + append_dims(output_dims, this_shape); return result; } diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index d455cfc3f18e603cb678e1b9ffbd600ad40049ba..3170746157fbcfa7d0a7eaba6d226d46691105f9 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -25,6 +25,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -161,17 +162,14 @@ class HloPrintOptions { class HloInstruction { public: enum class FusionKind { - kLoop, // Fused into a loop. - kInput, // Op's input is fused into the op itself. - kOutput, // Op's output is fused into the op itself. - // REQUIRES: At least one operand buffer must be able - // to alias the output buffer. - kTransposeDot, // Fused into a dot with transposed operands. - kConvBackwardFilter, // Fused into a backward filter convolution. - kConvBackwardInput, // Fused into a backward input convolution. - - kCustom, // Custom category for backend-specific fusions that - // do not match any of the more specific ones. + kLoop, // Fused into a loop. + kInput, // Op's input is fused into the op itself. + kOutput, // Op's output is fused into the op itself. + // REQUIRES: At least one operand buffer must be able + // to alias the output buffer. + kTransposeDot, // Fused into a dot with transposed operands. + kCustom, // Custom category for backend-specific fusions that + // do not match any of the more specific ones. }; ~HloInstruction(); @@ -261,6 +259,11 @@ class HloInstruction { const Window& window, const ConvolutionDimensionNumbers& dimension_numbers); + // Creates an FFT op, of the type indicated by fft_type. + static std::unique_ptr CreateFft( + const Shape& shape, HloInstruction* operand, FftType fft_type, + tensorflow::gtl::ArraySlice fft_length); + // Creates a dot op with operands 'lhs' and 'rhs' with contracting and batch // dimensions specified in 'dimension_numbers'. static std::unique_ptr CreateDot( @@ -403,6 +406,20 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, tensorflow::gtl::ArraySlice broadcast_dimensions); + // Creates a sequence of instructions that performs an explicit broadcast of + // the operand to the target shape. + // + // Interior HLOs are passed to "adder", but the "root" HLO of the sequence is + // returned as a unique_ptr for API consistency with other factory methods in + // this interface. + // + // TODO(b/72173833) Ideally HloComputations would always be present, and so + // the adder being passed by the caller would not be necessary. + static std::unique_ptr CreateBroadcastSequence( + const Shape& output_shape, HloInstruction* operand, + const std::function)>& + adder); + // Creates a pad instruction, where the operand is padded on the edges and // between the elements with the given padding value. static std::unique_ptr CreatePad( @@ -446,14 +463,6 @@ class HloInstruction { tensorflow::gtl::ArraySlice operands, HloComputation* fusion_computation); - // Creates a fusion instruction that represents backward convolution. This is - // similar to CreateFusion, but with extra arguments indicating the window and - // dimemsion mapping of the backward convolution. - static std::unique_ptr CreateFusionForBackwardConvolution( - const Shape& shape, FusionKind fusion_kind, const Window& window, - const ConvolutionDimensionNumbers& conv_dnums, - HloInstruction* fused_root); - // Creates a call instruction that applies the given computation on the given // operands. "shape" is the resultant shape. static std::unique_ptr CreateCall( @@ -545,28 +554,42 @@ class HloInstruction { } // Returns true if "other" performs the same computation as this instruction. - // Layout of the instructions' output array is not considered. bool Identical( const HloInstruction& other, - std::function + const std::function& eq_operands = std::equal_to(), - std::function - eq_computations = std::equal_to()) const { + const std::function& + eq_computations = std::equal_to(), + bool layout_sensitive = true) const { // An instruction is always identical to itself. if (this == &other) { return true; } - // Identical instruction must have the same opcode and identical operands. - // In general, there is no need to check shape because shape is inferred - // from the shape of the operands. - if (opcode() != other.opcode() || - !ContainersEqual(operands(), other.operands(), - std::move(eq_operands))) { + // Identical instruction must have the same opcode, shape, and identical + // operands. + if (opcode() != other.opcode()) { return false; } + using EqShapeFuncType = bool (*)(const Shape&, const Shape&); + EqShapeFuncType eq_shapes = + layout_sensitive ? ShapeUtil::Equal : ShapeUtil::Compatible; + if (!eq_shapes(shape(), other.shape())) { + return false; + } + if (operands().size() != other.operands().size()) { + return false; + } + + // Use an explicit loop rather than ContainerEquals, because copying around + // std::functions may be too expensive in some cases. + for (size_t i = 0; i < operands().size(); ++i) { + if (!eq_operands(operand(i), other.operand(i))) { + return false; + } + } - return IdenticalSlowPath(other, eq_computations); + return IdenticalSlowPath(other, eq_computations, eq_shapes); } // Returns whether the instruction has a constant operand. @@ -857,8 +880,8 @@ class HloInstruction { // Returns true if this instruction is a fusion instruction that generates // multiple outputs. const bool IsMultiOutputFusion() const { - return (opcode() == HloOpcode::kFusion && - fused_expression_root()->opcode() == HloOpcode::kTuple); + return opcode() == HloOpcode::kFusion && + fused_expression_root()->opcode() == HloOpcode::kTuple; } FusionKind fusion_kind() const { @@ -1024,13 +1047,33 @@ class HloInstruction { return *padding_config_; } - // Returns data on the dimension numbers used for a convolution - // operation. + // Returns data on the dimension numbers used for a convolution operation, + // which may be a kConvolution instruction or a kCustomCall that implements a + // convolution. const ConvolutionDimensionNumbers& convolution_dimension_numbers() const { CHECK(convolution_dimension_numbers_ != nullptr); return *convolution_dimension_numbers_; } + // Sets the convolution dimension numbers on this instruction. In general you + // shouldn't need to call this; instead, specify the convolution dimension + // numbers when you create the instruction. + void set_convolution_dimension_numbers( + const ConvolutionDimensionNumbers& dnums) { + convolution_dimension_numbers_ = + MakeUnique(dnums); + } + + FftType fft_type() const { + CHECK_EQ(HloOpcode::kFft, opcode_); + return fft_type_; + } + + const std::vector& fft_length() const { + CHECK_EQ(HloOpcode::kFft, opcode_); + return fft_length_; + } + // Returns the dump string of the convolution dimension numbers. string ConvolutionDimensionNumbersToString() const; @@ -1195,10 +1238,14 @@ class HloInstruction { class FusionReusesParamElements; // See comments on Identical(). + // eq_shapes() is used to check shapes for equality, and would normally be + // expected to be ShapeUtil::Equals or ShapeUtil::Compatible, depending on + // whether we want a layout-sensitive check or not. bool IdenticalSlowPath( const HloInstruction& other, - std::function - eq_computations) const; + const std::function& + eq_computations, + const std::function& eq_shapes) const; // Creates an n-ary elementwise operation. static std::unique_ptr CreateNary( @@ -1303,6 +1350,12 @@ class HloInstruction { // Describes the dimension numbers used for a dot. std::unique_ptr dot_dimension_numbers_; + // Describes FFT type for an FFT instruction. + FftType fft_type_ = FftType::FFT; + + // Indicates the FFT length for an FFT instruction. + std::vector fft_length_; + // Describes the [begin, end) index range for a slice. std::vector slice_starts_; std::vector slice_limits_; @@ -1430,6 +1483,10 @@ template using ConstHloInstructionMap = std::map; +using HloInstructionSet = std::set; +using ConstHloInstructionSet = + std::set; + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INSTRUCTION_H_ diff --git a/tensorflow/compiler/xla/service/hlo_instruction_test.cc b/tensorflow/compiler/xla/service/hlo_instruction_test.cc index e762523a5e3941f535227e03d72d8f7751cd7b0b..94e9bfe56eb445ec0b459a55342cd3cc4c6f68ef 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction_test.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction_test.cc @@ -712,8 +712,8 @@ TEST_F(HloInstructionTest, PreserveOutfeedShapeThroughClone) { {1, 2}, {3, 4}, }))); - auto shape10 = ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {1, 0}); - auto shape01 = ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {0, 1}); + auto shape10 = ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {1, 0}); + auto shape01 = ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {0, 1}); auto outfeed10 = builder.AddInstruction( HloInstruction::CreateOutfeed(shape10, constant, "")); auto outfeed01 = builder.AddInstruction( @@ -825,17 +825,42 @@ TEST_F(HloInstructionTest, ComplexFusionOp) { EXPECT_THAT(c1->users(), ElementsAre(fusion)); } -// Convenience function for comparing two HloInstructions inside of -// std::unique_ptrs. -static bool Identical(std::unique_ptr instruction1, - std::unique_ptr instruction2) { +// Convenience function for comparing two HloInstructions. +static bool Identical(const HloInstruction& instruction1, + const HloInstruction& instruction2) { // Verify Identical is reflexive for both instructions. - EXPECT_TRUE(instruction1->Identical(*instruction1)); - EXPECT_TRUE(instruction2->Identical(*instruction2)); + EXPECT_TRUE(instruction1.Identical(instruction1)); + EXPECT_TRUE(instruction2.Identical(instruction2)); - bool is_equal = instruction1->Identical(*instruction2); + bool is_equal = instruction1.Identical(instruction2); // Verify Identical is symmetric. - EXPECT_EQ(is_equal, instruction2->Identical(*instruction1)); + EXPECT_EQ(is_equal, instruction2.Identical(instruction1)); + return is_equal; +} + +// Convenience function for comparing two HloInstructions for structural +// equality. +static bool StructuralEqual(const HloInstruction& instruction1, + const HloInstruction& instruction2) { + auto eq_operand_shapes = [](const HloInstruction* a, + const HloInstruction* b) { + return ShapeUtil::Equal(a->shape(), b->shape()); + }; + auto eq_computations = [](const HloComputation* a, const HloComputation* b) { + return *a == *b; + }; + + // Verify Identical is reflexive for both instructions. + EXPECT_TRUE( + instruction1.Identical(instruction1, eq_operand_shapes, eq_computations)); + EXPECT_TRUE( + instruction2.Identical(instruction2, eq_operand_shapes, eq_computations)); + + bool is_equal = + instruction1.Identical(instruction2, eq_operand_shapes, eq_computations); + // Verify Identical is symmetric. + EXPECT_EQ(is_equal, instruction2.Identical(instruction1, eq_operand_shapes, + eq_computations)); return is_equal; } @@ -858,42 +883,42 @@ TEST_F(HloInstructionTest, IdenticalInstructions) { // Operations which only depend on their operands and opcode. EXPECT_TRUE( - Identical(HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1), - HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1))); + Identical(*HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1), + *HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1))); EXPECT_FALSE( - Identical(HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1), - HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op2))); + Identical(*HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1), + *HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op2))); EXPECT_FALSE( - Identical(HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1), - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, op1))); + Identical(*HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1), + *HloInstruction::CreateUnary(shape, HloOpcode::kNegate, op1))); // Tuples. - EXPECT_TRUE(Identical(HloInstruction::CreateTuple({op1, op2}), - HloInstruction::CreateTuple({op1, op2}))); - EXPECT_FALSE(Identical(HloInstruction::CreateTuple({op1, op2}), - HloInstruction::CreateTuple({op2, op1}))); + EXPECT_TRUE(Identical(*HloInstruction::CreateTuple({op1, op2}), + *HloInstruction::CreateTuple({op1, op2}))); + EXPECT_FALSE(Identical(*HloInstruction::CreateTuple({op1, op2}), + *HloInstruction::CreateTuple({op2, op1}))); // Broadcasts. - EXPECT_TRUE(Identical(HloInstruction::CreateBroadcast(shape, op1, {0, 1}), - HloInstruction::CreateBroadcast(shape, op1, {0, 1}))); - EXPECT_FALSE(Identical(HloInstruction::CreateBroadcast(shape, op1, {0, 1}), - HloInstruction::CreateBroadcast(shape, op1, {1, 0}))); + EXPECT_TRUE(Identical(*HloInstruction::CreateBroadcast(shape, op1, {0, 1}), + *HloInstruction::CreateBroadcast(shape, op1, {0, 1}))); + EXPECT_FALSE(Identical(*HloInstruction::CreateBroadcast(shape, op1, {0, 1}), + *HloInstruction::CreateBroadcast(shape, op1, {1, 0}))); Shape bcast_shape1 = ShapeUtil::MakeShape(F32, {2, 2, 42}); Shape bcast_shape2 = ShapeUtil::MakeShape(F32, {2, 2, 123}); EXPECT_FALSE( - Identical(HloInstruction::CreateBroadcast(bcast_shape1, op1, {0, 1}), - HloInstruction::CreateBroadcast(bcast_shape2, op1, {0, 1}))); + Identical(*HloInstruction::CreateBroadcast(bcast_shape1, op1, {0, 1}), + *HloInstruction::CreateBroadcast(bcast_shape2, op1, {0, 1}))); // Binary operands. EXPECT_TRUE(Identical( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op1, op2), - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op1, op2))); + *HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op1, op2), + *HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op1, op2))); EXPECT_FALSE(Identical( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op1, op2), - HloInstruction::CreateBinary(shape, HloOpcode::kDivide, op2, op1))); + *HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op1, op2), + *HloInstruction::CreateBinary(shape, HloOpcode::kDivide, op2, op1))); EXPECT_FALSE(Identical( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op1, op2), - HloInstruction::CreateBinary(shape, HloOpcode::kDivide, op1, op2))); + *HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op1, op2), + *HloInstruction::CreateBinary(shape, HloOpcode::kDivide, op1, op2))); } TEST_F(HloInstructionTest, FunctionVisitor) { @@ -1089,6 +1114,70 @@ TEST_F(HloInstructionTest, CloneOfFusionPreservesShape) { ShapeUtil::Equal(root->operand(1)->shape(), root2->operand(1)->shape())); EXPECT_TRUE(ShapeUtil::Equal(root->operand(1)->operand(0)->shape(), root2->operand(1)->operand(0)->shape())); + EXPECT_TRUE(StructuralEqual(*fusion, *fusion2)); +} + +TEST_F(HloInstructionTest, FusionEquality) { + HloModule module(TestName()); + HloComputation::Builder builder(TestName()); + + // Create two fusion instructions containing a single unary operation. + auto parameter = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "x")); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, parameter)); + auto neg = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, parameter)); + auto* computation = module.AddEntryComputation(builder.Build()); + auto* fusion = computation->CreateFusionInstruction( + {exp}, HloInstruction::FusionKind::kLoop); + auto* fusion2 = computation->CreateFusionInstruction( + {neg}, HloInstruction::FusionKind::kLoop); + EXPECT_FALSE(StructuralEqual(*fusion, *fusion2)); + + auto clone = fusion->Clone(); + EXPECT_TRUE(StructuralEqual(*fusion, *clone)); +} + +TEST_F(HloInstructionTest, NestedFusionEquality) { + HloModule module(TestName()); + HloComputation::Builder builder(TestName()); + + // Build a nested fusion computation. + Shape data_shape = ShapeUtil::MakeShape(F32, {2, 2}); + auto a = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{1.0, 0.0}, {0.0, 1.0}}))); + auto b = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{2.0, 2.0}, {2.0, 2.0}}))); + auto b_t = builder.AddInstruction( + HloInstruction::CreateTranspose(data_shape, b, {1, 0})); + DotDimensionNumbers dot_dnums; + dot_dnums.add_lhs_contracting_dimensions(1); + dot_dnums.add_rhs_contracting_dimensions(0); + auto dot = builder.AddInstruction( + HloInstruction::CreateDot(data_shape, a, b_t, dot_dnums)); + auto one = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto add_operand = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape, one, {1})); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + data_shape, HloOpcode::kAdd, dot, add_operand)); + auto sub = builder.AddInstruction(HloInstruction::CreateBinary( + data_shape, HloOpcode::kSubtract, dot, add_operand)); + builder.AddInstruction( + HloInstruction::CreateBinary(data_shape, HloOpcode::kMultiply, add, sub)); + auto computation = module.AddEntryComputation(builder.Build()); + + auto nested_fusion = computation->CreateFusionInstruction( + {dot, b_t}, HloInstruction::FusionKind::kTransposeDot); + + auto fusion = computation->CreateFusionInstruction( + {add, nested_fusion}, HloInstruction::FusionKind::kOutput); + auto fusion2 = computation->CreateFusionInstruction( + {sub, nested_fusion}, HloInstruction::FusionKind::kOutput); + auto clone = fusion->Clone(); + EXPECT_TRUE(StructuralEqual(*fusion, *clone)); + EXPECT_FALSE(StructuralEqual(*fusion, *fusion2)); } TEST_F(HloInstructionTest, CloneSuffixNames) { @@ -1162,7 +1251,7 @@ TEST_F(HloInstructionTest, Stringification) { EXPECT_EQ( fusion->ToString(options), - "%fusion = f32[5,20]{1,0} fusion(f32[5,10]{1,0} %x, " + "%dot_fusion = f32[5,20]{1,0} fusion(f32[5,10]{1,0} %x, " "f32[20,10]{1,0} %y), kind=kTransposeDot, calls=%fused_computation"); HloInstruction* loop = builder.AddInstruction( diff --git a/tensorflow/compiler/xla/service/hlo_matchers.cc b/tensorflow/compiler/xla/service/hlo_matchers.cc index 4255d6086625dfb9a045e4431e968a5ee0106ac7..bc74c4bc10cad20eab20b5caf8550b17048a5276 100644 --- a/tensorflow/compiler/xla/service/hlo_matchers.cc +++ b/tensorflow/compiler/xla/service/hlo_matchers.cc @@ -102,6 +102,36 @@ bool HloGetTupleElementMatcher::MatchAndExplain( return true; } +void HloCustomCallMatcher::DescribeTo(std::ostream* os) const { + HloMatcher::DescribeTo(os); + *os << " with call target that "; + call_target_matcher_.DescribeTo(os); +} + +bool HloCustomCallMatcher::MatchAndExplain( + const HloInstruction* instruction, + ::testing::MatchResultListener* listener) const { + if (!HloMatcher::MatchAndExplain(instruction, listener)) { + return false; + } + ::testing::StringMatchResultListener sub_listener; + bool result = ExplainMatchResult( + call_target_matcher_, instruction->custom_call_target(), &sub_listener); + if (sub_listener.str().empty()) { + sub_listener << " that "; + + std::stringstream desc_stream; + if (result) { + call_target_matcher_.DescribeTo(&desc_stream); + } else { + call_target_matcher_.DescribeNegationTo(&desc_stream); + } + sub_listener << desc_stream.str(); + } + *listener << "custom-call with call target" << sub_listener.str(); + return result; +} + } // namespace testing void PrintTo(const HloInstruction* inst, ::std::ostream* os) { diff --git a/tensorflow/compiler/xla/service/hlo_matchers.h b/tensorflow/compiler/xla/service/hlo_matchers.h index 992f55788b4900949f4994ba5b7be015bcd0d3de..103f04a2cb7a1a5ae877d8bf259692f7cbed3408 100644 --- a/tensorflow/compiler/xla/service/hlo_matchers.h +++ b/tensorflow/compiler/xla/service/hlo_matchers.h @@ -56,8 +56,8 @@ class HloParameterMatcher : public HloMatcher { // index to match. class HloGetTupleElementMatcher : public HloMatcher { public: - explicit HloGetTupleElementMatcher( - ::testing::Matcher operand, int64 tuple_index) + HloGetTupleElementMatcher(::testing::Matcher operand, + int64 tuple_index) : HloMatcher(HloOpcode::kGetTupleElement, /*operands=*/{operand}), tuple_index_(tuple_index) {} @@ -68,6 +68,24 @@ class HloGetTupleElementMatcher : public HloMatcher { int64 tuple_index_; }; +// Custom matcher for custom-call instructions, which accepts a matcher for its +// call target. +class HloCustomCallMatcher : public HloMatcher { + public: + HloCustomCallMatcher( + ::testing::Matcher call_target_matcher, + std::vector<::testing::Matcher> operands) + : HloMatcher(HloOpcode::kCustomCall, operands), + call_target_matcher_(call_target_matcher) {} + + bool MatchAndExplain(const HloInstruction* instruction, + ::testing::MatchResultListener* listener) const override; + void DescribeTo(std::ostream* os) const override; + + private: + ::testing::Matcher call_target_matcher_; +}; + // HloInstruction* matchers for opcode and operands. Example: // namespace op = xla::opcode_matchers; // EXPECT_THAT(instruction, @@ -83,6 +101,7 @@ HLO_MATCHER(Abs); HLO_MATCHER(Add); HLO_MATCHER(Bitcast); HLO_MATCHER(Broadcast); +HLO_MATCHER(BatchNormGrad); HLO_MATCHER(Call); HLO_MATCHER(Ceil); HLO_MATCHER(Clamp); @@ -93,7 +112,6 @@ HLO_MATCHER(Convert); HLO_MATCHER(Convolution); HLO_MATCHER(Copy); HLO_MATCHER(CrossReplicaSum); -HLO_MATCHER(CustomCall); HLO_MATCHER(Divide); HLO_MATCHER(Dot); HLO_MATCHER(DynamicSlice); @@ -183,6 +201,36 @@ inline ::testing::Matcher GetTupleElement() { new ::xla::testing::HloMatcher(HloOpcode::kGetTupleElement, {})); } +// - CustomCall(T, operand1, ..., operandN) matches a CustomCall with call +// target T and the given operands. +// +// - CustomCall(operand1, ..., operandN) matches any CustomCall HLO with the +// given operands. +// +// - CustomCall() matches any CustomCall HLO at all. +template +inline ::testing::Matcher CustomCall( + ::testing::Matcher call_target_matcher, M... operands) { + return ::testing::MakeMatcher(new ::xla::testing::HloCustomCallMatcher( + call_target_matcher, {operands...})); +} +// This overload of CustomCall(A, B, C, ...) exists iff A is not convertible to +// ::testing::Matcher. In that case, we want to prefer the overload +// above. +template >::value, + void>::type*> +inline ::testing::Matcher CustomCall( + FirstM operands_first, M... operands_rest) { + return ::testing::MakeMatcher(new ::xla::testing::HloMatcher( + HloOpcode::kCustomCall, {operands_first, operands_rest...})); +} +inline ::testing::Matcher CustomCall() { + return ::testing::MakeMatcher( + new ::xla::testing::HloMatcher(HloOpcode::kCustomCall, {})); +} + #undef HLO_MATCHER } // namespace opcode_matchers diff --git a/tensorflow/compiler/xla/service/hlo_matchers_test.cc b/tensorflow/compiler/xla/service/hlo_matchers_test.cc index 1465d1cacdc971a04c620bc48bed33239a67a955..1c21703a45e11914854153bc14fabd85e9ea57f2 100644 --- a/tensorflow/compiler/xla/service/hlo_matchers_test.cc +++ b/tensorflow/compiler/xla/service/hlo_matchers_test.cc @@ -23,6 +23,12 @@ using ::testing::Eq; namespace xla { namespace { +string DescribeHloMatcher(const ::testing::Matcher& m) { + std::stringstream ss; + m.DescribeTo(&ss); + return ss.str(); +} + template string Explain(const T& t, const M& m) { ::testing::StringMatchResultListener listener; @@ -67,5 +73,32 @@ TEST(HloMatchersTest, Test) { "add")); } +TEST(HloMatchersTest, CustomCallMatcher) { + auto c1 = HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3})); + auto c2 = HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3})); + auto call = HloInstruction::CreateCustomCall( + ShapeUtil::MakeShape(F32, {1}), {c1.get(), c2.get()}, "foo_target"); + + EXPECT_THAT(call.get(), op::CustomCall()); + EXPECT_THAT(call.get(), op::CustomCall(c1.get(), c2.get())); + EXPECT_THAT(call.get(), op::CustomCall("foo_target")); + EXPECT_THAT(call.get(), op::CustomCall("foo_target", c1.get(), c2.get())); + EXPECT_THAT(call.get(), op::CustomCall(::testing::StartsWith("foo"))); + EXPECT_THAT(call.get(), + op::CustomCall(::testing::Not(::testing::StartsWith("bar")))); + + // Wrong number of operands. + EXPECT_THAT(call.get(), ::testing::Not(op::CustomCall(c1.get()))); + + // Call target does not match. + EXPECT_THAT(call.get(), + ::testing::Not(op::CustomCall(::testing::StartsWith("bar")))); + + EXPECT_THAT(Explain(call.get(), op::CustomCall("bar")), + R"(custom-call with call target that isn't equal to "bar")"); + EXPECT_THAT(DescribeHloMatcher(op::CustomCall("foo_target")), + R"(custom-call with call target that is equal to "foo_target")"); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc index 6103cab3e7e73079ef9e65b4ada181aa088c4541..60270b0595dcfca8f1fcea5ab0914428880f35b5 100644 --- a/tensorflow/compiler/xla/service/hlo_module.cc +++ b/tensorflow/compiler/xla/service/hlo_module.cc @@ -38,12 +38,16 @@ HloModule::HloModule(const string& name, : name_(NameUniquer::GetSanitizedName(name)), config_(config), has_entry_computation_handle_(true), - entry_computation_handle_(entry_computation_handle) {} + entry_computation_handle_(entry_computation_handle), + unique_id_(next_unique_module_id_++) {} HloModule::HloModule(const string& name) - : name_(NameUniquer::GetSanitizedName(name)) {} + : name_(NameUniquer::GetSanitizedName(name)), + unique_id_(next_unique_module_id_++) {} HloModule::HloModule(const string& name, const HloModuleConfig& config) - : name_(NameUniquer::GetSanitizedName(name)), config_(config) {} + : name_(NameUniquer::GetSanitizedName(name)), + config_(config), + unique_id_(next_unique_module_id_++) {} HloComputation* HloModule::AddComputationInternal( std::unique_ptr computation, bool is_entry, @@ -457,6 +461,14 @@ HloInstruction* HloModule::OutlineExpressionFromComputation( return call; } +int64 HloModule::instruction_count() const { + int64 n = 0; + for (const auto& computation : computations_) { + n += computation->instruction_count(); + } + return n; +} + std::list HloModule::MakeComputationPostOrder() const { // First determine all root computations by building a set of nonroot // computations (computations which are called by an instruction in the @@ -515,7 +527,15 @@ std::unique_ptr HloModule::Clone(const string& suffix) const { std::unordered_map clone_map; for (auto& computation : computations_) { - auto cloned_computation = computation->Clone(suffix); + if (computation->IsFusionComputation()) { + // Cloning of a fused computation is handled by its fusion instruction. + continue; + } + + // When cloning a computation, pass in the new module, so that for any + // fusion instruction in this computation, the fused computation will be + // deep cloned to the new module. + auto cloned_computation = computation->Clone(suffix, module.get()); InsertOrDie(&clone_map, computation.get(), cloned_computation.get()); if (entry_computation_ == computation.get()) { @@ -529,8 +549,15 @@ std::unique_ptr HloModule::Clone(const string& suffix) const { for (auto* instruction : cloned_computation->instructions()) { // Rewrite instruction's called_computation to point to the cloned // computations. - instruction->ReplaceCalledComputations( - [&](HloComputation* hlo) { return FindOrDie(clone_map, hlo); }); + instruction->ReplaceCalledComputations([&](HloComputation* hlo) { + if (hlo->IsFusionComputation()) { + // Cloning of a fused computation has already been handled when its + // fusion instruction is cloned. So this hlo computation is already + // the cloned one. + return hlo; + } + return FindOrDie(clone_map, hlo); + }); } } return module; @@ -541,4 +568,6 @@ uint64 HloModule::RandomNew64() const { return rng_(); } +/* static */ std::atomic HloModule::next_unique_module_id_(0); + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h index d3bb46bffca15549ef22e2908f129efd8586fa67..4bfe8d89ce0a285de6d05d4867aaa6b266d78d12 100644 --- a/tensorflow/compiler/xla/service/hlo_module.h +++ b/tensorflow/compiler/xla/service/hlo_module.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_H_ +#include #include #include #include @@ -129,6 +130,9 @@ class HloModule { // Gets the number of computations in this module. int64 computation_count() const { return computations_.size(); } + // Gets the number of instructions in this module. + int64 instruction_count() const; + // Compute and return a post order of all computations in the module. The sort // is defined like so: if computation A has an instruction which calls // computation B, then A will appear after B in the sort. @@ -198,6 +202,10 @@ class HloModule { // this point are guaranteed to be in the range [0..NumUniqueInstructionIds()) int NumUniqueInstructionIds() const { return next_unique_id_; } + // Returns an id that is unique to this module across all modules created over + // the lifetime of this process. + int unique_id() const { return unique_id_; } + private: HloComputation* AddComputationInternal( std::unique_ptr computation, bool is_entry, @@ -224,6 +232,11 @@ class HloModule { NameUniquer computation_name_uniquer_{/*separator=*/"."}; NameUniquer instruction_name_uniquer_{/*separator=*/"."}; int next_unique_id_ = 0; + + // Used to keep track of the next unique module id that should be assigned. + static std::atomic next_unique_module_id_; + // A unique id to label modules with. + int unique_id_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module_test.cc b/tensorflow/compiler/xla/service/hlo_module_test.cc index 0f5d3dccb74e6e3c88e51685392171f940c03596..7f28a804bfec9c2f1bbb5fa08f7dd4e68be14d35 100644 --- a/tensorflow/compiler/xla/service/hlo_module_test.cc +++ b/tensorflow/compiler/xla/service/hlo_module_test.cc @@ -105,6 +105,48 @@ TEST_F(HloModuleTest, CloneTest) { } } +TEST_F(HloModuleTest, CloneHasFusion) { + auto module = CreateNewModule(); + + // Create the fused computation. + HloComputation* fused_computation; + { + auto b = HloComputation::Builder("Fused"); + auto x = b.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "x")); + b.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, x, x)); + fused_computation = module->AddEmbeddedComputation(b.Build()); + } + + // Create the entry computation. + { + auto b = HloComputation::Builder("Entry"); + auto input = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); + b.AddInstruction( + HloInstruction::CreateFusion(r0f32_, HloInstruction::FusionKind::kInput, + /*operands=*/{input}, fused_computation)); + module->AddEntryComputation(b.Build()); + } + + auto post_order = module->MakeComputationPostOrder(); + auto cloned_module = module->Clone("copy"); + auto post_order_copied = cloned_module->MakeComputationPostOrder(); + + EXPECT_EQ(post_order.size(), post_order_copied.size()); + for (auto origin = post_order.begin(), copied = post_order_copied.begin(); + origin != post_order.end() && copied != post_order_copied.end(); + ++origin, ++copied) { + if ((*origin)->name() == "Fused") { + // Clone of the fused computation is handled when its fusion instruction + // is cloned, which always use suffix ".clone". + EXPECT_EQ((*origin)->name() + ".clone", (*copied)->name()); + } else { + EXPECT_EQ((*origin)->name() + ".copy", (*copied)->name()); + } + } +} + TEST_F(HloModuleTest, DiamondComputationsPostOrder) { // Create a module with a diamond call graph of computations. auto module = CreateNewModule(); @@ -146,6 +188,12 @@ TEST_F(HloModuleTest, LargeConstantToString) { module->ToString(HloPrintOptions().set_print_large_constants(true))); } +TEST_F(HloModuleTest, UniqueModuleId) { + auto module_a = CreateNewModule(); + auto module_b = CreateNewModule(); + EXPECT_NE(module_a->unique_id(), module_b->unique_id()); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_opcode.h b/tensorflow/compiler/xla/service/hlo_opcode.h index f3f79357582ac7661a532e94031acdbca0b86784..3d64523a79fc50638fdf378b5d521a5cd4482b90 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.h +++ b/tensorflow/compiler/xla/service/hlo_opcode.h @@ -73,6 +73,7 @@ namespace xla { V(kDynamicUpdateSlice, "dynamic-update-slice") \ V(kEq, "equal-to", kHloOpcodeIsComparison) \ V(kExp, "exponential") \ + V(kFft, "fft") \ V(kFloor, "floor") \ V(kFusion, "fusion", kHloOpcodeIsVariadic) \ V(kGe, "greater-than-or-equal-to", kHloOpcodeIsComparison) \ diff --git a/tensorflow/compiler/xla/service/hlo_ordering.cc b/tensorflow/compiler/xla/service/hlo_ordering.cc index 6f6e679a21870e46da85963c3b2998465ac43420..68e3c9618c1fe9daacb0aee3ee98862c8b9e4bc4 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering.cc @@ -249,7 +249,7 @@ bool PredecessorHloOrdering::ExecutesBeforeInSameComputation( string PredecessorHloOrdering::ToStringHelper(const string& name) const { std::vector pieces; pieces.push_back(name); - for (auto* computation : module_->computations()) { + for (auto* computation : module_->MakeNonfusionComputations()) { pieces.push_back(tensorflow::strings::Printf("computation %s:", computation->name().c_str())); const auto all = computation->MakeInstructionPostOrder(); diff --git a/tensorflow/compiler/xla/service/hlo_ordering_test.cc b/tensorflow/compiler/xla/service/hlo_ordering_test.cc index 33bafd05c15c47abaa313f92eb53a791de43d7d9..aba66114de649ce7667ae77174e9c4073b010b90 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering_test.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_scheduling.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -310,5 +311,56 @@ TEST_F(HloOrderingTest, ValuesInWhileComputations) { *dataflow)); } +// Regression test for HloOrdering::ToString() crashing when fed a computation +// containing a fusion node. +TEST_F(HloOrderingTest, ToStringDoesNotCrash) { + const char* module_str = R"( +HloModule test_module + +body.v8 { + prev.1 = (s32[], f32[3]{0}, f32[3]{0}, f32[3]{0}) parameter(0) + get-tuple-element.4 = s32[] get-tuple-element(prev.1), index=0 + constant.1 = s32[] constant(1) + add = s32[] add(get-tuple-element.4, constant.1) + get-tuple-element.5 = f32[3]{0} get-tuple-element(prev.1), index=3 + get-tuple-element.6 = f32[3]{0} get-tuple-element(prev.1), index=1 + get-tuple-element.7 = f32[3]{0} get-tuple-element(prev.1), index=2 + ROOT tuple = (s32[], f32[3]{0}, f32[3]{0}, f32[3]{0}) tuple(add, get-tuple-element.5, get-tuple-element.6, get-tuple-element.7) +} + +condition.v4 { + constant.2 = s32[] constant(2) + prev.2 = (s32[], f32[3]{0}, f32[3]{0}, f32[3]{0}) parameter(0) + get-tuple-element.8 = s32[] get-tuple-element(prev.2), index=0 + ROOT greater-than = pred[] greater-than(constant.2, get-tuple-element.8) +} + +fused_computation { + get-tuple-element.5.param_1 = f32[3]{0} parameter(1) + get-tuple-element.6.param_2 = f32[3]{0} parameter(2) + add.4 = f32[3]{0} add(get-tuple-element.5.param_1, get-tuple-element.6.param_2) + get-tuple-element.7.param_1.1 = f32[3]{0} parameter(0) + ROOT add.5 = f32[3]{0} add(add.4, get-tuple-element.7.param_1.1) +} + +ENTRY while.v11 { + constant.5 = s32[] constant(0) + constant.6 = f32[3]{0} constant({1, 1, 1}) + constant.7 = f32[3]{0} constant({2, 2, 2}) + constant.8 = f32[3]{0} constant({3, 3, 3}) + tuple.1 = (s32[], f32[3]{0}, f32[3]{0}, f32[3]{0}) tuple(constant.5, constant.6, constant.7, constant.8) + while = (s32[], f32[3]{0}, f32[3]{0}, f32[3]{0}) while(tuple.1), condition=condition.v4, body=body.v8 + get-tuple-element.9 = f32[3]{0} get-tuple-element(while), index=3 + get-tuple-element.10 = f32[3]{0} get-tuple-element(while), index=1 + get-tuple-element.11 = f32[3]{0} get-tuple-element(while), index=2 + ROOT fusion = f32[3]{0} fusion(get-tuple-element.9, get-tuple-element.10, get-tuple-element.11), kind=kLoop, calls=fused_computation +})"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + tools::Parse(module_str)); + DependencyHloOrdering ordering(module.get()); + ordering.ToString(); // Shouldn't crash. +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc index 53bd46a641afcba1b9551895955742e74a9f374b..5120775737bfa32bbb656421216f2b3fbef590ea 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc +++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" +#include "tensorflow/compiler/xla/service/hlo_proto_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" @@ -32,12 +33,28 @@ using ::tensorflow::strings::StrCat; namespace xla { namespace { -void DumpModule(const HloModule& module, - const string& message) { +void DumpModuleGraph(const HloModule& module, const string& message) { hlo_graph_dumper::MaybeDumpHloModule(module, message); VLOG(3) << "HLO " << message << ":"; XLA_VLOG_LINES(3, module.ToString()); } + +void DumpModuleProto(const HloModule& module, const string& dump_to, + const string& pipeline_name, const string& pass_name) { + static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED); + static auto* const module_id_to_pass_number = + new tensorflow::gtl::FlatMap(); + + tensorflow::mutex_lock lock(mu); + const int64 pass_number = (*module_id_to_pass_number)[module.unique_id()]++; + + const string mod_name = SanitizeFileName(tensorflow::strings::Printf( + "module_%04d.%04lld.%s.after_%s", module.unique_id(), pass_number, + pipeline_name.c_str(), pass_name.c_str())); + + TF_QCHECK_OK(protobuf_util::DumpProtoToDirectory(MakeHloProto(module), + dump_to, mod_name)); +} } // namespace StatusOr HloPassPipeline::Run(HloModule* module) { @@ -78,6 +95,13 @@ StatusOr HloPassPipeline::Run(HloModule* module) { string message; TF_RETURN_IF_ERROR( run_invariant_checkers(StrCat("before running pipeline: ", name()))); + const string xla_dump_per_pass_hlo_proto_to = + module->config().debug_options().xla_dump_per_pass_hlo_proto_to(); + if (!xla_dump_per_pass_hlo_proto_to.empty()) { + DumpModuleProto(*module, xla_dump_per_pass_hlo_proto_to, name().ToString(), + "pipeline_start"); + } + for (auto& pass : passes_) { if (disabled_passes.count(pass->name().ToString()) > 0) { VLOG(1) << " Skipping HLO pass " << pass->name() @@ -90,17 +114,21 @@ StatusOr HloPassPipeline::Run(HloModule* module) { // Emit label containing: "after foo-pass, before bar-pass". message.clear(); StrAppend(&message, prefix, ", before ", pass->name()); - DumpModule(*module, message); + DumpModuleGraph(*module, message); TF_ASSIGN_OR_RETURN(bool changed_this_pass, pass->Run(module)); TF_RETURN_IF_ERROR( run_invariant_checkers(StrCat("after running pass: ", pass->name()))); + if (!xla_dump_per_pass_hlo_proto_to.empty()) { + DumpModuleProto(*module, xla_dump_per_pass_hlo_proto_to, + name().ToString(), pass->name().ToString()); + } changed |= changed_this_pass; prefix.clear(); StrAppend(&prefix, name(), ": after ", pass->name()); } - DumpModule(*module, prefix + ", pipeline end"); + DumpModuleGraph(*module, prefix + ", pipeline end"); return changed; } diff --git a/tensorflow/compiler/xla/service/hlo_profile_printer.cc b/tensorflow/compiler/xla/service/hlo_profile_printer.cc index e944ad15139af0d2f98e8e68d3d48303f47ecf1c..dcc22793015147aaf3229875078b2989e4ef7559 100644 --- a/tensorflow/compiler/xla/service/hlo_profile_printer.cc +++ b/tensorflow/compiler/xla/service/hlo_profile_printer.cc @@ -18,20 +18,20 @@ limitations under the License. #include "tensorflow/compiler/xla/service/human_readable_profile_builder.h" namespace xla { -string HloProfilePrinter::ToString(const int64* counters, - double clock_rate_ghz) const { +string PrintHloProfile(const HloProfilePrinterData& hlo_profile_printer_data, + const int64* counters, double clock_rate_ghz) { + using HloComputationInfo = HloProfilePrinterData::HloComputationInfo; + using HloInstructionInfo = HloProfilePrinterData::HloInstructionInfo; + string result; - for (int computation_idx = 0; computation_idx < computation_infos_size_; - computation_idx++) { - const HloComputationInfo& computation = computation_infos_[computation_idx]; - const HloInstructionInfo* instructions_begin = computation.instructions; - const HloInstructionInfo* instructions_end = - computation.instructions + computation.instructions_size; + for (const HloComputationInfo& computation_info : + hlo_profile_printer_data.computation_infos()) { + const auto& instruction_infos = computation_info.instruction_infos(); bool any_instruction_profiled = - std::any_of(instructions_begin, instructions_end, + std::any_of(instruction_infos.begin(), instruction_infos.end(), [&](const HloInstructionInfo& instruction_info) { - return counters[instruction_info.profile_index] != 0; + return counters[instruction_info.profile_index()] != 0; }); if (!any_instruction_profiled) { @@ -41,16 +41,19 @@ string HloProfilePrinter::ToString(const int64* counters, // Once we start using this in AOT for real, we will probably need a more // minimal version of HumanReadableProfileBuilder. HumanReadableProfileBuilder builder( - computation.name, counters[computation.profile_index], clock_rate_ghz); + computation_info.name(), counters[computation_info.profile_index()], + clock_rate_ghz); - for (const auto* instruction = instructions_begin; - instruction != instructions_end; instruction++) { + for (const auto& instruction_info : instruction_infos) { builder.AddOp( - /*op_name=*/instruction->long_name, - /*short_name=*/instruction->short_name, instruction->category, - counters[instruction->profile_index], instruction->flop_count, - instruction->transcendental_count, instruction->bytes_accessed, - instruction->optimal_seconds); + /*op_name=*/instruction_info.long_name(), + /*short_name=*/instruction_info.short_name(), + instruction_info.category(), + counters[instruction_info.profile_index()], + instruction_info.flop_count(), + instruction_info.transcendental_count(), + instruction_info.bytes_accessed(), + instruction_info.optimal_seconds()); } result += builder.ToString(); @@ -58,10 +61,4 @@ string HloProfilePrinter::ToString(const int64* counters, return result; } - -HloProfilePrinter::~HloProfilePrinter() { - if (deleter_) { - deleter_(computation_infos_, computation_infos_size_); - } -} } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_profile_printer.h b/tensorflow/compiler/xla/service/hlo_profile_printer.h index 2f056490ae027872570f7a0821ee63114f49fab8..b72325c7554acad258c2da55a18e5e18ec1b06a6 100644 --- a/tensorflow/compiler/xla/service/hlo_profile_printer.h +++ b/tensorflow/compiler/xla/service/hlo_profile_printer.h @@ -13,91 +13,20 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROFILE_PRINTER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROFILE_PRINTER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROFILE_PRINTER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROFILE_PRINTER_H_ #include #include #include +#include "tensorflow/compiler/xla/service/hlo_profile_printer_data.pb.h" #include "tensorflow/compiler/xla/types.h" namespace xla { -// Instances of this class can pretty-print profile counters gathered from -// running an XLA computation without having access to the backing module. -class HloProfilePrinter { - public: - // Holds meta information about an HloInstruction. - // - // The pointer-typed fields can be owning or non-owning -- this decision is - // manifested as the deleter_ function in the containing HloProfilePrinter. - struct HloInstructionInfo { - // Textual information for pretty printing. - const char* long_name; - const char* short_name; - const char* category; - - // Metrics computed by HloCostAnalysis. - float flop_count; - float transcendental_count; - float bytes_accessed; - float optimal_seconds; - - // The index into the profile counters array for the HloInstruction - // corresponding to this HloInstructionInfo. - int64 profile_index; - }; - - // Holds meta information about an HloComputation. - // - // The pointer-typed fields can be owning or non-owning -- this decision is - // manifested as the deleter_ function in the containing HloProfilePrinter. - struct HloComputationInfo { - const char* name; - - // The index into the profile counters array for the HloInstruction - // corresponding to this HloComputationInfo. - int64 profile_index; - - HloInstructionInfo* instructions; - int64 instructions_size; - }; - - HloProfilePrinter( - HloComputationInfo* computation_infos, int64 computation_infos_size, - int64 profile_counters_size, - std::function deleter = nullptr) - : computation_infos_(computation_infos), - computation_infos_size_(computation_infos_size), - profile_counters_size_(profile_counters_size), - deleter_(std::move(deleter)) {} - - HloProfilePrinter(HloProfilePrinter&& other) { - std::swap(other.computation_infos_, computation_infos_); - std::swap(other.computation_infos_size_, computation_infos_size_); - std::swap(other.deleter_, deleter_); - } - - HloProfilePrinter(const HloProfilePrinter&) = delete; - HloProfilePrinter& operator=(const HloProfilePrinter&) = delete; - - // Converts the profile counter sequence `counters` to a human readable string - // representation. - string ToString(const int64* counters, double clock_rate_ghz) const; - - // Returns the size of the profile buffer expected by this printer. - int64 profile_counters_size() const { return profile_counters_size_; } - - ~HloProfilePrinter(); - - private: - // The `computation_infos_` field can be owning or non-owning -- this decision - // is manifested as the deleter_ function. - HloComputationInfo* computation_infos_ = nullptr; - int64 computation_infos_size_ = 0; - int64 profile_counters_size_ = 0; - std::function deleter_; -}; +// Pretty-print an array of profile counters using hlo_profile_printer_data. +string PrintHloProfile(const HloProfilePrinterData& hlo_profile_printer_data, + const int64* counters, double clock_rate_ghz); } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROFILE_PRINTER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROFILE_PRINTER_H_ diff --git a/tensorflow/compiler/xla/service/hlo_profile_printer_data.proto b/tensorflow/compiler/xla/service/hlo_profile_printer_data.proto new file mode 100644 index 0000000000000000000000000000000000000000..9f22b733fe1d676b177039a9d7a3064b8638d7bc --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_profile_printer_data.proto @@ -0,0 +1,60 @@ +/* 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. +==============================================================================*/ + +syntax = "proto3"; + +package xla; + +option cc_enable_arenas = true; + +// Describes how to pretty-print a profile counter array gathered for a specific +// HloModule. +message HloProfilePrinterData { + // Pretty-printer information about an HloInstruction. + message HloInstructionInfo { + string long_name = 1; + string short_name = 2; + string category = 3; + + // Metrics computed by HloCostAnalysis. + float flop_count = 4; + float transcendental_count = 5; + float bytes_accessed = 6; + float optimal_seconds = 7; + + // The index into the profile counters array for the HloInstruction + // corresponding to this HloInstructionInfo. + int64 profile_index = 8; + } + + // Pretty-printer information about an HloComputation. + message HloComputationInfo { + string name = 1; + + // The index into the profile counters array for the HloComputation + // corresponding to this HloComputationInfo. + int64 profile_index = 2; + + // HloInstructionInfos for every HloInstruction in the HloComputation for + // corresponding to this HloComputattionInfo. + repeated HloInstructionInfo instruction_infos = 3; + } + + // HloComputationInfos for every HloComputation in the HloModule. + repeated HloComputationInfo computation_infos = 1; + + // The size of the profile counters array we will pretty-print. + int64 profile_counters_size = 2; +} diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.cc b/tensorflow/compiler/xla/service/hlo_proto_util.cc index 727ad0178c6227cd2e64c31a4618e781671b9393..78e6a101c10a1e812e3e2631d520139fd0bc425c 100644 --- a/tensorflow/compiler/xla/service/hlo_proto_util.cc +++ b/tensorflow/compiler/xla/service/hlo_proto_util.cc @@ -19,15 +19,20 @@ namespace xla { HloProto MakeHloProto(const HloModule& module, const BufferAssignment& assignment) { - HloModuleProto proto_module = module.ToProto(); HloOrderingProto proto_ordering = assignment.liveness().hlo_ordering().ToProto(); BufferAssignmentProto proto_assignment = assignment.ToProto(); - HloProto proto; - proto.mutable_hlo_module()->Swap(&proto_module); + HloProto proto = MakeHloProto(module); proto.mutable_hlo_ordering()->Swap(&proto_ordering); proto.mutable_buffer_assignment()->Swap(&proto_assignment); return proto; } +HloProto MakeHloProto(const HloModule& module) { + HloModuleProto proto_module = module.ToProto(); + HloProto proto; + proto.mutable_hlo_module()->Swap(&proto_module); + return proto; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.h b/tensorflow/compiler/xla/service/hlo_proto_util.h index 603259a11fcdca59f58653d9a7a164c983711a57..320288fdb9aa0810b306b1d78bd1ff4cfc366ed2 100644 --- a/tensorflow/compiler/xla/service/hlo_proto_util.h +++ b/tensorflow/compiler/xla/service/hlo_proto_util.h @@ -31,6 +31,10 @@ namespace xla { HloProto MakeHloProto(const HloModule& module, const BufferAssignment& assignment); +// Returns a serialized representation of the HLO state, but buffer assignment +// will not be included in the output. +HloProto MakeHloProto(const HloModule& module); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROTO_UTIL_H_ diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc b/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc index 216825959a560bd5baa4b49d1a3cace277e16098..1b7d26dde501a6a0955d62ea0938e0683a32d49d 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc @@ -345,7 +345,7 @@ TEST_F(HloRematerializationTest, RngNotRematerialized) { auto param = builder.AddInstruction( HloInstruction::CreateParameter(0, scalar_shape_, "param")); auto rng = builder.AddInstruction(HloInstruction::CreateRng( - vec1024_shape_, RandomDistribution::RNG_BERNOULLI, {param})); + vec1024_shape_, RandomDistribution::RNG_UNIFORM, {param, param})); auto tanh = builder.AddInstruction( HloInstruction::CreateUnary(vec1024_shape_, HloOpcode::kTanh, rng)); auto exp = builder.AddInstruction( diff --git a/tensorflow/compiler/xla/service/hlo_runner.cc b/tensorflow/compiler/xla/service/hlo_runner.cc index 7b3a8cef97b5670b1ab753cee14203a58c1e5c27..41b079eb799d06321a31f7d7ae0630dc8d58c46b 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.cc +++ b/tensorflow/compiler/xla/service/hlo_runner.cc @@ -47,22 +47,11 @@ HloRunner::CreateModuleFromString(const tensorflow::StringPiece hlo_string, return tools::Parse(hlo_string, config); } -/*static*/ StatusOr> -HloRunner::ReadModuleFromHloProtoFile(const std::string& filename, - const DebugOptions& debug_options) { - HloProto proto; - - const Status s = - tensorflow::ReadBinaryProto(tensorflow::Env::Default(), filename, &proto); - - if (!s.ok()) { - const Status s2 = - tensorflow::ReadTextProto(tensorflow::Env::Default(), filename, &proto); - if (!s2.ok()) { - return Status(s2.code(), s.error_message() + "\n" + s2.error_message()); - } - } +namespace { +// Creates an HloModule from the given proto. +StatusOr> HloProtoToModule( + const HloProto& proto, const DebugOptions& debug_options) { TF_ASSIGN_OR_RETURN( HloModuleConfig config, HloModule::CreateModuleConfigFromProto(proto.hlo_module())); @@ -72,9 +61,29 @@ HloRunner::ReadModuleFromHloProtoFile(const std::string& filename, return std::move(module); } +} // namespace + /*static*/ StatusOr> -HloRunner::ReadModuleFromHloTextDumpFile(const std::string& filename, +HloRunner::ReadModuleFromBinaryProtoFile(const std::string& filename, const DebugOptions& debug_options) { + HloProto proto; + TF_RETURN_IF_ERROR(tensorflow::ReadBinaryProto(tensorflow::Env::Default(), + filename, &proto)); + return HloProtoToModule(proto, debug_options); +} + +/*static*/ StatusOr> +HloRunner::ReadModuleFromTextProtoFile(const std::string& filename, + const DebugOptions& debug_options) { + HloProto proto; + TF_RETURN_IF_ERROR( + tensorflow::ReadTextProto(tensorflow::Env::Default(), filename, &proto)); + return HloProtoToModule(proto, debug_options); +} + +/*static*/ StatusOr> +HloRunner::ReadModuleFromHloTextFile(const std::string& filename, + const DebugOptions& debug_options) { string hlo_string; TF_RETURN_IF_ERROR(tensorflow::ReadFileToString(tensorflow::Env::Default(), filename, &hlo_string)); @@ -83,19 +92,6 @@ HloRunner::ReadModuleFromHloTextDumpFile(const std::string& filename, return tools::Parse(hlo_string, config); } -/*static*/ StatusOr> HloRunner::ReadModule( - const std::string& filename, const DebugOptions& debug_options) { - auto module = HloRunner::ReadModuleFromHloProtoFile(filename, debug_options); - if (module.ok()) { - return module; - } - const std::string e = module.status().error_message(); - module = HloRunner::ReadModuleFromHloTextDumpFile(filename, debug_options); - return module.ok() ? std::move(module) - : Status(module.status().code(), - e + "\n" + module.status().error_message()); -} - // Define this in .cc file to avoid having to include eigen or forward declare // these types in the header. struct HloRunner::EigenThreadPoolWrapper { @@ -121,12 +117,14 @@ StatusOr> HloRunner::ExecuteInternal( if (run_hlo_passes) { TF_ASSIGN_OR_RETURN( module, backend().compiler()->RunHloPasses( - std::move(module), backend().default_stream_executor())); + std::move(module), backend().default_stream_executor(), + /*device_allocator=*/nullptr)); } TF_ASSIGN_OR_RETURN( std::unique_ptr executable, backend().compiler()->RunBackend(std::move(module), - backend().default_stream_executor())); + backend().default_stream_executor(), + /*device_allocator=*/nullptr)); se::Stream stream(backend().default_stream_executor()); stream.Init(); @@ -169,8 +167,16 @@ StatusOr> HloRunner::ExecuteInternal( std::unique_ptr scoped_result, ScopedShapedBuffer::MakeScoped(result.get(), run_options.allocator())); - return backend().transfer_manager()->TransferLiteralFromDevice( + auto result_literal = backend().transfer_manager()->TransferLiteralFromDevice( stream.parent(), *scoped_result); + if (result_literal.ok()) { + VLOG(4) << "Executed binary and got result: " + << result_literal.ValueOrDie()->ToString(); + } else { + VLOG(4) << "Executed binary and got status: " + << result_literal.status().ToString(); + } + return result_literal; } Backend& HloRunner::backend() { diff --git a/tensorflow/compiler/xla/service/hlo_runner.h b/tensorflow/compiler/xla/service/hlo_runner.h index d4b221fb52dff64dda264a931df6fd19b86e5260..cbaebc68bee708090b8ccb2eae19b556c4d6d453 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.h +++ b/tensorflow/compiler/xla/service/hlo_runner.h @@ -52,21 +52,15 @@ class HloRunner { const DebugOptions& debug_options); // Reads the proto file in xla.HloProto format, creates and returns the - // HloModule. Will try to parse the filename as binary proto, then try as - // text proto if that fails. - static StatusOr> ReadModuleFromHloProtoFile( + // HloModule. + static StatusOr> ReadModuleFromBinaryProtoFile( + const std::string& filename, const DebugOptions& debug_options); + static StatusOr> ReadModuleFromTextProtoFile( const std::string& filename, const DebugOptions& debug_options); // Reads the hlo text dump file in HloModule::ToString format, creates and // returns the HloModule. - static StatusOr> ReadModuleFromHloTextDumpFile( - const std::string& filename, const DebugOptions& debug_options); - - // Tries to parse the filename specified first as binary proto format, then - // as a textual proto format, then textual IR, then gives up if both fail. - // ReadModuleFromHloProtoFile or ReadModuleFromHloTextDumpFile should be used - // explicitly when you know the format, this if you don't. - static StatusOr> ReadModule( + static StatusOr> ReadModuleFromHloTextFile( const std::string& filename, const DebugOptions& debug_options); // Executes the given module with given literals as input and returns the diff --git a/tensorflow/compiler/xla/service/hlo_scheduling.cc b/tensorflow/compiler/xla/service/hlo_scheduling.cc index 2594c29efd717b3bead34d326c28c7efdf093c50..8dc4d4f7bac1b2007f2b9f60d126fa07e314dac9 100644 --- a/tensorflow/compiler/xla/service/hlo_scheduling.cc +++ b/tensorflow/compiler/xla/service/hlo_scheduling.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_scheduling.h" +#include #include #include @@ -100,7 +101,7 @@ class ListScheduler { // LogicalBuffer is in an operand of the instruction as indicated by // points-to analysis. for (auto* instruction : computation.instructions()) { - std::unordered_set instr_uses; + tensorflow::gtl::FlatSet instr_uses; for (auto* operand : instruction->operands()) { for (const LogicalBuffer* buffer : points_to_analysis.GetBuffersDefinedByInstruction(operand)) { @@ -150,10 +151,8 @@ class ListScheduler { int64 bytes_defined; // For each buffer B used by this instruction, we keep a pair (B, U), where - // U is the number of uses of B that have not yet been scheduled. This pair - // is a pointer into the unscheduled_use_count_ map, so it gets updated for - // free when we update counts in the map. - std::vector*> + // U is the number of uses of B that have not yet been scheduled. + std::vector> used_buffer_unscheduled_use_counts; }; @@ -176,8 +175,8 @@ class ListScheduler { } auto unscheduled_use_count_it = unscheduled_use_count_.find(buffer); CHECK(unscheduled_use_count_it != unscheduled_use_count_.end()); - entry.used_buffer_unscheduled_use_counts.push_back( - &*unscheduled_use_count_it); + entry.used_buffer_unscheduled_use_counts.emplace_back( + unscheduled_use_count_it->first, unscheduled_use_count_it->second); } return entry; } @@ -186,8 +185,8 @@ class ListScheduler { int64 BytesFreedIfScheduled(const ReadyListEntry& entry) { int64 freed_bytes = 0; for (const auto& kv : entry.used_buffer_unscheduled_use_counts) { - auto buffer = kv->first; - auto use_count = kv->second; + auto buffer = kv.first; + auto use_count = kv.second; if (use_count == 1) { freed_bytes += size_function_(*buffer); } @@ -205,7 +204,8 @@ class ListScheduler { // Populate the ready list with instructions which have no operands or // control predecessors. - std::unordered_map unscheduled_pred_count; + tensorflow::gtl::FlatMap + unscheduled_pred_count; for (auto* instruction : computation_.instructions()) { // TODO(b/34466113): Replace this and above with successors() or // predecessors() when these methods are added to HloInstruction. @@ -217,39 +217,57 @@ class ListScheduler { } } - std::list ready_list; + auto priority_comparator = + [this](const std::pair& lhs, + const std::pair& rhs) { + return lhs.first < rhs.first; + }; + std::priority_queue, + std::vector>, + decltype(priority_comparator)> + ready_queue(priority_comparator); + + // Set of instructions in the ready list. + tensorflow::gtl::FlatSet ready_instructions; + + auto add_to_ready_queue = [&](HloInstruction* inst) { + auto entry = MakeReadyListEntry(inst); + ready_queue.emplace(GetPriority(entry), std::move(entry)); + ready_instructions.insert(inst); + }; + for (auto* instruction : computation_.instructions()) { // Instruction with no operands or control predecessors will // not be in the map. if (unscheduled_pred_count.count(instruction) == 0) { - ready_list.push_back(MakeReadyListEntry(instruction)); + add_to_ready_queue(instruction); } } - while (!ready_list.empty()) { - // Select the highest priority HLO instruction from the ready list. - auto best_it = ready_list.begin(); - Priority best_priority = GetPriority(*best_it); - for (auto ready_it = std::next(ready_list.begin()); - ready_it != ready_list.end(); ++ready_it) { - Priority priority = GetPriority(*ready_it); - if (priority > best_priority) { - best_it = ready_it; - best_priority = priority; - } - } - + while (!ready_queue.empty()) { // Remove the selected instruction from the ready list and add it to the // schedule. - const HloInstruction* best = best_it->instruction; - ready_list.erase(best_it); + const HloInstruction* best = ready_queue.top().second.instruction; + ready_queue.pop(); + // We may have duplicates in the priority queue, because when a ready + // instruction's priority goes up, we reinsert it to the priority queue. + // Skip the duplicate. + if (scheduled_instructions_.find(best) != scheduled_instructions_.end()) { + continue; + } + ready_instructions.erase(best); schedule.push_back(best); scheduled_instructions_.insert(best); + bool adjust_ready_queue = false; // Update the unscheduled uses of the logical buffers. for (const LogicalBuffer* buffer : buffer_uses_.at(best)) { - CHECK_GT(unscheduled_use_count_.at(buffer), 0); - --unscheduled_use_count_[buffer]; + int64& count = unscheduled_use_count_[buffer]; + CHECK_GT(count, 0); + --count; + if (count == 1) { + adjust_ready_queue = true; + } } // Add new instructions to ready list. @@ -257,7 +275,7 @@ class ListScheduler { int64 pred_count = --unscheduled_pred_count.at(inst); CHECK_GE(pred_count, 0); if (pred_count == 0) { - ready_list.push_back(MakeReadyListEntry(inst)); + add_to_ready_queue(inst); } }; // TODO(b/34466113): Replace this and above with successors() or @@ -268,6 +286,20 @@ class ListScheduler { for (HloInstruction* succ : best->control_successors()) { update_pred_count(succ); } + // The unscheduled use count for a buffer has changed to 1, so the + // priorities of some ready instructions may go up. We reinsert them to + // the priority queue, so that they can appear earlier. The old entries + // will become duplicates and will be skipped. + if (adjust_ready_queue) { + for (HloInstruction* operand : best->operands()) { + for (HloInstruction* operand_user : operand->users()) { + if (ready_instructions.find(operand_user) != + ready_instructions.end()) { + add_to_ready_queue(operand_user); + } + } + } + } } CHECK_EQ(schedule.size(), computation_.instruction_count()); CHECK_EQ(scheduled_instructions_.size(), computation_.instruction_count()); @@ -280,15 +312,16 @@ class ListScheduler { const LogicalBuffer::SizeFunction& size_function_; // A map containing the LogicalBuffers that each instruction uses. - std::unordered_map> + tensorflow::gtl::FlatMap> buffer_uses_; // A map containing the count of unscheduled HLOs which using a particular // LogicalBuffer. We rely on iterator stability in this map. - std::unordered_map unscheduled_use_count_; + tensorflow::gtl::FlatMap unscheduled_use_count_; // Set of instructions which have been scheduled. - std::unordered_set scheduled_instructions_; + tensorflow::gtl::FlatSet scheduled_instructions_; }; int64 SumLogicalBufferSizes( diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder.h b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.h index 9aa3e501d5f85e3b61b20555e3d13c5687f33f2f..c4876b852e32d34693202f4023aa20ad2b301ffd 100644 --- a/tensorflow/compiler/xla/service/hlo_tfgraph_builder.h +++ b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/xla.pb.h" @@ -56,4 +56,4 @@ class HloTfGraphBuilder { } // namespace hlo_graph_dumper } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index d963a8a2f4fac563f7e8d4e9d4dc3d6e761d40de..e2b3bb9d71497c352b0b92add2d2f6b4b777bee8 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -13,418 +13,529 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include + #include "tensorflow/compiler/xla/service/hlo_verifier.h" -#include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { -namespace { +Status ShapeVerifier::HandleElementwiseUnary(HloInstruction* hlo) { + return CheckUnaryShape(hlo); +} -// Visitor which verifies that the output shape is correctly set. Verifies -// against the inferred shape for the instruction. -// TODO(b/26024837): Check output shape for all instruction types. -class ShapeVerifier : public DfsHloVisitor { - public: - explicit ShapeVerifier( - const std::function& shape_size_fn) - : shape_size_fn_(shape_size_fn) {} - - Status HandleElementwiseUnary(HloInstruction* hlo) override { - return CheckUnaryShape(hlo); - } +Status ShapeVerifier::HandleElementwiseBinary(HloInstruction* hlo) { + return CheckBinaryShape(hlo); +} - Status HandleElementwiseBinary(HloInstruction* hlo) override { - return CheckBinaryShape(hlo); - } +Status ShapeVerifier::HandleClamp(HloInstruction* clamp) { + return CheckTernaryShape(clamp); +} - Status HandleClamp(HloInstruction* clamp) override { - return CheckTernaryShape(clamp); - } +Status ShapeVerifier::HandleSelect(HloInstruction* select) { + return CheckTernaryShape(select); +} - Status HandleSelect(HloInstruction* select) override { - return CheckTernaryShape(select); +Status ShapeVerifier::HandleConcatenate(HloInstruction* concatenate) { + std::vector operand_shapes; + for (const HloInstruction* operand : concatenate->operands()) { + operand_shapes.push_back(&operand->shape()); } + return CheckShape(concatenate, + ShapeInference::InferConcatOpShape( + operand_shapes, concatenate->concatenate_dimension())); +} - Status HandleConcatenate(HloInstruction* concatenate) override { - std::vector operand_shapes; - for (const HloInstruction* operand : concatenate->operands()) { - operand_shapes.push_back(&operand->shape()); - } - return CheckShape( - concatenate, ShapeInference::InferConcatOpShape( - operand_shapes, concatenate->concatenate_dimension())); - } +Status ShapeVerifier::HandleConvert(HloInstruction* convert) { + return CheckShape(convert, ShapeInference::InferConvertShape( + convert->operand(0)->shape(), + convert->shape().element_type())); +} - Status HandleConvert(HloInstruction* convert) override { - return CheckShape(convert, ShapeInference::InferConvertShape( - convert->operand(0)->shape(), - convert->shape().element_type())); - } +Status ShapeVerifier::HandleBitcastConvert(HloInstruction* convert) { + return CheckShape(convert, ShapeInference::InferBitcastConvertShape( + convert->operand(0)->shape(), + convert->shape().element_type())); +} - Status HandleBitcastConvert(HloInstruction* convert) override { - return CheckShape(convert, ShapeInference::InferBitcastConvertShape( - convert->operand(0)->shape(), - convert->shape().element_type())); - } +Status ShapeVerifier::HandleCopy(HloInstruction* copy) { + return CheckUnaryShape(copy); +} - Status HandleCopy(HloInstruction* copy) override { - return CheckUnaryShape(copy); - } +Status ShapeVerifier::HandleDot(HloInstruction* dot) { + TF_ASSIGN_OR_RETURN(const Shape expected, + ShapeInference::InferDotOpShape( + dot->operand(0)->shape(), dot->operand(1)->shape(), + dot->dot_dimension_numbers())); + return CheckShape(dot, expected); +} - Status HandleDot(HloInstruction* dot) override { - TF_ASSIGN_OR_RETURN(const Shape expected, - ShapeInference::InferDotOpShape( - dot->operand(0)->shape(), dot->operand(1)->shape(), - dot->dot_dimension_numbers())); - return CheckShape(dot, expected); - } +Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) { + TF_ASSIGN_OR_RETURN( + const Shape expected, + ShapeInference::InferConvolveShape( + convolution->operand(0)->shape(), convolution->operand(1)->shape(), + convolution->window(), convolution->convolution_dimension_numbers())); + return CheckShape(convolution, expected); +} - Status HandleConvolution(HloInstruction* convolution) override { - TF_ASSIGN_OR_RETURN( - const Shape expected, - ShapeInference::InferConvolveShape( - convolution->operand(0)->shape(), convolution->operand(1)->shape(), - convolution->window(), - convolution->convolution_dimension_numbers())); - return CheckShape(convolution, expected); - } +Status ShapeVerifier::HandleFft(HloInstruction* fft) { + TF_ASSIGN_OR_RETURN( + const Shape expected, + ShapeInference::InferFftShape(fft->operand(0)->shape(), fft->fft_type(), + fft->fft_length())); + return CheckShape(fft, expected); +} - Status HandleCrossReplicaSum(HloInstruction* crs) override { - std::vector operand_shapes; - for (const HloInstruction* operand : crs->operands()) { - operand_shapes.push_back(&operand->shape()); - } - return CheckShape( - crs, ShapeInference::InferCrossReplicaSumShape(operand_shapes)); +Status ShapeVerifier::HandleCrossReplicaSum(HloInstruction* crs) { + std::vector operand_shapes; + for (const HloInstruction* operand : crs->operands()) { + operand_shapes.push_back(&operand->shape()); } + return CheckShape(crs, + ShapeInference::InferCrossReplicaSumShape(operand_shapes)); +} - Status HandleReducePrecision(HloInstruction* reduce_precision) override { - return CheckShape(reduce_precision, - ShapeInference::InferReducePrecisionShape( - reduce_precision->operand(0)->shape(), - reduce_precision->exponent_bits(), - reduce_precision->mantissa_bits())); - } +Status ShapeVerifier::HandleReducePrecision(HloInstruction* reduce_precision) { + return CheckShape(reduce_precision, ShapeInference::InferReducePrecisionShape( + reduce_precision->operand(0)->shape(), + reduce_precision->exponent_bits(), + reduce_precision->mantissa_bits())); +} - Status HandleInfeed(HloInstruction*) override { - return tensorflow::Status::OK(); - } +Status ShapeVerifier::HandleInfeed(HloInstruction*) { + return tensorflow::Status::OK(); +} - Status HandleOutfeed(HloInstruction*) override { - return tensorflow::Status::OK(); - } +Status ShapeVerifier::HandleOutfeed(HloInstruction* outfeed) { + // Outfeed has a separate shape field for the value which is outfed to the + // host. The shape of the instruction itself is always nil because the outfeed + // produces no HLO value in the graph. + if (!ShapeUtil::Compatible(outfeed->outfeed_shape(), + outfeed->operand(0)->shape())) { + return InvalidArgument( + "Expected outfeed to have shape compatible with operand's shape %s, " + "actual shape is %s:\n%s", + ShapeUtil::HumanString(outfeed->operand(0)->shape()).c_str(), + ShapeUtil::HumanString(outfeed->outfeed_shape()).c_str(), + outfeed->ToString().c_str()); + } + return CheckShape(outfeed, ShapeUtil::MakeNil()); +} - Status HandleRng(HloInstruction*) override { - return tensorflow::Status::OK(); - } +Status ShapeVerifier::HandleRng(HloInstruction*) { + return tensorflow::Status::OK(); +} - Status HandleReverse(HloInstruction* reverse) override { - return CheckShape( - reverse, ShapeInference::InferReverseShape(reverse->operand(0)->shape(), - reverse->dimensions())); - } +Status ShapeVerifier::HandleReverse(HloInstruction* reverse) { + return CheckShape( + reverse, ShapeInference::InferReverseShape(reverse->operand(0)->shape(), + reverse->dimensions())); +} - Status HandleSort(HloInstruction* sort) override { - return CheckUnaryShape(sort); - } +Status ShapeVerifier::HandleSort(HloInstruction* sort) { + return CheckUnaryShape(sort); +} - Status HandleConstant(HloInstruction* constant) override { - return CheckShape(constant, constant->literal().shape()); - } +Status ShapeVerifier::HandleConstant(HloInstruction* constant) { + return CheckShape(constant, constant->literal().shape()); +} - Status HandleGetTupleElement(HloInstruction* get_tuple_element) override { - return CheckShape(get_tuple_element, - ShapeInference::InferGetTupleElementShape( - get_tuple_element->operand(0)->shape(), - get_tuple_element->tuple_index())); - } +Status ShapeVerifier::HandleGetTupleElement(HloInstruction* get_tuple_element) { + return CheckShape(get_tuple_element, + ShapeInference::InferGetTupleElementShape( + get_tuple_element->operand(0)->shape(), + get_tuple_element->tuple_index())); +} - Status HandleReduce(HloInstruction* reduce) override { - return CheckShape( - reduce, - ShapeInference::InferReduceShape( - reduce->operand(0)->shape(), reduce->operand(1)->shape(), - reduce->dimensions(), reduce->to_apply()->ComputeProgramShape())); - } +Status ShapeVerifier::HandleReduce(HloInstruction* reduce) { + return CheckShape( + reduce, + ShapeInference::InferReduceShape( + reduce->operand(0)->shape(), reduce->operand(1)->shape(), + reduce->dimensions(), reduce->to_apply()->ComputeProgramShape())); +} - Status HandleBitcast(HloInstruction* bitcast) override { - return tensorflow::Status::OK(); - } +Status ShapeVerifier::HandleBitcast(HloInstruction* bitcast) { + return tensorflow::Status::OK(); +} - Status HandleBroadcast(HloInstruction* broadcast) override { - // HLO broadcast has no exact analog at the proto level so there is no - // ShapeInference method. Check the output shape explicitly. - const Shape& operand_shape = broadcast->operand(0)->shape(); - TF_RET_CHECK(ShapeUtil::Rank(operand_shape) == - broadcast->dimensions().size()); - for (int64 operand_dimension = 0; - operand_dimension < ShapeUtil::Rank(operand_shape); - ++operand_dimension) { - int64 output_dimension = broadcast->dimensions()[operand_dimension]; - TF_RET_CHECK(broadcast->shape().dimensions(output_dimension) == - operand_shape.dimensions(operand_dimension)); - } - return tensorflow::Status::OK(); +Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) { + // HLO broadcast has no exact analog at the proto level so there is no + // ShapeInference method. Check the output shape explicitly. + const Shape& operand_shape = broadcast->operand(0)->shape(); + // Check for mixed precision. + TF_RETURN_IF_ERROR(CheckShape(broadcast, broadcast->shape())); + TF_RET_CHECK(ShapeUtil::Rank(operand_shape) == + broadcast->dimensions().size()); + for (int64 operand_dimension = 0; + operand_dimension < ShapeUtil::Rank(operand_shape); + ++operand_dimension) { + int64 output_dimension = broadcast->dimensions()[operand_dimension]; + TF_RET_CHECK(broadcast->shape().dimensions(output_dimension) == + operand_shape.dimensions(operand_dimension)) + << broadcast->ToString() << " operand shape " << operand_shape; } + return tensorflow::Status::OK(); +} - Status HandleReshape(HloInstruction* reshape) override { - TF_RET_CHECK(ShapeUtil::ElementsIn(reshape->shape()) == - ShapeUtil::ElementsIn(reshape->operand(0)->shape())); - return tensorflow::Status::OK(); - } +Status ShapeVerifier::HandleReshape(HloInstruction* reshape) { + // Check for mixed precision. + TF_RETURN_IF_ERROR(CheckShape(reshape, reshape->shape())); + TF_RET_CHECK(ShapeUtil::ElementsIn(reshape->shape()) == + ShapeUtil::ElementsIn(reshape->operand(0)->shape())); + return tensorflow::Status::OK(); +} - Status HandleTranspose(HloInstruction* transpose) override { - return CheckShape(transpose, ShapeInference::InferTransposeShape( - transpose->operand(0)->shape(), - transpose->dimensions())); - } +Status ShapeVerifier::HandleTranspose(HloInstruction* transpose) { + return CheckShape( + transpose, ShapeInference::InferTransposeShape( + transpose->operand(0)->shape(), transpose->dimensions())); +} - Status HandleParameter(HloInstruction*) override { - return tensorflow::Status::OK(); - } +Status ShapeVerifier::HandleParameter(HloInstruction*) { + return tensorflow::Status::OK(); +} - Status HandleFusion(HloInstruction*) override { - return tensorflow::Status::OK(); - } +Status ShapeVerifier::HandleFusion(HloInstruction*) { + return tensorflow::Status::OK(); +} - Status HandleCall(HloInstruction* call) override { - // The shape of kCall should match the shape of the computation it calls. - return CheckShape(call, call->to_apply()->ComputeProgramShape().result()); - } +Status ShapeVerifier::HandleCall(HloInstruction* call) { + // The shape of kCall should match the shape of the computation it calls. + return CheckShape(call, call->to_apply()->ComputeProgramShape().result()); +} - Status HandleCustomCall(HloInstruction*) override { - return tensorflow::Status::OK(); - } +Status ShapeVerifier::HandleCustomCall(HloInstruction*) { + return tensorflow::Status::OK(); +} - Status HandleSlice(HloInstruction* slice) override { - return CheckShape(slice, - ShapeInference::InferSliceShape( - slice->operand(0)->shape(), slice->slice_starts(), - slice->slice_limits(), slice->slice_strides())); - } +Status ShapeVerifier::HandleSlice(HloInstruction* slice) { + return CheckShape(slice, + ShapeInference::InferSliceShape( + slice->operand(0)->shape(), slice->slice_starts(), + slice->slice_limits(), slice->slice_strides())); +} - Status HandleDynamicSlice(HloInstruction* dynamic_slice) override { - return CheckShape(dynamic_slice, ShapeInference::InferDynamicSliceShape( - dynamic_slice->operand(0)->shape(), - dynamic_slice->operand(1)->shape(), - dynamic_slice->dynamic_slice_sizes())); - } +Status ShapeVerifier::HandleDynamicSlice(HloInstruction* dynamic_slice) { + return CheckShape(dynamic_slice, ShapeInference::InferDynamicSliceShape( + dynamic_slice->operand(0)->shape(), + dynamic_slice->operand(1)->shape(), + dynamic_slice->dynamic_slice_sizes())); +} - Status HandleDynamicUpdateSlice( - HloInstruction* dynamic_update_slice) override { - return CheckShape(dynamic_update_slice, - ShapeInference::InferDynamicUpdateSliceShape( - dynamic_update_slice->operand(0)->shape(), - dynamic_update_slice->operand(1)->shape(), - dynamic_update_slice->operand(2)->shape())); - } +Status ShapeVerifier::HandleDynamicUpdateSlice( + HloInstruction* dynamic_update_slice) { + return CheckShape(dynamic_update_slice, + ShapeInference::InferDynamicUpdateSliceShape( + dynamic_update_slice->operand(0)->shape(), + dynamic_update_slice->operand(1)->shape(), + dynamic_update_slice->operand(2)->shape())); +} - Status HandleTuple(HloInstruction* tuple) override { - return CheckVariadicShape(tuple); - } +Status ShapeVerifier::HandleTuple(HloInstruction* tuple) { + return CheckVariadicShape(tuple); +} - Status HandleMap(HloInstruction* map) override { - std::vector operand_shapes; - int64 max_operand_rank = 0; - for (const HloInstruction* operand : map->operands()) { - operand_shapes.push_back(&operand->shape()); - max_operand_rank = - std::max(max_operand_rank, ShapeUtil::Rank(operand->shape())); - } - // TODO(b/65689298) Remove code below once Map is generalized to accept - // arbitrary map dimensions. - std::vector map_dims(max_operand_rank); - std::iota(map_dims.begin(), map_dims.end(), 0); - return CheckShape( - map, - ShapeInference::InferMapShape( - operand_shapes, map->to_apply()->ComputeProgramShape(), map_dims)); - } +Status ShapeVerifier::HandleMap(HloInstruction* map) { + std::vector operand_shapes; + int64 max_operand_rank = 0; + for (const HloInstruction* operand : map->operands()) { + operand_shapes.push_back(&operand->shape()); + max_operand_rank = + std::max(max_operand_rank, ShapeUtil::Rank(operand->shape())); + } + // TODO(b/65689298) Remove code below once Map is generalized to accept + // arbitrary map dimensions. + std::vector map_dims(max_operand_rank); + std::iota(map_dims.begin(), map_dims.end(), 0); + return CheckShape(map, ShapeInference::InferMapShape( + operand_shapes, + map->to_apply()->ComputeProgramShape(), map_dims)); +} - Status HandleReduceWindow(HloInstruction* reduce_window) override { - return CheckShape( - reduce_window, - ShapeInference::InferReduceWindowShape( - reduce_window->operand(0)->shape(), - reduce_window->operand(1)->shape(), reduce_window->window(), - reduce_window->to_apply()->ComputeProgramShape())); - } +Status ShapeVerifier::HandleReduceWindow(HloInstruction* reduce_window) { + return CheckShape( + reduce_window, + ShapeInference::InferReduceWindowShape( + reduce_window->operand(0)->shape(), + reduce_window->operand(1)->shape(), reduce_window->window(), + reduce_window->to_apply()->ComputeProgramShape())); +} - Status HandleSelectAndScatter(HloInstruction* instruction) override { - return CheckShape( - instruction, - ShapeInference::InferSelectAndScatterShape( - instruction->operand(0)->shape(), - instruction->select()->ComputeProgramShape(), instruction->window(), - instruction->operand(1)->shape(), instruction->operand(2)->shape(), - instruction->scatter()->ComputeProgramShape())); - } +Status ShapeVerifier::HandleSelectAndScatter(HloInstruction* instruction) { + return CheckShape( + instruction, + ShapeInference::InferSelectAndScatterShape( + instruction->operand(0)->shape(), + instruction->select()->ComputeProgramShape(), instruction->window(), + instruction->operand(1)->shape(), instruction->operand(2)->shape(), + instruction->scatter()->ComputeProgramShape())); +} - Status HandleWhile(HloInstruction* xla_while) override { - // The shape of kWhile should match the shape of the body computation it - // calls. - return CheckShape(xla_while, - xla_while->while_body()->ComputeProgramShape().result()); - } +Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) { + // The shape of kWhile should match the shape of the body computation it + // calls. + return CheckShape(xla_while, + xla_while->while_body()->ComputeProgramShape().result()); +} - Status HandleConditional(HloInstruction* conditional) override { - TF_RETURN_IF_ERROR(CheckShape( - conditional, - conditional->true_computation()->ComputeProgramShape().result())); - return CheckShape( - conditional, - conditional->false_computation()->ComputeProgramShape().result()); - } +Status ShapeVerifier::HandleConditional(HloInstruction* conditional) { + TF_RETURN_IF_ERROR(CheckShape( + conditional, + conditional->true_computation()->ComputeProgramShape().result())); + return CheckShape( + conditional, + conditional->false_computation()->ComputeProgramShape().result()); +} - Status HandlePad(HloInstruction* pad) override { - return CheckShape(pad, - ShapeInference::InferPadShape(pad->operand(0)->shape(), - pad->operand(1)->shape(), - pad->padding_config())); - } +Status ShapeVerifier::HandlePad(HloInstruction* pad) { + return CheckShape(pad, ShapeInference::InferPadShape(pad->operand(0)->shape(), + pad->operand(1)->shape(), + pad->padding_config())); +} - Status HandleSend(HloInstruction* send) override { - TF_RET_CHECK(send->users().size() == 1); - const HloInstruction* send_done = send->users().front(); - TF_RET_CHECK(send_done->opcode() == HloOpcode::kSendDone); - TF_RETURN_IF_ERROR(CheckSameChannel(send, send_done)); - return CheckShape( - send, ShapeUtil::MakeTupleShape( - {send->operand(0)->shape(), ShapeUtil::MakeShape(U32, {})})); - } +Status ShapeVerifier::HandleSend(HloInstruction* send) { + TF_RET_CHECK(send->users().size() == 1); + const HloInstruction* send_done = send->users().front(); + TF_RET_CHECK(send_done->opcode() == HloOpcode::kSendDone); + TF_RETURN_IF_ERROR(CheckSameChannel(send, send_done)); + return CheckShape( + send, ShapeUtil::MakeTupleShape( + {send->operand(0)->shape(), ShapeUtil::MakeShape(U32, {})})); +} - Status HandleSendDone(HloInstruction* send_done) override { - TF_RET_CHECK(send_done->operands().size() == 1); - const HloInstruction* send = send_done->operand(0); - TF_RET_CHECK(send->opcode() == HloOpcode::kSend); - TF_RETURN_IF_ERROR(CheckSameChannel(send, send_done)); - return CheckShape(send_done, ShapeUtil::MakeNil()); - } +Status ShapeVerifier::HandleSendDone(HloInstruction* send_done) { + TF_RET_CHECK(send_done->operands().size() == 1); + const HloInstruction* send = send_done->operand(0); + TF_RET_CHECK(send->opcode() == HloOpcode::kSend); + TF_RETURN_IF_ERROR(CheckSameChannel(send, send_done)); + return CheckShape(send_done, ShapeUtil::MakeNil()); +} - Status HandleRecv(HloInstruction* recv) override { - TF_RET_CHECK(recv->users().size() == 1); - const HloInstruction* recv_done = recv->users().front(); - TF_RET_CHECK(recv_done->opcode() == HloOpcode::kRecvDone); - TF_RETURN_IF_ERROR(CheckSameChannel(recv, recv_done)); - return CheckShape(recv, - ShapeUtil::MakeTupleShape( - {recv_done->shape(), ShapeUtil::MakeShape(U32, {})})); - } +Status ShapeVerifier::HandleRecv(HloInstruction* recv) { + TF_RET_CHECK(recv->users().size() == 1); + const HloInstruction* recv_done = recv->users().front(); + TF_RET_CHECK(recv_done->opcode() == HloOpcode::kRecvDone); + TF_RETURN_IF_ERROR(CheckSameChannel(recv, recv_done)); + return CheckShape(recv, + ShapeUtil::MakeTupleShape( + {recv_done->shape(), ShapeUtil::MakeShape(U32, {})})); +} - Status HandleRecvDone(HloInstruction* recv_done) override { - TF_RET_CHECK(recv_done->operands().size() == 1); - const HloInstruction* recv = recv_done->operand(0); - TF_RET_CHECK(recv->opcode() == HloOpcode::kRecv); - TF_RETURN_IF_ERROR(CheckSameChannel(recv, recv_done)); - return CheckShape(recv_done, recv->shape().tuple_shapes(0)); - } +Status ShapeVerifier::HandleRecvDone(HloInstruction* recv_done) { + TF_RET_CHECK(recv_done->operands().size() == 1); + const HloInstruction* recv = recv_done->operand(0); + TF_RET_CHECK(recv->opcode() == HloOpcode::kRecv); + TF_RETURN_IF_ERROR(CheckSameChannel(recv, recv_done)); + return CheckShape(recv_done, recv->shape().tuple_shapes(0)); +} - Status HandleBatchNormTraining(HloInstruction* batch_norm_training) override { - return CheckShape(batch_norm_training, - ShapeInference::InferBatchNormTrainingShape( - batch_norm_training->operand(0)->shape(), - batch_norm_training->operand(1)->shape(), - batch_norm_training->operand(2)->shape(), - batch_norm_training->feature_index())); - } +Status ShapeVerifier::HandleBatchNormTraining( + HloInstruction* batch_norm_training) { + return CheckShape(batch_norm_training, + ShapeInference::InferBatchNormTrainingShape( + batch_norm_training->operand(0)->shape(), + batch_norm_training->operand(1)->shape(), + batch_norm_training->operand(2)->shape(), + batch_norm_training->feature_index())); +} - Status HandleBatchNormInference( - HloInstruction* batch_norm_inference) override { - return CheckShape(batch_norm_inference, - ShapeInference::InferBatchNormInferenceShape( - batch_norm_inference->operand(0)->shape(), - batch_norm_inference->operand(1)->shape(), - batch_norm_inference->operand(2)->shape(), - batch_norm_inference->operand(3)->shape(), - batch_norm_inference->operand(4)->shape(), - batch_norm_inference->feature_index())); - } +Status ShapeVerifier::HandleBatchNormInference( + HloInstruction* batch_norm_inference) { + return CheckShape(batch_norm_inference, + ShapeInference::InferBatchNormInferenceShape( + batch_norm_inference->operand(0)->shape(), + batch_norm_inference->operand(1)->shape(), + batch_norm_inference->operand(2)->shape(), + batch_norm_inference->operand(3)->shape(), + batch_norm_inference->operand(4)->shape(), + batch_norm_inference->feature_index())); +} - Status HandleBatchNormGrad(HloInstruction* batch_norm_grad) override { - return CheckShape(batch_norm_grad, ShapeInference::InferBatchNormGradShape( - batch_norm_grad->operand(0)->shape(), - batch_norm_grad->operand(1)->shape(), - batch_norm_grad->operand(2)->shape(), - batch_norm_grad->operand(3)->shape(), - batch_norm_grad->operand(4)->shape(), - batch_norm_grad->feature_index())); - } +Status ShapeVerifier::HandleBatchNormGrad(HloInstruction* batch_norm_grad) { + return CheckShape(batch_norm_grad, ShapeInference::InferBatchNormGradShape( + batch_norm_grad->operand(0)->shape(), + batch_norm_grad->operand(1)->shape(), + batch_norm_grad->operand(2)->shape(), + batch_norm_grad->operand(3)->shape(), + batch_norm_grad->operand(4)->shape(), + batch_norm_grad->feature_index())); +} - Status FinishVisit(HloInstruction*) override { - return tensorflow::Status::OK(); - } +namespace { - private: - // Check the instruction's shape against the given expected shape and return - // an appropriate error if there is a mismatch. - Status CheckShape(const HloInstruction* instruction, - const Shape& expected_shape) { - if (!ShapeUtil::Compatible(instruction->shape(), expected_shape)) { - return InvalidArgument( - "Expected instruction to have shape compatible with %s, actual " - "shape is %s:\n%s", - ShapeUtil::HumanString(expected_shape).c_str(), - ShapeUtil::HumanString(instruction->shape()).c_str(), - instruction->ToString().c_str()); +// Checks that the instruction does not have mixed precision floating point +// inputs. +Status CheckMixedPrecisionOperands(const HloInstruction* instruction) { + switch (instruction->opcode()) { + // White list the following opcodes for mixed-precision check, because they + // involve data pass through or grouping via tuples, where the precisions + // of buffers can be different. + case HloOpcode::kCall: + case HloOpcode::kConditional: + case HloOpcode::kConstant: + case HloOpcode::kCrossReplicaSum: + case HloOpcode::kCustomCall: + case HloOpcode::kFusion: + case HloOpcode::kGetTupleElement: + case HloOpcode::kInfeed: + case HloOpcode::kOutfeed: + case HloOpcode::kParameter: + case HloOpcode::kRecv: + case HloOpcode::kRecvDone: + case HloOpcode::kReducePrecision: + case HloOpcode::kSelect: + case HloOpcode::kSend: + case HloOpcode::kSendDone: + case HloOpcode::kTuple: + case HloOpcode::kWhile: + break; + default: { + PrimitiveType fp_type = PRIMITIVE_TYPE_INVALID; + for (auto operand : instruction->operands()) { + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + operand->shape(), + [&](const Shape& subshape, const ShapeIndex& index) { + if (!ShapeUtil::ElementIsFloating(subshape)) { + return Status::OK(); + } + if (fp_type == PRIMITIVE_TYPE_INVALID) { + fp_type = subshape.element_type(); + } else if (fp_type != subshape.element_type()) { + return FailedPrecondition( + "Seen floating point types of different precisions in " + "%s, but mixed precision is disallowed.", + instruction->ToString().c_str()); + } + return Status::OK(); + })); + } } - return tensorflow::Status::OK(); } + return Status::OK(); +} - // Overload which takes a StatusOr to reduce boilerplate in the caller. - Status CheckShape(const HloInstruction* instruction, - const StatusOr& expected_shape_status) { - if (!expected_shape_status.ok()) { - Status s = expected_shape_status.status(); - tensorflow::errors::AppendToMessage(&s, ", for instruction ", - instruction->ToString()); - return s; - } - return CheckShape(instruction, expected_shape_status.ValueOrDie()); - } +} // namespace - // Check a unary (binary, etc) instruction's shape against the inferred shape. - Status CheckUnaryShape(const HloInstruction* instruction) { - return CheckShape(instruction, - ShapeInference::InferUnaryOpShape( - instruction->opcode(), instruction->operand(0))); - } - Status CheckBinaryShape(const HloInstruction* instruction) { - return CheckShape(instruction, - ShapeInference::InferBinaryOpShape( - instruction->opcode(), instruction->operand(0), - instruction->operand(1))); - } - Status CheckTernaryShape(const HloInstruction* instruction) { - return CheckShape(instruction, - ShapeInference::InferTernaryOpShape( - instruction->opcode(), instruction->operand(0), - instruction->operand(1), instruction->operand(2))); +Status ShapeVerifier::CheckShape(const HloInstruction* instruction, + const Shape& inferred_shape) { + // If allow_mixed_precision_ is false, check if there are operands with + // different precisions. We need this check because ShapeInference allows + // mixed precision inputs. + if (!allow_mixed_precision_) { + TF_RETURN_IF_ERROR(CheckMixedPrecisionOperands(instruction)); + } + + // Check if the output shape matches the expected shape. + bool compatible; + // We treat BF16 and F32 as compatible types if mixed precision is allowed, + // but only when the instruction defines the BF16/F32 buffer. + switch (instruction->opcode()) { + case HloOpcode::kSelect: + if (ShapeUtil::IsTuple(inferred_shape) || !allow_mixed_precision_) { + // Select only defines the top-level buffer, which in this case is the + // tuple, so we cannot allow mixed precision. + compatible = + ShapeUtil::Compatible(instruction->shape(), inferred_shape); + } else { + compatible = ShapeUtil::CompatibleIgnoringFpPrecision( + instruction->shape(), inferred_shape); + } + break; + case HloOpcode::kGetTupleElement: + case HloOpcode::kTuple: + // Tuple and GetTupleElement do not define BF16/F32 buffers, so mixed + // precision is disallowed. + case HloOpcode::kConstant: + case HloOpcode::kBitcast: + case HloOpcode::kBitcastConvert: + case HloOpcode::kCall: + case HloOpcode::kConditional: + case HloOpcode::kConvert: + case HloOpcode::kCustomCall: + case HloOpcode::kInfeed: + case HloOpcode::kOutfeed: + case HloOpcode::kParameter: + case HloOpcode::kRecv: + case HloOpcode::kRecvDone: + case HloOpcode::kSend: + case HloOpcode::kSendDone: + case HloOpcode::kWhile: + // The above opcodes should match the expected shapes exactly. + compatible = ShapeUtil::Compatible(instruction->shape(), inferred_shape); + break; + default: + if (allow_mixed_precision_) { + compatible = ShapeUtil::CompatibleIgnoringFpPrecision( + instruction->shape(), inferred_shape); + } else { + compatible = + ShapeUtil::Compatible(instruction->shape(), inferred_shape); + } } - Status CheckVariadicShape(const HloInstruction* instruction) { - return CheckShape(instruction, - ShapeInference::InferVariadicOpShape( - instruction->opcode(), instruction->operands())); + if (!compatible) { + return InvalidArgument( + "Expected instruction to have shape compatible with %s, actual " + "shape is %s:\n%s", + ShapeUtil::HumanString(inferred_shape).c_str(), + ShapeUtil::HumanString(instruction->shape()).c_str(), + instruction->ToString().c_str()); } + return tensorflow::Status::OK(); +} - // Checks if the given two instructions shares the same channel id. - Status CheckSameChannel(const HloInstruction* instr1, - const HloInstruction* instr2) { - if (instr1->channel_id() != instr2->channel_id()) { - return FailedPrecondition( - "Expected to have the same channel id, actual channel ids are: %s " - "(%lld), %s (%lld)", - instr1->ToString().c_str(), instr1->channel_id(), - instr2->ToString().c_str(), instr2->channel_id()); - } - return tensorflow::Status::OK(); +Status ShapeVerifier::CheckShape(const HloInstruction* instruction, + const StatusOr& inferred_shape_status) { + if (!inferred_shape_status.ok()) { + Status s = inferred_shape_status.status(); + tensorflow::errors::AppendToMessage(&s, ", for instruction ", + instruction->ToString()); + return s; } + return CheckShape(instruction, inferred_shape_status.ValueOrDie()); +} + +Status ShapeVerifier::CheckUnaryShape(const HloInstruction* instruction) { + return CheckShape(instruction, + ShapeInference::InferUnaryOpShape(instruction->opcode(), + instruction->operand(0))); +} + +Status ShapeVerifier::CheckBinaryShape(const HloInstruction* instruction) { + return CheckShape( + instruction, ShapeInference::InferBinaryOpShape(instruction->opcode(), + instruction->operand(0), + instruction->operand(1))); +} + +Status ShapeVerifier::CheckTernaryShape(const HloInstruction* instruction) { + return CheckShape(instruction, + ShapeInference::InferTernaryOpShape( + instruction->opcode(), instruction->operand(0), + instruction->operand(1), instruction->operand(2))); +} - // Returns the size of a Shape in bytes. - const std::function shape_size_fn_; -}; +Status ShapeVerifier::CheckVariadicShape(const HloInstruction* instruction) { + return CheckShape(instruction, + ShapeInference::InferVariadicOpShape( + instruction->opcode(), instruction->operands())); +} + +// Checks if the given two instructions shares the same channel id. +Status ShapeVerifier::CheckSameChannel(const HloInstruction* instr1, + const HloInstruction* instr2) { + if (instr1->channel_id() != instr2->channel_id()) { + return FailedPrecondition( + "Expected to have the same channel id, actual channel ids are: %s " + "(%lld), %s (%lld)", + instr1->ToString().c_str(), instr1->channel_id(), + instr2->ToString().c_str(), instr2->channel_id()); + } + return tensorflow::Status::OK(); +} string ComputationsToString( tensorflow::gtl::ArraySlice computations) { @@ -491,8 +602,6 @@ Status VerifyHloStructure(HloModule* module) { return tensorflow::Status::OK(); } -} // namespace - Status HloVerifier::CheckFusionInstruction(HloInstruction* fusion) const { // The parent fusion instruction of the fusion computation must be 'fusion'. HloComputation* fused_computation = fusion->fused_instructions_computation(); @@ -614,7 +723,6 @@ StatusOr HloVerifier::Run(HloModule* module) { TF_RETURN_IF_ERROR(VerifyHloStructure(module)); tensorflow::gtl::FlatMap instructions; - ShapeVerifier shape_verifier(shape_size_fn_); for (auto* computation : module->computations()) { for (const auto& instruction : computation->instructions()) { @@ -694,7 +802,8 @@ StatusOr HloVerifier::Run(HloModule* module) { instructions[instruction->name()] = instruction; } - TF_RETURN_IF_ERROR(computation->Accept(&shape_verifier)); + std::unique_ptr shape_verifier = shape_verifier_factory_(); + TF_RETURN_IF_ERROR(computation->Accept(shape_verifier.get())); } return false; diff --git a/tensorflow/compiler/xla/service/hlo_verifier.h b/tensorflow/compiler/xla/service/hlo_verifier.h index e35a7f3642ccf91df37f69a3a11bd8c8e428b846..7eccf834bbd3ac6af0d5762a7241758b416a3523 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.h +++ b/tensorflow/compiler/xla/service/hlo_verifier.h @@ -13,19 +13,124 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VERIFIER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VERIFIER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VERIFIER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VERIFIER_H_ #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/service/shape_inference.h" + namespace xla { +// Visitor which verifies that the output shape is correctly set. Verifies +// against the inferred shape for the instruction. +// TODO(b/26024837): Check output shape for all instruction types. +class ShapeVerifier : public DfsHloVisitor { + public: + explicit ShapeVerifier() : allow_mixed_precision_(false) {} + explicit ShapeVerifier(bool allow_mixed_precision) + : allow_mixed_precision_(allow_mixed_precision) {} + + Status HandleElementwiseUnary(HloInstruction* hlo) override; + Status HandleElementwiseBinary(HloInstruction* hlo) override; + Status HandleClamp(HloInstruction* clamp) override; + Status HandleSelect(HloInstruction* select) override; + Status HandleConcatenate(HloInstruction* concatenate) override; + Status HandleConvert(HloInstruction* convert) override; + Status HandleBitcastConvert(HloInstruction* convert) override; + Status HandleCopy(HloInstruction* copy) override; + Status HandleDot(HloInstruction* dot) override; + Status HandleConvolution(HloInstruction* convolution) override; + Status HandleFft(HloInstruction* fft) override; + Status HandleCrossReplicaSum(HloInstruction* crs) override; + Status HandleReducePrecision(HloInstruction* reduce_precision) override; + Status HandleInfeed(HloInstruction*) override; + Status HandleOutfeed(HloInstruction*) override; + Status HandleRng(HloInstruction*) override; + Status HandleReverse(HloInstruction* reverse) override; + Status HandleSort(HloInstruction* sort) override; + Status HandleConstant(HloInstruction* constant) override; + Status HandleGetTupleElement(HloInstruction* get_tuple_element) override; + Status HandleReduce(HloInstruction* reduce) override; + Status HandleBitcast(HloInstruction* bitcast) override; + Status HandleBroadcast(HloInstruction* broadcast) override; + Status HandleReshape(HloInstruction* reshape) override; + Status HandleTranspose(HloInstruction* transpose) override; + Status HandleParameter(HloInstruction*) override; + Status HandleFusion(HloInstruction*) override; + Status HandleCall(HloInstruction* call) override; + Status HandleCustomCall(HloInstruction*) override; + Status HandleSlice(HloInstruction* slice) override; + Status HandleDynamicSlice(HloInstruction* dynamic_slice) override; + Status HandleDynamicUpdateSlice( + HloInstruction* dynamic_update_slice) override; + Status HandleTuple(HloInstruction* tuple) override; + Status HandleMap(HloInstruction* map) override; + Status HandleReduceWindow(HloInstruction* reduce_window) override; + Status HandleSelectAndScatter(HloInstruction* instruction) override; + Status HandleWhile(HloInstruction* xla_while) override; + Status HandleConditional(HloInstruction* conditional) override; + Status HandlePad(HloInstruction* pad) override; + Status HandleSend(HloInstruction* send) override; + Status HandleSendDone(HloInstruction* send_done) override; + Status HandleRecv(HloInstruction* recv) override; + Status HandleRecvDone(HloInstruction* recv_done) override; + Status HandleBatchNormTraining(HloInstruction* batch_norm_training) override; + Status HandleBatchNormInference( + HloInstruction* batch_norm_inference) override; + Status HandleBatchNormGrad(HloInstruction* batch_norm_grad) override; + + Status FinishVisit(HloInstruction*) override { + return tensorflow::Status::OK(); + } + + protected: + // Check the instruction's shape against the shape given by ShapeInference + // and return an appropriate error if there is a mismatch. + Status CheckShape(const HloInstruction* instruction, + const Shape& inferred_shape); + + // Overload which takes a StatusOr to reduce boilerplate in the caller. + Status CheckShape(const HloInstruction* instruction, + const StatusOr& inferred_shape_status); + + // Check a unary (binary, etc) instruction's shape against the inferred shape. + Status CheckUnaryShape(const HloInstruction* instruction); + Status CheckBinaryShape(const HloInstruction* instruction); + Status CheckTernaryShape(const HloInstruction* instruction); + Status CheckVariadicShape(const HloInstruction* instruction); + + // Checks if the given two instructions shares the same channel id. + Status CheckSameChannel(const HloInstruction* instr1, + const HloInstruction* instr2); + + private: + // Whether the inputs and output of an instruction can contain both F32s and + // BF16s. Tuples that include both F32s and BF16s are allowed regardless of + // this flag. + bool allow_mixed_precision_; +}; + // HLO pass that verifies invariants of HLO instructions for each computation in // the module. class HloVerifier : public HloPassInterface { public: - explicit HloVerifier(const std::function& shape_size_fn) - : shape_size_fn_(shape_size_fn) {} + using ShapeVerifierFactory = std::function()>; + + // Uses standard shape inference. + explicit HloVerifier() + : shape_verifier_factory_( + [] { return MakeUnique(false); }) {} + + explicit HloVerifier(bool allow_mixed_precision) + : shape_verifier_factory_([allow_mixed_precision] { + return MakeUnique(allow_mixed_precision); + }) {} + + // Uses custom shape verification. + explicit HloVerifier(ShapeVerifierFactory shape_verifier_factory) + : shape_verifier_factory_(std::move(shape_verifier_factory)) {} + ~HloVerifier() override = default; tensorflow::StringPiece name() const override { return "verifier"; } @@ -37,10 +142,13 @@ class HloVerifier : public HloPassInterface { // CHECKs various invariants of a fusion instruction. Status CheckFusionInstruction(HloInstruction* fusion) const; - // Returns the size of a Shape in bytes. - const std::function shape_size_fn_; + // Creates a ShapeVerifier that checks that shapes match inferred + // expectations. This is a factory function because ShapeVerifier, Note that + // ShapeVerifier, being a DfsHloVisitor, is stateful. We want a clean object + // for each run of the verifier. + ShapeVerifierFactory shape_verifier_factory_; }; } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VERIFIER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VERIFIER_H_ diff --git a/tensorflow/compiler/xla/service/hlo_verifier_test.cc b/tensorflow/compiler/xla/service/hlo_verifier_test.cc index 2a3b55decc5289e7e576d3c5897b333c0b1bc922..c92db0be14dceb32ea86521dcc99b8f63738e4a5 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier_test.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier_test.cc @@ -97,5 +97,31 @@ TEST_F(HloVerifierTest, DifferentOperandParents) { HasSubstr("is in a different computation")); } +TEST_F(HloVerifierTest, ResetsShapeVerifierState) { + HloComputation::Builder builder(TestName()); + Shape s1 = ShapeUtil::MakeShape(F32, {1}); + Shape s2 = ShapeUtil::MakeShape(F32, {2}); + + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter(0, s1, "param")); + + // Create an add instruction with the incorrect shape. + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(s2, HloOpcode::kAdd, param, param)); + + // In order to trigger the bug we're checking for, the instruction with the + // bad shape can't be the root of the computation. + builder.AddInstruction( + HloInstruction::CreateBinary(s2, HloOpcode::kMultiply, add, add)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + // Run the verifier twice. It should fail both times, because it shouldn't + // carry state in its DFS visitor between runs. + EXPECT_FALSE(verifier().Run(module.get()).status().ok()); + EXPECT_FALSE(verifier().Run(module.get()).status().ok()); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/human_readable_profile_builder.cc b/tensorflow/compiler/xla/service/human_readable_profile_builder.cc index b7c40fdeeb157fc74900bd9cf9d68a06a2cb1d56..13e4557317f74b3fb46f07fb91c339fd2f34752f 100644 --- a/tensorflow/compiler/xla/service/human_readable_profile_builder.cc +++ b/tensorflow/compiler/xla/service/human_readable_profile_builder.cc @@ -25,6 +25,7 @@ namespace xla { using tensorflow::strings::Appendf; using tensorflow::strings::HumanReadableElapsedTime; using tensorflow::strings::HumanReadableNumBytes; +using tensorflow::strings::Printf; using tensorflow::strings::StrAppend; string HumanReadableProfileBuilder::ToString() const { @@ -43,7 +44,12 @@ string HumanReadableProfileBuilder::ToString() const { } else { bytes_per_sec = HumanReadableNumBytes(op.bytes_accessed / CyclesToSeconds(op.cycles)); - bytes_per_cycle = HumanReadableNumBytes(op.bytes_accessed / op.cycles); + if (op.bytes_accessed > op.cycles) { + bytes_per_cycle = HumanReadableNumBytes(op.bytes_accessed / op.cycles); + } else { + bytes_per_cycle = + Printf("%.3fB", static_cast(op.bytes_accessed) / op.cycles); + } } double cycles_percent = 0; diff --git a/tensorflow/compiler/xla/service/implicit_broadcast_remover.cc b/tensorflow/compiler/xla/service/implicit_broadcast_remover.cc new file mode 100644 index 0000000000000000000000000000000000000000..ada21345014dac70d61129aaf7bbc7466a7db914 --- /dev/null +++ b/tensorflow/compiler/xla/service/implicit_broadcast_remover.cc @@ -0,0 +1,124 @@ +/* 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/compiler/xla/service/implicit_broadcast_remover.h" + +#include +#include +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_dce.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +namespace { + +// Visitor for removing implicit broadcasts. +class ImplicitBroadcastVisitor : public DfsHloVisitorWithDefault { + public: + Status DefaultAction(HloInstruction* hlo_instruction) override { + return Status::OK(); + } + + Status HandleElementwiseBinary(HloInstruction* hlo) override { + return ReplaceImplicitBroadcastOperands(hlo); + } + + Status HandleClamp(HloInstruction* hlo) override { + // Clamp is the only element-wise ternary operation. + return ReplaceImplicitBroadcastOperands(hlo); + } + + // Returns whether any modification has been made to any visited instruction. + bool changed() const { return changed_; } + + private: + // Iterates through the operands of 'hlo' and replace any operands which are + // implicitly broadcast with the equivalent sequence of broadcast and reshape + // instructions. An operand is considered to be implicitly broadcast if the + // operand shape does have the same dimensions as the shape of 'hlo'. + Status ReplaceImplicitBroadcastOperands(HloInstruction* hlo) { + auto fadd = [hlo](std::unique_ptr x) { + return hlo->parent()->AddInstruction(std::move(x)); + }; + std::vector operands; + bool operands_changed = false; + for (int i = 0; i < hlo->operand_count(); ++i) { + HloInstruction* operand = hlo->mutable_operand(i); + if (!ShapeUtil::SameDimensions(hlo->shape(), operand->shape())) { + HloInstruction* new_operand = hlo->parent()->AddInstruction( + HloInstruction::CreateBroadcastSequence(hlo->shape(), operand, + fadd)); + operands.push_back(new_operand); + operands_changed = true; + } else { + operands.push_back(operand); + } + } + if (operands_changed) { + // Create a new HLO instruction because the HloInstruction::Replace* + // methods check that the shape does not change with the replacement. + HloInstruction* new_hlo = hlo->parent()->AddInstruction( + hlo->CloneWithNewOperands(hlo->shape(), operands)); + TF_RETURN_IF_ERROR(hlo->ReplaceAllUsesWith(new_hlo)); + changed_ = true; + } + return Status::OK(); + } + + bool changed_ = false; +}; + +} // namespace + +StatusOr ImplicitBroadcastRemover::Run(HloModule* module) { + VLOG(1) << "Removing implicit broadcast from module " << module->name(); + XLA_VLOG_LINES(2, + "Before removing implicit broadcasts:\n" + module->ToString()); + + ImplicitBroadcastVisitor visitor; + for (HloComputation* computation : module->computations()) { + TF_RETURN_IF_ERROR(computation->Accept(&visitor)); + } + + if (visitor.changed()) { + // HLO instructions with implicitly broadcast operands are cloned and left + // for dead. Remove them. + HloDCE dce; + TF_RETURN_IF_ERROR(dce.Run(module).status()); + } + + XLA_VLOG_LINES(2, + "After removing implicit broadcasts:\n" + module->ToString()); + + return visitor.changed(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/implicit_broadcast_remover.h b/tensorflow/compiler/xla/service/implicit_broadcast_remover.h new file mode 100644 index 0000000000000000000000000000000000000000..aa325dc8a353c5bfbfded0c2774c66bfcc71c9cb --- /dev/null +++ b/tensorflow/compiler/xla/service/implicit_broadcast_remover.h @@ -0,0 +1,42 @@ +/* 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_COMPILER_XLA_SERVICE_IMPLICIT_BROADCAST_REMOVER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_IMPLICIT_BROADCAST_REMOVER_H_ + +#include + +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// Pass which replaces all implicit broadcasts with their equivalent sequence of +// explicit broadcast and reshape instructions. +class ImplicitBroadcastRemover : public HloPassInterface { + public: + ImplicitBroadcastRemover() {} + ~ImplicitBroadcastRemover() override {} + + tensorflow::StringPiece name() const override { + return "implicit-broadcast-remover"; + } + + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_IMPLICIT_BROADCAST_REMOVER_H_ diff --git a/tensorflow/compiler/xla/service/implicit_broadcast_remover_test.cc b/tensorflow/compiler/xla/service/implicit_broadcast_remover_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..8c7b38dd1bf73e0be7b669d7215812aaef1cee17 --- /dev/null +++ b/tensorflow/compiler/xla/service/implicit_broadcast_remover_test.cc @@ -0,0 +1,176 @@ +/* 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/compiler/xla/service/implicit_broadcast_remover.h" + +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" + +namespace op = xla::testing::opcode_matchers; + +namespace xla { +namespace { + +class ImplicitBroadcastRemoverTest : public HloVerifiedTestBase { + protected: + ImplicitBroadcastRemover remover_; +}; + +TEST_F(ImplicitBroadcastRemoverTest, NoImplicitBroadcast) { + auto builder = HloComputation::Builder(TestName()); + + const Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + auto param0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + auto param1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + + HloComputation* computation = module().AddEntryComputation(builder.Build()); + + EXPECT_FALSE(remover_.Run(&module()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Parameter(), op::Parameter())); +} + +TEST_F(ImplicitBroadcastRemoverTest, ScalarBroadcast) { + auto builder = HloComputation::Builder(TestName()); + + const Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "scalar_param")); + auto param1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kPower, param0, param1)); + + HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + + EXPECT_FALSE(ShapeUtil::Compatible(root->shape(), root->operand(0)->shape())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(1)->shape())); + + EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + root = computation->root_instruction(); + + EXPECT_THAT(root, op::Power(op::Broadcast(op::Parameter()), op::Parameter())); + + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(0)->shape())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(1)->shape())); +} + +TEST_F(ImplicitBroadcastRemoverTest, DegenerateDimensionBroadcast) { + auto builder = HloComputation::Builder(TestName()); + + const Shape shape = ShapeUtil::MakeShape(F32, {2, 4, 6}); + auto param0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + auto param1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {1, 4, 1}), "p1")); + builder.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kSubtract, param0, param1)); + + HloComputation* computation = module().AddEntryComputation(builder.Build()); + + EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + + HloInstruction* root = computation->root_instruction(); + EXPECT_THAT(root, op::Subtract(op::Parameter(), + op::Broadcast(op::Reshape(op::Parameter())))); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(0)->shape())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(1)->shape())); +} + +TEST_F(ImplicitBroadcastRemoverTest, ScalarBroadcastToDegenerateDimensions) { + auto builder = HloComputation::Builder(TestName()); + + const Shape shape = ShapeUtil::MakeShape(F32, {1, 4, 1}); + auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "scalar_param")); + auto param1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + builder.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kSubtract, param0, param1)); + + HloComputation* computation = module().AddEntryComputation(builder.Build()); + + EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + + HloInstruction* root = computation->root_instruction(); + EXPECT_THAT(root, + op::Subtract(op::Broadcast(op::Parameter()), op::Parameter())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(0)->shape())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(1)->shape())); +} + +TEST_F(ImplicitBroadcastRemoverTest, TernaryDegenerateDimensionBroadcast) { + auto builder = HloComputation::Builder(TestName()); + + const Shape shape = ShapeUtil::MakeShape(F32, {2, 4, 6, 8}); + auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 4, 1, 8}), "p0")); + auto param1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {1, 1, 6, 8}), "p1")); + auto param2 = builder.AddInstruction(HloInstruction::CreateParameter( + 2, ShapeUtil::MakeShape(F32, {2, 1, 6, 8}), "p2")); + builder.AddInstruction(HloInstruction::CreateTernary(shape, HloOpcode::kClamp, + param0, param1, param2)); + + HloComputation* computation = module().AddEntryComputation(builder.Build()); + + EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + + HloInstruction* root = computation->root_instruction(); + EXPECT_THAT(root, op::Clamp(op::Broadcast(op::Reshape(op::Parameter())), + op::Broadcast(op::Reshape(op::Parameter())), + op::Broadcast(op::Reshape(op::Parameter())))); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(0)->shape())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(1)->shape())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(2)->shape())); +} + +TEST_F(ImplicitBroadcastRemoverTest, + TernaryScalarAndDegenerateDimensionBroadcast) { + auto builder = HloComputation::Builder(TestName()); + + const Shape shape = ShapeUtil::MakeShape(F32, {2, 4, 6}); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p0")); + auto param1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {1, 4, 6}), "p1")); + auto param2 = + builder.AddInstruction(HloInstruction::CreateParameter(2, shape, "p2")); + builder.AddInstruction(HloInstruction::CreateTernary(shape, HloOpcode::kClamp, + param0, param1, param2)); + + HloComputation* computation = module().AddEntryComputation(builder.Build()); + + EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + + HloInstruction* root = computation->root_instruction(); + EXPECT_THAT(root, op::Clamp(op::Broadcast(op::Parameter()), + op::Broadcast(op::Reshape(op::Parameter())), + op::Parameter())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(0)->shape())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(1)->shape())); + EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(2)->shape())); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc index ba901b99e4f3c72c84c1ecdf4e19e58ad9ab6506..90e1f0acdc4cdeda280dabaab2df66b181d0f407 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion.cc @@ -100,6 +100,7 @@ namespace xla { case HloOpcode::kDivide: case HloOpcode::kDot: case HloOpcode::kExp: + case HloOpcode::kFft: case HloOpcode::kFusion: case HloOpcode::kLog: case HloOpcode::kMap: diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.cc b/tensorflow/compiler/xla/service/interpreter/compiler.cc index dc63a2224d659fa427d4d1a30c5dc0f94d643b36..9171e859c6f84ceef9664aa1eb90a07c87dfab40 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.cc +++ b/tensorflow/compiler/xla/service/interpreter/compiler.cc @@ -44,41 +44,26 @@ namespace interpreter { namespace se = ::perftools::gputools; namespace sep = ::perftools::gputools::interpreter; -/* - * Run optimization passes on the module. The graph is transformed by - * each pass in the optimization pipeline. The service subdirectory - * contains useful optimization passes. - */ Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) { HloPassPipeline pipeline("Interpreter"); - pipeline.AddPass(); - pipeline.AddPass(); - pipeline.AddPass(false); - - pipeline.AddPass>( - false, [](const Shape&, const Shape&) { return false; }); - pipeline.AddPass(); - pipeline.AddPass(); - pipeline.AddPass(); - pipeline.AddPass(true); + pipeline.AddPass( hlo_module->mutable_entry_computation_layout()); - pipeline.AddPass(); - pipeline.AddPass(); return pipeline.Run(hlo_module).status(); } StatusOr> InterpreterCompiler::RunHloPasses( - std::unique_ptr hlo_module, - se::StreamExecutor* /*stream_exec*/) { + std::unique_ptr hlo_module, se::StreamExecutor* /*stream_exec*/, + DeviceMemoryAllocator* /*device_allocator*/) { VLOG(1) << "Run hlo passes on graph " << hlo_module->name(); TF_RETURN_IF_ERROR(RunHloOptimization(hlo_module.get())); return std::move(hlo_module); } StatusOr> InterpreterCompiler::RunBackend( - std::unique_ptr hlo_module, se::StreamExecutor* stream_exec) { + std::unique_ptr hlo_module, se::StreamExecutor* stream_exec, + DeviceMemoryAllocator* /*device_allocator*/) { TF_RET_CHECK(stream_exec != nullptr); VLOG(1) << "Run backend " << hlo_module->name(); @@ -96,7 +81,8 @@ StatusOr> InterpreterCompiler::RunBackend( StatusOr>> InterpreterCompiler::Compile( std::vector> /*hlo_modules*/, - std::vector> /*stream_execs*/) { + std::vector> /*stream_execs*/, + DeviceMemoryAllocator* /*device_allocator*/) { return tensorflow::errors::Unimplemented( "Compilation of multiple HLO modules is not supported on Interpreter."); } diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.h b/tensorflow/compiler/xla/service/interpreter/compiler.h index 278cf5184227ae25518b1d46c0e16e4cce7bd1a8..c8660c04d86a82e7dfcfd1658310c2a0e4fa0083 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.h +++ b/tensorflow/compiler/xla/service/interpreter/compiler.h @@ -45,16 +45,19 @@ class InterpreterCompiler : public Compiler { StatusOr> RunHloPasses( std::unique_ptr hlo_module, - perftools::gputools::StreamExecutor* stream_exec) override; + perftools::gputools::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) override; StatusOr> RunBackend( std::unique_ptr hlo_module, - perftools::gputools::StreamExecutor* stream_exec) override; + perftools::gputools::StreamExecutor* stream_exec, + DeviceMemoryAllocator* device_allocator) override; StatusOr>> Compile( std::vector> hlo_modules, std::vector> - stream_exec) override; + stream_exec, + DeviceMemoryAllocator* device_allocator) override; StatusOr>> CompileAheadOfTime(std::vector> hlo_modules, diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc index 51572f621ab9a8dacaac074c5cc35b004b54d041..fce135ef61a7868386b869def1a79167c428d928 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.cc +++ b/tensorflow/compiler/xla/service/layout_assignment.cc @@ -61,8 +61,8 @@ std::ostream& operator<<(std::ostream& out, BufferLayoutConstraint::BufferLayoutConstraint(const Layout& layout, const LogicalBuffer& buffer, - bool mandatory) - : LayoutConstraint(mandatory), layout_(layout), buffer_(&buffer) { + bool mandatory, bool dfs) + : LayoutConstraint(mandatory, dfs), layout_(layout), buffer_(&buffer) { CHECK(LayoutUtil::ValidateLayoutForShape(layout, buffer.shape()).ok()); } @@ -74,14 +74,17 @@ string BufferLayoutConstraint::ToString() const { OperandLayoutConstraint::OperandLayoutConstraint( const ShapeLayout& shape_layout, const HloInstruction* instruction, - int64 operand_no, bool mandatory) - : LayoutConstraint(mandatory), + int64 operand_no, bool mandatory, bool dfs) + : LayoutConstraint(mandatory, dfs), shape_layout_(shape_layout), instruction_(instruction), operand_no_(operand_no) { CHECK(shape_layout_.LayoutIsSet()); CHECK(ShapeUtil::Compatible(shape_layout.shape(), - instruction->operand(operand_no)->shape())); + instruction->operand(operand_no)->shape())) + << shape_layout.shape() << " is not compatible with " + << instruction->operand(operand_no)->shape() << " (for operand " + << operand_no << " of instruction " << instruction->ToString() << ")"; } string OperandLayoutConstraint::ToString() const { @@ -131,7 +134,7 @@ bool LayoutConstraints::OperandBufferForwarded( Status LayoutConstraints::SetBufferLayout(const Layout& layout, const LogicalBuffer& buffer, - bool mandatory) { + bool mandatory, bool dfs) { VLOG(3) << "SetBufferLayout : " << buffer << " : " << LayoutUtil::HumanString(layout); @@ -168,10 +171,11 @@ Status LayoutConstraints::SetBufferLayout(const Layout& layout, if (!overwrite) { iter = buffer_constraints_ .insert(std::make_pair( - &buffer, BufferLayoutConstraint(layout, buffer, mandatory))) + &buffer, + BufferLayoutConstraint(layout, buffer, mandatory, dfs))) .first; } else { - iter->second = BufferLayoutConstraint(layout, buffer, /*mandatory=*/true); + iter->second = BufferLayoutConstraint(layout, buffer, mandatory, dfs); } added_constraints_.push_back(&iter->second); @@ -185,7 +189,8 @@ Status LayoutConstraints::SetBufferLayout(const Layout& layout, Status LayoutConstraints::SetOperandLayout(const Shape& shape_with_layout, const HloInstruction* instruction, - int64 operand_no, bool mandatory) { + int64 operand_no, bool mandatory, + bool dfs) { VLOG(3) << "SetOperandLayout : " << instruction->name() << ", operand " << operand_no << " : " << ShapeUtil::HumanStringWithLayout(shape_with_layout); @@ -223,12 +228,12 @@ Status LayoutConstraints::SetOperandLayout(const Shape& shape_with_layout, if (iter == operand_constraints_.end()) { auto pair = std::make_pair( key, OperandLayoutConstraint(ShapeLayout(shape_with_layout), - instruction, operand_no, mandatory)); + instruction, operand_no, mandatory, dfs)); iter = operand_constraints_.insert(pair).first; } else { iter->second = OperandLayoutConstraint(ShapeLayout(shape_with_layout), instruction, - operand_no, /*mandatory=*/true); + operand_no, mandatory, dfs); } added_constraints_.push_back(&iter->second); @@ -237,16 +242,17 @@ Status LayoutConstraints::SetOperandLayout(const Shape& shape_with_layout, Status LayoutConstraints::SetArrayOperandLayout( const Layout& layout, const HloInstruction* instruction, int64 operand_no, - bool mandatory) { + bool mandatory, bool dfs) { const HloInstruction* operand = instruction->operand(operand_no); TF_RET_CHECK(ShapeUtil::IsArray(operand->shape())); Shape shape(operand->shape()); *shape.mutable_layout() = layout; TF_RETURN_IF_ERROR(LayoutUtil::ValidateLayoutInShape(shape)); - return SetOperandLayout(shape, instruction, operand_no, mandatory); + return SetOperandLayout(shape, instruction, operand_no, mandatory, dfs); } -Status LayoutConstraints::SetResultLayout(const Shape& shape_with_layout) { +Status LayoutConstraints::SetResultLayout(const Shape& shape_with_layout, + bool dfs) { VLOG(3) << "SetResultLayout : " << ShapeUtil::HumanStringWithLayout(shape_with_layout); @@ -264,14 +270,15 @@ Status LayoutConstraints::SetResultLayout(const Shape& shape_with_layout) { } result_constraint_.reset( - new ResultLayoutConstraint(ShapeLayout(shape_with_layout))); + new ResultLayoutConstraint(ShapeLayout(shape_with_layout), dfs)); added_constraints_.push_back(result_constraint_.get()); return Status::OK(); } Status LayoutConstraints::SetInstructionLayout( - const Shape& shape_with_layout, const HloInstruction* instruction) { + const Shape& shape_with_layout, const HloInstruction* instruction, + bool mandatory, bool dfs) { VLOG(3) << "SetInstructionLayout : " << instruction->name() << ", " << ShapeUtil::HumanStringWithLayout(shape_with_layout); @@ -287,8 +294,8 @@ Status LayoutConstraints::SetInstructionLayout( // instruction. return ShapeUtil::ForEachSubshapeWithStatus( shape_with_layout, - [this, instruction](const Shape& subshape, - const ShapeIndex& index) -> Status { + [this, instruction, mandatory](const Shape& subshape, + const ShapeIndex& index) -> Status { // The precondition for this method is that the instruction defines all // buffers in its output. auto buffers = @@ -297,7 +304,7 @@ Status LayoutConstraints::SetInstructionLayout( CHECK_EQ(buffers[0]->instruction(), instruction); if (ShapeUtil::IsArray(subshape)) { - return SetBufferLayout(subshape.layout(), *buffers[0]); + return SetBufferLayout(subshape.layout(), *buffers[0], mandatory); } else { return Status::OK(); } @@ -369,8 +376,9 @@ string LayoutConstraints::ToString() const { } Status LayoutAssignment::AddMandatoryConstraints( - const ComputationLayout& computation_layout, HloComputation* computation, - LayoutConstraints* constraints) { + const ComputationLayout& computation_layout, + const ChannelLayoutConstraints* channel_constraints, + HloComputation* computation, LayoutConstraints* constraints) { VLOG(3) << "Adding mandatory layout constraints to computation " << computation->name(); @@ -390,8 +398,7 @@ Status LayoutAssignment::AddMandatoryConstraints( // Constrain the input to the Outfeed instruction to be the expected // layout of the Outfeed. TF_RETURN_IF_ERROR(constraints->SetOperandLayout( - instruction->outfeed_shape(), instruction, 0, - /*mandatory=*/true)); + instruction->outfeed_shape(), instruction, 0)); } else if (instruction->opcode() == HloOpcode::kParameter) { // Parameter layouts must match the respective layout in // ComputationLayout. @@ -403,6 +410,37 @@ Status LayoutAssignment::AddMandatoryConstraints( TF_RETURN_IF_ERROR( constraints->SetInstructionLayout(*shape_with_layout, instruction)); } + + if (instruction->opcode() == HloOpcode::kSend || + instruction->opcode() == HloOpcode::kRecv) { + CHECK(channel_constraints) + << "Multi-module layout assignment requires ChannelLayoutConstraints"; + int64 channel_id = instruction->channel_id(); + if (!channel_constraints->IsChannelConstrained(channel_id)) { + continue; + } + if (instruction->opcode() == HloOpcode::kSend) { + // TODO(b/68493863): Change to use SetOperandLayout(). + const Shape send_buffer_shape = instruction->operand(0)->shape(); + TF_RET_CHECK(ShapeUtil::IsArray(send_buffer_shape)); + Shape new_buffer_shape = channel_constraints->LayoutShapeForChannel( + send_buffer_shape, instruction->channel_id()); + TF_RETURN_IF_ERROR(constraints->SetInstructionLayout( + new_buffer_shape, instruction->operand(0))); + } else { + const Shape recv_buffer_shape = + ShapeUtil::GetTupleElementShape(instruction->shape(), 0); + TF_RET_CHECK(ShapeUtil::IsArray(recv_buffer_shape)); + TF_ASSIGN_OR_RETURN( + const LogicalBuffer* buffer, + constraints->points_to_analysis().GetBufferDefinedAt(instruction, + {0})); + Shape new_shape = channel_constraints->LayoutShapeForChannel( + recv_buffer_shape, instruction->channel_id()); + TF_RETURN_IF_ERROR( + constraints->SetBufferLayout(new_shape.layout(), *buffer)); + } + } } // Constrain layouts of instructions which call computations which have @@ -422,7 +460,7 @@ Status LayoutAssignment::AddMandatoryConstraints( for (int64 i = 0; i < instruction->operand_count(); ++i) { TF_RETURN_IF_ERROR(constraints->SetOperandLayout( called_computation_layout.parameter_layout(i).shape(), instruction, - i, /*mandatory=*/true)); + i)); } } else if (instruction->opcode() == HloOpcode::kWhile) { // Layout of input and output of kWhile instruction must be equal and must @@ -473,8 +511,7 @@ Status LayoutAssignment::AddMandatoryConstraints( TF_RETURN_IF_ERROR(constraints->SetInstructionLayout( body_layout.result_shape(), instruction)); TF_RETURN_IF_ERROR(constraints->SetOperandLayout( - body_layout.result_shape(), instruction, 0, - /*mandatory=*/true)); + body_layout.result_shape(), instruction, 0)); } else if (instruction->opcode() == HloOpcode::kCustomCall) { if (!CustomCallRequiresMajorFirstLayout(instruction)) { continue; @@ -498,7 +535,7 @@ Status LayoutAssignment::AddMandatoryConstraints( operand_shape.element_type(), AsInt64Slice(operand_shape.dimensions())); TF_RETURN_IF_ERROR(constraints->SetOperandLayout( - row_major_operand_shape, instruction, i, /*mandatory=*/true)); + row_major_operand_shape, instruction, i)); } } } @@ -690,8 +727,11 @@ Status LayoutAssignment::CheckLayouts(HloModule* module) { return Status::OK(); } -LayoutAssignment::LayoutAssignment(ComputationLayout* entry_computation_layout) - : entry_computation_layout_(entry_computation_layout) { +LayoutAssignment::LayoutAssignment( + ComputationLayout* entry_computation_layout, + ChannelLayoutConstraints* channel_constraints) + : entry_computation_layout_(entry_computation_layout), + channel_layout_constraints_(channel_constraints) { VLOG(1) << "entry computation layout given to layout assignment: " << entry_computation_layout_->ToString(); // Layouts of all parameter instructions must be set. @@ -869,7 +909,11 @@ Status LayoutAssignment::PropagateConstraints(LayoutConstraints* constraints) { auto add_new_constraints_to_worklist = [constraints, &worklist]() { // Add constraints to the front of the deque for DFS ordering. for (auto* constraint : constraints->ConsumeAddedConstraints()) { - worklist.push_front(constraint); + if (constraint->dfs()) { + worklist.push_front(constraint); + } else { + worklist.push_back(constraint); + } } }; add_new_constraints_to_worklist(); @@ -1198,7 +1242,8 @@ Status CopyOperandIfLayoutsDiffer(const ShapeLayout& operand_layout, // instruction itself. Status SetFusionLayouts(HloInstruction* fusion) { TF_RET_CHECK(fusion->opcode() == HloOpcode::kFusion); - for (auto* fused_instruction : fusion->fused_instructions()) { + for (auto* fused_instruction : + fusion->fused_instructions_computation()->MakeInstructionPostOrder()) { if (fused_instruction->opcode() == HloOpcode::kParameter) { const HloInstruction* fusion_operand = fusion->operand(fused_instruction->parameter_number()); @@ -1213,11 +1258,22 @@ Status SetFusionLayouts(HloInstruction* fusion) { ShapeUtil::Compatible(fusion->shape(), fused_instruction->shape())); TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( fusion->shape(), fused_instruction->mutable_shape())); - } else if (fused_instruction->opcode() != HloOpcode::kConstant && - fused_instruction->opcode() != HloOpcode::kGetTupleElement && - fused_instruction->opcode() != HloOpcode::kInfeed) { - // Internal fused instructions with the exception of constants - // and infeed need no layout. + } else if (fused_instruction->opcode() == HloOpcode::kGetTupleElement) { + // A GTE inherits its layout from its operand (which should ultimately be + // a parameter). + TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( + fused_instruction->operand(0)->shape().tuple_shapes( + fused_instruction->tuple_index()), + fused_instruction->mutable_shape())); + } else if (fused_instruction->opcode() == HloOpcode::kConstant) { + // Give constants the layout of their literal. + TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( + fused_instruction->literal().shape(), + fused_instruction->mutable_shape())); + } else if (fused_instruction->opcode() == HloOpcode::kInfeed) { + // Nop; leave the infeed layout alone. + } else { + // Other instructions don't have layouts inside of fusion nodes. LayoutUtil::ClearLayout(fused_instruction->mutable_shape()); } } @@ -1321,39 +1377,26 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints, Status LayoutAssignment::RunOnComputation( const ComputationLayout& computation_layout, const TuplePointsToAnalysis& points_to_analysis, - HloComputation* computation) { + HloComputation* computation, + ChannelLayoutConstraints* channel_constraints) { DCHECK(computation_layout.LayoutIsSet()); InsertOrDie(&computation_layouts_, computation, computation_layout); VLOG(2) << "LayoutAssignment::RunOnComputation(" << computation->name() << ")"; VLOG(2) << " ComputationLayout = " << computation_layout.ToString(); - // Clear existing layouts of the instructions. All layouts must be assigned by - // the LayoutAssignment pass, except for Infeed, Outfeed, Parameters and the - // computation result. The latter two are specified in computation_layout, so - // we only need to keep the existing layouts for Infeed and Outfeed. Clearing - // the layouts here avoids hiding potential bugs in the layout assignment pass - // that may accidently use the existing layout. - for (HloInstruction* instruction : computation->instructions()) { - if (instruction->opcode() == HloOpcode::kInfeed || - instruction->opcode() == HloOpcode::kOutfeed) { - continue; - } - LayoutUtil::ClearLayout(instruction->mutable_shape()); - } - // Construct LayoutConstraints with all layout constraints of the computation. LayoutConstraints constraints(points_to_analysis, computation); // Add constraints required for correctness on all backends (eg, entry // parameter layout constraints). - TF_RETURN_IF_ERROR( - AddMandatoryConstraints(computation_layout, computation, &constraints)); + TF_RETURN_IF_ERROR(AddMandatoryConstraints( + computation_layout, channel_constraints, computation, &constraints)); // Add any backend-specific constraints. TF_RETURN_IF_ERROR(AddBackendConstraints(&constraints)); - // Propagates layouts from an HLO to its neighbors. + // Propagates layouts from mandatory and backend constraints. TF_RETURN_IF_ERROR(PropagateConstraints(&constraints)); // While any unconstrained buffers remain, pick an arbitrary buffer, give it a @@ -1387,7 +1430,20 @@ Status LayoutAssignment::RunOnComputation( // All logical buffers should have constraints at this point. All that // remains is assign the constraints to the buffers and infer layouts for // aliased buffers. - return AssignLayouts(constraints, computation); + TF_RETURN_IF_ERROR(AssignLayouts(constraints, computation)); + + // Record the layouts assigned for any communication ops in + // channel_constraints so that they are constrained for future modules. + for (HloInstruction* instruction : computation->instructions()) { + if (instruction->opcode() == HloOpcode::kSend) { + channel_constraints->ConstrainChannel( + instruction->channel_id(), instruction->operand(0)->shape().layout()); + } else if (instruction->opcode() == HloOpcode::kRecvDone) { + channel_constraints->ConstrainChannel(instruction->channel_id(), + instruction->shape().layout()); + } + } + return Status::OK(); } StatusOr LayoutAssignment::Run(HloModule* module) { @@ -1405,20 +1461,35 @@ StatusOr LayoutAssignment::Run(HloModule* module) { // Assign layouts to computations in an order such that a callee computation // is handled before its caller computation. This ensures that the layout of // all callers of a computation will agree. + std::list computation_post_order = + module->MakeComputationPostOrder(); for (auto* computation : module->MakeComputationPostOrder()) { - if (computation == module->entry_computation()) { - TF_RETURN_IF_ERROR(RunOnComputation(*entry_computation_layout_, - *points_to_analysis, - module->entry_computation())); - } else if (computation->IsFusionComputation()) { + if (computation->IsFusionComputation()) { continue; + } + // Clear existing layouts of the instructions. All layouts must be assigned + // by the LayoutAssignment pass, except for those on infeeds, parameters, + // and the computation result. The latter two are specified in + // computation_layout, so we only need to keep the existing layouts for + // infeeds. Clearing the layouts here avoids hiding potential bugs in the + // layout assignment pass that may accidently use the existing layout. + for (HloInstruction* instruction : computation->instructions()) { + if (instruction->opcode() != HloOpcode::kInfeed) { + LayoutUtil::ClearLayout(instruction->mutable_shape()); + } + } + if (computation == module->entry_computation()) { + TF_RETURN_IF_ERROR(RunOnComputation( + *entry_computation_layout_, *points_to_analysis, + module->entry_computation(), channel_layout_constraints_)); } else { ComputationLayout computation_layout(computation->ComputeProgramShape()); // Setting all embedded computations to the default layout is potentially // suboptimal. computation_layout.SetToDefaultLayout(); TF_RETURN_IF_ERROR(RunOnComputation(computation_layout, - *points_to_analysis, computation)); + *points_to_analysis, computation, + channel_layout_constraints_)); } } diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h index 77cd454e52d48cfb9537d4d3f4879ff235849f99..29018584487cabfd740d7914625c2a50f552d6ff 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.h +++ b/tensorflow/compiler/xla/service/layout_assignment.h @@ -46,7 +46,8 @@ namespace xla { // gathered together in LayoutConstraints object. class LayoutConstraint { public: - LayoutConstraint(bool mandatory) : mandatory_(mandatory) {} + LayoutConstraint(bool mandatory, bool dfs) + : mandatory_(mandatory), dfs_(dfs) {} virtual ~LayoutConstraint() = default; virtual string ToString() const = 0; @@ -54,8 +55,12 @@ class LayoutConstraint { // True if this constraint cannot be overwritten by a different constraint. bool mandatory() const { return mandatory_; } + // When true, propagate in DFS. When false, constraint will propagate in BFS. + bool dfs() const { return dfs_; } + private: bool mandatory_; + bool dfs_; }; std::ostream& operator<<(std::ostream& out, const LayoutConstraint& constraint); @@ -65,7 +70,7 @@ std::ostream& operator<<(std::ostream& out, const LayoutConstraint& constraint); class BufferLayoutConstraint : public LayoutConstraint { public: BufferLayoutConstraint(const Layout& layout, const LogicalBuffer& buffer, - bool mandatory); + bool mandatory, bool dfs); const LogicalBuffer& buffer() const { return *buffer_; } const Layout& layout() const { return layout_; } @@ -86,7 +91,7 @@ class OperandLayoutConstraint : public LayoutConstraint { public: OperandLayoutConstraint(const ShapeLayout& shape_layout, const HloInstruction* instruction, int64 operand_no, - bool mandatory); + bool mandatory, bool dfs); const ShapeLayout& shape_layout() const { return shape_layout_; } const HloInstruction* instruction() const { return instruction_; } @@ -106,8 +111,10 @@ class OperandLayoutConstraint : public LayoutConstraint { // Constraint on the layout of the result of the entry computation. class ResultLayoutConstraint : public LayoutConstraint { public: - explicit ResultLayoutConstraint(const ShapeLayout& shape_layout) - : LayoutConstraint(/*mandatory=*/true), shape_layout_(shape_layout) {} + explicit ResultLayoutConstraint(const ShapeLayout& shape_layout, + bool dfs = false) + : LayoutConstraint(/*mandatory=*/true, dfs), + shape_layout_(shape_layout) {} const ShapeLayout& shape_layout() const { return shape_layout_; } string ToString() const override; @@ -157,23 +164,25 @@ class LayoutConstraints { // operand of the instruction, or the layout of the result of the computation, // respectively. Status SetBufferLayout(const Layout& layout, const LogicalBuffer& buffer, - bool mandatory = true); + bool mandatory = true, bool dfs = true); Status SetOperandLayout(const Shape& shape_with_layout, const HloInstruction* instruction, int64 operand_no, - bool mandatory = true); - Status SetResultLayout(const Shape& shape_with_layout); + bool mandatory = true, bool dfs = true); + Status SetResultLayout(const Shape& shape_with_layout, bool dfs = true); // Convenience wrapper around SetOperandLayout for setting the layout of a // operand using a Layout object. The operand must be array-shaped. Status SetArrayOperandLayout(const Layout& layout, const HloInstruction* instruction, - int64 operand_no, bool mandatory = true); + int64 operand_no, bool mandatory = true, + bool dfs = true); // Convenience wrapper around SetBufferLayout. Sets the layouts of all buffers // created by the instruction to the layouts in the given shape. The // instruction must define every logical buffer in its output. Status SetInstructionLayout(const Shape& shape_with_layout, - const HloInstruction* instruction); + const HloInstruction* instruction, + bool mandatory = true, bool dfs = true); // Returns true if any buffer in the given operand is forwarded to the output // of the given instruction. For example, the Tuple instruction forwards the @@ -215,13 +224,62 @@ class LayoutConstraints { HloComputation* computation_; }; +// Contains constraints on the layout of channels; sends and recvs. +class ChannelLayoutConstraints { + public: + // Construct an empty constraint set. + ChannelLayoutConstraints() {} + + // Returns true if channel_id has a layout constraint. + bool IsChannelConstrained(int64 channel_id) const { + return constraints_.count(channel_id) > 0; + } + + // Given `shape`, apply the layout for `channel_id`. `channel_id` must already + // be constrained. + Shape LayoutShapeForChannel(Shape shape, int64 channel_id) const { + CHECK(IsChannelConstrained(channel_id)); + *shape.mutable_layout() = constraints_.at(channel_id); + return shape; + } + + // Returns the layout constraint for `channel_id`, which must already be + // constrained. + Layout LayoutForChannel(int64 channel_id) const { + CHECK(IsChannelConstrained(channel_id)); + return constraints_.at(channel_id); + } + + // Adds a new layout constraint for `channel_id`. If a constraint for + // `channel_id` already exists, this operation requires that the new layout is + // the same as the previously constrained layout. + void ConstrainChannel(int64 channel_id, const Layout& layout) { + CHECK(!IsChannelConstrained(channel_id) || + LayoutUtil::Equal(layout, constraints_[channel_id])); + constraints_[channel_id] = layout; + } + + private: + std::unordered_map constraints_; +}; + // HLO pass which assigns layouts to all instructions in the HLO module while // satisfying all necessary invariants and minimizing cost. class LayoutAssignment : public HloPassInterface { public: // entry_computation_layout is modified to populate a layout for the result in // the case that no particular layout is requested. - explicit LayoutAssignment(ComputationLayout* entry_computation_layout); + // + // channel_constraints is both an input and output. Any sends or recvs that + // are present in channel_constraints will be layed out as constrained. Any + // unconstrained sends or recvs will be layed out as locally optimal and their + // layout will be added as a constraint to channel_constraints. + // + // If channel_constraints is nullptr, no kSend or kRecvs must be contained + // within any module passed to `Run`. + explicit LayoutAssignment( + ComputationLayout* entry_computation_layout, + ChannelLayoutConstraints* channel_constraints = nullptr); ~LayoutAssignment() override {} tensorflow::StringPiece name() const override { return "layout-assignment"; } @@ -247,7 +305,7 @@ class LayoutAssignment : public HloPassInterface { const ResultLayoutConstraint& layout_constraint, LayoutConstraints* constraints); - // By default LayoutAssignment ensures that inputs and ouptuts of CustomCalls + // By default LayoutAssignment ensures that inputs and outputs of CustomCalls // have the "major-first" layout (i.e. {n, n-1, ..., 0}). // // If this function returns true, LayoutAssignment does not set a layout for @@ -296,9 +354,10 @@ class LayoutAssignment : public HloPassInterface { private: // Adds constraints which must be satisfied for correctness on all // backends. Called once prior to propagating constraints. - Status AddMandatoryConstraints(const ComputationLayout& computation_layout, - HloComputation* computation, - LayoutConstraints* constraints); + Status AddMandatoryConstraints( + const ComputationLayout& computation_layout, + const ChannelLayoutConstraints* channel_constraints, + HloComputation* computation, LayoutConstraints* constraints); // This method can be overridden to add backend-specific constraints to the // layout of the instructions of a computation. This method is called after @@ -314,7 +373,8 @@ class LayoutAssignment : public HloPassInterface { // constrained. Status RunOnComputation(const ComputationLayout& computation_layout, const TuplePointsToAnalysis& points_to_analysis, - HloComputation* computation); + HloComputation* computation, + ChannelLayoutConstraints* channel_constraints); // Assign layouts to the instructions of a computation which satisfy the given // layout constraints. Copies may be added to satisfy the constraints. The @@ -333,6 +393,7 @@ class LayoutAssignment : public HloPassInterface { Status CheckLayouts(HloModule* module); ComputationLayout* entry_computation_layout_; + ChannelLayoutConstraints* channel_layout_constraints_; protected: // Map containing the layouts of all computations assigned so diff --git a/tensorflow/compiler/xla/service/layout_assignment_test.cc b/tensorflow/compiler/xla/service/layout_assignment_test.cc index d51c0d1dfb727801d6d2a8328eba60838373479f..e269a13459f1146f1d2952870399827d9e705e38 100644 --- a/tensorflow/compiler/xla/service/layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/layout_assignment_test.cc @@ -35,9 +35,11 @@ limitations under the License. #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/tests/test_utils.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/lib/gtl/array_slice.h" namespace op = xla::testing::opcode_matchers; @@ -587,5 +589,74 @@ TEST_F(LayoutAssignmentTest, TransposeToBitcastToUser) { EXPECT_TRUE(ShapeUtil::TransposeIsBitcast(transpose->operand(0)->shape(), transpose->shape(), {2, 3, 0, 1})); } + +// A GTE inside of a fusion node inherits the layout of its operand (which +// should, if we keep following operands, eventually be a parameter). +TEST_F(LayoutAssignmentTest, GTEInheritsLayoutFromOperand) { + const char* module_str = R"( + HloModule test_module + + fused_computation { + fparam = (f32[2,2,2], (f32[2,2,2], f32[2,2,2])) parameter(0) + gte0 = f32[2,2,2] get-tuple-element(fparam), index=0 + gte1 = (f32[2,2,2], f32[2,2,2]) get-tuple-element(fparam), index=1 + gte1a = f32[2,2,2] get-tuple-element(gte1), index=0 + gte1b = f32[2,2,2] get-tuple-element(gte1), index=1 + add = f32[2,2,2] add(gte1a, gte1b) + ROOT fresult = f32[2,2,2] add(gte0, add) + } + + ENTRY entry_computation { + param = (f32[2,2,2], (f32[2,2,2], f32[2,2,2])) parameter(0) + ROOT fusion = + f32[2,2,2] fusion(param), kind=kLoop, calls=fused_computation + } + )"; + + auto module = tools::Parse(module_str).ValueOrDie(); + ComputationLayout computation_layout( + module->entry_computation()->ComputeProgramShape()); + Shape param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShapeWithLayout(F32, {2, 2, 2}, {0, 1, 2}), + ShapeUtil::MakeTupleShape({ + ShapeUtil::MakeShapeWithLayout(F32, {2, 2, 2}, {1, 2, 0}), + ShapeUtil::MakeShapeWithLayout(F32, {2, 2, 2}, {2, 0, 1}), + })}); + TF_ASSERT_OK( + computation_layout.mutable_parameter_layout(0)->CopyLayoutFromShape( + param_shape)); + computation_layout.mutable_result_layout()->ResetLayout( + LayoutUtil::MakeLayout({2, 1, 0})); + AssignLayouts(module.get(), &computation_layout); + + HloComputation* fused_computation = *std::find_if( + module->computations().begin(), module->computations().end(), + [](const HloComputation* c) { return c->name() == "fused_computation"; }); + + auto fused_instr = [&](const string& name) { + auto it = std::find_if( + fused_computation->instructions().begin(), + fused_computation->instructions().end(), + [&](const HloInstruction* i) { return i->name() == name; }); + CHECK(it != fused_computation->instructions().end()); + return *it; + }; + + EXPECT_THAT(fused_instr("gte0")->shape().layout().minor_to_major(), + ElementsAre(0, 1, 2)); + EXPECT_THAT( + fused_instr("gte1")->shape().tuple_shapes(0).layout().minor_to_major(), + ElementsAre(1, 2, 0)); + EXPECT_THAT( + fused_instr("gte1")->shape().tuple_shapes(1).layout().minor_to_major(), + ElementsAre(2, 0, 1)); + EXPECT_THAT(fused_instr("gte1a")->shape().layout().minor_to_major(), + ElementsAre(1, 2, 0)); + EXPECT_THAT(fused_instr("gte1b")->shape().layout().minor_to_major(), + ElementsAre(2, 0, 1)); + EXPECT_THAT(fused_instr("fresult")->shape().layout().minor_to_major(), + ElementsAre(2, 1, 0)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/llvm_compiler.cc b/tensorflow/compiler/xla/service/llvm_compiler.cc index 34f3419269abbc73cd0ddb13c723a8da38ab19ff..68c35c0c1f8717957d5ab8cb067b8f8892aa426d 100644 --- a/tensorflow/compiler/xla/service/llvm_compiler.cc +++ b/tensorflow/compiler/xla/service/llvm_compiler.cc @@ -15,11 +15,15 @@ limitations under the License. #include "tensorflow/compiler/xla/service/llvm_compiler.h" +#ifdef __FAST_MATH__ +#error "Don't build XLA with -ffast-math" +#endif + namespace xla { StatusOr>> LLVMCompiler::Compile( std::vector> modules, - std::vector> - stream_execs) { + std::vector> stream_execs, + DeviceMemoryAllocator* device_allocator) { std::vector> result; for (size_t i = 0; i < modules.size(); i++) { if (stream_execs[i].size() != 1) { @@ -27,10 +31,12 @@ StatusOr>> LLVMCompiler::Compile( "Model partitioning not implemented for the CPU/GPU compilers!"); } - TF_ASSIGN_OR_RETURN( - modules[i], RunHloPasses(std::move(modules[i]), stream_execs[i][0])); + TF_ASSIGN_OR_RETURN(modules[i], + RunHloPasses(std::move(modules[i]), stream_execs[i][0], + device_allocator)); TF_ASSIGN_OR_RETURN(std::unique_ptr executable, - RunBackend(std::move(modules[i]), stream_execs[i][0])); + RunBackend(std::move(modules[i]), stream_execs[i][0], + device_allocator)); result.push_back(std::move(executable)); } diff --git a/tensorflow/compiler/xla/service/llvm_compiler.h b/tensorflow/compiler/xla/service/llvm_compiler.h index c5393cef4f961c5d04c32d0d4291732b8ec702f1..d74e81bb7f622ac5e89203a3d02ca5ad839da07e 100644 --- a/tensorflow/compiler/xla/service/llvm_compiler.h +++ b/tensorflow/compiler/xla/service/llvm_compiler.h @@ -60,17 +60,20 @@ class LLVMCompiler : public Compiler { // Bring in // StatusOr> RunBackend( // std::unique_ptr module, - // perftools::gputools::StreamExecutor* stream_exec) + // perftools::gputools::StreamExecutor* stream_exec, + // DeviceMemoryAllocator* device_allocator) // StatusOr> RunHloPasses( // std::unique_ptr module, - // perftools::gputools::StreamExecutor* stream_exec) + // perftools::gputools::StreamExecutor* stream_exec, + // DeviceMemoryAllocator* device_allocator) using Compiler::RunBackend; using Compiler::RunHloPasses; StatusOr>> Compile( std::vector> modules, std::vector> - stream_execs) override; + stream_execs, + DeviceMemoryAllocator* device_allocator) override; protected: ModuleHook user_pre_optimization_hook_; diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD index d878061f724de1c82f8285b0f082d0be4d5778df..37261ed1e665ebed9685751161a412ad114a9e96 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/BUILD +++ b/tensorflow/compiler/xla/service/llvm_ir/BUILD @@ -48,11 +48,13 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service:name_uniquer", "//tensorflow/core:lib", "@llvm//:core", "@llvm//:support", "@llvm//:target", + "@llvm//:transform_utils", ], ) @@ -156,18 +158,6 @@ cc_library( ], ) -cc_library( - name = "vector_support_library", - srcs = ["vector_support_library.cc"], - hdrs = ["vector_support_library.h"], - deps = [ - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", - "@llvm//:core", - ], -) - cc_library( name = "kernel_support_library", srcs = ["kernel_support_library.cc"], diff --git a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h index 9ad7cd82cb8ca862fd7acec3dfb12c9fd61f6e27..b3b6026ef17daa184c0a015fdea618597ef068b3 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h +++ b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h @@ -32,8 +32,23 @@ limitations under the License. namespace xla { -// Unlike IrEmitter, this creates host functions which emit IR to generate the -// output element at the given index. It is used to generate fused operations. +// FusedIrEmitter is used to generate code for fusion nodes. +// +// Unlike IrEmitter and its ilk, which directly create LLVM IR in an LLVM +// Module, FusedIrEmitter is better understood as "IR generator generator". +// FusedIrEmitter recursively creates a generator (a host function) which the +// compiler can invoke at a later time. Invoking the generator emits LLVM IR +// that, when run, produces the value at a particular index of the output. +// +// After building this generator, the compiler creates a loop (or its moral +// equivalent, e.g. a GPU kernel) and calls the generator from within the loop. +// This generates code that produces each element of the output. +// +// This class handles both vanilla fusion and multi-output fusion. In the MOF +// case, the fusion node ends with a kTuple instruction, and the generator +// created produces an LLVM struct with N elements, one for each element of the +// arrays in the tuple. It follows that the arrays in the tuple must have the +// same length. class FusedIrEmitter : public DfsHloVisitorWithDefault { public: using Generator = llvm_ir::ElementGenerator; diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc index c558f7388cab587b5858d0594cdb2f3c41d75562..6384c7f46f5ebbedaeda232b40095611a5d738a4 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc @@ -39,13 +39,27 @@ IrArray::Index::Index(llvm::Value* linear, const Shape& shape, << "Shape " << ShapeUtil::HumanStringWithLayout(shape) << " should have a layout."; int64 divisor = 1; - for (int64 dimension : LayoutUtil::MinorToMajor(layout_)) { + for (int64 i = 0; i < layout_.minor_to_major_size(); ++i) { + int64 dimension = layout_.minor_to_major(i); int64 size_of_current_dimension = shape.dimensions(dimension); - // Emit IR instructions that compute - // (linear_index / divisor) % current_dimension - multidim_[dimension] = ir_builder->CreateURem( - ir_builder->CreateUDiv(linear, ir_builder->getInt64(divisor)), - ir_builder->getInt64(size_of_current_dimension)); + + // If i is not the last dimension, compute + // (linear_index / divisor) % current_dimension. + // If i is the last dimension, we can skip the mod, because we assume that + // linear is in bounds. + // + // TODO(jlebar): We could add bounds checks here and elsewhere in this file, + // guarded under some sort of xla-memcheck flag. This might be particularly + // useful because cuda-memcheck can't help us much in XLA: Most of our + // memory lives in one big allocation, so cuda-memcheck can't detect + // out-of-bounds accesses. + auto* quot = ir_builder->CreateUDiv(linear, ir_builder->getInt64(divisor)); + if (i < layout_.minor_to_major_size() - 1) { + multidim_[dimension] = ir_builder->CreateURem( + quot, ir_builder->getInt64(size_of_current_dimension)); + } else { + multidim_[dimension] = quot; + } divisor *= size_of_current_dimension; } } diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h index 827e092a3fa9116c461716b27c309033f7988745..1c00b2aabd182da72e78d2c9c01cbe70cfd8e33c 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h +++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_KERNEL_SUPPORT_LIBRARY_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_KERNEL_SUPPORT_LIBRARY_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_KERNEL_SUPPORT_LIBRARY_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_KERNEL_SUPPORT_LIBRARY_H_ #include @@ -179,4 +179,4 @@ class KernelSupportLibrary { }; } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_KERNEL_SUPPORT_LIBRARY_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_KERNEL_SUPPORT_LIBRARY_H_ diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc index 61c47a0b6eca38db5d78dc622a8cf909f6cf14ee..5c1866311d1ae1e0c33ab061ee326d86d647a908 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc @@ -20,9 +20,11 @@ limitations under the License. #include #include "llvm/IR/DerivedTypes.h" +#include "llvm/IR/GlobalValue.h" #include "llvm/IR/MDBuilder.h" #include "llvm/IR/Operator.h" #include "llvm/Target/TargetOptions.h" +#include "llvm/Transforms/Utils/Cloning.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/name_uniquer.h" @@ -61,6 +63,16 @@ llvm::StringRef AsStringRef(tensorflow::StringPiece str) { return llvm::StringRef(str.data(), str.size()); } +std::unique_ptr DropConstantInitializers( + const llvm::Module& module) { + std::unique_ptr cloned_module = CloneModule(module); + for (llvm::GlobalVariable& global_var : cloned_module->globals()) { + global_var.setInitializer(nullptr); + global_var.setLinkage(llvm::GlobalValue::LinkageTypes::ExternalLinkage); + } + return cloned_module; +} + string DumpModuleToString(const llvm::Module& module) { std::string buffer_string; llvm::raw_string_ostream ostream(buffer_string); @@ -150,6 +162,8 @@ llvm::Type* PrimitiveTypeToIrType(PrimitiveType element_type, // addition to an addition on this type (int16) - this is just the type // used for storage. return llvm::Type::getInt16Ty(module->getContext()); + case F16: + return llvm::Type::getHalfTy(module->getContext()); case S32: case U32: return llvm::Type::getInt32Ty(module->getContext()); @@ -292,6 +306,11 @@ llvm::Constant* LiteralToConstant(const Literal& literal, int64 dimension_index, ir_element_type, tensorflow::bit_cast(literal.Get(*multi_index))); break; + case F16: + value = llvm::ConstantFP::get( + ir_element_type, + static_cast(literal.Get(*multi_index))); + break; case F64: value = llvm::ConstantFP::get(ir_element_type, literal.Get(*multi_index)); @@ -665,6 +684,19 @@ static string GetProcessUniqueIrFileName(tensorflow::StringPiece prefix) { return uniquer->GetUniqueName(prefix); } +static Status CreateAndWriteStringToFile(const string& directory_name, + const string& file_name, + const string& text) { + std::unique_ptr f; + TF_RETURN_IF_ERROR( + tensorflow::Env::Default()->RecursivelyCreateDir(directory_name)); + TF_RETURN_IF_ERROR( + tensorflow::Env::Default()->NewWritableFile(file_name, &f)); + TF_RETURN_IF_ERROR(f->Append(text)); + TF_RETURN_IF_ERROR(f->Close()); + return Status::OK(); +} + Status DumpIRToDirectory(const string& directory_name, const string& hlo_module_name, const llvm::Module& llvm_module, bool optimized) { @@ -679,13 +711,17 @@ Status DumpIRToDirectory(const string& directory_name, directory_name, tensorflow::strings::StrCat(unique_and_safe_file_name, ".ll")); - std::unique_ptr f; - TF_RETURN_IF_ERROR( - tensorflow::Env::Default()->RecursivelyCreateDir(directory_name)); - TF_RETURN_IF_ERROR( - tensorflow::Env::Default()->NewWritableFile(ir_file_name, &f)); - TF_RETURN_IF_ERROR(f->Append(DumpModuleToString(llvm_module))); - return f->Close(); + // For some models the embedded constants can be huge, so also dump the module + // with the constants stripped to get IR that is easier to manipulate. + string ir_no_constant_initializers_file_name = tensorflow::io::JoinPath( + directory_name, + tensorflow::strings::StrCat(unique_and_safe_file_name, "-noconst.ll")); + + TF_RETURN_IF_ERROR(CreateAndWriteStringToFile( + directory_name, ir_file_name, DumpModuleToString(llvm_module))); + return CreateAndWriteStringToFile( + directory_name, ir_no_constant_initializers_file_name, + DumpModuleToString(*DropConstantInitializers(llvm_module))); } llvm::Function* CreateFunction(llvm::FunctionType* function_type, @@ -715,5 +751,31 @@ llvm::Function* CreateFunction(llvm::FunctionType* function_type, return function; } +void InitializeLLVMCommandLineOptions(const HloModuleConfig& config) { + auto options = config.debug_options().xla_backend_extra_options(); + if (!options.empty()) { + std::vector fake_argv_storage; + fake_argv_storage.push_back(""); + for (const auto& it : options) { + // Skip options the XLA backend itself consumes. + if (!tensorflow::StringPiece(it.first).starts_with("xla_")) { + if (it.second.empty()) { + fake_argv_storage.push_back(it.first); + } else { + fake_argv_storage.push_back(it.first + "=" + it.second); + } + } + } + + VLOG(2) << "Passing argv to LLVM:"; + std::vector fake_argv; + for (const auto& s : fake_argv_storage) { + fake_argv.push_back(s.c_str()); + VLOG(2) << s; + } + llvm::cl::ParseCommandLineOptions(fake_argv.size(), &fake_argv[0]); + } +} + } // namespace llvm_ir } // namespace xla diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h index 6bdc6a01a2b487df3dd80a02e67f5bcf62dead31..4a10ec466dae6fdb56546fb8d8b353dcff6a5b8d 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h @@ -29,6 +29,7 @@ limitations under the License. #include "llvm/Support/raw_ostream.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/stringpiece.h" @@ -287,6 +288,10 @@ llvm::Function* CreateFunction(llvm::FunctionType* function_type, tensorflow::StringPiece name, llvm::Module* module); +// Extracts the xla_backend_extra_options from `config` and passes those that +// don't start with xla_ to LLVM. +void InitializeLLVMCommandLineOptions(const HloModuleConfig& config); + } // namespace llvm_ir } // namespace xla diff --git a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc index a5f7c850c33757fe8d48567ade35544d81224e46..b6b918ec78a27b90325f72eea14b97f9aee43c54 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc @@ -51,37 +51,40 @@ LoopEmitter::LoopEmitter(const ElementGenerator& target_element_generator, shape_(target_array.GetShape()), ir_builder_(ir_builder) {} +static LoopEmitter::BodyEmitter MakeBodyEmitterForMultiOutputFusion( + const ElementGenerator& target_element_generator, + const std::vector& target_arrays, llvm::IRBuilder<>* ir_builder) { + return [=](const llvm_ir::IrArray::Index array_index) { + TF_ASSIGN_OR_RETURN(llvm::Value * target_element, + target_element_generator(array_index)); + CHECK(target_element->getType()->isStructTy()) + << "This BodyEmitter is for multi-output fusion, but target element " + "generator does not produce values of struct type."; + CHECK_EQ(target_element->getType()->getStructNumElements(), + target_arrays.size()); + + for (int64 i = 0; i < target_arrays.size(); ++i) { + target_arrays[i].EmitWriteArrayElement( + array_index, ir_builder->CreateExtractValue(target_element, i), + ir_builder); + } + return Status::OK(); + }; +} + LoopEmitter::LoopEmitter(const ElementGenerator& target_element_generator, tensorflow::gtl::ArraySlice target_arrays, llvm::IRBuilder<>* ir_builder) - : body_emitter_([=](const llvm_ir::IrArray::Index array_index) - -> ::tensorflow::Status { - // Convert target_element_generator to a BodyEmitter. - TF_ASSIGN_OR_RETURN(llvm::Value * target_element, - target_element_generator(array_index)); - if (target_arrays.size() == 1) { - target_arrays[0].EmitWriteArrayElement(array_index, target_element, - ir_builder); - return tensorflow::Status::OK(); - } - - for (int64 i = 0; i < target_arrays.size(); ++i) { - target_arrays[i].EmitWriteArrayElement( - array_index, ir_builder_->CreateExtractValue(target_element, i), - ir_builder); - } - return tensorflow::Status::OK(); - }), + : body_emitter_(MakeBodyEmitterForMultiOutputFusion( + target_element_generator, + std::vector(target_arrays.begin(), target_arrays.end()), + ir_builder)), + shape_(target_arrays[0].GetShape()), ir_builder_(ir_builder) { - if (target_arrays.size() > 1) { - // The sanity check for multiple outputs. - shape_ = target_arrays[0].GetShape(); - for (int64 i = 1; i < target_arrays.size(); ++i) { - const Shape& element_shape = target_arrays[i].GetShape(); - CHECK(ShapeUtil::SameDimensions(shape_, element_shape)); - } - } else { - shape_ = target_arrays[0].GetShape(); + // Sanity check: In multi-output fusion, all shapes produced must have the + // same dimensions. + for (const IrArray& array : target_arrays) { + CHECK(ShapeUtil::SameDimensions(shape_, array.GetShape())); } } diff --git a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h index 1ef1dc246442041698d96f6aff48794c8788f1d1..0fc528439a0d5bf8382dfcf2d8b3051f8900bf1d 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h +++ b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h @@ -47,10 +47,16 @@ class LoopEmitter { // element of the given target array. LoopEmitter(const ElementGenerator& target_element_generator, const IrArray& target_array, llvm::IRBuilder<>* ir_builder); - // Same as previous method except emits multiple targets in an array. + + // Constructs a LoopEmitter that emits one element into each of N separate + // arrays on each iteration of the loop. + // + // This is used for multi-output fusion. target_element_generator must + // produce an LLVM struct with N elements. LoopEmitter(const ElementGenerator& target_element_generator, tensorflow::gtl::ArraySlice target_arrays, llvm::IRBuilder<>* ir_builder); + LoopEmitter(const LoopEmitter&) = delete; LoopEmitter& operator=(const LoopEmitter&) = delete; virtual ~LoopEmitter() = default; diff --git a/tensorflow/compiler/xla/service/llvm_ir/ops.h b/tensorflow/compiler/xla/service/llvm_ir/ops.h index f72f482e3128c61e53cc454e7da8b5795ba6f695..175b081e84d31779b15560cb0998011fe046ca01 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ops.h +++ b/tensorflow/compiler/xla/service/llvm_ir/ops.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_OPS_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_OPS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_OPS_H_ #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/elemental_ir_emitter.h" @@ -90,4 +90,4 @@ Status EmitParallelFusedDynamicUpdateSliceInPlace( } // namespace llvm_ir } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_OPS_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_OPS_H_ diff --git a/tensorflow/compiler/xla/service/local_service.cc b/tensorflow/compiler/xla/service/local_service.cc index 4071b948a5f94bcc2e87d8bb3b9533fb3b1d2cb1..07f989d4faea199e812e54d2ae74d3ff9e7fa19a 100644 --- a/tensorflow/compiler/xla/service/local_service.cc +++ b/tensorflow/compiler/xla/service/local_service.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/client/executable_build_options.h" #include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/backend.h" @@ -71,7 +72,7 @@ LocalService::LocalService(const ServiceOptions& options, StatusOr> LocalService::CompileExecutable( const ComputationHandle& computation, const tensorflow::gtl::ArraySlice argument_layouts, - const Shape* result_layout, int device_ordinal) { + const ExecutableBuildOptions& build_options) { TF_ASSIGN_OR_RETURN(UserComputation * user_computation, computation_tracker_.Resolve(computation)); VersionedComputationHandle versioned_handle = @@ -84,27 +85,47 @@ StatusOr> LocalService::CompileExecutable( // Validate incoming layouts. if (argument_layouts.size() != program_shape->parameters_size()) { return InvalidArgument( - "invalid number of arguments for computation: expected %d, got %zu", + "Invalid number of arguments for computation: expected %d, got %zu.", program_shape->parameters_size(), argument_layouts.size()); } for (int i = 0; i < argument_layouts.size(); ++i) { const Shape& argument_shape = *argument_layouts[i]; TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(argument_shape)); if (!ShapeUtil::Compatible(argument_shape, program_shape->parameters(i))) { + tensorflow::gtl::optional metadata = + user_computation->ParameterMetadata(i); + auto metadata_string = [&metadata]() -> string { + if (!metadata.has_value()) { + return ""; + } + CHECK(metadata.value() != nullptr); + const OpMetadata& m = *metadata.value(); + if (!m.source_file().empty()) { + return tensorflow::strings::Printf( + " (%s:%d)", m.source_file().c_str(), m.source_line()); + } + return ""; + }; return InvalidArgument( - "invalid argument shape for argument %d, expected %s, got %s", i, + "Invalid argument shape for argument %d%s, expected %s, got %s.", i, + metadata_string().c_str(), ShapeUtil::HumanString(program_shape->parameters(i)).c_str(), ShapeUtil::HumanString(argument_shape).c_str()); } } - if (result_layout != nullptr) { - TF_RETURN_IF_ERROR( - ValidateResultShapeWithLayout(*result_layout, program_shape->result())); + if (build_options.result_layout() != nullptr) { + TF_RETURN_IF_ERROR(ValidateResultShapeWithLayout( + *build_options.result_layout(), program_shape->result())); } ExecutionOptions execution_options = CreateDefaultExecutionOptions(); - if (result_layout != nullptr) { - *execution_options.mutable_shape_with_output_layout() = *result_layout; + if (build_options.generate_hlo_graph().has_value()) { + execution_options.mutable_debug_options()->set_xla_generate_hlo_graph( + build_options.generate_hlo_graph().value()); + } + if (build_options.result_layout() != nullptr) { + *execution_options.mutable_shape_with_output_layout() = + *build_options.result_layout(); } else { *execution_options.mutable_shape_with_output_layout() = program_shape->result(); @@ -113,13 +134,22 @@ StatusOr> LocalService::CompileExecutable( } TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(*program_shape, argument_layouts, &execution_options)); + CreateModuleConfig(*program_shape, argument_layouts, &execution_options, + *user_computation)); - TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor, - execute_backend_->stream_executor(device_ordinal)); + TF_ASSIGN_OR_RETURN( + se::StreamExecutor * executor, + execute_backend_->stream_executor(build_options.device_ordinal())); return BuildExecutable(versioned_handle, std::move(module_config), - execute_backend_.get(), executor); + execute_backend_.get(), executor, + build_options.device_allocator()); +} + +StatusOr LocalService::ReplicaNumberToDeviceOrdinal(int replica_number) { + return backend().computation_placer()->DeviceId( + replica_number, /*computation=*/0, options_.number_of_replicas(), + /*computation_count=*/1); } } // namespace xla diff --git a/tensorflow/compiler/xla/service/local_service.h b/tensorflow/compiler/xla/service/local_service.h index 52c4346385eb663baa6e7579d7b3883ba084205b..15e120685e1be9190d49fdaf5ed6706bdf991a6c 100644 --- a/tensorflow/compiler/xla/service/local_service.h +++ b/tensorflow/compiler/xla/service/local_service.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/client/executable_build_options.h" #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" @@ -41,11 +42,20 @@ class LocalService : public Service { // Builds an Executable with the given argument layouts and options. If // result_layout is non-null, then the executable is compiled to produce a - // result of the given layout. + // result of the given layout. If device_allocator is non-null, then the + // compiler may use it to allocate temp space on the device. The compiler is + // responsible for freeing any memory it allocates this way. StatusOr> CompileExecutable( const ComputationHandle& computation, const tensorflow::gtl::ArraySlice argument_layouts, - const Shape* result_layout, int device_ordinal); + const ExecutableBuildOptions& options); + + // Returns the device ordinal that corresponds to the given replica number. + // + // This returns an error if there is not a one-to-one correspondence of + // replicas to device ordinals, but is useful as a short term mechanism for + // the "easy" case where a single replica is a single device. + StatusOr ReplicaNumberToDeviceOrdinal(int replica_number); private: explicit LocalService(const ServiceOptions& options, diff --git a/tensorflow/compiler/xla/service/logical_buffer_analysis.h b/tensorflow/compiler/xla/service/logical_buffer_analysis.h index 598d08b7203b25b194dfc3b3125ec58c96b2cd4c..f4c63dd86b4d8a6f598d46047012e4e5bc7b3d7e 100644 --- a/tensorflow/compiler/xla/service/logical_buffer_analysis.h +++ b/tensorflow/compiler/xla/service/logical_buffer_analysis.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" @@ -90,4 +90,4 @@ class LogicalBufferAnalysis : public DfsHloVisitorWithDefault { } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc index 0b987141682558b88e846d2d4fdeb3c45fbc447d..98dfc89867ab33788c4cc837a66d6751a1ef2507 100644 --- a/tensorflow/compiler/xla/service/service.cc +++ b/tensorflow/compiler/xla/service/service.cc @@ -34,8 +34,10 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/service/hlo_proto_util.h" #include "tensorflow/compiler/xla/service/platform_util.h" #include "tensorflow/compiler/xla/service/session.pb.h" +#include "tensorflow/compiler/xla/service/source_map_util.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" #include "tensorflow/compiler/xla/shape_layout.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -55,6 +57,7 @@ namespace se = ::perftools::gputools; using ::tensorflow::strings::Printf; using ::tensorflow::strings::StrCat; +using ::xla::source_map_util::InvalidParameterArgument; namespace xla { @@ -260,7 +263,8 @@ StatusOr> Service::ResolveAndValidateArguments( StatusOr> Service::CreateModuleConfig( const ProgramShape& program_shape, tensorflow::gtl::ArraySlice argument_shapes, - const ExecutionOptions* execution_options) { + const ExecutionOptions* execution_options, + const UserComputation& user_computation) { auto config = MakeUnique(program_shape); auto* computation_layout = config->mutable_entry_computation_layout(); @@ -274,8 +278,10 @@ StatusOr> Service::CreateModuleConfig( // ProgramShape. if (!ShapeUtil::Compatible(*argument_shapes[i], program_shape.parameters(i))) { - return InvalidArgument( - "computation expects parameter %d to have shape %s, given shape %s", + return InvalidParameterArgument( + *user_computation.ParameterMetadata(i).value(), + "Argument does not match shape of computation parameter %d: want %s, " + "got %s", i, ShapeUtil::HumanString(program_shape.parameters(i)).c_str(), ShapeUtil::HumanString(*argument_shapes[i]).c_str()); } @@ -317,19 +323,22 @@ StatusOr> Service::CreateModuleConfig( StatusOr> Service::CreateModuleConfig( const ProgramShape& program_shape, tensorflow::gtl::ArraySlice arguments, - const ExecutionOptions& execution_options) { + const ExecutionOptions& execution_options, + const UserComputation& user_computation) { std::vector argument_shapes; for (const auto* arg : arguments) { argument_shapes.push_back(&arg->on_host_shape()); } - return CreateModuleConfig(program_shape, argument_shapes, &execution_options); + return CreateModuleConfig(program_shape, argument_shapes, &execution_options, + user_computation); } StatusOr>> Service::BuildExecutables( std::vector versioned_handles, std::vector> module_configs, Backend* backend, - std::vector> executors) { + std::vector> executors, + DeviceMemoryAllocator* device_allocator) { VLOG(1) << Printf("BuildExecutable on service %p", this); // Dump computation proto state if flag is set. @@ -375,7 +384,8 @@ StatusOr>> Service::BuildExecutables( TF_ASSIGN_OR_RETURN( std::vector> executables, - backend->compiler()->Compile(std::move(modules), std::move(executors))); + backend->compiler()->Compile(std::move(modules), std::move(executors), + device_allocator)); for (size_t i = 0; i < versioned_handles.size(); ++i) { if (!module_configs[i]->debug_options().xla_dump_executions_to().empty()) { @@ -388,8 +398,8 @@ StatusOr>> Service::BuildExecutables( StatusOr> Service::BuildExecutable( const VersionedComputationHandle& versioned_handle, - std::unique_ptr module_config, - Backend* backend, se::StreamExecutor* executor) { + std::unique_ptr module_config, Backend* backend, + se::StreamExecutor* executor, DeviceMemoryAllocator* device_allocator) { VLOG(1) << Printf("BuildExecutable on service %p with handle %s", this, versioned_handle.ToString().c_str()); @@ -419,12 +429,15 @@ StatusOr> Service::BuildExecutable( /*include_unreachable_instructions=*/ true)); - TF_ASSIGN_OR_RETURN( - module, backend->compiler()->RunHloPasses(std::move(module), executor)); + TF_RETURN_IF_ERROR(MaybeDumpHloModule(*module)); TF_ASSIGN_OR_RETURN( - std::unique_ptr executable, - backend->compiler()->RunBackend(std::move(module), executor)); + module, backend->compiler()->RunHloPasses(std::move(module), executor, + device_allocator)); + + TF_ASSIGN_OR_RETURN(std::unique_ptr executable, + backend->compiler()->RunBackend( + std::move(module), executor, device_allocator)); if (!other_directory_path.empty()) { executable->set_session_module(std::move(session_module)); @@ -435,9 +448,9 @@ StatusOr> Service::BuildExecutable( StatusOr> Service::BuildAndCacheExecutable( const VersionedComputationHandle& versioned_handle, - std::unique_ptr module_config, - Backend* backend, perftools::gputools::StreamExecutor* executor, - ExecutionProfile* profile) { + std::unique_ptr module_config, Backend* backend, + perftools::gputools::StreamExecutor* executor, ExecutionProfile* profile, + DeviceMemoryAllocator* device_allocator) { std::shared_ptr executable = compilation_cache_.LookUp(versioned_handle, *module_config); @@ -459,7 +472,7 @@ StatusOr> Service::BuildAndCacheExecutable( TF_ASSIGN_OR_RETURN( std::unique_ptr executable_unique_ptr, BuildExecutable(versioned_handle, std::move(module_config), backend, - executor)); + executor, device_allocator)); if (profile != nullptr) { uint64 end_micros = tensorflow::Env::Default()->NowMicros(); @@ -566,7 +579,7 @@ Service::ExecuteParallelAndRegisterResult( se::Stream* stream = index_to_profiled_stream.second; Executable* executable = executables[device]; const HloModule& module = executable->module(); - HloExecutionProfile hlo_profile(&executable->hlo_profile_printer(), + HloExecutionProfile hlo_profile(&executable->hlo_profile_printer_data(), &executable->hlo_profile_index_map()); TF_RETURN_IF_ERROR( executable->PopulateExecutionProfile(&hlo_profile, stream->parent())); @@ -739,9 +752,10 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, // Create an HloModuleConfig object for the computation, given the shape of // the program and the argument allocations. - TF_ASSIGN_OR_RETURN(std::unique_ptr module_config, - CreateModuleConfig(*program_shape, arguments, - request.execution_options())); + TF_ASSIGN_OR_RETURN( + std::unique_ptr module_config, + CreateModuleConfig(*program_shape, arguments, + request.execution_options(), *user_computation)); VLOG(3) << "ExecuteParallel created HloModuleConfig computation layout: " << module_config->entry_computation_layout().ToString(); @@ -760,10 +774,14 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, // Build the user computations into HloModules and compile to generate the // executables. + // + // TODO(jlebar): There's currently no way to pass a device allocator to + // ExecuteParallel, so we have to pass a null device_allocator below. TF_ASSIGN_OR_RETURN( std::vector> executables, BuildExecutables(versioned_handles, std::move(module_configs), - execute_backend_.get(), all_executors)); + execute_backend_.get(), all_executors, + /*device_allocator=*/nullptr)); std::vector executable_ptrs; executable_ptrs.reserve(executables.size()); for (const auto& executable : executables) { @@ -849,7 +867,8 @@ tensorflow::Status Service::Execute(const ExecuteRequest* arg, TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(*program_shape, arguments, arg->execution_options())); + CreateModuleConfig(*program_shape, arguments, arg->execution_options(), + *user_computation)); VLOG(3) << "Execute created HloModuleConfig computation layout: " << module_config->entry_computation_layout().ToString(); @@ -913,7 +932,8 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(*program_shape, arguments, arg->execution_options())); + CreateModuleConfig(*program_shape, arguments, arg->execution_options(), + *user_computation)); VLOG(3) << "ExecuteAsync created HloModuleConfig computation layout: " << module_config->entry_computation_layout().ToString(); @@ -1040,8 +1060,9 @@ std::unique_ptr CloneShapedBufferOnDevice( tensorflow::Status Service::TransferToServer(const TransferToServerRequest* arg, TransferToServerResponse* result) { - Literal literal = Literal(arg->literal()); - const Shape& shape = literal.shape(); + TF_ASSIGN_OR_RETURN(std::unique_ptr literal, + Literal::CreateFromProto(arg->literal())); + const Shape& shape = literal->shape(); std::vector replicas; if (arg->has_device_handle()) { @@ -1065,7 +1086,7 @@ tensorflow::Status Service::TransferToServer(const TransferToServerRequest* arg, if (executor->device_ordinal() == master_device_ordinal) { TF_RETURN_IF_ERROR( execute_backend_->transfer_manager()->TransferLiteralToDevice( - executor, literal, *shaped_buffer)); + executor, *literal, *shaped_buffer)); } else { // The replica is not the master. Create an cloned shaped buffer with // the replica's device ordinal. This is required because @@ -1075,7 +1096,7 @@ tensorflow::Status Service::TransferToServer(const TransferToServerRequest* arg, CloneShapedBufferOnDevice(*shaped_buffer, executor->device_ordinal()); TF_RETURN_IF_ERROR( execute_backend_->transfer_manager()->TransferLiteralToDevice( - executor, literal, *clone)); + executor, *literal, *clone)); } } TF_ASSIGN_OR_RETURN( @@ -1111,8 +1132,10 @@ tensorflow::Status Service::TransferToInfeed(const TransferToInfeedRequest* arg, executor = replicas[arg->replica_id()]; } + TF_ASSIGN_OR_RETURN(std::unique_ptr literal, + Literal::CreateFromProto(arg->literal())); return execute_backend_->transfer_manager()->TransferLiteralToInfeed( - executor, Literal(arg->literal())); + executor, *literal); } tensorflow::Status Service::TransferFromOutfeed( @@ -1187,7 +1210,22 @@ tensorflow::Status Service::ComputeConstant(const ComputeConstantRequest* arg, bool is_constant, user_computation->IsConstant(arg->operand(), arg->parameters_size())); if (!is_constant) { - return InvalidArgument("Operand to ComputeConstant depends on parameter."); + StatusOr op_request_status = + user_computation->LookUpRequestForErrorReporting(arg->operand()); + string op_request_string = ""; + if (op_request_status.ok()) { + op_request_string = op_request_status.ValueOrDie()->ShortDebugString(); + } + return InvalidArgument( + "Operand to ComputeConstant depends on a parameter.\n\n" + " op requested for constant evaluation: %s\n\n" + "This is an internal error that typically happens when the XLA user " + "(e.g. TensorFlow) is attempting to determine a value that must be a " + "compile-time constant (e.g. an array dimension) but it is not capable " + "of being evaluated at XLA compile time.\n\n" + "Please file a usability bug with the framework being used (e.g. " + "TensorFlow).", + op_request_string.c_str()); } // We can't use ComputeProgramShape because it checks that all parameter @@ -1215,7 +1253,8 @@ tensorflow::Status Service::ComputeConstant(const ComputeConstantRequest* arg, } TF_ASSIGN_OR_RETURN(std::unique_ptr module_config, - CreateModuleConfig(program_shape, {}, execution_options)); + CreateModuleConfig(program_shape, {}, execution_options, + *user_computation)); // Exclude dead parameter instructions for the purpose of computing constants. TF_ASSIGN_OR_RETURN( @@ -1224,18 +1263,15 @@ tensorflow::Status Service::ComputeConstant(const ComputeConstantRequest* arg, /*include_unreachable_instructions=*/ false)); - std::vector parameters(arg->parameters_size()); + std::vector> parameters(arg->parameters_size()); for (int64 i = 0; i < arg->parameters_size(); ++i) { - parameters[i] = Literal(arg->parameters(i)); + TF_ASSIGN_OR_RETURN(parameters[i], + Literal::CreateFromProto(arg->parameters(i))); } - std::vector parameter_ptrs; - std::transform(parameters.begin(), parameters.end(), - std::back_inserter(parameter_ptrs), - [](const Literal& literal) { return &literal; }); - HloEvaluator evaluator; - TF_ASSIGN_OR_RETURN(auto result_literal, evaluator.Evaluate( - *module, parameter_ptrs)); + TF_ASSIGN_OR_RETURN( + auto result_literal, + evaluator.Evaluate>(*module, parameters)); // Since the shape_with_output_layout option in ExecutionOption is // non-effective to the Evaluator results, explicit relayout here. @@ -1406,6 +1442,9 @@ tensorflow::Status Service::Op(const OpRequest* arg, OpResponse* result) { handle_status = computation->AddDynamicUpdateSliceInstruction( arg->dynamic_update_slice_request()); break; + case OpRequest::kFftRequest: + handle_status = computation->AddFftInstruction(arg->fft_request()); + break; case OpRequest::kGetTupleElementRequest: handle_status = computation->AddGetTupleElementInstruction( arg->get_tuple_element_request()); @@ -1414,9 +1453,9 @@ tensorflow::Status Service::Op(const OpRequest* arg, OpResponse* result) { handle_status = computation->AddInfeedInstruction(arg->infeed_request()); break; case OpRequest::kOutfeedRequest: - TF_RETURN_IF_ERROR( - computation->AddOutfeedInstruction(arg->outfeed_request())); - return tensorflow::Status::OK(); + handle_status = + computation->AddOutfeedInstruction(arg->outfeed_request()); + break; case OpRequest::kMapRequest: { TF_ASSIGN_OR_RETURN( UserComputation * to_apply, @@ -1518,8 +1557,6 @@ tensorflow::Status Service::Op(const OpRequest* arg, OpResponse* result) { handle_status = computation->AddRecvInstruction(arg->recv_request()); break; } - case OpRequest::kFftRequest: - return Unimplemented("FftRequest not implemented in XLA service."); case OpRequest::OP_NOT_SET: return InvalidArgument("XLA service received OpRequest with OP_NOT_SET"); default: @@ -1581,4 +1618,15 @@ StatusOr> Service::Replicas( return replicas; } +Status Service::MaybeDumpHloModule(const HloModule& module) const { + const string xla_dump_unoptimized_hlo_proto_to = + module.config().debug_options().xla_dump_unoptimized_hlo_proto_to(); + if (xla_dump_unoptimized_hlo_proto_to.empty()) { + return Status::OK(); + } + HloProto proto = MakeHloProto(module); + return protobuf_util::DumpProtoToDirectory( + proto, xla_dump_unoptimized_hlo_proto_to, module.name()); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/service.h b/tensorflow/compiler/xla/service/service.h index f962d0cdc7d41e1aeab55da5abcb1b40215b4144..6ce241971156599aaa25aea1b0caac0e1bd5379c 100644 --- a/tensorflow/compiler/xla/service/service.h +++ b/tensorflow/compiler/xla/service/service.h @@ -251,7 +251,8 @@ class Service : public ServiceInterface { StatusOr> CreateModuleConfig( const ProgramShape& program_shape, tensorflow::gtl::ArraySlice arguments, - const ExecutionOptions& execution_options); + const ExecutionOptions& execution_options, + const UserComputation& user_computation); protected: friend class LocalExecutable; @@ -275,13 +276,19 @@ class Service : public ServiceInterface { StatusOr> CreateModuleConfig( const ProgramShape& program_shape, tensorflow::gtl::ArraySlice argument_shapes, - const ExecutionOptions* execution_options); + const ExecutionOptions* execution_options, + const UserComputation& user_computation); // Builds an Executable for the given parameters. + // + // If device_allocator is not null, the compiler may use it to allocate temp + // buffers, which the compiler is responsible for freeing. The allocator + // given here need not match the allocator used when running the executable. StatusOr> BuildExecutable( const VersionedComputationHandle& versioned_handle, - std::unique_ptr module_config, - Backend* backend, perftools::gputools::StreamExecutor* executor); + std::unique_ptr module_config, Backend* backend, + perftools::gputools::StreamExecutor* executor, + DeviceMemoryAllocator* device_allocator = nullptr); // Same as BuildExecutable() above, but builds a list of Executables for the // given computations that may interact with each other. @@ -289,16 +296,17 @@ class Service : public ServiceInterface { std::vector versioned_handles, std::vector> module_configs, Backend* backend, - std::vector> executors); + std::vector> executors, + DeviceMemoryAllocator* device_allocator); // Similar to BuildExecutable, but look in the compilation cache for the // executable first. If the executable is not in the cache, it is built and // inserted into the cache. StatusOr> BuildAndCacheExecutable( const VersionedComputationHandle& versioned_handle, - std::unique_ptr module_config, - Backend* backend, perftools::gputools::StreamExecutor* executor, - ExecutionProfile* profile); + std::unique_ptr module_config, Backend* backend, + perftools::gputools::StreamExecutor* executor, ExecutionProfile* profile, + DeviceMemoryAllocator* device_allocator = nullptr); // Runs the given executable with the given arguments and register the result // in the allocation tracker. The handle of the result from the tracker is @@ -340,6 +348,8 @@ class Service : public ServiceInterface { StatusOr> Replicas( const Backend& backend, const DeviceHandle& device_handle) const; + Status MaybeDumpHloModule(const HloModule& module) const; + // Returns the device handle that represents the replicated device for a // single computation that is not model-parallelized. DeviceHandle SingleComputationDeviceHandle() const; diff --git a/tensorflow/compiler/xla/service/service_executable_run_options.h b/tensorflow/compiler/xla/service/service_executable_run_options.h index 017e5ef09ed2f52b862821e9408540d188a1edf5..6c1f8feac7ed4423051cf2737be57dcfab508671 100644 --- a/tensorflow/compiler/xla/service/service_executable_run_options.h +++ b/tensorflow/compiler/xla/service/service_executable_run_options.h @@ -30,6 +30,9 @@ class ServiceExecutableRunOptions { using StreamBorrower = std::function::SmartPtr>(int)>; + ServiceExecutableRunOptions() + : ServiceExecutableRunOptions(ExecutableRunOptions()) {} + explicit ServiceExecutableRunOptions( ExecutableRunOptions run_options, StreamBorrower borrow_stream = nullptr, tensorflow::thread::ThreadPool* xla_intra_op_thread_pool = nullptr) diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc index 9c1b951d017569a6dc89bc6583c72b5e42f0c07c..004889b5f216015ee1e1308702b2bf4cb0deb344 100644 --- a/tensorflow/compiler/xla/service/shape_inference.cc +++ b/tensorflow/compiler/xla/service/shape_inference.cc @@ -37,6 +37,9 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/protobuf.h" +using tensorflow::str_util::Join; +using tensorflow::strings::Printf; + namespace xla { namespace { @@ -206,7 +209,8 @@ tensorflow::Status VerifyReducerShape(const ProgramShape& reducer_shape, } // Check that init_value's shape is suitable for reducer_shape. - if (!ShapeUtil::Compatible(accumulator_shape, init_value_shape)) { + if (!ShapeUtil::CompatibleIgnoringFpPrecision(accumulator_shape, + init_value_shape)) { return InvalidArgument( "Reduction function's accumulator shape differs from the " "init_value shape: %s vs %s", @@ -217,8 +221,8 @@ tensorflow::Status VerifyReducerShape(const ProgramShape& reducer_shape, // Check that the inputs can be passed in as the second argument. const Shape& input_element_shape = ShapeUtil::MakeShape(input_element_type, {}); - if (!ShapeUtil::Compatible(input_element_shape, - reducer_shape.parameters(1))) { + if (!ShapeUtil::CompatibleIgnoringFpPrecision(input_element_shape, + reducer_shape.parameters(1))) { return InvalidArgument( "Reduction function's second parameter shape differs from the " "input type element type: %s vs %s", @@ -228,7 +232,8 @@ tensorflow::Status VerifyReducerShape(const ProgramShape& reducer_shape, // Currently the accumulator and inputs must be the same type, // though that restriction could be relaxed. - if (!ShapeUtil::Compatible(accumulator_shape, reducer_shape.parameters(1))) { + if (!ShapeUtil::CompatibleIgnoringFpPrecision(accumulator_shape, + reducer_shape.parameters(1))) { return InvalidArgument( "Reduction function's second parameter shape currently must " "match the result shape. Got %s vs %s", @@ -391,11 +396,13 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, dimension); } const Shape* arg_shape = nullptr; + PrimitiveType element_type = PRIMITIVE_TYPE_INVALID; for (const Shape* shape : arg_shapes) { TF_RETURN_IF_ERROR( ExpectNotTupleOrOpaque(*shape, "operand of concatenation")); if (!arg_shape) { arg_shape = shape; + element_type = arg_shape->element_type(); continue; } if (ShapeUtil::Rank(*arg_shape) != ShapeUtil::Rank(*shape)) { @@ -406,7 +413,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, ShapeUtil::HumanString(*arg_shape).c_str(), ShapeUtil::Rank(*shape), ShapeUtil::HumanString(*shape).c_str()); } - if (arg_shape->element_type() != shape->element_type()) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(*arg_shape, *shape)) { return InvalidArgument( "cannot concatenate arrays with different element types: %s vs %s", PrimitiveType_Name(arg_shape->element_type()).c_str(), @@ -428,6 +435,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, ShapeUtil::HumanString(*shape).c_str(), dimension); } } + element_type = ShapeUtil::HigherPrecisionElementType(*shape, *arg_shape); } std::vector new_dimensions(arg_shape->dimensions().begin(), @@ -435,7 +443,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, for (size_t i = 1; i < arg_shapes.size(); ++i) { new_dimensions[dimension] += arg_shapes[i]->dimensions(dimension); } - return ShapeUtil::MakeShape(arg_shape->element_type(), new_dimensions); + return ShapeUtil::MakeShape(element_type, new_dimensions); } /* static */ StatusOr ShapeInference::InferConvertShape( @@ -533,7 +541,8 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, ShapeUtil::HumanString(operand_shape).c_str(), padding_config.ShortDebugString().c_str()); } - if (operand_shape.element_type() != padding_value_shape.element_type()) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(operand_shape, + padding_value_shape)) { return InvalidArgument( "the element types of the operands to pad do not match"); } @@ -545,7 +554,9 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, std::max(operand_shape.dimensions(i) - 1, 0LL) * padding_config.dimensions(i).interior_padding(); } - return ShapeUtil::MakeShape(operand_shape.element_type(), dimensions); + return ShapeUtil::MakeShape( + ShapeUtil::HigherPrecisionElementType(operand_shape, padding_value_shape), + dimensions); } // Current DotDimensionNumbers Requirements: @@ -670,7 +681,7 @@ Status ValidateDotDimensionNumbers( }; // Check if both element types are the same. - if (lhs.element_type() != rhs.element_type()) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(lhs, rhs)) { return fail("element types do not match"); } @@ -733,7 +744,8 @@ Status ValidateDotDimensionNumbers( dimensions.push_back(rhs.dimensions(i)); } } - Shape result = ShapeUtil::MakeShape(lhs.element_type(), dimensions); + Shape result = ShapeUtil::MakeShape( + ShapeUtil::HigherPrecisionElementType(lhs, rhs), dimensions); TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(result)); VLOG(2) << "inferred dot shape: " << ShapeUtil::HumanString(result); @@ -764,7 +776,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( ShapeUtil::HumanString(rhs).c_str()); } } - return ShapeUtil::MakeShape(lhs.element_type(), output_dimensions); + return ShapeUtil::MakeShape(ShapeUtil::HigherPrecisionElementType(lhs, rhs), + output_dimensions); } /* static */ StatusOr ShapeInference::InferInDimBroadcastShape( @@ -826,6 +839,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( // specified in broadcast_dimensions are then changed to match the // corresponding dimension size in smaller_shape. Shape output_shape(larger_shape); + output_shape.set_element_type( + ShapeUtil::HigherPrecisionElementType(larger_shape, smaller_shape)); for (int i = 0; i < smaller_shape.dimensions_size(); ++i) { int64 dimension_to_match = broadcast_dimensions.at(i); @@ -875,7 +890,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( TF_RETURN_IF_ERROR( ExpectNotTupleOrOpaque(rhs, "rhs of elementwise binary operation")); - if (!ShapeUtil::SameElementType(lhs, rhs)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(lhs, rhs)) { return InvalidArgument( "binary op %s with different element types: %s and %s", BinaryOperation_Name(operation).c_str(), @@ -894,10 +909,11 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } } - if (ShapeUtil::Compatible(lhs, rhs)) { + if (ShapeUtil::CompatibleIgnoringFpPrecision(lhs, rhs)) { // If the shapes are the same other than layout, the output shape is the // same (elementwise op). - return lhs; + return ShapeUtil::ChangeElementType( + lhs, ShapeUtil::HigherPrecisionElementType(lhs, rhs)); } if (ShapeUtil::Rank(lhs) == ShapeUtil::Rank(rhs)) { @@ -934,7 +950,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( "inferring shape for <%s>(%s, %s) with broadcast_dimensions={%s}", BinaryOperation_Name(operation).c_str(), ShapeUtil::HumanString(lhs).c_str(), ShapeUtil::HumanString(rhs).c_str(), - tensorflow::str_util::Join(broadcast_dimensions, ", ").c_str()); + Join(broadcast_dimensions, ", ").c_str()); TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(lhs)); TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(rhs)); @@ -970,7 +986,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( TF_ASSIGN_OR_RETURN(const Shape& shape, InferElementwiseBinaryOpShape(operation, lhs, rhs, broadcast_dimensions)); - if (lhs.element_type() == F32) { + if (lhs.element_type() == F32 && rhs.element_type() == F32) { return ShapeUtil::ChangeElementType(shape, C64); } else { return Unimplemented("complex component type not supported"); @@ -1075,12 +1091,13 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( TF_RETURN_IF_ERROR( ExpectNotTupleOrOpaque(*arg_shapes[i], "operand of map")); - if (ShapeUtil::Compatible(*arg_shapes[i], *arg_shape)) { + if (ShapeUtil::CompatibleIgnoringFpPrecision(*arg_shapes[i], *arg_shape)) { continue; } if (!ShapeUtil::IsTuple(*arg_shapes[i]) && !ShapeUtil::IsTuple(*arg_shape) && - ShapeUtil::SameElementType(*arg_shapes[i], *arg_shape)) { + ShapeUtil::SameElementTypeIgnoringFpPrecision(*arg_shapes[i], + *arg_shape)) { if (ShapeUtil::IsScalar(*arg_shapes[i])) { continue; } @@ -1097,7 +1114,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "Map operation requires all operands to have the same shape; got: " "%s", - tensorflow::str_util::Join(pieces, ", ").c_str()); + Join(pieces, ", ").c_str()); } // Check that dimensions.size == arg_shape.dimensions_size() (we currently @@ -1114,7 +1131,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( if (dimensions[i] != i) { return InvalidArgument( "Map requires monotonically increasing dimension numbers, found: %s ", - tensorflow::str_util::Join(dimensions, ", ").c_str()); + Join(dimensions, ", ").c_str()); } } @@ -1145,7 +1162,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( i, ShapeUtil::HumanString(parameter_shape).c_str()); } - if (parameter_shape.element_type() != arg_shape->element_type()) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(parameter_shape, + *arg_shape)) { return InvalidArgument( "mapped computation's parameter type has to match argument element " "type; got parameter %d shape: %s, argument shape: %s", @@ -1218,7 +1236,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(offset_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(offset_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for batch-norm-training, " "but the shape of offset factor is %s " @@ -1227,7 +1246,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(scale_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(scale_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for batch-norm-training, " "but the shape of scale factor is %s " @@ -1326,7 +1346,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(offset_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(offset_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for " "batch-norm-inference, " @@ -1336,7 +1357,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(scale_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(scale_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for " "batch-norm-inference, " @@ -1346,7 +1368,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(mean_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(mean_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for " "batch-norm-inference, " @@ -1356,7 +1379,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(variance_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(variance_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for " "batch-norm-inference, " @@ -1478,7 +1502,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(output_grad_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(output_grad_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(output_grad_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for batch-norm-grad, " "but the element type of output_grad is %s " @@ -1487,7 +1512,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(scale_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(scale_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for batch-norm-grad, " "but the element type of scale factor is %s " @@ -1496,7 +1522,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(mean_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(mean_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for batch-norm-grad, " "but the element type of mean is %s " @@ -1505,7 +1532,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( PrimitiveType_Name(operand_shape.element_type()).c_str()); } - if (!ShapeUtil::SameElementType(var_shape, operand_shape)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(var_shape, + operand_shape)) { return InvalidArgument( "The inputs should have the same element type for batch-norm-grad, " "but the element type of mean is %s " @@ -1566,7 +1594,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(lhs, "lhs of convolution")); TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(rhs, "rhs of convolution")); - if (!ShapeUtil::SameElementType(lhs, rhs)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(lhs, rhs)) { return InvalidArgument( "Convolution with different element types: %s and %s", ShapeUtil::HumanString(lhs).c_str(), @@ -1711,8 +1739,91 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( dimensions[dnums.output_spatial_dimensions(i)] = window_output_shape.dimensions(i); } + return ShapeUtil::MakeShape(ShapeUtil::HigherPrecisionElementType(lhs, rhs), + dimensions); +} - return ShapeUtil::MakeShape(lhs.element_type(), dimensions); +/* static */ StatusOr ShapeInference::InferFftShape( + const Shape& in, const FftType fft_type, + const tensorflow::gtl::ArraySlice fft_length) { + const int64 fft_rank = fft_length.size(); + if (fft_rank < 1 || fft_rank > 3) { + return InvalidArgument("FFT only supports ranks 1-3, but got %lld", + fft_rank); + } +#define RET_CHECK_RANK(x) \ + if (x.dimensions_size() < fft_rank) { \ + return InvalidArgument( \ + "FFT of rank %lld requires input of at least " \ + "same rank; got input of rank %d", \ + fft_rank, x.dimensions_size()); \ + } + switch (fft_type) { + case FFT: + case IFFT: + if (in.element_type() != C64) { + return InvalidArgument("%s requires C64 input type, found %s", + FftType_Name(fft_type).c_str(), + PrimitiveType_Name(in.element_type()).c_str()); + } + RET_CHECK_RANK(in); + return in; + case RFFT: { + if (in.element_type() != F32) { + return InvalidArgument("RFFT requires F32 input type, found %s", + PrimitiveType_Name(in.element_type()).c_str()); + } + RET_CHECK_RANK(in); + for (int i = 0; i < fft_rank; i++) { + if (in.dimensions(in.dimensions_size() - fft_rank + i) != + fft_length[i]) { + return InvalidArgument( + "RFFT requires innermost dimensions match fft_length but " + "dimension %lld is %lld and should be %lld", + in.dimensions_size() - fft_rank + i, + in.dimensions(in.dimensions_size() - fft_rank + i), + fft_length[i]); + } + } + Shape result = ShapeUtil::ChangeElementType(in, C64); + result.set_dimensions(result.dimensions_size() - 1, + fft_length[fft_rank - 1] / 2 + 1); + return result; + } + case IRFFT: { + if (in.element_type() != C64) { + return InvalidArgument("IRFFT requires C64 input type, found %s", + PrimitiveType_Name(in.element_type()).c_str()); + } + RET_CHECK_RANK(in); + Shape result = ShapeUtil::ComplexComponentShape(in); + for (int i = 0; i < fft_rank - 1; i++) { + if (in.dimensions(in.dimensions_size() - fft_rank + i) != + fft_length[i]) { + return InvalidArgument( + "IRFFT requires all but one innermost dimensions match " + "fft_length, but dimension %lld is %lld and should be %lld", + in.dimensions_size() - fft_rank + i, + in.dimensions(in.dimensions_size() - fft_rank + i), + fft_length[i]); + } + } + if (in.dimensions(in.dimensions_size() - 1) != + fft_length[fft_rank - 1] / 2 + 1) { + return InvalidArgument( + "IRFFT requires innermost dimension matches fft_length/2+1, but " + "dimension %d is %lld and should be %lld", + in.dimensions_size() - 1, in.dimensions(in.dimensions_size() - 1), + fft_length[fft_rank - 1] / 2 + 1); + } + result.set_dimensions(result.dimensions_size() - 1, + fft_length[fft_rank - 1]); + return result; + } + default: + LOG(FATAL) << "Unexpected fft_type: " << fft_type; + } +#undef RET_CHECK_RANK } /* static */ StatusOr ShapeInference::InferCrossReplicaSumShape( @@ -1791,16 +1902,16 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } const Shape& operand_element_shape = ShapeUtil::MakeShape(operand_shape.element_type(), {}); - if (!ShapeUtil::Compatible(operand_element_shape, - select_shape.parameters(0))) { + if (!ShapeUtil::CompatibleIgnoringFpPrecision(operand_element_shape, + select_shape.parameters(0))) { return InvalidArgument( "select function's first parameter shape currently must " "match the operand element shape. Got %s vs %s", ShapeUtil::HumanString(select_shape.parameters(0)).c_str(), ShapeUtil::HumanString(operand_element_shape).c_str()); } - if (!ShapeUtil::Compatible(operand_element_shape, - select_shape.parameters(1))) { + if (!ShapeUtil::CompatibleIgnoringFpPrecision(operand_element_shape, + select_shape.parameters(1))) { return InvalidArgument( "select function's second parameter shape currently must " "match the operand element shape. Got %s vs %s", @@ -1817,7 +1928,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( InferWindowOutputShape(operand_shape, window, operand_shape.element_type(), /*allow_negative_padding=*/false)); - if (!ShapeUtil::Compatible(source_shape, window_result_shape)) { + if (!ShapeUtil::CompatibleIgnoringFpPrecision(source_shape, + window_result_shape)) { return InvalidArgument( "source shape does not match the shape of window-reduced operand: " "source(%s), window-reduced operand(%s)", @@ -1831,21 +1943,28 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const Shape& arg, tensorflow::gtl::ArraySlice starts, tensorflow::gtl::ArraySlice limits, tensorflow::gtl::ArraySlice strides) { + auto error = [&](const string& message) { + return InvalidArgument( + "%s in slice operation; argument shape: %s; starts: {%s}; limits: " + "{%s}; strides: {%s}", + message.c_str(), ShapeUtil::HumanString(arg).c_str(), + Join(starts, ",").c_str(), Join(limits, ",").c_str(), + Join(strides, ",").c_str()); + }; TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(arg, "operand of slice")); VLOG(2) << tensorflow::strings::Printf( "slicing shape %s starts={%s} limits={%s}", - ShapeUtil::HumanString(arg).c_str(), - tensorflow::str_util::Join(starts, ", ").c_str(), - tensorflow::str_util::Join(limits, ", ").c_str()); + ShapeUtil::HumanString(arg).c_str(), Join(starts, ", ").c_str(), + Join(limits, ", ").c_str()); if (starts.size() != limits.size()) { - return InvalidArgument("slice start and limit sizes differ: %zu vs %zu", - starts.size(), limits.size()); + return error(Printf("slice start and limit sizes differ: %zu vs %zu", + starts.size(), limits.size())); } if (starts.size() != strides.size()) { - return InvalidArgument("slice start and strides sizes differ: %zu vs %zu", - starts.size(), strides.size()); + return error(Printf("slice start and strides sizes differ: %zu vs %zu", + starts.size(), strides.size())); } if (starts.size() != ShapeUtil::Rank(arg)) { @@ -1864,20 +1983,20 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( start_index); } if (limit_index > arg.dimensions(dimension)) { - return InvalidArgument( - "limit index (%lld) must be less than or equal to dimension " - "size (%lld)", - limit_index, arg.dimensions(dimension)); + return error( + Printf("limit index (%lld) must be less than or equal to dimension " + "size (%lld)", + limit_index, arg.dimensions(dimension))); } VLOG(2) << tensorflow::strings::Printf("starts[%lld] = %lld", dimension, start_index); VLOG(2) << tensorflow::strings::Printf("limits[%lld] = %lld", dimension, limit_index); if (start_index > limit_index) { - return InvalidArgument( - "limit index (%lld) must be greater or equal to " - "start index (%lld) in slice with positive stride", - limit_index, start_index); + return error( + Printf("limit index (%lld) must be greater or equal to " + "start index (%lld) in slice with positive stride", + limit_index, start_index)); } if (stride <= 0) { return InvalidArgument("stride (%lld) must be positive", stride); @@ -1900,7 +2019,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( "slicing shape %s at dynamic start_indices %s with slice_sizes={%s}", ShapeUtil::HumanString(operand_shape).c_str(), ShapeUtil::HumanString(start_indices_shape).c_str(), - tensorflow::str_util::Join(slice_sizes, ", ").c_str()); + Join(slice_sizes, ", ").c_str()); if (ShapeUtil::Rank(start_indices_shape) != 1) { return InvalidArgument( @@ -1993,7 +2112,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( ShapeUtil::Rank(update_shape), ShapeUtil::Rank(operand_shape)); } - if (operand_shape.element_type() != update_shape.element_type()) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(operand_shape, + update_shape)) { return InvalidArgument( "dynamic update slice update element type does not match argument. " "operand.element_type: %s vs update.element_type: %s", @@ -2197,8 +2317,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( return InvalidArgument( "Reshape dimensions [%s] are not a permutation of the operand " "dimensions (operand shape is %s).", - tensorflow::str_util::Join(dimensions, ",").c_str(), - ShapeUtil::HumanString(operand).c_str()); + Join(dimensions, ",").c_str(), ShapeUtil::HumanString(operand).c_str()); } return inferred_shape; @@ -2230,24 +2349,26 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(min, "clamp min")); TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(operand, "clamp operand")); TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(max, "clamp max")); - if (!ShapeUtil::SameElementType(min, operand) || - !ShapeUtil::SameElementType(max, operand)) { + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(min, operand) || + !ShapeUtil::SameElementTypeIgnoringFpPrecision(max, operand)) { return InvalidArgument("clamp op with different operand types: %s, %s, %s", ShapeUtil::HumanString(min).c_str(), ShapeUtil::HumanString(operand).c_str(), ShapeUtil::HumanString(max).c_str()); } - if (((ShapeUtil::Compatible(min, operand) || ShapeUtil::IsScalar(min)) && - (ShapeUtil::Compatible(max, operand) || ShapeUtil::IsScalar(max)))) { + if (((ShapeUtil::CompatibleIgnoringFpPrecision(min, operand) || + ShapeUtil::IsScalar(min)) && + (ShapeUtil::CompatibleIgnoringFpPrecision(max, operand) || + ShapeUtil::IsScalar(max)))) { return operand; } if (ShapeUtil::IsScalar(operand)) { - if (ShapeUtil::Compatible(min, max)) { - return min; + if (ShapeUtil::CompatibleIgnoringFpPrecision(min, max)) { + return ShapeUtil::ChangeElementType(min, operand.element_type()); } else if (ShapeUtil::IsScalar(min)) { - return max; + return ShapeUtil::ChangeElementType(max, operand.element_type()); } else if (ShapeUtil::IsScalar(max)) { - return min; + return ShapeUtil::ChangeElementType(min, operand.element_type()); } } return Unimplemented( @@ -2260,7 +2381,15 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( // broadcast from all operands, not just the predicate. /* static */ StatusOr ShapeInference::InferSelectShape( const Shape& pred, const Shape& on_true, const Shape& on_false) { - if (!ShapeUtil::Compatible(on_true, on_false)) { + bool compatible; + if (ShapeUtil::IsTuple(on_true)) { + // Select only defines the top-level buffer, so if it's a tuple, the two + // input must match exactly. + compatible = ShapeUtil::Compatible(on_true, on_false); + } else { + compatible = ShapeUtil::CompatibleIgnoringFpPrecision(on_true, on_false); + } + if (!compatible) { return InvalidArgument( "operands to select must be the same shape; got %s and %s", ShapeUtil::HumanString(on_true).c_str(), @@ -2275,7 +2404,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( // By this stage we know that pred's element type is PRED. Therefore, this // check restricts pred to be a PRED scalar, or a PRED array with the same // dimensions as on_true and on_false. - return on_true; + return ShapeUtil::ChangeElementType( + on_true, ShapeUtil::HigherPrecisionElementType(on_true, on_false)); } else { return Unimplemented( "select operation with non-scalar predicate with dimensionality " @@ -2290,8 +2420,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( // The applied function's arity equals the number of arguments. if (arg_shapes.size() != to_apply.parameters_size()) { string computation_signature = ShapeUtil::HumanString(to_apply); - string argument_shapes = tensorflow::str_util::Join( - arg_shapes, ", ", [](string* out, const Shape* shape) { + string argument_shapes = + Join(arg_shapes, ", ", [](string* out, const Shape* shape) { tensorflow::strings::StrAppend(out, ShapeUtil::HumanString(*shape)); }); return InvalidArgument( diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h index c06340d2d5df239642eb0af4836df64a898a1eaf..b39151ebbc19f5d0b702a80da5069f58c8dfb07d 100644 --- a/tensorflow/compiler/xla/service/shape_inference.h +++ b/tensorflow/compiler/xla/service/shape_inference.h @@ -109,6 +109,11 @@ class ShapeInference { const Shape& lhs, const Shape& rhs, const Window& window, const ConvolutionDimensionNumbers& dimension_numbers); + // Infers the shape produced by the given FFT type on the given operand. + static StatusOr InferFftShape( + const Shape& in, FftType fft_type, + tensorflow::gtl::ArraySlice fft_length); + // Infers the shape produced a cross replica sum with the given operand // shapes. static StatusOr InferCrossReplicaSumShape( diff --git a/tensorflow/compiler/xla/service/shape_inference_test.cc b/tensorflow/compiler/xla/service/shape_inference_test.cc index 99d87f3b550ae72befe254f23fad080dd210aaf4..026c021165785bd3945d6a846dae446ad45da9b7 100644 --- a/tensorflow/compiler/xla/service/shape_inference_test.cc +++ b/tensorflow/compiler/xla/service/shape_inference_test.cc @@ -1512,5 +1512,20 @@ TEST_F(ShapeInferenceTest, Conditional) { "must have the same shape")); } +TEST_F(ShapeInferenceTest, BadSlice) { + auto arg = ShapeUtil::MakeShape(F32, {4}); + StatusOr statusor = + ShapeInference::InferSliceShape(arg, {0}, {5}, {1}); + ASSERT_FALSE(statusor.ok()); + + LOG(INFO) << statusor.status(); + + EXPECT_THAT(statusor.status().error_message(), + HasSubstr("less than or equal to dimension size")) + << statusor.status(); + EXPECT_THAT(statusor.status().error_message(), HasSubstr("argument shape")) + << statusor.status(); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/shaped_buffer.cc b/tensorflow/compiler/xla/service/shaped_buffer.cc index c679d401c3691b14a43ce77cbe953cd4c64a9e92..6e9986165f7eaf71a964b42b734a5ae5db5e45d7 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.cc +++ b/tensorflow/compiler/xla/service/shaped_buffer.cc @@ -41,7 +41,32 @@ ShapedBuffer::ShapedBuffer(const Shape& on_host_shape, on_device_shape_(on_device_shape), platform_(platform), device_ordinal_(device_ordinal), - buffers_(on_device_shape) {} + buffers_(&on_device_shape_) {} + +ShapedBuffer::ShapedBuffer(ShapedBuffer&& s) + : on_host_shape_(std::move(s.on_host_shape_)), + on_device_shape_(std::move(s.on_device_shape_)), + platform_(s.platform_), + device_ordinal_(s.device_ordinal_), + buffers_(std::move(s.buffers_)) { + // s.buffers_ has a pointer to s.on_device_shape_. When we move s.buffers_ + // into buffers_, we also need to update this pointer so that buffers_ doesn't + // point into s. + buffers_.replace_shape_ptr(&on_device_shape_); +} + +ShapedBuffer& ShapedBuffer::operator=(ShapedBuffer&& s) { + on_host_shape_ = std::move(s.on_host_shape_); + on_device_shape_ = std::move(s.on_device_shape_); + platform_ = s.platform_; + device_ordinal_ = s.device_ordinal_; + buffers_ = std::move(s.buffers_); + // buffers_ has a pointer to its on_device_shape_. When we move s.buffers_ + // into buffers_, we also need to update this pointer so that buffers_ doesn't + // point into s. + buffers_.replace_shape_ptr(&on_device_shape_); + return *this; +} void ShapedBuffer::clear() { for (auto& pair : buffers_) { @@ -99,6 +124,10 @@ ScopedShapedBuffer::ScopedShapedBuffer(const Shape& on_host_shape, device_ordinal), allocator_(allocator) {} +ScopedShapedBuffer::ScopedShapedBuffer(ShapedBuffer shaped_buffer, + DeviceMemoryAllocator* allocator) + : ShapedBuffer(std::move(shaped_buffer)), allocator_(allocator) {} + ScopedShapedBuffer::~ScopedShapedBuffer() { // Deallocate all non-null buffers. A buffer may appear in more than one spot // in the shape (eg, a tuple with a repeated element) so keep track of what @@ -116,12 +145,8 @@ ScopedShapedBuffer::~ScopedShapedBuffer() { } std::unique_ptr ScopedShapedBuffer::release() { - auto shaped_buffer = MakeUnique( - on_host_shape(), on_device_shape(), platform(), device_ordinal()); - - shaped_buffer->buffers() = buffers(); - clear(); - + auto shaped_buffer = MakeUnique(std::move(*this)); + buffers_ = ShapeTree(); return shaped_buffer; } diff --git a/tensorflow/compiler/xla/service/shaped_buffer.h b/tensorflow/compiler/xla/service/shaped_buffer.h index d397e47d2ca734458c7dc99baa5c81b16d0fd72b..b816df8385ef65b0b69ede1d6e65a1991b4bd7c6 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.h +++ b/tensorflow/compiler/xla/service/shaped_buffer.h @@ -87,18 +87,24 @@ class ShapedBuffer { string ToString() const; + ShapedBuffer(ShapedBuffer&& s); + ShapedBuffer& operator=(ShapedBuffer&&); + protected: + ShapedBuffer(const ShapedBuffer&) = delete; + ShapedBuffer& operator=(const ShapedBuffer&) = delete; + // The shape of the data when represented on the host. - const Shape on_host_shape_; + Shape on_host_shape_; // The shape of the data on the device. - const Shape on_device_shape_; + Shape on_device_shape_; // The platform the memory is allocated on. const perftools::gputools::Platform* platform_; // The device the memory is allocated on. - const int device_ordinal_; + int device_ordinal_; // The tree of device buffers. Its shape is on_device_shape(). ShapeTree buffers_; @@ -121,14 +127,20 @@ class ScopedShapedBuffer : public ShapedBuffer { ScopedShapedBuffer(const Shape& on_host_shape, const Shape& on_device_shape, DeviceMemoryAllocator* allocator, int device_ordinal); + // Create a ScopedShapedBuffer by taking over the memory from the incoming + // ShapedBuffer. + ScopedShapedBuffer(ShapedBuffer shaped_buffer, + DeviceMemoryAllocator* allocator); + // Return the allocator used to allocate the device memory held in this // ScopedShapedBuffer. DeviceMemoryAllocator* memory_allocator() const { return allocator_; } - // Release all device memory owned by this ScopedShapedBuffer and return the - // device memory pointers in the form of a ShapedBuffer. Device memory - // pointers in this ScopedShapedBuffer object are set to null. This method is - // analogous to std::unique_ptr::release(). + // Release all device memory owned by this ScopedShapedBuffer and + // return the device memory pointers in the form of a + // ShapedBuffer. The returned ShapedBuffer takes over the memory + // from the ScopedShapedBuffer. The resulting ScopedShapedBuffer can + // only be destroyed. std::unique_ptr release(); // All buffers in the shape are deallocated on destruction. diff --git a/tensorflow/compiler/xla/service/source_map_util.cc b/tensorflow/compiler/xla/service/source_map_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..8cbaac7b3760717bcacb57adc8782a5755c0aa6d --- /dev/null +++ b/tensorflow/compiler/xla/service/source_map_util.cc @@ -0,0 +1,66 @@ +/* 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/compiler/xla/service/source_map_util.h" + +#include "tensorflow/compiler/xla/util.h" + +namespace xla { +namespace source_map_util { +namespace { + +Status InvalidParameterArgumentV(const OpMetadata& op_metadata, + const char* format, va_list args) { + string message; + tensorflow::strings::Appendv(&message, format, args); + if (!op_metadata.source_file().empty()) { + tensorflow::strings::Appendf(&message, " (%s:%d)", + op_metadata.source_file().c_str(), + op_metadata.source_line()); + } + return InvalidArgument("%s", message.c_str()); +} + +} // namespace + +Status InvalidParameterArgument(const OpMetadata& op_metadata, + const char* format, ...) { + va_list args; + va_start(args, format); + Status result = InvalidParameterArgumentV(op_metadata, format, args); + va_end(args); + return result; +} + +Status InvalidParameterArgument(Executable* executable, int parameter_number, + const char* format, ...) { + va_list args; + va_start(args, format); + if (executable != nullptr && executable->has_module()) { + const HloModule& module = executable->module(); + const HloComputation& computation = *module.entry_computation(); + HloInstruction* param = computation.parameter_instruction(parameter_number); + const OpMetadata& metadata = param->metadata(); + Status result = InvalidParameterArgumentV(metadata, format, args); + va_end(args); + return result; + } + Status result = InvalidArgumentV(format, args); + va_end(args); + return result; +} + +} // namespace source_map_util +} // namespace xla diff --git a/tensorflow/compiler/xla/service/source_map_util.h b/tensorflow/compiler/xla/service/source_map_util.h new file mode 100644 index 0000000000000000000000000000000000000000..a776d745f4e56ca4f3d2480740259832bbc85011 --- /dev/null +++ b/tensorflow/compiler/xla/service/source_map_util.h @@ -0,0 +1,46 @@ +/* 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_COMPILER_XLA_SOURCE_MAP_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_SOURCE_MAP_UTIL_H_ + +#include "tensorflow/compiler/xla/service/executable.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { +namespace source_map_util { + +// Creates an INVALID_ARUGMENT status with the given format string. +// +// Also, attempts to extract the OpMetadata for parameter_number on executable +// and append it to the status message for source mapping to user code. +// +// executable may be nullptr, but parameter_number should not be out of bounds +// or a CHECK-failure may occur. +Status InvalidParameterArgument(Executable* executable, int parameter_number, + const char* format, ...) + TF_PRINTF_ATTRIBUTE(3, 4); + +// As above, but takes the parameter metadata directly instead of extracting it +// from the executable. +Status InvalidParameterArgument(const OpMetadata& op_metadata, + const char* format, ...) + TF_PRINTF_ATTRIBUTE(2, 3); + +} // namespace source_map_util +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SOURCE_MAP_UTIL_H_ diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h index 5ca6ccb5c93d4eb7b5b2058848bdd369c1fde35d..c3743b150168ebcf1051050dc511e50c43108c4f 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h @@ -199,12 +199,10 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault { StatusOr GetBufferDefinedAt( const HloInstruction* instruction, const ShapeIndex& index) const; - // Return a vector containing all BufferAliases of the given logical buffer - // This trivially includes the BufferAlias with same instruction and index as - // the logical buffer itself, so the returned vector is never empty. The - // buffer alias set is the inverse of the points-to set. That is, - // LogicalBuffer B is in the points-to set of instruction I at index N iff - // instruction I, index N is a BufferAlias of B. + // Return a (possibly empty) vector containing all BufferAliases of the given + // logical buffer The buffer alias set is the inverse of the points-to set. + // That is, LogicalBuffer B is in the points-to set of instruction I at index + // N iff instruction I, index N is a BufferAlias of B. using BufferAliasVector = tensorflow::gtl::InlinedVector; const BufferAliasVector& GetBufferAliases(const LogicalBuffer& buffer) const; diff --git a/tensorflow/compiler/xla/service/tuple_util.cc b/tensorflow/compiler/xla/service/tuple_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..4a530bb0b20582b303f4af969514748b46fd5064 --- /dev/null +++ b/tensorflow/compiler/xla/service/tuple_util.cc @@ -0,0 +1,61 @@ +/* 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/compiler/xla/service/tuple_util.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/core/lib/gtl/array_slice.h" + +namespace xla { + +/*static*/ HloInstruction* TupleUtil::ExtractPrefix(HloInstruction* input_tuple, + int64 elements) { + CHECK(ShapeUtil::IsTuple(input_tuple->shape())); + + HloComputation* computation = input_tuple->parent(); + const Shape& input_shape = input_tuple->shape(); + + std::vector tuple_elements; + tuple_elements.reserve(elements); + for (int i = 0; i < elements; i++) { + tuple_elements.push_back( + computation->AddInstruction(HloInstruction::CreateGetTupleElement( + input_shape.tuple_shapes(i), input_tuple, i))); + } + + return computation->AddInstruction( + HloInstruction::CreateTuple(tuple_elements)); +} + +/*static*/ HloInstruction* TupleUtil::AppendSuffix( + HloInstruction* input_tuple, + tensorflow::gtl::ArraySlice trailing_values) { + CHECK(ShapeUtil::IsTuple(input_tuple->shape())); + + HloComputation* computation = input_tuple->parent(); + const Shape& input_shape = input_tuple->shape(); + std::vector tuple_elements; + tuple_elements.reserve(input_shape.tuple_shapes_size()); + for (int i = 0; i < input_shape.tuple_shapes_size(); i++) { + tuple_elements.push_back( + computation->AddInstruction(HloInstruction::CreateGetTupleElement( + input_shape.tuple_shapes(i), input_tuple, i))); + } + tuple_elements.insert(tuple_elements.end(), trailing_values.begin(), + trailing_values.end()); + return computation->AddInstruction( + HloInstruction::CreateTuple(tuple_elements)); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/tuple_util.h b/tensorflow/compiler/xla/service/tuple_util.h new file mode 100644 index 0000000000000000000000000000000000000000..e5ff9aaa8357fe8e4777d6dee37bbec72e144c06 --- /dev/null +++ b/tensorflow/compiler/xla/service/tuple_util.h @@ -0,0 +1,45 @@ +/* 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_COMPILER_XLA_SERVICE_TUPLE_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_TUPLE_UTIL_H_ + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" + +namespace xla { +class TupleUtil { + public: + // Generates HLO instructions to get a prefix tuple from `input_tuple` (which + // must be of tuple shape) of length `elements`. Returns the root of the + // graph of instructions generated. + // + // The instructions are generated into the computation containing + // `input_tuple`. + static HloInstruction* ExtractPrefix(HloInstruction* input_tuple, + int64 elements); + + // Generates HLO instructions to create a tuple that consists of the values in + // `trailing_values` appended to `input_tuple` (which must be of tuple shape). + // Returns the root of the graph of instructions generated. + // + // The instructions are generated into the computation containing + // `input_tuple`. + static HloInstruction* AppendSuffix( + HloInstruction* input_tuple, + tensorflow::gtl::ArraySlice trailing_values); +}; +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_TUPLE_UTIL_H_ diff --git a/tensorflow/compiler/xla/service/tuple_util_test.cc b/tensorflow/compiler/xla/service/tuple_util_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..754fd8ef169231827eeb5bfd72aeb596644ca767 --- /dev/null +++ b/tensorflow/compiler/xla/service/tuple_util_test.cc @@ -0,0 +1,81 @@ +/* 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/compiler/xla/service/tuple_util.h" + +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" + +namespace xla { +namespace { + +namespace op = ::xla::testing::opcode_matchers; + +StatusOr> GetParsedModule( + HloComputation** entry_computation, HloInstruction** param0, + HloInstruction** param1) { + const char* const hlo_string = R"( +HloModule Module + +ENTRY entry { + p0 = (f32[32,32]{1,0},f32[32,32]{1,0},f32[32,32]{1,0}) parameter(0) + ROOT p1 = f32[32,32]{1,0} parameter(1) +} +)"; + + TF_ASSIGN_OR_RETURN(std::unique_ptr module, + tools::Parse(hlo_string)); + + *entry_computation = module->entry_computation(); + *param0 = (*entry_computation)->parameter_instruction(0); + *param1 = (*entry_computation)->parameter_instruction(1); + + return std::move(module); +} + +TEST(TupleUtilTest, ExtractPrefix) { + HloInstruction *param0, *param1; + HloComputation* entry_computation; + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + GetParsedModule(&entry_computation, ¶m0, ¶m1)); + + HloInstruction* prefix = TupleUtil::ExtractPrefix(param0, 2); + + EXPECT_THAT(prefix, op::Tuple(op::GetTupleElement(op::Parameter(0), 0), + op::GetTupleElement(op::Parameter(0), 1))); +} + +TEST(TupleUtilTest, AppendSuffix) { + HloInstruction *param0, *param1; + HloComputation* entry_computation; + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + GetParsedModule(&entry_computation, ¶m0, ¶m1)); + + HloInstruction* with_suffix = + TupleUtil::AppendSuffix(param0, {param1, param1}); + + EXPECT_THAT(with_suffix, op::Tuple(op::GetTupleElement(op::Parameter(0), 0), + op::GetTupleElement(op::Parameter(0), 1), + op::GetTupleElement(op::Parameter(0), 2), + op::Parameter(1), op::Parameter(1))); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/user_computation.cc b/tensorflow/compiler/xla/service/user_computation.cc index e42cbfa97690eb3f8df38404659f2c478ab4353f..fead9b92362bcd1974f2dff6e030bc47dfc5aa85 100644 --- a/tensorflow/compiler/xla/service/user_computation.cc +++ b/tensorflow/compiler/xla/service/user_computation.cc @@ -369,14 +369,6 @@ StatusOr UserComputation::AddRngInstruction( // Check the number of parameters per RNG distribution. switch (rng_request.distribution()) { - case RandomDistribution::RNG_BERNOULLI: - if (rng_request.parameter_size() != 1) { - return InvalidArgument( - "RNG distribution (%s) expects 1 parameters, but got %d", - RandomDistribution_Name(rng_request.distribution()).c_str(), - rng_request.parameter_size()); - } - break; case RandomDistribution::RNG_NORMAL: case RandomDistribution::RNG_UNIFORM: if (rng_request.parameter_size() != 2) { @@ -1121,6 +1113,31 @@ StatusOr UserComputation::AddConvolveInstruction( return handle; } +StatusOr UserComputation::AddFftInstruction( + const FftRequest& fft_request) { + tensorflow::mutex_lock lock(mutex_); + + TF_ASSIGN_OR_RETURN(const OperationRequest* operand, + LookUpRequest(fft_request.operand())); + TF_ASSIGN_OR_RETURN(Shape shape, + ShapeInference::InferFftShape( + operand->output_shape(), fft_request.fft_type(), + AsInt64Slice(fft_request.fft_length()))); + + const ComputationDataHandle handle = CreateComputationDataHandle(); + + OperationRequest& request = + (*session_computation_.mutable_requests())[handle.handle()]; + *request.mutable_output_handle() = handle; + *request.mutable_output_shape() = shape; + *request.mutable_request()->mutable_fft_request() = fft_request; + + VLOG(1) << "AddFftInstruction (" << GetVersionedHandleInternal() + << "), data handle " << handle.handle() << ": " + << fft_request.ShortDebugString(); + return handle; +} + StatusOr UserComputation::AddCrossReplicaSumInstruction( const CrossReplicaSumRequest& cross_replica_sum_request) { tensorflow::mutex_lock lock(mutex_); @@ -1168,7 +1185,7 @@ StatusOr UserComputation::AddInfeedInstruction( return handle; } -Status UserComputation::AddOutfeedInstruction( +StatusOr UserComputation::AddOutfeedInstruction( const OutfeedRequest& outfeed_request) { tensorflow::mutex_lock lock(mutex_); @@ -1180,8 +1197,6 @@ Status UserComputation::AddOutfeedInstruction( // Verify that operand is valid. TF_RETURN_IF_ERROR(LookUpRequest(outfeed_request.operand()).status()); - // No handle is returned, but a handle must be assigned to this instruction - // for computation versioning. ComputationDataHandle handle = CreateComputationDataHandle(); OperationRequest& request = (*session_computation_.mutable_requests())[handle.handle()]; @@ -1192,7 +1207,7 @@ Status UserComputation::AddOutfeedInstruction( VLOG(1) << "AddOutfeedInstruction (" << GetVersionedHandleInternal() << "), data handle " << handle.handle() << ": " << outfeed_request.ShortDebugString(); - return Status::OK(); + return handle; } StatusOr UserComputation::AddCallInstruction( @@ -1675,6 +1690,13 @@ void PureFunctionalVisitor(const SessionComputation& session_computation, break; } + case OpRequest::kFftRequest: { + const FftRequest& fft_request = request.request().fft_request(); + PureFunctionalVisitor(session_computation, fft_request.operand(), + num_parameters, visited, is_functional); + break; + } + case OpRequest::kCrossReplicaSumRequest: { // TODO(b/33009255): Implmement constant folding for cross replica sum. *is_functional = false; @@ -1975,6 +1997,9 @@ void PureFunctionalVisitor(const SessionComputation& session_computation, default: LOG(FATAL) << "Unexpected request type: " << request.request().op_case(); } + if (!*is_functional) { + VLOG(1) << "Non-functional: " << request.request().DebugString(); + } visited->insert(handle.handle()); } @@ -2122,6 +2147,24 @@ UserComputation::GetEmbeddedComputations( return computations; } +StatusOr +UserComputation::LookUpRequestForErrorReporting( + const ComputationDataHandle& handle) const { + tensorflow::mutex_lock lock(mutex_); + return LookUpRequest(handle); +} + +tensorflow::gtl::optional UserComputation::ParameterMetadata( + int parameter_number) const { + tensorflow::mutex_lock lock(mutex_); + auto it = parameters_.find(parameter_number); + if (it == parameters_.end()) { + return tensorflow::gtl::nullopt; + } + OperationRequest* op = it->second; + return &op->request().metadata(); +} + Status UserComputation::RemapEmbeddedComputations( const std::map& old_to_new) { auto update = [&old_to_new](ComputationHandle* to_update) -> Status { @@ -2406,6 +2449,12 @@ static void ForEachOperand( break; } + case OpRequest::kFftRequest: { + const FftRequest& fft_request = request.request().fft_request(); + apply(fft_request.operand()); + break; + } + case OpRequest::kBatchNormTrainingRequest: { const BatchNormTrainingRequest& batch_norm_training_request = request.request().batch_norm_training_request(); @@ -2719,48 +2768,11 @@ HloComputation* ComputationLowerer::ResolveComputation( HloInstruction* ComputationLowerer::ImplicitBroadcastToExplicitBroadcast( HloInstruction* operand, const Shape& output_shape) { - CHECK(ShapeUtil::IsScalar(operand->shape()) || - ShapeUtil::Rank(operand->shape()) == ShapeUtil::Rank(output_shape)); - Shape broadcast_shape = ShapeUtil::MakeShape( - operand->shape().element_type(), AsInt64Slice(output_shape.dimensions())); - // Do explicit broadcast for scalar. - if (ShapeUtil::IsScalar(operand->shape())) { - HloInstruction* broadcast = hlo_builder_.AddInstruction( - HloInstruction::CreateBroadcast(broadcast_shape, operand, {})); - broadcast->set_metadata(operand->metadata()); - if (operand->has_sharding()) { - broadcast->set_sharding(operand->sharding()); - } - return broadcast; - } - // Do explicit broadcast for degenerate broadcast. - std::vector broadcast_dimensions; - std::vector reshaped_dimensions; - for (int i = 0; i < ShapeUtil::Rank(operand->shape()); i++) { - if (operand->shape().dimensions(i) == output_shape.dimensions(i)) { - broadcast_dimensions.push_back(i); - reshaped_dimensions.push_back(operand->shape().dimensions(i)); - } - } - // Eliminate the size one dimensions. - HloInstruction* reshaped_operand = - hlo_builder_.AddInstruction(HloInstruction::CreateReshape( - ShapeUtil::MakeShape(operand->shape().element_type(), - reshaped_dimensions), - operand)); - reshaped_operand->set_metadata(operand->metadata()); - if (operand->has_sharding()) { - reshaped_operand->set_sharding(operand->sharding()); - } - // Broadcast 'reshape' up to the larger size. - HloInstruction* broadcast = - hlo_builder_.AddInstruction(HloInstruction::CreateBroadcast( - broadcast_shape, reshaped_operand, broadcast_dimensions)); - broadcast->set_metadata(operand->metadata()); - if (operand->has_sharding()) { - broadcast->set_sharding(operand->sharding()); - } - return broadcast; + auto fadd = [this](std::unique_ptr x) { + return hlo_builder_.AddInstruction(std::move(x)); + }; + return fadd( + HloInstruction::CreateBroadcastSequence(output_shape, operand, fadd)); } void ComputationLowerer::Visit( @@ -2801,7 +2813,8 @@ void ComputationLowerer::Visit( const ConstantRequest& constant_request = request.request().constant_request(); hlo_instruction = add_instruction(HloInstruction::CreateConstant( - Literal(constant_request.literal()).CloneToUnique())); + Literal::CreateFromProto(constant_request.literal()) + .ConsumeValueOrDie())); break; } @@ -2880,6 +2893,15 @@ void ComputationLowerer::Visit( break; } + case OpRequest::kFftRequest: { + const FftRequest& fft_request = request.request().fft_request(); + HloInstruction* operand = lookup_instruction(fft_request.operand()); + hlo_instruction = add_instruction(HloInstruction::CreateFft( + request.output_shape(), operand, fft_request.fft_type(), + AsInt64Slice(fft_request.fft_length()))); + break; + } + case OpRequest::kDotRequest: { const DotRequest& dot_request = request.request().dot_request(); HloInstruction* lhs = lookup_instruction(dot_request.lhs()); diff --git a/tensorflow/compiler/xla/service/user_computation.h b/tensorflow/compiler/xla/service/user_computation.h index 8a78d520e19024f5e397d6e0c2f4e0523264e176..54bb24d6d7fe7aa8cc7c684795e40464e4eb6614 100644 --- a/tensorflow/compiler/xla/service/user_computation.h +++ b/tensorflow/compiler/xla/service/user_computation.h @@ -133,6 +133,10 @@ class UserComputation { StatusOr AddConvolveInstruction( const ConvolveRequest& convolve_request); + // Enqueues an FFT instruction onto this user computation. + StatusOr AddFftInstruction( + const FftRequest& fft_request); + // Enqueues a cross replica sum instruction onto this user computation. StatusOr AddCrossReplicaSumInstruction( const CrossReplicaSumRequest& cross_replica_sum_request); @@ -142,7 +146,8 @@ class UserComputation { const InfeedRequest& infeed_request); // Enqueues an outfeed instruction onto this user computation. - Status AddOutfeedInstruction(const OutfeedRequest& outfeed_request); + StatusOr AddOutfeedInstruction( + const OutfeedRequest& outfeed_request); // Enqueues a call instruction onto this user computation. StatusOr AddCallInstruction( @@ -317,6 +322,23 @@ class UserComputation { SessionComputation CloneSessionComputation( VersionedComputationHandle::Version version) const; + // Warning: typically we don't want to look up computation data handles until + // the computation is finished being built, for consistency purposes. We + // expose this routine for error reporting purposes so that we can provide + // more meaningful error messages from the XLA service layer. + // + // Returns the operation request that the handle comes from. + StatusOr LookUpRequestForErrorReporting( + const ComputationDataHandle& handle) const; + + // Retrieves the parameter metadata for the given parameter number. + // + // If the parameter number is invalid for this computation, nullopt is + // returned. When the return value has_value(), nullptr will never be + // the held value. + tensorflow::gtl::optional ParameterMetadata( + int parameter_number) const; + private: // Warning: dangerous mutating operation that doesn't respect versioning. // This is only used at initialization time when constructing from a diff --git a/tensorflow/compiler/xla/service/user_computation_test.cc b/tensorflow/compiler/xla/service/user_computation_test.cc index ca02115863e6906ef709ba63259024877e0dcef4..2fa163953f638c0038e9f6bb11ce2a3742e0558c 100644 --- a/tensorflow/compiler/xla/service/user_computation_test.cc +++ b/tensorflow/compiler/xla/service/user_computation_test.cc @@ -67,7 +67,8 @@ TEST_F(UserComputationTest, SimpleComputation) { *outfeed_request.mutable_operand() = constant_handle; *outfeed_request.mutable_shape() = kVectorShape; outfeed_request.set_outfeed_config("abc"); - TF_ASSERT_OK(computation.AddOutfeedInstruction(outfeed_request)); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle outfeed_handle, + computation.AddOutfeedInstruction(outfeed_request)); auto hlo_resolver = [](const VersionedComputationHandle& handle) { return nullptr; diff --git a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc new file mode 100644 index 0000000000000000000000000000000000000000..a5f9b01f011ce04f1114c74391a967c62f015221 --- /dev/null +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc @@ -0,0 +1,296 @@ +/* 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/compiler/xla/service/while_loop_invariant_code_motion.h" +#include "tensorflow/compiler/xla/service/tuple_util.h" +#include "tensorflow/compiler/xla/service/while_util.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/lib/gtl/flatset.h" +#include "tensorflow/core/lib/gtl/inlined_vector.h" + +namespace xla { + +using tensorflow::gtl::FlatMap; +using tensorflow::gtl::FlatSet; +using tensorflow::gtl::InlinedVector; + +// Copies `to_hoist` to the computation containing `while_instr`, hoisting its +// operands as needed. All of its transitive operands are expected to be either +// in `hoisted_instructions` or `unhoisted_invariant_instructions`. This +// function hoists the operands in `unhoisted_invariant_instructions` and moves +// them into `hoisted_instructions`. +static void CreateLoopInvariantCopy( + FlatMap* hoisted_instructions, + FlatSet* unhoisted_invariant_instructions, + HloInstruction* while_instr, HloInstruction* to_hoist) { + HloComputation* parent_of_while = while_instr->parent(); + HloComputation* while_body = while_instr->while_body(); + + struct DFSFrame { + HloInstruction* instruction; + int64 operand_index; + }; + + InlinedVector dfs_stack; + dfs_stack.push_back({to_hoist, 0}); + + HloInstruction* while_body_param = while_body->parameter_instruction(0); + HloInstruction* while_operand = while_instr->mutable_operand(0); + + do { + DFSFrame* frame = &dfs_stack.back(); + if (frame->operand_index == frame->instruction->operand_count()) { + HloInstruction* old_instruction = frame->instruction; + + // All of the operands for old_instruction have been cloned, so it is + // time to clone old_instruction itself. + + auto get_new_operand = [&](HloInstruction* old_operand) { + return old_operand == while_body_param + ? while_operand + : FindOrDie(*hoisted_instructions, old_operand); + }; + + InlinedVector new_operands; + c_transform(old_instruction->operands(), std::back_inserter(new_operands), + get_new_operand); + + HloInstruction* new_instruction = + parent_of_while->AddInstruction(old_instruction->CloneWithNewOperands( + old_instruction->shape(), new_operands)); + + InsertOrDie(hoisted_instructions, old_instruction, new_instruction); + + // Approximately half of the instructions that would normally be present + // in unhoisted_invariant_instructions are constants. We save a bit of + // compile time by not putting these in the hashtable. + CHECK_EQ(unhoisted_invariant_instructions->erase(old_instruction), + to_hoist != old_instruction && + old_instruction->opcode() != HloOpcode::kConstant); + dfs_stack.pop_back(); + continue; + } + + HloInstruction* next_operand = + frame->instruction->mutable_operand(frame->operand_index++); + if (hoisted_instructions->count(next_operand) || + next_operand == while_body_param) { + continue; + } + + dfs_stack.push_back({next_operand, 0}); + } while (!dfs_stack.empty()); +} + +// Returns true if `instruction` is worth hoisting only if it lets us hoist some +// instruction using it. The rationale is that hoisting these instructions will +// prevent simplification and fusion in the while body. +static bool NotWorthHoistingIndividually(const HloInstruction& instruction) { + switch (instruction.opcode()) { + default: + return false; + + case HloOpcode::kBitcast: + case HloOpcode::kBroadcast: + case HloOpcode::kConstant: + case HloOpcode::kReverse: + case HloOpcode::kSlice: + case HloOpcode::kTuple: + return true; + + case HloOpcode::kTranspose: + return ShapeUtil::TransposeIsBitcast( + /*input_shape=*/instruction.operand(0)->shape(), + /*output_shape=*/instruction.shape(), instruction.dimensions()); + + case HloOpcode::kReshape: + return ShapeUtil::ReshapeIsBitcast( + /*input_shape=*/instruction.operand(0)->shape(), + /*output_shape=*/instruction.shape()); + } +} + +// Populates `gte_set` with the GetTupleElement instructions in `while_body` +// that access elements in the parameter tuple that don't change across +// iterations. Assumes `while_body` is the body computation of the while loop +// in question. +static void GatherInvariantGTEs(HloComputation* while_body, + FlatSet* gte_set) { + const HloInstruction::InstructionVector root_operands = + while_body->root_instruction()->operands(); + for (int i = 0; i < root_operands.size(); i++) { + HloInstruction* instr = root_operands[i]; + if (instr->opcode() == HloOpcode::kGetTupleElement && + instr->tuple_index() == i && + instr->operand(0) == while_body->parameter_instruction(0) && + ShapeUtil::IsArray(instr->shape())) { + InsertOrDie(gte_set, instr); + } + } +} + +static StatusOr TryHoistingInvariantInstructionsFromWhileBody( + HloInstruction* while_instr) { + auto print_no_metadata = HloPrintOptions{}.set_print_metadata(false); + + if (!ShapeUtil::IsTuple(while_instr->shape())) { + // This restriction leaves one interesting pattern on the table: + // + // while_body(f32[1024, 1024] %param) { + // %value = expensive_op(%param) + // outfeed(%value) + // ROOT = %param + // } + // + // If we see that pattern in the while, instead of generalizing this + // algorithm to work with non-tuples, we should instead add a pass that + // canonicalizes while loops like the above to use a tuple state. + return false; + } + + string while_instr_name = while_instr->ToString(print_no_metadata); + VLOG(2) << "Trying to hoist from " << while_instr_name; + + HloComputation* while_body = while_instr->while_body(); + + // Maps instructions in the while body to instructions hoisted outside the + // while that compute the same value. + FlatMap hoisted_instructions; + + // Contains instructions that can be legally hoisted, but were deemed to be + // unprofitable to be hoisted alone by NotWorthHoistingIndividually. When we + // hoist an instruction in this set, we move it from + // unhoisted_invariant_instructions to hoisted_instructions. + FlatSet unhoisted_invariant_instructions; + + // Invariant GTE's axiomatically satisfy the constraints for + // unhoisted_invariant_instructions -- they can be legally hoisted, but there + // is no benefit to hoisting them unless something that uses it is also + // hoisted. + GatherInvariantGTEs(while_body, &unhoisted_invariant_instructions); + + if (unhoisted_invariant_instructions.empty()) { + // There are no obviously loop invariant elements in the state being + // threaded through the while loop so give up. In theory this precondition + // is too strong -- we could have code that e.g. permutes the elements in + // the while state but uses a select to pick the same value on every + // iteration. + return false; + } + + // instructions_to_replace[i] is hoisted into a loop invariant instruction + // replacement_instructions[i]. + std::vector instructions_to_replace; + std::vector replacement_instructions; + + for (auto* instruction : while_body->MakeInstructionPostOrder()) { + if (instruction->HasSideEffect() || + instruction->opcode() == HloOpcode::kParameter || + !instruction->control_predecessors().empty() || + !instruction->control_successors().empty()) { + continue; + } + + auto is_invariant = [&](HloInstruction* op) { + return hoisted_instructions.find(op) != hoisted_instructions.end() || + unhoisted_invariant_instructions.count(op) || + op->opcode() == HloOpcode::kConstant; + }; + + if (!c_all_of(instruction->operands(), is_invariant)) { + continue; + } + + if (NotWorthHoistingIndividually(*instruction)) { + VLOG(2) << "Adding " << instruction->ToString(print_no_metadata) + << " to unhoisted invariant set."; + // Approximately half of the instructions that reach this point are + // constants. We save a bit of compile time by not putting these in the + // hashtable. + if (instruction->opcode() != HloOpcode::kConstant) { + InsertOrDie(&unhoisted_invariant_instructions, instruction); + } + continue; + } + + VLOG(2) << "Hoisting " << instruction->ToString(print_no_metadata); + + CreateLoopInvariantCopy(&hoisted_instructions, + &unhoisted_invariant_instructions, while_instr, + instruction); + + instructions_to_replace.push_back(instruction); + replacement_instructions.push_back( + FindOrDie(hoisted_instructions, instruction)); + } + + if (instructions_to_replace.empty()) { + return false; + } + + TF_ASSIGN_OR_RETURN( + WhileUtil::MakeInstructionsLiveInResult live_in_instructions_result, + WhileUtil::MakeInstructionsLiveIn(while_instr, replacement_instructions)); + + HloComputation* new_while_body = + live_in_instructions_result.new_while_instr->while_body(); + + for (int i = 0; i < instructions_to_replace.size(); i++) { + HloInstruction* instruction_to_replace_in_new_while = + FindOrDie(live_in_instructions_result.while_body_instruction_map, + instructions_to_replace[i]); + TF_RETURN_IF_ERROR(new_while_body->ReplaceInstruction( + instruction_to_replace_in_new_while, + live_in_instructions_result.while_body_live_in_values[i])); + } + + VLOG(1) << "Hoisted " << instructions_to_replace.size() + << " instructions from " << while_instr_name; + + return true; +} + +StatusOr WhileLoopInvariantCodeMotion::Run(HloModule* module) { + bool changed = false; + std::vector while_instrs; + for (auto* comp : module->computations()) { + c_copy_if(comp->instructions(), std::back_inserter(while_instrs), + [](const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kWhile; + }); + } + + for (HloInstruction* while_instr : while_instrs) { + // Right now we only hoist computations from the while body, but + // TryHoistingInvariantInstructionsFromWhileBody can be generalized to + // optimize the condition computation too, if needed. + // + // The transform we do here is a pessmization for while loops that execute + // zero times*, but at this time we expect those to be rare. If this + // becomes a problem we can consider using the conditional HLO to avoid + // doing extra work for while loops with zero trip count. + // + // * We delete while loops that have a zero trip count, so this would have + // to be a while loop with a somewhat opaque condition expression. + + TF_ASSIGN_OR_RETURN( + bool result, + TryHoistingInvariantInstructionsFromWhileBody(while_instr)); + changed |= result; + } + return changed; +} +} // namespace xla diff --git a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h new file mode 100644 index 0000000000000000000000000000000000000000..8c4b765b0003c48cfacb9d28e7c8259ac0927d66 --- /dev/null +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h @@ -0,0 +1,39 @@ +/* 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_COMPILER_XLA_SERVICE_WHILE_LOOP_INVARIANT_CODE_MOTION_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_INVARIANT_CODE_MOTION_H_ + +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +// HLO pass that rewrites while loops to hoist loop invariant instructions in +// the while body into the computation that contains the while instruction. + +class WhileLoopInvariantCodeMotion : public HloPassInterface { + public: + ~WhileLoopInvariantCodeMotion() override = default; + + tensorflow::StringPiece name() const override { + return "while-loop-invariant-code-motion"; + } + StatusOr Run(HloModule* module) override; +}; +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_INVARIANT_CODE_MOTION_H_ diff --git a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion_test.cc b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..799340fda905fb7d40b19b4cb79bb0fcb5629fd3 --- /dev/null +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion_test.cc @@ -0,0 +1,442 @@ +/* Copyright 2017 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/compiler/xla/service/while_loop_invariant_code_motion.h" + +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +namespace op = xla::testing::opcode_matchers; + +class WhileLoopInvariantCodeMotionTest : public HloVerifiedTestBase { + public: + // Makes a computation which has one parameter, of the given shape, and always + // returns PRED[]{true}. This is useful as a dummy loop condition. + HloComputation* MakeAlwaysTrueComputation(const Shape& param_shape, + HloModule* module); +}; + +static void FindOnlyWhileInstruction(HloComputation* computation, + HloInstruction** while_instruction) { + *while_instruction = nullptr; + for (auto* instr : computation->instructions()) { + if (instr->opcode() == HloOpcode::kWhile) { + ASSERT_EQ(*while_instruction, nullptr); + *while_instruction = instr; + } + } + + ASSERT_NE(*while_instruction, nullptr); +} + +HloComputation* WhileLoopInvariantCodeMotionTest::MakeAlwaysTrueComputation( + const Shape& param_shape, HloModule* module) { + HloComputation::Builder builder(TestName() + ".always_true"); + builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(true))); + return module->AddEmbeddedComputation(builder.Build()); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, HoistOneInvariantOperation) { + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + Shape while_shape = + ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloInstruction* gte_0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); + HloInstruction* gte_1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + HloInstruction* add_result = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kAdd, gte_0, gte_1)); + builder.AddInstruction( + HloInstruction::CreateTuple({gte_0, gte_1, add_result})); + + return module().AddEmbeddedComputation(builder.Build()); + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + HloComputation* entry_computation = + module().AddEntryComputation(builder.Build()); + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_TRUE(simplified_loop); + + HloInstruction* transformed_while; + FindOnlyWhileInstruction(entry_computation, &transformed_while); + + EXPECT_THAT(entry_computation->instructions(), Contains(op::Add())); + EXPECT_THAT(transformed_while->while_body()->instructions(), + Each(Not(op::Add()))); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, HoistInvariantOperationTree) { + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + Shape while_shape = + ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloInstruction* gte_0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); + HloInstruction* gte_1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + HloInstruction* gte_2_loop_variant = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 2)); + + HloInstruction* add_result = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kAdd, gte_0, gte_1)); + HloInstruction* mul_result = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kMultiply, add_result, gte_1)); + HloInstruction* negate_result = + builder.AddInstruction(HloInstruction::CreateUnary( + scalar_s32, HloOpcode::kNegate, mul_result)); + HloInstruction* constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(4))); + HloInstruction* sub_result = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kSubtract, negate_result, constant)); + HloInstruction* divide_result = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kDivide, sub_result, gte_2_loop_variant)); + builder.AddInstruction( + HloInstruction::CreateTuple({gte_0, gte_1, divide_result})); + + return module().AddEmbeddedComputation(builder.Build()); + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + HloComputation* entry_computation = + module().AddEntryComputation(builder.Build()); + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_TRUE(simplified_loop); + + HloInstruction* transformed_while; + FindOnlyWhileInstruction(entry_computation, &transformed_while); + + EXPECT_THAT(entry_computation->instructions(), + AllOf(Contains(op::Add()), Contains(op::Multiply()), + Contains(op::Negate()), Contains(op::Subtract()), + Contains(op::Constant()), + + // The division had a loop varying operand so that better + // not be hoisted. + Not(Contains(op::Divide())))); + + EXPECT_THAT(transformed_while->while_body()->instructions(), + Each(Not(AnyOf(op::Add(), op::Multiply(), op::Negate(), + op::Subtract(), op::Constant())))); + + EXPECT_THAT(transformed_while->while_body()->instructions(), + Contains(op::Divide())); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, + DontHoistTriviallyLoopVaryingComputation) { + // Basic negative test: the add expression is not loop invariant. + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloInstruction* gte_0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); + HloInstruction* gte_1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + HloInstruction* add_result = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kAdd, gte_0, gte_1)); + builder.AddInstruction(HloInstruction::CreateTuple({gte_0, add_result})); + + return module().AddEmbeddedComputation(builder.Build()); + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + auto* while_inst = builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + + module().AddEntryComputation(builder.Build()); + + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_FALSE(simplified_loop); + + EXPECT_THAT(while_inst->while_body()->instructions(), Contains(op::Add())); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, + DontHoistLoopVaryingComputationWithAlternatingTuples) { + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + Shape while_shape = + ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloInstruction* gte_0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); + HloInstruction* gte_1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + HloInstruction* add_result = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kAdd, gte_0, gte_1)); + builder.AddInstruction( + HloInstruction::CreateTuple({gte_1, gte_0, add_result})); + + return module().AddEmbeddedComputation(builder.Build()); + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + auto* while_inst = builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + + module().AddEntryComputation(builder.Build()); + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_FALSE(simplified_loop); + + EXPECT_THAT(while_inst->while_body()->instructions(), Contains(op::Add())); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistInstructionWithSideEffects) { + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloInstruction* gte_0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); + HloInstruction* gte_1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + builder.AddInstruction( + HloInstruction::CreateOutfeed(scalar_s32, gte_0, "")); + builder.AddInstruction(HloInstruction::CreateTuple({gte_0, gte_1})); + + return module().AddEmbeddedComputation(builder.Build()); + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + auto* while_inst = builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + + module().AddEntryComputation(builder.Build()); + + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_FALSE(simplified_loop); + + EXPECT_THAT(while_inst->while_body()->instructions(), + Contains(op::Outfeed())); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistBitcastAlone) { + // The bitcast's user, an outfeed, can't be hoisted, so don't hoist the + // bitcast either. + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + auto scalar_f32 = ShapeUtil::MakeShape(F32, {}); + Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloInstruction* gte_0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); + HloInstruction* gte_1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + HloInstruction* bitcast_inst = builder.AddInstruction( + HloInstruction::CreateUnary(scalar_f32, HloOpcode::kBitcast, gte_0)); + builder.AddInstruction( + HloInstruction::CreateOutfeed(scalar_f32, bitcast_inst, "")); + builder.AddInstruction(HloInstruction::CreateTuple({gte_0, gte_1})); + + return module().AddEmbeddedComputation(builder.Build()); + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + auto* while_inst = builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + + module().AddEntryComputation(builder.Build()); + + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_FALSE(simplified_loop); + + EXPECT_THAT(while_inst->while_body()->instructions(), + Contains(op::Outfeed())); + EXPECT_THAT(while_inst->while_body()->instructions(), + Contains(op::Bitcast())); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, HoistBitcastIfNeeded) { + // The bitcast's user can be hoisted, so hoist the bitcast too. + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + auto scalar_f32 = ShapeUtil::MakeShape(F32, {}); + Shape while_shape = + ShapeUtil::MakeTupleShape({scalar_s32, scalar_f32, scalar_f32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloInstruction* gte_0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); + HloInstruction* gte_1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_f32, param, 1)); + HloInstruction* bitcast_inst = builder.AddInstruction( + HloInstruction::CreateUnary(scalar_f32, HloOpcode::kBitcast, gte_0)); + HloInstruction* add_inst = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_f32, HloOpcode::kAdd, bitcast_inst, gte_1)); + builder.AddInstruction( + HloInstruction::CreateTuple({gte_0, gte_1, add_inst})); + + return module().AddEmbeddedComputation(builder.Build()); + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + + HloComputation* entry_computation = + module().AddEntryComputation(builder.Build()); + + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_TRUE(simplified_loop); + + HloInstruction* transformed_while; + FindOnlyWhileInstruction(entry_computation, &transformed_while); + + EXPECT_THAT(transformed_while->while_body()->instructions(), + Each(Not(op::Add()))); + EXPECT_THAT(transformed_while->while_body()->instructions(), + Each(Not(op::Bitcast()))); + EXPECT_THAT(entry_computation->instructions(), Contains(op::Add())); + EXPECT_THAT(entry_computation->instructions(), Contains(op::Bitcast())); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistControlDependencies) { + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + Shape while_shape = + ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32}); + + HloComputation* while_body; + { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloInstruction* gte_0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); + HloInstruction* gte_1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + HloInstruction* add_result = + builder.AddInstruction(HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kAdd, gte_0, gte_1)); + TF_ASSERT_OK(param->AddControlDependencyTo(add_result)); + builder.AddInstruction( + HloInstruction::CreateTuple({gte_0, gte_1, add_result})); + + while_body = module().AddEmbeddedComputation(builder.Build()); + } + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + module().AddEntryComputation(builder.Build()); + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_FALSE(simplified_loop); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, BodyHasNonTupleRoot) { + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".passthrough"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloComputation* result = module().AddEmbeddedComputation(builder.Build()); + + result->AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + return result; + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + module().AddEntryComputation(builder.Build()); + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(&module())); + EXPECT_FALSE(simplified_loop); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.cc b/tensorflow/compiler/xla/service/while_loop_simplifier.cc index fb0e6f7ce00cff48727dc55bf45c07994643331d..981de9b2200a9ae8938db21299580f510834d2f0 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier.cc @@ -236,7 +236,7 @@ static optional GetLoopTripCount(HloInstruction* while_op) { VLOG(2) << "Couldn't evaluate while cond: " << result.status(); return nullopt; } - return result.ValueOrDie()->GetArraySlice() == + return result.ValueOrDie()->data() == tensorflow::gtl::ArraySlice{true}; }; @@ -306,6 +306,11 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { return false; } + if (while_body_root->opcode() != HloOpcode::kTuple) { + VLOG(2) << "While body's root is not a tuple(...) instruction."; + return false; + } + auto print_no_metadata = HloPrintOptions().set_print_metadata(false); // Bail if param0 of while_cond or while_body has users which aren't of type @@ -559,9 +564,11 @@ static StatusOr TryRemoveWhileLoop(HloInstruction* while_op) { // // This is not a fundamental limitation. The control operands can be moved // onto the new HLOs after simplification, and any side-effecting ops inside - // the loop aren't removed, just cloned and added back to the loop. - // Nevertheless our infrastructure sees loop simplification as removal of - // these nodes and currently doesn't allow it. + // the loop aren't removed, just cloned and added back to the loop. But + // moving an op out of the loop also removes implicit control dependencies + // between the op and the ops outside the loop, so we'd have to add those back + // for things like infeed/outfeed. It gets complicated. So for now we just + // avoid it. if (!while_op->parent()->IsRemovable(while_op) || while_op->HasSideEffect()) { VLOG(2) << "Not attempting to remove while loop it is not removable: " << while_op->ToShortString(); @@ -590,7 +597,9 @@ static StatusOr TryRemoveWhileLoop(HloInstruction* while_op) { auto call_op = computation->AddInstruction(HloInstruction::CreateCall( while_op->shape(), while_op->operands(), while_op->while_body())); TF_RETURN_IF_ERROR(computation->ReplaceInstruction(while_op, call_op)); - TF_RETURN_IF_ERROR(CallInliner::Inline(call_op)); + TF_ASSIGN_OR_RETURN(auto inlined_instructions_map, + CallInliner::Inline(call_op)); + (void)inlined_instructions_map; return true; } return false; diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.h b/tensorflow/compiler/xla/service/while_loop_simplifier.h index 50dac32a4ab0a5de756c1ddf5e62c3560e54a079..d3d55634c97bbdf3f81321d8089bb808c411340b 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier.h +++ b/tensorflow/compiler/xla/service/while_loop_simplifier.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_SIMPLIFIER_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_SIMPLIFIER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_SIMPLIFIER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_SIMPLIFIER_H_ #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" @@ -41,4 +41,4 @@ class WhileLoopSimplifier : public HloPassInterface { } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_SIMPLIFIER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_SIMPLIFIER_H_ diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc b/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc index d99b31dc0037968bc88d5f22d53309a6a4546963..c5183f8d3aee99696ed4114c3f7e451888222137 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc @@ -418,5 +418,32 @@ TEST_F(WhileLoopSimplifierTest, RemoveUnusedOperand) { op::GetTupleElement(op::Parameter(0), /*tuple_index=*/1))); } +TEST_F(WhileLoopSimplifierTest, BodyHasNonTupleRoot) { + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32}); + + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".passthrough"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + HloComputation* result = module().AddEmbeddedComputation(builder.Build()); + + result->AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); + return result; + }(); + + HloComputation::Builder builder(TestName()); + auto* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, &module()), + while_body, init_value)); + module().AddEntryComputation(builder.Build()); + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopSimplifier{}.Run(&module())); + EXPECT_FALSE(simplified_loop); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/while_util.cc b/tensorflow/compiler/xla/service/while_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..e20b25e4a08a946f6b58575a4d4e557744f8035c --- /dev/null +++ b/tensorflow/compiler/xla/service/while_util.cc @@ -0,0 +1,140 @@ +/* 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/compiler/xla/service/while_util.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/tuple_util.h" + +namespace xla { + +static StatusOr WidenWhileCondition( + HloComputation* narrow_condition, const Shape& wide_shape) { + const Shape& narrow_shape = + narrow_condition->parameter_instruction(0)->shape(); + + HloComputation* wide_while_cond = [&]() { + HloComputation::Builder builder( + tensorflow::strings::StrCat("wide.", narrow_condition->name())); + builder.AddInstruction( + HloInstruction::CreateParameter(0, wide_shape, "wide_param")); + + // This is needed so that the root instruction is shaped as a PRED[] -- we + // need to get this right to begin with since we can't mutate the type of + // the root instruction later. We later change the root instruction to + // something more appropriate. + builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + return narrow_condition->parent()->AddEmbeddedComputation(builder.Build()); + }(); + + HloInstruction* truncated_parameter = + TupleUtil::ExtractPrefix(wide_while_cond->parameter_instruction(0), + narrow_shape.tuple_shapes_size()); + HloInstruction* call_narrow_cond = wide_while_cond->AddInstruction( + HloInstruction::CreateCall(ShapeUtil::MakeShape(PRED, {}), + {truncated_parameter}, narrow_condition)); + + wide_while_cond->set_root_instruction(call_narrow_cond); + + TF_RETURN_IF_ERROR(CallInliner::Inline(call_narrow_cond).status()); + return wide_while_cond; +} + +static StatusOr> +WidenWhileBody(HloComputation* narrow_body, const Shape& wide_shape) { + const Shape& narrow_shape = narrow_body->parameter_instruction(0)->shape(); + + HloComputation* wide_while_body = [&]() { + HloComputation::Builder builder( + tensorflow::strings::StrCat("wide.", narrow_body->name())); + builder.AddInstruction( + HloInstruction::CreateParameter(0, wide_shape, "wide_param")); + return narrow_body->parent()->AddEmbeddedComputation(builder.Build()); + }(); + + HloInstruction* wide_parameter = wide_while_body->parameter_instruction(0); + HloInstruction* truncated_parameter = TupleUtil::ExtractPrefix( + wide_parameter, narrow_shape.tuple_shapes_size()); + HloInstruction* call_narrow_body = + wide_while_body->AddInstruction(HloInstruction::CreateCall( + narrow_shape, {truncated_parameter}, narrow_body)); + + std::vector live_through_values; + for (int i = narrow_shape.tuple_shapes_size(); + i < wide_shape.tuple_shapes_size(); i++) { + live_through_values.push_back( + wide_while_body->AddInstruction(HloInstruction::CreateGetTupleElement( + wide_shape.tuple_shapes(i), wide_parameter, i))); + } + + wide_while_body->set_root_instruction( + TupleUtil::AppendSuffix(call_narrow_body, live_through_values)); + + TF_ASSIGN_OR_RETURN(auto inlined_instructions_map, + CallInliner::Inline(call_narrow_body)); + return {{wide_while_body, std::move(inlined_instructions_map)}}; +} + +/*static*/ StatusOr +WhileUtil::MakeInstructionsLiveIn( + HloInstruction* while_instr, + tensorflow::gtl::ArraySlice instructions) { + CHECK(ShapeUtil::IsTuple(while_instr->shape())); + + int64 elements_in_old_while_shape = while_instr->shape().tuple_shapes_size(); + Shape new_while_shape = while_instr->shape(); + for (auto* instruction : instructions) { + *new_while_shape.add_tuple_shapes() = instruction->shape(); + } + + TF_ASSIGN_OR_RETURN( + HloComputation * new_while_condition, + WidenWhileCondition(while_instr->while_condition(), new_while_shape)); + + HloComputation* new_while_body; + CallInliner::InlinedInstructionMap inlined_instructions_map; + TF_ASSIGN_OR_RETURN( + std::tie(new_while_body, inlined_instructions_map), + WidenWhileBody(while_instr->while_body(), new_while_shape)); + + HloInstruction* new_while_init = + TupleUtil::AppendSuffix(while_instr->mutable_operand(0), instructions); + HloComputation* containing_computation = while_instr->parent(); + HloInstruction* new_while = containing_computation->AddInstruction( + HloInstruction::CreateWhile(new_while_shape, new_while_condition, + new_while_body, new_while_init)); + TF_RETURN_IF_ERROR(containing_computation->ReplaceInstruction( + while_instr, TupleUtil::ExtractPrefix( + new_while, while_instr->shape().tuple_shapes_size()))); + + HloInstruction* while_body_param = new_while_body->parameter_instruction(0); + std::vector live_in_instructions; + for (int64 i = elements_in_old_while_shape; + i < new_while_shape.tuple_shapes_size(); i++) { + live_in_instructions.push_back( + new_while_body->AddInstruction(HloInstruction::CreateGetTupleElement( + instructions[i - elements_in_old_while_shape]->shape(), + while_body_param, i))); + } + + WhileUtil::MakeInstructionsLiveInResult result; + + result.new_while_instr = new_while; + result.while_body_live_in_values = std::move(live_in_instructions); + result.while_body_instruction_map = std::move(inlined_instructions_map); + + return std::move(result); +} +} // namespace xla diff --git a/tensorflow/compiler/xla/service/while_util.h b/tensorflow/compiler/xla/service/while_util.h new file mode 100644 index 0000000000000000000000000000000000000000..3600b5a80d26e37fdb7d5173c3b8743734306390 --- /dev/null +++ b/tensorflow/compiler/xla/service/while_util.h @@ -0,0 +1,58 @@ +/* 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_COMPILER_XLA_SERVICE_WHILE_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_UTIL_H_ + +#include "tensorflow/compiler/xla/service/call_inliner.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" + +namespace xla { +class WhileUtil { + public: + // Holds a return value from MakeInstructionsLiveIn. + struct MakeInstructionsLiveInResult { + // The new while operation that has the requested values live in. + HloInstruction* new_while_instr; + + // The i'th element of `while_body_live_in_values` is an instruction in the + // while body that holds the i'th *newly added* live in value at runtime. + std::vector while_body_live_in_values; + + // `while_body_instruction_map` maps instructions in the original while body + // to the corresponding instructions in the body for the newly created while + // operation. + CallInliner::InlinedInstructionMap while_body_instruction_map; + }; + + // Replaces `while_instr` with a new while instruction that is equivalent to + // `while_instr`, except that it has all of the HLO instructions in + // `instructions` as live-in, loop invariant values. These new live in values + // are represented as new elements appended to the parameter of the while + // loop, which must be of tuple shape. GetTupleElement instructions computing + // each new live in value is returned in the `while_body_live_in_values` + // vector. + // + // Precondition: `while_instr` must have a tuple shaped state. + // + // Every instruction in `instructions` must be contained in the computation + // that contains `while_instr`. + static StatusOr MakeInstructionsLiveIn( + HloInstruction* while_instr, + tensorflow::gtl::ArraySlice instructions); +}; +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_UTIL_H_ diff --git a/tensorflow/compiler/xla/service/while_util_test.cc b/tensorflow/compiler/xla/service/while_util_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..cf0d0db99bd92b6b364b4e28e56a0902d4065963 --- /dev/null +++ b/tensorflow/compiler/xla/service/while_util_test.cc @@ -0,0 +1,130 @@ +/* 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/compiler/xla/service/while_util.h" + +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" + +namespace xla { +namespace { + +namespace op = ::xla::testing::opcode_matchers; + +StatusOr> GetParsedModule( + HloComputation** entry_computation, HloInstruction** param0, + HloInstruction** param1, HloInstruction** param2) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +while_body { + ROOT p_body = (f32[32,32]{1,0}, f32[32,32]{1,0}) parameter(0) +} + +while_condition { + p_cond = f32[32,32]{1,0} parameter(0) + ROOT result = pred[] constant(true) +} + +ENTRY entry { + p_entry_0 = f32[32,32]{1,0} parameter(0) + p_entry_1 = s32[32,32]{1,0} parameter(1) + p_entry_2 = s64[32,32]{1,0} parameter(2) + while_init = (f32[32,32]{1,0}, f32[32,32]{1,0}) tuple(p_entry_0, p_entry_0) + ROOT while = (f32[32,32]{1,0}, f32[32,32]{1,0}) while(while_init), condition=while_condition, body=while_body +} +)"; + + TF_ASSIGN_OR_RETURN(std::unique_ptr module, + tools::Parse(hlo_string)); + + *entry_computation = module->entry_computation(); + *param0 = (*entry_computation)->parameter_instruction(0); + *param1 = (*entry_computation)->parameter_instruction(1); + *param2 = (*entry_computation)->parameter_instruction(2); + + return std::move(module); +} + +TEST(WhileUtil, MakeZeroInstructionsLiveOp) { + HloInstruction *param0, *param1, *param2; + HloComputation* entry_computation; + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + GetParsedModule(&entry_computation, ¶m0, ¶m1, ¶m2)); + + HloInstruction* while_instr = entry_computation->root_instruction(); + ASSERT_EQ(while_instr->opcode(), HloOpcode::kWhile); + + TF_ASSERT_OK_AND_ASSIGN( + WhileUtil::MakeInstructionsLiveInResult make_live_in_result, + WhileUtil::MakeInstructionsLiveIn(while_instr, /*instructions=*/{})); + + HloInstruction* new_while_instr = make_live_in_result.new_while_instr; + + EXPECT_THAT( + entry_computation->root_instruction(), + op::Tuple(op::GetTupleElement(::testing::Eq(new_while_instr), 0), + op::GetTupleElement(::testing::Eq(new_while_instr), 1))); + + auto param_reconstructed = + op::Tuple(op::GetTupleElement(op::Parameter(0), 0), + op::GetTupleElement(op::Parameter(0), 1)); + + EXPECT_THAT(new_while_instr->while_body()->root_instruction(), + op::Tuple(op::GetTupleElement(param_reconstructed, 0), + op::GetTupleElement(param_reconstructed, 1))); +} + +TEST(WhileUtilTest, MakeTwoInstructionsLive) { + HloInstruction *param0, *param1, *param2; + HloComputation* entry_computation; + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + GetParsedModule(&entry_computation, ¶m0, ¶m1, ¶m2)); + + HloInstruction* while_instr = entry_computation->root_instruction(); + ASSERT_EQ(while_instr->opcode(), HloOpcode::kWhile); + + TF_ASSERT_OK_AND_ASSIGN( + WhileUtil::MakeInstructionsLiveInResult make_live_in_result, + WhileUtil::MakeInstructionsLiveIn(while_instr, + /*instructions=*/{param0, param1})); + + HloInstruction* new_while_instr = make_live_in_result.new_while_instr; + + XLA_VLOG_LINES(3, module->ToString()); + + EXPECT_THAT( + entry_computation->root_instruction(), + op::Tuple(op::GetTupleElement(::testing::Eq(new_while_instr), 0), + op::GetTupleElement(::testing::Eq(new_while_instr), 1))); + + auto first_half_param_reconstructed = + op::Tuple(op::GetTupleElement(op::Parameter(0), 0), + op::GetTupleElement(op::Parameter(0), 1)); + + EXPECT_THAT(new_while_instr->while_body()->root_instruction(), + op::Tuple(op::GetTupleElement(first_half_param_reconstructed, 0), + op::GetTupleElement(first_half_param_reconstructed, 1), + op::GetTupleElement(op::Parameter(0), 2), + op::GetTupleElement(op::Parameter(0), 3))); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.h b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.h index 63afab4206eb072e84745ced3307295c0516da7b..063e312df66ce9cba0fa9f49c2fc6026ba6b74aa 100644 --- a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.h +++ b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_ZERO_SIZED_HLO_ELIMINATION_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_ZERO_SIZED_HLO_ELIMINATION_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_ZERO_SIZED_HLO_ELIMINATION_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_ZERO_SIZED_HLO_ELIMINATION_H_ #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" @@ -29,4 +29,4 @@ class ZeroSizedHloElimination : public HloPassInterface { } }; } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_ZERO_SIZED_HLO_ELIMINATION_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_ZERO_SIZED_HLO_ELIMINATION_H_ diff --git a/tensorflow/compiler/xla/shape_tree.h b/tensorflow/compiler/xla/shape_tree.h index d752619bd65751779c24f061e44e206d66b01465..280f02e88675381bd75108bfae0dd22c462ba718 100644 --- a/tensorflow/compiler/xla/shape_tree.h +++ b/tensorflow/compiler/xla/shape_tree.h @@ -143,6 +143,18 @@ class ShapeTree { // Return the shape represented with this ShapeTree. const Shape& shape() const { return *shape_; } + // Replaces *only* the underlying shape of this ShapeTree. The caller must own + // the Shape object and hence shape_storage_ is not updated. + // + // Only safe to use this if the ShapeTree was constructed with 'explicit + // ShapeTree(const Shape* shape)' or is moved from one such ShapeTree. The + // caller must ensure that the input shape is consistent with the underlying + // tree. + void replace_shape_ptr(const Shape* shape) { + CHECK(shape_storage_.get() == nullptr); + shape_ = shape; + } + // Returns true if the node at the given index is a leaf node (an array // shape). bool IsLeaf(const ShapeIndex& index) const { diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc index 48b7515ecff6604c6901826282968ba55a2af313..604e0173e789348923316174873f58058eaf2815 100644 --- a/tensorflow/compiler/xla/shape_util.cc +++ b/tensorflow/compiler/xla/shape_util.cc @@ -59,6 +59,11 @@ std::ostream& operator<<(std::ostream& out, const ShapeIndex& shape_index) { return out; } +std::ostream& operator<<(std::ostream& out, const ShapeIndexView& shape_index) { + out << shape_index.ToString(); + return out; +} + namespace { // Recursive helper for comparing the equality of two shapes. Returns true if @@ -79,7 +84,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { if (lhs.layout().format() != rhs.layout().format()) { return false; } - if (LayoutUtil::IsDense(lhs)) { + if (LayoutUtil::IsDenseArray(lhs)) { if (!ContainersEqual(LayoutUtil::MinorToMajor(lhs), LayoutUtil::MinorToMajor(rhs))) { VLOG(3) << "CompareShapes: lhs layout != rhs layout"; @@ -148,7 +153,8 @@ StatusOr MakeShapeWithLayoutInternal( } /* static */ int64 ShapeUtil::Rank(const Shape& shape) { - CHECK(!ShapeUtil::IsTuple(shape)) << "Tuples do not have a rank"; + CHECK(!ShapeUtil::IsTuple(shape)) + << "Tuples do not have a rank, shape: " << shape; return shape.dimensions_size(); } @@ -196,6 +202,17 @@ StatusOr MakeShapeWithLayoutInternal( return MakeShapeWithLayout(element_type, dimensions, layout); } +/* static */ Shape ShapeUtil::MakeShapeWithSparseLayout( + PrimitiveType element_type, tensorflow::gtl::ArraySlice dimensions, + int64 max_sparse_elements) { + DCHECK_NE(TUPLE, element_type); + DCHECK_NE(OPAQUE, element_type); + Shape shape = ShapeUtil::MakeShape(element_type, dimensions); + *shape.mutable_layout() = LayoutUtil::MakeSparseLayout(max_sparse_elements); + TF_DCHECK_OK(ShapeUtil::ValidateShape(shape)); + return shape; +} + /* static */ Shape ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( const Shape& shape) { @@ -243,7 +260,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ void ShapeUtil::AppendMajorDimension(int bound, Shape* shape) { - CHECK(LayoutUtil::IsDense(*shape)); + CHECK(LayoutUtil::IsDenseArray(*shape)); shape->mutable_layout()->add_minor_to_major(Rank(*shape)); shape->add_dimensions(bound); TF_DCHECK_OK(ValidateShape(*shape)); @@ -353,7 +370,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ int64 ShapeUtil::ElementsIn(const Shape& shape) { - CHECK(!IsTuple(shape)); + CHECK(!IsTuple(shape)) << ShapeUtil::HumanString(shape); CHECK_EQ(shape.dimensions_size(), Rank(shape)); return std::accumulate( shape.dimensions().begin(), shape.dimensions().end(), 1LL, @@ -369,7 +386,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ string ShapeUtil::HumanString(const Shape& shape) { - if (shape.element_type() == TUPLE) { + if (IsTuple(shape)) { string text = "("; const char* prefix = ""; for (const Shape& elem_shape : shape.tuple_shapes()) { @@ -436,7 +453,7 @@ StatusOr StringToPrimitiveType(const string& name) { } // namespace /* static */ string ShapeUtil::HumanStringWithLayout(const Shape& shape) { - if (shape.element_type() == TUPLE) { + if (IsTuple(shape)) { string text = "("; const char* prefix = ""; for (const Shape& elem_shape : shape.tuple_shapes()) { @@ -458,8 +475,6 @@ StatusOr StringToPrimitiveType(const string& name) { if (LayoutUtil::HasLayout(shape)) { tensorflow::strings::StrAppend(&result, LayoutUtil::HumanString(shape.layout())); - } else { - tensorflow::strings::StrAppend(&result, "{no layout}"); } } return result; @@ -507,26 +522,35 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { string element_type_string; string dimensions_string; + string format_string; string layout_string; // tensorflow::StringPiece is not compatible with internal RE2 StringPiece, so // we convert in to the RE2-consumable type and then consume the corresponding // amount from our StringPiece type. tensorflow::RegexpStringPiece s_consumable(s->data(), s->size()); - if (RE2::Consume(&s_consumable, - "^(\\w*\\d*)\\[([\\d,]*)\\](?:\\s*{([\\d,]*)})?", - &element_type_string, &dimensions_string, &layout_string)) { + if (RE2::Consume( + &s_consumable, + "^(\\w*\\d*)\\[([\\d,]*)\\](?:\\s*(dense|sparse)?\\s*{([\\d,]+)})?", + &element_type_string, &dimensions_string, &format_string, + &layout_string)) { size_t consumed = s->size() - s_consumable.size(); s->remove_prefix(consumed); + auto string_to_int64 = [&s](const string& input) -> StatusOr { + int64 element; + if (!tensorflow::strings::safe_strto64(input.c_str(), &element)) { + return InvalidArgument( + "Invalid s64 value in parsed shape string: \"%s\" in \"%s\"", + input.c_str(), s->ToString().c_str()); + } + return element; + }; + auto comma_list_to_int64s = - [&s](const string& input) -> StatusOr> { + [&s, + string_to_int64](const string& input) -> StatusOr> { std::vector results; for (const string& piece : tensorflow::str_util::Split(input, ',')) { - int64 element; - if (!tensorflow::strings::safe_strto64(piece.c_str(), &element)) { - return InvalidArgument( - "Invalid s64 value in parsed shape string: \"%s\" in \"%s\"", - piece.c_str(), s->ToString().c_str()); - } + TF_ASSIGN_OR_RETURN(int64 element, string_to_int64(piece)); results.push_back(element); } return results; @@ -546,15 +570,23 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { } Shape result; - if (layout_string.empty()) { + if (format_string.empty() && layout_string.empty()) { // Create a shape without a layout set. result = ShapeUtil::MakeShape(primitive_type, dimensions); - } else { + } else if (format_string == "sparse") { + TF_ASSIGN_OR_RETURN(int64 max_elements, string_to_int64(layout_string)); + result = ShapeUtil::MakeShapeWithSparseLayout(primitive_type, dimensions, + max_elements); + } else if (format_string.empty() || format_string == "dense") { // Extract the layout minor-to-major and set it. TF_ASSIGN_OR_RETURN(std::vector min2maj, comma_list_to_int64s(layout_string)); TF_ASSIGN_OR_RETURN(result, MakeShapeWithLayoutInternal( primitive_type, dimensions, min2maj)); + } else { + // This should not be reached. + LOG(FATAL) << "Unhandled condition when parsing shape; format: \"" + << format_string << "\", layout: \"" << layout_string << "\""; } TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(result)); return std::move(result); @@ -567,7 +599,12 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { /* static */ StatusOr ShapeUtil::ParseShapeString( tensorflow::StringPiece s) { - return ParseShapeStringInternal(&s); + TF_ASSIGN_OR_RETURN(Shape shape, ParseShapeStringInternal(&s)); + if (!s.empty()) { + return InvalidArgument("Invalid shape string to parse: \"%s\"", + s.ToString().c_str()); + } + return shape; } /* static */ bool ShapeUtil::SameDimensions(const Shape& lhs, @@ -593,6 +630,19 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { return SameDimensions(lhs, rhs); } +/* static */ bool ShapeUtil::CompatibleIgnoringFpPrecision(const Shape& lhs, + const Shape& rhs) { + if (lhs.element_type() == TUPLE) { + return rhs.element_type() == TUPLE && + ContainersEqual(lhs.tuple_shapes(), rhs.tuple_shapes(), + CompatibleIgnoringFpPrecision); + } + if (SameElementTypeIgnoringFpPrecision(lhs, rhs)) { + return CompatibleIgnoringElementType(lhs, rhs); + } + return false; +} + /* static */ int64 ShapeUtil::GetDimension(const Shape& shape, int64 dimension_number) { return shape.dimensions(GetDimensionNumber(shape, dimension_number)); @@ -652,23 +702,55 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { TF_DCHECK_OK(ValidateShape(shape)); DCHECK_NE(OPAQUE, shape.element_type()); if (shape.element_type() == TUPLE) { - CHECK_GT(pointer_size, 0); - return pointer_size * shape.tuple_shapes_size(); + return ByteSizeOfTupleIndexTable(shape, pointer_size); + } + int64 byte_size = ByteSizeOfElements(shape); + if (LayoutUtil::IsSparseArray(shape)) { + byte_size += ByteSizeOfSparseIndices(shape); } + return byte_size; +} + +/* static */ int64 ShapeUtil::ByteSizeOfTupleIndexTable(const Shape& shape, + int64 pointer_size) { + TF_DCHECK_OK(ValidateShape(shape)); + DCHECK_EQ(TUPLE, shape.element_type()); + CHECK_GT(pointer_size, 0); + return pointer_size * shape.tuple_shapes_size(); +} + +/* static */ int64 ShapeUtil::ByteSizeOfElements(const Shape& shape) { + TF_DCHECK_OK(ValidateShape(shape)); + DCHECK(ShapeUtil::IsArray(shape)); int64 allocated_element_count; - if (shape.layout().padded_dimensions_size() > 0) { - CHECK_EQ(Rank(shape), shape.layout().padded_dimensions_size()); - allocated_element_count = 1; - for (int64 dimension_size : shape.layout().padded_dimensions()) { - allocated_element_count *= dimension_size; - } + + if (LayoutUtil::IsSparseArray(shape)) { + allocated_element_count = LayoutUtil::MaxSparseElements(shape.layout()); } else { - allocated_element_count = ElementsIn(shape); + CHECK(LayoutUtil::IsDenseArray(shape)); + tensorflow::gtl::ArraySlice padded_dimensions = + LayoutUtil::PaddedDimensions(shape); + if (!padded_dimensions.empty()) { + CHECK_EQ(Rank(shape), padded_dimensions.size()); + allocated_element_count = 1; + for (int64 dimension_size : padded_dimensions) { + allocated_element_count *= dimension_size; + } + } else { + allocated_element_count = ElementsIn(shape); + } } return allocated_element_count * ByteSizeOfPrimitiveType(shape.element_type()); } +/* static */ int64 ShapeUtil::ByteSizeOfSparseIndices(const Shape& shape) { + TF_DCHECK_OK(ValidateShape(shape)); + DCHECK(LayoutUtil::IsSparseArray(shape)); + return LayoutUtil::MaxSparseElements(shape.layout()) * + ShapeUtil::Rank(shape) * sizeof(int64); +} + /* static */ Status ShapeUtil::ValidateShapeWithOptionalLayoutInternal( const Shape& shape) { if (shape.element_type() == TUPLE) { @@ -735,7 +817,8 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { ShapeIndexView index) { const Shape* return_shape = &shape; for (auto i : index) { - CHECK(IsTuple(*return_shape)); + CHECK(IsTuple(*return_shape)) + << "Invalid index " << index << " for shape " << shape; return_shape = &return_shape->tuple_shapes(i); } return *return_shape; @@ -893,7 +976,7 @@ Status ForEachMutableSubshapeHelper( new_shape.add_dimensions(dim); } if (shape.has_layout()) { - CHECK(LayoutUtil::IsDense(shape)); + CHECK(LayoutUtil::IsDenseArray(shape)); Layout* new_layout = new_shape.mutable_layout(); new_layout->set_format(DENSE); new_layout->clear_minor_to_major(); @@ -1352,4 +1435,9 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, return shape; } +std::ostream& operator<<(std::ostream& out, const Shape& shape) { + out << ShapeUtil::HumanString(shape); + return out; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/shape_util.h b/tensorflow/compiler/xla/shape_util.h index a2043eff1ed64104cfad69b271edbdb8cebdb02b..19b1aa93bd373ebd5f502d0dca56c9b31ab4fd7f 100644 --- a/tensorflow/compiler/xla/shape_util.h +++ b/tensorflow/compiler/xla/shape_util.h @@ -23,6 +23,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -62,6 +63,9 @@ class ShapeIndex { void push_back(int64 value) { indices_.push_back(value); } void pop_back() { indices_.pop_back(); } + // push_front is O(n^2), but shapes don't usually have a ton of dimensions. + void push_front(int64 value) { indices_.insert(indices_.begin(), value); } + std::vector::const_iterator begin() const { return indices_.begin(); } std::vector::const_iterator end() const { return indices_.end(); } std::vector::iterator begin() { return indices_.begin(); } @@ -134,6 +138,7 @@ class ShapeIndexView { }; std::ostream& operator<<(std::ostream& out, const ShapeIndex& shape_index); +std::ostream& operator<<(std::ostream& out, const ShapeIndexView& shape_index); // Namespaced collection of (static) shape utilities. // @@ -142,7 +147,10 @@ std::ostream& operator<<(std::ostream& out, const ShapeIndex& shape_index); class ShapeUtil { public: // Returns the number of elements are contained within the provided shape; - // e.g. for rank 0 (scalars) the result is always 1. + // e.g. for rank 0 (scalars) the result is always 1. Note that sparse shapes + // may not actually be able to store this number of elements. See + // LayoutUtil::MaxSparseElements(shape) to obtain the maximum number of + // elements that can be stored in a sparse shape. // Precondition: !IsTuple(shape) static int64 ElementsIn(const Shape& shape); @@ -163,6 +171,27 @@ class ShapeUtil { // Precondition: !ShapeUtil::IsOpaque(shape) && !ShapeUtil::IsTuple(shape) static int64 ByteSizeOfPrimitiveType(PrimitiveType primitive_type); + // Returns the number of bytes required to store the tuple member pointers for + // a allocation of shape. The `shape` must be a TUPLE shape, and + // `pointer_size` must be larger than zero. + static int64 ByteSizeOfTupleIndexTable(const Shape& shape, + int64 pointer_size); + + // Returns the number of bytes required for the elements in an allocation of + // `shape`, which must be an array shape. The return value does not include + // the bytes needed to store sparse indices. Dense shapes use a separate + // memory location for each element, and so for these shapes, + // `ByteSizeOf(shape) == ByteSizeOfElements(shape)`. For dense shapes, this + // size also includes padding if present in the layout. For sparse shapes, + // `ByteSizeOf(shape) == ByteSizeOfElements(shape) + + // ByteSizeOfSparseindices(shape)`. + static int64 ByteSizeOfElements(const Shape& shape); + + // Returns the number of bytes required for the sparse indices in an + // allocation of shape. The shape must be an array shape. The return value + // does not include the bytes needed to store sparse indices. + static int64 ByteSizeOfSparseIndices(const Shape& shape); + // Returns a human-readable string that represents the given shape, with or // without layout. e.g. "f32[42x12] {0, 1}" or "f32[64]". static string HumanString(const Shape& shape); @@ -186,6 +215,31 @@ class ShapeUtil { return lhs.element_type() == rhs.element_type(); } + // As SameElementType, but allows floating point types to have different + // precisions. + static bool SameElementTypeIgnoringFpPrecision(const Shape& a, + const Shape& b) { + if (ElementIsFloating(a) && ElementIsFloating(b)) { + return true; + } + return ShapeUtil::SameElementType(a, b); + } + + // Returns the higher-precision element type if a and b are both floating + // point types; otherwise, checks that that they have the same element type + // and returns it. + static PrimitiveType HigherPrecisionElementType(const Shape& a, + const Shape& b) { + if (SameElementType(a, b)) { + return a.element_type(); + } + CHECK(SameElementTypeIgnoringFpPrecision(a, b)); + return primitive_util::BitWidth(a.element_type()) < + primitive_util::BitWidth(b.element_type()) + ? b.element_type() + : a.element_type(); + } + // Returns true if the rank, dimension sizes, and element type are // identical. Layout is ignored. Tuple elements are compared recursively for // compatibility. @@ -196,6 +250,10 @@ class ShapeUtil { // compatibility. static bool CompatibleIgnoringElementType(const Shape& lhs, const Shape& rhs); + // As Compatible, but allow one of lhs and rhs to be BF16 while the other + // being F32. Tuple elements are compared recursively for compatibility. + static bool CompatibleIgnoringFpPrecision(const Shape& lhs, const Shape& rhs); + // Returns whether the lhs and rhs shapes are identical protobufs. static bool Equal(const Shape& lhs, const Shape& rhs); @@ -268,6 +326,10 @@ class ShapeUtil { PrimitiveType element_type, tensorflow::gtl::ArraySlice dimensions, tensorflow::gtl::ArraySlice minor_to_major); + static Shape MakeShapeWithSparseLayout( + PrimitiveType element_type, tensorflow::gtl::ArraySlice dimensions, + int64 max_sparse_elements); + // Constructs a new shape with major-first layout (i.e. {n, n-1, ..., 0}). static Shape MakeShapeWithDescendingLayout( PrimitiveType element_type, @@ -538,6 +600,8 @@ class ShapeUtil { TF_DISALLOW_COPY_AND_ASSIGN(ShapeUtil); }; +std::ostream& operator<<(std::ostream& out, const Shape& shape); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SHAPE_UTIL_H_ diff --git a/tensorflow/compiler/xla/shape_util_test.cc b/tensorflow/compiler/xla/shape_util_test.cc index 3be6d6c4299aff62582c1b9fdc46fb78712f95c8..4db97d45b20b86dc60531845c6e28a223203ff7f 100644 --- a/tensorflow/compiler/xla/shape_util_test.cc +++ b/tensorflow/compiler/xla/shape_util_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/types.h" @@ -71,7 +72,8 @@ TEST(ShapeUtilTest, Rank4DimensionIndexing) { TEST(ShapeUtilTest, ParseShapeStringR2F32) { string shape_string = "f32[123,456]"; - Shape actual = ShapeUtil::ParseShapeString(shape_string).ValueOrDie(); + TF_ASSERT_OK_AND_ASSIGN(Shape actual, + ShapeUtil::ParseShapeString(shape_string)); Shape expected = ShapeUtil::MakeShape(F32, {123, 456}); ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) << "expected: " << ShapeUtil::HumanString(expected) @@ -80,7 +82,8 @@ TEST(ShapeUtilTest, ParseShapeStringR2F32) { TEST(ShapeUtilTest, ParseShapeStringTupleOfArrays) { string shape_string = "(f32[1572864],s8[5120,1024])"; - Shape actual = ShapeUtil::ParseShapeString(shape_string).ValueOrDie(); + TF_ASSERT_OK_AND_ASSIGN(Shape actual, + ShapeUtil::ParseShapeString(shape_string)); Shape expected = ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {1572864}), ShapeUtil::MakeShape(S8, {5120, 1024})}); @@ -91,7 +94,8 @@ TEST(ShapeUtilTest, ParseShapeStringTupleOfArrays) { TEST(ShapeUtilTest, ParseShapeStringNestedTuple) { string shape_string = "(f32[1],(f32[2]), f32[3])"; - Shape actual = ShapeUtil::ParseShapeString(shape_string).ValueOrDie(); + TF_ASSERT_OK_AND_ASSIGN(Shape actual, + ShapeUtil::ParseShapeString(shape_string)); Shape expected = ShapeUtil::MakeTupleShape({ ShapeUtil::MakeShape(F32, {1}), ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {2})}), @@ -102,6 +106,47 @@ TEST(ShapeUtilTest, ParseShapeStringNestedTuple) { << "actual: " << ShapeUtil::HumanString(actual); } +TEST(ShapeUtilTest, ParseShapeStringWithLayout) { + string shape_string = "f32[123,456]{0,1}"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, + ShapeUtil::ParseShapeString(shape_string)); + Shape expected = ShapeUtil::MakeShapeWithLayout(F32, {123, 456}, {0, 1}); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST(ShapeUtilTest, ParseShapeStringWithExplicitDenseLayout) { + string shape_string = "f32[123,456]dense{0,1}"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, + ShapeUtil::ParseShapeString(shape_string)); + Shape expected = ShapeUtil::MakeShapeWithLayout(F32, {123, 456}, {0, 1}); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST(ShapeUtilTest, ParseShapeStringWithSparseLayout) { + string shape_string = "f32[123,456]sparse{10}"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, + ShapeUtil::ParseShapeString(shape_string)); + Shape expected = ShapeUtil::MakeShapeWithSparseLayout(F32, {123, 456}, 10); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST(ShapeUtilTest, ParseInvalidShapeString) { + string shape_strings[] = { + "f32[123,456]foobar{0,1}", "f32[123,456]sparse{0,1}", "f32[123,456]{foo}", + "f32[123,456]dense{foo}", "f32[123,456]sparse{foo}", + }; + for (const string& shape_string : shape_strings) { + StatusOr result = ShapeUtil::ParseShapeString(shape_string); + ASSERT_FALSE(result.ok()) << "shape: " << shape_string; + } +} + TEST(ShapeUtilTest, CompatibleIdenticalShapes) { Shape shape1 = ShapeUtil::MakeShape(F32, {3, 2}); Shape shape2 = ShapeUtil::MakeShape(F32, {3, 2}); @@ -125,6 +170,18 @@ TEST(ShapeUtilTest, CompatibleNotIdenticalShapes) { EXPECT_TRUE(ShapeUtil::Compatible(shape_1, shape_2)); } +TEST(ShapeUtilTest, CompatibleIgnoringFpPrecision) { + Shape shape1 = ShapeUtil::MakeShape(BF16, {3, 2}); + Shape shape2 = ShapeUtil::MakeShape(F32, {3, 2}); + ASSERT_TRUE(ShapeUtil::CompatibleIgnoringFpPrecision(shape1, shape2)); +} + +TEST(ShapeUtilTest, IncompatibleIgnoringFpPrecision) { + Shape shape1 = ShapeUtil::MakeShape(BF16, {3, 2}); + Shape shape2 = ShapeUtil::MakeShape(F32, {2, 2}); + ASSERT_FALSE(ShapeUtil::CompatibleIgnoringFpPrecision(shape1, shape2)); +} + TEST(ShapeUtilTest, IncompatibleDifferentElementShapes) { Shape shape_1 = ShapeUtil::MakeShape(F32, {3, 2}); Shape shape_2 = ShapeUtil::MakeShape(PRED, {3, 2}); @@ -139,6 +196,14 @@ TEST(ShapeUtilTest, CompatibleTuples) { EXPECT_TRUE(ShapeUtil::Compatible(tuple1, tuple2)); } +TEST(ShapeUtilTest, CompatibleTuplesIgnoringFpPrecision) { + Shape tuple1 = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(BF16, {3, 2}), ShapeUtil::MakeShape(F32, {4, 5})}); + Shape tuple2 = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F64, {3, 2}), ShapeUtil::MakeShape(BF16, {4, 5})}); + EXPECT_TRUE(ShapeUtil::CompatibleIgnoringFpPrecision(tuple1, tuple2)); +} + TEST(ShapeUtilTest, IncompatibleTuplesWithSwappedElements) { Shape tuple1 = ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShape(PRED, {4, 5}), ShapeUtil::MakeShape(F32, {3, 2})}); @@ -148,6 +213,14 @@ TEST(ShapeUtilTest, IncompatibleTuplesWithSwappedElements) { EXPECT_FALSE(ShapeUtil::CompatibleIgnoringElementType(tuple1, tuple2)); } +TEST(ShapeUtilTest, IncompatibleTuplesIgnoringFpPrecision) { + Shape tuple1 = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(BF16, {4, 5}), ShapeUtil::MakeShape(F32, {3, 2})}); + Shape tuple2 = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {3, 2}), ShapeUtil::MakeShape(BF16, {4, 5})}); + EXPECT_FALSE(ShapeUtil::CompatibleIgnoringFpPrecision(tuple1, tuple2)); +} + TEST(ShapeUtilTest, IncompatibleTuplesWithDifferentPrimitiveType) { Shape tuple1 = ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShape(PRED, {4, 5}), ShapeUtil::MakeShape(F32, {3, 2})}); diff --git a/tensorflow/compiler/xla/sparse_index_array.cc b/tensorflow/compiler/xla/sparse_index_array.cc new file mode 100644 index 0000000000000000000000000000000000000000..31844abd89a020c87c403353374a80fb639a3244 --- /dev/null +++ b/tensorflow/compiler/xla/sparse_index_array.cc @@ -0,0 +1,110 @@ +/* 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/compiler/xla/sparse_index_array.h" + +#include "tensorflow/compiler/xla/index_util.h" +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/shape_util.h" + +namespace xla { + +SparseIndexArray::SparseIndexArray() : rank_(0), max_indices_(0) {} + +SparseIndexArray::SparseIndexArray(int64 max_indices, int64 rank, + std::vector indices) + : indices_(std::move(indices)), rank_(rank), max_indices_(max_indices) { + CHECK_GT(rank_, 0); + CHECK_EQ(indices_.size() % rank_, 0) + << "indices_.size(): " << indices_.size() << ", rank_: " << rank_; + CHECK_LT(index_count(), max_indices_); +} + +SparseIndexArray::SparseIndexArray(int64 max_indices, int64 rank, + tensorflow::gtl::ArraySlice indices) + : SparseIndexArray(max_indices, rank, + std::vector(indices.begin(), indices.end())) {} + +SparseIndexArray::SparseIndexArray(int64 max_indices, + const Array2D& indices) + : SparseIndexArray(max_indices, indices.n2(), + std::vector(indices.begin(), indices.end())) {} + +int64 SparseIndexArray::index_count() const { + CHECK_GT(rank_, 0); + CHECK_EQ(indices_.size() % rank_, 0); + return indices_.size() / rank_; +} + +tensorflow::gtl::ArraySlice SparseIndexArray::At( + int64 sparse_element_number) const { + CHECK_GT(rank_, 0); + CHECK_GE(sparse_element_number, 0); + CHECK_LE(rank_ * sparse_element_number + rank_, indices_.size()); + return tensorflow::gtl::ArraySlice( + indices_.data() + rank_ * sparse_element_number, rank_); +} + +tensorflow::gtl::MutableArraySlice SparseIndexArray::At( + int64 sparse_element_number) { + CHECK_GT(rank_, 0); + CHECK_GE(sparse_element_number, 0); + CHECK_LE(rank_ * sparse_element_number + rank_, indices_.size()); + return tensorflow::gtl::MutableArraySlice( + indices_.data() + rank_ * sparse_element_number, rank_); +} + +void SparseIndexArray::Append(tensorflow::gtl::ArraySlice index) { + CHECK_GT(rank_, 0); + CHECK_EQ(index.size(), rank_); + indices_.insert(indices_.end(), index.begin(), index.end()); +} + +void SparseIndexArray::Clear() { indices_.clear(); } + +void SparseIndexArray::Resize(int64 num_indices) { + CHECK_GT(rank_, 0); + indices_.resize(rank_ * num_indices); +} + +bool SparseIndexArray::Validate(const Shape& shape) const { + if (rank_ == 0 || rank_ != ShapeUtil::Rank(shape)) { + return false; + } + int64 num_indices = index_count(); + if (num_indices > LayoutUtil::MaxSparseElements(shape.layout())) { + return false; + } + if (num_indices < 2) { + return true; + } + tensorflow::gtl::ArraySlice last = At(0); + if (!IndexUtil::IndexInBounds(shape, last)) { + return false; + } + for (int64 n = 1; n < num_indices; ++n) { + tensorflow::gtl::ArraySlice next = At(n); + if (!IndexUtil::IndexInBounds(shape, next)) { + return false; + } + if (IndexUtil::CompareIndices(last, next) >= 0) { + return false; + } + last = next; + } + return true; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/sparse_index_array.h b/tensorflow/compiler/xla/sparse_index_array.h new file mode 100644 index 0000000000000000000000000000000000000000..f2ce22d6721ff8da46f741ccedc2a63dea5994c8 --- /dev/null +++ b/tensorflow/compiler/xla/sparse_index_array.h @@ -0,0 +1,176 @@ +/* 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. +==============================================================================*/ + +// Utility class for managing sparse array indices. + +#ifndef TENSORFLOW_COMPILER_XLA_SPARSE_INDEX_ARRAY_H_ +#define TENSORFLOW_COMPILER_XLA_SPARSE_INDEX_ARRAY_H_ + +#include + +#include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/index_util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" + +namespace xla { + +// Encapsulates the array of indices for a sparse array. A SparseIndexArray +// contain indices for up to `max_indices` elements of a sparse array. Each +// sparse index is an array of `rank` int64 value that gives the location of a +// value within a sparse array. Note that the dimensions of the array are not +// checked (except for the rank). To avoid confusion, we refer to the position +// of an index within a SparseIndexArray as a sparse index number. +class SparseIndexArray { + public: + SparseIndexArray(); + SparseIndexArray(const SparseIndexArray&) = default; + SparseIndexArray(SparseIndexArray&&) = default; + SparseIndexArray& operator=(const SparseIndexArray&) = default; + SparseIndexArray& operator=(SparseIndexArray&&) = default; + + // Constructs a SparseIndexArray that can hold up to `max_indices` sparse + // indices, with an initial contents obtained from the given array. The rank + // is taken from the minor dimension of the array. The major dimension of the + // array must not exceed `max_indices`. + SparseIndexArray(int64 max_indices, const Array2D& indices); + + // Like above, but the array is flattened. For example, the following are + // equivalent: + // + // SparseIndexArray(10, 3, + // Array2D{ + // {0, 1, 2}, + // {3, 4, 5}, + // {6, 7, 8}, + // {9, 10, 11}, + // }) + // + // SparseIndexArray(10, 3, + // {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}) + // + SparseIndexArray(int64 max_indices, int64 rank, + std::vector indices = {}); + SparseIndexArray(int64 max_indices, int64 rank, + tensorflow::gtl::ArraySlice indices); + + // Returns the number of elements represented by the indices stored in the + // array. + int64 index_count() const; + + // Returns a slice that refers to the given sparse index number. The argument + // must be in the range [0, element_count()). + tensorflow::gtl::ArraySlice At(int64 sparse_element_number) const; + tensorflow::gtl::MutableArraySlice At(int64 sparse_element_number); + + // Adds the given index at the end of the array. The new size of the + // SparseIndexArray must not exceed `max_indices`. + void Append(tensorflow::gtl::ArraySlice index); + + // Removes all indices from the array. + void Clear(); + + // Resizes the array to contain the given number of sparse indices. The new + // size must be smaller than `max_indices`. If the new size is larger than + // the old size, the value of the new indices is not specified. + void Resize(int64 num_indices); + + // Returns true iff all indices are unique and occur in sorted order, and are + // valid for the given shape. + bool Validate(const Shape& shape) const; + + int64 rank() const { return rank_; } + int64 max_indices() const { return max_indices_; } + + // Returns a pointer to the int64 array that holds the sparse indices. + tensorflow::gtl::MutableArraySlice mutable_data() { return &indices_; } + tensorflow::gtl::ArraySlice data() const { return indices_; } + + // Sorts this sparse index array along with the set of corresponding values. + // The indices and values are sorted in the lexicographic order of the + // indices, from smallest to largest. + // + // For example: + // + // std::vector v{10.0, 11.0, 12.0}; + // SparseIndexArray a(10, 3, + // {{3, 4, 5}, + // {1, 2, 3}, + // {2, 3, 4}}); + // a.SortWithValues(&v); + // // Prints "11.0, 12.0, 10.0": + // std::cout << v[0] << ", " << v[1] << ", " << v[2] << std::endl; + // + template + void SortWithValues(tensorflow::gtl::MutableArraySlice values); + + private: + std::vector indices_; + int64 rank_; + int64 max_indices_; +}; + +template +void SparseIndexArray::SortWithValues( + tensorflow::gtl::MutableArraySlice values) { + int64 num_elements = index_count(); + CHECK_EQ(values.size(), num_elements); + std::vector sort_order; + sort_order.reserve(num_elements); + for (int64 i = 0; i < num_elements; ++i) { + sort_order.push_back(i); + } + auto sort_order_less = [this](int64 lhs, int64 rhs) { + return IndexUtil::CompareIndices(At(lhs), At(rhs)) < 0; + }; + std::sort(sort_order.begin(), sort_order.end(), sort_order_less); + + // Reorder the array elements according to sort_order. Work through the array + // and follow cycles so we can do the reorder in-place. + tensorflow::gtl::InlinedVector saved_index(rank()); + for (int64 i = 0; i < num_elements; ++i) { + // sort_order[i] == -1 indicates the element has already been copied. + if (sort_order[i] < 0) { + continue; + } else if (i == sort_order[i]) { + // The element is already in sorted order. + sort_order[i] = -1; + continue; + } + + std::copy_n(At(i).begin(), rank(), saved_index.begin()); + NativeT saved_value = values[i]; + int64 j = i; + for (;;) { + if (sort_order[j] == i) { + std::copy_n(saved_index.begin(), rank(), At(j).begin()); + values[j] = saved_value; + sort_order[j] = -1; + break; + } + + std::copy_n(At(sort_order[j]).begin(), rank(), At(j).begin()); + values[j] = values[sort_order[j]]; + + int64 k = sort_order[j]; + sort_order[j] = -1; + j = k; + } + } +} + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SPARSE_INDEX_ARRAY_H_ diff --git a/tensorflow/compiler/xla/sparse_index_array_test.cc b/tensorflow/compiler/xla/sparse_index_array_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..7377f88958dcb7daf3d3f4f0e07966fdc9294580 --- /dev/null +++ b/tensorflow/compiler/xla/sparse_index_array_test.cc @@ -0,0 +1,43 @@ +/* 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/compiler/xla/sparse_index_array.h" + +#include + +#include "tensorflow/compiler/xla/test.h" + +namespace xla { +namespace { + +TEST(SparseIndexArrayTest, Sort) { + SparseIndexArray a(10, 3); + a.Append({2, 3, 4}); + a.Append({3, 4, 5}); + a.Append({1, 2, 3}); + a.Append({5, 6, 7}); + a.Append({4, 5, 6}); + a.Append({6, 7, 8}); + std::vector values = { + 12.0, 13.0, 11.0, 15.0, 14.0, 16.0, + }; + a.SortWithValues(&values); + ASSERT_EQ(a.data(), std::vector({1, 2, 3, 2, 3, 4, 3, 4, 5, 4, 5, 6, 5, + 6, 7, 6, 7, 8})); + ASSERT_EQ(values, std::vector({11.0, 12.0, 13.0, 14.0, 15.0, 16.0})); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/statusor_internals.h b/tensorflow/compiler/xla/statusor_internals.h index a2fda5bb3c6f11c20fc45c57885b1ce7523db81d..14636bd144bc0a155fc96c5a350c658fd2dadfe6 100644 --- a/tensorflow/compiler/xla/statusor_internals.h +++ b/tensorflow/compiler/xla/statusor_internals.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ +#ifndef TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ +#define TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ #include "tensorflow/compiler/xla/status.h" #include "tensorflow/core/platform/macros.h" @@ -242,4 +242,4 @@ struct TraitsBase { } // namespace internal_statusor } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ +#endif // TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD index 45689976b0da2de86b0b037127882180cac06d0e..8339d08ef4d7455f9739b80074ab0405a404e8e8 100644 --- a/tensorflow/compiler/xla/tests/BUILD +++ b/tensorflow/compiler/xla/tests/BUILD @@ -345,6 +345,24 @@ xla_test( ], ) +xla_test( + name = "xla_hlo_profile_test", + srcs = ["xla_hlo_profile_test.cc"], + deps = [ + "//tensorflow/compiler/xla:array2d", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/service:platform_util", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/core:lib", + "//tensorflow/core:regexp_internal", + "//tensorflow/core:test", + ], +) + xla_test( name = "axpy_simple_test", srcs = ["axpy_simple_test.cc"], @@ -557,9 +575,31 @@ xla_test( ], ) +xla_test( + name = "exhaustive_f32_elementwise_op_test", + srcs = ["exhaustive_f32_elementwise_op_test.cc"], + backends = [ + "cpu", + "gpu", + ], + shard_count = 48, + tags = [ + "enormous", + "manual", + "notap", + ], + deps = [ + ":client_library_test_base", + ":literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", + ], +) + xla_test( name = "reduce_precision_test", srcs = ["reduce_precision_test.cc"], + tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal_util", @@ -582,6 +622,9 @@ xla_test( xla_test( name = "dot_operation_test", srcs = ["dot_operation_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -604,6 +647,9 @@ xla_test( xla_test( name = "dot_operation_runtime_test", srcs = ["dot_operation_test.cc"], + tags = [ + "enable_for_xla_interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -798,9 +844,6 @@ xla_test( xla_test( name = "bfloat16_test", srcs = ["bfloat16_test.cc"], - blacklisted_backends = [ - "gpu", - ], shard_count = 40, deps = [ ":test_utils", @@ -830,6 +873,31 @@ xla_test( ], ) +xla_test( + name = "half_test", + srcs = ["half_test.cc"], + backends = [ + # TODO(b/72509305): Flaky (fails with SEGV) as of 2018-01-25 + # "cpu", + "gpu", + ], + deps = [ + ":test_utils", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla/client:computation", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + xla_test( name = "slice_test", srcs = ["slice_test.cc"], @@ -984,7 +1052,10 @@ xla_test( name = "reduce_window_test", timeout = "long", srcs = [], - tags = ["optonly"], + tags = [ + "enable_for_xla_interpreter", + "optonly", + ], xla_test_library_deps = [":reduce_window_test_library"], deps = [], ) @@ -993,6 +1064,10 @@ xla_test( name = "select_and_scatter_test", timeout = "long", srcs = ["select_and_scatter_test.cc"], + tags = [ + "enable_for_xla_interpreter", + "optonly", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal_util", @@ -1031,6 +1106,19 @@ xla_test( ], ) +xla_test( + name = "reduce_hlo_test", + srcs = ["reduce_hlo_test.cc"], + deps = [ + ":client_library_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + xla_test( name = "call_test", srcs = ["call_test.cc"], @@ -1388,6 +1476,31 @@ xla_test( ], ) +xla_test( + name = "execution_profile_test", + srcs = ["execution_profile_test.cc"], + deps = [ + ":client_library_test_base", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:global_data", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) + +xla_test( + name = "execution_profile_test_with_xla_hlo_profile", + srcs = ["execution_profile_test.cc"], + args = ["--xla_hlo_profile"], + deps = [ + ":client_library_test_base", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:global_data", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) + xla_test( name = "replay_test", srcs = ["replay_test.cc"], @@ -1480,6 +1593,7 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_runner", "//tensorflow/compiler/xla/service:platform_util", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:hlo_test_base", @@ -1506,6 +1620,7 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_runner", "//tensorflow/compiler/xla/service:platform_util", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:hlo_test_base", diff --git a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc index c6e8b24d1211743d07878d388522feacf9c0e7f1..7e9005001db34d403ea923eb9c152d114bf32803 100644 --- a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc +++ b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc @@ -1879,20 +1879,73 @@ XLA_TEST_F(ArrayElementwiseOpTest, ClampF32ScalarVector) { auto min_scalar = builder.ConstantR0(0.0f); auto min_vector = builder.ConstantR1({1.0f, -6.5f, 1.0f, 2.25f, 0.0f}); auto arg_vector = builder.ConstantR1({2.0f, 10.0f, -5.0f, 1.0f, 4.0f}); - auto arg_scalar = builder.ConstantR1({2.0f, 10.0f, -5.0f, 1.0f, 4.0f}); auto max_scalar = builder.ConstantR0(3.0f); auto max_vector = builder.ConstantR1({3.0f, 0.5f, 25.5f, 5.0f, 123.0}); // Perform clamp with broadcasted scalar and vector. auto clamp = builder.Add( builder.Add(builder.Clamp(min_vector, arg_vector, max_scalar), builder.Clamp(min_scalar, arg_vector, max_vector)), - builder.Add(builder.Clamp(min_vector, arg_scalar, max_vector), - builder.Clamp(min_scalar, arg_scalar, max_vector))); + builder.Add(builder.Clamp(min_vector, arg_vector, max_vector), + builder.Clamp(min_scalar, arg_vector, max_scalar))); - ComputeAndCompareR1(&builder, {8.0f, 4.5f, 2.0f, 6.5f, 15.0f}, {}, + ComputeAndCompareR1(&builder, {8.0f, 7.0f, 2.0f, 6.5f, 14.0f}, {}, error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, ClampS32Vector) { + ComputationBuilder builder(client_, TestName()); + auto min_vector = builder.ConstantR1({1, -6, 1, 2, 0, -5}); + auto arg_vector = builder.ConstantR1({2, 10, -5, 1, 4, 10}); + auto max_vector = builder.ConstantR1({3, 0, 25, 5, 123, -1}); + auto clamp = builder.Clamp(min_vector, arg_vector, max_vector); + + ComputeAndCompareR1(&builder, {2, 0, 1, 2, 4, -1}, {}); +} + +XLA_TEST_F(ArrayElementwiseOpTest, ClampS32ScalarVector) { + ComputationBuilder builder(client_, TestName()); + auto min_scalar = builder.ConstantR0(0); + auto min_vector = builder.ConstantR1({1, -6, 1, 2, 0}); + auto arg_vector = builder.ConstantR1({2, 10, -5, 1, 4}); + auto max_scalar = builder.ConstantR0(3); + auto max_vector = builder.ConstantR1({3, 1, 25, 5, 123}); + // Perform clamp with broadcasted scalar and vector. + auto clamp = builder.Add( + builder.Add(builder.Clamp(min_vector, arg_vector, max_scalar), + builder.Clamp(min_scalar, arg_vector, max_vector)), + builder.Add(builder.Clamp(min_vector, arg_vector, max_vector), + builder.Clamp(min_scalar, arg_vector, max_scalar))); + + ComputeAndCompareR1(&builder, {8, 8, 2, 6, 14}, {}); +} + +XLA_TEST_F(ArrayElementwiseOpTest, ClampU32Vector) { + ComputationBuilder builder(client_, TestName()); + auto min_vector = builder.ConstantR1({1, 2, 1, 2, 0, ~0u - 4}); + auto arg_vector = builder.ConstantR1({2, 10, 5, 1, 4, 10}); + auto max_vector = builder.ConstantR1({3, 5, 25, 5, 123, ~0u}); + auto clamp = builder.Clamp(min_vector, arg_vector, max_vector); + + ComputeAndCompareR1(&builder, {2, 5, 5, 2, 4, ~0u - 4}, {}); +} + +XLA_TEST_F(ArrayElementwiseOpTest, ClampU32ScalarVector) { + ComputationBuilder builder(client_, TestName()); + auto min_scalar = builder.ConstantR0(0); + auto min_vector = builder.ConstantR1({1, 0, 1, 2, 0}); + auto arg_vector = builder.ConstantR1({2, 10, 0, 1, 4}); + auto max_scalar = builder.ConstantR0(3); + auto max_vector = builder.ConstantR1({3, 1, 25, 5, 123}); + // Perform clamp with broadcasted scalar and vector. + auto clamp = builder.Add( + builder.Add(builder.Clamp(min_vector, arg_vector, max_scalar), + builder.Clamp(min_scalar, arg_vector, max_vector)), + builder.Add(builder.Clamp(min_vector, arg_vector, max_vector), + builder.Clamp(min_scalar, arg_vector, max_scalar))); + + ComputeAndCompareR1(&builder, {8, 8, 2, 6, 14}, {}); +} + XLA_TEST_F(ArrayElementwiseOpTest, AddTwoParametersF32s) { ComputationBuilder builder(client_, TestName()); @@ -1971,6 +2024,18 @@ XLA_TEST_F(ArrayElementwiseOpTest, SinF32s) { error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, Atan2F32s) { + ComputationBuilder builder(client_, TestName()); + auto a = builder.ConstantR1({0.0f, 5.0f, 0.0f, -3.0f, 2.0f, -8.0f}); + auto b = builder.ConstantR1({6.0f, 0.0f, -4.0f, 0.0f, 2.0f, 8.0f}); + auto atan = builder.Atan2(a, b); + + ComputeAndCompareR1( + &builder, + {0.0f, 1.57079633f, 3.14159265f, -1.57079633f, 0.78539816f, -0.78539816f}, + {}, error_spec_); +} + XLA_TEST_F(ArrayElementwiseOpTest, TanhF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-2.5f, 3.14f, 2.25f}); @@ -1983,47 +2048,117 @@ XLA_TEST_F(ArrayElementwiseOpTest, TanhF32s) { XLA_TEST_F(ArrayElementwiseOpTest, TanhF32sVector) { // This is like the test ArrayElementwiseOpTest.TanhF32s above, except that // the input tensor is large enough to exercise the vectorized tanh - // implementation. - ComputationBuilder builder(client_, TestName()); - auto input_literal = Literal::CreateR2( - {{1.02, -0.32, 0.85, 0.90, 1.23, -0.91, -0.49, 0.80}, - {-0.67, 0.16, -0.07, 0.39, -0.41, 0.04, 1.36, 1.25}, - {0.41, 0.65, -1.08, 0.32, -1.45, -0.77, -1.09, 0.91}, - {-1.03, -0.30, -1.11, -1.17, 1.50, -0.85, 0.04, 1.02}, - {0.34, -0.61, 0.41, 0.07, -0.02, 1.42, -0.62, 0.81}, - {0.08, 0.81, -0.30, 1.17, -0.65, -0.44, 0.92, 1.26}, - {-1.29, 1.35, 0.08, -1.24, -0.92, 0.49, 1.17, -0.45}, - {-1.31, -1.44, -0.13, -1.31, -0.79, 1.41, 1.21, 1.05}}); - auto input_data = - client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + // implementation on XLA CPU. + ComputationBuilder builder(client_, TestName()); + auto input_literal = Literal::CreateR1( + {1.02, -0.32, 0.85, 0.90, 1.23, -0.91, -0.49, 0.80, -0.67, 0.16, + -0.07, 0.39, -0.41, 0.04, 1.36, 1.25, 0.41, 0.65, -1.08, 0.32, + -1.45, -0.77, -1.09, 0.91, -1.03, -0.30, -1.11, -1.17, 1.50, -0.85, + 0.04, 1.02, 0.34, -0.61, 0.41, 0.07, -0.02, 1.42, -0.62, 0.81, + 0.08, 0.81, -0.30, 1.17, -0.65, -0.44, 0.92, 1.26, -1.29, 1.35, + 0.08, -1.24, -0.92, 0.49, 1.17, -0.45, -1.31, -1.44, -0.13, -1.31, + -0.79, 1.41, 1.21, 1.05}); + TF_ASSERT_OK_AND_ASSIGN(auto input_data, + client_->TransferToServer(*input_literal)); auto input = builder.Parameter(0, input_literal->shape(), "input"); builder.Tanh(input); - ComputeAndCompareR2( + ComputeAndCompareR1( &builder, - {{0.77009583, -0.30665702, 0.69070244, 0.71401149, 0.84400684, - -0.71985596, -0.45764771, 0.66664988}, - {-0.58278900, 0.16050975, -0.06770509, 0.36843640, -0.38476998, - 0.04018109, 0.87562293, 0.84788644}, - {0.38603750, 0.57294142, -0.79140943, 0.31032649, -0.89590985, - -0.64770776, -0.79625875, 0.72234446}, - {-0.77389336, -0.28871772, -0.80428445, -0.82541436, 0.90456349, - -0.68856895, 0.03877772, 0.76877952}, - {0.32561871, -0.54546672, 0.39072621, 0.07273290, -0.01924866, - 0.88924897, -0.55283129, 0.67183107}, - {0.08006320, 0.66944766, -0.29068485, 0.82573754, -0.57170743, - -0.41581789, 0.72739530, 0.85025692}, - {-0.85931867, 0.87357593, 0.07782833, -0.84597743, -0.72748238, - 0.45396307, 0.82449573, -0.42462519}, - {-0.86363792, -0.89368379, -0.12621804, -0.86445558, -0.65565848, - 0.88789743, 0.83566397, 0.78287679}}, + {0.77009583, -0.30665702, 0.69070244, 0.71401149, 0.84400684, + -0.71985596, -0.45764771, 0.66664988, -0.58278900, 0.16050975, + -0.06770509, 0.36843640, -0.38476998, 0.04018109, 0.87562293, + 0.84788644, 0.38603750, 0.57294142, -0.79140943, 0.31032649, + -0.89590985, -0.64770776, -0.79625875, 0.72234446, -0.77389336, + -0.28871772, -0.80428445, -0.82541436, 0.90456349, -0.68856895, + 0.03877772, 0.76877952, 0.32561871, -0.54546672, 0.39072621, + 0.07273290, -0.01924866, 0.88924897, -0.55283129, 0.67183107, + 0.08006320, 0.66944766, -0.29068485, 0.82573754, -0.57170743, + -0.41581789, 0.72739530, 0.85025692, -0.85931867, 0.87357593, + 0.07782833, -0.84597743, -0.72748238, 0.45396307, 0.82449573, + -0.42462519, -0.86363792, -0.89368379, -0.12621804, -0.86445558, + -0.65565848, 0.88789743, 0.83566397, 0.78287679}, {input_data.get()}, // The error spec is unusually high here to account for the fact that we // use a rational interpolant to approximate tanh. ErrorSpec(0.004, 0.004)); } +XLA_TEST_F(ArrayElementwiseOpTest, ExpF32sVector) { + // The input tensor is large enough to exercise the vectorized exp + // implementation on XLA CPU. + ComputationBuilder builder(client_, TestName()); + + // Just to help make sense of the scales here -- exp(89) saturates float32 and + // exp(-10) is smaller than our error spec. + std::unique_ptr input_literal = Literal::CreateR1( + {1.02, -0.32, 0.85, 0.9, 1.23, -0.91, -0.49, 0.8, -1.31, + -1.44, -0.13, -1.31, -0.79, 1.41, 1.21, 1.05, -195.6, -194.5, + -193.4, -192.3, -191.2, -190.1, -189.0, -187.9, -19.6, -18.5, -17.4, + -16.3, -15.2, -14.1, -13.0, -11.9, -10.8, -9.7, -8.6, -7.5, + -6.4, -5.3, -4.2, -3.1, -2.0, -0.9, 0.2, 1.3, 2.4, + 3.5, 4.6, 5.7, 6.8, 7.9, 9.0, 10.1, 11.2, 12.3, + 13.4, 14.5, 15.6, 16.7, 17.8, 18.9, 20.0, 21.1, 22.2, + 23.3, 24.4, 25.5, 26.6, 27.7, 28.8, 29.9, 31.0, 32.1, + 68.4, 69.5, 70.6, 71.7, 72.8, 73.9, 75.0, 76.1, 77.2, + 78.3, 79.4, 80.5, 81.6, 82.7, 83.8, 84.9, 85.2, 86.3, + 86.4, 86.5, 87.6, 87.7, 87.8, 87.9}); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, + client_->TransferToServer(*input_literal)); + + auto input = builder.Parameter(0, input_literal->shape(), "input"); + builder.Exp(input); + + std::vector expected_result; + int64 input_size = input_literal->shape().dimensions(0); + expected_result.reserve(input_size); + for (int64 i = 0; i < input_size; i++) { + expected_result.push_back(std::exp(input_literal->Get({i}))); + } + + ComputeAndCompareR1(&builder, expected_result, {input_data.get()}, + error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, LogF32sVector) { + // The input tensor is large enough to exercise the vectorized exp + // implementation on XLA CPU. + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr input_literal = Literal::CreateR1( + {-1.29, -1.41, -1.25, -13.5, -11.7, -17.9, -198, + -167, 1.29, 1.41, 1.25, 13.5, 11.7, 17.9, + 198, 167, 1.27e+03, 1.33e+03, 1.74e+03, 1.6e+04, 1.84e+04, + 1.74e+04, 1.89e+05, 1.9e+05, 1.93e+06, 1.98e+06, 1.65e+06, 1.97e+07, + 1.66e+07, 1e+07, 1.98e+08, 1.96e+08, 1.64e+09, 1.58e+09, 1.64e+09, + 1.44e+10, 1.5e+10, 1.99e+10, 1.17e+11, 1.08e+11, 1.08e+12, 1.38e+12, + 1.4e+12, 1.03e+13, 1.6e+13, 1.99e+13, 1.26e+14, 1.51e+14, 1.33e+15, + 1.41e+15, 1.63e+15, 1.39e+16, 1.21e+16, 1.27e+16, 1.28e+17, 1.62e+17, + 2e+18, 1.96e+18, 1.81e+18, 1.99e+19, 1.86e+19, 1.61e+19, 1.71e+20, + 1.47e+20, 1.83e+21, 1.33e+21, 1.3e+21, 1.35e+22, 1.84e+22, 1.02e+22, + 1.81e+23, 1.02e+23, 1.89e+24, 1.49e+24, 1.08e+24, 1.95e+25, 1.1e+25, + 1.62e+25, 1.2e+26, 1.41e+26, 1.93e+27, 1.66e+27, 1.62e+27, 1.05e+28, + 1.5e+28, 1.79e+28, 1.36e+29, 1.95e+29, 1.5e+30, 1.81e+30, 1.34e+30, + 1.7e+31, 1.44e+31, 1.1e+31, 1.4e+32, 1.67e+32, 1.96e+33, 1.11e+33, + 1.19e+33, 1.61e+34, 1.05e+34, 1.88e+34, 1.67e+35, 1.7e+35}); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, + client_->TransferToServer(*input_literal)); + + auto input = builder.Parameter(0, input_literal->shape(), "input"); + builder.Log(input); + + std::vector expected_result; + int64 input_size = input_literal->shape().dimensions(0); + expected_result.reserve(input_size); + for (int64 i = 0; i < input_size; i++) { + expected_result.push_back(std::log(input_literal->Get({i}))); + } + + ComputeAndCompareR1(&builder, expected_result, {input_data.get()}, + error_spec_); +} + XLA_TEST_F(ArrayElementwiseOpTest, AddChainFoldLeft) { // a ------ (add) --------- (add) // / / @@ -2520,9 +2655,8 @@ XLA_TEST_F(ArrayElementwiseOpTest, R4_16x16x2x2_Plus_R1_16) { std::iota(r1.begin(), r1.end(), 1.0); ComputationBuilder builder(client_, TestName()); - std::unique_ptr a_literal = Literal::CreateR4FromArray4D(r4); - *a_literal->mutable_shape()->mutable_layout() = - LayoutUtil::MakeLayout({0, 1, 2, 3}); + std::unique_ptr a_literal = Literal::CreateR4FromArray4DWithLayout( + r4, LayoutUtil::MakeLayout({0, 1, 2, 3})); auto a = builder.ConstantLiteral(*a_literal); auto b = builder.ConstantR1(r1); builder.Add(a, b, {1}); diff --git a/tensorflow/compiler/xla/tests/axpy_simple_test.cc b/tensorflow/compiler/xla/tests/axpy_simple_test.cc index 627a9c3e7d9f6eb8d360228362ea5adf12c6c798..3f6fd7c65d3360a622dbf754833009fb20410535 100644 --- a/tensorflow/compiler/xla/tests/axpy_simple_test.cc +++ b/tensorflow/compiler/xla/tests/axpy_simple_test.cc @@ -62,6 +62,10 @@ TEST_F(AxpySimpleTest, AxpyTenValues) { auto ax = builder.Mul(alpha, x); auto axpy = builder.Add(ax, y); + TF_ASSERT_OK_AND_ASSIGN(ProgramShape shape, builder.GetProgramShape()); + + EXPECT_EQ("() -> f32[10]", ShapeUtil::HumanString(shape)); + std::vector expected = { 1.85840735, -1.85840735, 2.28318531, -2.28318531, -6.42477796, 6.42477796, 10.56637061, -10.56637061, -14.70796327, 14.70796327}; diff --git a/tensorflow/compiler/xla/tests/batch_normalization_test.cc b/tensorflow/compiler/xla/tests/batch_normalization_test.cc index dbbecc58fb65fbf9fcb9229a00fe78229d8fee54..28ab9654997728fbafd6610af840e721e72cce5a 100644 --- a/tensorflow/compiler/xla/tests/batch_normalization_test.cc +++ b/tensorflow/compiler/xla/tests/batch_normalization_test.cc @@ -62,7 +62,7 @@ class BatchNormalizationTest {5.0f, 4.4f}, // p2 }); input_array_.FillWithPZ(pz); - input_literal_ = *Literal::CreateR4FromArray4D(input_array_); + input_literal_ = std::move(*Literal::CreateR4FromArray4D(input_array_)); CHECK_EQ(kSamples, input_array_.planes()); CHECK_EQ(kZ, input_array_.depth()); CHECK_EQ(kY, input_array_.height()); @@ -231,14 +231,14 @@ XLA_TEST_P(BatchNormalizationTest, BasicTraining) { auto tuple = builder.BatchNormTraining(operand, scale, offset, /*epsilon=*/0.001, kFeatureIndex); - auto expected = *Literal::MakeTuple( + auto expected = Literal::MakeTuple( {Literal::CreateR4({{{{-1.6f, -2.0f}}, {{0.1f, 0.6f}}}, {{{1.9f, 3.3f}}, {{3.7f, 6.0f}}}}) .get(), Literal::CreateR1({4, 5}).get(), Literal::CreateR1({5, 5}).get()}); - ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.1)); + ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.1)); } XLA_TEST_P(BatchNormalizationTest, BasicTrainingOnSublane) { @@ -255,14 +255,14 @@ XLA_TEST_P(BatchNormalizationTest, BasicTrainingOnSublane) { auto tuple = builder.BatchNormTraining(operand, scale, offset, /*epsilon=*/0.001, kFeatureIndex); - auto expected = *Literal::MakeTuple( + auto expected = Literal::MakeTuple( {Literal::CreateR4({{{{-1.6f}, {-2.0f}}, {{0.1f}, {0.6f}}}, {{{1.9f}, {3.3f}}, {{3.7f}, {6.0f}}}}) .get(), Literal::CreateR1({4, 5}).get(), Literal::CreateR1({5, 5}).get()}); - ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.1)); + ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.1)); } XLA_TEST_P(BatchNormalizationTest, TrainingWithFeatureOnLowDimension) { @@ -286,13 +286,13 @@ XLA_TEST_P(BatchNormalizationTest, TrainingWithFeatureOnLowDimension) { auto tuple = builder.BatchNormTraining(h0, h1, h2, /*epsilon=*/1, kFeatureIndex); - auto expected = *Literal::MakeTuple( + auto expected = Literal::MakeTuple( {Literal::CreateR3FromArray3D(Array3D(260, 2, 2, 1.0f)) .get(), Literal::CreateR1(std::vector(260, 1.0f)).get(), Literal::CreateR1(std::vector(260, 0.0f)).get()}); - ComputeAndCompareTuple(&builder, expected, + ComputeAndCompareTuple(&builder, *expected, {operand.get(), scale.get(), offset.get()}, ErrorSpec(0.1)); } @@ -319,13 +319,13 @@ XLA_TEST_P(BatchNormalizationTest, LargeEpsilonTest) { auto tuple = builder.BatchNormTraining(h0, h1, h2, /*epsilon=*/-100, kFeatureIndex); - auto expected = *Literal::MakeTuple( + auto expected = Literal::MakeTuple( {Literal::CreateR3FromArray3D({{{-3.0f}, {-1.0f}, {1.0f}, {3.0f}}}) .get(), Literal::CreateR1(std::vector(1, 15.0f)).get(), Literal::CreateR1(std::vector(1, 125.0f)).get()}); - ComputeAndCompareTuple(&builder, expected, + ComputeAndCompareTuple(&builder, *expected, {operand.get(), scale.get(), offset.get()}, ErrorSpec(0.1)); } @@ -349,14 +349,14 @@ XLA_TEST_P(BatchNormalizationTest, BatchNormGradBasic) { builder.BatchNormGrad(operand, scale, mean, var, grad_output, /*epsilon=*/0.0, kFeatureIndex); - auto expected = *Literal::MakeTuple( + auto expected = Literal::MakeTuple( {Literal::CreateR4({{{{-3.f}, {-3.f}}, {{-1.f}, {-1.f}}}, {{{1.f}, {1.f}}, {{3.f}, {3.f}}}}) .get(), Literal::CreateR1({0, 0}).get(), Literal::CreateR1({16, 20}).get()}); - ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.1)); + ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.1)); } struct BatchNormTestParam { @@ -513,9 +513,9 @@ XLA_TEST_P(BatchNormTestManySizes, RandomizedTrainingTests) { auto offset_activations = builder.Parameter(2, offset_literal->shape(), "scale"); - auto expected = *Literal::MakeTuple({expected_normalized.get(), - Literal::CreateR1(mean).get(), - Literal::CreateR1(var).get()}); + auto expected = Literal::MakeTuple({expected_normalized.get(), + Literal::CreateR1(mean).get(), + Literal::CreateR1(var).get()}); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -532,7 +532,7 @@ XLA_TEST_P(BatchNormTestManySizes, RandomizedTrainingTests) { // testcase. execution_options_.mutable_debug_options()->clear_xla_disable_hlo_passes(); ComputeAndCompareTuple( - &builder, expected, + &builder, *expected, {input_data.get(), scale_data.get(), offset_data.get()}, ErrorSpec(0.01, 1)); } @@ -819,16 +819,16 @@ XLA_TEST_P(BatchNormTestManySizes, RandomizedGradTests) { grad_output_parameter, epsilon, feature_index); auto expected = - *Literal::MakeTuple({expected_grad_activation.get(), - Literal::CreateR1(grad_scale).get(), - Literal::CreateR1(grad_offset).get()}); + Literal::MakeTuple({expected_grad_activation.get(), + Literal::CreateR1(grad_scale).get(), + Literal::CreateR1(grad_offset).get()}); // Run all HLO passes during this test. In particular, ClientLibraryTestBase // disables constant folding, but we want it enabled for our zero-sized tensor // testcase. execution_options_.mutable_debug_options()->clear_xla_disable_hlo_passes(); - ComputeAndCompareTuple(&builder, expected, + ComputeAndCompareTuple(&builder, *expected, {input_data.get(), scale_data.get(), mean_data.get(), var_data.get(), grad_output_data.get()}, ErrorSpec(0.01, 1)); diff --git a/tensorflow/compiler/xla/tests/bfloat16_test.cc b/tensorflow/compiler/xla/tests/bfloat16_test.cc index ac3f3f4c9ddb03d003a44f5abd7a2e26c42f490d..b853dfaa15d7ff2e21048a5a6a486d22c5a05416 100644 --- a/tensorflow/compiler/xla/tests/bfloat16_test.cc +++ b/tensorflow/compiler/xla/tests/bfloat16_test.cc @@ -97,10 +97,11 @@ XLA_TEST_F(Bfloat16Test, BatchNormTraining) { auto tuple = builder.BatchNormTraining(operand, scale, offset, /*epsilon=*/0.001, kFeatureIndex); - auto expected = *Literal::MakeTuple( + auto expected = Literal::MakeTuple( {Literal::CreateR4( - {{{{static_cast(-1.7f)}, {static_cast(-2.04f)}}, - {{static_cast(0.105f)}, {static_cast(0.65f)}}}, + {{{{static_cast(-1.6875f)}, + {static_cast(-2.04f)}}, + {{static_cast(0.105f)}, {static_cast(0.66f)}}}, {{{static_cast(1.89f)}, {static_cast(3.35f)}}, {{static_cast(3.7f)}, {static_cast(6.04f)}}}}) .get(), @@ -111,7 +112,7 @@ XLA_TEST_F(Bfloat16Test, BatchNormTraining) { {static_cast(5), static_cast(5)}) .get()}); - ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.01)); + ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.01)); } XLA_TEST_F(Bfloat16Test, BatchNormGrad) { @@ -139,7 +140,7 @@ XLA_TEST_F(Bfloat16Test, BatchNormGrad) { builder.BatchNormGrad(operand, scale, mean, var, grad_output, /*epsilon=*/0.0, kFeatureIndex); - auto expected = *Literal::MakeTuple( + auto expected = Literal::MakeTuple( {Literal::CreateR4( {{{{static_cast(-3.f)}, {static_cast(-3.f)}}, {{static_cast(-1.f)}, {static_cast(-1.f)}}}, @@ -153,7 +154,7 @@ XLA_TEST_F(Bfloat16Test, BatchNormGrad) { {static_cast(16), static_cast(20)}) .get()}); - ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.01)); + ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.01)); } } // namespace diff --git a/tensorflow/compiler/xla/tests/broadcast_test.cc b/tensorflow/compiler/xla/tests/broadcast_test.cc index 0294628a127c9d506e6387d0b80f3da583c5a174..6ebbf7191833ef85ee4a48cc96c0a3be38c71228 100644 --- a/tensorflow/compiler/xla/tests/broadcast_test.cc +++ b/tensorflow/compiler/xla/tests/broadcast_test.cc @@ -87,11 +87,11 @@ XLA_TEST_F(BroadcastTest, BroadcastVectorTo2D) { LiteralTestUtil::ExpectNear( *Literal::CreateR2({{1.0, 1.0}, {2.0, 2.0}, {3.0, 3.0}}), - result->tuple_literals(0), error_spec_); + LiteralView::Create(*result, {0}), error_spec_); LiteralTestUtil::ExpectNear( *Literal::CreateR2({{1.0, 2.0, 3.0}, {1.0, 2.0, 3.0}}), - result->tuple_literals(1), error_spec_); + LiteralView::Create(*result, {1}), error_spec_); } XLA_TEST_F(BroadcastTest, Broadcast2DTo2D) { diff --git a/tensorflow/compiler/xla/tests/check_execution_arity_test.cc b/tensorflow/compiler/xla/tests/check_execution_arity_test.cc index 659660d91e519b428d28ced8591d05b4e4d45f53..f594cc10ac6496f710d03f0b0b134e6dd3b6d38f 100644 --- a/tensorflow/compiler/xla/tests/check_execution_arity_test.cc +++ b/tensorflow/compiler/xla/tests/check_execution_arity_test.cc @@ -104,7 +104,8 @@ XLA_TEST_F(CheckExecutionArityTest, CheckArgumentShapes) { ASSERT_FALSE(status.ok()); ASSERT_EQ(status.status().code(), tensorflow::error::INVALID_ARGUMENT); ASSERT_THAT(status.status().error_message(), - ContainsRegex("expects parameter 0")); + ContainsRegex( + "Argument does not match shape of computation parameter 0")); // Shape mismatch in parameter 1 (rank) status = client_->Execute(computation, {f32_data.get(), f32_data.get()}, @@ -112,7 +113,8 @@ XLA_TEST_F(CheckExecutionArityTest, CheckArgumentShapes) { ASSERT_FALSE(status.ok()); ASSERT_EQ(status.status().code(), tensorflow::error::INVALID_ARGUMENT); ASSERT_THAT(status.status().error_message(), - ContainsRegex("expects parameter 1")); + ContainsRegex( + "Argument does not match shape of computation parameter 1")); // Shape mismatch in parameter 1 (element type) status = client_->Execute(computation, {f32_data.get(), u8_4_data.get()}, @@ -120,7 +122,8 @@ XLA_TEST_F(CheckExecutionArityTest, CheckArgumentShapes) { ASSERT_FALSE(status.ok()); ASSERT_EQ(status.status().code(), tensorflow::error::INVALID_ARGUMENT); ASSERT_THAT(status.status().error_message(), - ContainsRegex("expects parameter 1")); + ContainsRegex( + "Argument does not match shape of computation parameter 1")); } } // namespace diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.cc b/tensorflow/compiler/xla/tests/client_library_test_base.cc index 50bf185936808fbd9c49f7fbd5ab0c0b4a76504b..a677986cd926cc0054d8f36abc98ccac33dc043d 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.cc +++ b/tensorflow/compiler/xla/tests/client_library_test_base.cc @@ -251,8 +251,17 @@ ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts( tensorflow::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus( ComputationBuilder* builder, const Literal& expected, - tensorflow::gtl::ArraySlice arguments, + tensorflow::gtl::ArraySlice arguments_passed_in, const Shape* shape_with_layout) { + std::vector arguments(arguments_passed_in.begin(), + arguments_passed_in.end()); + if (!arguments_.empty()) { + CHECK(arguments.empty()); + for (const auto& argument : arguments_) { + arguments.push_back(argument.get()); + } + } + TF_ASSIGN_OR_RETURN(auto computation, builder->Build()); if (ShapeUtil::ElementIsFloating(expected.shape()) || ShapeUtil::ElementIsComplex(expected.shape())) { @@ -300,8 +309,17 @@ tensorflow::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus( tensorflow::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus( ComputationBuilder* builder, const Literal& expected, - tensorflow::gtl::ArraySlice arguments, ErrorSpec error, - const Shape* shape_with_layout) { + tensorflow::gtl::ArraySlice arguments_passed_in, + ErrorSpec error, const Shape* shape_with_layout) { + std::vector arguments(arguments_passed_in.begin(), + arguments_passed_in.end()); + if (!arguments_.empty()) { + CHECK(arguments.empty()); + for (const auto& argument : arguments_) { + arguments.push_back(argument.get()); + } + } + TF_RET_CHECK(ShapeUtil::ElementIsFloating(expected.shape()) || ShapeUtil::ElementIsComplex(expected.shape())); TF_ASSIGN_OR_RETURN(auto computation, builder->Build()); @@ -357,7 +375,7 @@ void ClientLibraryTestBase::ComputeAndCompareR1U8( VLOG(1) << "expected: " << expected_literal->ToString(); VLOG(1) << "actual: " << actual->ToString(); - EXPECT_EQ(expected, actual->u8s_string()); + EXPECT_EQ(expected, actual->GetR1U8AsString()); } void ClientLibraryTestBase::ComputeAndCompareTuple( @@ -369,7 +387,7 @@ void ClientLibraryTestBase::ComputeAndCompareTuple( return; } auto actual = actual_status.ConsumeValueOrDie(); - LiteralTestUtil::ExpectEqualTuple(expected, *actual); + LiteralTestUtil::ExpectEqual(expected, *actual); } void ClientLibraryTestBase::ComputeAndCompareTuple( @@ -381,7 +399,7 @@ void ClientLibraryTestBase::ComputeAndCompareTuple( return; } auto actual = actual_status.ConsumeValueOrDie(); - LiteralTestUtil::ExpectNearTuple(expected, *actual, error); + LiteralTestUtil::ExpectNear(expected, *actual, error); } void ClientLibraryTestBase::ComputeAndCompare( @@ -508,6 +526,15 @@ std::unique_ptr ClientLibraryTestBase::CreateParameterAndTransferLiteral( int64 parameter_number, const Literal& literal, const string& name, ComputationBuilder* builder, ComputationDataHandle* data_handle) { + return CreateParameterAndTransferLiteral(parameter_number, literal, name, + nullptr, builder, data_handle); +} + +std::unique_ptr +ClientLibraryTestBase::CreateParameterAndTransferLiteral( + int64 parameter_number, const Literal& literal, const string& name, + const DeviceHandle* device_handle, ComputationBuilder* builder, + ComputationDataHandle* data_handle) { const Literal* param_literal = &literal; std::unique_ptr converted_literal; if (use_bfloat16_) { @@ -515,12 +542,21 @@ ClientLibraryTestBase::CreateParameterAndTransferLiteral( param_literal = converted_literal.get(); } std::unique_ptr data = - client_->TransferToServer(*param_literal).ConsumeValueOrDie(); + client_->TransferToServer(*param_literal, device_handle) + .ConsumeValueOrDie(); *data_handle = builder->Parameter(parameter_number, param_literal->shape(), name); return data; } +ComputationDataHandle ClientLibraryTestBase::AddParam( + const Literal& argument, ComputationBuilder* builder) { + ComputationDataHandle data_handle; + arguments_.push_back(CreateParameterAndTransferLiteral( + arguments_.size(), argument, "", builder, &data_handle)); + return data_handle; +} + ComputationDataHandle ClientLibraryTestBase::CreateConstantFromLiteral( const Literal& literal, ComputationBuilder* builder) { return builder->ConstantLiteral( diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.h b/tensorflow/compiler/xla/tests/client_library_test_base.h index 4d0cf8bf71cf22d7c046bb22754a8d4e299ed9db..ba0319990bc04196386e6812b0a03671676698ec 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.h +++ b/tensorflow/compiler/xla/tests/client_library_test_base.h @@ -43,6 +43,23 @@ limitations under the License. namespace xla { +// Sets the use_bfloat16 on a container of test cases according to the values in +// use_bfloat16_params. Generates one set of test cases for each values in +// use_bfloat16_params with that value. Returns the result. +template +std::vector ExpandUseBfloat16( + tensorflow::gtl::ArraySlice use_bfloat16_params, + tensorflow::gtl::ArraySlice specs) { + std::vector expanded; + for (bool use_bfloat16 : use_bfloat16_params) { + for (const auto& spec : specs) { + expanded.push_back(spec); + expanded.back().use_bfloat16 = use_bfloat16; + } + } + return expanded; +} + // A client library test establishes an in-process XLA client connection. class ClientLibraryTestBase : public ::testing::Test { protected: @@ -253,6 +270,28 @@ class ClientLibraryTestBase : public ::testing::Test { int64 parameter_number, const Literal& literal, const string& name, ComputationBuilder* builder, ComputationDataHandle* data_handle); + // As above, but the caller can specify the device that the literal is + // transferred to. If device_handle is nullptr, the literal will be + // transferred to the default device. + std::unique_ptr CreateParameterAndTransferLiteral( + int64 parameter_number, const Literal& literal, const string& name, + const DeviceHandle* device_handle, ComputationBuilder* builder, + ComputationDataHandle* data_handle); + + // Creates a parameter instruction and sets the value that will be passed to + // the computation as specified. This function must be used for all parameters + // or none and no parameters must be passed when invoking the computation if + // using this mechanism. If using this mechanism, then each parameter must be + // set exactly once. The first added parameter gets index 0, then 1 and so on. + ComputationDataHandle AddParam(const Literal& argument, + ComputationBuilder* builder); + + template + ComputationDataHandle AddParam(const Array& argument, + ComputationBuilder* builder) { + return AddParam(*Literal::CreateFromArray(argument), builder); + } + // Creates a constant instruction with the given literal. When the // use_bfloat16 flag is set but the literal has F32 elements, the elements // will be converted to BF16s. @@ -370,6 +409,9 @@ class ClientLibraryTestBase : public ::testing::Test { // Whether to run tests with all float-type input/output converted to // bfloat16. bool use_bfloat16_ = false; + + // Arguments to be passed to the computation when it runs. + std::vector> arguments_; }; template @@ -389,6 +431,7 @@ void ClientLibraryTestBase::ComputeAndCompareR0( static_assert(std::is_same::value || std::is_same::value || std::is_same::value || + std::is_same::value || std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = @@ -414,6 +457,7 @@ void ClientLibraryTestBase::ComputeAndCompareR1( static_assert(std::is_same::value || std::is_same::value || std::is_same::value || + std::is_same::value || std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = @@ -439,6 +483,7 @@ void ClientLibraryTestBase::ComputeAndCompareR2( static_assert(std::is_same::value || std::is_same::value || std::is_same::value || + std::is_same::value || std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = @@ -464,6 +509,7 @@ void ClientLibraryTestBase::ComputeAndCompareR3( static_assert(std::is_same::value || std::is_same::value || std::is_same::value || + std::is_same::value || std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = @@ -489,6 +535,7 @@ void ClientLibraryTestBase::ComputeAndCompareR4( static_assert(std::is_same::value || std::is_same::value || std::is_same::value || + std::is_same::value || std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = diff --git a/tensorflow/compiler/xla/tests/client_test.cc b/tensorflow/compiler/xla/tests/client_test.cc index 8853ed9e5780672d4006c326291767b8b5253f56..045148cdd11da94ae4789a753efca95c6aaa1f27 100644 --- a/tensorflow/compiler/xla/tests/client_test.cc +++ b/tensorflow/compiler/xla/tests/client_test.cc @@ -36,7 +36,7 @@ namespace { class ClientTest : public ClientLibraryTestBase {}; -TEST_F(ClientTest, ExecuteWithLayout) { +XLA_TEST_F(ClientTest, ExecuteWithLayout) { ComputationBuilder b(client_, TestName()); std::vector> layouts = {{0, 1}, {1, 0}}; @@ -68,7 +68,7 @@ TEST_F(ClientTest, ExecuteWithLayout) { } } -TEST_F(ClientTest, ExecuteWithTupleLayout) { +XLA_TEST_F(ClientTest, ExecuteWithTupleLayout) { ComputationBuilder b(client_, TestName()); b.Tuple({b.ConstantR2({{1, 2}, {3, 4}}), @@ -90,9 +90,9 @@ TEST_F(ClientTest, ExecuteWithTupleLayout) { auto result, client_->ExecuteAndTransfer(computation, {}, &execution_options)); LiteralTestUtil::ExpectR2Equal({{1, 2}, {3, 4}}, - result->tuple_literals(0)); + LiteralView::Create(*result, {0})); LiteralTestUtil::ExpectR2Equal({{10, 20}, {30, 40}}, - result->tuple_literals(1)); + LiteralView::Create(*result, {1})); EXPECT_TRUE(ShapeUtil::IsTuple(result->shape())); EXPECT_EQ(2, ShapeUtil::TupleElementCount(result->shape())); @@ -107,7 +107,8 @@ TEST_F(ClientTest, ExecuteWithTupleLayout) { /*minor_to_major=*/{1, 0}))); } -TEST_F(ClientTest, DISABLED_ON_CPU_PARALLEL(DISABLED_ON_GPU(ExecuteParallel))) { +XLA_TEST_F(ClientTest, + DISABLED_ON_CPU_PARALLEL(DISABLED_ON_GPU(ExecuteParallel))) { Computation add_with_one_arg, mul_with_two_args, dot_with_one_arg; Shape shape = ShapeUtil::MakeShape(S32, {2, 2}); diff --git a/tensorflow/compiler/xla/tests/codegen_test_base.cc b/tensorflow/compiler/xla/tests/codegen_test_base.cc index e472408dcf7ed5fec74e886fd0092ce47ee2e7eb..022641394f113ef28e7c53058385d77572822213 100644 --- a/tensorflow/compiler/xla/tests/codegen_test_base.cc +++ b/tensorflow/compiler/xla/tests/codegen_test_base.cc @@ -21,9 +21,11 @@ StatusOr> CodegenTestBase::CompileToExecutable( std::unique_ptr hlo_module) { TF_ASSIGN_OR_RETURN(hlo_module, backend().compiler()->RunHloPasses( std::move(hlo_module), - backend().default_stream_executor())); + backend().default_stream_executor(), + /*device_allocator=*/nullptr)); return backend().compiler()->RunBackend(std::move(hlo_module), - backend().default_stream_executor()); + backend().default_stream_executor(), + /*device_allocator=*/nullptr); } StatusOr> diff --git a/tensorflow/compiler/xla/tests/compute_constant_test.cc b/tensorflow/compiler/xla/tests/compute_constant_test.cc index 5226a78386824a94572d3e5cc3329677108a910a..ec2c580670cfac14ba42e8c9a836c86551af4b89 100644 --- a/tensorflow/compiler/xla/tests/compute_constant_test.cc +++ b/tensorflow/compiler/xla/tests/compute_constant_test.cc @@ -149,7 +149,7 @@ TEST_F(ComputeConstantTest, Param) { auto computation = b.Add(param, b.ConstantR0(1.5f)); std::vector arguments; - arguments.emplace_back(*Literal::CreateR0(42.5f)); + arguments.push_back(std::move(*Literal::CreateR0(42.5f))); EXPECT_TRUE(IsConstant(computation, &b, arguments.size())); auto value = @@ -168,7 +168,7 @@ TEST_F(ComputeConstantTest, DirectParamMissing) { auto value = ComputeConstantScalar(client, computation, &b); EXPECT_TRUE(tensorflow::StringPiece(value.status().ToString()) - .contains("depends on parameter")) + .contains("depends on a parameter")) << value.status(); } } @@ -184,7 +184,7 @@ TEST_F(ComputeConstantTest, IndirectParamMissing) { auto value = ComputeConstantScalar(client, computation, &b); EXPECT_TRUE(tensorflow::StringPiece(value.status().ToString()) - .contains("depends on parameter")) + .contains("depends on a parameter")) << value.status(); } } diff --git a/tensorflow/compiler/xla/tests/conditional_test.cc b/tensorflow/compiler/xla/tests/conditional_test.cc index c8c4932be821e410e25c41741df436544ab876f0..bc821674820fb128823786d7149037fc59b22ab6 100644 --- a/tensorflow/compiler/xla/tests/conditional_test.cc +++ b/tensorflow/compiler/xla/tests/conditional_test.cc @@ -23,7 +23,7 @@ namespace { class ConditionalOpTest : public ClientLibraryTestBase { protected: - Computation CreateR0F32ConstantComputation(float value) { + Computation CreateR0ConstantComputation(float value) { ComputationBuilder builder(client_, "Constant"); builder.Parameter(0, empty_tuple_, "tuple"); builder.ConstantR0(value); @@ -32,7 +32,7 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0F32IdentityComputation() { + Computation CreateR0IdentityComputation() { ComputationBuilder builder(client_, "Identity"); builder.Parameter(0, r0f32_, "x"); auto build_status = builder.Build(); @@ -40,25 +40,85 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0F32CeilComputation() { + Computation CreateCeilComputation(const Shape& shape) { ComputationBuilder builder(client_, "Ceil"); - auto param = builder.Parameter(0, r0f32_, "param"); + auto param = builder.Parameter(0, shape, "param"); builder.Ceil(param); auto build_status = builder.Build(); EXPECT_IS_OK(build_status.status()); return build_status.ConsumeValueOrDie(); } - Computation CreateR0F32FloorComputation() { - ComputationBuilder builder(client_, "Ceil"); - auto param = builder.Parameter(0, r0f32_, "param"); + Computation CreateR0CeilComputation() { + return CreateCeilComputation(r0f32_); + } + + Computation CreateR1CeilComputation() { + return CreateCeilComputation(r1s2f32_); + } + + Computation CreateFloorComputation(const Shape& shape) { + ComputationBuilder builder(client_, "Floor"); + auto param = builder.Parameter(0, shape, "param"); builder.Floor(param); auto build_status = builder.Build(); EXPECT_IS_OK(build_status.status()); return build_status.ConsumeValueOrDie(); } - Computation CreateAddTupleComputation(const string& computation_name, + Computation CreateR0FloorComputation() { + return CreateFloorComputation(r0f32_); + } + + Computation CreateR1FloorComputation() { + return CreateFloorComputation(r1s2f32_); + } + + Computation CreateTupleCeilComputation(const string& computation_name, + const Shape& tuple_shape) { + ComputationBuilder builder(client_, computation_name); + auto tuple = builder.Parameter(0, tuple_shape, "tuple"); + auto x = builder.GetTupleElement(tuple, 0); + auto y = builder.GetTupleElement(tuple, 1); + auto x_ceil = builder.Ceil(x); + auto y_ceil = builder.Ceil(y); + builder.Tuple({x_ceil, y_ceil}); + auto build_status = builder.Build(); + EXPECT_IS_OK(build_status.status()); + return build_status.ConsumeValueOrDie(); + } + + Computation CreateR0TupleCeilComputation() { + return CreateTupleCeilComputation("CeilR0", tuple_2_r0f32_); + } + + Computation CreateR1TupleCeilComputation() { + return CreateTupleCeilComputation("CeilR1", tuple_2_r1s2f32_); + } + + Computation CreateTupleFloorComputation(const string& computation_name, + const Shape& tuple_shape) { + ComputationBuilder builder(client_, computation_name); + auto tuple = builder.Parameter(0, tuple_shape, "tuple"); + auto x = builder.GetTupleElement(tuple, 0); + auto y = builder.GetTupleElement(tuple, 1); + auto x_floor = builder.Floor(x); + auto y_floor = builder.Floor(y); + builder.Tuple({x_floor, y_floor}); + auto build_status = builder.Build(); + EXPECT_IS_OK(build_status.status()); + return build_status.ConsumeValueOrDie(); + } + + Computation CreateR0TupleFloorComputation() { + return CreateTupleFloorComputation("FloorR0", tuple_2_r0f32_); + } + + Computation CreateR1TupleFloorComputation() { + return CreateTupleFloorComputation("FloorR1", tuple_2_r1s2f32_); + } + + Computation CreateTupleAddComputation(const string& computation_name, const Shape& tuple_shape) { ComputationBuilder builder(client_, computation_name); auto tuple = builder.Parameter(0, tuple_shape, "tuple"); @@ -70,15 +130,15 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateAddR0Computation() { - return CreateAddTupleComputation("AddR0", tuple_2_r0f32_); + Computation CreateR0TupleAddComputation() { + return CreateTupleAddComputation("AddR0", tuple_2_r0f32_); } - Computation CreateAddR1Computation() { - return CreateAddTupleComputation("AddR1", tuple_2_r1s2f32_); + Computation CreateR1TupleAddComputation() { + return CreateTupleAddComputation("AddR1", tuple_2_r1s2f32_); } - Computation CreateSubTupleComputation(const string& computation_name, + Computation CreateTupleSubComputation(const string& computation_name, const Shape& tuple_shape) { ComputationBuilder builder(client_, computation_name); auto tuple = builder.Parameter(0, tuple_shape, "tuple"); @@ -90,15 +150,16 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateSubR0Computation() { - return CreateSubTupleComputation("SubR0", tuple_2_r0f32_); + Computation CreateR0TupleSubComputation() { + return CreateTupleSubComputation("SubR0", tuple_2_r0f32_); } - Computation CreateSubR1Computation() { - return CreateSubTupleComputation("SubR1", tuple_2_r1s2f32_); + Computation CreateR1TupleSubComputation() { + return CreateTupleSubComputation("SubR1", tuple_2_r1s2f32_); } Shape r0f32_ = ShapeUtil::MakeShape(F32, {}); + Shape r1s2f32_ = ShapeUtil::MakeShape(F32, {2}); Shape tuple_2_r0f32_ = ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShape(F32, {}), ShapeUtil::MakeShape(F32, {})}); Shape tuple_2_r1s2f32_ = ShapeUtil::MakeTupleShape( @@ -112,8 +173,8 @@ XLA_TEST_F(ConditionalOpTest, Parameters0) { ComputationBuilder builder(client_, TestName()); auto pred = builder.ConstantR0(true); auto operands = builder.Tuple({}); - auto true_computation = CreateR0F32ConstantComputation(56.0f); - auto false_computation = CreateR0F32ConstantComputation(12.0f); + auto true_computation = CreateR0ConstantComputation(56.0f); + auto false_computation = CreateR0ConstantComputation(12.0f); auto result = builder.Conditional(pred, operands, true_computation, operands, false_computation); @@ -126,7 +187,7 @@ XLA_TEST_F(ConditionalOpTest, Parameters1) { auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.0f); auto operand2 = builder.ConstantR0(12.0f); - auto identity = CreateR0F32IdentityComputation(); + auto identity = CreateR0IdentityComputation(); auto result = builder.Conditional(pred, operand1, identity, operand2, identity); @@ -140,9 +201,8 @@ XLA_TEST_F(ConditionalOpTest, DiffComputationsDiffArgs) { auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.4f); auto operand2 = builder.ConstantR0(12.6f); - auto result = - builder.Conditional(pred, operand1, CreateR0F32CeilComputation(), - operand2, CreateR0F32FloorComputation()); + auto result = builder.Conditional(pred, operand1, CreateR0CeilComputation(), + operand2, CreateR0FloorComputation()); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -153,8 +213,8 @@ XLA_TEST_F(ConditionalOpTest, DiffComputationsSameArg) { ComputationBuilder builder(client_, TestName()); auto pred = builder.ConstantR0(false); auto operand = builder.ConstantR0(12.6f); - auto result = builder.Conditional(pred, operand, CreateR0F32CeilComputation(), - operand, CreateR0F32FloorComputation()); + auto result = builder.Conditional(pred, operand, CreateR0CeilComputation(), + operand, CreateR0FloorComputation()); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -166,7 +226,7 @@ XLA_TEST_F(ConditionalOpTest, SameComputationDiffArgs) { auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.4f); auto operand2 = builder.ConstantR0(12.6f); - auto floor = CreateR0F32FloorComputation(); + auto floor = CreateR0FloorComputation(); auto result = builder.Conditional(pred, operand1, floor, operand2, floor); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); @@ -178,7 +238,7 @@ XLA_TEST_F(ConditionalOpTest, SameComputationSameArg) { ComputationBuilder builder(client_, TestName()); auto pred = builder.ConstantR0(false); auto operand = builder.ConstantR0(12.6f); - auto floor = CreateR0F32FloorComputation(); + auto floor = CreateR0FloorComputation(); auto result = builder.Conditional(pred, operand, floor, operand, floor); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); @@ -191,9 +251,8 @@ XLA_TEST_F(ConditionalOpTest, SameComputationDiffInstances) { auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.4f); auto operand2 = builder.ConstantR0(12.6f); - auto result = - builder.Conditional(pred, operand1, CreateR0F32FloorComputation(), - operand2, CreateR0F32FloorComputation()); + auto result = builder.Conditional(pred, operand1, CreateR0FloorComputation(), + operand2, CreateR0FloorComputation()); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -205,9 +264,8 @@ XLA_TEST_F(ConditionalOpTest, ConditionalWithCall) { auto pred_cond = inner_builder.Parameter(0, r0bool, "param0"); auto true_operand = inner_builder.Parameter(1, r0f32_, "param1"); auto false_operand = inner_builder.Parameter(2, r0f32_, "param2"); - inner_builder.Conditional(pred_cond, true_operand, - CreateR0F32CeilComputation(), false_operand, - CreateR0F32FloorComputation()); + inner_builder.Conditional(pred_cond, true_operand, CreateR0CeilComputation(), + false_operand, CreateR0FloorComputation()); auto inner_builder_result = inner_builder.Build(); ComputationBuilder builder(client_, TestName()); @@ -228,8 +286,9 @@ XLA_TEST_F(ConditionalOpTest, Parameters2TrueBranch) { auto operand1 = builder.ConstantR0(56.0f); auto operand2 = builder.ConstantR0(12.0f); auto operands = builder.Tuple({operand1, operand2}); - auto result = builder.Conditional(pred, operands, CreateAddR0Computation(), - operands, CreateSubR0Computation()); + auto result = + builder.Conditional(pred, operands, CreateR0TupleAddComputation(), + operands, CreateR0TupleSubComputation()); ComputeAndCompareR0(&builder, 68.0f, {}, error_spec_); } @@ -242,8 +301,9 @@ XLA_TEST_F(ConditionalOpTest, Parameters2FalseBranch) { auto operand1 = builder.ConstantR0(56.0f); auto operand2 = builder.ConstantR0(12.0f); auto operands = builder.Tuple({operand1, operand2}); - auto result = builder.Conditional(pred, operands, CreateAddR0Computation(), - operands, CreateSubR0Computation()); + auto result = + builder.Conditional(pred, operands, CreateR0TupleAddComputation(), + operands, CreateR0TupleSubComputation()); ComputeAndCompareR0(&builder, 44.0f, {}, error_spec_); } @@ -256,8 +316,9 @@ XLA_TEST_F(ConditionalOpTest, Parameters2ArrayTrueBranch) { auto operand1 = builder.ConstantR1({24.0f, 56.0f}); auto operand2 = builder.ConstantR1({10.0f, 11.0f}); auto operands = builder.Tuple({operand1, operand2}); - auto result = builder.Conditional(pred, operands, CreateAddR1Computation(), - operands, CreateSubR1Computation()); + auto result = + builder.Conditional(pred, operands, CreateR1TupleAddComputation(), + operands, CreateR1TupleSubComputation()); ComputeAndCompareR1(&builder, {34.0f, 67.0f}, {}, error_spec_); } @@ -270,25 +331,188 @@ XLA_TEST_F(ConditionalOpTest, Parameters2ArrayFalseBranch) { auto operand1 = builder.ConstantR1({24.0f, 56.0f}); auto operand2 = builder.ConstantR1({10.0f, 11.0f}); auto operands = builder.Tuple({operand1, operand2}); - auto result = builder.Conditional(pred, operands, CreateAddR1Computation(), - operands, CreateSubR1Computation()); + auto result = + builder.Conditional(pred, operands, CreateR1TupleAddComputation(), + operands, CreateR1TupleSubComputation()); ComputeAndCompareR1(&builder, {14.0f, 45.0f}, {}, error_spec_); } +// Test true and false computations that return a tuple of scalars. +XLA_TEST_F(ConditionalOpTest, ReturnTupleOfScalars) { + ComputationBuilder builder(client_, TestName()); + auto pred = builder.ConstantR0(false); + auto operands = builder.Tuple( + {builder.ConstantR0(12.2f), builder.ConstantR0(25.6f)}); + builder.Conditional(pred, operands, CreateR0TupleCeilComputation(), operands, + CreateR0TupleFloorComputation()); + + ComputeAndCompareTuple( + &builder, + *Literal::MakeTuple({Literal::CreateR0(12.0f).get(), + Literal::CreateR0(25.0f).get()}), + {}, error_spec_); +} + +// Test true and false computations that return a tuple of arrays. +XLA_TEST_F(ConditionalOpTest, ReturnTupleOfArrays) { + ComputationBuilder builder(client_, TestName()); + auto pred = builder.ConstantR0(true); + auto operands = builder.Tuple({builder.ConstantR1({12.2f, 15.8f}), + builder.ConstantR1({25.6f, 29.2f})}); + builder.Conditional(pred, operands, CreateR1TupleCeilComputation(), operands, + CreateR1TupleFloorComputation()); + + ComputeAndCompareTuple( + &builder, + *Literal::MakeTuple({Literal::CreateR1({13.0f, 16.0f}).get(), + Literal::CreateR1({26.0f, 30.0f}).get()}), + {}, error_spec_); +} + +// Test true and false computations that return a tuple of a predicate, a +// scalar, and an array. +XLA_TEST_F(ConditionalOpTest, ReturnTupleofPredicateScalarArray) { + ComputationBuilder true_builder(client_, TestName() + ".true"); + { + true_builder.Parameter(0, empty_tuple_, "tuple"); + auto true_pred = true_builder.ConstantR0(true); + auto true_scalar = true_builder.ConstantR0(12.2f); + auto true_array = true_builder.ConstantR1({12.8f, 14.6f}); + true_builder.Tuple({true_pred, true_scalar, true_array}); + } + auto true_builder_result = true_builder.Build(); + EXPECT_IS_OK(true_builder_result.status()); + + ComputationBuilder false_builder(client_, TestName() + ".false"); + { + false_builder.Parameter(0, empty_tuple_, "tuple"); + auto false_pred = false_builder.ConstantR0(false); + auto false_scalar = false_builder.ConstantR0(25.6f); + auto false_array = false_builder.ConstantR1({26.4f, 32.6f}); + false_builder.Tuple({false_pred, false_scalar, false_array}); + } + auto false_builder_result = false_builder.Build(); + EXPECT_IS_OK(false_builder_result.status()); + + ComputationBuilder builder(client_, TestName()); + auto pred = builder.ConstantR0(true); + auto operands = builder.Tuple({}); + builder.Conditional(pred, operands, true_builder_result.ConsumeValueOrDie(), + operands, false_builder_result.ConsumeValueOrDie()); + + ComputeAndCompareTuple( + &builder, + *Literal::MakeTuple({Literal::CreateR0(true).get(), + Literal::CreateR0(12.2f).get(), + Literal::CreateR1({12.8f, 14.6f}).get()}), + {}, error_spec_); +} + +// Test true and false computations that return a nested tuple. +XLA_TEST_F(ConditionalOpTest, ReturnNestedTuple) { + ComputationBuilder true_builder(client_, TestName() + ".true"); + { + true_builder.Parameter(0, empty_tuple_, "tuple"); + auto true_constant1 = true_builder.ConstantR0(12.2f); + auto true_constant2 = true_builder.ConstantR1({12.8f, 14.6f}); + auto true_constant3 = true_builder.ConstantR1({25.4f, 29.8f}); + auto true_constant4 = true_builder.ConstantR0(35.6f); + true_builder.Tuple({true_builder.Tuple({true_constant1, true_constant2}), + true_builder.Tuple({true_constant3, true_constant4})}); + } + auto true_builder_result = true_builder.Build(); + EXPECT_IS_OK(true_builder_result.status()); + + ComputationBuilder false_builder(client_, TestName() + ".false"); + { + false_builder.Parameter(0, empty_tuple_, "tuple"); + auto false_constant1 = false_builder.ConstantR0(46.6f); + auto false_constant2 = false_builder.ConstantR1({54.4f, 58.4f}); + auto false_constant3 = false_builder.ConstantR1({62.1f, 67.4f}); + auto false_constant4 = false_builder.ConstantR0(9.3f); + false_builder.Tuple( + {false_builder.Tuple({false_constant1, false_constant2}), + false_builder.Tuple({false_constant3, false_constant4})}); + } + auto false_builder_result = false_builder.Build(); + EXPECT_IS_OK(false_builder_result.status()); + + ComputationBuilder builder(client_, TestName()); + auto pred = builder.ConstantR0(false); + auto operands = builder.Tuple({}); + builder.Conditional(pred, operands, true_builder_result.ConsumeValueOrDie(), + operands, false_builder_result.ConsumeValueOrDie()); + + ComputeAndCompareTuple( + &builder, + *Literal::MakeTuple( + {Literal::MakeTuple({Literal::CreateR0(46.6f).get(), + Literal::CreateR1({54.4f, 58.4f}).get()}) + .get(), + Literal::MakeTuple({Literal::CreateR1({62.1f, 67.4f}).get(), + Literal::CreateR0(9.3f).get()}) + .get()}), + {}, error_spec_); +} + +// Test conditional that takes in scalar operands in the form of external +// params. +XLA_TEST_F(ConditionalOpTest, ScalarOperandsFromExternalParams) { + Shape r0bool = ShapeUtil::MakeShape(PRED, {}); + ComputationBuilder builder(client_, TestName()); + + ComputationDataHandle pred, operand1, operand2; + auto pred_arg = CreateR0Parameter(true, 0, "pred", &builder, &pred); + auto operand1_param = + CreateR0Parameter(56.3f, 1, "operand1", &builder, &operand1); + auto operand2_param = + CreateR0Parameter(12.7f, 2, "operand2", &builder, &operand2); + auto result = builder.Conditional(pred, operand1, CreateR0CeilComputation(), + operand2, CreateR0FloorComputation()); + + ComputeAndCompareR0( + &builder, 57.0f, + {pred_arg.get(), operand1_param.get(), operand2_param.get()}, + error_spec_); +} + +// Test conditional that takes in array operands in the form of external params. +XLA_TEST_F(ConditionalOpTest, ArrayOperandsFromExternalParams) { + Shape r0bool = ShapeUtil::MakeShape(PRED, {}); + ComputationBuilder builder(client_, TestName()); + + ComputationDataHandle pred, operand1, operand2; + auto pred_arg = CreateR0Parameter(false, 0, "pred", &builder, &pred); + auto operand1_param = CreateR1Parameter({24.3f, 56.7f}, 1, "operand1", + &builder, &operand1); + auto operand2_param = CreateR1Parameter({10.2f, 11.6f}, 2, "operand2", + &builder, &operand2); + auto result = builder.Conditional(pred, operand1, CreateR1CeilComputation(), + operand2, CreateR1FloorComputation()); + + ComputeAndCompareR1( + &builder, {10.0f, 11.0f}, + {pred_arg.get(), operand1_param.get(), operand2_param.get()}, + error_spec_); +} + // Test the case where one conditional is nested within another. XLA_TEST_F(ConditionalOpTest, NestedConditionals) { - Shape r0bool = ShapeUtil::MakeShape(PRED, {}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({r0bool, r0f32_, r0f32_}); ComputationBuilder inner_builder(client_, TestName() + ".inner_conditional"); - auto param0 = inner_builder.Parameter(0, tuple_shape, "param0"); - auto pred_cond = inner_builder.GetTupleElement(param0, 0); - auto true_operand = inner_builder.GetTupleElement(param0, 1); - auto false_operand = inner_builder.GetTupleElement(param0, 2); - inner_builder.Conditional(pred_cond, true_operand, - CreateR0F32CeilComputation(), false_operand, - CreateR0F32FloorComputation()); + { + Shape r0bool = ShapeUtil::MakeShape(PRED, {}); + Shape tuple_shape = ShapeUtil::MakeTupleShape({r0bool, r0f32_, r0f32_}); + auto param0 = inner_builder.Parameter(0, tuple_shape, "param0"); + auto pred_cond = inner_builder.GetTupleElement(param0, 0); + auto true_operand = inner_builder.GetTupleElement(param0, 1); + auto false_operand = inner_builder.GetTupleElement(param0, 2); + inner_builder.Conditional(pred_cond, true_operand, + CreateR0CeilComputation(), false_operand, + CreateR0FloorComputation()); + } auto inner_builder_result = inner_builder.Build(); + EXPECT_IS_OK(inner_builder_result.status()); ComputationBuilder builder(client_, TestName()); auto pred1 = builder.ConstantR0(true); @@ -299,7 +523,33 @@ XLA_TEST_F(ConditionalOpTest, NestedConditionals) { auto tuple_operand = builder.Tuple({pred2, operand1, operand2}); builder.Conditional(pred1, tuple_operand, inner_builder_result.ConsumeValueOrDie(), operand3, - CreateR0F32IdentityComputation()); + CreateR0IdentityComputation()); + + ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); +} + +XLA_TEST_F(ConditionalOpTest, ConditionalInNestedComputation) { + ComputationBuilder inner_builder(client_, TestName() + ".inner_conditional"); + { + Shape r0bool = ShapeUtil::MakeShape(PRED, {}); + Shape tuple_shape = ShapeUtil::MakeTupleShape({r0bool, r0f32_, r0f32_}); + auto param0 = inner_builder.Parameter(0, tuple_shape, "param0"); + auto pred_cond = inner_builder.GetTupleElement(param0, 0); + auto true_operand = inner_builder.GetTupleElement(param0, 1); + auto false_operand = inner_builder.GetTupleElement(param0, 2); + inner_builder.Conditional(pred_cond, true_operand, + CreateR0CeilComputation(), false_operand, + CreateR0FloorComputation()); + } + auto inner_builder_result = inner_builder.Build(); + EXPECT_IS_OK(inner_builder_result.status()); + + ComputationBuilder builder(client_, TestName()); + auto pred2 = builder.ConstantR0(false); + auto operand1 = builder.ConstantR0(1.1f); + auto operand2 = builder.ConstantR0(12.2f); + auto tuple_operand = builder.Tuple({pred2, operand1, operand2}); + builder.Call(inner_builder_result.ConsumeValueOrDie(), {tuple_operand}); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -311,8 +561,8 @@ XLA_TEST_F(ConditionalOpTest, ShapeMismatch) { auto operand1 = builder.ConstantR0(56.0f); auto operand2 = builder.ConstantR0(12.0f); auto operands = builder.Tuple({operand1, operand2}); - builder.Conditional(pred, operands, CreateAddR1Computation(), operands, - CreateSubR0Computation()); + builder.Conditional(pred, operands, CreateR1TupleAddComputation(), operands, + CreateR0TupleSubComputation()); auto result = builder.Build(); EXPECT_FALSE(result.ok()); diff --git a/tensorflow/compiler/xla/tests/constants_test.cc b/tensorflow/compiler/xla/tests/constants_test.cc index 97bd1553664a6c0fcb097b441ec42efb4eaa9cc2..35aa3f6d696297efb7d95d826ed75a504a24529d 100644 --- a/tensorflow/compiler/xla/tests/constants_test.cc +++ b/tensorflow/compiler/xla/tests/constants_test.cc @@ -141,11 +141,12 @@ TEST_F(ConstantsTest, Small_3x2x1x1) { {5.0f, 4.4f}, // p2 }); input_array.FillWithPZ(pz); - Literal input_literal = *Literal::CreateR4FromArray4D(input_array); + std::unique_ptr input_literal = + Literal::CreateR4FromArray4D(input_array); { ComputationBuilder builder(client_, TestName()); - builder.ConstantLiteral(input_literal); + builder.ConstantLiteral(*input_literal); ComputeAndCompareR4(&builder, input_array, {}, error_spec_); } @@ -165,10 +166,10 @@ TEST_F(ConstantsTest, DISABLED_TupleConstant) { std::unique_ptr result = ExecuteAndTransferOrDie(&builder, {}); - LiteralTestUtil::ExpectR2Near({{1.0}, {2.0}}, - result->tuple_literals(0), error_spec_); - LiteralTestUtil::ExpectR1Near({2.0, 42.0}, result->tuple_literals(1), - error_spec_); + LiteralTestUtil::ExpectR2Near( + {{1.0}, {2.0}}, LiteralView::Create(*result, {0}), error_spec_); + LiteralTestUtil::ExpectR1Near( + {2.0, 42.0}, LiteralView::Create(*result, {1}), error_spec_); } } // namespace diff --git a/tensorflow/compiler/xla/tests/convolution_test.cc b/tensorflow/compiler/xla/tests/convolution_test.cc index 2924c08615fa706bb19addf04bf58e1d5dd5a659..1385b437fc47fe5289c401581fab8b5278872382 100644 --- a/tensorflow/compiler/xla/tests/convolution_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_test.cc @@ -53,157 +53,200 @@ class ConvolutionTest : public ClientLibraryTestBase { #endif }; -XLA_TEST_F(ConvolutionTest, ForwardPassConvolution_3x3x256_256_OutputZ_Iota) { - const int kInputActivationSizeY = 3; - const int kInputActivationSizeX = 3; - const int kInputActivationSizeZ = 256; - const int kKernelSizeX = 2; - const int kKernelSizeY = 2; - const int kOutputActivationSizeZ = 256; - const int kMiniBatchSize = 4; - auto alhs = - MakeUnique>(kMiniBatchSize, kInputActivationSizeZ, - kInputActivationSizeY, kInputActivationSizeX); - alhs->FillWithMultiples(1.0f); - ASSERT_EQ(3, alhs->width()); - ASSERT_EQ(3, alhs->height()); - - auto arhs = - MakeUnique>(kOutputActivationSizeZ, kInputActivationSizeZ, - kKernelSizeY, kKernelSizeX); - Array2D rhs_raster({ - {1.0f, 0.0f}, // row 0 - {0.0f, 0.0f}, // row 1 - }); - arhs->FillWithYX(rhs_raster); - ASSERT_EQ(2, arhs->width()); - ASSERT_EQ(2, arhs->height()); +// TODO(b/72509305): Enable half data type tests for CPU +#if (XLA_TEST_BACKEND_GPU) +using TestTypes = ::testing::Types; +#else +using TestTypes = ::testing::Types; +#endif - ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR4FromArray4D(*alhs); - auto rhs = builder.ConstantR4FromArray4D(*arhs); - auto conv = builder.Conv(lhs, rhs, {1, 1}, Padding::kValid); +template +Shape MakeShapeWrapper(tensorflow::gtl::ArraySlice dimensions); - ComputeAndCompare(&builder, conv, {}, error_spec_); +template <> +Shape MakeShapeWrapper(tensorflow::gtl::ArraySlice dimensions) { + return ShapeUtil::MakeShape(F32, dimensions); } -TEST_F(ConvolutionTest, Convolve_1x1x1x2_1x1x1x2_Valid) { - ComputationBuilder builder(client_, TestName()); - Shape input_shape = ShapeUtil::MakeShape(F32, {1, 1, 1, 2}); - Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 1, 1, 2}); - auto input = builder.Parameter(0, input_shape, "input"); - auto filter = builder.Parameter(1, filter_shape, "filter"); - auto conv = builder.Conv(input, filter, {1, 1}, Padding::kValid); +template <> +Shape MakeShapeWrapper( + tensorflow::gtl::ArraySlice dimensions) { + return ShapeUtil::MakeShape(F16, dimensions); +} - Array4D input_data(1, 1, 1, 2); - input_data.FillWithYX(Array2D({ - {1, 2}, - })); - Array4D filter_data(1, 1, 1, 2); - filter_data.FillWithYX(Array2D({ - {5, 6}, - })); +template +class ForwardPassConvolution_3x3x256_256_OutputZ_Iota : public ConvolutionTest { + public: + void RunTest() { + const int kInputActivationSizeY = 3; + const int kInputActivationSizeX = 3; + const int kInputActivationSizeZ = 256; + const int kKernelSizeX = 2; + const int kKernelSizeY = 2; + const int kOutputActivationSizeZ = 256; + const int kMiniBatchSize = 4; + auto alhs = + MakeUnique>(kMiniBatchSize, kInputActivationSizeZ, + kInputActivationSizeY, kInputActivationSizeX); + alhs->FillWithMultiples(static_cast(1.0f)); + ASSERT_EQ(3, alhs->width()); + ASSERT_EQ(3, alhs->height()); + + auto arhs = + MakeUnique>(kOutputActivationSizeZ, kInputActivationSizeZ, + kKernelSizeY, kKernelSizeX); + Array2D rhs_raster({ + {1.0f, 0.0f}, // row 0 + {0.0f, 0.0f}, // row 1 + }); + arhs->FillWithYX(rhs_raster); + ASSERT_EQ(2, arhs->width()); + ASSERT_EQ(2, arhs->height()); + + ComputationBuilder builder(client_, TestName()); + auto lhs = builder.ConstantR4FromArray4D(*alhs); + auto rhs = builder.ConstantR4FromArray4D(*arhs); + auto conv = builder.Conv(lhs, rhs, {1, 1}, Padding::kValid); + + ComputeAndCompare(&builder, conv, {}, error_spec_); + } +}; - ComputeAndCompare(&builder, conv, - {*Literal::CreateFromArray(input_data), - *Literal::CreateFromArray(filter_data)}, - error_spec_); +TYPED_TEST_CASE(ForwardPassConvolution_3x3x256_256_OutputZ_Iota, TestTypes); +XLA_TYPED_TEST(ForwardPassConvolution_3x3x256_256_OutputZ_Iota, Types) { + this->RunTest(); } +template +class Convolve_1x1x1x2_1x1x1x2_Valid : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + Shape input_shape = MakeShapeWrapper({1, 1, 1, 2}); + Shape filter_shape = MakeShapeWrapper({1, 1, 1, 2}); + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + auto conv = builder.Conv(input, filter, {1, 1}, Padding::kValid); + + Array4D input_data(1, 1, 1, 2); + input_data.FillWithYX(Array2D({ + {1.0f, 2.0f}, + })); + Array4D filter_data(1, 1, 1, 2); + filter_data.FillWithYX(Array2D({ + {5.0f, 6.0f}, + })); + + ComputeAndCompare(&builder, conv, + {std::move(*Literal::CreateFromArray(input_data)), + std::move(*Literal::CreateFromArray(filter_data))}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve_1x1x1x2_1x1x1x2_Valid, TestTypes); +TYPED_TEST(Convolve_1x1x1x2_1x1x1x2_Valid, Types) { this->RunTest(); } + // Tests valid padding for 2D convolution in raster space. -TEST_F(ConvolutionTest, Convolve_1x1x4x4_1x1x2x2_Valid) { - ComputationBuilder builder(client_, TestName()); - Shape input_shape = ShapeUtil::MakeShape(F32, {1, 1, 4, 4}); - Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 1, 2, 2}); - auto input = builder.Parameter(0, input_shape, "input"); - auto filter = builder.Parameter(1, filter_shape, "filter"); - auto conv = builder.Conv(input, filter, {1, 1}, Padding::kValid); +template +class Convolve_1x1x4x4_1x1x2x2_Valid : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + Shape input_shape = MakeShapeWrapper({1, 1, 4, 4}); + Shape filter_shape = MakeShapeWrapper({1, 1, 2, 2}); + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + auto conv = builder.Conv(input, filter, {1, 1}, Padding::kValid); + + Array4D input_data(1, 1, 4, 4); + input_data.FillWithYX(Array2D({ + {1.0f, 2.0f, 3.0f, 4.0f}, + {5.0f, 6.0f, 7.0f, 8.0f}, + {9.0f, 10.0f, 11.0f, 12.0f}, + {13.0f, 14.0f, 15.0f, 16.0f}, + })); + Array4D filter_data(1, 1, 2, 2); + filter_data.FillWithYX(Array2D({ + {5.0f, 6.0f}, + {7.0f, 8.0f}, + })); + ComputeAndCompare(&builder, conv, + {std::move(*Literal::CreateFromArray(input_data)), + std::move(*Literal::CreateFromArray(filter_data))}, + error_spec_); + } +}; - Array4D input_data(1, 1, 4, 4); - // clang-format off - input_data.FillWithYX(Array2D({ - {1, 2, 3, 4 }, - {5, 6, 7, 8 }, - {9, 10, 11, 12}, - {13, 14, 15, 16}, - })); - // clang-format on - Array4D filter_data(1, 1, 2, 2); - // clang-format off - filter_data.FillWithYX(Array2D({ - {5, 6}, - {7, 8}, - })); - // clang-format on - ComputeAndCompare(&builder, conv, - {*Literal::CreateFromArray(input_data), - *Literal::CreateFromArray(filter_data)}, - error_spec_); -} +TYPED_TEST_CASE(Convolve_1x1x4x4_1x1x2x2_Valid, TestTypes); +TYPED_TEST(Convolve_1x1x4x4_1x1x2x2_Valid, Types) { this->RunTest(); } // Tests same padding for 2D convolution in raster space. -TEST_F(ConvolutionTest, Convolve_1x1x4x4_1x1x2x2_Same) { - ComputationBuilder builder(client_, TestName()); - Shape input_shape = ShapeUtil::MakeShape(F32, {1, 1, 4, 4}); - Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 1, 2, 2}); - auto input = builder.Parameter(0, input_shape, "input"); - auto filter = builder.Parameter(1, filter_shape, "filter"); - auto conv = builder.Conv(input, filter, {1, 1}, Padding::kSame); - - Array4D input_data(1, 1, 4, 4); - // clang-format off - input_data.FillWithYX(Array2D({ - {1, 2, 3, 4 }, - {5, 6, 7, 8 }, - {9, 10, 11, 12}, - {13, 14, 15, 16}, - })); - // clang-format on - Array4D filter_data(1, 1, 2, 2); - // clang-format off - filter_data.FillWithYX(Array2D({ - {5, 6}, - {7, 8}, - })); - // clang-format on - ComputeAndCompare(&builder, conv, - {*Literal::CreateFromArray(input_data), - *Literal::CreateFromArray(filter_data)}, - error_spec_); -} +template +class Convolve_1x1x4x4_1x1x2x2_Same : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + Shape input_shape = MakeShapeWrapper({1, 1, 4, 4}); + Shape filter_shape = MakeShapeWrapper({1, 1, 2, 2}); + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + auto conv = builder.Conv(input, filter, {1, 1}, Padding::kSame); + + Array4D input_data(1, 1, 4, 4); + input_data.FillWithYX(Array2D({ + {1.0f, 2.0f, 3.0f, 4.0f}, + {5.0f, 6.0f, 7.0f, 8.0f}, + {9.0f, 10.0f, 11.0f, 12.0f}, + {13.0f, 14.0f, 15.0f, 16.0f}, + })); + Array4D filter_data(1, 1, 2, 2); + filter_data.FillWithYX(Array2D({ + {5.0f, 6.0f}, + {7.0f, 8.0f}, + })); + + ComputeAndCompare(&builder, conv, + {std::move(*Literal::CreateFromArray(input_data)), + std::move(*Literal::CreateFromArray(filter_data))}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve_1x1x4x4_1x1x2x2_Same, TestTypes); +TYPED_TEST(Convolve_1x1x4x4_1x1x2x2_Same, Types) { this->RunTest(); } // Tests same padding for 2D convolution in raster space with an odd sized // kernel. -TEST_F(ConvolutionTest, Convolve_1x1x4x4_1x1x3x3_Same) { - ComputationBuilder builder(client_, TestName()); - Shape input_shape = ShapeUtil::MakeShape(F32, {1, 1, 4, 4}); - Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 1, 3, 3}); - auto input = builder.Parameter(0, input_shape, "input"); - auto filter = builder.Parameter(1, filter_shape, "filter"); - auto conv = builder.Conv(input, filter, {1, 1}, Padding::kSame); - - Array4D input_data(1, 1, 4, 4); - // clang-format off - input_data.FillWithYX(Array2D({ - {1, 2, 3, 4 }, - {5, 6, 7, 8 }, - {9, 10, 11, 12}, - {13, 14, 15, 16}, - })); - // clang-format on - Array4D filter_data(1, 1, 3, 3); - // clang-format off - filter_data.FillWithYX(Array2D({ - { 5, 6, 7}, - { 8, 9, 10}, - {11, 12, 13}, - })); - // clang-format on - ComputeAndCompare(&builder, conv, - {*Literal::CreateFromArray(input_data), - *Literal::CreateFromArray(filter_data)}, - error_spec_); -} +template +class Convolve_1x1x4x4_1x1x3x3_Same : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + Shape input_shape = MakeShapeWrapper({1, 1, 4, 4}); + Shape filter_shape = MakeShapeWrapper({1, 1, 3, 3}); + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + auto conv = builder.Conv(input, filter, {1, 1}, Padding::kSame); + + Array4D input_data(1, 1, 4, 4); + input_data.FillWithYX(Array2D({{1.0f, 2.0f, 3.0f, 4.0f}, + {5.0f, 6.0f, 7.0f, 8.0f}, + {9.0f, 10.0f, 11.0f, 12.0f}, + {13.0f, 14.0f, 15.0f, 16.0f}})); + Array4D filter_data(1, 1, 3, 3); + filter_data.FillWithYX(Array2D( + {{5.0f, 6.0f, 7.0f}, {8.0f, 9.0f, 10.0f}, {11.0f, 12.0f, 13.0f}})); + // clang-format on + ComputeAndCompare(&builder, conv, + {std::move(*Literal::CreateFromArray(input_data)), + std::move(*Literal::CreateFromArray(filter_data))}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve_1x1x4x4_1x1x3x3_Same, TestTypes); +TYPED_TEST(Convolve_1x1x4x4_1x1x3x3_Same, Types) { this->RunTest(); } XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_Valid) { ComputationBuilder builder(client_, TestName()); @@ -232,36 +275,44 @@ XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_Valid) { error_spec_); } -XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_WithRHSDilation) { - ComputationBuilder builder(client_, TestName()); - { - Shape input_shape = ShapeUtil::MakeShape(F32, {1, 2, 5}); - Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); - auto input = builder.Parameter(0, input_shape, "input"); - auto filter = builder.Parameter(1, filter_shape, "filter"); - // Convolution dimensions are bf0_oi0->bo0. - builder.ConvGeneralDilated( - input, filter, /*window_strides=*/{1}, /*padding=*/{{0, 0}}, - /*lhs_dilation=*/{1}, /*rhs_dilation=*/{2}, - /*dimension_numbers=*/builder.CreateDefaultConvDimensionNumbers(1)); +template +class Convolve1D_1x2x5_1x2x2_WithRHSDilation : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + { + Shape input_shape = MakeShapeWrapper({1, 2, 5}); + Shape filter_shape = MakeShapeWrapper({1, 2, 2}); + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + // Convolution dimensions are bf0_oi0->bo0. + builder.ConvGeneralDilated( + input, filter, /*window_strides=*/{1}, /*padding=*/{{0, 0}}, + /*lhs_dilation=*/{1}, /*rhs_dilation=*/{2}, + /*dimension_numbers=*/builder.CreateDefaultConvDimensionNumbers(1)); + } + + Array3D input( + {{{1.0f, 2.0f, 3.0f, 4.0f, 5.0f}, {6.0f, 7.0f, 8.0f, 9.0f, 10.0f}}}); + Array3D filter({{{10.0f, 20.0f}, {30.0f, 40.0f}}}); + + Array3D expected({{{570.0f, 670.0f, 770.0f}}}); + + auto input_literal = + client_->TransferToServer(*Literal::CreateR3FromArray3D(input)) + .ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(*Literal::CreateR3FromArray3D(filter)) + .ConsumeValueOrDie(); + + ComputeAndCompareR3(&builder, expected, + {input_literal.get(), filter_literal.get()}, + error_spec_); } +}; // namespace - Array3D input({{{1, 2, 3, 4, 5}, {6, 7, 8, 9, 10}}}); - Array3D filter({{{10, 20}, {30, 40}}}); - - Array3D expected({{{570, 670, 770}}}); - - auto input_literal = - client_->TransferToServer(*Literal::CreateR3FromArray3D(input)) - .ConsumeValueOrDie(); - auto filter_literal = - client_->TransferToServer(*Literal::CreateR3FromArray3D(filter)) - .ConsumeValueOrDie(); - - ComputeAndCompareR3(&builder, expected, - {input_literal.get(), filter_literal.get()}, - error_spec_); -} +TYPED_TEST_CASE(Convolve1D_1x2x5_1x2x2_WithRHSDilation, TestTypes); +TYPED_TEST(Convolve1D_1x2x5_1x2x2_WithRHSDilation, Types) { this->RunTest(); } XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_WithLHSDilation) { ComputationBuilder builder(client_, TestName()); @@ -325,36 +376,45 @@ XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_WithLHSAndRHSDilation) { error_spec_); } -XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_WithPadding) { - ComputationBuilder builder(client_, TestName()); - { - Shape input_shape = ShapeUtil::MakeShape(F32, {1, 2, 5}); - Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); - auto input = builder.Parameter(0, input_shape, "input"); - auto filter = builder.Parameter(1, filter_shape, "filter"); - // Convolution dimensions are bf0_oi0->bo0. - builder.ConvGeneralDilated( - input, filter, /*window_strides=*/{1}, /*padding=*/{{2, 2}}, - /*lhs_dilation=*/{1}, /*rhs_dilation=*/{1}, - /*dimension_numbers=*/builder.CreateDefaultConvDimensionNumbers(1)); +template +class Convolve1D_1x2x5_1x2x2_WithPadding : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + { + Shape input_shape = MakeShapeWrapper({1, 2, 5}); + Shape filter_shape = MakeShapeWrapper({1, 2, 2}); + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + // Convolution dimensions are bf0_oi0->bo0. + builder.ConvGeneralDilated( + input, filter, /*window_strides=*/{1}, /*padding=*/{{2, 2}}, + /*lhs_dilation=*/{1}, /*rhs_dilation=*/{1}, + /*dimension_numbers=*/builder.CreateDefaultConvDimensionNumbers(1)); + } + + Array3D input( + {{{1.0f, 2.0f, 3.0f, 4.0f, 5.0f}, {6.0f, 7.0f, 8.0f, 9.0f, 10.0f}}}); + Array3D filter({{{10.0f, 20.0f}, {30.0f, 40.0f}}}); + + Array3D expected( + {{{0.0f, 260.0f, 510.0f, 610.0f, 710.0f, 810.0f, 350.0f, 0.0f}}}); + + auto input_literal = + client_->TransferToServer(*Literal::CreateR3FromArray3D(input)) + .ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(*Literal::CreateR3FromArray3D(filter)) + .ConsumeValueOrDie(); + + ComputeAndCompareR3(&builder, expected, + {input_literal.get(), filter_literal.get()}, + error_spec_); } +}; - Array3D input({{{1, 2, 3, 4, 5}, {6, 7, 8, 9, 10}}}); - Array3D filter({{{10, 20}, {30, 40}}}); - - Array3D expected({{{0, 260, 510, 610, 710, 810, 350, 0}}}); - - auto input_literal = - client_->TransferToServer(*Literal::CreateR3FromArray3D(input)) - .ConsumeValueOrDie(); - auto filter_literal = - client_->TransferToServer(*Literal::CreateR3FromArray3D(filter)) - .ConsumeValueOrDie(); - - ComputeAndCompareR3(&builder, expected, - {input_literal.get(), filter_literal.get()}, - error_spec_); -} +TYPED_TEST_CASE(Convolve1D_1x2x5_1x2x2_WithPadding, TestTypes); +TYPED_TEST(Convolve1D_1x2x5_1x2x2_WithPadding, Types) { this->RunTest(); } XLA_TEST_F(ConvolutionTest, Convolve3D_1x4x2x3x3_2x2x2x3x3_Valid) { ComputationBuilder builder(client_, TestName()); @@ -389,12 +449,12 @@ XLA_TEST_F(ConvolutionTest, Convolve3D_1x4x2x3x3_2x2x2x3x3_Valid) { } std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); - std::iota(input_elems.begin(), input_elems.end(), 1.0f); + iota(input_elems.begin(), input_elems.end(), 1.0f); auto input_r1 = Literal::CreateR1(input_elems); auto input_r5 = input_r1->Reshape(input_dims).ConsumeValueOrDie(); std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); - std::iota(filter_elems.begin(), filter_elems.end(), 1.0f); + iota(filter_elems.begin(), filter_elems.end(), 1.0f); auto filter_r1 = Literal::CreateR1(filter_elems); auto filter_r5 = filter_r1->Reshape(filter_dims).ConsumeValueOrDie(); @@ -412,56 +472,73 @@ XLA_TEST_F(ConvolutionTest, Convolve3D_1x4x2x3x3_2x2x2x3x3_Valid) { error_spec_); } -XLA_TEST_F(ConvolutionTest, Convolve2D_1x3x3x5_3x3x5x5_Valid) { - ComputationBuilder builder(client_, TestName()); - std::vector input_dims = {1, 3, 3, 5}; - std::vector filter_dims = {3, 3, 5, 3}; - Shape input_shape = ShapeUtil::MakeShape(F32, input_dims); - Shape filter_shape = ShapeUtil::MakeShape(F32, filter_dims); - { - auto input = builder.Parameter(0, input_shape, "input"); - auto filter = builder.Parameter(1, filter_shape, "filter"); - - // Tensorflow dimension numbers for 2D convolution. - ConvolutionDimensionNumbers dnums; - dnums.set_input_batch_dimension(0); - dnums.set_output_batch_dimension(0); - dnums.add_input_spatial_dimensions(1); - dnums.add_output_spatial_dimensions(1); - dnums.add_input_spatial_dimensions(2); - dnums.add_output_spatial_dimensions(2); - dnums.set_input_feature_dimension(3); - dnums.set_output_feature_dimension(3); - dnums.add_kernel_spatial_dimensions(0); - dnums.add_kernel_spatial_dimensions(1); - dnums.set_kernel_input_feature_dimension(2); - dnums.set_kernel_output_feature_dimension(3); +// std::iota doesn't work when init_value has a type Eigen::half in some build +// servers. The error message is missing the operator ++. +template +void iota_int_init_value(std::vector& values, int init_value) { + std::for_each(values.begin(), values.end(), + [&](T& value) { value = static_cast(init_value++); }); +} - builder.ConvWithGeneralDimensions(input, filter, {1, 1}, Padding::kValid, - dnums); +template +class Convolve2D_1x3x3x5_3x3x5x5_Valid : public ConvolutionTest { + public: + void RunTest() { + ComputationBuilder builder(client_, TestName()); + std::vector input_dims = {1, 3, 3, 5}; + std::vector filter_dims = {3, 3, 5, 3}; + Shape input_shape = MakeShapeWrapper(input_dims); + Shape filter_shape = MakeShapeWrapper(filter_dims); + { + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + ConvolutionDimensionNumbers dnums; + dnums.set_input_batch_dimension(0); + dnums.set_output_batch_dimension(0); + dnums.add_input_spatial_dimensions(1); + dnums.add_output_spatial_dimensions(1); + dnums.add_input_spatial_dimensions(2); + dnums.add_output_spatial_dimensions(2); + dnums.set_input_feature_dimension(3); + dnums.set_output_feature_dimension(3); + dnums.add_kernel_spatial_dimensions(0); + dnums.add_kernel_spatial_dimensions(1); + dnums.set_kernel_input_feature_dimension(2); + dnums.set_kernel_output_feature_dimension(3); + + builder.ConvWithGeneralDimensions(input, filter, {1, 1}, Padding::kValid, + dnums); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); + iota_int_init_value(input_elems, 1); + auto input_r1 = Literal::CreateR1(input_elems); + auto input_r4 = input_r1->Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); + iota_int_init_value(filter_elems, 1); + auto filter_r1 = Literal::CreateR1(filter_elems); + auto filter_r4 = filter_r1->Reshape(filter_dims).ConsumeValueOrDie(); + + auto expected_r1 = Literal::CreateR1( + {static_cast(92115), static_cast(93150), static_cast(94185)}); + auto expected_r4 = expected_r1->Reshape({1, 1, 1, 3}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(*input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(*filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, *expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); } +}; - std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); - std::iota(input_elems.begin(), input_elems.end(), 1.0f); - auto input_r1 = Literal::CreateR1(input_elems); - auto input_r4 = input_r1->Reshape(input_dims).ConsumeValueOrDie(); - - std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); - std::iota(filter_elems.begin(), filter_elems.end(), 1.0f); - auto filter_r1 = Literal::CreateR1(filter_elems); - auto filter_r4 = filter_r1->Reshape(filter_dims).ConsumeValueOrDie(); - - auto expected_r1 = Literal::CreateR1({92115, 93150, 94185}); - auto expected_r4 = expected_r1->Reshape({1, 1, 1, 3}).ConsumeValueOrDie(); - - auto input_literal = client_->TransferToServer(*input_r4).ConsumeValueOrDie(); - auto filter_literal = - client_->TransferToServer(*filter_r4).ConsumeValueOrDie(); - - ComputeAndCompareLiteral(&builder, *expected_r4, - {input_literal.get(), filter_literal.get()}, - error_spec_); -} +TYPED_TEST_CASE(Convolve2D_1x3x3x5_3x3x5x5_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x3x3x5_3x3x5x5_Valid, Types) { this->RunTest(); } // Test fixture to run convolution tests with and without convolution // canonicalization enabled. @@ -501,10 +578,10 @@ XLA_TEST_P(ConvolveWithAndWithoutCanonicalization, Array2D expected_result(29, 10); expected_result.Fill(0); - ComputeAndCompare( - &builder, conv, - {*Literal::CreateFromArray(param0), *Literal::CreateFromArray(param1)}, - error_spec_); + ComputeAndCompare(&builder, conv, + {std::move(*Literal::CreateFromArray(param0)), + std::move(*Literal::CreateFromArray(param1))}, + error_spec_); } INSTANTIATE_TEST_CASE_P(ConvolveWithAndWithoutCanonicalization_Instantiation, @@ -519,67 +596,117 @@ struct Convolve1DTestParam { int64 num_windows; }; -class Convolve1D1WindowTest +class Convolve1D1WindowTestBase : public ConvolutionTest, - public ::testing::WithParamInterface {}; - -XLA_TEST_P(Convolve1D1WindowTest, Convolve1D1Window) { - ComputationBuilder builder(client_, TestName()); - int64 input_feature = GetParam().input_feature; - int64 output_feature = GetParam().output_feature; - int64 batch = GetParam().batch; - int64 num_windows = GetParam().num_windows; - int64 window_size = GetParam().window_size; - std::vector input_dims = {batch, window_size + num_windows - 1, - input_feature}; - std::vector filter_dims = {window_size, input_feature, output_feature}; - Shape input_shape = ShapeUtil::MakeShape(F32, input_dims); - Shape filter_shape = ShapeUtil::MakeShape(F32, filter_dims); - { - auto input = builder.Parameter(0, input_shape, "input"); - auto filter = builder.Parameter(1, filter_shape, "filter"); - - // Tensorflow dimension numbers for 1D convolution. - ConvolutionDimensionNumbers dnums; - dnums.set_input_batch_dimension(0); - dnums.set_output_batch_dimension(0); - dnums.add_input_spatial_dimensions(1); - dnums.add_output_spatial_dimensions(1); - dnums.set_input_feature_dimension(2); - dnums.set_output_feature_dimension(2); - dnums.add_kernel_spatial_dimensions(0); - dnums.set_kernel_input_feature_dimension(1); - dnums.set_kernel_output_feature_dimension(2); - - builder.ConvWithGeneralDimensions(input, filter, {1}, Padding::kValid, - dnums); + public ::testing::WithParamInterface { + protected: + template + void TestImpl() { + ComputationBuilder builder(client_, TestName()); + int64 input_feature = GetParam().input_feature; + int64 output_feature = GetParam().output_feature; + int64 batch = GetParam().batch; + int64 num_windows = GetParam().num_windows; + int64 window_size = GetParam().window_size; + std::vector input_dims = {batch, window_size + num_windows - 1, + input_feature}; + std::vector filter_dims = {window_size, input_feature, + output_feature}; + Shape input_shape = MakeShapeWrapper(input_dims); + Shape filter_shape = MakeShapeWrapper(filter_dims); + { + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 1D convolution. + ConvolutionDimensionNumbers dnums; + dnums.set_input_batch_dimension(0); + dnums.set_output_batch_dimension(0); + dnums.add_input_spatial_dimensions(1); + dnums.add_output_spatial_dimensions(1); + dnums.set_input_feature_dimension(2); + dnums.set_output_feature_dimension(2); + dnums.add_kernel_spatial_dimensions(0); + dnums.set_kernel_input_feature_dimension(1); + dnums.set_kernel_output_feature_dimension(2); + + builder.ConvWithGeneralDimensions(input, filter, {1}, Padding::kValid, + dnums); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1.0f)); + auto input_r1 = Literal::CreateR1(input_elems); + auto input_r3 = input_r1->Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(1.0f)); + + auto filter_r1 = Literal::CreateR1(filter_elems); + auto filter_r3 = filter_r1->Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector expect_elems(batch * output_feature * num_windows, + static_cast(window_size * input_feature)); + auto expected_r1 = Literal::CreateR1(expect_elems); + auto expected_r3 = + expected_r1->Reshape({batch, num_windows, output_feature}) + .ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(*input_r3).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(*filter_r3).ConsumeValueOrDie(); + ComputeAndCompareLiteral(&builder, *expected_r3, + {input_literal.get(), filter_literal.get()}, + error_spec_); } +}; - std::vector input_elems(ShapeUtil::ElementsIn(input_shape), 1.0); - auto input_r1 = Literal::CreateR1(input_elems); - auto input_r3 = input_r1->Reshape(input_dims).ConsumeValueOrDie(); +class Convolve1D1WindowTestFloat : public Convolve1D1WindowTestBase {}; - std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), 1.0); +XLA_TEST_P(Convolve1D1WindowTestFloat, Convolve1D1Window) { TestImpl(); } - auto filter_r1 = Literal::CreateR1(filter_elems); - auto filter_r3 = filter_r1->Reshape(filter_dims).ConsumeValueOrDie(); +INSTANTIATE_TEST_CASE_P( + Convolve1D1WindowTest_Instantiation, Convolve1D1WindowTestFloat, + ::testing::Values(Convolve1DTestParam{1, 1, 1, 1, 2}, + Convolve1DTestParam{160, 1, 1, 5, 1}, + Convolve1DTestParam{24, 1, 1, 20, 1}, + Convolve1DTestParam{30, 1, 1, 20, 1}, + Convolve1DTestParam{23, 1, 1, 20, 20}, + Convolve1DTestParam{25, 1, 1, 20, 1}, + Convolve1DTestParam{24, 1, 1, 10, 5}, + Convolve1DTestParam{160, 1, 1, 10, 1}, + Convolve1DTestParam{255, 1, 1, 3, 1}, + Convolve1DTestParam{130, 1, 1, 1, 3}, + Convolve1DTestParam{64, 1, 1, 1, 1}, + Convolve1DTestParam{128, 1, 1, 1, 1}, + Convolve1DTestParam{139, 1, 1, 128, 1}, + Convolve1DTestParam{1, 10, 10, 1, 10}, + Convolve1DTestParam{1, 10, 130, 1, 2}, + Convolve1DTestParam{1, 10, 130, 1, 1}, + Convolve1DTestParam{1, 64, 64, 1, 10}, + Convolve1DTestParam{1, 65, 65, 1, 1}, + Convolve1DTestParam{1, 128, 128, 1, 1}, + Convolve1DTestParam{128, 128, 128, 128, 1}, + Convolve1DTestParam{1, 128, 128, 1, 1}, + Convolve1DTestParam{2, 2, 2, 2, 1}, + Convolve1DTestParam{161, 1, 1, 10, 1}, + Convolve1DTestParam{900, 1, 1, 10, 1}, + Convolve1DTestParam{640, 3, 3, 128, 1}) - std::vector expect_elems(batch * output_feature * num_windows, - window_size * input_feature); - auto expected_r1 = Literal::CreateR1(expect_elems); - auto expected_r3 = expected_r1->Reshape({batch, num_windows, output_feature}) - .ConsumeValueOrDie(); +); - auto input_literal = client_->TransferToServer(*input_r3).ConsumeValueOrDie(); - auto filter_literal = - client_->TransferToServer(*filter_r3).ConsumeValueOrDie(); - ComputeAndCompareLiteral(&builder, *expected_r3, - {input_literal.get(), filter_literal.get()}, - error_spec_); +#if (XLA_TEST_BACKEND_GPU || XLA_TEST_BACKEND_CPU) +class Convolve1D1WindowTestHalf : public Convolve1D1WindowTestBase {}; + +// TODO(b/72509305): Enable half data type tests for CPU. +XLA_TEST_P(Convolve1D1WindowTestHalf, + DISABLED_ON_CPU_PARALLEL(DISABLED_ON_CPU(Convolve1D1Window))) { + TestImpl(); } INSTANTIATE_TEST_CASE_P( - Convolve1D1WindowTest_Instantiation, Convolve1D1WindowTest, + Convolve1D1WindowTest_Instantiation, Convolve1D1WindowTestHalf, ::testing::Values(Convolve1DTestParam{1, 1, 1, 1, 2}, Convolve1DTestParam{160, 1, 1, 5, 1}, Convolve1DTestParam{24, 1, 1, 20, 1}, @@ -592,7 +719,11 @@ INSTANTIATE_TEST_CASE_P( Convolve1DTestParam{130, 1, 1, 1, 3}, Convolve1DTestParam{64, 1, 1, 1, 1}, Convolve1DTestParam{128, 1, 1, 1, 1}, + // TODO(b/72566306): the following three tests fail on CPU + // backend due to result miscompare. Convolve1DTestParam{139, 1, 1, 128, 1}, + Convolve1DTestParam{640, 3, 3, 128, 1}, + Convolve1DTestParam{900, 1, 1, 10, 1}, Convolve1DTestParam{1, 10, 10, 1, 10}, Convolve1DTestParam{1, 10, 130, 1, 2}, Convolve1DTestParam{1, 10, 130, 1, 1}, @@ -602,11 +733,33 @@ INSTANTIATE_TEST_CASE_P( Convolve1DTestParam{128, 128, 128, 128, 1}, Convolve1DTestParam{1, 128, 128, 1, 1}, Convolve1DTestParam{2, 2, 2, 2, 1}, - Convolve1DTestParam{161, 1, 1, 10, 1}, - Convolve1DTestParam{900, 1, 1, 10, 1}, - Convolve1DTestParam{640, 3, 3, 128, 1}) + Convolve1DTestParam{161, 1, 1, 10, 1}) ); +#endif + +TEST_F(ConvolutionTest, Convolve_bf16_1x1x1x2_1x1x1x2_Valid) { + ComputationBuilder builder(client_, TestName()); + Shape input_shape = ShapeUtil::MakeShape(BF16, {1, 1, 1, 2}); + Shape filter_shape = ShapeUtil::MakeShape(BF16, {1, 1, 1, 2}); + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + auto conv = builder.Conv(input, filter, {1, 1}, Padding::kValid); + + Array4D input_data(1, 1, 1, 2); + input_data.FillWithYX(Array2D({ + {bfloat16(1), bfloat16(2)}, + })); + Array4D filter_data(1, 1, 1, 2); + filter_data.FillWithYX(Array2D({ + {bfloat16(5), bfloat16(6)}, + })); + + ComputeAndCompare(&builder, conv, + {std::move(*Literal::CreateFromArray(input_data)), + std::move(*Literal::CreateFromArray(filter_data))}, + error_spec_); +} } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/copy_test.cc b/tensorflow/compiler/xla/tests/copy_test.cc index d64bf0aa5bd5e9d6213ea07b3da3305a9c621c65..ece7c3b05e7fafa299db7f9cbf50610c8204f95e 100644 --- a/tensorflow/compiler/xla/tests/copy_test.cc +++ b/tensorflow/compiler/xla/tests/copy_test.cc @@ -40,7 +40,7 @@ class CopyOpTest : public HloTestBase { void TestCopyOp(const Literal& literal) { auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(MakeUnique(literal))); + HloInstruction::CreateConstant(literal.CloneToUnique())); builder.AddInstruction(HloInstruction::CreateUnary( constant->shape(), HloOpcode::kCopy, constant)); auto computation = builder.Build(); @@ -132,7 +132,8 @@ XLA_TEST_F(CopyOpTest, CopyConstantR2DifferentLayouts) { std::unique_ptr literal = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); // Reverse the minor-to-major order of the literal. - Layout* literal_layout = literal->mutable_shape()->mutable_layout(); + Layout* literal_layout = + literal->mutable_shape_do_not_use()->mutable_layout(); ASSERT_EQ(2, literal_layout->minor_to_major_size()); literal_layout->mutable_minor_to_major()->SwapElements(0, 1); diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc index cc683701e6305510d202721fe645310f1009081c..6b0c04c2c083bbfce267dd92d24ef15c06186d26 100644 --- a/tensorflow/compiler/xla/tests/dot_operation_test.cc +++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc @@ -520,9 +520,39 @@ XLA_TEST_F(DotOperationTest, BatchMatMul) { ComputeAndCompareR4( &builder, - /*expected=*/{{{{1300, 2400}, {13, 24}}, {{11400, 13600}, {114, 136}}}, - {{{42900, 79200}, {429, 792}}, - {{250800, 299200}, {2508, 2992}}}}, + /*expected=*/ + {{{{1300, 2400}, {13, 24}}, {{11400, 13600}, {114, 136}}}, + {{{42900, 79200}, {429, 792}}, {{250800, 299200}, {2508, 2992}}}}, + {x_data.get(), y_data.get()}, error_spec_); +} + +XLA_TEST_F(DotOperationTest, GeneralMatMul) { + ComputationBuilder builder(client_, TestName()); + auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {2, 2, 2}), "x"); + auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {2, 2, 2}), "y"); + + DotDimensionNumbers dnums; + dnums.add_lhs_contracting_dimensions(2); + dnums.add_rhs_contracting_dimensions(1); + dnums.add_lhs_batch_dimensions(0); + dnums.add_rhs_batch_dimensions(0); + + auto out = builder.DotGeneral(x, y, dnums); + + auto x_data = client_ + ->TransferToServer(*Literal::CreateR3( + {{{1.0, 2.0}, {3.0, 4.0}}, {{5.0, 6.0}, {7.0, 8.0}}})) + .ConsumeValueOrDie(); + + auto y_data = client_ + ->TransferToServer(*Literal::CreateR3( + {{{1.0, 0.0}, {0.0, 1.0}}, {{1.0, 0.0}, {0.0, 1.0}}})) + .ConsumeValueOrDie(); + + ComputeAndCompareR3( + &builder, + /*expected=*/ + {{{1.0, 2.0}, {3.0, 4.0}}, {{5.0, 6.0}, {7.0, 8.0}}}, {x_data.get(), y_data.get()}, error_spec_); } diff --git a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc index 59be32a8ff584a6189302a0835ba74b2e08956b1..877dc7db0eec229a7119b3627f177a33ed0d971b 100644 --- a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc +++ b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc @@ -51,12 +51,16 @@ class DynamicSliceTest : public ClientLibraryTestBase { RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {2}, {3}, {2, 3, 4}); // Slice at dimension boundaries. RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {5}, {3}, {5, 6, 7}); - // Slice at dimension boundaries, but with sizes that cause indices to wrap. - RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {6}, {4}, {6, 7, 0, 1}); // Zero element slice. RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {2}, {0}, {}); } + template + void TestR1Wrap() { + // Slice at dimension boundaries, but with sizes that cause indices to wrap. + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {6}, {4}, {6, 7, 0, 1}); + } + template void TestR2() { // Slice at dimension start. @@ -68,15 +72,19 @@ class DynamicSliceTest : public ClientLibraryTestBase { // Slice at dimension boundaries. RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {1, 1}, {2, 1}, {{5}, {8}}); - // Slice at dimension boundaries, but with sizes that cause indices to wrap. - RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {1, 1}, {3, 3}, - {{5, 6, 4}, {8, 9, 7}, {2, 3, 1}}); // Zero element slice: 2x0. RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {0, 0}, {2, 0}, {{}, {}}); // Zero element slice: 0x2. RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {0, 0}, {0, 2}, - Array2D(0, 2)); + Array2D(0, 2)); + } + + template + void TestR2Wrap() { + // Slice at dimension boundaries, but with sizes that cause indices to wrap. + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {1, 1}, {3, 3}, + {{5, 6, 4}, {8, 9, 7}, {2, 3, 1}}); } template @@ -97,85 +105,119 @@ class DynamicSliceTest : public ClientLibraryTestBase { {{7, 8}, {9, 10}, {11, 12}}}, {0, 1, 1}, {2, 2, 1}, {{{4}, {6}}, {{10}, {12}}}); + // clang-format on + } + template + void TestR3Wrap() { // Slice at dimension boundaries, but with sizes that cause indices to wrap. RunR3( {{{1, 2}, {3, 4}, {5, 6}}, {{7, 8}, {9, 10}, {11, 12}}}, {0, 2, 1}, {2, 1, 2}, {{{6, 5}}, {{12, 11}}}); - - // clang-format on } template - void RunR1(tensorflow::gtl::ArraySlice input_values, + void RunR1(tensorflow::gtl::ArraySlice input_values_int, const std::vector slice_starts, const std::vector& slice_sizes, - tensorflow::gtl::ArraySlice expected_values) { + tensorflow::gtl::ArraySlice expected_values_int) { + // bfloat16 has explicit constructors, so it does not implicitly convert the + // way built-in types do, which is why we can't take the parameter as an + // ArraySlice. We also can't convert it to a vector, because + // vector is special so that it cannot be an ArraySlice, which + // is what the code below wants. So instead we do this. + Literal input_values = + std::move(*Literal::CreateR1(input_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal expected_values = + std::move(*Literal::CreateR1(expected_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR1(input_values); + auto input = builder.ConstantLiteral(input_values); builder.DynamicSlice(input, starts, slice_sizes); // Run computation and compare against expected values. - ComputeAndCompareR1(&builder, expected_values, {start_data.get()}); + ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); } template - void RunR2(const Array2D& input_values, + void RunR2(const Array2D& input_values_int, const std::vector slice_starts, const std::vector& slice_sizes, - const Array2D& expected_values) { + const Array2D& expected_values_int) { + Literal input_values = + std::move(*Literal::CreateR2FromArray2D(input_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal expected_values = + std::move(*Literal::CreateR2FromArray2D(expected_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR2FromArray2D(input_values); + auto input = builder.ConstantLiteral(input_values); builder.DynamicSlice(input, starts, slice_sizes); // Run computation and compare against expected values. - ComputeAndCompareR2(&builder, expected_values, {start_data.get()}); + ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); } template - void RunR3(const Array3D& input_values, + void RunR3(const Array3D& input_values_int, const std::vector slice_starts, const std::vector& slice_sizes, - const Array3D& expected_values) { + const Array3D& expected_values_int) { + Literal input_values = + std::move(*Literal::CreateR3FromArray3D(input_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal expected_values = + std::move(*Literal::CreateR3FromArray3D(expected_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR3FromArray3D(input_values); + auto input = builder.ConstantLiteral(input_values); builder.DynamicSlice(input, starts, slice_sizes); // Run computation and compare against expected values. - ComputeAndCompareR3(&builder, expected_values, {start_data.get()}); + ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); } }; +XLA_TEST_F(DynamicSliceTest, Int32R1BF16) { TestR1(); } XLA_TEST_F(DynamicSliceTest, Int32R1) { TestR1(); } - +XLA_TEST_F(DynamicSliceTest, Int32R1Wrap) { TestR1Wrap(); } XLA_TEST_F(DynamicSliceTest, Int64R1) { TestR1(); } - XLA_TEST_F(DynamicSliceTest, UInt64R1) { TestR1(); } -XLA_TEST_F(DynamicSliceTest, Int32R2) { TestR2(); } - +XLA_TEST_F(DynamicSliceTest, Int32R2BF16) { TestR2(); } +XLA_TEST_F(DynamicSliceTest, Int32R2) { TestR2(); } +XLA_TEST_F(DynamicSliceTest, Int32R2Wrap) { TestR2Wrap(); } XLA_TEST_F(DynamicSliceTest, Int64R2) { TestR2(); } - XLA_TEST_F(DynamicSliceTest, UInt64R2) { TestR2(); } -XLA_TEST_F(DynamicSliceTest, Int32R3) { TestR3(); } - +XLA_TEST_F(DynamicSliceTest, Int32R3BF16) { TestR3(); } +XLA_TEST_F(DynamicSliceTest, Int32R3) { TestR3(); } +XLA_TEST_F(DynamicSliceTest, Int32R3Wrap) { TestR3Wrap(); } XLA_TEST_F(DynamicSliceTest, Int64R3) { TestR3(); } - XLA_TEST_F(DynamicSliceTest, UInt64R3) { TestR3(); } XLA_TEST_F(DynamicSliceTest, Int32R1Pred) { @@ -213,7 +255,7 @@ XLA_TEST_F(DynamicSliceTest, Int32R2Pred) { // Zero element slice: 0x2. RunR2( {{true, false, true}, {false, false, true}, {true, true, false}}, {0, 0}, - {0, 2}, Array2D(0, 2)); + {0, 2}, Array2D(0, 2)); } XLA_TEST_F(DynamicSliceTest, Int32R3Pred) { @@ -300,107 +342,154 @@ class DynamicUpdateSliceTest : public ClientLibraryTestBase { } template - void RunR1(tensorflow::gtl::ArraySlice input_values, - tensorflow::gtl::ArraySlice update_values, + void RunR1(tensorflow::gtl::ArraySlice input_values_int, + tensorflow::gtl::ArraySlice update_values_int, const std::vector slice_starts, - tensorflow::gtl::ArraySlice expected_values) { + tensorflow::gtl::ArraySlice expected_values_int) { + Literal input_values = + std::move(*Literal::CreateR1(input_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal update_values = + std::move(*Literal::CreateR1(update_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal expected_values = + std::move(*Literal::CreateR1(expected_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR1(input_values); - auto update = builder.ConstantR1(update_values); + auto input = builder.ConstantLiteral(input_values); + auto update = builder.ConstantLiteral(update_values); builder.DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareR1(&builder, expected_values, {start_data.get()}); + ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); } template - void RunR2(const Array2D& input_values, - const Array2D& update_values, + void RunR2(const Array2D& input_values_int, + const Array2D& update_values_int, const std::vector slice_starts, - const Array2D& expected_values) { + const Array2D& expected_values_int) { + Literal input_values = + std::move(*Literal::CreateR2FromArray2D(input_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal update_values = + std::move(*Literal::CreateR2FromArray2D(update_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal expected_values = + std::move(*Literal::CreateR2FromArray2D(expected_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR2FromArray2D(input_values); - auto update = builder.ConstantR2FromArray2D(update_values); + auto input = builder.ConstantLiteral(input_values); + auto update = builder.ConstantLiteral(update_values); builder.DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareR2(&builder, expected_values, {start_data.get()}); + ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); } template - void RunR3(const Array3D& input_values, - const Array3D& update_values, + void RunR3(const Array3D& input_values_int, + const Array3D& update_values_int, const std::vector slice_starts, - const Array3D& expected_values) { + const Array3D& expected_values_int) { + Literal input_values = + std::move(*Literal::CreateR3FromArray3D(input_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal update_values = + std::move(*Literal::CreateR3FromArray3D(update_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + Literal expected_values = + std::move(*Literal::CreateR3FromArray3D(expected_values_int) + ->Convert(primitive_util::NativeToPrimitiveType()) + .ValueOrDie()); + ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR3FromArray3D(input_values); - auto update = builder.ConstantR3FromArray3D(update_values); + auto input = builder.ConstantLiteral(input_values); + auto update = builder.ConstantLiteral(update_values); builder.DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareR3(&builder, expected_values, {start_data.get()}); + ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); } + template void RunR3Contiguous(std::vector operand_shape, int32 index, int32 size) { +#ifdef XLA_TEST_BACKEND_CPU_PARALLEL + // TODO(b/71820067): The CPU parallel backend failed for this on 2018-01-10. + if (std::is_same::value) { + return; + } +#endif + const int32 kSeq = operand_shape[0]; const int32 kBatch = operand_shape[1]; const int32 kDim = operand_shape[2]; - Array3D input_values(kSeq, kBatch, kDim); - Array3D update_values(size, kBatch, kDim); - Array3D expected_values(kSeq, kBatch, kDim); + Array3D input_values(kSeq, kBatch, kDim); + Array3D update_values(size, kBatch, kDim); + Array3D expected_values(kSeq, kBatch, kDim); - input_values.FillIota(0); - float val = 1000; - update_values.FillIota(val); + input_values.FillIota(static_cast(0)); + T value = static_cast(10); + update_values.FillIota(static_cast(value)); // TODO(b/34128753) Expected values may vary depending on backend when // the update wraps. According to documentation, the results are technically // implementation specific where the update is out of bounds, and hence // we don't really know what to pass into ComputeAndCompareR3. - expected_values.FillIota(0); + expected_values.FillIota(static_cast(0)); for (int i = 0; i < size; i++) { for (int j = 0; j < kBatch; j++) { for (int k = 0; k < kDim; k++) { - expected_values((index + i) % kSeq, j, k) = val++; + expected_values((index + i) % kSeq, j, k) = value++; } } } if (VLOG_IS_ON(1)) { - DumpArray("input", input_values); - DumpArray("update", update_values); - DumpArray("expected", expected_values); + DumpArray("input", input_values); + DumpArray("update", update_values); + DumpArray("expected", expected_values); } // Build dynamic slice computation. ComputationBuilder builder(client_, TestName()); // Initialize and transfer input parameter. ComputationDataHandle input; - std::unique_ptr input_data = CreateR3Parameter( - input_values, 0, "input_values", &builder, &input); + std::unique_ptr input_data = + CreateR3Parameter(input_values, 0, "input_values", &builder, &input); // Initialize and transfer update parameter. ComputationDataHandle update; - std::unique_ptr update_data = CreateR3Parameter( + std::unique_ptr update_data = CreateR3Parameter( update_values, 1, "update_values", &builder, &update); auto starts = builder.ConstantR1({index, 0, 0}); builder.DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareR3(&builder, expected_values, - {input_data.get(), update_data.get()}, - ErrorSpec(0.000001)); + ComputeAndCompareR3(&builder, expected_values, + {input_data.get(), update_data.get()}, + ErrorSpec(0.000001)); } template @@ -411,28 +500,35 @@ class DynamicUpdateSliceTest : public ClientLibraryTestBase { } }; +// TODO(b/71820067): The CPU parallel backend failed for this on 2018-01-10. +XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_CPU_PARALLEL(Int32R1BF16)) { + TestR1(); +} XLA_TEST_F(DynamicUpdateSliceTest, Int32R1) { TestR1(); } - XLA_TEST_F(DynamicUpdateSliceTest, Int64R1) { TestR1(); } - XLA_TEST_F(DynamicUpdateSliceTest, UInt64R1) { TestR1(); } +// TODO(b/71820067): The CPU parallel backend failed for this on 2018-01-10. +XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_CPU_PARALLEL(Int32R2BF16)) { + TestR2(); +} XLA_TEST_F(DynamicUpdateSliceTest, Int32R2) { TestR2(); } - XLA_TEST_F(DynamicUpdateSliceTest, Int64R2) { TestR2(); } - XLA_TEST_F(DynamicUpdateSliceTest, UInt64R2) { TestR2(); } +// TODO(b/71820067): The CPU parallel backend failed for this on 2018-01-10. +XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_CPU_PARALLEL(Int32R3BF16)) { + TestR3(); +} XLA_TEST_F(DynamicUpdateSliceTest, Int32R3) { TestR3(); } - XLA_TEST_F(DynamicUpdateSliceTest, Int64R3) { TestR3(); } - XLA_TEST_F(DynamicUpdateSliceTest, UInt64R3) { TestR3(); } +XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_CPU_PARALLEL(Int32WrapBF16)) { + TestWrap(); +} XLA_TEST_F(DynamicUpdateSliceTest, Int32Wrap) { TestWrap(); } - XLA_TEST_F(DynamicUpdateSliceTest, Int64Wrap) { TestWrap(); } - XLA_TEST_F(DynamicUpdateSliceTest, UInt64Wrap) { TestWrap(); } XLA_TEST_F(DynamicUpdateSliceTest, Int32R1Pred) { @@ -498,36 +594,70 @@ XLA_TEST_F(DynamicUpdateSliceTest, Int32R3Pred) { XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousSingleElement) { // Single element, no wrap. std::vector operand_shape({4, 5, 2}); - RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/1); + RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/1); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousSingleElementBF16) { + // Single element, no wrap. + std::vector operand_shape({4, 5, 2}); + RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/1); } XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousMultipleElements) { // Multiple element, no wrap. std::vector operand_shape({4, 5, 2}); - RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/2); + RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/2); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousMultipleElementsBF16) { + // Multiple element, no wrap. + std::vector operand_shape({4, 5, 2}); + RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/2); } XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousMultipleWrapping) { // Multiple element, wrapping. std::vector operand_shape({4, 5, 2}); - RunR3Contiguous(operand_shape, /*index=*/3, /*size=*/2); + RunR3Contiguous(operand_shape, /*index=*/3, /*size=*/2); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousMultipleWrappingBF16) { + // Multiple element, wrapping. + std::vector operand_shape({4, 5, 2}); + RunR3Contiguous(operand_shape, /*index=*/3, /*size=*/2); } XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousTooLarge) { // Multiple element, update size larger than operand. std::vector operand_shape({4, 5, 2}); - RunR3Contiguous(operand_shape, /*index=*/5, /*size=*/2); + RunR3Contiguous(operand_shape, /*index=*/5, /*size=*/2); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousTooLargeBF16) { + // Multiple element, update size larger than operand. + std::vector operand_shape({4, 5, 2}); + RunR3Contiguous(operand_shape, /*index=*/5, /*size=*/2); } XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousUnaligned) { std::vector operand_shape({3, 123, 247}); - RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/1); + RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/1); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousUnalignedBF16) { + std::vector operand_shape({3, 123, 247}); + RunR3Contiguous(operand_shape, /*index=*/1, /*size=*/1); } // TODO(b/34134076) Disabled on GPU 2016-01-06 due to out-of-memory error. XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_GPU(R3ContiguousLarger)) { std::vector operand_shape({32, 128, 1024}); - RunR3Contiguous(operand_shape, /*index=*/7, /*size=*/1); + RunR3Contiguous(operand_shape, /*index=*/7, /*size=*/1); +} + +XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_GPU(R3ContiguousLargerBF16)) { + std::vector operand_shape({32, 128, 1024}); + RunR3Contiguous(operand_shape, /*index=*/7, /*size=*/1); } void BM_DynamicSlice(int num_iters) { diff --git a/tensorflow/compiler/xla/tests/execution_profile_test.cc b/tensorflow/compiler/xla/tests/execution_profile_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..644cbbf40f296eb2a574ae568b4f32aa3d0bd12f --- /dev/null +++ b/tensorflow/compiler/xla/tests/execution_profile_test.cc @@ -0,0 +1,71 @@ +/* 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/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/client/global_data.h" +#include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { + +class ExecutionProfileTest : public ClientLibraryTestBase {}; + +XLA_TEST_F(ExecutionProfileTest, + DISABLED_ON_CPU_PARALLEL(ExecuteWithExecutionProfile)) { + Shape shape = ShapeUtil::MakeShape(F32, {256, 256}); + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr input, + client_->TransferToServer( + *Literal::CreateR2F32Linspace(1e0, 1e5, 256, 256))); + + ComputationBuilder b(client_, TestName() + ".add"); + b.Dot(b.Parameter(0, shape, "param_0"), b.Parameter(1, shape, "param_1")); + TF_ASSERT_OK_AND_ASSIGN(Computation dot_product, b.Build()); + + ExecutionProfile execution_profile; + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr data, + client_->Execute(dot_product, {input.get(), input.get()}, + &execution_options_, &execution_profile)); + + VLOG(3) << "execution_profile.compute_cycle_count() = " + << execution_profile.compute_cycle_count(); + VLOG(3) << "execution_profile.compute_and_transfer_time_ns() = " + << execution_profile.compute_and_transfer_time_ns(); + VLOG(3) << "execution_profile.compute_time_ns() = " + << execution_profile.compute_time_ns(); + + bool hlo_profiling_enabled = + execution_options_.debug_options().xla_hlo_profile(); + + // If HLO profiling is enabled we always expect cycle count to be populated. + // If HLO profiling is disabled then depending on the backend the cycle count + // may or may not be populated. + if (hlo_profiling_enabled) { + EXPECT_GT(execution_profile.compute_cycle_count(), 0); + } + + EXPECT_GT(execution_profile.compute_and_transfer_time_ns(), 0); + EXPECT_GT(execution_profile.compute_time_ns(), 0); + + TF_ASSERT_OK_AND_ASSIGN(auto computed, client_->Transfer(*data, &shape)); + (void)computed; +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc b/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..6fe7737de7af349dca2931b52d62dbc03b14e0b3 --- /dev/null +++ b/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc @@ -0,0 +1,128 @@ +/* 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/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/tests/literal_test_util.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/core/lib/core/casts.h" + +namespace xla { +namespace { +class ExhaustiveF32ElementwiseOpTest + : public ClientLibraryTestBase, + public ::testing::WithParamInterface> { + protected: + ErrorSpec error_spec_{0.0001, 0.0001, /*relaxed_nans=*/true}; + + template + void ExhaustivelyTestF32Op(EnqueueOpTy enqueue_op, + float (*evaluate_op)(float), + std::pair known_incorrect_range) { + int64 begin, end; + std::tie(begin, end) = GetParam(); + int64 input_size = end - begin; + LOG(INFO) << "Checking range [" << begin << ", " << end << ")"; + + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr input_literal = + Literal::CreateFromDimensions(F32, {input_size}); + for (int64 i = begin; i < end; i++) { + if (i >= known_incorrect_range.first && + i < known_incorrect_range.second) { + // If the operation is known to be buggy on a specific input clamp that + // input to 0 under the assumption that the op is at least correct on 0. + input_literal->Set({i - begin}, 0.0f); + } else { + input_literal->Set({i - begin}, tensorflow::bit_cast(i)); + } + } + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, + client_->TransferToServer(*input_literal)); + + auto input = builder.Parameter(0, input_literal->shape(), "input"); + enqueue_op(&builder, input); + + std::vector expected_result; + expected_result.reserve(input_size); + for (int64 i = 0; i < input_size; i++) { + expected_result.push_back(evaluate_op(input_literal->Get({i}))); + } + + ComputeAndCompareR1(&builder, expected_result, {input_data.get()}, + error_spec_); + } +}; + +XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, LogF32) { +#ifdef XLA_TEST_BACKEND_CPU + // TODO(b/73141998): The vectorized Log implementation gives results outside + // our error spec in this range (these numbers are bitwise representations of + // floats expressed as a zero extended int64): + std::pair known_incorrect_range = {1, 8315654}; +#else + std::pair known_incorrect_range = {0, 0}; +#endif + + ExhaustivelyTestF32Op( + [](ComputationBuilder* builder, const ComputationDataHandle& input) { + builder->Log(input); + }, + std::log, known_incorrect_range); +} + +XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, ExpF32) { +#ifdef XLA_TEST_BACKEND_CPU + // TODO(b/73142289): The vectorized Exp implementation gives results outside + // our error spec in this range (these numbers are bitwise representations of + // floats expressed as a zero extended int64): + std::pair known_incorrect_range = {1107296256 + 11583654, + 1107296256 + 11629080}; +#else + std::pair known_incorrect_range = {0, 0}; +#endif + + ExhaustivelyTestF32Op( + [](ComputationBuilder* builder, const ComputationDataHandle& input) { + builder->Exp(input); + }, + std::exp, known_incorrect_range); +} + +XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, TanhF32) { + ExhaustivelyTestF32Op( + [](ComputationBuilder* builder, const ComputationDataHandle& input) { + builder->Tanh(input); + }, + std::tanh, /*known_incorrect_range=*/{0, 0}); +} + +std::vector> CreateExhaustiveParameters() { + // We break up the 2^32-element space into small'ish chunks to keep peak + // memory usage low. + std::vector> result; + const int64 step = 1 << 25; + for (int64 i = 0; i < (1l << 32); i += step) { + result.push_back({i, i + step}); + } + return result; +} + +INSTANTIATE_TEST_CASE_P(ExhaustiveF32ElementwiseOpTestInstance, + ExhaustiveF32ElementwiseOpTest, + ::testing::ValuesIn(CreateExhaustiveParameters())); +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/filecheck.h b/tensorflow/compiler/xla/tests/filecheck.h index 493ff7414bde31b18a39a5098925d9c991529b00..3830d5a44d2ca483fbe839231b0136d13033b48b 100644 --- a/tensorflow/compiler/xla/tests/filecheck.h +++ b/tensorflow/compiler/xla/tests/filecheck.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ -#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ +#ifndef TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ +#define TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ #include @@ -30,4 +30,4 @@ StatusOr RunFileCheck(const string& input, const string& pattern); } // namespace xla -#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ +#endif // TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ diff --git a/tensorflow/compiler/xla/tests/half_test.cc b/tensorflow/compiler/xla/tests/half_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..ec2f49d43bd8cee84c6b0abe1892e8b2278eefeb --- /dev/null +++ b/tensorflow/compiler/xla/tests/half_test.cc @@ -0,0 +1,257 @@ +/* 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 +#include + +#include "tensorflow/compiler/xla/client/computation.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/compiler/xla/tests/test_utils.h" + +// Tests the handling of the basic mathematics operations with F16 operands. + +namespace xla { +namespace { + +class HalfTestBase : public ClientLibraryTestBase { + protected: + const ErrorSpec error_spec_{0.001, 0.001}; + // Number of elements in the input buffers. + static const int kNumElements = 4; +}; + +using UnaryBuildFuncTy = + std::function; + +struct UnaryOpTestParam { + std::function compute_func; + UnaryBuildFuncTy build_func; +}; + +class UnaryOpTest : public HalfTestBase, + public ::testing::WithParamInterface {}; + +XLA_TEST_P(UnaryOpTest, Ops) { + std::vector x({half(1.4), half(-2.3), half(3.2), half(-4.1)}); + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle x_opnd; + auto x_data = CreateR1Parameter(x, /*parameter_number=*/0, "x", + &builder, &x_opnd); + + std::function compute_func = GetParam().compute_func; + std::vector expected; + for (int64 i = 0; i < x.size(); ++i) { + expected.push_back(compute_func(x[i])); + } + + UnaryBuildFuncTy build_func = GetParam().build_func; + build_func(&builder, x_opnd); + + ComputeAndCompareR1(&builder, expected, {x_data.get()}, error_spec_); +} + +half sign_imp(half value) { + const float x(std::move(value)); + return half((x < .0) ? -1 : (x > .0)); +} + +half round_imp(half value) { + return half(round(static_cast(std::move(value)))); +} + +INSTANTIATE_TEST_CASE_P( + half, UnaryOpTest, + ::testing::Values(UnaryOpTestParam{[](half x) { return abs(x); }, + &ComputationBuilder::Abs}, + UnaryOpTestParam{[](half x) { return round_imp(x); }, + &ComputationBuilder::Round}, + UnaryOpTestParam{[](half x) { return ceil(x); }, + &ComputationBuilder::Ceil}, + UnaryOpTestParam{[](half x) { return cos(x); }, + &ComputationBuilder::Cos}, + UnaryOpTestParam{[](half x) { return exp(x); }, + &ComputationBuilder::Exp}, + UnaryOpTestParam{[](half x) { return floor(x); }, + &ComputationBuilder::Floor}, + UnaryOpTestParam{[](half x) { return log(x); }, + &ComputationBuilder::Log}, + UnaryOpTestParam{[](half x) { return -x; }, + &ComputationBuilder::Neg}, + UnaryOpTestParam{[](half x) { return sign_imp(x); }, + &ComputationBuilder::Sign}, + UnaryOpTestParam{[](half x) { return sin(x); }, + &ComputationBuilder::Sin}, + UnaryOpTestParam{[](half x) { return tanh(x); }, + &ComputationBuilder::Tanh} + + )); + +struct UnaryPredTestParam { + std::function compute_func; + UnaryBuildFuncTy build_func; +}; + +class UnaryPredTest : public HalfTestBase, + public ::testing::WithParamInterface { +}; + +XLA_TEST_P(UnaryPredTest, Ops) { + std::vector x({half(1.4), half(-2.3), half(3.2), half(-4.1)}); + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle x_opnd; + auto x_data = CreateR1Parameter(x, /*parameter_number=*/0, "x", + &builder, &x_opnd); + + std::function compute_func = GetParam().compute_func; + CHECK_EQ(kNumElements, x.size()); + bool expected[kNumElements]; + for (int64 i = 0; i < x.size(); ++i) { + expected[i] = compute_func(x[i]); + } + + UnaryBuildFuncTy build_func = GetParam().build_func; + build_func(&builder, x_opnd); + + ComputeAndCompareR1(&builder, expected, {x_data.get()}); +} + +INSTANTIATE_TEST_CASE_P(half, UnaryPredTest, + ::testing::Values(UnaryPredTestParam{ + [](half x) { return isfinite(x); }, + &ComputationBuilder::IsFinite})); + +using BinaryBuildFuncTy = std::function)>; + +struct BinaryOpTestParam { + std::function compute_func; + BinaryBuildFuncTy build_func; +}; + +class BinaryOpTest : public HalfTestBase, + public ::testing::WithParamInterface {}; + +XLA_TEST_P(BinaryOpTest, Ops) { + std::vector x({half(1.0), half(2.0), half(3.0), half(-4.0)}); + std::vector y({half(0.4), half(-0.3), half(0.2), half(0.1)}); + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle x_opnd; + auto x_data = CreateR1Parameter(x, /*parameter_number=*/0, "x", + &builder, &x_opnd); + + ComputationDataHandle y_opnd; + auto y_data = CreateR1Parameter(y, /*parameter_number=*/1, "y", + &builder, &y_opnd); + + std::function compute_func = GetParam().compute_func; + std::vector expected; + for (int64 i = 0; i < x.size(); ++i) { + expected.push_back(compute_func(x[i], y[i])); + } + + BinaryBuildFuncTy build_func = GetParam().build_func; + build_func(&builder, x_opnd, y_opnd, {}); + + ComputeAndCompareR1(&builder, expected, {x_data.get(), y_data.get()}, + error_spec_); +} + +half atan2_imp(half x, half y) { + return half(atan2(static_cast(std::move(x)), + static_cast(std::move(y)))); +} + +INSTANTIATE_TEST_CASE_P( + half, BinaryOpTest, + ::testing::Values( + BinaryOpTestParam{[](half x, half y) { return x + y; }, + &ComputationBuilder::Add}, + BinaryOpTestParam{[](half x, half y) { return atan2_imp(x, y); }, + &ComputationBuilder::Atan2}, + BinaryOpTestParam{[](half x, half y) { return x / y; }, + &ComputationBuilder::Div}, + BinaryOpTestParam{[](half x, half y) { return max(x, y); }, + &ComputationBuilder::Max}, + BinaryOpTestParam{[](half x, half y) { return min(x, y); }, + &ComputationBuilder::Min}, + BinaryOpTestParam{[](half x, half y) { return x * y; }, + &ComputationBuilder::Mul}, + BinaryOpTestParam{[](half x, half y) { return pow(x, y); }, + &ComputationBuilder::Pow}, + BinaryOpTestParam{[](half x, half y) { return x - y; }, + &ComputationBuilder::Sub} + + )); + +struct BinaryPredTestParam { + std::function compute_func; + BinaryBuildFuncTy build_func; +}; + +class BinaryPredTest + : public HalfTestBase, + public ::testing::WithParamInterface {}; + +XLA_TEST_P(BinaryPredTest, Ops) { + std::vector x({half(1.0), half(2.0), half(0.2), half(-4.0)}); + std::vector y({half(0.4), half(-0.3), half(0.2), half(0.1)}); + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle x_opnd; + auto x_data = CreateR1Parameter(x, /*parameter_number=*/0, "x", + &builder, &x_opnd); + + ComputationDataHandle y_opnd; + auto y_data = CreateR1Parameter(y, /*parameter_number=*/1, "y", + &builder, &y_opnd); + + std::function compute_func = GetParam().compute_func; + CHECK_EQ(kNumElements, x.size()); + bool expected[kNumElements]; + for (int64 i = 0; i < x.size(); ++i) { + expected[i] = compute_func(x[i], y[i]); + } + + BinaryBuildFuncTy build_func = GetParam().build_func; + build_func(&builder, x_opnd, y_opnd, {}); + + ComputeAndCompareR1(&builder, expected, {x_data.get(), y_data.get()}); +} + +INSTANTIATE_TEST_CASE_P( + half, BinaryPredTest, + ::testing::Values(BinaryPredTestParam{[](half x, half y) { return x == y; }, + &ComputationBuilder::Eq}, + BinaryPredTestParam{[](half x, half y) { return x != y; }, + &ComputationBuilder::Ne}, + BinaryPredTestParam{[](half x, half y) { return x >= y; }, + &ComputationBuilder::Ge}, + BinaryPredTestParam{[](half x, half y) { return x > y; }, + &ComputationBuilder::Gt}, + BinaryPredTestParam{[](half x, half y) { return x <= y; }, + &ComputationBuilder::Le}, + BinaryPredTestParam{[](half x, half y) { return x < y; }, + &ComputationBuilder::Lt} + + )); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc index a27e0f2c106c2ffa2ba108e1963e7111fd347482..9f5806c5e16c30cf198027cffab5f78c315cb957 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc @@ -91,9 +91,7 @@ HloTestBase::HloTestBase() HloTestBase::HloTestBase(se::Platform* test_platform, se::Platform* reference_platform) : test_runner_(test_platform), reference_runner_(reference_platform) { - hlo_verifier_ = MakeUnique([this](const Shape& shape) { - return backend().transfer_manager()->GetByteSizeRequirement(shape); - }); + hlo_verifier_ = MakeUnique(); } /* static */ @@ -232,7 +230,7 @@ template const string& filename, const tensorflow::gtl::optional& error, const std::function& reference_preprocessor) { auto module_or_status = - HloRunner::ReadModule(filename, GetDebugOptionsForTest()); + HloRunner::ReadModuleFromHloTextFile(filename, GetDebugOptionsForTest()); if (!module_or_status.ok()) { return ::testing::AssertionFailure() << "failed reading hlo module from file"; @@ -260,7 +258,7 @@ template const string& filename, const tensorflow::gtl::optional& error, const std::function& reference_preprocessor) { auto module_or_status = - HloRunner::ReadModule(filename, GetDebugOptionsForTest()); + HloRunner::ReadModuleFromHloTextFile(filename, GetDebugOptionsForTest()); if (!module_or_status.ok()) { return ::testing::AssertionFailure() << "failed reading hlo module from file"; diff --git a/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc b/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc index 31060b9e80fcd50aefdedca27c70ec8a9b8be743..506091ddd8d1d8e6519525bb7031f4e8b296b5fb 100644 --- a/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc @@ -23,15 +23,8 @@ limitations under the License. namespace xla { -/*static*/ int64 HloVerifiedTestBase::DefaultShapeSize(const Shape& shape) { - constexpr int64 kPointerSize = sizeof(void*); - if (ShapeUtil::IsOpaque(shape)) { - return kPointerSize; - } - return ShapeUtil::ByteSizeOf(shape, kPointerSize); -} - -HloVerifiedTestBase::HloVerifiedTestBase() : shape_size_fn_(DefaultShapeSize) {} +HloVerifiedTestBase::HloVerifiedTestBase() + : shape_verifier_(MakeUnique()) {} HloVerifiedTestBase::~HloVerifiedTestBase() { // We can't call the ASSERT or EXPECT test macros in destructors, so we @@ -47,7 +40,7 @@ void HloVerifiedTestBase::TearDown() { << "TearDown called more than once; it should be called exactly once."; tear_down_called_ = true; if (module_) { - HloVerifier verifier(shape_size_fn_); + HloVerifier verifier; xla::StatusOr mutated = verifier.Run(module_.get()); if (!mutated.ok()) { ADD_FAILURE() << "HloVerifier failed: " << mutated.status(); diff --git a/tensorflow/compiler/xla/tests/hlo_verified_test_base.h b/tensorflow/compiler/xla/tests/hlo_verified_test_base.h index b3d6b5af3b46f932707abf309669d23c327d1334..492688bf7d682cf991cb8c09399492a0437f651b 100644 --- a/tensorflow/compiler/xla/tests/hlo_verified_test_base.h +++ b/tensorflow/compiler/xla/tests/hlo_verified_test_base.h @@ -28,14 +28,13 @@ namespace xla { // A base class for HLO tests that stores a default HloModule, and automatically // performs verification on that module on tear-down. class HloVerifiedTestBase : public HloTestBase { - public: - // Returns the size in bytes of the given shape, using a default pointer size. - static int64 DefaultShapeSize(const Shape& shape); - protected: HloVerifiedTestBase(); ~HloVerifiedTestBase() override; + // Constructs a default shape verifier. + std::unique_ptr MakeShapeVerifier(); + // Performs verification on the default HloModule returned by module(). // Automatically called by the testing framework for each test. // @@ -47,14 +46,14 @@ class HloVerifiedTestBase : public HloTestBase { HloModule& module(); // Sets the shape-size function used during hlo verification. If this isn't - // called, DefaultShapeSize is used instead. - void SetShapeSizeFn(std::function shape_size_fn) { - shape_size_fn_ = std::move(shape_size_fn); + // called, a default ShapeVerifier is used instead. + void SetShapeVerifier(std::unique_ptr shape_verifier) { + shape_verifier_ = std::move(shape_verifier); } private: std::unique_ptr module_; // Lazily populated. Access via module(). - std::function shape_size_fn_; + std::unique_ptr shape_verifier_; bool tear_down_called_ = false; }; diff --git a/tensorflow/compiler/xla/tests/literal_test_util.cc b/tensorflow/compiler/xla/tests/literal_test_util.cc index fb425fe6f3cfbb35d7824f3dd1b7d3a2f869313f..5aa71a9261dbd414d1499f15c9b83cd63b634b49 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util.cc @@ -101,56 +101,57 @@ namespace xla { ASSERT_EQ(expected.ShortDebugString(), actual.ShortDebugString()); } +namespace { + +// Return a literal with all arrays of type FromNativeT converted to type +// ToNativeT in the given literal. +template +std::unique_ptr ConvertType(const Literal& literal) { + // First construct shape of the result. + Shape result_shape(literal.shape()); + ShapeUtil::ForEachMutableSubshape( + &result_shape, [](Shape* subshape, const ShapeIndex&) { + if (subshape->element_type() == + primitive_util::NativeToPrimitiveType()) { + subshape->set_element_type( + primitive_util::NativeToPrimitiveType()); + } + }); + auto result = MakeUnique(result_shape); + + // Then copy over the data from 'literal' converting FromNativeT values to + // ToNativeT values as necessary. + ShapeUtil::ForEachSubshape( + literal.shape(), + [&](const Shape& subshape, const ShapeIndex& shape_index) { + if (ShapeUtil::IsArray(subshape)) { + if (subshape.element_type() == + primitive_util::NativeToPrimitiveType()) { + auto src = literal.data(shape_index); + auto dest = result->data(shape_index); + for (int64 i = 0; i < src.size(); ++i) { + dest[i] = static_cast(src[i]); + } + } else { + TF_CHECK_OK(result->CopyFrom(literal, + /*dest_shape_index=*/shape_index, + /*src_shape_index=*/shape_index)); + } + } + }); + return result; +} + +} // namespace + /* static */ std::unique_ptr LiteralTestUtil::ConvertBF16ToF32( const Literal& literal) { - if (ShapeUtil::IsTuple(literal.shape())) { - std::vector> converted_elements; - for (const auto& element : literal.tuple_literals()) { - converted_elements.push_back(ConvertBF16ToF32(element)); - } - return Literal::MakeTupleOwned(std::move(converted_elements)); - } - - if (literal.shape().element_type() != BF16) { - return MakeUnique(literal); - } - Shape converted_shape = literal.shape(); - converted_shape.set_element_type(F32); - auto converted = Literal::CreateFromShape(converted_shape); - if (!ShapeUtil::HasZeroElements(converted_shape)) { - std::vector index(converted_shape.dimensions_size(), 0); - do { - converted->Set(index, - static_cast(literal.Get(index))); - } while (IndexUtil::BumpIndices(converted_shape, &index)); - } - return converted; + return ConvertType(literal); } /* static */ std::unique_ptr LiteralTestUtil::ConvertF32ToBF16( const Literal& literal) { - if (ShapeUtil::IsTuple(literal.shape())) { - std::vector> converted_elements; - for (const auto& element : literal.tuple_literals()) { - converted_elements.push_back(ConvertF32ToBF16(element)); - } - return Literal::MakeTupleOwned(std::move(converted_elements)); - } - - if (literal.shape().element_type() != F32) { - return MakeUnique(literal); - } - Shape converted_shape = literal.shape(); - converted_shape.set_element_type(BF16); - auto converted = Literal::CreateFromShape(converted_shape); - if (!ShapeUtil::HasZeroElements(converted_shape)) { - std::vector index(converted_shape.dimensions_size(), 0); - do { - converted->Set( - index, static_cast(literal.Get(index))); - } while (IndexUtil::BumpIndices(converted_shape, &index)); - } - return converted; + return ConvertType(literal); } namespace { @@ -300,6 +301,9 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, case BF16: match = ExpectLiteralsEqual(expected, actual, &multi_index, 0); break; + case F16: + match = ExpectLiteralsEqual(expected, actual, &multi_index, 0); + break; case F32: match = ExpectLiteralsEqual(expected, actual, &multi_index, 0); break; @@ -311,9 +315,14 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, break; case TUPLE: { bool tuple_match = true; - for (int i = 0; i < actual.tuple_literals_size(); ++i) { - auto result = - Equal(expected.tuple_literals(i), actual.tuple_literals(i)); + for (int i = 0; i < ShapeUtil::TupleElementCount(expected.shape()); ++i) { + SCOPED_TRACE(tensorflow::strings::StrCat( + "Tuple index ", i, " in ", + ShapeUtil::HumanString(expected.shape()))); + + // Create LiteralViews of the expected and actual elements. + auto result = Equal(LiteralView::Create(expected, {i}), + LiteralView::Create(actual, {i})); tuple_match = tuple_match ? !!result : false; } match = tuple_match; @@ -334,47 +343,6 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, return result; } -/* static */ ::testing::AssertionResult LiteralTestUtil::EqualTuple( - const Literal& expected, const Literal& actual) { - VLOG(1) << "expected: " << expected.ToString(); - VLOG(1) << "actual: " << actual.ToString(); - - if (!ShapeUtil::IsTuple(expected.shape()) || - !ShapeUtil::IsTuple(actual.shape())) { - return ::testing::AssertionFailure() - << "tuples expected shape = " << expected.shape().ShortDebugString() - << " actual shape = " << actual.shape().ShortDebugString(); - } - AssertEqualShapes(expected.shape(), actual.shape()); - - ::testing::AssertionResult err = ::testing::AssertionSuccess(); - for (uint64 i = 0; i < expected.tuple_literals_size(); ++i) { - SCOPED_TRACE(tensorflow::strings::StrCat( - "Tuple index ", i, " in ", ShapeUtil::HumanString(expected.shape()))); - const auto& expected_element = expected.tuple_literals(i); - const auto& actual_element = actual.tuple_literals(i); - - ::testing::AssertionResult res = [&] { - if (ShapeUtil::IsTuple(expected_element.shape())) { - return EqualTuple(expected_element, actual_element); - } else { - return Equal(expected_element, actual_element); - } - }(); - - if (!res && err) { - err = res; - } - } - - return err; -} - -/* static */ void LiteralTestUtil::ExpectEqualTuple(const Literal& expected, - const Literal& actual) { - EXPECT_TRUE(EqualTuple(expected, actual)); -} - namespace { // Helper class for comparing floating-point literals within an error bound. @@ -387,9 +355,9 @@ class NearComparator { // temporary files on failure. Returns true if literals match. bool ExpectNear(const Literal& expected, const Literal& actual) { VLOG(1) << "expected:"; - XLA_VLOG_LINES(1, expected.ToString()); + XLA_VLOG_LINES(1, TruncateHugeLiteral(expected)); VLOG(1) << "actual:"; - XLA_VLOG_LINES(1, actual.ToString()); + XLA_VLOG_LINES(1, TruncateHugeLiteral(actual)); // If the shapes mismatch, we simply fail the expectation instead of // printing out data, as it's a type error rather than a value error. @@ -408,16 +376,21 @@ class NearComparator { abs_expected_miscompare_sum_ = 0.0; max_rel_err_ = 0.0; max_abs_err_ = 0.0; - *miscompares_.mutable_shape() = - ShapeUtil::ChangeElementType(actual.shape(), PRED); - miscompares_.mutable_preds()->resize( - ShapeUtil::ElementsIn(miscompares_.shape()), false); + first_linear_index_ = -1; + last_linear_index_ = -1; + max_rel_linear_index_ = -1; + max_abs_linear_index_ = -1; + miscompares_ = Literal(ShapeUtil::ChangeElementType(actual.shape(), PRED)); + miscompares_.PopulateWithValue(false); multi_index_.resize(expected.shape().dimensions_size(), 0); switch (expected.shape().element_type()) { case BF16: ExpectLiteralsNear(expected, actual, 0); break; + case F16: + ExpectLiteralsNear(expected, actual, 0); + break; case F32: ExpectLiteralsNear(expected, actual, 0); break; @@ -436,21 +409,33 @@ class NearComparator { if (num_miscompares_ > 0) { if (!VLOG_IS_ON(1)) { LOG(INFO) << "expected: " << ShapeUtil::HumanString(expected.shape()) - << " " << expected.ToString(); + << " " << TruncateHugeLiteral(expected); LOG(INFO) << "actual: " << ShapeUtil::HumanString(actual.shape()) - << " " << actual.ToString(); + << " " << TruncateHugeLiteral(actual); + LOG(INFO) << "Dumping literals to temp files..."; + WriteLiteralToTempFile(expected, "expected"); + WriteLiteralToTempFile(actual, "actual"); + WriteLiteralToTempFile(miscompares_, "miscompares"); } EXPECT_TRUE(num_miscompares_ == 0) << "\nmax relative mismatch at index " - << LiteralTestUtil::MultiIndexAsString(max_rel_multi_index_) + << LiteralTestUtil::MultiIndexAsString( + IndexUtil::LinearIndexToMultidimensionalIndex( + actual.shape(), max_rel_linear_index_)) << "\nmaximum relative error " << max_rel_err_ << "\nmax absolute mismatch at index " - << LiteralTestUtil::MultiIndexAsString(max_abs_multi_index_) + << LiteralTestUtil::MultiIndexAsString( + IndexUtil::LinearIndexToMultidimensionalIndex( + actual.shape(), max_abs_linear_index_)) << "\nmaximum absolute error " << max_abs_err_ << "\nfirst mismatch at index " - << LiteralTestUtil::MultiIndexAsString(first_multi_index_) + << LiteralTestUtil::MultiIndexAsString( + IndexUtil::LinearIndexToMultidimensionalIndex( + actual.shape(), first_linear_index_)) << "\nlast mismatch at index " - << LiteralTestUtil::MultiIndexAsString(last_multi_index_) + << LiteralTestUtil::MultiIndexAsString( + IndexUtil::LinearIndexToMultidimensionalIndex( + actual.shape(), last_linear_index_)) << "\ntotal absolute error " << abs_diff_sum_ << "\ntotal absolute error of miscompares " << abs_diff_miscompare_sum_ << "\ntotal relative error " @@ -458,18 +443,18 @@ class NearComparator { << "\ntotal relative error of miscompares " << (abs_diff_miscompare_sum_ / abs_expected_miscompare_sum_) << "\nfailure count " << num_miscompares_; - - WriteLiteralToTempFile(expected, "expected"); - WriteLiteralToTempFile(actual, "actual"); - WriteLiteralToTempFile(miscompares_, "miscompares"); } return num_miscompares_ == 0; } private: template - bool NanMismatch(NativeT lhs, NativeT rhs) { - return std::isnan(lhs) != std::isnan(rhs); + bool NanMismatch(NativeT expected, NativeT actual, bool relaxed_nans) { + if (relaxed_nans) { + return !std::isnan(expected) && std::isnan(actual); + } else { + return std::isnan(expected) != std::isnan(actual); + } } template @@ -489,57 +474,94 @@ class NearComparator { return true; } - float abs_diff = std::abs(actual - expected); - float rel_err = abs_diff / std::abs(expected); + const float abs_diff = std::abs(actual - expected); + const float rel_err = abs_diff / std::abs(expected); + const bool nan_mismatch = + NanMismatch(expected, actual, error_.relaxed_nans); + const bool mismatch = + (nan_mismatch || (abs_diff >= error_.abs && rel_err >= error_.rel)); + return !mismatch; + } + + // Assumes that expected vs actual fail ExpectValuesNear. + template + void UpdateAndLogMiscompares(const NativeT expected, const NativeT actual, + const Shape& shape, const int64 linear_index) { + const float abs_diff = std::abs(actual - expected); + const float rel_err = abs_diff / std::abs(expected); abs_diff_sum_ += abs_diff; abs_expected_sum_ += std::abs(expected); - if (rel_err > max_rel_err_) { + if (rel_err > max_rel_err_ || std::isnan(rel_err)) { max_rel_err_ = rel_err; - max_rel_multi_index_ = multi_index_; + max_rel_linear_index_ = linear_index; } - if (abs_diff > max_abs_err_) { + if (abs_diff > max_abs_err_ || std::isnan(abs_diff)) { max_abs_err_ = abs_diff; - max_abs_multi_index_ = multi_index_; + max_abs_linear_index_ = linear_index; } - VLOG(10) << tensorflow::strings::Printf( - "index %s abs_diff %f rel_err %f", - LiteralTestUtil::MultiIndexAsString(multi_index_).c_str(), abs_diff, - rel_err); - bool nan_mismatch = NanMismatch(expected, actual); - bool mismatch = - (nan_mismatch || (abs_diff >= error_.abs && rel_err >= error_.rel)); - if (mismatch) { - abs_diff_miscompare_sum_ += abs_diff; - abs_expected_miscompare_sum_ += std::abs(expected); - const int64 kMaxFailures = 2; - if (num_miscompares_ < kMaxFailures) { - ::testing::Message msg; - msg << "mismatch at index " - << LiteralTestUtil::MultiIndexAsString(multi_index_) << " abs diff " - << abs_diff << " rel err " << rel_err << " failure #" - << num_miscompares_; - ExpectNear(expected, actual, msg); - } else if (num_miscompares_ == kMaxFailures) { - LOG(ERROR) - << "reached max 'loud' failure count; silently proceeding..."; - } - if (num_miscompares_ == 0) { - first_multi_index_ = multi_index_; - } - num_miscompares_++; - last_multi_index_ = multi_index_; + if (VLOG_IS_ON(10)) { + VLOG(10) << tensorflow::strings::Printf( + "index %s abs_diff %f rel_err %f", + LiteralTestUtil::MultiIndexAsString( + IndexUtil::LinearIndexToMultidimensionalIndex(shape, + linear_index)) + .c_str(), + abs_diff, rel_err); } - return !mismatch; + abs_diff_miscompare_sum_ += abs_diff; + abs_expected_miscompare_sum_ += std::abs(expected); + const int64 kMaxFailures = 2; + if (num_miscompares_ < kMaxFailures) { + const auto multi_index = + IndexUtil::LinearIndexToMultidimensionalIndex(shape, linear_index); + ::testing::Message msg; + msg << "mismatch at index " + << LiteralTestUtil::MultiIndexAsString(multi_index) << " abs diff " + << abs_diff << " rel err " << rel_err << " failure #" + << num_miscompares_; + ExpectNear(expected, actual, msg); + } else if (num_miscompares_ == kMaxFailures) { + LOG(ERROR) << "reached max 'loud' failure count; silently proceeding..."; + } + if (num_miscompares_ == 0) { + first_linear_index_ = linear_index; + } + num_miscompares_++; + last_linear_index_ = linear_index; + miscompares_.data()[linear_index] = true; } // Recursive function which compares the two given literals elementwise. template void ExpectLiteralsNear(const Literal& expected, const Literal& actual, int64 dimension) { + // Fast path optimization for the case were layouts match. + if (LayoutUtil::Equal(actual.shape().layout(), expected.shape().layout())) { + tensorflow::gtl::ArraySlice expected_data = + expected.data(); + tensorflow::gtl::ArraySlice actual_data = + actual.data(); + const int64 len = expected_data.size(); + for (int64 i = 0; i < len; ++i) { + const bool near = ExpectValuesNear(expected_data[i], actual_data[i]); + if (!near) { + UpdateAndLogMiscompares(expected_data[i], actual_data[i], + actual.shape(), i); + } + } + return; + } + if (dimension == expected.shape().dimensions_size()) { bool near = ExpectValuesNear(expected.Get(multi_index_), actual.Get(multi_index_)); - miscompares_.Set(multi_index_, !near); + if (!near) { + UpdateAndLogMiscompares( + expected.Get(multi_index_), + actual.Get(multi_index_), actual.shape(), + IndexUtil::MultidimensionalIndexToLinearIndex(actual.shape(), + multi_index_)); + } } else { for (int64 i = 0; i < expected.shape().dimensions(dimension); ++i) { multi_index_[dimension] = i; @@ -560,6 +582,32 @@ class NearComparator { LOG(ERROR) << "wrote to " << name << " file: " << filename; } + // Gets the total element count. For tuples, this is not the count of tuple + // elements, but the sum of elements of each tuple element. + int64 RecursiveElementCount(const Shape& shape) { + if (ShapeUtil::IsTuple(shape)) { + const int64 tuple_elements = ShapeUtil::TupleElementCount(shape); + int64 total = 0; + for (int64 i = 0; i < tuple_elements; ++i) { + total += + RecursiveElementCount(ShapeUtil::GetTupleElementShape(shape, i)); + } + return total; + } else { + return ShapeUtil::ElementsIn(shape); + } + } + + // Calling ToString on a literal with over 100 million elements takes around + // 3 minutes. The utility of printing a literal with >1000 elements is + // questionable, especially when writing the Literal proto to disk is orders + // of magnitude faster. + string TruncateHugeLiteral(const Literal& literal) { + return RecursiveElementCount(literal.shape()) < 1000 + ? literal.ToString() + : "[TRUNCATED, Literal with more than 1000 values]"; + } + ErrorSpec error_; // Number of element miscomparisons encountered so far. @@ -580,16 +628,18 @@ class NearComparator { double abs_expected_miscompare_sum_; float max_rel_err_; float max_abs_err_; - std::vector first_multi_index_; - std::vector last_multi_index_; - std::vector max_rel_multi_index_; - std::vector max_abs_multi_index_; + int64 first_linear_index_; + int64 last_linear_index_; + int64 max_rel_linear_index_; + int64 max_abs_linear_index_; }; template <> -bool NearComparator::NanMismatch(complex64 lhs, complex64 rhs) { - return std::isnan(lhs.real()) != std::isnan(rhs.real()) || - std::isnan(lhs.imag()) != std::isnan(rhs.imag()); +bool NearComparator::NanMismatch(complex64 expected, + complex64 actual, + bool relaxed_nans) { + return NanMismatch(expected.real(), actual.real(), relaxed_nans) || + NanMismatch(expected.imag(), actual.imag(), relaxed_nans); } template <> @@ -610,14 +660,64 @@ bool NearComparator::ExpectValuesNear(bfloat16 expected, static_cast(actual)); } +template <> +bool NearComparator::ExpectValuesNear(half expected, half actual) { + return ExpectValuesNear(static_cast(std::move(expected)), + static_cast(std::move(actual))); +} + +template <> +void NearComparator::UpdateAndLogMiscompares( + const bfloat16 expected, const bfloat16 actual, const Shape& shape, + const int64 linear_index) { + UpdateAndLogMiscompares(static_cast(expected), + static_cast(actual), shape, linear_index); +} + +template <> +void NearComparator::UpdateAndLogMiscompares(half expected, half actual, + const Shape& shape, + const int64 linear_index) { + UpdateAndLogMiscompares(static_cast(std::move(expected)), + static_cast(std::move(actual)), shape, + linear_index); +} + } // namespace /* static */ ::testing::AssertionResult LiteralTestUtil::Near( const Literal& expected, const Literal& actual, const ErrorSpec& error) { - NearComparator comparator(error); - return comparator.ExpectNear(expected, actual) - ? ::testing::AssertionSuccess() - : ::testing::AssertionFailure() << "values were not near"; + ::testing::AssertionResult err = + EqualShapes(expected.shape(), actual.shape()); + if (!err) { + return err; + } + + if (ShapeUtil::IsTuple(expected.shape())) { + for (int64 i = 0; i < ShapeUtil::TupleElementCount(expected.shape()); ++i) { + SCOPED_TRACE(tensorflow::strings::StrCat( + "Tuple index ", i, " in ", ShapeUtil::HumanString(expected.shape()))); + const auto expected_element = LiteralView::Create(expected, {i}); + const auto actual_element = LiteralView::Create(actual, {i}); + + ::testing::AssertionResult res = + Near(expected_element, actual_element, error); + if (err && !res) { + err = res; + } + } + return err; + } + + if (ShapeUtil::ElementIsFloating(expected.shape()) || + ShapeUtil::ElementIsComplex(expected.shape())) { + NearComparator comparator(error); + return comparator.ExpectNear(expected, actual) + ? ::testing::AssertionSuccess() + : ::testing::AssertionFailure() << "values were not near"; + } + + return Equal(expected, actual); } /* static */ void LiteralTestUtil::ExpectNear(const Literal& expected, @@ -630,65 +730,13 @@ bool NearComparator::ExpectValuesNear(bfloat16 expected, : tensorflow::strings::StrCat("\nmessage: ", message)); } -/* static */ ::testing::AssertionResult LiteralTestUtil::NearTuple( - const Literal& expected, const Literal& actual, const ErrorSpec& error) { - VLOG(1) << "expected: " << expected.ToString(); - VLOG(1) << "actual: " << actual.ToString(); - - if (!ShapeUtil::IsTuple(expected.shape()) || - !ShapeUtil::IsTuple(actual.shape())) { - return ::testing::AssertionFailure() - << "tuples expected shape = " << expected.shape().ShortDebugString() - << " actual shape = " << actual.shape().ShortDebugString(); - } - AssertEqualShapes(expected.shape(), actual.shape()); - - ::testing::AssertionResult err = ::testing::AssertionSuccess(); - for (uint64 i = 0; i < expected.tuple_literals_size(); ++i) { - SCOPED_TRACE(tensorflow::strings::StrCat( - "Tuple index ", i, " in ", ShapeUtil::HumanString(expected.shape()))); - const auto& expected_element = expected.tuple_literals(i); - const auto& actual_element = actual.tuple_literals(i); - - ::testing::AssertionResult res = [&] { - if (ShapeUtil::IsTuple(expected_element.shape())) { - return NearTuple(expected_element, actual_element, error); - } else if (ShapeUtil::ElementIsFloating(expected_element.shape())) { - return Near(expected_element, actual_element, error); - } else { - return Equal(expected_element, actual_element); - } - }(); - - if (err && !res) { - err = res; - } - } - return err; -} - -/* static */ void LiteralTestUtil::ExpectNearTuple(const Literal& expected, - const Literal& actual, - const ErrorSpec& error) { - EXPECT_TRUE(NearTuple(expected, actual, error)); -} - /*static*/ ::testing::AssertionResult LiteralTestUtil::NearOrEqual( const Literal& expected, const Literal& actual, const tensorflow::gtl::optional& error) { - bool is_tuple = ShapeUtil::IsTuple(expected.shape()); if (error.has_value()) { - if (is_tuple) { - VLOG(1) << "Expects near tuple"; - return NearTuple(expected, actual, *error); - } VLOG(1) << "Expects near"; return Near(expected, actual, *error); } - if (is_tuple) { - VLOG(1) << "Expects equal tuple"; - return EqualTuple(expected, actual); - } VLOG(1) << "Expects equal"; return Equal(expected, actual); } @@ -713,10 +761,10 @@ bool NearComparator::ExpectValuesNear(bfloat16 expected, new_num_elements *= new_dimensions[i]; } CHECK_EQ(ShapeUtil::ElementsIn(literal.shape()), new_num_elements); + CHECK_EQ(new_dimensions.size(), minor_to_major.size()); - auto new_literal = MakeUnique(); - *new_literal->mutable_shape() = - ShapeUtil::MakeShape(literal.shape().element_type(), new_dimensions); + auto new_literal = MakeUnique( + ShapeUtil::MakeShape(literal.shape().element_type(), new_dimensions)); // Create a new shape with the given minor-to-major layout. This shape is used // solely for converting linear address to multi-dimensional addresses when @@ -724,9 +772,6 @@ bool NearComparator::ExpectValuesNear(bfloat16 expected, Shape shape_with_layout = new_literal->shape(); *shape_with_layout.mutable_layout() = LayoutUtil::MakeLayout(minor_to_major); - // Allocate space in the new literal. - new_literal->Reserve(ShapeUtil::ElementsIn(literal.shape())); - // Copy data into new literal, element-by-element. for (int64 i = 0; i < ShapeUtil::ElementsIn(literal.shape()); ++i) { std::vector from_multi_index = @@ -766,6 +811,10 @@ bool NearComparator::ExpectValuesNear(bfloat16 expected, new_literal->Set(to_multi_index, literal.Get(from_multi_index)); break; + case C64: + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); + break; default: LOG(FATAL) << "Unhandled primitive element type: " << PrimitiveType_Name(literal.shape().element_type()); diff --git a/tensorflow/compiler/xla/tests/literal_test_util.h b/tensorflow/compiler/xla/tests/literal_test_util.h index f53553c70170bdcda717e72ffd791016effd0774..7b757a4bd7e7592583b7596b4305ddb7e6c52d75 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.h +++ b/tensorflow/compiler/xla/tests/literal_test_util.h @@ -40,10 +40,16 @@ namespace xla { // Structure describing permissible absolute and relative error bounds. struct ErrorSpec { - explicit ErrorSpec(float aabs, float arel = 0) : abs(aabs), rel(arel) {} + explicit ErrorSpec(float aabs, float arel = 0, bool relaxed_nans = false) + : abs(aabs), rel(arel), relaxed_nans(relaxed_nans) {} float abs; // Absolute error bound. float rel; // Relative error bound. + + // If relaxed_nans is true then any result is valid if we are expecting NaNs. + // In effect, this allows the tested operation to produce incorrect results + // for inputs outside its mathematical domain. + bool relaxed_nans; }; // Utility class for making expectations/assertions related to XLA literals. @@ -111,17 +117,18 @@ class LiteralTestUtil { static void ExpectR4EqualArray4D(const Array4D& expected, const Literal& actual); - // Returns whether the two tuples are equal. - static ::testing::AssertionResult EqualTuple( - const Literal& expected, const Literal& actual) TF_MUST_USE_RESULT; - - // Expects that the values of the elements in the expected and actual tuples - // are equal. Tuples are matched recursively. - static void ExpectEqualTuple(const Literal& expected, const Literal& actual); - // Asserts that the expected and actual literals are within the given error // bound for all elements. Also, asserts that the rank, dimensions sizes, and - // bounds are equivalent. Only supported for floating point values. + // bounds are equivalent. + // + // Tuples are matched recursively. When comparing tensors of + // non-floating-point type, checks for exact equality, ignoring the ErroSpec. + // + // If the shape of the literals is neither a complex/floating-point tensor nor + // a tuple which contains a complex/floating-point tensor, Near() is + // equivalent to Equal(). We don't raise an error in this case, because we + // want to allow callers to call Near() even if they have no preconceptions + // about the shapes being compared. static ::testing::AssertionResult Near( const Literal& expected, const Literal& actual, const ErrorSpec& error) TF_MUST_USE_RESULT; @@ -170,18 +177,6 @@ class LiteralTestUtil { const Literal& actual, const ErrorSpec& error); - // Returns whether the values of the elements in the expected and actual - // tuples are within the given error bound. Tuples are matched recursively. - // If the elements of the tuple are not floating-point types, the error spec - // is ignored and exact equality is checked. - static ::testing::AssertionResult NearTuple( - const Literal& expected, const Literal& actual, - const ErrorSpec& error) TF_MUST_USE_RESULT; - - // Expects that the expected and actual values are near. - static void ExpectNearTuple(const Literal& expected, const Literal& actual, - const ErrorSpec& error); - // If the error spec is given, returns whether the expected and the actual are // within the error bound; otherwise, returns whether they are equal. Tuples // will be compared recursively. diff --git a/tensorflow/compiler/xla/tests/literal_test_util_test.cc b/tensorflow/compiler/xla/tests/literal_test_util_test.cc index 2acf27ed390b0732ba40fcf505c746bd7d8b651e..3a421f8458268a14dcdd84889bcae4990c095ea4 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util_test.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util_test.cc @@ -83,18 +83,43 @@ TEST(LiteralTestUtilTest, ExpectNearFailurePlacesResultsInTemporaryDirectory) { LiteralProto literal_proto; TF_CHECK_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(), result, &literal_proto)); - Literal literal(literal_proto); + std::unique_ptr literal = + Literal::CreateFromProto(literal_proto).ConsumeValueOrDie(); if (result.find("expected") != string::npos) { - EXPECT_EQ("2", literal.ToString()); + EXPECT_EQ("2", literal->ToString()); } else if (result.find("actual") != string::npos) { - EXPECT_EQ("4", literal.ToString()); + EXPECT_EQ("4", literal->ToString()); } else if (result.find("miscompares") != string::npos) { - EXPECT_EQ("true", literal.ToString()); + EXPECT_EQ("true", literal->ToString()); } else { FAIL() << "unknown file in temporary directory: " << result; } } } +TEST(LiteralTestUtilTest, NearComparatorR1) { + auto a = + Literal::CreateR1({0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}); + auto b = + Literal::CreateR1({0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}); + EXPECT_TRUE(LiteralTestUtil::Near(*a, *b, ErrorSpec{0.0001})); +} + +TEST(LiteralTestUtilTest, NearComparatorR1Nan) { + auto a = + Literal::CreateR1({0.0, 0.1, 0.2, 0.3, NAN, 0.5, 0.6, 0.7, 0.8}); + auto b = + Literal::CreateR1({0.0, 0.1, 0.2, 0.3, NAN, 0.5, 0.6, 0.7, 0.8}); + EXPECT_TRUE(LiteralTestUtil::Near(*a, *b, ErrorSpec{0.0001})); +} + +TEST(LiteralTestUtil, NearComparatorDifferentLengths) { + auto a = + Literal::CreateR1({0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}); + auto b = Literal::CreateR1({0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7}); + EXPECT_FALSE(LiteralTestUtil::Near(*a, *b, ErrorSpec{0.0001})); + EXPECT_FALSE(LiteralTestUtil::Near(*b, *a, ErrorSpec{0.0001})); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/llvm_compiler_test.cc b/tensorflow/compiler/xla/tests/llvm_compiler_test.cc index b5b95967ff9162301a092f3a57996e0f3f78658f..7e92439c494b677f718a63c71c20828d65bebef4 100644 --- a/tensorflow/compiler/xla/tests/llvm_compiler_test.cc +++ b/tensorflow/compiler/xla/tests/llvm_compiler_test.cc @@ -74,7 +74,8 @@ class LLVMCompilerTest : public ::testing::Test { ASSERT_TRUE(compiler ->RunBackend(std::move(hlo_module), - backend_->default_stream_executor()) + backend_->default_stream_executor(), + /*device_allocator=*/nullptr) .ok()); // Test that hooks were called. @@ -98,7 +99,8 @@ class LLVMCompilerTest : public ::testing::Test { executors.push_back({backend_->default_stream_executor()}); executors.push_back({backend_->default_stream_executor()}); - EXPECT_IS_OK(compiler->Compile(std::move(modules), std::move(executors))); + EXPECT_IS_OK(compiler->Compile(std::move(modules), std::move(executors), + /*device_allocator=*/nullptr)); } private: diff --git a/tensorflow/compiler/xla/tests/local_client_aot_test.cc b/tensorflow/compiler/xla/tests/local_client_aot_test.cc index 569d5944cab0ae8f6a7b58a651285d20d4f9d019..47cab796041e9669affaebd7866d0d80100730f1 100644 --- a/tensorflow/compiler/xla/tests/local_client_aot_test.cc +++ b/tensorflow/compiler/xla/tests/local_client_aot_test.cc @@ -44,8 +44,7 @@ TEST_F(LocalClientAotTest, Constant) { OpaqueData opaque_data{100, 20, 3}; void* parameters[] = {&opaque_data}; float out = 0; - char tmp[4] = {0}; - void* temporary_buffers[] = {nullptr, &out, &tmp}; + void* temporary_buffers[] = {nullptr, &out}; SumAndDouble(&out, &run_options, parameters, temporary_buffers); EXPECT_EQ(out, 246.0f); diff --git a/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc b/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc index 4d3b513b092e0b447a1452a3809fb7099e54dbb9..3704ddd8010bf727b75ff81b63605e8b7ffe2ca8 100644 --- a/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc +++ b/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc @@ -87,10 +87,9 @@ int main(int argc, char** argv) { // It's lame to hard-code the buffer assignments, but we need // local_client_aot_test.cc to be able to easily invoke the function. CHECK_EQ(result->result_buffer_index(), 1); - CHECK_EQ(result->buffer_sizes().size(), 3); + CHECK_EQ(result->buffer_sizes().size(), 2); CHECK_EQ(result->buffer_sizes()[0], -1); // param buffer CHECK_EQ(result->buffer_sizes()[1], sizeof(float)); // result buffer - CHECK_EQ(result->buffer_sizes()[2], sizeof(float)); // temp buffer if (triple.isOSBinFormatELF()) { // Check the ELF magic. CHECK_EQ(result->object_file_data()[0], 0x7F); diff --git a/tensorflow/compiler/xla/tests/local_client_execute_test.cc b/tensorflow/compiler/xla/tests/local_client_execute_test.cc index e3298e98c67969f97adfdf15d22dfe72592b56aa..2462ea39f914b1dbb525ea777a48d9ce66035638 100644 --- a/tensorflow/compiler/xla/tests/local_client_execute_test.cc +++ b/tensorflow/compiler/xla/tests/local_client_execute_test.cc @@ -217,12 +217,13 @@ XLA_TEST_F(LocalClientExecuteTest, TupleResult) { EXPECT_EQ(3, ShapeUtil::TupleElementCount(result->on_host_shape())); std::unique_ptr result_literal = ShapedBufferToLiteral(*result); - LiteralTestUtil::ExpectR2Equal({{1.0f, 2.0f}, {3.0f, 4.0f}}, - result_literal->tuple_literals(0)); - LiteralTestUtil::ExpectR2Equal({{10.0f, 20.0f}, {30.0f, 40.0f}}, - result_literal->tuple_literals(1)); - LiteralTestUtil::ExpectR2Equal({{1.0f, 2.0f}, {3.0f, 4.0f}}, - result_literal->tuple_literals(2)); + LiteralTestUtil::ExpectR2Equal( + {{1.0f, 2.0f}, {3.0f, 4.0f}}, LiteralView::Create(*result_literal, {0})); + LiteralTestUtil::ExpectR2Equal( + {{10.0f, 20.0f}, {30.0f, 40.0f}}, + LiteralView::Create(*result_literal, {1})); + LiteralTestUtil::ExpectR2Equal( + {{1.0f, 2.0f}, {3.0f, 4.0f}}, LiteralView::Create(*result_literal, {2})); } XLA_TEST_F(LocalClientExecuteTest, NestedTupleResult) { @@ -245,15 +246,17 @@ XLA_TEST_F(LocalClientExecuteTest, NestedTupleResult) { EXPECT_EQ(2, ShapeUtil::TupleElementCount(result->on_host_shape())); std::unique_ptr result_literal = ShapedBufferToLiteral(*result); - LiteralTestUtil::ExpectR2Equal({{1.0f, 2.0f}, {3.0f, 4.0f}}, - result_literal->tuple_literals(1)); - const Literal& inner_tuple_literal = result_literal->tuple_literals(0); - LiteralTestUtil::ExpectR2Equal({{1.0f, 2.0f}, {3.0f, 4.0f}}, - inner_tuple_literal.tuple_literals(0)); - LiteralTestUtil::ExpectR2Equal({{10.0f, 20.0f}, {30.0f, 40.0f}}, - inner_tuple_literal.tuple_literals(1)); - LiteralTestUtil::ExpectR2Equal({{1.0f, 2.0f}, {3.0f, 4.0f}}, - inner_tuple_literal.tuple_literals(2)); + LiteralTestUtil::ExpectR2Equal( + {{1.0f, 2.0f}, {3.0f, 4.0f}}, LiteralView::Create(*result_literal, {1})); + LiteralTestUtil::ExpectR2Equal( + {{1.0f, 2.0f}, {3.0f, 4.0f}}, + LiteralView::Create(*result_literal, {0, 0})); + LiteralTestUtil::ExpectR2Equal( + {{10.0f, 20.0f}, {30.0f, 40.0f}}, + LiteralView::Create(*result_literal, {0, 1})); + LiteralTestUtil::ExpectR2Equal( + {{1.0f, 2.0f}, {3.0f, 4.0f}}, + LiteralView::Create(*result_literal, {0, 2})); } XLA_TEST_F(LocalClientExecuteTest, TupleResultWithLayout) { @@ -278,10 +281,10 @@ XLA_TEST_F(LocalClientExecuteTest, TupleResultWithLayout) { DefaultExecutableRunOptions()); std::unique_ptr result_literal = ShapedBufferToLiteral(*result); - LiteralTestUtil::ExpectR2Equal({{1.0f, 2.0f}, {3.0f, 4.0f}}, - result_literal->tuple_literals(0)); - LiteralTestUtil::ExpectR2Equal({{1.0f, 2.0f}, {3.0f, 4.0f}}, - result_literal->tuple_literals(1)); + LiteralTestUtil::ExpectR2Equal( + {{1.0f, 2.0f}, {3.0f, 4.0f}}, LiteralView::Create(*result_literal, {0})); + LiteralTestUtil::ExpectR2Equal( + {{1.0f, 2.0f}, {3.0f, 4.0f}}, LiteralView::Create(*result_literal, {1})); } XLA_TEST_F(LocalClientExecuteTest, TupleArguments) { @@ -324,10 +327,11 @@ XLA_TEST_F(LocalClientExecuteTest, TupleArguments) { EXPECT_EQ(2, ShapeUtil::TupleElementCount(result->on_host_shape())); std::unique_ptr result_literal = ShapedBufferToLiteral(*result); - LiteralTestUtil::ExpectR2Equal({{56.0f, 46.0f}, {36.0f, 26.0f}}, - result_literal->tuple_literals(0)); - LiteralTestUtil::ExpectR1Equal({40.0f, 71.0f, 117.0f}, - result_literal->tuple_literals(1)); + LiteralTestUtil::ExpectR2Equal( + {{56.0f, 46.0f}, {36.0f, 26.0f}}, + LiteralView::Create(*result_literal, {0})); + LiteralTestUtil::ExpectR1Equal( + {40.0f, 71.0f, 117.0f}, LiteralView::Create(*result_literal, {1})); } XLA_TEST_F(LocalClientExecuteTest, NestedTupleArgument) { @@ -365,10 +369,10 @@ XLA_TEST_F(LocalClientExecuteTest, NestedTupleArgument) { ExecuteLocallyOrDie(computation, {arg_buffer.get()}); std::unique_ptr result_literal = ShapedBufferToLiteral(*result); - LiteralTestUtil::ExpectR2Equal({{-1.0, -2.0}, {-3.0, -4}}, - result_literal->tuple_literals(0)); - LiteralTestUtil::ExpectR1Equal({264.0, 73.0, 133.0}, - result_literal->tuple_literals(1)); + LiteralTestUtil::ExpectR2Equal( + {{-1.0, -2.0}, {-3.0, -4}}, LiteralView::Create(*result_literal, {0})); + LiteralTestUtil::ExpectR1Equal( + {264.0, 73.0, 133.0}, LiteralView::Create(*result_literal, {1})); } XLA_TEST_F(LocalClientExecuteTest, PassingTupleResultBackIntoComputation) { @@ -395,18 +399,19 @@ XLA_TEST_F(LocalClientExecuteTest, PassingTupleResultBackIntoComputation) { std::unique_ptr result_0 = ExecuteLocallyOrDie(computation, {arg_buffer.get()}); std::unique_ptr result_0_literal = ShapedBufferToLiteral(*result_0); - LiteralTestUtil::ExpectR2Equal({{-1.0, -2.0}, {-3.0, -4.0}}, - result_0_literal->tuple_literals(0)); - LiteralTestUtil::ExpectR2Equal({{22.0, 6.0}, {8.0, 10}}, - result_0_literal->tuple_literals(1)); + LiteralTestUtil::ExpectR2Equal( + {{-1.0, -2.0}, {-3.0, -4.0}}, + LiteralView::Create(*result_0_literal, {0})); + LiteralTestUtil::ExpectR2Equal( + {{22.0, 6.0}, {8.0, 10}}, LiteralView::Create(*result_0_literal, {1})); std::unique_ptr result_1 = ExecuteLocallyOrDie(computation, {result_0.get()}); std::unique_ptr result_1_literal = ShapedBufferToLiteral(*result_1); - LiteralTestUtil::ExpectR2Equal({{1.0, 2.0}, {3.0, 4.0}}, - result_1_literal->tuple_literals(0)); - LiteralTestUtil::ExpectR2Equal({{44.0, 12.0}, {16.0, 20}}, - result_1_literal->tuple_literals(1)); + LiteralTestUtil::ExpectR2Equal( + {{1.0, 2.0}, {3.0, 4.0}}, LiteralView::Create(*result_1_literal, {0})); + LiteralTestUtil::ExpectR2Equal( + {{44.0, 12.0}, {16.0, 20}}, LiteralView::Create(*result_1_literal, {1})); } XLA_TEST_F(LocalClientExecuteTest, LargeTuple) { @@ -455,7 +460,8 @@ XLA_TEST_F(LocalClientExecuteTest, LargeTuple) { for (int i = 0; i < kElementCount; ++i) { LiteralTestUtil::ExpectR1Near( - {2.0f * i, 0.0f}, result_literal->tuple_literals(i), error_spec_); + {2.0f * i, 0.0f}, LiteralView::Create(*result_literal, {i}), + error_spec_); } } @@ -512,8 +518,8 @@ XLA_TEST_F(LocalClientExecuteTest, DISABLED_ON_CPU_PARALLEL(LargeNestedTuple)) { for (int i = 0; i < kFanout; ++i) { for (int j = 0; j < kFanout; ++j) { LiteralTestUtil::ExpectR0Near( - i + j + i * kFanout + j, - result_literal->tuple_literals(i).tuple_literals(j), error_spec_); + i + j + i * kFanout + j, LiteralView::Create(*result_literal, {i, j}), + error_spec_); } } } @@ -554,11 +560,12 @@ XLA_TEST_F(LocalClientExecuteTest, DeepTuple) { ExecuteLocallyOrDie(computation, {arg_buffer.get()}); std::unique_ptr result_literal = ShapedBufferToLiteral(*result); - const Literal* result_element = result_literal.get(); + ShapeIndex index; for (int i = 0; i < kTupleDepth; ++i) { - result_element = &result_element->tuple_literals(0); + index.push_back(0); } - LiteralTestUtil::ExpectR0Equal(165.0, *result_element); + LiteralTestUtil::ExpectR0Equal( + 165.0, LiteralView::Create(*result_literal, index)); } XLA_TEST_F(LocalClientExecuteTest, InvalidNumberOfArguments) { @@ -575,7 +582,7 @@ XLA_TEST_F(LocalClientExecuteTest, InvalidNumberOfArguments) { EXPECT_FALSE(execute_status.ok()); EXPECT_THAT(execute_status.status().error_message(), - ContainsRegex("invalid number of arguments")); + ContainsRegex("Invalid number of arguments")); } XLA_TEST_F(LocalClientExecuteTest, IncorrectArgumentShape) { @@ -591,7 +598,7 @@ XLA_TEST_F(LocalClientExecuteTest, IncorrectArgumentShape) { EXPECT_FALSE(execute_status.ok()); EXPECT_THAT(execute_status.status().error_message(), - ContainsRegex("invalid argument shape")) + ContainsRegex("Invalid argument shape")) << execute_status.status(); } @@ -763,10 +770,10 @@ XLA_TEST_F(LocalClientExecuteTest, SelectBetweenTuples) { std::unique_ptr result = ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {}); std::unique_ptr tuple_literal = ShapedBufferToLiteral(*result); - LiteralTestUtil::ExpectR1Equal({2.0f, 4.0f, 6.0f}, - tuple_literal->tuple_literals(0)); - LiteralTestUtil::ExpectR1Equal({1.0f, 2.0f, 3.0f}, - tuple_literal->tuple_literals(1)); + LiteralTestUtil::ExpectR1Equal( + {2.0f, 4.0f, 6.0f}, LiteralView::Create(*tuple_literal, {0})); + LiteralTestUtil::ExpectR1Equal( + {1.0f, 2.0f, 3.0f}, LiteralView::Create(*tuple_literal, {1})); } XLA_TEST_F(LocalClientExecuteTest, CompileExecutable) { diff --git a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc index 0fb87c3c2ccbad387d46016cfad4e7d3cc537dcc..6c86dd5b9ef673c9facffafa37e00a859ce82010 100644 --- a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc +++ b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc @@ -221,5 +221,77 @@ INSTANTIATE_TEST_CASE_P(MatOpsDotAddTestInstances, MatOpsDotAddTest, ::testing::Combine(::testing::Bool(), ::testing::Bool(), ::testing::Bool())); +class MatOpsDotAddTest_bf16 + : public ClientLibraryTestBase, + public ::testing::WithParamInterface> {}; + +TEST_P(MatOpsDotAddTest_bf16, Dot_Add_2x2_2x2) { + bool row_major = std::get<0>(GetParam()); + bool add_lhs = std::get<1>(GetParam()); + bool transpose = std::get<2>(GetParam()); + Array2D lhs( + {{bfloat16(1.0f), bfloat16(2.0f)}, {bfloat16(3.0), bfloat16(4.0)}}); + Array2D rhs( + {{bfloat16(10.0f), bfloat16(11.0f)}, {bfloat16(12.0f), bfloat16(13.0f)}}); + + auto minor_to_major = [](bool row_major) -> std::vector { + return {row_major ? 1 : 0, row_major ? 0 : 1}; + }; + + auto prim_type = primitive_util::NativeToPrimitiveType(); + Shape lhs_shape = + ShapeUtil::MakeShape(prim_type, {lhs.height(), lhs.width()}); + Shape rhs_shape = + ShapeUtil::MakeShape(prim_type, {rhs.height(), rhs.width()}); + + TF_ASSERT_OK_AND_ASSIGN( + auto lhs_handle, + client_->TransferToServer( + *Literal::CreateR2FromArray2DWithLayout( + lhs, LayoutUtil::MakeLayout(minor_to_major(row_major))))); + TF_ASSERT_OK_AND_ASSIGN( + auto rhs_handle, + client_->TransferToServer( + *Literal::CreateR2FromArray2DWithLayout( + rhs, LayoutUtil::MakeLayout(minor_to_major(row_major))))); + + ComputationBuilder builder(client_, TestName()); + auto lhs_arg = builder.Parameter(0, lhs_shape, "lhs"); + auto lhs_mat_arg = lhs_arg; + if (transpose) { + lhs_mat_arg = builder.Transpose(lhs_mat_arg, {1, 0}); + } + auto rhs_arg = builder.Parameter(1, rhs_shape, "rhs"); + auto result = builder.Dot(lhs_mat_arg, rhs_arg); + Array2D expected; + if (add_lhs) { + result = builder.Add(result, lhs_arg); + if (transpose) { + expected = Array2D( + {{bfloat16(47), bfloat16(52)}, {bfloat16(71), bfloat16(78)}}); + } else { + expected = Array2D( + {{bfloat16(35), bfloat16(39)}, {bfloat16(81), bfloat16(89)}}); + } + } else { + result = builder.Add(result, rhs_arg); + if (transpose) { + expected = Array2D( + {{bfloat16(56), bfloat16(61)}, {bfloat16(80), bfloat16(87)}}); + } else { + expected = Array2D( + {{bfloat16(44), bfloat16(48)}, {bfloat16(90), bfloat16(98)}}); + } + } + + ComputeAndCompareR2(&builder, expected, + {lhs_handle.get(), rhs_handle.get()}, + ErrorSpec(1e-6)); +} + +INSTANTIATE_TEST_CASE_P(MatOpsDotAddTestInstances, MatOpsDotAddTest_bf16, + ::testing::Combine(::testing::Bool(), ::testing::Bool(), + ::testing::Bool())); + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc b/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc index 62d24a11fdb164ed6776d1e83877cf3acd319cc6..0a603f4954badd12adf3144320789a5edd0d9c6c 100644 --- a/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc +++ b/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc @@ -28,6 +28,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_runner.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" @@ -35,6 +36,7 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/tests/test_utils.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/protobuf.h" @@ -99,11 +101,11 @@ class MultiOutputFusionTest : public HloTestBase { nullptr); } - Literal arg1; - arg1.PopulateWithValue(2.5f, {size, size}); + Literal arg1(ShapeUtil::MakeShape(F32, {size, size})); + arg1.PopulateWithValue(2.5f); - Literal expect; - expect.PopulateWithValue(size * 1.5f * 3.5f, {size, size}); + Literal expect(ShapeUtil::MakeShape(F32, {size, size})); + expect.PopulateWithValue(size * 1.5f * 3.5f); auto actual = ExecuteAndTransfer( std::move(hlo_module), {Literal::CreateR0(-9.0f).get(), &arg1}); LiteralTestUtil::ExpectNear(expect, *actual, error_spec_); @@ -159,11 +161,12 @@ class MultiOutputFusionTest : public HloTestBase { nullptr); } - Literal input0, input1; - input0.PopulateWithValue(2.5f, {size}); - input1.PopulateWithValue(1, {size}); + Literal input0(ShapeUtil::MakeShape(F32, {size})); + input0.PopulateWithValue(2.5f); + Literal input1(ShapeUtil::MakeShape(F64, {size})); + input1.PopulateWithValue(1.); - Literal expect = *Literal::CreateR1({size * 1.5f * 3.5f}); + Literal expect = std::move(*Literal::CreateR1({size * 1.5f * 3.5f})); auto actual = ExecuteAndTransfer(std::move(hlo_module), {&input0, &input1}); LiteralTestUtil::ExpectNear(expect, *actual, error_spec_); } @@ -175,5 +178,38 @@ XLA_TEST_F(MultiOutputFusionTest, 2DFusionSize129) { RunTest2D(true, 129); } XLA_TEST_F(MultiOutputFusionTest, DiffentTypesNoFusion) { RunTest1D(false, 8); } XLA_TEST_F(MultiOutputFusionTest, DiffentTypesFusion) { RunTest1D(true, 8); } +XLA_TEST_F(MultiOutputFusionTest, FusionNodeIsRoot) { + const char* testcase = R"( + HloModule m + + fused_computation { + x.param_0 = (((s32[]), f32[]), (f32[], s32[])) parameter(0) + gte.3 = ((s32[]), f32[]) get-tuple-element(x.param_0), index=0 + gte.2 = (s32[]) get-tuple-element(gte.3), index=0 + gte.4 = s32[] get-tuple-element(gte.2), index=0 + copy = s32[] copy(gte.4) + ROOT tuple = (s32[]) tuple(copy) + } + + ENTRY thing.v3 { + x = (((s32[]), f32[]), (f32[], s32[])) parameter(0) + ROOT fusion = (s32[]) fusion(x), kind=kLoop, calls=fused_computation + } + )"; + auto module = + HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) + .ValueOrDie(); + auto param = Literal::MakeTupleOwned( + Literal::MakeTupleOwned( + Literal::MakeTupleOwned(Literal::CreateR0(42)), + Literal::CreateR0(1.0)), + Literal::MakeTupleOwned(Literal::CreateR0(3.0), + Literal::CreateR0(4))); + TF_ASSERT_OK_AND_ASSIGN(auto result, + Execute(std::move(module), {param.get()})); + EXPECT_TRUE(LiteralTestUtil::Equal( + *result, *Literal::MakeTupleOwned(Literal::CreateR0(42)))); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/pad_test.cc b/tensorflow/compiler/xla/tests/pad_test.cc index 3fd83a4c3b104831f03366339fb7b8b5d816a3f7..8cef8dd34dc7b16b1e58ded67d6b6a4ba79f20db 100644 --- a/tensorflow/compiler/xla/tests/pad_test.cc +++ b/tensorflow/compiler/xla/tests/pad_test.cc @@ -33,6 +33,14 @@ limitations under the License. namespace xla { namespace { +#ifdef XLA_BACKEND_SUPPORTS_BFLOAT16 +// Tests both F32 and BF16. +static std::array use_bfloat16_params{false, true}; +#else +// Only tests F32. +static std::array use_bfloat16_params{false}; +#endif + class PadTest : public ClientLibraryTestBase { protected: PadTest() { @@ -61,8 +69,22 @@ class PadTest : public ClientLibraryTestBase { PaddingConfig r4_padding_on_dim0_dim1_; }; +class PadTestFloat : public PadTest, + public ::testing::WithParamInterface { + protected: + PadTestFloat() { set_use_bfloat16(GetParam()); } + + ErrorSpec DefaultErrorSpec() const { + if (use_bfloat16()) { + return ErrorSpec(1e-3, 1e-3); + } else { + return ErrorSpec(1e-5, 1e-5); + } + } +}; + // Tests a Pad() with a zero-element input and output. -XLA_TEST_F(PadTest, Pad1DS0ToS0Array) { +XLA_TEST_P(PadTestFloat, Pad1DS0ToS0Array) { ComputationBuilder b(client_, TestName()); // Set up the padding configuration {low: 0, high: 0, interior: 0}. PaddingConfig padding_config; @@ -71,12 +93,13 @@ XLA_TEST_F(PadTest, Pad1DS0ToS0Array) { dimension->set_edge_padding_high(0); dimension->set_interior_padding(0); - b.Pad(b.ConstantR1({}), b.ConstantR0(0.1), padding_config); - ComputeAndCompareR1(&b, {}, {}, ErrorSpec(0.0001)); + b.Pad(AddParam(*Literal::CreateR1({}), &b), + AddParam(*Literal::CreateR0(0.1), &b), padding_config); + ComputeAndCompareR1(&b, {}, {}, DefaultErrorSpec()); } // Tests a Pad() with a zero-element input but a non-zero-element output. -XLA_TEST_F(PadTest, Pad1DS0ToS5Array) { +XLA_TEST_P(PadTestFloat, Pad1DS0ToS5Array) { ComputationBuilder b(client_, TestName()); // Set up the padding configuration {low: 3, high: 0, interior: 1}. PaddingConfig padding_config; @@ -85,12 +108,13 @@ XLA_TEST_F(PadTest, Pad1DS0ToS5Array) { dimension->set_edge_padding_high(4); dimension->set_interior_padding(7); - b.Pad(b.ConstantR1({}), b.ConstantR0(0.1), padding_config); + b.Pad(AddParam(*Literal::CreateR1({}), &b), + AddParam(*Literal::CreateR0(0.1), &b), padding_config); ComputeAndCompareR1(&b, std::vector(5, 0.1), {}, - ErrorSpec(0.0001)); + DefaultErrorSpec()); } -XLA_TEST_F(PadTest, Pad1DS3Array) { +XLA_TEST_P(PadTestFloat, Pad1DS3Array) { ComputationBuilder b(client_, TestName()); // Set up the padding configuration {low: 3, high: 0, interior: 1}. PaddingConfig padding_config; @@ -99,21 +123,21 @@ XLA_TEST_F(PadTest, Pad1DS3Array) { dimension->set_edge_padding_high(0); dimension->set_interior_padding(1); - b.Pad(b.ConstantR1({1, 2, 3}), b.ConstantR0(0.1), - padding_config); + b.Pad(AddParam(*Literal::CreateR1({1, 2, 3}), &b), + AddParam(*Literal::CreateR0(0.1), &b), padding_config); std::vector expected({0.1, 0.1, 0.1, 1, 0.1, 2, 0.1, 3}); - ComputeAndCompareR1(&b, expected, {}, ErrorSpec(0.0001)); + ComputeAndCompareR1(&b, expected, {}, DefaultErrorSpec()); } -XLA_TEST_F(PadTest, Pad4D_2x0x3x2_FloatArray) { +XLA_TEST_P(PadTestFloat, Pad4D_2x0x3x2_FloatArray) { ComputationBuilder b(client_, TestName()); - b.Pad(b.ConstantR4FromArray4D(Array4D(2, 0, 3, 2)), - b.ConstantR0(1.5), r4_padding_on_dim0_dim1_); + b.Pad(AddParam(Array4D(2, 0, 3, 2), &b), + AddParam(*Literal::CreateR0(1.5), &b), r4_padding_on_dim0_dim1_); ComputeAndCompareR4(&b, Array4D(5, 2, 3, 2, 1.5f), {}, - ErrorSpec(0.0001)); + DefaultErrorSpec()); } -TEST_F(PadTest, Pad4DFloat_1x1x3x2_Array) { +TEST_P(PadTestFloat, Pad4DFloat_1x1x3x2_Array) { ComputationBuilder b(client_, TestName()); auto input = MakeUnique>(1, 1, 3, 2); Array2D input_xy({ @@ -123,7 +147,7 @@ TEST_F(PadTest, Pad4DFloat_1x1x3x2_Array) { }); input->FillWithYX(input_xy); - b.Pad(b.ConstantR4FromArray4D(*input), b.ConstantR0(1.5), + b.Pad(AddParam(*input, &b), AddParam(*Literal::CreateR0(1.5), &b), r4_padding_on_dim0_dim1_); auto expected = MakeUnique>(2, 3, 3, 2); @@ -134,15 +158,15 @@ TEST_F(PadTest, Pad4DFloat_1x1x3x2_Array) { (*expected)(1, 0, 1, 1) = 4.0f; (*expected)(1, 0, 2, 0) = 5.0f; (*expected)(1, 0, 2, 1) = 6.0f; - ComputeAndCompareR4(&b, *expected, {}, ErrorSpec(0.0001)); + ComputeAndCompareR4(&b, *expected, {}, DefaultErrorSpec()); } -TEST_F(PadTest, Pad4DFloatArrayWithInteriorPadding) { +TEST_P(PadTestFloat, Pad4DFloatArrayWithInteriorPadding) { ComputationBuilder b(client_, TestName()); const float pad_value = 1.5f; Array4D input(3, 2, 1, 1, {1, 2, 3, 4, 5, 6}); - b.Pad(b.ConstantR4FromArray4D(input), b.ConstantR0(pad_value), + b.Pad(AddParam(input, &b), AddParam(*Literal::CreateR0(pad_value), &b), r4_padding_on_dim0_dim1_); auto expected = MakeUnique>(8, 5, 1, 1); @@ -156,7 +180,7 @@ TEST_F(PadTest, Pad4DFloatArrayWithInteriorPadding) { ComputeAndCompareR4(&b, *expected, {}, ErrorSpec(0.0001)); } -TEST_F(PadTest, Pad4DFloatArrayMinorFirstSmall) { +TEST_P(PadTestFloat, Pad4DFloatArrayMinorFirstSmall) { ComputationBuilder b(client_, TestName()); PaddingConfig padding_config; @@ -184,7 +208,8 @@ TEST_F(PadTest, Pad4DFloatArrayMinorFirstSmall) { auto input = Literal::CreateR4FromArray4D(input_array); input = input->Relayout(layout); - b.Pad(b.ConstantLiteral(*input), b.ConstantR0(pad_value), padding_config); + b.Pad(AddParam(*input, &b), + AddParam(*Literal::CreateR0(pad_value), &b), padding_config); Array4D expected_array(1, 1, 5, 8); expected_array.Fill(pad_value); @@ -197,7 +222,7 @@ TEST_F(PadTest, Pad4DFloatArrayMinorFirstSmall) { ComputeAndCompareR4(&b, expected_array, {}, ErrorSpec(0.0001)); } -XLA_TEST_F(PadTest, Pad4DFloatArrayMinorFirstNonTrivialMinorDimensions) { +XLA_TEST_P(PadTestFloat, Pad4DFloatArrayMinorFirstNonTrivialMinorDimensions) { ComputationBuilder b(client_, TestName()); PaddingConfig padding_config; @@ -229,7 +254,8 @@ XLA_TEST_F(PadTest, Pad4DFloatArrayMinorFirstNonTrivialMinorDimensions) { auto input = Literal::CreateR4FromArray4D(input_array); input = input->Relayout(layout); - b.Pad(b.ConstantLiteral(*input), b.ConstantR0(pad_value), padding_config); + b.Pad(AddParam(*input, &b), + AddParam(*Literal::CreateR0(pad_value), &b), padding_config); Array4D expected_array(1, 25, 17, 11); expected_array.Fill(pad_value); @@ -249,7 +275,7 @@ XLA_TEST_F(PadTest, Pad4DU8Array) { }); input->FillWithYX(input_xy); - b.Pad(b.ConstantR4FromArray4D(*input), b.ConstantR0(35), + b.Pad(AddParam(*input, &b), b.ConstantR0(35), r4_padding_on_dim0_dim1_); auto expected = MakeUnique>(2, 3, 3, 2); @@ -277,8 +303,7 @@ XLA_TEST_F(PadTest, Pad4DPredArray) { auto ones = MakeUnique>(2, 3, 3, 2); zeros->Fill(0); ones->Fill(1); - b.Select(padded, b.ConstantR4FromArray4D(*ones), - b.ConstantR4FromArray4D(*zeros)); + b.Select(padded, AddParam(*ones, &b), AddParam(*zeros, &b)); auto expected = MakeUnique>(2, 3, 3, 2); expected->Fill(0); @@ -291,10 +316,12 @@ XLA_TEST_F(PadTest, Pad4DPredArray) { ComputeAndCompareR4(&b, *expected, {}); } -XLA_TEST_F(PadTest, Large2DPad) { +XLA_TEST_P(PadTestFloat, Large2DPad) { ComputationBuilder b(client_, TestName()); - auto input = b.Parameter(0, ShapeUtil::MakeShape(F32, {4, 4}), "input"); + auto ones = MakeUnique>(4, 4); + ones->Fill(1.0f); + auto input = AddParam(*ones, &b); PaddingConfig padding_config = MakeNoPaddingConfig(2); for (int dim : {0, 1}) { padding_config.mutable_dimensions(dim)->set_edge_padding_low( @@ -302,25 +329,22 @@ XLA_TEST_F(PadTest, Large2DPad) { padding_config.mutable_dimensions(dim)->set_edge_padding_high(58 + 100 * dim); } - auto padded = b.Pad(input, b.ConstantR0(0.0f), padding_config); - - auto ones = MakeUnique>(4, 4); - ones->Fill(1.0f); - auto input_literal = Literal::CreateR2FromArray2D(*ones); - std::unique_ptr input_data = - client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + auto padded = b.Pad(input, AddParam(*Literal::CreateR0(0.0f), &b), + padding_config); auto expected = ReferenceUtil::PadArray2D(*ones, padding_config, 0.0f); - ComputeAndCompareR2(&b, *expected, {input_data.get()}); + ComputeAndCompareR2(&b, *expected, {}, DefaultErrorSpec()); } -XLA_TEST_F(PadTest, AllTypes2DPad) { +XLA_TEST_P(PadTestFloat, AllTypes2DPad) { ComputationBuilder b(client_, TestName()); constexpr int64 in_rows = 35; constexpr int64 in_cols = 35; - auto input = - b.Parameter(0, ShapeUtil::MakeShape(F32, {in_rows, in_cols}), "input"); + auto operand = MakeUnique>(in_rows, in_cols); + operand->FillUnique(0.0f); + auto input = AddParam(*operand, &b); + PaddingConfig padding_config = MakeNoPaddingConfig(2); padding_config.mutable_dimensions(0)->set_edge_padding_low(7); padding_config.mutable_dimensions(0)->set_edge_padding_high(5); @@ -328,20 +352,14 @@ XLA_TEST_F(PadTest, AllTypes2DPad) { padding_config.mutable_dimensions(1)->set_edge_padding_low(6); padding_config.mutable_dimensions(1)->set_edge_padding_high(4); padding_config.mutable_dimensions(1)->set_interior_padding(2); - auto padded = b.Pad(input, b.ConstantR0(3.14f), padding_config); - - auto operand = MakeUnique>(in_rows, in_cols); - operand->FillUnique(0.0f); - auto input_literal = Literal::CreateR2FromArray2D(*operand); - std::unique_ptr input_data = - client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + auto padded = b.Pad(input, AddParam(*Literal::CreateR0(3.14f), &b), + padding_config); auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 3.14f); - ComputeAndCompareR2(&b, *expected, {input_data.get()}, - ErrorSpec{0.0001}); + ComputeAndCompareR2(&b, *expected, {}, DefaultErrorSpec()); } -XLA_TEST_F(PadTest, High2DPad) { +XLA_TEST_P(PadTestFloat, High2DPad) { ComputationBuilder b(client_, TestName()); constexpr int64 in_rows = 129; @@ -349,8 +367,9 @@ XLA_TEST_F(PadTest, High2DPad) { constexpr int64 low_padding = 0; int64 high_padding[2] = {5, 7}; constexpr int64 interior_padding = 0; - auto input = - b.Parameter(0, ShapeUtil::MakeShape(F32, {in_rows, in_cols}), "input"); + auto operand = MakeUnique>(in_rows, in_cols); + operand->FillUnique(1.0f); + auto input = AddParam(*operand, &b); PaddingConfig padding_config = MakeNoPaddingConfig(2); for (int dim : {0, 1}) { padding_config.mutable_dimensions(dim)->set_edge_padding_low(low_padding); @@ -359,20 +378,15 @@ XLA_TEST_F(PadTest, High2DPad) { padding_config.mutable_dimensions(dim)->set_interior_padding( interior_padding); } - auto padded = b.Pad(input, b.ConstantR0(2.718f), padding_config); + auto padded = b.Pad(input, AddParam(*Literal::CreateR0(2.718f), &b), + padding_config); - auto operand = MakeUnique>(in_rows, in_cols); - operand->FillUnique(1.0f); - auto input_literal = Literal::CreateR2FromArray2D(*operand); auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f); - std::unique_ptr input_data = - client_->TransferToServer(*input_literal).ConsumeValueOrDie(); - ComputeAndCompareR2(&b, *expected, {input_data.get()}, - ErrorSpec(0.0001)); + ComputeAndCompareR2(&b, *expected, {}, DefaultErrorSpec()); } -XLA_TEST_F(PadTest, NegativePadding2D) { +XLA_TEST_P(PadTestFloat, NegativePadding2D) { ComputationBuilder b(client_, TestName()); constexpr int64 in_rows = 129; @@ -380,8 +394,9 @@ XLA_TEST_F(PadTest, NegativePadding2D) { int64 low_padding[2] = {-1, -2}; int64 high_padding[2] = {-3, 4}; constexpr int64 interior_padding = 0; - auto input = - b.Parameter(0, ShapeUtil::MakeShape(F32, {in_rows, in_cols}), "input"); + auto operand = MakeUnique>(in_rows, in_cols); + operand->FillUnique(1.0f); + auto input = AddParam(*operand, &b); PaddingConfig padding_config = MakeNoPaddingConfig(2); for (int dim : {0, 1}) { padding_config.mutable_dimensions(dim)->set_edge_padding_low( @@ -391,20 +406,15 @@ XLA_TEST_F(PadTest, NegativePadding2D) { padding_config.mutable_dimensions(dim)->set_interior_padding( interior_padding); } - auto padded = b.Pad(input, b.ConstantR0(2.718f), padding_config); + auto padded = b.Pad(input, AddParam(*Literal::CreateR0(2.718f), &b), + padding_config); - auto operand = MakeUnique>(in_rows, in_cols); - operand->FillUnique(1.0f); - auto input_literal = Literal::CreateR2FromArray2D(*operand); auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f); - std::unique_ptr input_data = - client_->TransferToServer(*input_literal).ConsumeValueOrDie(); - ComputeAndCompareR2(&b, *expected, {input_data.get()}, - ErrorSpec(0.0001)); + ComputeAndCompareR2(&b, *expected, {}, DefaultErrorSpec()); } -XLA_TEST_F(PadTest, NegativeAndInteriorPadding2D) { +XLA_TEST_P(PadTestFloat, NegativeAndInteriorPadding2D) { ComputationBuilder b(client_, TestName()); constexpr int64 in_rows = 8; @@ -412,8 +422,9 @@ XLA_TEST_F(PadTest, NegativeAndInteriorPadding2D) { int64 low_padding[2] = {4, -1}; int64 high_padding[2] = {-2, -4}; int64 interior_padding[2] = {1, 2}; - auto input = - b.Parameter(0, ShapeUtil::MakeShape(F32, {in_rows, in_cols}), "input"); + auto operand = MakeUnique>(in_rows, in_cols); + operand->FillUnique(1.0f); + auto input = AddParam(*operand, &b); PaddingConfig padding_config = MakeNoPaddingConfig(2); for (int dim : {0, 1}) { padding_config.mutable_dimensions(dim)->set_edge_padding_low( @@ -423,44 +434,40 @@ XLA_TEST_F(PadTest, NegativeAndInteriorPadding2D) { padding_config.mutable_dimensions(dim)->set_interior_padding( interior_padding[dim]); } - auto padded = b.Pad(input, b.ConstantR0(2.718f), padding_config); + auto padded = b.Pad(input, AddParam(*Literal::CreateR0(2.718f), &b), + padding_config); - auto operand = MakeUnique>(in_rows, in_cols); - operand->FillUnique(1.0f); - auto input_literal = Literal::CreateR2FromArray2D(*operand); auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f); - std::unique_ptr input_data = - client_->TransferToServer(*input_literal).ConsumeValueOrDie(); - ComputeAndCompareR2(&b, *expected, {input_data.get()}, - ErrorSpec(0.0001)); + ComputeAndCompareR2(&b, *expected, {}, DefaultErrorSpec()); } // Regression test for b/31827337. -XLA_TEST_F(PadTest, ReducePad) { +XLA_TEST_P(PadTestFloat, ReducePad) { ComputationBuilder b(client_, TestName()); - auto input = b.Parameter(0, ShapeUtil::MakeShape(F32, {2, 2, 2, 2}), "input"); + auto ones = MakeUnique>(2, 2, 2, 2); + ones->Fill(1.0); + auto input = AddParam(*ones, &b); - Computation add_f32 = CreateScalarAddComputation(F32, &b); - auto reduce = b.Reduce(input, b.ConstantR0(0.0), add_f32, {0}); + Computation add = CreateScalarAddComputation(FloatType(), &b); + auto reduce = + b.Reduce(input, AddParam(*Literal::CreateR0(0.0), &b), add, {0}); PaddingConfig padding_config = MakeNoPaddingConfig(3); padding_config.mutable_dimensions(0)->set_edge_padding_low(1); padding_config.mutable_dimensions(0)->set_edge_padding_high(1); - auto pad = b.Pad(reduce, b.ConstantR0(0.0), padding_config); - - auto ones = MakeUnique>(2, 2, 2, 2); - ones->Fill(1.0); - auto input_literal = Literal::CreateR4FromArray4D(*ones); - std::unique_ptr input_data = - client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + auto padded = b.Pad(reduce, AddParam(*Literal::CreateR0(0.0f), &b), + padding_config); Array3D expected({{{0.0, 0.0}, {0.0, 0.0}}, {{2.0, 2.0}, {2.0, 2.0}}, {{2.0, 2.0}, {2.0, 2.0}}, {{0.0, 0.0}, {0.0, 0.0}}}); - ComputeAndCompareR3(&b, expected, {input_data.get()}); + ComputeAndCompareR3(&b, expected, {}, DefaultErrorSpec()); } +INSTANTIATE_TEST_CASE_P(PadTestFloatInstantiation, PadTestFloat, + ::testing::ValuesIn(use_bfloat16_params)); + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/params_test.cc b/tensorflow/compiler/xla/tests/params_test.cc index c260258d6e9af6dee6075c92cf35dac4ed46abed..bb7e800df84121f2045141bc366c34b94ba694ea 100644 --- a/tensorflow/compiler/xla/tests/params_test.cc +++ b/tensorflow/compiler/xla/tests/params_test.cc @@ -436,8 +436,7 @@ XLA_TEST_F(ParamsTest, #endif -XLA_TEST_F(ParamsTest, - DISABLED_ON_CPU_PARALLEL(TupleOfR1ParametersAddedTogether)) { +XLA_TEST_F(ParamsTest, TupleOfR1ParametersAddedTogether) { ComputationBuilder builder(client_, TestName()); Shape r1f32_3 = ShapeUtil::MakeShape(F32, {3}); @@ -463,10 +462,8 @@ XLA_TEST_F(ParamsTest, // Verifies that passing a 2x2 with {0, 1} layout returns the same value back // when (transferred to the server and) passed through a parameter. XLA_TEST_F(ParamsTest, R2_2x2_Layout_01) { - std::unique_ptr literal = Literal::CreateR2({ - {1, 2}, {3, 4}, - }); - *literal->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); + std::unique_ptr literal = Literal::CreateR2WithLayout( + {{1, 2}, {3, 4}}, LayoutUtil::MakeLayout({0, 1})); ComputationBuilder builder(client_, TestName()); builder.Parameter(0, literal->shape(), "input"); @@ -477,10 +474,8 @@ XLA_TEST_F(ParamsTest, R2_2x2_Layout_01) { // As above, but for {1, 0} layout. XLA_TEST_F(ParamsTest, R2_2x2_Layout_10) { - std::unique_ptr literal = Literal::CreateR2({ - {1, 3}, {2, 4}, - }); - *literal->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({1, 0}); + std::unique_ptr literal = Literal::CreateR2WithLayout( + {{1, 3}, {2, 4}}, LayoutUtil::MakeLayout({1, 0})); ComputationBuilder builder(client_, TestName()); builder.Parameter(0, literal->shape(), "input"); @@ -501,7 +496,7 @@ XLA_TEST_F(ParamsTest, R2_2x2_TryToPassReverseLayoutToParameter) { original.layout().minor_to_major().begin(), original.layout().minor_to_major().end()); std::reverse(original_layout.begin(), original_layout.end()); - *literal->mutable_shape()->mutable_layout() = + *literal->mutable_shape_do_not_use()->mutable_layout() = LayoutUtil::MakeLayout(original_layout); ASSERT_EQ(2, literal->Get({0, 1})); } diff --git a/tensorflow/compiler/xla/tests/prng_test.cc b/tensorflow/compiler/xla/tests/prng_test.cc index 209f063cc5a34648453d12deae79f261b95dc3b4..6aafb9fa6cb2175c478f0e9a5e16f5808cbea590 100644 --- a/tensorflow/compiler/xla/tests/prng_test.cc +++ b/tensorflow/compiler/xla/tests/prng_test.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include "tensorflow/compiler/xla/client/computation_builder.h" @@ -36,65 +37,42 @@ namespace { class PrngTest : public ClientLibraryTestBase { protected: template - void UniformTest(T a, T b, tensorflow::gtl::ArraySlice dims); - void BernoulliTest(float p, tensorflow::gtl::ArraySlice dims); + std::unique_ptr UniformTest(T a, T b, + tensorflow::gtl::ArraySlice dims, + int64 seed = 42); // Computes the χ² statistic of a sample of the discrete uniform distribution // of the given range size. `expected_count` is the number of times each // possible value is expected to be generated. Thus, the sample size is // `range_size * expected_count`. - double UniformChiSquared(int32 range_size, int32 expected_count); + double UniformChiSquared(int32 range_size, int32 expected_count, + int64 seed = 42); }; template -void PrngTest::UniformTest(T a, T b, tensorflow::gtl::ArraySlice dims) { +std::unique_ptr PrngTest::UniformTest( + T a, T b, tensorflow::gtl::ArraySlice dims, int64 seed) { ComputationBuilder builder(client_, TestName()); builder.RngUniform( builder.ConstantR0(a), builder.ConstantR0(b), ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), dims)); - SetSeed(42); + SetSeed(seed); auto actual = ExecuteAndTransferOrDie(&builder, /*arguments=*/{}); EXPECT_THAT(dims, ::testing::ElementsAreArray(actual->shape().dimensions())); actual->EachCell([=](tensorflow::gtl::ArraySlice, T value) { EXPECT_LE(a, value); EXPECT_LT(value, b); }); -} - -void PrngTest::BernoulliTest(float p, tensorflow::gtl::ArraySlice dims) { - ComputationBuilder builder(client_, TestName()); - auto shape = ShapeUtil::MakeShape(U32, dims); - builder.RngBernoulli(builder.ConstantR0(p), shape); - - TF_ASSERT_OK_AND_ASSIGN(auto computation, builder.Build()); - ExecutionOptions execution_options = execution_options_; - execution_options.set_seed(42); - TF_ASSERT_OK_AND_ASSIGN( - auto actual, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, - &execution_options)); - EXPECT_THAT(dims, ::testing::ElementsAreArray(actual->shape().dimensions())); - int32 sum = 0; - actual->EachCell( - [&sum](tensorflow::gtl::ArraySlice, uint32 value) { - EXPECT_TRUE(value == 0 || value == 1); - sum += value; - }); - int32 total = ShapeUtil::ElementsIn(shape); - float p_tilde = sum / static_cast(total); - - // Test within expected range using normal approximation. The test uses a - // fixed seed and has a fixed output per p and backend. Using the normal - // approximation as this test is invoked for different `p` and the different - // backends could use different random number generators and produce different - // values. Choose 95% confidence level, so that z_{1-\alpha/2} = 1.96. - float normal_approximation_term = 1.96 * sqrt(p * (1 - p) / total); - EXPECT_GE(p_tilde, p - normal_approximation_term); - EXPECT_LE(p_tilde, p + normal_approximation_term); + return actual; } // Uniform random number generation tests XLA_TEST_F(PrngTest, ScalarU01) { UniformTest(0, 1, {}); } +XLA_TEST_F(PrngTest, ScalarU01limits) { + UniformTest(std::numeric_limits::min(), + std::numeric_limits::max(), {}); +} XLA_TEST_F(PrngTest, ZeroValuesU01) { UniformTest(0, 1, {0}); } XLA_TEST_F(PrngTest, TenValuesU01) { UniformTest(0, 1, {10}); } XLA_TEST_F(PrngTest, TenValuesU37) { UniformTest(3, 7, {10}); } @@ -102,6 +80,56 @@ XLA_TEST_F(PrngTest, ZeroValuesR2) { UniformTest(0, 1, {0, 20}); } XLA_TEST_F(PrngTest, LargeU01) { UniformTest(0, 1, {0x100, 0x100}); } XLA_TEST_F(PrngTest, TwelveValuesU524) { UniformTest(5, 24, {12}); } +// TODO(b/71543667): Fix Rng ops on LLVM backends. +XLA_TEST_F(PrngTest, DISABLED_ON_GPU(DISABLED_ON_CPU_PARALLEL( + DISABLED_ON_CPU(ScalarBF16Tests)))) { + for (int64 seed = 0; seed < 100; ++seed) { + // The largest negative number smaller than zero in bf16 that's not + // denormalized. + int32 low_raw = 0x80800000; + const float low = reinterpret_cast(low_raw); + float high = 0.0f; + UniformTest(static_cast(low), + static_cast(high), {}, /*seed=*/seed); + + // Test odd and even values. + UniformTest(static_cast(32.75), + static_cast(33), {}, /*seed=*/seed); + UniformTest(static_cast(32.50), + static_cast(32.75), {}, /*seed=*/seed); + UniformTest(static_cast(-33.00), + static_cast(-32.75), {}, /*seed=*/seed); + UniformTest(static_cast(-32.75), + static_cast(-32.50), {}, /*seed=*/seed); + } +} + +// TODO(b/71543667): Fix Rng ops on LLVM backends. +XLA_TEST_F(PrngTest, DISABLED_ON_GPU(DISABLED_ON_CPU( + DISABLED_ON_CPU_PARALLEL(ScalarBF16CountTests)))) { + // There are 3 BF16 values in the range of [32.25, 33): 32.25, 32.5, 32.75, + // they should get similar counts. + bfloat16 low = static_cast(32.25); + bfloat16 high = static_cast(33); + bfloat16 interval = static_cast(0.25); + std::vector counts(static_cast((high - low) / interval), 0); + + constexpr int64 count = 100; + for (int64 seed = 0; seed < count; ++seed) { + auto result = UniformTest(low, high, {}, /*seed=*/seed); + result->Literal::EachCell( + [&](tensorflow::gtl::ArraySlice, bfloat16 value) { + int64 index = static_cast((value - low) / interval); + counts[index]++; + }); + } + // Each bucket should have similar amount of counts. That is, not more than + // 10% of total counts. This mostly tests that we don't fall into a 1:2:2 + // distribution, which yields 20% expected difference. + EXPECT_LT(std::abs(counts[0] - counts[1]), count * 0.1); + EXPECT_LT(std::abs(counts[1] - counts[2]), count * 0.1); +} + namespace { template T Square(T x) { @@ -109,7 +137,8 @@ T Square(T x) { } } // namespace -double PrngTest::UniformChiSquared(int32 range_size, int32 expected_count) { +double PrngTest::UniformChiSquared(int32 range_size, int32 expected_count, + int64 seed) { int32 sample_size = range_size * expected_count; ComputationBuilder builder(client_, TestName()); @@ -117,7 +146,7 @@ double PrngTest::UniformChiSquared(int32 range_size, int32 expected_count) { builder.ConstantR0(range_size), ShapeUtil::MakeShape(S32, {sample_size})); - SetSeed(42); + SetSeed(seed); auto actual = ExecuteAndTransferOrDie(&builder, /*arguments=*/{}); std::vector counts(range_size, 0); actual->EachCell([&counts](tensorflow::gtl::ArraySlice, @@ -181,10 +210,12 @@ XLA_TEST_F(PrngTest, MapUsingRng) { computation, /*arguments=*/{param0_data.get()}, &execution_options)); - EXPECT_EQ(actual->f32s_size(), param0_literal->f32s_size()); - for (int i = 0; i < param0_literal->f32s_size(); ++i) { - EXPECT_GE(actual->f32s(i), param0_literal->f32s(i)); - EXPECT_LT(actual->f32s(i), param0_literal->f32s(i) + 1.0f); + EXPECT_EQ(ShapeUtil::ElementsIn(actual->shape()), + ShapeUtil::ElementsIn(param0_literal->shape())); + for (int i = 0; i < ShapeUtil::ElementsIn(actual->shape()); ++i) { + EXPECT_GE(actual->data()[i], param0_literal->data()[i]); + EXPECT_LT(actual->data()[i], + param0_literal->data()[i] + 1.0f); } } @@ -250,10 +281,6 @@ XLA_TEST_F(PrngTest, PassInGlobalRngSeed) { LiteralTestUtil::ExpectNotEqual(*result5, *result6); } -// Bernoulli random number generation tests -XLA_TEST_F(PrngTest, HundredValuesB10p5) { BernoulliTest(0.5, {100}); } -XLA_TEST_F(PrngTest, HundredValuesB10p1) { BernoulliTest(0.1, {100}); } - XLA_TEST_F(PrngTest, TenValuesN01) { ComputationBuilder builder(client_, TestName()); builder.RngNormal(builder.ConstantR0(0), builder.ConstantR0(1), diff --git a/tensorflow/compiler/xla/tests/reduce_hlo_test.cc b/tensorflow/compiler/xla/tests/reduce_hlo_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..c0a2c0ca4cb8414e0771a541b9f963f9aedc8376 --- /dev/null +++ b/tensorflow/compiler/xla/tests/reduce_hlo_test.cc @@ -0,0 +1,132 @@ +/* Copyright 2017 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 + +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/types.h" + +// Tests the Reduce HLO in ways that can't be done using the ComputationBuilder +// API. + +namespace xla { +namespace { + +namespace str_util = tensorflow::str_util; +namespace strings = tensorflow::strings; + +struct ReduceLayout { + std::array input_minor_to_major; + std::array output_minor_to_major; + + string ToString() const { + return strings::StrCat(str_util::Join(input_minor_to_major, "x"), "_", + str_util::Join(output_minor_to_major, "x")); + } +}; + +string PrintReduceLayout( + ::testing::TestParamInfo reduce_layout_param) { + return reduce_layout_param.param.ToString(); +} + +void PrintTo(const ReduceLayout& reduce_layout, ::std::ostream* os) { + *os << reduce_layout.ToString(); +} + +class ReduceWithLayoutTest + : public HloTestBase, + public ::testing::WithParamInterface {}; + +StatusOr> GetParsedModule() { + const char* const hlo_string = R"( +HloModule BadReduce + +Sum { + x.1 = f32[] parameter(0) + y.1 = f32[] parameter(1) + ROOT add.1 = f32[] add(x.1, y.1) +} + +ENTRY reduce.1 { + parameter = f32[2,2,2,3]{3,2,1,0} parameter(0) + init_value = f32[] constant(0) + reduce = f32[2,2,3]{2,1,0} reduce(parameter, init_value), dimensions={1}, to_apply=Sum + ROOT copy = f32[2,2,3]{2,1,0} copy(reduce) +} +)"; + + return tools::Parse(hlo_string); +} + +// TODO(b/72454718): XLA:GPU does not support executing code compiled without +// optimizations. +XLA_TEST_P(ReduceWithLayoutTest, DISABLED_ON_GPU(Reduce)) { + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, GetParsedModule()); + HloInstruction* reduce_instruction = + module->entry_computation()->root_instruction()->mutable_operand(0); + ASSERT_EQ(reduce_instruction->opcode(), HloOpcode::kReduce); + + const ReduceLayout& reduce_layout = GetParam(); + + Shape* reduce_output_shape = reduce_instruction->mutable_shape(); + *reduce_output_shape->mutable_layout() = + LayoutUtil::MakeLayout(reduce_layout.output_minor_to_major); + + Shape* reduce_input_shape = + reduce_instruction->mutable_operand(0)->mutable_shape(); + *reduce_input_shape->mutable_layout() = + LayoutUtil::MakeLayout(reduce_layout.input_minor_to_major); + + std::unique_ptr reduce_input = + Literal::CreateR4({{ /*i0=0*/ + {/*i1=0*/ + {-0.246092796, -0.179497838, -0.161181688}, + {-0.151643038, -0.240213156, -0.198156}}, + {/*i1=1*/ + {-0.14222312, -0.162200093, -0.193907976}, + {-0.239411, -0.198166847, -0.172471642}}}, + { /*i0=1*/ + {/*i1=0*/ + {-0.22965157, -0.218723893, -0.129257083}, + {-0.188762426, -0.16123569, -0.181166649}}, + {/*i1=1*/ + {-0.241772294, -0.245131493, -0.160247207}, + {-0.179881215, -0.23383224, -0.121976733}}}}); + + EXPECT_TRUE(RunAndCompareNoHloPasses(std::move(module), ErrorSpec(1e-5))); +} + +INSTANTIATE_TEST_CASE_P(ReduceWithLayoutTest_Instantiation, + ReduceWithLayoutTest, + ::testing::Values( // + ReduceLayout{{3, 2, 1, 0}, {0, 1, 2}}, // + ReduceLayout{{3, 2, 1, 0}, {0, 2, 1}}, // + ReduceLayout{{3, 2, 1, 0}, {1, 2, 0}}, // + ReduceLayout{{3, 2, 1, 0}, {1, 0, 2}}, // + ReduceLayout{{3, 2, 1, 0}, {2, 0, 1}}, // + ReduceLayout{{3, 2, 1, 0}, {2, 1, 0}}, // + ReduceLayout{{3, 1, 2, 0}, {1, 2, 0}}, // + ReduceLayout{{1, 2, 3, 0}, {1, 0, 2}}, // + ReduceLayout{{0, 2, 1, 3}, {2, 0, 1}}), // + PrintReduceLayout); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/reduce_precision_test.cc b/tensorflow/compiler/xla/tests/reduce_precision_test.cc index 4756ba096896806ece8fe35d18c4eaef041b8830..dc7ce3253cee255a7949326fa5b49fc8917432b8 100644 --- a/tensorflow/compiler/xla/tests/reduce_precision_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_precision_test.cc @@ -249,7 +249,9 @@ INSTANTIATE_TEST_CASE_P(ReducePrecisionAccuracyTest, // ReducePrecisionInsertion passes. class ReducePrecisionInsertionTest : public ClientLibraryTestBase {}; -XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionBeforeFusion) { +// The interpreter has no fusion pass, so skip this test. +XLA_TEST_F(ReducePrecisionInsertionTest, + DISABLED_ON_INTERPRETER(ReducePrecisionBeforeFusion)) { ComputationBuilder builder(client_, TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); @@ -276,7 +278,9 @@ XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionBeforeFusion) { ComputeAndCompareR1(&builder, {0.0f}, {a_data.get()}); } -XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionSkippedAfterFusion) { +// The interpreter has no fusion pass, so skip this test. +XLA_TEST_F(ReducePrecisionInsertionTest, + DISABLED_ON_INTERPRETER(ReducePrecisionSkippedAfterFusion)) { ComputationBuilder builder(client_, TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); @@ -300,7 +304,9 @@ XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionSkippedAfterFusion) { ComputeAndCompareR1(&builder, {-1.00001f}, {a_data.get()}); } -XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionAddedAfterFusion) { +// The interpreter has no fusion pass, so skip this test. +XLA_TEST_F(ReducePrecisionInsertionTest, + DISABLED_ON_INTERPRETER(ReducePrecisionAddedAfterFusion)) { ComputationBuilder builder(client_, TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); @@ -322,7 +328,9 @@ XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionAddedAfterFusion) { ComputeAndCompareR1(&builder, {-1.0f}, {a_data.get()}); } -XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionSkippedFusionContains) { +// The interpreter has no fusion pass, so skip this test. +XLA_TEST_F(ReducePrecisionInsertionTest, + DISABLED_ON_INTERPRETER(ReducePrecisionSkippedFusionContains)) { ComputationBuilder builder(client_, TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); @@ -345,7 +353,9 @@ XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionSkippedFusionContains) { ComputeAndCompareR1(&builder, {-1.00001f}, {a_data.get()}); } -XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionAddedFusionContains) { +// The interpreter has no fusion pass, so skip this test. +XLA_TEST_F(ReducePrecisionInsertionTest, + DISABLED_ON_INTERPRETER(ReducePrecisionAddedFusionContains)) { ComputationBuilder builder(client_, TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); diff --git a/tensorflow/compiler/xla/tests/reduce_test.cc b/tensorflow/compiler/xla/tests/reduce_test.cc index b09ccdd679b6c8f628e40f78f58dbd1734926af6..50d7b5074d201d2292cf90224ef4cd37efdbb8d3 100644 --- a/tensorflow/compiler/xla/tests/reduce_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_test.cc @@ -143,6 +143,55 @@ class ReduceTest : public ClientLibraryTestBase { ComputeAndCompareR0(&builder, expected, {input_global_data.get()}); } + // Reduce predicate tensor with dimension rows * cols to dimension cols, to + // test the implementation of atomic operations on misaligned small data + // types. + template + void RunR2ToR1PredTest(bool and_reduce, int64 rows, int64 minor = 1, + int64 major = 0) { + ComputationBuilder builder(client_, TestName()); + const Shape input_shape = ShapeUtil::MakeShape(U8, {rows, cols}); + auto input = builder.Parameter(0, input_shape, "input"); + auto input_pred = builder.Eq(input, builder.ConstantR0(1)); + + ComputationDataHandle init_value; + Computation reduce_op; + if (and_reduce) { + init_value = builder.ConstantR0(true); + reduce_op = CreateScalarAndComputation(&builder); + } else { + init_value = builder.ConstantR0(false); + reduce_op = CreateScalarOrComputation(&builder); + } + + builder.Reduce(input_pred, init_value, reduce_op, + /*dimensions_to_reduce=*/{0}); + + Array2D input_data(rows, cols); + input_data.FillRandom(0, 1); + std::unique_ptr input_literal = + Literal::CreateR2FromArray2D(input_data); + input_literal = + input_literal->Relayout(LayoutUtil::MakeLayout({minor, major})); + std::unique_ptr input_global_data = + client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + + std::array expected; + for (int64 colno = 0; colno < cols; ++colno) { + bool column_sum = and_reduce ? true : false; + for (int64 rowno = 0; rowno < rows; ++rowno) { + if (and_reduce) { + column_sum = column_sum && input_data(rowno, colno); + } else { + column_sum = column_sum || input_data(rowno, colno); + } + } + expected[colno] = column_sum; + } + + ComputeAndCompareR1(&builder, expected, {input_global_data.get()}); + } + // Runs an R2 => R0 reduction test with the given number of (rows, cols). void RunR2ToR0Test(int64 rows, int64 cols, int64 minor = 1, int64 major = 0) { ComputationBuilder builder(client_, TestName()); @@ -445,6 +494,26 @@ XLA_TEST_F(ReduceTest, TransposeAndReduceElementwiseR2_111x50_To_R1) { ErrorSpec(0.01, 1e-4)); } +// Test that algebraic simplifier does not incorrectly fold a transpose into a +// reduction operation. +XLA_TEST_F(ReduceTest, TransposeAndReduceR3_12x111x50_To_R2) { + ComputationBuilder builder(client_, TestName()); + Computation add_f32 = CreateScalarAddComputation(F32, &builder); + const Shape input_shape = ShapeUtil::MakeShape(F32, {12, 111, 50}); + ComputationDataHandle input = builder.Parameter(0, input_shape, "input"); + ComputationDataHandle zero = builder.ConstantR0(0.0); + ComputationDataHandle transpose = + builder.Transpose(input, /*permutation=*/{1, 0, 2}); + ComputationDataHandle reduce = + builder.Reduce(transpose, zero, add_f32, /*dimensions_to_reduce=*/{0}); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, + MakeFakeLiteral(input_shape)); + + ComputeAndCompare(&builder, reduce, {std::move(*input_data)}, + ErrorSpec(0.01, 1e-4)); +} + XLA_TEST_F(ReduceTest, Reshape_111x2x25Reduce_111x50_To_R1) { const int64 rows = 111, cols = 50; @@ -808,5 +877,12 @@ XLA_TEST_F(ReduceTest, DISABLED_ON_GPU(OperationOnConstantAsInitValue)) { ComputeAndCompareR0(&builder, 4.0f, {b_data.get()}); } +XLA_TEST_F(ReduceTest, ReduceAndPredR2_128x64_To_R1) { + RunR2ToR1PredTest(/*and_reduce=true*/ true, /*rows=128*/ 128); +} +XLA_TEST_F(ReduceTest, ReduceOrPredR2_64x32_To_R1) { + RunR2ToR1PredTest(/*and_reduce=false*/ false, /*rows=64*/ 64); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/reduce_window_test.cc b/tensorflow/compiler/xla/tests/reduce_window_test.cc index e857d5b50dbaf974bad2c7b167fff4ee16925788..b11b64e40a582150d6adf29e915cd70b4bcb982b 100644 --- a/tensorflow/compiler/xla/tests/reduce_window_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_window_test.cc @@ -138,6 +138,16 @@ TEST_P(ReduceWindowTest, Min3In5Stride2) { ErrorSpec(0.00001)); } +TEST_P(ReduceWindowTest, Min3In5Stride1WithSamePadding) { + const auto input = CreateConstantFromLiteral( + *Literal::CreateR1({10000, 1000, 100, 10, 1}), &builder_); + ReduceWindowMin(input, /*window_dimensions=*/{3}, /*window_strides=*/{1}, + Padding::kSame); + ComputeAndCompareLiteral(&builder_, + *Literal::CreateR1({1000, 100, 10, 1, 1}), {}, + ErrorSpec(0.00001)); +} + XLA_TEST_P(ReduceWindowTest, ZeroElementSmall) { Array4D input_array(1, 0, 2, 1); const auto input = CreateConstantFromArray(input_array, &builder_); @@ -242,65 +252,6 @@ TEST_P(ReduceWindowTest, AmongMajor2DimsMediumSize) { DefaultErrorSpec()); } -XLA_TEST_P(ReduceWindowTest, Add1x1x2In2x1x2) { - Array3D input_array(2, 1, 2); - input_array(0, 0, 0) = 1000; - input_array(0, 0, 1) = 100; - input_array(1, 0, 0) = 10; - input_array(1, 0, 1) = 1; - const auto input = CreateConstantFromArray(input_array, &builder_); - - ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 1}, Padding::kValid); - - Array3D expected(2, 1, 1); - expected(0, 0, 0) = 1100; - expected(1, 0, 0) = 11; - ComputeAndCompareLiteral(&builder_, *Literal::CreateFromArray(expected), {}, - DefaultErrorSpec()); -} - -XLA_TEST_P(ReduceWindowTest, Add1x1x2In2x1x3Stride1x1x2) { - Array3D input_array(2, 1, 3); - input_array(0, 0, 0) = 100; - input_array(0, 0, 1) = 10; - input_array(0, 0, 2) = 1; - input_array(1, 0, 0) = 500; - input_array(1, 0, 1) = 50; - input_array(1, 0, 2) = 5; - const auto input = CreateConstantFromArray(input_array, &builder_); - - ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 2}, Padding::kValid); - - Array3D expected(2, 1, 1); - expected(0, 0, 0) = 110; - expected(1, 0, 0) = 550; - ComputeAndCompareLiteral(&builder_, *Literal::CreateFromArray(expected), {}, - DefaultErrorSpec()); -} - -XLA_TEST_P(ReduceWindowTest, Add1x1x2In2x1x3SamePad) { - Array3D input_array(2, 1, 3); - input_array(0, 0, 0) = 100; - input_array(0, 0, 1) = 10; - input_array(0, 0, 2) = 1; - input_array(1, 0, 0) = 500; - input_array(1, 0, 1) = 50; - input_array(1, 0, 2) = 5; - const auto input = CreateConstantFromArray(input_array, &builder_); - - ReduceWindowAdd(input, {1, 1, 2}, {1, 1, 1}, Padding::kSame); - - Array3D expected(2, 1, 3); - expected(0, 0, 0) = 110; - expected(0, 0, 1) = 11; - expected(0, 0, 2) = 1; - expected(1, 0, 0) = 550; - expected(1, 0, 1) = 55; - expected(1, 0, 2) = 5; - ComputeAndCompareLiteral(&builder_, *Literal::CreateFromArray(expected), {}, - DefaultErrorSpec()); -} - // Tests a reduction function that is not a simple add/min/max/etc. XLA_TEST_P(ReduceWindowTest, NonstandardReduceFunction) { Array4D input_array(1, 2, 2, 1); @@ -321,7 +272,7 @@ XLA_TEST_P(ReduceWindowTest, NonstandardReduceFunction) { builder_.ReduceWindow( input, - CreateConstantFromLiteral(*Literal::CreateR0(3.0f), &builder_), + CreateConstantFromLiteral(*Literal::CreateR0(0.0f), &builder_), reduce_fn, /*window_dimensions=*/{1, 1, 2, 1}, /*window_strides=*/{1, 1, 1, 1}, padding); @@ -331,7 +282,7 @@ XLA_TEST_P(ReduceWindowTest, NonstandardReduceFunction) { }; auto expected = - ReferenceUtil::ReduceWindow4DGeneric(input_array, 3.0f, reduce_func, + ReferenceUtil::ReduceWindow4DGeneric(input_array, 0.0f, reduce_func, /*window=*/{1, 1, 2, 1}, /*stride=*/{1, 1, 1, 1}, padding); @@ -522,6 +473,23 @@ XLA_TEST_P(ReduceWindowTest, Add128In128Stride128) { DefaultErrorSpec()); } +XLA_TEST_P(ReduceWindowTest, Add128In128) { + std::vector input_vector{ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; + const auto input = CreateConstantFromLiteral( + *Literal::CreateR1(input_vector), &builder_); + ReduceWindowAdd(input, {128}, {1}, Padding::kValid); + ComputeAndCompareLiteral(&builder_, *Literal::CreateR1({1088}), {}, + DefaultErrorSpec()); +} + // Regression test for a bug that appeared in Inception (b/34784899). TEST_P(ReduceWindowTest, R2ReduceWindowInceptionFromBroadcast) { Array2D input_array(14, 14, 1.0f); @@ -565,6 +533,7 @@ struct R4ReduceWindowTestData { int64 strides[4]; int64 pad_low[4]; int64 pad_high[4]; + int64 layout[4]; Reducer reducer; }; @@ -580,7 +549,8 @@ string R4ReduceWindowTestDataToString( "__strides_", tensorflow::str_util::Join(param.strides, "x"), // "__pad_low_", tensorflow::str_util::Join(param.pad_low, "x"), // "__pad_high_", tensorflow::str_util::Join(param.pad_high, "x"), // - (param.reducer == kAdd) ? "add" : "max"); + "__layout_", tensorflow::str_util::Join(param.layout, "_"), // + (param.reducer == kAdd) ? "_add" : "_max"); CHECK(param.reducer == kAdd || param.reducer == kMax); // Test names are not allowed to contain the '-' character. @@ -607,7 +577,8 @@ class R4ReduceWindowTest : public ReduceWindowTestBase, param.base_bounds[2], param.base_bounds[3]); input.FillIota(1); std::unique_ptr input_literal = - Literal::CreateR4FromArray4D(input); + Literal::CreateR4FromArray4DWithLayout( + input, LayoutUtil::MakeLayout(param.layout)); ComputationDataHandle parameter; auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", &b, ¶meter); @@ -643,8 +614,13 @@ class R4ReduceWindowTest : public ReduceWindowTestBase, /*window=*/param.window_bounds, /*stride=*/param.strides, /*padding=*/padding); - ComputeAndCompareLiteral(&b, *Literal::CreateFromArray(*expected), - {input_arg.get()}, DefaultErrorSpec()); + std::unique_ptr expected_literal = + Literal::CreateFromArray(*expected); + const Shape& expected_shape_with_layout = ShapeUtil::MakeShapeWithLayout( + input_literal->shape().element_type(), + AsInt64Slice(expected_literal->shape().dimensions()), param.layout); + ComputeAndCompareLiteral(&b, *expected_literal, {input_arg.get()}, + DefaultErrorSpec(), &expected_shape_with_layout); } }; @@ -658,6 +634,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // Arbitrary padding (not kSame or kValid). @@ -666,6 +643,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{2, 2, 1, 1}, /*pad_low=*/{4, 4, 0, 0}, /*pad_high=*/{4, 4, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // Zero base bound edge case. @@ -674,6 +652,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With non-1x1 window. @@ -682,6 +661,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With max instead of add. @@ -690,6 +670,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kMax}, // With stride. @@ -698,6 +679,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{2, 4, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With low padding. @@ -706,6 +688,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{2, 2, 1, 1}, /*pad_low=*/{3, 2, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With high padding. @@ -714,6 +697,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{2, 2, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{2, 3, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // Window touches both sides of the padding simultaneously. @@ -722,6 +706,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{1, 1, 0, 0}, /*pad_high=*/{1, 1, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // Window is entirely in the padding for some positions. @@ -730,6 +715,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{4, 4, 0, 0}, /*pad_high=*/{4, 4, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // Zero base bound with padding edge case. @@ -738,6 +724,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{0, 1, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With stride, low padding and high padding. @@ -746,6 +733,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{3, 1, 1, 1}, /*pad_low=*/{10, 1, 0, 0}, /*pad_high=*/{2, 3, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With second minor dimension == 9. @@ -754,6 +742,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With minor dimension == 129. @@ -762,6 +751,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With minor dims reduction and non-overlapped stride. @@ -770,6 +760,7 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*strides=*/{1, 1, 2, 2}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, // With minor dims reduction and overlapped stride. @@ -777,7 +768,8 @@ const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { /*window_bounds=*/{1, 1, 4, 4}, /*strides=*/{1, 1, 2, 2}, /*pad_low=*/{0, 0, 0, 0}, - /*pad_high=*/{0, 0, 0, 0}, + /*pad_high=*/{1, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, }; @@ -789,15 +781,16 @@ INSTANTIATE_TEST_CASE_P( class R4ReduceWindowLargeTest : public R4ReduceWindowTest {}; -XLA_TEST_P(R4ReduceWindowLargeTest, DoIt) { DoIt(); } +XLA_TEST_P(R4ReduceWindowLargeTest, DISABLED_ON_INTERPRETER(DoIt)) { DoIt(); } // Test cases that are large/slow/failed. const R4ReduceWindowTestData kR4ReduceWindowLargeTestValues[] = { R4ReduceWindowTestData{/*base_bounds=*/{28, 28, 256, 128}, - /*window_bounds=*/{3, 3, 1, 1}, - /*strides=*/{1, 1, 1, 1}, + /*window_bounds=*/{3, 3, 1, 5}, + /*strides=*/{1, 1, 1, 5}, /*pad_low=*/{1, 1, 0, 0}, /*pad_high=*/{1, 1, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kMax}, R4ReduceWindowTestData{/*base_bounds=*/{112, 112, 64, 128}, @@ -805,7 +798,16 @@ const R4ReduceWindowTestData kR4ReduceWindowLargeTestValues[] = { /*strides=*/{2, 2, 1, 1}, /*pad_low=*/{0, 0, 0, 0}, /*pad_high=*/{1, 1, 0, 0}, + /*layout=*/{3, 2, 1, 0}, /*reducer=*/kAdd}, + + R4ReduceWindowTestData{/*base_bounds=*/{1, 1, 32768 - 3, 2}, + /*window_bounds=*/{1, 1, 4, 1}, + /*strides=*/{1, 1, 4, 1}, + /*pad_low=*/{0, 0, 1, 0}, + /*pad_high=*/{0, 0, 2, 0}, + /*layout=*/{3, 2, 1, 0}, + /*reducer=*/kMax}, }; INSTANTIATE_TEST_CASE_P( @@ -814,6 +816,146 @@ INSTANTIATE_TEST_CASE_P( ::testing::ValuesIn(use_bfloat16_params)), R4ReduceWindowTestDataToString); +class R4ReduceWindowAnyDimsTest : public R4ReduceWindowTest {}; + +// TODO(b/72234705): Fix the test cases failed on CPU and GPU. +XLA_TEST_P(R4ReduceWindowAnyDimsTest, + DISABLED_ON_CPU_PARALLEL(DISABLED_ON_CPU(DISABLED_ON_GPU(DoIt)))) { + DoIt(); +} + +const R4ReduceWindowTestData kR4ReduceWindowAnyDimsTestValues[] = { + R4ReduceWindowTestData{/*base_bounds=*/{4, 6, 17, 140}, + /*window_bounds=*/{2, 3, 4, 5}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, + /*reducer=*/kAdd}, + R4ReduceWindowTestData{/*base_bounds=*/{4, 6, 17, 140}, + /*window_bounds=*/{2, 3, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{3, 2, 1, 0}, + /*reducer=*/kMax}, + // With 0321 layout. + R4ReduceWindowTestData{/*base_bounds=*/{4, 6, 17, 140}, + /*window_bounds=*/{2, 3, 4, 5}, + /*strides=*/{1, 2, 3, 4}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{0, 3, 2, 1}, + /*reducer=*/kAdd}, + + // With 0123 layout. + R4ReduceWindowTestData{/*base_bounds=*/{4, 6, 17, 23}, + /*window_bounds=*/{2, 3, 7, 9}, + /*strides=*/{1, 2, 5, 8}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*layout=*/{0, 1, 2, 3}, + /*reducer=*/kAdd}, +}; + +INSTANTIATE_TEST_CASE_P( + R4ReduceWindowAnyDimsTestInstantiation, R4ReduceWindowAnyDimsTest, + ::testing::Combine(::testing::ValuesIn(kR4ReduceWindowAnyDimsTestValues), + ::testing::ValuesIn(use_bfloat16_params)), + R4ReduceWindowTestDataToString); + +struct R3ReduceWindowTestData { + int64 base_bounds[3]; + int64 window_bounds[3]; + int64 strides[3]; + int64 layout[3]; + Padding padding; + Reducer reducer; +} kR3TestCases[] = { + {/*base_bounds=*/{2, 1, 2}, /*window_bounds=*/{1, 1, 2}, + /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{4, 3, 3}, /*window_bounds=*/{2, 2, 2}, + /*strides=*/{2, 2, 2}, /*layout=*/{2, 1, 0}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{4, 3, 3}, /*window_bounds=*/{2, 2, 2}, + /*strides=*/{2, 2, 2}, /*layout=*/{2, 1, 0}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{6, 21, 3}, /*window_bounds=*/{2, 3, 2}, + /*strides=*/{1, 2, 2}, /*layout=*/{2, 1, 0}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{10, 21, 129}, /*window_bounds=*/{2, 9, 1}, + /*strides=*/{5, 2, 1}, /*layout=*/{2, 1, 0}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{6, 21, 3}, /*window_bounds=*/{2, 3, 2}, + /*strides=*/{1, 2, 2}, /*layout=*/{0, 1, 2}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{6, 21, 3}, /*window_bounds=*/{2, 3, 2}, + /*strides=*/{1, 2, 2}, /*layout=*/{1, 0, 2}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, +}; + +string R3ReduceWindowTestDataToString( + const ::testing::TestParamInfo< + ::testing::tuple>& data) { + const auto& param = ::testing::get<0>(data.param); + string str = tensorflow::strings::StrCat( + "base_bounds_", tensorflow::str_util::Join(param.base_bounds, "x"), + "__window_bounds_", tensorflow::str_util::Join(param.window_bounds, "x"), + "__strides_", tensorflow::str_util::Join(param.strides, "x"), + "__padding_", param.padding == Padding::kSame ? "same" : "valid", + "__layout_", param.layout[0], "_", param.layout[1], "_", param.layout[2], + "__reducer_", param.reducer == kAdd ? "add" : "max"); + if (::testing::get<1>(data.param)) { + str = tensorflow::strings::StrCat(str, "_bfloat16"); + } + return str; +} + +class R3ReduceWindowTest : public ReduceWindowTestBase, + public ::testing::WithParamInterface< + ::testing::tuple> { + protected: + R3ReduceWindowTest() { set_use_bfloat16(::testing::get<1>(GetParam())); } +}; + +TEST_P(R3ReduceWindowTest, Add) { + ComputationBuilder b(client_, TestName()); + const auto& param = ::testing::get<0>(GetParam()); + CHECK(param.reducer == kAdd); + + const float kInitValue = 0.0f; + Array3D input(param.base_bounds[0], param.base_bounds[1], + param.base_bounds[2], 1.0f); + std::unique_ptr input_literal = + Literal::CreateR3FromArray3DWithLayout( + input, LayoutUtil::MakeLayout(param.layout)); + + ComputationDataHandle parameter; + auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", + &b, ¶meter); + auto init_value = + CreateConstantFromLiteral(*Literal::CreateR0(kInitValue), &b); + b.ReduceWindow(/*operand=*/parameter, + /*init_value=*/init_value, + /*computation=*/CreateScalarAddComputation(FloatType(), &b), + /*window_dimensions=*/param.window_bounds, + /*window_strides=*/param.strides, /*padding=*/param.padding); + + auto expected = ReferenceUtil::ReduceWindow3DAdd( + /*operand=*/input, /*init=*/kInitValue, /*window=*/param.window_bounds, + /*stride=*/param.strides, /*padding=*/param.padding); + + ComputeAndCompareLiteral(&b, *Literal::CreateFromArray(*expected), + {input_arg.get()}, DefaultErrorSpec()); +} + +INSTANTIATE_TEST_CASE_P( + R3ReduceWindowTestInstantiation, R3ReduceWindowTest, + ::testing::Combine(::testing::ValuesIn(kR3TestCases), + ::testing::ValuesIn(use_bfloat16_params)), + R3ReduceWindowTestDataToString); + struct R2ReduceWindowTestData { int64 base_bounds[2]; int64 window_bounds[2]; @@ -882,37 +1024,39 @@ class R2ReduceWindowTest : public ReduceWindowTestBase, ::testing::tuple> { protected: R2ReduceWindowTest() { set_use_bfloat16(::testing::get<1>(GetParam())); } -}; -TEST_P(R2ReduceWindowTest, Add) { - ComputationBuilder b(client_, TestName()); - const auto& param = ::testing::get<0>(GetParam()); - CHECK(param.reducer == kAdd); + void DoIt() { + ComputationBuilder b(client_, TestName()); + const auto& param = ::testing::get<0>(GetParam()); + CHECK(param.reducer == kAdd); - const float kInitValue = 0.0f; - Array2D input(param.base_bounds[0], param.base_bounds[1], 1.0f); - std::unique_ptr input_literal = - Literal::CreateR2FromArray2DWithLayout( - input, LayoutUtil::MakeLayout(param.layout)); + const float kInitValue = 0.0f; + Array2D input(param.base_bounds[0], param.base_bounds[1], 1.0f); + std::unique_ptr input_literal = + Literal::CreateR2FromArray2DWithLayout( + input, LayoutUtil::MakeLayout(param.layout)); - ComputationDataHandle parameter; - auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", - &b, ¶meter); - auto init_value = - CreateConstantFromLiteral(*Literal::CreateR0(kInitValue), &b); - b.ReduceWindow(/*operand=*/parameter, - /*init_value=*/init_value, - /*computation=*/CreateScalarAddComputation(FloatType(), &b), - /*window_dimensions=*/param.window_bounds, - /*window_strides=*/param.strides, /*padding=*/param.padding); + ComputationDataHandle parameter; + auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", + &b, ¶meter); + auto init_value = + CreateConstantFromLiteral(*Literal::CreateR0(kInitValue), &b); + b.ReduceWindow(/*operand=*/parameter, + /*init_value=*/init_value, + /*computation=*/CreateScalarAddComputation(FloatType(), &b), + /*window_dimensions=*/param.window_bounds, + /*window_strides=*/param.strides, /*padding=*/param.padding); - auto expected = ReferenceUtil::ReduceWindow2DAdd( - /*operand=*/input, /*init=*/kInitValue, /*window=*/param.window_bounds, - /*stride=*/param.strides, /*padding=*/param.padding); + auto expected = ReferenceUtil::ReduceWindow2DAdd( + /*operand=*/input, /*init=*/kInitValue, /*window=*/param.window_bounds, + /*stride=*/param.strides, /*padding=*/param.padding); - ComputeAndCompareLiteral(&b, *Literal::CreateFromArray(*expected), - {input_arg.get()}, DefaultErrorSpec()); -} + ComputeAndCompareLiteral(&b, *Literal::CreateFromArray(*expected), + {input_arg.get()}, DefaultErrorSpec()); + } +}; + +TEST_P(R2ReduceWindowTest, DoIt) { DoIt(); } INSTANTIATE_TEST_CASE_P( R2ReduceWindowTestInstantiation, R2ReduceWindowTest, @@ -920,6 +1064,26 @@ INSTANTIATE_TEST_CASE_P( ::testing::ValuesIn(use_bfloat16_params)), R2ReduceWindowTestDataToString); +class R2ReduceWindowFailingCpuGpuBf16Test : public R2ReduceWindowTest {}; + +// TODO(b/72234705): Fix the test cases failed on CPU and GPU. +XLA_TEST_P(R2ReduceWindowFailingCpuGpuBf16Test, + DISABLED_ON_CPU_PARALLEL(DISABLED_ON_CPU(DISABLED_ON_GPU(DoIt)))) { + DoIt(); +} + +const R2ReduceWindowTestData kR2FailingValuesCpuGpuBf16Test[] = { + {/*base_bounds=*/{8, 128}, /*window_bounds=*/{8, 128}, + /*strides=*/{1, 1}, /*layout=*/{1, 0}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, +}; + +INSTANTIATE_TEST_CASE_P( + R2ReduceWindowFailingInstantiation, R2ReduceWindowFailingCpuGpuBf16Test, + ::testing::Combine(::testing::ValuesIn(kR2FailingValuesCpuGpuBf16Test), + ::testing::ValuesIn(use_bfloat16_params)), + R2ReduceWindowTestDataToString); + struct R1ReduceWindowTestData { int64 base_bounds[1]; int64 window_bounds[1]; @@ -1095,5 +1259,61 @@ INSTANTIATE_TEST_CASE_P( ::testing::Combine(::testing::ValuesIn(kR1TestCases), ::testing::ValuesIn(use_bfloat16_params)), R1ReduceWindowTestDataToString); + +// Test class for text-based test cases. Note that this compares with the +// results on the interpreter backend. +class ReduceWindowTextTest : public HloTestBase {}; + +TEST_F(ReduceWindowTextTest, R2General256x384) { + const string& hlo_string = R"( +HloModule R2Window +mul { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT mul = f32[] multiply(lhs, rhs) +} +ENTRY R2Window { + operand = f32[256,384]{1,0} parameter(0) + constant = f32[] constant(1) + ROOT reduce-window = f32[256,384]{1,0} reduce-window(operand, constant), window={size=2x3 pad=0_1x1_1}, to_apply=mul +} +)"; + EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{0.001})); +} + +TEST_F(ReduceWindowTextTest, R2General256x384Layout01) { + const string& hlo_string = R"( +HloModule R2Window +mul { +lhs = f32[] parameter(0) +rhs = f32[] parameter(1) +ROOT mul = f32[] multiply(lhs, rhs) +} +ENTRY R2Window { +operand = f32[256,384]{0,1} parameter(0) +constant = f32[] constant(1) +ROOT reduce-window = f32[256,384]{0,1} reduce-window(operand, constant), window={size=2x3 pad=0_1x1_1}, to_apply=mul +} +)"; + EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{0.001})); +} + +TEST_F(ReduceWindowTextTest, R2General2x5) { + const string& hlo_string = R"( +HloModule R2Window +mul { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT mul = f32[] multiply(lhs, rhs) +} +ENTRY R2Window { + operand = f32[2,5]{1,0} parameter(0) + constant = f32[] constant(1) + ROOT reduce-window = f32[3,5]{1,0} reduce-window(operand, constant), window={size=2x1 pad=0_2x0_0}, to_apply=mul +} +)"; + EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{0.001})); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/reshape_test.cc b/tensorflow/compiler/xla/tests/reshape_test.cc index ddd50d7a5864d73de7916ce736bb7cd40c1c4dc9..f7b04debd4f5c40a904e32c832b6fc384a03c33b 100644 --- a/tensorflow/compiler/xla/tests/reshape_test.cc +++ b/tensorflow/compiler/xla/tests/reshape_test.cc @@ -566,14 +566,15 @@ XLA_TEST_P(ReshapeTest, FullyConnectedCollapseDesugared) { XLA_TEST_P(ReshapeTest, ToScalar) { for (int rank = 0; rank < 8; ++rank) { ComputationBuilder b(client_, TestName()); - auto input_literal = Literal::CreateR1({83.0f}); std::vector ones(rank, 1); // this is {1, ..., 1}. std::vector dimensions(rank); std::iota(dimensions.begin(), dimensions.end(), 0); - *input_literal->mutable_shape() = ShapeUtil::MakeShape(F32, ones); + Literal input_literal(ShapeUtil::MakeShape(F32, ones)); + std::vector zeros(rank, 0); // this is {0, ..., 0}. + input_literal.Set(zeros, 83.0f); ComputationDataHandle parameter; - auto input = CreateParameterAndTransferLiteral(0, *input_literal, "input", + auto input = CreateParameterAndTransferLiteral(0, input_literal, "input", &b, ¶meter); b.Reshape(parameter, dimensions, {}); @@ -818,11 +819,9 @@ XLA_TEST_P(ReshapeTest, NoopReshape) { // data. if (use_bfloat16()) { auto expected = LiteralTestUtil::ConvertF32ToBF16(*input_literal); - EXPECT_EQ(tensorflow::gtl::ArraySlice(expected->bf16s()), - tensorflow::gtl::ArraySlice(output_literal->bf16s())); + EXPECT_EQ(expected->data(), output_literal->data()); } else { - EXPECT_EQ(tensorflow::gtl::ArraySlice(input_literal->f32s()), - tensorflow::gtl::ArraySlice(output_literal->f32s())); + EXPECT_EQ(input_literal->data(), output_literal->data()); } } diff --git a/tensorflow/compiler/xla/tests/reverse_test.cc b/tensorflow/compiler/xla/tests/reverse_test.cc index 1f6cfc85ccd25bb22db51411f7376489c14c3603..8fc841f14087cdea02fe44cdaea521ff92122aec 100644 --- a/tensorflow/compiler/xla/tests/reverse_test.cc +++ b/tensorflow/compiler/xla/tests/reverse_test.cc @@ -28,56 +28,89 @@ limitations under the License. namespace xla { namespace { -class ReverseTest : public ClientLibraryTestBase {}; - -// Tests the reverse operation on a scalar. -XLA_TEST_F(ReverseTest, ReverseScalar) { - ComputationBuilder b(client_, TestName()); - float input = 3.5f; - b.Rev(b.ConstantR0(input), {}); - ComputeAndCompareR0(&b, input, {}); -} - -// Tests the reverse operation on a 0x0 float array on both dimensions. -XLA_TEST_F(ReverseTest, Reverse0x0FloatArray) { - ComputationBuilder b(client_, TestName()); - b.Rev(b.ConstantR2FromArray2D(Array2D(0, 0)), {0, 1}); - ComputeAndCompareR2(&b, Array2D(0, 0), {}); -} - -// Tests the reverse operation on a 0x1 float array on both dimensions. -XLA_TEST_F(ReverseTest, Reverse0x1FloatArray) { - ComputationBuilder b(client_, TestName()); - b.Rev(b.ConstantR2FromArray2D(Array2D(0, 1)), {0, 1}); - ComputeAndCompareR2(&b, Array2D(0, 1), {}); +#ifdef XLA_BACKEND_SUPPORTS_BFLOAT16 +// Tests both F32 and BF16. +static std::array use_bfloat16_params{false, true}; +#else +// Only tests F32. +static std::array use_bfloat16_params{false}; +#endif + +struct ReverseSpec { + tensorflow::gtl::ArraySlice input_dims; + tensorflow::gtl::ArraySlice reversal; + bool use_bfloat16; + + string ToTestCaseName() const { + return tensorflow::strings::Printf( + "reverse_%s_in_dims_%s_%s", + tensorflow::str_util::Join(input_dims, "x").c_str(), + tensorflow::str_util::Join(reversal, "x").c_str(), + use_bfloat16 ? "bf16" : "f32"); + } +}; + +static std::vector GetTestCases() { + // clang-format off + return ExpandUseBfloat16( + use_bfloat16_params, + {{{}, {}}, + {{0, 0}, {0, 1}}, + {{0, 1}, {0, 1}}, + {{1, 0}, {0, 1}}, + {{1, 1}, {0, 1}}, + {{2, 0, 4, 3}, {0, 2}}, + {{2, 0, 4, 3}, {1, 3}}, + {{1, 2, 3, 4}, {0, 3}}, + {{4, 3, 2, 1}, {0, 1}}, + }); + // clang-format on } -// Tests the reverse operation on a 1x0 float array on both dimensions. -XLA_TEST_F(ReverseTest, Reverse1x0FloatArray) { - ComputationBuilder b(client_, TestName()); - b.Rev(b.ConstantR2FromArray2D(Array2D(1, 0)), {0, 1}); - ComputeAndCompareR2(&b, Array2D(1, 0), {}); +void PrintTo(const ReverseSpec& spec, std::ostream* os) { + *os << spec.ToTestCaseName(); } -// Tests the reverse operation on a 1x1 float array on both dimensions. -XLA_TEST_F(ReverseTest, Reverse1x1FloatArray) { - ComputationBuilder b(client_, TestName()); - Array2D input({{3.5f}}); - b.Rev(b.ConstantR2FromArray2D(input), {0, 1}); - ComputeAndCompareR2(&b, input, {}); +class FloatReverseTest : public ClientLibraryTestBase, + public ::testing::WithParamInterface { + public: + FloatReverseTest() { set_use_bfloat16(GetParam().use_bfloat16); } +}; + +TEST_P(FloatReverseTest, Reverses) { + const ReverseSpec& spec = GetParam(); + std::vector input_vector( + ShapeUtil::ElementsIn(ShapeUtil::MakeShape(F32, spec.input_dims))); + std::iota(input_vector.begin(), input_vector.end(), 0.0); + auto r1_literal = Literal::CreateR1(input_vector); + auto input_literal = r1_literal->Reshape(spec.input_dims).ConsumeValueOrDie(); + + ComputationBuilder builder(client_, TestName()); + auto a = AddParam(*input_literal, &builder); + builder.Rev(a, spec.reversal); + + std::unique_ptr expected = input_literal->CloneToUnique(); + std::vector output_indices(spec.input_dims.size()); + expected->EachCell( + [&](tensorflow::gtl::ArraySlice indices, float) { + for (int64 i = 0; i < indices.size(); ++i) { + output_indices[i] = indices[i]; + } + float value = input_literal->Get(indices); + for (int64 dim : spec.reversal) { + output_indices[dim] = (spec.input_dims[dim] - 1) - indices[dim]; + } + expected->Set(output_indices, value); + }); + ComputeAndCompareLiteral(&builder, *expected, {}); } -XLA_TEST_F(ReverseTest, Reverse2x0x4x3FloatArrayDim02) { - ComputationBuilder b(client_, TestName()); - b.Rev(b.ConstantR4FromArray4D(Array4D(2, 0, 4, 3)), {0, 2}); - ComputeAndCompareR4(&b, Array4D(2, 0, 4, 3), {}); -} +INSTANTIATE_TEST_CASE_P(FloatReverseInstance, FloatReverseTest, + ::testing::ValuesIn(GetTestCases()), + ::testing::PrintToStringParamName()); -XLA_TEST_F(ReverseTest, Reverse2x0x4x3FloatArrayDim13) { - ComputationBuilder b(client_, TestName()); - b.Rev(b.ConstantR4FromArray4D(Array4D(2, 0, 4, 3)), {1, 3}); - ComputeAndCompareR4(&b, Array4D(2, 0, 4, 3), {}); -} +// A simple test class which not templated by float precision. +class ReverseTest : public ClientLibraryTestBase {}; // Tests the reverse operation on a 4D U8 array on dimension 0 and 3. XLA_TEST_F(ReverseTest, Reverse4DU8ArrayOnDim23) { diff --git a/tensorflow/compiler/xla/tests/scalar_computations_test.cc b/tensorflow/compiler/xla/tests/scalar_computations_test.cc index b5e7570778ffeca66cc15d7cd2b153639637a647..4da6ee91607941b395b00befc98a10e7c17746ed 100644 --- a/tensorflow/compiler/xla/tests/scalar_computations_test.cc +++ b/tensorflow/compiler/xla/tests/scalar_computations_test.cc @@ -69,6 +69,13 @@ class ScalarComputationsTest : public ClientLibraryTestBase { } }; +XLA_TEST_F(ScalarComputationsTest, ReturnScalarF32) { + ComputationBuilder builder(client_, TestName()); + builder.ConstantR0(2.1f); + + ComputeAndCompareR0(&builder, 2.1f, {}, error_spec_); +} + XLA_TEST_F(ScalarComputationsTest, NegateScalarF32) { ComputationBuilder builder(client_, TestName()); builder.Neg(builder.ConstantR0(2.1f)); @@ -730,7 +737,61 @@ XLA_TEST_F(ScalarComputationsTest, PowScalar) { ComputeAndCompareR0(&builder, 8.0, {}, error_spec_); } -XLA_TEST_F(ScalarComputationsTest, ClampScalarHigh) { +XLA_TEST_F(ScalarComputationsTest, ClampScalarHighS32) { + ComputationBuilder builder(client_, TestName()); + builder.Clamp(builder.ConstantR0(-1), // The lower bound. + builder.ConstantR0(5), // The operand to be clamped. + builder.ConstantR0(3)); // The upper bound. + + ComputeAndCompareR0(&builder, 3, {}); +} + +XLA_TEST_F(ScalarComputationsTest, ClampScalarMiddleS32) { + ComputationBuilder builder(client_, TestName()); + builder.Clamp(builder.ConstantR0(-1), // The lower bound. + builder.ConstantR0(2), // The operand to be clamped. + builder.ConstantR0(3)); // The upper bound. + + ComputeAndCompareR0(&builder, 2, {}); +} + +XLA_TEST_F(ScalarComputationsTest, ClampScalarLowS32) { + ComputationBuilder builder(client_, TestName()); + builder.Clamp(builder.ConstantR0(-1), // The lower bound. + builder.ConstantR0(-5), // The operand to be clamped. + builder.ConstantR0(3)); // The upper bound. + + ComputeAndCompareR0(&builder, -1, {}); +} + +XLA_TEST_F(ScalarComputationsTest, ClampScalarHighU32) { + ComputationBuilder builder(client_, TestName()); + builder.Clamp(builder.ConstantR0(1), // The lower bound. + builder.ConstantR0(5), // The operand to be clamped. + builder.ConstantR0(3)); // The upper bound. + + ComputeAndCompareR0(&builder, 3, {}); +} + +XLA_TEST_F(ScalarComputationsTest, ClampScalarMiddleU32) { + ComputationBuilder builder(client_, TestName()); + builder.Clamp(builder.ConstantR0(1), // The lower bound. + builder.ConstantR0(2), // The operand to be clamped. + builder.ConstantR0(3)); // The upper bound. + + ComputeAndCompareR0(&builder, 2, {}); +} + +XLA_TEST_F(ScalarComputationsTest, ClampScalarLowU32) { + ComputationBuilder builder(client_, TestName()); + builder.Clamp(builder.ConstantR0(1), // The lower bound. + builder.ConstantR0(0), // The operand to be clamped. + builder.ConstantR0(3)); // The upper bound. + + ComputeAndCompareR0(&builder, 1, {}); +} + +XLA_TEST_F(ScalarComputationsTest, ClampScalarHighF32) { ComputationBuilder builder(client_, TestName()); builder.Clamp(builder.ConstantR0(2.0f), // The lower bound. builder.ConstantR0(5.0f), // The operand to be clamped. @@ -739,7 +800,7 @@ XLA_TEST_F(ScalarComputationsTest, ClampScalarHigh) { ComputeAndCompareR0(&builder, 3.0, {}, error_spec_); } -XLA_TEST_F(ScalarComputationsTest, ClampScalarMiddle) { +XLA_TEST_F(ScalarComputationsTest, ClampScalarMiddleF32) { ComputationBuilder builder(client_, TestName()); builder.Clamp(builder.ConstantR0(2.0f), // The lower bound. builder.ConstantR0(2.5f), // The operand to be clamped. @@ -748,7 +809,7 @@ XLA_TEST_F(ScalarComputationsTest, ClampScalarMiddle) { ComputeAndCompareR0(&builder, 2.5, {}, error_spec_); } -XLA_TEST_F(ScalarComputationsTest, ClampScalarLow) { +XLA_TEST_F(ScalarComputationsTest, ClampScalarLowF32) { ComputationBuilder builder(client_, TestName()); builder.Clamp(builder.ConstantR0(2.0f), // The lower bound. builder.ConstantR0(-5.0f), // The operand to be clamped. @@ -845,5 +906,12 @@ XLA_TEST_F(ScalarComputationsTest, SqrtF320) { ComputeAndCompareR0(&builder, 0.0f, {zero_data.get()}, error_spec_); } +XLA_TEST_F(ScalarComputationsTest, RoundScalar) { + ComputationBuilder builder(client_, TestName()); + builder.Round(builder.ConstantR0(1.4f)); + + ComputeAndCompareR0(&builder, 1.0f, {}, error_spec_); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/select_and_scatter_test.cc b/tensorflow/compiler/xla/tests/select_and_scatter_test.cc index 62ff349e9c011e0eb845192013a74aeb0956b791..9ee94b8571e5fc8789b60501462986967ce909a0 100644 --- a/tensorflow/compiler/xla/tests/select_and_scatter_test.cc +++ b/tensorflow/compiler/xla/tests/select_and_scatter_test.cc @@ -39,8 +39,8 @@ namespace xla { namespace { struct SelectAndScatterTestParam { - Array4D operand_shape; - Array4D source_shape; + std::vector operand_shape; + std::vector source_shape; Padding padding_type; tensorflow::gtl::ArraySlice window_dimensions; tensorflow::gtl::ArraySlice window_strides; @@ -69,83 +69,132 @@ class SelectAndScatterTest Computation min_f32_; }; -XLA_TEST_P(SelectAndScatterTest, R4Randomized) { - Array4D o(GetParam().operand_shape); +XLA_TEST_P(SelectAndScatterTest, ParamTest) { + auto operand_shape = GetParam().operand_shape; + Array o(operand_shape); o.FillRandom(1.5f); - auto operand = builder_.ConstantR4FromArray4D(o); + auto operand = builder_.ConstantFromArray(o); - Array4D s(GetParam().source_shape); + auto source_shape = GetParam().source_shape; + Array s(source_shape); s.FillRandom(12.0f); - auto source = builder_.ConstantR4FromArray4D(s); - - builder_.SelectAndScatter(operand, ge_f32_, GetParam().window_dimensions, - GetParam().window_strides, GetParam().padding_type, - source, builder_.ConstantR0(0.0f), add_f32_); + auto source = builder_.ConstantFromArray(s); - auto e = ReferenceUtil::SelectAndScatter4DGePlus( - o, s, 0.0f, GetParam().window_dimensions, GetParam().window_strides, - GetParam().padding_type == Padding::kSame); + auto select_and_scatter = builder_.SelectAndScatter( + operand, ge_f32_, GetParam().window_dimensions, GetParam().window_strides, + GetParam().padding_type, source, builder_.ConstantR0(0.0f), + add_f32_); - ComputeAndCompareR4(&builder_, *e, {}, ErrorSpec(1e-5)); + ComputeAndCompare(&builder_, select_and_scatter, {}, ErrorSpec(1e-5)); } INSTANTIATE_TEST_CASE_P( SelectAndScatterTest_Instantiation, SelectAndScatterTest, - ::testing::Values(SelectAndScatterTestParam{{6, 6, 256, 128}, - {3, 3, 256, 128}, - Padding::kSame, - {3, 3, 1, 1}, - {2, 2, 1, 1}}, - SelectAndScatterTestParam{{7, 7, 256, 128}, - {3, 3, 256, 128}, - Padding::kValid, - {3, 3, 1, 1}, - {2, 2, 1, 1}}, - SelectAndScatterTestParam{{6, 7, 256, 128}, - {3, 3, 256, 128}, - Padding::kValid, - {2, 3, 1, 1}, - {2, 2, 1, 1}}, - SelectAndScatterTestParam{{6, 7, 256, 128}, - {2, 3, 256, 128}, - Padding::kValid, - {2, 3, 1, 1}, - {3, 2, 1, 1}}, - SelectAndScatterTestParam{{9, 9, 16, 128}, - {3, 3, 16, 128}, - Padding::kValid, - {3, 3, 1, 1}, - {3, 3, 1, 1}}, - SelectAndScatterTestParam{{3, 3, 4, 4}, - {1, 1, 4, 4}, - Padding::kValid, - {3, 3, 1, 1}, - {3, 3, 1, 1}}, - SelectAndScatterTestParam{{3, 3, 4, 4}, - {1, 1, 4, 4}, - Padding::kValid, - {3, 3, 1, 1}, - {3, 3, 1, 1}}, - SelectAndScatterTestParam{{9, 3, 4, 4}, - {3, 1, 4, 4}, - Padding::kValid, - {3, 3, 1, 1}, - {3, 3, 1, 1}}, - SelectAndScatterTestParam{{7, 3, 4, 4}, - {3, 1, 4, 4}, - Padding::kValid, - {3, 3, 1, 1}, - {2, 3, 1, 1}}, - SelectAndScatterTestParam{{1, 1, 5, 5}, - {1, 1, 5, 5}, - Padding::kSame, - {3, 3, 1, 1}, - {3, 3, 1, 1}}, - SelectAndScatterTestParam{{7, 7, 8, 256}, - {4, 4, 8, 256}, - Padding::kSame, - {2, 2, 1, 1}, - {2, 2, 1, 1}})); + ::testing::Values( + SelectAndScatterTestParam{{6, 6, 6, 4, 4}, + {3, 3, 3, 4, 4}, + Padding::kSame, + {3, 3, 3, 1, 1}, + {2, 2, 2, 1, 1}}, + SelectAndScatterTestParam{{7, 7, 7, 4, 4}, + {3, 3, 3, 4, 4}, + Padding::kValid, + {3, 3, 3, 1, 1}, + {2, 2, 2, 1, 1}}, + + SelectAndScatterTestParam{{8, 8, 8, 4, 4}, + {1, 3, 3, 4, 4}, + Padding::kValid, + {8, 4, 4, 1, 1}, + {1, 2, 2, 1, 1}}, + SelectAndScatterTestParam{{6, 6, 256, 128}, + {3, 3, 256, 128}, + Padding::kSame, + {3, 3, 1, 1}, + {2, 2, 1, 1}}, + SelectAndScatterTestParam{{7, 7, 256, 128}, + {3, 3, 256, 128}, + Padding::kValid, + {3, 3, 1, 1}, + {2, 2, 1, 1}}, + SelectAndScatterTestParam{{6, 7, 256, 128}, + {3, 3, 256, 128}, + Padding::kValid, + {2, 3, 1, 1}, + {2, 2, 1, 1}}, + SelectAndScatterTestParam{{6, 7, 256, 128}, + {2, 3, 256, 128}, + Padding::kValid, + {2, 3, 1, 1}, + {3, 2, 1, 1}}, + SelectAndScatterTestParam{{9, 9, 16, 128}, + {3, 3, 16, 128}, + Padding::kValid, + {3, 3, 1, 1}, + {3, 3, 1, 1}}, + SelectAndScatterTestParam{{3, 3, 4, 4}, + {1, 1, 4, 4}, + Padding::kValid, + {3, 3, 1, 1}, + {3, 3, 1, 1}}, + SelectAndScatterTestParam{{3, 3, 4, 4}, + {1, 1, 4, 4}, + Padding::kValid, + {3, 3, 1, 1}, + {3, 3, 1, 1}}, + SelectAndScatterTestParam{{9, 3, 4, 4}, + {3, 1, 4, 4}, + Padding::kValid, + {3, 3, 1, 1}, + {3, 3, 1, 1}}, + SelectAndScatterTestParam{{7, 3, 4, 4}, + {3, 1, 4, 4}, + Padding::kValid, + {3, 3, 1, 1}, + {2, 3, 1, 1}}, + SelectAndScatterTestParam{{1, 1, 5, 5}, + {1, 1, 5, 5}, + Padding::kSame, + {3, 3, 1, 1}, + {3, 3, 1, 1}}, + SelectAndScatterTestParam{{7, 7, 8, 256}, + {4, 4, 8, 256}, + Padding::kSame, + {2, 2, 1, 1}, + {2, 2, 1, 1}}, + SelectAndScatterTestParam{ + {6, 4, 4}, {3, 4, 4}, Padding::kSame, {3, 1, 1}, {2, 1, 1}}, + SelectAndScatterTestParam{ + {6, 256, 128}, {3, 256, 128}, Padding::kSame, {3, 1, 1}, {2, 1, 1}}, + SelectAndScatterTestParam{{7, 256, 128}, + {3, 256, 128}, + Padding::kValid, + {3, 1, 1}, + {2, 1, 1}}, + SelectAndScatterTestParam{{6, 256, 128}, + {3, 256, 128}, + Padding::kValid, + {2, 1, 1}, + {2, 1, 1}}, + SelectAndScatterTestParam{{6, 256, 128}, + {2, 256, 128}, + Padding::kValid, + {2, 1, 1}, + {3, 1, 1}}, + SelectAndScatterTestParam{ + {9, 16, 128}, {3, 16, 128}, Padding::kValid, {3, 1, 1}, {3, 1, 1}}, + SelectAndScatterTestParam{ + {3, 4, 4}, {1, 4, 4}, Padding::kValid, {3, 1, 1}, {3, 1, 1}}, + SelectAndScatterTestParam{ + {3, 4, 4}, {1, 4, 4}, Padding::kValid, {3, 1, 1}, {3, 1, 1}}, + SelectAndScatterTestParam{ + {9, 4, 4}, {3, 4, 4}, Padding::kValid, {3, 1, 1}, {3, 1, 1}}, + SelectAndScatterTestParam{ + {7, 4, 4}, {3, 4, 4}, Padding::kValid, {3, 1, 1}, {2, 1, 1}}, + SelectAndScatterTestParam{ + {1, 5, 5}, {1, 5, 5}, Padding::kSame, {3, 1, 1}, {3, 1, 1}}, + SelectAndScatterTestParam{ + {7, 8, 256}, {4, 8, 256}, Padding::kSame, {2, 1, 1}, {2, 1, 1}})); // Test for F32 1D array, with a zero-element input. XLA_TEST_F(SelectAndScatterTest, R1S0F32) { diff --git a/tensorflow/compiler/xla/tests/slice_test.cc b/tensorflow/compiler/xla/tests/slice_test.cc index 4db566f7841829359ea06fe25408048418c547ad..ac163df127e0087c02777fa3d5ce7970c51b97b9 100644 --- a/tensorflow/compiler/xla/tests/slice_test.cc +++ b/tensorflow/compiler/xla/tests/slice_test.cc @@ -34,7 +34,6 @@ namespace xla { namespace { using ::tensorflow::str_util::Join; -using ::tensorflow::strings::StrCat; class SliceTest : public ClientLibraryTestBase {}; @@ -383,7 +382,7 @@ struct R4Spec { string R4SpecToString(const ::testing::TestParamInfo& data) { const R4Spec& spec = data.param; - return StrCat( // + return tensorflow::strings::StrCat( // "input_", Join(spec.input_dims, "x"), // "__layout_", Join(spec.input_layout, ""), // "__starts_", Join(spec.slice_starts, "x"), // diff --git a/tensorflow/compiler/xla/tests/test_macros.h b/tensorflow/compiler/xla/tests/test_macros.h index 28a2d0198a707cec1aa5e0fbed341ee9b2a927f7..e2d406f66d94f8ec76faa5b7d2d2e84dcaf6db57 100644 --- a/tensorflow/compiler/xla/tests/test_macros.h +++ b/tensorflow/compiler/xla/tests/test_macros.h @@ -36,6 +36,7 @@ limitations under the License. #define DISABLED_ON_CPU(X) X #define DISABLED_ON_CPU_PARALLEL(X) X #define DISABLED_ON_GPU(X) X +#define DISABLED_ON_INTERPRETER(X) X // We need this macro instead of pasting directly to support nesting // the DISABLED_ON_FOO macros, as in the definition of DISABLED_ON_CPU. @@ -62,6 +63,11 @@ limitations under the License. # define DISABLED_ON_GPU(X) XLA_TEST_PASTE(DISABLED_, X) #endif // XLA_TEST_BACKEND_GPU +#ifdef XLA_TEST_BACKEND_INTERPRETER +# undef DISABLED_ON_INTERPRETER +# define DISABLED_ON_INTERPRETER(X) XLA_TEST_PASTE(DISABLED_, X) +#endif // XLA_TEST_BACKEND_INTERPRETER + // clang-format on namespace xla { @@ -155,4 +161,31 @@ string PrependDisabledIfIndicated(const string& test_case_name, #define XLA_TEST_P(test_case_name, test_name) \ XLA_TEST_P_IMPL_(test_case_name, test_name) + +// This is identical to the TEST_F macro from "gtest", but it potentially +// disables the test based on an external manifest file, DISABLED_MANIFEST. +#define XLA_TYPED_TEST(CaseName, TestName) \ + template \ + class GTEST_TEST_CLASS_NAME_(CaseName, TestName) \ + : public CaseName { \ + private: \ + typedef CaseName TestFixture; \ + typedef gtest_TypeParam_ TypeParam; \ + virtual void TestBody(); \ + }; \ + bool gtest_##CaseName##_##TestName##_registered_ GTEST_ATTRIBUTE_UNUSED_ = \ + ::testing::internal::TypeParameterizedTest< \ + CaseName, \ + ::testing::internal::TemplateSel, \ + GTEST_TYPE_PARAMS_(CaseName)>:: \ + Register( \ + "", ::testing::internal::CodeLocation(__FILE__, __LINE__), \ + #CaseName, \ + ::xla::PrependDisabledIfIndicated(#CaseName, #TestName).c_str(), \ + 0); \ + template \ + void GTEST_TEST_CLASS_NAME_(CaseName, \ + TestName)::TestBody() + #endif // TENSORFLOW_COMPILER_XLA_TESTS_TEST_MACROS_H_ diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc index f9c62ec217d085e5c5a55f484c4bd712c6ccf05a..b060fb13b1451aab30cfca73bea0a4a598a9fa3a 100644 --- a/tensorflow/compiler/xla/tests/test_utils.cc +++ b/tensorflow/compiler/xla/tests/test_utils.cc @@ -24,44 +24,127 @@ namespace xla { namespace { template -void PopulateWithRandomFloatingPointData(Literal* literal) { - // TODO(b/69179121): Generate data that is less self-similar. +void PopulateWithRandomFloatingPointData(Literal* literal, + std::minstd_rand0* engine) { CHECK_EQ(literal->shape().element_type(), primitive_util::NativeToPrimitiveType()); - std::minstd_rand0 engine; - std::uniform_real_distribution generator(0.0f, 1.0f); + // Create uniform numbers between 1 and 1.125 to avoid creating denormal + // numbers. + std::uniform_real_distribution generator(1.0f, 1.125f); + const bool should_index_bias = ShapeUtil::ElementsIn(literal->shape()) > 1000; TF_CHECK_OK(literal->Populate( - [&](tensorflow::gtl::ArraySlice /*indices*/) { - return generator(engine); + [&](tensorflow::gtl::ArraySlice indices) { + // Generate a random uniform number from -0.0625 and 0.0625 and bias it + // with a position dependent number with mean 0.037109375. These number + // should allow for long chains of accumulation without being too close + // to zero or too large to accumulate all numbers accurately. Only do + // this for large literals where the number of elements is much greater + // than 47 otherwise only negative values are produced. + // + // The value is positionally biased using a product of the indices. Add + // one to each index value to avoid collapsing to zero if any of the + // indices are zero. + int64 index_product = 1; + for (int64 i : indices) { + index_product *= (1 + i); + } + const int64 negative_bias = should_index_bias ? 47 : 0; + FloatT index_bias = + static_cast(index_product % 113 - negative_bias) / + static_cast(256.0f); + return (generator(*engine) - 1.0625) + index_bias; })); } // The standard library does not have a case for bfloat16, unsurprisingly, so we // handle that one specially. template <> -void PopulateWithRandomFloatingPointData(Literal* literal) { +void PopulateWithRandomFloatingPointData(Literal* literal, + std::minstd_rand0* engine) { CHECK_EQ(literal->shape().element_type(), BF16); - std::minstd_rand0 engine; - std::uniform_real_distribution generator(0.0f, 1.0f); + std::uniform_real_distribution generator(-0.9f, 1.0f); TF_CHECK_OK(literal->Populate( [&](tensorflow::gtl::ArraySlice /*indices*/) { - return static_cast(generator(engine)); + return static_cast(generator(*engine)); })); } template -void PopulateWithRandomIntegralData(Literal* literal) { +void PopulateWithRandomIntegralData(Literal* literal, + std::minstd_rand0* engine) { CHECK_EQ(literal->shape().element_type(), primitive_util::NativeToPrimitiveType()); - std::minstd_rand0 engine; std::uniform_int_distribution generator( std::numeric_limits::lowest(), std::numeric_limits::max()); TF_CHECK_OK(literal->Populate( [&](tensorflow::gtl::ArraySlice /*indices*/) { - return generator(engine); + return generator(*engine); })); } +// Similar to MakeFakeLiteral but takes a random number generator engine to +// enable reusing the engine across randomly generated literals. +StatusOr> MakeFakeLiteralInternal( + const Shape& shape, std::minstd_rand0* engine) { + if (ShapeUtil::IsTuple(shape)) { + std::vector> elements; + for (const Shape& element_shape : shape.tuple_shapes()) { + TF_ASSIGN_OR_RETURN(std::unique_ptr element, + MakeFakeLiteralInternal(element_shape, engine)); + elements.push_back(std::move(element)); + } + return Literal::MakeTupleOwned(std::move(elements)); + } + std::unique_ptr literal = Literal::CreateFromShape(shape); + switch (shape.element_type()) { + case BF16: + PopulateWithRandomFloatingPointData(literal.get(), engine); + break; + case F32: + PopulateWithRandomFloatingPointData(literal.get(), engine); + break; + case F64: + PopulateWithRandomFloatingPointData(literal.get(), engine); + break; + case S8: + PopulateWithRandomIntegralData(literal.get(), engine); + break; + case U8: + PopulateWithRandomIntegralData(literal.get(), engine); + break; + case S16: + PopulateWithRandomIntegralData(literal.get(), engine); + break; + case U16: + PopulateWithRandomIntegralData(literal.get(), engine); + break; + case S32: + PopulateWithRandomIntegralData(literal.get(), engine); + break; + case U32: + PopulateWithRandomIntegralData(literal.get(), engine); + break; + case S64: + PopulateWithRandomIntegralData(literal.get(), engine); + break; + case U64: + PopulateWithRandomIntegralData(literal.get(), engine); + break; + case PRED: { + std::uniform_int_distribution generator(0, 1); + TF_CHECK_OK(literal->Populate( + [&](tensorflow::gtl::ArraySlice /*indices*/) { + return generator(*engine); + })); + break; + } + default: + return Unimplemented("Unsupported type for fake literal generation: %s", + ShapeUtil::HumanString(shape).c_str()); + } + return std::move(literal); +} + // Matches binary addition computations. bool LooksLikeSum(const HloComputation& computation) { const HloInstruction* const root = computation.root_instruction(); @@ -88,15 +171,15 @@ bool NeedsZeroInitValue(const HloUse& use) { // Generate random values that are constrained to the input_shape minus the // output_shape so as not to produce wrapping slices, for instance. std::unique_ptr MakeRandomNonwrappingSliceIndex( - const Shape& input_shape, const Shape& slice_shape) { + const Shape& input_shape, const Shape& slice_shape, + std::minstd_rand0* engine) { const int64 rank = ShapeUtil::Rank(input_shape); std::vector start_indices(rank); - std::minstd_rand0 engine; for (int i = 0; i < rank; ++i) { const int32 upper_bound = ShapeUtil::GetDimension(input_shape, i) - ShapeUtil::GetDimension(slice_shape, i); std::uniform_int_distribution generator(0, upper_bound); - start_indices[i] = generator(engine); + start_indices[i] = generator(*engine); } return Literal::CreateR1(start_indices); } @@ -126,6 +209,11 @@ std::vector FindConstrainedUses( fused_uses.end()); } else if (NeedsZeroInitValue(use)) { constrained_uses.push_back(instruction); + } else if (opcode == HloOpcode::kConvert || + opcode == HloOpcode::kReducePrecision) { + auto converted_uses = FindConstrainedUses(dataflow, *instruction); + constrained_uses.insert(constrained_uses.end(), converted_uses.begin(), + converted_uses.end()); } } } @@ -138,14 +226,13 @@ std::vector FindConstrainedUses( // zero in the case of init_values for reductions). StatusOr> CreateLiteralForConstrainedUses( const tensorflow::gtl::ArraySlice constrained_uses, - const HloInstruction& param) { + const HloInstruction& param, std::minstd_rand0* engine) { HloInstruction* needs_index = nullptr; HloInstruction* needs_zero = nullptr; for (HloInstruction* use : constrained_uses) { switch (use->opcode()) { case HloOpcode::kDynamicSlice: case HloOpcode::kDynamicUpdateSlice: - TF_RET_CHECK(ShapeUtil::Equal(param.shape(), use->operand(0)->shape())); if (needs_index != nullptr && !ShapeUtil::Equal(needs_index->shape(), use->shape())) { return Unimplemented( @@ -173,106 +260,47 @@ StatusOr> CreateLiteralForConstrainedUses( needs_index->ToString().c_str(), needs_zero->ToString().c_str()); } if (needs_index != nullptr) { - return MakeRandomNonwrappingSliceIndex(param.shape(), needs_index->shape()); + return MakeRandomNonwrappingSliceIndex(needs_index->operand(0)->shape(), + needs_index->shape(), engine); } else if (needs_zero != nullptr) { return Literal::CreateFromShape(param.shape()); } else { - return MakeFakeLiteral(param.shape()); + return MakeFakeLiteralInternal(param.shape(), engine); } } // Given a module entry parameter, use the dataflow analysis to see if a // special case literal must be created, or if we can generate fake data. StatusOr> MakeConstrainedArgument( - const HloDataflowAnalysis& dataflow, const HloInstruction& param) { + const HloDataflowAnalysis& dataflow, const HloInstruction& param, + std::minstd_rand0* engine) { const auto constrained_uses = FindConstrainedUses(dataflow, param); - return CreateLiteralForConstrainedUses(constrained_uses, param); + return CreateLiteralForConstrainedUses(constrained_uses, param, engine); } } // namespace StatusOr> MakeFakeLiteral(const Shape& shape) { - if (ShapeUtil::IsTuple(shape)) { - std::vector> elements; - for (const Shape& element_shape : shape.tuple_shapes()) { - TF_ASSIGN_OR_RETURN(std::unique_ptr element, - MakeFakeLiteral(element_shape)); - elements.push_back(std::move(element)); - } - return Literal::MakeTupleOwned(std::move(elements)); - } - std::unique_ptr literal = Literal::CreateFromShape(shape); - switch (shape.element_type()) { - case BF16: - PopulateWithRandomFloatingPointData(literal.get()); - break; - case F32: - PopulateWithRandomFloatingPointData(literal.get()); - break; - case F64: - PopulateWithRandomFloatingPointData(literal.get()); - break; - case S8: - PopulateWithRandomIntegralData(literal.get()); - break; - case U8: - PopulateWithRandomIntegralData(literal.get()); - break; - case S16: - PopulateWithRandomIntegralData(literal.get()); - break; - case U16: - PopulateWithRandomIntegralData(literal.get()); - break; - case S32: - PopulateWithRandomIntegralData(literal.get()); - break; - case U32: - PopulateWithRandomIntegralData(literal.get()); - break; - case S64: - PopulateWithRandomIntegralData(literal.get()); - break; - case U64: - PopulateWithRandomIntegralData(literal.get()); - break; - case PRED: { - std::uniform_int_distribution generator(0, 1); - std::minstd_rand0 engine; - TF_CHECK_OK(literal->Populate( - [&](tensorflow::gtl::ArraySlice /*indices*/) { - return generator(engine); - })); - break; - } - default: - return Unimplemented("Unsupported type for fake literal generation: %s", - ShapeUtil::HumanString(shape).c_str()); - } - return std::move(literal); + std::minstd_rand0 engine; + return MakeFakeLiteralInternal(shape, &engine); } StatusOr>> MakeFakeArguments( HloModule* const module) { TF_ASSIGN_OR_RETURN(auto dataflow, HloDataflowAnalysis::Run(module)); const auto params = module->entry_computation()->parameter_instructions(); + std::minstd_rand0 engine; std::vector> arguments(params.size()); for (int i = 0; i < params.size(); ++i) { - TF_ASSIGN_OR_RETURN(arguments[i], - MakeConstrainedArgument(*dataflow, *params[i])); + TF_ASSIGN_OR_RETURN( + arguments[i], MakeConstrainedArgument(*dataflow, *params[i], &engine)); } return std::move(arguments); } Status VerifyHloModule(const perftools::gputools::Platform& platform, HloModule* const module) { - return HloVerifier( - std::bind( - &TransferManager::GetByteSizeRequirement, - TransferManager::GetForPlatform(&platform).ConsumeValueOrDie(), - std::placeholders::_1)) - .Run(module) - .status(); + return HloVerifier().Run(module).status(); } } // namespace xla diff --git a/tensorflow/compiler/xla/tests/transfer_manager_test.cc b/tensorflow/compiler/xla/tests/transfer_manager_test.cc index ed556fafb17cb2d243141783f822400d3c54b459..268ba338f2e6740a1d1a046d5a85494f3cf2e9f8 100644 --- a/tensorflow/compiler/xla/tests/transfer_manager_test.cc +++ b/tensorflow/compiler/xla/tests/transfer_manager_test.cc @@ -119,7 +119,7 @@ XLA_TEST_F(TransferManagerTest, TransferR1U8) { transfer_manager_->TransferLiteralFromDevice( stream_executor_, *device_buffer)); - EXPECT_EQ(result->u8s_string(), test_string); + EXPECT_EQ(result->GetR1U8AsString(), test_string); } XLA_TEST_F(TransferManagerTest, TransferR2F32) { diff --git a/tensorflow/compiler/xla/tests/tuple_test.cc b/tensorflow/compiler/xla/tests/tuple_test.cc index 65489cfff19c8fecbdead8a7e295bf9cca56038f..2029312f94a14bc81706368b9ecfc2727fd9fe4c 100644 --- a/tensorflow/compiler/xla/tests/tuple_test.cc +++ b/tensorflow/compiler/xla/tests/tuple_test.cc @@ -57,6 +57,20 @@ XLA_TEST_F(TupleTest, TupleConstant) { ComputeAndCompareTuple(&builder, *value, {}, error_spec_); } +// Tests a tuple made of scalar constants. +XLA_TEST_F(TupleTest, TupleScalarConstant) { + ComputationBuilder builder(client_, TestName()); + + const float constant_scalar1 = 7.3f; + const float constant_scalar2 = 1.2f; + auto value = + Literal::MakeTuple({Literal::CreateR0(constant_scalar1).get(), + Literal::CreateR0(constant_scalar2).get()}); + + auto result = builder.ConstantLiteral(*value); + ComputeAndCompareTuple(&builder, *value, {}, error_spec_); +} + // Tests the creation of tuple data. XLA_TEST_F(TupleTest, TupleCreate) { ComputationBuilder builder(client_, TestName()); @@ -180,8 +194,7 @@ XLA_TEST_F(TupleTest, TupleGTEToTuple) { ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } -// TODO(b/68395210): GPU does not tolerate ambiguous top-level buffers. -XLA_TEST_F(TupleTest, DISABLED_ON_GPU(SelectBetweenPredTuples)) { +XLA_TEST_F(TupleTest, SelectBetweenPredTuples) { ComputationBuilder b(client_, TestName()); ComputationDataHandle v1, v2; diff --git a/tensorflow/compiler/xla/tests/unary_op_test.cc b/tensorflow/compiler/xla/tests/unary_op_test.cc index fa4192e9281784a4a3063601afe89fba6a9dac18..835e2d7e5594d7c8c6e523f9806e32dce23a87e9 100644 --- a/tensorflow/compiler/xla/tests/unary_op_test.cc +++ b/tensorflow/compiler/xla/tests/unary_op_test.cc @@ -215,5 +215,23 @@ XLA_TEST_F(UnaryOpTest, SignAbsTestR2) { ComputeAndCompareR2(&builder, {{0, 0}, {0, 0}}, {}); } +XLA_TEST_F(UnaryOpTest, ConvertElementTypePredToS32) { + ComputationBuilder builder(client_, TestName()); + auto lhs = builder.ConstantR1({0, 1}); + auto rhs = builder.ConstantR1({1, 1}); + builder.ConvertElementType(builder.Eq(lhs, rhs), S32); + + ComputeAndCompareR1(&builder, {0, 1}, {}); +} + +XLA_TEST_F(UnaryOpTest, ConvertElementTypePredToF32) { + ComputationBuilder builder(client_, TestName()); + auto lhs = builder.ConstantR1({0, 1}); + auto rhs = builder.ConstantR1({1, 1}); + builder.ConvertElementType(builder.Eq(lhs, rhs), F32); + + ComputeAndCompareR1(&builder, {0.0, 1.0}, {}); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/while_test.cc b/tensorflow/compiler/xla/tests/while_test.cc index 0b3430ee1ee515c2c98c64a947b7a7021c04f22b..52157b837c383205f77a030ef98b2fd03a41aff5 100644 --- a/tensorflow/compiler/xla/tests/while_test.cc +++ b/tensorflow/compiler/xla/tests/while_test.cc @@ -563,6 +563,53 @@ TEST_F(WhileTest, WhileWithPredicateTupleResult) { ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0)); } +TEST_F(WhileTest, WhileWithTupleConstantScalarResult) { + std::vector shape_elements = {ShapeUtil::MakeShape(S32, {}), + ShapeUtil::MakeShape(S32, {})}; + Shape result_shape = ShapeUtil::MakeTupleShape(shape_elements); + + // Create a computation for the condition. + // Repeat for 5 iterations. + Computation condition; + { + ComputationBuilder builder(client_, "condition"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + builder.Gt(builder.ConstantR0(5), iteration); + condition = builder.Build().ConsumeValueOrDie(); + } + + // Create a computation for the body. + // Add 1 to the iteration variable and set the other tuple element to a + // constant. + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + auto result = + builder.Tuple({builder.Add(iteration, builder.ConstantR0(1)), + builder.ConstantR0(7)}); + body = builder.Build().ConsumeValueOrDie(); + } + + // Create a While node with computations for the condition and the body. + ComputationBuilder builder(client_, "while"); + auto init = builder.Tuple( + {builder.ConstantR0(0), builder.ConstantR0(7)}); + auto result = builder.While(condition, body, init); + VLOG(2) << "while = " + << ShapeUtil::HumanString( + *builder.GetShape(result).ConsumeValueOrDie()); + + auto expected_counter = Literal::CreateR0(5); + auto expected_data = Literal::CreateR0(7); + auto expected = + Literal::MakeTuple({expected_counter.get(), expected_data.get()}); + VLOG(2) << "expected = " << ShapeUtil::HumanString(expected->shape()); + ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001)); +} + // Tests two while nodes when the result type T is a Tuple and the second // while node uses the result of the first while node which is used in two // nodes. @@ -1160,6 +1207,50 @@ TEST_F(WhileTest, WhileWithCallInsideCondition) { ComputeAndCompareR0(&builder, 5, {}); } +TEST_F(WhileTest, WhileWithLoopInvariantOperation) { + auto matrix_shape = ShapeUtil::MakeShape(F32, {2, 2}); + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + auto while_shape = ShapeUtil::MakeTupleShape( + {scalar_s32, matrix_shape, matrix_shape, matrix_shape}); + + // Create a computation for the condition: repeat for 5 iterations. + Computation condition; + { + ComputationBuilder builder(client_, "condition"); + auto state = builder.Parameter(0, while_shape, "state"); + builder.Gt(builder.ConstantR0(5), builder.GetTupleElement(state, 0)); + TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build()); + } + + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto state = builder.Parameter(0, while_shape, "state"); + auto indvar = builder.GetTupleElement(state, 0); + auto input_0 = builder.GetTupleElement(state, 1); + auto input_1 = builder.GetTupleElement(state, 2); + auto output = builder.Tanh(builder.Dot(input_0, input_1)); + auto indvar_next = builder.Add(indvar, builder.ConstantR0(1)); + auto tuple_result = builder.Tuple({indvar_next, input_0, input_1, output}); + TF_ASSERT_OK_AND_ASSIGN(body, builder.Build()); + } + + ComputationBuilder builder(client_, TestName()); + auto matrix_input = builder.Parameter(0, matrix_shape, "matrix"); + auto init = builder.Tuple( + {builder.ConstantR0(0), matrix_input, matrix_input, matrix_input}); + auto while_instruction = builder.While(condition, body, init); + builder.GetTupleElement(while_instruction, 3); + + TF_ASSERT_OK_AND_ASSIGN(auto param_value, + client_->TransferToServer(*Literal::CreateR2( + {{1.0, 2.0}, {-1.0, -2.0}}))); + + ComputeAndCompareR2( + &builder, {{-0.76159416, -0.96402758}, {0.76159416, 0.96402758}}, + {param_value.get()}, ErrorSpec(4e-5)); +} + void BM_WhileLoop(int num_iters) { // Benchmark a simple kernel to measure while loop overheads. tensorflow::testing::StopTiming(); diff --git a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..9ad2a1985331b80625dd0687ea052300bc99e440 --- /dev/null +++ b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc @@ -0,0 +1,362 @@ +/* Copyright 2017 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 +#include + +#include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/service/platform_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/compiler/xla/tests/test_utils.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/platform/regexp.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace { +namespace se = ::perftools::gputools; +namespace gtl = ::tensorflow::gtl; + +class HloProfileTest : public ClientLibraryTestBase {}; + +struct ParsedProfileOutputLine { + int64 cycles; + string cycles_percentage; + double usec; + string flops; + string trops; + string bytes_per_sec; + string bytes_per_cycle; + string opcode; +}; + +::testing::AssertionResult HasFlops( + const ParsedProfileOutputLine& parsed_line) { + if (RE2::FullMatch(parsed_line.flops, "[0-9.TGMk]+FLOP/s")) { + return ::testing::AssertionSuccess() + << "'flops' field present in " << parsed_line.opcode << ": '" + << parsed_line.flops << "'"; + } + + return ::testing::AssertionFailure() + << "'flops' field absent in " << parsed_line.opcode << ": '" + << parsed_line.flops << "'"; +} + +::testing::AssertionResult HasTrops( + const ParsedProfileOutputLine& parsed_line) { + if (RE2::FullMatch(parsed_line.trops, "[0-9.TGMk]+TROP/s")) { + return ::testing::AssertionSuccess() + << "'trops' field present in " << parsed_line.opcode << ": '" + << parsed_line.trops << "'"; + } + + return ::testing::AssertionFailure() + << "'trops' field absent in " << parsed_line.opcode << ": '" + << parsed_line.trops << "'"; +} + +Status ParseOneProfileOutputLine( + const string& line, bool expect_hlo, + gtl::FlatMap* parsed_results) { + string separator = "[^:]*:: +"; + string match_percentage = "\\d+\\.\\d\\d%"; + string match_cycles = "(\\d+) cycles +\\( *(" + match_percentage + ")\\)"; + string match_usecs = "([0-9.]+) usec"; + string match_flops = "([^ ]+)"; + string match_trops = "([^ ]+)"; + string match_bytes_per_sec = "([0-9.TGMKi]+)B/s"; + string match_bytes_per_cycle = "([0-9.TGMKi]+)B/cycle"; + + // The underlined part is what we're trying to match with match_opcode: + // + // %dot33 = f32[256,256]{1,0} dot(...) + // ^^^ + + string match_opcode = + expect_hlo ? "%[^=]+= [^ ]+ ([^(]+)\\(.*" : "(\\[total\\])"; + string regexp_pattern = tensorflow::strings::StrCat( + " +", match_cycles, separator, match_usecs, separator, match_flops, + separator, match_trops, separator, match_bytes_per_sec, separator, + match_bytes_per_cycle, separator, match_opcode); + + ParsedProfileOutputLine parsed_line; + bool matched = RE2::FullMatch( + line, regexp_pattern, &parsed_line.cycles, &parsed_line.cycles_percentage, + &parsed_line.usec, &parsed_line.flops, &parsed_line.trops, + &parsed_line.bytes_per_sec, &parsed_line.bytes_per_cycle, + &parsed_line.opcode); + if (!matched) { + return tensorflow::errors::InvalidArgument( + "Input did not match regexp. Input: ", line, + ", Regexp: ", regexp_pattern); + } + + InsertOrDie(parsed_results, parsed_line.opcode, parsed_line); + + return Status::OK(); +} + +// Returns void so that we can ASSERT. +void ExecuteAndFetchProfile(string* profile_output, LocalClient* client, + const Computation& computation, + const Shape& lhs_arg_shape, + const Shape& rhs_arg_shape) { + LocalService* service = ClientLibrary::GetXlaService(client->platform()); + Backend* backend = service->mutable_backend(); + se::StreamExecutor* executor = backend->default_stream_executor(); + DeviceMemoryAllocator* allocator = backend->memory_allocator(); + auto* transfer_manager = backend->transfer_manager(); + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr lhs_arg, + transfer_manager->AllocateScopedShapedBuffer( + lhs_arg_shape, allocator, backend->default_device_ordinal())); + TF_ASSERT_OK(transfer_manager->TransferLiteralToDevice( + executor, *Literal::CreateFromShape(lhs_arg_shape), *lhs_arg)); + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr rhs_arg, + transfer_manager->AllocateScopedShapedBuffer( + rhs_arg_shape, allocator, backend->default_device_ordinal())); + TF_ASSERT_OK(transfer_manager->TransferLiteralToDevice( + executor, *Literal::CreateFromShape(rhs_arg_shape), *rhs_arg)); + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr local_executable, + client->Compile(computation, {&lhs_arg_shape, &rhs_arg_shape}, + ExecutableBuildOptions())); + + Executable* executable = local_executable->executable(); + HloExecutionProfile hlo_execution_profile( + &executable->hlo_profile_printer_data(), + &executable->hlo_profile_index_map()); + + TF_ASSERT_OK_AND_ASSIGN( + Backend::StreamPtr stream_ptr, + backend->BorrowStream(backend->default_device_ordinal())); + ExecutableRunOptions exec_run_options; + exec_run_options.set_stream(stream_ptr.get()); + exec_run_options.set_allocator(backend->memory_allocator()); + exec_run_options.set_intra_op_thread_pool( + backend->eigen_intra_op_thread_pool_device()); + ServiceExecutableRunOptions run_options( + exec_run_options, /*borrow_stream=*/nullptr, + backend->eigen_intra_op_thread_pool()); + TF_ASSERT_OK_AND_ASSIGN( + auto execution_result, + executable->ExecuteOnStream(&run_options, {lhs_arg.get(), rhs_arg.get()}, + &hlo_execution_profile)); + (void)execution_result; + + *profile_output = + hlo_execution_profile.ToString(executor->GetDeviceDescription()); + + XLA_VLOG_LINES(4, *profile_output); +} + +// TODO(b/71364943): This test exposes a bug in the parallel CPU backend. +XLA_TEST_F(HloProfileTest, DISABLED_ON_CPU_PARALLEL(ProfileSingleComputation)) { + const int64 m = 256, k = 256, n = 256; + Shape lhs_shape = ShapeUtil::MakeShape(F32, {m, k}); + Shape rhs_shape = ShapeUtil::MakeShape(F32, {m, k}); + + TF_ASSERT_OK_AND_ASSIGN(se::Platform * platform, + PlatformUtil::GetDefaultPlatform()); + TF_ASSERT_OK_AND_ASSIGN(LocalClient * client, + ClientLibrary::GetOrCreateLocalClient(platform)); + + ComputationBuilder builder(client, TestName()); + auto result = builder.Tanh(builder.Add( + builder.Parameter(0, ShapeUtil::MakeShape(F32, {m, k}), "dot_lhs"), + builder.Parameter(1, ShapeUtil::MakeShape(F32, {k, n}), "dot_rhs"))); + + TF_ASSERT_OK_AND_ASSIGN(auto computation, builder.Build()); + + string profile_output; + ExecuteAndFetchProfile(&profile_output, client, computation, lhs_shape, + rhs_shape); + + std::vector profile_output_lines = + tensorflow::str_util::Split(profile_output, '\n'); + + gtl::FlatMap parsed_profile_lines; + + TF_ASSERT_OK(ParseOneProfileOutputLine( + profile_output_lines[1], /*expect_hlo=*/false, &parsed_profile_lines)); + + TF_ASSERT_OK(ParseOneProfileOutputLine( + profile_output_lines[2], /*expect_hlo=*/true, &parsed_profile_lines)); + + TF_ASSERT_OK(ParseOneProfileOutputLine( + profile_output_lines[3], /*expect_hlo=*/true, &parsed_profile_lines)); + + TF_ASSERT_OK_AND_ASSIGN(ParsedProfileOutputLine total_profile, + MaybeFind(parsed_profile_lines, "[total]")); + TF_ASSERT_OK_AND_ASSIGN(ParsedProfileOutputLine dot_profile, + MaybeFind(parsed_profile_lines, "add")); + TF_ASSERT_OK_AND_ASSIGN(ParsedProfileOutputLine tanh_profile, + MaybeFind(parsed_profile_lines, "tanh")); + + EXPECT_GT(total_profile.cycles, 0); + EXPECT_EQ(total_profile.cycles_percentage, "100.00%"); + + EXPECT_TRUE(HasFlops(total_profile)); + EXPECT_TRUE(HasTrops(total_profile)); + + EXPECT_GT(total_profile.cycles, dot_profile.cycles); + EXPECT_NE(dot_profile.cycles_percentage, "0.00%"); + EXPECT_NE(dot_profile.cycles_percentage, "100.00%"); + + EXPECT_TRUE(HasFlops(dot_profile)); + EXPECT_FALSE(HasTrops(dot_profile)); + + EXPECT_GT(total_profile.cycles, tanh_profile.cycles); + EXPECT_NE(tanh_profile.cycles_percentage, "0.00%"); + EXPECT_NE(tanh_profile.cycles_percentage, "100.00%"); + + EXPECT_FALSE(HasFlops(tanh_profile)); + EXPECT_TRUE(HasTrops(tanh_profile)); +} + +// TODO(b/71364943): This test exposes a bug in the parallel CPU backend. +// +// TODO(b/71544591): The GPU backend does not record cycles spent in on Hlo +// instructions "interior" to while nodes. +XLA_TEST_F(HloProfileTest, + DISABLED_ON_GPU(DISABLED_ON_CPU_PARALLEL(ProfileWhileComputation))) { + const int64 size = 256; + Shape matrix_shape = ShapeUtil::MakeShape(F32, {size, size}); + Shape while_result_shape = + ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(S32, {}), matrix_shape}); + + TF_ASSERT_OK_AND_ASSIGN(se::Platform * platform, + PlatformUtil::GetDefaultPlatform()); + TF_ASSERT_OK_AND_ASSIGN(LocalClient * client, + ClientLibrary::GetOrCreateLocalClient(platform)); + + Computation condition; + { + ComputationBuilder builder(client, "condition"); + auto state = builder.Parameter(0, while_result_shape, "state"); + auto iteration = builder.GetTupleElement(state, 0); + builder.Gt(builder.ConstantR0(5), iteration); + TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build()); + } + + Computation body; + { + ComputationBuilder builder(client, "body"); + auto state = builder.Parameter(0, while_result_shape, "state"); + auto matrix = builder.GetTupleElement(state, 1); + auto next_iteration = builder.Add(builder.GetTupleElement(state, 0), + builder.ConstantR0(1)); + builder.Tuple({next_iteration, builder.Add(matrix, matrix)}); + TF_ASSERT_OK_AND_ASSIGN(body, builder.Build()); + } + + ComputationBuilder builder(client, TestName()); + auto initial_while_state = + builder.Tuple({builder.ConstantR0(0), + builder.Parameter(0, matrix_shape, "initial_value")}); + auto while_result = builder.While(condition, body, initial_while_state); + builder.Add(builder.GetTupleElement(while_result, 1), + builder.Parameter(1, matrix_shape, "other_value")); + + TF_ASSERT_OK_AND_ASSIGN(auto computation, builder.Build()); + + string profile_output; + ExecuteAndFetchProfile(&profile_output, client, computation, matrix_shape, + matrix_shape); + + std::vector profile_output_lines = + tensorflow::str_util::Split(profile_output, '\n'); + + auto while_body_profile_start = + std::find_if(profile_output_lines.begin(), profile_output_lines.end(), + [](tensorflow::StringPiece s) { + return s.starts_with("Execution profile for body"); + }); + + ASSERT_NE(while_body_profile_start, profile_output_lines.end()); + + gtl::FlatMap parsed_profile_lines; + + TF_ASSERT_OK( + ParseOneProfileOutputLine(*std::next(while_body_profile_start, 1), + /*expect_hlo=*/false, &parsed_profile_lines)); + + TF_ASSERT_OK( + ParseOneProfileOutputLine(*std::next(while_body_profile_start, 2), + /*expect_hlo=*/true, &parsed_profile_lines)); + + TF_ASSERT_OK_AND_ASSIGN(ParsedProfileOutputLine total_while_body_profile, + MaybeFind(parsed_profile_lines, "[total]")); + TF_ASSERT_OK_AND_ASSIGN(ParsedProfileOutputLine dot_profile, + MaybeFind(parsed_profile_lines, "add")); + + EXPECT_GT(total_while_body_profile.cycles, 0); + EXPECT_EQ(total_while_body_profile.opcode, "[total]"); + EXPECT_EQ(total_while_body_profile.cycles_percentage, "100.00%"); + + EXPECT_GT(total_while_body_profile.cycles, dot_profile.cycles); + EXPECT_NE(dot_profile.cycles_percentage, "0.00%"); + EXPECT_NE(dot_profile.cycles_percentage, "100.00%"); +} +} // namespace +} // namespace xla + +static std::pair AddXlaHloProfileFlag(int argc, char** argv) { + // Intentional "leak". + char** new_argv = new char*[argc + 2]; + for (int i = 0; i < argc; i++) { + new_argv[i] = argv[i]; + } + + // We do it this way (as opposed to piping in a modified DebugOptions + // instance) for better end-to-end integration testing. + new_argv[argc] = strdup("--xla_hlo_profile"); + + // Fusion can change the Hlo instructions that show up in the final Hlo + // executable, so block it here. + new_argv[argc + 1] = strdup("--xla_disable_hlo_passes=fusion"); + return {argc + 2, new_argv}; +} + +GTEST_API_ int main(int argc, char** argv) { + std::vector flag_list; + xla::legacy_flags::AppendDebugOptionsFlags(&flag_list); + std::tie(argc, argv) = AddXlaHloProfileFlag(argc, argv); + + auto usage = tensorflow::Flags::Usage(argv[0], flag_list); + if (!tensorflow::Flags::Parse(&argc, argv, flag_list)) { + LOG(ERROR) << "\n" << usage; + return 2; + } + + testing::InitGoogleTest(&argc, argv); + if (argc > 1) { + LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; + return 2; + } + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/compiler/xla/text_literal_reader.cc b/tensorflow/compiler/xla/text_literal_reader.cc index 4d060895d357493327ec50b38016478c65fef94d..6fa4c48e11d1102367b21bc21d4734466495ef0e 100644 --- a/tensorflow/compiler/xla/text_literal_reader.cc +++ b/tensorflow/compiler/xla/text_literal_reader.cc @@ -102,9 +102,9 @@ StatusOr> TextLiteralReader::ReadAllLines() { ShapeUtil::HumanString(shape).c_str()); } - auto result = MakeUnique(); + auto result = MakeUnique(shape); const float fill = std::numeric_limits::quiet_NaN(); - result->PopulateWithValue(fill, AsInt64Slice(shape.dimensions())); + result->PopulateWithValue(fill); std::vector pieces; std::vector coordinates; std::vector coordinate_values; diff --git a/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc index 5ede37b8737bd4fa6235464ddeb6382af17c8a80..b82f1c81c84b487c1661af5267b9123da97bb107 100644 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc @@ -85,10 +85,12 @@ void RealMain(tensorflow::gtl::ArraySlice args) { for (int i = 0; i < program_shape->parameters_size(); ++i) { layouts.push_back(&program_shape->parameters(i)); } + ExecutableBuildOptions build_options; + build_options.set_device_ordinal(0); + build_options.set_result_layout(program_shape->result()); StatusOr> executable = local_service->CompileExecutable(computation.handle(), layouts, - &program_shape->result(), - /*device_ordinal=*/0); + build_options); const HloModule& module = executable.ValueOrDie()->module(); diff --git a/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc index 24417a0cb8212e59cc0af53bd5bb21afcf3e134b..05c0fdf97d27c09eb2bbb0f265b5b2a5982ca7b1 100644 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc @@ -60,10 +60,13 @@ void RealMain(tensorflow::gtl::ArraySlice args, bool compile) { for (int i = 0; i < program_shape->parameters_size(); ++i) { layouts.push_back(&program_shape->parameters(i)); } + + ExecutableBuildOptions build_options; + build_options.set_device_ordinal(0); + build_options.set_result_layout(program_shape->result()); StatusOr> executable = local_service->CompileExecutable(computation.handle(), layouts, - &program_shape->result(), - /*device_ordinal=*/0); + build_options); const HloModule& module = executable.ValueOrDie()->module(); diff --git a/tensorflow/compiler/xla/tools/hlo_proto_to_json.cc b/tensorflow/compiler/xla/tools/hlo_proto_to_json.cc index 4e02e17db65c0a4220672733be8319e1a0cc4f0f..8460ae3e4991ee091af72d2553a8491f627c722e 100644 --- a/tensorflow/compiler/xla/tools/hlo_proto_to_json.cc +++ b/tensorflow/compiler/xla/tools/hlo_proto_to_json.cc @@ -19,7 +19,7 @@ limitations under the License. // // Reads one serilized Hlo module, convert it into JSON format and dump into // some output directory. some_binaray_proto is obtained by serializing Hlo -// module to disk using --xla_dump_hlo_proto_to debug optoin. +// module to disk using --xla_dump_optimized_hlo_proto_to debug option. #include #include diff --git a/tensorflow/compiler/xla/tools/parser/README.md b/tensorflow/compiler/xla/tools/parser/README.md index 4b43810a68214982dce62addd9aed69f87dc871b..f0f3dd7785c13e505e1eb6d4c8cd4bad157c4993 100644 --- a/tensorflow/compiler/xla/tools/parser/README.md +++ b/tensorflow/compiler/xla/tools/parser/README.md @@ -116,7 +116,29 @@ non_tuple | rank2345 ; rank2345 - : shape nested_array + : shape sparse_or_nested_array + ; +sparse_or_nested_array + : sparse_array + | nested_array + ; +sparse_array + : '{' sparse_array1 '}' + ; +sparse_array1 + : sparse_array_item + | sparse_array1 ',' sparse_array_item + ; +sparse_array_item + : multi_index ':' scalar + ; +multi_index + : kInt + | '[' multi_index1 ']' + ; +multi_index1 + : kInt + | multi_index1 ',' kInt ; ``` diff --git a/tensorflow/compiler/xla/tools/parser/hlo_lexer.cc b/tensorflow/compiler/xla/tools/parser/hlo_lexer.cc index 6d1e4173d25a032970284fc7abbc3d2ec30b27cd..fc0e4444521247734fc240a03da669244fe1a6a4 100644 --- a/tensorflow/compiler/xla/tools/parser/hlo_lexer.cc +++ b/tensorflow/compiler/xla/tools/parser/hlo_lexer.cc @@ -166,7 +166,7 @@ TokKind HloLexer::LexIdentifier() { auto consumable = RegexpStringPieceFromPointers(token_start_, buf_.end()); // 'consumable' will be advanced iff its prefix matches the pattern. static LazyRE2 shape_pattern = { - R"(^(\w*\d*)\[([\d,]*)\](?:{([\d,]*)})?)"}; + R"(^(\w*\d*)\[([\d,]*)\](?:(dense|sparse)?{([\d,]+)})?)"}; if (RE2::Consume(&consumable, *shape_pattern)) { auto status_or_shape = ShapeUtil::ParseShapeString( StringPieceFromPointers(token_start_, consumable.begin())); diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc index 58139ee0c1443e4fad3b78c8e4cff963a264824d..89def5d5610cb9522a69297668b443b8c4e03fb5 100644 --- a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc +++ b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc @@ -68,6 +68,13 @@ class HloParser { bool ParseTupleLiteral(std::unique_ptr* literal, const Shape& shape); bool ParseNonTupleLiteral(std::unique_ptr* literal, const Shape& shape); + bool ParseDenseLiteral(std::unique_ptr* literal, const Shape& shape); + bool ParseSparseLiteral(std::unique_ptr* literal, + const Shape& shape); + template + bool ParseSparseLiteralHelper(std::unique_ptr* literal, + const Shape& shape); + // Sets the sub-value of literal at the given index to the given value. The // literal's shape must have the default layout. bool SetValueInLiteral(int64 value, int64 linear_index, Literal* literal); @@ -99,6 +106,7 @@ class HloParser { kString, kBracedInt64List, kHloComputation, + kFftType, kWindow, kConvolutionDimensionNumbers, kSharding, @@ -178,6 +186,7 @@ class HloParser { bool ParseString(string* result); bool ParseShape(Shape* result); bool ParseOpcode(HloOpcode* result); + bool ParseFftType(FftType* result); bool ParseFusionKind(HloInstruction::FusionKind* result); bool ParseRandomDistribution(RandomDistribution* result); bool ParseInt64(int64* result); @@ -211,10 +220,13 @@ class HloParser { bool AddComputation(const string& name, HloComputation* computation, LocTy name_loc); - // The map from the instruction name to the instruction. This does not own the - // instructions. - std::unordered_map instruction_pool_; - std::unordered_map computation_pool_; + // The map from the instruction/computation name to the + // instruction/computation itself and it's location. This does not own the + // pointers. + std::unordered_map> + instruction_pool_; + std::unordered_map> + computation_pool_; HloLexer lexer_; std::unique_ptr module_; @@ -331,15 +343,16 @@ bool HloParser::ParseComputation(HloComputation** entry_computation) { return false; } - HloInstruction* root = - tensorflow::gtl::FindPtrOrNull(instruction_pool_, root_name); + std::pair* root_node = + tensorflow::gtl::FindOrNull(instruction_pool_, root_name); // This means some instruction was marked as ROOT but we didn't find it in the // pool, which should not happen. - if (!root_name.empty() && root == nullptr) { + if (!root_name.empty() && root_node == nullptr) { LOG(FATAL) << "instruction " << root_name << " was marked as ROOT but the parser has not seen it before"; } + HloInstruction* root = root_node == nullptr ? nullptr : root_node->first; // Now root can be either an existing instruction or a nullptr. If it's a // nullptr, the implementation of Builder will set the last instruction as // root instruction. @@ -685,6 +698,20 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder, shape, /*lhs=*/operands[0], /*rhs=*/operands[1], *window, *dnums)); break; } + case HloOpcode::kFft: { + optional fft_type; + optional> fft_length; + attrs["fft_type"] = {/*required=*/true, AttrTy::kFftType, &fft_type}; + attrs["fft_length"] = {/*required=*/true, AttrTy::kBracedInt64List, + &fft_length}; + if (!ParseOperands(&operands, /*expected_size=*/1) || + !ParseAttributes(attrs)) { + return false; + } + instruction = builder->AddInstruction(HloInstruction::CreateFft( + shape, operands[0], *fft_type, *fft_length)); + break; + } case HloOpcode::kBroadcast: { optional> broadcast_dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, @@ -908,7 +935,7 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder, return false; } instruction = builder->AddInstruction(HloInstruction::CreateOutfeed( - shape, operands[0], config ? *config : "")); + operands[0]->shape(), operands[0], config ? *config : "")); break; } case HloOpcode::kRng: { @@ -1206,13 +1233,13 @@ bool HloParser::ParseInstructionNames( if (!ParseName(&name)) { return Error(loc, "expects a instruction name"); } - HloInstruction* instr = - tensorflow::gtl::FindPtrOrNull(instruction_pool_, name); + std::pair* instr = + tensorflow::gtl::FindOrNull(instruction_pool_, name); if (!instr) { return TokenError( Printf("instruction '%s' is not defined", name.c_str())); } - instructions->push_back(instr); + instructions->push_back(instr->first); } while (EatIfPresent(TokKind::kComma)); return ParseToken(TokKind::kRbrace, @@ -1308,7 +1335,7 @@ bool HloParser::SetValueInLiteralHelper(ParsedElemT value, int64 linear_index, PrimitiveType_Name(literal->shape().element_type()))); } - literal->GetMutableArraySlice().at(linear_index) = + literal->data().at(linear_index) = static_cast(value); return true; } @@ -1375,9 +1402,19 @@ bool HloParser::ParseTupleLiteral(std::unique_ptr* literal, // non_tuple // ::= rank01 // ::= rank2345 -// rank2345 ::= shape nested_array +// rank2345 ::= shape sparse_or_nested_array bool HloParser::ParseNonTupleLiteral(std::unique_ptr* literal, const Shape& shape) { + if (LayoutUtil::IsSparseArray(shape)) { + return ParseSparseLiteral(literal, shape); + } + + CHECK(LayoutUtil::IsDenseArray(shape)); + return ParseDenseLiteral(literal, shape); +} + +bool HloParser::ParseDenseLiteral(std::unique_ptr* literal, + const Shape& shape) { const int64 rank = ShapeUtil::Rank(shape); if (rank > 1 && !EatShapeAndCheckCompatible(shape)) { return false; @@ -1499,7 +1536,7 @@ bool HloParser::ParseNonTupleLiteral(std::unique_ptr* literal, return false; } } else { - return TokenError(StrCat("unsupported premitive type ", + return TokenError(StrCat("unsupported primitive type ", PrimitiveType_Name(shape.element_type()))); } break; @@ -1511,6 +1548,142 @@ bool HloParser::ParseNonTupleLiteral(std::unique_ptr* literal, return true; } +bool HloParser::ParseSparseLiteral(std::unique_ptr* literal, + const Shape& shape) { + if (!EatShapeAndCheckCompatible(shape)) { + return false; + } + + switch (shape.element_type()) { + case PRED: + return ParseSparseLiteralHelper(literal, shape); + case S8: + return ParseSparseLiteralHelper(literal, shape); + case S16: + return ParseSparseLiteralHelper(literal, shape); + case S32: + return ParseSparseLiteralHelper(literal, shape); + case S64: + return ParseSparseLiteralHelper(literal, shape); + case U8: + return ParseSparseLiteralHelper(literal, shape); + case U16: + return ParseSparseLiteralHelper(literal, shape); + case U32: + return ParseSparseLiteralHelper(literal, shape); + case U64: + return ParseSparseLiteralHelper(literal, shape); + case F16: + return ParseSparseLiteralHelper(literal, shape); + case F32: + return ParseSparseLiteralHelper(literal, shape); + case BF16: + return ParseSparseLiteralHelper(literal, shape); + case F64: + return ParseSparseLiteralHelper(literal, shape); + default: + return Error(lexer_.GetLoc(), + StrCat("invalid primitive type for sparse literal: ", + PrimitiveType_Name(shape.element_type()))); + } +} + +template +bool HloParser::ParseSparseLiteralHelper(std::unique_ptr* literal, + const Shape& shape) { + std::vector index; + + int64 rank = ShapeUtil::Rank(shape); + + *literal = MakeUnique(shape); + + if (!ParseToken(TokKind::kLbrace, + "expects '{' at the beginning of a sparse literal")) { + return false; + } + + for (;;) { + if (lexer_.GetKind() == TokKind::kRbrace) { + lexer_.Lex(); + break; + } + + LocTy index_loc = lexer_.GetLoc(); + index.clear(); + if (lexer_.GetKind() == TokKind::kInt) { + int64 single_index = lexer_.GetInt64Val(); + lexer_.Lex(); + if (rank != 1) { + return Error( + index_loc, + StrCat("invalid single-dimensional index for shape with rank ", + rank, ": ", single_index)); + } + index.push_back(single_index); + } else { + if (!ParseInt64List(TokKind::kLsquare, TokKind::kRsquare, TokKind::kComma, + &index)) { + return false; + } + if (index.size() != rank) { + return Error( + index_loc, + StrCat("invalid multi-dimension index for shape with rank ", rank, + ": [", tensorflow::str_util::Join(index, ", "), "]")); + } + } + if (!ParseToken(TokKind::kColon, + "expects ':' after after the sparse array index and before " + "the sparse array value")) { + return false; + } + LocTy value_loc = lexer_.GetLoc(); + LiteralNativeT value; + if (lexer_.GetKind() == TokKind::kw_true || + lexer_.GetKind() == TokKind::kw_false) { + value = static_cast(lexer_.GetKind() == TokKind::kw_true); + lexer_.Lex(); + } else if (primitive_util::IsIntegralType(shape.element_type())) { + int64 value_s64; + if (!ParseInt64(&value_s64)) { + return Error(value_loc, + StrCat("expects integer for primitive type: ", + PrimitiveType_Name(shape.element_type()))); + } + value = static_cast(value_s64); + } else if (primitive_util::IsFloatingPointType(shape.element_type())) { + double value_f64; + if (!ParseDouble(&value_f64)) { + return Error(value_loc, + StrCat("expects floating point value for primitive type: ", + PrimitiveType_Name(shape.element_type()))); + } + value = static_cast(value_f64); + } else { + LOG(FATAL) << "Unexpected element type: " + << PrimitiveType_Name(shape.element_type()); + } + if (lexer_.GetKind() != TokKind::kRbrace && + !ParseToken(TokKind::kComma, + "expects ',' separator between sparse array elements")) { + return false; + } + + if ((*literal)->sparse_element_count() + 1 == + LayoutUtil::MaxSparseElements(shape.layout())) { + return Error( + lexer_.GetLoc(), + StrCat("number of sparse elements exceeds maximum for layout: ", + ShapeUtil::HumanStringWithLayout(shape))); + } + + (*literal)->AppendSparseElement(index, value); + } + + (*literal)->SortSparseElements(); + return true; +} + // operands ::= '(' operands1 ')' // operands1 // ::= /*empty*/ @@ -1536,12 +1709,12 @@ bool HloParser::ParseOperands(std::vector* operands) { if (!ParseName(&name)) { return false; } - HloInstruction* instruction = - tensorflow::gtl::FindPtrOrNull(instruction_pool_, name); + std::pair* instruction = + tensorflow::gtl::FindOrNull(instruction_pool_, name); if (!instruction) { return Error(loc, StrCat("instruction does not exist: ", name)); } - operands->push_back(instruction); + operands->push_back(instruction->first); } while (EatIfPresent(TokKind::kComma)); } return ParseToken(TokKind::kRparen, "expects ')' at the end of operands"); @@ -1672,6 +1845,14 @@ bool HloParser::ParseAttributeHelper( static_cast*>(attr_out_ptr)->emplace(result); return true; } + case AttrTy::kFftType: { + FftType result; + if (!ParseFftType(&result)) { + return false; + } + static_cast*>(attr_out_ptr)->emplace(result); + return true; + } case AttrTy::kWindow: { Window result; if (!ParseWindow(&result)) { @@ -1780,10 +1961,12 @@ bool HloParser::ParseComputationName(HloComputation** value) { if (!ParseName(&name)) { return Error(loc, "expects computation name"); } - *value = tensorflow::gtl::FindPtrOrNull(computation_pool_, name); - if (*value == nullptr) { + std::pair* computation = + tensorflow::gtl::FindOrNull(computation_pool_, name); + if (computation == nullptr) { return Error(loc, StrCat("computation does not exist: ", name)); } + *value = computation->first; return true; } @@ -1827,7 +2010,7 @@ bool HloParser::ParseWindow(Window* window) { if (field_name == "rhs_reversal") { return ParseDxD("rhs_reversal", &rhs_reversal); } - return Error(loc, StrCat("unexpected attribute name: ", field_name)); + return Error(attr_loc, StrCat("unexpected attribute name: ", field_name)); }(); if (!ok) { return false; @@ -1996,7 +2179,7 @@ bool HloParser::ParseConvolutionDimensionNumbers( // // {[2:3:4], [5:6:7], [8:9]} // -// The the parsed result will be: +// The parsed result will be: // // {/*starts=*/{2, 5, 8}, /*limits=*/{3, 6, 9}, /*strides=*/{4, 7, 1}} // @@ -2289,6 +2472,19 @@ bool HloParser::ParseOpcode(HloOpcode* result) { return true; } +bool HloParser::ParseFftType(FftType* result) { + VLOG(1) << "ParseFftType"; + if (lexer_.GetKind() != TokKind::kIdent) { + return TokenError("expects fft type"); + } + string val = lexer_.GetStrVal(); + if (!FftType_Parse(val, result) || !FftType_IsValid(*result)) { + return TokenError(Printf("expects fft type but sees: %s", val.c_str())); + } + lexer_.Lex(); + return true; +} + bool HloParser::ParseFusionKind(HloInstruction::FusionKind* result) { VLOG(1) << "ParseFusionKind"; if (lexer_.GetKind() != TokKind::kIdent) { @@ -2386,18 +2582,22 @@ bool HloParser::EatIfPresent(TokKind kind) { bool HloParser::AddInstruction(const string& name, HloInstruction* instruction, LocTy name_loc) { - auto result = instruction_pool_.insert({name, instruction}); + auto result = instruction_pool_.insert({name, {instruction, name_loc}}); if (!result.second) { - return Error(name_loc, StrCat("instruction already exists: ", name)); + Error(name_loc, StrCat("instruction already exists: ", name)); + return Error(/*loc=*/result.first->second.second, + "instruction previously defined here"); } return true; } bool HloParser::AddComputation(const string& name, HloComputation* computation, LocTy name_loc) { - auto result = computation_pool_.insert({name, computation}); + auto result = computation_pool_.insert({name, {computation, name_loc}}); if (!result.second) { - return Error(name_loc, StrCat("computation already exists: ", name)); + Error(name_loc, StrCat("computation already exists: ", name)); + return Error(/*loc=*/result.first->second.second, + "computation previously defined here"); } return true; } diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc index 69c59ad6c70353c8dce4c0f2bfe7f718f2a2189e..b8c6b59204f897c7dc07b846370b5b776a19a808 100644 --- a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc @@ -25,7 +25,6 @@ namespace tools { namespace { using tensorflow::StringPiece; -using tensorflow::strings::StrCat; struct TestData { string test_name; @@ -580,6 +579,54 @@ ENTRY %BatchNormGrad.v4 (input: f32[2,2,2,2], scale: f32[2], mean: f32[2], varia ROOT %batch-norm-grad = (f32[2,2,2,2]{3,2,1,0}, f32[2]{0}, f32[2]{0}) batch-norm-grad(f32[2,2,2,2]{3,2,1,0} %input, f32[2]{0} %scale, f32[2]{0} %mean, f32[2]{0} %variance, f32[2,2,2,2]{3,2,1,0} %grad_output), epsilon=0.001, feature_index=0 } +)" +}, +// fft +{ +"Fft", +R"(HloModule Fft_module + +ENTRY %Fft (input: c64[8,32]) -> c64[8,32] { + %input = c64[8,32]{1,0} parameter(0) + ROOT %fft = c64[8,32]{1,0} fft(c64[8,32]{1,0} %input), fft_type=FFT, fft_length={32} +} + +)" +}, +// ifft +{ +"Ifft2d", +R"(HloModule Ifft2d_module + +ENTRY %Ifft2d (input: c64[5,8,32]) -> c64[5,8,32] { + %input = c64[5,8,32]{2,1,0} parameter(0) + ROOT %fft = c64[5,8,32]{2,1,0} fft(c64[5,8,32]{2,1,0} %input), fft_type=IFFT, fft_length={8,32} +} + +)" +}, +// rfft2d +{ +"Rfft2d", +R"(HloModule Rfft2d_module + +ENTRY %Rfft2d (input: f32[5,64,32]) -> c64[5,64,17] { + %input = f32[5,64,32]{2,1,0} parameter(0) + ROOT %fft = c64[5,64,17]{2,1,0} fft(f32[5,64,32]{2,1,0} %input), fft_type=RFFT, fft_length={64,32} +} + +)" +}, +// irfft3d +{ +"Irfft3d", +R"(HloModule Irfft3d_module + +ENTRY %Irfft3d (input: c64[5,64,128,33]) -> f32[5,64,128,64] { + %input = c64[5,64,128,33]{3,2,1,0} parameter(0) + ROOT %fft = f32[5,64,128,64]{3,2,1,0} fft(c64[5,64,128,33]{3,2,1,0} %input), fft_type=IRFFT, fft_length={64,128,64} +} + )" }, // pad @@ -640,7 +687,37 @@ ENTRY %fusion.v3 () -> f32[3,2,1,1] { } )" +}, +{ +"Sparse", +R"(HloModule sparse_f32 + +ENTRY %sparse () -> f32[2,3,4] { + ROOT %foo = f32[2,3,4]sparse{10} constant(f32[2,3,4]{[0, 1, 2]: 1, [1, 2, 3]: 2, [2, 3, 4]: 3}) } + +)" +}, +{ +"SparseEmpty", +R"(HloModule sparse_f32_empty + +ENTRY %sparse_f32_empty () -> f32[2,3,4] { + ROOT %foo = f32[2,3,4]sparse{10} constant(f32[2,3,4]{}) +} + +)" +}, +{ +"SparseR1", +R"(HloModule sparse_f32_r1 + +ENTRY %sparse_f32_r1 () -> f32[9] { + ROOT %foo = f32[9]sparse{10} constant(f32[9]{1: 2, 3: 4, 5: 6}) +} + +)" +}, }); // clang-format on } @@ -1010,12 +1087,14 @@ ENTRY %Convolve1D1Window_0.v3 (input: f32[1,2,1], filter: f32[1,1,1]) -> f32[1,2 )"; - ExpectHasSubstr(Parse(StrCat(prefix, ",dim_labels=00_01_10", suffix)) - .status() - .error_message(), - "expects dim labels pattern"); + ExpectHasSubstr( + Parse(tensorflow::strings::StrCat(prefix, ",dim_labels=00_01_10", suffix)) + .status() + .error_message(), + "expects dim labels pattern"); - ExpectHasSubstr(Parse(StrCat(prefix, ",dim_labels=010_1100->010", suffix)) + ExpectHasSubstr(Parse(tensorflow::strings::StrCat( + prefix, ",dim_labels=010_1100->010", suffix)) .status() .error_message(), "must have the same rank"); @@ -1196,6 +1275,35 @@ ENTRY consts { "one computation should have only one ROOT"); } +TEST_F(HloParserTest, InstructionExists) { + const string original = R"(HloModule comp_exists +c1 { + instr = f32[1]{0} constant({12345}) +} +c2 { + instr = f32[1]{0} constant({67890}) +})"; + + ExpectHasSubstr(Parse(original).status().error_message(), + R"(was parsing 3:3: error: instruction previously defined here + instr = f32[1]{0} constant({12345}) + ^)"); +} + +TEST_F(HloParserTest, ComputationExists) { + const string original = R"(HloModule comp_exists +comp { + const1 = f32[1]{0} constant({12345}) +} +comp { + const2 = f32[1]{0} constant({67890}) +})"; + ExpectHasSubstr(Parse(original).status().error_message(), + R"(was parsing 2:1: error: computation previously defined here +comp { +^)"); +} + } // namespace } // namespace tools } // namespace xla diff --git a/tensorflow/compiler/xla/tools/replay_computation.cc b/tensorflow/compiler/xla/tools/replay_computation.cc index a7dc5862057047f7c56faeb211cc0b13992caec7..eda5effbb92db92c9317a956497a00c0ec15c27c 100644 --- a/tensorflow/compiler/xla/tools/replay_computation.cc +++ b/tensorflow/compiler/xla/tools/replay_computation.cc @@ -82,9 +82,10 @@ StatusOr> ReplayComputation( arguments = MakeFakeArgumentsOrDie(computation, client); } else { // use recorded data if available for (const auto& proto : module.arguments()) { - Literal literal(proto); + TF_ASSIGN_OR_RETURN(std::unique_ptr literal, + Literal::CreateFromProto(proto)); TF_ASSIGN_OR_RETURN(std::unique_ptr data, - client->TransferToServer(literal)); + client->TransferToServer(*literal)); arguments.push_back(std::move(data)); } } @@ -162,12 +163,16 @@ int RealMain(tensorflow::gtl::ArraySlice args, const Options& opts) { ShapeUtil::HumanString(result->shape()).c_str(), result->ToString().c_str()); if (module.has_result()) { + std::unique_ptr literal = + Literal::CreateFromProto(module.result()).ConsumeValueOrDie(); fprintf(stdout, "was %s:%s\n", ShapeUtil::HumanString(module.result().shape()).c_str(), - Literal(module.result()).ToString().c_str()); + literal->ToString().c_str()); } } } + + ClientLibrary::DestroyLocalInstances(); return exit_status; } diff --git a/tensorflow/compiler/xla/tools/show_literal.cc b/tensorflow/compiler/xla/tools/show_literal.cc index b50cb5e28eac14ed99af566939f8bd64e393ff64..fe8e72ba32bb4493b2751cfdfeb977f271092f9c 100644 --- a/tensorflow/compiler/xla/tools/show_literal.cc +++ b/tensorflow/compiler/xla/tools/show_literal.cc @@ -40,7 +40,8 @@ int main(int argc, char **argv) { xla::LiteralProto literal_proto; TF_CHECK_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(), argv[1], &literal_proto)); - xla::Literal literal(literal_proto); + std::unique_ptr literal = + xla::Literal::CreateFromProto(literal_proto).ConsumeValueOrDie(); LOG(INFO) << "literal: " << literal_proto.ShortDebugString(); - fprintf(stderr, "%s\n", literal.ToString().c_str()); + fprintf(stderr, "%s\n", literal->ToString().c_str()); } diff --git a/tensorflow/compiler/xla/tools/show_text_literal.cc b/tensorflow/compiler/xla/tools/show_text_literal.cc index bbe9902aa17a585c4bad5b732330305dfdd45302..8525873e913185554d18df8c8c3584bfcdcdcabe 100644 --- a/tensorflow/compiler/xla/tools/show_text_literal.cc +++ b/tensorflow/compiler/xla/tools/show_text_literal.cc @@ -39,13 +39,13 @@ int main(int argc, char **argv) { std::unique_ptr literal = xla::TextLiteralReader::ReadPath(argv[1]).ConsumeValueOrDie(); - LOG(INFO) << "literal: " << literal->ShortDebugString(); + LOG(INFO) << "literal: " << *literal; fprintf(stderr, "%s\n", literal->ToString().c_str()); if (literal->shape().element_type() == xla::F32) { - float min = - *std::min_element(literal->f32s().begin(), literal->f32s().end()); - float max = - *std::max_element(literal->f32s().begin(), literal->f32s().end()); + float min = *std::min_element(literal->data().begin(), + literal->data().end()); + float max = *std::max_element(literal->data().begin(), + literal->data().end()); fprintf(stderr, "min: %a=%f\n", min, min); fprintf(stderr, "max: %a=%f\n", max, max); } diff --git a/tensorflow/compiler/xla/util.cc b/tensorflow/compiler/xla/util.cc index fe5d29a6b655a89d559eb1214c2b8dd54d34094c..1f0c626bbb2d64ef4e67c9ec51485ae96ae73d04 100644 --- a/tensorflow/compiler/xla/util.cc +++ b/tensorflow/compiler/xla/util.cc @@ -30,9 +30,7 @@ limitations under the License. #include "tensorflow/core/platform/stacktrace.h" namespace xla { -namespace { -// Logs the provided status message with a backtrace. Status WithLogBacktrace(const Status& status) { CHECK(!status.ok()); VLOG(1) << status.ToString(); @@ -40,8 +38,6 @@ Status WithLogBacktrace(const Status& status) { return status; } -} // namespace - ScopedLoggingTimer::ScopedLoggingTimer(const string& label, bool enabled) : enabled(enabled), label(label) { if (enabled) { @@ -74,13 +70,18 @@ Status AppendStatus(Status prior, tensorflow::StringPiece context) { // Implementation note: we can't common these out (without using macros) because // they all need to va_start/va_end their varargs in their frame. -Status InvalidArgument(const char* format, ...) { +Status InvalidArgumentV(const char* format, va_list args) { string message; + tensorflow::strings::Appendv(&message, format, args); + return WithLogBacktrace(tensorflow::errors::InvalidArgument(message)); +} + +Status InvalidArgument(const char* format, ...) { va_list args; va_start(args, format); - tensorflow::strings::Appendv(&message, format, args); + Status result = InvalidArgumentV(format, args); va_end(args); - return WithLogBacktrace(tensorflow::errors::InvalidArgument(message)); + return result; } Status Unimplemented(const char* format, ...) { @@ -338,7 +339,7 @@ std::vector> CommonFactors( string SanitizeFileName(string file_name) { for (char& c : file_name) { - if (c == '/' || c == '\\' || c == '[' || c == ']') { + if (c == '/' || c == '\\' || c == '[' || c == ']' || c == ' ') { c = '_'; } } diff --git a/tensorflow/compiler/xla/util.h b/tensorflow/compiler/xla/util.h index 277cc5ec86f23a10cdadf7fc3b89c96397bab704..08df5b12b3a53a138f56705531baa3333b23c5d8 100644 --- a/tensorflow/compiler/xla/util.h +++ b/tensorflow/compiler/xla/util.h @@ -40,6 +40,13 @@ limitations under the License. namespace xla { +// Logs the provided status message with a backtrace. +// +// For use by Status-factories, logs a backtrace at the point where the status +// is created, such that we can use --vmodule=util=1 to see all status +// creation backtraces. +Status WithLogBacktrace(const Status& status); + // Ranks greater than 8 are very rare, so use InlinedVector to store // the bounds and indices. And for the rare cases of ranks greater than 8, // the InlinedVector will just behave like an std::vector<> and allocate the @@ -207,6 +214,27 @@ Status ResourceExhausted(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); Status NotFound(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); Status Unavailable(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); +// Passed-varargs variant of the InvalidArgument factory above. +Status InvalidArgumentV(const char* format, va_list args); + +template +Status UnimplementedStrCat(Args&&... concat) { + return Unimplemented( + "%s", tensorflow::strings::StrCat(std::forward(concat)...).c_str()); +} + +template +Status InternalErrorStrCat(Args&&... concat) { + return InternalError( + "%s", tensorflow::strings::StrCat(std::forward(concat)...).c_str()); +} + +template +Status ResourceExhaustedStrCat(Args&&... concat) { + return ResourceExhausted( + "%s", tensorflow::strings::StrCat(std::forward(concat)...).c_str()); +} + // Splits the lines of the original, replaces leading whitespace with the prefix // given by "indentation", and returns the string joined by newlines again. As a // side effect, any additional trailing whitespace is removed. @@ -332,7 +360,7 @@ T CeilOfRatio(T dividend, T divisor) { } // Rounds the value up to a multiple of the divisor by first calling CeilOfRatio -// then multiplying by the divisor. For example: RoundUpToMultiple(13, 8) => 16 +// then multiplying by the divisor. For example: RoundUpToNearest(13, 8) => 16 template T RoundUpToNearest(T value, T divisor) { return CeilOfRatio(value, divisor) * divisor; @@ -340,7 +368,7 @@ T RoundUpToNearest(T value, T divisor) { // Rounds the value down to a multiple of the divisor by first calling // FloorOfRatio then multiplying by the divisor. For example: -// RoundUpToMultiple(13, 8) => 8 +// RoundDownToNearest(13, 8) => 8 template T RoundDownToNearest(T value, T divisor) { return FloorOfRatio(value, divisor) * divisor; @@ -398,6 +426,33 @@ std::vector> CommonFactors( // Removes illegal characters from filenames. string SanitizeFileName(string file_name); +template +bool c_all_of(Container container, Predicate predicate) { + return std::all_of(std::begin(container), std::end(container), predicate); +} + +template +OutputIterator c_transform(InputContainer input_container, + OutputIterator output_iterator, + UnaryOperation unary_op) { + return std::transform(std::begin(input_container), std::end(input_container), + output_iterator, unary_op); +} + +template +OutputIterator c_copy_if(InputContainer input_container, + OutputIterator output_iterator, + UnaryPredicate predicate) { + return std::copy_if(std::begin(input_container), std::end(input_container), + output_iterator, predicate); +} + +template +void c_sort(InputContainer& input_container, Comparator comparator) { + std::sort(input_container.begin(), input_container.end(), comparator); +} + } // namespace xla #define XLA_LOG_LINES(SEV, STRING) \ diff --git a/tensorflow/compiler/xla/window_util.cc b/tensorflow/compiler/xla/window_util.cc index 293f0781a203d092a7996d5548de1dbf5bf32e4c..93284b80f9e1f82c4b18dc7388754d5c01a7740c 100644 --- a/tensorflow/compiler/xla/window_util.cc +++ b/tensorflow/compiler/xla/window_util.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/lib/strings/stringprintf.h" @@ -25,6 +26,28 @@ limitations under the License. namespace xla { namespace window_util { +Window MakeWindow(tensorflow::gtl::ArraySlice sizes) { + Window window; + for (int64 size : sizes) { + auto* dimension = window.add_dimensions(); + dimension->set_size(size); + dimension->set_stride(1); + dimension->set_base_dilation(1); + dimension->set_window_dilation(1); + } + return window; +} + +PaddingConfig MakeSymmetricPadding(tensorflow::gtl::ArraySlice sizes) { + PaddingConfig config; + for (int64 size : sizes) { + auto* dimension = config.add_dimensions(); + dimension->set_edge_padding_low(size); + dimension->set_edge_padding_high(size); + } + return config; +} + /* static */ string ToString(const WindowDimension& dim) { using tensorflow::strings::StrAppend; using tensorflow::strings::StrCat; @@ -114,13 +137,21 @@ bool HasPadding(const Window& window) { return false; } -bool HasEvenPadding(const Window& window) { +bool HasSymmetricPadding(const Window& window) { return std::all_of(window.dimensions().begin(), window.dimensions().end(), [](const WindowDimension& dim) { return dim.padding_low() == dim.padding_high(); }); } +bool HasSymmetricPadding(const PaddingConfig& padding_config) { + return std::all_of(padding_config.dimensions().begin(), + padding_config.dimensions().end(), + [](const PaddingConfig::PaddingConfigDimension& dim) { + return dim.edge_padding_low() == dim.edge_padding_high(); + }); +} + bool HasNegativePadding(const Window& window) { return std::any_of(window.dimensions().begin(), window.dimensions().end(), [](const WindowDimension& dim) { @@ -159,6 +190,12 @@ bool HasDilation(const Window& window) { return HasBaseDilation(window) || HasWindowDilation(window); } +bool IsInactiveWindowDimension(const Window& window, int64 logical_dim) { + const WindowDimension& window_dim = window.dimensions(logical_dim); + return window_dim.size() == 1 && window_dim.stride() == 1 && + window_dim.padding_low() == 0 && window_dim.padding_high() == 0; +} + int64 DilatedBound(int64 bound, int64 dilation) { CHECK_GE(bound, 0); CHECK_GE(dilation, 1); diff --git a/tensorflow/compiler/xla/window_util.h b/tensorflow/compiler/xla/window_util.h index 125900dac0c5ab478b834c315b4a438c9238ef6d..ba473e2c8c35202865a9a4981da7653fe1d6f552 100644 --- a/tensorflow/compiler/xla/window_util.h +++ b/tensorflow/compiler/xla/window_util.h @@ -18,10 +18,21 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" namespace xla { namespace window_util { +// Creates a window with the given sizes in the dimensions and all strides set +// to 1. +Window MakeWindow(tensorflow::gtl::ArraySlice sizes); + +// Creates a padding config with symmetrical padding in each dimension, of value +// given by sizes; e.g. {0, 1, 2} would create a R3 padding config that had zero +// pixels of padding in dimension 0, one pixel of padding symmetrically, on each +// side of dimension 1, and two pixels of padding symmetrically on dimension 2. +PaddingConfig MakeSymmetricPadding(tensorflow::gtl::ArraySlice sizes); + string ToString(const WindowDimension& dim); string ToString(const Window& window); @@ -32,15 +43,24 @@ string ToString(const Window& window); bool HasStride(const Window& window); bool HasPadding(const Window& window); -bool HasEvenPadding(const Window& window); +bool HasSymmetricPadding(const Window& window); bool HasNegativePadding(const Window& window); +// As with HasSymmetricPadding(Window) above, returns whether the "padding low" +// is equivalent to the "padding high" for all dimensions, but works on a +// padding configuration. +bool HasSymmetricPadding(const PaddingConfig& padding_config); + bool HasBaseDilation(const Window& window); bool HasWindowDilation(const Window& window); bool HasDilation(const Window& window); bool HasWindowReversal(const Window& window); +// Returns true if the given logical dimension is inactive in the sense that it +// has window bound 1, no striding and no padding. +bool IsInactiveWindowDimension(const Window& window, int64 logical_dim); + // Returns the new bound after dilation. // // If a window with the given bound in some dimension is dilated with the given diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto index fda1a4c27b6dea1b7e4dee76de976f93ba61c007..56162ab44e2e0e3e4478fe631888f243332dc1d8 100644 --- a/tensorflow/compiler/xla/xla.proto +++ b/tensorflow/compiler/xla/xla.proto @@ -82,8 +82,9 @@ message DebugOptions { // Dump all HLO modules as text into the provided directory path. string xla_generate_hlo_text_to = 7; - // Dump compilation artifacts in binary proto into this directory. - string xla_dump_hlo_proto_to = 8; + // Dump Hlo after all hlo passes are executed as proto binary into this + // directory. + string xla_dump_optimized_hlo_proto_to = 8; // Instrument the computation to collect per-HLO cycle counts. bool xla_hlo_profile = 9; @@ -179,6 +180,14 @@ message DebugOptions { // ops. bool xla_gpu_use_cudnn_batchnorm = 94; + // Dump HLO before any hlo passes are executed as proto binary into this + // directory. + string xla_dump_unoptimized_hlo_proto_to = 95; + + // Dump HLO after each pass as an HloProto in binary file format into this + // directory. + string xla_dump_per_pass_hlo_proto_to = 96; + // Extra options to pass to the compilation backend; specific interpretation // of these values is left to the backend. map xla_backend_extra_options = 500; diff --git a/tensorflow/compiler/xla/xla_data.proto b/tensorflow/compiler/xla/xla_data.proto index 95045d5e28b96c8e9b31fccd62a24d5c83523092..3aea0217539b89b5d60ecfaf2605eee4b69af728 100644 --- a/tensorflow/compiler/xla/xla_data.proto +++ b/tensorflow/compiler/xla/xla_data.proto @@ -120,6 +120,9 @@ enum Format { // The default layout, with exactly one storage location per element (ignoring // padding). DENSE = 1; + // A sparsely encoded layout, providing only the index/value pairs of non-zero + // elements. + SPARSE = 2; } // A layout describes how the array is placed in (1D) memory space. This @@ -151,6 +154,11 @@ message Layout { // field must be unset unless the format is DENSE. PaddingValue padding_value = 3; + // The maximum number of elements that can be stored for SPARSE formats. This + // can be used to determine the maximum size in bytes of arrays stored in + // memory. This field must be unset unless the format is SPARSE. + int64 max_sparse_elements = 5; + // Important: if any field is added, be sure to modify ShapeUtil::Equal() // appropriately to account for the new field. } @@ -333,7 +341,8 @@ message LiteralProto { // The F16s and BF16s are encoded in little endian byte order bytes f16s = 11; bytes bf16s = 13; - // Next = 14 + repeated int64 sparse_indices = 14; + // Next = 15 } message WindowDimension { @@ -814,9 +823,7 @@ enum RandomDistribution { // parameter[0] and standard deviation parameter[1]. RNG_NORMAL = 2; - // Creates a Bernoulli-distribution-generated random number with mean - // parameter[0]. - RNG_BERNOULLI = 3; + // Next: 4 } message RngRequest { diff --git a/tensorflow/contrib/BUILD b/tensorflow/contrib/BUILD index 6e2320bd0d6376cfddb60f8069e141a88bc93563..bab37e8906e5c648acdc1556da7e5f4601776ff5 100644 --- a/tensorflow/contrib/BUILD +++ b/tensorflow/contrib/BUILD @@ -6,6 +6,8 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//tensorflow:__subpackages__"]) load("//third_party/mpi:mpi.bzl", "if_mpi") +load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") +load("@local_config_tensorrt//:build_defs.bzl", "if_tensorrt") py_library( name = "contrib_py", @@ -23,7 +25,9 @@ py_library( "//tensorflow/contrib/bayesflow:bayesflow_py", "//tensorflow/contrib/boosted_trees:init_py", "//tensorflow/contrib/cloud:cloud_py", + "//tensorflow/contrib/cluster_resolver:cluster_resolver_pip", "//tensorflow/contrib/cluster_resolver:cluster_resolver_py", + "//tensorflow/contrib/coder:coder_ops_py", "//tensorflow/contrib/compiler:compiler_py", "//tensorflow/contrib/copy_graph:copy_graph_py", "//tensorflow/contrib/crf:crf_py", @@ -34,6 +38,7 @@ py_library( "//tensorflow/contrib/eager/python:tfe", "//tensorflow/contrib/estimator:estimator_py", "//tensorflow/contrib/factorization:factorization_py", + "//tensorflow/contrib/feature_column:feature_column_py", "//tensorflow/contrib/ffmpeg:ffmpeg_ops_py", "//tensorflow/contrib/framework:framework_py", "//tensorflow/contrib/fused_conv:fused_conv_py", @@ -46,6 +51,7 @@ py_library( "//tensorflow/contrib/image:single_image_random_dot_stereograms_py", "//tensorflow/contrib/input_pipeline:input_pipeline_py", "//tensorflow/contrib/integrate:integrate_py", + "//tensorflow/contrib/kafka", "//tensorflow/contrib/keras", "//tensorflow/contrib/kernel_methods", "//tensorflow/contrib/kfac", @@ -66,7 +72,6 @@ py_library( "//tensorflow/contrib/metrics:metrics_py", "//tensorflow/contrib/model_pruning", "//tensorflow/contrib/nccl:nccl_py", - "//tensorflow/contrib/ndlstm", "//tensorflow/contrib/nearest_neighbor:nearest_neighbor_py", "//tensorflow/contrib/nn:nn_py", "//tensorflow/contrib/opt:opt_py", @@ -74,6 +79,8 @@ py_library( "//tensorflow/contrib/predictor", "//tensorflow/contrib/quantization:quantization_py", "//tensorflow/contrib/quantize:quantize_graph", + "//tensorflow/contrib/py2tf", + "//tensorflow/contrib/receptive_field:receptive_field_py", "//tensorflow/contrib/reduce_slice_ops:reduce_slice_ops_py", "//tensorflow/contrib/remote_fused_graph/pylib:remote_fused_graph_ops_py", "//tensorflow/contrib/resampler:resampler_py", @@ -101,20 +108,22 @@ py_library( "//tensorflow/contrib/training:training_py", "//tensorflow/contrib/util:util_py", "//tensorflow/python:util", - ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_ops_py"]), + ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]) + if_tensorrt([ + "//tensorflow/contrib/tensorrt:init_py", + ]), ) cc_library( name = "contrib_kernels", visibility = ["//visibility:public"], deps = [ - "//tensorflow/contrib/batching:batch_ops_kernels", "//tensorflow/contrib/boosted_trees:boosted_trees_kernels", + "//tensorflow/contrib/coder:all_kernels", "//tensorflow/contrib/cudnn_rnn:cudnn_rnn_kernels", + "//tensorflow/contrib/data/kernels:dataset_kernels", "//tensorflow/contrib/factorization/kernels:all_kernels", "//tensorflow/contrib/input_pipeline:input_pipeline_ops_kernels", "//tensorflow/contrib/layers:sparse_feature_cross_op_kernel", - "//tensorflow/contrib/nccl:nccl_kernels", "//tensorflow/contrib/nearest_neighbor:nearest_neighbor_ops_kernels", "//tensorflow/contrib/rnn:all_kernels", "//tensorflow/contrib/seq2seq:beam_search_ops_kernels", @@ -122,19 +131,23 @@ cc_library( "//tensorflow/contrib/tensor_forest:stats_ops_kernels", "//tensorflow/contrib/tensor_forest:tensor_forest_kernels", "//tensorflow/contrib/text:all_kernels", - ], + ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]) + if_cuda([ + "//tensorflow/contrib/nccl:nccl_kernels", + ]), ) cc_library( name = "contrib_ops_op_lib", visibility = ["//visibility:public"], deps = [ - "//tensorflow/contrib/batching:batch_ops_op_lib", "//tensorflow/contrib/boosted_trees:boosted_trees_ops_op_lib", + "//tensorflow/contrib/coder:all_ops", "//tensorflow/contrib/cudnn_rnn:cudnn_rnn_ops_op_lib", + "//tensorflow/contrib/data:dataset_ops_op_lib", "//tensorflow/contrib/factorization:all_ops", "//tensorflow/contrib/framework:all_ops", "//tensorflow/contrib/input_pipeline:input_pipeline_ops_op_lib", + "//tensorflow/contrib/kafka:kafka_ops_op_lib", "//tensorflow/contrib/layers:sparse_feature_cross_op_op_lib", "//tensorflow/contrib/nccl:nccl_ops_op_lib", "//tensorflow/contrib/nearest_neighbor:nearest_neighbor_ops_op_lib", diff --git a/tensorflow/contrib/__init__.py b/tensorflow/contrib/__init__.py index 08247c6b38a4df663ad28a6b4d3c41a1da41a020..4f6f539027b040de7554d09fe9118ff97aa006f8 100644 --- a/tensorflow/contrib/__init__.py +++ b/tensorflow/contrib/__init__.py @@ -19,9 +19,11 @@ from __future__ import division from __future__ import print_function # Add projects here, they will show up under tf.contrib. +from tensorflow.contrib import batching from tensorflow.contrib import bayesflow from tensorflow.contrib import cloud from tensorflow.contrib import cluster_resolver +from tensorflow.contrib import coder from tensorflow.contrib import compiler from tensorflow.contrib import copy_graph from tensorflow.contrib import crf @@ -31,6 +33,7 @@ from tensorflow.contrib import deprecated from tensorflow.contrib import distributions from tensorflow.contrib import estimator from tensorflow.contrib import factorization +from tensorflow.contrib import feature_column from tensorflow.contrib import framework from tensorflow.contrib import gan from tensorflow.contrib import graph_editor @@ -81,14 +84,14 @@ from tensorflow.contrib import training from tensorflow.contrib import util from tensorflow.contrib.eager.python import tfe as eager from tensorflow.contrib.lite.python import lite -from tensorflow.contrib.ndlstm import python as ndlstm +from tensorflow.contrib.receptive_field import receptive_field_api as receptive_field from tensorflow.contrib.remote_fused_graph import pylib as remote_fused_graph from tensorflow.contrib.specs import python as specs from tensorflow.contrib.summary import summary from tensorflow.python.util.lazy_loader import LazyLoader -ffmpeg = LazyLoader("ffmpeg", - globals(), "tensorflow.contrib.ffmpeg") +ffmpeg = LazyLoader("ffmpeg", globals(), + "tensorflow.contrib.ffmpeg") del LazyLoader del absolute_import diff --git a/tensorflow/contrib/all_reduce/python/all_reduce.py b/tensorflow/contrib/all_reduce/python/all_reduce.py index a5057da9fd43a88575813613d6ac9d17fd2b2e28..6658f0d9c13f6db17b25354cde2593d57f104f17 100644 --- a/tensorflow/contrib/all_reduce/python/all_reduce.py +++ b/tensorflow/contrib/all_reduce/python/all_reduce.py @@ -48,7 +48,7 @@ def _flatten_tensors(tensors): if shape.ndims is None: raise ValueError("At least one of the tensors in 'tensors' must have " "statically known rank.") - if len(shape) > 1: + if len(shape) != 1: reshaped = [] for t in tensors: with ops.colocate_with(t): @@ -289,7 +289,7 @@ def build_ring_all_reduce(input_tensors, num_workers, num_subchunks, chunks_by_dev) if pad_len > 0: output_tensors = _strip_padding(output_tensors, pad_len) - if len(shape) > 1: + if len(shape) != 1: output_tensors = _reshape_tensors(output_tensors, shape) return output_tensors @@ -466,7 +466,7 @@ def build_recursive_hd_all_reduce(input_tensors, red_op, un_op=None): if un_op: reduced_shards = [un_op(t) for t in reduced_shards] output_tensors = _build_recursive_hd_scatter(reduced_shards, devices) - if len(shape) > 1: + if len(shape) != 1: output_tensors = _reshape_tensors(output_tensors, shape) return output_tensors @@ -578,7 +578,7 @@ def build_shuffle_all_reduce(input_tensors, gather_devices, red_op, un_op=None): reduced_shards = _build_shuffle_gather(input_tensors, gather_devices, red_op, un_op) output_tensors = _build_shuffle_scatter(reduced_shards, dst_devices) - if len(shape) > 1: + if len(shape) != 1: output_tensors = _reshape_tensors(output_tensors, shape) return output_tensors @@ -744,21 +744,21 @@ def _build_nccl_hybrid(input_tensors, red_op, upper_level_f): level_2_output = upper_level_f(up_values) # Third stage: propagate within each worker using NCCL Broadcast for w in range(0, num_workers): - dst_devices = per_worker_devices[w][1:] - send_op, dst_tensors = nccl.broadcast(level_2_output[w], dst_devices) - # NOTE: need control dependency to ensure send_op executes - with ops.control_dependencies([send_op]): - with ops.device(per_worker_devices[w][0]): - dst_tensors.insert(0, array_ops.identity(level_2_output[w])) - down_values[w] = dst_tensors + dst_tensors = [] + with ops.device(per_worker_devices[w][0]): + broadcast_src = nccl.broadcast(array_ops.identity(level_2_output[w])) + for d in per_worker_devices[w]: + with ops.device(d): + dst_tensors.append(array_ops.identity(broadcast_src)) + down_values[w] = dst_tensors output_tensors = [v for sublist in down_values for v in sublist] - if len(shape) > 1: + if len(shape) != 1: output_tensors = _reshape_tensors(output_tensors, shape) return output_tensors def _reduce_non_singleton(input_tensors, red_f, un_op): - """If input_tenors has more than one element apply red_f, else apply un_op.""" + """If input_tensors has more than one element apply red_f, else apply un_op.""" if len(input_tensors) > 1: return red_f(input_tensors) else: @@ -831,7 +831,7 @@ def _build_shuffle_hybrid(input_tensors, gather_devices, red_op, upper_level_f): for w in range(0, num_workers): output_tensors += _build_shuffle_scatter( [level_2_output[w]], per_worker_devices[w]) - if len(shape) > 1: + if len(shape) != 1: output_tensors = _reshape_tensors(output_tensors, shape) return output_tensors diff --git a/tensorflow/contrib/all_reduce/python/all_reduce_test.py b/tensorflow/contrib/all_reduce/python/all_reduce_test.py index 0802b2736909c2a6f075ea2eac6d4dd3ab2918d8..47bab0a3670a90644972b2c961954a3036b8ecba 100644 --- a/tensorflow/contrib/all_reduce/python/all_reduce_test.py +++ b/tensorflow/contrib/all_reduce/python/all_reduce_test.py @@ -119,7 +119,7 @@ class AllReduceTest(test_util.TensorFlowTestCase): def _buildInitialVars(self, shape, dev_list): values = [] num_devices = len(dev_list) - dim = np.prod(shape) + dim = np.prod(shape) if shape else 1 for d in range(0, num_devices): with ops.device(dev_list[d]): npt = np.zeros(shape).astype(np.float32) @@ -164,6 +164,7 @@ class AllReduceTest(test_util.TensorFlowTestCase): (num_workers, num_gpus, shape, subdiv, elapsed)) def testRingAllReduce(self): + self._testRingAllReduce(1, 2, [], 1) self._testRingAllReduce(1, 2, [8], 1) self._testRingAllReduce(1, 2, [4, 4], 1) self._testRingAllReduce(6, 1, [8], 1) @@ -192,6 +193,7 @@ class AllReduceTest(test_util.TensorFlowTestCase): "elapsed=%f" % (num_workers, num_gpus, shape, elapsed)) def testShuffleAllReduce(self): + self._testShuffleAllReduce(1, 2, [], 1) self._testShuffleAllReduce(1, 2, [8], 1) self._testShuffleAllReduce(1, 2, [4, 4], 1) self._testShuffleAllReduce(1, 8, [32], 1) diff --git a/tensorflow/contrib/android/README.md b/tensorflow/contrib/android/README.md index b8d73bf24ce60e0b3850d4f39ac9e6d6c2194a02..db37bcf73d144eb81c32a461a276d10be7e2d193 100644 --- a/tensorflow/contrib/android/README.md +++ b/tensorflow/contrib/android/README.md @@ -81,6 +81,11 @@ For documentation on building a self-contained AAR file with cmake, see [tensorflow/contrib/android/cmake](cmake). +### Makefile + +For documentation on building native TF libraries with make, including a CUDA-enabled variant for devices like the Nvidia Shield TV, see [tensorflow/contrib/makefile/README.md](../makefile/README.md) + + ## AssetManagerFileSystem This directory also contains a TensorFlow filesystem supporting the Android diff --git a/tensorflow/contrib/android/asset_manager_filesystem.h b/tensorflow/contrib/android/asset_manager_filesystem.h index 2b43939f148e360945e5d488d148fcb2c13008a6..665304b5eef1f8a3633c8c522259e20d744b1808 100644 --- a/tensorflow/contrib/android/asset_manager_filesystem.h +++ b/tensorflow/contrib/android/asset_manager_filesystem.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_ANDROID_ASSET_MANAGER_FILESYSTEM_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_ANDROID_ASSET_MANAGER_FILESYSTEM_H_ +#ifndef TENSORFLOW_CONTRIB_ANDROID_ASSET_MANAGER_FILESYSTEM_H_ +#define TENSORFLOW_CONTRIB_ANDROID_ASSET_MANAGER_FILESYSTEM_H_ #include #include @@ -79,4 +79,4 @@ class AssetManagerFileSystem : public FileSystem { }; } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_ANDROID_ASSET_MANAGER_FILESYSTEM_H_ +#endif // TENSORFLOW_CONTRIB_ANDROID_ASSET_MANAGER_FILESYSTEM_H_ diff --git a/tensorflow/contrib/android/java/org/tensorflow/contrib/android/TensorFlowInferenceInterface.java b/tensorflow/contrib/android/java/org/tensorflow/contrib/android/TensorFlowInferenceInterface.java index dc5b9fb88742d78d0f40207b589e29451a6358dd..abddadac5bcace9b1f992b69bdcc69c24b29cd13 100644 --- a/tensorflow/contrib/android/java/org/tensorflow/contrib/android/TensorFlowInferenceInterface.java +++ b/tensorflow/contrib/android/java/org/tensorflow/contrib/android/TensorFlowInferenceInterface.java @@ -194,6 +194,11 @@ public class TensorFlowInferenceInterface { * @param outputNames A list of output nodes which should be filled by the inference pass. */ public void run(String[] outputNames, boolean enableStats) { + run(outputNames, enableStats, new String[] {}); + } + + /** An overloaded version of runInference that allows supplying targetNodeNames as well */ + public void run(String[] outputNames, boolean enableStats, String[] targetNodeNames) { // Release any Tensors from the previous run calls. closeFetches(); @@ -204,6 +209,11 @@ public class TensorFlowInferenceInterface { runner.fetch(tid.name, tid.outputIndex); } + // Add targets. + for (String t : targetNodeNames) { + runner.addTarget(t); + } + // Run the session. try { if (enableStats) { diff --git a/tensorflow/contrib/android/jni/run_stats_jni.cc b/tensorflow/contrib/android/jni/run_stats_jni.cc index 119fa9cd2c378d2ba2383ea8b0e09e1b6083d84e..707853b59befc2625145ad96952fbf9f66d62b43 100644 --- a/tensorflow/contrib/android/jni/run_stats_jni.cc +++ b/tensorflow/contrib/android/jni/run_stats_jni.cc @@ -21,8 +21,8 @@ limitations under the License. #include "tensorflow/core/protobuf/config.pb.h" #include "tensorflow/core/util/stat_summarizer.h" -using tensorflow::StatSummarizer; using tensorflow::RunMetadata; +using tensorflow::StatSummarizer; namespace { StatSummarizer* requireHandle(JNIEnv* env, jlong handle) { diff --git a/tensorflow/contrib/batching/BUILD b/tensorflow/contrib/batching/BUILD index ea8ac2c680e62ee03a45716aa1e0870d44495f1e..ee67909133fc26ba98355db05a4b90d3dfa6b97b 100644 --- a/tensorflow/contrib/batching/BUILD +++ b/tensorflow/contrib/batching/BUILD @@ -12,7 +12,7 @@ cc_library( name = "batch_scheduler_hdrs", hdrs = ["batch_scheduler.h"], deps = [ - "//tensorflow/core:framework_headers_lib", + "//tensorflow/core/kernels/batching_util:batch_scheduler_hdrs", ], ) @@ -20,18 +20,7 @@ cc_library( name = "batch_scheduler", hdrs = ["batch_scheduler.h"], deps = [ - "//tensorflow/core:lib", - ], -) - -tf_cc_test( - name = "batch_scheduler_test", - srcs = ["batch_scheduler_test.cc"], - deps = [ - ":batch_scheduler", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "//tensorflow/core:test_main", + "//tensorflow/core/kernels/batching_util:batch_scheduler", ], ) @@ -39,9 +28,7 @@ cc_library( name = "shared_batch_scheduler_hdrs", hdrs = ["shared_batch_scheduler.h"], deps = [ - ":batch_scheduler_hdrs", - "//tensorflow/contrib/batching/util:periodic_function_dynamic", - "//tensorflow/core:framework_headers_lib", + "//tensorflow/core/kernels/batching_util:shared_batch_scheduler_hdrs", ], ) @@ -49,49 +36,16 @@ cc_library( name = "shared_batch_scheduler", hdrs = ["shared_batch_scheduler.h"], deps = [ - ":batch_scheduler", - "//tensorflow/contrib/batching/util:periodic_function_dynamic", - "//tensorflow/core:lib", + "//tensorflow/core/kernels/batching_util:shared_batch_scheduler", ], alwayslink = 1, ) -tf_cc_test( - name = "shared_batch_scheduler_test", - srcs = ["shared_batch_scheduler_test.cc"], - deps = [ - ":shared_batch_scheduler", - "//tensorflow/contrib/batching/test_util:fake_clock_env", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - ], -) - cc_library( name = "adaptive_shared_batch_scheduler", hdrs = ["adaptive_shared_batch_scheduler.h"], deps = [ - ":batch_scheduler", - "//tensorflow/contrib/batching/util:periodic_function_dynamic", - "//tensorflow/core:lib", - ], -) - -tf_cc_test( - name = "adaptive_shared_batch_scheduler_test", - srcs = ["adaptive_shared_batch_scheduler_test.cc"], - tags = [ - "local", - "manual", - ], - deps = [ - ":adaptive_shared_batch_scheduler", - "//tensorflow/contrib/batching/test_util:fake_clock_env", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "//tensorflow/core:test_main", + "//tensorflow/core/kernels/batching_util:adaptive_shared_batch_scheduler", ], ) @@ -99,34 +53,7 @@ cc_library( name = "basic_batch_scheduler", hdrs = ["basic_batch_scheduler.h"], deps = [ - ":shared_batch_scheduler", - ], -) - -tf_cc_test( - name = "basic_batch_scheduler_test", - srcs = ["basic_batch_scheduler_test.cc"], - deps = [ - ":basic_batch_scheduler", - ":batch_scheduler", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - ], -) - -tf_cc_test( - name = "basic_batch_scheduler_benchmark", - srcs = ["basic_batch_scheduler_benchmark.cc"], - tags = [ - "local", - "manual", - ], - deps = [ - ":basic_batch_scheduler", - "//tensorflow/core:lib", - "//tensorflow/core:tensorflow", - "//tensorflow/core:test", + "//tensorflow/core/kernels/batching_util:basic_batch_scheduler", ], ) @@ -140,48 +67,14 @@ load( ) load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") -tf_custom_op_library( - name = "python/ops/_batch_ops.so", - srcs = ["ops/batch_ops.cc"], - deps = [ - "//tensorflow/contrib/batching/kernels:batch_kernels", - ], -) - -tf_gen_op_libs( - op_lib_names = ["batch_ops"], -) - -tf_gen_op_wrapper_py( - name = "batch_ops", - deps = [":batch_ops_op_lib"], -) - -tf_kernel_library( - name = "batch_ops_kernels", - deps = [ - "//tensorflow/contrib/batching/kernels:batch_kernels", - "//tensorflow/contrib/batching/util:periodic_function", - "//tensorflow/core/kernels:concat_lib", - "//tensorflow/core/kernels:ops_util", - "//tensorflow/core/kernels:split_lib", - ], - alwayslink = 1, -) - -tf_custom_op_py_library( +py_library( name = "batch_py", srcs = glob(["python/ops/*.py"]) + ["__init__.py"], - dso = [":python/ops/_batch_ops.so"], - kernels = [ - ":batch_ops_kernels", - ":batch_ops_op_lib", - ], srcs_version = "PY2AND3", deps = [ - ":batch_ops", "//tensorflow/contrib/util:util_py", "//tensorflow/python:array_ops", + "//tensorflow/python:batch_ops_gen", "//tensorflow/python:client_testlib", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:gradients", @@ -191,6 +84,14 @@ tf_custom_op_py_library( ], ) +cc_library( + name = "batch_ops_kernels", + deps = [ + "//tensorflow/core/kernels:batch_kernels", + ], + alwayslink = 1, +) + py_test( name = "batch_ops_test", size = "small", @@ -206,6 +107,7 @@ py_test( "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", + "//tensorflow/python:framework", "//tensorflow/python:gradients", "//tensorflow/python:script_ops", ], diff --git a/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h b/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h index 6773accc6fc813a81cba05b35c995aad4e801cb4..86250e6692004a12a1fa338767a5db1e4c2e4195 100644 --- a/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h +++ b/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h @@ -13,647 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_ +#ifndef TENSORFLOW_CONTRIB_BATCHING_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_ +#define TENSORFLOW_CONTRIB_BATCHING_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_ -#include -#include -#include -#include -#include -#include -#include +#include "tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h" -#include "tensorflow/contrib/batching/batch_scheduler.h" -#include "tensorflow/contrib/batching/util/periodic_function.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/core/threadpool.h" -#include "tensorflow/core/platform/cpu_info.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/thread_annotations.h" -#include "tensorflow/core/platform/types.h" - -namespace tensorflow { -namespace serving { -namespace internal { -template -class ASBSBatch; - -template -class ASBSQueue; -} // namespace internal - -// EXPERIMENTAL: API MAY BE SUBJECTED TO SUDDEN CHANGES. -// -// Shared batch scheduler designed to minimize latency. The scheduler keeps -// track of a number of queues (one per model or model version) which are -// continuously enqueuing requests. The scheduler groups the requests into -// batches which it periodically sends off for processing (see -// shared_batch_scheduler.h for more details). The AdaptiveSharedBatchScheduler -// prioritizes batches by age (i.e. the batch's oldest request) irrespective of -// queue or batch size. -// -// The scheduling decision currently exists in two flavors, controlled by the -// option use_in_flight_batches_implementation. It is expected that setting this -// option to true will give universally better results; after a period of -// testing to confirm, the old implementation will be removed. -// -// If use_in_flight_batches_implementation is set to true, the scheduler -// limits the number of batches which can be processed concurrently. If a new -// batch is created, and the number of in flight batches is below the limit, -// the next (i.e. oldest) batch is immediately scheduled. Similarly, when a -// batch finishes processing, the limit is rechecked, and another batch may be -// scheduled. To avoid the need to carefully tune the limit for workload, -// model type, platform, etc, it is dynamically adjusted in order to provide the -// lowest latency. -// -// If use_in_flight_batches_implementation is set to false, the scheduler will -// process the oldest batch at an adjustable rate, regardless of batch size. -// The user can provide feedback to help set this rate to achieve some goal -// (i.e. minimize overall latency, limit cpu usage, etc). The rate (or rather, -// the corresponding period) is adjusted each time a batch is processed, using -// an exponentially weighted moving average to smooth noisy feedback: -// ewma_feedback = ((N - 1) * ewma_feedback + feedback()) / N -// period *= (1 + K * emwa_feedback) -// -// Some potential use cases: -// Hardware Accelerators (GPUs & TPUs) - If some phase of batch processing -// involves serial processing by a device, from a latency perspective it is -// desirable to keep the device evenly loaded, avoiding the need to wait for -// the device to process prior batches. -// feedback = num_pending_on_device() - desired_pending. -// CPU utilization - If the batch processing is cpu dominated, you can reap -// latency gains when underutilized by increasing the processing rate, but -// back the rate off when the load increases to avoid overload. -// feedback = cpu_rate() - desired_cpu_rate. - -template -class AdaptiveSharedBatchScheduler - : public std::enable_shared_from_this< - AdaptiveSharedBatchScheduler> { - public: - ~AdaptiveSharedBatchScheduler() { - // Finish processing batches before destorying other class members. - batch_thread_pool_.reset(); - } - - struct Options { - // The name to use for the pool of batch threads. - string thread_pool_name = {"batch_threads"}; - // Number of batch processing threads; equivalently the maximum number of - // concurrently running batches. - int64 num_batch_threads = port::NumSchedulableCPUs(); - // The environment to use (typically only overridden by test code). - Env* env = Env::Default(); - // Which implementation to use (described in class comments above). - bool use_in_flight_batches_implementation = false; - // Initial limit for number of batches being concurrently processed. - // Non-integer values correspond to probabilistic limits - i.e. a value of - // 3.2 results in an actual cap of 3 80% of the time, and 4 20% of the time. - double initial_in_flight_batches_limit = 3; - // Number of batches between adjustments of in_flight_batches_limit. Larger - // numbers will give less noisy latency measurements, but will be less - // responsive to changes in workload. - int64 batches_to_average_over = 1000; - - // TODO(kte): remove the rate based implementation and corresponding options - // below once testing confirms the superiority of the in flight batches - // implementation. - // Initial batch scheduling period in microseconds. Will be altered for - // non-zero rate_feedback. - double initial_scheduling_period_micros = 500; - // Minimum batch scheduling period in microseconds. Recommend setting this - // value greater than 0, otherwise it may take a while to recover from a - // sustained time of negative scheduling_period_feedback (which may occur - // under low load). - double min_scheduling_period_micros = 100; - // Maximum batch scheduling period in microseconds. - double max_scheduling_period_micros = 10000; - // Feedback function used to modify the scheduling period each time a batch - // is scheduled. Should return values roughly O(1), with positive values - // resulting in an increased period. - std::function scheduling_period_feedback{[] { return 0.; }}; - // To handle potentially noisy scheduling_period_feedback, the period is - // adjusted using an exponentially weighted moving average over the previous - // feedback_smoothing_batches batches. Must be greater than 0. - int64 feedback_smoothing_batches = 10; - }; - - // Ownership is shared between the caller of Create() and any queues created - // via AddQueue(). - static Status Create( - const Options& options, - std::shared_ptr>* scheduler); - - struct QueueOptions { - // Maximum size of each batch. - int max_batch_size = 1000; - // Maximum number of enqueued (i.e. non-scheduled) batches. - int max_enqueued_batches = 10; - }; - - using BatchProcessor = std::function>)>; - - // Adds queue (and its callback) to be managed by this scheduler. - Status AddQueue(const QueueOptions& options, - BatchProcessor process_batch_callback, - std::unique_ptr>* queue); - - double in_flight_batches_limit() { - mutex_lock l(mu_); - return in_flight_batches_limit_; - } - - private: - // access to AddBatch, RemoveQueue, GetEnv. - friend class internal::ASBSQueue; - - explicit AdaptiveSharedBatchScheduler(const Options& options); - - // Batch scheduling function which runs every scheduling_period_ microseconds. - // Only used when options_.use_in_flight_batches_implementation == false. - void ProcessOneBatch(); - - // Tracks processing latency and adjusts in_flight_batches_limit to minimize. - // Only used when options_.use_in_flight_batches_implementation == true. - void CallbackWrapper(const internal::ASBSBatch* batch, - BatchProcessor callback); - - // Schedules batch if in_flight_batches_limit_ is not met. - // Only used when options_.use_in_flight_batches_implementation == true. - void MaybeScheduleNextBatch() EXCLUSIVE_LOCKS_REQUIRED(mu_); - - // Notifies scheduler of non-empty batch which is eligible for processing. - void AddBatch(const internal::ASBSBatch* batch); - - // Removes queue from scheduler. - void RemoveQueue(const internal::ASBSQueue* queue); - - Env* GetEnv() const { return options_.env; } - - const Options options_; - - struct BatchCompare { - bool operator()(const internal::ASBSBatch* a, - const internal::ASBSBatch* b); - }; - - // Collection of batches added by AddBatch, ordered by age. Owned by scheduler - // until they are released for processing. - std::priority_queue*, - std::vector*>, - BatchCompare> - batches_ GUARDED_BY(mu_); - - // Unowned queues and callbacks added by AddQueue. - std::unordered_map*, BatchProcessor> - queues_and_callbacks_ GUARDED_BY(mu_); - - mutex mu_; - - // Responsible for running ProcessOneBatch. PeriodicFunction was used in order - // to check for deletion so that the thread can be shut down. - // Only used when options_.use_in_flight_batches_implementation == false. - std::unique_ptr scheduling_thread_; - - // Responsible for running the batch processing callbacks. - std::unique_ptr batch_thread_pool_; - - // Time interval in microseconds between successive ProcessOneBatch calls. - // Only used when options_.use_in_flight_batches_implementation == false. - double scheduling_period_; - - // Exponentially weighted moving average of - // options_.scheduling_period_feedback() evaluated in each ProcessOneBatch - // call. - // Only used when options_.use_in_flight_batches_implementation == false. - double ewma_feedback_ = 0; - - // Limit on number of batches which can be concurrently processed. - // Non-integer values correspond to probabilistic limits - i.e. a value of 3.2 - // results in an actual cap of 3 80% of the time, and 4 20% of the time. - // Only used when options_.use_in_flight_batches_implementation == true. - double in_flight_batches_limit_ GUARDED_BY(mu_); - - // Number of batches currently being processed. - // Only used when options_.use_in_flight_batches_implementation == true. - int64 in_flight_batches_ GUARDED_BY(mu_) = 0; - - // RNG engine and distribution. - // Only used when options_.use_in_flight_batches_implementation == true. - std::default_random_engine rand_engine_; - std::uniform_real_distribution rand_double_; - - // Fields controlling the dynamic adjustment of in_flight_batches_limit_. - // Only used when options_.use_in_flight_batches_implementation == true. - // Number of batches since the last in_flight_batches_limit_ adjustment. - int64 batch_count_ GUARDED_BY(mu_) = 0; - // Sum of processing latency for batches counted by batch_count_. - int64 batch_latency_sum_ GUARDED_BY(mu_) = 0; - // Average batch latency for previous value of in_flight_batches_limit_. - double last_avg_latency_ms_ GUARDED_BY(mu_) = 0; - // Did last_avg_latency_ms_ decrease from the previous last_avg_latency_ms_? - bool last_latency_decreased_ GUARDED_BY(mu_) = false; - // Current direction (+-) to adjust in_flight_batches_limit_ - int step_direction_ GUARDED_BY(mu_) = 1; - // Max adjustment size (as a fraction of in_flight_batches_limit_). - constexpr static double kMaxStepSizeMultiplier = 0.125; // 1/8; - // Min adjustment size (as a fraction of in_flight_batches_limit_). - constexpr static double kMinStepSizeMultiplier = 0.0078125; // 1/128 - // Current adjustment size (as a fraction of in_flight_batches_limit_). - double step_size_multiplier_ GUARDED_BY(mu_) = kMaxStepSizeMultiplier; - - TF_DISALLOW_COPY_AND_ASSIGN(AdaptiveSharedBatchScheduler); -}; - -////////////////////////////////////////////////////////// -// Implementation details follow. API users need not read. - -namespace internal { -// Consolidates tasks into batches, passing them off to the -// AdaptiveSharedBatchScheduler for processing. -template -class ASBSQueue : public BatchScheduler { - public: - using QueueOptions = - typename AdaptiveSharedBatchScheduler::QueueOptions; - - ASBSQueue(std::shared_ptr> scheduler, - const QueueOptions& options); - - ~ASBSQueue() override; - - // Adds task to current batch. Fails if the task size is larger than the batch - // size or if the current batch is full and this queue's number of outstanding - // batches is at its maximum. - Status Schedule(std::unique_ptr* task) override; - - // Number of tasks waiting to be scheduled. - size_t NumEnqueuedTasks() const override; - - // Number of size 1 tasks which could currently be scheduled without failing. - size_t SchedulingCapacity() const override; - - // Notifies queue that a batch is about to be scheduled; the queue should not - // place any more tasks in this batch. - void ReleaseBatch(const ASBSBatch* batch); - - size_t max_task_size() const override { return options_.max_batch_size; } - - private: - std::shared_ptr> scheduler_; - const QueueOptions options_; - // Owned by scheduler_. - ASBSBatch* current_batch_ GUARDED_BY(mu_) = nullptr; - int64 num_enqueued_batches_ GUARDED_BY(mu_) = 0; - int64 num_enqueued_tasks_ GUARDED_BY(mu_) = 0; - mutable mutex mu_; - TF_DISALLOW_COPY_AND_ASSIGN(ASBSQueue); -}; - -// Batch which remembers when and by whom it was created. -template -class ASBSBatch : public Batch { - public: - ASBSBatch(ASBSQueue* queue, int64 creation_time_micros) - : queue_(queue), creation_time_micros_(creation_time_micros) {} - - ~ASBSBatch() override {} - - ASBSQueue* queue() const { return queue_; } - - int64 creation_time_micros() const { return creation_time_micros_; } - - private: - ASBSQueue* queue_; - const int64 creation_time_micros_; - TF_DISALLOW_COPY_AND_ASSIGN(ASBSBatch); -}; -} // namespace internal - -// ---------------- AdaptiveSharedBatchScheduler ---------------- - -template -constexpr double AdaptiveSharedBatchScheduler::kMaxStepSizeMultiplier; - -template -constexpr double AdaptiveSharedBatchScheduler::kMinStepSizeMultiplier; - -template -Status AdaptiveSharedBatchScheduler::Create( - const Options& options, - std::shared_ptr>* scheduler) { - if (options.num_batch_threads < 1) { - return errors::InvalidArgument("num_batch_threads must be positive; was ", - options.num_batch_threads); - } - if (options.min_scheduling_period_micros < 0) { - return errors::InvalidArgument( - "min_scheduling_period_micros must be >= 0; was ", - options.min_scheduling_period_micros); - } - if (options.min_scheduling_period_micros > - options.initial_scheduling_period_micros) { - return errors::InvalidArgument( - "initial_scheduling_period_micros (", - options.initial_scheduling_period_micros, - ") must be >= min_scheduling_period_micros (", - options.min_scheduling_period_micros, ")"); - } - if (options.initial_scheduling_period_micros > - options.max_scheduling_period_micros) { - return errors::InvalidArgument( - "initial_scheduling_period_micros (", - options.initial_scheduling_period_micros, - ") must be <= max_scheduling_period_micros (", - options.max_scheduling_period_micros, ")"); - } - if (options.feedback_smoothing_batches < 1) { - return errors::InvalidArgument( - "feedback_smoothing_batches must be positive; was ", - options.feedback_smoothing_batches); - } - if (options.initial_in_flight_batches_limit > options.num_batch_threads) { - return errors::InvalidArgument( - "initial_in_flight_batches_limit (", - options.initial_in_flight_batches_limit, - ") should not be larger than num_batch_threads (", - options.num_batch_threads, ")"); - } - if (options.initial_in_flight_batches_limit < 1) { - return errors::InvalidArgument( - "initial_in_flight_batches_limit should be " - "greater than or equal to 1; was ", - options.initial_in_flight_batches_limit); - } - if (options.batches_to_average_over < 1) { - return errors::InvalidArgument( - "batches_to_average_over should be " - "greater than or equal to 1; was ", - options.batches_to_average_over); - } - scheduler->reset(new AdaptiveSharedBatchScheduler(options)); - return Status::OK(); -} - -template -AdaptiveSharedBatchScheduler::AdaptiveSharedBatchScheduler( - const Options& options) - : options_(options), - scheduling_period_(options.initial_scheduling_period_micros), - in_flight_batches_limit_(options.initial_in_flight_batches_limit), - rand_double_(0.0, 1.0) { - std::random_device device; - rand_engine_.seed(device()); - PeriodicFunction::Options opts; - opts.thread_name_prefix = "scheduling_thread"; - opts.env = GetEnv(); - batch_thread_pool_.reset(new thread::ThreadPool( - GetEnv(), options.thread_pool_name, options.num_batch_threads)); - if (!options.use_in_flight_batches_implementation) { - scheduling_thread_.reset( - new PeriodicFunction([this] { ProcessOneBatch(); }, 0, opts)); - } -} - -template -Status AdaptiveSharedBatchScheduler::AddQueue( - const QueueOptions& options, BatchProcessor process_batch_callback, - std::unique_ptr>* queue) { - if (options.max_batch_size <= 0) { - return errors::InvalidArgument("max_batch_size must be positive; was ", - options.max_batch_size); - } - if (options.max_enqueued_batches <= 0) { - return errors::InvalidArgument( - "max_enqueued_batches must be positive; was ", - options.max_enqueued_batches); - } - internal::ASBSQueue* asbs_queue_raw; - queue->reset(asbs_queue_raw = new internal::ASBSQueue( - this->shared_from_this(), options)); - mutex_lock l(mu_); - queues_and_callbacks_[asbs_queue_raw] = process_batch_callback; - return Status::OK(); -} - -template -void AdaptiveSharedBatchScheduler::AddBatch( - const internal::ASBSBatch* batch) { - mutex_lock l(mu_); - batches_.push(batch); - if (options_.use_in_flight_batches_implementation) { - MaybeScheduleNextBatch(); - } -} - -template -void AdaptiveSharedBatchScheduler::RemoveQueue( - const internal::ASBSQueue* queue) { - mutex_lock l(mu_); - queues_and_callbacks_.erase(queue); -} - -template -void AdaptiveSharedBatchScheduler::MaybeScheduleNextBatch() { - if (batches_.empty() || in_flight_batches_ >= in_flight_batches_limit_) - return; - // Non-integer limit handled probabilistially. - if (in_flight_batches_limit_ - in_flight_batches_ < 1 && - rand_double_(rand_engine_) > - (in_flight_batches_limit_ - in_flight_batches_)) - return; - const internal::ASBSBatch* batch = batches_.top(); - batches_.pop(); - // Queue may destroy itself after ReleaseBatch is called. - batch->queue()->ReleaseBatch(batch); - batch_thread_pool_->Schedule( - std::bind(&AdaptiveSharedBatchScheduler::CallbackWrapper, this, - batch, queues_and_callbacks_[batch->queue()])); - in_flight_batches_++; -} - -template -void AdaptiveSharedBatchScheduler::CallbackWrapper( - const internal::ASBSBatch* batch, - AdaptiveSharedBatchScheduler::BatchProcessor callback) { - int64 start_time = batch->creation_time_micros(); - callback(std::unique_ptr>( - const_cast*>(batch))); - int64 end_time = GetEnv()->NowMicros(); - mutex_lock l(mu_); - in_flight_batches_--; - batch_count_++; - batch_latency_sum_ += end_time - start_time; - // Occasionally adjust in_flight_batches_limit_ to minimize average latency. - // Although the optimal value may depend on the workload, the latency should - // be a simple convex function of in_flight_batches_limit_, allowing us to - // locate the global minimum relatively quickly. - if (batch_count_ == options_.batches_to_average_over) { - double current_avg_latency_ms = (batch_latency_sum_ / 1000.) / batch_count_; - bool current_latency_decreased = - current_avg_latency_ms < last_avg_latency_ms_; - if (current_latency_decreased) { - // If latency improvement was because we're moving in the correct - // direction, increase step_size so that we can get to the minimum faster. - // If latency improvement was due to backtracking from a previous failure, - // decrease step_size in order to refine our location. - step_size_multiplier_ *= (last_latency_decreased_ ? 2 : 0.5); - step_size_multiplier_ = - std::min(step_size_multiplier_, kMaxStepSizeMultiplier); - step_size_multiplier_ = - std::max(step_size_multiplier_, kMinStepSizeMultiplier); - } else { - // Return (nearly) to previous position and confirm that latency is better - // there before decreasing step size. - step_direction_ = -step_direction_; - } - in_flight_batches_limit_ += - step_direction_ * in_flight_batches_limit_ * step_size_multiplier_; - in_flight_batches_limit_ = - std::min(in_flight_batches_limit_, - static_cast(options_.num_batch_threads)); - in_flight_batches_limit_ = std::max(in_flight_batches_limit_, 1.0); - last_avg_latency_ms_ = current_avg_latency_ms; - last_latency_decreased_ = current_latency_decreased; - batch_count_ = 0; - batch_latency_sum_ = 0; - } - MaybeScheduleNextBatch(); -} - -template -void AdaptiveSharedBatchScheduler::ProcessOneBatch() { - static const double kFeedbackMultiplier = .001; - const internal::ASBSBatch* batch = nullptr; - BatchProcessor callback; - const int64 start_time_micros = GetEnv()->NowMicros(); - { - mutex_lock l(mu_); - if (!batches_.empty()) { - batch = batches_.top(); - batches_.pop(); - callback = queues_and_callbacks_[batch->queue()]; - } - } - if (batch != nullptr) { - double feedback = options_.scheduling_period_feedback(); - const int64 N = options_.feedback_smoothing_batches; - ewma_feedback_ = ((N - 1) * ewma_feedback_ + feedback) / N; - scheduling_period_ *= (1 + kFeedbackMultiplier * ewma_feedback_); - if (scheduling_period_ < options_.min_scheduling_period_micros) { - scheduling_period_ = options_.min_scheduling_period_micros; - } else if (scheduling_period_ > options_.max_scheduling_period_micros) { - scheduling_period_ = options_.max_scheduling_period_micros; - } - // Queue may destroy itself after ReleaseBatch is called. - batch->queue()->ReleaseBatch(batch); - batch_thread_pool_->Schedule([callback, batch] { - callback(std::unique_ptr>( - const_cast*>(batch))); - }); - } - const int64 sleep_time = - scheduling_period_ - (GetEnv()->NowMicros() - start_time_micros); - if (sleep_time > 0) { - GetEnv()->SleepForMicroseconds(sleep_time); - } -} - -template -bool AdaptiveSharedBatchScheduler::BatchCompare::operator()( - const internal::ASBSBatch* a, - const internal::ASBSBatch* b) { - return a->creation_time_micros() > b->creation_time_micros(); -} - -// ---------------- ASBSQueue ---------------- - -namespace internal { -template -ASBSQueue::ASBSQueue( - std::shared_ptr> scheduler, - const QueueOptions& options) - : scheduler_(scheduler), options_(options) {} - -template -ASBSQueue::~ASBSQueue() { - // Wait until last batch has been scheduled. - const int kSleepMicros = 1000; - for (;;) { - { - mutex_lock l(mu_); - if (num_enqueued_batches_ == 0) { - break; - } - } - scheduler_->GetEnv()->SleepForMicroseconds(kSleepMicros); - } - scheduler_->RemoveQueue(this); -} - -template -Status ASBSQueue::Schedule(std::unique_ptr* task) { - ASBSBatch* new_batch = nullptr; - size_t size = (*task)->size(); - if (size > options_.max_batch_size) { - return errors::InvalidArgument("Task size ", size, - " is larger than maximum batch size ", - options_.max_batch_size); - } - { - mutex_lock l(mu_); - // Current batch is full, create another if allowed. - if (current_batch_ && - current_batch_->size() + size > options_.max_batch_size) { - if (num_enqueued_batches_ >= options_.max_enqueued_batches) { - return errors::Unavailable("The batch scheduling queue is full"); - } - current_batch_->Close(); - current_batch_ = nullptr; - } - if (!current_batch_) { - num_enqueued_batches_++; - current_batch_ = new_batch = - new ASBSBatch(this, scheduler_->GetEnv()->NowMicros()); - } - current_batch_->AddTask(std::move(*task)); - num_enqueued_tasks_++; - } - // AddBatch must be called outside of lock, since it may call ReleaseBatch. - if (new_batch != nullptr) scheduler_->AddBatch(new_batch); - return Status::OK(); -} - -template -void ASBSQueue::ReleaseBatch(const ASBSBatch* batch) { - mutex_lock l(mu_); - num_enqueued_batches_--; - num_enqueued_tasks_ -= batch->num_tasks(); - if (batch == current_batch_) { - current_batch_->Close(); - current_batch_ = nullptr; - } -} - -template -size_t ASBSQueue::NumEnqueuedTasks() const { - mutex_lock l(mu_); - return num_enqueued_tasks_; -} - -template -size_t ASBSQueue::SchedulingCapacity() const { - mutex_lock l(mu_); - const int current_batch_capacity = - current_batch_ ? options_.max_batch_size - current_batch_->size() : 0; - const int spare_batches = - options_.max_enqueued_batches - num_enqueued_batches_; - return spare_batches * options_.max_batch_size + current_batch_capacity; -} -} // namespace internal -} // namespace serving -} // namespace tensorflow - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_ +#endif // TENSORFLOW_CONTRIB_BATCHING_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_ diff --git a/tensorflow/contrib/batching/basic_batch_scheduler.h b/tensorflow/contrib/batching/basic_batch_scheduler.h index 91065db2499dffd2687a53bd6304d9b7593f7b3a..d9b37da6933aa0847c229607f43d1d5d121a928c 100644 --- a/tensorflow/contrib/batching/basic_batch_scheduler.h +++ b/tensorflow/contrib/batching/basic_batch_scheduler.h @@ -13,256 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_BASIC_BATCH_SCHEDULER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_BASIC_BATCH_SCHEDULER_H_ +#ifndef TENSORFLOW_CONTRIB_BATCHING_BASIC_BATCH_SCHEDULER_H_ +#define TENSORFLOW_CONTRIB_BATCHING_BASIC_BATCH_SCHEDULER_H_ -#include -#include -#include -#include -#include +#include "tensorflow/core/kernels/batching_util/basic_batch_scheduler.h" -#include "tensorflow/contrib/batching/shared_batch_scheduler.h" - -namespace tensorflow { -namespace serving { - -// A BatchScheduler implementation geared toward handling a single request type -// running on a specific set of hardware resources. A typical scenario is one in -// which all requests invoke the same machine-learned model on one GPU. -// -// If there are, say, two GPUs and two models each bound to one of the GPUs, one -// could use two BasicBatchScheduler instances to schedule the two model/GPU -// combinations independently. If multiple models must share a given GPU or -// other hardware resource, consider using SharedBatchScheduler instead. -// -// -// PARAMETERS AND BEHAVIOR: -// -// BasicBatchScheduler runs a fixed pool of threads, which it uses to process -// batches of tasks. It enforces a maximum batch size, and enqueues a bounded -// number of tasks. If the queue is nearly empty, such that a full batch cannot -// be formed, when a thread becomes free, it anyway schedules a batch -// immediately if a task has been in the queue for longer than a given timeout -// parameter. If the timeout parameter is set to 0, then the batch threads will -// always be kept busy (unless there are zero tasks waiting to be processed). -// -// For online serving, it is recommended to set the maximum number of enqueued -// batches worth of tasks equal to the number of batch threads, which allows -// enqueuing of enough tasks s.t. if every thread becomes available it can be -// kept busy, but no more. For bulk processing jobs and throughput-oriented -// benchmarks, you may want to set it much higher. -// -// When Schedule() is called, if the queue is full the call will fail with an -// UNAVAILABLE error (after which the client may retry again later). If the call -// succeeds, the maximum time the task will spend in the queue before being -// placed in a batch and assigned to a thread for processing, is the greater of: -// - the maximum time to process ceil(max_enqueued_batches/num_batch_threads) -// (1 in the recommended configuration) batches of previously-submitted tasks -// - the configured timeout parameter (which can be 0, as mentioned above) -// -// Unlike StreamingBatchScheduler, when BasicBatchScheduler assigns a batch to a -// thread, it closes the batch. The process-batch callback may assume that every -// batch it receives is closed at the outset. -// -// -// RECOMMENDED USE-CASES: -// -// BasicBatchScheduler is suitable for use-cases that feature a single kind of -// request (e.g. a server performing inference with a single machine-learned -// model, possibly evolving over time), with loose versioning semantics. -// Concretely, the following conditions should hold: -// -// A. All requests batched onto a given resource (e.g. a hardware accelerator, -// or a pool accelerators) are of the same type. For example, they all -// invoke the same machine-learned model. -// -// These variations are permitted: -// - The model may reside in a single servable, or it may be spread across -// multiple servables that are used in unison (e.g. a vocabulary lookup -// table servable and a tensorflow session servable). -// - The model's servable(s) may be static, or they may evolve over time -// (successive servable versions). -// - Zero or more of the servables are used in the request thread; the rest -// are used in the batch thread. In our running example, the vocabulary -// lookups and tensorflow runs may both be performed in the batch thread, -// or alternatively the vocabulary lookup may occur in the request thread -// with only the tensorflow run performed in the batch thread. -// -// In contrast, BasicBatchScheduler is not a good fit if the server -// hosts multiple distinct models running on a pool accelerators, with each -// request specifying which model it wants to use. BasicBatchScheduler -// has no facility to time-multiplex the batch threads across multiple -// models in a principled way. More basically, it cannot ensure that a given -// batch doesn't contain a mixture of requests for different models. -// -// B. Requests do not specify a particular version of the servable(s) that must -// be used. Instead, each request is content to use the "latest" version. -// -// BasicBatchScheduler does not constrain which requests get grouped -// together into a batch, so using this scheduler there is no way to achieve -// cohesion of versioned requests to version-specific batches. -// -// C. No servable version coordination needs to be performed between the -// request threads and the batch threads. Often, servables are only used in -// the batch threads, in which case this condition trivially holds. If -// servables are used in both threads, then the use-case must tolerate -// version skew across the servables used in the two kinds of threads. -// -// -// EXAMPLE USE-CASE FLOW: -// -// For such use-cases, request processing via BasicBatchScheduler generally -// follows this flow (given for illustration; variations are possible): -// 1. Optionally perform some pre-processing on each request in the request -// threads. -// 2. Route the requests to the batch scheduler, as batching::Task objects. -// (Since all requests are of the same type and are not versioned, the -// scheduler is free to group them into batches arbitrarily.) -// 3. Merge the requests into a single batched representation B. -// 4. Obtain handles to the servable(s) needed to process B. The simplest -// approach is to obtain the latest version of each servable. Alternatively, -// if cross-servable consistency is required (e.g. the vocabulary lookup -// table's version number must match that of the tensorflow session), -// identify an appropriate version number and obtain the servable handles -// accordingly. -// 5. Process B using the obtained servable handles, and split the result into -// individual per-request units. -// 6. Perform any post-processing in the batch thread and/or request thread. -// -// -// PERFORMANCE TUNING: See README.md. -// -template -class BasicBatchScheduler : public BatchScheduler { - public: - // TODO(b/25089730): Tune defaults based on best practices as they develop. - // (Keep them mirrored to the ones in SharedBatchScheduler::QueueOptions and - // SharedBatchScheduler::Options.) - struct Options { - // The maximum size of each batch. - // - // The scheduler may form batches of any size between 1 and this number - // (inclusive). If there is a need to quantize the batch sizes, i.e. only - // submit batches whose size is in a small set of allowed sizes, that can be - // done by adding padding in the process-batch callback. - int max_batch_size = 1000; - - // If a task has been enqueued for this amount of time (in microseconds), - // and a thread is available, the scheduler will immediately form a batch - // from enqueued tasks and assign the batch to the thread for processing, - // even if the batch's size is below 'max_batch_size'. - // - // This parameter offers a way to bound queue latency, so that a task isn't - // stuck in the queue indefinitely waiting for enough tasks to arrive to - // make a full batch. (The latency bound is given in the class documentation - // above.) - // - // The goal is to smooth out batch sizes under low request rates, and thus - // avoid latency spikes. - int64 batch_timeout_micros = 0; - - // The name to use for the pool of batch threads. - string thread_pool_name = {"batch_threads"}; - - // The number of threads to use to process batches. - // Must be >= 1, and should be tuned carefully. - int num_batch_threads = port::NumSchedulableCPUs(); - - // The maximum allowable number of enqueued (accepted by Schedule() but - // not yet being processed on a batch thread) tasks in terms of batches. - // If this limit is reached, Schedule() will return an UNAVAILABLE error. - // See the class documentation above for guidelines on how to tune this - // parameter. - int max_enqueued_batches = 10; - - // The following options are typically only overridden by test code. - - // The environment to use. - Env* env = Env::Default(); - }; - static Status Create(const Options& options, - std::function>)> - process_batch_callback, - std::unique_ptr* scheduler); - - ~BasicBatchScheduler() override = default; - - Status Schedule(std::unique_ptr* task) override; - size_t NumEnqueuedTasks() const override; - size_t SchedulingCapacity() const override; - - size_t max_task_size() const override { - return shared_scheduler_queue_->max_task_size(); - } - - private: - explicit BasicBatchScheduler( - std::unique_ptr> shared_scheduler_queue); - - // This class is merely a thin wrapper around a SharedBatchScheduler with a - // single queue. - std::unique_ptr> shared_scheduler_queue_; - - TF_DISALLOW_COPY_AND_ASSIGN(BasicBatchScheduler); -}; - -////////// -// Implementation details follow. API users need not read. - -template -Status BasicBatchScheduler::Create( - const Options& options, - std::function>)> - process_batch_callback, - std::unique_ptr* scheduler) { - typename SharedBatchScheduler::Options shared_scheduler_options; - shared_scheduler_options.thread_pool_name = options.thread_pool_name; - shared_scheduler_options.num_batch_threads = options.num_batch_threads; - shared_scheduler_options.env = options.env; - std::shared_ptr> shared_scheduler; - TF_RETURN_IF_ERROR(SharedBatchScheduler::Create( - shared_scheduler_options, &shared_scheduler)); - - typename SharedBatchScheduler::QueueOptions - shared_scheduler_queue_options; - shared_scheduler_queue_options.max_batch_size = options.max_batch_size; - shared_scheduler_queue_options.batch_timeout_micros = - options.batch_timeout_micros; - shared_scheduler_queue_options.max_enqueued_batches = - options.max_enqueued_batches; - std::unique_ptr> shared_scheduler_queue; - TF_RETURN_IF_ERROR(shared_scheduler->AddQueue(shared_scheduler_queue_options, - process_batch_callback, - &shared_scheduler_queue)); - - scheduler->reset( - new BasicBatchScheduler(std::move(shared_scheduler_queue))); - return Status::OK(); -} - -template -Status BasicBatchScheduler::Schedule( - std::unique_ptr* task) { - return shared_scheduler_queue_->Schedule(task); -} - -template -size_t BasicBatchScheduler::NumEnqueuedTasks() const { - return shared_scheduler_queue_->NumEnqueuedTasks(); -} - -template -size_t BasicBatchScheduler::SchedulingCapacity() const { - return shared_scheduler_queue_->SchedulingCapacity(); -} - -template -BasicBatchScheduler::BasicBatchScheduler( - std::unique_ptr> shared_scheduler_queue) - : shared_scheduler_queue_(std::move(shared_scheduler_queue)) {} - -} // namespace serving -} // namespace tensorflow - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_BASIC_BATCH_SCHEDULER_H_ +#endif // TENSORFLOW_CONTRIB_BATCHING_BASIC_BATCH_SCHEDULER_H_ diff --git a/tensorflow/contrib/batching/batch_scheduler.h b/tensorflow/contrib/batching/batch_scheduler.h index e18cf6c35059e4d720768e3b2c02b03727a6bac4..8e94e1fd8b969d4fef8dbc8c322557f9da3833e6 100644 --- a/tensorflow/contrib/batching/batch_scheduler.h +++ b/tensorflow/contrib/batching/batch_scheduler.h @@ -13,268 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// Abstractions for processing small tasks in a batched fashion, to reduce -// processing times and costs that can be amortized across multiple tasks. -// -// The core class is BatchScheduler, which groups tasks into batches. -// -// BatchScheduler encapsulates logic for aggregating multiple tasks into a -// batch, and kicking off processing of a batch on a thread pool it manages. -// -// This file defines an abstract BatchScheduler class. +#ifndef TENSORFLOW_CONTRIB_BATCHING_BATCH_SCHEDULER_H_ +#define TENSORFLOW_CONTRIB_BATCHING_BATCH_SCHEDULER_H_ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_BATCH_SCHEDULER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_BATCH_SCHEDULER_H_ +#include "tensorflow/core/kernels/batching_util/batch_scheduler.h" -#include -#include -#include -#include -#include -#include - -#include "tensorflow/core/lib/core/notification.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/macros.h" -#include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/thread_annotations.h" -#include "tensorflow/core/platform/types.h" - -namespace tensorflow { -namespace serving { - -// The abstract superclass for a unit of work to be done as part of a batch. -// -// An implementing subclass typically contains (or points to): -// (a) input data; -// (b) a thread-safe completion signal (e.g. a Notification); -// (c) a place to store the outcome (success, or some error), upon completion; -// (d) a place to store the output data, upon success. -// -// Items (b), (c) and (d) are typically non-owned pointers to data homed -// elsewhere, because a task's ownership gets transferred to a BatchScheduler -// (see below) and it may be deleted as soon as it is done executing. -class BatchTask { - public: - virtual ~BatchTask() = default; - - // Returns the size of the task, in terms of how much it contributes to the - // size of a batch. (A batch's size is the sum of its task sizes.) - virtual size_t size() const = 0; -}; - -// A thread-safe collection of BatchTasks, to be executed together in some -// fashion. -// -// At a given time, a batch is either "open" or "closed": an open batch can -// accept new tasks; a closed one cannot. A batch is monotonic: initially it is -// open and tasks can be added to it; then it is closed and its set of tasks -// remains fixed for the remainder of its life. A closed batch cannot be re- -// opened. Tasks can never be removed from a batch. -// -// Type parameter TaskType must be a subclass of BatchTask. -template -class Batch { - public: - Batch() = default; - virtual ~Batch(); // Blocks until the batch is closed. - - // Appends 'task' to the batch. After calling AddTask(), the newly-added task - // can be accessed via task(num_tasks()-1) or mutable_task(num_tasks()-1). - // Dies if the batch is closed. - void AddTask(std::unique_ptr task); - - // Removes the most recently added task. Returns nullptr if the batch is - // empty. - std::unique_ptr RemoveTask(); - - // Returns the number of tasks in the batch. - int num_tasks() const; - - // Returns true iff the batch contains 0 tasks. - bool empty() const; - - // Returns a reference to the ith task (in terms of insertion order). - const TaskType& task(int i) const; - - // Returns a pointer to the ith task (in terms of insertion order). - TaskType* mutable_task(int i); - - // Returns the sum of the task sizes. - size_t size() const; - - // Returns true iff the batch is currently closed. - bool IsClosed() const; - - // Blocks until the batch is closed. - void WaitUntilClosed() const; - - // Marks the batch as closed. Dies if called more than once. - void Close(); - - private: - mutable mutex mu_; - - // The tasks in the batch. - std::vector> tasks_ GUARDED_BY(mu_); - - // The sum of the sizes of the tasks in 'tasks_'. - size_t size_ GUARDED_BY(mu_) = 0; - - // Whether the batch has been closed. - Notification closed_; - - TF_DISALLOW_COPY_AND_ASSIGN(Batch); -}; - -// An abstract batch scheduler class. Collects individual tasks into batches, -// and processes each batch on a pool of "batch threads" that it manages. The -// actual logic for processing a batch is accomplished via a callback. -// -// Type parameter TaskType must be a subclass of BatchTask. -template -class BatchScheduler { - public: - virtual ~BatchScheduler() = default; - - // Submits a task to be processed as part of a batch. - // - // Ownership of '*task' is transferred to the callee iff the method returns - // Status::OK. In that case, '*task' is left as nullptr. Otherwise, '*task' is - // left as-is. - // - // If no batch processing capacity is available to process this task at the - // present time, and any task queue maintained by the implementing subclass is - // full, this method returns an UNAVAILABLE error code. The client may retry - // later. - // - // Other problems, such as the task size being larger than the maximum batch - // size, yield other, permanent error types. - // - // In all cases, this method returns "quickly" without blocking for any - // substantial amount of time. If the method returns Status::OK, the task is - // processed asynchronously, and any errors that occur during the processing - // of the batch that includes the task can be reported to 'task'. - virtual Status Schedule(std::unique_ptr* task) = 0; - - // Returns the number of tasks that have been scheduled (i.e. accepted by - // Schedule()), but have yet to be handed to a thread for execution as part of - // a batch. Note that this returns the number of tasks, not the aggregate task - // size (so if there is one task of size 3 and one task of size 5, this method - // returns 2 rather than 8). - virtual size_t NumEnqueuedTasks() const = 0; - - // Returns a guaranteed number of size 1 tasks that can be Schedule()d without - // getting an UNAVAILABLE error. In a typical implementation, returns the - // available space on a queue. - // - // There are two important caveats: - // 1. The guarantee does not extend to varying-size tasks due to possible - // internal fragmentation of batches. - // 2. The guarantee only holds in a single-thread environment or critical - // section, i.e. if an intervening thread cannot call Schedule(). - // - // This method is useful for monitoring, or for guaranteeing a future slot in - // the schedule (but being mindful about the caveats listed above). - virtual size_t SchedulingCapacity() const = 0; - - // Returns the maximum allowed size of tasks submitted to the scheduler. (This - // is typically equal to a configured maximum batch size.) - virtual size_t max_task_size() const = 0; -}; - -////////// -// Implementation details follow. API users need not read. - -template -Batch::~Batch() { - WaitUntilClosed(); -} - -template -void Batch::AddTask(std::unique_ptr task) { - DCHECK(!IsClosed()); - { - mutex_lock l(mu_); - size_ += task->size(); - tasks_.push_back(std::move(task)); - } -} - -template -std::unique_ptr Batch::RemoveTask() { - { - mutex_lock l(mu_); - if (tasks_.empty()) { - return nullptr; - } - std::unique_ptr task = std::move(tasks_.back()); - tasks_.pop_back(); - return task; - } -} - -template -int Batch::num_tasks() const { - { - mutex_lock l(mu_); - return tasks_.size(); - } -} - -template -bool Batch::empty() const { - { - mutex_lock l(mu_); - return tasks_.empty(); - } -} - -template -const TaskType& Batch::task(int i) const { - DCHECK_GE(i, 0); - { - mutex_lock l(mu_); - DCHECK_LT(i, tasks_.size()); - return *tasks_[i].get(); - } -} - -template -TaskType* Batch::mutable_task(int i) { - DCHECK_GE(i, 0); - { - mutex_lock l(mu_); - DCHECK_LT(i, tasks_.size()); - return tasks_[i].get(); - } -} - -template -size_t Batch::size() const { - { - mutex_lock l(mu_); - return size_; - } -} - -template -bool Batch::IsClosed() const { - return const_cast(&closed_)->HasBeenNotified(); -} - -template -void Batch::WaitUntilClosed() const { - const_cast(&closed_)->WaitForNotification(); -} - -template -void Batch::Close() { - closed_.Notify(); -} - -} // namespace serving -} // namespace tensorflow - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_BATCH_SCHEDULER_H_ +#endif // TENSORFLOW_CONTRIB_BATCHING_BATCH_SCHEDULER_H_ diff --git a/tensorflow/contrib/batching/kernels/BUILD b/tensorflow/contrib/batching/kernels/BUILD deleted file mode 100644 index 6e53dd9a5fc0201c5ed91d1eaf07f940e341fb5e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/batching/kernels/BUILD +++ /dev/null @@ -1,34 +0,0 @@ -# Description: -# Contains kernels for the batching ops. - -package(default_visibility = ["//tensorflow:__subpackages__"]) - -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - -cc_library( - name = "batch_kernels", - srcs = ["batch_kernels.cc"], - deps = [ - "//tensorflow/contrib/batching:shared_batch_scheduler_hdrs", - "//tensorflow/contrib/batching/util:periodic_function_dynamic", - "//tensorflow/core:framework_headers_lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core/kernels:concat_lib_hdrs", - "//tensorflow/core/kernels:ops_util_hdrs", - "//tensorflow/core/kernels:split_lib_hdrs", - ], - alwayslink = 1, -) - -filegroup( - name = "all_files", - srcs = glob( - ["**/*"], - exclude = [ - "**/METADATA", - "**/OWNERS", - ], - ), -) diff --git a/tensorflow/contrib/batching/ops/batch_ops.cc b/tensorflow/contrib/batching/ops/batch_ops.cc deleted file mode 100644 index 85e0ccba4aa372bdc21fb194263569b8b787bb6c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/batching/ops/batch_ops.cc +++ /dev/null @@ -1,164 +0,0 @@ -/* Copyright 2017 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/common_shape_fns.h" -#include "tensorflow/core/framework/op.h" -#include "tensorflow/core/framework/shape_inference.h" - -namespace tensorflow { - -REGISTER_OP("Batch") - .Input("in_tensors: T") - .Output("batched_tensors: T") - .Output("batch_index: int64") - .Output("id: int64") - .Attr("num_batch_threads: int") - .Attr("max_batch_size: int") - .Attr("batch_timeout_micros: int") - .Attr("allowed_batch_sizes: list(int) = []") - .Attr("grad_timeout_micros: int") - .Attr("container: string = ''") - .Attr("shared_name: string = ''") - .Attr("batching_queue: string = ''") - .Attr("T: list(type)") - .SetShapeFn([](shape_inference::InferenceContext* c) { - std::vector in_shapes; - TF_RETURN_IF_ERROR(c->input("in_tensors", &in_shapes)); - std::vector out_shapes(in_shapes.size()); - for (int i = 0; i < in_shapes.size(); ++i) { - TF_RETURN_IF_ERROR( - c->ReplaceDim(in_shapes[i], 0, c->UnknownDim(), &out_shapes[i])); - } - TF_RETURN_IF_ERROR(c->set_output("batched_tensors", out_shapes)); - TF_RETURN_IF_ERROR(c->set_output("id", {c->Scalar()})); - TF_RETURN_IF_ERROR(c->set_output( - "batch_index", - {c->MakeShape({shape_inference::DimensionOrConstant(c->UnknownDim()), - shape_inference::DimensionOrConstant(3)})})); - return Status::OK(); - }) - .Doc(R"doc( -Batches all input tensors nondeterministically. - -When many instances of this Op are being run concurrently with the same -container/shared_name in the same device, some will output zero-shaped Tensors -and others will output Tensors of size up to max_batch_size. - -All Tensors in in_tensors are batched together (so, for example, labels and -features should be batched with a single instance of this operation. - -Each invocation of batch emits an `id` scalar which will be used to identify -this particular invocation when doing unbatch or its gradient. - -Each op which emits a non-empty batch will also emit a non-empty batch_index -Tensor, which, is a [K, 3] matrix where each row contains the invocation's id, -start, and length of elements of each set of Tensors present in batched_tensors. - -Batched tensors are concatenated along the first dimension, and all tensors in -in_tensors must have the first dimension of the same size. - -in_tensors: The tensors to be batched. -num_batch_threads: Number of scheduling threads for processing batches of work. - Determines the number of batches processed in parallel. -max_batch_size: Batch sizes will never be bigger than this. -batch_timeout_micros: Maximum number of microseconds to wait before outputting - an incomplete batch. -allowed_batch_sizes: Optional list of allowed batch sizes. If left empty, does - nothing. Otherwise, supplies a list of batch sizes, causing the op to pad - batches up to one of those sizes. The entries must increase monotonically, and - the final entry must equal max_batch_size. -grad_timeout_micros: The timeout to use for the gradient. See Unbatch. -batched_tensors: Either empty tensors or a batch of concatenated Tensors. -batch_index: If out_tensors is non-empty, has information to invert it. -container: Controls the scope of sharing of this batch. -id: always contains a scalar with a unique ID for this invocation of Batch. -shared_name: Concurrently running instances of batch in the same device with the - same container and shared_name will batch their elements together. If left - empty, the op name will be used as the shared name. -T: the types of tensors to be batched. -)doc"); - -REGISTER_OP("Unbatch") - .Input("batched_tensor: T") - .Input("batch_index: int64") - .Input("id: int64") - .Output("unbatched_tensor: T") - .Attr("timeout_micros: int") - .Attr("container: string = ''") - .Attr("shared_name: string = ''") - .Attr("T: type") - .SetShapeFn([](shape_inference::InferenceContext* c) { - shape_inference::ShapeHandle out_shape; - TF_RETURN_IF_ERROR( - c->ReplaceDim(c->input(0), 0, c->UnknownDim(), &out_shape)); - c->set_output(0, out_shape); - return Status::OK(); - }) - .Doc(R"doc( -Reverses the operation of Batch for a single output Tensor. - -An instance of Unbatch either receives an empty batched_tensor, in which case it -asynchronously waits until the values become available from a concurrently -running instance of Unbatch with the same container and shared_name, or receives -a non-empty batched_tensor in which case it finalizes all other concurrently -running instances and outputs its own element from the batch. - -batched_tensor: The possibly transformed output of Batch. The size of the first - dimension should remain unchanged by the transformations for the operation to - work. -batch_index: The matching batch_index obtained from Batch. -id: The id scalar emitted by Batch. -unbatched_tensor: The Tensor corresponding to this execution. -timeout_micros: Maximum amount of time (in microseconds) to wait to receive the - batched input tensor associated with a given invocation of the op. -container: Container to control resource sharing. -shared_name: Instances of Unbatch with the same container and shared_name are - assumed to possibly belong to the same batch. If left empty, the op name will - be used as the shared name. -)doc"); - -REGISTER_OP("UnbatchGrad") - .Input("original_input: T") - .Input("batch_index: int64") - .Input("grad: T") - .Input("id: int64") - .Output("batched_grad: T") - .Attr("container: string = ''") - .Attr("shared_name: string = ''") - .Attr("T: type") - .SetShapeFn([](shape_inference::InferenceContext* c) { - c->set_output(0, c->UnknownShapeOfRank(c->Rank(c->input(2)))); - return Status::OK(); - }) - .Doc(R"doc( -Gradient of Unbatch. - -Acts like Batch but using the given batch_index index of batching things as they -become available. This ensures that the gradients are propagated back in the -same session which did the forward pass. - -original_input: The input to the Unbatch operation this is the gradient of. -batch_index: The batch_index given to the Unbatch operation this is the gradient -of. -grad: The downstream gradient. -id: The id scalar emitted by Batch. -batched_grad: The return value, either an empty tensor or the batched gradient. -container: Container to control resource sharing. -shared_name: Instances of UnbatchGrad with the same container and shared_name - are assumed to possibly belong to the same batch. If left empty, the op name - will be used as the shared name. - )doc"); - -} // namespace tensorflow diff --git a/tensorflow/contrib/batching/python/ops/batch_ops.py b/tensorflow/contrib/batching/python/ops/batch_ops.py index cee4d7b4a9710e285957f27ace7c2762c473c5c7..921d6917a4e478c3e60771fdc3ae99febc33d2e3 100644 --- a/tensorflow/contrib/batching/python/ops/batch_ops.py +++ b/tensorflow/contrib/batching/python/ops/batch_ops.py @@ -18,18 +18,12 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.batching.ops import gen_batch_ops +from tensorflow.python.framework import ops +from tensorflow.python.ops import gen_batch_ops # go/tf-wildcard-import # pylint: disable=wildcard-import -from tensorflow.contrib.batching.ops.gen_batch_ops import * +from tensorflow.python.ops.gen_batch_ops import * # pylint: enable=wildcard-import -from tensorflow.contrib.util import loader -from tensorflow.python.framework import ops -from tensorflow.python.platform import resource_loader - - -_batch_ops = loader.load_op_library( - resource_loader.get_path_to_datafile("_batch_ops.so")) @ops.RegisterGradient("Batch") @@ -59,10 +53,13 @@ def _UnbatchGrad(op, grad): # pylint: disable=invalid-name ] -def batch_function(num_batch_threads, max_batch_size, batch_timeout_micros, +def batch_function(num_batch_threads, + max_batch_size, + batch_timeout_micros, allowed_batch_sizes=None, grad_timeout_micros=60 * 1000 * 1000, - unbatch_timeout_micros=60 * 1000 * 1000): + unbatch_timeout_micros=60 * 1000 * 1000, + max_enqueued_batches=10): """Batches the computation done by the decorated function. So, for example, in the following code @@ -100,6 +97,7 @@ def batch_function(num_batch_threads, max_batch_size, batch_timeout_micros, documentation of the unbatch op for more details. Defaults to 60s. unbatch_timeout_micros: The timeout to use for unbatching. See the documentation of the unbatch op for more details. Defaults to 60s. + max_enqueued_batches: The maximum depth of the batch queue. Defaults to 10. Returns: The decorated function will return the unbatched computation output Tensors. @@ -117,6 +115,7 @@ def batch_function(num_batch_threads, max_batch_size, batch_timeout_micros, num_batch_threads=num_batch_threads, max_batch_size=max_batch_size, batch_timeout_micros=batch_timeout_micros, + max_enqueued_batches=max_enqueued_batches, allowed_batch_sizes=allowed_batch_sizes, grad_timeout_micros=grad_timeout_micros, shared_name=name) diff --git a/tensorflow/contrib/batching/shared_batch_scheduler.h b/tensorflow/contrib/batching/shared_batch_scheduler.h index 86c45bdc2e66e30fbde15f6cafe481cf969c14d0..83a59695d7db7e0a24fb437a3ea71a4d9e23c93f 100644 --- a/tensorflow/contrib/batching/shared_batch_scheduler.h +++ b/tensorflow/contrib/batching/shared_batch_scheduler.h @@ -13,693 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_SHARED_BATCH_SCHEDULER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_SHARED_BATCH_SCHEDULER_H_ +#ifndef TENSORFLOW_CONTRIB_BATCHING_SHARED_BATCH_SCHEDULER_H_ +#define TENSORFLOW_CONTRIB_BATCHING_SHARED_BATCH_SCHEDULER_H_ -#include -#include -#include -#include -#include -#include -#include -#include +#include "tensorflow/core/kernels/batching_util/shared_batch_scheduler.h" -#include "tensorflow/contrib/batching/batch_scheduler.h" -#include "tensorflow/contrib/batching/util/periodic_function.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/strings/strcat.h" -#include "tensorflow/core/platform/cpu_info.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/thread_annotations.h" -#include "tensorflow/core/platform/types.h" - -namespace tensorflow { -namespace serving { -namespace internal { -template -class Queue; -} // namespace internal -} // namespace serving -} // namespace tensorflow - -namespace tensorflow { -namespace serving { - -// A batch scheduler for server instances that service multiple request types -// (e.g. multiple machine-learned models, or multiple versions of a model served -// concurrently), or even multiple distinct tasks for a given request. The -// scheduler multiplexes batches of different kinds of tasks onto a fixed-size -// thread pool (each batch contains tasks of a single type), in a carefully -// controlled manner. A common configuration is to set the number of threads -// equal to the number of hardware accelerator units, in which case the -// scheduler takes care of multiplexing the task types onto the shared hardware, -// in a manner that is both fair and efficient. -// -// Semantically, SharedBatchScheduler behaves like having N instances of -// BasicBatchScheduler (see basic_batch_scheduler.h), one per task type. The -// difference is that under the covers there is a single shared thread pool, -// instead of N independent ones, with their sharing deliberately coordinated. -// -// SharedBatchScheduler does not implement the BatchScheduler API; rather, it -// presents an abstraction of "queues", where each queue corresponds to one type -// of task. Tasks submitted to a given queue are placed in their own batches, -// and cannot be mixed with other tasks. Queues can be added and deleted -// dynamically, to accommodate e.g. versions of a model being brought up and -// down over the lifetime of a server. -// -// The batch thread pool round-robins through the queues, running one batch -// from a queue and then moving to the next queue. Each queue behaves like a -// BasicBatchScheduler instance, in the sense that it has maximum batch size and -// timeout parameters, which govern when a batch is eligible to be processed. -// -// Each queue is independently configured with a maximum size (in terms of the -// maximum number of batches worth of enqueued tasks). For online serving, it is -// recommended that the queue sizes be configured such that the sum of the sizes -// of the active queues roughly equal the number of batch threads. (The idea is -// that if all threads become available at roughly the same time, there will be -// enough enqueued work for them to take on, but no more.) -// -// If queue sizes are configured in the manner suggested above, the maximum time -// a task can spend in a queue before being placed in a batch and assigned to a -// thread for processing, is the greater of: -// - the maximum time to process one batch of tasks from any active queue -// - the configured timeout parameter for the task's queue (which can be 0) -// -// For bulk processing jobs and throughput-oriented benchmarks, you may want to -// set the maximum queue size to a large value. -// -// TODO(b/26539183): Support queue servicing policies other than round-robin. -// E.g. let each queue specify a "share" (an int >= 1), so e.g. with queues A -// and B having shares 1 and 2 respectively, the servicing pattern is ABBABB... -// -// -// PERFORMANCE TUNING: See README.md. -// -template -class SharedBatchScheduler - : public std::enable_shared_from_this> { - public: - // TODO(b/25089730): Tune defaults based on best practices as they develop. - struct Options { - // The name to use for the pool of batch threads. - string thread_pool_name = {"batch_threads"}; - - // The number of threads to use to process batches. - // Must be >= 1, and should be tuned carefully. - int num_batch_threads = port::NumSchedulableCPUs(); - - // The environment to use. - // (Typically only overridden by test code.) - Env* env = Env::Default(); - }; - // Ownership is shared between the caller of Create() and any queues created - // via AddQueue(). - static Status Create( - const Options& options, - std::shared_ptr>* scheduler); - - ~SharedBatchScheduler(); - - // Adds a queue to which tasks may be submitted. The returned queue implements - // the BatchScheduler API. Each queue has its own set of scheduling options, - // and its own callback to process batches of tasks submitted to the queue. - // - // The returned queue's destructor blocks until all tasks submitted to it have - // been processed. - struct QueueOptions { - // The maximum size of each batch. - // - // The scheduler may form batches of any size between 1 and this number - // (inclusive). If there is a need to quantize the batch sizes, i.e. only - // submit batches whose size is in a small set of allowed sizes, that can be - // done by adding padding in the process-batch callback. - int max_batch_size = 1000; - - // If a task has been enqueued for this amount of time (in microseconds), - // and a thread is available, the scheduler will immediately form a batch - // from enqueued tasks and assign the batch to the thread for processing, - // even if the batch's size is below 'max_batch_size'. - // - // This parameter offers a way to bound queue latency, so that a task isn't - // stuck in the queue indefinitely waiting for enough tasks to arrive to - // make a full batch. (The latency bound is given in the class documentation - // above.) - // - // The goal is to smooth out batch sizes under low request rates, and thus - // avoid latency spikes. - int64 batch_timeout_micros = 0; - - // The maximum allowable number of enqueued (accepted by Schedule() but - // not yet being processed on a batch thread) tasks in terms of batches. - // If this limit is reached, Schedule() will return an UNAVAILABLE error. - // See the class documentation above for guidelines on how to tune this - // parameter. - int max_enqueued_batches = 10; - }; - Status AddQueue(const QueueOptions& options, - std::function>)> - process_batch_callback, - std::unique_ptr>* queue); - - private: - explicit SharedBatchScheduler(const Options& options); - - // The code executed in 'batch_threads_'. Obtains a batch to process from the - // queue pointed to by 'next_queue_to_schedule_', and processes it. If that - // queue declines to provide a batch to process, moves onto the next queue. If - // no queues provide a batch to process, just sleeps briefly and exits. - void ThreadLogic(); - - const Options options_; - - mutex mu_; - - // A list of queues. (We use std::list instead of std::vector to ensure that - // iterators are not invalidated by adding/removing elements. It also offers - // efficient removal of elements from the middle.) - using QueueList = std::list>>; - - // All "active" queues, i.e. ones that either: - // - have not been removed, or - // - have been removed but are not yet empty. - QueueList queues_ GUARDED_BY(mu_); - - // An iterator over 'queues_', pointing to the queue from which the next - // available batch thread should grab work. - typename QueueList::iterator next_queue_to_schedule_ GUARDED_BY(mu_); - - // Used by idle batch threads to wait for work to enter the system. Notified - // whenever a batch becomes schedulable. - condition_variable schedulable_batch_cv_; - - // Threads that process batches obtained from the queues. - std::vector> batch_threads_; - - TF_DISALLOW_COPY_AND_ASSIGN(SharedBatchScheduler); -}; - -////////// -// Implementation details follow. API users need not read. - -namespace internal { - -// A task queue for SharedBatchScheduler. Accepts tasks and accumulates them -// into batches, and dispenses those batches to be processed via a "pull" -// interface. The queue's behavior is governed by maximum batch size, timeout -// and maximum queue length parameters; see their documentation in -// SharedBatchScheduler. -// -// The queue is implemented as a deque of batches, with these invariants: -// - The number of batches is between 1 and 'options_.max_enqueued_batches'. -// - The back-most batch is open; the rest are closed. -// -// Submitted tasks are added to the open batch. If that batch doesn't have room -// but the queue isn't full, then that batch is closed and a new open batch is -// started. -// -// Batch pull requests are handled by dequeuing the front-most batch if it is -// closed. If the front-most batch is open (i.e. the queue contains only one -// batch) and has reached the timeout, it is immediately closed and returned; -// otherwise no batch is returned for the request. -template -class Queue { - public: - using ProcessBatchCallback = - std::function>)>; - using SchedulableBatchCallback = std::function; - Queue(const typename SharedBatchScheduler::QueueOptions& options, - Env* env, ProcessBatchCallback process_batch_callback, - SchedulableBatchCallback schdulable_batch_callback); - - // Illegal to destruct unless the queue is empty. - ~Queue(); - - // Submits a task to the queue, with the same semantics as - // BatchScheduler::Schedule(). - Status Schedule(std::unique_ptr* task); - - // Returns the number of enqueued tasks, with the same semantics as - // BatchScheduler::NumEnqueuedTasks(). - size_t NumEnqueuedTasks() const; - - // Returns the queue capacity, with the same semantics as - // BatchScheduler::SchedulingCapacity(). - size_t SchedulingCapacity() const; - - // Returns the maximum allowed size of tasks submitted to the queue. - size_t max_task_size() const { return options_.max_batch_size; } - - // Called by a thread that is ready to process a batch, to request one from - // this queue. Either returns a batch that is ready to be processed, or - // nullptr if the queue declines to schedule a batch at this time. If it - // returns a batch, the batch is guaranteed to be closed. - std::unique_ptr> ScheduleBatch(); - - // Processes a batch that has been returned earlier by ScheduleBatch(). - void ProcessBatch(std::unique_ptr> batch); - - // Determines whether the queue is empty, i.e. has no tasks waiting or being - // processed. - bool IsEmpty() const; - - // Marks the queue closed, and waits until it is empty. - void CloseAndWaitUntilEmpty(); - - bool closed() const { - mutex_lock l(mu_); - return closed_; - } - - private: - // Same as IsEmpty(), but assumes the caller already holds a lock on 'mu_'. - bool IsEmptyInternal() const EXCLUSIVE_LOCKS_REQUIRED(mu_); - - // Closes the open batch residing at the back of 'batches_', and inserts a - // fresh open batch behind it. - void StartNewBatch() EXCLUSIVE_LOCKS_REQUIRED(mu_); - - // Determines whether the open batch residing at the back of 'batches_' is - // currently schedulable. - bool IsOpenBatchSchedulable() const EXCLUSIVE_LOCKS_REQUIRED(mu_); - - const typename SharedBatchScheduler::QueueOptions options_; - - // The environment to use. - Env* env_; - - // A callback invoked to processes a batch of work units. Always invoked from - // a batch thread. - ProcessBatchCallback process_batch_callback_; - - // A callback invoked to notify the scheduler that a new batch has become - // schedulable. - SchedulableBatchCallback schedulable_batch_callback_; - - mutable mutex mu_; - - // Whether this queue can accept new tasks. This variable is monotonic: it - // starts as false, and then at some point gets set to true and remains true - // for the duration of this object's life. - bool closed_ GUARDED_BY(mu_) = false; - - // The enqueued batches. See the invariants in the class comments above. - std::deque>> batches_ GUARDED_BY(mu_); - - // The time at which the first task was added to the open (back-most) batch - // in 'batches_'. Valid iff that batch contains at least one task. - uint64 open_batch_start_time_micros_ GUARDED_BY(mu_); - - // Whether this queue contains a batch that is eligible to be scheduled. Used - // to keep track of when to call 'schedulable_batch_callback_'. - bool schedulable_batch_ GUARDED_BY(mu_) = false; - - // The number of batches currently being processed by batch threads. - // Incremented in ScheduleBatch() and decremented in ProcessBatch(). - int num_batches_being_processed_ GUARDED_BY(mu_) = 0; - - // Used by CloseAndWaitUntilEmpty() to wait until the queue is empty, for the - // case in which the queue is not empty when CloseAndWaitUntilEmpty() starts. - // When ProcessBatch() dequeues the last batch and makes the queue empty, if - // 'empty_notification_' is non-null it calls 'empty_notification_->Notify()'. - Notification* empty_notification_ GUARDED_BY(mu_) = nullptr; - - TF_DISALLOW_COPY_AND_ASSIGN(Queue); -}; - -// A RAII-style object that points to a Queue and implements -// the BatchScheduler API. To be handed out to clients who call AddQueue(). -template -class QueueHandle : public BatchScheduler { - public: - QueueHandle(std::shared_ptr> scheduler, - Queue* queue); - ~QueueHandle() override; - - Status Schedule(std::unique_ptr* task) override; - size_t NumEnqueuedTasks() const override; - size_t SchedulingCapacity() const override; - - size_t max_task_size() const override { return queue_->max_task_size(); } - - private: - // The scheduler that owns 'queue_'. - std::shared_ptr> scheduler_; - - // The queue this handle wraps. Owned by 'scheduler_', which keeps it alive at - // least until this class's destructor closes it. - Queue* queue_; - - TF_DISALLOW_COPY_AND_ASSIGN(QueueHandle); -}; - -} // namespace internal - -template -Status SharedBatchScheduler::Create( - const Options& options, - std::shared_ptr>* scheduler) { - if (options.num_batch_threads < 1) { - return errors::InvalidArgument("num_batch_threads must be positive; was ", - options.num_batch_threads); - } - scheduler->reset(new SharedBatchScheduler(options)); - return Status::OK(); -} - -template -SharedBatchScheduler::~SharedBatchScheduler() { - // Wait until the batch threads finish clearing out and deleting the closed - // queues. - for (;;) { - { - mutex_lock l(mu_); - if (queues_.empty()) { - break; - } - } - const int64 kSleepTimeMicros = 100; - options_.env->SleepForMicroseconds(kSleepTimeMicros); - } - // Delete the batch threads before allowing state the threads may access (e.g. - // 'mu_') to be deleted. - batch_threads_.clear(); -} - -template -Status SharedBatchScheduler::AddQueue( - const QueueOptions& options, - std::function>)> - process_batch_callback, - std::unique_ptr>* queue) { - if (options.max_batch_size <= 0) { - return errors::InvalidArgument("max_batch_size must be positive; was ", - options.max_batch_size); - } - if (options.batch_timeout_micros < 0) { - return errors::InvalidArgument( - "batch_timeout_micros must be non-negative; was ", - options.batch_timeout_micros); - } - if (options.max_enqueued_batches < 0) { - return errors::InvalidArgument( - "max_enqueued_batches must be non-negative; was ", - options.max_enqueued_batches); - } - - auto schedulable_batch_callback = [this] { - mutex_lock l(mu_); - schedulable_batch_cv_.notify_one(); - }; - auto internal_queue = - std::unique_ptr>(new internal::Queue( - options, options_.env, process_batch_callback, - schedulable_batch_callback)); - auto handle = std::unique_ptr>( - new internal::QueueHandle(this->shared_from_this(), - internal_queue.get())); - { - mutex_lock l(mu_); - queues_.push_back(std::move(internal_queue)); - if (next_queue_to_schedule_ == queues_.end()) { - next_queue_to_schedule_ = queues_.begin(); - } - } - *queue = std::move(handle); - return Status::OK(); -} - -template -SharedBatchScheduler::SharedBatchScheduler(const Options& options) - : options_(options), next_queue_to_schedule_(queues_.end()) { - // Kick off the batch threads. - PeriodicFunction::Options periodic_fn_options; - periodic_fn_options.thread_name_prefix = - strings::StrCat(options.thread_pool_name, "_"); - for (int i = 0; i < options.num_batch_threads; ++i) { - std::unique_ptr thread(new PeriodicFunction( - [this] { this->ThreadLogic(); }, - 0 /* function invocation interval time */, periodic_fn_options)); - batch_threads_.push_back(std::move(thread)); - } -} - -template -void SharedBatchScheduler::ThreadLogic() { - // A batch to process next (or nullptr if no work to do). - std::unique_ptr> batch_to_process; - // The queue with which 'batch_to_process' is associated. - internal::Queue* queue_for_batch = nullptr; - { - mutex_lock l(mu_); - - const int num_queues = queues_.size(); - for (int num_queues_tried = 0; - batch_to_process == nullptr && num_queues_tried < num_queues; - ++num_queues_tried) { - DCHECK(next_queue_to_schedule_ != queues_.end()); - - // If a closed queue responds to ScheduleBatch() with nullptr, the queue - // will never yield any further batches so we can drop it. To avoid a - // race, we take a snapshot of the queue's closedness state *before* - // calling ScheduleBatch(). - const bool queue_closed = (*next_queue_to_schedule_)->closed(); - - // Ask '*next_queue_to_schedule_' if it wants us to process a batch. - batch_to_process = (*next_queue_to_schedule_)->ScheduleBatch(); - if (batch_to_process != nullptr) { - queue_for_batch = next_queue_to_schedule_->get(); - } - - // Advance 'next_queue_to_schedule_'. - if (queue_closed && (*next_queue_to_schedule_)->IsEmpty() && - batch_to_process == nullptr) { - // We've encountered a closed queue with no work to do. Drop it. - DCHECK_NE(queue_for_batch, next_queue_to_schedule_->get()); - next_queue_to_schedule_ = queues_.erase(next_queue_to_schedule_); - } else { - ++next_queue_to_schedule_; - } - if (next_queue_to_schedule_ == queues_.end() && !queues_.empty()) { - // We've hit the end. Wrap to the first queue. - next_queue_to_schedule_ = queues_.begin(); - } - } - - if (batch_to_process == nullptr) { - // We couldn't find any work to do. Wait until a new batch becomes - // schedulable, or some time has elapsed, before checking again. - const int64 kTimeoutMillis = 1; // The smallest accepted granule of time. - WaitForMilliseconds(&l, &schedulable_batch_cv_, kTimeoutMillis); - return; - } - } - - queue_for_batch->ProcessBatch(std::move(batch_to_process)); -} - -namespace internal { - -template -Queue::Queue( - const typename SharedBatchScheduler::QueueOptions& options, - Env* env, ProcessBatchCallback process_batch_callback, - SchedulableBatchCallback schedulable_batch_callback) - : options_(options), - env_(env), - process_batch_callback_(process_batch_callback), - schedulable_batch_callback_(schedulable_batch_callback) { - // Create an initial, open batch. - batches_.emplace_back(new Batch); -} - -template -Queue::~Queue() { - mutex_lock l(mu_); - DCHECK(IsEmptyInternal()); - - // Close the (empty) open batch, so its destructor doesn't block. - batches_.back()->Close(); -} - -template -Status Queue::Schedule(std::unique_ptr* task) { - if ((*task)->size() > options_.max_batch_size) { - return errors::InvalidArgument("Task size ", (*task)->size(), - " is larger than maximum batch size ", - options_.max_batch_size); - } - - bool notify_of_schedulable_batch = false; - { - mutex_lock l(mu_); - - DCHECK(!closed_); - - if (batches_.back()->size() + (*task)->size() > options_.max_batch_size) { - if (batches_.size() >= options_.max_enqueued_batches) { - return errors::Unavailable( - "The batch scheduling queue to which this task was submitted is " - "full"); - } - StartNewBatch(); - } - if (batches_.back()->empty()) { - open_batch_start_time_micros_ = env_->NowMicros(); - } - batches_.back()->AddTask(std::move(*task)); - - if (!schedulable_batch_) { - if (batches_.size() > 1 || IsOpenBatchSchedulable()) { - schedulable_batch_ = true; - notify_of_schedulable_batch = true; - } - } - } - - if (notify_of_schedulable_batch) { - schedulable_batch_callback_(); - } - - return Status::OK(); -} - -template -size_t Queue::NumEnqueuedTasks() const { - mutex_lock l(mu_); - size_t num_enqueued_tasks = 0; - for (const auto& batch : batches_) { - num_enqueued_tasks += batch->num_tasks(); - } - return num_enqueued_tasks; -} - -template -size_t Queue::SchedulingCapacity() const { - mutex_lock l(mu_); - const int num_new_batches_schedulable = - options_.max_enqueued_batches - batches_.size(); - const int open_batch_capacity = - options_.max_batch_size - batches_.back()->size(); - return (num_new_batches_schedulable * options_.max_batch_size) + - open_batch_capacity; -} - -template -std::unique_ptr> Queue::ScheduleBatch() { - // The batch to schedule, which we may populate below. (If left as nullptr, - // that means we are electing not to schedule a batch at this time.) - std::unique_ptr> batch_to_schedule; - - { - mutex_lock l(mu_); - - // Consider closing the open batch at this time, to schedule it. - if (batches_.size() == 1 && IsOpenBatchSchedulable()) { - StartNewBatch(); - } - - if (batches_.size() >= 2) { - // There is at least one closed batch that is ready to be scheduled. - ++num_batches_being_processed_; - batch_to_schedule = std::move(batches_.front()); - batches_.pop_front(); - } else { - schedulable_batch_ = false; - } - } - - return batch_to_schedule; -} - -template -void Queue::ProcessBatch(std::unique_ptr> batch) { - process_batch_callback_(std::move(batch)); - - { - mutex_lock l(mu_); - --num_batches_being_processed_; - if (empty_notification_ != nullptr && IsEmptyInternal()) { - empty_notification_->Notify(); - } - } -} - -template -bool Queue::IsEmpty() const { - mutex_lock l(mu_); - return IsEmptyInternal(); -} - -template -void Queue::CloseAndWaitUntilEmpty() { - Notification empty; - { - mutex_lock l(mu_); - closed_ = true; - if (IsEmptyInternal()) { - empty.Notify(); - } else { - // Arrange for ProcessBatch() to notify when the queue becomes empty. - empty_notification_ = ∅ - } - } - empty.WaitForNotification(); -} - -template -bool Queue::IsEmptyInternal() const { - return num_batches_being_processed_ == 0 && batches_.size() == 1 && - batches_.back()->empty(); -} - -template -void Queue::StartNewBatch() { - batches_.back()->Close(); - batches_.emplace_back(new Batch); -} - -template -bool Queue::IsOpenBatchSchedulable() const { - Batch* open_batch = batches_.back().get(); - if (open_batch->empty()) { - return false; - } - return closed_ || open_batch->size() >= options_.max_batch_size || - env_->NowMicros() >= - open_batch_start_time_micros_ + options_.batch_timeout_micros; -} - -template -QueueHandle::QueueHandle( - std::shared_ptr> scheduler, - Queue* queue) - : scheduler_(scheduler), queue_(queue) {} - -template -QueueHandle::~QueueHandle() { - queue_->CloseAndWaitUntilEmpty(); -} - -template -Status QueueHandle::Schedule(std::unique_ptr* task) { - return queue_->Schedule(task); -} - -template -size_t QueueHandle::NumEnqueuedTasks() const { - return queue_->NumEnqueuedTasks(); -} - -template -size_t QueueHandle::SchedulingCapacity() const { - return queue_->SchedulingCapacity(); -} - -} // namespace internal - -} // namespace serving -} // namespace tensorflow - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_SHARED_BATCH_SCHEDULER_H_ +#endif // TENSORFLOW_CONTRIB_BATCHING_SHARED_BATCH_SCHEDULER_H_ diff --git a/tensorflow/contrib/batching/test_util/BUILD b/tensorflow/contrib/batching/test_util/BUILD index d1ced0d8c367f44b520a9bba2db8a3e0969bab4c..6db627faad1df4a4b73082e74e7754829ff2b514 100644 --- a/tensorflow/contrib/batching/test_util/BUILD +++ b/tensorflow/contrib/batching/test_util/BUILD @@ -22,11 +22,9 @@ filegroup( cc_library( name = "fake_clock_env", testonly = 1, - srcs = ["fake_clock_env.cc"], hdrs = ["fake_clock_env.h"], visibility = ["//visibility:public"], deps = [ - "//tensorflow/core:lib", - "//tensorflow/core:tensorflow", + "//tensorflow/core/kernels/batching_util:fake_clock_env", ], ) diff --git a/tensorflow/contrib/batching/test_util/fake_clock_env.h b/tensorflow/contrib/batching/test_util/fake_clock_env.h index 35cafcb73c51feb4e9e15a61d1830c8ef6bc3e0f..40a39a5569854350c72a47102f3dac07b362ce8e 100644 --- a/tensorflow/contrib/batching/test_util/fake_clock_env.h +++ b/tensorflow/contrib/batching/test_util/fake_clock_env.h @@ -13,64 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_TEST_UTIL_FAKE_CLOCK_ENV_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_TEST_UTIL_FAKE_CLOCK_ENV_H_ +#ifndef TENSORFLOW_CONTRIB_BATCHING_TEST_UTIL_FAKE_CLOCK_ENV_H_ +#define TENSORFLOW_CONTRIB_BATCHING_TEST_UTIL_FAKE_CLOCK_ENV_H_ -#include -#include -#include +#include "tensorflow/core/kernels/batching_util/fake_clock_env.h" -#include "tensorflow/core/lib/core/notification.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/macros.h" -#include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/thread_annotations.h" -#include "tensorflow/core/platform/types.h" - -namespace tensorflow { -namespace serving { -namespace test_util { - -// An Env implementation with a fake clock for NowMicros() and -// SleepForMicroseconds(). The clock doesn't advance on its own; it advances via -// an explicit Advance() method. -// All other Env virtual methods pass through to a wrapped Env. -class FakeClockEnv : public EnvWrapper { - public: - explicit FakeClockEnv(Env* wrapped); - ~FakeClockEnv() override = default; - - // Advance the clock by a certain number of microseconds. - void AdvanceByMicroseconds(int micros); - - // Blocks until there is a sleeping thread that is scheduled to wake up at - // the given (absolute) time. - void BlockUntilSleepingThread(uint64 wake_time); - - // Blocks until there are at least num_threads sleeping. - void BlockUntilThreadsAsleep(int num_threads); - - // Methods that this class implements. - uint64 NowMicros() override; - void SleepForMicroseconds(int64 micros) override; - - private: - mutex mu_; - - uint64 current_time_ GUARDED_BY(mu_) = 0; - - struct SleepingThread { - uint64 wake_time; - Notification* wake_notification; - }; - std::vector sleeping_threads_ GUARDED_BY(mu_); - - TF_DISALLOW_COPY_AND_ASSIGN(FakeClockEnv); -}; - -} // namespace test_util -} // namespace serving -} // namespace tensorflow - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_TEST_UTIL_FAKE_CLOCK_ENV_H_ +#endif // TENSORFLOW_CONTRIB_BATCHING_TEST_UTIL_FAKE_CLOCK_ENV_H_ diff --git a/tensorflow/contrib/batching/util/BUILD b/tensorflow/contrib/batching/util/BUILD index f33a08cb817e9f2832be953ef6ff1aba04c4c288..2a84a7712a8fa66e89db41ff4e7ebe4f620029ca 100644 --- a/tensorflow/contrib/batching/util/BUILD +++ b/tensorflow/contrib/batching/util/BUILD @@ -22,12 +22,11 @@ filegroup( cc_library( name = "periodic_function_dynamic", - srcs = ["periodic_function.cc"], hdrs = ["periodic_function.h"], visibility = ["//visibility:public"], deps = [ - "//tensorflow/core:framework_headers_lib", - "//tensorflow/core:protos_all_cc", + "//tensorflow/core/kernels/batching_util:periodic_function_dynamic", + "//third_party/eigen3", ], ) @@ -36,17 +35,6 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":periodic_function_dynamic", - "//tensorflow/core:lib", - ], -) - -tf_cc_test( - name = "periodic_function_test", - srcs = ["periodic_function_test.cc"], - deps = [ - ":periodic_function_dynamic", - "//tensorflow/contrib/batching/test_util:fake_clock_env", - "//tensorflow/core:test", - "//tensorflow/core:test_main", + "//tensorflow/core/kernels/batching_util:periodic_function", ], ) diff --git a/tensorflow/contrib/batching/util/periodic_function.h b/tensorflow/contrib/batching/util/periodic_function.h index 2c032d802fe5f23a267db28dc869a253f16afc34..aa2ed0a385125fa090a7a56b6339a87eb2d57b1f 100644 --- a/tensorflow/contrib/batching/util/periodic_function.h +++ b/tensorflow/contrib/batching/util/periodic_function.h @@ -12,121 +12,9 @@ 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_BATCHING_UTIL_PERIODIC_FUNCTION_H_ +#define TENSORFLOW_CONTRIB_BATCHING_UTIL_PERIODIC_FUNCTION_H_ -// PeriodicFunction will periodically call the given function with a specified -// period in a background thread. After Start() returns, the thread is -// guaranteed to have started. The destruction of the class causes the -// background thread to be destroyed as well. Start() should not be called more -// than once. -// -// PeriodicFunction runs the function as soon as any previous run both is -// complete and was started more than "interval_micros" earlier. Thus, runs are -// both serialized, and normally have a period of "interval_micros" if no run -// exceeds the time. -// -// Note that, if the function takes longer than two interval_micross to finish, -// then PeriodicFunction will "skip" at least one call to the function. For -// instance, if the period is 50ms and the function starts runs at time 0 for -// 150ms, then the function will immediately start executing again at time 150, -// but there will be no function runs corresponding to times 50 or 100. This is -// especially important to remember when using an environment with a simulated -// clock: advancing simulated time atomically over N interval_micross will not -// cause the function to be called N times. -// -// This object is thread-safe. -// -// Example: -// -// class Foo { -// public: -// Foo() : periodic_function_([this]() { Bar(); }, -// 1000 /* 1000us == 1ms*/) { -// } -// -// private: -// void Bar() { ... } -// -// PeriodicFunction periodic_function_; -// }; +#include "tensorflow/core/kernels/batching_util/periodic_function.h" -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_UTIL_PERIODIC_FUNCTION_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_UTIL_PERIODIC_FUNCTION_H_ - -#include -#include -#include - -#include "tensorflow/core/lib/core/notification.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/macros.h" -#include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/thread_annotations.h" -#include "tensorflow/core/platform/types.h" - -namespace tensorflow { -namespace serving { - -namespace internal { -class PeriodicFunctionTestAccess; -} - -class PeriodicFunction { - public: - // Provides the ability to customize several aspects of the PeriodicFunction. - // Passed to constructor of PeriodicFunction. - struct Options { - Options() {} - - // Any standard thread options, such as stack size, should - // be passed via "thread_options". - ThreadOptions thread_options; - - // Specifies the thread name prefix (see the description in class - // Thread). - string thread_name_prefix = "periodic_function"; - - // The environment to use. Does not take ownership, but must remain alive - // for as long as the PeriodicFunction exists. - Env* env = Env::Default(); - - // Specifies the length of sleep before the first invocation of the - // function. - // This can be used for adding a random jitter to avoid synchronous behavior - // across multiple periodic functions. - int64 startup_delay_micros = 0; - }; - - // Also starts the background thread which will be calling the function. - PeriodicFunction(const std::function& function, int64 interval_micros, - const Options& options = Options()); - - ~PeriodicFunction(); - - private: - friend class internal::PeriodicFunctionTestAccess; - - // Notifies the background thread to stop. - void NotifyStop(); - - // (Blocking.) Loops forever calling "function_" every "interval_micros_". - void RunLoop(int64 start) LOCKS_EXCLUDED(mutex_); - - const std::function function_; // Actual client function - const int64 interval_micros_; // Interval between calls. - const Options options_; - - // Protects state below. - mutable mutex mutex_; - // Used to notify the thread to stop. - Notification stop_thread_; - - // Thread for running "function_" - std::unique_ptr thread_ GUARDED_BY(mutex_) = nullptr; - - TF_DISALLOW_COPY_AND_ASSIGN(PeriodicFunction); -}; - -} // namespace serving -} // namespace tensorflow - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BATCHING_UTIL_PERIODIC_FUNCTION_H_ +#endif // TENSORFLOW_CONTRIB_BATCHING_UTIL_PERIODIC_FUNCTION_H_ diff --git a/tensorflow/contrib/bayesflow/BUILD b/tensorflow/contrib/bayesflow/BUILD index 11c3c037c4e8b4ba41eae60d28d6aac49f1488f2..74712aeb67c3f0a31def78f25a0298f9c02c9590 100644 --- a/tensorflow/contrib/bayesflow/BUILD +++ b/tensorflow/contrib/bayesflow/BUILD @@ -137,6 +137,26 @@ cuda_py_test( ], ) +cuda_py_test( + name = "mcmc_diagnostics_test", + size = "small", + srcs = ["python/kernel_tests/mcmc_diagnostics_test.py"], + additional_deps = [ + ":bayesflow_py", + "//third_party/py/numpy", + "//tensorflow/python:spectral_ops_test_util", + "//tensorflow/contrib/distributions:distributions_py", + "//tensorflow/python/ops/distributions", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + "//tensorflow/python:random_seed", + ], +) + cuda_py_test( name = "monte_carlo_test", size = "small", @@ -175,6 +195,7 @@ cuda_py_test( "//tensorflow/python:variable_scope", "//tensorflow/python:variables", ], + tags = ["no_mac"], # b/73192243 ) cuda_py_test( @@ -217,6 +238,24 @@ cuda_py_test( "//tensorflow/python:platform_test", "//tensorflow/python:random_seed", ], + tags = ["notsan"], +) + +cuda_py_test( + name = "variable_utils_test", + size = "small", + srcs = ["python/kernel_tests/variable_utils_test.py"], + additional_deps = [ + ":bayesflow_py", + "//third_party/py/numpy", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:gradients", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + ], ) cuda_py_test( diff --git a/tensorflow/contrib/bayesflow/__init__.py b/tensorflow/contrib/bayesflow/__init__.py index 95b9452b1ada60c44672f37800ced2133d2bd8b2..528c4fbacd06c7b0defa0e32bd24a98b2bc07b64 100644 --- a/tensorflow/contrib/bayesflow/__init__.py +++ b/tensorflow/contrib/bayesflow/__init__.py @@ -26,9 +26,11 @@ from tensorflow.contrib.bayesflow.python.ops import custom_grad from tensorflow.contrib.bayesflow.python.ops import halton_sequence from tensorflow.contrib.bayesflow.python.ops import hmc from tensorflow.contrib.bayesflow.python.ops import layers +from tensorflow.contrib.bayesflow.python.ops import mcmc_diagnostics from tensorflow.contrib.bayesflow.python.ops import metropolis_hastings from tensorflow.contrib.bayesflow.python.ops import monte_carlo from tensorflow.contrib.bayesflow.python.ops import optimizers +from tensorflow.contrib.bayesflow.python.ops import variable_utils # pylint: enable=unused-import,line-too-long from tensorflow.python.util.all_util import remove_undocumented @@ -42,10 +44,12 @@ _allowed_symbols = [ 'hmc', 'layers', 'metropolis_hastings', + 'mcmc_diagnostics', 'monte_carlo', 'optimizers', 'special_math', 'stochastic_variables', + 'variable_utils', 'variational_inference', ] diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/hmc_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/hmc_test.py index cbc66b6dc13db62c25952de6b6c13b2fdfe27f12..5bd834e56245ab4d874544cfd014fe59ae521ea8 100644 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/hmc_test.py +++ b/tensorflow/contrib/bayesflow/python/kernel_tests/hmc_test.py @@ -18,30 +18,40 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections + import numpy as np -from scipy import special from scipy import stats from tensorflow.contrib.bayesflow.python.ops import hmc +from tensorflow.contrib.bayesflow.python.ops.hmc_impl import _compute_energy_change +from tensorflow.contrib.bayesflow.python.ops.hmc_impl import _leapfrog_integrator +from tensorflow.contrib.distributions.python.ops import independent as independent_lib +from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed from tensorflow.python.ops import array_ops -from tensorflow.python.ops import gradients_impl +from tensorflow.python.ops import gen_linalg_ops +from tensorflow.python.ops import gradients_impl as gradients_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops +from tensorflow.python.ops.distributions import gamma as gamma_lib +from tensorflow.python.ops.distributions import normal as normal_lib from tensorflow.python.platform import test -from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.platform import tf_logging as logging_ops + + +def _reduce_variance(x, axis=None, keepdims=False): + sample_mean = math_ops.reduce_mean(x, axis, keepdims=True) + return math_ops.reduce_mean( + math_ops.squared_difference(x, sample_mean), axis, keepdims) -# TODO(b/66964210): Test float16. class HMCTest(test.TestCase): def setUp(self): self._shape_param = 5. self._rate_param = 10. - self._expected_x = (special.digamma(self._shape_param) - - np.log(self._rate_param)) - self._expected_exp_x = self._shape_param / self._rate_param random_seed.set_random_seed(10003) np.random.seed(10003) @@ -63,63 +73,46 @@ class HMCTest(test.TestCase): self._rate_param * math_ops.exp(x), event_dims) - def _log_gamma_log_prob_grad(self, x, event_dims=()): - """Computes log-pdf and gradient of a log-gamma random variable. - - Args: - x: Value of the random variable. - event_dims: Dimensions not to treat as independent. Default is (), - i.e., all dimensions are independent. - - Returns: - log_prob: The log-pdf up to a normalizing constant. - grad: The gradient of the log-pdf with respect to x. - """ - return (math_ops.reduce_sum(self._shape_param * x - - self._rate_param * math_ops.exp(x), - event_dims), - self._shape_param - self._rate_param * math_ops.exp(x)) - - def _n_event_dims(self, x_shape, event_dims): - return np.prod([int(x_shape[i]) for i in event_dims]) - - def _integrator_conserves_energy(self, x, event_dims, sess, + def _integrator_conserves_energy(self, x, independent_chain_ndims, sess, feed_dict=None): - def potential_and_grad(x): - log_prob, grad = self._log_gamma_log_prob_grad(x, event_dims) - return -log_prob, -grad - - step_size = array_ops.placeholder(np.float32, [], name='step_size') - hmc_lf_steps = array_ops.placeholder(np.int32, [], name='hmc_lf_steps') + step_size = array_ops.placeholder(np.float32, [], name="step_size") + hmc_lf_steps = array_ops.placeholder(np.int32, [], name="hmc_lf_steps") if feed_dict is None: feed_dict = {} feed_dict[hmc_lf_steps] = 1000 - m = random_ops.random_normal(array_ops.shape(x)) - potential_0, grad_0 = potential_and_grad(x) - old_energy = potential_0 + 0.5 * math_ops.reduce_sum(m * m, - event_dims) - - _, new_m, potential_1, _ = ( - hmc.leapfrog_integrator(step_size, hmc_lf_steps, x, - m, potential_and_grad, grad_0)) + event_dims = math_ops.range(independent_chain_ndims, + array_ops.rank(x)) - new_energy = potential_1 + 0.5 * math_ops.reduce_sum(new_m * new_m, + m = random_ops.random_normal(array_ops.shape(x)) + log_prob_0 = self._log_gamma_log_prob(x, event_dims) + grad_0 = gradients_ops.gradients(log_prob_0, x) + old_energy = -log_prob_0 + 0.5 * math_ops.reduce_sum(m**2., event_dims) + + new_m, _, log_prob_1, _ = _leapfrog_integrator( + current_momentums=[m], + target_log_prob_fn=lambda x: self._log_gamma_log_prob(x, event_dims), + current_state_parts=[x], + step_sizes=[step_size], + num_leapfrog_steps=hmc_lf_steps, + current_target_log_prob=log_prob_0, + current_grads_target_log_prob=grad_0) + new_m = new_m[0] + + new_energy = -log_prob_1 + 0.5 * math_ops.reduce_sum(new_m * new_m, event_dims) x_shape = sess.run(x, feed_dict).shape - n_event_dims = self._n_event_dims(x_shape, event_dims) - feed_dict[step_size] = 0.1 / n_event_dims - old_energy_val, new_energy_val = sess.run([old_energy, new_energy], - feed_dict) - logging.vlog(1, 'average energy change: {}'.format( - abs(old_energy_val - new_energy_val).mean())) - - self.assertAllEqual(np.ones_like(new_energy_val, dtype=np.bool), - abs(old_energy_val - new_energy_val) < 1.) - - def _integrator_conserves_energy_wrapper(self, event_dims): + event_size = np.prod(x_shape[independent_chain_ndims:]) + feed_dict[step_size] = 0.1 / event_size + old_energy_, new_energy_ = sess.run([old_energy, new_energy], + feed_dict) + logging_ops.vlog(1, "average energy relative change: {}".format( + (1. - new_energy_ / old_energy_).mean())) + self.assertAllClose(old_energy_, new_energy_, atol=0., rtol=0.02) + + def _integrator_conserves_energy_wrapper(self, independent_chain_ndims): """Tests the long-term energy conservation of the leapfrog integrator. The leapfrog integrator is symplectic, so for sufficiently small step @@ -127,135 +120,310 @@ class HMCTest(test.TestCase): the energy of the system blowing up or collapsing. Args: - event_dims: A tuple of dimensions that should not be treated as - independent. This allows for multiple chains to be run independently - in parallel. Default is (), i.e., all dimensions are independent. + independent_chain_ndims: Python `int` scalar representing the number of + dims associated with independent chains. """ - with self.test_session() as sess: - x_ph = array_ops.placeholder(np.float32, name='x_ph') - - feed_dict = {x_ph: np.zeros([50, 10, 2])} - self._integrator_conserves_energy(x_ph, event_dims, sess, feed_dict) + with self.test_session(graph=ops.Graph()) as sess: + x_ph = array_ops.placeholder(np.float32, name="x_ph") + feed_dict = {x_ph: np.random.rand(50, 10, 2)} + self._integrator_conserves_energy(x_ph, independent_chain_ndims, + sess, feed_dict) def testIntegratorEnergyConservationNullShape(self): - self._integrator_conserves_energy_wrapper([]) + self._integrator_conserves_energy_wrapper(0) def testIntegratorEnergyConservation1(self): - self._integrator_conserves_energy_wrapper([1]) + self._integrator_conserves_energy_wrapper(1) def testIntegratorEnergyConservation2(self): - self._integrator_conserves_energy_wrapper([2]) - - def testIntegratorEnergyConservation12(self): - self._integrator_conserves_energy_wrapper([1, 2]) - - def testIntegratorEnergyConservation012(self): - self._integrator_conserves_energy_wrapper([0, 1, 2]) - - def _chain_gets_correct_expectations(self, x, event_dims, sess, - feed_dict=None): + self._integrator_conserves_energy_wrapper(2) + + def testIntegratorEnergyConservation3(self): + self._integrator_conserves_energy_wrapper(3) + + def testSampleChainSeedReproducibleWorksCorrectly(self): + with self.test_session(graph=ops.Graph()) as sess: + num_results = 10 + independent_chain_ndims = 1 + + def log_gamma_log_prob(x): + event_dims = math_ops.range(independent_chain_ndims, + array_ops.rank(x)) + return self._log_gamma_log_prob(x, event_dims) + + kwargs = dict( + target_log_prob_fn=log_gamma_log_prob, + current_state=np.random.rand(4, 3, 2), + step_size=0.1, + num_leapfrog_steps=2, + num_burnin_steps=150, + seed=52, + ) + + samples0, kernel_results0 = hmc.sample_chain( + **dict(list(kwargs.items()) + list(dict( + num_results=2 * num_results, + num_steps_between_results=0).items()))) + + samples1, kernel_results1 = hmc.sample_chain( + **dict(list(kwargs.items()) + list(dict( + num_results=num_results, + num_steps_between_results=1).items()))) + + [ + samples0_, + samples1_, + target_log_prob0_, + target_log_prob1_, + ] = sess.run([ + samples0, + samples1, + kernel_results0.current_target_log_prob, + kernel_results1.current_target_log_prob, + ]) + self.assertAllClose(samples0_[::2], samples1_, + atol=1e-5, rtol=1e-5) + self.assertAllClose(target_log_prob0_[::2], target_log_prob1_, + atol=1e-5, rtol=1e-5) + + def _chain_gets_correct_expectations(self, x, independent_chain_ndims, + sess, feed_dict=None): + counter = collections.Counter() def log_gamma_log_prob(x): + counter["target_calls"] += 1 + event_dims = math_ops.range(independent_chain_ndims, + array_ops.rank(x)) return self._log_gamma_log_prob(x, event_dims) - step_size = array_ops.placeholder(np.float32, [], name='step_size') - hmc_lf_steps = array_ops.placeholder(np.int32, [], name='hmc_lf_steps') - hmc_n_steps = array_ops.placeholder(np.int32, [], name='hmc_n_steps') + num_results = array_ops.placeholder( + np.int32, [], name="num_results") + step_size = array_ops.placeholder( + np.float32, [], name="step_size") + num_leapfrog_steps = array_ops.placeholder( + np.int32, [], name="num_leapfrog_steps") if feed_dict is None: feed_dict = {} - feed_dict.update({step_size: 0.1, - hmc_lf_steps: 2, - hmc_n_steps: 300}) - - sample_chain, acceptance_prob_chain = hmc.chain([hmc_n_steps], - step_size, - hmc_lf_steps, - x, log_gamma_log_prob, - event_dims) - - acceptance_probs, samples = sess.run([acceptance_prob_chain, sample_chain], - feed_dict) - samples = samples[feed_dict[hmc_n_steps] // 2:] - expected_x_est = samples.mean() - expected_exp_x_est = np.exp(samples).mean() - - logging.vlog(1, 'True E[x, exp(x)]: {}\t{}'.format( - self._expected_x, self._expected_exp_x)) - logging.vlog(1, 'Estimated E[x, exp(x)]: {}\t{}'.format( - expected_x_est, expected_exp_x_est)) - self.assertNear(expected_x_est, self._expected_x, 2e-2) - self.assertNear(expected_exp_x_est, self._expected_exp_x, 2e-2) - self.assertTrue((acceptance_probs > 0.5).all()) - self.assertTrue((acceptance_probs <= 1.0).all()) - - def _chain_gets_correct_expectations_wrapper(self, event_dims): - with self.test_session() as sess: - x_ph = array_ops.placeholder(np.float32, name='x_ph') - - feed_dict = {x_ph: np.zeros([50, 10, 2])} - self._chain_gets_correct_expectations(x_ph, event_dims, sess, - feed_dict) + feed_dict.update({num_results: 150, + step_size: 0.05, + num_leapfrog_steps: 2}) + + samples, kernel_results = hmc.sample_chain( + num_results=num_results, + target_log_prob_fn=log_gamma_log_prob, + current_state=x, + step_size=step_size, + num_leapfrog_steps=num_leapfrog_steps, + num_burnin_steps=150, + seed=42) + + self.assertAllEqual(dict(target_calls=2), counter) + + expected_x = (math_ops.digamma(self._shape_param) + - np.log(self._rate_param)) + + expected_exp_x = self._shape_param / self._rate_param + + acceptance_probs_, samples_, expected_x_ = sess.run( + [kernel_results.acceptance_probs, samples, expected_x], + feed_dict) + + actual_x = samples_.mean() + actual_exp_x = np.exp(samples_).mean() + + logging_ops.vlog(1, "True E[x, exp(x)]: {}\t{}".format( + expected_x_, expected_exp_x)) + logging_ops.vlog(1, "Estimated E[x, exp(x)]: {}\t{}".format( + actual_x, actual_exp_x)) + self.assertNear(actual_x, expected_x_, 2e-2) + self.assertNear(actual_exp_x, expected_exp_x, 2e-2) + self.assertAllEqual(np.ones_like(acceptance_probs_, np.bool), + acceptance_probs_ > 0.5) + self.assertAllEqual(np.ones_like(acceptance_probs_, np.bool), + acceptance_probs_ <= 1.) + + def _chain_gets_correct_expectations_wrapper(self, independent_chain_ndims): + with self.test_session(graph=ops.Graph()) as sess: + x_ph = array_ops.placeholder(np.float32, name="x_ph") + feed_dict = {x_ph: np.random.rand(50, 10, 2)} + self._chain_gets_correct_expectations(x_ph, independent_chain_ndims, + sess, feed_dict) def testHMCChainExpectationsNullShape(self): - self._chain_gets_correct_expectations_wrapper([]) + self._chain_gets_correct_expectations_wrapper(0) def testHMCChainExpectations1(self): - self._chain_gets_correct_expectations_wrapper([1]) + self._chain_gets_correct_expectations_wrapper(1) def testHMCChainExpectations2(self): - self._chain_gets_correct_expectations_wrapper([2]) + self._chain_gets_correct_expectations_wrapper(2) - def testHMCChainExpectations12(self): - self._chain_gets_correct_expectations_wrapper([1, 2]) - - def _kernel_leaves_target_invariant(self, initial_draws, event_dims, + def testKernelResultsUsingTruncatedDistribution(self): + def log_prob(x): + return array_ops.where( + x >= 0., + -x - x**2, # Non-constant gradient. + array_ops.fill(x.shape, math_ops.cast(-np.inf, x.dtype))) + # This log_prob has the property that it is likely to attract + # the HMC flow toward, and below, zero...but for x <=0, + # log_prob(x) = -inf, which should result in rejection, as well + # as a non-finite log_prob. Thus, this distribution gives us an opportunity + # to test out the kernel results ability to correctly capture rejections due + # to finite AND non-finite reasons. + # Why use a non-constant gradient? This ensures the leapfrog integrator + # will not be exact. + + num_results = 1000 + # Large step size, will give rejections due to integration error in addition + # to rejection due to going into a region of log_prob = -inf. + step_size = 0.1 + num_leapfrog_steps = 5 + num_chains = 2 + + with self.test_session(graph=ops.Graph()) as sess: + + # Start multiple independent chains. + initial_state = ops.convert_to_tensor([0.1] * num_chains) + + states, kernel_results = hmc.sample_chain( + num_results=num_results, + target_log_prob_fn=log_prob, + current_state=initial_state, + step_size=step_size, + num_leapfrog_steps=num_leapfrog_steps, + seed=42) + + states_, kernel_results_ = sess.run([states, kernel_results]) + pstates_ = kernel_results_.proposed_state + + neg_inf_mask = np.isneginf(kernel_results_.proposed_target_log_prob) + + # First: Test that the mathematical properties of the above log prob + # function in conjunction with HMC show up as expected in kernel_results_. + + # We better have log_prob = -inf some of the time. + self.assertLess(0, neg_inf_mask.sum()) + # We better have some rejections due to something other than -inf. + self.assertLess(neg_inf_mask.sum(), (~kernel_results_.is_accepted).sum()) + # We better have been accepted a decent amount, even near the end of the + # chain, or else this HMC run just got stuck at some point. + self.assertLess( + 0.1, kernel_results_.is_accepted[int(0.9 * num_results):].mean()) + # We better not have any NaNs in proposed state or log_prob. + # We may have some NaN in grads, which involve multiplication/addition due + # to gradient rules. This is the known "NaN grad issue with tf.where." + self.assertAllEqual(np.zeros_like(states_), + np.isnan(kernel_results_.proposed_target_log_prob)) + self.assertAllEqual(np.zeros_like(states_), + np.isnan(states_)) + # We better not have any +inf in states, grads, or log_prob. + self.assertAllEqual(np.zeros_like(states_), + np.isposinf(kernel_results_.proposed_target_log_prob)) + self.assertAllEqual( + np.zeros_like(states_), + np.isposinf(kernel_results_.proposed_grads_target_log_prob[0])) + self.assertAllEqual(np.zeros_like(states_), + np.isposinf(states_)) + + # Second: Test that kernel_results is congruent with itself and + # acceptance/rejection of states. + + # Proposed state is negative iff proposed target log prob is -inf. + np.testing.assert_array_less(pstates_[neg_inf_mask], 0.) + np.testing.assert_array_less(0., pstates_[~neg_inf_mask]) + + # Acceptance probs are zero whenever proposed state is negative. + self.assertAllEqual( + np.zeros_like(pstates_[neg_inf_mask]), + kernel_results_.acceptance_probs[neg_inf_mask]) + + # The move is accepted ==> state = proposed state. + self.assertAllEqual( + states_[kernel_results_.is_accepted], + pstates_[kernel_results_.is_accepted], + ) + # The move was rejected <==> state[t] == state[t - 1]. + for t in range(1, num_results): + for i in range(num_chains): + if kernel_results_.is_accepted[t, i]: + self.assertNotEqual(states_[t, i], states_[t - 1, i]) + else: + self.assertEqual(states_[t, i], states_[t - 1, i]) + + def _kernel_leaves_target_invariant(self, initial_draws, + independent_chain_ndims, sess, feed_dict=None): def log_gamma_log_prob(x): + event_dims = math_ops.range(independent_chain_ndims, array_ops.rank(x)) return self._log_gamma_log_prob(x, event_dims) def fake_log_prob(x): """Cooled version of the target distribution.""" return 1.1 * log_gamma_log_prob(x) - step_size = array_ops.placeholder(np.float32, [], name='step_size') + step_size = array_ops.placeholder(np.float32, [], name="step_size") if feed_dict is None: feed_dict = {} feed_dict[step_size] = 0.4 - sample, acceptance_probs, _, _ = hmc.kernel(step_size, 5, initial_draws, - log_gamma_log_prob, event_dims) - bad_sample, bad_acceptance_probs, _, _ = hmc.kernel( - step_size, 5, initial_draws, fake_log_prob, event_dims) - (acceptance_probs_val, bad_acceptance_probs_val, initial_draws_val, - updated_draws_val, fake_draws_val) = sess.run([acceptance_probs, - bad_acceptance_probs, - initial_draws, sample, - bad_sample], feed_dict) + sample, kernel_results = hmc.kernel( + target_log_prob_fn=log_gamma_log_prob, + current_state=initial_draws, + step_size=step_size, + num_leapfrog_steps=5, + seed=43) + + bad_sample, bad_kernel_results = hmc.kernel( + target_log_prob_fn=fake_log_prob, + current_state=initial_draws, + step_size=step_size, + num_leapfrog_steps=5, + seed=44) + + [ + acceptance_probs_, + bad_acceptance_probs_, + initial_draws_, + updated_draws_, + fake_draws_, + ] = sess.run([ + kernel_results.acceptance_probs, + bad_kernel_results.acceptance_probs, + initial_draws, + sample, + bad_sample, + ], feed_dict) + # Confirm step size is small enough that we usually accept. - self.assertGreater(acceptance_probs_val.mean(), 0.5) - self.assertGreater(bad_acceptance_probs_val.mean(), 0.5) + self.assertGreater(acceptance_probs_.mean(), 0.5) + self.assertGreater(bad_acceptance_probs_.mean(), 0.5) + # Confirm step size is large enough that we sometimes reject. - self.assertLess(acceptance_probs_val.mean(), 0.99) - self.assertLess(bad_acceptance_probs_val.mean(), 0.99) - _, ks_p_value_true = stats.ks_2samp(initial_draws_val.flatten(), - updated_draws_val.flatten()) - _, ks_p_value_fake = stats.ks_2samp(initial_draws_val.flatten(), - fake_draws_val.flatten()) - logging.vlog(1, 'acceptance rate for true target: {}'.format( - acceptance_probs_val.mean())) - logging.vlog(1, 'acceptance rate for fake target: {}'.format( - bad_acceptance_probs_val.mean())) - logging.vlog(1, 'K-S p-value for true target: {}'.format(ks_p_value_true)) - logging.vlog(1, 'K-S p-value for fake target: {}'.format(ks_p_value_fake)) + self.assertLess(acceptance_probs_.mean(), 0.99) + self.assertLess(bad_acceptance_probs_.mean(), 0.99) + + _, ks_p_value_true = stats.ks_2samp(initial_draws_.flatten(), + updated_draws_.flatten()) + _, ks_p_value_fake = stats.ks_2samp(initial_draws_.flatten(), + fake_draws_.flatten()) + + logging_ops.vlog(1, "acceptance rate for true target: {}".format( + acceptance_probs_.mean())) + logging_ops.vlog(1, "acceptance rate for fake target: {}".format( + bad_acceptance_probs_.mean())) + logging_ops.vlog(1, "K-S p-value for true target: {}".format( + ks_p_value_true)) + logging_ops.vlog(1, "K-S p-value for fake target: {}".format( + ks_p_value_fake)) # Make sure that the MCMC update hasn't changed the empirical CDF much. self.assertGreater(ks_p_value_true, 1e-3) # Confirm that targeting the wrong distribution does # significantly change the empirical CDF. self.assertLess(ks_p_value_fake, 1e-6) - def _kernel_leaves_target_invariant_wrapper(self, event_dims): + def _kernel_leaves_target_invariant_wrapper(self, independent_chain_ndims): """Tests that the kernel leaves the target distribution invariant. Draws some independent samples from the target distribution, @@ -267,86 +435,160 @@ class HMCTest(test.TestCase): does change the target distribution. (And that we can detect that.) Args: - event_dims: A tuple of dimensions that should not be treated as - independent. This allows for multiple chains to be run independently - in parallel. Default is (), i.e., all dimensions are independent. + independent_chain_ndims: Python `int` scalar representing the number of + dims associated with independent chains. """ - with self.test_session() as sess: + with self.test_session(graph=ops.Graph()) as sess: initial_draws = np.log(np.random.gamma(self._shape_param, size=[50000, 2, 2])) initial_draws -= np.log(self._rate_param) - x_ph = array_ops.placeholder(np.float32, name='x_ph') + x_ph = array_ops.placeholder(np.float32, name="x_ph") feed_dict = {x_ph: initial_draws} - self._kernel_leaves_target_invariant(x_ph, event_dims, sess, - feed_dict) - - def testKernelLeavesTargetInvariantNullShape(self): - self._kernel_leaves_target_invariant_wrapper([]) + self._kernel_leaves_target_invariant(x_ph, independent_chain_ndims, + sess, feed_dict) def testKernelLeavesTargetInvariant1(self): - self._kernel_leaves_target_invariant_wrapper([1]) + self._kernel_leaves_target_invariant_wrapper(1) def testKernelLeavesTargetInvariant2(self): - self._kernel_leaves_target_invariant_wrapper([2]) + self._kernel_leaves_target_invariant_wrapper(2) + + def testKernelLeavesTargetInvariant3(self): + self._kernel_leaves_target_invariant_wrapper(3) - def testKernelLeavesTargetInvariant12(self): - self._kernel_leaves_target_invariant_wrapper([1, 2]) + def _ais_gets_correct_log_normalizer(self, init, independent_chain_ndims, + sess, feed_dict=None): + counter = collections.Counter() - def _ais_gets_correct_log_normalizer(self, init, event_dims, sess, - feed_dict=None): def proposal_log_prob(x): - return math_ops.reduce_sum(-0.5 * x * x - 0.5 * np.log(2*np.pi), - event_dims) + counter["proposal_calls"] += 1 + event_dims = math_ops.range(independent_chain_ndims, array_ops.rank(x)) + return -0.5 * math_ops.reduce_sum(x**2. + np.log(2 * np.pi), + axis=event_dims) def target_log_prob(x): + counter["target_calls"] += 1 + event_dims = math_ops.range(independent_chain_ndims, array_ops.rank(x)) return self._log_gamma_log_prob(x, event_dims) if feed_dict is None: feed_dict = {} - w, _, _ = hmc.ais_chain(200, 0.5, 2, init, target_log_prob, - proposal_log_prob, event_dims) + num_steps = 200 + + _, ais_weights, _ = hmc.sample_annealed_importance_chain( + proposal_log_prob_fn=proposal_log_prob, + num_steps=num_steps, + target_log_prob_fn=target_log_prob, + step_size=0.5, + current_state=init, + num_leapfrog_steps=2, + seed=45) + + # We have three calls because the calculation of `ais_weights` entails + # another call to the `convex_combined_log_prob_fn`. We could refactor + # things to avoid this, if needed (eg, b/72994218). + self.assertAllEqual(dict(target_calls=3, proposal_calls=3), counter) + + event_shape = array_ops.shape(init)[independent_chain_ndims:] + event_size = math_ops.reduce_prod(event_shape) + + log_true_normalizer = ( + -self._shape_param * math_ops.log(self._rate_param) + + math_ops.lgamma(self._shape_param)) + log_true_normalizer *= math_ops.cast(event_size, log_true_normalizer.dtype) + + log_estimated_normalizer = (math_ops.reduce_logsumexp(ais_weights) + - np.log(num_steps)) + + ratio_estimate_true = math_ops.exp(ais_weights - log_true_normalizer) + ais_weights_size = array_ops.size(ais_weights) + standard_error = math_ops.sqrt( + _reduce_variance(ratio_estimate_true) + / math_ops.cast(ais_weights_size, ratio_estimate_true.dtype)) + + [ + ratio_estimate_true_, + log_true_normalizer_, + log_estimated_normalizer_, + standard_error_, + ais_weights_size_, + event_size_, + ] = sess.run([ + ratio_estimate_true, + log_true_normalizer, + log_estimated_normalizer, + standard_error, + ais_weights_size, + event_size, + ], feed_dict) + + logging_ops.vlog(1, " log_true_normalizer: {}\n" + " log_estimated_normalizer: {}\n" + " ais_weights_size: {}\n" + " event_size: {}\n".format( + log_true_normalizer_, + log_estimated_normalizer_, + ais_weights_size_, + event_size_)) + self.assertNear(ratio_estimate_true_.mean(), 1., 4. * standard_error_) + + def _ais_gets_correct_log_normalizer_wrapper(self, independent_chain_ndims): + """Tests that AIS yields reasonable estimates of normalizers.""" + with self.test_session(graph=ops.Graph()) as sess: + x_ph = array_ops.placeholder(np.float32, name="x_ph") + initial_draws = np.random.normal(size=[30, 2, 1]) + self._ais_gets_correct_log_normalizer( + x_ph, + independent_chain_ndims, + sess, + feed_dict={x_ph: initial_draws}) - w_val = sess.run(w, feed_dict) - init_shape = sess.run(init, feed_dict).shape - normalizer_multiplier = np.prod([init_shape[i] for i in event_dims]) + def testAIS1(self): + self._ais_gets_correct_log_normalizer_wrapper(1) - true_normalizer = -self._shape_param * np.log(self._rate_param) - true_normalizer += special.gammaln(self._shape_param) - true_normalizer *= normalizer_multiplier + def testAIS2(self): + self._ais_gets_correct_log_normalizer_wrapper(2) - n_weights = np.prod(w_val.shape) - normalized_w = np.exp(w_val - true_normalizer) - standard_error = np.std(normalized_w) / np.sqrt(n_weights) - logging.vlog(1, 'True normalizer {}, estimated {}, n_weights {}'.format( - true_normalizer, np.log(normalized_w.mean()) + true_normalizer, - n_weights)) - self.assertNear(normalized_w.mean(), 1.0, 4.0 * standard_error) + def testAIS3(self): + self._ais_gets_correct_log_normalizer_wrapper(3) - def _ais_gets_correct_log_normalizer_wrapper(self, event_dims): - """Tests that AIS yields reasonable estimates of normalizers.""" - with self.test_session() as sess: - x_ph = array_ops.placeholder(np.float32, name='x_ph') + def testSampleAIChainSeedReproducibleWorksCorrectly(self): + with self.test_session(graph=ops.Graph()) as sess: + independent_chain_ndims = 1 + x = np.random.rand(4, 3, 2) - initial_draws = np.random.normal(size=[30, 2, 1]) - feed_dict = {x_ph: initial_draws} + def proposal_log_prob(x): + event_dims = math_ops.range(independent_chain_ndims, array_ops.rank(x)) + return -0.5 * math_ops.reduce_sum(x**2. + np.log(2 * np.pi), + axis=event_dims) - self._ais_gets_correct_log_normalizer(x_ph, event_dims, sess, - feed_dict) + def target_log_prob(x): + event_dims = math_ops.range(independent_chain_ndims, array_ops.rank(x)) + return self._log_gamma_log_prob(x, event_dims) - def testAISNullShape(self): - self._ais_gets_correct_log_normalizer_wrapper([]) + ais_kwargs = dict( + proposal_log_prob_fn=proposal_log_prob, + num_steps=200, + target_log_prob_fn=target_log_prob, + step_size=0.5, + current_state=x, + num_leapfrog_steps=2, + seed=53) - def testAIS1(self): - self._ais_gets_correct_log_normalizer_wrapper([1]) + _, ais_weights0, _ = hmc.sample_annealed_importance_chain( + **ais_kwargs) - def testAIS2(self): - self._ais_gets_correct_log_normalizer_wrapper([2]) + _, ais_weights1, _ = hmc.sample_annealed_importance_chain( + **ais_kwargs) + + [ais_weights0_, ais_weights1_] = sess.run([ + ais_weights0, ais_weights1]) - def testAIS12(self): - self._ais_gets_correct_log_normalizer_wrapper([1, 2]) + self.assertAllClose(ais_weights0_, ais_weights1_, + atol=1e-5, rtol=1e-5) def testNanRejection(self): """Tests that an update that yields NaN potentials gets rejected. @@ -359,86 +601,263 @@ class HMCTest(test.TestCase): """ def _unbounded_exponential_log_prob(x): """An exponential distribution with log-likelihood NaN for x < 0.""" - per_element_potentials = array_ops.where(x < 0, - np.nan * array_ops.ones_like(x), - -x) + per_element_potentials = array_ops.where( + x < 0., + array_ops.fill(array_ops.shape(x), x.dtype.as_numpy_dtype(np.nan)), + -x) return math_ops.reduce_sum(per_element_potentials) - with self.test_session() as sess: + with self.test_session(graph=ops.Graph()) as sess: initial_x = math_ops.linspace(0.01, 5, 10) - updated_x, acceptance_probs, _, _ = hmc.kernel( - 2., 5, initial_x, _unbounded_exponential_log_prob, [0]) - initial_x_val, updated_x_val, acceptance_probs_val = sess.run( - [initial_x, updated_x, acceptance_probs]) - - logging.vlog(1, 'initial_x = {}'.format(initial_x_val)) - logging.vlog(1, 'updated_x = {}'.format(updated_x_val)) - logging.vlog(1, 'acceptance_probs = {}'.format(acceptance_probs_val)) - - self.assertAllEqual(initial_x_val, updated_x_val) - self.assertEqual(acceptance_probs_val, 0.) + updated_x, kernel_results = hmc.kernel( + target_log_prob_fn=_unbounded_exponential_log_prob, + current_state=initial_x, + step_size=2., + num_leapfrog_steps=5, + seed=46) + initial_x_, updated_x_, acceptance_probs_ = sess.run( + [initial_x, updated_x, kernel_results.acceptance_probs]) + + logging_ops.vlog(1, "initial_x = {}".format(initial_x_)) + logging_ops.vlog(1, "updated_x = {}".format(updated_x_)) + logging_ops.vlog(1, "acceptance_probs = {}".format(acceptance_probs_)) + + self.assertAllEqual(initial_x_, updated_x_) + self.assertEqual(acceptance_probs_, 0.) def testNanFromGradsDontPropagate(self): """Test that update with NaN gradients does not cause NaN in results.""" def _nan_log_prob_with_nan_gradient(x): return np.nan * math_ops.reduce_sum(x) - with self.test_session() as sess: + with self.test_session(graph=ops.Graph()) as sess: initial_x = math_ops.linspace(0.01, 5, 10) - updated_x, acceptance_probs, new_log_prob, new_grad = hmc.kernel( - 2., 5, initial_x, _nan_log_prob_with_nan_gradient, [0]) - initial_x_val, updated_x_val, acceptance_probs_val = sess.run( - [initial_x, updated_x, acceptance_probs]) - - logging.vlog(1, 'initial_x = {}'.format(initial_x_val)) - logging.vlog(1, 'updated_x = {}'.format(updated_x_val)) - logging.vlog(1, 'acceptance_probs = {}'.format(acceptance_probs_val)) - - self.assertAllEqual(initial_x_val, updated_x_val) - self.assertEqual(acceptance_probs_val, 0.) + updated_x, kernel_results = hmc.kernel( + target_log_prob_fn=_nan_log_prob_with_nan_gradient, + current_state=initial_x, + step_size=2., + num_leapfrog_steps=5, + seed=47) + initial_x_, updated_x_, acceptance_probs_ = sess.run( + [initial_x, updated_x, kernel_results.acceptance_probs]) + + logging_ops.vlog(1, "initial_x = {}".format(initial_x_)) + logging_ops.vlog(1, "updated_x = {}".format(updated_x_)) + logging_ops.vlog(1, "acceptance_probs = {}".format(acceptance_probs_)) + + self.assertAllEqual(initial_x_, updated_x_) + self.assertEqual(acceptance_probs_, 0.) self.assertAllFinite( - gradients_impl.gradients(updated_x, initial_x)[0].eval()) - self.assertTrue( - gradients_impl.gradients(new_grad, initial_x)[0] is None) + gradients_ops.gradients(updated_x, initial_x)[0].eval()) + self.assertAllEqual([True], [g is None for g in gradients_ops.gradients( + kernel_results.proposed_grads_target_log_prob, initial_x)]) + self.assertAllEqual([False], [g is None for g in gradients_ops.gradients( + kernel_results.proposed_grads_target_log_prob, + kernel_results.proposed_state)]) # Gradients of the acceptance probs and new log prob are not finite. - _ = new_log_prob # Prevent unused arg error. # self.assertAllFinite( - # gradients_impl.gradients(acceptance_probs, initial_x)[0].eval()) + # gradients_ops.gradients(acceptance_probs, initial_x)[0].eval()) # self.assertAllFinite( - # gradients_impl.gradients(new_log_prob, initial_x)[0].eval()) + # gradients_ops.gradients(new_log_prob, initial_x)[0].eval()) + + def _testChainWorksDtype(self, dtype): + with self.test_session(graph=ops.Graph()) as sess: + states, kernel_results = hmc.sample_chain( + num_results=10, + target_log_prob_fn=lambda x: -math_ops.reduce_sum(x**2., axis=-1), + current_state=np.zeros(5).astype(dtype), + step_size=0.01, + num_leapfrog_steps=10, + seed=48) + states_, acceptance_probs_ = sess.run( + [states, kernel_results.acceptance_probs]) + self.assertEqual(dtype, states_.dtype) + self.assertEqual(dtype, acceptance_probs_.dtype) def testChainWorksIn64Bit(self): - def log_prob(x): - return - math_ops.reduce_sum(x * x, axis=-1) - states, acceptance_probs = hmc.chain( - n_iterations=10, - step_size=np.float64(0.01), - n_leapfrog_steps=10, - initial_x=np.zeros(5).astype(np.float64), - target_log_prob_fn=log_prob, - event_dims=[-1]) - with self.test_session() as sess: - states_, acceptance_probs_ = sess.run([states, acceptance_probs]) - self.assertEqual(np.float64, states_.dtype) - self.assertEqual(np.float64, acceptance_probs_.dtype) + self._testChainWorksDtype(np.float64) def testChainWorksIn16Bit(self): - def log_prob(x): - return - math_ops.reduce_sum(x * x, axis=-1) - states, acceptance_probs = hmc.chain( - n_iterations=10, - step_size=np.float16(0.01), - n_leapfrog_steps=10, - initial_x=np.zeros(5).astype(np.float16), - target_log_prob_fn=log_prob, - event_dims=[-1]) - with self.test_session() as sess: - states_, acceptance_probs_ = sess.run([states, acceptance_probs]) - self.assertEqual(np.float16, states_.dtype) - self.assertEqual(np.float16, acceptance_probs_.dtype) - - -if __name__ == '__main__': + self._testChainWorksDtype(np.float16) + + def testChainWorksCorrelatedMultivariate(self): + dtype = np.float32 + true_mean = dtype([0, 0]) + true_cov = dtype([[1, 0.5], + [0.5, 1]]) + num_results = 2000 + counter = collections.Counter() + with self.test_session(graph=ops.Graph()) as sess: + def target_log_prob(x, y): + counter["target_calls"] += 1 + # Corresponds to unnormalized MVN. + # z = matmul(inv(chol(true_cov)), [x, y] - true_mean) + z = array_ops.stack([x, y], axis=-1) - true_mean + z = array_ops.squeeze( + gen_linalg_ops.matrix_triangular_solve( + np.linalg.cholesky(true_cov), + z[..., array_ops.newaxis]), + axis=-1) + return -0.5 * math_ops.reduce_sum(z**2., axis=-1) + states, _ = hmc.sample_chain( + num_results=num_results, + target_log_prob_fn=target_log_prob, + current_state=[dtype(-2), dtype(2)], + step_size=[0.5, 0.5], + num_leapfrog_steps=2, + num_burnin_steps=200, + num_steps_between_results=1, + seed=54) + self.assertAllEqual(dict(target_calls=2), counter) + states = array_ops.stack(states, axis=-1) + self.assertEqual(num_results, states.shape[0].value) + sample_mean = math_ops.reduce_mean(states, axis=0) + x = states - sample_mean + sample_cov = math_ops.matmul(x, x, transpose_a=True) / dtype(num_results) + [sample_mean_, sample_cov_] = sess.run([ + sample_mean, sample_cov]) + self.assertAllClose(true_mean, sample_mean_, + atol=0.05, rtol=0.) + self.assertAllClose(true_cov, sample_cov_, + atol=0., rtol=0.1) + + +class _EnergyComputationTest(object): + + def testHandlesNanFromPotential(self): + with self.test_session(graph=ops.Graph()) as sess: + x = [1, np.inf, -np.inf, np.nan] + target_log_prob, proposed_target_log_prob = [ + self.dtype(x.flatten()) for x in np.meshgrid(x, x)] + num_chains = len(target_log_prob) + dummy_momentums = [-1, 1] + momentums = [self.dtype([dummy_momentums] * num_chains)] + proposed_momentums = [self.dtype([dummy_momentums] * num_chains)] + + target_log_prob = ops.convert_to_tensor(target_log_prob) + momentums = [ops.convert_to_tensor(momentums[0])] + proposed_target_log_prob = ops.convert_to_tensor(proposed_target_log_prob) + proposed_momentums = [ops.convert_to_tensor(proposed_momentums[0])] + + energy = _compute_energy_change( + target_log_prob, + momentums, + proposed_target_log_prob, + proposed_momentums, + independent_chain_ndims=1) + grads = gradients_ops.gradients(energy, momentums) + + [actual_energy, grads_] = sess.run([energy, grads]) + + # Ensure energy is `inf` (note: that's positive inf) in weird cases and + # finite otherwise. + expected_energy = self.dtype([0] + [np.inf]*(num_chains - 1)) + self.assertAllEqual(expected_energy, actual_energy) + + # Ensure gradient is finite. + self.assertAllEqual(np.ones_like(grads_).astype(np.bool), + np.isfinite(grads_)) + + def testHandlesNanFromKinetic(self): + with self.test_session(graph=ops.Graph()) as sess: + x = [1, np.inf, -np.inf, np.nan] + momentums, proposed_momentums = [ + [np.reshape(self.dtype(x), [-1, 1])] + for x in np.meshgrid(x, x)] + num_chains = len(momentums[0]) + target_log_prob = np.ones(num_chains, self.dtype) + proposed_target_log_prob = np.ones(num_chains, self.dtype) + + target_log_prob = ops.convert_to_tensor(target_log_prob) + momentums = [ops.convert_to_tensor(momentums[0])] + proposed_target_log_prob = ops.convert_to_tensor(proposed_target_log_prob) + proposed_momentums = [ops.convert_to_tensor(proposed_momentums[0])] + + energy = _compute_energy_change( + target_log_prob, + momentums, + proposed_target_log_prob, + proposed_momentums, + independent_chain_ndims=1) + grads = gradients_ops.gradients(energy, momentums) + + [actual_energy, grads_] = sess.run([energy, grads]) + + # Ensure energy is `inf` (note: that's positive inf) in weird cases and + # finite otherwise. + expected_energy = self.dtype([0] + [np.inf]*(num_chains - 1)) + self.assertAllEqual(expected_energy, actual_energy) + + # Ensure gradient is finite. + g = grads_[0].reshape([len(x), len(x)])[:, 0] + self.assertAllEqual(np.ones_like(g).astype(np.bool), np.isfinite(g)) + + # The remaining gradients are nan because the momentum was itself nan or + # inf. + g = grads_[0].reshape([len(x), len(x)])[:, 1:] + self.assertAllEqual(np.ones_like(g).astype(np.bool), np.isnan(g)) + + +class EnergyComputationTest16(test.TestCase, _EnergyComputationTest): + dtype = np.float16 + + +class EnergyComputationTest32(test.TestCase, _EnergyComputationTest): + dtype = np.float32 + + +class EnergyComputationTest64(test.TestCase, _EnergyComputationTest): + dtype = np.float64 + + +class _HMCHandlesLists(object): + + def testStateParts(self): + with self.test_session(graph=ops.Graph()) as sess: + dist_x = normal_lib.Normal(loc=self.dtype(0), scale=self.dtype(1)) + dist_y = independent_lib.Independent( + gamma_lib.Gamma(concentration=self.dtype([1, 2]), + rate=self.dtype([0.5, 0.75])), + reinterpreted_batch_ndims=1) + def target_log_prob(x, y): + return dist_x.log_prob(x) + dist_y.log_prob(y) + x0 = [dist_x.sample(seed=1), dist_y.sample(seed=2)] + samples, _ = hmc.sample_chain( + num_results=int(2e3), + target_log_prob_fn=target_log_prob, + current_state=x0, + step_size=0.85, + num_leapfrog_steps=3, + num_burnin_steps=int(250), + seed=49) + actual_means = [math_ops.reduce_mean(s, axis=0) for s in samples] + actual_vars = [_reduce_variance(s, axis=0) for s in samples] + expected_means = [dist_x.mean(), dist_y.mean()] + expected_vars = [dist_x.variance(), dist_y.variance()] + [ + actual_means_, + actual_vars_, + expected_means_, + expected_vars_, + ] = sess.run([ + actual_means, + actual_vars, + expected_means, + expected_vars, + ]) + self.assertAllClose(expected_means_, actual_means_, atol=0.05, rtol=0.16) + self.assertAllClose(expected_vars_, actual_vars_, atol=0., rtol=0.25) + + +class HMCHandlesLists32(_HMCHandlesLists, test.TestCase): + dtype = np.float32 + + +class HMCHandlesLists64(_HMCHandlesLists, test.TestCase): + dtype = np.float64 + + +if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/layers_conv_variational_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/layers_conv_variational_test.py index 57f44aef1a198f62cd8a715472a68a3d889ec3ac..750afb6654311fea30a1dc6b31b20aa3b4160ae2 100644 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/layers_conv_variational_test.py +++ b/tensorflow/contrib/bayesflow/python/kernel_tests/layers_conv_variational_test.py @@ -18,15 +18,21 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numpy as np + from tensorflow.contrib.bayesflow.python.ops import layers_conv_variational as prob_layers_lib from tensorflow.contrib.bayesflow.python.ops import layers_util as prob_layers_util from tensorflow.contrib.distributions.python.ops import independent as independent_lib +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 math_ops from tensorflow.python.ops import nn from tensorflow.python.ops import nn_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops.distributions import normal as normal_lib +from tensorflow.python.ops.distributions import util as distribution_util from tensorflow.python.platform import test @@ -46,7 +52,7 @@ class Counter(object): class MockDistribution(independent_lib.Independent): - """Monitors DenseVariational calls to the underlying distribution.""" + """Monitors layer calls to the underlying distribution.""" def __init__(self, result_sample, result_log_prob, loc=None, scale=None): self.result_sample = result_sample @@ -104,11 +110,14 @@ class ConvVariational(test.TestCase): def _testKLPenaltyKernel(self, layer_class): with self.test_session(): layer = layer_class(filters=2, kernel_size=3) - if layer_class == prob_layers_lib.Conv1DVariational: + if layer_class in (prob_layers_lib.Conv1DReparameterization, + prob_layers_lib.Conv1DFlipout): inputs = random_ops.random_uniform([2, 3, 1], seed=1) - elif layer_class == prob_layers_lib.Conv2DVariational: + elif layer_class in (prob_layers_lib.Conv2DReparameterization, + prob_layers_lib.Conv2DFlipout): inputs = random_ops.random_uniform([2, 3, 3, 1], seed=1) - elif layer_class == prob_layers_lib.Conv3DVariational: + elif layer_class in (prob_layers_lib.Conv3DReparameterization, + prob_layers_lib.Conv3DFlipout): inputs = random_ops.random_uniform([2, 3, 3, 3, 1], seed=1) # No keys. @@ -133,11 +142,14 @@ class ConvVariational(test.TestCase): kernel_size=3, bias_posterior_fn=prob_layers_util.default_mean_field_normal_fn(), bias_prior_fn=_make_normal) - if layer_class == prob_layers_lib.Conv1DVariational: + if layer_class in (prob_layers_lib.Conv1DReparameterization, + prob_layers_lib.Conv1DFlipout): inputs = random_ops.random_uniform([2, 3, 1], seed=1) - elif layer_class == prob_layers_lib.Conv2DVariational: + elif layer_class in (prob_layers_lib.Conv2DReparameterization, + prob_layers_lib.Conv2DFlipout): inputs = random_ops.random_uniform([2, 3, 3, 1], seed=1) - elif layer_class == prob_layers_lib.Conv3DVariational: + elif layer_class in (prob_layers_lib.Conv3DReparameterization, + prob_layers_lib.Conv3DFlipout): inputs = random_ops.random_uniform([2, 3, 3, 3, 1], seed=1) # No keys. @@ -152,42 +164,78 @@ class ConvVariational(test.TestCase): self.assertEqual(len(losses), 2) self.assertListEqual(layer.losses, losses) - def _testConvVariational(self, layer_class): + def _testConvSetUp(self, layer_class, batch_size, depth=None, + height=None, width=None, channels=None, filters=None, + **kwargs): + seed = Counter() + if layer_class in (prob_layers_lib.Conv1DReparameterization, + prob_layers_lib.Conv1DFlipout): + inputs = random_ops.random_uniform( + [batch_size, width, channels], seed=seed()) + kernel_size = (2,) + elif layer_class in (prob_layers_lib.Conv2DReparameterization, + prob_layers_lib.Conv2DFlipout): + inputs = random_ops.random_uniform( + [batch_size, height, width, channels], seed=seed()) + kernel_size = (2, 2) + elif layer_class in (prob_layers_lib.Conv3DReparameterization, + prob_layers_lib.Conv3DFlipout): + inputs = random_ops.random_uniform( + [batch_size, depth, height, width, channels], seed=seed()) + kernel_size = (2, 2, 2) + + kernel_shape = kernel_size + (channels, filters) + kernel_posterior = MockDistribution( + loc=random_ops.random_uniform(kernel_shape, seed=seed()), + scale=random_ops.random_uniform(kernel_shape, seed=seed()), + result_log_prob=random_ops.random_uniform(kernel_shape, seed=seed()), + result_sample=random_ops.random_uniform(kernel_shape, seed=seed())) + kernel_prior = MockDistribution( + result_log_prob=random_ops.random_uniform(kernel_shape, seed=seed()), + result_sample=random_ops.random_uniform(kernel_shape, seed=seed())) + kernel_divergence = MockKLDivergence( + result=random_ops.random_uniform(kernel_shape, seed=seed())) + + bias_size = (filters,) + bias_posterior = MockDistribution( + result_log_prob=random_ops.random_uniform(bias_size, seed=seed()), + result_sample=random_ops.random_uniform(bias_size, seed=seed())) + bias_prior = MockDistribution( + result_log_prob=random_ops.random_uniform(bias_size, seed=seed()), + result_sample=random_ops.random_uniform(bias_size, seed=seed())) + bias_divergence = MockKLDivergence( + result=random_ops.random_uniform(bias_size, seed=seed())) + + layer = layer_class( + filters=filters, + kernel_size=kernel_size, + padding="SAME", + kernel_posterior_fn=lambda *args: kernel_posterior, + kernel_posterior_tensor_fn=lambda d: d.sample(seed=42), + kernel_prior_fn=lambda *args: kernel_prior, + kernel_divergence_fn=kernel_divergence, + bias_posterior_fn=lambda *args: bias_posterior, + bias_posterior_tensor_fn=lambda d: d.sample(seed=43), + bias_prior_fn=lambda *args: bias_prior, + bias_divergence_fn=bias_divergence, + **kwargs) + + outputs = layer(inputs) + + kl_penalty = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) + return (kernel_posterior, kernel_prior, kernel_divergence, + bias_posterior, bias_prior, bias_divergence, + layer, inputs, outputs, kl_penalty, kernel_shape) + + def _testConvReparameterization(self, layer_class): batch_size, depth, height, width, channels, filters = 2, 4, 4, 4, 3, 5 with self.test_session() as sess: - seed = Counter() - if layer_class == prob_layers_lib.Conv1DVariational: - inputs = random_ops.random_uniform( - [batch_size, width, channels], seed=seed()) - kernel_size = (2,) - elif layer_class == prob_layers_lib.Conv2DVariational: - inputs = random_ops.random_uniform( - [batch_size, height, width, channels], seed=seed()) - kernel_size = (2, 2) - elif layer_class == prob_layers_lib.Conv3DVariational: - inputs = random_ops.random_uniform( - [batch_size, depth, height, width, channels], seed=seed()) - kernel_size = (2, 2, 2) - - kernel_shape = kernel_size + (channels, filters) - kernel_posterior = MockDistribution( - result_log_prob=random_ops.random_uniform(kernel_shape, seed=seed()), - result_sample=random_ops.random_uniform(kernel_shape, seed=seed())) - kernel_prior = MockDistribution( - result_log_prob=random_ops.random_uniform(kernel_shape, seed=seed()), - result_sample=random_ops.random_uniform(kernel_shape, seed=seed())) - kernel_divergence = MockKLDivergence( - result=random_ops.random_uniform(kernel_shape, seed=seed())) - - bias_size = (filters,) - bias_posterior = MockDistribution( - result_log_prob=random_ops.random_uniform(bias_size, seed=seed()), - result_sample=random_ops.random_uniform(bias_size, seed=seed())) - bias_prior = MockDistribution( - result_log_prob=random_ops.random_uniform(bias_size, seed=seed()), - result_sample=random_ops.random_uniform(bias_size, seed=seed())) - bias_divergence = MockKLDivergence( - result=random_ops.random_uniform(bias_size, seed=seed())) + (kernel_posterior, kernel_prior, kernel_divergence, + bias_posterior, bias_prior, bias_divergence, layer, inputs, + outputs, kl_penalty, kernel_shape) = self._testConvSetUp( + layer_class, batch_size, + depth=depth, height=height, width=width, channels=channels, + filters=filters) convolution_op = nn_ops.Convolution( tensor_shape.TensorShape(inputs.shape), @@ -198,23 +246,6 @@ class ConvVariational(test.TestCase): bias_posterior.result_sample, data_format="NHWC") - layer = layer_class( - filters=filters, - kernel_size=kernel_size, - padding="SAME", - kernel_posterior_fn=lambda *args: kernel_posterior, - kernel_posterior_tensor_fn=lambda d: d.sample(seed=42), - kernel_prior_fn=lambda *args: kernel_prior, - kernel_divergence_fn=kernel_divergence, - bias_posterior_fn=lambda *args: bias_posterior, - bias_posterior_tensor_fn=lambda d: d.sample(seed=43), - bias_prior_fn=lambda *args: bias_prior, - bias_divergence_fn=bias_divergence) - - outputs = layer(inputs) - - kl_penalty = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES) - [ expected_outputs_, actual_outputs_, expected_kernel_, actual_kernel_, @@ -257,32 +288,233 @@ class ConvVariational(test.TestCase): bias_posterior.result_sample]], bias_divergence.args) - def testKLPenaltyKernelConv1DVariational(self): - self._testKLPenaltyKernel(prob_layers_lib.Conv1DVariational) + def _testConvFlipout(self, layer_class): + batch_size, depth, height, width, channels, filters = 2, 4, 4, 4, 3, 5 + with self.test_session() as sess: + (kernel_posterior, kernel_prior, kernel_divergence, + bias_posterior, bias_prior, bias_divergence, layer, inputs, + outputs, kl_penalty, kernel_shape) = self._testConvSetUp( + layer_class, batch_size, + depth=depth, height=height, width=width, channels=channels, + filters=filters, seed=44) + + convolution_op = nn_ops.Convolution( + tensor_shape.TensorShape(inputs.shape), + filter_shape=tensor_shape.TensorShape(kernel_shape), + padding="SAME") + + expected_kernel_posterior_affine = normal_lib.Normal( + loc=array_ops.zeros_like(kernel_posterior.result_loc), + scale=kernel_posterior.result_scale) + expected_kernel_posterior_affine_tensor = ( + expected_kernel_posterior_affine.sample(seed=42)) + + expected_outputs = convolution_op( + inputs, kernel_posterior.distribution.loc) + + input_shape = array_ops.shape(inputs) + output_shape = array_ops.shape(expected_outputs) + batch_shape = array_ops.expand_dims(input_shape[0], 0) + channels = input_shape[-1] + rank = len(inputs.get_shape()) - 2 + + sign_input = random_ops.random_uniform( + array_ops.concat([batch_shape, + array_ops.expand_dims(channels, 0)], 0), + minval=0, + maxval=2, + dtype=dtypes.int32, + seed=layer.seed) + sign_input = math_ops.cast(2 * sign_input - 1, inputs.dtype) + sign_output = random_ops.random_uniform( + array_ops.concat([batch_shape, + array_ops.expand_dims(filters, 0)], 0), + minval=0, + maxval=2, + dtype=dtypes.int32, + seed=distribution_util.gen_new_seed( + layer.seed, salt="conv_flipout")) + sign_output = math_ops.cast(2 * sign_output - 1, inputs.dtype) + for _ in range(rank): + sign_input = array_ops.expand_dims(sign_input, 1) # 2D ex: (B, 1, 1, C) + sign_output = array_ops.expand_dims(sign_output, 1) + + sign_input = array_ops.tile( # tile for element-wise op broadcasting + sign_input, + [1] + [input_shape[i + 1] for i in range(rank)] + [1]) + sign_output = array_ops.tile( + sign_output, + [1] + [output_shape[i + 1] for i in range(rank)] + [1]) + + perturbed_inputs = convolution_op( + inputs * sign_input, expected_kernel_posterior_affine_tensor) + perturbed_inputs *= sign_output + + expected_outputs += perturbed_inputs + expected_outputs = nn.bias_add(expected_outputs, + bias_posterior.result_sample, + data_format="NHWC") + + [ + expected_outputs_, actual_outputs_, + expected_kernel_divergence_, actual_kernel_divergence_, + expected_bias_, actual_bias_, + expected_bias_divergence_, actual_bias_divergence_, + ] = sess.run([ + expected_outputs, outputs, + kernel_divergence.result, kl_penalty[0], + bias_posterior.result_sample, layer.bias_posterior_tensor, + bias_divergence.result, kl_penalty[1], + ]) + + self.assertAllClose( + expected_bias_, actual_bias_, + rtol=1e-6, atol=0.) + self.assertAllClose( + expected_outputs_, actual_outputs_, + rtol=1e-6, atol=0.) + self.assertAllClose( + expected_kernel_divergence_, actual_kernel_divergence_, + rtol=1e-6, atol=0.) + self.assertAllClose( + expected_bias_divergence_, actual_bias_divergence_, + rtol=1e-6, atol=0.) + + self.assertAllEqual( + [[kernel_posterior.distribution, kernel_prior.distribution, None]], + kernel_divergence.args) + + self.assertAllEqual( + [[bias_posterior.distribution, + bias_prior.distribution, + bias_posterior.result_sample]], + bias_divergence.args) + + def _testRandomConvFlipout(self, layer_class): + batch_size, depth, height, width, channels, filters = 2, 4, 4, 4, 3, 5 + with self.test_session() as sess: + seed = Counter() + if layer_class in (prob_layers_lib.Conv1DReparameterization, + prob_layers_lib.Conv1DFlipout): + inputs = random_ops.random_uniform( + [batch_size, width, channels], seed=seed()) + kernel_size = (2,) + elif layer_class in (prob_layers_lib.Conv2DReparameterization, + prob_layers_lib.Conv2DFlipout): + inputs = random_ops.random_uniform( + [batch_size, height, width, channels], seed=seed()) + kernel_size = (2, 2) + elif layer_class in (prob_layers_lib.Conv3DReparameterization, + prob_layers_lib.Conv3DFlipout): + inputs = random_ops.random_uniform( + [batch_size, depth, height, width, channels], seed=seed()) + kernel_size = (2, 2, 2) + + kernel_shape = kernel_size + (channels, filters) + bias_size = (filters,) + + kernel_posterior = MockDistribution( + loc=random_ops.random_uniform( + kernel_shape, seed=seed()), + scale=random_ops.random_uniform( + kernel_shape, seed=seed()), + result_log_prob=random_ops.random_uniform( + kernel_shape, seed=seed()), + result_sample=random_ops.random_uniform( + kernel_shape, seed=seed())) + bias_posterior = MockDistribution( + loc=random_ops.random_uniform( + bias_size, seed=seed()), + scale=random_ops.random_uniform( + bias_size, seed=seed()), + result_log_prob=random_ops.random_uniform( + bias_size, seed=seed()), + result_sample=random_ops.random_uniform( + bias_size, seed=seed())) + layer_one = layer_class( + filters=filters, + kernel_size=kernel_size, + padding="SAME", + kernel_posterior_fn=lambda *args: kernel_posterior, + kernel_posterior_tensor_fn=lambda d: d.sample(seed=42), + bias_posterior_fn=lambda *args: bias_posterior, + bias_posterior_tensor_fn=lambda d: d.sample(seed=43), + seed=44) + layer_two = layer_class( + filters=filters, + kernel_size=kernel_size, + padding="SAME", + kernel_posterior_fn=lambda *args: kernel_posterior, + kernel_posterior_tensor_fn=lambda d: d.sample(seed=42), + bias_posterior_fn=lambda *args: bias_posterior, + bias_posterior_tensor_fn=lambda d: d.sample(seed=43), + seed=45) + + outputs_one = layer_one(inputs) + outputs_two = layer_two(inputs) + + outputs_one_, outputs_two_ = sess.run([ + outputs_one, outputs_two]) + + self.assertLess(np.sum(np.isclose(outputs_one_, outputs_two_)), + np.prod(outputs_one_.shape)) + + def testKLPenaltyKernelConv1DReparameterization(self): + self._testKLPenaltyKernel(prob_layers_lib.Conv1DReparameterization) + + def testKLPenaltyKernelConv2DReparameterization(self): + self._testKLPenaltyKernel(prob_layers_lib.Conv2DReparameterization) + + def testKLPenaltyKernelConv3DReparameterization(self): + self._testKLPenaltyKernel(prob_layers_lib.Conv3DReparameterization) + + def testKLPenaltyKernelConv1DFlipout(self): + self._testKLPenaltyKernel(prob_layers_lib.Conv1DFlipout) + + def testKLPenaltyKernelConv2DFlipout(self): + self._testKLPenaltyKernel(prob_layers_lib.Conv2DFlipout) + + def testKLPenaltyKernelConv3DFlipout(self): + self._testKLPenaltyKernel(prob_layers_lib.Conv3DFlipout) + + def testKLPenaltyBothConv1DReparameterization(self): + self._testKLPenaltyBoth(prob_layers_lib.Conv1DReparameterization) + + def testKLPenaltyBothConv2DReparameterization(self): + self._testKLPenaltyBoth(prob_layers_lib.Conv2DReparameterization) + + def testKLPenaltyBothConv3DReparameterization(self): + self._testKLPenaltyBoth(prob_layers_lib.Conv3DReparameterization) + + def testKLPenaltyBothConv1DFlipout(self): + self._testKLPenaltyBoth(prob_layers_lib.Conv1DFlipout) + + def testKLPenaltyBothConv2DFlipout(self): + self._testKLPenaltyBoth(prob_layers_lib.Conv2DFlipout) - def testKLPenaltyKernelConv2DVariational(self): - self._testKLPenaltyKernel(prob_layers_lib.Conv2DVariational) + def testKLPenaltyBothConv3DFlipout(self): + self._testKLPenaltyBoth(prob_layers_lib.Conv3DFlipout) - def testKLPenaltyKernelConv3DVariational(self): - self._testKLPenaltyKernel(prob_layers_lib.Conv3DVariational) + def testConv1DReparameterization(self): + self._testConvReparameterization(prob_layers_lib.Conv1DReparameterization) - def testKLPenaltyBothConv1DVariational(self): - self._testKLPenaltyBoth(prob_layers_lib.Conv1DVariational) + def testConv2DReparameterization(self): + self._testConvReparameterization(prob_layers_lib.Conv2DReparameterization) - def testKLPenaltyBothConv2DVariational(self): - self._testKLPenaltyBoth(prob_layers_lib.Conv2DVariational) + def testConv3DReparameterization(self): + self._testConvReparameterization(prob_layers_lib.Conv3DReparameterization) - def testKLPenaltyBothConv3DVariational(self): - self._testKLPenaltyBoth(prob_layers_lib.Conv3DVariational) + def testConv1DFlipout(self): + self._testConvFlipout(prob_layers_lib.Conv1DFlipout) - def testConv1DVariational(self): - self._testConvVariational(prob_layers_lib.Conv1DVariational) + def testConv2DFlipout(self): + self._testConvFlipout(prob_layers_lib.Conv2DFlipout) - def testConv2DVariational(self): - self._testConvVariational(prob_layers_lib.Conv2DVariational) + def testConv3DFlipout(self): + self._testConvFlipout(prob_layers_lib.Conv3DFlipout) - def testConv3DVariational(self): - self._testConvVariational(prob_layers_lib.Conv3DVariational) + def testRandomConv1DFlipout(self): + self._testRandomConvFlipout(prob_layers_lib.Conv1DFlipout) if __name__ == "__main__": diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/layers_dense_variational_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/layers_dense_variational_test.py index 4e9f1193511c35beead85914ca988fde69b3afde..342f38ccec7ec74db1b393d6cdc22300205cc547 100644 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/layers_dense_variational_test.py +++ b/tensorflow/contrib/bayesflow/python/kernel_tests/layers_dense_variational_test.py @@ -20,7 +20,7 @@ from __future__ import print_function import numpy as np -from tensorflow.contrib.bayesflow.python.ops import layers_dense_variational_impl as prob_layers_lib +from tensorflow.contrib.bayesflow.python.ops import layers_dense_variational as prob_layers_lib from tensorflow.contrib.bayesflow.python.ops import layers_util as prob_layers_util from tensorflow.contrib.distributions.python.ops import independent as independent_lib from tensorflow.python.framework import dtypes diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/mcmc_diagnostics_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/mcmc_diagnostics_test.py new file mode 100644 index 0000000000000000000000000000000000000000..52e36e135d95c1ec919c710f35d59073c2134d05 --- /dev/null +++ b/tensorflow/contrib/bayesflow/python/kernel_tests/mcmc_diagnostics_test.py @@ -0,0 +1,445 @@ +# 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 MCMC diagnostic utilities.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.bayesflow.python.ops import mcmc_diagnostics_impl as mcmc_diagnostics +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import spectral_ops_test_util +from tensorflow.python.platform import test + +rng = np.random.RandomState(42) + + +class _EffectiveSampleSizeTest(object): + + @property + def use_static_shape(self): + raise NotImplementedError( + "Subclass failed to implement `use_static_shape`.") + + def _check_versus_expected_effective_sample_size(self, + x_, + expected_ess, + sess, + atol=1e-2, + rtol=1e-2, + filter_threshold=None, + filter_beyond_lag=None): + x = array_ops.placeholder_with_default( + input=x_, shape=x_.shape if self.use_static_shape else None) + ess = mcmc_diagnostics.effective_sample_size( + x, + filter_threshold=filter_threshold, + filter_beyond_lag=filter_beyond_lag) + if self.use_static_shape: + self.assertAllEqual(x.shape[1:], ess.shape) + + ess_ = sess.run(ess) + + self.assertAllClose( + np.ones_like(ess_) * expected_ess, ess_, atol=atol, rtol=rtol) + + def testIidRank1NormalHasFullEssMaxLags10(self): + # With a length 5000 iid normal sequence, and filter_beyond_lag = 10, we + # should have a good estimate of ESS, and it should be close to the full + # sequence length of 5000. + # The choice of filter_beyond_lag = 10 is a short cutoff, reasonable only + # since we know the correlation length should be zero right away. + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + self._check_versus_expected_effective_sample_size( + x_=rng.randn(5000).astype(np.float32), + expected_ess=5000, + sess=sess, + filter_beyond_lag=10, + filter_threshold=None, + rtol=0.3) + + def testIidRank2NormalHasFullEssMaxLags10(self): + # See similar test for Rank1Normal for reasoning. + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + self._check_versus_expected_effective_sample_size( + x_=rng.randn(5000, 2).astype(np.float32), + expected_ess=5000, + sess=sess, + filter_beyond_lag=10, + filter_threshold=None, + rtol=0.3) + + def testIidRank1NormalHasFullEssMaxLagThresholdZero(self): + # With a length 5000 iid normal sequence, and filter_threshold = 0, + # we should have a super-duper estimate of ESS, and it should be very close + # to the full sequence length of 5000. + # The choice of filter_beyond_lag = 0 means we cutoff as soon as the + # auto-corris below zero. This should happen very quickly, due to the fact + # that the theoretical auto-corr is [1, 0, 0,...] + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + self._check_versus_expected_effective_sample_size( + x_=rng.randn(5000).astype(np.float32), + expected_ess=5000, + sess=sess, + filter_beyond_lag=None, + filter_threshold=0., + rtol=0.1) + + def testIidRank2NormalHasFullEssMaxLagThresholdZero(self): + # See similar test for Rank1Normal for reasoning. + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + self._check_versus_expected_effective_sample_size( + x_=rng.randn(5000, 2).astype(np.float32), + expected_ess=5000, + sess=sess, + filter_beyond_lag=None, + filter_threshold=0., + rtol=0.1) + + def testLength10CorrelationHasEssOneTenthTotalLengthUsingMaxLags50(self): + # Create x_, such that + # x_[i] = iid_x_[0], i = 0,...,9 + # x_[i] = iid_x_[1], i = 10,..., 19, + # and so on. + iid_x_ = rng.randn(5000, 1).astype(np.float32) + x_ = (iid_x_ * np.ones((5000, 10)).astype(np.float32)).reshape((50000,)) + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + self._check_versus_expected_effective_sample_size( + x_=x_, + expected_ess=50000 // 10, + sess=sess, + filter_beyond_lag=50, + filter_threshold=None, + rtol=0.2) + + def testLength10CorrelationHasEssOneTenthTotalLengthUsingMaxLagsThresholdZero( + self): + # Create x_, such that + # x_[i] = iid_x_[0], i = 0,...,9 + # x_[i] = iid_x_[1], i = 10,..., 19, + # and so on. + iid_x_ = rng.randn(5000, 1).astype(np.float32) + x_ = (iid_x_ * np.ones((5000, 10)).astype(np.float32)).reshape((50000,)) + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + self._check_versus_expected_effective_sample_size( + x_=x_, + expected_ess=50000 // 10, + sess=sess, + filter_beyond_lag=None, + filter_threshold=0., + rtol=0.1) + + def testListArgs(self): + # x_ has correlation length 10 ==> ESS = N / 10 + # y_ has correlation length 1 ==> ESS = N + iid_x_ = rng.randn(5000, 1).astype(np.float32) + x_ = (iid_x_ * np.ones((5000, 10)).astype(np.float32)).reshape((50000,)) + y_ = rng.randn(50000).astype(np.float32) + states = [x_, x_, y_, y_] + filter_threshold = [0., None, 0., None] + filter_beyond_lag = [None, 5, None, 5] + + # See other tests for reasoning on tolerance. + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + ess = mcmc_diagnostics.effective_sample_size( + states, + filter_threshold=filter_threshold, + filter_beyond_lag=filter_beyond_lag) + ess_ = sess.run(ess) + self.assertAllEqual(4, len(ess_)) + + self.assertAllClose(50000 // 10, ess_[0], rtol=0.3) + self.assertAllClose(50000 // 10, ess_[1], rtol=0.3) + self.assertAllClose(50000, ess_[2], rtol=0.1) + self.assertAllClose(50000, ess_[3], rtol=0.1) + + def testMaxLagsThresholdLessThanNeg1SameAsNone(self): + # Setting both means we filter out items R_k from the auto-correlation + # sequence if k > filter_beyond_lag OR k >= j where R_j < filter_threshold. + + # x_ has correlation length 10. + iid_x_ = rng.randn(500, 1).astype(np.float32) + x_ = (iid_x_ * np.ones((500, 10)).astype(np.float32)).reshape((5000,)) + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + x = array_ops.placeholder_with_default( + input=x_, shape=x_.shape if self.use_static_shape else None) + + ess_none_none = mcmc_diagnostics.effective_sample_size( + x, filter_threshold=None, filter_beyond_lag=None) + ess_none_200 = mcmc_diagnostics.effective_sample_size( + x, filter_threshold=None, filter_beyond_lag=200) + ess_neg2_200 = mcmc_diagnostics.effective_sample_size( + x, filter_threshold=-2., filter_beyond_lag=200) + ess_neg2_none = mcmc_diagnostics.effective_sample_size( + x, filter_threshold=-2., filter_beyond_lag=None) + ess_none_none_, ess_none_200_, ess_neg2_200_, ess_neg2_none_ = sess.run( + [ess_none_none, ess_none_200, ess_neg2_200, ess_neg2_none]) + + # filter_threshold=-2 <==> filter_threshold=None. + self.assertAllClose(ess_none_none_, ess_neg2_none_) + self.assertAllClose(ess_none_200_, ess_neg2_200_) + + def testMaxLagsArgsAddInAnOrManner(self): + # Setting both means we filter out items R_k from the auto-correlation + # sequence if k > filter_beyond_lag OR k >= j where R_j < filter_threshold. + + # x_ has correlation length 10. + iid_x_ = rng.randn(500, 1).astype(np.float32) + x_ = (iid_x_ * np.ones((500, 10)).astype(np.float32)).reshape((5000,)) + with self.test_session() as sess: + with spectral_ops_test_util.fft_kernel_label_map(): + x = array_ops.placeholder_with_default( + input=x_, shape=x_.shape if self.use_static_shape else None) + + ess_1_9 = mcmc_diagnostics.effective_sample_size( + x, filter_threshold=1., filter_beyond_lag=9) + ess_1_none = mcmc_diagnostics.effective_sample_size( + x, filter_threshold=1., filter_beyond_lag=None) + ess_none_9 = mcmc_diagnostics.effective_sample_size( + x, filter_threshold=1., filter_beyond_lag=9) + ess_1_9_, ess_1_none_, ess_none_9_ = sess.run( + [ess_1_9, ess_1_none, ess_none_9]) + + # Since R_k = 1 for k < 10, and R_k < 1 for k >= 10, + # filter_threshold = 1 <==> filter_beyond_lag = 9. + self.assertAllClose(ess_1_9_, ess_1_none_) + self.assertAllClose(ess_1_9_, ess_none_9_) + + +class EffectiveSampleSizeStaticTest(test.TestCase, _EffectiveSampleSizeTest): + + @property + def use_static_shape(self): + return True + + +class EffectiveSampleSizeDynamicTest(test.TestCase, _EffectiveSampleSizeTest): + + @property + def use_static_shape(self): + return False + + +class _PotentialScaleReductionTest(object): + + @property + def use_static_shape(self): + raise NotImplementedError( + "Subclass failed to impliment `use_static_shape`.") + + def testListOfStatesWhereFirstPassesSecondFails(self): + """Simple test showing API with two states. Read first!.""" + n_samples = 1000 + + # state_0 is two scalar chains taken from iid Normal(0, 1). Will pass. + state_0 = rng.randn(n_samples, 2) + + # state_1 is three 4-variate chains taken from Normal(0, 1) that have been + # shifted. Since every chain is shifted, they are not the same, and the + # test should fail. + offset = np.array([1., -1., 2.]).reshape(3, 1) + state_1 = rng.randn(n_samples, 3, 4) + offset + + rhat = mcmc_diagnostics.potential_scale_reduction( + chains_states=[state_0, state_1], independent_chain_ndims=1) + + self.assertIsInstance(rhat, list) + with self.test_session() as sess: + rhat_0_, rhat_1_ = sess.run(rhat) + + # r_hat_0 should be close to 1, meaning test is passed. + self.assertAllEqual((), rhat_0_.shape) + self.assertAllClose(1., rhat_0_, rtol=0.02) + + # r_hat_1 should be greater than 1.2, meaning test has failed. + self.assertAllEqual((4,), rhat_1_.shape) + self.assertAllEqual(np.ones_like(rhat_1_).astype(bool), rhat_1_ > 1.2) + + def check_results(self, state_, independent_chain_shape, should_pass): + sample_ndims = 1 + independent_chain_ndims = len(independent_chain_shape) + with self.test_session(): + state = array_ops.placeholder_with_default( + input=state_, shape=state_.shape if self.use_static_shape else None) + + rhat = mcmc_diagnostics.potential_scale_reduction( + state, independent_chain_ndims=independent_chain_ndims) + + if self.use_static_shape: + self.assertAllEqual( + state_.shape[sample_ndims + independent_chain_ndims:], rhat.shape) + + rhat_ = rhat.eval() + if should_pass: + self.assertAllClose(np.ones_like(rhat_), rhat_, atol=0, rtol=0.02) + else: + self.assertAllEqual(np.ones_like(rhat_).astype(bool), rhat_ > 1.2) + + def iid_normal_chains_should_pass_wrapper(self, + sample_shape, + independent_chain_shape, + other_shape, + dtype=np.float32): + """Check results with iid normal chains.""" + + state_shape = sample_shape + independent_chain_shape + other_shape + state_ = rng.randn(*state_shape).astype(dtype) + + # The "other" dimensions do not have to be identical, just independent, so + # force them to not be identical. + if other_shape: + state_ *= rng.rand(*other_shape).astype(dtype) + + self.check_results(state_, independent_chain_shape, should_pass=True) + + def testPassingIIDNdimsAreIndependentOneOtherZero(self): + self.iid_normal_chains_should_pass_wrapper( + sample_shape=[10000], independent_chain_shape=[4], other_shape=[]) + + def testPassingIIDNdimsAreIndependentOneOtherOne(self): + self.iid_normal_chains_should_pass_wrapper( + sample_shape=[10000], independent_chain_shape=[3], other_shape=[7]) + + def testPassingIIDNdimsAreIndependentOneOtherTwo(self): + self.iid_normal_chains_should_pass_wrapper( + sample_shape=[10000], independent_chain_shape=[2], other_shape=[5, 7]) + + def testPassingIIDNdimsAreIndependentTwoOtherTwo64Bit(self): + self.iid_normal_chains_should_pass_wrapper( + sample_shape=[10000], + independent_chain_shape=[2, 3], + other_shape=[5, 7], + dtype=np.float64) + + def offset_normal_chains_should_fail_wrapper( + self, sample_shape, independent_chain_shape, other_shape): + """Check results with normal chains that are offset from each other.""" + + state_shape = sample_shape + independent_chain_shape + other_shape + state_ = rng.randn(*state_shape) + + # Add a significant offset to the different (formerly iid) chains. + offset = np.linspace( + 0, 2, num=np.prod(independent_chain_shape)).reshape([1] * len( + sample_shape) + independent_chain_shape + [1] * len(other_shape)) + state_ += offset + + self.check_results(state_, independent_chain_shape, should_pass=False) + + def testFailingOffsetNdimsAreSampleOneIndependentOneOtherOne(self): + self.offset_normal_chains_should_fail_wrapper( + sample_shape=[10000], independent_chain_shape=[2], other_shape=[5]) + + +class PotentialScaleReductionStaticTest(test.TestCase, + _PotentialScaleReductionTest): + + @property + def use_static_shape(self): + return True + + def testIndependentNdimsLessThanOneRaises(self): + with self.assertRaisesRegexp(ValueError, "independent_chain_ndims"): + mcmc_diagnostics.potential_scale_reduction( + rng.rand(2, 3, 4), independent_chain_ndims=0) + + +class PotentialScaleReductionDynamicTest(test.TestCase, + _PotentialScaleReductionTest): + + @property + def use_static_shape(self): + return False + + +class _ReduceVarianceTest(object): + + @property + def use_static_shape(self): + raise NotImplementedError( + "Subclass failed to impliment `use_static_shape`.") + + def check_versus_numpy(self, x_, axis, biased, keepdims): + with self.test_session(): + x_ = np.asarray(x_) + x = array_ops.placeholder_with_default( + input=x_, shape=x_.shape if self.use_static_shape else None) + var = mcmc_diagnostics._reduce_variance( + x, axis=axis, biased=biased, keepdims=keepdims) + np_var = np.var(x_, axis=axis, ddof=0 if biased else 1, keepdims=keepdims) + + if self.use_static_shape: + self.assertAllEqual(np_var.shape, var.shape) + + var_ = var.eval() + # We will mask below, which changes shape, so check shape explicitly here. + self.assertAllEqual(np_var.shape, var_.shape) + + # We get NaN when we divide by zero due to the size being the same as ddof + nan_mask = np.isnan(np_var) + if nan_mask.any(): + self.assertTrue(np.isnan(var_[nan_mask]).all()) + self.assertAllClose(np_var[~nan_mask], var_[~nan_mask], atol=0, rtol=0.02) + + def testScalarBiasedTrue(self): + self.check_versus_numpy(x_=-1.234, axis=None, biased=True, keepdims=False) + + def testScalarBiasedFalse(self): + # This should result in NaN. + self.check_versus_numpy(x_=-1.234, axis=None, biased=False, keepdims=False) + + def testShape2x3x4AxisNoneBiasedFalseKeepdimsFalse(self): + self.check_versus_numpy( + x_=rng.randn(2, 3, 4), axis=None, biased=True, keepdims=False) + + def testShape2x3x4Axis1BiasedFalseKeepdimsTrue(self): + self.check_versus_numpy( + x_=rng.randn(2, 3, 4), axis=1, biased=True, keepdims=True) + + def testShape2x3x4x5Axis13BiasedFalseKeepdimsTrue(self): + self.check_versus_numpy( + x_=rng.randn(2, 3, 4, 5), axis=1, biased=True, keepdims=True) + + def testShape2x3x4x5Axis13BiasedFalseKeepdimsFalse(self): + self.check_versus_numpy( + x_=rng.randn(2, 3, 4, 5), axis=1, biased=False, keepdims=False) + + +class ReduceVarianceTestStaticShape(test.TestCase, _ReduceVarianceTest): + + @property + def use_static_shape(self): + return True + + +class ReduceVarianceTestDynamicShape(test.TestCase, _ReduceVarianceTest): + + @property + def use_static_shape(self): + return False + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/variable_utils_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/variable_utils_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f978cf86417dc5ff5412a3eee584330a266e0964 --- /dev/null +++ b/tensorflow/contrib/bayesflow/python/kernel_tests/variable_utils_test.py @@ -0,0 +1,135 @@ +# 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 utility functions related to managing `tf.Variable`s.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import warnings + +import numpy as np + +from tensorflow.contrib.bayesflow.python.ops import variable_utils + +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.ops import variable_scope as varscope_ops +from tensorflow.python.ops import variables as variables_ops +from tensorflow.python.platform import test + + +def test_fn(x): + x = ops.convert_to_tensor(x, name="x") + dtype = x.dtype.as_numpy_dtype + s = x.shape.as_list() + z = varscope_ops.get_variable( + name="z", + dtype=dtype, + initializer=np.arange(np.prod(s)).reshape(s).astype(dtype)) + y = varscope_ops.get_variable( + name="y", + dtype=dtype, + initializer=np.arange(np.prod(s)).reshape(s).astype(dtype)**2) + return x + y + z + + +class _WrapCallableTest(object): + + def testDefaultArgsWorkCorrectly(self): + with self.test_session(): + x = constant_op.constant(self.dtype([0.1, 0.2])) + wrapped_fn, vars_args = variable_utils.externalize_variables_as_args( + test_fn, [x]) + + varscope_ops.get_variable_scope().reuse_variables() + + result = wrapped_fn(self.dtype(2), [3, 4, 5], 0.5) + + y_actual = varscope_ops.get_variable("y", dtype=self.dtype) + z_actual = varscope_ops.get_variable("z", dtype=self.dtype) + + variables_ops.global_variables_initializer().run() + result_ = result.eval() + + self.assertEqual(self.dtype, result_.dtype) + self.assertAllEqual([5.5, 6.5, 7.5], result_) + self.assertAllEqual([y_actual, z_actual], vars_args) + + def testNonDefaultArgsWorkCorrectly(self): + with self.test_session(): + x = constant_op.constant(self.dtype([0.1, 0.2])) + + _ = test_fn(self.dtype([0., 0.])) # Needed to create vars. + varscope_ops.get_variable_scope().reuse_variables() + + y_actual = varscope_ops.get_variable("y", dtype=self.dtype) + + wrapped_fn, vars_args = variable_utils.externalize_variables_as_args( + test_fn, [x], possible_ancestor_vars=[y_actual]) + + result = wrapped_fn(self.dtype([2, 3]), 0.5) # x, y + + variables_ops.global_variables_initializer().run() + result_ = result.eval() + + self.assertEqual(self.dtype, result_.dtype) + self.assertAllEqual([2.5, 4.5], result_) + self.assertAllEqual([y_actual], vars_args) + + def testWarnings(self): + with self.test_session(): + x = constant_op.constant(self.dtype([0.1, 0.2])) + wrapped_fn, _ = variable_utils.externalize_variables_as_args( + test_fn, [x], possible_ancestor_vars=[]) + varscope_ops.get_variable_scope().reuse_variables() + with warnings.catch_warnings(record=True) as w: + wrapped_fn(self.dtype(2)) + w = sorted(w, key=lambda w: str(w.message)) + self.assertEqual(2, len(w)) + self.assertRegexpMatches( + str(w[0].message), + r"Variable .* 'y:0' .* not found in bypass dict.") + self.assertRegexpMatches( + str(w[1].message), + r"Variable .* 'z:0' .* not found in bypass dict.") + + def testExceptions(self): + with self.test_session(): + x = constant_op.constant(self.dtype([0.1, 0.2])) + wrapped_fn, _ = variable_utils.externalize_variables_as_args( + test_fn, + [x], + possible_ancestor_vars=[], + assert_variable_override=True) + varscope_ops.get_variable_scope().reuse_variables() + with self.assertRaisesRegexp(ValueError, r"not found"): + wrapped_fn(self.dtype(2)) + + +class WrapCallableTest16(test.TestCase, _WrapCallableTest): + dtype = np.float16 + + +class WrapCallableTest32(test.TestCase, _WrapCallableTest): + dtype = np.float32 + + +class WrapCallableTest64(test.TestCase, _WrapCallableTest): + dtype = np.float64 + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py b/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py index ee3719232d8796c338247320fd8ef832a41df12b..d44fe6529a7ff0da0c6747e193fdb98a272a8da3 100644 --- a/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py +++ b/tensorflow/contrib/bayesflow/python/ops/custom_grad_impl.py @@ -31,8 +31,7 @@ __all__ = [ ] -def custom_gradient(fx, gx, x, axis=(), - fx_gx_manually_stopped=False, +def custom_gradient(fx, gx, x, axis=(), fx_gx_manually_stopped=False, name=None): """Enables specifying a custom gradient. @@ -43,7 +42,8 @@ def custom_gradient(fx, gx, x, axis=(), h(x) = x * stop_gradient(g(x)) + stop_gradient(f(x) - x * g(x)) ``` - is such that `h(x) = stop(f(x))` and `grad[h(x), x] = stop_gradient(g(x)).` + is such that `h(x) = stop_gradient(f(x))` and `grad[h(x), x] = + stop_gradient(g(x)).` In addition to scalar-domain/scalar-range functions, this function also supports tensor-domain/scalar-range functions. However, in the latter case it diff --git a/tensorflow/contrib/bayesflow/python/ops/hmc.py b/tensorflow/contrib/bayesflow/python/ops/hmc.py index 977d42fc16bb91777a76c45ac24f3c5dc587f5fe..7fd5652c5c3e085b23c05baef6e3a42b7a42e08f 100644 --- a/tensorflow/contrib/bayesflow/python/ops/hmc.py +++ b/tensorflow/contrib/bayesflow/python/ops/hmc.py @@ -12,8 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Hamiltonian Monte Carlo, a gradient-based MCMC algorithm. -""" +"""Hamiltonian Monte Carlo, a gradient-based MCMC algorithm.""" from __future__ import absolute_import from __future__ import division @@ -24,11 +23,9 @@ from tensorflow.contrib.bayesflow.python.ops.hmc_impl import * # pylint: disabl from tensorflow.python.util import all_util _allowed_symbols = [ - 'chain', - 'kernel', - 'leapfrog_integrator', - 'leapfrog_step', - 'ais_chain' + "sample_chain", + "sample_annealed_importance_chain", + "kernel", ] all_util.remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/bayesflow/python/ops/hmc_impl.py b/tensorflow/contrib/bayesflow/python/ops/hmc_impl.py index 5685a942e98800a39ec718adc67bcfd43aeafd52..f724910c59315867a42a56fab3deb36f5d3adb7a 100644 --- a/tensorflow/contrib/bayesflow/python/ops/hmc_impl.py +++ b/tensorflow/contrib/bayesflow/python/ops/hmc_impl.py @@ -14,17 +14,16 @@ # ============================================================================== """Hamiltonian Monte Carlo, a gradient-based MCMC algorithm. -@@chain -@@update -@@leapfrog_integrator -@@leapfrog_step -@@ais_chain +@@sample_chain +@@sample_annealed_importance_chain +@@kernel """ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections import numpy as np from tensorflow.python.framework import dtypes @@ -32,168 +31,326 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import functional_ops -from tensorflow.python.ops import gradients_impl +from tensorflow.python.ops import gradients_impl as gradients_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops -from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.ops.distributions import util as distributions_util __all__ = [ - 'chain', - 'kernel', - 'leapfrog_integrator', - 'leapfrog_step', - 'ais_chain' + "sample_chain", + "sample_annealed_importance_chain", + "kernel", ] -def _make_potential_and_grad(target_log_prob_fn): - def potential_and_grad(x): - log_prob_result = -target_log_prob_fn(x) - grad_result = gradients_impl.gradients(math_ops.reduce_sum(log_prob_result), - x)[0] - return log_prob_result, grad_result - return potential_and_grad - - -def chain(n_iterations, step_size, n_leapfrog_steps, initial_x, - target_log_prob_fn, event_dims=(), name=None): +KernelResults = collections.namedtuple( + "KernelResults", + [ + "acceptance_probs", + "current_grads_target_log_prob", # "Current result" means "accepted". + "current_target_log_prob", # "Current result" means "accepted". + "energy_change", + "is_accepted", + "proposed_grads_target_log_prob", + "proposed_state", + "proposed_target_log_prob", + "random_positive", + ]) + + +def _make_dummy_kernel_results( + dummy_state, + dummy_target_log_prob, + dummy_grads_target_log_prob): + return KernelResults( + acceptance_probs=dummy_target_log_prob, + current_grads_target_log_prob=dummy_grads_target_log_prob, + current_target_log_prob=dummy_target_log_prob, + energy_change=dummy_target_log_prob, + is_accepted=array_ops.ones_like(dummy_target_log_prob, dtypes.bool), + proposed_grads_target_log_prob=dummy_grads_target_log_prob, + proposed_state=dummy_state, + proposed_target_log_prob=dummy_target_log_prob, + random_positive=dummy_target_log_prob, + ) + + +def sample_chain( + num_results, + target_log_prob_fn, + current_state, + step_size, + num_leapfrog_steps, + num_burnin_steps=0, + num_steps_between_results=0, + seed=None, + current_target_log_prob=None, + current_grads_target_log_prob=None, + name=None): """Runs multiple iterations of one or more Hamiltonian Monte Carlo chains. - Hamiltonian Monte Carlo (HMC) is a Markov chain Monte Carlo (MCMC) - algorithm that takes a series of gradient-informed steps to produce - a Metropolis proposal. This function samples from an HMC Markov - chain whose initial state is `initial_x` and whose stationary - distribution has log-density `target_log_prob_fn()`. - - This function can update multiple chains in parallel. It assumes - that all dimensions of `initial_x` not specified in `event_dims` are - independent, and should therefore be updated independently. The - output of `target_log_prob_fn()` should sum log-probabilities across - all event dimensions. Slices along dimensions not in `event_dims` - may have different target distributions; this is up to + Hamiltonian Monte Carlo (HMC) is a Markov chain Monte Carlo (MCMC) algorithm + that takes a series of gradient-informed steps to produce a Metropolis + proposal. This function samples from an HMC Markov chain at `current_state` + and whose stationary distribution has log-unnormalized-density `target_log_prob_fn()`. - This function basically just wraps `hmc.kernel()` in a tf.scan() loop. + This function samples from multiple chains in parallel. It assumes that the + the leftmost dimensions of (each) `current_state` (part) index an independent + chain. The function `target_log_prob_fn()` sums log-probabilities across + event dimensions (i.e., current state (part) rightmost dimensions). Each + element of the output of `target_log_prob_fn()` represents the (possibly + unnormalized) log-probability of the joint distribution over (all) the current + state (parts). - Args: - n_iterations: Integer number of Markov chain updates to run. - step_size: Scalar step size or array of step sizes for the - leapfrog integrator. Broadcasts to the shape of - `initial_x`. Larger step sizes lead to faster progress, but - too-large step sizes make rejection exponentially more likely. - When possible, it's often helpful to match per-variable step - sizes to the standard deviations of the target distribution in - each variable. - n_leapfrog_steps: Integer number of steps to run the leapfrog - integrator for. Total progress per HMC step is roughly - proportional to step_size * n_leapfrog_steps. - initial_x: Tensor of initial state(s) of the Markov chain(s). - target_log_prob_fn: Python callable which takes an argument like `initial_x` - and returns its (possibly unnormalized) log-density under the target - distribution. - event_dims: List of dimensions that should not be treated as - independent. This allows for multiple chains to be run independently - in parallel. Default is (), i.e., all dimensions are independent. - name: Python `str` name prefixed to Ops created by this function. + The `current_state` can be represented as a single `Tensor` or a `list` of + `Tensors` which collectively represent the current state. When specifying a + `list`, one must also specify a list of `step_size`s. - Returns: - acceptance_probs: Tensor with the acceptance probabilities for each - iteration. Has shape matching `target_log_prob_fn(initial_x)`. - chain_states: Tensor with the state of the Markov chain at each iteration. - Has shape `[n_iterations, initial_x.shape[0],...,initial_x.shape[-1]`. + Note: `target_log_prob_fn` is called exactly twice. + + Only one out of every `num_steps_between_samples + 1` steps is included in the + returned results. This "thinning" comes at a cost of reduced statistical + power, while reducing memory requirements and autocorrelation. For more + discussion see [1]. + + [1]: "Statistically efficient thinning of a Markov chain sampler." + Art B. Owen. April 2017. + http://statweb.stanford.edu/~owen/reports/bestthinning.pdf #### Examples: - ```python - # Sampling from a standard normal (note `log_joint()` is unnormalized): - def log_joint(x): - return tf.reduce_sum(-0.5 * tf.square(x)) - chain, acceptance_probs = hmc.chain(1000, 0.5, 2, tf.zeros(10), log_joint, - event_dims=[0]) - # Discard first half of chain as warmup/burn-in - warmed_up = chain[500:] - mean_est = tf.reduce_mean(warmed_up, 0) - var_est = tf.reduce_mean(tf.square(warmed_up), 0) - tf.square(mean_est) - ``` + ##### Sample from a diagonal-variance Gaussian. ```python - # Sampling from a diagonal-variance Gaussian: - variances = tf.linspace(1., 3., 10) - def log_joint(x): - return tf.reduce_sum(-0.5 / variances * tf.square(x)) - chain, acceptance_probs = hmc.chain(1000, 0.5, 2, tf.zeros(10), log_joint, - event_dims=[0]) - # Discard first half of chain as warmup/burn-in - warmed_up = chain[500:] - mean_est = tf.reduce_mean(warmed_up, 0) - var_est = tf.reduce_mean(tf.square(warmed_up), 0) - tf.square(mean_est) + tfd = tf.contrib.distributions + + def make_likelihood(true_variances): + return tfd.MultivariateNormalDiag( + scale_diag=tf.sqrt(true_variances)) + + dims = 10 + dtype = np.float32 + true_variances = tf.linspace(dtype(1), dtype(3), dims) + likelihood = make_likelihood(true_variances) + + states, kernel_results = hmc.sample_chain( + num_results=1000, + target_log_prob_fn=likelihood.log_prob, + current_state=tf.zeros(dims), + step_size=0.5, + num_leapfrog_steps=2, + num_burnin_steps=500) + + # Compute sample stats. + sample_mean = tf.reduce_mean(states, axis=0) + sample_var = tf.reduce_mean( + tf.squared_difference(states, sample_mean), + axis=0) ``` - ```python - # Sampling from factor-analysis posteriors with known factors W: - # mu[i, j] ~ Normal(0, 1) - # x[i] ~ Normal(matmul(mu[i], W), I) - def log_joint(mu, x, W): - prior = -0.5 * tf.reduce_sum(tf.square(mu), 1) - x_mean = tf.matmul(mu, W) - likelihood = -0.5 * tf.reduce_sum(tf.square(x - x_mean), 1) - return prior + likelihood - chain, acceptance_probs = hmc.chain(1000, 0.1, 2, - tf.zeros([x.shape[0], W.shape[0]]), - lambda mu: log_joint(mu, x, W), - event_dims=[1]) - # Discard first half of chain as warmup/burn-in - warmed_up = chain[500:] - mean_est = tf.reduce_mean(warmed_up, 0) - var_est = tf.reduce_mean(tf.square(warmed_up), 0) - tf.square(mean_est) + ##### Sampling from factor-analysis posteriors with known factors. + + I.e., + + ```none + for i=1..n: + w[i] ~ Normal(0, eye(d)) # prior + x[i] ~ Normal(loc=matmul(w[i], F)) # likelihood ``` + where `F` denotes factors. + ```python - # Sampling from the posterior of a Bayesian regression model.: - - # Run 100 chains in parallel, each with a different initialization. - initial_beta = tf.random_normal([100, x.shape[1]]) - chain, acceptance_probs = hmc.chain(1000, 0.1, 10, initial_beta, - log_joint_partial, event_dims=[1]) - # Discard first halves of chains as warmup/burn-in - warmed_up = chain[500:] - # Averaging across samples within a chain and across chains - mean_est = tf.reduce_mean(warmed_up, [0, 1]) - var_est = tf.reduce_mean(tf.square(warmed_up), [0, 1]) - tf.square(mean_est) + tfd = tf.contrib.distributions + + def make_prior(dims, dtype): + return tfd.MultivariateNormalDiag( + loc=tf.zeros(dims, dtype)) + + def make_likelihood(weights, factors): + return tfd.MultivariateNormalDiag( + loc=tf.tensordot(weights, factors, axes=[[0], [-1]])) + + # Setup data. + num_weights = 10 + num_factors = 4 + num_chains = 100 + dtype = np.float32 + + prior = make_prior(num_weights, dtype) + weights = prior.sample(num_chains) + factors = np.random.randn(num_factors, num_weights).astype(dtype) + x = make_likelihood(weights, factors).sample(num_chains) + + def target_log_prob(w): + # Target joint is: `f(w) = p(w, x | factors)`. + return prior.log_prob(w) + make_likelihood(w, factors).log_prob(x) + + # Get `num_results` samples from `num_chains` independent chains. + chains_states, kernels_results = hmc.sample_chain( + num_results=1000, + target_log_prob_fn=target_log_prob, + current_state=tf.zeros([num_chains, dims], dtype), + step_size=0.1, + num_leapfrog_steps=2, + num_burnin_steps=500) + + # Compute sample stats. + sample_mean = tf.reduce_mean(chains_states, axis=[0, 1]) + sample_var = tf.reduce_mean( + tf.squared_difference(chains_states, sample_mean), + axis=[0, 1]) ``` + + Args: + num_results: Integer number of Markov chain draws. + target_log_prob_fn: Python callable which takes an argument like + `current_state` (or `*current_state` if it's a list) and returns its + (possibly unnormalized) log-density under the target distribution. + current_state: `Tensor` or Python `list` of `Tensor`s representing the + current state(s) of the Markov chain(s). The first `r` dimensions index + independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`. + step_size: `Tensor` or Python `list` of `Tensor`s representing the step size + for the leapfrog integrator. Must broadcast with the shape of + `current_state`. Larger step sizes lead to faster progress, but too-large + step sizes make rejection exponentially more likely. When possible, it's + often helpful to match per-variable step sizes to the standard deviations + of the target distribution in each variable. + num_leapfrog_steps: Integer number of steps to run the leapfrog integrator + for. Total progress per HMC step is roughly proportional to `step_size * + num_leapfrog_steps`. + num_burnin_steps: Integer number of chain steps to take before starting to + collect results. + Default value: 0 (i.e., no burn-in). + num_steps_between_results: Integer number of chain steps between collecting + a result. Only one out of every `num_steps_between_samples + 1` steps is + included in the returned results. This "thinning" comes at a cost of + reduced statistical power, while reducing memory requirements and + autocorrelation. For more discussion see [1]. + Default value: 0 (i.e., no subsampling). + seed: Python integer to seed the random number generator. + current_target_log_prob: (Optional) `Tensor` representing the value of + `target_log_prob_fn` at the `current_state`. The only reason to specify + this argument is to reduce TF graph size. + Default value: `None` (i.e., compute as needed). + current_grads_target_log_prob: (Optional) Python list of `Tensor`s + representing gradient of `target_log_prob` at the `current_state` and wrt + the `current_state`. Must have same shape as `current_state`. The only + reason to specify this argument is to reduce TF graph size. + Default value: `None` (i.e., compute as needed). + name: Python `str` name prefixed to Ops created by this function. + Default value: `None` (i.e., "hmc_sample_chain"). + + Returns: + accepted_states: Tensor or Python list of `Tensor`s representing the + state(s) of the Markov chain(s) at each result step. Has same shape as + input `current_state` but with a prepended `num_results`-size dimension. + kernel_results: `collections.namedtuple` of internal calculations used to + advance the chain. """ - with ops.name_scope(name, 'hmc_chain', [n_iterations, step_size, - n_leapfrog_steps, initial_x]): - initial_x = ops.convert_to_tensor(initial_x, name='initial_x') - non_event_shape = array_ops.shape(target_log_prob_fn(initial_x)) - - def body(a, _): - updated_x, acceptance_probs, log_prob, grad = kernel( - step_size, n_leapfrog_steps, a[0], target_log_prob_fn, event_dims, - a[2], a[3]) - return updated_x, acceptance_probs, log_prob, grad - - potential_and_grad = _make_potential_and_grad(target_log_prob_fn) - potential, grad = potential_and_grad(initial_x) - return functional_ops.scan( - body, array_ops.zeros(n_iterations, dtype=initial_x.dtype), - (initial_x, - array_ops.zeros(non_event_shape, dtype=initial_x.dtype), - -potential, -grad))[:2] - - -def ais_chain(n_iterations, step_size, n_leapfrog_steps, initial_x, - target_log_prob_fn, proposal_log_prob_fn, event_dims=(), - name=None): + with ops.name_scope( + name, "hmc_sample_chain", + [num_results, current_state, step_size, num_leapfrog_steps, + num_burnin_steps, num_steps_between_results, seed, + current_target_log_prob, current_grads_target_log_prob]): + with ops.name_scope("initialize"): + [ + current_state, + step_size, + current_target_log_prob, + current_grads_target_log_prob, + ] = _prepare_args( + target_log_prob_fn, + current_state, + step_size, + current_target_log_prob, + current_grads_target_log_prob) + num_results = ops.convert_to_tensor( + num_results, + dtype=dtypes.int32, + name="num_results") + num_leapfrog_steps = ops.convert_to_tensor( + num_leapfrog_steps, + dtype=dtypes.int32, + name="num_leapfrog_steps") + num_burnin_steps = ops.convert_to_tensor( + num_burnin_steps, + dtype=dtypes.int32, + name="num_burnin_steps") + num_steps_between_results = ops.convert_to_tensor( + num_steps_between_results, + dtype=dtypes.int32, + name="num_steps_between_results") + + def _run_chain(num_steps, current_state, kernel_results): + """Runs the chain(s) for `num_steps`.""" + def _loop_body(iter_, current_state, kernel_results): + return [iter_ + 1] + list(kernel( + target_log_prob_fn, + current_state, + step_size, + num_leapfrog_steps, + seed, + kernel_results.current_target_log_prob, + kernel_results.current_grads_target_log_prob)) + while_loop_kwargs = dict( + cond=lambda iter_, *args: iter_ < num_steps, + body=_loop_body, + loop_vars=[ + np.int32(0), + current_state, + kernel_results, + ], + ) + if seed is not None: + while_loop_kwargs["parallel_iterations"] = 1 + return control_flow_ops.while_loop( + **while_loop_kwargs)[1:] # Lop-off "iter_". + + def _scan_body(args_list, iter_): + """Closure which implements `tf.scan` body.""" + current_state, kernel_results = args_list + return _run_chain( + 1 + array_ops.where(math_ops.equal(iter_, 0), + num_burnin_steps, + num_steps_between_results), + current_state, + kernel_results) + + scan_kwargs = dict( + fn=_scan_body, + elems=math_ops.range(num_results), # iter_: used to choose burnin. + initializer=[ + current_state, + _make_dummy_kernel_results( + current_state, + current_target_log_prob, + current_grads_target_log_prob), + ]) + if seed is not None: + scan_kwargs["parallel_iterations"] = 1 + return functional_ops.scan(**scan_kwargs) + + +def sample_annealed_importance_chain( + proposal_log_prob_fn, + num_steps, + target_log_prob_fn, + current_state, + step_size, + num_leapfrog_steps, + seed=None, + name=None): """Runs annealed importance sampling (AIS) to estimate normalizing constants. - This routine uses Hamiltonian Monte Carlo to sample from a series of + This function uses Hamiltonian Monte Carlo to sample from a series of distributions that slowly interpolates between an initial "proposal" - distribution + distribution: `exp(proposal_log_prob_fn(x) - proposal_log_normalizer)` - and the target distribution + and the target distribution: `exp(target_log_prob_fn(x) - target_log_normalizer)`, @@ -202,113 +359,203 @@ def ais_chain(n_iterations, step_size, n_leapfrog_steps, initial_x, normalizing constants of the initial distribution and the target distribution: - E[exp(w)] = exp(target_log_normalizer - proposal_log_normalizer). + `E[exp(ais_weights)] = exp(target_log_normalizer - proposal_log_normalizer)`. - Args: - n_iterations: Integer number of Markov chain updates to run. More - iterations means more expense, but smoother annealing between q - and p, which in turn means exponentially lower variance for the - normalizing constant estimator. - step_size: Scalar step size or array of step sizes for the - leapfrog integrator. Broadcasts to the shape of - `initial_x`. Larger step sizes lead to faster progress, but - too-large step sizes make rejection exponentially more likely. - When possible, it's often helpful to match per-variable step - sizes to the standard deviations of the target distribution in - each variable. - n_leapfrog_steps: Integer number of steps to run the leapfrog - integrator for. Total progress per HMC step is roughly - proportional to step_size * n_leapfrog_steps. - initial_x: Tensor of initial state(s) of the Markov chain(s). Must - be a sample from q, or results will be incorrect. - target_log_prob_fn: Python callable which takes an argument like `initial_x` - and returns its (possibly unnormalized) log-density under the target - distribution. - proposal_log_prob_fn: Python callable that returns the log density of the - initial distribution. - event_dims: List of dimensions that should not be treated as - independent. This allows for multiple chains to be run independently - in parallel. Default is (), i.e., all dimensions are independent. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - ais_weights: Tensor with the estimated weight(s). Has shape matching - `target_log_prob_fn(initial_x)`. - chain_states: Tensor with the state(s) of the Markov chain(s) the final - iteration. Has shape matching `initial_x`. - acceptance_probs: Tensor with the acceptance probabilities for the final - iteration. Has shape matching `target_log_prob_fn(initial_x)`. + Note: `proposal_log_prob_fn` and `target_log_prob_fn` are called exactly three + times (although this may be reduced to two times, in the future). #### Examples: + ##### Estimate the normalizing constant of a log-gamma distribution. + ```python - # Estimating the normalizing constant of a log-gamma distribution: - def proposal_log_prob(x): - # Standard normal log-probability. This is properly normalized. - return tf.reduce_sum(-0.5 * tf.square(x) - 0.5 * np.log(2 * np.pi), 1) - def target_log_prob(x): - # Unnormalized log-gamma(2, 3) distribution. - # True normalizer is (lgamma(2) - 2 * log(3)) * x.shape[1] - return tf.reduce_sum(2. * x - 3. * tf.exp(x), 1) + tfd = tf.contrib.distributions + # Run 100 AIS chains in parallel - initial_x = tf.random_normal([100, 20]) - w, _, _ = hmc.ais_chain(1000, 0.2, 2, initial_x, target_log_prob, - proposal_log_prob, event_dims=[1]) - log_normalizer_estimate = tf.reduce_logsumexp(w) - np.log(100) + num_chains = 100 + dims = 20 + dtype = np.float32 + + proposal = tfd.MultivatiateNormalDiag( + loc=tf.zeros([dims], dtype=dtype)) + + target = tfd.TransformedDistribution( + distribution=tfd.Gamma(concentration=dtype(2), + rate=dtype(3)), + bijector=tfd.bijectors.Invert(tfd.bijectors.Exp()), + event_shape=[dims]) + + chains_state, ais_weights, kernels_results = ( + hmc.sample_annealed_importance_chain( + proposal_log_prob_fn=proposal.log_prob, + num_steps=1000, + target_log_prob_fn=target.log_prob, + step_size=0.2, + current_state=proposal.sample(num_chains), + num_leapfrog_steps=2)) + + log_estimated_normalizer = (tf.reduce_logsumexp(ais_weights) + - np.log(num_chains)) + log_true_normalizer = tf.lgamma(2.) - 2. * tf.log(3.) ``` + ##### Estimate marginal likelihood of a Bayesian regression model. + ```python - # Estimating the marginal likelihood of a Bayesian regression model: - base_measure = -0.5 * np.log(2 * np.pi) - def proposal_log_prob(x): - # Standard normal log-probability. This is properly normalized. - return tf.reduce_sum(-0.5 * tf.square(x) + base_measure, 1) - def regression_log_joint(beta, x, y): - # This function returns a vector whose ith element is log p(beta[i], y | x). - # Each row of beta corresponds to the state of an independent Markov chain. - log_prior = tf.reduce_sum(-0.5 * tf.square(beta) + base_measure, 1) - means = tf.matmul(beta, x, transpose_b=True) - log_likelihood = tf.reduce_sum(-0.5 * tf.square(y - means) + - base_measure, 1) - return log_prior + log_likelihood - def log_joint_partial(beta): - return regression_log_joint(beta, x, y) + tfd = tf.contrib.distributions + + def make_prior(dims, dtype): + return tfd.MultivariateNormalDiag( + loc=tf.zeros(dims, dtype)) + + def make_likelihood(weights, x): + return tfd.MultivariateNormalDiag( + loc=tf.tensordot(weights, x, axes=[[0], [-1]])) + # Run 100 AIS chains in parallel - initial_beta = tf.random_normal([100, x.shape[1]]) - w, beta_samples, _ = hmc.ais_chain(1000, 0.1, 2, initial_beta, - log_joint_partial, proposal_log_prob, - event_dims=[1]) - log_normalizer_estimate = tf.reduce_logsumexp(w) - np.log(100) + num_chains = 100 + dims = 10 + dtype = np.float32 + + # Make training data. + x = np.random.randn(num_chains, dims).astype(dtype) + true_weights = np.random.randn(dims).astype(dtype) + y = np.dot(x, true_weights) + np.random.randn(num_chains) + + # Setup model. + prior = make_prior(dims, dtype) + def target_log_prob_fn(weights): + return prior.log_prob(weights) + make_likelihood(weights, x).log_prob(y) + + proposal = tfd.MultivariateNormalDiag( + loc=tf.zeros(dims, dtype)) + + weight_samples, ais_weights, kernel_results = ( + hmc.sample_annealed_importance_chain( + num_steps=1000, + proposal_log_prob_fn=proposal.log_prob, + target_log_prob_fn=target_log_prob_fn + current_state=tf.zeros([num_chains, dims], dtype), + step_size=0.1, + num_leapfrog_steps=2)) + log_normalizer_estimate = (tf.reduce_logsumexp(ais_weights) + - np.log(num_chains)) ``` + + Args: + proposal_log_prob_fn: Python callable that returns the log density of the + initial distribution. + num_steps: Integer number of Markov chain updates to run. More + iterations means more expense, but smoother annealing between q + and p, which in turn means exponentially lower variance for the + normalizing constant estimator. + target_log_prob_fn: Python callable which takes an argument like + `current_state` (or `*current_state` if it's a list) and returns its + (possibly unnormalized) log-density under the target distribution. + current_state: `Tensor` or Python `list` of `Tensor`s representing the + current state(s) of the Markov chain(s). The first `r` dimensions index + independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`. + step_size: `Tensor` or Python `list` of `Tensor`s representing the step size + for the leapfrog integrator. Must broadcast with the shape of + `current_state`. Larger step sizes lead to faster progress, but too-large + step sizes make rejection exponentially more likely. When possible, it's + often helpful to match per-variable step sizes to the standard deviations + of the target distribution in each variable. + num_leapfrog_steps: Integer number of steps to run the leapfrog integrator + for. Total progress per HMC step is roughly proportional to `step_size * + num_leapfrog_steps`. + seed: Python integer to seed the random number generator. + name: Python `str` name prefixed to Ops created by this function. + Default value: `None` (i.e., "hmc_sample_annealed_importance_chain"). + + Returns: + accepted_state: `Tensor` or Python list of `Tensor`s representing the + state(s) of the Markov chain(s) at the final iteration. Has same shape as + input `current_state`. + ais_weights: Tensor with the estimated weight(s). Has shape matching + `target_log_prob_fn(current_state)`. + kernel_results: `collections.namedtuple` of internal calculations used to + advance the chain. """ - with ops.name_scope(name, 'hmc_ais_chain', - [n_iterations, step_size, n_leapfrog_steps, initial_x]): - non_event_shape = array_ops.shape(target_log_prob_fn(initial_x)) - - beta_series = math_ops.linspace(0., 1., n_iterations+1)[1:] - def _body(a, beta): # pylint: disable=missing-docstring - def log_prob_beta(x): - return ((1 - beta) * proposal_log_prob_fn(x) + - beta * target_log_prob_fn(x)) - last_x = a[0] - w = a[2] - w += (1. / n_iterations) * (target_log_prob_fn(last_x) - - proposal_log_prob_fn(last_x)) - # TODO(b/66917083): There's an opportunity for gradient reuse here. - updated_x, acceptance_probs, _, _ = kernel(step_size, n_leapfrog_steps, - last_x, log_prob_beta, - event_dims) - return updated_x, acceptance_probs, w - - x, acceptance_probs, w = functional_ops.scan( - _body, beta_series, - (initial_x, array_ops.zeros(non_event_shape, dtype=initial_x.dtype), - array_ops.zeros(non_event_shape, dtype=initial_x.dtype))) - return w[-1], x[-1], acceptance_probs[-1] - - -def kernel(step_size, n_leapfrog_steps, x, target_log_prob_fn, event_dims=(), - x_log_prob=None, x_grad=None, name=None): + def make_convex_combined_log_prob_fn(iter_): + def _fn(*args): + p = proposal_log_prob_fn(*args) + t = target_log_prob_fn(*args) + dtype = p.dtype.base_dtype + beta = (math_ops.cast(iter_ + 1, dtype) + / math_ops.cast(num_steps, dtype)) + return (1. - beta) * p + beta * t + return _fn + + with ops.name_scope( + name, "hmc_sample_annealed_importance_chain", + [num_steps, current_state, step_size, num_leapfrog_steps, seed]): + with ops.name_scope("initialize"): + [ + current_state, + step_size, + current_log_prob, + current_grads_log_prob, + ] = _prepare_args( + make_convex_combined_log_prob_fn(iter_=0), + current_state, + step_size, + description="convex_combined_log_prob") + num_steps = ops.convert_to_tensor( + num_steps, + dtype=dtypes.int32, + name="num_steps") + num_leapfrog_steps = ops.convert_to_tensor( + num_leapfrog_steps, + dtype=dtypes.int32, + name="num_leapfrog_steps") + def _loop_body(iter_, ais_weights, current_state, kernel_results): + """Closure which implements `tf.while_loop` body.""" + current_state_parts = (list(current_state) + if _is_list_like(current_state) + else [current_state]) + # TODO(b/72994218): Consider refactoring things to avoid this unecessary + # call. + ais_weights += ((target_log_prob_fn(*current_state_parts) + - proposal_log_prob_fn(*current_state_parts)) + / math_ops.cast(num_steps, ais_weights.dtype)) + return [iter_ + 1, ais_weights] + list(kernel( + make_convex_combined_log_prob_fn(iter_), + current_state, + step_size, + num_leapfrog_steps, + seed, + kernel_results.current_target_log_prob, + kernel_results.current_grads_target_log_prob)) + + while_loop_kwargs = dict( + cond=lambda iter_, *args: iter_ < num_steps, + body=_loop_body, + loop_vars=[ + np.int32(0), # iter_ + array_ops.zeros_like(current_log_prob), # ais_weights + current_state, + _make_dummy_kernel_results(current_state, + current_log_prob, + current_grads_log_prob), + ]) + if seed is not None: + while_loop_kwargs["parallel_iterations"] = 1 + + [ais_weights, current_state, kernel_results] = control_flow_ops.while_loop( + **while_loop_kwargs)[1:] # Lop-off "iter_". + + return [current_state, ais_weights, kernel_results] + + +def kernel(target_log_prob_fn, + current_state, + step_size, + num_leapfrog_steps, + seed=None, + current_target_log_prob=None, + current_grads_target_log_prob=None, + name=None): """Runs one iteration of Hamiltonian Monte Carlo. Hamiltonian Monte Carlo (HMC) is a Markov chain Monte Carlo (MCMC) @@ -316,334 +563,623 @@ def kernel(step_size, n_leapfrog_steps, x, target_log_prob_fn, event_dims=(), a Metropolis proposal. This function applies one step of HMC to randomly update the variable `x`. - This function can update multiple chains in parallel. It assumes - that all dimensions of `x` not specified in `event_dims` are - independent, and should therefore be updated independently. The - output of `target_log_prob_fn()` should sum log-probabilities across - all event dimensions. Slices along dimensions not in `event_dims` - may have different target distributions; for example, if - `event_dims == (1,)`, then `x[0, :]` could have a different target - distribution from x[1, :]. This is up to `target_log_prob_fn()`. - - Args: - step_size: Scalar step size or array of step sizes for the - leapfrog integrator. Broadcasts to the shape of - `x`. Larger step sizes lead to faster progress, but - too-large step sizes make rejection exponentially more likely. - When possible, it's often helpful to match per-variable step - sizes to the standard deviations of the target distribution in - each variable. - n_leapfrog_steps: Integer number of steps to run the leapfrog - integrator for. Total progress per HMC step is roughly - proportional to step_size * n_leapfrog_steps. - x: Tensor containing the value(s) of the random variable(s) to update. - target_log_prob_fn: Python callable which takes an argument like `initial_x` - and returns its (possibly unnormalized) log-density under the target - distribution. - event_dims: List of dimensions that should not be treated as - independent. This allows for multiple chains to be run independently - in parallel. Default is (), i.e., all dimensions are independent. - x_log_prob (optional): Tensor containing the cached output of a previous - call to `target_log_prob_fn()` evaluated at `x` (such as that provided by - a previous call to `kernel()`). Providing `x_log_prob` and - `x_grad` saves one gradient computation per call to `kernel()`. - x_grad (optional): Tensor containing the cached gradient of - `target_log_prob_fn()` evaluated at `x` (such as that provided by - a previous call to `kernel()`). Providing `x_log_prob` and - `x_grad` saves one gradient computation per call to `kernel()`. - name: Python `str` name prefixed to Ops created by this function. - - Returns: - updated_x: The updated variable(s) x. Has shape matching `initial_x`. - acceptance_probs: Tensor with the acceptance probabilities for the final - iteration. This is useful for diagnosing step size problems etc. Has - shape matching `target_log_prob_fn(initial_x)`. - new_log_prob: The value of `target_log_prob_fn()` evaluated at `updated_x`. - new_grad: The value of the gradient of `target_log_prob_fn()` evaluated at - `updated_x`. + This function can update multiple chains in parallel. It assumes that all + leftmost dimensions of `current_state` index independent chain states (and are + therefore updated independently). The output of `target_log_prob_fn()` should + sum log-probabilities across all event dimensions. Slices along the rightmost + dimensions may have different target distributions; for example, + `current_state[0, :]` could have a different target distribution from + `current_state[1, :]`. This is up to `target_log_prob_fn()`. (The number of + independent chains is `tf.size(target_log_prob_fn(*current_state))`.) #### Examples: + ##### Simple chain with warm-up. + ```python + tfd = tf.contrib.distributions + # Tuning acceptance rates: + dtype = np.float32 target_accept_rate = 0.631 - def target_log_prob(x): - # Standard normal - return tf.reduce_sum(-0.5 * tf.square(x)) - initial_x = tf.zeros([10]) - initial_log_prob = target_log_prob(initial_x) - initial_grad = tf.gradients(initial_log_prob, initial_x)[0] - # Algorithm state - x = tf.Variable(initial_x, name='x') - step_size = tf.Variable(1., name='step_size') - last_log_prob = tf.Variable(initial_log_prob, name='last_log_prob') - last_grad = tf.Variable(initial_grad, name='last_grad') - # Compute updates - new_x, acceptance_prob, log_prob, grad = hmc.kernel(step_size, 3, x, - target_log_prob, - event_dims=[0], - x_log_prob=last_log_prob) - x_update = tf.assign(x, new_x) - log_prob_update = tf.assign(last_log_prob, log_prob) - grad_update = tf.assign(last_grad, grad) - step_size_update = tf.assign(step_size, - tf.where(acceptance_prob > target_accept_rate, - step_size * 1.01, step_size / 1.01)) - adaptive_updates = [x_update, log_prob_update, grad_update, step_size_update] - sampling_updates = [x_update, log_prob_update, grad_update] - - sess = tf.Session() - sess.run(tf.global_variables_initializer()) + num_warmup_iter = 500 + num_chain_iter = 500 + + x = tf.get_variable(name="x", initializer=dtype(1)) + step_size = tf.get_variable(name="step_size", initializer=dtype(1)) + + target = tfd.Normal(loc=dtype(0), scale=dtype(1)) + + new_x, other_results = hmc.kernel( + target_log_prob_fn=target.log_prob, + current_state=x, + step_size=step_size, + num_leapfrog_steps=3)[:4] + + x_update = x.assign(new_x) + + step_size_update = step_size.assign_add( + step_size * tf.where( + other_results.acceptance_probs > target_accept_rate, + 0.01, -0.01)) + + warmup = tf.group([x_update, step_size_update]) + + tf.global_variables_initializer().run() + + sess.graph.finalize() # No more graph building. + # Warm up the sampler and adapt the step size - for i in xrange(500): - sess.run(adaptive_updates) - # Collect samples without adapting step size - samples = np.zeros([500, 10]) - for i in xrange(500): - x_val, _ = sess.run([new_x, sampling_updates]) - samples[i] = x_val - ``` + for _ in xrange(num_warmup_iter): + sess.run(warmup) - ```python - # Empirical-Bayes estimation of a hyperparameter by MCMC-EM: - - # Problem setup - N = 150 - D = 10 - x = np.random.randn(N, D).astype(np.float32) - true_sigma = 0.5 - true_beta = true_sigma * np.random.randn(D).astype(np.float32) - y = x.dot(true_beta) + np.random.randn(N).astype(np.float32) - - def log_prior(beta, log_sigma): - return tf.reduce_sum(-0.5 / tf.exp(2 * log_sigma) * tf.square(beta) - - log_sigma) - def regression_log_joint(beta, log_sigma, x, y): - # This function returns log p(beta | log_sigma) + log p(y | x, beta). - means = tf.matmul(tf.expand_dims(beta, 0), x, transpose_b=True) - means = tf.squeeze(means) - log_likelihood = tf.reduce_sum(-0.5 * tf.square(y - means)) - return log_prior(beta, log_sigma) + log_likelihood - def log_joint_partial(beta): - return regression_log_joint(beta, log_sigma, x, y) - # Our estimate of log(sigma) - log_sigma = tf.Variable(0., name='log_sigma') - # The state of the Markov chain - beta = tf.Variable(tf.random_normal([x.shape[1]]), name='beta') - new_beta, _, _, _ = hmc.kernel(0.1, 5, beta, log_joint_partial, - event_dims=[0]) - beta_update = tf.assign(beta, new_beta) - optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01) - with tf.control_dependencies([beta_update]): - log_sigma_update = optimizer.minimize(-log_prior(beta, log_sigma), - var_list=[log_sigma]) - - sess = tf.Session() - sess.run(tf.global_variables_initializer()) - log_sigma_history = np.zeros(1000) - for i in xrange(1000): - log_sigma_val, _ = sess.run([log_sigma, log_sigma_update]) - log_sigma_history[i] = log_sigma_val - # Should converge to something close to true_sigma - plt.plot(np.exp(log_sigma_history)) + # Collect samples without adapting step size + samples = np.zeros([num_chain_iter]) + for i in xrange(num_chain_iter): + _, x_, target_log_prob_, grad_ = sess.run([ + x_update, + x, + other_results.target_log_prob, + other_results.grads_target_log_prob]) + samples[i] = x_ + + print(samples.mean(), samples.std()) ``` - """ - with ops.name_scope(name, 'hmc_kernel', [step_size, n_leapfrog_steps, x]): - potential_and_grad = _make_potential_and_grad(target_log_prob_fn) - x = ops.convert_to_tensor(x, name='x') - - x_shape = array_ops.shape(x) - m = random_ops.random_normal(x_shape, dtype=x.dtype) - - kinetic_0 = 0.5 * math_ops.reduce_sum(math_ops.square(m), event_dims) - - if (x_log_prob is not None) and (x_grad is not None): - log_potential_0, grad_0 = -x_log_prob, -x_grad # pylint: disable=invalid-unary-operand-type - else: - if x_log_prob is not None: - logging.warn('x_log_prob was provided, but x_grad was not,' - ' so x_log_prob was not used.') - if x_grad is not None: - logging.warn('x_grad was provided, but x_log_prob was not,' - ' so x_grad was not used.') - log_potential_0, grad_0 = potential_and_grad(x) - - new_x, new_m, log_potential_1, grad_1 = leapfrog_integrator( - step_size, n_leapfrog_steps, x, m, potential_and_grad, grad_0) - - kinetic_1 = 0.5 * math_ops.reduce_sum(math_ops.square(new_m), event_dims) - - energy_change = log_potential_1 - log_potential_0 + kinetic_1 - kinetic_0 - # Treat NaN as infinite energy (and therefore guaranteed rejection). - energy_change = array_ops.where( - math_ops.is_nan(energy_change), - array_ops.fill(array_ops.shape(energy_change), - energy_change.dtype.as_numpy_dtype(np.inf)), - energy_change) - acceptance_probs = math_ops.exp(math_ops.minimum(-energy_change, 0.)) - accepted = ( - random_ops.random_uniform( - array_ops.shape(acceptance_probs), dtype=x.dtype) - < acceptance_probs) - new_log_prob = -array_ops.where(accepted, log_potential_1, log_potential_0) - - # TODO(b/65738010): This should work, but it doesn't for now. - # reduced_shape = math_ops.reduced_shape(x_shape, event_dims) - reduced_shape = array_ops.shape(math_ops.reduce_sum(x, event_dims, - keep_dims=True)) - accepted = array_ops.reshape(accepted, reduced_shape) - accepted = math_ops.logical_or( - accepted, math_ops.cast(array_ops.zeros_like(x), dtypes.bool)) - new_x = array_ops.where(accepted, new_x, x) - new_grad = -array_ops.where(accepted, grad_1, grad_0) - - # TODO(langmore) Gradients of acceptance_probs and new_log_prob with respect - # to initial_x will propagate NaNs (see testNanFromGradsDontPropagate). This - # should be fixed. - return new_x, acceptance_probs, new_log_prob, new_grad - - -def leapfrog_integrator(step_size, n_steps, initial_position, initial_momentum, - potential_and_grad, initial_grad, name=None): - """Applies `n_steps` steps of the leapfrog integrator. - - This just wraps `leapfrog_step()` in a `tf.while_loop()`, reusing - gradient computations where possible. - Args: - step_size: Scalar step size or array of step sizes for the - leapfrog integrator. Broadcasts to the shape of - `initial_position`. Larger step sizes lead to faster progress, but - too-large step sizes lead to larger discretization error and - worse energy conservation. - n_steps: Number of steps to run the leapfrog integrator. - initial_position: Tensor containing the value(s) of the position variable(s) - to update. - initial_momentum: Tensor containing the value(s) of the momentum variable(s) - to update. - potential_and_grad: Python callable that takes a position tensor like - `initial_position` and returns the potential energy and its gradient at - that position. - initial_grad: Tensor with the value of the gradient of the potential energy - at `initial_position`. - name: Python `str` name prefixed to Ops created by this function. + ##### Sample from more complicated posterior. - Returns: - updated_position: Updated value of the position. - updated_momentum: Updated value of the momentum. - new_potential: Potential energy of the new position. Has shape matching - `potential_and_grad(initial_position)`. - new_grad: Gradient from potential_and_grad() evaluated at the new position. - Has shape matching `initial_position`. + I.e., - Example: Simple quadratic potential. + ```none + W ~ MVN(loc=0, scale=sigma * eye(dims)) + for i=1...num_samples: + X[i] ~ MVN(loc=0, scale=eye(dims)) + eps[i] ~ Normal(loc=0, scale=1) + Y[i] = X[i].T * W + eps[i] + ``` ```python - def potential_and_grad(position): - return tf.reduce_sum(0.5 * tf.square(position)), position - position = tf.placeholder(np.float32) - momentum = tf.placeholder(np.float32) - potential, grad = potential_and_grad(position) - new_position, new_momentum, new_potential, new_grad = hmc.leapfrog_integrator( - 0.1, 3, position, momentum, potential_and_grad, grad) - - sess = tf.Session() - position_val = np.random.randn(10) - momentum_val = np.random.randn(10) - potential_val, grad_val = sess.run([potential, grad], - {position: position_val}) - positions = np.zeros([100, 10]) - for i in xrange(100): - position_val, momentum_val, potential_val, grad_val = sess.run( - [new_position, new_momentum, new_potential, new_grad], - {position: position_val, momentum: momentum_val}) - positions[i] = position_val - # Should trace out sinusoidal dynamics. - plt.plot(positions[:, 0]) - ``` - """ - def leapfrog_wrapper(step_size, x, m, grad, l): - x, m, _, grad = leapfrog_step(step_size, x, m, potential_and_grad, grad) - return step_size, x, m, grad, l + 1 + tfd = tf.contrib.distributions + + def make_training_data(num_samples, dims, sigma): + dt = np.asarray(sigma).dtype + zeros = tf.zeros(dims, dtype=dt) + x = tfd.MultivariateNormalDiag( + loc=zeros).sample(num_samples, seed=1) + w = tfd.MultivariateNormalDiag( + loc=zeros, + scale_identity_multiplier=sigma).sample(seed=2) + noise = tfd.Normal( + loc=dt(0), + scale=dt(1)).sample(num_samples, seed=3) + y = tf.tensordot(x, w, axes=[[1], [0]]) + noise + return y, x, w + + def make_prior(sigma, dims): + # p(w | sigma) + return tfd.MultivariateNormalDiag( + loc=tf.zeros([dims], dtype=sigma.dtype), + scale_identity_multiplier=sigma) + + def make_likelihood(x, w): + # p(y | x, w) + return tfd.MultivariateNormalDiag( + loc=tf.tensordot(x, w, axes=[[1], [0]])) + + # Setup assumptions. + dtype = np.float32 + num_samples = 150 + dims = 10 + num_iters = int(5e3) + + true_sigma = dtype(0.5) + y, x, true_weights = make_training_data(num_samples, dims, true_sigma) + + # Estimate of `log(true_sigma)`. + log_sigma = tf.get_variable(name="log_sigma", initializer=dtype(0)) + sigma = tf.exp(log_sigma) + + # State of the Markov chain. + weights = tf.get_variable( + name="weights", + initializer=np.random.randn(dims).astype(dtype)) + + prior = make_prior(sigma, dims) + + def joint_log_prob_fn(w): + # f(w) = log p(w, y | x) + return prior.log_prob(w) + make_likelihood(x, w).log_prob(y) + + weights_update = weights.assign( + hmc.kernel(target_log_prob_fn=joint_log_prob, + current_state=weights, + step_size=0.1, + num_leapfrog_steps=5)[0]) + + with tf.control_dependencies([weights_update]): + loss = -prior.log_prob(weights) + + optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01) + log_sigma_update = optimizer.minimize(loss, var_list=[log_sigma]) + + sess.graph.finalize() # No more graph building. - def counter_fn(a, b, c, d, counter): # pylint: disable=unused-argument - return counter < n_steps + tf.global_variables_initializer().run() - with ops.name_scope(name, 'leapfrog_integrator', - [step_size, n_steps, initial_position, initial_momentum, - initial_grad]): - _, new_x, new_m, new_grad, _ = control_flow_ops.while_loop( - counter_fn, leapfrog_wrapper, [step_size, initial_position, - initial_momentum, initial_grad, - array_ops.constant(0)], back_prop=False) - # We're counting on the runtime to eliminate this redundant computation. - new_potential, new_grad = potential_and_grad(new_x) - return new_x, new_m, new_potential, new_grad + sigma_history = np.zeros(num_iters, dtype) + weights_history = np.zeros([num_iters, dims], dtype) + for i in xrange(num_iters): + _, sigma_, weights_, _ = sess.run([log_sigma_update, sigma, weights]) + weights_history[i, :] = weights_ + sigma_history[i] = sigma_ -def leapfrog_step(step_size, position, momentum, potential_and_grad, grad, - name=None): - """Applies one step of the leapfrog integrator. + true_weights_ = sess.run(true_weights) - Assumes a simple quadratic kinetic energy function: 0.5 * ||momentum||^2. + # Should converge to something close to true_sigma. + plt.plot(sigma_history); + plt.ylabel("sigma"); + plt.xlabel("iteration"); + ``` Args: - step_size: Scalar step size or array of step sizes for the - leapfrog integrator. Broadcasts to the shape of - `position`. Larger step sizes lead to faster progress, but - too-large step sizes lead to larger discretization error and - worse energy conservation. - position: Tensor containing the value(s) of the position variable(s) - to update. - momentum: Tensor containing the value(s) of the momentum variable(s) - to update. - potential_and_grad: Python callable that takes a position tensor like - `position` and returns the potential energy and its gradient at that - position. - grad: Tensor with the value of the gradient of the potential energy - at `position`. + target_log_prob_fn: Python callable which takes an argument like + `current_state` (or `*current_state` if it's a list) and returns its + (possibly unnormalized) log-density under the target distribution. + current_state: `Tensor` or Python `list` of `Tensor`s representing the + current state(s) of the Markov chain(s). The first `r` dimensions index + independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`. + step_size: `Tensor` or Python `list` of `Tensor`s representing the step size + for the leapfrog integrator. Must broadcast with the shape of + `current_state`. Larger step sizes lead to faster progress, but too-large + step sizes make rejection exponentially more likely. When possible, it's + often helpful to match per-variable step sizes to the standard deviations + of the target distribution in each variable. + num_leapfrog_steps: Integer number of steps to run the leapfrog integrator + for. Total progress per HMC step is roughly proportional to `step_size * + num_leapfrog_steps`. + seed: Python integer to seed the random number generator. + current_target_log_prob: (Optional) `Tensor` representing the value of + `target_log_prob_fn` at the `current_state`. The only reason to + specify this argument is to reduce TF graph size. + Default value: `None` (i.e., compute as needed). + current_grads_target_log_prob: (Optional) Python list of `Tensor`s + representing gradient of `current_target_log_prob` at the `current_state` + and wrt the `current_state`. Must have same shape as `current_state`. The + only reason to specify this argument is to reduce TF graph size. + Default value: `None` (i.e., compute as needed). name: Python `str` name prefixed to Ops created by this function. + Default value: `None` (i.e., "hmc_kernel"). Returns: - updated_position: Updated value of the position. - updated_momentum: Updated value of the momentum. - new_potential: Potential energy of the new position. Has shape matching - `potential_and_grad(position)`. - new_grad: Gradient from potential_and_grad() evaluated at the new position. - Has shape matching `position`. + accepted_state: Tensor or Python list of `Tensor`s representing the state(s) + of the Markov chain(s) at each result step. Has same shape as + `current_state`. + kernel_results: `collections.namedtuple` of internal calculations used to + advance the chain. + + Raises: + ValueError: if there isn't one `step_size` or a list with same length as + `current_state`. + """ + with ops.name_scope( + name, "hmc_kernel", + [current_state, step_size, num_leapfrog_steps, seed, + current_target_log_prob, current_grads_target_log_prob]): + with ops.name_scope("initialize"): + [current_state_parts, step_sizes, current_target_log_prob, + current_grads_target_log_prob] = _prepare_args( + target_log_prob_fn, current_state, step_size, + current_target_log_prob, current_grads_target_log_prob, + maybe_expand=True) + independent_chain_ndims = distributions_util.prefer_static_rank( + current_target_log_prob) + current_momentums = [] + for s in current_state_parts: + current_momentums.append(random_ops.random_normal( + shape=array_ops.shape(s), + dtype=s.dtype.base_dtype, + seed=seed)) + seed = distributions_util.gen_new_seed( + seed, salt="hmc_kernel_momentums") + + num_leapfrog_steps = ops.convert_to_tensor( + num_leapfrog_steps, + dtype=dtypes.int32, + name="num_leapfrog_steps") + [ + proposed_momentums, + proposed_state_parts, + proposed_target_log_prob, + proposed_grads_target_log_prob, + ] = _leapfrog_integrator(current_momentums, + target_log_prob_fn, + current_state_parts, + step_sizes, + num_leapfrog_steps, + current_target_log_prob, + current_grads_target_log_prob) + + energy_change = _compute_energy_change(current_target_log_prob, + current_momentums, + proposed_target_log_prob, + proposed_momentums, + independent_chain_ndims) + + # u < exp(min(-energy, 0)), where u~Uniform[0,1) + # ==> -log(u) >= max(e, 0) + # ==> -log(u) >= e + # (Perhaps surprisingly, we don't have a better way to obtain a random + # uniform from positive reals, i.e., `tf.random_uniform(minval=0, + # maxval=np.inf)` won't work.) + random_uniform = random_ops.random_uniform( + shape=array_ops.shape(energy_change), + dtype=energy_change.dtype, + seed=seed) + random_positive = -math_ops.log(random_uniform) + is_accepted = random_positive >= energy_change + + accepted_target_log_prob = array_ops.where(is_accepted, + proposed_target_log_prob, + current_target_log_prob) + + accepted_state_parts = [_choose(is_accepted, + proposed_state_part, + current_state_part, + independent_chain_ndims) + for current_state_part, proposed_state_part + in zip(current_state_parts, proposed_state_parts)] + + accepted_grads_target_log_prob = [ + _choose(is_accepted, + proposed_grad, + grad, + independent_chain_ndims) + for proposed_grad, grad + in zip(proposed_grads_target_log_prob, current_grads_target_log_prob)] + + maybe_flatten = lambda x: x if _is_list_like(current_state) else x[0] + return [ + maybe_flatten(accepted_state_parts), + KernelResults( + acceptance_probs=math_ops.exp(math_ops.minimum(-energy_change, 0.)), + current_grads_target_log_prob=accepted_grads_target_log_prob, + current_target_log_prob=accepted_target_log_prob, + energy_change=energy_change, + is_accepted=is_accepted, + proposed_grads_target_log_prob=proposed_grads_target_log_prob, + proposed_state=maybe_flatten(proposed_state_parts), + proposed_target_log_prob=proposed_target_log_prob, + random_positive=random_positive, + ), + ] + + +def _leapfrog_integrator(current_momentums, + target_log_prob_fn, + current_state_parts, + step_sizes, + num_leapfrog_steps, + current_target_log_prob=None, + current_grads_target_log_prob=None, + name=None): + """Applies `num_leapfrog_steps` of the leapfrog integrator. + + Assumes a simple quadratic kinetic energy function: `0.5 ||momentum||**2`. + + #### Examples: - Example: Simple quadratic potential. + ##### Simple quadratic potential. ```python - def potential_and_grad(position): - # Simple quadratic potential - return tf.reduce_sum(0.5 * tf.square(position)), position + tfd = tf.contrib.distributions + + dims = 10 + num_iter = int(1e3) + dtype = np.float32 + position = tf.placeholder(np.float32) momentum = tf.placeholder(np.float32) - potential, grad = potential_and_grad(position) - new_position, new_momentum, new_potential, new_grad = hmc.leapfrog_step( - 0.1, position, momentum, potential_and_grad, grad) - - sess = tf.Session() - position_val = np.random.randn(10) - momentum_val = np.random.randn(10) - potential_val, grad_val = sess.run([potential, grad], - {position: position_val}) - positions = np.zeros([100, 10]) - for i in xrange(100): - position_val, momentum_val, potential_val, grad_val = sess.run( - [new_position, new_momentum, new_potential, new_grad], - {position: position_val, momentum: momentum_val}) - positions[i] = position_val - # Should trace out sinusoidal dynamics. - plt.plot(positions[:, 0]) + + [ + new_momentums, + new_positions, + ] = hmc._leapfrog_integrator( + current_momentums=[momentum], + target_log_prob_fn=tfd.MultivariateNormalDiag( + loc=tf.zeros(dims, dtype)).log_prob, + current_state_parts=[position], + step_sizes=0.1, + num_leapfrog_steps=3)[:2] + + sess.graph.finalize() # No more graph building. + + momentum_ = np.random.randn(dims).astype(dtype) + position_ = np.random.randn(dims).astype(dtype) + + positions = np.zeros([num_iter, dims], dtype) + for i in xrange(num_iter): + position_, momentum_ = sess.run( + [new_momentums[0], new_position[0]], + feed_dict={position: position_, momentum: momentum_}) + positions[i] = position_ + + plt.plot(positions[:, 0]); # Sinusoidal. ``` + + Args: + current_momentums: Tensor containing the value(s) of the momentum + variable(s) to update. + target_log_prob_fn: Python callable which takes an argument like + `*current_state_parts` and returns its (possibly unnormalized) log-density + under the target distribution. + current_state_parts: Python `list` of `Tensor`s representing the current + state(s) of the Markov chain(s). The first `independent_chain_ndims` of + the `Tensor`(s) index different chains. + step_sizes: Python `list` of `Tensor`s representing the step size for the + leapfrog integrator. Must broadcast with the shape of + `current_state_parts`. Larger step sizes lead to faster progress, but + too-large step sizes make rejection exponentially more likely. When + possible, it's often helpful to match per-variable step sizes to the + standard deviations of the target distribution in each variable. + num_leapfrog_steps: Integer number of steps to run the leapfrog integrator + for. Total progress per HMC step is roughly proportional to `step_size * + num_leapfrog_steps`. + current_target_log_prob: (Optional) `Tensor` representing the value of + `target_log_prob_fn(*current_state_parts)`. The only reason to specify + this argument is to reduce TF graph size. + Default value: `None` (i.e., compute as needed). + current_grads_target_log_prob: (Optional) Python list of `Tensor`s + representing gradient of `target_log_prob_fn(*current_state_parts`) wrt + `current_state_parts`. Must have same shape as `current_state_parts`. The + only reason to specify this argument is to reduce TF graph size. + Default value: `None` (i.e., compute as needed). + name: Python `str` name prefixed to Ops created by this function. + Default value: `None` (i.e., "hmc_leapfrog_integrator"). + + Returns: + proposed_momentums: Updated value of the momentum. + proposed_state_parts: Tensor or Python list of `Tensor`s representing the + state(s) of the Markov chain(s) at each result step. Has same shape as + input `current_state_parts`. + proposed_target_log_prob: `Tensor` representing the value of + `target_log_prob_fn` at `accepted_state`. + proposed_grads_target_log_prob: Gradient of `proposed_target_log_prob` wrt + `accepted_state`. + + Raises: + ValueError: if `len(momentums) != len(state_parts)`. + ValueError: if `len(state_parts) != len(step_sizes)`. + ValueError: if `len(state_parts) != len(grads_target_log_prob)`. + TypeError: if `not target_log_prob.dtype.is_floating`. """ - with ops.name_scope(name, 'leapfrog_step', [step_size, position, momentum, - grad]): - momentum -= 0.5 * step_size * grad - position += step_size * momentum - potential, grad = potential_and_grad(position) - momentum -= 0.5 * step_size * grad - - return position, momentum, potential, grad + def _loop_body(step, + current_momentums, + current_state_parts, + ignore_current_target_log_prob, # pylint: disable=unused-argument + current_grads_target_log_prob): + return [step + 1] + list(_leapfrog_step(current_momentums, + target_log_prob_fn, + current_state_parts, + step_sizes, + current_grads_target_log_prob)) + + with ops.name_scope( + name, "hmc_leapfrog_integrator", + [current_momentums, current_state_parts, step_sizes, num_leapfrog_steps, + current_target_log_prob, current_grads_target_log_prob]): + if len(current_momentums) != len(current_state_parts): + raise ValueError("`momentums` must be in one-to-one correspondence " + "with `state_parts`") + num_leapfrog_steps = ops.convert_to_tensor(num_leapfrog_steps, + name="num_leapfrog_steps") + current_target_log_prob, current_grads_target_log_prob = ( + _maybe_call_fn_and_grads( + target_log_prob_fn, + current_state_parts, + current_target_log_prob, + current_grads_target_log_prob)) + return control_flow_ops.while_loop( + cond=lambda iter_, *args: iter_ < num_leapfrog_steps, + body=_loop_body, + loop_vars=[ + np.int32(0), # iter_ + current_momentums, + current_state_parts, + current_target_log_prob, + current_grads_target_log_prob, + ], + back_prop=False)[1:] # Lop-off "iter_". + + +def _leapfrog_step(current_momentums, + target_log_prob_fn, + current_state_parts, + step_sizes, + current_grads_target_log_prob, + name=None): + """Applies one step of the leapfrog integrator.""" + with ops.name_scope( + name, "_leapfrog_step", + [current_momentums, current_state_parts, step_sizes, + current_grads_target_log_prob]): + proposed_momentums = [m + 0.5 * ss * g for m, ss, g + in zip(current_momentums, + step_sizes, + current_grads_target_log_prob)] + proposed_state_parts = [x + ss * m for x, ss, m + in zip(current_state_parts, + step_sizes, + proposed_momentums)] + proposed_target_log_prob = target_log_prob_fn(*proposed_state_parts) + if not proposed_target_log_prob.dtype.is_floating: + raise TypeError("`target_log_prob_fn` must produce a `Tensor` " + "with `float` `dtype`.") + proposed_grads_target_log_prob = gradients_ops.gradients( + proposed_target_log_prob, proposed_state_parts) + if any(g is None for g in proposed_grads_target_log_prob): + raise ValueError( + "Encountered `None` gradient. Does your target `target_log_prob_fn` " + "access all `tf.Variable`s via `tf.get_variable`?\n" + " current_state_parts: {}\n" + " proposed_state_parts: {}\n" + " proposed_grads_target_log_prob: {}".format( + current_state_parts, + proposed_state_parts, + proposed_grads_target_log_prob)) + proposed_momentums = [m + 0.5 * ss * g for m, ss, g + in zip(proposed_momentums, + step_sizes, + proposed_grads_target_log_prob)] + return [ + proposed_momentums, + proposed_state_parts, + proposed_target_log_prob, + proposed_grads_target_log_prob, + ] + + +def _compute_energy_change(current_target_log_prob, + current_momentums, + proposed_target_log_prob, + proposed_momentums, + independent_chain_ndims, + name=None): + """Helper to `kernel` which computes the energy change.""" + with ops.name_scope( + name, "compute_energy_change", + ([current_target_log_prob, proposed_target_log_prob, + independent_chain_ndims] + + current_momentums + proposed_momentums)): + # Abbreviate lk0=log_kinetic_energy and lk1=proposed_log_kinetic_energy + # since they're a mouthful and lets us inline more. + lk0, lk1 = [], [] + for current_momentum, proposed_momentum in zip(current_momentums, + proposed_momentums): + axis = math_ops.range(independent_chain_ndims, + array_ops.rank(current_momentum)) + lk0.append(_log_sum_sq(current_momentum, axis)) + lk1.append(_log_sum_sq(proposed_momentum, axis)) + + lk0 = -np.log(2.) + math_ops.reduce_logsumexp(array_ops.stack(lk0, axis=-1), + axis=-1) + lk1 = -np.log(2.) + math_ops.reduce_logsumexp(array_ops.stack(lk1, axis=-1), + axis=-1) + lp0 = -current_target_log_prob # log_potential + lp1 = -proposed_target_log_prob # proposed_log_potential + x = array_ops.stack([lp1, math_ops.exp(lk1), -lp0, -math_ops.exp(lk0)], + axis=-1) + + # The sum is NaN if any element is NaN or we see both +Inf and -Inf. + # Thus we will replace such rows with infinite energy change which implies + # rejection. Recall that float-comparisons with NaN are always False. + is_sum_determinate = ( + math_ops.reduce_all(math_ops.is_finite(x) | (x >= 0.), axis=-1) & + math_ops.reduce_all(math_ops.is_finite(x) | (x <= 0.), axis=-1)) + is_sum_determinate = array_ops.tile( + is_sum_determinate[..., array_ops.newaxis], + multiples=array_ops.concat([ + array_ops.ones(array_ops.rank(is_sum_determinate), + dtype=dtypes.int32), + [4], + ], axis=0)) + x = array_ops.where(is_sum_determinate, + x, + array_ops.fill(array_ops.shape(x), + value=x.dtype.as_numpy_dtype(np.inf))) + + return math_ops.reduce_sum(x, axis=-1) + + +def _choose(is_accepted, + accepted, + rejected, + independent_chain_ndims, + name=None): + """Helper to `kernel` which expand_dims `is_accepted` to apply tf.where.""" + def _expand_is_accepted_like(x): + with ops.name_scope("_choose"): + expand_shape = array_ops.concat([ + array_ops.shape(is_accepted), + array_ops.ones([array_ops.rank(x) - array_ops.rank(is_accepted)], + dtype=dtypes.int32), + ], axis=0) + multiples = array_ops.concat([ + array_ops.ones([array_ops.rank(is_accepted)], dtype=dtypes.int32), + array_ops.shape(x)[independent_chain_ndims:], + ], axis=0) + m = array_ops.tile(array_ops.reshape(is_accepted, expand_shape), + multiples) + m.set_shape(x.shape) + return m + with ops.name_scope(name, "_choose", values=[ + is_accepted, accepted, rejected, independent_chain_ndims]): + return array_ops.where(_expand_is_accepted_like(accepted), + accepted, + rejected) + + +def _maybe_call_fn_and_grads(fn, + fn_arg_list, + fn_result=None, + grads_fn_result=None, + description="target_log_prob"): + """Helper which computes `fn_result` and `grads` if needed.""" + fn_arg_list = (list(fn_arg_list) if _is_list_like(fn_arg_list) + else [fn_arg_list]) + if fn_result is None: + fn_result = fn(*fn_arg_list) + if not fn_result.dtype.is_floating: + raise TypeError("`{}` must be a `Tensor` with `float` `dtype`.".format( + description)) + if grads_fn_result is None: + grads_fn_result = gradients_ops.gradients( + fn_result, fn_arg_list) + if len(fn_arg_list) != len(grads_fn_result): + raise ValueError("`{}` must be in one-to-one correspondence with " + "`grads_{}`".format(*[description]*2)) + if any(g is None for g in grads_fn_result): + raise ValueError("Encountered `None` gradient.") + return fn_result, grads_fn_result + + +def _prepare_args(target_log_prob_fn, state, step_size, + target_log_prob=None, grads_target_log_prob=None, + maybe_expand=False, description="target_log_prob"): + """Helper which processes input args to meet list-like assumptions.""" + state_parts = list(state) if _is_list_like(state) else [state] + state_parts = [ops.convert_to_tensor(s, name="state") + for s in state_parts] + target_log_prob, grads_target_log_prob = _maybe_call_fn_and_grads( + target_log_prob_fn, + state_parts, + target_log_prob, + grads_target_log_prob, + description) + step_sizes = list(step_size) if _is_list_like(step_size) else [step_size] + step_sizes = [ + ops.convert_to_tensor( + s, name="step_size", dtype=target_log_prob.dtype) + for s in step_sizes] + if len(step_sizes) == 1: + step_sizes *= len(state_parts) + if len(state_parts) != len(step_sizes): + raise ValueError("There should be exactly one `step_size` or it should " + "have same length as `current_state`.") + maybe_flatten = lambda x: x if maybe_expand or _is_list_like(state) else x[0] + return [ + maybe_flatten(state_parts), + maybe_flatten(step_sizes), + target_log_prob, + grads_target_log_prob, + ] + + +def _is_list_like(x): + """Helper which returns `True` if input is `list`-like.""" + return isinstance(x, (tuple, list)) + + +def _log_sum_sq(x, axis=None): + """Computes log(sum(x**2)).""" + return math_ops.reduce_logsumexp(2. * math_ops.log(math_ops.abs(x)), axis) diff --git a/tensorflow/contrib/bayesflow/python/ops/layers.py b/tensorflow/contrib/bayesflow/python/ops/layers.py index 93412afae738564d440065f230c9df0036591467..a742b7c1aa593d6c08bf9d8d597c99c9fc4e7aed 100644 --- a/tensorflow/contrib/bayesflow/python/ops/layers.py +++ b/tensorflow/contrib/bayesflow/python/ops/layers.py @@ -24,24 +24,36 @@ from __future__ import print_function # go/tf-wildcard-import # pylint: disable=wildcard-import from tensorflow.contrib.bayesflow.python.ops.layers_conv_variational import * -from tensorflow.contrib.bayesflow.python.ops.layers_dense_variational_impl import * +from tensorflow.contrib.bayesflow.python.ops.layers_dense_variational import * from tensorflow.contrib.bayesflow.python.ops.layers_util import * # pylint: enable=wildcard-import from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ - 'Convolution1DVariational', - 'Convolution2DVariational', - 'Convolution3DVariational', - 'Conv1DVariational', - 'Conv2DVariational', - 'Conv3DVariational', - 'convolution1d_variational', - 'convolution2d_variational', - 'convolution3d_variational', - 'conv1d_variational', - 'conv2d_variational', - 'conv3d_variational', + 'Convolution1DReparameterization', + 'Convolution2DReparameterization', + 'Convolution3DReparameterization', + 'Convolution1DFlipout', + 'Convolution2DFlipout', + 'Convolution3DFlipout', + 'Conv1DReparameterization', + 'Conv2DReparameterization', + 'Conv3DReparameterization', + 'Conv1DFlipout', + 'Conv2DFlipout', + 'Conv3DFlipout', + 'convolution1d_reparameterization', + 'convolution2d_reparameterization', + 'convolution3d_reparameterization', + 'convolution1d_flipout', + 'convolution2d_flipout', + 'convolution3d_flipout', + 'conv1d_reparameterization', + 'conv2d_reparameterization', + 'conv3d_reparameterization', + 'conv1d_flipout', + 'conv2d_flipout', + 'conv3d_flipout', 'DenseReparameterization', 'DenseLocalReparameterization', 'DenseFlipout', diff --git a/tensorflow/contrib/bayesflow/python/ops/layers_conv_variational.py b/tensorflow/contrib/bayesflow/python/ops/layers_conv_variational.py index 6ffb55feb1ad754bf96473c075ad6fd38d4e8be9..7723cfb442712626ff415f1412e3362f2392ce9f 100644 --- a/tensorflow/contrib/bayesflow/python/ops/layers_conv_variational.py +++ b/tensorflow/contrib/bayesflow/python/ops/layers_conv_variational.py @@ -32,6 +32,7 @@ from tensorflow.python.ops import nn_ops from tensorflow.python.ops import standard_ops from tensorflow.python.ops.distributions import kullback_leibler as kl_lib from tensorflow.python.ops.distributions import normal as normal_lib +from tensorflow.python.ops.distributions import util as distribution_util class _ConvVariational(layers_lib.Layer): @@ -123,8 +124,6 @@ class _ConvVariational(layers_lib.Layer): dilation_rate: Dilation rate for an atrous convolution. activation: Activation function (`callable`). activity_regularizer: Regularizer function for the output. - kernel_use_local_reparameterization: Python `bool` indicating whether - `kernel` calculation should employ the Local Reparameterization Trick. kernel_posterior_fn: `callable` returning posterior. kernel_posterior_tensor_fn: `callable` operating on posterior. kernel_prior_fn: `callable` returning prior. @@ -271,12 +270,6 @@ class _ConvVariational(layers_lib.Layer): self._built_bias_divergence = True return outputs - def _apply_variational_kernel(self, inputs): - self.kernel_posterior_tensor = self.kernel_posterior_tensor_fn( - self.kernel_posterior) - outputs = self._convolution_op(inputs, self.kernel_posterior_tensor) - return outputs - def _apply_variational_bias(self, inputs): if self.bias_posterior is None: self.bias_posterior_tensor = None @@ -356,7 +349,165 @@ class _ConvVariational(layers_lib.Layer): new_space) -class Conv1DVariational(_ConvVariational): +class _ConvReparameterization(_ConvVariational): + """Abstract nD convolution layer (private, used as implementation base). + + This layer creates a convolution kernel that is convolved + (actually cross-correlated) with the layer input to produce a tensor of + outputs. It may also include a bias addition and activation function + on the outputs. It assumes the `kernel` and/or `bias` are drawn from + distributions. + + By default, the layer implements a stochastic forward pass via + sampling from the kernel and bias posteriors, + ```none + outputs = f(inputs; kernel, bias), kernel, bias ~ posterior + ``` + where f denotes the layer's calculation. It uses the reparameterization + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. + + The arguments permit separate specification of the surrogate posterior + (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` + distributions. + + Arguments: + rank: An integer, the rank of the convolution, e.g. "2" for 2D convolution. + filters: Integer, the dimensionality of the output space (i.e. the number + of filters in the convolution). + kernel_size: An integer or tuple/list of n integers, specifying the + length of the convolution window. + strides: An integer or tuple/list of n integers, + specifying the stride length of the convolution. + Specifying any stride value != 1 is incompatible with specifying + any `dilation_rate` value != 1. + padding: One of `"valid"` or `"same"` (case-insensitive). + data_format: A string, one of `channels_last` (default) or `channels_first`. + The ordering of the dimensions in the inputs. + `channels_last` corresponds to inputs with shape + `(batch, ..., channels)` while `channels_first` corresponds to + inputs with shape `(batch, channels, ...)`. + dilation_rate: An integer or tuple/list of n integers, specifying + the dilation rate to use for dilated convolution. + Currently, specifying any `dilation_rate` value != 1 is + incompatible with specifying any `strides` value != 1. + activation: Activation function. Set it to None to maintain a + linear activation. + activity_regularizer: Optional regularizer function for the output. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + kernel_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `kernel` parameter. Default value: + `default_mean_field_normal_fn()`. + kernel_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + kernel_prior_fn: Python `callable` which creates `tf.distributions` + instance. See `default_mean_field_normal_fn` docstring for required + parameter signature. + Default value: `tf.distributions.Normal(loc=0., scale=1.)`. + kernel_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + bias_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `bias` parameter. Default value: + `default_mean_field_normal_fn(is_singular=True)` (which creates an + instance of `tf.distributions.Deterministic`). + bias_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + bias_prior_fn: Python `callable` which creates `tf.distributions` instance. + See `default_mean_field_normal_fn` docstring for required parameter + signature. Default value: `None` (no prior, no variational inference) + bias_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + name: A string, the name of the layer. + + Properties: + rank: Python integer, dimensionality of convolution. + filters: Python integer, dimensionality of the output space. + kernel_size: Size of the convolution window. + strides: Stride length of convolution. + padding: Python string describing padding approach. + data_format: Python string describing input data's dimensions. + dilation_rate: Dilation rate for an atrous convolution. + activation: Activation function (`callable`). + activity_regularizer: Regularizer function for the output. + kernel_posterior_fn: `callable` returning posterior. + kernel_posterior_tensor_fn: `callable` operating on posterior. + kernel_prior_fn: `callable` returning prior. + kernel_divergence_fn: `callable` returning divergence. + bias_posterior_fn: `callable` returning posterior. + bias_posterior_tensor_fn: `callable` operating on posterior. + bias_prior_fn: `callable` returning prior. + bias_divergence_fn: `callable` returning divergence. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. + """ + + def __init__( + self, + rank, + filters, + kernel_size, + strides=1, + padding="valid", + data_format="channels_last", + dilation_rate=1, + activation=None, + activity_regularizer=None, + trainable=True, + kernel_posterior_fn=layers_util.default_mean_field_normal_fn(), + kernel_posterior_tensor_fn=lambda d: d.sample(), + kernel_prior_fn=lambda dtype, *args: normal_lib.Normal( # pylint: disable=g-long-lambda + loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)), + kernel_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + bias_posterior_fn=layers_util.default_mean_field_normal_fn(is_singular=True), # pylint: disable=line-too-long + bias_posterior_tensor_fn=lambda d: d.sample(), + bias_prior_fn=None, + bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + name=None, + **kwargs): + super(_ConvReparameterization, self).__init__( + rank=rank, + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activation, + activity_regularizer=activity_regularizer, + trainable=trainable, + kernel_posterior_fn=kernel_posterior_fn, + kernel_posterior_tensor_fn=kernel_posterior_tensor_fn, + kernel_prior_fn=kernel_prior_fn, + kernel_divergence_fn=kernel_divergence_fn, + bias_posterior_fn=bias_posterior_fn, + bias_posterior_tensor_fn=bias_posterior_tensor_fn, + bias_prior_fn=bias_prior_fn, + bias_divergence_fn=bias_divergence_fn, + name=name, **kwargs) + + def _apply_variational_kernel(self, inputs): + self.kernel_posterior_tensor = self.kernel_posterior_tensor_fn( + self.kernel_posterior) + self.kernel_posterior_affine = None + self.kernel_posterior_affine_tensor = None + outputs = self._convolution_op(inputs, self.kernel_posterior_tensor) + return outputs + + +class Conv1DReparameterization(_ConvReparameterization): """1D convolution layer (e.g. temporal convolution). This layer creates a convolution kernel that is convolved @@ -370,7 +521,9 @@ class Conv1DVariational(_ConvVariational): ```none outputs = f(inputs; kernel, bias), kernel, bias ~ posterior ``` - where f denotes the layer's calculation. + where f denotes the layer's calculation. It uses the reparameterization + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` @@ -442,8 +595,6 @@ class Conv1DVariational(_ConvVariational): dilation_rate: Dilation rate for an atrous convolution. activation: Activation function (`callable`). activity_regularizer: Regularizer function for the output. - kernel_use_local_reparameterization: Python `bool` indicating whether - `kernel` calculation should employ the Local Reparameterization Trick. kernel_posterior_fn: `callable` returning posterior. kernel_posterior_tensor_fn: `callable` operating on posterior. kernel_prior_fn: `callable` returning prior. @@ -463,12 +614,12 @@ class Conv1DVariational(_ConvVariational): tfp = tf.contrib.bayesflow net = tf.reshape(features, [-1, 128, 1]) - net = tfp.layers.Conv1DVariational(64, - kernel_size=5, - padding="SAME", - activation=tf.nn.relu)(net) + net = tfp.layers.Conv1DReparameterization(64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu)(net) net = tf.reshape(net, [-1, 128 * 64]) - logits = tfp.layers.DenseVariational(10)(net) + logits = tfp.layers.DenseReparameterization(10)(net) neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( labels=labels, logits=logits) kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) @@ -482,6 +633,10 @@ class Conv1DVariational(_ConvVariational): the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. """ def __init__( @@ -506,7 +661,7 @@ class Conv1DVariational(_ConvVariational): bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), name=None, **kwargs): - super(Conv1DVariational, self).__init__( + super(Conv1DReparameterization, self).__init__( rank=1, filters=filters, kernel_size=kernel_size, @@ -528,7 +683,7 @@ class Conv1DVariational(_ConvVariational): name=name, **kwargs) -def conv1d_variational( +def conv1d_reparameterization( inputs, filters, kernel_size, @@ -563,7 +718,9 @@ def conv1d_variational( ```none outputs = f(inputs; kernel, bias), kernel, bias ~ posterior ``` - where f denotes the layer's calculation. + where f denotes the layer's calculation. It uses the reparameterization + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` @@ -645,13 +802,13 @@ def conv1d_variational( tfp = tf.contrib.bayesflow net = tf.reshape(features, [-1, 128, 1]) - net = tfp.layers.conv1d_variational(net, - 64, - kernel_size=5, - padding="SAME", - activation=tf.nn.relu) + net = tfp.layers.conv1d_reparameterization(net, + filters=64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu) net = tf.reshape(net, [-1, 128 * 64]) - logits = tfp.layers.dense_variational(net, 10) + logits = tfp.layers.dense_reparameterization(net, 10) neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( labels=labels, logits=logits) kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) @@ -665,8 +822,12 @@ def conv1d_variational( the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. """ - layer = Conv1DVariational( + layer = Conv1DReparameterization( filters=filters, kernel_size=kernel_size, strides=strides, @@ -691,7 +852,7 @@ def conv1d_variational( return layer.apply(inputs) -class Conv2DVariational(_ConvVariational): +class Conv2DReparameterization(_ConvReparameterization): """2D convolution layer (e.g. spatial convolution over images). This layer creates a convolution kernel that is convolved @@ -705,7 +866,9 @@ class Conv2DVariational(_ConvVariational): ```none outputs = f(inputs; kernel, bias), kernel, bias ~ posterior ``` - where f denotes the layer's calculation. + where f denotes the layer's calculation. It uses the reparameterization + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` @@ -784,8 +947,6 @@ class Conv2DVariational(_ConvVariational): dilation_rate: Dilation rate for an atrous convolution. activation: Activation function (`callable`). activity_regularizer: Regularizer function for the output. - kernel_use_local_reparameterization: Python `bool` indicating whether - `kernel` calculation should employ the Local Reparameterization Trick. kernel_posterior_fn: `callable` returning posterior. kernel_posterior_tensor_fn: `callable` operating on posterior. kernel_prior_fn: `callable` returning prior. @@ -805,15 +966,15 @@ class Conv2DVariational(_ConvVariational): tfp = tf.contrib.bayesflow net = tf.reshape(features, [-1, 32, 32, 3]) - net = tfp.layers.Conv2DVariational(64, - kernel_size=5, - padding="SAME", - activation=tf.nn.relu)(net) + net = tfp.layers.Conv2DReparameterization(64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu)(net) net = tf.layers.MaxPooling2D(pool_size=2, strides=2, padding="SAME")(net) net = tf.reshape(net, [-1, 8 * 8 * 64]) - logits = tfp.layers.DenseVariational(10)(net) + logits = tfp.layers.DenseReparameterization(10)(net) neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( labels=labels, logits=logits) kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) @@ -827,6 +988,10 @@ class Conv2DVariational(_ConvVariational): the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. """ def __init__( @@ -851,7 +1016,7 @@ class Conv2DVariational(_ConvVariational): bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), name=None, **kwargs): - super(Conv2DVariational, self).__init__( + super(Conv2DReparameterization, self).__init__( rank=2, filters=filters, kernel_size=kernel_size, @@ -873,7 +1038,7 @@ class Conv2DVariational(_ConvVariational): name=name, **kwargs) -def conv2d_variational( +def conv2d_reparameterization( inputs, filters, kernel_size, @@ -908,7 +1073,9 @@ def conv2d_variational( ```none outputs = f(inputs; kernel, bias), kernel, bias ~ posterior ``` - where f denotes the layer's calculation. + where f denotes the layer's calculation. It uses the reparameterization + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` @@ -997,17 +1164,17 @@ def conv2d_variational( tfp = tf.contrib.bayesflow net = tf.reshape(features, [-1, 32, 32, 3]) - net = tfp.layers.conv2d_variational(net, - 64, - kernel_size=5, - padding="SAME", - activation=tf.nn.relu) + net = tfp.layers.conv2d_reparameterization(net, + filters=64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu) net = tf.layers.max_pooling2d(net, pool_size=2, strides=2, padding="SAME") net = tf.reshape(net, [-1, 8 * 8 * 64]) - logits = tfp.layers.dense_variational(net, 10) + logits = tfp.layers.dense_reparameterization(net, 10) neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( labels=labels, logits=logits) kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) @@ -1021,8 +1188,12 @@ def conv2d_variational( the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. """ - layer = Conv2DVariational( + layer = Conv2DReparameterization( filters=filters, kernel_size=kernel_size, strides=strides, @@ -1047,7 +1218,7 @@ def conv2d_variational( return layer.apply(inputs) -class Conv3DVariational(_ConvVariational): +class Conv3DReparameterization(_ConvReparameterization): """3D convolution layer (e.g. spatial convolution over volumes). This layer creates a convolution kernel that is convolved @@ -1061,7 +1232,9 @@ class Conv3DVariational(_ConvVariational): ```none outputs = f(inputs; kernel, bias), kernel, bias ~ posterior ``` - where f denotes the layer's calculation. + where f denotes the layer's calculation. It uses the reparameterization + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` @@ -1141,8 +1314,6 @@ class Conv3DVariational(_ConvVariational): dilation_rate: Dilation rate for an atrous convolution. activation: Activation function (`callable`). activity_regularizer: Regularizer function for the output. - kernel_use_local_reparameterization: Python `bool` indicating whether - `kernel` calculation should employ the Local Reparameterization Trick. kernel_posterior_fn: `callable` returning posterior. kernel_posterior_tensor_fn: `callable` operating on posterior. kernel_prior_fn: `callable` returning prior. @@ -1162,15 +1333,15 @@ class Conv3DVariational(_ConvVariational): tfp = tf.contrib.bayesflow net = tf.reshape(features, [-1, 256, 32, 32, 3]) - net = tfp.layers.Conv3DVariational(64, - kernel_size=5, - padding="SAME", - activation=tf.nn.relu)(net) + net = tfp.layers.Conv3DReparameterization(64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu)(net) net = tf.layers.MaxPooling2D(pool_size=2, strides=2, padding="SAME")(net) net = tf.reshape(net, [-1, 256 * 8 * 8 * 64]) - logits = tfp.layers.DenseVariational(10)(net) + logits = tfp.layers.DenseReparameterization(10)(net) neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( labels=labels, logits=logits) kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) @@ -1184,6 +1355,10 @@ class Conv3DVariational(_ConvVariational): the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. """ def __init__( @@ -1208,7 +1383,7 @@ class Conv3DVariational(_ConvVariational): bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), name=None, **kwargs): - super(Conv3DVariational, self).__init__( + super(Conv3DReparameterization, self).__init__( rank=3, filters=filters, kernel_size=kernel_size, @@ -1230,7 +1405,7 @@ class Conv3DVariational(_ConvVariational): name=name, **kwargs) -def conv3d_variational( +def conv3d_reparameterization( inputs, filters, kernel_size, @@ -1265,7 +1440,9 @@ def conv3d_variational( ```none outputs = f(inputs; kernel, bias), kernel, bias ~ posterior ``` - where f denotes the layer's calculation. + where f denotes the layer's calculation. It uses the reparameterization + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` @@ -1355,17 +1532,17 @@ def conv3d_variational( tfp = tf.contrib.bayesflow net = tf.reshape(features, [-1, 256, 32, 32, 3]) - net = tfp.layers.conv3d_variational(net, - 64, - kernel_size=5, - padding="SAME", - activation=tf.nn.relu) + net = tfp.layers.conv3d_reparameterization(net, + filters=64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu) net = tf.layers.max_pooling2d(net, pool_size=2, strides=2, padding="SAME") net = tf.reshape(net, [-1, 256 * 8 * 8 * 64]) - logits = tfp.layers.dense_variational(net, 10) + logits = tfp.layers.dense_reparameterization(net, 10) neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( labels=labels, logits=logits) kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) @@ -1379,8 +1556,1352 @@ def conv3d_variational( the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. + """ + layer = Conv3DReparameterization( + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activation, + activity_regularizer=activity_regularizer, + trainable=trainable, + kernel_posterior_fn=kernel_posterior_fn, + kernel_posterior_tensor_fn=kernel_posterior_tensor_fn, + kernel_prior_fn=kernel_prior_fn, + kernel_divergence_fn=kernel_divergence_fn, + bias_posterior_fn=bias_posterior_fn, + bias_posterior_tensor_fn=bias_posterior_tensor_fn, + bias_prior_fn=bias_prior_fn, + bias_divergence_fn=bias_divergence_fn, + name=name, + dtype=inputs.dtype.base_dtype, + _scope=name, + _reuse=reuse) + return layer.apply(inputs) + + +class _ConvFlipout(_ConvVariational): + """Abstract nD convolution layer (private, used as implementation base). + + This layer creates a convolution kernel that is convolved + (actually cross-correlated) with the layer input to produce a tensor of + outputs. It may also include a bias addition and activation function + on the outputs. It assumes the `kernel` and/or `bias` are drawn from + distributions. + + By default, the layer implements a stochastic forward pass via + sampling from the kernel and bias posteriors, + ```none + outputs = f(inputs; kernel, bias), kernel, bias ~ posterior + ``` + where f denotes the layer's calculation. It uses the Flipout + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. Flipout uses + roughly twice as many floating point operations as the + reparameterization estimator but has the advantage of significantly + lower variance. + + The arguments permit separate specification of the surrogate posterior + (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` + distributions. + + Arguments: + rank: An integer, the rank of the convolution, e.g. "2" for 2D convolution. + filters: Integer, the dimensionality of the output space (i.e. the number + of filters in the convolution). + kernel_size: An integer or tuple/list of n integers, specifying the + length of the convolution window. + strides: An integer or tuple/list of n integers, + specifying the stride length of the convolution. + Specifying any stride value != 1 is incompatible with specifying + any `dilation_rate` value != 1. + padding: One of `"valid"` or `"same"` (case-insensitive). + data_format: A string, one of `channels_last` (default) or `channels_first`. + The ordering of the dimensions in the inputs. + `channels_last` corresponds to inputs with shape + `(batch, ..., channels)` while `channels_first` corresponds to + inputs with shape `(batch, channels, ...)`. + dilation_rate: An integer or tuple/list of n integers, specifying + the dilation rate to use for dilated convolution. + Currently, specifying any `dilation_rate` value != 1 is + incompatible with specifying any `strides` value != 1. + activation: Activation function. Set it to None to maintain a + linear activation. + activity_regularizer: Optional regularizer function for the output. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + kernel_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `kernel` parameter. Default value: + `default_mean_field_normal_fn()`. + kernel_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + kernel_prior_fn: Python `callable` which creates `tf.distributions` + instance. See `default_mean_field_normal_fn` docstring for required + parameter signature. + Default value: `tf.distributions.Normal(loc=0., scale=1.)`. + kernel_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + bias_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `bias` parameter. Default value: + `default_mean_field_normal_fn(is_singular=True)` (which creates an + instance of `tf.distributions.Deterministic`). + bias_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + bias_prior_fn: Python `callable` which creates `tf.distributions` instance. + See `default_mean_field_normal_fn` docstring for required parameter + signature. Default value: `None` (no prior, no variational inference) + bias_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + seed: Python scalar `int` which initializes the random number + generator. Default value: `None` (i.e., use global seed). + name: A string, the name of the layer. + + Properties: + rank: Python integer, dimensionality of convolution. + filters: Python integer, dimensionality of the output space. + kernel_size: Size of the convolution window. + strides: Stride length of convolution. + padding: Python string describing padding approach. + data_format: Python string describing input data's dimensions. + dilation_rate: Dilation rate for an atrous convolution. + activation: Activation function (`callable`). + activity_regularizer: Regularizer function for the output. + kernel_posterior_fn: `callable` returning posterior. + kernel_posterior_tensor_fn: `callable` operating on posterior. + kernel_prior_fn: `callable` returning prior. + kernel_divergence_fn: `callable` returning divergence. + bias_posterior_fn: `callable` returning posterior. + bias_posterior_tensor_fn: `callable` operating on posterior. + bias_prior_fn: `callable` returning prior. + bias_divergence_fn: `callable` returning divergence. + seed: Python integer, used to create random seeds. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb + """ + + def __init__( + self, + rank, + filters, + kernel_size, + strides=1, + padding="valid", + data_format="channels_last", + dilation_rate=1, + activation=None, + activity_regularizer=None, + trainable=True, + kernel_posterior_fn=layers_util.default_mean_field_normal_fn(), + kernel_posterior_tensor_fn=lambda d: d.sample(), + kernel_prior_fn=lambda dtype, *args: normal_lib.Normal( # pylint: disable=g-long-lambda + loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)), + kernel_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + bias_posterior_fn=layers_util.default_mean_field_normal_fn(is_singular=True), # pylint: disable=line-too-long + bias_posterior_tensor_fn=lambda d: d.sample(), + bias_prior_fn=None, + bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + seed=None, + name=None, + **kwargs): + super(_ConvFlipout, self).__init__( + rank=rank, + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activation, + activity_regularizer=activity_regularizer, + trainable=trainable, + kernel_posterior_fn=kernel_posterior_fn, + kernel_posterior_tensor_fn=kernel_posterior_tensor_fn, + kernel_prior_fn=kernel_prior_fn, + kernel_divergence_fn=kernel_divergence_fn, + bias_posterior_fn=bias_posterior_fn, + bias_posterior_tensor_fn=bias_posterior_tensor_fn, + bias_prior_fn=bias_prior_fn, + bias_divergence_fn=bias_divergence_fn, + name=name, **kwargs) + self.seed = seed + + def _apply_variational_kernel(self, inputs): + if (not isinstance(self.kernel_posterior, independent_lib.Independent) or + not isinstance(self.kernel_posterior.distribution, normal_lib.Normal)): + raise TypeError( + "`{}` requires " + "`kernel_posterior_fn` produce an instance of " + "`tf.distributions.Independent(tf.distributions.Normal)` " + "(saw: \"{}\").".format( + type(self).__name__, self.kernel_posterior.name)) + self.kernel_posterior_affine = normal_lib.Normal( + loc=array_ops.zeros_like(self.kernel_posterior.distribution.loc), + scale=self.kernel_posterior.distribution.scale) + self.kernel_posterior_affine_tensor = ( + self.kernel_posterior_tensor_fn(self.kernel_posterior_affine)) + self.kernel_posterior_tensor = None + + outputs = self._convolution_op( + inputs, self.kernel_posterior.distribution.loc) + + input_shape = array_ops.shape(inputs) + output_shape = array_ops.shape(outputs) + batch_shape = array_ops.expand_dims(input_shape[0], 0) + channels = input_shape[-1] + + sign_input = layers_util.random_sign( + array_ops.concat([batch_shape, + array_ops.expand_dims(channels, 0)], 0), + dtype=inputs.dtype, + seed=self.seed) + sign_output = layers_util.random_sign( + array_ops.concat([batch_shape, + array_ops.expand_dims(self.filters, 0)], 0), + dtype=inputs.dtype, + seed=distribution_util.gen_new_seed( + self.seed, salt="conv_flipout")) + for _ in range(self.rank): + sign_input = array_ops.expand_dims(sign_input, 1) # 2D ex: (B, 1, 1, C) + sign_output = array_ops.expand_dims(sign_output, 1) + + sign_input = array_ops.tile( # tile for element-wise op broadcasting + sign_input, + [1] + [input_shape[i + 1] for i in range(self.rank)] + [1]) + sign_output = array_ops.tile( + sign_output, + [1] + [output_shape[i + 1] for i in range(self.rank)] + [1]) + + perturbed_inputs = self._convolution_op( + inputs * sign_input, self.kernel_posterior_affine_tensor) * sign_output + + outputs += perturbed_inputs + return outputs + + +class Conv1DFlipout(_ConvFlipout): + """1D convolution layer (e.g. temporal convolution) with Flipout. + + This layer creates a convolution kernel that is convolved + (actually cross-correlated) with the layer input to produce a tensor of + outputs. It may also include a bias addition and activation function + on the outputs. It assumes the `kernel` and/or `bias` are drawn from + distributions. + + By default, the layer implements a stochastic forward pass via + sampling from the kernel and bias posteriors, + ```none + outputs = f(inputs; kernel, bias), kernel, bias ~ posterior + ``` + where f denotes the layer's calculation. It uses the Flipout + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. Flipout uses + roughly twice as many floating point operations as the + reparameterization estimator but has the advantage of significantly + lower variance. + + The arguments permit separate specification of the surrogate posterior + (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` + distributions. + + Arguments: + filters: Integer, the dimensionality of the output space (i.e. the number + of filters in the convolution). + kernel_size: An integer or tuple/list of a single integer, specifying the + length of the 1D convolution window. + strides: An integer or tuple/list of a single integer, + specifying the stride length of the convolution. + Specifying any stride value != 1 is incompatible with specifying + any `dilation_rate` value != 1. + padding: One of `"valid"` or `"same"` (case-insensitive). + data_format: A string, one of `channels_last` (default) or `channels_first`. + The ordering of the dimensions in the inputs. + `channels_last` corresponds to inputs with shape + `(batch, length, channels)` while `channels_first` corresponds to + inputs with shape `(batch, channels, length)`. + dilation_rate: An integer or tuple/list of a single integer, specifying + the dilation rate to use for dilated convolution. + Currently, specifying any `dilation_rate` value != 1 is + incompatible with specifying any `strides` value != 1. + activation: Activation function. Set it to None to maintain a + linear activation. + activity_regularizer: Optional regularizer function for the output. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + kernel_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `kernel` parameter. Default value: + `default_mean_field_normal_fn()`. + kernel_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + kernel_prior_fn: Python `callable` which creates `tf.distributions` + instance. See `default_mean_field_normal_fn` docstring for required + parameter signature. + Default value: `tf.distributions.Normal(loc=0., scale=1.)`. + kernel_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + bias_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `bias` parameter. Default value: + `default_mean_field_normal_fn(is_singular=True)` (which creates an + instance of `tf.distributions.Deterministic`). + bias_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + bias_prior_fn: Python `callable` which creates `tf.distributions` instance. + See `default_mean_field_normal_fn` docstring for required parameter + signature. Default value: `None` (no prior, no variational inference) + bias_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + seed: Python scalar `int` which initializes the random number + generator. Default value: `None` (i.e., use global seed). + name: A string, the name of the layer. + + Properties: + filters: Python integer, dimensionality of the output space. + kernel_size: Size of the convolution window. + strides: Stride length of convolution. + padding: Python string describing padding approach. + data_format: Python string describing input data's dimensions. + dilation_rate: Dilation rate for an atrous convolution. + activation: Activation function (`callable`). + activity_regularizer: Regularizer function for the output. + kernel_posterior_fn: `callable` returning posterior. + kernel_posterior_tensor_fn: `callable` operating on posterior. + kernel_prior_fn: `callable` returning prior. + kernel_divergence_fn: `callable` returning divergence. + bias_posterior_fn: `callable` returning posterior. + bias_posterior_tensor_fn: `callable` operating on posterior. + bias_prior_fn: `callable` returning prior. + bias_divergence_fn: `callable` returning divergence. + seed: Python integer, used to create random seeds. + + #### Examples + + We illustrate a Bayesian neural network with [variational inference]( + https://en.wikipedia.org/wiki/Variational_Bayesian_methods), + assuming a dataset of `features` and `labels`. + + ```python + tfp = tf.contrib.bayesflow + + net = tf.reshape(features, [-1, 128, 1]) + net = tfp.layers.Conv1DFlipout(64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu)(net) + net = tf.reshape(net, [-1, 128 * 64]) + logits = tfp.layers.DenseFlipout(10)(net) + neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( + labels=labels, logits=logits) + kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) + loss = neg_log_likelihood + kl + train_op = tf.train.AdamOptimizer().minimize(loss) + ``` + + It uses the Flipout gradient estimator to minimize the + Kullback-Leibler divergence up to a constant, also known as the + negative Evidence Lower Bound. It consists of the sum of two terms: + the expected negative log-likelihood, which we approximate via + Monte Carlo; and the KL divergence, which is added via regularizer + terms which are arguments to the layer. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb + """ + + def __init__( + self, + filters, + kernel_size, + strides=1, + padding="valid", + data_format="channels_last", + dilation_rate=1, + activation=None, + activity_regularizer=None, + trainable=True, + kernel_posterior_fn=layers_util.default_mean_field_normal_fn(), + kernel_posterior_tensor_fn=lambda d: d.sample(), + kernel_prior_fn=lambda dtype, *args: normal_lib.Normal( # pylint: disable=g-long-lambda + loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)), + kernel_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + bias_posterior_fn=layers_util.default_mean_field_normal_fn(is_singular=True), # pylint: disable=line-too-long + bias_posterior_tensor_fn=lambda d: d.sample(), + bias_prior_fn=None, + bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + seed=None, + name=None, + **kwargs): + super(Conv1DFlipout, self).__init__( + rank=1, + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activation, + activity_regularizer=activity_regularizer, + trainable=trainable, + kernel_posterior_fn=kernel_posterior_fn, + kernel_posterior_tensor_fn=kernel_posterior_tensor_fn, + kernel_prior_fn=kernel_prior_fn, + kernel_divergence_fn=kernel_divergence_fn, + bias_posterior_fn=bias_posterior_fn, + bias_posterior_tensor_fn=bias_posterior_tensor_fn, + bias_prior_fn=bias_prior_fn, + bias_divergence_fn=bias_divergence_fn, + seed=seed, + name=name, **kwargs) + + +def conv1d_flipout( + inputs, + filters, + kernel_size, + strides=1, + padding="valid", + data_format="channels_last", + dilation_rate=1, + activation=None, + activity_regularizer=None, + trainable=True, + kernel_posterior_fn=layers_util.default_mean_field_normal_fn(), + kernel_posterior_tensor_fn=lambda d: d.sample(), + kernel_prior_fn=lambda dtype, *args: normal_lib.Normal( # pylint: disable=g-long-lambda + loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)), + kernel_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + bias_posterior_fn=layers_util.default_mean_field_normal_fn(is_singular=True), # pylint: disable=line-too-long + bias_posterior_tensor_fn=lambda d: d.sample(), + bias_prior_fn=None, + bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + seed=None, + name=None, + reuse=None): + """Functional interface for 1D convolution layer (e.g. temporal convolution). + + This layer creates a convolution kernel that is convolved + (actually cross-correlated) with the layer input to produce a tensor of + outputs. It may also include a bias addition and activation function + on the outputs. It assumes the `kernel` and/or `bias` are drawn from + distributions. + + By default, the layer implements a stochastic forward pass via + sampling from the kernel and bias posteriors, + ```none + outputs = f(inputs; kernel, bias), kernel, bias ~ posterior + ``` + where f denotes the layer's calculation. It uses the Flipout + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. Flipout uses + roughly twice as many floating point operations as the + reparameterization estimator but has the advantage of significantly + lower variance. + + The arguments permit separate specification of the surrogate posterior + (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` + distributions. + + Arguments: + inputs: Tensor input. + filters: Integer, the dimensionality of the output space (i.e. the number + of filters in the convolution). + kernel_size: An integer or tuple/list of a single integer, specifying the + length of the 1D convolution window. + strides: An integer or tuple/list of a single integer, + specifying the stride length of the convolution. + Specifying any stride value != 1 is incompatible with specifying + any `dilation_rate` value != 1. + padding: One of `"valid"` or `"same"` (case-insensitive). + data_format: A string, one of `channels_last` (default) or `channels_first`. + The ordering of the dimensions in the inputs. + `channels_last` corresponds to inputs with shape + `(batch, length, channels)` while `channels_first` corresponds to + inputs with shape `(batch, channels, length)`. + dilation_rate: An integer or tuple/list of a single integer, specifying + the dilation rate to use for dilated convolution. + Currently, specifying any `dilation_rate` value != 1 is + incompatible with specifying any `strides` value != 1. + activation: Activation function. Set it to None to maintain a + linear activation. + activity_regularizer: Optional regularizer function for the output. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + kernel_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `kernel` parameter. Default value: + `default_mean_field_normal_fn()`. + kernel_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + kernel_prior_fn: Python `callable` which creates `tf.distributions` + instance. See `default_mean_field_normal_fn` docstring for required + parameter signature. + Default value: `tf.distributions.Normal(loc=0., scale=1.)`. + kernel_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + bias_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `bias` parameter. Default value: + `default_mean_field_normal_fn(is_singular=True)` (which creates an + instance of `tf.distributions.Deterministic`). + bias_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + bias_prior_fn: Python `callable` which creates `tf.distributions` instance. + See `default_mean_field_normal_fn` docstring for required parameter + signature. Default value: `None` (no prior, no variational inference) + bias_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + seed: Python scalar `int` which initializes the random number + generator. Default value: `None` (i.e., use global seed). + name: A string, the name of the layer. + reuse: Boolean, whether to reuse the weights of a previous layer + by the same name. + + Returns: + Output tensor. + + Raises: + ValueError: if eager execution is enabled. + + #### Examples + + We illustrate a Bayesian neural network with [variational inference]( + https://en.wikipedia.org/wiki/Variational_Bayesian_methods), + assuming a dataset of `features` and `labels`. + + ```python + tfp = tf.contrib.bayesflow + + net = tf.reshape(features, [-1, 128, 1]) + net = tfp.layers.conv1d_flipout(net, + filters=64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu) + net = tf.reshape(net, [-1, 128 * 64]) + logits = tfp.layers.dense_flipout(net, 10) + neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( + labels=labels, logits=logits) + kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) + loss = neg_log_likelihood + kl + train_op = tf.train.AdamOptimizer().minimize(loss) + ``` + + It uses the Flipout gradient estimator to minimize the + Kullback-Leibler divergence up to a constant, also known as the + negative Evidence Lower Bound. It consists of the sum of two terms: + the expected negative log-likelihood, which we approximate via + Monte Carlo; and the KL divergence, which is added via regularizer + terms which are arguments to the layer. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb + """ + layer = Conv1DFlipout( + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activation, + activity_regularizer=activity_regularizer, + trainable=trainable, + kernel_posterior_fn=kernel_posterior_fn, + kernel_posterior_tensor_fn=kernel_posterior_tensor_fn, + kernel_prior_fn=kernel_prior_fn, + kernel_divergence_fn=kernel_divergence_fn, + bias_posterior_fn=bias_posterior_fn, + bias_posterior_tensor_fn=bias_posterior_tensor_fn, + bias_prior_fn=bias_prior_fn, + bias_divergence_fn=bias_divergence_fn, + seed=seed, + name=name, + dtype=inputs.dtype.base_dtype, + _scope=name, + _reuse=reuse) + return layer.apply(inputs) + + +class Conv2DFlipout(_ConvFlipout): + """2D convolution layer (e.g. spatial convolution over images) with Flipout. + + This layer creates a convolution kernel that is convolved + (actually cross-correlated) with the layer input to produce a tensor of + outputs. It may also include a bias addition and activation function + on the outputs. It assumes the `kernel` and/or `bias` are drawn from + distributions. + + By default, the layer implements a stochastic forward pass via + sampling from the kernel and bias posteriors, + ```none + outputs = f(inputs; kernel, bias), kernel, bias ~ posterior + ``` + where f denotes the layer's calculation. It uses the Flipout + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. Flipout uses + roughly twice as many floating point operations as the + reparameterization estimator but has the advantage of significantly + lower variance. + + The arguments permit separate specification of the surrogate posterior + (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` + distributions. + + Arguments: + filters: Integer, the dimensionality of the output space (i.e. the number + of filters in the convolution). + kernel_size: An integer or tuple/list of 2 integers, specifying the + height and width of the 2D convolution window. + Can be a single integer to specify the same value for + all spatial dimensions. + strides: An integer or tuple/list of 2 integers, + specifying the strides of the convolution along the height and width. + Can be a single integer to specify the same value for + all spatial dimensions. + Specifying any stride value != 1 is incompatible with specifying + any `dilation_rate` value != 1. + padding: One of `"valid"` or `"same"` (case-insensitive). + data_format: A string, one of `channels_last` (default) or `channels_first`. + The ordering of the dimensions in the inputs. + `channels_last` corresponds to inputs with shape + `(batch, height, width, channels)` while `channels_first` corresponds to + inputs with shape `(batch, channels, height, width)`. + + dilation_rate: An integer or tuple/list of 2 integers, specifying + the dilation rate to use for dilated convolution. + Can be a single integer to specify the same value for + all spatial dimensions. + Currently, specifying any `dilation_rate` value != 1 is + incompatible with specifying any stride value != 1. + activation: Activation function. Set it to None to maintain a + linear activation. + activity_regularizer: Optional regularizer function for the output. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + kernel_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `kernel` parameter. Default value: + `default_mean_field_normal_fn()`. + kernel_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + kernel_prior_fn: Python `callable` which creates `tf.distributions` + instance. See `default_mean_field_normal_fn` docstring for required + parameter signature. + Default value: `tf.distributions.Normal(loc=0., scale=1.)`. + kernel_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + bias_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `bias` parameter. Default value: + `default_mean_field_normal_fn(is_singular=True)` (which creates an + instance of `tf.distributions.Deterministic`). + bias_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + bias_prior_fn: Python `callable` which creates `tf.distributions` instance. + See `default_mean_field_normal_fn` docstring for required parameter + signature. Default value: `None` (no prior, no variational inference) + bias_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + seed: Python scalar `int` which initializes the random number + generator. Default value: `None` (i.e., use global seed). + name: A string, the name of the layer. + + Properties: + filters: Python integer, dimensionality of the output space. + kernel_size: Size of the convolution window. + strides: Stride length of convolution. + padding: Python string describing padding approach. + data_format: Python string describing input data's dimensions. + dilation_rate: Dilation rate for an atrous convolution. + activation: Activation function (`callable`). + activity_regularizer: Regularizer function for the output. + kernel_posterior_fn: `callable` returning posterior. + kernel_posterior_tensor_fn: `callable` operating on posterior. + kernel_prior_fn: `callable` returning prior. + kernel_divergence_fn: `callable` returning divergence. + bias_posterior_fn: `callable` returning posterior. + bias_posterior_tensor_fn: `callable` operating on posterior. + bias_prior_fn: `callable` returning prior. + bias_divergence_fn: `callable` returning divergence. + seed: Python integer, used to create random seeds. + + #### Examples + + We illustrate a Bayesian neural network with [variational inference]( + https://en.wikipedia.org/wiki/Variational_Bayesian_methods), + assuming a dataset of `features` and `labels`. + + ```python + tfp = tf.contrib.bayesflow + + net = tf.reshape(features, [-1, 32, 32, 3]) + net = tfp.layers.Conv2DFlipout(64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu)(net) + net = tf.layers.MaxPooling2D(pool_size=2, + strides=2, + padding="SAME")(net) + net = tf.reshape(net, [-1, 8 * 8 * 64]) + logits = tfp.layers.DenseFlipout(10)(net) + neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( + labels=labels, logits=logits) + kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) + loss = neg_log_likelihood + kl + train_op = tf.train.AdamOptimizer().minimize(loss) + ``` + + It uses the Flipout gradient estimator to minimize the + Kullback-Leibler divergence up to a constant, also known as the + negative Evidence Lower Bound. It consists of the sum of two terms: + the expected negative log-likelihood, which we approximate via + Monte Carlo; and the KL divergence, which is added via regularizer + terms which are arguments to the layer. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb + """ + + def __init__( + self, + filters, + kernel_size, + strides=(1, 1), + padding="valid", + data_format="channels_last", + dilation_rate=(1, 1), + activation=None, + activity_regularizer=None, + trainable=True, + kernel_posterior_fn=layers_util.default_mean_field_normal_fn(), + kernel_posterior_tensor_fn=lambda d: d.sample(), + kernel_prior_fn=lambda dtype, *args: normal_lib.Normal( # pylint: disable=g-long-lambda + loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)), + kernel_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + bias_posterior_fn=layers_util.default_mean_field_normal_fn(is_singular=True), # pylint: disable=line-too-long + bias_posterior_tensor_fn=lambda d: d.sample(), + bias_prior_fn=None, + bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + seed=None, + name=None, + **kwargs): + super(Conv2DFlipout, self).__init__( + rank=2, + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activation, + activity_regularizer=activity_regularizer, + trainable=trainable, + kernel_posterior_fn=kernel_posterior_fn, + kernel_posterior_tensor_fn=kernel_posterior_tensor_fn, + kernel_prior_fn=kernel_prior_fn, + kernel_divergence_fn=kernel_divergence_fn, + bias_posterior_fn=bias_posterior_fn, + bias_posterior_tensor_fn=bias_posterior_tensor_fn, + bias_prior_fn=bias_prior_fn, + bias_divergence_fn=bias_divergence_fn, + seed=seed, + name=name, **kwargs) + + +def conv2d_flipout( + inputs, + filters, + kernel_size, + strides=(1, 1), + padding="valid", + data_format="channels_last", + dilation_rate=(1, 1), + activation=None, + activity_regularizer=None, + trainable=True, + kernel_posterior_fn=layers_util.default_mean_field_normal_fn(), + kernel_posterior_tensor_fn=lambda d: d.sample(), + kernel_prior_fn=lambda dtype, *args: normal_lib.Normal( # pylint: disable=g-long-lambda + loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)), + kernel_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + bias_posterior_fn=layers_util.default_mean_field_normal_fn(is_singular=True), # pylint: disable=line-too-long + bias_posterior_tensor_fn=lambda d: d.sample(), + bias_prior_fn=None, + bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + seed=None, + name=None, + reuse=None): + """Functional interface for the 2D convolution layer. + + This layer creates a convolution kernel that is convolved + (actually cross-correlated) with the layer input to produce a tensor of + outputs. It may also include a bias addition and activation function + on the outputs. It assumes the `kernel` and/or `bias` are drawn from + distributions. + + By default, the layer implements a stochastic forward pass via + sampling from the kernel and bias posteriors, + ```none + outputs = f(inputs; kernel, bias), kernel, bias ~ posterior + ``` + where f denotes the layer's calculation. It uses the Flipout + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. Flipout uses + roughly twice as many floating point operations as the + reparameterization estimator but has the advantage of significantly + lower variance. + + The arguments permit separate specification of the surrogate posterior + (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` + distributions. + + Arguments: + inputs: Tensor input. + filters: Integer, the dimensionality of the output space (i.e. the number + of filters in the convolution). + kernel_size: An integer or tuple/list of 2 integers, specifying the + height and width of the 2D convolution window. + Can be a single integer to specify the same value for + all spatial dimensions. + strides: An integer or tuple/list of 2 integers, + specifying the strides of the convolution along the height and width. + Can be a single integer to specify the same value for + all spatial dimensions. + Specifying any stride value != 1 is incompatible with specifying + any `dilation_rate` value != 1. + padding: One of `"valid"` or `"same"` (case-insensitive). + data_format: A string, one of `channels_last` (default) or `channels_first`. + The ordering of the dimensions in the inputs. + `channels_last` corresponds to inputs with shape + `(batch, height, width, channels)` while `channels_first` corresponds to + inputs with shape `(batch, channels, height, width)`. + + dilation_rate: An integer or tuple/list of 2 integers, specifying + the dilation rate to use for dilated convolution. + Can be a single integer to specify the same value for + all spatial dimensions. + Currently, specifying any `dilation_rate` value != 1 is + incompatible with specifying any stride value != 1. + activation: Activation function. Set it to None to maintain a + linear activation. + activity_regularizer: Optional regularizer function for the output. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + kernel_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `kernel` parameter. Default value: + `default_mean_field_normal_fn()`. + kernel_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + kernel_prior_fn: Python `callable` which creates `tf.distributions` + instance. See `default_mean_field_normal_fn` docstring for required + parameter signature. + Default value: `tf.distributions.Normal(loc=0., scale=1.)`. + kernel_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + bias_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `bias` parameter. Default value: + `default_mean_field_normal_fn(is_singular=True)` (which creates an + instance of `tf.distributions.Deterministic`). + bias_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + bias_prior_fn: Python `callable` which creates `tf.distributions` instance. + See `default_mean_field_normal_fn` docstring for required parameter + signature. Default value: `None` (no prior, no variational inference) + bias_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + seed: Python scalar `int` which initializes the random number + generator. Default value: `None` (i.e., use global seed). + name: A string, the name of the layer. + reuse: Boolean, whether to reuse the weights of a previous layer + by the same name. + + Returns: + Output tensor. + + Raises: + ValueError: if eager execution is enabled. + + #### Examples + + We illustrate a Bayesian neural network with [variational inference]( + https://en.wikipedia.org/wiki/Variational_Bayesian_methods), + assuming a dataset of `features` and `labels`. + + ```python + tfp = tf.contrib.bayesflow + + net = tf.reshape(features, [-1, 32, 32, 3]) + net = tfp.layers.conv2d_flipout(net, + filters=64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu) + net = tf.layers.max_pooling2d(net, + pool_size=2, + strides=2, + padding="SAME") + net = tf.reshape(net, [-1, 8 * 8 * 64]) + logits = tfp.layers.dense_flipout(net, 10) + neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( + labels=labels, logits=logits) + kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) + loss = neg_log_likelihood + kl + train_op = tf.train.AdamOptimizer().minimize(loss) + ``` + + It uses the Flipout gradient estimator to minimize the + Kullback-Leibler divergence up to a constant, also known as the + negative Evidence Lower Bound. It consists of the sum of two terms: + the expected negative log-likelihood, which we approximate via + Monte Carlo; and the KL divergence, which is added via regularizer + terms which are arguments to the layer. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb + """ + layer = Conv2DFlipout( + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activation, + activity_regularizer=activity_regularizer, + trainable=trainable, + kernel_posterior_fn=kernel_posterior_fn, + kernel_posterior_tensor_fn=kernel_posterior_tensor_fn, + kernel_prior_fn=kernel_prior_fn, + kernel_divergence_fn=kernel_divergence_fn, + bias_posterior_fn=bias_posterior_fn, + bias_posterior_tensor_fn=bias_posterior_tensor_fn, + bias_prior_fn=bias_prior_fn, + bias_divergence_fn=bias_divergence_fn, + seed=seed, + name=name, + dtype=inputs.dtype.base_dtype, + _scope=name, + _reuse=reuse) + return layer.apply(inputs) + + +class Conv3DFlipout(_ConvFlipout): + """3D convolution layer (e.g. spatial convolution over volumes) with Flipout. + + This layer creates a convolution kernel that is convolved + (actually cross-correlated) with the layer input to produce a tensor of + outputs. It may also include a bias addition and activation function + on the outputs. It assumes the `kernel` and/or `bias` are drawn from + distributions. + + By default, the layer implements a stochastic forward pass via + sampling from the kernel and bias posteriors, + ```none + outputs = f(inputs; kernel, bias), kernel, bias ~ posterior + ``` + where f denotes the layer's calculation. It uses the Flipout + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. Flipout uses + roughly twice as many floating point operations as the + reparameterization estimator but has the advantage of significantly + lower variance. + + The arguments permit separate specification of the surrogate posterior + (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` + distributions. + + Arguments: + filters: Integer, the dimensionality of the output space (i.e. the number + of filters in the convolution). + kernel_size: An integer or tuple/list of 3 integers, specifying the + depth, height and width of the 3D convolution window. + Can be a single integer to specify the same value for + all spatial dimensions. + strides: An integer or tuple/list of 3 integers, + specifying the strides of the convolution along the depth, + height and width. + Can be a single integer to specify the same value for + all spatial dimensions. + Specifying any stride value != 1 is incompatible with specifying + any `dilation_rate` value != 1. + padding: One of `"valid"` or `"same"` (case-insensitive). + data_format: A string, one of `channels_last` (default) or `channels_first`. + The ordering of the dimensions in the inputs. + `channels_last` corresponds to inputs with shape + `(batch, depth, height, width, channels)` while `channels_first` + corresponds to inputs with shape + `(batch, channels, depth, height, width)`. + dilation_rate: An integer or tuple/list of 3 integers, specifying + the dilation rate to use for dilated convolution. + Can be a single integer to specify the same value for + all spatial dimensions. + Currently, specifying any `dilation_rate` value != 1 is + incompatible with specifying any stride value != 1. + activation: Activation function. Set it to None to maintain a + linear activation. + activity_regularizer: Optional regularizer function for the output. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + kernel_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `kernel` parameter. Default value: + `default_mean_field_normal_fn()`. + kernel_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + kernel_prior_fn: Python `callable` which creates `tf.distributions` + instance. See `default_mean_field_normal_fn` docstring for required + parameter signature. + Default value: `tf.distributions.Normal(loc=0., scale=1.)`. + kernel_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + bias_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `bias` parameter. Default value: + `default_mean_field_normal_fn(is_singular=True)` (which creates an + instance of `tf.distributions.Deterministic`). + bias_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + bias_prior_fn: Python `callable` which creates `tf.distributions` instance. + See `default_mean_field_normal_fn` docstring for required parameter + signature. Default value: `None` (no prior, no variational inference) + bias_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + seed: Python scalar `int` which initializes the random number + generator. Default value: `None` (i.e., use global seed). + name: A string, the name of the layer. + + Properties: + filters: Python integer, dimensionality of the output space. + kernel_size: Size of the convolution window. + strides: Stride length of convolution. + padding: Python string describing padding approach. + data_format: Python string describing input data's dimensions. + dilation_rate: Dilation rate for an atrous convolution. + activation: Activation function (`callable`). + activity_regularizer: Regularizer function for the output. + kernel_posterior_fn: `callable` returning posterior. + kernel_posterior_tensor_fn: `callable` operating on posterior. + kernel_prior_fn: `callable` returning prior. + kernel_divergence_fn: `callable` returning divergence. + bias_posterior_fn: `callable` returning posterior. + bias_posterior_tensor_fn: `callable` operating on posterior. + bias_prior_fn: `callable` returning prior. + bias_divergence_fn: `callable` returning divergence. + seed: Python integer, used to create random seeds. + + #### Examples + + We illustrate a Bayesian neural network with [variational inference]( + https://en.wikipedia.org/wiki/Variational_Bayesian_methods), + assuming a dataset of `features` and `labels`. + + ```python + tfp = tf.contrib.bayesflow + + net = tf.reshape(features, [-1, 256, 32, 32, 3]) + net = tfp.layers.Conv3DFlipout(64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu)(net) + net = tf.layers.MaxPooling2D(pool_size=2, + strides=2, + padding="SAME")(net) + net = tf.reshape(net, [-1, 256 * 8 * 8 * 64]) + logits = tfp.layers.DenseFlipout(10)(net) + neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( + labels=labels, logits=logits) + kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) + loss = neg_log_likelihood + kl + train_op = tf.train.AdamOptimizer().minimize(loss) + ``` + + It uses the Flipout gradient estimator to minimize the + Kullback-Leibler divergence up to a constant, also known as the + negative Evidence Lower Bound. It consists of the sum of two terms: + the expected negative log-likelihood, which we approximate via + Monte Carlo; and the KL divergence, which is added via regularizer + terms which are arguments to the layer. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb + """ + + def __init__( + self, + filters, + kernel_size, + strides=(1, 1, 1), + padding="valid", + data_format="channels_last", + dilation_rate=(1, 1, 1), + activation=None, + activity_regularizer=None, + trainable=True, + kernel_posterior_fn=layers_util.default_mean_field_normal_fn(), + kernel_posterior_tensor_fn=lambda d: d.sample(), + kernel_prior_fn=lambda dtype, *args: normal_lib.Normal( # pylint: disable=g-long-lambda + loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)), + kernel_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + bias_posterior_fn=layers_util.default_mean_field_normal_fn(is_singular=True), # pylint: disable=line-too-long + bias_posterior_tensor_fn=lambda d: d.sample(), + bias_prior_fn=None, + bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + seed=None, + name=None, + **kwargs): + super(Conv3DFlipout, self).__init__( + rank=3, + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activation, + activity_regularizer=activity_regularizer, + trainable=trainable, + kernel_posterior_fn=kernel_posterior_fn, + kernel_posterior_tensor_fn=kernel_posterior_tensor_fn, + kernel_prior_fn=kernel_prior_fn, + kernel_divergence_fn=kernel_divergence_fn, + bias_posterior_fn=bias_posterior_fn, + bias_posterior_tensor_fn=bias_posterior_tensor_fn, + bias_prior_fn=bias_prior_fn, + bias_divergence_fn=bias_divergence_fn, + seed=seed, + name=name, **kwargs) + + +def conv3d_flipout( + inputs, + filters, + kernel_size, + strides=(1, 1, 1), + padding="valid", + data_format="channels_last", + dilation_rate=(1, 1, 1), + activation=None, + activity_regularizer=None, + trainable=True, + kernel_posterior_fn=layers_util.default_mean_field_normal_fn(), + kernel_posterior_tensor_fn=lambda d: d.sample(), + kernel_prior_fn=lambda dtype, *args: normal_lib.Normal( # pylint: disable=g-long-lambda + loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.)), + kernel_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + bias_posterior_fn=layers_util.default_mean_field_normal_fn(is_singular=True), # pylint: disable=line-too-long + bias_posterior_tensor_fn=lambda d: d.sample(), + bias_prior_fn=None, + bias_divergence_fn=lambda q, p, ignore: kl_lib.kl_divergence(q, p), + seed=None, + name=None, + reuse=None): + """Functional interface for the 3D convolution layer. + + This layer creates a convolution kernel that is convolved + (actually cross-correlated) with the layer input to produce a tensor of + outputs. It may also include a bias addition and activation function + on the outputs. It assumes the `kernel` and/or `bias` are drawn from + distributions. + + By default, the layer implements a stochastic forward pass via + sampling from the kernel and bias posteriors, + ```none + outputs = f(inputs; kernel, bias), kernel, bias ~ posterior + ``` + where f denotes the layer's calculation. It uses the Flipout + estimator [1], which performs a Monte Carlo approximation of the + distribution integrating over the `kernel` and `bias`. Flipout uses + roughly twice as many floating point operations as the + reparameterization estimator but has the advantage of significantly + lower variance. + + The arguments permit separate specification of the surrogate posterior + (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` + distributions. + + Arguments: + inputs: Tensor input. + filters: Integer, the dimensionality of the output space (i.e. the number + of filters in the convolution). + kernel_size: An integer or tuple/list of 3 integers, specifying the + depth, height and width of the 3D convolution window. + Can be a single integer to specify the same value for + all spatial dimensions. + strides: An integer or tuple/list of 3 integers, + specifying the strides of the convolution along the depth, + height and width. + Can be a single integer to specify the same value for + all spatial dimensions. + Specifying any stride value != 1 is incompatible with specifying + any `dilation_rate` value != 1. + padding: One of `"valid"` or `"same"` (case-insensitive). + data_format: A string, one of `channels_last` (default) or `channels_first`. + The ordering of the dimensions in the inputs. + `channels_last` corresponds to inputs with shape + `(batch, depth, height, width, channels)` while `channels_first` + corresponds to inputs with shape + `(batch, channels, depth, height, width)`. + dilation_rate: An integer or tuple/list of 3 integers, specifying + the dilation rate to use for dilated convolution. + Can be a single integer to specify the same value for + all spatial dimensions. + Currently, specifying any `dilation_rate` value != 1 is + incompatible with specifying any stride value != 1. + activation: Activation function. Set it to None to maintain a + linear activation. + activity_regularizer: Optional regularizer function for the output. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + kernel_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `kernel` parameter. Default value: + `default_mean_field_normal_fn()`. + kernel_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + kernel_prior_fn: Python `callable` which creates `tf.distributions` + instance. See `default_mean_field_normal_fn` docstring for required + parameter signature. + Default value: `tf.distributions.Normal(loc=0., scale=1.)`. + kernel_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + sample is a `Tensor`. + bias_posterior_fn: Python `callable` which creates + `tf.distributions.Distribution` instance representing the surrogate + posterior of the `bias` parameter. Default value: + `default_mean_field_normal_fn(is_singular=True)` (which creates an + instance of `tf.distributions.Deterministic`). + bias_posterior_tensor_fn: Python `callable` which takes a + `tf.distributions.Distribution` instance and returns a representative + value. Default value: `lambda d: d.sample()`. + bias_prior_fn: Python `callable` which creates `tf.distributions` instance. + See `default_mean_field_normal_fn` docstring for required parameter + signature. Default value: `None` (no prior, no variational inference) + bias_divergence_fn: Python `callable` which takes the surrogate posterior + distribution, prior distribution and random variate sample(s) from the + surrogate posterior and computes or approximates the KL divergence. The + distributions are `tf.distributions.Distribution`-like instances and the + seed: Python scalar `int` which initializes the random number + generator. Default value: `None` (i.e., use global seed). + name: A string, the name of the layer. + reuse: Boolean, whether to reuse the weights of a previous layer + by the same name. + + Returns: + Output tensor. + + Raises: + ValueError: if eager execution is enabled. + + #### Examples + + We illustrate a Bayesian neural network with [variational inference]( + https://en.wikipedia.org/wiki/Variational_Bayesian_methods), + assuming a dataset of `features` and `labels`. + + ```python + tfp = tf.contrib.bayesflow + + net = tf.reshape(features, [-1, 256, 32, 32, 3]) + net = tfp.layers.conv3d_flipout(net, + filters=64, + kernel_size=5, + padding="SAME", + activation=tf.nn.relu) + net = tf.layers.max_pooling2d(net, + pool_size=2, + strides=2, + padding="SAME") + net = tf.reshape(net, [-1, 256 * 8 * 8 * 64]) + logits = tfp.layers.dense_flipout(net, 10) + neg_log_likelihood = tf.nn.softmax_cross_entropy_with_logits( + labels=labels, logits=logits) + kl = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) + loss = neg_log_likelihood + kl + train_op = tf.train.AdamOptimizer().minimize(loss) + ``` + + It uses the Flipout gradient estimator to minimize the + Kullback-Leibler divergence up to a constant, also known as the + negative Evidence Lower Bound. It consists of the sum of two terms: + the expected negative log-likelihood, which we approximate via + Monte Carlo; and the KL divergence, which is added via regularizer + terms which are arguments to the layer. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb """ - layer = Conv3DVariational( + layer = Conv3DFlipout( filters=filters, kernel_size=kernel_size, strides=strides, @@ -1398,6 +2919,7 @@ def conv3d_variational( bias_posterior_tensor_fn=bias_posterior_tensor_fn, bias_prior_fn=bias_prior_fn, bias_divergence_fn=bias_divergence_fn, + seed=seed, name=name, dtype=inputs.dtype.base_dtype, _scope=name, @@ -1407,9 +2929,15 @@ def conv3d_variational( # Aliases -Convolution1DVariational = Conv1DVariational -Convolution2DVariational = Conv2DVariational -Convolution3DVariational = Conv3DVariational -convolution1d_variational = conv1d_variational -convolution2d_variational = conv2d_variational -convolution3d_variational = conv3d_variational +Convolution1DReparameterization = Conv1DReparameterization +Convolution2DReparameterization = Conv2DReparameterization +Convolution3DReparameterization = Conv3DReparameterization +convolution1d_reparameterization = conv1d_reparameterization +convolution2d_reparameterization = conv2d_reparameterization +convolution3d_reparameterization = conv3d_reparameterization +Convolution1DFlipout = Conv1DFlipout +Convolution2DFlipout = Conv2DFlipout +Convolution3DFlipout = Conv3DFlipout +convolution1d_flipout = conv1d_flipout +convolution2d_flipout = conv2d_flipout +convolution3d_flipout = conv3d_flipout diff --git a/tensorflow/contrib/bayesflow/python/ops/layers_dense_variational_impl.py b/tensorflow/contrib/bayesflow/python/ops/layers_dense_variational.py similarity index 95% rename from tensorflow/contrib/bayesflow/python/ops/layers_dense_variational_impl.py rename to tensorflow/contrib/bayesflow/python/ops/layers_dense_variational.py index a749a396f15188ef345b4ae7c53017b6004c5e71..591a8e553de0c194786c7ee8693665f762711b2d 100644 --- a/tensorflow/contrib/bayesflow/python/ops/layers_dense_variational_impl.py +++ b/tensorflow/contrib/bayesflow/python/ops/layers_dense_variational.py @@ -13,13 +13,6 @@ # limitations under the License. # ============================================================================== """Dense Bayesian layer using KL-divergence based variational inference. - -@@DenseReparameterization -@@DenseLocalReparameterization -@@DenseFlipout -@@dense_reparameterization -@@dense_local_reparameterization -@@dense_flipout """ from __future__ import absolute_import @@ -33,25 +26,13 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.layers import base as layers_lib from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn -from tensorflow.python.ops import random_ops from tensorflow.python.ops import standard_ops from tensorflow.python.ops.distributions import kullback_leibler as kl_lib from tensorflow.python.ops.distributions import normal as normal_lib from tensorflow.python.ops.distributions import util as distribution_util -__all__ = [ - "DenseReparameterization", - "DenseLocalReparameterization", - "DenseFlipout", - "dense_reparameterization", - "dense_local_reparameterization", - "dense_flipout", -] - - class _DenseVariational(layers_lib.Layer): """Abstract densely-connected class (private, used as implementation base). @@ -285,6 +266,10 @@ class DenseReparameterization(_DenseVariational): outputs = activation(matmul(inputs, kernel) + bias) ``` + It uses the reparameterization estimator [1], which performs a Monte Carlo + approximation of the distribution integrating over the `kernel` and + `bias`. + The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` distributions. @@ -372,6 +357,10 @@ class DenseReparameterization(_DenseVariational): the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. """ def __init__( @@ -445,6 +434,10 @@ def dense_reparameterization( outputs = activation(matmul(inputs, kernel) + bias) ``` + It uses the reparameterization estimator [1], which performs a Monte Carlo + approximation of the distribution integrating over the `kernel` and + `bias`. + The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` distributions. @@ -524,6 +517,10 @@ def dense_reparameterization( the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Auto-Encoding Variational Bayes." + Diederik P. Kingma, Max Welling. + International Conference on Learning Representations, 2014. """ layer = DenseReparameterization( units, @@ -558,6 +555,10 @@ class DenseLocalReparameterization(_DenseVariational): outputs = activation(matmul(inputs, kernel) + bias) ``` + It uses the local reparameterization estimator [1], which performs a + Monte Carlo approximation of the distribution on the hidden units + induced by the `kernel` and `bias`. + The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` distributions. @@ -645,6 +646,10 @@ class DenseLocalReparameterization(_DenseVariational): the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Variational Dropout and the Local Reparameterization Trick." + Diederik P. Kingma, Tim Salimans, Max Welling. + Neural Information Processing Systems, 2015. """ def __init__( @@ -688,7 +693,7 @@ class DenseLocalReparameterization(_DenseVariational): "`DenseLocalReparameterization` requires " "`kernel_posterior_fn` produce an instance of " "`tf.distributions.Independent(tf.distributions.Normal)` " - "(saw: \"{}\").".format(type(self.kernel_posterior).__name__)) + "(saw: \"{}\").".format(self.kernel_posterior.name)) self.kernel_posterior_affine = normal_lib.Normal( loc=self._matmul(inputs, self.kernel_posterior.distribution.loc), scale=standard_ops.sqrt(self._matmul( @@ -730,6 +735,10 @@ def dense_local_reparameterization( outputs = activation(matmul(inputs, kernel) + bias) ``` + It uses the local reparameterization estimator [1], which performs a + Monte Carlo approximation of the distribution on the hidden units + induced by the `kernel` and `bias`. + The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` distributions. @@ -809,6 +818,10 @@ def dense_local_reparameterization( the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Variational Dropout and the Local Reparameterization Trick." + Diederik P. Kingma, Tim Salimans, Max Welling. + Neural Information Processing Systems, 2015. """ layer = DenseLocalReparameterization( units, @@ -843,6 +856,12 @@ class DenseFlipout(_DenseVariational): outputs = activation(matmul(inputs, kernel) + bias) ``` + It uses the Flipout estimator [1], which performs a Monte Carlo + approximation of the distribution integrating over the `kernel` and + `bias`. Flipout uses roughly twice as many floating point operations + as the reparameterization estimator but has the advantage of + significantly lower variance. + The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` distributions. @@ -933,6 +952,11 @@ class DenseFlipout(_DenseVariational): the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb """ def __init__( @@ -978,7 +1002,7 @@ class DenseFlipout(_DenseVariational): "`DenseFlipout` requires " "`kernel_posterior_fn` produce an instance of " "`tf.distributions.Independent(tf.distributions.Normal)` " - "(saw: \"{}\").".format(type(self.kernel_posterior).__name__)) + "(saw: \"{}\").".format(self.kernel_posterior.name)) self.kernel_posterior_affine = normal_lib.Normal( loc=array_ops.zeros_like(self.kernel_posterior.distribution.loc), scale=self.kernel_posterior.distribution.scale) @@ -989,8 +1013,11 @@ class DenseFlipout(_DenseVariational): input_shape = array_ops.shape(inputs) batch_shape = input_shape[:-1] - sign_input = random_sign(input_shape, dtype=inputs.dtype, seed=self.seed) - sign_output = random_sign( + sign_input = layers_util.random_sign( + input_shape, + dtype=inputs.dtype, + seed=self.seed) + sign_output = layers_util.random_sign( array_ops.concat([batch_shape, array_ops.expand_dims(self.units, 0)], 0), dtype=inputs.dtype, @@ -1035,6 +1062,12 @@ def dense_flipout( outputs = activation(matmul(inputs, kernel) + bias) ``` + It uses the Flipout estimator [1], which performs a Monte Carlo + approximation of the distribution integrating over the `kernel` and + `bias`. Flipout uses roughly twice as many floating point operations + as the reparameterization estimator but has the advantage of + significantly lower variance. + The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` distributions. @@ -1116,6 +1149,11 @@ def dense_flipout( the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer. + + [1]: "Flipout: Efficient Pseudo-Independent Weight Perturbations on + Mini-Batches." + Anonymous. OpenReview, 2017. + https://openreview.net/forum?id=rJnpifWAb """ layer = DenseFlipout( units, @@ -1136,11 +1174,3 @@ def dense_flipout( _scope=name, _reuse=reuse) return layer.apply(inputs) - - -def random_sign(shape, dtype=dtypes.float32, seed=None): - """Draw values from {-1, 1} uniformly, i.e., Rademacher distribution.""" - random_bernoulli = random_ops.random_uniform(shape, minval=0, maxval=2, - dtype=dtypes.int32, - seed=seed) - return math_ops.cast(2 * random_bernoulli - 1, dtype) diff --git a/tensorflow/contrib/bayesflow/python/ops/layers_util.py b/tensorflow/contrib/bayesflow/python/ops/layers_util.py index 9a4fecf4e5dcb1e1008303b07b4f76d5e5ce557f..8c1fb203f7328e8260e49b4326d813fbe133613e 100644 --- a/tensorflow/contrib/bayesflow/python/ops/layers_util.py +++ b/tensorflow/contrib/bayesflow/python/ops/layers_util.py @@ -23,9 +23,12 @@ import numpy as np from tensorflow.contrib.distributions.python.ops import deterministic as deterministic_lib from tensorflow.contrib.distributions.python.ops import independent as independent_lib +from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops +from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops +from tensorflow.python.ops import random_ops from tensorflow.python.ops.distributions import normal as normal_lib @@ -178,3 +181,11 @@ def default_mean_field_normal_fn( return independent_lib.Independent( dist, reinterpreted_batch_ndims=reinterpreted_batch_ndims) return _fn + + +def random_sign(shape, dtype=dtypes.float32, seed=None): + """Draw values from {-1, 1} uniformly, i.e., Rademacher distribution.""" + random_bernoulli = random_ops.random_uniform(shape, minval=0, maxval=2, + dtype=dtypes.int32, + seed=seed) + return math_ops.cast(2 * random_bernoulli - 1, dtype) diff --git a/tensorflow/contrib/bayesflow/python/ops/mcmc_diagnostics.py b/tensorflow/contrib/bayesflow/python/ops/mcmc_diagnostics.py new file mode 100644 index 0000000000000000000000000000000000000000..f3a645eafc249d1c39e0d4a238ae7ec8755c78d8 --- /dev/null +++ b/tensorflow/contrib/bayesflow/python/ops/mcmc_diagnostics.py @@ -0,0 +1,32 @@ +# 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. +# ============================================================================== +"""Utilities for Markov Chain Monte Carlo (MCMC) sampling.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# go/tf-wildcard-import +# pylint: disable=wildcard-import +from tensorflow.contrib.bayesflow.python.ops.mcmc_diagnostics_impl import * +# pylint: enable=wildcard-import +from tensorflow.python.util.all_util import remove_undocumented + +_allowed_symbols = [ + "effective_sample_size", + "potential_scale_reduction", +] + +remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/bayesflow/python/ops/mcmc_diagnostics_impl.py b/tensorflow/contrib/bayesflow/python/ops/mcmc_diagnostics_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..0424b6952bc89ce7fe5b00b0135c9a5fe1faa8cf --- /dev/null +++ b/tensorflow/contrib/bayesflow/python/ops/mcmc_diagnostics_impl.py @@ -0,0 +1,400 @@ +# 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. +# ============================================================================== +"""Utilities for Markov Chain Monte Carlo (MCMC) sampling. + +@@effective_sample_size +@@potential_scale_reduction +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.distributions.python.ops import sample_stats +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops + +__all__ = [ + "effective_sample_size", + "potential_scale_reduction", +] + + +def effective_sample_size(states, + filter_threshold=0., + filter_beyond_lag=None, + name=None): + """Estimate a lower bound on effective sample size for each independent chain. + + Roughly speaking, "effective sample size" (ESS) is the size of an iid sample + with the same variance as `state`. + + More precisely, given a stationary sequence of possibly correlated random + variables `X_1, X_2,...,X_N`, each identically distributed ESS is the number + such that + + ```Variance{ N**-1 * Sum{X_i} } = ESS**-1 * Variance{ X_1 }.``` + + If the sequence is uncorrelated, `ESS = N`. In general, one should expect + `ESS <= N`, with more highly correlated sequences having smaller `ESS`. + + #### Example of using ESS to estimate standard error. + + ``` + tfd = tf.contrib.distributions + tfb = tf.contrib.bayesflow + + target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.]) + + # Get 1000 states from one chain. + states = tfb.hmc.sample_chain( + num_results=1000, + target_log_prob_fn=target.log_prob, + current_state=tf.constant([0., 0.]), + step_size=0.05, + num_leapfrog_steps=20, + num_burnin_steps=200) + states.shape + ==> (1000, 2) + + ess = effective_sample_size(states) + ==> Shape (2,) Tensor + + mean, variance = tf.nn.moments(states, axis=0) + standard_error = tf.sqrt(variance / ess) + ``` + + Some math shows that, with `R_k` the auto-correlation sequence, + `R_k := Covariance{X_1, X_{1+k}} / Variance{X_1}`, we have + + ```ESS(N) = N / [ 1 + 2 * ( (N - 1) / N * R_1 + ... + 1 / N * R_{N-1} ) ]``` + + This function estimates the above by first estimating the auto-correlation. + Since `R_k` must be estimated using only `N - k` samples, it becomes + progressively noisier for larger `k`. For this reason, the summation over + `R_k` should be truncated at some number `filter_beyond_lag < N`. Since many + MCMC methods generate chains where `R_k > 0`, a reasonable critera is to + truncate at the first index where the estimated auto-correlation becomes + negative. + + The arguments `filter_beyond_lag`, `filter_threshold` are filters intended to + remove noisy tail terms from `R_k`. They combine in an "OR" manner meaning + terms are removed if they were to be filtered under the `filter_beyond_lag` OR + `filter_threshold` criteria. + + Args: + states: `Tensor` or list of `Tensor` objects. Dimension zero should index + identically distributed states. + filter_threshold: `Tensor` or list of `Tensor` objects. + Must broadcast with `state`. The auto-correlation sequence is truncated + after the first appearance of a term less than `filter_threshold`. + Setting to `None` means we use no threshold filter. Since `|R_k| <= 1`, + setting to any number less than `-1` has the same effect. + filter_beyond_lag: `Tensor` or list of `Tensor` objects. Must be + `int`-like and scalar valued. The auto-correlation sequence is truncated + to this length. Setting to `None` means we do not filter based on number + of lags. + name: `String` name to prepend to created ops. + + Returns: + ess: `Tensor` or list of `Tensor` objects. The effective sample size of + each component of `states`. Shape will be `states.shape[1:]`. + + Raises: + ValueError: If `states` and `filter_threshold` or `states` and + `filter_beyond_lag` are both lists with different lengths. + """ + states_was_list = _is_list_like(states) + + # Convert all args to lists. + if not states_was_list: + states = [states] + + filter_beyond_lag = _broadcast_maybelist_arg(states, filter_beyond_lag, + "filter_beyond_lag") + filter_threshold = _broadcast_maybelist_arg(states, filter_threshold, + "filter_threshold") + + # Process items, one at a time. + with ops.name_scope(name, "effective_sample_size"): + ess_list = [ + _effective_sample_size_single_state(s, ml, mlt) + for (s, ml, mlt) in zip(states, filter_beyond_lag, filter_threshold) + ] + + if states_was_list: + return ess_list + return ess_list[0] + + +def _effective_sample_size_single_state(states, filter_beyond_lag, + filter_threshold): + """ESS computation for one single Tensor argument.""" + + with ops.name_scope( + "effective_sample_size_single_state", + values=[states, filter_beyond_lag, filter_threshold]): + + states = ops.convert_to_tensor(states, name="states") + dt = states.dtype + + # filter_beyond_lag == None ==> auto_corr is the full sequence. + auto_corr = sample_stats.auto_correlation( + states, axis=0, max_lags=filter_beyond_lag) + if filter_threshold is not None: + filter_threshold = ops.convert_to_tensor( + filter_threshold, dtype=dt, name="filter_threshold") + # Get a binary mask to zero out values of auto_corr below the threshold. + # mask[i, ...] = 1 if auto_corr[j, ...] > threshold for all j <= i, + # mask[i, ...] = 0, otherwise. + # So, along dimension zero, the mask will look like [1, 1, ..., 0, 0,...] + # Building step by step, + # Assume auto_corr = [1, 0.5, 0.0, 0.3], and filter_threshold = 0.2. + # Step 1: mask = [False, False, True, False] + mask = auto_corr < filter_threshold + # Step 2: mask = [0, 0, 1, 1] + mask = math_ops.cast(mask, dtype=dt) + # Step 3: mask = [0, 0, 1, 2] + mask = math_ops.cumsum(mask, axis=0) + # Step 4: mask = [1, 1, 0, 0] + mask = math_ops.maximum(1. - mask, 0.) + auto_corr *= mask + + # With R[k] := auto_corr[k, ...], + # ESS = N / {1 + 2 * Sum_{k=1}^N (N - k) / N * R[k]} + # = N / {-1 + 2 * Sum_{k=0}^N (N - k) / N * R[k]} (since R[0] = 1) + # approx N / {-1 + 2 * Sum_{k=0}^M (N - k) / N * R[k]} + # where M is the filter_beyond_lag truncation point chosen above. + + # Get the factor (N - k) / N, and give it shape [M, 1,...,1], having total + # ndims the same as auto_corr + n = _axis_size(states, axis=0) + k = math_ops.range(0., _axis_size(auto_corr, axis=0)) + nk_factor = (n - k) / n + if auto_corr.shape.ndims is not None: + new_shape = [-1] + [1] * (auto_corr.shape.ndims - 1) + else: + new_shape = array_ops.concat( + ([-1], + array_ops.ones([array_ops.rank(auto_corr) - 1], dtype=dtypes.int32)), + axis=0) + nk_factor = array_ops.reshape(nk_factor, new_shape) + + return n / (-1 + 2 * math_ops.reduce_sum(nk_factor * auto_corr, axis=0)) + + +def potential_scale_reduction(chains_states, + independent_chain_ndims=1, + name=None): + """Gelman and Rubin's potential scale reduction factor for chain convergence. + + Given `N > 1` states from each of `C > 1` independent chains, the potential + scale reduction factor, commonly referred to as R-hat, measures convergence of + the chains (to the same target) by testing for equality of means. + Specifically, R-hat measures the degree to which variance (of the means) + between chains exceeds what one would expect if the chains were identically + distributed. See [1], [2]. + + Some guidelines: + + * The initial state of the chains should be drawn from a distribution + overdispersed with respect to the target. + * If all chains converge to the target, then as `N --> infinity`, R-hat --> 1. + Before that, R-hat > 1 (except in pathological cases, e.g. if the chain + paths were identical). + * The above holds for any number of chains `C > 1`. Increasing `C` does + improves effectiveness of the diagnostic. + * Sometimes, R-hat < 1.2 is used to indicate approximate convergence, but of + course this is problem depedendent. See [2]. + * R-hat only measures non-convergence of the mean. If higher moments, or other + statistics are desired, a different diagnostic should be used. See [2]. + + #### Examples + + Diagnosing convergence by monitoring 10 chains that each attempt to + sample from a 2-variate normal. + + ```python + tfd = tf.contrib.distributions + tfb = tf.contrib.bayesflow + + target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.]) + + # Get 10 (2x) overdispersed initial states. + initial_state = target.sample(10) * 2. + ==> (10, 2) + + # Get 1000 samples from the 10 independent chains. + chains_states, _ = tfb.hmc.sample_chain( + num_results=1000, + target_log_prob_fn=target.log_prob, + current_state=initial_state, + step_size=0.05, + num_leapfrog_steps=20, + num_burnin_steps=200) + chains_states.shape + ==> (1000, 10, 2) + + rhat = tfb.mcmc_diagnostics.potential_scale_reduction( + chains_states, independent_chain_ndims=1) + + # The second dimension needed a longer burn-in. + rhat.eval() + ==> [1.05, 1.3] + ``` + + To see why R-hat is reasonable, let `X` be a random variable drawn uniformly + from the combined states (combined over all chains). Then, in the limit + `N, C --> infinity`, with `E`, `Var` denoting expectation and variance, + + ```R-hat = ( E[Var[X | chain]] + Var[E[X | chain]] ) / E[Var[X | chain]].``` + + Using the law of total variance, the numerator is the variance of the combined + states, and the denominator is the total variance minus the variance of the + the individual chain means. If the chains are all drawing from the same + distribution, they will have the same mean, and thus the ratio should be one. + + [1] "Inference from Iterative Simulation Using Multiple Sequences" + Andrew Gelman and Donald B. Rubin + Statist. Sci. Volume 7, Number 4 (1992), 457-472. + [2] "General Methods for Monitoring Convergence of Iterative Simulations" + Stephen P. Brooks and Andrew Gelman + Journal of Computational and Graphical Statistics, 1998. Vol 7, No. 4. + + Args: + chains_states: `Tensor` or Python `list` of `Tensor`s representing the + state(s) of a Markov Chain at each result step. The `ith` state is + assumed to have shape `[Ni, Ci1, Ci2,...,CiD] + A`. + Dimension `0` indexes the `Ni > 1` result steps of the Markov Chain. + Dimensions `1` through `D` index the `Ci1 x ... x CiD` independent + chains to be tested for convergence to the same target. + The remaining dimensions, `A`, can have any shape (even empty). + independent_chain_ndims: Integer type `Tensor` with value `>= 1` giving the + number of giving the number of dimensions, from `dim = 1` to `dim = D`, + holding independent chain results to be tested for convergence. + name: `String` name to prepend to created ops. Default: + `potential_scale_reduction`. + + Returns: + `Tensor` or Python `list` of `Tensor`s representing the R-hat statistic for + the state(s). Same `dtype` as `state`, and shape equal to + `state.shape[1 + independent_chain_ndims:]`. + + Raises: + ValueError: If `independent_chain_ndims < 1`. + """ + chains_states_was_list = _is_list_like(chains_states) + if not chains_states_was_list: + chains_states = [chains_states] + + # tensor_util.constant_value returns None iff a constant value (as a numpy + # array) is not efficiently computable. Therefore, we try constant_value then + # check for None. + icn_const_ = tensor_util.constant_value( + ops.convert_to_tensor(independent_chain_ndims)) + if icn_const_ is not None: + independent_chain_ndims = icn_const_ + if icn_const_ < 1: + raise ValueError( + "Argument `independent_chain_ndims` must be `>= 1`, found: {}".format( + independent_chain_ndims)) + + with ops.name_scope(name, "potential_scale_reduction"): + rhat_list = [ + _potential_scale_reduction_single_state(s, independent_chain_ndims) + for s in chains_states + ] + + if chains_states_was_list: + return rhat_list + return rhat_list[0] + + +def _potential_scale_reduction_single_state(state, independent_chain_ndims): + """potential_scale_reduction for one single state `Tensor`.""" + with ops.name_scope( + "potential_scale_reduction_single_state", + values=[state, independent_chain_ndims]): + # We assume exactly one leading dimension indexes e.g. correlated samples + # from each Markov chain. + state = ops.convert_to_tensor(state, name="state") + sample_ndims = 1 + + sample_axis = math_ops.range(0, sample_ndims) + chain_axis = math_ops.range(sample_ndims, + sample_ndims + independent_chain_ndims) + sample_and_chain_axis = math_ops.range( + 0, sample_ndims + independent_chain_ndims) + + n = _axis_size(state, sample_axis) + m = _axis_size(state, chain_axis) + + # In the language of [2], + # B / n is the between chain variance, the variance of the chain means. + # W is the within sequence variance, the mean of the chain variances. + b_div_n = _reduce_variance( + math_ops.reduce_mean(state, sample_axis, keepdims=True), + sample_and_chain_axis, + biased=False) + w = math_ops.reduce_mean( + _reduce_variance(state, sample_axis, keepdims=True, biased=True), + sample_and_chain_axis) + + # sigma^2_+ is an estimate of the true variance, which would be unbiased if + # each chain was drawn from the target. c.f. "law of total variance." + sigma_2_plus = w + b_div_n + + return ((m + 1.) / m) * sigma_2_plus / w - (n - 1.) / (m * n) + + +# TODO(b/72873233) Move some variant of this to sample_stats. +def _reduce_variance(x, axis=None, biased=True, keepdims=False): + with ops.name_scope("reduce_variance"): + x = ops.convert_to_tensor(x, name="x") + mean = math_ops.reduce_mean(x, axis=axis, keepdims=True) + biased_var = math_ops.reduce_mean( + math_ops.squared_difference(x, mean), axis=axis, keepdims=keepdims) + if biased: + return biased_var + n = _axis_size(x, axis) + return (n / (n - 1.)) * biased_var + + +def _axis_size(x, axis=None): + """Get number of elements of `x` in `axis`, as type `x.dtype`.""" + if axis is None: + return math_ops.cast(array_ops.size(x), x.dtype) + return math_ops.cast( + math_ops.reduce_prod(array_ops.gather(array_ops.shape(x), axis)), x.dtype) + + +def _is_list_like(x): + """Helper which returns `True` if input is `list`-like.""" + return isinstance(x, (tuple, list)) + + +def _broadcast_maybelist_arg(states, secondary_arg, name): + """Broadcast a listable secondary_arg to that of states.""" + if _is_list_like(secondary_arg): + if len(secondary_arg) != len(states): + raise ValueError("Argument `%s` was a list of different length ({}) than " + "`states` ({})".format(name, len(states))) + else: + secondary_arg = [secondary_arg] * len(states) + + return secondary_arg diff --git a/tensorflow/contrib/bayesflow/python/ops/variable_utils.py b/tensorflow/contrib/bayesflow/python/ops/variable_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eadf6f4d5fa1c776e2c71c66c4b64b8f5ac98359 --- /dev/null +++ b/tensorflow/contrib/bayesflow/python/ops/variable_utils.py @@ -0,0 +1,29 @@ +# 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. +# ============================================================================== +"""Utility functions related to managing `tf.Variable`s.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# go/tf-wildcard-import +from tensorflow.contrib.bayesflow.python.ops.variable_utils_impl import * # pylint: disable=wildcard-import,unused-wildcard-import,g-importing-member +from tensorflow.python.util import all_util + +_allowed_symbols = [ + "externalize_variables_as_args", +] + +all_util.remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/bayesflow/python/ops/variable_utils_impl.py b/tensorflow/contrib/bayesflow/python/ops/variable_utils_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..ca3d75b5bfee093449026c7d1d62e3bdeff6b096 --- /dev/null +++ b/tensorflow/contrib/bayesflow/python/ops/variable_utils_impl.py @@ -0,0 +1,157 @@ +# 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. +# ============================================================================== +"""Utility functions related to managing `tf.Variable`s. + +@@externalize_variables_as_args +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import warnings + +from tensorflow.python.framework import ops +from tensorflow.python.ops import gradients_impl as gradients_ops +from tensorflow.python.ops import variable_scope as varscope_ops +from tensorflow.python.ops import variables as variables_ops + +__all__ = [ + "externalize_variables_as_args", +] + + +# Cause all warnings to always be triggered. +# Not having this means subsequent calls wont trigger the warning. +warnings.simplefilter("always") + + +def externalize_variables_as_args(fn, + fn_args=(), + ancestor_variables=None, + possible_ancestor_vars=None, + assert_variable_override=False, + name=None): + """"Converts variables within a callable into explicit args. + + Makes a new callable from `fn` which has arguments `list(fn_args) + + list(ancestor_variables)`. If `ancestor_variables` is not specified, it is + inferred by checking which of `possible_ancestor_vars` actually influences the + return value of `fn` (concretely, gradient of `fn(*fn_args)` is not `None`). + By default `possible_ancestor_vars` is `tf.trainable_variables() + + tf.get_collection(tf.GraphKeys.TRAINABLE_RESOURCE_VARIABLES)`. + + #### Examples: + + ```python + num_samples = 2 + num_dims = 1 + dtype = np.float32 + + def foo(x): + x = tf.convert_to_tensor(x, dtype=dtype, name="x") + s = x.shape.as_list() + y = tf.get_variable( + name="y", + dtype=dtype, + initializer=np.arange(np.prod(s)).reshape(s).astype(dtype)) + return x + y + + x = tf.constant(dtype([0.1, 0.2])) + + wrapped_foo, discovered_ancestor_variables = ( + externalize_variables_as_args(foo, [x])) + + new_x = dtype([[1.], [2.]]) + new_y = dtype([[3.], [4.]]) + new_result = wrapped_foo(new_x, new_y) + # ==> [[4.], [6.]] + + discovered_ancestor_variables == [tf.get_variable("y", dtype)] + # ==> [True] + ``` + + Args: + fn: Python callable which returns a `Tensor` and accepts `*fn_args`. + fn_args: Python list of args to `fn`. Represents dummy arguments passed to + `fn` to trace its execution; actual values are unimportant. These args are + only used to construct the output of `fn` and to resolve the ancestor + `tf.Variable`s. + Default value: `()` (i.e., `fn` takes no args). + ancestor_variables: Python list of `tf.Variable`s. When `None` the list is + expanded to non-`None` gradients of `fn(*fn_args)`. By directly providing + the `ancestor_variables` the internal call to `fn` is avoided. + Default value: `None` (i.e., `tf.Variable` dependencies are discovered). + possible_ancestor_vars: Python list of possible `tf.Variable`s which might + be a dependency of computing `fn(*fn_args)`. + Default value: `None` (i.e., expanded as described above). + assert_variable_override: Python `bool` indicating that not finding a + `tf.Variable` in the override list is an exception. + Default value: `False` (i.e., missing a `Variable` triggers a `warning`). + name: Python `str` name prefixed to Ops created by this function. + Default value: `None` (i.e., "externalize_variables_as_args"). + + Returns: + wrapped_fn: Python callable taking arguments like + `*(list(fn_args) + discovered_ancestor_variables)`. + discovered_ancestor_variables: Python list of `tf.Variable`s known to be a + dependency of `fn(*fn_args)`. + + Raises: + ValueError: if `assert_variable_override` is `True` and `Variable` is + requested but not overridden. + """ + def _make_bypassing_custom_getter_fn(new_var_dict): + """Return dict value rather than what would otherwise be dict key.""" + def _custom_getter(getter, *args, **kwargs): + v = getter(*args, **kwargs) + new_v = new_var_dict.get(v, None) + if new_v is None: + msg = "Variable \"{}\" not found in bypass dict.".format(v) + if assert_variable_override: + raise ValueError(msg) + warnings.warn(msg) + return v + return new_v + return _custom_getter + + with ops.name_scope(name, "externalize_variables_as_args"): + if ancestor_variables is not None and not ancestor_variables: + return fn, () + if ancestor_variables is None: + y = fn(*fn_args) # Side-effect: adds trainable vars. + if possible_ancestor_vars is None: + possible_ancestor_vars = ( + variables_ops.trainable_variables() + + ops.get_collection(ops.GraphKeys.TRAINABLE_RESOURCE_VARIABLES)) + # TODO(b/72873296): Add a dedicated op for identifying ancestors. + ancestors = [v for g, v + in zip(gradients_ops.gradients(y, possible_ancestor_vars), + possible_ancestor_vars) + if g is not None] + ancestor_variables = sorted(ancestors, key=lambda v: v.name) + n = len(fn_args) + def _fn(*args): + with ops.name_scope("wrapped_fn"): + vars_dict = dict( + (k, ops.convert_to_tensor( + v, dtype=k.dtype.base_dtype, name=k.op.name)) + for k, v in zip(ancestor_variables, args[n:])) + with varscope_ops.variable_scope( + varscope_ops.get_variable_scope(), + reuse=True, + custom_getter=_make_bypassing_custom_getter_fn(vars_dict)): + return fn(*args[:n]) + return _fn, ancestor_variables diff --git a/tensorflow/contrib/boosted_trees/BUILD b/tensorflow/contrib/boosted_trees/BUILD index 392ac7fa1ce600a64ee3b941b70b01447645e4aa..6fdcd0f996ee011842a5add79f06264a28a2145c 100644 --- a/tensorflow/contrib/boosted_trees/BUILD +++ b/tensorflow/contrib/boosted_trees/BUILD @@ -196,6 +196,7 @@ py_test( name = "quantile_ops_test", size = "small", srcs = ["python/kernel_tests/quantile_ops_test.py"], + shard_count = 3, srcs_version = "PY2AND3", deps = [ ":quantile_ops_py", diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/BUILD b/tensorflow/contrib/boosted_trees/estimator_batch/BUILD index 2c11dd981506334bcc98de6695d57c42a58f902a..289f5bb3140974d8c37f4938ceef27275b099f9a 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/BUILD +++ b/tensorflow/contrib/boosted_trees/estimator_batch/BUILD @@ -152,7 +152,11 @@ py_test( size = "small", srcs = ["dnn_tree_combined_estimator_test.py"], srcs_version = "PY2AND3", - tags = ["notsan"], + tags = [ + "no_gpu", + "no_pip_gpu", + "notsan", + ], deps = [ ":dnn_tree_combined_estimator", "//tensorflow/contrib/boosted_trees:gbdt_batch", diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py b/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py index 6ebc7d7911df878ec91701db8b75feb9a27d18a2..31f5c444817b9b82723c86bea3504d4934e57eb8 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py @@ -155,6 +155,9 @@ def convert_to_universal_format(dtec, sorted_feature_names, inequality_test.feature_id.id.value = ( _SPARSE_FLOAT_FEATURE_NAME_TEMPLATE % (sorted_feature_names[feature_id], split.dimension_id)) + model_and_features.features.pop(sorted_feature_names[feature_id]) + (model_and_features.features[inequality_test.feature_id.id.value] + .SetInParent()) inequality_test.type = ( generic_tree_model_pb2.InequalityTest.LESS_OR_EQUAL) inequality_test.threshold.float_value = split.threshold @@ -169,6 +172,9 @@ def convert_to_universal_format(dtec, sorted_feature_names, inequality_test.feature_id.id.value = ( _SPARSE_FLOAT_FEATURE_NAME_TEMPLATE % (sorted_feature_names[feature_id], split.dimension_id)) + model_and_features.features.pop(sorted_feature_names[feature_id]) + (model_and_features.features[inequality_test.feature_id.id.value] + .SetInParent()) inequality_test.type = ( generic_tree_model_pb2.InequalityTest.LESS_OR_EQUAL) inequality_test.threshold.float_value = split.threshold diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy_test.py b/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy_test.py index 492d9ca40c5cfa84e186020605429aacc02af6a6..67ec0e16bf815e9dbea6567cc87c3980a825a004 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy_test.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy_test.py @@ -150,8 +150,8 @@ class ConvertModelTest(test_util.TensorFlowTestCase): dtec, feature_columns, 1, 2, 1) # Features a and a_m are sparse float features, a_m is multidimensional. expected_tree = """ - features { key: "feature_a" } - features { key: "feature_a_m" } + features { key: "feature_a_0" } + features { key: "feature_a_m_3" } features { key: "feature_b" } features { key: "feature_d" } model { diff --git a/tensorflow/contrib/boosted_trees/kernels/model_ops.cc b/tensorflow/contrib/boosted_trees/kernels/model_ops.cc index 4b5d5ba0de6c3995ee2da7a44ab0ba099cbf1b35..754b7bc3270d647fc381033b769eadd7b791771e 100644 --- a/tensorflow/contrib/boosted_trees/kernels/model_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/model_ops.cc @@ -48,8 +48,9 @@ class CreateTreeEnsembleVariableOp : public OpKernel { if (!result->InitFromSerialized(tree_ensemble_config_t->scalar()(), stamp_token)) { result->Unref(); - OP_REQUIRES(context, false, errors::InvalidArgument( - "Unable to parse tree ensemble config.")); + OP_REQUIRES( + context, false, + errors::InvalidArgument("Unable to parse tree ensemble config.")); } // Only create one, if one does not exist already. Report status for all diff --git a/tensorflow/contrib/boosted_trees/kernels/prediction_ops.cc b/tensorflow/contrib/boosted_trees/kernels/prediction_ops.cc index f8086b0c2bb93eae6af0336bbe33fc23f8fcde22..b3fe38614e05801b223f0c96f7a70ce7e432a70b 100644 --- a/tensorflow/contrib/boosted_trees/kernels/prediction_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/prediction_ops.cc @@ -47,8 +47,8 @@ namespace boosted_trees { using boosted_trees::learner::LearnerConfig; using boosted_trees::learner::LearningRateConfig; using boosted_trees::learner::LearningRateDropoutDrivenConfig; -using boosted_trees::models::MultipleAdditiveTrees; using boosted_trees::models::DecisionTreeEnsembleResource; +using boosted_trees::models::MultipleAdditiveTrees; using boosted_trees::utils::DropoutUtils; using boosted_trees::utils::TensorUtils; diff --git a/tensorflow/contrib/boosted_trees/kernels/quantile_ops.cc b/tensorflow/contrib/boosted_trees/kernels/quantile_ops.cc index 8600c8c53caa5fd4274ba6730fc764d8315d680c..0f4c2298f56be48bb32f52d5d44cff8afe284f1e 100644 --- a/tensorflow/contrib/boosted_trees/kernels/quantile_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/quantile_ops.cc @@ -36,22 +36,21 @@ namespace tensorflow { using ::boosted_trees::QuantileConfig; -using boosted_trees::utils::TensorUtils; using boosted_trees::QuantileStreamResource; +using boosted_trees::utils::TensorUtils; namespace { const char* const kExampleWeightsName = "example_weights"; const char* const kMaxElementsName = "max_elements"; -const char* const kHandleName = "handle"; const char* const kNextStampTokenName = "next_stamp_token"; const char* const kStampTokenName = "stamp_token"; const char* const kAreBucketsReadyName = "are_buckets_ready"; +const char* const kGenerateQuantiles = "generate_quantiles"; // Names for sparse arguments. const char* const kNumSparseFeaturesName = "num_sparse_features"; const char* const kSparseBucketsName = "sparse_buckets"; const char* const kSparseValuesName = "sparse_values"; const char* const kSparseIndicesName = "sparse_indices"; -const char* const kSparseStreamsStateName = "sparse_streams_state"; const char* const kSparseSummariesName = "sparse_summaries"; const char* const kSparseConfigName = "sparse_config"; const char* const kSparseOutputTensorName = "sparse_quantiles"; @@ -59,7 +58,6 @@ const char* const kSparseOutputTensorName = "sparse_quantiles"; const char* const kDenseBucketsName = "dense_buckets"; const char* const kDenseConfigName = "dense_config"; const char* const kDenseOutputTensorName = "dense_quantiles"; -const char* const kDenseStreamsStateName = "dense_streams_state"; const char* const kDenseSummariesName = "dense_summaries"; const char* const kDenseValuesName = "dense_values"; const char* const kNumDenseFeaturesName = "num_dense_features"; @@ -182,6 +180,16 @@ std::vector GenerateBoundaries(const QuantileStream& stream, return boundaries; } +// Generates quantiles on a finalized QuantileStream. +std::vector GenerateQuantiles(const QuantileStream& stream, + int num_quantiles) { + // Do not de-dup boundaries. Exactly num_quantiles+1 boundary values + // will be returned. + std::vector boundaries = stream.GenerateQuantiles(num_quantiles); + CHECK_EQ(boundaries.size(), num_quantiles + 1); + return boundaries; +} + // Copies quantiles to output list. void CopyBoundaries(OpKernelContext* const context, const std::vector& boundaries, const int64 index, @@ -224,6 +232,8 @@ class CreateQuantileAccumulatorOp : public OpKernel { OP_REQUIRES_OK(context, context->GetAttr(kNumQuantilesName, &num_quantiles_)); OP_REQUIRES_OK(context, context->GetAttr(kMaxElementsName, &max_elements_)); + OP_REQUIRES_OK(context, + context->GetAttr(kGenerateQuantiles, &generate_quantiles_)); } void Compute(OpKernelContext* context) override { @@ -231,9 +241,9 @@ class CreateQuantileAccumulatorOp : public OpKernel { // other exceptions. If one already exists, it unrefs the new one. const Tensor* stamp_token_t; OP_REQUIRES_OK(context, context->input(kStampTokenName, &stamp_token_t)); - auto result = - new QuantileStreamResource(epsilon_, num_quantiles_, max_elements_, - stamp_token_t->scalar()()); + auto result = new QuantileStreamResource(epsilon_, num_quantiles_, + max_elements_, generate_quantiles_, + stamp_token_t->scalar()()); auto status = CreateResource(context, HandleFromInput(context, 0), result); if (!status.ok() && status.code() != tensorflow::error::ALREADY_EXISTS) { OP_REQUIRES(context, false, status); @@ -246,6 +256,7 @@ class CreateQuantileAccumulatorOp : public OpKernel { // An upperbound on the number of enteries that the summaries might have // for a feature. int64 max_elements_; + bool generate_quantiles_; }; REGISTER_KERNEL_BUILDER(Name("CreateQuantileAccumulator").Device(DEVICE_CPU), @@ -373,7 +384,7 @@ class MakeQuantileSummariesOp : public OpKernel { protobuf::Arena arena; ::boosted_trees::QuantileSummaryState* summary_proto = protobuf::Arena::CreateMessage< - ::boosted_trees::QuantileSummaryState>(&arena); + ::boosted_trees::QuantileSummaryState>(&arena); const auto& summary = stream.GetFinalSummary(); CopySummaryToProto(summary, summary_proto); // Output to tensor. @@ -597,10 +608,15 @@ class QuantileAccumulatorFlushOp : public OpKernel { << "Passed stamp token: " << stamp_token << " " << "Current token: " << streams_resource->stamp(); QuantileStream* stream = streams_resource->stream(stamp_token); + bool generate_quantiles = streams_resource->generate_quantiles(); stream->Finalize(); + streams_resource->set_boundaries( stamp_token, - GenerateBoundaries(*stream, streams_resource->num_quantiles())); + generate_quantiles + ? GenerateQuantiles(*stream, streams_resource->num_quantiles()) + : GenerateBoundaries(*stream, streams_resource->num_quantiles())); + streams_resource->Reset(next_stamp_token); } }; diff --git a/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc b/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc index 18b4abd654ea3541d646a43ac901aca1a678446f..44a8ffaf4b2f5a9c11b3abc46ce55a18c80ad318 100644 --- a/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc @@ -34,10 +34,10 @@ namespace tensorflow { +using boosted_trees::learner::LearnerConfig_MultiClassStrategy; using boosted_trees::learner::SplitInfo; using boosted_trees::learner::stochastic::GradientStats; using boosted_trees::learner::stochastic::NodeStats; -using boosted_trees::learner::LearnerConfig_MultiClassStrategy; namespace { const int32 DUMMY_FEATURE_DIMENSION = -1; @@ -47,9 +47,8 @@ class BaseBuildSplitOp : public OpKernel { public: explicit BaseBuildSplitOp(OpKernelConstruction* const context) : OpKernel(context) { - OP_REQUIRES_OK( - context, - context->GetAttr("feature_column_group_id", &feature_column_group_id_)); + OP_REQUIRES_OK(context, context->GetAttr("feature_column_group_id", + &feature_column_group_id_)); OP_REQUIRES_OK(context, context->GetAttr("l1_regularization", &l1_regularization_)); OP_REQUIRES_OK(context, diff --git a/tensorflow/contrib/boosted_trees/kernels/stats_accumulator_ops.cc b/tensorflow/contrib/boosted_trees/kernels/stats_accumulator_ops.cc index a9a229c8ae0c26bba5f0a684dad7e546298577bb..90a0655201f8cb8df6fc6417cb51216dec91b4d7 100644 --- a/tensorflow/contrib/boosted_trees/kernels/stats_accumulator_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/stats_accumulator_ops.cc @@ -134,10 +134,9 @@ void SerializeScalarAccumulatorToOutput( OpKernelContext* context) { int64 num_slots = accumulator_resource.values().size(); Tensor* partition_ids_t = nullptr; - OP_REQUIRES_OK( - context, - context->allocate_output("output_partition_ids", TensorShape({num_slots}), - &partition_ids_t)); + OP_REQUIRES_OK(context, context->allocate_output("output_partition_ids", + TensorShape({num_slots}), + &partition_ids_t)); auto partition_ids = partition_ids_t->vec(); // Feature ids tensor has ids of feature columns and their dimensions. @@ -149,15 +148,14 @@ void SerializeScalarAccumulatorToOutput( Tensor* gradients_t = nullptr; OP_REQUIRES_OK( - context, - context->allocate_output("output_gradients", TensorShape({num_slots}), - &gradients_t)); + context, context->allocate_output( + "output_gradients", TensorShape({num_slots}), &gradients_t)); auto gradients = gradients_t->vec(); Tensor* hessians_t = nullptr; - OP_REQUIRES_OK(context, - context->allocate_output( - "output_hessians", TensorShape({num_slots}), &hessians_t)); + OP_REQUIRES_OK( + context, context->allocate_output("output_hessians", + TensorShape({num_slots}), &hessians_t)); auto hessians = hessians_t->vec(); int i = 0; @@ -177,10 +175,9 @@ void SerializeTensorAccumulatorToOutput( OpKernelContext* context) { int64 num_slots = accumulator_resource.values().size(); Tensor* partition_ids_t = nullptr; - OP_REQUIRES_OK( - context, - context->allocate_output("output_partition_ids", TensorShape({num_slots}), - &partition_ids_t)); + OP_REQUIRES_OK(context, context->allocate_output("output_partition_ids", + TensorShape({num_slots}), + &partition_ids_t)); auto partition_ids = partition_ids_t->vec(); Tensor* feature_ids_t = nullptr; @@ -202,9 +199,8 @@ void SerializeTensorAccumulatorToOutput( int64 num_hessian_elements = hessian_shape.num_elements(); hessian_shape.InsertDim(0, num_slots); Tensor* hessians_t = nullptr; - OP_REQUIRES_OK( - context, - context->allocate_output("output_hessians", hessian_shape, &hessians_t)); + OP_REQUIRES_OK(context, context->allocate_output("output_hessians", + hessian_shape, &hessians_t)); auto hessians = hessians_t->flat_outer_dims(); int i = 0; diff --git a/tensorflow/contrib/boosted_trees/kernels/training_ops.cc b/tensorflow/contrib/boosted_trees/kernels/training_ops.cc index c77d90e243c304ec8e9a10a0b63401f9bd825c3e..7f8dea1d3c2a04b725843f6e2932a0cdfbc7733c 100644 --- a/tensorflow/contrib/boosted_trees/kernels/training_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/training_ops.cc @@ -361,10 +361,27 @@ class GrowTreeEnsembleOp : public OpKernel { // Increment attempt stats. ensemble_resource->IncrementAttempts(); + // In case we want to do feature selection and we have reached the limit, + // build a list of handlers used so far to avoid adding new features. + std::vector allowed_handlers; + if (learner_config_.constraints().max_number_of_unique_feature_columns() > + 0) { + allowed_handlers = ensemble_resource->GetUsedHandlers(); + // TODO(soroush): We can disable handlers that are not going to be used to + // avoid unnecessary computations. + if (allowed_handlers.size() < + learner_config_.constraints() + .max_number_of_unique_feature_columns()) { + // We have not reached the limit yet. Empty the list of allow features + // which means we can keep adding new features. + allowed_handlers.clear(); + } + } + // Find best splits for each active partition. std::map best_splits; - FindBestSplitsPerPartition(context, partition_ids_list, gains_list, - splits_list, &best_splits); + FindBestSplitsPerPartition(context, allowed_handlers, partition_ids_list, + gains_list, splits_list, &best_splits); // No-op if no new splits can be considered. if (best_splits.empty()) { @@ -381,7 +398,8 @@ class GrowTreeEnsembleOp : public OpKernel { // Split tree nodes. for (auto& split_entry : best_splits) { - SplitTreeNode(split_entry.first, &split_entry.second, tree_config); + SplitTreeNode(split_entry.first, &split_entry.second, tree_config, + ensemble_resource); } // Post-prune finalized tree if needed. @@ -403,12 +421,20 @@ class GrowTreeEnsembleOp : public OpKernel { // Helper method which effectively does a reduce over all split candidates // and finds the best split for each partition. void FindBestSplitsPerPartition( - OpKernelContext* const context, const OpInputList& partition_ids_list, - const OpInputList& gains_list, const OpInputList& splits_list, + OpKernelContext* const context, + const std::vector& allowed_handlers, // Empty means all handlers. + const OpInputList& partition_ids_list, const OpInputList& gains_list, + const OpInputList& splits_list, std::map* best_splits) { // Find best split per partition going through every feature candidate. // TODO(salehay): Is this worth parallelizing? for (int64 handler_id = 0; handler_id < num_handlers_; ++handler_id) { + if (!allowed_handlers.empty()) { + if (!std::binary_search(allowed_handlers.begin(), + allowed_handlers.end(), handler_id)) { + continue; + } + } const auto& partition_ids = partition_ids_list[handler_id].vec(); const auto& gains = gains_list[handler_id].vec(); const auto& splits = splits_list[handler_id].vec(); @@ -592,8 +618,10 @@ class GrowTreeEnsembleOp : public OpKernel { // Helper method to split a tree node and append its respective // leaf children given the split candidate. - void SplitTreeNode(const int32 node_id, SplitCandidate* split, - boosted_trees::trees::DecisionTreeConfig* tree_config) { + void SplitTreeNode( + const int32 node_id, SplitCandidate* split, + boosted_trees::trees::DecisionTreeConfig* tree_config, + boosted_trees::models::DecisionTreeEnsembleResource* ensemble_resource) { // No-op if we have no real node. CHECK(node_id < tree_config->nodes_size()) << "Invalid node " << node_id << " to split."; @@ -633,6 +661,9 @@ class GrowTreeEnsembleOp : public OpKernel { // Replace node in tree. (*tree_config->mutable_nodes(node_id)) = *split->split_info.mutable_split_node(); + if (learner_config_.constraints().max_number_of_unique_feature_columns()) { + ensemble_resource->MaybeAddUsedHandler(split->handler_id); + } } void PruneTree(boosted_trees::trees::DecisionTreeConfig* tree_config) { diff --git a/tensorflow/contrib/boosted_trees/lib/learner/common/accumulators/class-partition-key.h b/tensorflow/contrib/boosted_trees/lib/learner/common/accumulators/class-partition-key.h index e1bef0278846e7ff6abc91e8c57f780af45e8b41..3c54868951a6db93a8b685c8da4dfc78996b7b1f 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/common/accumulators/class-partition-key.h +++ b/tensorflow/contrib/boosted_trees/lib/learner/common/accumulators/class-partition-key.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_CLASS_PARTITION_KEY_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_CLASS_PARTITION_KEY_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_CLASS_PARTITION_KEY_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_CLASS_PARTITION_KEY_H_ #include "tensorflow/core/lib/hash/hash.h" @@ -58,4 +58,4 @@ struct ClassPartitionKey { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_CLASS_PARTITION_KEY_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_CLASS_PARTITION_KEY_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/learner/common/accumulators/feature-stats-accumulator.h b/tensorflow/contrib/boosted_trees/lib/learner/common/accumulators/feature-stats-accumulator.h index 3814edb5675be74794a08e00becb649f8fc53fdb..ec4e7c52bb5f4536a50192e1b5fcc019dd7b2511 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/common/accumulators/feature-stats-accumulator.h +++ b/tensorflow/contrib/boosted_trees/lib/learner/common/accumulators/feature-stats-accumulator.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_FEATURE_STATS_ACCUMULATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_FEATURE_STATS_ACCUMULATOR_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_FEATURE_STATS_ACCUMULATOR_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_FEATURE_STATS_ACCUMULATOR_H_ #include #include @@ -79,4 +79,4 @@ class FeatureStatsAccumulator { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_FEATURE_STATS_ACCUMULATOR_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_ACCUMULATORS_FEATURE_STATS_ACCUMULATOR_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/learner/common/partitioners/example_partitioner.h b/tensorflow/contrib/boosted_trees/lib/learner/common/partitioners/example_partitioner.h index aed0d9fdac108dff4576cc1563dae420340387be..37a71037041445e6a6fcf6290015b93cffef1618 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/common/partitioners/example_partitioner.h +++ b/tensorflow/contrib/boosted_trees/lib/learner/common/partitioners/example_partitioner.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_PARTITIONERS_EXAMPLE_PARTITIONER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_PARTITIONERS_EXAMPLE_PARTITIONER_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_PARTITIONERS_EXAMPLE_PARTITIONER_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_PARTITIONERS_EXAMPLE_PARTITIONER_H_ #include #include "tensorflow/contrib/boosted_trees/lib/trees/decision_tree.h" @@ -50,4 +50,4 @@ class ExamplePartitioner { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_PARTITIONERS_EXAMPLE_PARTITIONER_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_PARTITIONERS_EXAMPLE_PARTITIONER_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/feature-split-candidate.h b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/feature-split-candidate.h index 339c2e0fded10e6a7b140da62e152e2868ffd164..382b85cf0b2c146f82fa79551c569b9c70d9b7a6 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/feature-split-candidate.h +++ b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/feature-split-candidate.h @@ -13,8 +13,8 @@ // limitations under the License. // // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_FEATURE_SPLIT_CANDIDATE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_FEATURE_SPLIT_CANDIDATE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_FEATURE_SPLIT_CANDIDATE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_FEATURE_SPLIT_CANDIDATE_H_ #include "tensorflow/contrib/boosted_trees/lib/learner/common/stats/split-stats.h" #include "tensorflow/contrib/boosted_trees/proto/tree_config.pb.h" @@ -58,4 +58,4 @@ struct FeatureSplitCandidate { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_FEATURE_SPLIT_CANDIDATE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_FEATURE_SPLIT_CANDIDATE_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/gradient-stats.h b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/gradient-stats.h index 34e3ddb777242553d62035a51f1aec33d0f9ba54..3dd03215d88abc223a2d081d11901ffd3fb7aaa9 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/gradient-stats.h +++ b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/gradient-stats.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_GRADIENT_STATS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_GRADIENT_STATS_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_GRADIENT_STATS_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_GRADIENT_STATS_H_ #include @@ -190,4 +190,4 @@ inline GradientStats operator-(const GradientStats& a, const GradientStats& b) { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_GRADIENT_STATS_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_GRADIENT_STATS_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/node-stats.h b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/node-stats.h index 642a183aec5c7e591579fa5ee91d45729bfb624d..cd925f6b65e569538212e9c26aef0abc8482960b 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/node-stats.h +++ b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/node-stats.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_NODE_STATS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_NODE_STATS_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_NODE_STATS_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_NODE_STATS_H_ #include "third_party/eigen3/Eigen/Core" #include "third_party/eigen3/Eigen/Eigenvalues" @@ -298,4 +298,4 @@ struct NodeStats { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_NODE_STATS_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_NODE_STATS_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/node-stats_test.cc b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/node-stats_test.cc index f867e77d3ef0609774628b2a9c36ca52bcf2a957..8bca132acfde9397942b198db9a8d4c0e4d74897 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/node-stats_test.cc +++ b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/node-stats_test.cc @@ -17,8 +17,8 @@ #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/platform/test.h" -using tensorflow::test::AsTensor; using std::vector; +using tensorflow::test::AsTensor; namespace tensorflow { namespace boosted_trees { diff --git a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/split-stats.h b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/split-stats.h index 054ccd9a8cd0be0c48b14cca013f15677deba900..81ee2774bdab91f492064455055181c56ef6a065 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/common/stats/split-stats.h +++ b/tensorflow/contrib/boosted_trees/lib/learner/common/stats/split-stats.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_SPLIT_STATS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_SPLIT_STATS_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_SPLIT_STATS_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_SPLIT_STATS_H_ #include @@ -81,4 +81,4 @@ struct SplitStats { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_SPLIT_STATS_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_LEARNER_COMMON_STATS_SPLIT_STATS_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/models/multiple_additive_trees.h b/tensorflow/contrib/boosted_trees/lib/models/multiple_additive_trees.h index ee29a8aa797b96d41ec2d77bf831ee287d5443e7..cc3dc226cdbc88fc7010ada1e7f0e6c0a3913c5f 100644 --- a/tensorflow/contrib/boosted_trees/lib/models/multiple_additive_trees.h +++ b/tensorflow/contrib/boosted_trees/lib/models/multiple_additive_trees.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_MODELS_MULTIPLE_ADDITIVE_TREES_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_MODELS_MULTIPLE_ADDITIVE_TREES_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_MODELS_MULTIPLE_ADDITIVE_TREES_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_MODELS_MULTIPLE_ADDITIVE_TREES_H_ #include @@ -45,4 +45,4 @@ class MultipleAdditiveTrees { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_MODELS_MULTIPLE_ADDITIVE_TREES_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_MODELS_MULTIPLE_ADDITIVE_TREES_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_buffer.h b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_buffer.h index 70037d5bd8f446bdbbfcc468edb8a76c05e4fab7..804b218f1c08338df80f8dd2e6135f5d92b9928e 100644 --- a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_buffer.h +++ b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_buffer.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_BUFFER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_BUFFER_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_BUFFER_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_BUFFER_H_ #include #include @@ -129,4 +129,4 @@ constexpr decltype(CompareFn()) } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_BUFFER_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_BUFFER_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream.h b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream.h index fd577ad712f228fa8016a48942511a3263aae5da..8ad97fedc923ac50bcaad86e0ba2c2e46df6821b 100644 --- a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream.h +++ b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream.h @@ -12,12 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_STREAM_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_STREAM_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_STREAM_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_STREAM_H_ +#include #include #include -#include #include "tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_buffer.h" #include "tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h" @@ -322,4 +322,4 @@ WeightedQuantilesStream::GetQuantileSpecs( } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_STREAM_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_STREAM_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h index c329c6d4f7363a7738b06648943fe1dbd065cce5..aec232f3cbb096f0aa51e4362a821882391f8027 100644 --- a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h +++ b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_SUMMARY_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_SUMMARY_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_SUMMARY_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_SUMMARY_H_ #include #include @@ -334,4 +334,4 @@ constexpr decltype(CompareFn()) } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_SUMMARY_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_QUANTILES_WEIGHTED_QUANTILES_SUMMARY_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/testutil/batch_features_testutil.h b/tensorflow/contrib/boosted_trees/lib/testutil/batch_features_testutil.h index d95878ec87b9e903930d2016bb573eee2573f776..b98190b10dc88d5bba9023e771844a2bd6c9a45d 100644 --- a/tensorflow/contrib/boosted_trees/lib/testutil/batch_features_testutil.h +++ b/tensorflow/contrib/boosted_trees/lib/testutil/batch_features_testutil.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_BATCH_FEATURES_TESTUTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_BATCH_FEATURES_TESTUTIL_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_BATCH_FEATURES_TESTUTIL_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_BATCH_FEATURES_TESTUTIL_H_ #include "tensorflow/contrib/boosted_trees/lib/utils/batch_features.h" #include "tensorflow/core/framework/tensor.h" @@ -42,4 +42,4 @@ void RandomlyInitializeBatchFeatures( } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_BATCH_FEATURES_TESTUTIL_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_BATCH_FEATURES_TESTUTIL_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/testutil/random_tree_gen.cc b/tensorflow/contrib/boosted_trees/lib/testutil/random_tree_gen.cc index cbe26ba918d384ad903fb854ca3e88e84d16a923..705b65e9db9f1aed9af1be153240d57e163c2d5b 100644 --- a/tensorflow/contrib/boosted_trees/lib/testutil/random_tree_gen.cc +++ b/tensorflow/contrib/boosted_trees/lib/testutil/random_tree_gen.cc @@ -22,9 +22,9 @@ namespace tensorflow { namespace boosted_trees { namespace testutil { +using boosted_trees::trees::DenseFloatBinarySplit; using tensorflow::boosted_trees::trees::DecisionTreeConfig; using tensorflow::boosted_trees::trees::TreeNode; -using boosted_trees::trees::DenseFloatBinarySplit; namespace { diff --git a/tensorflow/contrib/boosted_trees/lib/testutil/random_tree_gen.h b/tensorflow/contrib/boosted_trees/lib/testutil/random_tree_gen.h index 5e12429ba778344edda623d149e017661f1e0222..1838b4cee21afb5df72a9b902f0ec0ce6f7ac627 100644 --- a/tensorflow/contrib/boosted_trees/lib/testutil/random_tree_gen.h +++ b/tensorflow/contrib/boosted_trees/lib/testutil/random_tree_gen.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_ #include @@ -72,4 +72,4 @@ class RandomTreeGen { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TESTUTIL_RANDOM_TREE_GEN_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.h b/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.h index 604ff02744b25b136bd935bf85635731730effe8..43526c229a65d45a2b0ced4aa1262d489526fc7b 100644 --- a/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.h +++ b/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TREES_DECISION_TREE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TREES_DECISION_TREE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TREES_DECISION_TREE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TREES_DECISION_TREE_H_ #include "tensorflow/contrib/boosted_trees/lib/utils/example.h" #include "tensorflow/contrib/boosted_trees/proto/tree_config.pb.h" // NOLINT @@ -46,4 +46,4 @@ class DecisionTree { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TREES_DECISION_TREE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_TREES_DECISION_TREE_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/batch_features.h b/tensorflow/contrib/boosted_trees/lib/utils/batch_features.h index badc629a118f768d5aa25ef1b94b8190e6910c7f..da5e7448519cb7f4092f7bbbe1b526271008ec22 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/batch_features.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/batch_features.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_BATCH_FEATURES_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_BATCH_FEATURES_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_BATCH_FEATURES_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_BATCH_FEATURES_H_ #include #include "tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h" @@ -92,4 +92,4 @@ class BatchFeatures { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_BATCH_FEATURES_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_BATCH_FEATURES_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/batch_features_test.cc b/tensorflow/contrib/boosted_trees/lib/utils/batch_features_test.cc index 9de3e32b097a151b3bd6f5c30df2db0938b65e9c..609519e8b1153a27d987c5f9ca9bfcc9ee6717d6 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/batch_features_test.cc +++ b/tensorflow/contrib/boosted_trees/lib/utils/batch_features_test.cc @@ -25,8 +25,8 @@ namespace boosted_trees { namespace utils { namespace { -using test::AsTensor; using errors::InvalidArgument; +using test::AsTensor; class BatchFeaturesTest : public ::testing::Test {}; diff --git a/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.cc b/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.cc index 38f0151255bbf4fcd87f1d0d76fd111649ee4a12..db34db998a7442c69f2ab468f4557d991429f4ee 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.cc +++ b/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.cc @@ -23,10 +23,10 @@ #include "tensorflow/core/lib/random/simple_philox.h" #include "tensorflow/core/platform/logging.h" +using tensorflow::Status; using tensorflow::boosted_trees::learner::LearningRateDropoutDrivenConfig; using tensorflow::random::PhiloxRandom; using tensorflow::random::SimplePhilox; -using tensorflow::Status; namespace tensorflow { namespace boosted_trees { diff --git a/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.h b/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.h index c3f1c918ca5f603cf9470071017d8ee384dc9320..928bfbfe5c9394ab4083aabced4c8e1149bb10aa 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_DROPOUT_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_DROPOUT_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_DROPOUT_UTILS_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_DROPOUT_UTILS_H_ #include #include @@ -74,4 +74,4 @@ class DropoutUtils { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_DROPOUT_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_DROPOUT_UTILS_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils_test.cc b/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils_test.cc index ce7632e58987f5890beaded5dd305724f950e1e8..02f972c8e00e8229426ac53d8f20765484787b6e 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils_test.cc +++ b/tensorflow/contrib/boosted_trees/lib/utils/dropout_utils_test.cc @@ -26,9 +26,9 @@ #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/env.h" +using std::unordered_set; using tensorflow::boosted_trees::learner::LearningRateDropoutDrivenConfig; using tensorflow::boosted_trees::trees::DecisionTreeEnsembleConfig; -using std::unordered_set; namespace tensorflow { namespace boosted_trees { diff --git a/tensorflow/contrib/boosted_trees/lib/utils/example.h b/tensorflow/contrib/boosted_trees/lib/utils/example.h index 54f60e1dee49a4a40b84fcc6e042fac1858aa187..1371ff337f78dd1c38f2bd0ba86911642f3aeb3e 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/example.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/example.h @@ -13,8 +13,8 @@ // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLE_H_ #include #include @@ -131,4 +131,4 @@ struct Example { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLE_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h b/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h index 5b33c8158879ec65425ac77b5338ee98fbdf07db..1b654e1c44e545fb97216ad950f3cd2d3240ffd0 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h @@ -13,8 +13,8 @@ // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLES_ITERABLE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLES_ITERABLE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLES_ITERABLE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLES_ITERABLE_H_ #include @@ -205,4 +205,4 @@ class ExamplesIterable { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLES_ITERABLE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_EXAMPLES_ITERABLE_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/macros.h b/tensorflow/contrib/boosted_trees/lib/utils/macros.h index 28ea0a4dc191af66ced574d78d9873cc8335f491..9a53fb2ef7d0581986885f3bc8233d91b67c0166 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/macros.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/macros.h @@ -13,8 +13,8 @@ // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_MACROS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_MACROS_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_MACROS_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_MACROS_H_ #include "tensorflow/core/platform/macros.h" @@ -23,4 +23,4 @@ return (STATUS); \ } -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_MACROS_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_MACROS_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/optional_value.h b/tensorflow/contrib/boosted_trees/lib/utils/optional_value.h index c141fe059d48072c6c4495535eafec9633616d21..b2166f53d7a037fb8ec53d5295b98bb82b17d4c7 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/optional_value.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/optional_value.h @@ -13,8 +13,8 @@ // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_OPTIONAL_VALUE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_OPTIONAL_VALUE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_OPTIONAL_VALUE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_OPTIONAL_VALUE_H_ #include "tensorflow/core/platform/logging.h" @@ -44,4 +44,4 @@ class OptionalValue { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_OPTIONAL_VALUE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_OPTIONAL_VALUE_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/parallel_for.h b/tensorflow/contrib/boosted_trees/lib/utils/parallel_for.h index c80431b5587cecc0bce22f6150a69d30397529da..ec06787e1db69514c9e60f6d152f3b0c7de23842 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/parallel_for.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/parallel_for.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LIB_UTILS_PARALLEL_FOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LIB_UTILS_PARALLEL_FOR_H_ +#ifndef TENSORFLOW_CONTRIB_LIB_UTILS_PARALLEL_FOR_H_ +#define TENSORFLOW_CONTRIB_LIB_UTILS_PARALLEL_FOR_H_ #include "tensorflow/core/lib/core/threadpool.h" @@ -30,4 +30,4 @@ void ParallelFor(int64 batch_size, int64 desired_parallelism, } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LIB_UTILS_PARALLEL_FOR_H_ +#endif // TENSORFLOW_CONTRIB_LIB_UTILS_PARALLEL_FOR_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/random.h b/tensorflow/contrib/boosted_trees/lib/utils/random.h index 6dd55fcacc42b88116737ab6fb413852ffc1473d..546d344f5585458f10699a644621f0adf26b6446 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/random.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/random.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LIB_UTILS_RANDOM_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LIB_UTILS_RANDOM_H_ +#ifndef TENSORFLOW_CONTRIB_LIB_UTILS_RANDOM_H_ +#define TENSORFLOW_CONTRIB_LIB_UTILS_RANDOM_H_ #include "tensorflow/core/lib/random/simple_philox.h" @@ -36,4 +36,4 @@ inline int32 PoissonBootstrap(random::SimplePhilox* rng) { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LIB_UTILS_RANDOM_H_ +#endif // TENSORFLOW_CONTRIB_LIB_UTILS_RANDOM_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/sparse_column_iterable.cc b/tensorflow/contrib/boosted_trees/lib/utils/sparse_column_iterable.cc index ccee9530b6897924453461c13b1238402c0f6cfa..1297aa884938f2f099a32568acc80c6cd8162651 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/sparse_column_iterable.cc +++ b/tensorflow/contrib/boosted_trees/lib/utils/sparse_column_iterable.cc @@ -51,7 +51,7 @@ class IndicesRowIterator return tmp; } - reference operator*() { return iter_->ix()(row_idx_, 0); } + reference operator*() const { return iter_->ix()(row_idx_, 0); } pointer operator->() { return &iter_->ix()(row_idx_, 0); } diff --git a/tensorflow/contrib/boosted_trees/lib/utils/sparse_column_iterable.h b/tensorflow/contrib/boosted_trees/lib/utils/sparse_column_iterable.h index 9664c9d1c6a0c0c8b1bbd1506944c54d2310c611..87fb1fbf5ae3cc6bcf25f68a180d1d9b21ef4d6f 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/sparse_column_iterable.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/sparse_column_iterable.h @@ -13,8 +13,8 @@ // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_SPARSE_COLUMN_ITERABLE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_SPARSE_COLUMN_ITERABLE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_SPARSE_COLUMN_ITERABLE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_SPARSE_COLUMN_ITERABLE_H_ #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" @@ -127,4 +127,4 @@ class SparseColumnIterable { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_SPARSE_COLUMN_ITERABLE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_SPARSE_COLUMN_ITERABLE_H_ diff --git a/tensorflow/contrib/boosted_trees/lib/utils/tensor_utils.h b/tensorflow/contrib/boosted_trees/lib/utils/tensor_utils.h index 58f5e5a0d18788375cd8166d1fcbdc7c294ba5e2..475d3718eccc2b23260b7cf5286abdd31ef1bad6 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/tensor_utils.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/tensor_utils.h @@ -13,8 +13,8 @@ // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_TENSOR_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_TENSOR_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_TENSOR_UTILS_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_TENSOR_UTILS_H_ #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" @@ -57,4 +57,4 @@ class TensorUtils { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_TENSOR_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_LIB_UTILS_TENSOR_UTILS_H_ diff --git a/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc b/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc index 1fa70bafddb0c94f47d006d5694bea941edaddf9..ae99d53a2cf805d70d60746cd44f73f7fd9dc6e2 100644 --- a/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc +++ b/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc @@ -19,8 +19,8 @@ namespace tensorflow { namespace boosted_trees { -using shape_inference::InferenceContext; using shape_inference::DimensionHandle; +using shape_inference::InferenceContext; using shape_inference::ShapeHandle; REGISTER_RESOURCE_HANDLE_OP(QuantileStreamResource); @@ -39,6 +39,7 @@ REGISTER_OP("CreateQuantileAccumulator") .Attr("max_elements: int = 1099511627776") // 1 << 40 .Attr("epsilon: float") .Attr("num_quantiles: int") + .Attr("generate_quantiles: bool=False") .Input("quantile_accumulator_handle: resource") .Input("stamp_token: int64") .SetShapeFn([](shape_inference::InferenceContext* c) { diff --git a/tensorflow/contrib/boosted_trees/ops/split_handler_ops.cc b/tensorflow/contrib/boosted_trees/ops/split_handler_ops.cc index 0d27ddaf3a1d540efee268c2bcca217077ff5871..5d0ebbf73ce1272b51a475f67984db3a181b7130 100644 --- a/tensorflow/contrib/boosted_trees/ops/split_handler_ops.cc +++ b/tensorflow/contrib/boosted_trees/ops/split_handler_ops.cc @@ -18,9 +18,9 @@ namespace tensorflow { +using shape_inference::DimensionHandle; using shape_inference::InferenceContext; using shape_inference::ShapeHandle; -using shape_inference::DimensionHandle; REGISTER_OP("BuildDenseInequalitySplits") .Attr("feature_column_group_id: int") diff --git a/tensorflow/contrib/boosted_trees/ops/stats_accumulator_ops.cc b/tensorflow/contrib/boosted_trees/ops/stats_accumulator_ops.cc index 0354f7853cbedf22d0a299273b4dbd225b3121ab..179505eef01f79bb149137400468b84285fe478a 100644 --- a/tensorflow/contrib/boosted_trees/ops/stats_accumulator_ops.cc +++ b/tensorflow/contrib/boosted_trees/ops/stats_accumulator_ops.cc @@ -19,9 +19,9 @@ namespace tensorflow { namespace boosted_trees { +using shape_inference::DimensionHandle; using shape_inference::InferenceContext; using shape_inference::ShapeHandle; -using shape_inference::DimensionHandle; REGISTER_RESOURCE_HANDLE_OP(StatsAccumulatorScalarResource); diff --git a/tensorflow/contrib/boosted_trees/proto/learner.proto b/tensorflow/contrib/boosted_trees/proto/learner.proto index 919e7cd81427c27cf892bc77998f52406d2bcf15..d84ba7438e7f03685d5bafca52ff8283f0fce898 100644 --- a/tensorflow/contrib/boosted_trees/proto/learner.proto +++ b/tensorflow/contrib/boosted_trees/proto/learner.proto @@ -22,6 +22,10 @@ message TreeConstraintsConfig { // Min hessian weight per node. float min_node_weight = 2; + + // Maximum number of unique features used in the tree. Zero means there is no + // limit. + int64 max_number_of_unique_feature_columns = 3; } // LearningRateConfig describes all supported learning rate tuners. diff --git a/tensorflow/contrib/boosted_trees/proto/tree_config.proto b/tensorflow/contrib/boosted_trees/proto/tree_config.proto index fc570c1083d01a65760a456c109dad93afd9f62a..4407c4d981785a279b6296f4726a221cacb4c5b1 100644 --- a/tensorflow/contrib/boosted_trees/proto/tree_config.proto +++ b/tensorflow/contrib/boosted_trees/proto/tree_config.proto @@ -128,6 +128,10 @@ message GrowingMetadata { // Number of layers that we have attempted to build. After pruning, these // layers might have been removed. int64 num_layers_attempted = 2; + + // Sorted list of column handlers that have been used in at least one split + // so far. + repeated int64 used_handler_ids = 3; } // DecisionTreeEnsembleConfig describes an ensemble of decision trees. diff --git a/tensorflow/contrib/boosted_trees/python/kernel_tests/quantile_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/quantile_ops_test.py index 888d5c57ed33446c8b6f18d2d1e393647613d132..81f58de28cbe98bb996c6665114eeb0030ee52f9 100644 --- a/tensorflow/contrib/boosted_trees/python/kernel_tests/quantile_ops_test.py +++ b/tensorflow/contrib/boosted_trees/python/kernel_tests/quantile_ops_test.py @@ -106,9 +106,11 @@ class QuantileBucketsOpTest(test_util.TensorFlowTestCase): | 6 | 16 | [16, 17, 18, 19, 20, 21] """ + num_quantiles = 3 with self.test_session() as sess: accumulator = quantile_ops.QuantileAccumulator( - init_stamp_token=0, num_quantiles=3, epsilon=0.001, name="q1") + init_stamp_token=0, num_quantiles=num_quantiles, + epsilon=0.001, name="q1") resources.initialize_resources(resources.shared_resources()).run() input_column = array_ops.placeholder(dtypes.float32) weights = array_ops.placeholder(dtypes.float32) @@ -131,8 +133,128 @@ class QuantileBucketsOpTest(test_util.TensorFlowTestCase): buckets, are_ready_flush = (sess.run( [buckets, are_ready_flush])) self.assertEqual(True, are_ready_flush) + self.assertEqual(num_quantiles + 1, len(buckets)) self.assertAllEqual([1, 86., 170., 253.], buckets) + def testStreamingQuantileBucketsLowPrecisionInput(self): + """Tests inputs that simulate low precision float16 values.""" + + num_quantiles = 3 + # set generate_quantiles to True since the test will generate fewer + # boundaries otherwise. + with self.test_session() as sess: + accumulator = quantile_ops.QuantileAccumulator( + init_stamp_token=0, num_quantiles=num_quantiles, + epsilon=0.001, name="q1", generate_quantiles=True) + resources.initialize_resources(resources.shared_resources()).run() + input_column = array_ops.placeholder(dtypes.float32) + weights = array_ops.placeholder(dtypes.float32) + update = accumulator.add_summary( + stamp_token=0, + column=input_column, + example_weights=weights) + + with self.test_session() as sess: + # This input is generated by integer in the range [2030, 2060] + # but represented by with float16 precision. Integers <= 2048 are + # exactly represented, whereas numbers > 2048 are rounded; and hence + # numbers > 2048 are repeated. For precision loss / rounding, see: + # https://en.wikipedia.org/wiki/Half-precision_floating-point_format. + # + # The intent of the test is not handling of float16 values, but to + # validate the number of buckets is returned, in cases where the input + # may contain repeated values. + inputs = [ + 2030.0, 2031.0, 2032.0, 2033.0, 2034.0, 2035.0, 2036.0, 2037.0, + 2038.0, 2039.0, 2040.0, 2041.0, 2042.0, 2043.0, 2044.0, 2045.0, + 2046.0, 2047.0, 2048.0, 2048.0, 2050.0, 2052.0, 2052.0, 2052.0, + 2054.0, 2056.0, 2056.0, 2056.0, 2058.0, 2060.0 + ] + sess.run(update, + {input_column: inputs, + weights: [1] * len(inputs)}) + + with self.test_session() as sess: + sess.run(accumulator.flush(stamp_token=0, next_stamp_token=1)) + are_ready_flush, buckets = (accumulator.get_buckets(stamp_token=1)) + buckets, are_ready_flush = (sess.run( + [buckets, are_ready_flush])) + self.assertEqual(True, are_ready_flush) + self.assertEqual(num_quantiles + 1, len(buckets)) + self.assertAllEqual([2030, 2040, 2050, 2060], buckets) + + def _testStreamingQuantileBucketsHelper( + self, inputs, num_quantiles=3, expected_buckets=None): + """Helper to test quantile buckets on different inputs.""" + + # set generate_quantiles to True since the test will generate fewer + # boundaries otherwise. + with self.test_session() as sess: + accumulator = quantile_ops.QuantileAccumulator( + init_stamp_token=0, num_quantiles=num_quantiles, + epsilon=0.001, name="q1", generate_quantiles=True) + resources.initialize_resources(resources.shared_resources()).run() + input_column = array_ops.placeholder(dtypes.float32) + weights = array_ops.placeholder(dtypes.float32) + update = accumulator.add_summary( + stamp_token=0, + column=input_column, + example_weights=weights) + + with self.test_session() as sess: + sess.run(update, + {input_column: inputs, + weights: [1] * len(inputs)}) + + with self.test_session() as sess: + sess.run(accumulator.flush(stamp_token=0, next_stamp_token=1)) + are_ready_flush, buckets = (accumulator.get_buckets(stamp_token=1)) + buckets, are_ready_flush = (sess.run( + [buckets, are_ready_flush])) + self.assertEqual(True, are_ready_flush) + # By default, use 3 quantiles, 4 boundaries for simplicity. + self.assertEqual(num_quantiles + 1, len(buckets)) + if expected_buckets: + self.assertAllEqual(buckets, expected_buckets) + + def testStreamingQuantileBucketsRepeatedSingleValue(self): + inputs = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1] + self._testStreamingQuantileBucketsHelper(inputs) + + def testStreamingQ2antileBucketsRepeatedTwoValues(self): + inputs = [1, 1, 1, 2, 2, 2, 2, 2, 1, 1] + self._testStreamingQuantileBucketsHelper(inputs) + + def testStreamingQ2antileBucketsRepeatedTwoValuesUnbalanced(self): + inputs = [7, 7, 7, 2, 7, 7, 2, 2, 7, 7] + self._testStreamingQuantileBucketsHelper(inputs) + + def testStreamingQuantileBucketsFewerInputstThanBuckets(self): + inputs = [5] + self._testStreamingQuantileBucketsHelper(inputs) + + def testStreamingQuantileBucketsEqualDistributionInSequence(self): + # Input pattern is of the form [1, 1, 1, 2, 2, 2, 3, 3, 3, ...] + ones = 100 * [1] + inputs = [] + for i in range(1, 101): + inputs += [i * k for k in ones] + # Expect 100 equally spaced buckets. + expected_buckets = range(1, 101) + self._testStreamingQuantileBucketsHelper( + inputs, num_quantiles=99, expected_buckets=expected_buckets) + + def testStreamingQuantileBucketsEqualDistributionInterleaved(self): + # Input pattern is of the form [1, 2, 3, 1, 2, 3, 1, 2, 3, ...] + sequence = range(1, 101) + inputs = [] + for _ in range(1, 101): + inputs += sequence + # Expect 100 equally spaced buckets. + expected_buckets = range(1, 101) + self._testStreamingQuantileBucketsHelper( + inputs, num_quantiles=99, expected_buckets=expected_buckets) + def testStreamingQuantileBuckets(self): """Sets up the quantile summary op test as follows. diff --git a/tensorflow/contrib/boosted_trees/python/kernel_tests/training_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/training_ops_test.py index c2e65b643df90e88aadb0bb9acaf692da35b1a16..8ca1aabacaf53b66aaba184962922294427d6803 100644 --- a/tensorflow/contrib/boosted_trees/python/kernel_tests/training_ops_test.py +++ b/tensorflow/contrib/boosted_trees/python/kernel_tests/training_ops_test.py @@ -63,7 +63,7 @@ def _gen_learner_config(num_classes, if dropout_prob_of_skipping is not None: config.learning_rate_tuner.dropout.dropout_prob_of_skipping = ( dropout_prob_of_skipping) - return config.SerializeToString() + return config def _gen_dense_split_info(fc, threshold, left_weight, right_weight): @@ -145,7 +145,7 @@ class CenterTreeEnsembleBiasOpTest(test_util.TensorFlowTestCase): pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE, # Dropout does not change anything here. - dropout_probability=0.5) + dropout_probability=0.5).SerializeToString() # Center bias for the initial step. grads = constant_op.constant([0.4, -0.3]) @@ -296,7 +296,7 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE, # Dropout does not change anything here, tree is not finalized. - dropout_probability=0.5) + dropout_probability=0.5).SerializeToString() # Prepare handler inputs. # Note that handlers 1 & 3 have the same gain but different splits. @@ -443,7 +443,7 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE, # Dropout does not change anything here - tree is not finalized. - dropout_probability=0.5) + dropout_probability=0.5).SerializeToString() # Prepare handler inputs. # Handler 1 only has a candidate for partition 1, handler 2 has candidates @@ -632,7 +632,8 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): max_depth=1, min_node_weight=0, pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, - growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE) + growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE).SerializeToString( + ) # Prepare handler inputs. handler1_partitions = np.array([0], dtype=np.int32) @@ -772,7 +773,8 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): max_depth=1, min_node_weight=0, pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, - growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE) + growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE).SerializeToString( + ) # Prepare handler inputs. # All handlers have negative gain. @@ -837,7 +839,8 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): max_depth=1, min_node_weight=0, pruning_mode=learner_pb2.LearnerConfig.POST_PRUNE, - growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE) + growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE).SerializeToString( + ) # Prepare handler inputs. # Note that handlers 1 & 3 have the same gain but different splits. @@ -943,7 +946,8 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): max_depth=2, min_node_weight=0, pruning_mode=learner_pb2.LearnerConfig.POST_PRUNE, - growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE) + growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE).SerializeToString( + ) # Prepare handler inputs. # All handlers have negative gain. @@ -1090,7 +1094,8 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): max_depth=2, min_node_weight=0, pruning_mode=learner_pb2.LearnerConfig.POST_PRUNE, - growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE) + growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE).SerializeToString( + ) # Prepare handler inputs. # Second handler has positive gain. @@ -1330,7 +1335,7 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, growing_mode=learner_pb2.LearnerConfig.LAYER_BY_LAYER, # Dropout will have no effect, since the tree will not be fully grown. - dropout_probability=1.0) + dropout_probability=1.0).SerializeToString() # Prepare handler inputs. # Handler 1 only has a candidate for partition 1, handler 2 has candidates @@ -1538,7 +1543,7 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): min_node_weight=0, pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE, - dropout_probability=1.0) + dropout_probability=1.0).SerializeToString() # Prepare handler inputs. handler1_partitions = np.array([0], dtype=np.int32) @@ -1583,6 +1588,301 @@ class GrowTreeEnsembleOpTest(test_util.TensorFlowTestCase): self.assertEqual( 2, tree_ensemble_config.tree_metadata[2].num_tree_weight_updates) + def testGrowExistingEnsembleTreeWithFeatureSelectionCanStillGrow(self): + """Test growing a tree with feature selection.""" + with self.test_session() as session: + # Create existing ensemble with one root split and one bias tree. + tree_ensemble_config = tree_config_pb2.DecisionTreeEnsembleConfig() + text_format.Merge(""" + trees { + nodes { + leaf { + vector { + value: -0.32 + value: 0.28 + } + } + } + } + trees { + nodes { + categorical_id_binary_split { + feature_column: 3 + feature_id: 7 + left_id: 1 + right_id: 2 + } + node_metadata { + gain: 1.3 + } + } + nodes { + leaf { + sparse_vector { + index: 0 + value: 2.3 + } + } + } + nodes { + leaf { + sparse_vector { + index: 0 + value: -0.9 + } + } + } + } + tree_weights: 0.7 + tree_weights: 1 + tree_metadata { + num_tree_weight_updates: 1 + num_layers_grown: 1 + is_finalized: true + } + tree_metadata { + num_tree_weight_updates: 5 + num_layers_grown: 1 + is_finalized: true + } + growing_metadata { + num_trees_attempted: 2 + num_layers_attempted: 2 + used_handler_ids: 2 + used_handler_ids: 5 + } + """, tree_ensemble_config) + tree_ensemble_handle = model_ops.tree_ensemble_variable( + stamp_token=0, + tree_ensemble_config=tree_ensemble_config.SerializeToString(), + name="tree_ensemble") + resources.initialize_resources(resources.shared_resources()).run() + + # Prepare learner config. + learner_config = _gen_learner_config( + num_classes=2, + l1_reg=0, + l2_reg=0, + tree_complexity=0, + max_depth=1, + min_node_weight=0, + pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, + growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE) + # There are 2 handler_ids in used_handler_ids already but one of them + # is handler 2, so we can still grow trees. + learner_config.constraints.max_number_of_unique_feature_columns = 2 + learner_config = learner_config.SerializeToString() + # Prepare handler inputs. + handler1_partitions = np.array([0], dtype=np.int32) + handler1_gains = np.array([7.62], dtype=np.float32) + handler1_split = [_gen_dense_split_info(5, 0.52, -4.375, 7.143)] + handler2_partitions = np.array([0], dtype=np.int32) + handler2_gains = np.array([0.63], dtype=np.float32) + handler2_split = [_gen_dense_split_info(2, 0.23, -0.6, 0.24)] + handler3_partitions = np.array([0], dtype=np.int32) + handler3_gains = np.array([7.62], dtype=np.float32) + handler3_split = [_gen_categorical_split_info(8, 7, -4.375, 7.143)] + + # Grow tree ensemble. + grow_op = training_ops.grow_tree_ensemble( + tree_ensemble_handle, + stamp_token=0, + next_stamp_token=1, + learning_rate=1, + partition_ids=[ + handler1_partitions, handler2_partitions, handler3_partitions + ], + gains=[handler1_gains, handler2_gains, handler3_gains], + splits=[handler1_split, handler2_split, handler3_split], + learner_config=learner_config, + dropout_seed=123, + center_bias=True) + session.run(grow_op) + + # Expect a new tree to be added with the split from handler 1. + _, serialized = session.run( + model_ops.tree_ensemble_serialize(tree_ensemble_handle)) + tree_ensemble_config.ParseFromString(serialized) + self.assertEqual(3, len(tree_ensemble_config.trees)) + self.assertEqual( + 2, len(tree_ensemble_config.growing_metadata.used_handler_ids)) + + def testGrowExistingEnsembleTreeWithFeatureSelectionEmptyEnsemble(self): + """Test growing a tree with feature selection with empty ensemble.""" + with self.test_session() as session: + # Create existing ensemble with one root split and one bias tree. + tree_ensemble_config = tree_config_pb2.DecisionTreeEnsembleConfig() + tree_ensemble_handle = model_ops.tree_ensemble_variable( + stamp_token=0, + tree_ensemble_config=tree_ensemble_config.SerializeToString(), + name="tree_ensemble") + resources.initialize_resources(resources.shared_resources()).run() + + # Prepare learner config. + learner_config = _gen_learner_config( + num_classes=2, + l1_reg=0, + l2_reg=0, + tree_complexity=0, + max_depth=1, + min_node_weight=0, + pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, + growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE) + learner_config.constraints.max_number_of_unique_feature_columns = 2 + learner_config = learner_config.SerializeToString() + # Prepare handler inputs. + handler1_partitions = np.array([0], dtype=np.int32) + handler1_gains = np.array([7.62], dtype=np.float32) + handler1_split = [_gen_dense_split_info(5, 0.52, -4.375, 7.143)] + handler2_partitions = np.array([0], dtype=np.int32) + handler2_gains = np.array([0.63], dtype=np.float32) + handler2_split = [_gen_dense_split_info(2, 0.23, -0.6, 0.24)] + handler3_partitions = np.array([0], dtype=np.int32) + handler3_gains = np.array([7.62], dtype=np.float32) + handler3_split = [_gen_categorical_split_info(8, 7, -4.375, 7.143)] + + # Grow tree ensemble. + grow_op = training_ops.grow_tree_ensemble( + tree_ensemble_handle, + stamp_token=0, + next_stamp_token=1, + learning_rate=1, + partition_ids=[ + handler1_partitions, handler2_partitions, handler3_partitions + ], + gains=[handler1_gains, handler2_gains, handler3_gains], + splits=[handler1_split, handler2_split, handler3_split], + learner_config=learner_config, + dropout_seed=123, + center_bias=True) + session.run(grow_op) + + _, serialized = session.run( + model_ops.tree_ensemble_serialize(tree_ensemble_handle)) + tree_ensemble_config.ParseFromString(serialized) + self.assertEqual(1, len(tree_ensemble_config.trees)) + self.assertEqual( + 1, len(tree_ensemble_config.growing_metadata.used_handler_ids)) + + def testGrowExistingEnsembleTreeWithFeatureSelectionCantGrow(self): + """Test growing a tree with feature selection with empty ensemble.""" + with self.test_session() as session: + # Create existing ensemble with one root split and one bias tree. + tree_ensemble_config = tree_config_pb2.DecisionTreeEnsembleConfig() + text_format.Merge(""" + trees { + nodes { + leaf { + vector { + value: -0.32 + value: 0.28 + } + } + } + } + trees { + nodes { + categorical_id_binary_split { + feature_column: 3 + feature_id: 7 + left_id: 1 + right_id: 2 + } + node_metadata { + gain: 1.3 + } + } + nodes { + leaf { + sparse_vector { + index: 0 + value: 2.3 + } + } + } + nodes { + leaf { + sparse_vector { + index: 0 + value: -0.9 + } + } + } + } + tree_weights: 0.7 + tree_weights: 1 + tree_metadata { + num_tree_weight_updates: 1 + num_layers_grown: 1 + is_finalized: true + } + tree_metadata { + num_tree_weight_updates: 5 + num_layers_grown: 1 + is_finalized: true + } + growing_metadata { + num_trees_attempted: 2 + num_layers_attempted: 2 + used_handler_ids: 4 + used_handler_ids: 5 + } + """, tree_ensemble_config) + tree_ensemble_handle = model_ops.tree_ensemble_variable( + stamp_token=0, + tree_ensemble_config=tree_ensemble_config.SerializeToString(), + name="tree_ensemble") + resources.initialize_resources(resources.shared_resources()).run() + + # Prepare learner config. + learner_config = _gen_learner_config( + num_classes=2, + l1_reg=0, + l2_reg=0, + tree_complexity=0, + max_depth=1, + min_node_weight=0, + pruning_mode=learner_pb2.LearnerConfig.PRE_PRUNE, + growing_mode=learner_pb2.LearnerConfig.WHOLE_TREE) + learner_config.constraints.max_number_of_unique_feature_columns = 2 + learner_config = learner_config.SerializeToString() + # Prepare handler inputs. + handler1_partitions = np.array([0], dtype=np.int32) + handler1_gains = np.array([7.62], dtype=np.float32) + handler1_split = [_gen_dense_split_info(5, 0.52, -4.375, 7.143)] + handler2_partitions = np.array([0], dtype=np.int32) + handler2_gains = np.array([0.63], dtype=np.float32) + handler2_split = [_gen_dense_split_info(2, 0.23, -0.6, 0.24)] + handler3_partitions = np.array([0], dtype=np.int32) + handler3_gains = np.array([7.62], dtype=np.float32) + handler3_split = [_gen_categorical_split_info(8, 7, -4.375, 7.143)] + + # Grow tree ensemble. + grow_op = training_ops.grow_tree_ensemble( + tree_ensemble_handle, + stamp_token=0, + next_stamp_token=1, + learning_rate=1, + partition_ids=[ + handler1_partitions, handler2_partitions, handler3_partitions + ], + gains=[handler1_gains, handler2_gains, handler3_gains], + splits=[handler1_split, handler2_split, handler3_split], + learner_config=learner_config, + dropout_seed=123, + center_bias=True) + session.run(grow_op) + + _, serialized = session.run( + model_ops.tree_ensemble_serialize(tree_ensemble_handle)) + tree_ensemble_config.ParseFromString(serialized) + # We can't grow a tree since we have reached the limit of 2 unique + # features [4, 5] and the only available splits are from + # handlers [0, 1, 2]. + self.assertEqual(2, len(tree_ensemble_config.trees)) + self.assertEqual( + 2, len(tree_ensemble_config.growing_metadata.used_handler_ids)) + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/contrib/boosted_trees/python/ops/batch_ops_utils.py b/tensorflow/contrib/boosted_trees/python/ops/batch_ops_utils.py index 23168bf4935e92bcb5072348361ae04861641b6d..7a5f329b7ab3216972180ccbb4c85f2537175422 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/batch_ops_utils.py +++ b/tensorflow/contrib/boosted_trees/python/ops/batch_ops_utils.py @@ -81,32 +81,32 @@ def _scheduled_stamp_resource_op_runner(batch, stamp): if not batch: return arg_keys = set(batch[0].args.keys()) - grouped_args = collections.defaultdict(list) + grouped_args = collections.OrderedDict() resource_handles = [] # Check that the set of arguments is the same across all the scheduled ops. for op in batch: if set(op.args.keys()) != arg_keys: raise ValueError("Mismatching arguments: %s, %s.", op.args, arg_keys) for key in arg_keys: - grouped_args[key].append(op.args[key]) + grouped_args.setdefault(key, []).append(op.args[key]) resource_handles.append(op.resource_handle) # Move all the inputs to the op device in one RPC. - grouped_args = { - k: _move_tensors(v, resource_handles[0].device) - for k, v in grouped_args.items() - } + grouped_args = collections.OrderedDict( + (k, _move_tensors(v, resource_handles[0].device)) + for k, v in sorted(grouped_args.items())) with ops.device(resource_handles[0].device): return batch[0].op(resource_handles, stamp, **grouped_args) def run_handler_scheduled_ops(per_handler_ops, stamp, worker_device): """Given a dictionary of ops for each handler, runs them in batch.""" - batched_ops = collections.defaultdict(list) + batched_ops = collections.OrderedDict() # Group the ops by their batching_key. Ops that share the same batching key # can be executed together. for handler in per_handler_ops.keys(): for op in per_handler_ops[handler]: - batched_ops[(op.batching_key(), op.batch_runner_fn())].append(op) + key = (op.batching_key(), op.batch_runner_fn()) + batched_ops.setdefault(key, []).append(op) op_results = {} for batch in batched_ops.values(): # Run each of the batched ops using its runner. diff --git a/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py b/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py index 294e04002adac62fc123a3242a05a1b36f422433..97d57e8b23608d4c3a8719426a75056fc6417d1d 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py @@ -47,7 +47,8 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): num_quantiles, max_elements=None, name=None, - container=None): + container=None, + generate_quantiles=False): """Creates a QuantileAccumulator object. Args: @@ -57,8 +58,11 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): max_elements: Maximum number of elements added to the accumulator. name: the name to save the accumulator under. container: An optional `string`. Defaults to `""` + generate_quantiles: Generate quantiles instead of approximate boundaries. + If true, exactly `num_quantiles` will be produced in the final summary. """ self._epsilon = epsilon + self._generate_quantiles = generate_quantiles name = _PATTERN.sub("", name) with ops.name_scope(name, "QuantileAccumulator") as name: @@ -70,7 +74,8 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): init_stamp_token, epsilon=epsilon, max_elements=max_elements, - num_quantiles=num_quantiles) + num_quantiles=num_quantiles, + generate_quantiles=generate_quantiles) is_initialized_op = gen_quantile_ops.quantile_accumulator_is_initialized( self._quantile_accumulator_handle) resources.register_resource(self._quantile_accumulator_handle, @@ -176,7 +181,14 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): summaries=summary) def flush(self, stamp_token, next_stamp_token): - """Finalizes quantile summary stream and resets it for next iteration.""" + """Finalizes quantile summary stream and resets it for next iteration. + + Args: + stamp_token: Exepcted current token. + next_stamp_token: Next value for the token. + Returns: + A list of quantiles or approximate boundaries. + """ return gen_quantile_ops.quantile_accumulator_flush( quantile_accumulator_handle=self._quantile_accumulator_handle, stamp_token=stamp_token, diff --git a/tensorflow/contrib/boosted_trees/python/training/__init__.py b/tensorflow/contrib/boosted_trees/python/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b569ac5fdb60e0907c322ad73aca65645e548d94 --- /dev/null +++ b/tensorflow/contrib/boosted_trees/python/training/__init__.py @@ -0,0 +1,18 @@ +# 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. +# ============================================================================== +"""training module under boosted_trees.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function diff --git a/tensorflow/contrib/boosted_trees/python/training/functions/__init__.py b/tensorflow/contrib/boosted_trees/python/training/functions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c1750117cd7c311515b4bca6882d55f496daac0e --- /dev/null +++ b/tensorflow/contrib/boosted_trees/python/training/functions/__init__.py @@ -0,0 +1,18 @@ +# 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. +# ============================================================================== +"""functions module under boosted_trees.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function diff --git a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py index b95956dae2a62b28643cd31815c5f5650eca337b..f0b66dcbbe1c5167b9993e66b30b1dc8a839c380 100644 --- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py +++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections import copy from tensorflow.contrib import learn @@ -163,7 +164,7 @@ def extract_features(features, feature_columns): scope = "gbdt" with variable_scope.variable_scope(scope): feature_columns = list(feature_columns) - transformed_features = {} + transformed_features = collections.OrderedDict() for fc in feature_columns: # pylint: disable=protected-access if isinstance(fc, feature_column_lib._EmbeddingColumn): @@ -681,13 +682,13 @@ class GradientBoostedDecisionTreeModel(object): control_flow_ops.no_op)) # Update handler stats. - handler_reads = {} + handler_reads = collections.OrderedDict() for handler in handlers: handler_reads[handler] = handler.scheduled_reads() handler_results = batch_ops_utils.run_handler_scheduled_ops( handler_reads, ensemble_stamp, worker_device) - per_handler_updates = {} + per_handler_updates = collections.OrderedDict() # Two values per handler. First one is if the handler is active for the # current layer. The second one is if the handler is going to be active # for the next layer. diff --git a/tensorflow/contrib/boosted_trees/python/utils/__init__.py b/tensorflow/contrib/boosted_trees/python/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6ceb150c26552584d631948f5eef2fedfa690894 --- /dev/null +++ b/tensorflow/contrib/boosted_trees/python/utils/__init__.py @@ -0,0 +1,18 @@ +# 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. +# ============================================================================== +"""utils module under boosted_trees.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function diff --git a/tensorflow/contrib/boosted_trees/python/utils/losses.py b/tensorflow/contrib/boosted_trees/python/utils/losses.py index 1e8b3ac08a74a94a0e5729e42ace91398a7b5c94..ab7ac2aba605db22a8ed370049b27d55cf1d413a 100644 --- a/tensorflow/contrib/boosted_trees/python/utils/losses.py +++ b/tensorflow/contrib/boosted_trees/python/utils/losses.py @@ -78,7 +78,7 @@ def per_example_maxent_loss(labels, weights, logits, num_classes, eps=1e-15): # Calculate softmax probabilities for each class. unnormalized_probs = math_ops.exp(logits) - normalizers = math_ops.reduce_sum(unnormalized_probs, 1, keep_dims=True) + normalizers = math_ops.reduce_sum(unnormalized_probs, 1, keepdims=True) softmax_predictions = math_ops.divide(unnormalized_probs, math_ops.add(normalizers, eps)) @@ -120,7 +120,7 @@ def per_example_squared_loss(labels, weights, predictions): update_op: An update operation to update the loss's internal state. """ unweighted_loss = math_ops.reduce_sum( - math_ops.square(predictions - labels), 1, keep_dims=True) + math_ops.square(predictions - labels), 1, keepdims=True) return unweighted_loss * weights, control_flow_ops.no_op() diff --git a/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h b/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h index 284ad5cdb9abf374650940ade7bb36663d72c0dd..3ebf28ea442edf87815c39971ae9e01a2a8aae9a 100644 --- a/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h +++ b/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_DECISION_TREE_ENSEMBLE_RESOURCE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_DECISION_TREE_ENSEMBLE_RESOURCE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_DECISION_TREE_ENSEMBLE_RESOURCE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_DECISION_TREE_ENSEMBLE_RESOURCE_H_ #include "tensorflow/contrib/boosted_trees/lib/trees/decision_tree.h" #include "tensorflow/contrib/boosted_trees/resources/stamped_resource.h" @@ -111,6 +111,35 @@ class DecisionTreeEnsembleResource : public StampedResource { return decision_tree_ensemble_->tree_weights(index); } + void MaybeAddUsedHandler(const int32 handler_id) { + protobuf::RepeatedField* used_ids = + decision_tree_ensemble_->mutable_growing_metadata() + ->mutable_used_handler_ids(); + protobuf::RepeatedField::iterator first = + std::lower_bound(used_ids->begin(), used_ids->end(), handler_id); + if (first == used_ids->end()) { + used_ids->Add(handler_id); + return; + } + if (handler_id == *first) { + // It is a duplicate entry. + return; + } + used_ids->Add(handler_id); + std::rotate(first, used_ids->end() - 1, used_ids->end()); + } + + std::vector GetUsedHandlers() const { + std::vector result; + result.reserve( + decision_tree_ensemble_->growing_metadata().used_handler_ids().size()); + for (int64 h : + decision_tree_ensemble_->growing_metadata().used_handler_ids()) { + result.push_back(h); + } + return result; + } + // Sets the weight of i'th tree, and increment num_updates in tree_metadata. void SetTreeWeight(const int32 index, const float weight, const int32 increment_num_updates) { @@ -150,4 +179,4 @@ class DecisionTreeEnsembleResource : public StampedResource { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_DECISION_TREE_ENSEMBLE_RESOURCE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_DECISION_TREE_ENSEMBLE_RESOURCE_H_ diff --git a/tensorflow/contrib/boosted_trees/resources/quantile_stream_resource.h b/tensorflow/contrib/boosted_trees/resources/quantile_stream_resource.h index fb29f79e578e8e52b67de631c527be35b7772b41..fdaaae7f472c8f564ab45a8366d3746cbf1158ee 100644 --- a/tensorflow/contrib/boosted_trees/resources/quantile_stream_resource.h +++ b/tensorflow/contrib/boosted_trees/resources/quantile_stream_resource.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_QUANTILE_STREAM_RESOURCE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_QUANTILE_STREAM_RESOURCE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_QUANTILE_STREAM_RESOURCE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_QUANTILE_STREAM_RESOURCE_H_ #include "tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream.h" #include "tensorflow/contrib/boosted_trees/proto/quantiles.pb.h" // NOLINT @@ -32,12 +32,14 @@ using QuantileStream = class QuantileStreamResource : public StampedResource { public: QuantileStreamResource(const float epsilon, const int32 num_quantiles, - const int64 max_elements, int64 stamp_token) + const int64 max_elements, bool generate_quantiles, + int64 stamp_token) : stream_(epsilon, max_elements), are_buckets_ready_(false), epsilon_(epsilon), num_quantiles_(num_quantiles), - max_elements_(max_elements) { + max_elements_(max_elements), + generate_quantiles_(generate_quantiles) { set_stamp(stamp_token); } @@ -74,6 +76,11 @@ class QuantileStreamResource : public StampedResource { are_buckets_ready_ = are_buckets_ready; } + bool generate_quantiles() const { return generate_quantiles_; } + void set_generate_quantiles(bool generate_quantiles) { + generate_quantiles_ = generate_quantiles; + } + private: ~QuantileStreamResource() override {} @@ -95,10 +102,15 @@ class QuantileStreamResource : public StampedResource { const int32 num_quantiles_; // An upper-bound for the number of elements. int64 max_elements_; + + // Generate quantiles instead of approximate boundaries. + // If true, exactly `num_quantiles` will be produced in the final summary. + bool generate_quantiles_; + TF_DISALLOW_COPY_AND_ASSIGN(QuantileStreamResource); }; } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_QUANTILE_STREAM_RESOURCE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_QUANTILE_STREAM_RESOURCE_H_ diff --git a/tensorflow/contrib/boosted_trees/resources/stamped_resource.h b/tensorflow/contrib/boosted_trees/resources/stamped_resource.h index aabeeb98516eda6f7e8e7e296d6860fe5d8d5ec3..957bbe8d61d3dd32adba1a7f0cf840c69bce6273 100644 --- a/tensorflow/contrib/boosted_trees/resources/stamped_resource.h +++ b/tensorflow/contrib/boosted_trees/resources/stamped_resource.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_STAMPED_RESOURCE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_STAMPED_RESOURCE_H_ +#ifndef TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_STAMPED_RESOURCE_H_ +#define TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_STAMPED_RESOURCE_H_ #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/platform/mutex.h" @@ -39,4 +39,4 @@ class StampedResource : public ResourceBase { } // namespace boosted_trees } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_STAMPED_RESOURCE_H_ +#endif // TENSORFLOW_CONTRIB_BOOSTED_TREES_RESOURCES_STAMPED_RESOURCE_H_ diff --git a/tensorflow/contrib/cloud/kernels/bigquery_table_accessor.cc b/tensorflow/contrib/cloud/kernels/bigquery_table_accessor.cc index deb324634b6edc17c9725996115d80c5bd11cbde..1bfd27305d569668a0bd67d876e59eec082296b3 100644 --- a/tensorflow/contrib/cloud/kernels/bigquery_table_accessor.cc +++ b/tensorflow/contrib/cloud/kernels/bigquery_table_accessor.cc @@ -18,7 +18,6 @@ limitations under the License. #include "tensorflow/core/lib/strings/numbers.h" namespace tensorflow { - namespace { constexpr size_t kBufferSize = 1024 * 1024; // In bytes. @@ -40,33 +39,6 @@ Status ParseJson(StringPiece json, Json::Value* result) { return Status::OK(); } -string ColumnTypeToString(BigQueryTableAccessor::ColumnType enum_type) { - switch (enum_type) { - case BigQueryTableAccessor::ColumnType::kRecord: - return "RECORD"; - case BigQueryTableAccessor::ColumnType::kString: - return "STRING"; - case BigQueryTableAccessor::ColumnType::kBytes: - return "BYTES"; - case BigQueryTableAccessor::ColumnType::kInteger: - return "INTEGER"; - case BigQueryTableAccessor::ColumnType::kFloat: - return "FLOAT"; - case BigQueryTableAccessor::ColumnType::kBoolean: - return "BOOLEAN"; - case BigQueryTableAccessor::ColumnType::kTimestamp: - return "TIMESTAMP"; - case BigQueryTableAccessor::ColumnType::kDate: - return "DATE"; - case BigQueryTableAccessor::ColumnType::kTime: - return "TIME"; - case BigQueryTableAccessor::ColumnType::kDatetime: - return "DATETIME"; - case BigQueryTableAccessor::ColumnType::kNone: - return "NONE"; - } -} - Status ParseColumnType(const string& type, BigQueryTableAccessor::ColumnType* enum_type) { if (type == "RECORD") { diff --git a/tensorflow/contrib/cloud/kernels/bigquery_table_accessor.h b/tensorflow/contrib/cloud/kernels/bigquery_table_accessor.h index 7d0eee59ae2f47503c4f8994ef356ce0dc336733..b349063715c903c982cfe2fb116b6525e35ff63b 100644 --- a/tensorflow/contrib/cloud/kernels/bigquery_table_accessor.h +++ b/tensorflow/contrib/cloud/kernels/bigquery_table_accessor.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_PARTITION_ACCESSOR_H_ -#define THIRD_PARTY_TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_PARTITION_ACCESSOR_H_ +#ifndef TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_PARTITION_ACCESSOR_H_ +#define TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_PARTITION_ACCESSOR_H_ #include #include @@ -205,4 +205,4 @@ class BigQueryTableAccessor { }; } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_PARTITION_ACCESSOR_H_ +#endif // TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_PARTITION_ACCESSOR_H_ diff --git a/tensorflow/contrib/cloud/kernels/bigquery_table_accessor_test_data.h b/tensorflow/contrib/cloud/kernels/bigquery_table_accessor_test_data.h index b2b11f4f57800d55ebc86273fcda71e673ff143a..fea6b15640ded74432f35112bc5d5d68e641c9dc 100644 --- a/tensorflow/contrib/cloud/kernels/bigquery_table_accessor_test_data.h +++ b/tensorflow/contrib/cloud/kernels/bigquery_table_accessor_test_data.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_TABLE_ACCESSOR_TEST_DATA_H_ -#define THIRD_PARTY_TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_TABLE_ACCESSOR_TEST_DATA_H_ +#ifndef TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_TABLE_ACCESSOR_TEST_DATA_H_ +#define TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_TABLE_ACCESSOR_TEST_DATA_H_ #include @@ -399,6 +399,6 @@ const string kTestEmptyRow = R"({ }]}]})"; } // namespace -} // namepsace tensorflow +} // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_TABLE_ACCESSOR_TEST_DATA_H_ +#endif // TENSORFLOW_CORE_KERNELS_CLOUD_BIGQUERY_TABLE_ACCESSOR_TEST_DATA_H_ diff --git a/tensorflow/contrib/cluster_resolver/BUILD b/tensorflow/contrib/cluster_resolver/BUILD index 15abd2be0385eb776ff4f76484133efb6e34f076..80e18a43a71cc9d6c9e2ccf5836e50c6427a30f6 100644 --- a/tensorflow/contrib/cluster_resolver/BUILD +++ b/tensorflow/contrib/cluster_resolver/BUILD @@ -34,6 +34,7 @@ py_library( ":cluster_resolver_py", ":gce_cluster_resolver_py", ":tpu_cluster_resolver_py", + "//tensorflow/python:util", ], ) diff --git a/tensorflow/contrib/cluster_resolver/__init__.py b/tensorflow/contrib/cluster_resolver/__init__.py index d17501e87e79158b1602ac6ddecc091bd86f2c2d..b4d8cd4a7cf42e910e7506dbeec8656a2cef62eb 100644 --- a/tensorflow/contrib/cluster_resolver/__init__.py +++ b/tensorflow/contrib/cluster_resolver/__init__.py @@ -26,3 +26,15 @@ from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import from tensorflow.contrib.cluster_resolver.python.training.gce_cluster_resolver import GceClusterResolver from tensorflow.contrib.cluster_resolver.python.training.tpu_cluster_resolver import TPUClusterResolver # pylint: enable=wildcard-import,unused-import + +from tensorflow.python.util.all_util import remove_undocumented + +_allowed_symbols = [ + 'ClusterResolver', + 'SimpleClusterResolver', + 'UnionClusterResolver', + 'GceClusterResolver', + 'TPUClusterResolver', +] + +remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py index c74da9cabd6816bc9c7891e32937534cff2d677d..a6a6e642e4e4c721b94821a70d55d6fe931347d6 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py @@ -18,6 +18,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function + +from six.moves.urllib.request import Request +from six.moves.urllib.request import urlopen + from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import ClusterResolver from tensorflow.python.training.server_lib import ClusterSpec @@ -38,10 +42,16 @@ class TPUClusterResolver(ClusterResolver): Cloud Platform project. """ + def _requestComputeMetadata(self, path): + req = Request('http://metadata/computeMetadata/v1/%s' % path, + headers={'Metadata-Flavor': 'Google'}) + resp = urlopen(req) + return resp.read() + def __init__(self, - project, - zone, tpu_names, + zone=None, + project=None, job_name='tpu_worker', credentials='default', service=None): @@ -51,9 +61,13 @@ class TPUClusterResolver(ClusterResolver): for the IP addresses and ports of each Cloud TPU listed. Args: - project: Name of the GCP project containing Cloud TPUs - zone: Zone where the TPUs are located tpu_names: A list of names of the target Cloud TPUs. + zone: Zone where the TPUs are located. If omitted or empty, we will assume + that the zone of the TPU is the same as the zone of the GCE VM, which we + will try to discover from the GCE metadata service. + project: Name of the GCP project containing Cloud TPUs. If omitted or + empty, we will try to discover the project name of the GCE VM from the + GCE metadata service. job_name: Name of the TensorFlow job the TPUs belong to. credentials: GCE Credentials. If None, then we use default credentials from the oauth2client @@ -65,6 +79,13 @@ class TPUClusterResolver(ClusterResolver): ImportError: If the googleapiclient is not installed. """ + if not project: + project = self._requestComputeMetadata('/project/project-id') + + if not zone: + zone_path = self._requestComputeMetadata('/instance/zone') + zone = zone_path.split('/')[-1] + self._project = project self._zone = zone self._tpu_names = tpu_names @@ -122,7 +143,8 @@ class TPUClusterResolver(ClusterResolver): request = self._service.projects().locations().nodes().get(name=full_name) response = request.execute() - instance_url = '%s:%s' % (response['ipAddress'], response['port']) - worker_list.append(instance_url) + if 'health' in response and response['health'] == 'HEALTHY': + instance_url = '%s:%s' % (response['ipAddress'], response['port']) + worker_list.append(instance_url) return ClusterSpec({self._job_name: worker_list}) diff --git a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver_test.py b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver_test.py index db7419be06b58e1c5737f69f2c7fd9fee44b9d95..4fd34629cf74f90869c77b8cb098d3c585a49404 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver_test.py +++ b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver_test.py @@ -48,6 +48,15 @@ class MockNodeClass(object): return MockRequestClass(name, self._tpu_map) +def mock_request_compute_metadata(cls, *args, **kwargs): + del cls, kwargs # Unused. + if args[0] == '/project/project-id': + return 'test-project' + elif args[0] == '/instance/zone': + return 'projects/test-project/locations/us-central1-c' + return '' + + class TPUClusterResolverTest(test.TestCase): def _verifyClusterSpecEquality(self, cluster_spec, expected_proto): @@ -89,11 +98,37 @@ class TPUClusterResolverTest(test.TestCase): return mock_client + @mock.patch.object(TPUClusterResolver, + '_requestComputeMetadata', + mock_request_compute_metadata) + def testRetrieveProjectAndZoneFromMetadata(self): + tpu_map = { + 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { + 'ipAddress': '10.1.2.3', + 'port': '8470', + 'health': 'HEALTHY' + } + } + + tpu_cluster_resolver = TPUClusterResolver( + project=None, + zone=None, + tpu_names=['test-tpu-1'], + credentials=None, + service=self.mock_service_client(tpu_map=tpu_map)) + + actual_cluster_spec = tpu_cluster_resolver.cluster_spec() + expected_proto = """ + job { name: 'tpu_worker' tasks { key: 0 value: '10.1.2.3:8470' } } + """ + self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) + def testSimpleSuccessfulRetrieval(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { 'ipAddress': '10.1.2.3', - 'port': '8470' + 'port': '8470', + 'health': 'HEALTHY' } } @@ -114,11 +149,13 @@ class TPUClusterResolverTest(test.TestCase): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { 'ipAddress': '10.1.2.3', - 'port': '8470' + 'port': '8470', + 'health': 'HEALTHY' }, 'projects/test-project/locations/us-central1-c/nodes/test-tpu-2': { 'ipAddress': '10.4.5.6', - 'port': '8470' + 'port': '8470', + 'health': 'HEALTHY' } } @@ -136,15 +173,54 @@ class TPUClusterResolverTest(test.TestCase): """ self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) + def testHealthyTpuNodeRetrieval(self): + tpu_map = { + 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { + 'ipAddress': '10.1.2.3', + 'port': '8470', + 'health': 'HEALTHY' + }, + 'projects/test-project/locations/us-central1-c/nodes/test-tpu-2': { + 'ipAddress': '10.4.5.6', + 'port': '8470', + }, + 'projects/test-project/locations/us-central1-c/nodes/test-tpu-3': { + 'ipAddress': '10.7.8.9', + 'port': '8470', + 'health': 'UNHEALTHY' + } + } + + tpu_cluster_resolver = TPUClusterResolver( + project='test-project', + zone='us-central1-c', + tpu_names=['test-tpu-2', 'test-tpu-1', 'test-tpu-3'], + credentials=None, + service=self.mock_service_client(tpu_map=tpu_map)) + + actual_cluster_spec = tpu_cluster_resolver.cluster_spec() + expected_proto = """ + job { + name: 'tpu_worker' + tasks { + key: 0 + value: '10.1.2.3:8470' + } + } + """ + self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) + def testGetMasterMultipleEntries(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { 'ipAddress': '10.1.2.3', - 'port': '8470' + 'port': '8470', + 'health': 'HEALTHY' }, 'projects/test-project/locations/us-central1-c/nodes/test-tpu-2': { 'ipAddress': '10.4.5.6', - 'port': '8470' + 'port': '8470', + 'health': 'HEALTHY' } } diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt index 8d023cc81dd73751f0b5690f3649ded3fc038155..23b31ae1dcc83d8a7152354ac147de9ada320429 100644 --- a/tensorflow/contrib/cmake/CMakeLists.txt +++ b/tensorflow/contrib/cmake/CMakeLists.txt @@ -52,6 +52,7 @@ if (NOT WIN32) # for targets that link ${CMAKE_THREAD_LIBS_INIT}. find_package (Threads) + # Options for linking CUDA/CUDNN libraries option(tensorflow_PATH_STATIC_LIB "Additional library search path for libcudnn_static.a, libnccl_static.a, libculibos.a" /usr/local/cuda/lib64/) option(tensorflow_CUDNN_INCLUDE "cudnn.h header install path" /usr/include/) if (NOT tensorflow_CUDNN_INCLUDE) @@ -73,6 +74,14 @@ if (NOT WIN32) # option's default value is OFF. Fill it with real default values set(tensorflow_CUDA_LIBRARY_PATH /usr/local/cuda/lib64) endif (NOT tensorflow_CUDA_LIBRARY_PATH) + + # Options for linking other libraries + option(systemlib_ZLIB "Use the system installed library as shared objects instead of downloading ZLIB and statically linking to it: ZLIB" OFF) + + option(systemlib_ALL "Turn on every possible systemlib_* options" OFF) + if (systemlib_ALL) + set (systmelib_ZLIB ON) + endif (systemlib_ALL) endif() if (WIN32) @@ -106,6 +115,9 @@ else() set(CMAKE_POSITION_INDEPENDENT_CODE OFF) endif() +# TODO(jart): We should make this only apply to snapfn.cc +add_definitions(-DSQLITE_OMIT_LOAD_EXTENSION) + if (tensorflow_DISABLE_EIGEN_FORCEINLINE) add_definitions(-DEIGEN_STRONG_INLINE=inline) endif() @@ -131,6 +143,9 @@ if(WIN32) set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /D_ITERATOR_DEBUG_LEVEL=0") set(CMAKE_CXX_FLAGS_MINSIZEREL "${CMAKE_CXX_FLAGS_MINSIZEREL} /D_ITERATOR_DEBUG_LEVEL=0") set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} /D_ITERATOR_DEBUG_LEVEL=0") + + # Try to avoid flaky failures due to failed generation of generate.stamp files. + set(CMAKE_SUPPRESS_REGENERATION ON) endif() if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU") @@ -182,8 +197,10 @@ if (tensorflow_BUILD_CC_TESTS) include(googletest) endif() +add_definitions(${ADD_CFLAGS}) +link_directories(${ADD_LINK_DIRECTORY}) + set(tensorflow_EXTERNAL_LIBRARIES - ${zlib_STATIC_LIBRARIES} ${gif_STATIC_LIBRARIES} ${png_STATIC_LIBRARIES} ${jpeg_STATIC_LIBRARIES} @@ -197,6 +214,15 @@ set(tensorflow_EXTERNAL_LIBRARIES ${re2_STATIC_LIBRARIES} ${sqlite_STATIC_LIBRARIES} ) + +if (systemlib_ZLIB) + set(tensorflow_EXTERNAL_LIBRARIES ${tensorflow_EXTERNAL_LIBRARIES} + ${ZLIB_LIBRARIES}) +else (systemlib_ZLIB) + set(tensorflow_EXTERNAL_LIBRARIES ${tensorflow_EXTERNAL_LIBRARIES} + ${zlib_STATIC_LIBRARIES}) +endif (systemlib_ZLIB) + set(tensorflow_EXTERNAL_DEPENDENCIES zlib_copy_headers_to_destination gif_copy_headers_to_destination @@ -280,7 +306,21 @@ if (tensorflow_ENABLE_GPU) list(APPEND CMAKE_LIBRARY_PATH "${tensorflow_CUDA_LIBRARY_PATH}/stubs") endif (NOT WIN32) - find_package(CUDA ${tensorflow_CUDA_VERSION} REQUIRED) + # later command will make use of the value in tensorflow_CUDA_VERSION + find_package(CUDA ${tensorflow_CUDA_VERSION} REQUIRED EXACT) + + # Test compatibility of compiler on CUDA + try_compile(CUDA_TEST_COMPILE_C + ${CMAKE_CURRENT_BINARY_DIR}/tests/cuda + ${CMAKE_CURRENT_SOURCE_DIR}/tests/cuda/compatibility_test.c + CMAKE_FLAGS -DINCLUDE_DIRECTORIES=${CUDA_INCLUDE_DIRS}) + try_compile(CUDA_TEST_COMPILE_CXX + ${CMAKE_CURRENT_BINARY_DIR}/tests/cuda + ${CMAKE_CURRENT_SOURCE_DIR}/tests/cuda/compatibility_test.cc + CMAKE_FLAGS -DINCLUDE_DIRECTORIES=${CUDA_INCLUDE_DIRS}) + if(NOT (CUDA_TEST_COMPILE_C AND CUDA_TEST_COMPILE_CXX)) + message(FATAL_ERROR "Selected compiler (or version) is not supported for CUDA") + endif() # by default we assume compute cabability 3.5 and 5.2. If you change this change it in # CUDA_NVCC_FLAGS and cuda_config.h below @@ -301,7 +341,8 @@ if (tensorflow_ENABLE_GPU) if(NOT CUDNN_HOME) set(CUDNN_HOME ${CUDA_TOOLKIT_TARGET_DIR}) endif(NOT CUDNN_HOME) - include_directories(${CUDNN_HOME}) + set(CUDNN_INCLUDE "${CUDNN_HOME}/include") + set(CUDA_LIBRARIES ${CUDA_LIBRARIES} ${CUDA_CUDA_LIBRARY} ${CUDA_CUBLAS_LIBRARIES} ${CUDA_CUFFT_LIBRARIES} ${CUDA_curand_LIBRARY} ${CUDA_cupti_LIBRARY} ${CUDA_cusolver_LIBRARY} ${CUDNN_HOME}/lib/x64/cudnn.lib) else (WIN32) @@ -329,10 +370,10 @@ if (tensorflow_ENABLE_GPU) message("culibos-static: ${culibos_STATIC_LIBRARY}") endif (NOT culibos_STATIC_LIBRARY) - include_directories(${CUDNN_INCLUDE}) set(CUDA_LIBRARIES ${CUDA_LIBRARIES} ${CUDA_CUDA_LIBRARY} ${CUDA_CUBLAS_LIBRARIES} ${CUDA_CUFFT_LIBRARIES} ${CUDA_curand_LIBRARY} ${CUDA_cupti_LIBRARY} ${CUDA_cusolver_LIBRARY} ${cudnn_STATIC_LIBRARY} ${culibos_STATIC_LIBRARY} ${nccl_STATIC_LIBRARY}) endif (WIN32) + include_directories(${CUDNN_INCLUDE}) # Remove "." from CUDA version variable. string(REPLACE "." "" short_CUDA_VER ${tensorflow_CUDA_VERSION}) @@ -348,29 +389,22 @@ if (tensorflow_ENABLE_GPU) "#endif // CUDA_CUDA_CONFIG_H_\n" ) - if (WIN32) - # tf assumes in various places header files to be in cuda/include. On windows the cuda sdk - # installs them under cuda/version/include and to avoid that we need to change tf we copy a - # few files to cuda/include - FILE(COPY - ${CUDA_TOOLKIT_TARGET_DIR}/include/cuda.h ${CUDA_TOOLKIT_TARGET_DIR}/include/cuComplex.h - ${CUDA_TOOLKIT_TARGET_DIR}/include/cublas_v2.h ${CUDNN_HOME}/include/cudnn.h - ${CUDA_TOOLKIT_TARGET_DIR}/include/cufft.h ${CUDA_TOOLKIT_TARGET_DIR}/include/curand.h - ${CUDA_TOOLKIT_TARGET_DIR}/include/cuda_runtime_api.h - ${CUDA_TOOLKIT_TARGET_DIR}/include/cusolverDn.h - DESTINATION ${tensorflow_source_dir}/third_party/gpus/cuda/include - ) - else(WIN32) - # Linux has slightly differnt install paths than Windows - FILE(COPY - ${CUDA_TOOLKIT_TARGET_DIR}/include/cuda.h ${CUDA_TOOLKIT_TARGET_DIR}/include/cuComplex.h - ${CUDA_TOOLKIT_TARGET_DIR}/include/cublas_v2.h ${CUDNN_INCLUDE}/cudnn.h - ${CUDA_TOOLKIT_TARGET_DIR}/include/cufft.h ${CUDA_TOOLKIT_TARGET_DIR}/include/curand.h - ${CUDA_TOOLKIT_TARGET_DIR}/include/cuda_runtime_api.h - ${CUDA_TOOLKIT_TARGET_DIR}/include/cusolverDn.h - DESTINATION ${tensorflow_source_dir}/third_party/gpus/cuda/include - ) - endif(WIN32) + # tf assumes in various places header files to be in cuda/include. On windows the cuda sdk + # installs them under cuda/version/include and to avoid that we need to change tf we copy a + # few files to cuda/include + FILE(COPY + ${CUDA_TOOLKIT_TARGET_DIR}/include/cuda.h + ${CUDA_TOOLKIT_TARGET_DIR}/include/cuComplex.h + ${CUDA_TOOLKIT_TARGET_DIR}/include/cublas_v2.h + ${CUDA_TOOLKIT_TARGET_DIR}/include/cusolverDn.h + ${CUDA_TOOLKIT_TARGET_DIR}/include/cuda_fp16.h + ${CUDA_TOOLKIT_TARGET_DIR}/include/device_functions.h + ${CUDA_TOOLKIT_TARGET_DIR}/include/cufft.h + ${CUDA_TOOLKIT_TARGET_DIR}/include/curand.h + ${CUDA_TOOLKIT_TARGET_DIR}/include/cuda_runtime_api.h + ${CUDNN_INCLUDE}/cudnn.h + DESTINATION ${tensorflow_source_dir}/third_party/gpus/cuda/include + ) include_directories(${tensorflow_source_dir}/third_party/gpus) # add cuda libraries to tensorflow_EXTERNAL_LIBRARIES @@ -404,7 +438,9 @@ endif() # Let's get to work! include(tf_core_framework.cmake) -include(tf_stream_executor.cmake) +if (tensorflow_ENABLE_GPU) + include(tf_stream_executor.cmake) +endif() include(tf_core_cpu.cmake) include(tf_core_ops.cmake) diff --git a/tensorflow/contrib/cmake/README.md b/tensorflow/contrib/cmake/README.md index 4be733a2809f366a214fa2bb853bccffb10ecaba..8f85a75ee466dbac524a1266dc2522109ca77cd5 100644 --- a/tensorflow/contrib/cmake/README.md +++ b/tensorflow/contrib/cmake/README.md @@ -30,7 +30,7 @@ bindings. * CMake version 3.5 or later. -* [Git](http://git-scm.com) +* [Git](https://git-scm.com) * [SWIG](http://www.swig.org/download.html) @@ -48,7 +48,7 @@ bindings. * Microsoft Windows 10 - Microsoft Visual Studio Enterprise 2015 with Visual C++ 2015 - - [Anaconda 4.1.1 (Python 3.5 64-bit)](https://www.continuum.io/downloads) + - [Anaconda 4.1.1 (Python 3.5 64-bit)](https://www.anaconda.com/download/) - [Git for Windows version 2.9.2.windows.1](https://git-scm.com/download/win) - [swigwin-3.0.10](http://www.swig.org/download.html) - [NVidia CUDA Toolkit 8.0](https://developer.nvidia.com/cuda-downloads) diff --git a/tensorflow/contrib/cmake/external/boringssl.cmake b/tensorflow/contrib/cmake/external/boringssl.cmake index cca8444e2ae9952ea7c69a9392580ead715d363b..3c4bb01e24fd121c9d0fc3594cc25de37af0e8a1 100644 --- a/tensorflow/contrib/cmake/external/boringssl.cmake +++ b/tensorflow/contrib/cmake/external/boringssl.cmake @@ -37,13 +37,10 @@ ExternalProject_Add(boringssl GIT_TAG ${boringssl_TAG} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" # BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${boringssl_STATIC_LIBRARIES} INSTALL_COMMAND "" CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF ) diff --git a/tensorflow/contrib/cmake/external/farmhash.cmake b/tensorflow/contrib/cmake/external/farmhash.cmake index 0cd0c1030c73d5218411f281d2b077af217e8275..d51569bc213f2bd354571a00910714e787120951 100644 --- a/tensorflow/contrib/cmake/external/farmhash.cmake +++ b/tensorflow/contrib/cmake/external/farmhash.cmake @@ -33,6 +33,7 @@ if(WIN32) URL_HASH ${farmhash_HASH} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${farmhash_STATIC_LIBRARIES} PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_CURRENT_SOURCE_DIR}/patches/farmhash/CMakeLists.txt ${farmhash_BUILD} INSTALL_DIR ${farmhash_INSTALL} CMAKE_CACHE_ARGS diff --git a/tensorflow/contrib/cmake/external/fft2d.cmake b/tensorflow/contrib/cmake/external/fft2d.cmake index d3af2a46761c0f7f0b5db134af8400fc93f2f095..a7bc50d5bcd4384d5c943d681fd7cd6fa1ffa796 100644 --- a/tensorflow/contrib/cmake/external/fft2d.cmake +++ b/tensorflow/contrib/cmake/external/fft2d.cmake @@ -29,6 +29,7 @@ if(WIN32) URL_HASH ${fft2d_HASH} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${fft2d_STATIC_LIBRARIES} PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_CURRENT_SOURCE_DIR}/patches/fft2d/CMakeLists.txt ${fft2d_BUILD}/src/fft2d/CMakeLists.txt INSTALL_DIR ${fft2d_INSTALL} CMAKE_CACHE_ARGS diff --git a/tensorflow/contrib/cmake/external/gif.cmake b/tensorflow/contrib/cmake/external/gif.cmake index 3d53c51fffcec1602a3b5553cdf3b225e3b0ae46..e1f8d13f8ea47b83e4a1840afac7398ef226eb45 100644 --- a/tensorflow/contrib/cmake/external/gif.cmake +++ b/tensorflow/contrib/cmake/external/gif.cmake @@ -33,6 +33,7 @@ if(WIN32) PREFIX gif URL ${gif_URL} URL_HASH ${gif_HASH} + BUILD_BYPRODUCTS ${gif_STATIC_LIBRARIES} PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_SOURCE_DIR}/patches/gif/CMakeLists.txt ${gif_BUILD} INSTALL_DIR ${gif_INSTALL} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" diff --git a/tensorflow/contrib/cmake/external/googletest.cmake b/tensorflow/contrib/cmake/external/googletest.cmake index d09bb02890f25a0312e62c876c1729e57a059e82..7cc5ae6390934773635cf7a4dff77a3cbfb41ba1 100644 --- a/tensorflow/contrib/cmake/external/googletest.cmake +++ b/tensorflow/contrib/cmake/external/googletest.cmake @@ -20,8 +20,13 @@ set(googletest_BUILD ${CMAKE_CURRENT_BINARY_DIR}/googletest/) set(googletest_TAG ec44c6c1675c25b9827aacd08c02433cccde7780) if(WIN32) - set(googletest_STATIC_LIBRARIES - ${CMAKE_CURRENT_BINARY_DIR}/googletest/src/googletest/googletest/$(Configuration)/gtest.lib) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(googletest_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/googletest/src/googletest/googletest/$(Configuration)/gtest.lib) + else() + set(googletest_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/googletest/src/googletest/googletest/gtest.lib) + endif() else() set(googletest_STATIC_LIBRARIES ${CMAKE_CURRENT_BINARY_DIR}/googletest/src/googletest/googletest/${CMAKE_BUILD_TYPE}/gtest.a) @@ -33,6 +38,7 @@ ExternalProject_Add(googletest GIT_TAG ${googletest_TAG} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${googletest_STATIC_LIBRARIES} #PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_SOURCE_DIR}/patches/grpc/CMakeLists.txt ${GRPC_BUILD} INSTALL_COMMAND "" CMAKE_CACHE_ARGS diff --git a/tensorflow/contrib/cmake/external/grpc.cmake b/tensorflow/contrib/cmake/external/grpc.cmake index 41ea0b48a4600d7ca2dd2f4a61c14ec0cc5b4734..a9f43a3ecba4830533efcc13f8c4c1c61fe1ef78 100644 --- a/tensorflow/contrib/cmake/external/grpc.cmake +++ b/tensorflow/contrib/cmake/external/grpc.cmake @@ -17,13 +17,20 @@ include (ExternalProject) set(GRPC_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/include) set(GRPC_URL https://github.com/grpc/grpc.git) set(GRPC_BUILD ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc) -set(GRPC_TAG 54e8f37e537794c2d814c1604c1282125f64f093) +set(GRPC_TAG 730b778632e79cc3c96ad237f282d687ee325ce7) if(WIN32) - set(grpc_STATIC_LIBRARIES - ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/grpc++_unsecure.lib - ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/grpc_unsecure.lib - ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/gpr.lib) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(grpc_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/grpc++_unsecure.lib + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/grpc_unsecure.lib + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/gpr.lib) + else() + set(grpc_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/grpc++_unsecure.lib + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/grpc_unsecure.lib + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/gpr.lib) + endif() else() set(grpc_STATIC_LIBRARIES ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/libgrpc++_unsecure.a @@ -40,6 +47,7 @@ ExternalProject_Add(grpc GIT_TAG ${GRPC_TAG} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${grpc_STATIC_LIBRARIES} BUILD_COMMAND ${CMAKE_COMMAND} --build . --config Release --target grpc++_unsecure COMMAND ${CMAKE_COMMAND} --build . --config Release --target grpc_cpp_plugin INSTALL_COMMAND "" diff --git a/tensorflow/contrib/cmake/external/highwayhash.cmake b/tensorflow/contrib/cmake/external/highwayhash.cmake index 2c23bef8a331de356c93dbf9d0e91d8bb13bd6c8..a6e8a38d8c2ee3deb5453c264e0c5eb23248301f 100644 --- a/tensorflow/contrib/cmake/external/highwayhash.cmake +++ b/tensorflow/contrib/cmake/external/highwayhash.cmake @@ -42,6 +42,7 @@ ExternalProject_Add(highwayhash GIT_TAG ${highwayhash_TAG} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${highwayhash_STATIC_LIBRARIES} PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_CURRENT_SOURCE_DIR}/patches/highwayhash/CMakeLists.txt ${highwayhash_BUILD} INSTALL_DIR ${highwayhash_INSTALL} CMAKE_CACHE_ARGS diff --git a/tensorflow/contrib/cmake/external/jemalloc.cmake b/tensorflow/contrib/cmake/external/jemalloc.cmake index 198ba13e64e4b6df57c4325a0104b1a6745d173a..afadcc007d66414be3306e91e7186a00b6e587ce 100644 --- a/tensorflow/contrib/cmake/external/jemalloc.cmake +++ b/tensorflow/contrib/cmake/external/jemalloc.cmake @@ -24,8 +24,11 @@ if (WIN32) ${jemalloc_INCLUDE_DIRS} ${CMAKE_CURRENT_BINARY_DIR}/jemalloc/src/jemalloc/include/msvc_compat ) - set(jemalloc_ADDITIONAL_CMAKE_OPTIONS -A x64) - set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/Release/jemalloc.lib) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/Release/jemalloc.lib) + else() + set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/jemalloc.lib) + endif() else() set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/Release/jemalloc.a) endif() @@ -36,12 +39,12 @@ ExternalProject_Add(jemalloc URL_HASH ${jemalloc_HASH} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 - CONFIGURE_COMMAND ${CMAKE_COMMAND} + BUILD_BYPRODUCTS ${jemalloc_STATIC_LIBRARIES} + BUILD_COMMAND ${CMAKE_COMMAND} --build . --config Release --target jemalloc + INSTALL_COMMAND ${CMAKE_COMMAND} -E echo "Skipping install step." + CMAKE_CACHE_ARGS -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF -Dwith-jemalloc-prefix:STRING=jemalloc_ -Dwithout-export:BOOL=ON - ${jemalloc_ADDITIONAL_CMAKE_OPTIONS} - BUILD_COMMAND ${CMAKE_COMMAND} --build . --config Release --target jemalloc - INSTALL_COMMAND ${CMAKE_COMMAND} -E echo "Skipping install step." ) diff --git a/tensorflow/contrib/cmake/external/jpeg.cmake b/tensorflow/contrib/cmake/external/jpeg.cmake index d9a165e856c588880ebdf996666d70c9e7f53da8..c1c5842aa4454f1c95ec284392194a89d47ee8d5 100644 --- a/tensorflow/contrib/cmake/external/jpeg.cmake +++ b/tensorflow/contrib/cmake/external/jpeg.cmake @@ -46,6 +46,7 @@ if (WIN32) PREFIX jpeg URL ${jpeg_URL} URL_HASH ${jpeg_HASH} + BUILD_BYPRODUCTS ${jpeg_STATIC_LIBRARIES} PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_CURRENT_SOURCE_DIR}/patches/jpeg/CMakeLists.txt ${jpeg_BUILD} INSTALL_DIR ${jpeg_INSTALL} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" diff --git a/tensorflow/contrib/cmake/external/jsoncpp.cmake b/tensorflow/contrib/cmake/external/jsoncpp.cmake index d2ae4c76e8cd175cdc3ba41fdf4e4009f8237309..84c52e3652ff935c287d32c0c80fd407e1213f29 100644 --- a/tensorflow/contrib/cmake/external/jsoncpp.cmake +++ b/tensorflow/contrib/cmake/external/jsoncpp.cmake @@ -23,7 +23,11 @@ set(jsoncpp_LIBRARIES ${jsoncpp_BUILD}/obj/so/libjsoncpp.so) set(jsoncpp_INCLUDES ${jsoncpp_BUILD}) if(WIN32) - set(jsoncpp_STATIC_LIBRARIES ${jsoncpp_BUILD}/$(Configuration)/jsoncpp.lib) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(jsoncpp_STATIC_LIBRARIES ${jsoncpp_BUILD}/$(Configuration)/jsoncpp.lib) + else() + set(jsoncpp_STATIC_LIBRARIES ${jsoncpp_BUILD}/jsoncpp.lib) + endif() else() set(jsoncpp_STATIC_LIBRARIES ${jsoncpp_BUILD}/libjsoncpp.a) endif() @@ -40,13 +44,10 @@ ExternalProject_Add(jsoncpp GIT_TAG ${jsoncpp_TAG} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${jsoncpp_STATIC_LIBRARIES} INSTALL_COMMAND "" CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF ) diff --git a/tensorflow/contrib/cmake/external/lmdb.cmake b/tensorflow/contrib/cmake/external/lmdb.cmake index e41384f023ca9fc4cba697917b491af5a9db92bc..ed5ab788acc5625b9c8020fce15f027d98433096 100644 --- a/tensorflow/contrib/cmake/external/lmdb.cmake +++ b/tensorflow/contrib/cmake/external/lmdb.cmake @@ -20,31 +20,28 @@ set(lmdb_HASH SHA256=108532fb94c6f227558d45be3f3347b52539f0f58290a7bb31ec06c462d set(lmdb_BUILD ${CMAKE_BINARY_DIR}/lmdb/src/lmdb) set(lmdb_INSTALL ${CMAKE_BINARY_DIR}/lmdb/install) +if(WIN32) + set(lmdb_STATIC_LIBRARIES ${lmdb_INSTALL}/lib/lmdb.lib) +else() + set(lmdb_STATIC_LIBRARIES ${lmdb_INSTALL}/lib/liblmdb.a) +endif() + ExternalProject_Add(lmdb PREFIX lmdb URL ${lmdb_URL} URL_HASH ${lmdb_HASH} + BUILD_BYPRODUCTS ${lmdb_STATIC_LIBRARIES} PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_CURRENT_SOURCE_DIR}/patches/lmdb/CMakeLists.txt ${lmdb_BUILD} INSTALL_DIR ${lmdb_INSTALL} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF -DCMAKE_INSTALL_PREFIX:STRING=${lmdb_INSTALL} ) -if(WIN32) - set(lmdb_STATIC_LIBRARIES ${lmdb_INSTALL}/lib/lmdb.lib) -else() - set(lmdb_STATIC_LIBRARIES ${lmdb_INSTALL}/lib/liblmdb.a) -endif() - set(lmdb_HEADERS "${lmdb_INSTALL}/include/lmdb.h" "${lmdb_INSTALL}/include/midl.h" diff --git a/tensorflow/contrib/cmake/external/nsync.cmake b/tensorflow/contrib/cmake/external/nsync.cmake index 05080060479b6240edb8ab9f65160b3dd182feb9..f3a37ff5088e3f9e54e38c0edb5777c27b26969f 100644 --- a/tensorflow/contrib/cmake/external/nsync.cmake +++ b/tensorflow/contrib/cmake/external/nsync.cmake @@ -42,6 +42,7 @@ ExternalProject_Add(nsync GIT_TAG ${nsync_TAG} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${nsync_STATIC_LIBRARIES} PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_CURRENT_SOURCE_DIR}/patches/nsync/CMakeLists.txt ${nsync_BUILD} INSTALL_DIR ${nsync_INSTALL} CMAKE_CACHE_ARGS diff --git a/tensorflow/contrib/cmake/external/png.cmake b/tensorflow/contrib/cmake/external/png.cmake index aad6618f52f909096fd2388e867ef3a965d033cb..6cd66a65990e7a2b963b52b310061b551752cd4d 100644 --- a/tensorflow/contrib/cmake/external/png.cmake +++ b/tensorflow/contrib/cmake/external/png.cmake @@ -21,9 +21,19 @@ set(png_BUILD ${CMAKE_BINARY_DIR}/png/src/png) set(png_INSTALL ${CMAKE_BINARY_DIR}/png/install) if(WIN32) - set(png_STATIC_LIBRARIES - debug ${CMAKE_BINARY_DIR}/png/install/lib/libpng12_staticd.lib - optimized ${CMAKE_BINARY_DIR}/png/install/lib/libpng12_static.lib) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(png_STATIC_LIBRARIES + debug ${CMAKE_BINARY_DIR}/png/install/lib/libpng12_staticd.lib + optimized ${CMAKE_BINARY_DIR}/png/install/lib/libpng12_static.lib) + else() + if(CMAKE_BUILD_TYPE EQUAL Debug) + set(png_STATIC_LIBRARIES + ${CMAKE_BINARY_DIR}/png/install/lib/libpng12_staticd.lib) + else() + set(png_STATIC_LIBRARIES + ${CMAKE_BINARY_DIR}/png/install/lib/libpng12_static.lib) + endif() + endif() else() set(png_STATIC_LIBRARIES ${CMAKE_BINARY_DIR}/png/install/lib/libpng12.a) endif() @@ -38,14 +48,11 @@ ExternalProject_Add(png DEPENDS zlib URL ${png_URL} URL_HASH ${png_HASH} + BUILD_BYPRODUCTS ${png_STATIC_LIBRARIES} INSTALL_DIR ${png_INSTALL} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF -DCMAKE_INSTALL_PREFIX:STRING=${png_INSTALL} diff --git a/tensorflow/contrib/cmake/external/protobuf.cmake b/tensorflow/contrib/cmake/external/protobuf.cmake index b53857a47bfbf797af02fe7f69474263119161cd..aba8a5244e17d717293deec6d9b6e8e725ef010e 100644 --- a/tensorflow/contrib/cmake/external/protobuf.cmake +++ b/tensorflow/contrib/cmake/external/protobuf.cmake @@ -16,14 +16,37 @@ include (ExternalProject) set(PROTOBUF_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/src) set(PROTOBUF_URL https://github.com/google/protobuf.git) -set(PROTOBUF_TAG b04e5cba356212e4e8c66c61bbe0c3a20537c5b9) +set(PROTOBUF_TAG 396336eb961b75f03b25824fe86cf6490fb75e3a) if(WIN32) - set(protobuf_STATIC_LIBRARIES - debug ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/$(Configuration)/libprotobufd.lib - optimized ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/$(Configuration)/libprotobuf.lib) - set(PROTOBUF_PROTOC_EXECUTABLE ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/$(Configuration)/protoc.exe) - set(PROTOBUF_ADDITIONAL_CMAKE_OPTIONS -Dprotobuf_MSVC_STATIC_RUNTIME:BOOL=OFF -A x64) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(protobuf_STATIC_LIBRARIES + debug ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/$(Configuration)/libprotobufd.lib + optimized ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/$(Configuration)/libprotobuf.lib) + set(PROTOBUF_PROTOC_EXECUTABLE ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/$(Configuration)/protoc.exe) + else() + if(CMAKE_BUILD_TYPE EQUAL Debug) + set(protobuf_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/libprotobufd.lib) + else() + set(protobuf_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/libprotobuf.lib) + endif() + set(PROTOBUF_PROTOC_EXECUTABLE ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/protoc.exe) + endif() + + # This section is to make sure CONFIGURE_COMMAND use the same generator settings + set(PROTOBUF_GENERATOR_PLATFORM) + if (CMAKE_GENERATOR_PLATFORM) + set(PROTOBUF_GENERATOR_PLATFORM -A ${CMAKE_GENERATOR_PLATFORM}) + endif() + set(PROTOBUF_GENERATOR_TOOLSET) + if (CMAKE_GENERATOR_TOOLSET) + set(PROTOBUF_GENERATOR_TOOLSET -T ${CMAKE_GENERATOR_TOOLSET}) + endif() + set(PROTOBUF_ADDITIONAL_CMAKE_OPTIONS -Dprotobuf_MSVC_STATIC_RUNTIME:BOOL=OFF + -G${CMAKE_GENERATOR} ${PROTOBUF_GENERATOR_PLATFORM} ${PROTOBUF_GENERATOR_TOOLSET}) + # End of section else() set(protobuf_STATIC_LIBRARIES ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/libprotobuf.a) set(PROTOBUF_PROTOC_EXECUTABLE ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/protoc) @@ -36,20 +59,23 @@ ExternalProject_Add(protobuf GIT_TAG ${PROTOBUF_TAG} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${PROTOBUF_PROTOC_EXECUTABLE} ${protobuf_STATIC_LIBRARIES} SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf + # SOURCE_SUBDIR cmake/ # Requires CMake 3.7, this will allow removal of CONFIGURE_COMMAND + # CONFIGURE_COMMAND resets some settings made in CMAKE_CACHE_ARGS and the generator used CONFIGURE_COMMAND ${CMAKE_COMMAND} cmake/ - -Dprotobuf_BUILD_TESTS=OFF - -DCMAKE_POSITION_INDEPENDENT_CODE=ON + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} + -DCMAKE_BUILD_TYPE:STRING=Release + -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF + -Dprotobuf_BUILD_TESTS:BOOL=OFF -DZLIB_ROOT=${ZLIB_INSTALL} ${PROTOBUF_ADDITIONAL_CMAKE_OPTIONS} INSTALL_COMMAND "" CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF + -Dprotobuf_BUILD_TESTS:BOOL=OFF + -Dprotobuf_MSVC_STATIC_RUNTIME:BOOL=OFF -DZLIB_ROOT:STRING=${ZLIB_INSTALL} ) diff --git a/tensorflow/contrib/cmake/external/re2.cmake b/tensorflow/contrib/cmake/external/re2.cmake index d10f5959f71dd350e6e2bcb81be8882b203fb231..c4bc0b1707bf9e86ea41234c8155fd6321c4c33b 100644 --- a/tensorflow/contrib/cmake/external/re2.cmake +++ b/tensorflow/contrib/cmake/external/re2.cmake @@ -21,7 +21,11 @@ set(re2_INSTALL ${CMAKE_CURRENT_BINARY_DIR}/re2/install) set(re2_TAG e7efc48) if(WIN32) - set(re2_STATIC_LIBRARIES ${re2_BUILD}/$(Configuration)/re2.lib) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(re2_STATIC_LIBRARIES ${re2_BUILD}/$(Configuration)/re2.lib) + else() + set(re2_STATIC_LIBRARIES ${re2_BUILD}/re2.lib) + endif() else() set(re2_STATIC_LIBRARIES ${re2_BUILD}/libre2.a) endif() @@ -36,13 +40,10 @@ ExternalProject_Add(re2 GIT_TAG ${re2_TAG} INSTALL_DIR ${re2_INSTALL} BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${re2_STATIC_LIBRARIES} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_INSTALL_PREFIX:STRING=${re2_INSTALL} -DRE2_BUILD_TESTING:BOOL=OFF diff --git a/tensorflow/contrib/cmake/external/snappy.cmake b/tensorflow/contrib/cmake/external/snappy.cmake index 926c271fd9ea6e2a30251aa408bd49859ae95070..f54197643b06781dad35b40f526f28d301047299 100644 --- a/tensorflow/contrib/cmake/external/snappy.cmake +++ b/tensorflow/contrib/cmake/external/snappy.cmake @@ -20,7 +20,11 @@ set(snappy_BUILD ${CMAKE_CURRENT_BINARY_DIR}/snappy/src/snappy) set(snappy_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/snappy/src/snappy) if(WIN32) - set(snappy_STATIC_LIBRARIES ${snappy_BUILD}/$(Configuration)/snappy.lib) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(snappy_STATIC_LIBRARIES ${snappy_BUILD}/$(Configuration)/snappy.lib) + else() + set(snappy_STATIC_LIBRARIES ${snappy_BUILD}/snappy.lib) + endif() else() set(snappy_STATIC_LIBRARIES ${snappy_BUILD}/libsnappy.a) endif() @@ -35,20 +39,17 @@ ExternalProject_Add(snappy GIT_TAG ${snappy_TAG} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${snappy_STATIC_LIBRARIES} INSTALL_COMMAND "" LOG_DOWNLOAD ON LOG_CONFIGURE ON LOG_BUILD ON CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF -DSNAPPY_BUILD_TESTS:BOOL=OFF ) # actually enables snappy in the source code -add_definitions(-DTF_USE_SNAPPY) \ No newline at end of file +add_definitions(-DTF_USE_SNAPPY) diff --git a/tensorflow/contrib/cmake/external/sqlite.cmake b/tensorflow/contrib/cmake/external/sqlite.cmake index 785039a46983747557607562675349c150e064ad..57c4ae76517e4d7247093edd5e5bd95a83258d87 100644 --- a/tensorflow/contrib/cmake/external/sqlite.cmake +++ b/tensorflow/contrib/cmake/external/sqlite.cmake @@ -28,6 +28,7 @@ endif() set(sqlite_HEADERS "${sqlite_BUILD}/sqlite3.h" + "${sqlite_BUILD}/sqlite3ext.h" ) if (WIN32) @@ -35,6 +36,7 @@ if (WIN32) PREFIX sqlite URL ${sqlite_URL} URL_HASH ${sqlite_HASH} + BUILD_BYPRODUCTS ${sqlite_STATIC_LIBRARIES} PATCH_COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_CURRENT_SOURCE_DIR}/patches/sqlite/CMakeLists.txt ${sqlite_BUILD} INSTALL_DIR ${sqlite_INSTALL} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" @@ -53,11 +55,7 @@ else() INSTALL_DIR ${sqlite_INSTALL} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} -DCMAKE_BUILD_TYPE:STRING=Release -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF -DCMAKE_INSTALL_PREFIX:STRING=${sqlite_INSTALL} diff --git a/tensorflow/contrib/cmake/external/zlib.cmake b/tensorflow/contrib/cmake/external/zlib.cmake index f10f84336e8b1c0a2c7de7ea1f8b8af7c21f8b51..116d42309394b92407cef79c9d3a975f494bc3ff 100644 --- a/tensorflow/contrib/cmake/external/zlib.cmake +++ b/tensorflow/contrib/cmake/external/zlib.cmake @@ -12,54 +12,75 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -include (ExternalProject) +if (systemlib_ZLIB) + find_package(PkgConfig) + pkg_search_module(ZLIB REQUIRED zlib) + set(zlib_INCLUDE_DIR ${ZLIB_INCLUDE_DIRS}) + set(ADD_LINK_DIRECTORY ${ADD_LINK_DIRECTORY} ${ZLIB_LIBRARY_DIRS}) + set(ADD_CFLAGS ${ADD_CFLAGS} ${ZLIB_CFLAGS_OTHER}) -set(zlib_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/external/zlib_archive) -set(ZLIB_URL https://github.com/madler/zlib) -set(ZLIB_BUILD ${CMAKE_CURRENT_BINARY_DIR}/zlib/src/zlib) -set(ZLIB_INSTALL ${CMAKE_CURRENT_BINARY_DIR}/zlib/install) -set(ZLIB_TAG 50893291621658f355bc5b4d450a8d06a563053d) + # To meet DEPENDS zlib from other projects. + # If we hit this line, zlib is already built and installed to the system. + add_custom_target(zlib) + add_custom_target(zlib_copy_headers_to_destination) -if(WIN32) - set(zlib_STATIC_LIBRARIES - debug ${CMAKE_CURRENT_BINARY_DIR}/zlib/install/lib/zlibstaticd.lib - optimized ${CMAKE_CURRENT_BINARY_DIR}/zlib/install/lib/zlibstatic.lib) -else() - set(zlib_STATIC_LIBRARIES - ${CMAKE_CURRENT_BINARY_DIR}/zlib/install/lib/libz.a) -endif() +else (systemlib_ZLIB) + include (ExternalProject) -set(ZLIB_HEADERS - "${ZLIB_INSTALL}/include/zconf.h" - "${ZLIB_INSTALL}/include/zlib.h" -) + set(zlib_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/external/zlib_archive) + set(ZLIB_URL https://github.com/madler/zlib) + set(ZLIB_BUILD ${CMAKE_CURRENT_BINARY_DIR}/zlib/src/zlib) + set(ZLIB_INSTALL ${CMAKE_CURRENT_BINARY_DIR}/zlib/install) + set(ZLIB_TAG 50893291621658f355bc5b4d450a8d06a563053d) -ExternalProject_Add(zlib - PREFIX zlib - GIT_REPOSITORY ${ZLIB_URL} - GIT_TAG ${ZLIB_TAG} - INSTALL_DIR ${ZLIB_INSTALL} - BUILD_IN_SOURCE 1 - DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" - CMAKE_CACHE_ARGS - if(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE) - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON - else() - -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=OFF - endif() - -DCMAKE_BUILD_TYPE:STRING=Release - -DCMAKE_INSTALL_PREFIX:STRING=${ZLIB_INSTALL} -) + if(WIN32) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(zlib_STATIC_LIBRARIES + debug ${CMAKE_CURRENT_BINARY_DIR}/zlib/install/lib/zlibstaticd.lib + optimized ${CMAKE_CURRENT_BINARY_DIR}/zlib/install/lib/zlibstatic.lib) + else() + if(CMAKE_BUILD_TYPE EQUAL Debug) + set(zlib_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/zlib/install/lib/zlibstaticd.lib) + else() + set(zlib_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/zlib/install/lib/zlibstatic.lib) + endif() + endif() + else() + set(zlib_STATIC_LIBRARIES + ${CMAKE_CURRENT_BINARY_DIR}/zlib/install/lib/libz.a) + endif() -# put zlib includes in the directory where they are expected -add_custom_target(zlib_create_destination_dir - COMMAND ${CMAKE_COMMAND} -E make_directory ${zlib_INCLUDE_DIR} - DEPENDS zlib) + set(ZLIB_HEADERS + "${ZLIB_INSTALL}/include/zconf.h" + "${ZLIB_INSTALL}/include/zlib.h" + ) -add_custom_target(zlib_copy_headers_to_destination - DEPENDS zlib_create_destination_dir) + ExternalProject_Add(zlib + PREFIX zlib + GIT_REPOSITORY ${ZLIB_URL} + GIT_TAG ${ZLIB_TAG} + INSTALL_DIR ${ZLIB_INSTALL} + BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${zlib_STATIC_LIBRARIES} + DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" + CMAKE_CACHE_ARGS + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} + -DCMAKE_BUILD_TYPE:STRING=Release + -DCMAKE_INSTALL_PREFIX:STRING=${ZLIB_INSTALL} + ) -foreach(header_file ${ZLIB_HEADERS}) - add_custom_command(TARGET zlib_copy_headers_to_destination PRE_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different ${header_file} ${zlib_INCLUDE_DIR}) -endforeach() + # put zlib includes in the directory where they are expected + add_custom_target(zlib_create_destination_dir + COMMAND ${CMAKE_COMMAND} -E make_directory ${zlib_INCLUDE_DIR} + DEPENDS zlib) + + add_custom_target(zlib_copy_headers_to_destination + DEPENDS zlib_create_destination_dir) + + foreach(header_file ${ZLIB_HEADERS}) + add_custom_command(TARGET zlib_copy_headers_to_destination PRE_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different ${header_file} ${zlib_INCLUDE_DIR}) + endforeach() +endif (systemlib_ZLIB) diff --git a/tensorflow/tools/ci_build/install/install_cmake_for_clang.sh b/tensorflow/contrib/cmake/make.sh similarity index 83% rename from tensorflow/tools/ci_build/install/install_cmake_for_clang.sh rename to tensorflow/contrib/cmake/make.sh index 3e626a69ab5e6b7f8d1b4997b459301606501a8e..eed3c34aba1f0326ec741169a187eb2982f253a3 100755 --- a/tensorflow/tools/ci_build/install/install_cmake_for_clang.sh +++ b/tensorflow/contrib/cmake/make.sh @@ -14,6 +14,13 @@ # limitations under the License. # ============================================================================== -CMAKE_URL="https://cmake.org/files/v3.7/cmake-3.7.2-Linux-x86_64.tar.gz" +( +cd "$(dirname "$0")" +mkdir -p _build -wget -O - "${CMAKE_URL}" | tar xzf - -C /usr/local --strip-components=1 +( +cd _build +rm -rf -- * +cmake .. +) +) diff --git a/tensorflow/contrib/cmake/python_modules.txt b/tensorflow/contrib/cmake/python_modules.txt index 92edce77dff699e75d1873ad0f56c6c489fbc571..bfe53c01b3b5fb9db8a5d8fa280d1d7f98974882 100644 --- a/tensorflow/contrib/cmake/python_modules.txt +++ b/tensorflow/contrib/cmake/python_modules.txt @@ -1,9 +1,12 @@ +# python_sanity_test.py will complain about invalid or missing entries +# problematic entries can be commented for temporary whitelisting tensorflow tensorflow/core tensorflow/core/example tensorflow/core/framework tensorflow/core/lib tensorflow/core/lib/core +tensorflow/core/profiler tensorflow/core/protobuf tensorflow/core/util tensorflow/examples @@ -31,9 +34,11 @@ tensorflow/python/grappler tensorflow/python/keras tensorflow/python/keras/activations tensorflow/python/keras/applications +tensorflow/python/keras/applications/densenet tensorflow/python/keras/applications/inception_resnet_v2 tensorflow/python/keras/applications/inception_v3 tensorflow/python/keras/applications/mobilenet +tensorflow/python/keras/applications/nasnet tensorflow/python/keras/applications/resnet50 tensorflow/python/keras/applications/vgg16 tensorflow/python/keras/applications/vgg19 @@ -86,9 +91,7 @@ tensorflow/python/ops/distributions tensorflow/python/ops/linalg tensorflow/python/ops/losses tensorflow/python/platform -tensorflow/python/platform/default -tensorflow/python/platform/summary -tensorflow/python/profiler/ +tensorflow/python/profiler tensorflow/python/profiler/internal tensorflow/python/saved_model tensorflow/python/summary @@ -111,12 +114,9 @@ tensorflow/contrib/android/java/org/tensorflow/contrib tensorflow/contrib/android/java/org/tensorflow/contrib/android tensorflow/contrib/android/jni tensorflow/contrib/batching -tensorflow/contrib/batching/kernels tensorflow/contrib/batching/python tensorflow/contrib/batching/python/ops tensorflow/contrib/bayesflow -tensorflow/contrib/bayesflow/examples -tensorflow/contrib/bayesflow/examples/reinforce_simple tensorflow/contrib/bayesflow/python tensorflow/contrib/bayesflow/python/ops tensorflow/contrib/boosted_trees @@ -134,6 +134,11 @@ tensorflow/contrib/cloud/python/ops tensorflow/contrib/cluster_resolver tensorflow/contrib/cluster_resolver/python tensorflow/contrib/cluster_resolver/python/training +tensorflow/contrib/coder +tensorflow/contrib/coder/kernels +tensorflow/contrib/coder/ops +tensorflow/contrib/coder/python +tensorflow/contrib/coder/python/ops tensorflow/contrib/compiler tensorflow/contrib/copy_graph tensorflow/contrib/copy_graph/python @@ -170,6 +175,9 @@ tensorflow/contrib/factorization/kernels tensorflow/contrib/factorization/ops tensorflow/contrib/factorization/python tensorflow/contrib/factorization/python/ops +tensorflow/contrib/feature_column +tensorflow/contrib/feature_column/python +tensorflow/contrib/feature_column/python/feature_column tensorflow/contrib/ffmpeg tensorflow/contrib/ffmpeg/default tensorflow/contrib/framework @@ -212,16 +220,8 @@ tensorflow/contrib/input_pipeline/python/ops tensorflow/contrib/integrate tensorflow/contrib/integrate/python tensorflow/contrib/integrate/python/ops -tensorflow/contrib/ios_examples -tensorflow/contrib/ios_examples/benchmark -tensorflow/contrib/ios_examples/benchmark/benchmark.xcodeproj -tensorflow/contrib/ios_examples/benchmark/data -tensorflow/contrib/ios_examples/camera -tensorflow/contrib/ios_examples/camera/camera_example.xcodeproj -tensorflow/contrib/ios_examples/camera/en.lproj -tensorflow/contrib/ios_examples/simple -tensorflow/contrib/ios_examples/simple/data -tensorflow/contrib/ios_examples/simple/tf_ios_makefile_example.xcodeproj +tensorflow/contrib/kafka/python +tensorflow/contrib/kafka/python/ops tensorflow/contrib/keras tensorflow/contrib/keras/api tensorflow/contrib/keras/api/keras @@ -276,9 +276,6 @@ tensorflow/contrib/layers/python/ops tensorflow/contrib/learn tensorflow/contrib/learn/python tensorflow/contrib/learn/python/learn -tensorflow/contrib/learn/python/learn/dataframe -tensorflow/contrib/learn/python/learn/dataframe/queues -tensorflow/contrib/learn/python/learn/dataframe/transforms tensorflow/contrib/learn/python/learn/datasets tensorflow/contrib/learn/python/learn/datasets/data tensorflow/contrib/learn/python/learn/estimators @@ -301,6 +298,11 @@ tensorflow/contrib/linear_optimizer/kernels tensorflow/contrib/linear_optimizer/kernels/g3doc tensorflow/contrib/linear_optimizer/python tensorflow/contrib/linear_optimizer/python/ops +# TODO(drpngx): Fix failing imports +# tensorflow/contrib/lite +# tensorflow/contrib/lite/python +# tensorflow/contrib/lite/toco +# tensorflow/contrib/lite/toco/python tensorflow/contrib/lookup tensorflow/contrib/losses tensorflow/contrib/losses/python @@ -314,10 +316,11 @@ tensorflow/contrib/memory_stats/python tensorflow/contrib/memory_stats/python/ops tensorflow/contrib/meta_graph_transform tensorflow/contrib/metrics -tensorflow/contrib/metrics/ops tensorflow/contrib/metrics/python tensorflow/contrib/metrics/python/metrics tensorflow/contrib/metrics/python/ops +tensorflow/contrib/mpi_collectives/python +tensorflow/contrib/mpi_collectives/python/ops tensorflow/contrib/model_pruning tensorflow/contrib/model_pruning/examples tensorflow/contrib/model_pruning/examples/cifar10 @@ -328,8 +331,6 @@ tensorflow/contrib/nccl/kernels tensorflow/contrib/nccl/ops tensorflow/contrib/nccl/python tensorflow/contrib/nccl/python/ops -tensorflow/contrib/ndlstm -tensorflow/contrib/ndlstm/python tensorflow/contrib/nearest_neighbor/kernels tensorflow/contrib/nearest_neighbor/ops tensorflow/contrib/nearest_neighbor/python @@ -346,7 +347,6 @@ tensorflow/contrib/pi_examples/label_image tensorflow/contrib/pi_examples/label_image/data tensorflow/contrib/periodic_resample tensorflow/contrib/periodic_resample/python -tensorflow/contrib/periodic_resample/python/kernels tensorflow/contrib/periodic_resample/python/ops tensorflow/contrib/predictor tensorflow/contrib/quantization @@ -355,11 +355,14 @@ tensorflow/contrib/quantize tensorflow/contrib/quantize/python tensorflow/contrib/receptive_field tensorflow/contrib/receptive_field/python +tensorflow/contrib/receptive_field/python/util +tensorflow/contrib/receptive_field/python/util/examples tensorflow/contrib/reduce_slice_ops tensorflow/contrib/reduce_slice_ops/kernels tensorflow/contrib/reduce_slice_ops/ops tensorflow/contrib/reduce_slice_ops/python tensorflow/contrib/reduce_slice_ops/python/ops +tensorflow/contrib/remote_fused_graph tensorflow/contrib/remote_fused_graph/pylib tensorflow/contrib/remote_fused_graph/pylib/python tensorflow/contrib/remote_fused_graph/pylib/python/ops @@ -409,20 +412,21 @@ tensorflow/contrib/summary tensorflow/contrib/tensorboard tensorflow/contrib/tensorboard/plugins tensorflow/contrib/tensorboard/plugins/projector +tensorflow/contrib/tensorboard/plugins/trace +# TODO(sami): Add cmake implementations. +# tensorflow/contrib/tensorrt/python +# tensorflow/contrib/tensorrt/python/ops tensorflow/contrib/tensor_forest tensorflow/contrib/tensor_forest/client -tensorflow/contrib/tensor_forest/core -tensorflow/contrib/tensor_forest/core/ops -tensorflow/contrib/tensor_forest/data tensorflow/contrib/tensor_forest/hybrid tensorflow/contrib/tensor_forest/hybrid/core tensorflow/contrib/tensor_forest/hybrid/core/ops -tensorflow/contrib/tensor_forest/hybrid/ops tensorflow/contrib/tensor_forest/hybrid/python tensorflow/contrib/tensor_forest/hybrid/python/layers tensorflow/contrib/tensor_forest/hybrid/python/models tensorflow/contrib/tensor_forest/hybrid/python/ops tensorflow/contrib/tensor_forest/kernels +tensorflow/contrib/tensor_forest/proto tensorflow/contrib/tensor_forest/python tensorflow/contrib/tensor_forest/python/ops tensorflow/contrib/testing @@ -443,6 +447,7 @@ tensorflow/contrib/timeseries/python/timeseries/state_space_models tensorflow/contrib/tpu tensorflow/contrib/tpu/ops tensorflow/contrib/tpu/profiler +tensorflow/contrib/tpu/proto tensorflow/contrib/tpu/python tensorflow/contrib/tpu/python/ops tensorflow/contrib/tpu/python/profiler diff --git a/tensorflow/contrib/cmake/python_sanity_test.py b/tensorflow/contrib/cmake/python_sanity_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e0056823a80833329bcb1f275a3384a33127bb40 --- /dev/null +++ b/tensorflow/contrib/cmake/python_sanity_test.py @@ -0,0 +1,128 @@ +# Copyright 2017 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. +# ============================================================================== +"""Complain about invalid or missing entries in python_*.txt files. + +Problematic entries can be commented for temporary whitelisting. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import unittest + + +def abs_path(path): + root = os.path.dirname(__file__) + + for _ in range(3): + root = os.path.join(root, os.pardir) + + path = os.path.join(root, path) + path = os.path.abspath(path) + return path + + +def read_entries(test): + with open(abs_path(test.entries_file), "r") as f: + lines = f.readlines() + + lines = [line.strip() for line in lines] + lines = [line for line in lines if line] + + test.entries = [] + test.whitelist = [] + + for line in lines: + # line is comment + if line.startswith("#"): + line = line[1:].strip() + # whitelist entry + if line.startswith("tensorflow/"): + test.whitelist.append(line) + # line has comment -> strip comment + elif line.find("#") != -1: + line = line[:line.find("#")].strip() + test.entries.append(line) + else: + test.entries.append(line) + + +def test_invalid_directories(test): + for entry in test.entries: + if not os.path.isdir(abs_path(entry)): + problem = "'" + test.entries_file + "' contains invalid '" + entry + "'" + solution = ("Please remove the invalid entry (or add the missing " + "directory).") + raise AssertionError(problem + "\n" + solution) + + +def test_missing_directory(test, path): + if path in test.whitelist: + return + + dir_exists = os.path.isdir(abs_path(path)) + entry_exists = path in test.entries + + if dir_exists and not entry_exists: + problem = "'" + test.entries_file + "' is missing '" + path + "'" + solution = "Please add the missing entry (comment to whitelist if needed)." + raise AssertionError(problem + "\n" + solution) + + +class PythonModuleTest(unittest.TestCase): + + def setUp(self): + self.entries_file = "tensorflow/contrib/cmake/python_modules.txt" + read_entries(self) + + def testInvalidEntries(self): + test_invalid_directories(self) + + def testMissingModules(self): + module_names = next(os.walk(abs_path("tensorflow/contrib")))[1] + + for module_name in module_names: + path = "tensorflow/contrib/" + module_name + + test_missing_directory(self, path + "/python") + test_missing_directory(self, path + "/python/ops") + test_missing_directory(self, path + "/python/kernels") + test_missing_directory(self, path + "/python/layers") + + +class PythonProtoTest(unittest.TestCase): + + def setUp(self): + self.entries_file = "tensorflow/contrib/cmake/python_protos.txt" + read_entries(self) + + def testInvalidEntries(self): + test_invalid_directories(self) + + +class PythonProtoCCTest(unittest.TestCase): + + def setUp(self): + self.entries_file = "tensorflow/contrib/cmake/python_protos_cc.txt" + read_entries(self) + + def testInvalidEntries(self): + test_invalid_directories(self) + + +if __name__ == "__main__": + unittest.main() diff --git a/tensorflow/contrib/cmake/tests/cuda/compatibility_test.c b/tensorflow/contrib/cmake/tests/cuda/compatibility_test.c new file mode 100644 index 0000000000000000000000000000000000000000..968ab13a0c43793341431248713f81010c87f148 --- /dev/null +++ b/tensorflow/contrib/cmake/tests/cuda/compatibility_test.c @@ -0,0 +1,7 @@ +// This is a program to test if compiler is compatible with CUDA. +#define __CUDACC__ +#include "crt/host_config.h" + +int main(void) { + return 0; +} diff --git a/tensorflow/contrib/cmake/tests/cuda/compatibility_test.cc b/tensorflow/contrib/cmake/tests/cuda/compatibility_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..968ab13a0c43793341431248713f81010c87f148 --- /dev/null +++ b/tensorflow/contrib/cmake/tests/cuda/compatibility_test.cc @@ -0,0 +1,7 @@ +// This is a program to test if compiler is compatible with CUDA. +#define __CUDACC__ +#include "crt/host_config.h" + +int main(void) { + return 0; +} diff --git a/tensorflow/contrib/cmake/tf_cc_ops.cmake b/tensorflow/contrib/cmake/tf_cc_ops.cmake index 6e2ac203f9a7f96cb14752a91483840a9eb6b451..f73da0b8ab18af1eca4c2bd577604595f8b8ec6d 100644 --- a/tensorflow/contrib/cmake/tf_cc_ops.cmake +++ b/tensorflow/contrib/cmake/tf_cc_ops.cmake @@ -83,7 +83,7 @@ foreach(tf_cc_op_lib_name ${tf_cc_op_lib_names}) ${cc_ops_target_dir}/${tf_cc_op_lib_name}.cc ${cc_ops_target_dir}/${tf_cc_op_lib_name}_internal.h ${cc_ops_target_dir}/${tf_cc_op_lib_name}_internal.cc - COMMAND ${tf_cc_op_lib_name}_gen_cc ${cc_ops_target_dir}/${tf_cc_op_lib_name}.h ${cc_ops_target_dir}/${tf_cc_op_lib_name}.cc ${tensorflow_source_dir}/tensorflow/cc/ops/op_gen_overrides.pbtxt ${cc_ops_include_internal} ${tensorflow_source_dir}/tensorflow/core/api_def/base_api + COMMAND ${tf_cc_op_lib_name}_gen_cc ${cc_ops_target_dir}/${tf_cc_op_lib_name}.h ${cc_ops_target_dir}/${tf_cc_op_lib_name}.cc ${cc_ops_include_internal} ${tensorflow_source_dir}/tensorflow/core/api_def/base_api DEPENDS ${tf_cc_op_lib_name}_gen_cc create_cc_ops_header_dir ) @@ -149,7 +149,11 @@ add_library(tf_cc OBJECT ${tf_cc_srcs}) add_dependencies(tf_cc tf_cc_framework tf_cc_ops) if (WIN32) - set (pywrap_tensorflow_lib "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}/pywrap_tensorflow_internal.lib") + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set (pywrap_tensorflow_lib "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}/pywrap_tensorflow_internal.lib") + else() + set (pywrap_tensorflow_lib "${CMAKE_CURRENT_BINARY_DIR}/pywrap_tensorflow_internal.lib") + endif() else (WIN32) set (pywrap_tensorflow_lib "${CMAKE_CURRENT_BINARY_DIR}/libpywrap_tensorflow_internal.so") endif (WIN32) diff --git a/tensorflow/contrib/cmake/tf_core_cpu.cmake b/tensorflow/contrib/cmake/tf_core_cpu.cmake index e4213ea2a47da2a7381cccd0504235ad62018d4e..96ac60d095dbc84470ff1be92f4bf52bb420fc52 100644 --- a/tensorflow/contrib/cmake/tf_core_cpu.cmake +++ b/tensorflow/contrib/cmake/tf_core_cpu.cmake @@ -50,6 +50,12 @@ file(GLOB_RECURSE tf_core_cpu_exclude_srcs "${tensorflow_source_dir}/tensorflow/core/graph/edgeset.cc" "${tensorflow_source_dir}/tensorflow/core/graph/graph.h" "${tensorflow_source_dir}/tensorflow/core/graph/graph.cc" + "${tensorflow_source_dir}/tensorflow/core/graph/graph_def_builder.h" + "${tensorflow_source_dir}/tensorflow/core/graph/graph_def_builder.cc" + "${tensorflow_source_dir}/tensorflow/core/graph/node_builder.h" + "${tensorflow_source_dir}/tensorflow/core/graph/node_builder.cc" + "${tensorflow_source_dir}/tensorflow/core/graph/tensor_id.h" + "${tensorflow_source_dir}/tensorflow/core/graph/tensor_id.cc" "${tensorflow_source_dir}/tensorflow/core/graph/while_context.h" "${tensorflow_source_dir}/tensorflow/core/graph/while_context.cc" "${tensorflow_source_dir}/tensorflow/core/grappler/clusters/single_machine.h" diff --git a/tensorflow/contrib/cmake/tf_core_framework.cmake b/tensorflow/contrib/cmake/tf_core_framework.cmake index 5ec1a8d04fa41c6b36400fc0998af77592866150..a1c320347fe60f87806736befc677541a93e7e93 100644 --- a/tensorflow/contrib/cmake/tf_core_framework.cmake +++ b/tensorflow/contrib/cmake/tf_core_framework.cmake @@ -126,7 +126,9 @@ endfunction() file(GLOB_RECURSE tf_protos_cc_srcs RELATIVE ${tensorflow_source_dir} "${tensorflow_source_dir}/tensorflow/core/*.proto" "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/proto/*.proto" + "${tensorflow_source_dir}/tensorflow/contrib/tpu/proto/*.proto" ) + RELATIVE_PROTOBUF_GENERATE_CPP(PROTO_SRCS PROTO_HDRS ${tensorflow_source_dir} ${tf_protos_cc_srcs} ) @@ -191,10 +193,6 @@ file(GLOB_RECURSE tf_core_lib_srcs "${tensorflow_source_dir}/tensorflow/core/lib/*.h" "${tensorflow_source_dir}/tensorflow/core/lib/*.cc" "${tensorflow_source_dir}/tensorflow/core/public/*.h" - # TODO(@jart): Move StatusOr into core. - "${tensorflow_source_dir}/tensorflow/compiler/xla/statusor.cc" - "${tensorflow_source_dir}/tensorflow/compiler/xla/statusor.h" - "${tensorflow_source_dir}/tensorflow/compiler/xla/statusor_internals.h" ) file(GLOB tf_core_platform_srcs @@ -294,6 +292,12 @@ file(GLOB_RECURSE tf_core_framework_srcs "${tensorflow_source_dir}/tensorflow/core/graph/edgeset.cc" "${tensorflow_source_dir}/tensorflow/core/graph/graph.h" "${tensorflow_source_dir}/tensorflow/core/graph/graph.cc" + "${tensorflow_source_dir}/tensorflow/core/graph/graph_def_builder.h" + "${tensorflow_source_dir}/tensorflow/core/graph/graph_def_builder.cc" + "${tensorflow_source_dir}/tensorflow/core/graph/node_builder.h" + "${tensorflow_source_dir}/tensorflow/core/graph/node_builder.cc" + "${tensorflow_source_dir}/tensorflow/core/graph/tensor_id.h" + "${tensorflow_source_dir}/tensorflow/core/graph/tensor_id.cc" "${tensorflow_source_dir}/tensorflow/core/graph/while_context.h" "${tensorflow_source_dir}/tensorflow/core/graph/while_context.cc" "${tensorflow_source_dir}/tensorflow/core/util/*.h" @@ -317,8 +321,15 @@ file(GLOB_RECURSE tf_core_framework_exclude_srcs "${tensorflow_source_dir}/tensorflow/core/util/*test*.cc" "${tensorflow_source_dir}/tensorflow/core/util/*main.cc" "${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/*test*.cc" + "${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/loader.cc" + "${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/vacuum.cc" ) +# TODO(jart): Why doesn't this work? +# set_source_files_properties( +# ${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/snapfn.cc +# PROPERTIES COMPILE_FLAGS -DSQLITE_OMIT_LOAD_EXTENSION) + list(REMOVE_ITEM tf_core_framework_srcs ${tf_core_framework_exclude_srcs}) add_library(tf_core_framework OBJECT diff --git a/tensorflow/contrib/cmake/tf_core_kernels.cmake b/tensorflow/contrib/cmake/tf_core_kernels.cmake index d3b6c0bdd385432dc469133c00960ebba0dbeec5..f219d5eb577afa9edaadca09aef9869c81d2bd87 100644 --- a/tensorflow/contrib/cmake/tf_core_kernels.cmake +++ b/tensorflow/contrib/cmake/tf_core_kernels.cmake @@ -63,10 +63,15 @@ if(tensorflow_BUILD_CONTRIB_KERNELS) "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/split_handler_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/stats_accumulator_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/training_ops.cc" + "${tensorflow_source_dir}/tensorflow/contrib/coder/kernels/range_coder.cc" + "${tensorflow_source_dir}/tensorflow/contrib/coder/kernels/range_coder_ops.cc" + "${tensorflow_source_dir}/tensorflow/contrib/coder/kernels/range_coder_ops_util.cc" + "${tensorflow_source_dir}/tensorflow/contrib/coder/ops/coder_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc" + "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc" "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/prefetching_kernels.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/ops/prefetching_ops.cc" + "${tensorflow_source_dir}/tensorflow/contrib/data/ops/dataset_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/factorization/kernels/clustering_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/factorization/kernels/wals_solver_ops.cc" @@ -79,6 +84,7 @@ if(tensorflow_BUILD_CONTRIB_KERNELS) "${tensorflow_source_dir}/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.cc" "${tensorflow_source_dir}/tensorflow/contrib/image/kernels/bipartite_match_op.cc" "${tensorflow_source_dir}/tensorflow/contrib/image/kernels/image_ops.cc" + "${tensorflow_source_dir}/tensorflow/contrib/image/kernels/segmentation_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/image/kernels/single_image_random_dot_stereograms_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/image/ops/distort_image_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/image/ops/image_ops.cc" diff --git a/tensorflow/contrib/cmake/tf_core_ops.cmake b/tensorflow/contrib/cmake/tf_core_ops.cmake index e8c2cd347327843d10d13c1d24a800ff776aa8c1..59e094812aaf4da2549d96314fc550e5635f9de8 100644 --- a/tensorflow/contrib/cmake/tf_core_ops.cmake +++ b/tensorflow/contrib/cmake/tf_core_ops.cmake @@ -15,6 +15,7 @@ set(tf_op_lib_names "audio_ops" "array_ops" + "batch_ops" "bitwise_ops" "candidate_sampling_ops" "checkpoint_ops" @@ -26,8 +27,10 @@ set(tf_op_lib_names "image_ops" "io_ops" "linalg_ops" + "list_ops" "lookup_ops" "logging_ops" + "manip_ops" "math_ops" "nn_ops" "no_op" @@ -80,8 +83,9 @@ GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_training "${tensorflow_source_dir}/ten GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_prediction "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/prediction_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_quantiles "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_stats_accumulator "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/stats_accumulator_ops.cc") +GENERATE_CONTRIB_OP_LIBRARY(coder "${tensorflow_source_dir}/tensorflow/contrib/coder/ops/coder_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(cudnn_rnn "${tensorflow_source_dir}/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc") -GENERATE_CONTRIB_OP_LIBRARY(data_prefetching "${tensorflow_source_dir}/tensorflow/contrib/data/ops/prefetching_ops.cc") +GENERATE_CONTRIB_OP_LIBRARY(data_dataset "${tensorflow_source_dir}/tensorflow/contrib/data/ops/dataset_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(factorization_clustering "${tensorflow_source_dir}/tensorflow/contrib/factorization/ops/clustering_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(factorization_factorization "${tensorflow_source_dir}/tensorflow/contrib/factorization/ops/factorization_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(framework_variable "${tensorflow_source_dir}/tensorflow/contrib/framework/ops/variable_ops.cc") diff --git a/tensorflow/contrib/cmake/tf_core_profiler.cmake b/tensorflow/contrib/cmake/tf_core_profiler.cmake index 61ed6a1e145299125d037b48b8b644cae1ce96e7..b91a7f43e5c03e933d10572e54e0c8c914c55f71 100644 --- a/tensorflow/contrib/cmake/tf_core_profiler.cmake +++ b/tensorflow/contrib/cmake/tf_core_profiler.cmake @@ -17,6 +17,8 @@ ######################################################## file(GLOB_RECURSE tf_core_profiler_srcs "${tensorflow_source_dir}/tensorflow/core/profiler/*.proto" + "${tensorflow_source_dir}/tensorflow/core/profiler/tfprof_options.h" + "${tensorflow_source_dir}/tensorflow/core/profiler/tfprof_options.cc" "${tensorflow_source_dir}/tensorflow/core/profiler/internal/*.h" "${tensorflow_source_dir}/tensorflow/core/profiler/internal/*.cc" "${tensorflow_source_dir}/tensorflow/core/profiler/internal/advisor/*.h" diff --git a/tensorflow/contrib/cmake/tf_python.cmake b/tensorflow/contrib/cmake/tf_python.cmake index 8db6929e31a1a5f5c793721f455a664bd6741b06..b730ebd3baacafe8ae401e8987104f3062372954 100755 --- a/tensorflow/contrib/cmake/tf_python.cmake +++ b/tensorflow/contrib/cmake/tf_python.cmake @@ -126,10 +126,16 @@ STRING(REGEX REPLACE ";" "\\\\;" python_protos "${python_protos}") STRING(REGEX REPLACE "\n" ";" python_protos "${python_protos}") foreach(python_proto ${python_protos}) - file(GLOB_RECURSE tf_python_protos_src RELATIVE ${tensorflow_source_dir} - "${tensorflow_source_dir}/${python_proto}/*.proto" - ) - list(APPEND tf_python_protos_srcs ${tf_python_protos_src}) + if(NOT python_proto MATCHES "^\#") + STRING(REGEX REPLACE " *\#.*" "" python_proto "${python_proto}") + if(NOT EXISTS "${tensorflow_source_dir}/${python_proto}") + message(SEND_ERROR "Python proto directory not found: ${python_proto}") + endif() + file(GLOB_RECURSE tf_python_protos_src RELATIVE ${tensorflow_source_dir} + "${tensorflow_source_dir}/${python_proto}/*.proto" + ) + list(APPEND tf_python_protos_srcs ${tf_python_protos_src}) + endif() endforeach(python_proto) RELATIVE_PROTOBUF_GENERATE_PYTHON( @@ -142,10 +148,16 @@ STRING(REGEX REPLACE ";" "\\\\;" python_protos_cc "${python_protos_cc}") STRING(REGEX REPLACE "\n" ";" python_protos_cc "${python_protos_cc}") foreach(python_proto_cc ${python_protos_cc}) - file(GLOB_RECURSE tf_python_protos_cc_src RELATIVE ${tensorflow_source_dir} - "${tensorflow_source_dir}/${python_proto_cc}/*.proto" - ) - list(APPEND tf_python_protos_cc_srcs ${tf_python_protos_cc_src}) + if(NOT python_proto_cc MATCHES "^\#") + STRING(REGEX REPLACE " *\#.*" "" python_proto_cc "${python_proto_cc}") + if(NOT EXISTS "${tensorflow_source_dir}/${python_proto_cc}") + message(SEND_ERROR "Python proto CC directory not found: ${python_proto_cc}") + endif() + file(GLOB_RECURSE tf_python_protos_cc_src RELATIVE ${tensorflow_source_dir} + "${tensorflow_source_dir}/${python_proto_cc}/*.proto" + ) + list(APPEND tf_python_protos_cc_srcs ${tf_python_protos_cc_src}) + endif() endforeach(python_proto_cc) RELATIVE_PROTOBUF_GENERATE_CPP(PROTO_SRCS PROTO_HDRS @@ -199,7 +211,13 @@ STRING(REGEX REPLACE ";" "\\\\;" python_modules "${python_modules}") STRING(REGEX REPLACE "\n" ";" python_modules "${python_modules}") foreach(python_module ${python_modules}) - add_python_module(${python_module}) + if(NOT python_module MATCHES "^\#") + STRING(REGEX REPLACE " *\#.*" "" python_module "${python_module}") + if(NOT EXISTS "${tensorflow_source_dir}/${python_module}") + message(SEND_ERROR "Python module not found: ${python_module}") + endif() + add_python_module(${python_module}) + endif() endforeach(python_module) add_custom_command(TARGET tf_python_touchup_modules PRE_BUILD @@ -289,7 +307,7 @@ function(GENERATE_PYTHON_OP_LIB tf_python_op_lib_name) # containing the wrappers. add_custom_command( OUTPUT ${GENERATE_PYTHON_OP_LIB_DESTINATION} - COMMAND ${tf_python_op_lib_name}_gen_python ${tensorflow_source_dir}/tensorflow/core/api_def/base_api,${tensorflow_source_dir}/tensorflow/core/api_def/python_api @${tensorflow_source_dir}/tensorflow/python/ops/hidden_ops.txt ${require_shape_fn} > ${GENERATE_PYTHON_OP_LIB_DESTINATION} + COMMAND ${tf_python_op_lib_name}_gen_python ${tensorflow_source_dir}/tensorflow/core/api_def/base_api,${tensorflow_source_dir}/tensorflow/core/api_def/python_api ${require_shape_fn} > ${GENERATE_PYTHON_OP_LIB_DESTINATION} DEPENDS ${tf_python_op_lib_name}_gen_python ) @@ -299,6 +317,7 @@ endfunction() GENERATE_PYTHON_OP_LIB("audio_ops") GENERATE_PYTHON_OP_LIB("array_ops") +GENERATE_PYTHON_OP_LIB("batch_ops") GENERATE_PYTHON_OP_LIB("bitwise_ops") GENERATE_PYTHON_OP_LIB("math_ops") GENERATE_PYTHON_OP_LIB("functional_ops") @@ -312,9 +331,11 @@ GENERATE_PYTHON_OP_LIB("dataset_ops") GENERATE_PYTHON_OP_LIB("image_ops") GENERATE_PYTHON_OP_LIB("io_ops") GENERATE_PYTHON_OP_LIB("linalg_ops") +GENERATE_PYTHON_OP_LIB("list_ops") GENERATE_PYTHON_OP_LIB("logging_ops") GENERATE_PYTHON_OP_LIB("lookup_ops") GENERATE_PYTHON_OP_LIB("nn_ops") +GENERATE_PYTHON_OP_LIB("manip_ops") GENERATE_PYTHON_OP_LIB("parsing_ops") GENERATE_PYTHON_OP_LIB("random_ops") GENERATE_PYTHON_OP_LIB("remote_fused_graph_ops" @@ -343,10 +364,12 @@ GENERATE_PYTHON_OP_LIB("contrib_boosted_trees_quantiles_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/boosted_trees/python/ops/gen_quantile_ops.py) GENERATE_PYTHON_OP_LIB("contrib_boosted_trees_stats_accumulator_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/boosted_trees/python/ops/gen_stats_accumulator_ops.py) +GENERATE_PYTHON_OP_LIB("contrib_coder_ops" + DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/coder/python/ops/gen_coder_ops.py) GENERATE_PYTHON_OP_LIB("contrib_cudnn_rnn_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/cudnn_rnn/ops/gen_cudnn_rnn_ops.py) -GENERATE_PYTHON_OP_LIB("contrib_data_prefetching_ops" - DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/data/python/ops/gen_prefetching_ops.py) +GENERATE_PYTHON_OP_LIB("contrib_data_dataset_ops" + DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/data/python/ops/gen_dataset_ops.py) GENERATE_PYTHON_OP_LIB("contrib_factorization_clustering_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/factorization/python/ops/gen_clustering_ops.py) GENERATE_PYTHON_OP_LIB("contrib_factorization_factorization_ops" @@ -518,7 +541,11 @@ if(WIN32) ${nsync_STATIC_LIBRARIES} ) - set(pywrap_tensorflow_deffile "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}/pywrap_tensorflow.def") + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(pywrap_tensorflow_deffile "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}/pywrap_tensorflow.def") + else() + set(pywrap_tensorflow_deffile "${CMAKE_CURRENT_BINARY_DIR}/pywrap_tensorflow.def") + endif() set_source_files_properties(${pywrap_tensorflow_deffile} PROPERTIES GENERATED TRUE) add_custom_command(TARGET pywrap_tensorflow_internal_static POST_BUILD @@ -526,6 +553,7 @@ if(WIN32) --input "${pywrap_tensorflow_internal_static_dependencies}" --output "${pywrap_tensorflow_deffile}" --target _pywrap_tensorflow_internal.pyd + BYPRODUCTS ${pywrap_tensorflow_deffile} # Required for Ninja ) endif(WIN32) @@ -679,11 +707,19 @@ add_custom_command(TARGET tf_python_copy_scripts_to_destination PRE_BUILD ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/testing/python/framework/) if(WIN32) - add_custom_command(TARGET tf_python_build_pip_package POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/$(Configuration)/pywrap_tensorflow_internal.dll - ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/_pywrap_tensorflow_internal.pyd - COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/$(Configuration)/pywrap_tensorflow_internal.lib - ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + add_custom_command(TARGET tf_python_build_pip_package POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/$(Configuration)/pywrap_tensorflow_internal.dll + ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/_pywrap_tensorflow_internal.pyd + COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/$(Configuration)/pywrap_tensorflow_internal.lib + ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/) + else() + add_custom_command(TARGET tf_python_build_pip_package POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/pywrap_tensorflow_internal.dll + ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/_pywrap_tensorflow_internal.pyd + COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/pywrap_tensorflow_internal.lib + ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/) + endif() else() add_custom_command(TARGET tf_python_build_pip_package POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/libpywrap_tensorflow_internal.so diff --git a/tensorflow/contrib/cmake/tf_shared_lib.cmake b/tensorflow/contrib/cmake/tf_shared_lib.cmake index 571d2b0decb5e9afcec2314f9837546f0974e90d..6d36d5fc5c2854b2d7d2542a3cb12e033e193b88 100644 --- a/tensorflow/contrib/cmake/tf_shared_lib.cmake +++ b/tensorflow/contrib/cmake/tf_shared_lib.cmake @@ -46,7 +46,11 @@ if(WIN32) $ ) - set(tensorflow_deffile "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}/tensorflow.def") + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(tensorflow_deffile "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}/tensorflow.def") + else() + set(tensorflow_deffile "${CMAKE_CURRENT_BINARY_DIR}/tensorflow.def") + endif() set_source_files_properties(${tensorflow_deffile} PROPERTIES GENERATED TRUE) add_custom_command(TARGET tensorflow_static POST_BUILD diff --git a/tensorflow/contrib/cmake/tf_tests.cmake b/tensorflow/contrib/cmake/tf_tests.cmake index 8fb19c055e1172a47f2947575707515496cc24ad..1c4ebd7f0c1113bcd0857fb0858df2248499f920 100644 --- a/tensorflow/contrib/cmake/tf_tests.cmake +++ b/tensorflow/contrib/cmake/tf_tests.cmake @@ -153,8 +153,10 @@ if (tensorflow_BUILD_PYTHON_TESTS) "${tensorflow_source_dir}/tensorflow/python/profiler/internal/*_test.py" "${tensorflow_source_dir}/tensorflow/python/saved_model/*_test.py" "${tensorflow_source_dir}/tensorflow/python/training/*_test.py" + "${tensorflow_source_dir}/tensorflow/contrib/coder/*_test.py" "${tensorflow_source_dir}/tensorflow/contrib/data/*_test.py" "${tensorflow_source_dir}/tensorflow/contrib/factorization/*_test.py" + "${tensorflow_source_dir}/tensorflow/contrib/feature_column/python/feature_column/*_test.py" "${tensorflow_source_dir}/tensorflow/contrib/image/*_test.py" "${tensorflow_source_dir}/tensorflow/python/keras/_impl/keras/*_test.py" "${tensorflow_source_dir}/tensorflow/contrib/periodic_resample/python/kernel_tests/*_test.py" @@ -190,6 +192,7 @@ if (tensorflow_BUILD_PYTHON_TESTS) "${tensorflow_source_dir}/tensorflow/python/profiler/pprof_profiler_test.py" # flaky test "${tensorflow_source_dir}/tensorflow/python/profiler/internal/run_metadata_test.py" + "${tensorflow_source_dir}/tensorflow/python/profiler/model_analyzer_test.py" # Fails because uses data dependencies with bazel "${tensorflow_source_dir}/tensorflow/python/saved_model/saved_model_test.py" # requires scipy @@ -225,6 +228,7 @@ if (tensorflow_BUILD_PYTHON_TESTS) "${tensorflow_source_dir}/tensorflow/python/kernel_tests/as_string_op_test.py" "${tensorflow_source_dir}/tensorflow/python/kernel_tests/string_to_number_op_test.py" "${tensorflow_source_dir}/tensorflow/python/kernel_tests/clip_ops_test.py" + "${tensorflow_source_dir}/tensorflow/python/kernel_tests/list_ops_test.py" # Needs portpicker. "${tensorflow_source_dir}/tensorflow/python/kernel_tests/tensor_array_ops_test.py" # Needs portpicker. # Numerical issues, calculations off. "${tensorflow_source_dir}/tensorflow/python/kernel_tests/concat_op_test.py" @@ -272,8 +276,7 @@ if (tensorflow_BUILD_PYTHON_TESTS) "${tensorflow_source_dir}/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py" # Segfaults on windows "${tensorflow_source_dir}/tensorflow/contrib/data/python/kernel_tests/dataset_constructor_op_test.py" # Segfaults on Windows. "${tensorflow_source_dir}/tensorflow/python/data/kernel_tests/iterator_ops_cluster_test.py" - # Broken tensorboard test due to cmake issues. - "${tensorflow_source_dir}/tensorflow/contrib/data/python/kernel_tests/iterator_ops_cluster_test.py" # Needs portpicker + "${tensorflow_source_dir}/tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py" # Deadlocks "${tensorflow_source_dir}/tensorflow/contrib/data/python/kernel_tests/sloppy_transformation_dataset_op_test.py" # b/65430561 # tensor_forest tests (also note that we exclude the hybrid tests for now) "${tensorflow_source_dir}/tensorflow/contrib/tensor_forest/python/kernel_tests/count_extremely_random_stats_op_test.py" # Results in wrong order. @@ -305,6 +308,10 @@ if (tensorflow_BUILD_PYTHON_TESTS) # Depends on python/framework/test_ops "${tensorflow_source_dir}/tensorflow/python/kernel_tests/array_ops_test.py" "${tensorflow_source_dir}/tensorflow/python/kernel_tests/control_flow_util_test.py" + # Flaky replicate_model_fn_test + "${tensorflow_source_dir}/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py" # b/71901810 + # Broken io_utils_test + "${tensorflow_source_dir}/tensorflow/python/keras/_impl/keras/utils/io_utils_test.py" # b/72894325 ) endif() list(REMOVE_ITEM tf_test_src_py ${tf_test_src_py_exclude}) diff --git a/tensorflow/contrib/cmake/tf_tools.cmake b/tensorflow/contrib/cmake/tf_tools.cmake index cb58a2e7df85b2f214654eff5547c5788592f208..58c7df95c821b4d1aa2cc63c8aaf4039518b83ca 100644 --- a/tensorflow/contrib/cmake/tf_tools.cmake +++ b/tensorflow/contrib/cmake/tf_tools.cmake @@ -48,9 +48,6 @@ file(GLOB_RECURSE tf_tools_transform_graph_lib_exclude_srcs "${tensorflow_source_dir}/tensorflow/tools/graph_transforms/compare_graphs.cc" "${tensorflow_source_dir}/tensorflow/tools/graph_transforms/summarize_graph_main.cc" "${tensorflow_source_dir}/tensorflow/tools/graph_transforms/transform_graph_main.cc" - "${tensorflow_source_dir}/tensorflow/tools/graph_transforms/quantize_nodes.cc" - "${tensorflow_source_dir}/tensorflow/tools/graph_transforms/quantize_weights.cc" - "${tensorflow_source_dir}/tensorflow/tools/graph_transforms/round_weights.cc" ) list(REMOVE_ITEM tf_tools_transform_graph_lib_srcs ${tf_tools_transform_graph_lib_exclude_srcs}) diff --git a/tensorflow/contrib/cmake/tools/create_def_file.py b/tensorflow/contrib/cmake/tools/create_def_file.py index f67698eb99a38eae307b52e55de748a67b798cbd..53c2285699a6ca94e1e6b147080338b507f4d768 100644 --- a/tensorflow/contrib/cmake/tools/create_def_file.py +++ b/tensorflow/contrib/cmake/tools/create_def_file.py @@ -31,7 +31,7 @@ from __future__ import division from __future__ import print_function import argparse -import io +import codecs import os import re import subprocess @@ -103,7 +103,7 @@ def main(): for lib_path in args.input: proc = subprocess.Popen([DUMPBIN, "/nologo", "/linkermember:1", lib_path], stdout=subprocess.PIPE) - for line in io.TextIOWrapper(proc.stdout, encoding="utf-8"): + for line in codecs.getreader("utf-8")(proc.stdout): cols = line.split() if len(cols) < 2: continue @@ -131,7 +131,7 @@ def main(): # We compare on undname but use the decorated name from candidates. dupes = 0 proc = subprocess.Popen([UNDNAME, tmpfile.name], stdout=subprocess.PIPE) - for idx, line in enumerate(io.TextIOWrapper(proc.stdout, encoding="utf-8")): + for idx, line in enumerate(codecs.getreader("utf-8")(proc.stdout)): decorated = candidates[idx] if decorated in taken: # Symbol is already in output, done. diff --git a/tensorflow/contrib/coder/BUILD b/tensorflow/contrib/coder/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..ec3d550b70d2aaa23b989c44f3d86fa87cffb335 --- /dev/null +++ b/tensorflow/contrib/coder/BUILD @@ -0,0 +1,167 @@ +# Description: +# Contains entropy coding related modules. + +package(default_visibility = [ + "//learning/brain:__subpackages__", + "//tensorflow:__subpackages__", +]) + +licenses(["notice"]) # Apache 2.0 + +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_test", + "tf_custom_op_library", + "tf_custom_op_py_library", + "tf_gen_op_libs", + "tf_gen_op_wrapper_py", + "tf_kernel_library", + "tf_py_test", +) + +cc_library( + name = "range_coder", + srcs = [ + "kernels/range_coder.cc", + ], + hdrs = [ + "kernels/range_coder.h", + ], + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/core:framework", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "range_coder_test", + size = "small", + srcs = ["kernels/range_coder_test.cc"], + deps = [ + ":range_coder", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +tf_gen_op_libs( + op_lib_names = ["coder_ops"], + deps = [ + "//tensorflow/core:lib", + ], +) + +tf_kernel_library( + name = "range_coder_ops", + srcs = [ + "kernels/range_coder_ops.cc", + "kernels/range_coder_ops_util.cc", + ], + hdrs = [ + "kernels/range_coder_ops_util.h", + ], + visibility = ["//visibility:public"], + deps = [ + ":coder_ops_op_lib", + ":range_coder", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + ], + alwayslink = 1, +) + +tf_cc_test( + name = "range_coder_ops_test", + size = "small", + srcs = ["kernels/range_coder_ops_test.cc"], + deps = [ + ":range_coder", + ":range_coder_ops", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + "//tensorflow/core/kernels:ops_testutil", + ], +) + +cc_library( + name = "all_ops", + deps = [":coder_ops_op_lib"], +) + +cc_library( + name = "all_kernels", + deps = [":range_coder_ops"], +) + +tf_custom_op_library( + name = "python/ops/_coder_ops.so", + srcs = [ + "kernels/range_coder.cc", + "kernels/range_coder.h", + "kernels/range_coder_ops.cc", + "kernels/range_coder_ops_util.cc", + "kernels/range_coder_ops_util.h", + "ops/coder_ops.cc", + ], +) + +tf_gen_op_wrapper_py( + name = "gen_coder_ops", + out = "python/ops/gen_coder_ops.py", + deps = [":coder_ops_op_lib"], +) + +tf_custom_op_py_library( + name = "coder_ops_py", + srcs = [ + "__init__.py", + "python/ops/coder_ops.py", + ], + dso = [ + ":python/ops/_coder_ops.so", + ], + kernels = [ + ":all_kernels", + ], + srcs_version = "PY2AND3", + deps = [ + ":gen_coder_ops", + "//tensorflow/contrib/util:util_py", + ], +) + +tf_py_test( + name = "coder_ops_py_test", + srcs = [ + "python/ops/coder_ops_test.py", + ], + additional_deps = [ + ":coder_ops_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dtypes", + "//tensorflow/python:math_ops", + "//tensorflow/python:random_ops", + ], + main = "python/ops/coder_ops_test.py", +) + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), +) diff --git a/tensorflow/contrib/coder/README.md b/tensorflow/contrib/coder/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c6c379c458893551b765327c0c1cbfff7f24f9c3 --- /dev/null +++ b/tensorflow/contrib/coder/README.md @@ -0,0 +1,73 @@ +# Entropy coder + +This module contains range encoder and range decoder which can encode integer +data into string with cumulative distribution functions (CDF). + +## Data and CDF values + +The data to be encoded should be non-negative integers in half-open interval +`[0, m)`. Then a CDF is represented as an integral vector of length `m + 1` +where `CDF(i) = f(Pr(X < i) * 2^precision)` for i = 0,1,...,m, and `precision` +is an attribute in range `0 < precision <= 16`. The function `f` maps real +values into integers, e.g., round or floor. It is important that to encode a +number `i`, `CDF(i + 1) - CDF(i)` cannot be zero. + +Note that we used `Pr(X < i)` not `Pr(X <= i)`, and therefore CDF(0) = 0 always. + +## RangeEncode: data shapes and CDF shapes + +For each data element, its CDF has to be provided. Therefore if the shape of CDF +should be `data.shape + (m + 1,)` in NumPy-like notation. For example, if `data` +is a 2-D tensor of shape (10, 10) and its elements are in `[0, 64)`, then the +CDF tensor should have shape (10, 10, 65). + +This may make CDF tensor too large, and in many applications all data elements +may have the same probability distribution. To handle this, `RangeEncode` +supports limited broadcasting CDF into data. Broadcasting is limited in the +following sense: + +- All CDF axes but the last one is broadcasted into data but not the other way + around, +- The number of CDF axes does not extend, i.e., `CDF.ndim == data.ndim + 1`. + +In the previous example where data has shape (10, 10), the following are +acceptable CDF shapes: + +- (10, 10, 65) +- (1, 10, 65) +- (10, 1, 65) +- (1, 1, 65) + +## RangeDecode + +`RangeEncode` encodes neither data shape nor termination character. Therefore +the decoder should know how many characters are encoded into the string, and +`RangeDecode` takes the encoded data shape as the second argument. The same +shape restrictions as `RangeEncode` inputs apply here. + +## Example + +```python +data = tf.random_uniform((128, 128), 0, 10, dtype=tf.int32) + +histogram = tf.bincount(data, minlength=10, maxlength=10) +cdf = tf.cumsum(histogram, exclusive=False) +# CDF should have length m + 1. +cdf = tf.pad(cdf, [[1, 0]]) +# CDF axis count must be one more than data. +cdf = tf.reshape(cdf, [1, 1, -1]) + +# Note that data has 2^14 elements, and therefore the sum of CDF is 2^14. +data = tf.cast(data, tf.int16) +encoded = coder.range_encode(data, cdf, precision=14) +decoded = coder.range_decode(encoded, tf.shape(data), cdf, precision=14) + +# data and decoded should be the same. +sess = tf.Session() +x, y = sess.run((data, decoded)) +assert np.all(x == y) +``` + +## Authors +Sung Jin Hwang (github: [ssjhv](https://github.com/ssjhv)) and Nick Johnston +(github: [nmjohn](https://github.com/nmjohn)) diff --git a/tensorflow/contrib/coder/__init__.py b/tensorflow/contrib/coder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7e663e6f1359f399cdaa80e037635a8f7546b37 --- /dev/null +++ b/tensorflow/contrib/coder/__init__.py @@ -0,0 +1,26 @@ +# 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. +# ============================================================================== +"""Entropy code operations.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# pylint: disable=wildcard-import +from tensorflow.contrib.coder.python.ops.coder_ops import * +# pylint: enable=wildcard-import + +from tensorflow.python.util.all_util import remove_undocumented +remove_undocumented(__name__) diff --git a/tensorflow/contrib/coder/kernels/range_coder.cc b/tensorflow/contrib/coder/kernels/range_coder.cc new file mode 100644 index 0000000000000000000000000000000000000000..21b35155ff317c6afbb1b86745f05385726505b6 --- /dev/null +++ b/tensorflow/contrib/coder/kernels/range_coder.cc @@ -0,0 +1,374 @@ +/* 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. +==============================================================================*/ + +// Range coder implementation, based on [1]. +// +// [1] G. N. N. Martin, "Range coding: an algorithm for removing redundancy from +// a digitised message", presented to the Video & Data Recording Conference, +// held in Southampton, July 24-27, 1979. +// +#include "tensorflow/contrib/coder/kernels/range_coder.h" + +#include +#include + +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +RangeEncoder::RangeEncoder(int precision) : precision_(precision) { + CHECK_GT(precision, 0); + CHECK_LE(precision, 16); +} + +void RangeEncoder::Encode(int32 lower, int32 upper, string* sink) { + // Input requirement: 0 <= lower < upper <= 2^precision. + DCHECK_LE(0, lower); + DCHECK_LT(lower, upper); + DCHECK_LE(upper, 1 << precision_); + + // `base` and `size` represent a half-open interval [base, base + size). + // Loop invariant: 2^16 <= size <= 2^32. + // + // Note that keeping size above 2^16 is important. Since the interval sizes + // are quantized to up to 16 bits, the smallest interval size the encode may + // handle is 2^-16. If size is smaller than 2^16, a small interval input may + // collapse the encoder range into an empty interval. + const uint64 size = static_cast(size_minus1_) + 1; + DCHECK_NE(size >> 16, 0); + + // For short notation, let u := lower and v := upper. + // + // The input u, v represents a half-open interval [u, v) / 2^precision. + // This narrows the current interval roughly to + // [base + (size * u) / 2^precision, base + (size * v) / 2^precision). + // + // TODO(sjhwang): Try rounding if it helps improve compression ratio, at the + // expense of more operations. In the test using Zipf distribution, the + // overhead over the theoretical compression ratio was ~0.01%. + // NOTE: The max value of `size` is 2^32 and size > 0. Therefore `size * u` + // can be rewritten as `(size - 1) * u + u` and all the computation can be + // done in 32-bit mode. If 32-bit multiply is faster, then rewrite. + const uint32 a = (size * static_cast(lower)) >> precision_; + const uint32 b = ((size * static_cast(upper)) >> precision_) - 1; + DCHECK_LE(a, b); + + // Let's confirm the RHS of a, b fit in uint32 type. + // Recall that 0 <= u < 2^precision, and size <= 2^32. Therefore + // (size * u) / 2^precision < size <= 2^32, + // and the value of a fits in uint32 type. Similarly, since v <= 2^precision, + // (size * v) / 2^precision - 1 <= size - 1 < 2^32. + // For lower bound of b, note that 1 <= v, 2^16 <= size, and 16 <= precision. + // Therefore (size * v) / 2^precision - 1 >= 2^16 / 2^precision - 1 >= 0. + + // The new interval is [base + a, base + b] = [base + a, base + b + 1). + base_ += a; // May overflow. + size_minus1_ = b - a; + const bool base_overflow = (base_ < a); + + // The encoder has two states. Let's call them state 0 and state 1. + // State 0 is when base < base + size <= 2^32. + // State 1 is when base < 2^32 < base + size. + // + // The encoder initially starts in state 0, with base = 0, size = 2^32. + // + // TODO(sjhwang): Requires some profiling, but the encoder stays in state 0 + // most of the time. Should optimize code for state 0. + // + // Each Encode() has up to two places where the interval changes: + // #1. Refine the interval. [base, base + size) -> [base + a, base + b + 1). + // #2. Expand interval if the new size is too small, + // and each change may cause a state transition. + // + // First, consider when the current state is 0. + // + // In this case, the next state after #1 is always state 0, since refining + // interval only shrinks the interval, therefore new_base + new_size <= 2^32. + // + // Let us explain #2. + // + // Recall that at the beginning of each Encode(), the encoder requires + // 2^16 < size <= 2^32. As precision <= 16, the new interval size can be as + // small as 1, but never zero. + // + // To keep size above 2^16, if new size is smaller than or equal to 2^16, the + // encoder would left-shift base and size by 16 bits: size' <- size * 2^16. + // Note that new size' is now in the range [2^16, 2^32]. + // + // Since size is left-shifted, the same should be applied to base as well. + // However, after the left-shift, base will then contain 48 bits instead of 32 + // bits. Therefore prior to the shift, The upper 16 bits in base should be + // stored somewhere else. + // + // If the upper 16 bits of all values in the interval were the same, i.e., if + // base[32:16] == (base + size - 1)[32:16], then base[32:16] can be written + // out to `output` string, since any further Encode() only narrows down the + // interval and that 16 bits would never change. + // + // If the upper 16 bits were not all the same, since this happens only when + // size <= 2^16, the upper 16 bits may differ only by one, i.e., + // base[32:16] + 1 == (base + size - 1)[32:16]. At this stage, it is not + // determined yet whether base[32:16] should be written to the output or + // (base[32:16] + 1) should be written to the output. In this case, + // (base[32:16] + 1) is temporarily stored in `delay`, and base is + // left-shifted by 16 bits. + // + // In the latter case, the condition implies that (base // 2^16) and + // ((base + size - 1) // 2^16) were different. Therefore after left-shift by + // 16 bits, the new (base + size) is greater than 2^32, i.e., the encoder + // transition to state 1. + // + // ==== Summary ==== + // To detect the current encoder state, + // state 0: delay == 0 iff (base mod 2^32) < (base + size) mod 2^32, + // state 1: delay != 0 iff (base + size) mod 2^32 <= base mod 2^32, + // because size <= 2^32. + // + // ==== Summary for state 0 ==== + // 1. Interval refinement does not cause state transition. + // 2. Interval expansion may cause state transition, depending on the upper 16 + // bits of base and base + size - 1. + // + // Now suppose the previous state was 1. This means that + // base <= 2^32 < base + size. + // + // When in state 1, an interval refinement may trigger state transition. + // After Encode() refines the interval, there are three possibilities: + // #1. base <= 2^32 < base + size (unchanged), + // #2. 2^32 <= base < base + size (base overflowed), + // #3. base < base + size <= 2^32 (base + size - 1 underflowed). + // + // In case #1, the encoder remains in state 1. + // In case #2 or #3, the encoder state changes to state 0. + // + // ==== State transition for interval refinement ==== + // 1. state 0 -> state 0, + // 2. state 1 -> state 0 or state 1. + // + // Therefore if the new state is 1, then the previous state must have been + // state 1. + if (base_ + size_minus1_ < base_) { + // If statement checked if 2^32 < base + size. The new state is 1, hence the + // previous state was also state 1. + DCHECK_NE(((base_ - a) + size) >> 32, 0); + DCHECK_NE(delay_ & 0xFFFF, 0); + + // Like in state 0, if the new size is <= 2^16, then base and size should + // be left-shifted by 16 bits. Combine the conditions + // base <= 2^32 < base + size and size <= 2^16 to conclude that + // base[32:16] >= 0xFFFF and (base + size - 1)[32:16] = 0x0000. + // + // Note that 2^32 - base < size, and since base is at least 0xFFFF0000, + // 2^16 - base[16:0] < size. Let base' and size' be the new base and size + // after the bit-shift. Then 2^32 - base' < size' => 2^32 < base' + size'. + // Therefore the encoder remains in state 1. + // + // Lastly, `delay` is modified. Conceptually, delay has to be changed to + // delay' <- delay * 2^16 + (base + size - 1)[32:16]. + // Since we know above that (base + size - 1)[32:16] = 0x0000, there is no + // need to explicitly do the computation above, but rather store how many + // trailing zeros there were. For this reason, the lower 16 bits of + // `delay` stores the delayed value when state changed from 0 to 1, and + // delay[32:16] stores the # of trailing zeros (in bytes). + // + // ==== State transition for interval expansion ==== + // 1. state 0 -> state 0 or state 1, + // 2. state 1 -> state 1. + if (size_minus1_ >> 16 == 0) { + DCHECK_EQ(base_ >> 16, 0xFFFF); + base_ <<= 16; + size_minus1_ <<= 16; + size_minus1_ |= 0xFFFF; + // TODO(sjhwang): It is possible that for very long input, delay + // overflow during below. If overflow is detected, this delay is too + // long the encoder should forcefully move to state 0. In such case, + // base can be raised to 2^32 (force case #2), or (base + size) can be + // lowered to 2^32 (force case #3), depending on which transition + // keeps size larger. + CHECK_LT(delay_, static_cast(1) << 62); + delay_ += 0x20000; // Two more bytes of zeros. Check overflow? + } + return; + } + + // If reached here, the current state is 0. + // First handle the case when the previous state was state 1. + if (delay_ != 0) { + // In case #2 or #3, the encoder state changes to state 0. Recall that when + // the encoder state changed from state 0 to state 1, the top 16 bits of + // (base + size - 1) was temporarily stored in `delay`, because the output + // could be either (delay - 1) or (delay). + // + // And from above, the delayed value encoded in `delay` is + // delay' <- delay[16:0] * 2^(8 * delay[MAX:16]) + // + // In case #2, the interval moved below 2^32. So (delay' - 1) is the + // converged value after interval refinements. Write out + // (delay[16:0] - 1) and write (8 * delay[MAX:16]) bytes of 0xFF. + // + // In case #3, the interval moved above 2^32. So delay' is the converged + // value after interval refinement. Write out delay[16:0] and write + // (8 * delay[MAX:16]) bytes of 0x00. + if (base_overflow) { + // Case #2. + DCHECK_NE((static_cast(base_ - a) + a) >> 32, 0); + sink->push_back(static_cast(delay_ >> 8)); + sink->push_back(static_cast(delay_ >> 0)); + sink->append(delay_ >> 16, static_cast(0)); + } else { + // Case #3. + DCHECK_EQ(static_cast(base_ + size_minus1_) >> 32, 0); + --delay_; + sink->push_back(static_cast(delay_ >> 8)); + sink->push_back(static_cast(delay_ >> 0)); + sink->append(delay_ >> 16, static_cast(0xFF)); + } + // Reset to state 0. + delay_ = 0; + } + + if (size_minus1_ >> 16 == 0) { + const uint32 top = base_ >> 16; + + base_ <<= 16; + size_minus1_ <<= 16; + size_minus1_ |= 0xFFFF; + + if (base_ <= base_ + size_minus1_) { + // Still in state 0. Write the top 16 bits. + sink->push_back(static_cast(top >> 8)); + sink->push_back(static_cast(top)); + } else { + // New state is 1. + DCHECK_LT(top, 0xFFFF); + delay_ = top + 1; + } + } +} + +void RangeEncoder::Finalize(string* sink) { + // Finalize the encode by writing out any number in the interval + // [base, base + size). + // + // Trailing zeros are not explicitly written out as decoder can fill in zeros + // by default. + if (delay_ != 0) { + // The last state was state 1. Since base < 2^32 < base + size, pick 2^32 + // (state 1, case #3). + // NOTE: It is a bit difficult to trigger this code path on purpose. + // TODO(sjhwang): Find a way to trigger this code path for test coverage. + sink->push_back(static_cast(delay_ >> 8)); + if ((delay_ & 0xFF) != 0) { + sink->push_back(static_cast(delay_)); + } + } else if (base_ != 0) { + // If base == 0, then pick 0 from [base, base + size) and no zeros are + // explicitly written. + // + // Otherwise, pick (base + (2^16 - base[16:0])), i.e., round up base to the + // next multiple of 2^16. As 2^16 < size, this value should be in the + // interval [base, base + size). + const uint32 mid = ((base_ - 1) >> 16) + 1; + DCHECK_EQ(mid & 0xFFFF, mid); + sink->push_back(static_cast(mid >> 8)); + if ((mid & 0xFF) != 0) { + sink->push_back(static_cast(mid >> 0)); + } + } + + base_ = 0; + size_minus1_ = std::numeric_limits::max(); + delay_ = 0; +} + +RangeDecoder::RangeDecoder(const string& source, int precision) + : current_(source.begin()), + begin_(source.begin()), + end_(source.end()), + precision_(precision) { + CHECK_LE(precision, 16); + + Read16BitValue(); + Read16BitValue(); +} + +int32 RangeDecoder::Decode(tensorflow::gtl::ArraySlice cdf) { + const uint64 size = static_cast(size_minus1_) + 1; + const uint64 offset = + ((static_cast(value_ - base_) + 1) << precision_) - 1; + + // This is similar to std::lower_range() with std::less_equal as comparison. + // After the binary search, `pv` points to the smallest number v that + // satisfies offset < (size * v) / 2^precision. + + // Assumes that cdf[0] == 0. Therefore (size * cdf[0]) / 2^precision is always + // less than or equal to offset. + const int32* pv = cdf.data() + 1; + // `len` can be cdf.size() - 2 if there is guarantee that the last element of + // cdf is 2^precision. + auto len = cdf.size() - 1; + DCHECK_GT(len, 0); + + do { + const auto half = len / 2; + const int32* mid = pv + half; + DCHECK_GE(*mid, 0); + DCHECK_LE(*mid, 1 << precision_); + if (size * static_cast(*mid) <= offset) { + pv = mid + 1; + len -= half + 1; + } else { + len = half; + } + } while (len > 0); + + // If (size * v) / 2^precision <= offset for all v in cdf, then pv points to + // one after the last element of cdf. That is a decoding error. + // + // TODO(sjhwang): Consider returning -1 to indicate error. Or start len = + // cdf.size() - 2 instead and give up detecting this error. + CHECK_LT(pv, cdf.data() + cdf.size()); + + const uint32 a = (size * static_cast(*(pv - 1))) >> precision_; + const uint32 b = ((size * static_cast(*pv)) >> precision_) - 1; + DCHECK_LE(a, offset >> precision_); + DCHECK_LE(offset >> precision_, b); + + base_ += a; + size_minus1_ = b - a; + + if (size_minus1_ >> 16 == 0) { + base_ <<= 16; + size_minus1_ <<= 16; + size_minus1_ |= 0xFFFF; + + Read16BitValue(); + } + + return pv - cdf.data() - 1; +} + +void RangeDecoder::Read16BitValue() { + value_ <<= 8; + if (current_ != end_) { + value_ |= static_cast(*current_++); + } + value_ <<= 8; + if (current_ != end_) { + value_ |= static_cast(*current_++); + } +} +} // namespace tensorflow diff --git a/tensorflow/contrib/coder/kernels/range_coder.h b/tensorflow/contrib/coder/kernels/range_coder.h new file mode 100644 index 0000000000000000000000000000000000000000..f46413072e34a55128d7854b9c312dfdde457d85 --- /dev/null +++ b/tensorflow/contrib/coder/kernels/range_coder.h @@ -0,0 +1,109 @@ +/* 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_CODER_KERNELS_RANGE_CODER_H_ +#define TENSORFLOW_CONTRIB_CODER_KERNELS_RANGE_CODER_H_ + +#include +#include + +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +class RangeEncoder { + public: + // `precision` determines the granularity of probability masses passed to + // Encode() function below. + // + // REQUIRES: 0 < precision <= 16. + explicit RangeEncoder(int precision); + + // Encodes a half-open interval [lower / 2^precision, upper / 2^precision). + // Suppose each character to be encoded is from an integer-valued + // distribution. When encoding a random character x0, the arguments lower and + // upper represent + // Pr(X < x0) = lower / 2^precision, + // Pr(X < x0 + 1) = upper / 2^precision, + // where X is a random variable following the distribution. + // + // For example, assume that the distribution has possible outputs 0, 1, 2, ... + // To encode value 0, lower = 0 and upper = Pr(X = 0). + // To encode value 1, lower = Pr(X = 0) and upper = Pr(X = 0 or 1). + // To encode value 2, lower = Pr(X = 0 or 1) and upper = Pr(X = 0, 1, or 2). + // ... + // + // REQUIRES: 0 <= lower < upper <= 2^precision. + void Encode(int32 lower, int32 upper, string* sink); + + // The encode may contain some under-determined values from previous encoding. + // After Encode() calls, Finalize() must be called. Otherwise the encoded + // string may not be decoded. + void Finalize(string* sink); + + private: + uint32 base_ = 0; + uint32 size_minus1_ = std::numeric_limits::max(); + uint64 delay_ = 0; + + const int precision_; +}; + +class RangeDecoder { + public: + // Holds a reference to `source`. The caller has to make sure that `source` + // outlives the decoder object. + // + // REQUIRES: `precision` must be the same as the encoder's precision. + // REQUIRES: 0 < precision <= 16. + RangeDecoder(const string& source, int precision); + + // Decodes a character from `source` using CDF. The size of `cdf` should be + // one more than the number of the character in the alphabet. + // + // If x0, x1, x2, ... are the possible characters (in increasing order) from + // the distribution, then + // cdf[0] = 0 + // cdf[1] = Pr(X <= x0), + // cdf[2] = Pr(X <= x1), + // cdf[3] = Pr(X <= x2), + // ... + // + // The returned value is an index to `cdf` where the decoded character + // corresponds to. + // + // REQUIRES: cdf.size() > 1. + // REQUIRES: cdf[i] <= cdf[i + 1] for i = 0, 1, ..., cdf.size() - 2. + // REQUIRES: cdf[cdf.size() - 1] <= 2^precision. + // + // In practice the last element of `cdf` should equal to 2^precision. + int32 Decode(gtl::ArraySlice cdf); + + private: + void Read16BitValue(); + + uint32 base_ = 0; + uint32 size_minus1_ = std::numeric_limits::max(); + uint32 value_ = 0; + + string::const_iterator current_; + const string::const_iterator begin_; + const string::const_iterator end_; + + const int precision_; +}; +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_CODER_KERNELS_RANGE_CODER_H_ diff --git a/tensorflow/contrib/coder/kernels/range_coder_ops.cc b/tensorflow/contrib/coder/kernels/range_coder_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..cde7982530fea6407aaf074f7af4a22263d50da3 --- /dev/null +++ b/tensorflow/contrib/coder/kernels/range_coder_ops.cc @@ -0,0 +1,307 @@ +/* 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. +==============================================================================*/ + +#define EIGEN_USE_THREADS + +#include +#include +#include +#include +#include + +#include "tensorflow/contrib/coder/kernels/range_coder.h" +#include "tensorflow/contrib/coder/kernels/range_coder_ops_util.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_types.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +namespace { +// A helper class to iterate over data and cdf simultaneously, while cdf is +// broadcasted to data. +// NOTE: Moving this class out of anonymous namespace impacts compiler +// optimization and affects performance. When moving this code around (e.g., +// into a library header), be sure to check the benchmark tests. +template +class BroadcastRange { + public: + BroadcastRange(T* data_pointer, gtl::ArraySlice data_shape, + const U* cdf_pointer, gtl::ArraySlice cdf_shape) + : data_pointer_(data_pointer), cdf_pointer_(cdf_pointer) { + CHECK(!data_shape.empty()); + CHECK_EQ(data_shape.size(), N); + CHECK_EQ(cdf_shape.size(), N + 1); + + std::copy(data_shape.begin(), data_shape.end(), &data_shape_[0]); + data_index_.fill(0); + + const int64 innermost_stride = cdf_shape[N]; + cdf_displace_.fill(innermost_stride); + + // Pre-compute the pointer displacement for cdf. + int64 stride = innermost_stride; + for (int i = N - 1; i >= 0; --i) { + const bool broadcasting = (cdf_shape[i] <= 1); + + // When the data linear index advances by one, the cdf linear index + // advances by `innermost_stride`. + // + // Suppose that the i-th axis coordinate of data increased by one, and + // that i-th axis is broadcasting. The cdf linear index should be wound + // back by i-th axis stride, so that i-th axis coordinate of cdf is + // effectively kept at 0. + if (broadcasting) { + cdf_displace_[i] -= stride; + } + stride *= cdf_shape[i]; + } + } + + // Returns the pointers to the current iterating locations to data and cdf + // tensors. + // + // Note that this function does not track whether data pointer is running past + // the end of data buffer. The caller has to make sure Next() is called no + // more than that. + std::pair Next() { + std::pair return_value = {data_pointer_, cdf_pointer_}; + + int i = N - 1; + for (; i > 0; --i) { + ++data_index_[i]; + if (data_index_[i] < data_shape_[i]) { + break; + } + data_index_[i] = 0; + } + + // Advance data pointer by one. + data_pointer_ += 1; + + // For cdf pointer, it's more complicated because of broadcasting. When i-th + // coordinate increase by one, and if i-th axis is broadcasting, then we + // need to rewind back the pointer so that the effective i-th axis + // coordinate for cdf is always 0. This value is precomputed as + // cdf_displace_. + cdf_pointer_ += cdf_displace_[i]; + return return_value; + } + + private: + std::array data_shape_; + std::array cdf_displace_; + std::array data_index_; + + T* data_pointer_; + const U* cdf_pointer_; +}; + +Status CheckCdfShape(const TensorShape& data_shape, + const TensorShape& cdf_shape) { + if (TF_PREDICT_FALSE(cdf_shape.dims() != data_shape.dims() + 1)) { + return errors::InvalidArgument( + "`cdf` should have one more axis than `data`: data shape=", + data_shape.DebugString(), ", cdf shape=", cdf_shape.DebugString()); + } + + if (TF_PREDICT_FALSE(cdf_shape.dim_size(cdf_shape.dims() - 1) <= 1)) { + return errors::InvalidArgument( + "The last dimension of `cdf` should be > 1: ", cdf_shape.DebugString()); + } + + return Status::OK(); +} + +// Non-incremental encoder op ------------------------------------------------- +class RangeEncodeOp : public OpKernel { + public: + explicit RangeEncodeOp(OpKernelConstruction* context) : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("precision", &precision_)); + OP_REQUIRES(context, 0 < precision_ && precision_ <= 16, + errors::InvalidArgument("`precision` must be in [1, 16]: ", + precision_)); + } + + void Compute(OpKernelContext* context) override { + const Tensor& data = context->input(0); + const Tensor& cdf = context->input(1); + + OP_REQUIRES_OK(context, CheckCdfShape(data.shape(), cdf.shape())); + + std::vector data_shape, cdf_shape; + OP_REQUIRES_OK( + context, MergeAxes(data.shape(), cdf.shape(), &data_shape, &cdf_shape)); + + Tensor* output_tensor; + OP_REQUIRES_OK(context, + context->allocate_output(0, TensorShape{}, &output_tensor)); + string* output = &output_tensor->scalar()(); + + switch (data_shape.size()) { +#define RANGE_ENCODE_CASE(dims) \ + case dims: { \ + RangeEncodeImpl(data.flat(), data_shape, \ + cdf.flat_inner_dims(), cdf_shape, output); \ + } break + RANGE_ENCODE_CASE(1); + RANGE_ENCODE_CASE(2); + RANGE_ENCODE_CASE(3); + RANGE_ENCODE_CASE(4); + RANGE_ENCODE_CASE(5); + RANGE_ENCODE_CASE(6); +#undef RANGE_ENCODE_CASE + default: + context->CtxFailure(errors::InvalidArgument( + "Irregular broadcast pattern: ", data.shape().DebugString(), ", ", + cdf.shape().DebugString())); + return; + } + } + + private: + template + void RangeEncodeImpl(TTypes::ConstFlat data, + gtl::ArraySlice data_shape, + TTypes::ConstMatrix cdf, + gtl::ArraySlice cdf_shape, string* output) const { + const int64 data_size = data.size(); + const int64 cdf_size = cdf.size(); + const int64 chip_size = cdf.dimension(1); + + BroadcastRange view{data.data(), data_shape, + cdf.data(), cdf_shape}; + RangeEncoder encoder{precision_}; + for (int64 linear = 0; linear < data_size; ++linear) { + const auto pair = view.Next(); + + const int64 index = *pair.first; + DCHECK_GE(index, 0); + DCHECK_LT(index + 1, chip_size); + + const int32* cdf_slice = pair.second; + DCHECK_LE(cdf_slice + chip_size, cdf.data() + cdf_size); + + const int32 lower = cdf_slice[index]; + const int32 upper = cdf_slice[index + 1]; + encoder.Encode(lower, upper, output); + } + + encoder.Finalize(output); + } + + int precision_; +}; + +REGISTER_KERNEL_BUILDER(Name("RangeEncode").Device(DEVICE_CPU), RangeEncodeOp); + +// Non-incremental decoder op ------------------------------------------------- +class RangeDecodeOp : public OpKernel { + public: + explicit RangeDecodeOp(OpKernelConstruction* context) : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("precision", &precision_)); + OP_REQUIRES(context, 0 < precision_ && precision_ <= 16, + errors::InvalidArgument("`precision` must be in [1, 16]: ", + precision_)); + } + + void Compute(OpKernelContext* context) override { + const Tensor& encoded_tensor = context->input(0); + const Tensor& shape = context->input(1); + const Tensor& cdf = context->input(2); + + OP_REQUIRES(context, TensorShapeUtils::IsScalar(encoded_tensor.shape()), + errors::InvalidArgument("Invalid `encoded` shape: ", + encoded_tensor.shape().DebugString())); + OP_REQUIRES(context, TensorShapeUtils::IsVector(shape.shape()), + errors::InvalidArgument("Invalid `shape` shape: ", + shape.shape().DebugString())); + TensorShape output_shape; + OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(shape.vec(), + &output_shape)); + OP_REQUIRES_OK(context, CheckCdfShape(output_shape, cdf.shape())); + + std::vector data_shape, cdf_shape; + OP_REQUIRES_OK( + context, MergeAxes(output_shape, cdf.shape(), &data_shape, &cdf_shape)); + + const string& encoded = encoded_tensor.scalar()(); + + Tensor* output; + OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); + + switch (data_shape.size()) { +#define RANGE_DECODE_CASE(dim) \ + case dim: { \ + RangeDecodeImpl(output->flat(), data_shape, \ + cdf.flat_inner_dims(), cdf_shape, encoded); \ + } break + RANGE_DECODE_CASE(1); + RANGE_DECODE_CASE(2); + RANGE_DECODE_CASE(3); + RANGE_DECODE_CASE(4); + RANGE_DECODE_CASE(5); + RANGE_DECODE_CASE(6); +#undef RANGE_DECODE_CASE + default: + context->CtxFailure(errors::InvalidArgument( + "Irregular broadcast pattern: ", output_shape.DebugString(), ", ", + cdf.shape().DebugString())); + return; + } + } + + private: + template + void RangeDecodeImpl(TTypes::Flat output, + gtl::ArraySlice output_shape, + TTypes::ConstMatrix cdf, + gtl::ArraySlice cdf_shape, + const string& encoded) const { + BroadcastRange view{output.data(), output_shape, + cdf.data(), cdf_shape}; + + RangeDecoder decoder{encoded, precision_}; + + const int64 output_size = output.size(); + const int64 cdf_size = cdf.size(); + const auto chip_size = + static_cast::size_type>(cdf.dimension(1)); + + for (int64 i = 0; i < output_size; ++i) { + const auto pair = view.Next(); + + int16* data = pair.first; + DCHECK_LT(data, output.data() + output_size); + + const int32* cdf_slice = pair.second; + DCHECK_LE(cdf_slice + chip_size, cdf.data() + cdf_size); + + *data = decoder.Decode(gtl::ArraySlice{cdf_slice, chip_size}); + } + } + + int precision_; +}; + +REGISTER_KERNEL_BUILDER(Name("RangeDecode").Device(DEVICE_CPU), RangeDecodeOp); +} // namespace +} // namespace tensorflow diff --git a/tensorflow/contrib/coder/kernels/range_coder_ops_test.cc b/tensorflow/contrib/coder/kernels/range_coder_ops_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..ae4d9d2836a0f89a9765004a85bc3c292b0e484f --- /dev/null +++ b/tensorflow/contrib/coder/kernels/range_coder_ops_test.cc @@ -0,0 +1,521 @@ +/* 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 +#include + +#include "tensorflow/contrib/coder/kernels/range_coder.h" +#include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h" +#include "tensorflow/core/common_runtime/shape_refiner.h" +#include "tensorflow/core/framework/fake_input.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/node_def_builder.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/tensor_types.h" +#include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/framework/versions.pb.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/graph/testlib.h" +#include "tensorflow/core/kernels/ops_testutil.h" +#include "tensorflow/core/lib/core/bits.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/random/random.h" +#include "tensorflow/core/lib/random/simple_philox.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/test_benchmark.h" +#include "tensorflow/core/public/session.h" +#include "tensorflow/core/public/session_options.h" + +namespace tensorflow { +namespace { +int LogUniform(random::SimplePhilox* gen, uint32 n) { + CHECK_GT(n, 0); + + // Split [0, n) into {0}, [1, 2), [2, 4), [4, 8), ..., [2^(m-1), n). + const int m = Log2Ceiling(n); + + int outcome; + do { + // Uniform() consumes at least 32 bits per call, therefore this is somewhat + // wasteful implementation. Since this is used only for test, we do not + // refine this implementation further. + const int k = gen->Uniform(m + 1) - 1; + // If k == -1, then sample from {0}. + // If k == 0, then sample from [1, 2). + // If k == 1, then sample from [2, 4), ... and so on. + if (k < 1) { + outcome = k + 1; + } else { + outcome = (1 << k) + gen->Uniform(1 << k); + } + } while (n <= outcome); + return outcome; +} + +std::vector ComputeStrides(const TensorShape& shape) { + std::vector stride(shape.dims()); + int64 current = 1; + for (int i = shape.dims() - 1; i >= 0; --i) { + stride[i] = current; + current *= shape.dim_size(i); + } + return stride; +} + +class RangeCoderOpsTest : public OpsTestBase { + protected: + Status RunEncodeOp(int precision, gtl::ArraySlice input, + Tensor* output) { + TF_RETURN_IF_ERROR(NodeDefBuilder("encode", "RangeEncode") + .Input(tensorflow::FakeInput(DT_INT16)) + .Input(tensorflow::FakeInput(DT_INT32)) + .Attr("precision", precision) + .Finalize(node_def())); + TF_RETURN_IF_ERROR(InitOp()); + + inputs_.clear(); + std::vector copies(input.size()); + for (int i = 0; i < input.size(); ++i) { + copies[i] = input[i]; + inputs_.emplace_back(&copies[i]); + } + + TF_RETURN_IF_ERROR(RunOpKernel()); + + *output = *GetOutput(0); + inputs_.clear(); + + return Status::OK(); + } + + Status RunDecodeOp(int precision, gtl::ArraySlice input, + Tensor* output) { + TF_RETURN_IF_ERROR(NodeDefBuilder("decode", "RangeDecode") + .Input(tensorflow::FakeInput(DT_STRING)) + .Input(tensorflow::FakeInput(DT_INT32)) + .Input(tensorflow::FakeInput(DT_INT32)) + .Attr("precision", precision) + .Finalize(node_def())); + TF_RETURN_IF_ERROR(InitOp()); + + inputs_.clear(); + std::vector copies(input.size()); + for (int i = 0; i < input.size(); ++i) { + copies[i] = input[i]; + inputs_.emplace_back(&copies[i]); + } + + TF_RETURN_IF_ERROR(RunOpKernel()); + + *output = *GetOutput(0); + inputs_.clear(); + + return Status::OK(); + } + + void TestEncodeAndDecode(int precision, const Tensor& data, + const Tensor& cdf) { + Tensor encoded; + TF_ASSERT_OK(RunEncodeOp(precision, {data, cdf}, &encoded)); + + const TensorShape& data_shape = data.shape(); + Tensor shape{DT_INT32, {data_shape.dims()}}; + for (int i = 0; i < data_shape.dims(); ++i) { + shape.flat()(i) = data_shape.dim_size(i); + } + + Tensor decoded; + TF_ASSERT_OK(RunDecodeOp(precision, {encoded, shape, cdf}, &decoded)); + + EXPECT_EQ(decoded.dtype(), data.dtype()); + EXPECT_EQ(decoded.shape(), data.shape()); + EXPECT_EQ(decoded.tensor_data(), data.tensor_data()); + } + + void PopulateMaxValues(random::SimplePhilox* gen, Tensor* maxvalue_tensor, + int min_maxvalue, int max_maxvalue) { + const int range = max_maxvalue - min_maxvalue; + TTypes::Flat flat = maxvalue_tensor->flat(); + + for (int64 i = 0; i < flat.size(); ++i) { + flat(i) = min_maxvalue + gen->Uniform(range); + } + } + + void BuildCdf(random::SimplePhilox* gen, Tensor* data_tensor, + Tensor* cdf_tensor, const Tensor& maxvalue_tensor) { + CHECK(TensorShapeUtils::StartsWith(cdf_tensor->shape(), + maxvalue_tensor.shape())); + CHECK_EQ(cdf_tensor->dims(), maxvalue_tensor.dims() + 1); + const int64 chip_size = cdf_tensor->dim_size(cdf_tensor->dims() - 1); + + std::vector data_stride = ComputeStrides(data_tensor->shape()); + std::vector cdf_stride = ComputeStrides(cdf_tensor->shape()); + + for (int i = 0; i < cdf_tensor->dims(); ++i) { + if (cdf_tensor->dim_size(i) == 1) { + cdf_stride[i] = 0; + } + } + + Tensor histogram_tensor{DT_INT32, cdf_tensor->shape()}; + TTypes::Flat data = data_tensor->flat(); + TTypes::Flat histogram = histogram_tensor.flat(); + TTypes::ConstFlat maxvalue = maxvalue_tensor.flat(); + histogram.setZero(); + + for (int64 index = 0; index < data.size(); ++index) { + int64 temp = index; + int64 offset = 0; + for (int dim = 0; dim < data_stride.size(); ++dim) { + const int64 coord = temp / data_stride[dim]; + offset += coord * cdf_stride[dim]; + temp -= coord * data_stride[dim]; + } + ASSERT_EQ(temp, 0); + + const int64 maxvalue_offset = offset / chip_size; + CHECK_EQ(maxvalue_offset * chip_size, offset); + CHECK_LT(maxvalue(maxvalue_offset) + 1, chip_size); + const int value = LogUniform(gen, maxvalue(maxvalue_offset)); + data(index) = value; + histogram(offset + value + 1) += 1; + } + + cdf_tensor->flat_inner_dims() = + histogram_tensor.flat_inner_dims().cumsum(1); + } +}; + +TEST_F(RangeCoderOpsTest, NoBroadcast) { + constexpr int kPrecision = 14; + constexpr int kMaxValue = 10; + + Tensor data{DT_INT16, {1, 32, 32, 16}}; + Tensor temp{DT_INT32, {1, 1, 1, 1, kMaxValue + 2}}; + Tensor maxvalue{DT_INT16, {1, 1, 1, 1}}; + maxvalue.flat()(0) = kMaxValue; + + ASSERT_LE(data.shape().num_elements(), 1 << kPrecision); + + random::PhiloxRandom philox(random::New64(), random::New64()); + random::SimplePhilox gen(&philox); + BuildCdf(&gen, &data, &temp, maxvalue); + + const Eigen::array broadcast = {1, 32, 32, 16, 1}; + + Tensor cdf{DT_INT32, {1, 32, 32, 16, kMaxValue + 2}}; + cdf.tensor() = temp.tensor().broadcast(broadcast); + + TestEncodeAndDecode(kPrecision, data, cdf); +} + +TEST_F(RangeCoderOpsTest, Broadcast1Axis) { + constexpr int kPrecision = 9; + constexpr int kDimensionSize = 1 << kPrecision; + constexpr int kMinMaxValue = 10; + constexpr int kMaxMaxValue = 64; + + random::PhiloxRandom philox(random::New64(), random::New64()); + random::SimplePhilox gen(&philox); + Tensor data{DT_INT16, {1, kDimensionSize, kDimensionSize}}; + + Tensor maxvalue{DT_INT16, {kDimensionSize}}; + PopulateMaxValues(&gen, &maxvalue, kMinMaxValue, kMaxMaxValue); + + { + // Axis 1. + Tensor maxvalue1; + ASSERT_TRUE(maxvalue1.CopyFrom(maxvalue, {1, 1, kDimensionSize})); + + Tensor cdf{DT_INT32, {1, 1, kDimensionSize, kMaxMaxValue + 2}}; + BuildCdf(&gen, &data, &cdf, maxvalue1); + TestEncodeAndDecode(kPrecision, data, cdf); + } + + { + // Axis 2. + Tensor maxvalue2; + ASSERT_TRUE(maxvalue2.CopyFrom(maxvalue, {1, kDimensionSize, 1})); + + Tensor cdf{DT_INT32, {1, kDimensionSize, 1, kMaxMaxValue + 2}}; + BuildCdf(&gen, &data, &cdf, maxvalue2); + TestEncodeAndDecode(kPrecision, data, cdf); + } +} + +TEST_F(RangeCoderOpsTest, Broadcast2Axes) { + constexpr int kPrecision = 13; + constexpr int kDimensionSize1 = 1 << (kPrecision / 2); + constexpr int kDimensionSize2 = 1 << (kPrecision - kPrecision / 2); + constexpr int kMinMaxValue = 10; + constexpr int kMaxMaxValue = 64; + + random::PhiloxRandom philox(random::New64(), random::New64()); + random::SimplePhilox gen(&philox); + Tensor maxvalue{DT_INT16, {2, 1, 1, 7}}; + PopulateMaxValues(&gen, &maxvalue, kMinMaxValue, kMaxMaxValue); + + Tensor data{DT_INT16, {2, kDimensionSize1, kDimensionSize2, 7}}; + Tensor cdf{DT_INT32, {2, 1, 1, 7, kMaxMaxValue + 2}}; + BuildCdf(&gen, &data, &cdf, maxvalue); + TestEncodeAndDecode(kPrecision, data, cdf); +} + +TEST_F(RangeCoderOpsTest, InvalidCdfShape) { + Tensor data{DT_INT16, {3, 3}}; + Tensor cdf{DT_INT32, {3, 3}}; + + Tensor unused; + { + const Status status = RunEncodeOp(10, {data, cdf}, &unused); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.error_message().find("`cdf` should have one more axis"), + string::npos); + } + + Tensor empty{DT_STRING, {}}; + Tensor shape{DT_INT32, {2}}; + shape.vec().setValues({3, 3}); + { + const Status status = RunDecodeOp(10, {empty, shape, cdf}, &unused); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.error_message().find("`cdf` should have one more axis"), + string::npos); + } + + cdf = Tensor{DT_INT32, {3, 3, 1}}; + { + const Status status = RunEncodeOp(10, {data, cdf}, &unused); + EXPECT_FALSE(status.ok()); + EXPECT_NE( + status.error_message().find("last dimension of `cdf` should be > 1"), + string::npos); + } + { + const Status status = RunDecodeOp(10, {empty, shape, cdf}, &unused); + EXPECT_FALSE(status.ok()); + EXPECT_NE( + status.error_message().find("last dimension of `cdf` should be > 1"), + string::npos); + } +} + +TEST_F(RangeCoderOpsTest, DecoderShapeFn) { + Tensor encoded_tensor{DT_STRING, {}}; + Tensor shape_tensor{DT_INT32, {3}}; + Tensor cdf_tensor{DT_INT32, {4, 6, 8, 2}}; + + shape_tensor.flat().setValues({4, 6, 8}); + + Graph g{OpRegistry::Global()}; + Node* encoded = test::graph::Constant(&g, encoded_tensor); + Node* shape = test::graph::Constant(&g, shape_tensor); + Node* cdf = test::graph::Constant(&g, cdf_tensor); + Node* decode; + TF_ASSERT_OK(NodeBuilder("range_decode", "RangeDecode", g.op_registry()) + .Input(encoded) + .Input(shape) + .Input(cdf) + .Attr("precision", 10) + .Finalize(&g, &decode)); + + ShapeRefiner refiner{g.versions().producer(), g.op_registry()}; + TF_ASSERT_OK(refiner.AddNode(encoded)); + TF_ASSERT_OK(refiner.AddNode(shape)); + TF_ASSERT_OK(refiner.AddNode(cdf)); + TF_ASSERT_OK(refiner.AddNode(decode)); + + auto* context = refiner.GetContext(decode); + ASSERT_NE(context, nullptr); + + ASSERT_EQ(context->num_outputs(), 1); + auto shape_handle = context->output(0); + + ASSERT_EQ(context->Rank(shape_handle), 3); + EXPECT_EQ(context->Value(context->Dim(shape_handle, 0)), 4); + EXPECT_EQ(context->Value(context->Dim(shape_handle, 1)), 6); + EXPECT_EQ(context->Value(context->Dim(shape_handle, 2)), 8); +} + +TEST_F(RangeCoderOpsTest, InvalidBroadcast) { + Tensor data{DT_INT16, {3, 3}}; + Tensor cdf{DT_INT32, {3, 2, 2}}; + + Tensor unused; + { + const Status status = RunEncodeOp(10, {data, cdf}, &unused); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.error_message().find("Cannot broadcast shape"), + string::npos); + } + + data = Tensor{DT_INT16, {3, 1}}; + cdf = Tensor{DT_INT32, {3, 3, 2}}; + Tensor empty{DT_STRING, {}}; + Tensor shape{DT_INT32, {2}}; + shape.vec().setValues({3, 1}); + { + const Status status = RunDecodeOp(10, {empty, shape, cdf}, &unused); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.error_message().find("Cannot broadcast shape"), + string::npos); + } + + std::vector shape_vector = {2, 2, 2, 2, 2, 2, 2, 2, 2}; + data = Tensor{DT_INT16, TensorShape{shape_vector}}; + cdf = Tensor{DT_INT32, {2, 1, 2, 1, 2, 1, 2, 1, 2, 2}}; + { + const Status status = RunEncodeOp(10, {data, cdf}, &unused); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.error_message().find("Irregular broadcast"), string::npos); + } + + shape = Tensor{DT_INT32, {static_cast(shape_vector.size())}}; + for (int i = 0; i < shape_vector.size(); ++i) { + shape.flat()(i) = shape_vector[i]; + } + { + const Status status = RunDecodeOp(10, {empty, shape, cdf}, &unused); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.error_message().find("Irregular broadcast"), string::npos); + } +} + +// Benchmark ------------------------------------------------------------- + +// This function creates RangeEncode graph with CDF built from a separate data +// sample. +Graph* CreateRangeEncodeFullBroadcastGraph(const TensorShape& shape, + int precision) { + CHECK_EQ(shape.dims(), 4); + + constexpr int kAlphabetSize = 70; + + Tensor histogram{DT_INT32, {kAlphabetSize + 1}}; + TTypes::Vec h = histogram.vec(); + h.setConstant(1); + h(0) = 0; + + random::PhiloxRandom philox(random::New64(), random::New64()); + random::SimplePhilox gen(&philox); + for (int i = 0; i < (1 << precision) - kAlphabetSize; ++i) { + const int value = LogUniform(&gen, kAlphabetSize - 1); + h(value + 1) += 1; + } + + Tensor cdf{DT_INT32, {1, 1, 1, 1, kAlphabetSize + 1}}; + cdf.flat() = h.cumsum(0); + + Tensor data{DT_INT16, shape}; + TTypes::Flat d = data.flat(); + for (int64 i = 0; i < d.size(); ++i) { + d(i) = LogUniform(&gen, kAlphabetSize - 1); + } + + Graph* g = new Graph(OpRegistry::Global()); + TF_CHECK_OK(NodeBuilder("range_encode", "RangeEncode", g->op_registry()) + .Input(test::graph::Constant(g, data)) + .Input(test::graph::Constant(g, cdf)) + .Attr("precision", precision) + .Finalize(g, nullptr)); + return g; +} + +// This function creates RangeDecode graph with CDF built from a separate data +// sample. +Graph* CreateRangeDecodeFullBroadcastGraph(const TensorShape& shape, + int precision) { + CHECK_EQ(shape.dims(), 4); + + constexpr int kAlphabetSize = 200; + const int64 num_elements = shape.num_elements(); + + Tensor histogram{DT_INT32, {kAlphabetSize + 1}}; + TTypes::Vec h = histogram.vec(); + h.setConstant(1); + h(0) = 0; + + random::PhiloxRandom philox(random::New64(), random::New64()); + random::SimplePhilox gen(&philox); + for (int i = 0; i < (1 << precision) - kAlphabetSize; ++i) { + const int value = LogUniform(&gen, kAlphabetSize - 1); + h(value + 1) += 1; + } + + Tensor cdf_tensor{DT_INT32, {1, 1, 1, 1, kAlphabetSize + 1}}; + TTypes::Flat cdf = cdf_tensor.flat(); + cdf = h.cumsum(0); + + Tensor string_tensor{DT_STRING, TensorShape{}}; + string& sink = string_tensor.scalar()(); + + RangeEncoder encoder{precision}; + for (int64 i = 0; i < num_elements; ++i) { + const int value = LogUniform(&gen, kAlphabetSize - 1); + encoder.Encode(cdf(value), cdf(value + 1), &sink); + } + encoder.Finalize(&sink); + + Tensor shape_tensor{DT_INT32, {shape.dims()}}; + for (int i = 0; i < shape.dims(); ++i) { + shape_tensor.flat()(i) = shape.dim_size(i); + } + + Graph* g = new Graph(OpRegistry::Global()); + TF_CHECK_OK(NodeBuilder("range_decode", "RangeDecode", g->op_registry()) + .Input(test::graph::Constant(g, string_tensor)) + .Input(test::graph::Constant(g, shape_tensor)) + .Input(test::graph::Constant(g, cdf_tensor)) + .Attr("precision", precision) + .Finalize(g, nullptr)); + return g; +} + +void RunTensorFlowBenchmark(int iters, Graph* g, int64 num_elements) { + SessionOptions opts; + opts.config.set_intra_op_parallelism_threads(1); + opts.config.set_inter_op_parallelism_threads(1); + + testing::UseRealTime(); + test::Benchmark("cpu", g, &opts).Run(iters); + + const int64 num_items = static_cast(iters) * num_elements; + testing::ItemsProcessed(num_items); +} + +void BM_RangeEncodeFullBroadcast(int iters, int code_size) { + constexpr int kPrecision = 14; + const TensorShape shape = {1, code_size, code_size, 256}; + Graph* g = CreateRangeEncodeFullBroadcastGraph(shape, kPrecision); + RunTensorFlowBenchmark(iters, g, shape.num_elements()); +} + +BENCHMARK(BM_RangeEncodeFullBroadcast)->Arg(32)->Arg(64); + +void BM_RangeDecodeFullBroadcast(int iters, int code_size) { + constexpr int kPrecision = 14; + const TensorShape shape = {1, code_size, code_size, 256}; + Graph* g = CreateRangeDecodeFullBroadcastGraph(shape, kPrecision); + RunTensorFlowBenchmark(iters, g, shape.num_elements()); +} + +BENCHMARK(BM_RangeDecodeFullBroadcast)->Arg(32)->Arg(64); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/contrib/coder/kernels/range_coder_ops_util.cc b/tensorflow/contrib/coder/kernels/range_coder_ops_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..d66730cb4881ea92b5477047c500291fa9c0c290 --- /dev/null +++ b/tensorflow/contrib/coder/kernels/range_coder_ops_util.cc @@ -0,0 +1,85 @@ +/* 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/coder/kernels/range_coder_ops_util.h" + +#include + +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/types.h" + +using tensorflow::errors::InvalidArgument; + +namespace tensorflow { +Status MergeAxes(const TensorShape& broadcast_shape, + const TensorShape& storage_shape, + std::vector* merged_broadcast_shape_pointer, + std::vector* merged_storage_shape_pointer) { + CHECK_EQ(storage_shape.dims(), broadcast_shape.dims() + 1); + + std::vector& merged_broadcast_shape = *merged_broadcast_shape_pointer; + std::vector& merged_storage_shape = *merged_storage_shape_pointer; + + // The shapes are simplified so that the conversions between linear index + // and coordinates takes less CPU cycles. Two adjacent dimensions are + // merged if they both are broadcasting dimensions or if they both are + // non-broadcasting dimensions. + merged_broadcast_shape.resize(1); + merged_broadcast_shape[0] = 1; + merged_storage_shape.resize(1); + merged_storage_shape[0] = 1; + + for (int i = 0, j = 0; j < broadcast_shape.dims(); ++j) { + if (TF_PREDICT_FALSE( + (broadcast_shape.dim_size(j) != storage_shape.dim_size(j)) && + (storage_shape.dim_size(j) != 1))) { + return InvalidArgument("Cannot broadcast shape ", + storage_shape.DebugString(), " to ", + broadcast_shape.DebugString()); + } + + const bool was_broadcasting = (merged_storage_shape[i] == 1); + const bool is_broadcasting = (storage_shape.dim_size(j) == 1); + + // Merge two adjacent axes if they both are broadcasting or both are + // non-broadcasting axes. The second and the third conditions in the if + // clause below are when the previously merged axis or the next j-th axis + // may be interpreted as either a broadcasting or a non-broadcasting axis. + const bool merge = (was_broadcasting == is_broadcasting) || + (broadcast_shape.dim_size(j) <= 1) || + (merged_broadcast_shape[i] <= 1); + + if (merge) { + merged_broadcast_shape[i] *= broadcast_shape.dim_size(j); + merged_storage_shape[i] *= storage_shape.dim_size(j); + } else { + // Move to the next axis. + merged_broadcast_shape.push_back(broadcast_shape.dim_size(j)); + merged_storage_shape.push_back(storage_shape.dim_size(j)); + ++i; + } + } + + int64 storage_stride = 1; + for (int i = broadcast_shape.dims(); i < storage_shape.dims(); ++i) { + storage_stride *= storage_shape.dim_size(i); + } + merged_storage_shape.push_back(storage_stride); + + return Status::OK(); +} +} // namespace tensorflow diff --git a/tensorflow/contrib/coder/kernels/range_coder_ops_util.h b/tensorflow/contrib/coder/kernels/range_coder_ops_util.h new file mode 100644 index 0000000000000000000000000000000000000000..b8aabcef62e9de53810397960f871abc4adc0cf9 --- /dev/null +++ b/tensorflow/contrib/coder/kernels/range_coder_ops_util.h @@ -0,0 +1,33 @@ +/* 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_CODER_KERNELS_RANGE_CODER_OPS_UTIL_H_ +#define TENSORFLOW_CONTRIB_CODER_KERNELS_RANGE_CODER_OPS_UTIL_H_ + +#include + +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +// The shapes are simplified to reduce indexing cost. +Status MergeAxes(const TensorShape& broadcast_shape, + const TensorShape& storage_shape, + std::vector* merged_broadcast_shape_pointer, + std::vector* merged_storage_shape_pointer); +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_CODER_KERNELS_RANGE_CODER_OPS_UTIL_H_ diff --git a/tensorflow/contrib/coder/kernels/range_coder_test.cc b/tensorflow/contrib/coder/kernels/range_coder_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..442994bf7c7566c1cbe1c439050a69e5b9a4208e --- /dev/null +++ b/tensorflow/contrib/coder/kernels/range_coder_test.cc @@ -0,0 +1,116 @@ +/* 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/coder/kernels/range_coder.h" + +#include + +#include "tensorflow/core/lib/random/distribution_sampler.h" +#include "tensorflow/core/lib/random/philox_random.h" +#include "tensorflow/core/lib/random/random.h" +#include "tensorflow/core/lib/random/simple_philox.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace { +void RangeEncodeDecodeTest(int precision, random::SimplePhilox* gen) { + constexpr int kAlphabetSize = 256; + + std::vector distribution_weight; + distribution_weight.reserve(kAlphabetSize); + for (int i = 1; i <= kAlphabetSize; ++i) { + distribution_weight.push_back(std::pow(static_cast(i), -2.0f)); + } + + random::DistributionSampler sampler(distribution_weight); + + const int multiplier = (precision > 7) ? 32 : 1; + std::vector histogram(kAlphabetSize, multiplier - 1); + + const int data_size = + (multiplier << precision) - histogram.size() * (multiplier - 1); + CHECK_GE(data_size, 0); + std::vector data(data_size); + for (uint8& x : data) { + x = sampler.Sample(gen); + ++histogram[x]; + } + + std::vector cdf(histogram.size() + 1, 0); + int partial_sum = 0; + for (int i = 0; i < histogram.size(); ++i) { + partial_sum += histogram[i]; + cdf[i + 1] = partial_sum / multiplier; + } + + ASSERT_EQ(cdf.front(), 0); + ASSERT_EQ(cdf.back(), 1 << precision); + + std::vector ideal_code_length(histogram.size()); + const double normalizer = static_cast(1 << precision); + for (int i = 0; i < ideal_code_length.size(); ++i) { + ideal_code_length[i] = -std::log2((cdf[i + 1] - cdf[i]) / normalizer); + } + + RangeEncoder encoder(precision); + string encoded; + double ideal_length = 0.0; + for (uint8 x : data) { + encoder.Encode(cdf[x], cdf[x + 1], &encoded); + ideal_length += ideal_code_length[x]; + } + encoder.Finalize(&encoded); + + LOG(INFO) << "Encoded string length (bits): " << 8 * encoded.size() + << ", whereas ideal " << ideal_length << " (" + << (8 * encoded.size()) / ideal_length << " of ideal) " + << " (ideal compression rate " << ideal_length / (8 * data.size()) + << ")"; + + RangeDecoder decoder(encoded, precision); + for (int i = 0; i < data.size(); ++i) { + const int32 decoded = decoder.Decode(cdf); + ASSERT_EQ(decoded, static_cast(data[i])) << i; + } +} + +TEST(RangeCoderTest, Precision1To11) { + random::PhiloxRandom gen(random::New64(), random::New64()); + random::SimplePhilox rand(&gen); + const int precision = 1 + rand.Uniform(11); + RangeEncodeDecodeTest(precision, &rand); +} + +TEST(RangeCoderTest, Precision12To16) { + random::PhiloxRandom gen(random::New64(), random::New64()); + random::SimplePhilox rand(&gen); + for (int precision = 12; precision < 17; ++precision) { + RangeEncodeDecodeTest(precision, &rand); + } +} + +TEST(RangeCoderTest, FinalizeState0) { + constexpr int kPrecision = 2; + + string output; + RangeEncoder encoder(kPrecision); + encoder.Encode(0, 2, &output); + encoder.Finalize(&output); + + RangeDecoder decoder(output, kPrecision); + EXPECT_EQ(decoder.Decode({0, 2, 4}), 0); +} +} // namespace +} // namespace tensorflow diff --git a/tensorflow/contrib/coder/ops/coder_ops.cc b/tensorflow/contrib/coder/ops/coder_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..9056d1a6963d7be92f499db31385fb6afe2dc515 --- /dev/null +++ b/tensorflow/contrib/coder/ops/coder_ops.cc @@ -0,0 +1,119 @@ +/* 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/common_shape_fns.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { +using shape_inference::InferenceContext; +using shape_inference::ShapeHandle; + +// clang-format off +REGISTER_OP("RangeEncode") + .Input("data: int16") + .Input("cdf: int32") + .Output("encoded: string") + .Attr("precision: int >= 1") + .SetShapeFn(shape_inference::ScalarShape) + .Doc(R"doc( +Using the provided cumulative distribution functions (CDF) inside `cdf`, returns +a range-code of `data`. + +The shape of `cdf` should have one more axis than the shape of `data`, and the +prefix `cdf.shape[:-1]` should be broadcastable to `data.shape`. That is, for +every `i = 0,...,rank(data) - 1`, the op requires that either +`cdf.shape[i] == 1` or `cdf.shape[i] == data.shape[i]`. Note that this +broadcasting is limited in the sense that the number of axes must match, and +broadcasts only `cdf` but not `data`. + +`data` should have an upper bound `m > 0` such that each element is an integer +in range `[0, m)`. Then the last dimension size of `cdf` must be `m + 1`. For +each element of `data`, the innermost strip of `cdf` is a vector representing a +CDF. For each k = 0,...,m, `cdf[..., k] / 2^precision` is the probability that +an outcome is less than `k` (not less than or equal to). + +``` + cdf[..., 0] / 2^precision = Pr(data[...] < 0) + cdf[..., 1] / 2^precision = Pr(data[...] < 1) = Pr(data[...] <= 0) + cdf[..., 2] / 2^precision = Pr(data[...] < 2) = Pr(data[...] <= 1) + ... + cdf[..., m] / 2^precision = Pr(data[...] < m) = 1 +``` + +Therefore each element of `cdf` must be in `[0, 2^precision]`. + +Ideally `cdf[..., m]` should equal to `2^precision` but this is not a hard +requirement as long as `cdf[..., m] <= 2^precision`. + +The encoded string neither contains the shape information of the encoded data +nor a termination symbol. Therefore the shape of the encoded data must be +explicitly provided to the decoder. + +Implementation notes: + +- Because of potential performance issues, the op does not check whether +elements of `data` is in the correct range `[0, m)`, or if `cdf` satisfies +monotonic increase property. + +- For the range coder to decode the encoded string correctly, the decoder should +be able to reproduce the internal states of the encoder precisely. Otherwise, +the decoding would fail and once an error occur, all subsequent decoded values +are incorrect. For this reason, the range coder uses integer arithmetics and +avoids using any floating point operations internally, and `cdf` should contain +integers representing quantized probability mass rather than floating points. + +data: An int32 tensor. +cdf: An int32 tensor representing the CDF's of `data`. Each integer is divided + by `2^precision` to represent a fraction. +encoded: A range-coded scalar string. +precision: The number of bits for probability quantization. Must be <= 16. +)doc"); + + +REGISTER_OP("RangeDecode") + .Input("encoded: string") + .Input("shape: int32") + .Input("cdf: int32") + .Output("decoded: int16") + .Attr("precision: int >= 1") + .SetShapeFn([] (InferenceContext* c) { + ShapeHandle out; + TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(1, &out)); + c->set_output(0, out); + return Status::OK(); + }) + .Doc(R"doc( +Decodes a range-coded `code` into an int32 tensor of shape `shape`. + +This is the reverse op of RangeEncode. The shape of the tensor that was encoded +should be known by the caller. + +Implementation notes: + +- If wrong input was given (e.g., corrupt `encoded` string, or `cdf` or +`precision` do not match encoder), the decode is unsuccessful. Because of +potential performance issues, the decoder does not return error status. + +encoded: A scalar string tensor from RangeEncode. +shape: An int32 1-D tensor representing the shape of the data encoded by + RangeEncode. +decoded: An int32 tensor with shape equal to `shape`. +precision: The number of bits for probability quantization. Must be <= 16, and + must match the precision used by RangeEncode that produced `encoded`. +)doc"); +// clang-format on +} // namespace tensorflow diff --git a/tensorflow/contrib/coder/python/ops/coder_ops.py b/tensorflow/contrib/coder/python/ops/coder_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..bb262e338baf1d9c3c043f03a02c2d2851e22b49 --- /dev/null +++ b/tensorflow/contrib/coder/python/ops/coder_ops.py @@ -0,0 +1,30 @@ +# 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. +# ============================================================================== +"""Range coder operations.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# pylint: disable=wildcard-import,unused-import +from tensorflow.contrib.coder.python.ops import gen_coder_ops +from tensorflow.contrib.coder.python.ops.gen_coder_ops import * +# pylint: enable=wildcard-import,unused-import +from tensorflow.contrib.util import loader +from tensorflow.python.platform import resource_loader + + +_coder_ops = loader.load_op_library( + resource_loader.get_path_to_datafile("_coder_ops.so")) diff --git a/tensorflow/contrib/coder/python/ops/coder_ops_test.py b/tensorflow/contrib/coder/python/ops/coder_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..d5e14e7a641b5673e97882daf2b5a1796ee1bbef --- /dev/null +++ b/tensorflow/contrib/coder/python/ops/coder_ops_test.py @@ -0,0 +1,53 @@ +# 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. +# ============================================================================== +"""Coder operations tests.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.coder.python.ops import coder_ops +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import random_ops +from tensorflow.python.platform import test + + +class CoderOpsTest(test.TestCase): + """Coder ops test. + + Coder ops have C++ tests. Python test just ensures that Python binding is not + broken. + """ + + def testReadmeExample(self): + data = random_ops.random_uniform((128, 128), 0, 10, dtype=dtypes.int32) + histogram = math_ops.bincount(data, minlength=10, maxlength=10) + cdf = math_ops.cumsum(histogram, exclusive=False) + cdf = array_ops.pad(cdf, [[1, 0]]) + cdf = array_ops.reshape(cdf, [1, 1, -1]) + + data = math_ops.cast(data, dtypes.int16) + encoded = coder_ops.range_encode(data, cdf, precision=14) + decoded = coder_ops.range_decode( + encoded, array_ops.shape(data), cdf, precision=14) + + with self.test_session() as sess: + self.assertAllEqual(*sess.run((data, decoded))) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/compiler/jit_test.py b/tensorflow/contrib/compiler/jit_test.py index 2108e42bce4eba1eed158fe85888f1699a69ba7e..29a593f6bcfa05dcafcdb2f94087380ad720dba1 100644 --- a/tensorflow/contrib/compiler/jit_test.py +++ b/tensorflow/contrib/compiler/jit_test.py @@ -24,6 +24,7 @@ from tensorflow.python.framework import function from tensorflow.python.framework import op_def_registry from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed +from tensorflow.python.framework import test_util from tensorflow.python.ops import gradients from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops @@ -169,6 +170,7 @@ class JITTest(test.TestCase): self.assertEqual(b"jit_scope_0", func_attrs["_XlaScope"].s) +@test_util.with_c_api class CompilationEnabledInGradientTest(test.TestCase): def testCompilationInGradient(self): @@ -188,7 +190,7 @@ class CompilationEnabledInGradientTest(test.TestCase): for cg in c_grad_ops: self.assertTrue(cg.get_attr("_XlaCompile")) for ncg in nc_grad_ops: - with self.assertRaisesRegexp(ValueError, "No attr named"): + with self.assertRaisesRegexp(ValueError, "[Nn]o attr named"): ncg.get_attr("_XlaCompile") # d/dx (x ** 4) = 4 * (x ** 3) diff --git a/tensorflow/contrib/copy_graph/python/util/copy_elements.py b/tensorflow/contrib/copy_graph/python/util/copy_elements.py index bae66ffd4289308f2cbfc730ec50d057b13923fb..b806799202bff4f2f6dbf717fbeea74a04b8cd6e 100644 --- a/tensorflow/contrib/copy_graph/python/util/copy_elements.py +++ b/tensorflow/contrib/copy_graph/python/util/copy_elements.py @@ -35,10 +35,10 @@ from tensorflow.python.ops.variables import Variable from tensorflow.python.client.session import Session from tensorflow.python.framework import ops -__all__ = ["copy_op_to_graph", "copy_variable_to_graph", "get_copied_op"] +__all__ = ['copy_op_to_graph', 'copy_variable_to_graph', 'get_copied_op'] -def copy_variable_to_graph(org_instance, to_graph, scope=""): +def copy_variable_to_graph(org_instance, to_graph, scope=''): """Given a `Variable` instance from one `Graph`, initializes and returns a copy of it from another `Graph`, under the specified scope (default `""`). @@ -56,12 +56,11 @@ def copy_variable_to_graph(org_instance, to_graph, scope=""): """ if not isinstance(org_instance, Variable): - raise TypeError(str(org_instance) + " is not a Variable") + raise TypeError(str(org_instance) + ' is not a Variable') #The name of the new variable - if scope != "": - new_name = (scope + '/' + - org_instance.name[:org_instance.name.index(':')]) + if scope != '': + new_name = (scope + '/' + org_instance.name[:org_instance.name.index(':')]) else: new_name = org_instance.name[:org_instance.name.index(':')] @@ -73,15 +72,15 @@ def copy_variable_to_graph(org_instance, to_graph, scope=""): for name, collection in org_instance.graph._collections.items(): if org_instance in collection: if (name == ops.GraphKeys.GLOBAL_VARIABLES or - name == ops.GraphKeys.TRAINABLE_VARIABLES or - scope == ''): + name == ops.GraphKeys.TRAINABLE_VARIABLES or scope == ''): collections.append(name) else: collections.append(scope + '/' + name) #See if its trainable. - trainable = (org_instance in org_instance.graph.get_collection( - ops.GraphKeys.TRAINABLE_VARIABLES)) + trainable = ( + org_instance in org_instance.graph.get_collection( + ops.GraphKeys.TRAINABLE_VARIABLES)) #Get the initial value with org_instance.graph.as_default(): temp_session = Session() @@ -89,17 +88,17 @@ def copy_variable_to_graph(org_instance, to_graph, scope=""): #Initialize the new variable with to_graph.as_default(): - new_var = Variable(init_value, - trainable, - name=new_name, - collections=collections, - validate_shape=False) + new_var = Variable( + init_value, + trainable, + name=new_name, + collections=collections, + validate_shape=False) return new_var -def copy_op_to_graph(org_instance, to_graph, variables, - scope=""): +def copy_op_to_graph(org_instance, to_graph, variables, scope=''): """Returns a copy of an operation from another Graph under a specified scope. Given an `Operation` `org_instance` from one `Graph`, @@ -139,14 +138,12 @@ def copy_op_to_graph(org_instance, to_graph, variables, #If a variable by the new name already exists, return the #correspondng tensor that will act as an input if new_name in copied_variables: - return to_graph.get_tensor_by_name( - copied_variables[new_name].name) + return to_graph.get_tensor_by_name(copied_variables[new_name].name) #If an instance of the same name exists, return appropriately try: - already_present = to_graph.as_graph_element(new_name, - allow_tensor=True, - allow_operation=True) + already_present = to_graph.as_graph_element( + new_name, allow_tensor=True, allow_operation=True) return already_present except: pass @@ -184,20 +181,21 @@ def copy_op_to_graph(org_instance, to_graph, variables, #If it has an original_op parameter, copy it if op._original_op is not None: - new_original_op = copy_op_to_graph(op._original_op, to_graph, - variables, scope) + new_original_op = copy_op_to_graph(op._original_op, to_graph, variables, + scope) else: new_original_op = None #If it has control inputs, call this function recursively on each. - new_control_inputs = [copy_op_to_graph(x, to_graph, variables, - scope) - for x in op.control_inputs] + new_control_inputs = [ + copy_op_to_graph(x, to_graph, variables, scope) + for x in op.control_inputs + ] #If it has inputs, call this function recursively on each. - new_inputs = [copy_op_to_graph(x, to_graph, variables, - scope) - for x in op.inputs] + new_inputs = [ + copy_op_to_graph(x, to_graph, variables, scope) for x in op.inputs + ] #Make a new node_def based on that of the original. #An instance of tensorflow.core.framework.node_def_pb2.NodeDef, it @@ -216,13 +214,8 @@ def copy_op_to_graph(org_instance, to_graph, variables, op_def = deepcopy(op._op_def) #Initialize a new Operation instance - new_op = ops.Operation(new_node_def, - to_graph, - new_inputs, - output_types, - new_control_inputs, - input_types, - new_original_op, + new_op = ops.Operation(new_node_def, to_graph, new_inputs, output_types, + new_control_inputs, input_types, new_original_op, op_def) #Use Graph's hidden methods to add the op to_graph._add_op(new_op) # pylint: disable=protected-access @@ -233,10 +226,10 @@ def copy_op_to_graph(org_instance, to_graph, variables, return new_op else: - raise TypeError("Could not copy instance: " + str(org_instance)) + raise TypeError('Could not copy instance: ' + str(org_instance)) -def get_copied_op(org_instance, graph, scope=""): +def get_copied_op(org_instance, graph, scope=''): """Given an `Operation` instance from some `Graph`, returns its namesake from `graph`, under the specified scope (default `""`). @@ -259,5 +252,5 @@ def get_copied_op(org_instance, graph, scope=""): else: new_name = org_instance.name - return graph.as_graph_element(new_name, allow_tensor=True, - allow_operation=True) + return graph.as_graph_element( + new_name, allow_tensor=True, allow_operation=True) diff --git a/tensorflow/contrib/copy_graph/python/util/copy_test.py b/tensorflow/contrib/copy_graph/python/util/copy_test.py index 2798d31229d048561f8ebd9b63d3df94a44c45c7..05744bec4e05405c04b5ec442e72e4495737ab5b 100644 --- a/tensorflow/contrib/copy_graph/python/util/copy_test.py +++ b/tensorflow/contrib/copy_graph/python/util/copy_test.py @@ -17,9 +17,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np from tensorflow.contrib.copy_graph.python.util import copy_elements -from tensorflow.contrib.framework.python.framework import tensor_util from tensorflow.python.client import session as session_lib from tensorflow.python.framework import constant_op from tensorflow.python.framework import ops diff --git a/tensorflow/contrib/crf/python/ops/crf.py b/tensorflow/contrib/crf/python/ops/crf.py index 62708636c6181ca63cddf2b2e7c84d3da740282a..faa78769b98699af59047aed2865771120110fc2 100644 --- a/tensorflow/contrib/crf/python/ops/crf.py +++ b/tensorflow/contrib/crf/python/ops/crf.py @@ -166,8 +166,8 @@ def crf_log_likelihood(inputs, sequence_lengths: A [batch_size] vector of true sequence lengths. transition_params: A [num_tags, num_tags] transition matrix, if available. Returns: - log_likelihood: A scalar containing the log-likelihood of the given sequence - of tag indices. + log_likelihood: A [batch_size] `Tensor` containing the log-likelihood of + each example, given the sequence of tag indices. transition_params: A [num_tags, num_tags] transition matrix. This is either provided by the caller or created in this function. """ @@ -182,7 +182,7 @@ def crf_log_likelihood(inputs, transition_params) log_norm = crf_log_norm(inputs, sequence_lengths, transition_params) - # Normalize the scores to get the log-likelihood. + # Normalize the scores to get the log-likelihood per example. log_likelihood = sequence_scores - log_norm return log_likelihood, transition_params diff --git a/tensorflow/contrib/cudnn_rnn/BUILD b/tensorflow/contrib/cudnn_rnn/BUILD index 0751624bc4b7fbf413c342db3e5c440c9d572cd4..fec358c4e1067dc8dc8173d1b9d05dc90b90ca05 100644 --- a/tensorflow/contrib/cudnn_rnn/BUILD +++ b/tensorflow/contrib/cudnn_rnn/BUILD @@ -25,6 +25,7 @@ tf_custom_op_library( ], deps = [ "//tensorflow/core/kernels:bounds_check_lib", + "@farmhash_archive//:farmhash", ], ) @@ -39,6 +40,7 @@ tf_kernel_library( "//tensorflow/core:stream_executor", "//tensorflow/core/kernels:bounds_check_lib", "//third_party/eigen3", + "@farmhash_archive//:farmhash", ], ) diff --git a/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc b/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc index 5d5f593d016a3bb9f7b5ea8f5cd40c29268dc4f5..ba9686e94ee7072cc485c955decb2287bd4a56f3 100644 --- a/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc +++ b/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc @@ -39,6 +39,7 @@ limitations under the License. #include "tensorflow/core/lib/gtl/inlined_vector.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/platform/fingerprint.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/core/util/env_var.h" @@ -369,6 +370,27 @@ struct CudnnModelShapes { } }; +// Utility class for using CudnnModelShapes as a hash table key. +struct CudnnModelShapesHasher { + uint64 operator()(const CudnnModelShapes& to_hash) const { + uint64 hash = static_cast(to_hash.num_layers); + hash = tensorflow::FingerprintCat64( + hash, static_cast(to_hash.input_size)); + hash = tensorflow::FingerprintCat64(hash, + static_cast(to_hash.num_units)); + return tensorflow::FingerprintCat64(hash, + static_cast(to_hash.dir_count)); + } +}; + +// Utility class for using CudnnModelShapes as a hash table key. +struct CudnnModelShapesComparator { + bool operator()(const CudnnModelShapes& first, + const CudnnModelShapes& second) const { + return first.IsCompatibleWith(second); + } +}; + // Extract and checks the forward input tensors, parameters, and shapes from the // OpKernelContext. Status ExtractForwardInput(OpKernelContext* context, @@ -627,7 +649,7 @@ class CudnnRNNParamsToCanonical : public CudnnRNNKernelCommon { } const int num_params_per_layer = num_params_ / num_layers / num_dirs; // Number of params applied on inputs. The rest are applied on recurrent - // hiddden states. + // hidden states. const int num_params_input_state = num_params_per_layer / 2; CHECK(num_params_ % (num_layers * num_dirs) == 0) << "Number of params is not a multiple of num_layers * num_dirs."; @@ -764,6 +786,13 @@ TF_CALL_float(REGISTER_GPU); TF_CALL_double(REGISTER_GPU); #undef REGISTER_GPU +// Pointers to RNN scratch space for a specific set of shape parameters (used as +// a hash table value in CudnnRNNForwardOp and CudnnRNNBackwardOp). +struct RnnScratchSpace { + std::unique_ptr rnn_desc; + std::unique_ptr dropout_state_allocator; +}; + // Run the forward operation of the RNN model. template class CudnnRNNForwardOp : public CudnnRNNKernelCommon { @@ -808,32 +837,7 @@ class CudnnRNNForwardOp : public CudnnRNNKernelCommon { OP_REQUIRES_OK(context, ToRNNInputMode(rnn_input_mode(), model_shapes.num_units, model_shapes.input_size, &input_mode)); - // TODO(zhengxq): cache the descriptor so we don't have to create them all - // the time. auto data_type = ToDataType::value; - { - mutex_lock l(mu_); - if (model_shapes_ == nullptr) { - model_shapes_.reset(new CudnnModelShapes(model_shapes)); - } else { - OP_REQUIRES(context, model_shapes_->IsCompatibleWith(model_shapes), - errors::InvalidArgument( - "Incompatible rnn model shapes inferred: expecting ", - model_shapes_->RnnDescDebugString(), ", getting ", - model_shapes.RnnDescDebugString(), ".")); - } - if (rnn_desc_ == nullptr || ResetRndGenState()) { - dropout_state_allocator_.reset( - new CudnnRNNPersistentSpaceAllocator(context)); - auto rnn_desc_s = executor->createRnnDescriptor( - model_shapes_->num_layers, model_shapes_->num_units, - model_shapes_->input_size, input_mode, rnn_direction_mode(), - rnn_mode(), data_type, dropout(), seed(), - dropout_state_allocator_.get()); - OP_REQUIRES_OK(context, FromExecutorStatus(rnn_desc_s)); - rnn_desc_ = std::move(rnn_desc_s.ConsumeValueOrDie()); - } - } auto input_desc_s = executor->createRnnSequenceTensorDescriptor( input_shape.dim_size(0), input_shape.dim_size(1), @@ -882,14 +886,27 @@ class CudnnRNNForwardOp : public CudnnRNNKernelCommon { bool launch_status = false; { mutex_lock l(mu_); + RnnScratchSpace& rnn_state = rnn_state_cache_[model_shapes]; + if (rnn_state.rnn_desc == nullptr || ResetRndGenState()) { + CudnnRNNPersistentSpaceAllocator* dropout_state_allocator = + new CudnnRNNPersistentSpaceAllocator(context); + rnn_state.dropout_state_allocator.reset(dropout_state_allocator); + auto rnn_desc_s = executor->createRnnDescriptor( + model_shapes.num_layers, model_shapes.num_units, + model_shapes.input_size, input_mode, rnn_direction_mode(), + rnn_mode(), data_type, dropout(), seed(), dropout_state_allocator); + OP_REQUIRES_OK(context, FromExecutorStatus(rnn_desc_s)); + rnn_state.rnn_desc = std::move(rnn_desc_s.ConsumeValueOrDie()); + } launch_status = stream - ->ThenRnnForward( - *rnn_desc_, *input_desc, input_data, *hidden_state_desc, - input_h_data, *hidden_state_desc, input_c_data, params_data, - *output_desc, &output_data, *hidden_state_desc, - &output_h_data, *hidden_state_desc, &output_c_data, - is_training_, &reserve_space_allocator, &workspace_allocator) + ->ThenRnnForward(*rnn_state.rnn_desc, *input_desc, input_data, + *hidden_state_desc, input_h_data, + *hidden_state_desc, input_c_data, params_data, + *output_desc, &output_data, *hidden_state_desc, + &output_h_data, *hidden_state_desc, + &output_c_data, is_training_, + &reserve_space_allocator, &workspace_allocator) .ok(); } OP_REQUIRES(context, launch_status, @@ -899,10 +916,9 @@ class CudnnRNNForwardOp : public CudnnRNNKernelCommon { private: mutex mu_; bool is_training_; - std::unique_ptr model_shapes_ GUARDED_BY(mu_); - std::unique_ptr rnn_desc_ GUARDED_BY(mu_); - std::unique_ptr dropout_state_allocator_ - GUARDED_BY(mu_); + std::unordered_map + rnn_state_cache_ GUARDED_BY(mu_); }; #define REGISTER_GPU(T) \ @@ -1022,32 +1038,6 @@ class CudnnRNNBackwardOp : public CudnnRNNKernelCommon { OP_REQUIRES_OK(context, ToRNNInputMode(rnn_input_mode(), model_shapes.num_units, model_shapes.input_size, &input_mode)); - // TODO(zhengxq): cache the descriptor so we don't have to create them all - // the time. - { - mutex_lock l(mu_); - if (model_shapes_ == nullptr) { - model_shapes_.reset(new CudnnModelShapes(model_shapes)); - } else { - OP_REQUIRES(context, model_shapes_->IsCompatibleWith(model_shapes), - errors::InvalidArgument( - "Incompatible rnn model shapes inferred: expecting ", - model_shapes_->RnnDescDebugString(), ", getting ", - model_shapes.RnnDescDebugString(), ".")); - } - - if (rnn_desc_ == nullptr || ResetRndGenState()) { - dropout_state_allocator_.reset( - new CudnnRNNPersistentSpaceAllocator(context)); - auto rnn_desc_s = executor->createRnnDescriptor( - model_shapes.num_layers, model_shapes.num_units, - model_shapes.input_size, input_mode, rnn_direction_mode(), - rnn_mode(), data_type, dropout(), seed(), - dropout_state_allocator_.get()); - OP_REQUIRES_OK(context, FromExecutorStatus(rnn_desc_s)); - rnn_desc_ = std::move(rnn_desc_s.ConsumeValueOrDie()); - } - } auto input_desc_s = executor->createRnnSequenceTensorDescriptor( input_shape.dim_size(0), input_shape.dim_size(1), @@ -1100,17 +1090,30 @@ class CudnnRNNBackwardOp : public CudnnRNNKernelCommon { bool launch_status = false; { mutex_lock l(mu_); + RnnScratchSpace& rnn_state = rnn_state_cache_[model_shapes]; + if (rnn_state.rnn_desc == nullptr || ResetRndGenState()) { + CudnnRNNPersistentSpaceAllocator* dropout_state_allocator = + new CudnnRNNPersistentSpaceAllocator(context); + rnn_state.dropout_state_allocator.reset(dropout_state_allocator); + auto rnn_desc_s = executor->createRnnDescriptor( + model_shapes.num_layers, model_shapes.num_units, + model_shapes.input_size, input_mode, rnn_direction_mode(), + rnn_mode(), data_type, dropout(), seed(), dropout_state_allocator); + OP_REQUIRES_OK(context, FromExecutorStatus(rnn_desc_s)); + rnn_state.rnn_desc = std::move(rnn_desc_s.ConsumeValueOrDie()); + } launch_status = stream - ->ThenRnnBackward( - *rnn_desc_, *input_desc, input_data, *hidden_state_desc, - input_h_data, *hidden_state_desc, input_c_data, params_data, - *output_desc, output_data, *hidden_state_desc, output_h_data, - *hidden_state_desc, output_c_data, output_backprop_data, - output_h_backprop_data, output_c_backprop_data, - &input_backprop_data, &input_h_backprop_data, - &input_c_backprop_data, ¶ms_backprop_data, - &reserve_space_uint8, &workspace_allocator) + ->ThenRnnBackward(*rnn_state.rnn_desc, *input_desc, input_data, + *hidden_state_desc, input_h_data, + *hidden_state_desc, input_c_data, params_data, + *output_desc, output_data, *hidden_state_desc, + output_h_data, *hidden_state_desc, + output_c_data, output_backprop_data, + output_h_backprop_data, output_c_backprop_data, + &input_backprop_data, &input_h_backprop_data, + &input_c_backprop_data, ¶ms_backprop_data, + &reserve_space_uint8, &workspace_allocator) .ok(); } OP_REQUIRES(context, launch_status, @@ -1119,10 +1122,9 @@ class CudnnRNNBackwardOp : public CudnnRNNKernelCommon { private: mutex mu_; - std::unique_ptr model_shapes_ GUARDED_BY(mu_); - std::unique_ptr rnn_desc_ GUARDED_BY(mu_); - std::unique_ptr dropout_state_allocator_ - GUARDED_BY(mu_); + std::unordered_map + rnn_state_cache_ GUARDED_BY(mu_); }; #define REGISTER_GPU(T) \ diff --git a/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc b/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc index 9e41e67857101534e8bfef8d5d0b8a45ed8f1f76..1a79bf066c3a27e040099729fb079ee963f59270 100644 --- a/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc +++ b/tensorflow/contrib/cudnn_rnn/ops/cudnn_rnn_ops.cc @@ -251,9 +251,8 @@ REGISTER_OP("CudnnRNNParamsToCanonical") TF_RETURN_IF_ERROR(c->GetAttr("num_params", &num_params)); // Set shape for weight matrices for (int i = 0; i < num_params; i++) { - c->set_output(i, - c->Matrix(InferenceContext::kUnknownDim, - InferenceContext::kUnknownDim)); + c->set_output(i, c->Matrix(InferenceContext::kUnknownDim, + InferenceContext::kUnknownDim)); } // Set shape for bias vectors for (int i = 0; i < num_params; i++) { @@ -300,6 +299,7 @@ upcoming training or inferences. num_params: number of parameter sets for all layers. Each layer may contain multiple parameter sets, with each set consisting of a weight matrix and a bias vector. -)doc", kCudnnRNNCommonAttrs)); +)doc", + kCudnnRNNCommonAttrs)); } // namespace tensorflow diff --git a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_benchmark.py b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_benchmark.py index 4fc5ff1bd1887c4532e95fcf0e791d72b20471b0..933df6d71dd7c972efe63d54fa7344ecfc39b0a7 100644 --- a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_benchmark.py +++ b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_benchmark.py @@ -20,6 +20,7 @@ from __future__ import print_function import time +from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.contrib import rnn as contrib_rnn from tensorflow.contrib.cudnn_rnn.python.ops import cudnn_rnn_ops from tensorflow.contrib.rnn.python.ops import lstm_ops diff --git a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py index 4ddd75de60016bf37f8b12d10a9abe47f0e33852..9897c31a98e0b335c18a84825fc518ed1fc310a2 100644 --- a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py +++ b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py @@ -29,6 +29,8 @@ import numpy as np from tensorflow.contrib.cudnn_rnn.python.layers import cudnn_rnn from tensorflow.contrib.cudnn_rnn.python.ops import cudnn_rnn_ops from tensorflow.contrib.rnn.python.ops import rnn as contrib_rnn_lib +from tensorflow.python.eager import backprop +from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed @@ -49,7 +51,11 @@ from tensorflow.python.ops.losses import losses from tensorflow.python.platform import googletest from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import adagrad +from tensorflow.python.training import adam from tensorflow.python.training import gradient_descent +from tensorflow.python.training import momentum +from tensorflow.python.training import rmsprop from tensorflow.python.training import saver as saver_lib @@ -314,6 +320,55 @@ class CudnnRNNTestBasic(TensorFlowTestCase): self.assertEqual(0, total_sum2_v) self.assertEqual(0, total_sum3_v) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def testOptimizersSupport(self): + for opt in ("adagrad", "adam", "rmsprop", "momentum", "sgd"): + self._TestOptimizerSupportHelper(opt) + + def _GetOptimizer(self, opt): + if opt == "adagrad": + return adagrad.AdagradOptimizer(learning_rate=1e-2) + elif opt == "adam": + return adam.AdamOptimizer(learning_rate=1e-2) + elif opt == "rmsprop": + return rmsprop.RMSPropOptimizer(learning_rate=1e-2) + elif opt == "momentum": + return momentum.MomentumOptimizer(learning_rate=1e-2, momentum=0.9) + elif opt == "sgd": + return gradient_descent.GradientDescentOptimizer(learning_rate=1e-2) + else: + raise ValueError("Unsupported optimizer: %s" % opt) + + def _TestOptimizerSupportHelper(self, opt): + num_layers = 4 + num_units = 2 + batch_size = 8 + direction = CUDNN_RNN_UNIDIRECTION + dir_count = 1 + + with ops.Graph().as_default() as g: + kernel_initializer = init_ops.constant_initializer(0.) + bias_initializer = init_ops.constant_initializer(0.) + inputs = random_ops.random_uniform([ + num_layers * dir_count, batch_size, num_units], dtype=dtypes.float32) + + lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units, + direction=direction, + kernel_initializer=kernel_initializer, + bias_initializer=bias_initializer, + name="awesome_lstm") + outputs, _ = lstm(inputs) + loss = math_ops.reduce_sum(outputs) + optimizer = self._GetOptimizer(opt) + train_op = optimizer.minimize(loss) + + with self.test_session(use_gpu=True, graph=g) as sess: + sess.run(variables.global_variables_initializer()) + sess.run(train_op) + + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") def testSaveableGraphDeviceAssignment(self): num_layers = 4 num_units = 2 @@ -355,6 +410,60 @@ class CudnnRNNTestBasic(TensorFlowTestCase): saver.restore(sess, save_path) sess.run(outputs) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def testDifferentShapesEager(self): + # Checks that kernel caching does not cause sharing of temporary storage + # across different input shapes when executing eagerly. + with context.eager_mode(): + with ops.device("gpu:0"): + first_output, _ = cudnn_rnn.CudnnGRU(1, 100)( + array_ops.zeros([28, 100, 28])) + second_output, _ = cudnn_rnn.CudnnGRU(1, 100)( + array_ops.zeros([28, 100, 100])) + self.assertAllEqual([28, 100, 100], first_output.shape) + self.assertAllEqual([28, 100, 100], second_output.shape) + + def _LossFunc(): + first_output, _ = cudnn_rnn.CudnnGRU(1, 100)( + array_ops.zeros([28, 100, 28])) + second_output, _ = cudnn_rnn.CudnnGRU(1, 100)( + array_ops.zeros([28, 100, 100])) + return (math_ops.reduce_sum(first_output) + + math_ops.reduce_sum(second_output)) + + backprop.implicit_grad(_LossFunc)() + + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def testDifferentShapesGraph(self): + # Tests that a single kernel instance presented with multiple input shapes + # does not crash with graph execution. + with ops.device("gpu:0"): + layer = cudnn_rnn.CudnnGRU(1, 100) + layer(array_ops.zeros([28, 100, 100])) + + def _Cond(index, accumulation): + del accumulation # unused + return math_ops.less(index, 4) + + def _Body(index, accumulation): + layer_input = accumulation[:, :, 10 * (1 + index % 2):] + output, _ = layer(layer_input) + return index + 1, accumulation + output + + original_input = array_ops.zeros([28, 100, 100]) + _, accumulation = control_flow_ops.while_loop(_Cond, _Body, + [0, original_input]) + grad, = gradients.gradients( + math_ops.reduce_sum(accumulation), (original_input,)) + init_op = variables.global_variables_initializer() + with self.test_session() as sess: + sess.run(init_op) + accumulation_eval, grad_eval = sess.run((accumulation, grad)) + self.assertAllEqual([28, 100, 100], accumulation_eval.shape) + self.assertAllEqual([28, 100, 100], grad_eval.shape) + # TODO(jamesqin): Transform to parameterized test after it is included in the # TF open source codebase. diff --git a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py index 3a21914041e7b121cdc0c5c0cdbf848bb188cbfa..36fba917a8f56c26fd5b4c3468d1d980a8ba2ba5 100644 --- a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py +++ b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py @@ -176,8 +176,9 @@ class _CudnnRNN(base_layer.Layer): otherwise, it implies 'linear_input'. direction: the direction model that the model operates. Can be either 'unidirectional' or 'bidirectional' - dropout: dropout rate, a number between [0, 1]. Dropout is applied on - inputs of each layer. When set to 0, dropout is disabled. + dropout: dropout rate, a number between [0, 1]. Dropout is applied between + each layer (no dropout is applied for a model with a single layer). + When set to 0, dropout is disabled. seed: the op seed used for initializing dropout. See @{tf.set_random_seed} for behavior. dtype: tf.float16, tf.float32 or tf.float64 @@ -358,7 +359,7 @@ class _CudnnRNN(base_layer.Layer): # Create saveable in the outer scope of the cudnn subgraph, such that # alternative subgraph with platform-independent rnn cells can load the # checkpoints directly. - if not (self.built or vs.get_variable_scope().reuse): + if not (self.built or vs.get_variable_scope().reuse is True): self._create_saveable() self.built = True @@ -497,7 +498,7 @@ class _CudnnRNN(base_layer.Layer): input_mode=self.input_mode, direction=self.direction, scope=vs.get_variable_scope(), - name="%s_saveable" % self.trainable_variables[0].op.name) + name="%s_saveable" % self.trainable_variables[0].name.split(":")[0]) ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self._saveable) diff --git a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py index 45a3674beda55854ef303d4c7e7534cb9fe81ed3..e87162f0ee9cc4eed795555171f55a93639e83cf 100644 --- a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py +++ b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py @@ -1354,7 +1354,7 @@ class _CudnnRNN(object): params: the parameter buffer created for this model. is_training: whether this operation will be used in training or inference. Returns: - output: the output sequuence. + output: the output sequence. output_h: the final state for h. output_c: the final state for c. This is only relevant for LSTM. """ @@ -1472,7 +1472,7 @@ class CudnnLSTM(_CudnnRNN): params: the parameter buffer created for this model. is_training: whether this operation will be used in training or inference. Returns: - output: the output sequuence. + output: the output sequence. output_h: the final state for h. output_c: the final state for c. """ @@ -1542,7 +1542,7 @@ class _CudnnRNNNoInputC(_CudnnRNN): params: the parameter buffer created for this model. is_training: whether this operation will be used in training or inference. Returns: - output: the output sequuence. + output: the output sequence. output_h: the final state for h. """ return _cudnn_rnn_no_input_c( diff --git a/tensorflow/contrib/data/BUILD b/tensorflow/contrib/data/BUILD index 8ecc003348d70379ee48d050e63e93d0dd38efaa..0458199ff771bc45603106411550a39448e515b8 100644 --- a/tensorflow/contrib/data/BUILD +++ b/tensorflow/contrib/data/BUILD @@ -27,13 +27,13 @@ py_library( ) tf_custom_op_library( - name = "_prefetching_ops.so", - srcs = ["ops/prefetching_ops.cc"], - deps = ["//tensorflow/contrib/data/kernels:prefetching_kernels"], + name = "_dataset_ops.so", + srcs = ["ops/dataset_ops.cc"], + deps = ["//tensorflow/contrib/data/kernels:dataset_kernels"], ) tf_gen_op_libs( - op_lib_names = ["prefetching_ops"], + op_lib_names = ["dataset_ops"], ) filegroup( diff --git a/tensorflow/contrib/data/__init__.py b/tensorflow/contrib/data/__init__.py index 7a6cbe2aebc65aa47aae24b9d339ea1e19e8a579..fcdccdd26ca1824bf13f1fd0cfd80b20ca8a10c3 100644 --- a/tensorflow/contrib/data/__init__.py +++ b/tensorflow/contrib/data/__init__.py @@ -12,16 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""`tf.contrib.data.Dataset` API for input pipelines. +"""Experimental API for building input pipelines. + +This module contains experimental `Dataset` sources and transformations that can +be used in conjunction with the @{tf.data.Dataset} API. Note that the +`tf.contrib.data` API is not subject to the same backwards compatibility +guarantees as `tf.data`, but we will provide deprecation advice in advance of +removing existing functionality. See the @{$datasets$Importing Data} Programmer's Guide for an overview. -@@Dataset @@Counter -@@Iterator -@@TFRecordDataset -@@FixedLengthRecordDataset -@@TextLineDataset @@batch_and_drop_remainder @@dense_to_sparse_batch @@ -40,8 +41,6 @@ See the @{$datasets$Importing Data} Programmer's Guide for an overview. @@unbatch @@get_single_element - -@@parse_single_example """ from __future__ import absolute_import @@ -56,24 +55,18 @@ from tensorflow.contrib.data.python.ops.batching import map_and_batch from tensorflow.contrib.data.python.ops.batching import padded_batch_and_drop_remainder from tensorflow.contrib.data.python.ops.batching import unbatch from tensorflow.contrib.data.python.ops.counter import Counter -from tensorflow.contrib.data.python.ops.dataset_ops import Dataset -from tensorflow.contrib.data.python.ops.dataset_ops import get_single_element from tensorflow.contrib.data.python.ops.enumerate_ops import enumerate_dataset from tensorflow.contrib.data.python.ops.error_ops import ignore_errors +from tensorflow.contrib.data.python.ops.get_single_element import get_single_element from tensorflow.contrib.data.python.ops.grouping import group_by_window from tensorflow.contrib.data.python.ops.interleave_ops import parallel_interleave from tensorflow.contrib.data.python.ops.interleave_ops import sloppy_interleave from tensorflow.contrib.data.python.ops.iterator_ops import make_saveable_from_iterator -from tensorflow.contrib.data.python.ops.readers import FixedLengthRecordDataset from tensorflow.contrib.data.python.ops.readers import read_batch_features from tensorflow.contrib.data.python.ops.readers import SqlDataset -from tensorflow.contrib.data.python.ops.readers import TextLineDataset -from tensorflow.contrib.data.python.ops.readers import TFRecordDataset from tensorflow.contrib.data.python.ops.resampling import rejection_resample from tensorflow.contrib.data.python.ops.scan_ops import scan from tensorflow.contrib.data.python.ops.shuffle_ops import shuffle_and_repeat -from tensorflow.python.data.ops.iterator_ops import Iterator -from tensorflow.python.ops.parsing_ops import parse_single_example_v2 as parse_single_example # pylint: enable=unused-import from tensorflow.python.util.all_util import remove_undocumented diff --git a/tensorflow/contrib/data/kernels/BUILD b/tensorflow/contrib/data/kernels/BUILD index 4cb53741ebf8cd0db41b382c878bd2ccd1dcf7f1..56471911c5c0d1c1825955c67997b5bbc0786463 100644 --- a/tensorflow/contrib/data/kernels/BUILD +++ b/tensorflow/contrib/data/kernels/BUILD @@ -17,6 +17,28 @@ cc_library( alwayslink = 1, ) +cc_library( + name = "ignore_errors_dataset_op", + srcs = ["ignore_errors_dataset_op.cc"], + deps = [ + "//tensorflow/core:framework_headers_lib", + "//third_party/eigen3", + "@protobuf_archive//:protobuf_headers", + ], + alwayslink = 1, +) + +cc_library( + name = "dataset_kernels", + deps = [ + ":ignore_errors_dataset_op", + ":prefetching_kernels", + "//tensorflow/core:framework_headers_lib", + "//third_party/eigen3", + "@protobuf_archive//:protobuf_headers", + ], +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/core/kernels/data/ignore_errors_dataset_op.cc b/tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc similarity index 97% rename from tensorflow/core/kernels/data/ignore_errors_dataset_op.cc rename to tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc index beedc7c6777bfa5240bd380081359adfdc3e3aca..bb29df60e8f114aaa50f578c43e73874f72ab0a3 100644 --- a/tensorflow/core/kernels/data/ignore_errors_dataset_op.cc +++ b/tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc @@ -12,9 +12,9 @@ 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/dataset.h" #include "tensorflow/core/framework/partial_tensor_shape.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/kernels/data/dataset.h" #include "tensorflow/core/lib/random/random.h" namespace tensorflow { @@ -108,7 +108,7 @@ class IgnoreErrorsDatasetOp : public UnaryDatasetOpKernel { return Status::OK(); } - Status RestoreInternal(OpKernelContext* ctx, + Status RestoreInternal(IteratorContext* ctx, IteratorStateReader* reader) override { mutex_lock l(mu_); if (reader->Contains(full_name("input_impls_empty"))) diff --git a/tensorflow/contrib/data/kernels/prefetching_kernels.cc b/tensorflow/contrib/data/kernels/prefetching_kernels.cc index c9a3537c70c711290fb1111a1594e6dea3bc07a9..d3df14bdd03476e9ee4015b374512e5bb9893a63 100644 --- a/tensorflow/contrib/data/kernels/prefetching_kernels.cc +++ b/tensorflow/contrib/data/kernels/prefetching_kernels.cc @@ -83,11 +83,10 @@ class FunctionBufferingResource : public ResourceBase { return Status::OK(); } AttrValueMap attr_values = func_.attr(); - AttrValue v; - v.set_s(target_device_); - AddAttr("_target", v, &attr_values); - - return lib_->Instantiate(func_.name(), AttrSlice(&attr_values), &handle_); + FunctionLibraryRuntime::InstantiateOptions opts; + opts.target = target_device_; + return lib_->Instantiate(func_.name(), AttrSlice(&attr_values), opts, + &handle_); } // Returns true if we've got to the end of the sequence and exhausted the diff --git a/tensorflow/contrib/data/ops/prefetching_ops.cc b/tensorflow/contrib/data/ops/dataset_ops.cc similarity index 86% rename from tensorflow/contrib/data/ops/prefetching_ops.cc rename to tensorflow/contrib/data/ops/dataset_ops.cc index 23cb62b6f0dbfed15667dd00ae0039b33aa944d4..289ffa1d9c29092cdf434e86ed5553ff9644d43e 100644 --- a/tensorflow/contrib/data/ops/prefetching_ops.cc +++ b/tensorflow/contrib/data/ops/dataset_ops.cc @@ -17,6 +17,16 @@ limitations under the License. namespace tensorflow { +REGISTER_OP("IgnoreErrorsDataset") + .Input("input_dataset: variant") + .Output("handle: variant") + .Attr("output_types: list(type) >= 1") + .Attr("output_shapes: list(shape) >= 1") + .SetShapeFn(shape_inference::ScalarShape) + .Doc(R"doc( +Creates a dataset that contains the elements of `input_dataset` ignoring errors. +)doc"); + REGISTER_OP("FunctionBufferingResource") .Input("string_arg: string") .Input("target_device: string") diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD index cc511f72047841c70df64dbab4af44259740285f..e51d57cc896dc32d8e11912cd89f34a04a858c78 100644 --- a/tensorflow/contrib/data/python/kernel_tests/BUILD +++ b/tensorflow/contrib/data/python/kernel_tests/BUILD @@ -36,6 +36,7 @@ py_test( srcs = ["bucketing_test.py"], srcs_version = "PY2AND3", deps = [ + ":dataset_serialization_test", "//tensorflow/contrib/data/python/ops:dataset_ops", "//tensorflow/contrib/data/python/ops:transformation_ops", "//tensorflow/python:array_ops", @@ -51,37 +52,17 @@ py_test( ], ) -py_test( - name = "cache_dataset_op_test", - size = "small", - srcs = ["cache_dataset_op_test.py"], - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:constant_op", - "//tensorflow/python:dtypes", - "//tensorflow/python:errors", - "//tensorflow/python:variables", - "//tensorflow/python/data/ops:iterator_ops", - "//third_party/py/numpy", - ], -) - py_test( name = "concatenate_dataset_op_test", size = "small", srcs = ["concatenate_dataset_op_test.py"], srcs_version = "PY2AND3", deps = [ + ":dataset_serialization_test", "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:iterator_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:errors", - "//tensorflow/python:framework_ops", "//tensorflow/python:tensor_shape", - "//tensorflow/python:training", "//tensorflow/python/data/util:nest", "//third_party/py/numpy", ], @@ -89,7 +70,7 @@ py_test( py_test( name = "dataset_constructor_op_test", - size = "small", + size = "medium", srcs = ["dataset_constructor_op_test.py"], srcs_version = "PY2AND3", tags = [ @@ -118,7 +99,6 @@ py_test( py_library( name = "dataset_serialization_test", - testonly = 1, srcs = [ "dataset_serialization_test_base.py", ], @@ -128,6 +108,7 @@ py_library( "//tensorflow/python:client_testlib", "//tensorflow/python:errors", "//tensorflow/python:framework_ops", + "//tensorflow/python:lookup_ops", "//tensorflow/python:platform", "//tensorflow/python:sparse_tensor", "//tensorflow/python:training", @@ -159,7 +140,7 @@ py_test( tf_py_test( name = "flat_map_dataset_op_test", - size = "small", + size = "medium", srcs = ["flat_map_dataset_op_test.py"], additional_deps = [ ":dataset_serialization_test", @@ -183,11 +164,12 @@ tf_py_test( py_test( name = "interleave_dataset_op_test", - size = "small", + size = "medium", srcs = ["interleave_dataset_op_test.py"], srcs_version = "PY2AND3", tags = [ - "manual", # b/67958761 + "no_oss", + "no_pip", ], deps = [ ":dataset_serialization_test", @@ -208,73 +190,18 @@ py_test( ) tf_py_test( - name = "iterator_ops_cluster_test", - size = "small", - srcs = ["iterator_ops_cluster_test.py"], - additional_deps = [ - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/core:protos_all_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:dtypes", - "//tensorflow/python:errors", - "//tensorflow/python:framework_ops", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:function", - "//tensorflow/python:functional_ops", - "//tensorflow/python:session", - "//tensorflow/python/data/ops:iterator_ops", - ], - grpc_enabled = True, - tags = [ - "no_windows", - "oss_serial", - ], -) - -tf_py_test( - name = "iterator_ops_test", + name = "get_single_element_test", size = "small", - srcs = ["iterator_ops_test.py"], + srcs = ["get_single_element_test.py"], additional_deps = [ "//third_party/py/numpy", "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:readers", - "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", - "//tensorflow/python:dataset_ops_gen", "//tensorflow/python:dtypes", "//tensorflow/python:errors", - "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", - "//tensorflow/python:function", - "//tensorflow/python:functional_ops", - "//tensorflow/python:gradients", - "//tensorflow/python:io_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:parsing_ops", - "//tensorflow/python:script_ops", - "//tensorflow/python:session", - "//tensorflow/python:training", - "//tensorflow/python/data/ops:iterator_ops", - ], - grpc_enabled = True, -) - -py_test( - name = "list_files_dataset_op_test", - size = "small", - srcs = ["list_files_dataset_op_test.py"], - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:dtypes", - "//tensorflow/python:errors", - "//tensorflow/python:util", ], ) @@ -307,6 +234,7 @@ py_test( "//tensorflow/python:string_ops", "//tensorflow/python:util", "//tensorflow/python:variable_scope", + "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", ], ) @@ -330,8 +258,8 @@ py_test( srcs = ["range_dataset_op_test.py"], srcs_version = "PY2AND3", deps = [ + ":dataset_serialization_test", "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:iterator_ops", "//tensorflow/contrib/data/python/ops:transformation_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", @@ -342,11 +270,8 @@ py_test( "//tensorflow/python:framework_ops", "//tensorflow/python:io_ops", "//tensorflow/python:parsing_ops", - "//tensorflow/python:platform", "//tensorflow/python:tensor_shape", - "//tensorflow/python:training", "//tensorflow/python:variables", - "//tensorflow/python/data/ops:iterator_ops", ], ) @@ -428,14 +353,17 @@ py_test( ) py_test( - name = "shard_dataset_op_test", + name = "serialization_integration_test", size = "small", - srcs = ["shard_dataset_op_test.py"], + srcs = ["serialization_integration_test.py"], srcs_version = "PY2AND3", + tags = ["no_pip"], deps = [ - "//tensorflow/contrib/data/python/ops:dataset_ops", + "//tensorflow/contrib/data/python/ops:iterator_ops", "//tensorflow/python:client_testlib", - "//tensorflow/python:errors", + "//tensorflow/python:framework_ops", + "//tensorflow/python:training", + "//tensorflow/python/data/ops:dataset_ops", ], ) @@ -448,7 +376,6 @@ py_test( deps = [ ":dataset_serialization_test", "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:iterator_ops", "//tensorflow/contrib/data/python/ops:shuffle_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", @@ -456,8 +383,6 @@ py_test( "//tensorflow/python:dtypes", "//tensorflow/python:errors", "//tensorflow/python:framework_ops", - "//tensorflow/python:platform", - "//tensorflow/python:training", "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python/data/ops:iterator_ops", "//third_party/py/numpy", @@ -475,12 +400,13 @@ py_test( "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", "//tensorflow/python:errors", + "@org_sqlite//:python", ], ) py_test( name = "stats_dataset_ops_test", - size = "small", + size = "medium", srcs = ["stats_dataset_ops_test.py"], srcs_version = "PY2AND3", tags = ["no_pip"], @@ -539,7 +465,7 @@ py_test( "no_oss", # b/68785503 ], deps = [ - "//tensorflow/contrib/data/python/ops:prefetching_py", + "//tensorflow/contrib/data/python/ops:prefetching_ops", "//tensorflow/core:protos_all_py", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", diff --git a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py index 015f69c5673f185c53e61a5df2636333699ae203..71dc1c1172c9d515d4c85f85257c952135098329 100644 --- a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py @@ -23,283 +23,24 @@ import numpy as np from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base from tensorflow.contrib.data.python.ops import batching -from tensorflow.contrib.data.python.ops import dataset_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import sparse_tensor -from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import string_ops from tensorflow.python.platform import test -from tensorflow.python.util import compat class BatchDatasetTest(test.TestCase): - def testBatchDataset(self): - """Test an dataset that maps a TF function across its input elements.""" - # The pipeline is TensorSliceDataset -> MapDataset(square_3) -> - # RepeatDataset(count) -> BatchDataset(batch_size). - components = (np.arange(7), - np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], - np.array(37.0) * np.arange(7)) - - count = array_ops.placeholder(dtypes.int64, shape=[]) - batch_size = array_ops.placeholder(dtypes.int64, shape=[]) - - def _map_fn(x, y, z): - return math_ops.square(x), math_ops.square(y), math_ops.square(z) - - iterator = ( - dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) - .repeat(count).batch(batch_size).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([[None] + list(c.shape[1:]) for c in components], - [t.shape.as_list() for t in get_next]) - - with self.test_session() as sess: - # Batch of a finite input, where the batch_size divides the - # total number of elements. - sess.run(init_op, feed_dict={count: 28, batch_size: 14}) - num_batches = (28 * 7) // 14 - for i in range(num_batches): - result = sess.run(get_next) - for component, result_component in zip(components, result): - for j in range(14): - self.assertAllEqual(component[(i * 14 + j) % 7]**2, - result_component[j]) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Batch of a finite input, where the batch_size does not - # divide the total number of elements. - sess.run(init_op, feed_dict={count: 14, batch_size: 8}) - - # We expect (num_batches - 1) full-sized batches. - num_batches = int(math.ceil((14 * 7) / 8)) - for i in range(num_batches - 1): - result = sess.run(get_next) - for component, result_component in zip(components, result): - for j in range(8): - self.assertAllEqual(component[(i * 8 + j) % 7]**2, - result_component[j]) - result = sess.run(get_next) - for component, result_component in zip(components, result): - for j in range((14 * 7) % 8): - self.assertAllEqual(component[((num_batches - 1) * 8 + j) % 7]**2, - result_component[j]) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Batch of an empty input should fail straight away. - sess.run(init_op, feed_dict={count: 0, batch_size: 8}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Empty batch should be an initialization time error. - with self.assertRaises(errors.InvalidArgumentError): - sess.run(init_op, feed_dict={count: 14, batch_size: 0}) - def assertSparseValuesEqual(self, a, b): self.assertAllEqual(a.indices, b.indices) self.assertAllEqual(a.values, b.values) self.assertAllEqual(a.dense_shape, b.dense_shape) - def testBatchSparse(self): - - def _sparse(i): - return sparse_tensor.SparseTensorValue( - indices=[[0]], values=(i * [1]), dense_shape=[1]) - - iterator = dataset_ops.Dataset.range(10).map(_sparse).batch( - 5).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(2): - actual = sess.run(get_next) - expected = sparse_tensor.SparseTensorValue( - indices=[[0, 0], [1, 0], [2, 0], [3, 0], [4, 0]], - values=[i * 5, i * 5 + 1, i * 5 + 2, i * 5 + 3, i * 5 + 4], - dense_shape=[5, 1]) - self.assertTrue(sparse_tensor.is_sparse(actual)) - self.assertSparseValuesEqual(actual, expected) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testNestedBatchSparse(self): - - def _sparse(i): - return sparse_tensor.SparseTensorValue( - indices=[[0]], values=(i * [1]), dense_shape=[1]) - - iterator = dataset_ops.Dataset.range(10).map(_sparse).batch(5).batch( - 2).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - actual = sess.run(get_next) - expected = sparse_tensor.SparseTensorValue( - indices=[[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 3, 0], [0, 4, 0], - [1, 0, 0], [1, 1, 0], [1, 2, 0], [1, 3, 0], [1, 4, 0]], - values=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], - dense_shape=[2, 5, 1]) - self.assertTrue(sparse_tensor.is_sparse(actual)) - self.assertSparseValuesEqual(actual, expected) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testPaddedBatchDataset(self): - seq_lens = array_ops.placeholder(dtypes.int32, shape=[None]) - padded_shape = array_ops.placeholder(dtypes.int64, shape=[1]) - - iterator = ( - dataset_ops.Dataset.from_tensor_slices(seq_lens) - .map(lambda x: array_ops.fill([x], x)).padded_batch( - 4, padded_shapes=padded_shape).make_initializable_iterator()) - - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - # Test with random sequence lengths, and max padding. - random_seq_lens = np.random.randint(20, size=(32,)).astype(np.int32) - sess.run( - init_op, feed_dict={ - padded_shape: [-1], - seq_lens: random_seq_lens - }) - for i in range(8): - result = sess.run(get_next) - padded_len = np.max(result) - self.assertEqual((4, padded_len), result.shape) - for j in range(4): - seq_len = random_seq_lens[(i * 4) + j] - self.assertAllEqual(result[j, :seq_len], [seq_len] * seq_len) - self.assertAllEqual(result[j, seq_len:], [0] * (padded_len - seq_len)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test with random sequence lengths, and constant padding. - sess.run( - init_op, feed_dict={ - padded_shape: [25], - seq_lens: random_seq_lens - }) - for i in range(8): - result = sess.run(get_next) - self.assertEqual((4, 25), result.shape) - for j in range(4): - seq_len = random_seq_lens[(i * 4) + j] - self.assertAllEqual(result[j, :seq_len], [seq_len] * seq_len) - self.assertAllEqual(result[j, seq_len:], [0] * (25 - seq_len)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test correct handling of empty tensors. - sess.run(init_op, feed_dict={padded_shape: [-1], seq_lens: [0, 0, 0, 0]}) - result = sess.run(get_next) - self.assertAllEqual([[], [], [], []], result) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test error handling with constant sequence lengths, and - # too-short padding. - sess.run(init_op, feed_dict={padded_shape: [5], seq_lens: [6, 5, 5, 5]}) - with self.assertRaises(errors.DataLossError): - result = sess.run(get_next) - - def testPaddedBatchDatasetNonDefaultPadding(self): - seq_lens = array_ops.placeholder(dtypes.int32, shape=[None]) - padded_shape = array_ops.placeholder(dtypes.int64, shape=[1]) - - def fill_tuple(x): - filled = array_ops.fill([x], x) - return (filled, string_ops.as_string(filled)) - - iterator = ( - dataset_ops.Dataset.from_tensor_slices(seq_lens).map(fill_tuple) - .padded_batch( - 4, - padded_shapes=(padded_shape, padded_shape), - padding_values=(-1, "")).make_initializable_iterator()) - - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - # Test with random sequence lengths, and max padding. - random_seq_lens = np.random.randint(20, size=(32,)).astype(np.int32) - sess.run( - init_op, feed_dict={ - padded_shape: [-1], - seq_lens: random_seq_lens - }) - for i in range(8): - result = sess.run(get_next) - padded_len = np.max(result[0]) - self.assertEqual((4, padded_len), result[0].shape) - self.assertEqual((4, padded_len), result[1].shape) - for j in range(4): - seq_len = random_seq_lens[(i * 4) + j] - self.assertAllEqual(result[0][j, :seq_len], [seq_len] * seq_len) - self.assertAllEqual(result[0][j, seq_len:], - [-1] * (padded_len - seq_len)) - self.assertAllEqual(result[1][j, :seq_len], - [compat.as_bytes(str(seq_len))] * seq_len) - self.assertAllEqual(result[1][j, seq_len:], - [b""] * (padded_len - seq_len)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testPaddedBatchDatasetShapeSpecifications(self): - int_placeholder = array_ops.placeholder(dtypes.int32) - float_placeholder = array_ops.placeholder(dtypes.float32) - string_placeholder = array_ops.placeholder(dtypes.string) - input_dataset = dataset_ops.Dataset.from_tensors( - (int_placeholder, float_placeholder, string_placeholder)) - - # Test different ways of specifying the `padded_shapes` argument. - dynamic_padding_from_tensor_shapes = input_dataset.padded_batch( - 32, - padded_shapes=(tensor_shape.TensorShape([None]), - tensor_shape.TensorShape([None, None]), - tensor_shape.TensorShape([37]))) - dynamic_padding_from_lists = input_dataset.padded_batch( - 32, padded_shapes=([None], [None, None], [37])) - dynamic_padding_from_lists_with_minus_one = input_dataset.padded_batch( - 32, padded_shapes=([-1], [-1, -1], [37])) - dynamic_padding_from_tensors = input_dataset.padded_batch( - 32, - padded_shapes=(constant_op.constant([-1], dtype=dtypes.int64), - constant_op.constant([-1, -1], dtype=dtypes.int64), - constant_op.constant([37], dtype=dtypes.int64))) - - for dataset in [ - dynamic_padding_from_tensor_shapes, dynamic_padding_from_lists, - dynamic_padding_from_lists_with_minus_one, dynamic_padding_from_tensors - ]: - self.assertEqual([None, None], dataset.output_shapes[0].as_list()) - self.assertEqual([None, None, None], dataset.output_shapes[1].as_list()) - self.assertEqual([None, 37], dataset.output_shapes[2].as_list()) - - def testPaddedBatchSparseError(self): - - def _map_fn(i): - return sparse_tensor.SparseTensorValue( - indices=[[0, 0]], values=(i * [1]), dense_shape=[1, 1]), i - - with self.assertRaises(TypeError): - _ = dataset_ops.Dataset.range(10).map(_map_fn).padded_batch(10) - def testDenseToSparseBatchDataset(self): components = np.random.randint(12, size=(100,)).astype(np.int32) iterator = ( @@ -744,6 +485,23 @@ class BatchDatasetSerializationTest( lambda: self._build_dataset_dense_to_sparse(diff_comp), num_outputs) + def _sparse(self, i): + return sparse_tensor.SparseTensorValue( + indices=[[0]], values=(i * [1]), dense_shape=[1]) + + def _build_dataset_sparse(self, batch_size=5): + return dataset_ops.Dataset.range(10).map(self._sparse).batch(batch_size) + + def testSparseCore(self): + self.run_core_tests(self._build_dataset_sparse, + lambda: self._build_dataset_sparse(2), 2) + + def _build_dataset_nested_sparse(self): + return dataset_ops.Dataset.range(10).map(self._sparse).batch(5).batch(2) + + def testNestedSparseCore(self): + self.run_core_tests(self._build_dataset_nested_sparse, None, 1) + class PaddedBatchDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): diff --git a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py b/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py index 765ed53618958a8c49b26e416c57be28ea3bba73..f1b494e1a620992365ed75613b508e32f94b40a4 100644 --- a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py @@ -19,8 +19,9 @@ from __future__ import print_function import numpy as np -from tensorflow.contrib.data.python.ops import dataset_ops +from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base from tensorflow.contrib.data.python.ops import grouping +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -40,8 +41,7 @@ class GroupByWindowTest(test.TestCase): dataset_ops.Dataset.from_tensor_slices(components).map(lambda x: x * x) .apply( grouping.group_by_window(lambda x: x % 2, lambda _, xs: xs.batch(4), - 4)) - .make_initializable_iterator()) + 4)).make_initializable_iterator()) init_op = iterator.initializer get_next = iterator.get_next() @@ -52,7 +52,8 @@ class GroupByWindowTest(test.TestCase): while True: result = sess.run(get_next) self.assertTrue( - all(x % 2 == 0 for x in result) or all(x % 2 == 1) + all(x % 2 == 0 + for x in result) or all(x % 2 == 1) for x in result) counts.append(result.shape[0]) @@ -115,8 +116,8 @@ class GroupByWindowTest(test.TestCase): iterator = ( dataset_ops.Dataset.from_tensor_slices(components) .map(lambda x: (x, ops.convert_to_tensor([x * x]))).apply( - grouping.group_by_window(lambda x, _: x % 2, reduce_func, 32)) - .make_initializable_iterator()) + grouping.group_by_window(lambda x, _: x % 2, reduce_func, + 32)).make_initializable_iterator()) init_op = iterator.initializer get_next = iterator.get_next() @@ -135,7 +136,8 @@ class GroupByWindowTest(test.TestCase): window.padded_batch( 4, padded_shapes=tensor_shape.TensorShape([None])), window.padded_batch( - 4, padded_shapes=ops.convert_to_tensor([(key + 1) * 10])),)) + 4, padded_shapes=ops.convert_to_tensor([(key + 1) * 10])), + )) iterator = ( dataset_ops.Dataset.from_tensor_slices(components) @@ -160,6 +162,34 @@ class GroupByWindowTest(test.TestCase): self.assertEqual(len(components), sum(counts)) +class GroupByWindowSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): + + def _build_dataset(self, components): + return dataset_ops.Dataset.from_tensor_slices(components).repeat(-1).apply( + grouping.group_by_window(lambda x: x % 3, lambda _, xs: xs.batch(4), 4)) + + def testCoreGroupByWindow(self): + components = np.array( + [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 0, 0, 2, 2, 0, 0], dtype=np.int64) + self.verify_unused_iterator( + lambda: self._build_dataset(components), 12, verify_exhausted=False) + self.verify_init_before_restore( + lambda: self._build_dataset(components), 12, verify_exhausted=False) + self.verify_multiple_breaks( + lambda: self._build_dataset(components), 12, verify_exhausted=False) + self.verify_reset_restored_iterator( + lambda: self._build_dataset(components), 12, verify_exhausted=False) + self.verify_restore_in_empty_graph( + lambda: self._build_dataset(components), 12, verify_exhausted=False) + diff_components = np.array([0, 0, 0, 1, 1, 1], dtype=np.int64) + self.verify_restore_in_modified_graph( + lambda: self._build_dataset(components), + lambda: self._build_dataset(diff_components), + 12, + verify_exhausted=False) + + # NOTE(mrry): These tests are based on the tests in bucket_ops_test.py. # Currently, they use a constant batch size, though should be made to use a # different batch size per key. @@ -171,9 +201,10 @@ class BucketTest(test.TestCase): # dynamically and does not rely on static shape information about # the arguments. return dataset_ops.Dataset.zip( - (dataset_ops.Dataset.from_tensors(bucket), window.padded_batch( - 32, (tensor_shape.TensorShape([]), tensor_shape.TensorShape([None]), - tensor_shape.TensorShape([3]))))) + (dataset_ops.Dataset.from_tensors(bucket), + window.padded_batch( + 32, (tensor_shape.TensorShape([]), tensor_shape.TensorShape( + [None]), tensor_shape.TensorShape([3]))))) def testSingleBucket(self): @@ -278,12 +309,13 @@ class BucketTest(test.TestCase): def _dynamic_pad_fn(bucket, window, _): return dataset_ops.Dataset.zip( - (dataset_ops.Dataset.from_tensors(bucket), window.padded_batch( - 32, { - "x": tensor_shape.TensorShape([]), - "y": tensor_shape.TensorShape([None]), - "z": tensor_shape.TensorShape([3]) - }))) + (dataset_ops.Dataset.from_tensors(bucket), + window.padded_batch( + 32, { + "x": tensor_shape.TensorShape([]), + "y": tensor_shape.TensorShape([None]), + "z": tensor_shape.TensorShape([3]) + }))) input_dataset = ( dataset_ops.Dataset.from_tensor_slices(math_ops.range(128)).map(_map_fn) diff --git a/tensorflow/contrib/data/python/kernel_tests/cache_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/cache_dataset_op_test.py deleted file mode 100644 index 9818020680afb9d0f0197d272ec5339c6358db36..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/data/python/kernel_tests/cache_dataset_op_test.py +++ /dev/null @@ -1,300 +0,0 @@ -# Copyright 2017 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 the experimental input pipeline ops.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from os import path -import shutil -import tempfile - -import numpy as np - -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.python.data.ops import iterator_ops -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import variables -from tensorflow.python.platform import test - - -class FilesystemCacheDatasetTest(test.TestCase): - - def setUp(self): - self.tmp_dir = tempfile.mkdtemp() - self.cache_prefix = path.join(self.tmp_dir, "cache") - - def tearDown(self): - if self.tmp_dir: - shutil.rmtree(self.tmp_dir, ignore_errors=True) - - def testCacheDatasetPassthrough(self): - components = (np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]), - np.array([9.0, 10.0, 11.0, 12.0])) - count_placeholder = array_ops.placeholder_with_default( - constant_op.constant(5, dtypes.int64), shape=[]) - filename_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - - repeat_dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .repeat(count_placeholder)) - - cache_dataset = repeat_dataset.cache(filename_placeholder) - - self.assertEqual( - tuple([c.shape[1:] for c in components]), cache_dataset.output_shapes) - - # Create initialization ops for iterators without and with - # caching, respectively. - iterator = iterator_ops.Iterator.from_structure(cache_dataset.output_types, - cache_dataset.output_shapes) - init_fifo_op = iterator.make_initializer(repeat_dataset) - init_cache_op = iterator.make_initializer(cache_dataset) - - get_next = iterator.get_next() - - with self.test_session() as sess: - # First run without caching to collect the "ground truth". - sess.run(init_fifo_op) - elements = [] - for _ in range(20): - elements.append(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Assert that the cached dataset has the same elements as the - # "ground truth". - sess.run( - init_cache_op, feed_dict={filename_placeholder: self.cache_prefix}) - cached_elements = [] - for _ in range(20): - cached_elements.append(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - self.assertAllEqual(elements, cached_elements) - - # Re-initialize with an empty upstream (to throw errors.OutOfRangeError - # if we didn't use the cache). - sess.run( - init_cache_op, - feed_dict={ - count_placeholder: 0, - filename_placeholder: self.cache_prefix - }) - replayed_elements = [] - for _ in range(20): - replayed_elements.append(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - self.assertEqual(cached_elements, replayed_elements) - - # Re-initialize with an empty upstream and a missing cache file (should - # throw errors.OutOfRangeError immediately). - sess.run( - init_cache_op, - feed_dict={ - count_placeholder: 0, - filename_placeholder: self.cache_prefix + "nonsense" - }) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testConcurrentWriters(self): - components = (np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]), - np.array([9.0, 10.0, 11.0, 12.0])) - filename_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - - cache_dataset1 = (dataset_ops.Dataset.from_tensor_slices(components) - .cache(filename_placeholder)) - cache_dataset2 = (dataset_ops.Dataset.from_tensor_slices(components) - .cache(filename_placeholder)) - - iterator1 = cache_dataset1.make_initializable_iterator() - iterator2 = cache_dataset2.make_initializable_iterator() - init_cache_op1 = iterator1.initializer - init_cache_op2 = iterator2.initializer - - get_next1 = iterator1.get_next() - get_next2 = iterator2.get_next() - - with self.test_session() as sess: - sess.run( - init_cache_op1, feed_dict={filename_placeholder: self.cache_prefix}) - sess.run(get_next1) # this should succeed - - sess.run( - init_cache_op2, feed_dict={filename_placeholder: self.cache_prefix}) - with self.assertRaises(errors.AlreadyExistsError): - sess.run(get_next2) - - sess.run(get_next1) # this should continue to succeed - - def testConcurrentReaders(self): - components = (np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]), - np.array([9.0, 10.0, 11.0, 12.0])) - filename_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - - cache_dataset1 = (dataset_ops.Dataset.from_tensor_slices(components) - .cache(filename_placeholder)) - cache_dataset2 = (dataset_ops.Dataset.from_tensor_slices(components) - .cache(filename_placeholder)) - - iterator1 = cache_dataset1.make_initializable_iterator() - iterator2 = cache_dataset2.make_initializable_iterator() - init_cache_op1 = iterator1.initializer - init_cache_op2 = iterator2.initializer - - get_next1 = iterator1.get_next() - get_next2 = iterator2.get_next() - - with self.test_session() as sess: - sess.run( - init_cache_op1, feed_dict={filename_placeholder: self.cache_prefix}) - elements = [] - for _ in range(4): - elements.append(sess.run(get_next1)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next1) - - # Re-initialize - sess.run( - init_cache_op1, feed_dict={filename_placeholder: self.cache_prefix}) - sess.run( - init_cache_op2, feed_dict={filename_placeholder: self.cache_prefix}) - - # Reading concurrently should succeed. - elements_itr1 = [] - elements_itr2 = [] - elements_itr2.append(sess.run(get_next2)) - elements_itr1.append(sess.run(get_next1)) - elements_itr2.append(sess.run(get_next2)) - elements_itr1.append(sess.run(get_next1)) - # Intentionally reversing the order - elements_itr1.append(sess.run(get_next1)) - elements_itr2.append(sess.run(get_next2)) - elements_itr1.append(sess.run(get_next1)) - elements_itr2.append(sess.run(get_next2)) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next2) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next1) - - self.assertAllEqual(elements, elements_itr1) - self.assertAllEqual(elements, elements_itr2) - - -class MemoryCacheDatasetTest(test.TestCase): - - def testCacheDatasetPassthrough(self): - repeat_count = variables.Variable(constant_op.constant(10, dtypes.int64)) - dataset = dataset_ops.Dataset.range(3).flat_map( - lambda x: dataset_ops.Dataset.from_tensors(x).repeat(repeat_count)) - - cached_dataset = dataset.cache().repeat(2) - uncached_dataset = dataset.repeat(2) - - # Needs to be initializable to capture the variable. - cached_iterator = cached_dataset.make_initializable_iterator() - cached_next = cached_iterator.get_next() - uncached_iterator = uncached_dataset.make_initializable_iterator() - uncached_next = uncached_iterator.get_next() - - with self.test_session() as sess: - - sess.run(repeat_count.initializer) - sess.run(cached_iterator.initializer) - sess.run(uncached_iterator.initializer) - - for i in range(3): - for _ in range(10): - self.assertEqual(sess.run(cached_next), i) - self.assertEqual(sess.run(uncached_next), i) - - sess.run(repeat_count.assign(0)) - - # The uncached iterator should now be empty. - with self.assertRaises(errors.OutOfRangeError): - sess.run(uncached_next) - - # The cached iterator replays from cache. - for i in range(3): - for _ in range(10): - self.assertEqual(sess.run(cached_next), i) - - # The cached iterator should now be empty. - with self.assertRaises(errors.OutOfRangeError): - sess.run(cached_next) - - def testEmptyCacheReading(self): - components = (np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]), - np.array([9.0, 10.0, 11.0, 12.0])) - count_placeholder = array_ops.placeholder_with_default( - constant_op.constant(5, dtypes.int64), shape=[]) - - repeat_dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .repeat(count_placeholder)) - - cache_dataset = repeat_dataset.cache() - - # Create initialization ops for iterators without and with - # caching, respectively. - iterator = cache_dataset.make_initializable_iterator() - init_cache_op = iterator.initializer - - get_next = iterator.get_next() - - with self.test_session() as sess: - # Initialize with an empty upstream and a missing cache file (should - # throw errors.OutOfRangeError immediately). - sess.run(init_cache_op, feed_dict={count_placeholder: 0}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testConcurrentReaders(self): - count_placeholder = array_ops.placeholder_with_default( - constant_op.constant(5, dtypes.int64), shape=[]) - dataset = dataset_ops.Dataset.range(count_placeholder).cache() - d1 = dataset.map(lambda x: x + 1) - d2 = dataset.map(lambda x: x + 6) - - i1 = d1.make_initializable_iterator() - i2 = d2.make_initializable_iterator() - - with self.test_session() as sess: - sess.run(i1.initializer) - - self.assertEqual(1, sess.run(i1.get_next())) - self.assertEqual(2, sess.run(i1.get_next())) - self.assertEqual(3, sess.run(i1.get_next())) - - sess.run(i2.initializer, feed_dict={count_placeholder: 3}) - - self.assertEqual(6, sess.run(i2.get_next())) - self.assertEqual(7, sess.run(i2.get_next())) - self.assertEqual(4, sess.run(i1.get_next())) # interleave execution - self.assertEqual([8, 5], sess.run([i2.get_next(), i1.get_next()])) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(i1.get_next()) - with self.assertRaises(errors.OutOfRangeError): - sess.run(i2.get_next()) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/concatenate_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/concatenate_dataset_op_test.py index 870352209a08e6bc08bcca227ba455ad1851e8bf..17f2980157ddd0350dafd1d745cbb9b64e65f7c5 100644 --- a/tensorflow/contrib/data/python/kernel_tests/concatenate_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/concatenate_dataset_op_test.py @@ -17,255 +17,32 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os import numpy as np -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.contrib.data.python.ops import iterator_ops -from tensorflow.python.data.util import nest -from tensorflow.python.framework import errors -from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape +from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.platform import test -from tensorflow.python.training import saver as saver_lib -class ConcatenateDatasetTest(test.TestCase): +class ConcatenateDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): - def testConcatenateDataset(self): - input_components = ( - np.tile(np.array([[1], [2], [3], [4]]), 20), - np.tile(np.array([[12], [13], [14], [15]]), 15), - np.array([37.0, 38.0, 39.0, 40.0])) - to_concatenate_components = ( - np.tile(np.array([[1], [2], [3], [4], [5]]), 20), - np.tile(np.array([[12], [13], [14], [15], [16]]), 15), - np.array([37.0, 38.0, 39.0, 40.0, 41.0])) - - input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components) - dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices( - to_concatenate_components) - concatenated = input_dataset.concatenate(dataset_to_concatenate) - self.assertEqual(concatenated.output_shapes, (tensor_shape.TensorShape( - [20]), tensor_shape.TensorShape([15]), tensor_shape.TensorShape([]))) - - iterator = concatenated.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(9): - result = sess.run(get_next) - if i < 4: - for component, result_component in zip(input_components, result): - self.assertAllEqual(component[i], result_component) - else: - for component, result_component in zip(to_concatenate_components, - result): - self.assertAllEqual(component[i - 4], result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testConcatenateDatasetDifferentShape(self): - input_components = ( - np.tile(np.array([[1], [2], [3], [4]]), 20), - np.tile(np.array([[12], [13], [14], [15]]), 4)) - to_concatenate_components = ( - np.tile(np.array([[1], [2], [3], [4], [5]]), 20), - np.tile(np.array([[12], [13], [14], [15], [16]]), 15)) - - input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components) - dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices( - to_concatenate_components) - concatenated = input_dataset.concatenate(dataset_to_concatenate) - self.assertEqual( - [ts.as_list() - for ts in nest.flatten(concatenated.output_shapes)], [[20], [None]]) - - iterator = concatenated.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(9): - result = sess.run(get_next) - if i < 4: - for component, result_component in zip(input_components, result): - self.assertAllEqual(component[i], result_component) - else: - for component, result_component in zip(to_concatenate_components, - result): - self.assertAllEqual(component[i - 4], result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testConcatenateDatasetDifferentStructure(self): - input_components = ( - np.tile(np.array([[1], [2], [3], [4]]), 5), - np.tile(np.array([[12], [13], [14], [15]]), 4)) - to_concatenate_components = ( - np.tile(np.array([[1], [2], [3], [4], [5]]), 20), - np.tile(np.array([[12], [13], [14], [15], [16]]), 15), - np.array([37.0, 38.0, 39.0, 40.0, 41.0])) - - input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components) - dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices( - to_concatenate_components) - - with self.assertRaisesRegexp(ValueError, - "don't have the same number of elements"): - input_dataset.concatenate(dataset_to_concatenate) - - def testConcatenateDatasetDifferentType(self): - input_components = ( - np.tile(np.array([[1], [2], [3], [4]]), 5), - np.tile(np.array([[12], [13], [14], [15]]), 4)) - to_concatenate_components = ( - np.tile(np.array([[1.0], [2.0], [3.0], [4.0]]), 5), - np.tile(np.array([[12], [13], [14], [15]]), 15)) - - input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components) - dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices( - to_concatenate_components) - - with self.assertRaisesRegexp(TypeError, "have different types"): - input_dataset.concatenate(dataset_to_concatenate) - - def _iterator_checkpoint_prefix(self): - return os.path.join(self.get_temp_dir(), "iterator") - - def _build_graph(self, input_components, to_concatenate_components): - input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components) - dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices( - to_concatenate_components) - iterator = input_dataset.concatenate( - dataset_to_concatenate).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - saveable = iterator_ops.make_saveable_from_iterator(iterator) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) - # TODO(shivaniagrawal) : non-intuitive way, add support in mata_graph - for t in nest.flatten(get_next): - ops.add_to_collection("get_next", t) - return init_op, get_next - - def _testSaveRestoreUtility(self, start, break_range, stop): - path = self._iterator_checkpoint_prefix() - step = 0 - meta_filename = path + "-%d.meta" % step - - input_components = (np.tile(np.array([[1], [2], [3], [4]]), 20), np.tile( - np.array([[12], [13], [14], [15]]), 4)) - to_concatenate_components = (np.tile( - np.array([[5], [6], [7], [8], [9]]), 20), np.tile( - np.array([[16], [17], [18], [19], [20]]), 15)) - - with ops.Graph().as_default() as g: - init_op, get_next = self._build_graph(input_components, - to_concatenate_components) - saver = saver_lib.Saver() - with self.test_session(graph=g) as sess: - sess.run(init_op) - for i in range(start, break_range): - result = sess.run(get_next) - if i < 4: - for component, result_component in zip(input_components, result): - self.assertAllEqual(component[i], result_component) - else: - for component, result_component in zip(to_concatenate_components, - result): - self.assertAllEqual(component[i - 4], result_component) - saver.save(sess, path, step) - - with ops.Graph().as_default() as g: - saver = saver_lib.import_meta_graph(meta_filename) - with self.test_session(graph=g) as sess: - get_next = nest.pack_sequence_as(("a", "b"), - ops.get_collection("get_next")) - saver.restore(sess, saver_lib.latest_checkpoint(self.get_temp_dir())) - for i in range(break_range, stop): - result = sess.run(get_next) - if i < 4: - for component, result_component in zip(input_components, result): - self.assertAllEqual(component[i], result_component) - else: - for component, result_component in zip(to_concatenate_components, - result): - self.assertAllEqual(component[i - 4], result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testRestoreAtFirstDataset(self): - start = 0 - stop = 9 - break_range = 3 - self._testSaveRestoreUtility(start, break_range, stop) - - def testRestoreAtSecondDataset(self): - start = 0 - stop = 9 - break_range = 6 - self._testSaveRestoreUtility(start, break_range, stop) - - def testRestoreAtBetweenDatasets(self): - start = 0 - stop = 9 - break_range = 4 - self._testSaveRestoreUtility(start, break_range, stop) - - def testRestoreExhaustedIterator(self): - start = 0 - stop = 9 - break_range = 9 - self._testSaveRestoreUtility(start, break_range, stop) - - def testRestoreInModifiedGraph(self): - start = 0 - stop = 9 - break_range = 6 - path = self._iterator_checkpoint_prefix() - step = 0 - - input_components = (np.tile(np.array([[1], [2], [3], [4]]), 20), np.tile( - np.array([[12], [13], [14], [15]]), 4)) + def _build_concatenate_dataset(self, var_array): + input_components = (np.tile(np.array([[1], [2], [3], [4]]), 20), + np.tile(np.array([[12], [13], [14], [15]]), 4)) to_concatenate_components = (np.tile( - np.array([[5], [6], [7], [8], [9]]), 20), np.tile( - np.array([[16], [17], [18], [19], [20]]), 15)) - - with ops.Graph().as_default() as g: - init_op, get_next = self._build_graph(input_components, - to_concatenate_components) - saver = saver_lib.Saver(allow_empty=True) - with self.test_session(graph=g) as sess: - sess.run(init_op) - for i in range(start, break_range): - result = sess.run(get_next) - if i < 4: - for component, result_component in zip(input_components, result): - self.assertAllEqual(component[i], result_component) - else: - for component, result_component in zip(to_concatenate_components, - result): - self.assertAllEqual(component[i - 4], result_component) - saver.save(sess, path, step) - - new_to_concatenate_components = (np.array([[5], [6], [7], [8], [9]]), - np.array([[16], [17], [18], [19], [20]])) - with ops.Graph().as_default() as g: - init_op, get_next = self._build_graph(input_components, - new_to_concatenate_components) - saver = saver_lib.Saver() - with self.test_session(graph=g) as sess: - saver.restore(sess, saver_lib.latest_checkpoint(self.get_temp_dir())) - for i in range(break_range, stop): - result = sess.run(get_next) - for component, result_component in zip(to_concatenate_components, - result): - self.assertAllEqual(component[i - 4], result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + np.array([[5], [6], [7], [8], [9]]), 20), var_array) + + return dataset_ops.Dataset.from_tensor_slices(input_components).concatenate( + dataset_ops.Dataset.from_tensor_slices(to_concatenate_components)) + + def testConcatenateCore(self): + num_outputs = 9 + array = np.tile(np.array([[16], [17], [18], [19], [20]]), 15) + diff_array = np.array([[1], [2], [3], [4], [5]]) + self.run_core_tests(lambda: self._build_concatenate_dataset(array), + lambda: self._build_concatenate_dataset(diff_array), + num_outputs) if __name__ == "__main__": diff --git a/tensorflow/contrib/data/python/kernel_tests/dataset_constructor_op_test.py b/tensorflow/contrib/data/python/kernel_tests/dataset_constructor_op_test.py index 55a1d3b95b212466b262ad3c26f1efd7ed0e067e..a842502cc6fe3605dde0be5f50cf46e3e37d7ed4 100644 --- a/tensorflow/contrib/data/python/kernel_tests/dataset_constructor_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/dataset_constructor_op_test.py @@ -17,712 +17,20 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import threading - import numpy as np from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base from tensorflow.contrib.data.python.ops import batching -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.core.protobuf import config_pb2 -from tensorflow.python.client import session +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.util import nest 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.framework import tensor_shape from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import resource_variable_ops from tensorflow.python.platform import test class DatasetConstructorTest(test.TestCase): - def testFromTensors(self): - """Test an dataset that represents a single tuple of tensors.""" - components = (np.array(1), np.array([1, 2, 3]), np.array(37.0)) - - iterator = (dataset_ops.Dataset.from_tensors(components) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - sess.run(init_op) - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertAllEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def assertSparseValuesEqual(self, a, b): - self.assertAllEqual(a.indices, b.indices) - self.assertAllEqual(a.values, b.values) - self.assertAllEqual(a.dense_shape, b.dense_shape) - - def testFromTensorsSparse(self): - """Test an dataset that represents a single tuple of tensors.""" - components = (sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([0]), - dense_shape=np.array([1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0], [1, 1]]), - values=np.array([-1, 1]), - dense_shape=np.array([2, 2]))) - - iterator = ( - dataset_ops.Dataset.from_tensors(components) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual( - [tensor_shape.TensorShape(c.dense_shape) for c in components], - [shape for shape in iterator.output_shapes]) - - with self.test_session() as sess: - sess.run(init_op) - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertSparseValuesEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromTensorsMixed(self): - """Test an dataset that represents a single tuple of tensors.""" - components = (np.array(1), np.array([1, 2, 3]), np.array(37.0), - sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([0]), - dense_shape=np.array([1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0], [1, 1]]), - values=np.array([-1, 1]), - dense_shape=np.array([2, 2]))) - - iterator = ( - dataset_ops.Dataset.from_tensors(components) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([ - tensor_shape.TensorShape(c.dense_shape) - if sparse_tensor.is_sparse(c) else c.shape for c in components - ], [shape for shape in iterator.output_shapes]) - - with self.test_session() as sess: - sess.run(init_op) - results = sess.run(get_next) - for component, result_component in zip(components, results): - if sparse_tensor.is_sparse(component): - self.assertSparseValuesEqual(component, result_component) - else: - self.assertAllEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromTensorSlices(self): - """Test an dataset that represents the slices from a tuple of tensors.""" - components = ( - np.tile(np.array([[1], [2], [3], [4]]), 20), np.tile( - np.array([[12], [13], [14], [15]]), 22), - np.array([37.0, 38.0, 39.0, 40.0]) - ) - - iterator = (dataset_ops.Dataset.from_tensor_slices(components) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape[1:] for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - sess.run(init_op) - for i in range(4): - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertAllEqual(component[i], result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromTensorSlicesSparse(self): - """Test an dataset that represents the slices from a tuple of tensors.""" - components = (sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0], [1, 0], [2, 0]]), - values=np.array([0, 0, 0]), - dense_shape=np.array([3, 1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0], [1, 1], [2, 2]]), - values=np.array([1, 2, 3]), - dense_shape=np.array([3, 3]))) - - iterator = ( - dataset_ops.Dataset.from_tensor_slices(components) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual( - [tensor_shape.TensorShape(c.dense_shape[1:]) for c in components], - [shape for shape in iterator.output_shapes]) - - with self.test_session() as sess: - sess.run(init_op) - expected = [ - (sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([0]), - dense_shape=np.array([1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([1]), - dense_shape=np.array([3]))), - (sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([0]), - dense_shape=np.array([1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[1]]), - values=np.array([2]), - dense_shape=np.array([3]))), - (sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([0]), - dense_shape=np.array([1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[2]]), - values=np.array([3]), - dense_shape=np.array([3]))), - ] - for i in range(3): - results = sess.run(get_next) - for component, result_component in zip(expected[i], results): - self.assertSparseValuesEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromTensorSlicesMixed(self): - """Test an dataset that represents the slices from a tuple of tensors.""" - components = (np.tile(np.array([[1], [2], [3]]), 20), - np.tile(np.array([[12], [13], [14]]), 22), - np.array([37.0, 38.0, 39.0]), - sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0], [1, 0], [2, 0]]), - values=np.array([0, 0, 0]), - dense_shape=np.array([3, 1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0], [1, 1], [2, 2]]), - values=np.array([1, 2, 3]), - dense_shape=np.array([3, 3]))) - - iterator = ( - dataset_ops.Dataset.from_tensor_slices(components) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([ - tensor_shape.TensorShape(c.dense_shape[1:]) - if sparse_tensor.is_sparse(c) else c.shape[1:] for c in components - ], [shape for shape in iterator.output_shapes]) - - with self.test_session() as sess: - sess.run(init_op) - expected = [ - (sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([0]), - dense_shape=np.array([1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([1]), - dense_shape=np.array([3]))), - (sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([0]), - dense_shape=np.array([1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[1]]), - values=np.array([2]), - dense_shape=np.array([3]))), - (sparse_tensor.SparseTensorValue( - indices=np.array([[0]]), - values=np.array([0]), - dense_shape=np.array([1])), - sparse_tensor.SparseTensorValue( - indices=np.array([[2]]), - values=np.array([3]), - dense_shape=np.array([3]))), - ] - for i in range(3): - results = sess.run(get_next) - for component, result_component in zip( - (zip(*components[:3])[i] + expected[i]), results): - if sparse_tensor.is_sparse(component): - self.assertSparseValuesEqual(component, result_component) - else: - self.assertAllEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromTensorSlicesWithDict(self): - components = {"foo": [1, 2, 3], "bar": [[4.0], [5.0], [6.0]]} - iterator = (dataset_ops.Dataset.from_tensor_slices(components) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual(dtypes.int32, iterator.output_types["foo"]) - self.assertEqual(dtypes.float32, iterator.output_types["bar"]) - self.assertEqual((), iterator.output_shapes["foo"]) - self.assertEqual((1,), iterator.output_shapes["bar"]) - - with self.test_session() as sess: - sess.run(init_op) - for i in range(3): - results = sess.run(get_next) - self.assertEqual(components["foo"][i], results["foo"]) - self.assertEqual(components["bar"][i], results["bar"]) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromSparseTensorSlices(self): - """Test a dataset based on slices of a `tf.SparseTensor`.""" - st = array_ops.sparse_placeholder(dtypes.float64) - iterator = (dataset_ops.Dataset.from_sparse_tensor_slices(st) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = sparse_tensor.SparseTensor(*iterator.get_next()) - - with self.test_session() as sess: - slices = [[1., 2., 3.], [1.], [1.], [1., 2.], [], [1., 2.], [], [], []] - - # Test with sparse tensor in the appropriate order. - indices = np.array( - [[i, j] for i in range(len(slices)) for j in range(len(slices[i]))]) - values = np.array([val for s in slices for val in s]) - dense_shape = np.array([len(slices), max(len(s) for s in slices) + 1]) - sparse_feed = sparse_tensor.SparseTensorValue(indices, values, - dense_shape) - sess.run(init_op, feed_dict={st: sparse_feed}) - for i, s in enumerate(slices): - results = sess.run(get_next) - self.assertAllEqual(s, results.values) - expected_indices = np.array( - [[j] for j in range(len(slices[i]))]).reshape([-1, 1]) - self.assertAllEqual(expected_indices, results.indices) - self.assertAllEqual(dense_shape[1:], results.dense_shape) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test with sparse tensor in the reverse order, which is not - # currently supported. - reverse_order_indices = indices[::-1, :] - reverse_order_values = values[::-1] - sparse_feed = sparse_tensor.SparseTensorValue( - reverse_order_indices, reverse_order_values, dense_shape) - with self.assertRaises(errors.UnimplementedError): - sess.run(init_op, feed_dict={st: sparse_feed}) - - # Test with an empty sparse tensor. - empty_indices = np.empty((0, 4), dtype=np.int64) - empty_values = np.empty((0,), dtype=np.float64) - empty_dense_shape = [0, 4, 37, 9] - sparse_feed = sparse_tensor.SparseTensorValue(empty_indices, empty_values, - empty_dense_shape) - sess.run(init_op, feed_dict={st: sparse_feed}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # pylint: disable=g-long-lambda,unnecessary-lambda - def testNestedStructure(self): - components = (np.array([1, 2, 3]), (np.array([4., 5.]), np.array([6., 7.])), - np.array([8, 9, 10])) - - dataset = dataset_ops.Dataset.from_tensors(components) - self.assertEquals((dtypes.int64, (dtypes.float64, dtypes.float64), - dtypes.int64), dataset.output_types) - self.assertEquals(([3], ([2], [2]), [3]), dataset.output_shapes) - - dataset = dataset.shuffle(10, 10) - self.assertEquals((dtypes.int64, (dtypes.float64, dtypes.float64), - dtypes.int64), dataset.output_types) - self.assertEquals(([3], ([2], [2]), [3]), dataset.output_shapes) - - dataset = dataset.repeat(-1) - self.assertEquals((dtypes.int64, (dtypes.float64, dtypes.float64), - dtypes.int64), dataset.output_types) - self.assertEquals(([3], ([2], [2]), [3]), dataset.output_shapes) - - dataset = dataset.filter(lambda x, y, z: True) - self.assertEquals((dtypes.int64, (dtypes.float64, dtypes.float64), - dtypes.int64), dataset.output_types) - self.assertEquals(([3], ([2], [2]), [3]), dataset.output_shapes) - - dataset = dataset.take(5) - self.assertEquals((dtypes.int64, (dtypes.float64, dtypes.float64), - dtypes.int64), dataset.output_types) - self.assertEquals(([3], ([2], [2]), [3]), dataset.output_shapes) - - dataset = dataset.map(lambda x, y, z: ((x, z), (y[0], y[1]))) - self.assertEquals(((dtypes.int64, dtypes.int64), - (dtypes.float64, dtypes.float64)), dataset.output_types) - self.assertEquals((([3], [3]), ([2], [2])), dataset.output_shapes) - - dataset = dataset.flat_map( - lambda x, y: dataset_ops.Dataset.from_tensors(((x[0], x[1]), - (y[0], y[1]))) - ) - self.assertEquals(((dtypes.int64, dtypes.int64), - (dtypes.float64, dtypes.float64)), dataset.output_types) - self.assertEquals((([3], [3]), ([2], [2])), dataset.output_shapes) - - dataset = dataset.batch(32) - self.assertEquals(((dtypes.int64, dtypes.int64), - (dtypes.float64, dtypes.float64)), dataset.output_types) - self.assertEquals((([None, 3], [None, 3]), ([None, 2], [None, 2])), - nest.pack_sequence_as(dataset.output_shapes, [ - s.as_list() - for s in nest.flatten(dataset.output_shapes) - ])) - - iterator = dataset.make_one_shot_iterator() - (w, x), (y, z) = iterator.get_next() - self.assertEquals(dtypes.int64, w.dtype) - self.assertEquals(dtypes.int64, x.dtype) - self.assertEquals(dtypes.float64, y.dtype) - self.assertEquals(dtypes.float64, z.dtype) - self.assertEquals([None, 3], w.shape.as_list()) - self.assertEquals([None, 3], x.shape.as_list()) - self.assertEquals([None, 2], y.shape.as_list()) - self.assertEquals([None, 2], z.shape.as_list()) - - iterator = dataset.make_initializable_iterator() - (w, x), (y, z) = iterator.get_next() - self.assertEquals(dtypes.int64, w.dtype) - self.assertEquals(dtypes.int64, x.dtype) - self.assertEquals(dtypes.float64, y.dtype) - self.assertEquals(dtypes.float64, z.dtype) - self.assertEquals([None, 3], w.shape.as_list()) - self.assertEquals([None, 3], x.shape.as_list()) - self.assertEquals([None, 2], y.shape.as_list()) - self.assertEquals([None, 2], z.shape.as_list()) - - # Define a separate set of components with matching leading - # dimension for the from-slices constructor. - components_for_slices = (np.array([1, 2, 3]), (np.array( - [4., 5., 6.]), np.array([7., 8., 9.])), np.array([10, 11, 12])) - - dataset = dataset_ops.Dataset.from_tensor_slices(components_for_slices) - self.assertEquals((dtypes.int64, (dtypes.float64, dtypes.float64), - dtypes.int64), dataset.output_types) - self.assertEquals(([], ([], []), []), dataset.output_shapes) - - def testNestedDict(self): - components = {"a": {"aa": 1, "ab": [2.0, 2.0]}, "b": [3, 3, 3]} - dataset = dataset_ops.Dataset.from_tensors(components) - self.assertEquals(dtypes.int32, dataset.output_types["a"]["aa"]) - self.assertEquals(dtypes.float32, dataset.output_types["a"]["ab"]) - self.assertEquals(dtypes.int32, dataset.output_types["b"]) - self.assertEquals([], dataset.output_shapes["a"]["aa"]) - self.assertEquals([2], dataset.output_shapes["a"]["ab"]) - self.assertEquals([3], dataset.output_shapes["b"]) - - def testNonSequenceNestedStructure(self): - components = np.array([1, 2, 3]) - - dataset = dataset_ops.Dataset.from_tensors(components) - self.assertEquals(dtypes.int64, dataset.output_types) - self.assertEquals([3], dataset.output_shapes) - - dataset = dataset.filter( - lambda x: math_ops.reduce_all(math_ops.equal(x, components))) - self.assertEquals(dtypes.int64, dataset.output_types) - self.assertEquals([3], dataset.output_shapes) - - dataset = dataset.map(lambda x: array_ops.stack([x, x])) - self.assertEquals(dtypes.int64, dataset.output_types) - self.assertEquals([2, 3], dataset.output_shapes) - - dataset = dataset.flat_map( - lambda x: dataset_ops.Dataset.from_tensor_slices(x)) - self.assertEquals(dtypes.int64, dataset.output_types) - self.assertEquals([3], dataset.output_shapes) - - iterator = dataset.make_one_shot_iterator() - get_next = iterator.get_next() - self.assertEquals(dtypes.int64, get_next.dtype) - self.assertEquals([3], get_next.shape) - - def _testFromGenerator(self, generator, elem_sequence, num_repeats): - iterator = ( - dataset_ops.Dataset.from_generator(generator, output_types=dtypes.int64) - .repeat(num_repeats) - .prefetch(5) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - for _ in range(2): # Run twice to test reinitialization. - sess.run(init_op) - for _ in range(num_repeats): - for elem in elem_sequence: - self.assertAllEqual(elem, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def _testFromGeneratorOneShot(self, generator, elem_sequence, num_repeats): - iterator = ( - dataset_ops.Dataset.from_generator(generator, output_types=dtypes.int64) - .repeat(num_repeats) - .prefetch(5) - .make_one_shot_iterator()) - get_next = iterator.get_next() - - with self.test_session() as sess: - for _ in range(num_repeats): - for elem in elem_sequence: - self.assertAllEqual(elem, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromGeneratorUsingFunction(self): - def generator(): - for i in range(1, 100): - yield [i] * i - elem_sequence = list(generator()) - self._testFromGenerator(generator, elem_sequence, 1) - self._testFromGenerator(generator, elem_sequence, 5) - self._testFromGeneratorOneShot(generator, elem_sequence, 1) - self._testFromGeneratorOneShot(generator, elem_sequence, 5) - - def testFromGeneratorUsingList(self): - generator = lambda: [[i] * i for i in range(1, 100)] - elem_sequence = list(generator()) - self._testFromGenerator(generator, elem_sequence, 1) - self._testFromGenerator(generator, elem_sequence, 5) - - def testFromGeneratorUsingNdarray(self): - generator = lambda: np.arange(100, dtype=np.int64) - elem_sequence = list(generator()) - self._testFromGenerator(generator, elem_sequence, 1) - self._testFromGenerator(generator, elem_sequence, 5) - - def testFromGeneratorUsingGeneratorExpression(self): - # NOTE(mrry): Generator *expressions* are not repeatable (or in - # general reusable), because they eagerly evaluate the `for` - # expression as `iter(range(1, 100))` and discard the means of - # reconstructing `range(1, 100)`. Wrapping the generator - # expression in a `lambda` makes it repeatable. - generator = lambda: ([i] * i for i in range(1, 100)) - elem_sequence = list(generator()) - self._testFromGenerator(generator, elem_sequence, 1) - self._testFromGenerator(generator, elem_sequence, 5) - - def testFromMultipleConcurrentGenerators(self): - num_inner_repeats = 5 - num_outer_repeats = 100 - - def generator(): - for i in range(1, 10): - yield ([i] * i, [i, i ** 2, i ** 3]) - input_list = list(generator()) - - # The interleave transformation is essentially a flat map that - # draws from multiple input datasets concurrently (in a cyclic - # fashion). By placing `Datsaet.from_generator()` inside an - # interleave, we test its behavior when multiple iterators are - # active at the same time; by additionally prefetching inside the - # interleave, we create the possibility of parallel (modulo GIL) - # invocations to several iterators created by the same dataset. - def interleave_fn(_): - return (dataset_ops.Dataset.from_generator( - generator, output_types=(dtypes.int64, dtypes.int64), - output_shapes=([None], [3])) - .repeat(num_inner_repeats).prefetch(5)) - - iterator = ( - dataset_ops.Dataset.range(num_outer_repeats) - .interleave(interleave_fn, cycle_length=10, - block_length=len(input_list)) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for _ in range(num_inner_repeats * num_outer_repeats): - for elem in input_list: - val0, val1 = sess.run(get_next) - self.assertAllEqual(elem[0], val0) - self.assertAllEqual(elem[1], val1) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromGeneratorsRunningInParallel(self): - num_parallel_iterators = 3 - - # Define shared state that multiple iterator instances will access to - # demonstrate their concurrent activity. - lock = threading.Lock() - condition = threading.Condition(lock) - next_ticket = [0] # GUARDED_BY(lock) - - def generator(): - # NOTE(mrry): We yield one element before the barrier, because - # the current implementation of `Dataset.interleave()` must - # fetch one element from each incoming dataset to start the - # prefetching. - yield 0 - - # Define a barrier that `num_parallel_iterators` iterators must enter - # before any can proceed. Demonstrates that multiple iterators may be - # active at the same time. - condition.acquire() - ticket = next_ticket[0] - next_ticket[0] += 1 - if ticket == num_parallel_iterators - 1: - # The last iterator to join the barrier notifies the others. - condition.notify_all() - else: - # Wait until the last iterator enters the barrier. - while next_ticket[0] < num_parallel_iterators: - condition.wait() - condition.release() - - yield 1 - - # As in `testFromMultipleConcurrentGenerators()`, we use a combination of - # `Dataset.interleave()` and `Dataset.prefetch()` to cause multiple - # iterators to be active concurrently. - def interleave_fn(_): - return dataset_ops.Dataset.from_generator( - generator, output_types=dtypes.int64, output_shapes=[]).prefetch(2) - - iterator = ( - dataset_ops.Dataset.range(num_parallel_iterators) - .interleave( - interleave_fn, cycle_length=num_parallel_iterators, block_length=1) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for elem in [0, 1]: - for _ in range(num_parallel_iterators): - self.assertAllEqual(elem, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromGeneratorImplicitConversion(self): - def generator(): - yield [1] - yield [2] - yield [3] - - for dtype in [dtypes.int8, dtypes.int32, dtypes.int64]: - iterator = (dataset_ops.Dataset.from_generator( - generator, output_types=dtype, output_shapes=[1]) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual(dtype, get_next.dtype) - - with self.test_session() as sess: - sess.run(init_op) - for expected in [[1], [2], [3]]: - next_val = sess.run(get_next) - self.assertEqual(dtype.as_numpy_dtype, next_val.dtype) - self.assertAllEqual(expected, next_val) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromGeneratorTypeError(self): - def generator(): - yield np.array([1, 2, 3], dtype=np.int64) - yield np.array([4, 5, 6], dtype=np.int64) - yield "ERROR" - yield np.array([7, 8, 9], dtype=np.int64) - - iterator = (dataset_ops.Dataset.from_generator( - generator, output_types=dtypes.int64, output_shapes=[3]) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - self.assertAllEqual([1, 2, 3], sess.run(get_next)) - self.assertAllEqual([4, 5, 6], sess.run(get_next)) - with self.assertRaisesOpError(r"invalid literal for long\(\)"): - sess.run(get_next) - self.assertAllEqual([7, 8, 9], sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFromGeneratorShapeError(self): - def generator(): - yield np.array([1, 2, 3], dtype=np.int64) - yield np.array([4, 5, 6], dtype=np.int64) - yield np.array([7, 8, 9, 10], dtype=np.int64) - yield np.array([11, 12, 13], dtype=np.int64) - - iterator = (dataset_ops.Dataset.from_generator( - generator, output_types=dtypes.int64, output_shapes=[3]) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - self.assertAllEqual([1, 2, 3], sess.run(get_next)) - self.assertAllEqual([4, 5, 6], sess.run(get_next)) - with self.assertRaisesOpError(r"element of shape \(3,\) was expected"): - sess.run(get_next) - self.assertAllEqual([11, 12, 13], sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testSplitPipelineFailsWithPlacementError(self): - with session.Session( - target="", - config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess: - - dataset = dataset_ops.Dataset.from_tensors(0) - - # Define a pipeline that attempts to use variables on two - # different devices. - # - # Initialize the variables before creating to iterator, to avoid the - # placement algorithm overriding the DT_RESOURCE colocation constraints. - with ops.device("/cpu:0"): - var_0 = resource_variable_ops.ResourceVariable(initial_value=0) - dataset = dataset.map(lambda x: x + var_0.read_value()) - sess.run(var_0.initializer) - - with ops.device("/cpu:1"): - var_1 = resource_variable_ops.ResourceVariable(initial_value=0) - dataset = dataset.map(lambda x: x + var_1.read_value()) - sess.run(var_1.initializer) - - iterator = dataset.make_initializable_iterator() - - with self.assertRaisesRegexp( - errors.InvalidArgumentError, - "Trying to access resource located in device"): - sess.run(iterator.initializer) - def testRestructureDataset(self): components = (array_ops.placeholder(dtypes.int32), (array_ops.placeholder(dtypes.int32, shape=[None]), diff --git a/tensorflow/contrib/data/python/kernel_tests/dataset_serialization_test_base.py b/tensorflow/contrib/data/python/kernel_tests/dataset_serialization_test_base.py index 7cde6e05b244773966fd7c1bd4ca1e95abf7fd5e..dbc35097ddda9f0375060d43aeb43efa8107f929 100644 --- a/tensorflow/contrib/data/python/kernel_tests/dataset_serialization_test_base.py +++ b/tensorflow/contrib/data/python/kernel_tests/dataset_serialization_test_base.py @@ -24,9 +24,11 @@ import numpy as np from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops from tensorflow.python.data.ops import iterator_ops +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.ops import lookup_ops from tensorflow.python.ops import variables from tensorflow.python.platform import gfile from tensorflow.python.platform import test @@ -34,14 +36,29 @@ from tensorflow.python.training import saver as saver_lib from tensorflow.python.util import nest +def remove_variants(get_next_op): + # TODO(b/72408568): Remove this once session.run can get + # variant tensors. + """Remove variants from a nest structure, so sess.run will execute.""" + + def _remove_variant(x): + if isinstance(x, ops.Tensor) and x.dtype == dtypes.variant: + return () + else: + return x + + return nest.map_structure(_remove_variant, get_next_op) + + class DatasetSerializationTestBase(test.TestCase): """Base class for testing serializable datasets.""" def tearDown(self): self._delete_ckpt() - # TODO(b/70988345): Support native `tf.SparseTensor` objects and get rid of - # `sparse_tensors` argument. + # TODO(b/72657739): Remove sparse_tensor argument, which is to test the + # (deprecated) saveable `SparseTensorSliceDataset`, once the API + # `from_sparse_tensor_slices()`and related tests are deleted. def run_core_tests(self, ds_fn1, ds_fn2, num_outputs, sparse_tensors=False): """Runs the core tests. @@ -233,10 +250,10 @@ class DatasetSerializationTestBase(test.TestCase): saver = self._import_meta_graph() init_op, get_next_op = self._get_iterator_ops_from_collection( ds_fn, sparse_tensors=sparse_tensors) + get_next_op = remove_variants(get_next_op) with self.test_session(graph=g) as sess: self._restore(saver, sess) - sess.run(variables.global_variables_initializer()) - sess.run(init_op) + self._initialize(init_op, sess) for _ in range(num_outputs): actual.append(sess.run(get_next_op)) if verify_exhausted: @@ -296,6 +313,7 @@ class DatasetSerializationTestBase(test.TestCase): with ops.Graph().as_default() as g: _, get_next_op, saver = self._build_graph( ds_fn2, sparse_tensors=sparse_tensors) + get_next_op = remove_variants(get_next_op) with self.test_session(graph=g) as sess: self._restore(saver, sess) for _ in range(num_outputs - break_point): @@ -356,6 +374,7 @@ class DatasetSerializationTestBase(test.TestCase): with ops.Graph().as_default() as g: get_next_op, saver = self._build_empty_graph( ds_fn, sparse_tensors=sparse_tensors) + get_next_op = remove_variants(get_next_op) with self.test_session(graph=g) as sess: self._restore(saver, sess) for _ in range(num_outputs - break_point): @@ -389,9 +408,9 @@ class DatasetSerializationTestBase(test.TestCase): with ops.Graph().as_default() as g: init_op, get_next_op, saver = self._build_graph( ds_fn, sparse_tensors=sparse_tensors) + get_next_op = remove_variants(get_next_op) with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) + self._initialize(init_op, sess) for _ in range(break_point): sess.run(get_next_op) with self.assertRaises(error): @@ -485,20 +504,20 @@ class DatasetSerializationTestBase(test.TestCase): else: init_op, get_next_op, saver = self._build_graph( ds_fn, sparse_tensors=sparse_tensors) + get_next_op = remove_variants(get_next_op) return init_op, get_next_op, saver for i in range(len(break_points) + 1): with ops.Graph().as_default() as g: init_op, get_next_op, saver = get_ops() + get_next_op = remove_variants(get_next_op) with self.test_session(graph=g) as sess: if ckpt_saved: if init_before_restore: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) + self._initialize(init_op, sess) self._restore(saver, sess) else: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) + self._initialize(init_op, sess) start = break_points[i - 1] if i > 0 else 0 end = break_points[i] if i < len(break_points) else num_outputs num_iters = end - start @@ -562,13 +581,16 @@ class DatasetSerializationTestBase(test.TestCase): get_next = sparse_tensor.SparseTensor(*iterator.get_next()) else: get_next = iterator.get_next() - self._add_iterator_ops_to_collection(init_op, get_next, sparse_tensors) + self._add_iterator_ops_to_collection(init_op, get_next, ds_fn, + sparse_tensors) saver = saver_lib.Saver(allow_empty=True) return init_op, get_next, saver def _build_empty_graph(self, ds_fn, sparse_tensors=False): iterator = iterator_ops.Iterator.from_structure( - self._get_output_types(ds_fn), self._get_output_shapes(ds_fn)) + self._get_output_types(ds_fn), + output_shapes=self._get_output_shapes(ds_fn), + output_classes=self._get_output_classes(ds_fn)) saveable = contrib_iterator_ops.make_saveable_from_iterator(iterator) ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) if sparse_tensors: @@ -581,12 +603,19 @@ class DatasetSerializationTestBase(test.TestCase): def _add_iterator_ops_to_collection(self, init_op, get_next, + ds_fn, sparse_tensors=False): ops.add_to_collection("iterator_ops", init_op) # `get_next` may be a tuple e.g. in TensorSliceDataset. Since Collections # do not support tuples we flatten the tensors and restore the shape in # `_get_iterator_ops_from_collection`. - if sparse_tensors: + + # TODO(shivaniagrwal): `output_classes` is a nested structure of classes, + # this base class is specific to current test cases. Update when tests are + # added with `output_classes` as a nested structure with at least one of the + # component being `tf.SparseTensor`. + if (sparse_tensors or + self._get_output_classes(ds_fn) is sparse_tensor.SparseTensor): ops.add_to_collection("iterator_ops", get_next.indices) ops.add_to_collection("iterator_ops", get_next.values) ops.add_to_collection("iterator_ops", get_next.dense_shape) @@ -596,7 +625,8 @@ class DatasetSerializationTestBase(test.TestCase): def _get_iterator_ops_from_collection(self, ds_fn, sparse_tensors=False): all_ops = ops.get_collection("iterator_ops") - if sparse_tensors: + if (sparse_tensors or + self._get_output_classes(ds_fn) is sparse_tensor.SparseTensor): init_op, indices, values, dense_shape = all_ops return init_op, sparse_tensor.SparseTensor(indices, values, dense_shape) else: @@ -611,6 +641,10 @@ class DatasetSerializationTestBase(test.TestCase): with ops.Graph().as_default(): return ds_fn().output_shapes + def _get_output_classes(self, ds_fn): + with ops.Graph().as_default(): + return ds_fn().output_classes + def _ckpt_path(self): return os.path.join(self.get_temp_dir(), "iterator") @@ -621,8 +655,14 @@ class DatasetSerializationTestBase(test.TestCase): saver.save(sess, self._ckpt_path()) def _restore(self, saver, sess): + sess.run(lookup_ops.tables_initializer()) saver.restore(sess, self._latest_ckpt()) + def _initialize(self, init_op, sess): + sess.run(variables.global_variables_initializer()) + sess.run(lookup_ops.tables_initializer()) + sess.run(init_op) + def _import_meta_graph(self): meta_file_path = self._ckpt_path() + ".meta" return saver_lib.import_meta_graph(meta_file_path) diff --git a/tensorflow/contrib/data/python/kernel_tests/filter_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/filter_dataset_op_test.py index 5921be2ae89ba1bbbb8d6e3a509cf49c65949544..b572d6ed770fc0fe0f852359baf343c55966eddd 100644 --- a/tensorflow/contrib/data/python/kernel_tests/filter_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/filter_dataset_op_test.py @@ -20,144 +20,12 @@ from __future__ import print_function import numpy as np from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import sparse_tensor -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import functional_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import test -class FilterDatasetTest(test.TestCase): - - def testFilterDataset(self): - components = ( - np.arange(7, dtype=np.int64), - np.array([[1, 2, 3]], dtype=np.int64) * np.arange( - 7, dtype=np.int64)[:, np.newaxis], - np.array(37.0, dtype=np.float64) * np.arange(7) - ) - count = array_ops.placeholder(dtypes.int64, shape=[]) - modulus = array_ops.placeholder(dtypes.int64) - - def _map_fn(x, y, z): - return math_ops.square(x), math_ops.square(y), math_ops.square(z) - - iterator = ( - dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) - .repeat(count) - .filter(lambda x, _y, _z: math_ops.equal(math_ops.mod(x, modulus), 0)) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape[1:] for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - # Test that we can dynamically feed a different modulus value for each - # iterator. - def do_test(count_val, modulus_val): - sess.run(init_op, feed_dict={count: count_val, modulus: modulus_val}) - for _ in range(count_val): - for i in [x for x in range(7) if x**2 % modulus_val == 0]: - result = sess.run(get_next) - for component, result_component in zip(components, result): - self.assertAllEqual(component[i]**2, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - do_test(14, 2) - do_test(4, 18) - - # Test an empty dataset. - do_test(0, 1) - - def testFilterRange(self): - dataset = dataset_ops.Dataset.range(100).filter( - lambda x: math_ops.not_equal(math_ops.mod(x, 3), 2)) - iterator = dataset.make_one_shot_iterator() - get_next = iterator.get_next() - - with self.test_session() as sess: - self.assertEqual(0, sess.run(get_next)) - self.assertEqual(1, sess.run(get_next)) - self.assertEqual(3, sess.run(get_next)) - - def testFilterDict(self): - iterator = (dataset_ops.Dataset.range(10) - .map(lambda x: {"foo": x * 2, "bar": x ** 2}) - .filter(lambda d: math_ops.equal(d["bar"] % 2, 0)) - .map(lambda d: d["foo"] + d["bar"]) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - if (i ** 2) % 2 == 0: - self.assertEqual(i * 2 + i ** 2, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testUseStepContainerInFilter(self): - input_data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64) - - # Define a predicate that returns true for the first element of - # the sequence and not the second, and uses `tf.map_fn()`. - def _predicate(xs): - squared_xs = functional_ops.map_fn(lambda x: x * x, xs) - summed = math_ops.reduce_sum(squared_xs) - return math_ops.equal(summed, 1 + 4 + 9) - - iterator = ( - dataset_ops.Dataset.from_tensor_slices([[1, 2, 3], [4, 5, 6]]) - .filter(_predicate) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - self.assertAllEqual(input_data[0], sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def assertSparseValuesEqual(self, a, b): - self.assertAllEqual(a.indices, b.indices) - self.assertAllEqual(a.values, b.values) - self.assertAllEqual(a.dense_shape, b.dense_shape) - - def testSparse(self): - - def _map_fn(i): - return sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0]]), - values=(i * np.array([1])), - dense_shape=np.array([1, 1])), i - - def _filter_fn(_, i): - return math_ops.equal(i % 2, 0) - - iterator = ( - dataset_ops.Dataset.range(10).map(_map_fn).filter(_filter_fn).map( - lambda x, i: x).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(5): - actual = sess.run(get_next) - self.assertTrue(isinstance(actual, sparse_tensor.SparseTensorValue)) - self.assertSparseValuesEqual(actual, _map_fn(i * 2)[0]) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - class FilterDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): @@ -194,6 +62,10 @@ class FilterDatasetSerializationTest( return dataset_ops.Dataset.range(10).map(_map_fn).filter(_filter_fn).map( lambda x, i: x) + def testSparseCore(self): + num_outputs = 5 + self.run_core_tests(self._build_sparse_filter, None, num_outputs) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/flat_map_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/flat_map_dataset_op_test.py index d4fbaa5cdcdd315aa0524134b48eb0515169722c..f3feecef32e587045be25056815315136a883ca7 100644 --- a/tensorflow/contrib/data/python/kernel_tests/flat_map_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/flat_map_dataset_op_test.py @@ -17,13 +17,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import random - -import numpy as np - from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.python.client import session +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -34,124 +29,6 @@ from tensorflow.python.ops import random_ops from tensorflow.python.ops import sparse_ops from tensorflow.python.ops import variable_scope from tensorflow.python.platform import test -from tensorflow.python.training import server_lib - - -class FlatMapDatasetTest(test.TestCase): - - # pylint: disable=g-long-lambda - def testFlatMapDataset(self): - repeats = [1, 2, 3, 4, 5, 0, 1] - components = np.array(repeats, dtype=np.int64) - iterator = ( - dataset_ops.Dataset.from_tensor_slices(components) - .flat_map(lambda x: dataset_ops.Dataset.from_tensors([x]).repeat(x)) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in repeats: - for _ in range(i): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testNestedFlatMapDataset(self): - repeats = [[1, 2], [3, 4], [5, 0], [1, 7]] - components = np.array(repeats, dtype=np.int64) - iterator = ( - dataset_ops.Dataset.from_tensor_slices(components) - .flat_map(lambda x: dataset_ops.Dataset.from_tensor_slices(x) - .flat_map(lambda y: dataset_ops.Dataset.from_tensors(y) - .repeat(y))).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for row in repeats: - for i in row: - for _ in range(i): - self.assertEqual(i, sess.run(get_next)) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testSharedResourceNestedFlatMapDataset(self): - repeats = [[1, 2], [3, 4], [5, 0], [1, 7]] - components = np.array(repeats, dtype=np.int64) - iterator = ( - dataset_ops.Dataset.from_tensor_slices(components) - .flat_map(lambda x: dataset_ops.Dataset.from_tensor_slices(x) - .flat_map(lambda y: dataset_ops.Dataset.from_tensors(y) - .repeat(y))).make_initializable_iterator( - shared_name="shared_flat_map_iterator")) - init_op = iterator.initializer - get_next = iterator.get_next() - - # Create two concurrent sessions that share the same iterator - # resource on the same server, and verify that a random - # interleaving of `Session.run(get_next)` calls on the two - # sessions yields the expected result. - server = server_lib.Server.create_local_server() - with session.Session(server.target) as sess1: - with session.Session(server.target) as sess2: - for _ in range(3): - sess = random.choice([sess1, sess2]) - sess.run(init_op) - for row in repeats: - for i in row: - for _ in range(i): - sess = random.choice([sess1, sess2]) - self.assertEqual(i, sess.run(get_next)) - - with self.assertRaises(errors.OutOfRangeError): - sess = random.choice([sess1, sess2]) - sess.run(get_next) - - def testMapDict(self): - iterator = (dataset_ops.Dataset.range(10) - .map(lambda x: {"foo": x * 2, "bar": x ** 2}) - .flat_map(lambda d: dataset_ops.Dataset.from_tensors(d["foo"]) - .repeat(d["bar"])) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - for _ in range(i ** 2): - self.assertEqual(i * 2, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - # pylint: enable=g-long-lambda - - def testSparse(self): - def _map_fn(i): - return sparse_tensor.SparseTensorValue( - indices=[[0, 0], [1, 1]], values=(i * [1, -1]), dense_shape=[2, 2]) - - def _flat_map_fn(x): - return dataset_ops.Dataset.from_tensor_slices( - sparse_ops.sparse_to_dense(x.indices, x.dense_shape, x.values)) - - iterator = ( - dataset_ops.Dataset.range(10).map(_map_fn).flat_map(_flat_map_fn) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - for j in range(2): - expected = [i, 0] if j % 2 == 0 else [0, -i] - self.assertAllEqual(expected, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) class FlatMapDatasetSerializationTest( @@ -225,6 +102,21 @@ class FlatMapDatasetSerializationTest( self.verify_error_on_save(build_ds, 500, errors.InvalidArgumentError) + def testSparseCore(self): + + def _map_fn(i): + return sparse_tensor.SparseTensorValue( + indices=[[0, 0], [1, 1]], values=(i * [1, -1]), dense_shape=[2, 2]) + + def _flat_map_fn(x): + return dataset_ops.Dataset.from_tensor_slices( + sparse_ops.sparse_to_dense(x.indices, x.dense_shape, x.values)) + + def _build_ds(): + return dataset_ops.Dataset.range(10).map(_map_fn).flat_map(_flat_map_fn) + + self.run_core_tests(_build_ds, None, 20) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/get_single_element_test.py b/tensorflow/contrib/data/python/kernel_tests/get_single_element_test.py new file mode 100644 index 0000000000000000000000000000000000000000..32ea44f7c7ba329dc253bb9fbbcac0a1ed16aec7 --- /dev/null +++ b/tensorflow/contrib/data/python/kernel_tests/get_single_element_test.py @@ -0,0 +1,58 @@ +# Copyright 2017 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 the experimental input pipeline ops.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.data.python.ops import get_single_element +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class GetSingleElementTest(test.TestCase): + + def testGetSingleElement(self): + skip_value = array_ops.placeholder(dtypes.int64, shape=[]) + take_value = array_ops.placeholder_with_default( + constant_op.constant(1, dtype=dtypes.int64), shape=[]) + + dataset = (dataset_ops.Dataset.range(100) + .skip(skip_value) + .map(lambda x: x * x) + .take(take_value)) + + element = get_single_element.get_single_element(dataset) + + with self.test_session() as sess: + self.assertEqual(0, sess.run(element, feed_dict={skip_value: 0})) + self.assertEqual(25, sess.run(element, feed_dict={skip_value: 5})) + self.assertEqual(100, sess.run(element, feed_dict={skip_value: 10})) + + with self.assertRaisesRegexp(errors.InvalidArgumentError, + "Dataset was empty."): + sess.run(element, feed_dict={skip_value: 100}) + + with self.assertRaisesRegexp(errors.InvalidArgumentError, + "Dataset had more than one element."): + sess.run(element, feed_dict={skip_value: 0, take_value: 2}) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py index e13c60c9a71d248ba7140feda40139ad65965aeb..256ad8d94dc1a7c2b26df3f1ebf8e8e321882c15 100644 --- a/tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py @@ -26,8 +26,8 @@ import numpy as np from six.moves import zip_longest from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base -from tensorflow.contrib.data.python.ops import dataset_ops from tensorflow.contrib.data.python.ops import interleave_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import sparse_tensor @@ -38,182 +38,7 @@ from tensorflow.python.ops import sparse_ops from tensorflow.python.platform import test -class InterleaveDatasetTest(test.TestCase): - - def _interleave(self, lists, cycle_length, block_length): - # TODO(b/69678297): Consolidate python interleave implementations. - num_open = 0 - - # `all_iterators` acts as a queue of iterators over each element of `lists`. - all_iterators = [iter(l) for l in lists] - - # `open_iterators` are the iterators whose elements are currently being - # interleaved. - open_iterators = [] - for i in range(cycle_length): - if all_iterators: - open_iterators.append(all_iterators.pop(0)) - num_open += 1 - else: - open_iterators.append(None) - - while num_open or all_iterators: - for i in range(cycle_length): - if open_iterators[i] is None: - if all_iterators: - open_iterators[i] = all_iterators.pop(0) - num_open += 1 - else: - continue - for _ in range(block_length): - try: - yield next(open_iterators[i]) - except StopIteration: - open_iterators[i] = None - num_open -= 1 - break - - def testPythonImplementation(self): - input_lists = [[4, 4, 4, 4], [5, 5, 5, 5, 5], [6, 6, 6, 6, 6, 6], - [4, 4, 4, 4], [5, 5, 5, 5, 5], [6, 6, 6, 6, 6, 6]] - - # Cycle length 1 acts like `Dataset.flat_map()`. - expected_elements = itertools.chain(*input_lists) - for expected, produced in zip( - expected_elements, self._interleave(input_lists, 1, 1)): - self.assertEqual(expected, produced) - - # Cycle length > 1. - expected_elements = [4, 5, 4, 5, 4, 5, 4, - 5, 5, 6, 6, # NOTE(mrry): When we cycle back - # to a list and are already at - # the end of that list, we move - # on to the next element. - 4, 6, 4, 6, 4, 6, 4, 6, 5, 6, 5, 6, 5, 6, 5, 6, 5] - for expected, produced in zip( - expected_elements, self._interleave(input_lists, 2, 1)): - self.assertEqual(expected, produced) - - # Cycle length > 1 and block length > 1. - expected_elements = [4, 4, 4, 5, 5, 5, 4, 5, 5, 6, 6, 6, 4, 4, 4, 6, 6, 6, - 4, 5, 5, 5, 6, 6, 6, 5, 5, 6, 6, 6] - for expected, produced in zip( - expected_elements, self._interleave(input_lists, 2, 3)): - self.assertEqual(expected, produced) - - # Cycle length > len(input_values). - expected_elements = [4, 4, 5, 5, 6, 6, 4, 4, 5, 5, 6, 6, 4, 4, 5, 5, 6, 6, - 4, 4, 5, 5, 6, 6, 5, 6, 6, 5, 6, 6] - for expected, produced in zip( - expected_elements, self._interleave(input_lists, 7, 2)): - self.assertEqual(expected, produced) - - def testInterleaveDataset(self): - input_values = array_ops.placeholder(dtypes.int64, shape=[None]) - cycle_length = array_ops.placeholder(dtypes.int64, shape=[]) - block_length = array_ops.placeholder(dtypes.int64, shape=[]) - - repeat_count = 2 - - dataset = ( - dataset_ops.Dataset.from_tensor_slices(input_values) - .repeat(repeat_count) - .interleave(lambda x: dataset_ops.Dataset.from_tensors(x).repeat(x), - cycle_length, block_length)) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - next_element = iterator.get_next() - - with self.test_session() as sess: - # Cycle length 1 acts like `Dataset.flat_map()`. - sess.run(init_op, feed_dict={input_values: [4, 5, 6], - cycle_length: 1, block_length: 3}) - - for expected_element in self._interleave( - [[4] * 4, [5] * 5, [6] * 6] * repeat_count, 1, 3): - self.assertEqual(expected_element, sess.run(next_element)) - - # Cycle length > 1. - # expected: [4, 5, 4, 5, 4, 5, 4, 5, 5, 6, 6, 4, 6, 4, 6, 4, 6, 4, 6, 5, - # 6, 5, 6, 5, 6, 5, 6, 5] - sess.run(init_op, feed_dict={input_values: [4, 5, 6], - cycle_length: 2, block_length: 1}) - for expected_element in self._interleave( - [[4] * 4, [5] * 5, [6] * 6] * repeat_count, 2, 1): - self.assertEqual(expected_element, sess.run(next_element)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - # Cycle length > 1 and block length > 1. - # expected: [4, 4, 4, 5, 5, 5, 4, 5, 5, 6, 6, 6, 4, 4, 4, 6, 6, 6, 4, 5, - # 5, 5, 6, 6, 6, 5, 5, 6, 6, 6] - sess.run(init_op, feed_dict={input_values: [4, 5, 6], - cycle_length: 2, block_length: 3}) - for expected_element in self._interleave( - [[4] * 4, [5] * 5, [6] * 6] * repeat_count, 2, 3): - self.assertEqual(expected_element, sess.run(next_element)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - # Cycle length > len(input_values) * repeat_count. - # expected: [4, 4, 5, 5, 6, 6, 4, 4, 5, 5, 6, 6, 4, 4, 5, 5, 6, 6, 4, 4, - # 5, 5, 6, 6, 5, 6, 6, 5, 6, 6] - sess.run(init_op, feed_dict={input_values: [4, 5, 6], - cycle_length: 7, block_length: 2}) - for expected_element in self._interleave( - [[4] * 4, [5] * 5, [6] * 6] * repeat_count, 7, 2): - self.assertEqual(expected_element, sess.run(next_element)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - # Empty input. - sess.run(init_op, feed_dict={input_values: [], - cycle_length: 2, block_length: 3}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - # Non-empty input leading to empty output. - sess.run(init_op, feed_dict={input_values: [0, 0, 0], - cycle_length: 2, block_length: 3}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - # Mixture of non-empty and empty interleaved datasets. - sess.run(init_op, feed_dict={input_values: [4, 0, 6], - cycle_length: 2, block_length: 3}) - for expected_element in self._interleave( - [[4] * 4, [], [6] * 6] * repeat_count, 2, 3): - self.assertEqual(expected_element, sess.run(next_element)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - def testSparse(self): - - def _map_fn(i): - return sparse_tensor.SparseTensorValue( - indices=[[0, 0], [1, 1]], values=(i * [1, -1]), dense_shape=[2, 2]) - - def _interleave_fn(x): - return dataset_ops.Dataset.from_tensor_slices( - sparse_ops.sparse_to_dense(x.indices, x.dense_shape, x.values)) - - iterator = ( - dataset_ops.Dataset.range(10).map(_map_fn).interleave( - _interleave_fn, cycle_length=1).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - for j in range(2): - expected = [i, 0] if j % 2 == 0 else [0, -i] - self.assertAllEqual(expected, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - -class InterleaveDatasetSeriazationTest( +class InterleaveDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): def _build_iterator_graph(self, input_values, cycle_length, block_length): @@ -252,6 +77,22 @@ class InterleaveDatasetSeriazationTest( None, num_outputs) # pylint: enable=g-long-lambda + def testSparseCore(self): + + def _map_fn(i): + return sparse_tensor.SparseTensorValue( + indices=[[0, 0], [1, 1]], values=(i * [1, -1]), dense_shape=[2, 2]) + + def _interleave_fn(x): + return dataset_ops.Dataset.from_tensor_slices( + sparse_ops.sparse_to_dense(x.indices, x.dense_shape, x.values)) + + def _build_dataset(): + return dataset_ops.Dataset.range(10).map(_map_fn).interleave( + _interleave_fn, cycle_length=1) + + self.run_core_tests(_build_dataset, None, 20) + class ParallelInterleaveDatasetTest(test.TestCase): @@ -524,7 +365,7 @@ class ParallelInterleaveDatasetTest(test.TestCase): self.read_coordination_events[expected_element].acquire() else: self.write_coordination_events[expected_element].set() - time.sleep(0.1) # Sleep to consistently "avoid" the race condition. + time.sleep(0.5) # Sleep to consistently "avoid" the race condition. actual_element = sess.run(self.next_element) if not done_first_event: done_first_event = True @@ -611,7 +452,7 @@ class ParallelInterleaveDatasetTest(test.TestCase): self.read_coordination_events[expected_element].acquire() else: self.write_coordination_events[expected_element].set() - time.sleep(0.1) # Sleep to consistently "avoid" the race condition. + time.sleep(0.5) # Sleep to consistently "avoid" the race condition. actual_element = sess.run(self.next_element) if not done_first_event: done_first_event = True diff --git a/tensorflow/contrib/data/python/kernel_tests/iterator_ops_cluster_test.py b/tensorflow/contrib/data/python/kernel_tests/iterator_ops_cluster_test.py deleted file mode 100644 index 02379d064d4ab857ce9c7d13881a3ae37eea0980..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/data/python/kernel_tests/iterator_ops_cluster_test.py +++ /dev/null @@ -1,108 +0,0 @@ -# Copyright 2017 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 the experimental input pipeline ops that need test_util.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.core.protobuf import config_pb2 -from tensorflow.python.client import session -from tensorflow.python.data.ops import iterator_ops -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.framework import function -from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import functional_ops -from tensorflow.python.platform import test - - -class IteratorClusterTest(test.TestCase): - - def testRemoteIteratorWithoutRemoteCallFail(self): - worker_config = config_pb2.ConfigProto() - worker_config.device_count["CPU"] = 2 - worker, _ = test_util.create_local_cluster( - 1, 1, worker_config=worker_config) - - with ops.device("/job:worker/replica:0/task:0/cpu:1"): - dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) - iterator_3 = dataset_3.make_one_shot_iterator() - iterator_3_handle = iterator_3.string_handle() - - with ops.device("/job:worker/replica:0/task:0/cpu:0"): - remote_it = iterator_ops.Iterator.from_string_handle( - iterator_3_handle, dataset_3.output_types, dataset_3.output_shapes) - get_next_op = remote_it.get_next() - - with session.Session(worker[0].target) as sess: - with self.assertRaises(errors.InvalidArgumentError): - sess.run(get_next_op) - - def _testRemoteIteratorHelper(self, device0, device1, target): - with ops.device(device1): - dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) - iterator_3 = dataset_3.make_one_shot_iterator() - iterator_3_handle = iterator_3.string_handle() - - @function.Defun(dtypes.string) - def _remote_fn(h): - remote_iterator = iterator_ops.Iterator.from_string_handle( - h, dataset_3.output_types, dataset_3.output_shapes) - return remote_iterator.get_next() - - with ops.device(device0): - target_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - remote_op = functional_ops.remote_call( - args=[iterator_3_handle], - Tout=[dtypes.int32], - f=_remote_fn, - target=target_placeholder) - - with session.Session(target) as sess: - elem = sess.run(remote_op, feed_dict={target_placeholder: device1}) - self.assertEqual(elem, [1]) - # Fails when target is cpu:0 where the resource is not located. - with self.assertRaises(errors.InvalidArgumentError): - sess.run(remote_op, feed_dict={target_placeholder: device0}) - elem = sess.run(iterator_3.get_next()) - self.assertEqual(elem, [2]) - elem = sess.run(remote_op, feed_dict={target_placeholder: device1}) - self.assertEqual(elem, [3]) - with self.assertRaises(errors.OutOfRangeError): - sess.run(remote_op, feed_dict={target_placeholder: device1}) - - def testRemoteIteratorUsingRemoteCallOp(self): - worker_config = config_pb2.ConfigProto() - worker_config.device_count["CPU"] = 2 - worker, _ = test_util.create_local_cluster( - 1, 1, worker_config=worker_config) - - self._testRemoteIteratorHelper("/job:worker/replica:0/task:0/cpu:0", - "/job:worker/replica:0/task:0/cpu:1", - worker[0].target) - - def testRemoteIteratorUsingRemoteCallOpCrossProcess(self): - workers, _ = test_util.create_local_cluster(2, 1) - - self._testRemoteIteratorHelper("/job:worker/replica:0/task:0/cpu:0", - "/job:worker/replica:0/task:1/cpu:0", - workers[0].target) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/iterator_ops_test.py b/tensorflow/contrib/data/python/kernel_tests/iterator_ops_test.py deleted file mode 100644 index bda9a2a4a37e9c3d35ff99041d1150ffc43f4c43..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/data/python/kernel_tests/iterator_ops_test.py +++ /dev/null @@ -1,625 +0,0 @@ -# Copyright 2017 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 the experimental input pipeline ops.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import os -import numpy as np - -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.contrib.data.python.ops import readers -from tensorflow.core.protobuf import config_pb2 -from tensorflow.python.client import session -from tensorflow.python.data.ops import iterator_ops -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.framework import function -from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import functional_ops -from tensorflow.python.ops import gen_dataset_ops -from tensorflow.python.ops import gradients_impl -from tensorflow.python.ops import io_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import parsing_ops -from tensorflow.python.ops import script_ops -from tensorflow.python.platform import test -from tensorflow.python.training import server_lib - - -class IteratorTest(test.TestCase): - - def testAttemptingGradientsRaiseExceptions(self): - component = constant_op.constant([1]) - side = constant_op.constant(0) - add = lambda x: x + side - dataset = dataset_ops.Dataset.from_tensor_slices(component).map(add) - value = dataset.make_one_shot_iterator().get_next() - with self.assertRaisesRegexp(LookupError, "No gradient defined"): - gradients_impl.gradients(value, component) - with self.assertRaisesRegexp(LookupError, "No gradient defined"): - gradients_impl.gradients(value, side) - with self.assertRaisesRegexp(LookupError, "No gradient defined"): - gradients_impl.gradients(value, [component, side]) - - def testOneShotIterator(self): - components = (np.arange(7), - np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], - np.array(37.0) * np.arange(7)) - - def _map_fn(x, y, z): - return math_ops.square(x), math_ops.square(y), math_ops.square(z) - - iterator = (dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) - .repeat(14).make_one_shot_iterator()) - get_next = iterator.get_next() - - self.assertEqual([c.shape[1:] for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - for _ in range(14): - for i in range(7): - result = sess.run(get_next) - for component, result_component in zip(components, result): - self.assertAllEqual(component[i]**2, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testOneShotIteratorCaptureByValue(self): - components = (np.arange(7), - np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], - np.array(37.0) * np.arange(7)) - tensor_components = tuple([ops.convert_to_tensor(c) for c in components]) - - def _map_fn(x, y, z): - return math_ops.square(x), math_ops.square(y), math_ops.square(z) - - iterator = (dataset_ops.Dataset.from_tensor_slices(tensor_components) - .map(_map_fn).repeat(14).make_one_shot_iterator()) - get_next = iterator.get_next() - - self.assertEqual([c.shape[1:] for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - for _ in range(14): - for i in range(7): - result = sess.run(get_next) - for component, result_component in zip(components, result): - self.assertAllEqual(component[i]**2, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testOneShotIteratorInsideContainer(self): - components = (np.arange(7), - np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], - np.array(37.0) * np.arange(7)) - - def within_container(): - def _map_fn(x, y, z): - return math_ops.square(x), math_ops.square(y), math_ops.square(z) - iterator = (dataset_ops.Dataset.from_tensor_slices(components) - .map(_map_fn).repeat(14).make_one_shot_iterator()) - return iterator.get_next() - - server = server_lib.Server.create_local_server() - - # Create two iterators within unique containers, and run them to - # make sure that the resources aren't shared. - # - # The test below would fail if cname were the same across both - # sessions. - for i in range(2): - with session.Session(server.target) as sess: - cname = "iteration%d" % i - with ops.container(cname): - get_next = within_container() - - for _ in range(14): - for i in range(7): - result = sess.run(get_next) - for component, result_component in zip(components, result): - self.assertAllEqual(component[i]**2, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testOneShotIteratorNonBlocking(self): - dataset = dataset_ops.Dataset.from_tensors([1, 2, 3]).map(lambda x: x * x) - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() - - # Create a session with a single thread to ensure that the - # one-shot iterator initializer does not deadlock. - config = config_pb2.ConfigProto(inter_op_parallelism_threads=1, - use_per_session_threads=True) - with session.Session(config=config) as sess: - self.assertAllEqual([1, 4, 9], sess.run(next_element)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - # Test with multiple threads invoking the one-shot iterator concurrently. - with session.Session(config=config) as sess: - results = [] - def consumer_thread(): - try: - results.append(sess.run(next_element)) - except errors.OutOfRangeError: - results.append(None) - - num_threads = 8 - threads = [ - self.checkedThread(consumer_thread) for _ in range(num_threads)] - for t in threads: - t.start() - for t in threads: - t.join() - - self.assertEqual(num_threads, len(results)) - self.assertEqual(num_threads - 1, - len([None for r in results if r is None])) - self.assertAllEqual([[1, 4, 9]], [r for r in results if r is not None]) - - def testOneShotIteratorInitializerFails(self): - # Define a dataset whose initialization will always fail. - dataset = dataset_ops.Dataset.from_tensors( - array_ops.check_numerics( - constant_op.constant(1.0) / constant_op.constant(0.0), "oops")) - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() - - with self.test_session() as sess: - with self.assertRaisesRegexp(errors.InvalidArgumentError, "oops"): - sess.run(next_element) - - # Test that subsequent attempts to use the iterator also fail. - with self.assertRaisesRegexp(errors.InvalidArgumentError, "oops"): - sess.run(next_element) - - with self.test_session() as sess: - def consumer_thread(): - with self.assertRaisesRegexp(errors.InvalidArgumentError, "oops"): - sess.run(next_element) - - num_threads = 8 - threads = [ - self.checkedThread(consumer_thread) for _ in range(num_threads)] - for t in threads: - t.start() - for t in threads: - t.join() - - def testSimpleSharedResource(self): - components = ( - np.array(1, dtype=np.int64), - np.array([1, 2, 3], dtype=np.int64), - np.array(37.0, dtype=np.float64) - ) - - server = server_lib.Server.create_local_server() - - # Create two non-overlapping sessions that share the same iterator - # resource on the same server, and verify that an action of the - # first session (initializing the iterator) is visible in the - # second session. - with ops.Graph().as_default(): - iterator = (dataset_ops.Dataset.from_tensors(components) - .map(lambda x, y, z: (x, y, z)).make_initializable_iterator( - shared_name="shared_iterator")) - init_op = iterator.initializer - get_next = iterator.get_next() - - with session.Session(server.target) as sess: - sess.run(init_op) - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertAllEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Re-initialize the iterator in the first session. - sess.run(init_op) - - with ops.Graph().as_default(): - # Re-define the iterator manually, without defining any of the - # functions in this graph, to ensure that we are not - # accidentally redefining functions with the same names in the - # new graph. - iterator = iterator_ops.Iterator.from_structure( - shared_name="shared_iterator", - output_types=(dtypes.int64, dtypes.int64, dtypes.float64), - output_shapes=([], [3], [])) - get_next = iterator.get_next() - - with session.Session(server.target) as sess: - # Use the iterator without re-initializing in the second session. - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertAllEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testNotInitializedError(self): - components = (np.array(1), np.array([1, 2, 3]), np.array(37.0)) - iterator = (dataset_ops.Dataset.from_tensors(components) - .make_initializable_iterator()) - get_next = iterator.get_next() - - with self.test_session() as sess: - with self.assertRaisesRegexp(errors.FailedPreconditionError, - "iterator has not been initialized"): - sess.run(get_next) - - def testReinitializableIterator(self): - dataset_3 = dataset_ops.Dataset.from_tensors( - constant_op.constant([1, 2, 3])) - dataset_4 = dataset_ops.Dataset.from_tensors( - constant_op.constant([4, 5, 6, 7])) - iterator = iterator_ops.Iterator.from_structure(dataset_3.output_types, - [None]) - - dataset_3_init_op = iterator.make_initializer(dataset_3) - dataset_4_init_op = iterator.make_initializer(dataset_4) - get_next = iterator.get_next() - - self.assertEqual(dataset_3.output_types, iterator.output_types) - self.assertEqual(dataset_4.output_types, iterator.output_types) - self.assertEqual([None], iterator.output_shapes.as_list()) - - with self.test_session() as sess: - # The iterator is initially uninitialized. - with self.assertRaises(errors.FailedPreconditionError): - sess.run(get_next) - - # Initialize with one dataset. - sess.run(dataset_3_init_op) - self.assertAllEqual([1, 2, 3], sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Initialize with a different dataset. - sess.run(dataset_4_init_op) - self.assertAllEqual([4, 5, 6, 7], sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Reinitialize with the first dataset. - sess.run(dataset_3_init_op) - self.assertAllEqual([1, 2, 3], sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testReinitializableIteratorStaticErrors(self): - # Non-matching structure for types and shapes. - with self.assertRaises(TypeError): - iterator = iterator_ops.Iterator.from_structure((dtypes.int64, - dtypes.float64), [None]) - - # Test validation of dataset argument. - iterator = iterator_ops.Iterator.from_structure((dtypes.int64, - dtypes.float64)) - - # Incompatible structure. - with self.assertRaises(ValueError): - iterator.make_initializer( - dataset_ops.Dataset.from_tensors(((constant_op.constant( - [1, 2, 3], dtype=dtypes.int64),), (constant_op.constant( - [4., 5., 6., 7.], dtype=dtypes.float64),)))) - - # Incompatible types. - with self.assertRaises(TypeError): - iterator.make_initializer( - dataset_ops.Dataset.from_tensors((constant_op.constant( - [1, 2, 3], dtype=dtypes.int32), constant_op.constant( - [4., 5., 6., 7.], dtype=dtypes.float32)))) - - # Incompatible shapes. - iterator = iterator_ops.Iterator.from_structure( - (dtypes.int64, dtypes.float64), ([None], [])) - with self.assertRaises(TypeError): - iterator.make_initializer( - dataset_ops.Dataset.from_tensors((constant_op.constant( - [1, 2, 3], dtype=dtypes.int64), constant_op.constant( - [4., 5., 6., 7.], dtype=dtypes.float64)))) - - def testIteratorStringHandle(self): - dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) - dataset_4 = dataset_ops.Dataset.from_tensor_slices([10, 20, 30, 40]) - - iterator_3 = dataset_3.make_one_shot_iterator() - iterator_4 = dataset_4.make_one_shot_iterator() - - handle_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - feedable_iterator = iterator_ops.Iterator.from_string_handle( - handle_placeholder, dataset_3.output_types, dataset_3.output_shapes) - next_element = feedable_iterator.get_next() - - self.assertEqual(dataset_3.output_types, feedable_iterator.output_types) - self.assertEqual(dataset_4.output_types, feedable_iterator.output_types) - self.assertEqual([], feedable_iterator.output_shapes) - - with self.test_session() as sess: - iterator_3_handle = sess.run(iterator_3.string_handle()) - iterator_4_handle = sess.run(iterator_4.string_handle()) - - self.assertEqual( - 10, sess.run(next_element, - feed_dict={handle_placeholder: iterator_4_handle})) - self.assertEqual( - 1, sess.run(next_element, - feed_dict={handle_placeholder: iterator_3_handle})) - self.assertEqual( - 20, sess.run(next_element, - feed_dict={handle_placeholder: iterator_4_handle})) - self.assertEqual( - 2, sess.run(next_element, - feed_dict={handle_placeholder: iterator_3_handle})) - self.assertEqual( - 30, sess.run(next_element, - feed_dict={handle_placeholder: iterator_4_handle})) - self.assertEqual( - 3, sess.run(next_element, - feed_dict={handle_placeholder: iterator_3_handle})) - self.assertEqual( - 40, sess.run(next_element, - feed_dict={handle_placeholder: iterator_4_handle})) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element, - feed_dict={handle_placeholder: iterator_3_handle}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element, - feed_dict={handle_placeholder: iterator_4_handle}) - - def testIteratorStringHandleError(self): - dataset_int_scalar = (dataset_ops.Dataset.from_tensor_slices([1, 2, - 3]).repeat()) - dataset_float_vector = (dataset_ops.Dataset.from_tensors([1.0, 2.0, 3.0])) - - handle_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - - feedable_int_scalar = iterator_ops.Iterator.from_string_handle( - handle_placeholder, dtypes.int32, []) - feedable_int_vector = iterator_ops.Iterator.from_string_handle( - handle_placeholder, dtypes.int32, [None]) - feedable_int_any = iterator_ops.Iterator.from_string_handle( - handle_placeholder, dtypes.int32) - - with self.test_session() as sess: - handle_int_scalar = sess.run( - dataset_int_scalar.make_one_shot_iterator().string_handle()) - handle_float_vector = sess.run( - dataset_float_vector.make_one_shot_iterator().string_handle()) - - self.assertEqual(1, - sess.run( - feedable_int_scalar.get_next(), - feed_dict={handle_placeholder: handle_int_scalar})) - - self.assertEqual(2, - sess.run( - feedable_int_any.get_next(), - feed_dict={handle_placeholder: handle_int_scalar})) - - with self.assertRaises(errors.InvalidArgumentError): - print(sess.run( - feedable_int_vector.get_next(), - feed_dict={handle_placeholder: handle_int_scalar})) - - with self.assertRaises(errors.InvalidArgumentError): - print(sess.run( - feedable_int_vector.get_next(), - feed_dict={handle_placeholder: handle_float_vector})) - - def testRemoteIteratorUsingRemoteCallOpDirectSession(self): - worker_config = config_pb2.ConfigProto() - worker_config.device_count["CPU"] = 3 - - with ops.device("/job:localhost/replica:0/task:0/cpu:1"): - dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) - iterator_3 = dataset_3.make_one_shot_iterator() - iterator_3_handle = iterator_3.string_handle() - - @function.Defun(dtypes.string) - def _remote_fn(h): - remote_iterator = iterator_ops.Iterator.from_string_handle( - h, dataset_3.output_types, dataset_3.output_shapes) - return remote_iterator.get_next() - - with ops.device("/job:localhost/replica:0/task:0/cpu:0"): - target_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - remote_op = functional_ops.remote_call( - args=[iterator_3_handle], - Tout=[dtypes.int32], - f=_remote_fn, - target=target_placeholder) - - with self.test_session(config=worker_config) as sess: - elem = sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:1" - }) - self.assertEqual(elem, [1]) - # Fails when target is cpu:2 where the resource is not located. - with self.assertRaises(errors.InvalidArgumentError): - sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:2" - }) - elem = sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:1" - }) - self.assertEqual(elem, [2]) - elem = sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:1" - }) - self.assertEqual(elem, [3]) - with self.assertRaises(errors.OutOfRangeError): - sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:1" - }) - - def testRemoteIteratorUsingRemoteCallOpDirectSessionGPUCPU(self): - if not test_util.is_gpu_available(): - self.skipTest("No GPU available") - - with ops.device("/job:localhost/replica:0/task:0/cpu:0"): - dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) - iterator_3 = dataset_3.make_one_shot_iterator() - iterator_3_handle = iterator_3.string_handle() - - def _encode_raw(byte_array): - return bytes(bytearray(byte_array)) - - @function.Defun(dtypes.uint8) - def _remote_fn(h): - handle = script_ops.py_func(_encode_raw, [h], dtypes.string) - remote_iterator = iterator_ops.Iterator.from_string_handle( - handle, dataset_3.output_types, dataset_3.output_shapes) - return remote_iterator.get_next() - - with ops.device("/job:localhost/replica:0/task:0/device:GPU:0"): - target_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - iterator_3_handle_uint8 = parsing_ops.decode_raw( - bytes=iterator_3_handle, out_type=dtypes.uint8) - remote_op = functional_ops.remote_call( - args=[iterator_3_handle_uint8], - Tout=[dtypes.int32], - f=_remote_fn, - target=target_placeholder) - - with self.test_session() as sess: - elem = sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:0" - }) - self.assertEqual(elem, [1]) - elem = sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:0" - }) - self.assertEqual(elem, [2]) - elem = sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:0" - }) - self.assertEqual(elem, [3]) - with self.assertRaises(errors.OutOfRangeError): - sess.run( - remote_op, - feed_dict={ - target_placeholder: "/job:localhost/replica:0/task:0/cpu:0" - }) - - def testIncorrectIteratorRestore(self): - - def _path(): - return os.path.join(self.get_temp_dir(), "iterator") - - def _save_op(iterator_resource): - iterator_state_variant = gen_dataset_ops.serialize_iterator( - iterator_resource) - save_op = io_ops.write_file( - _path(), parsing_ops.serialize_tensor(iterator_state_variant)) - return save_op - - def _restore_op(iterator_resource): - iterator_state_variant = parsing_ops.parse_tensor( - io_ops.read_file(_path()), dtypes.variant) - restore_op = gen_dataset_ops.deserialize_iterator(iterator_resource, - iterator_state_variant) - return restore_op - - def _build_range_dataset_graph(): - start = 1 - stop = 10 - iterator = dataset_ops.Dataset.range(start, - stop).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - save_op = _save_op(iterator._iterator_resource) - restore_op = _restore_op(iterator._iterator_resource) - return init_op, get_next, save_op, restore_op - - def _build_reader_dataset_graph(): - filenames = ["test"] # Does not exist but we don't care in this test. - iterator = readers.FixedLengthRecordDataset( - filenames, 1, 0, 0).make_initializable_iterator() - init_op = iterator.initializer - get_next_op = iterator.get_next() - save_op = _save_op(iterator._iterator_resource) - restore_op = _restore_op(iterator._iterator_resource) - return init_op, get_next_op, save_op, restore_op - - # Saving iterator for RangeDataset graph. - with ops.Graph().as_default() as g: - init_op, _, save_op, _ = _build_range_dataset_graph() - with self.test_session(graph=g) as sess: - sess.run(init_op) - sess.run(save_op) - - # Attempt to restore the saved iterator into an IteratorResource of - # incompatible type. An iterator of RangeDataset has output type int64, - # while an iterator of FixedLengthRecordDataset has output type string. - # So an InvalidArgumentError should be raised by - # IteratorResource::set_iterator. - with ops.Graph().as_default() as g: - _, _, _, restore_op = _build_reader_dataset_graph() - with self.test_session(graph=g) as sess: - with self.assertRaises(errors.InvalidArgumentError): - sess.run(restore_op) - - def testToSingleElement(self): - skip_value = array_ops.placeholder(dtypes.int64, shape=[]) - take_value = array_ops.placeholder_with_default( - constant_op.constant(1, dtype=dtypes.int64), shape=[]) - - dataset = (dataset_ops.Dataset.range(100) - .skip(skip_value) - .map(lambda x: x * x) - .take(take_value)) - - element = dataset_ops.get_single_element(dataset) - - with self.test_session() as sess: - self.assertEqual(0, sess.run(element, feed_dict={skip_value: 0})) - self.assertEqual(25, sess.run(element, feed_dict={skip_value: 5})) - self.assertEqual(100, sess.run(element, feed_dict={skip_value: 10})) - - with self.assertRaisesRegexp(errors.InvalidArgumentError, - "Dataset was empty."): - sess.run(element, feed_dict={skip_value: 100}) - - with self.assertRaisesRegexp(errors.InvalidArgumentError, - "Dataset had more than one element."): - sess.run(element, feed_dict={skip_value: 0, take_value: 2}) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/list_files_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/list_files_dataset_op_test.py deleted file mode 100644 index 27298de65f90c627e5eb638385bfe0478ef74fca..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/data/python/kernel_tests/list_files_dataset_op_test.py +++ /dev/null @@ -1,159 +0,0 @@ -# Copyright 2017 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 the experimental input pipeline ops.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from os import path -import shutil -import tempfile - -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.ops import array_ops -from tensorflow.python.platform import test -from tensorflow.python.util import compat - - -class ListFilesDatasetOpTest(test.TestCase): - - def setUp(self): - self.tmp_dir = tempfile.mkdtemp() - - def tearDown(self): - shutil.rmtree(self.tmp_dir, ignore_errors=True) - - def _touchTempFiles(self, filenames): - for filename in filenames: - open(path.join(self.tmp_dir, filename), 'a').close() - - def testEmptyDirectory(self): - dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*')) - with self.test_session() as sess: - itr = dataset.make_one_shot_iterator() - with self.assertRaises(errors.OutOfRangeError): - sess.run(itr.get_next()) - - def testSimpleDirectory(self): - filenames = ['a', 'b', 'c'] - self._touchTempFiles(filenames) - - dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*')) - with self.test_session() as sess: - itr = dataset.make_one_shot_iterator() - - full_filenames = [] - produced_filenames = [] - for filename in filenames: - full_filenames.append( - compat.as_bytes(path.join(self.tmp_dir, filename))) - produced_filenames.append(compat.as_bytes(sess.run(itr.get_next()))) - self.assertItemsEqual(full_filenames, produced_filenames) - with self.assertRaises(errors.OutOfRangeError): - sess.run(itr.get_next()) - - def testEmptyDirectoryInitializer(self): - filename_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - dataset = dataset_ops.Dataset.list_files(filename_placeholder) - - with self.test_session() as sess: - itr = dataset.make_initializable_iterator() - sess.run( - itr.initializer, - feed_dict={filename_placeholder: path.join(self.tmp_dir, '*')}) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(itr.get_next()) - - def testSimpleDirectoryInitializer(self): - filenames = ['a', 'b', 'c'] - self._touchTempFiles(filenames) - - filename_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - dataset = dataset_ops.Dataset.list_files(filename_placeholder) - - with self.test_session() as sess: - itr = dataset.make_initializable_iterator() - sess.run( - itr.initializer, - feed_dict={filename_placeholder: path.join(self.tmp_dir, '*')}) - - full_filenames = [] - produced_filenames = [] - for filename in filenames: - full_filenames.append( - compat.as_bytes(path.join(self.tmp_dir, filename))) - produced_filenames.append(compat.as_bytes(sess.run(itr.get_next()))) - - self.assertItemsEqual(full_filenames, produced_filenames) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(itr.get_next()) - - def testFileSuffixes(self): - filenames = ['a.txt', 'b.py', 'c.py', 'd.pyc'] - self._touchTempFiles(filenames) - - filename_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - dataset = dataset_ops.Dataset.list_files(filename_placeholder) - - with self.test_session() as sess: - itr = dataset.make_initializable_iterator() - sess.run( - itr.initializer, - feed_dict={filename_placeholder: path.join(self.tmp_dir, '*.py')}) - - full_filenames = [] - produced_filenames = [] - for filename in filenames[1:-1]: - full_filenames.append( - compat.as_bytes(path.join(self.tmp_dir, filename))) - produced_filenames.append(compat.as_bytes(sess.run(itr.get_next()))) - self.assertItemsEqual(full_filenames, produced_filenames) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(itr.get_next()) - - def testFileMiddles(self): - filenames = ['a.txt', 'b.py', 'c.pyc'] - self._touchTempFiles(filenames) - - filename_placeholder = array_ops.placeholder(dtypes.string, shape=[]) - dataset = dataset_ops.Dataset.list_files(filename_placeholder) - - with self.test_session() as sess: - itr = dataset.make_initializable_iterator() - sess.run( - itr.initializer, - feed_dict={filename_placeholder: path.join(self.tmp_dir, '*.py*')}) - - full_filenames = [] - produced_filenames = [] - for filename in filenames[1:]: - full_filenames.append( - compat.as_bytes(path.join(self.tmp_dir, filename))) - produced_filenames.append(compat.as_bytes(sess.run(itr.get_next()))) - - self.assertItemsEqual(full_filenames, produced_filenames) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(itr.get_next()) - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/map_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/map_dataset_op_test.py index e9a07da84a8c80c09ebd4dab0b1d69febe1c9790..8d4042927970cab2f5a518fc0da49b38444dbcdf 100644 --- a/tensorflow/contrib/data/python/kernel_tests/map_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/map_dataset_op_test.py @@ -16,16 +16,14 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from collections import namedtuple import os -import threading import numpy as np from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base -from tensorflow.contrib.data.python.ops import dataset_ops from tensorflow.contrib.data.python.ops import error_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -33,15 +31,9 @@ from tensorflow.python.framework import function from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import array_ops -from tensorflow.python.ops import data_flow_ops -from tensorflow.python.ops import functional_ops from tensorflow.python.ops import io_ops -from tensorflow.python.ops import lookup_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops -from tensorflow.python.ops import script_ops -from tensorflow.python.ops import sparse_ops -from tensorflow.python.ops import string_ops from tensorflow.python.ops import variable_scope from tensorflow.python.platform import test from tensorflow.python.util import compat @@ -49,227 +41,13 @@ from tensorflow.python.util import compat class MapDatasetTest(test.TestCase): - def _buildMapDataset(self, components, count): - def _map_fn(x, y, z): - return math_ops.square(x), math_ops.square(y), math_ops.square(z) - return (dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) - .repeat(count)) - - def testMapDataset(self): - """Test an dataset that maps a TF function across its input elements.""" - # The pipeline is TensorSliceDataset -> MapDataset(square_3) -> - # RepeatDataset(count). - components = (np.arange(7), - np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], - np.array(37.0) * np.arange(7)) - count = array_ops.placeholder(dtypes.int64, shape=[]) - - dataset = self._buildMapDataset(components, count) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape[1:] for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - # Test single-threaded access to the iterator. - sess.run(init_op, feed_dict={count: 14}) - for _ in range(14): - for i in range(7): - result = sess.run(get_next) - for component, result_component in zip(components, result): - self.assertAllEqual(component[i]**2, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test multi-threaded access to the same iterator. - sess.run(init_op, feed_dict={count: 18}) - results = [] - def iterator_thread(): - while True: - try: - results.append(sess.run(get_next)) - except errors.OutOfRangeError: - return - threads = [self.checkedThread(target=iterator_thread) for _ in range(8)] - for t in threads: - t.start() - for t in threads: - t.join() - - # `results` will contain the same elements components**2 - # repeated 18 times, but in a non-deterministic order. Sort the - # results, and assert that each element of components**2 is - # produced 18 times. - results.sort(key=lambda x: x[0]) - for i in range(7): - for j in range(18): - for component, result_component in zip(components, - results[i * 18 + j]): - self.assertAllEqual(component[i]**2, result_component) - - def _buildParallelMapDataset(self, components, count, num_threads, - output_buffer_size): - def _map_fn(x, y, z): - return math_ops.square(x), math_ops.square(y), math_ops.square(z) - return (dataset_ops.Dataset.from_tensor_slices(components).map( - _map_fn, num_threads=num_threads, output_buffer_size=output_buffer_size) - .repeat(count)) - - def testParallelMapDataset(self): - """Test an dataset that maps a TF function across its input elements.""" - # The pipeline is TensorSliceDataset -> ParallelMapDataset(square_3) -> - # RepeatDataset(count). - components = (np.arange(7), - np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], - np.array(37.0) * np.arange(7)) - count = array_ops.placeholder(dtypes.int64, shape=[]) - num_threads = array_ops.placeholder(dtypes.int32, shape=[]) - output_buffer_size = array_ops.placeholder(dtypes.int64, shape=[]) - - dataset = self._buildParallelMapDataset(components, count, num_threads, - output_buffer_size) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape[1:] for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - def do_test(num_threads_val, output_buffer_size_val): - # Test single-threaded access to the iterator. - sess.run(init_op, feed_dict={ - count: 14, - num_threads: num_threads_val, - output_buffer_size: output_buffer_size_val}) - for _ in range(14): - for i in range(7): - result = sess.run(get_next) - for component, result_component in zip(components, result): - self.assertAllEqual(component[i]**2, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test multi-threaded access to the same iterator. - sess.run(init_op, feed_dict={ - count: 18, - num_threads: num_threads_val, - output_buffer_size: output_buffer_size_val}) - results = [] - def iterator_thread(): - while True: - try: - results.append(sess.run(get_next)) - except errors.OutOfRangeError: - return - threads = [self.checkedThread(target=iterator_thread) - for _ in range(64)] - for t in threads: - t.start() - for t in threads: - t.join() - - # `results` will contain the same elements components**2 - # repeated 18 times, but in a non-deterministic order. Sort the - # results, and assert that each element of components**2 is - # produced 18 times. - results.sort(key=lambda x: x[0]) - for i in range(7): - for j in range(18): - for component, result_component in zip(components, - results[i * 18 + j]): - self.assertAllEqual(component[i]**2, result_component) - - for num_threads_val, output_buffer_size_val in [ - (1, 1), (1, 2), (2, 2), (2, 4), (8, 8), (8, 16)]: - do_test(num_threads_val, output_buffer_size_val) - - def testImplicitDisposeParallelMapDataset(self): - # Tests whether a parallel map dataset will be cleaned up correctly when - # the pipeline does not run it until exhaustion. - # The pipeline is TensorSliceDataset -> MapDataset(square_3) -> - # RepeatDataset(1000). - components = (np.arange(1000), - np.array([[1, 2, 3]]) * np.arange(1000)[:, np.newaxis], - np.array(37.0) * np.arange(1000)) - - dataset = self._buildParallelMapDataset(components, 1000, 100, 100) - # NOTE(mrry): Also test that the prefetching thread is cancelled correctly. - dataset = dataset.prefetch(100) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for _ in range(3): - sess.run(get_next) - - def testParallelMapUnspecifiedOutputSize(self): - components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) - - dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .map(lambda x: array_ops.check_numerics(x, "message"), - num_threads=2)) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for _ in range(3): - sess.run(get_next) - - def testParallelMapError(self): - components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) - - dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .map(lambda x: array_ops.check_numerics(x, "message"), - num_threads=2, output_buffer_size=2)) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for _ in range(3): - sess.run(get_next) - # The 4th element is NaN, so `array_ops.check_numerics()` should fail. - with self.assertRaises(errors.InvalidArgumentError): - sess.run(get_next) - sess.run(get_next) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testPrefetchError(self): - components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) - - dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .map(lambda x: array_ops.check_numerics(x, "message")) - .prefetch(2)) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for _ in range(3): - sess.run(get_next) - # The 4th element is NaN, so `array_ops.check_numerics()` should fail. - with self.assertRaises(errors.InvalidArgumentError): - sess.run(get_next) - sess.run(get_next) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - def testMapIgnoreError(self): components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) - dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .map(lambda x: array_ops.check_numerics(x, "message")).apply( - error_ops.ignore_errors())) + dataset = ( + dataset_ops.Dataset.from_tensor_slices(components) + .map(lambda x: array_ops.check_numerics(x, "message")).apply( + error_ops.ignore_errors())) iterator = dataset.make_initializable_iterator() init_op = iterator.initializer get_next = iterator.get_next() @@ -284,10 +62,10 @@ class MapDatasetTest(test.TestCase): def testParallelMapIgnoreError(self): components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) - dataset = (dataset_ops.Dataset.from_tensor_slices(components).map( - lambda x: array_ops.check_numerics(x, "message"), - num_threads=2, - output_buffer_size=2).apply(error_ops.ignore_errors())) + dataset = ( + dataset_ops.Dataset.from_tensor_slices(components).map( + lambda x: array_ops.check_numerics(x, "message"), + num_parallel_calls=2).prefetch(2).apply(error_ops.ignore_errors())) iterator = dataset.make_initializable_iterator() init_op = iterator.initializer get_next = iterator.get_next() @@ -308,9 +86,10 @@ class MapDatasetTest(test.TestCase): for filename in filenames: write_string_to_file(filename, filename) - dataset = (dataset_ops.Dataset.from_tensor_slices(filenames).map( - io_ops.read_file, num_threads=2, output_buffer_size=2).apply( - error_ops.ignore_errors())) + dataset = ( + dataset_ops.Dataset.from_tensor_slices(filenames).map( + io_ops.read_file, num_parallel_calls=2).prefetch(2).apply( + error_ops.ignore_errors())) iterator = dataset.make_initializable_iterator() init_op = iterator.initializer get_next = iterator.get_next() @@ -334,321 +113,125 @@ class MapDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - def testCaptureHashTable(self): - # NOTE(mrry): We must use the V2 variants of `HashTable` - # etc. because these produce a `tf.resource`-typed output that is - # compatible with the in-graph function implementation. - default_val = -1 - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup_ops.HashTable( - lookup_ops.KeyValueTensorInitializer(keys, values), default_val) - - input_sentences = dataset_ops.Dataset.from_tensor_slices( - ["brain brain tank salad surgery", "surgery brain"]) - - iterator = (input_sentences - .map(lambda x: string_ops.string_split([x]).values) - .map(table.lookup) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(table.init) - sess.run(init_op) - - print(sess.run(get_next)) - print(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testCaptureQueue(self): - elements = np.random.randint(100, size=[200]) - queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[]) - enqueue_op = queue.enqueue_many(elements) - close_op = queue.close() - iterator = (dataset_ops.Dataset.from_tensors(0).repeat(-1) - .map(lambda _: queue.dequeue()).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(enqueue_op) - sess.run(close_op) - sess.run(init_op) - for element in elements: - self.assertEqual(element, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testCaptureSameResourceMultipleTimes(self): - elements = np.random.randint(100, size=[200]) - queue = data_flow_ops.FIFOQueue( - 200, dtypes.int64, shapes=[], shared_name="shared_queue") - queue_2 = data_flow_ops.FIFOQueue( - 200, dtypes.int64, shapes=[], shared_name="shared_queue") - - enqueue_op = queue.enqueue_many(elements) - close_op = queue.close() + def testCaptureResourceInMapFn(self): - iterator = (dataset_ops.Dataset.from_tensors(0).repeat(-1) - .map(lambda _: (queue.dequeue(), queue_2.dequeue())) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() + def _build_ds(iterator): - with self.test_session() as sess: - sess.run(enqueue_op) - sess.run(close_op) - sess.run(init_op) - for i in range(100): - self.assertEqual(sorted([elements[i * 2], elements[i * 2 + 1]]), - sorted(sess.run(get_next))) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + def _map_fn(x): + get_next = iterator.get_next() + return x * get_next - def testCaptureVariable(self): - counter_var = variable_scope.get_variable( - "counter", (), dtypes.int32, use_resource=True) - iterator = (dataset_ops.Dataset.from_tensors(0).repeat(10) - .map(lambda _: counter_var.assign_add(1)) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() + return dataset_ops.Dataset.range(10).map(_map_fn) - with self.test_session() as sess: - sess.run(counter_var.initializer) - sess.run(init_op) - for i in range(10): - self.assertEqual(i, sess.run(counter_var)) - self.assertEqual(i + 1, sess.run(get_next)) - self.assertEqual(10, sess.run(counter_var)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - self.assertEqual(10, sess.run(counter_var)) - - def testCaptureUninitializedVariableError(self): - counter_var = variable_scope.get_variable( - "counter", (), dtypes.int32, use_resource=True) - iterator = (dataset_ops.Dataset.from_tensors(0).repeat(10) - .map(lambda _: counter_var.assign_add(1)) - .make_initializable_iterator()) - init_op = iterator.initializer + def _build_graph(): + captured_iterator = dataset_ops.Dataset.range( + 10).make_initializable_iterator() + ds = _build_ds(captured_iterator) + iterator = ds.make_initializable_iterator() + init_op = iterator.initializer + get_next = iterator.get_next() + return captured_iterator.initializer, init_op, get_next - with self.test_session() as sess: - with self.assertRaisesRegexp(errors.FailedPreconditionError, - "Failed to capture resource"): + with ops.Graph().as_default() as g: + captured_init_op, init_op, get_next = _build_graph() + with self.test_session(graph=g) as sess: + sess.run(captured_init_op) sess.run(init_op) + for i in range(10): + self.assertEquals(i * i, sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) - def testSeededStatefulOperatorIsProperlyStateful(self): - iterator = (dataset_ops.Dataset.from_tensors(0).repeat(10) - .map(lambda _: random_ops.random_uniform((), seed=11)).batch(2) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - with self.test_session() as sess: - sess.run(init_op) - random_values = [] - with self.assertRaises(errors.OutOfRangeError): - while True: - random_values.extend(sess.run(get_next)) - self.assertEqual(10, len(random_values)) - self.assertGreater(np.abs(np.diff(random_values)).max(), 1e-6) - sess.run(init_op) - random_values_2 = [] - with self.assertRaises(errors.OutOfRangeError): - while True: - random_values_2.extend(sess.run(get_next)) +class MapDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): - # Randomness is repeatable given same seed - self.assertAllClose(random_values, random_values_2) + def setUp(self): + self._tensor_slice_len = 7 + self._num_epochs = 14 + self._num_outputs = self._tensor_slice_len * self._num_epochs - def testMapDict(self): - iterator = (dataset_ops.Dataset.range(10) - .map(lambda x: {"foo": x * 2, "bar": x ** 2}) - .map(lambda d: d["foo"] + d["bar"]) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() + def _build_ds(self, multiplier=37.0): + components = (np.arange(self._tensor_slice_len), np.array([[1, 2, 3]]) * + np.arange(self._tensor_slice_len)[:, np.newaxis], + np.array(multiplier) * np.arange(self._tensor_slice_len)) - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - self.assertEqual(i * 2 + i ** 2, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + def _map_fn(x, y, z): + return math_ops.square(x), math_ops.square(y), math_ops.square(z) - def testMapNamedtuple(self, count=10): - # construct dataset of tuples - labels = dataset_ops.Dataset.range(count) - images = labels.map(lambda l: -l) - dataset_tuple = dataset_ops.Dataset.zip((labels, images)) + return ( + dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) + .repeat(self._num_epochs)) - # convert dataset of tuples to dataset of namedtuples - example = namedtuple("Example", ["label", "image"]) - dataset_namedtuple = dataset_tuple.map(example) + def testSaveRestoreCore(self): + self.run_core_tests( + self._build_ds, + lambda: self._build_ds(multiplier=15.0), + self._num_outputs) - def preprocess_tuple(label, image): - image = 2 * image - return label, image + def testSaveStatefulFunction(self): - def preprocess_namedtuple(example): - return example._replace(image=2 * example.image) + def _build_ds(): - # preprocess both datasets - dataset_tuple = dataset_tuple.map(preprocess_tuple) - dataset_namedtuple = dataset_namedtuple.map(preprocess_namedtuple) + def _map_fn(x): + return random_ops.random_uniform( + (), 0, 10, dtype=dtypes.int32) * math_ops.to_int32(x) - next_tuple = dataset_tuple.make_one_shot_iterator().get_next() - next_namedtuple = dataset_namedtuple.make_one_shot_iterator().get_next() + return dataset_ops.Dataset.range(100).map(_map_fn) - # make sure both datasets contain the same data - with self.test_session() as sess: - for i in range(count): - tuple_, namedtuple_ = sess.run([next_tuple, next_namedtuple]) - self.assertEqual(tuple_, namedtuple_) - self.assertEqual(tuple_, (i, -2 * i)) + self.verify_error_on_save(_build_ds, 15, errors.InvalidArgumentError) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_namedtuple) - - def testUseStepContainerInMap(self): - row = np.arange(6) - iterator = ( - dataset_ops.Dataset.from_tensors(row) - .map(lambda elems: functional_ops.map_fn(lambda x: x * x, elems)) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() + def testCaptureVariableInMapFn(self): - with self.test_session() as sess: - sess.run(init_op) - self.assertAllEqual(row ** 2, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + def _build_ds(): + counter_var = variable_scope.get_variable( + "counter", (), dtypes.int32, use_resource=True) + return (dataset_ops.Dataset.from_tensors(0).repeat(10).map( + lambda _: counter_var.assign_add(1))) - def testPrefetch(self): - # We will use this event to test that `_map_py_func()` has been - # invoked a certain number of times (6 times, to be exact) after - # consuming fewer elements from the iterator. - ev = threading.Event() + self.verify_error_on_save(_build_ds, 15, errors.InvalidArgumentError) - set_event_during_invocation = 5 + def testCaptureConstantInMapFn(self): - def _map_py_func(x): - if x == set_event_during_invocation: - ev.set() - return x * x + def _build_ds(): + constant_var = constant_op.constant(5) + return (dataset_ops.Dataset.from_tensors(0).repeat(10).map( + lambda x: x + constant_var)) - def _map_fn(x): - return script_ops.py_func(_map_py_func, [x], x.dtype) + self.run_core_tests(_build_ds, None, 10) - buffer_size_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = ( - dataset_ops.Dataset.range(100) - .map(_map_fn) - .prefetch(buffer_size_placeholder) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() + def testCaptureDefunInMapFn(self): + num_outputs = 100 - with self.test_session() as sess: - # Simple test that prefetch yields the expected values in the - # expected order. - for buffer_size in [1, 10, 100, 1000]: - sess.run(init_op, feed_dict={buffer_size_placeholder: buffer_size}) - for i in range(100): - self.assertEqual(i * i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + def _build_ds(): - # We can indirectly observe that varying the buffer size has the - # intended effect by observing when `ev` is set (on the 6th - # invocation of `_map_py_func()`). - # NOTE(mrry): We do not test with `buffer_size == - # set_event_during_invocation`, because we must consume at least - # one element to start the prefetching. - for buffer_size in range(1, set_event_during_invocation): - event_will_be_set_after_consuming = ( - set_event_during_invocation - buffer_size + 1) - - ev.clear() - sess.run(init_op, feed_dict={buffer_size_placeholder: buffer_size}) - for i in range(event_will_be_set_after_consuming): - self.assertFalse(ev.is_set()) - self.assertEqual(i * i, sess.run(get_next)) - ev.wait() - for i in range(event_will_be_set_after_consuming, 100): - self.assertEqual(i * i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + @function.Defun(dtypes.int64) + def defun_fn(x): + return constant_op.constant(1000) + math_ops.to_int32(x) - def testReturnList(self): - iterator = (dataset_ops.Dataset.range(10) - .map(lambda x: [x, constant_op.constant(37.0)]) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() + return dataset_ops.Dataset.range(num_outputs).map(defun_fn) - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - self.assertEqual((i, 37.0), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + self.run_core_tests(_build_ds, None, num_outputs) - def testMultiOutputPyFunc(self): - # The `tf.py_func()` op returns a list of tensors for its outputs. - def _map_fn(x_tensor): - def _map_py_func(x): - return x, np.array(37.0, dtype=np.float64) - return script_ops.py_func( - _map_py_func, [x_tensor], [dtypes.int64, dtypes.float64]) - - iterator = (dataset_ops.Dataset.range(10) - .map(_map_fn) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() + def testBuildDefunInMapFn(self): + num_outputs = 100 - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - self.assertEqual((i, 37.0), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + def _build_ds(): - def assertSparseValuesEqual(self, a, b): - self.assertAllEqual(a.indices, b.indices) - self.assertAllEqual(a.values, b.values) - self.assertAllEqual(a.dense_shape, b.dense_shape) + @function.Defun(dtypes.int64) + def defun_fn(x): - def testSparse(self): + @function.Defun(dtypes.int32) + def defun_fn_deep(x): + return constant_op.constant(1000) + math_ops.to_int32(x) - def _sparse(i): - return sparse_tensor.SparseTensorValue( - indices=np.array([[0, 0]]), - values=(i * np.array([1])), - dense_shape=np.array([1, 1])) + return constant_op.constant(11000) + defun_fn_deep(math_ops.to_int32(x)) - iterator = (dataset_ops.Dataset.range(10) - .map(_sparse) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() + return dataset_ops.Dataset.range(num_outputs).map(defun_fn) - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - actual = sess.run(get_next) - self.assertTrue(isinstance(actual, sparse_tensor.SparseTensorValue)) - self.assertSparseValuesEqual(actual, _sparse(i)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + self.run_core_tests(_build_ds, None, num_outputs) - def testSparseChain(self): + def testSparseCore(self): def _sparse(i): return sparse_tensor.SparseTensorValue( @@ -656,58 +239,20 @@ class MapDatasetTest(test.TestCase): values=(i * np.array([1])), dense_shape=np.array([1, 1])) - def _check(i): - self.assertTrue(sparse_tensor.is_sparse(i)) - return sparse_ops.sparse_concat(0, [i, i]) - - iterator = ( - dataset_ops.Dataset.range(10).map(_sparse).map(_check) - .make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - for i in range(10): - actual = sess.run(get_next) - self.assertTrue(isinstance(actual, sparse_tensor.SparseTensorValue)) - self.assertSparseValuesEqual(actual, _check(_sparse(i)).eval()) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testCaptureResourceInMapFn(self): - - def _build_ds(iterator): - - def _map_fn(x): - get_next = iterator.get_next() - return x * get_next - - return dataset_ops.Dataset.range(10).map(_map_fn) - - def _build_graph(): - captured_iterator = dataset_ops.Dataset.range( - 10).make_initializable_iterator() - ds = _build_ds(captured_iterator) - iterator = ds.make_initializable_iterator() - init_op = iterator.initializer - return captured_iterator.initializer, init_op + def _build_ds(num_outputs): + return dataset_ops.Dataset.range(num_outputs).map(_sparse) - with ops.Graph().as_default() as g: - captured_init_op, init_op = _build_graph() - with self.test_session(graph=g) as sess: - sess.run(captured_init_op) - with self.assertRaises(errors.UnimplementedError): - # CapturedFunction does not support capturing IteratorResource. - sess.run(init_op) + num_outputs = 10 + self.run_core_tests(lambda: _build_ds(num_outputs), + lambda: _build_ds(int(num_outputs / 2)), num_outputs) -class MapDatasetSerializationTest( +class ParallelMapDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): def setUp(self): self._tensor_slice_len = 7 - self._num_epochs = 14 + self._num_epochs = 1 self._num_outputs = self._tensor_slice_len * self._num_epochs def _build_ds(self, multiplier=37.0): @@ -718,14 +263,26 @@ class MapDatasetSerializationTest( def _map_fn(x, y, z): return math_ops.square(x), math_ops.square(y), math_ops.square(z) - return (dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) - .repeat(self._num_epochs)) + return (dataset_ops.Dataset.from_tensor_slices(components).map( + _map_fn, num_parallel_calls=3).repeat(self._num_epochs)) + + def _build_ds_with_prefetch(self, multiplier=37.0): + components = (np.arange(self._tensor_slice_len), np.array([[1, 2, 3]]) * + np.arange(self._tensor_slice_len)[:, np.newaxis], + np.array(multiplier) * np.arange(self._tensor_slice_len)) + + def _map_fn(x, y, z): + return math_ops.square(x), math_ops.square(y), math_ops.square(z) + + return (dataset_ops.Dataset.from_tensor_slices(components).map( + _map_fn, num_parallel_calls=3).repeat(self._num_epochs).prefetch(5)) def testSaveRestoreCore(self): - self.run_core_tests( - self._build_ds, - lambda: self._build_ds(multiplier=15.0), - self._num_outputs) + for ds_fn in [self._build_ds, self._build_ds_with_prefetch]: + self.run_core_tests( + ds_fn, + lambda: ds_fn(multiplier=15.0), + self._num_outputs) def testSaveStatefulFunction(self): @@ -735,7 +292,8 @@ class MapDatasetSerializationTest( return random_ops.random_uniform( (), 0, 10, dtype=dtypes.int32) * math_ops.to_int32(x) - return dataset_ops.Dataset.range(100).map(_map_fn) + return dataset_ops.Dataset.range(100).map( + _map_fn, num_parallel_calls=2).prefetch(2) self.verify_error_on_save(_build_ds, 15, errors.InvalidArgumentError) @@ -745,10 +303,20 @@ class MapDatasetSerializationTest( counter_var = variable_scope.get_variable( "counter", (), dtypes.int32, use_resource=True) return (dataset_ops.Dataset.from_tensors(0).repeat(10).map( - lambda _: counter_var.assign_add(1))) + lambda _: counter_var.assign_add(1), + num_parallel_calls=2).prefetch(2)) self.verify_error_on_save(_build_ds, 15, errors.InvalidArgumentError) + def testCaptureConstantInMapFn(self): + + def _build_ds(): + constant_var = constant_op.constant(5) + return (dataset_ops.Dataset.from_tensors(0).repeat(10).map( + lambda x: x + constant_var, num_parallel_calls=2).prefetch(2)) + + self.run_core_tests(_build_ds, None, 10) + def testCaptureDefunInMapFn(self): num_outputs = 100 @@ -758,7 +326,8 @@ class MapDatasetSerializationTest( def defun_fn(x): return constant_op.constant(1000) + math_ops.to_int32(x) - return dataset_ops.Dataset.range(num_outputs).map(defun_fn) + return dataset_ops.Dataset.range(num_outputs).map( + defun_fn, num_parallel_calls=2).prefetch(2) self.run_core_tests(_build_ds, None, num_outputs) @@ -776,7 +345,8 @@ class MapDatasetSerializationTest( return constant_op.constant(11000) + defun_fn_deep(math_ops.to_int32(x)) - return dataset_ops.Dataset.range(num_outputs).map(defun_fn) + return dataset_ops.Dataset.range(num_outputs).map( + defun_fn, num_parallel_calls=2).prefetch(2) self.run_core_tests(_build_ds, None, num_outputs) diff --git a/tensorflow/contrib/data/python/kernel_tests/range_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/range_dataset_op_test.py index 8e6ad061a11752ab7b1ffc13c90b4fa52f67d6aa..80e1cb0041024b68bd5268b5de5d69c88c839896 100644 --- a/tensorflow/contrib/data/python/kernel_tests/range_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/range_dataset_op_test.py @@ -19,160 +19,24 @@ from __future__ import print_function import os +from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base from tensorflow.contrib.data.python.ops import counter -from tensorflow.contrib.data.python.ops import dataset_ops from tensorflow.contrib.data.python.ops import enumerate_ops -from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops -from tensorflow.python.data.ops import iterator_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors 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 gen_dataset_ops from tensorflow.python.ops import io_ops from tensorflow.python.ops import parsing_ops from tensorflow.python.ops import variables -from tensorflow.python.platform import gfile from tensorflow.python.platform import test -from tensorflow.python.training import saver as saver_lib class RangeDatasetTest(test.TestCase): - def tearDown(self): - # Remove all checkpoint files. - prefix = self._iterator_checkpoint_prefix() - pattern = prefix + "*" - files = gfile.Glob(pattern) - map(gfile.Remove, files) - - def testStop(self): - stop = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = dataset_ops.Dataset.range(stop).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op, feed_dict={stop: 5}) - for i in range(5): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testStartStop(self): - start = array_ops.placeholder(dtypes.int64, shape=[]) - stop = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = dataset_ops.Dataset.range(start, - stop).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op, feed_dict={start: 2, stop: 5}) - for i in range(2, 5): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testStartStopStep(self): - start = array_ops.placeholder(dtypes.int64, shape=[]) - stop = array_ops.placeholder(dtypes.int64, shape=[]) - step = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = dataset_ops.Dataset.range(start, stop, - step).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op, feed_dict={start: 2, stop: 10, step: 2}) - for i in range(2, 10, 2): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testZeroStep(self): - start = array_ops.placeholder(dtypes.int64, shape=[]) - stop = array_ops.placeholder(dtypes.int64, shape=[]) - step = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = dataset_ops.Dataset.range(start, stop, - step).make_initializable_iterator() - init_op = iterator.initializer - - with self.test_session() as sess: - with self.assertRaises(errors.InvalidArgumentError): - sess.run(init_op, feed_dict={start: 2, stop: 10, step: 0}) - - def testNegativeStep(self): - start = array_ops.placeholder(dtypes.int64, shape=[]) - stop = array_ops.placeholder(dtypes.int64, shape=[]) - step = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = dataset_ops.Dataset.range(start, stop, - step).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op, feed_dict={start: 2, stop: 10, step: -1}) - # This for loop is a no-op but will ensure that the implementation is - # consistent with range if it ever changes. - for i in range(2, 10, -1): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testStopLessThanStart(self): - start = array_ops.placeholder(dtypes.int64, shape=[]) - stop = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = dataset_ops.Dataset.range(start, - stop).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op, feed_dict={start: 10, stop: 2}) - # This for loop is a no-op but will ensure that the implementation is - # consistent with range if it ever changes. - for i in range(10, 2): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testStopLessThanStartWithPositiveStep(self): - start = array_ops.placeholder(dtypes.int64, shape=[]) - stop = array_ops.placeholder(dtypes.int64, shape=[]) - step = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = dataset_ops.Dataset.range(start, stop, - step).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op, feed_dict={start: 10, stop: 2, step: 2}) - # This for loop is a no-op but will ensure that the implementation is - # consistent with range if it ever changes. - for i in range(10, 2, 2): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testStopLessThanStartWithNegativeStep(self): - start = array_ops.placeholder(dtypes.int64, shape=[]) - stop = array_ops.placeholder(dtypes.int64, shape=[]) - step = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = dataset_ops.Dataset.range(start, stop, - step).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op, feed_dict={start: 10, stop: 2, step: -1}) - for i in range(10, 2, -1): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - def testEnumerateDataset(self): components = (["a", "b"], [1, 2], [37.0, 38]) start = constant_op.constant(20, dtype=dtypes.int64) @@ -216,20 +80,25 @@ class RangeDatasetTest(test.TestCase): self.assertEqual(-1, sess.run(negative_get_next)) self.assertEqual(-2, sess.run(negative_get_next)) - def _iterator_checkpoint_prefix(self): + +class RangeDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): + + def _iterator_checkpoint_prefix_local(self): return os.path.join(self.get_temp_dir(), "iterator") def _save_op(self, iterator_resource): iterator_state_variant = gen_dataset_ops.serialize_iterator( iterator_resource) save_op = io_ops.write_file( - self._iterator_checkpoint_prefix(), + self._iterator_checkpoint_prefix_local(), parsing_ops.serialize_tensor(iterator_state_variant)) return save_op def _restore_op(self, iterator_resource): iterator_state_variant = parsing_ops.parse_tensor( - io_ops.read_file(self._iterator_checkpoint_prefix()), dtypes.variant) + io_ops.read_file(self._iterator_checkpoint_prefix_local()), + dtypes.variant) restore_op = gen_dataset_ops.deserialize_iterator(iterator_resource, iterator_state_variant) return restore_op @@ -283,382 +152,16 @@ class RangeDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - def testSaveRestoreUsingSaverFromMetaGraph(self): - - def _build_graph(start, stop): - iterator = dataset_ops.Dataset.range(start, - stop).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - ops.add_to_collection("iterator_ops", init_op) - ops.add_to_collection("iterator_ops", get_next) - saveable_obj = contrib_iterator_ops.make_saveable_from_iterator(iterator) - # Add the SaveableObject to the `SAVEABLE_OBJECTS` collection - # so that it can be automatically picked up by the Saver. - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable_obj) - saver = saver_lib.Saver() - return init_op, get_next, saver - - start = 2 - stop = 10 - break_point = 5 - path = self._iterator_checkpoint_prefix() - meta_filename = path + ".meta" - - # Execute input pipeline for a few steps and save iterator state. - with ops.Graph().as_default() as g: - init_op, get_next, saver = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - for i in range(start, break_point): - self.assertEqual(i, sess.run(get_next)) - saver.save(sess, path) - - # Build the saver from the MetaGraph using import_meta_graph and - # check that the iterator state is restored. - with ops.Graph().as_default() as g: - saver = saver_lib.import_meta_graph(meta_filename) - init_op, get_next = ops.get_collection("iterator_ops") - with self.test_session(graph=g) as sess: - saver.restore(sess, saver_lib.latest_checkpoint(self.get_temp_dir())) - for i in range(break_point, stop): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testSaveRestoreUsingBuiltSaver(self): - - def _build_graph(start, stop): - iterator = dataset_ops.Dataset.range(start, - stop).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - ops.add_to_collection("iterator_ops", init_op) - ops.add_to_collection("iterator_ops", get_next) - # Add the SaveableObject to the `SAVEABLE_OBJECTS` collection - # so that it can be automatically picked up by the Saver. - saveable_obj = contrib_iterator_ops.make_saveable_from_iterator(iterator) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable_obj) - saver = saver_lib.Saver() - return init_op, get_next, saver - - start = 2 - stop = 10 - stop_new = 15 - break_point = 5 - path = self._iterator_checkpoint_prefix() - - # Execute input pipeline for a few steps and save iterator state. - with ops.Graph().as_default() as g: - init_op, get_next, saver = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - for i in range(start, break_point): - self.assertEqual(i, sess.run(get_next)) - saver.save(sess, path) - - # Manually build a modified Graph and Saver instead of importing - # MetaGraph and verify that original iterator state gets restored. - with ops.Graph().as_default() as g: - init_op, get_next, saver = _build_graph(start, stop_new) - with self.test_session(graph=g) as sess: - saver.restore(sess, saver_lib.latest_checkpoint(self.get_temp_dir())) - for i in range(break_point, stop): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testSaveRestoreUsingSaverThenInit(self): - - def _build_graph(start, stop): - iterator = dataset_ops.Dataset.range(start, - stop).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - ops.add_to_collection("iterator_ops", init_op) - ops.add_to_collection("iterator_ops", get_next) - # Add the SaveableObject to the `SAVEABLE_OBJECTS` collection - # so that it can be automatically picked up by the Saver. - saveable_obj = contrib_iterator_ops.make_saveable_from_iterator(iterator) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable_obj) - saver = saver_lib.Saver() - return init_op, get_next, saver - - start = 2 - stop = 10 - stop_new = 15 - break_point = 5 - path = self._iterator_checkpoint_prefix() - - # Execute input pipeline for a few steps and save iterator state. - with ops.Graph().as_default() as g: - init_op, get_next, saver = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - for i in range(start, break_point): - self.assertEqual(i, sess.run(get_next)) - saver.save(sess, path) - - # Restore iterator state call and then call init_op for the iterator and - # verify that the new iterator hides the restored iterator. - with ops.Graph().as_default() as g: - init_op, get_next, saver = _build_graph(start, stop_new) - with self.test_session(graph=g) as sess: - saver.restore(sess, saver_lib.latest_checkpoint(self.get_temp_dir())) - sess.run(init_op) - for i in range(start, stop_new): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testRestoreWithoutBuildingDatasetGraph(self): - - def _build_graph(start, stop, num_epochs): - dataset = dataset_ops.Dataset.range(start, stop).repeat(num_epochs) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - save_op = self._save_op(iterator._iterator_resource) - restore_op = self._restore_op(iterator._iterator_resource) - return init_op, get_next, save_op, restore_op - - # Saving and restoring in different sessions. - start = 2 - stop = 10 - num_epochs = 5 - break_point = 5 - break_epoch = 3 - with ops.Graph().as_default() as g: - init_op, get_next, save_op, _ = _build_graph(start, stop, num_epochs) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - for _ in range(break_epoch): - for i in range(start, stop): - self.assertEqual(i, sess.run(get_next)) - for i in range(start, break_point): - self.assertEqual(i, sess.run(get_next)) - sess.run(save_op) - - with ops.Graph().as_default() as g: - # Create an empty IteratorResource and restore the Iterator into it. - output_types = dtypes.int64 - output_shapes = tensor_shape.scalar() - iterator = iterator_ops.Iterator.from_structure(output_types, - output_shapes) - restore_op = self._restore_op(iterator._iterator_resource) - get_next = iterator.get_next() - with self.test_session(graph=g) as sess: - sess.run(restore_op) - for i in range(break_point, stop): - self.assertEqual(i, sess.run(get_next)) - for _ in range(break_epoch + 1, num_epochs): - for i in range(start, stop): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testRestoreInModifiedGraph(self): + def _build_range_dataset(self, start, stop): + return dataset_ops.Dataset.range(start, stop) - def _build_graph(start, stop): - dataset = dataset_ops.Dataset.range(start, stop) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - save_op = self._save_op(iterator._iterator_resource) - restore_op = self._restore_op(iterator._iterator_resource) - return init_op, get_next, save_op, restore_op - - # Saving and restoring in different sessions. + def testRangeCore(self): start = 2 stop = 10 stop_1 = 8 - break_point = 5 - with ops.Graph().as_default() as g: - init_op, get_next, save_op, _ = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - for i in range(start, break_point): - self.assertEqual(i, sess.run(get_next)) - sess.run(save_op) - - with ops.Graph().as_default() as g: - # Intentionally build a graph with a different value for stop to make sure - # the original dataset graph is actually getting loaded. - init_op, get_next, _, restore_op = _build_graph(start, stop_1) - with self.test_session(graph=g) as sess: - sess.run(restore_op) - for i in range(break_point, stop): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testInitThenRestore(self): - # Note: Calling init_op before restore_op is redundant. This test just makes - # sure we do not fail if restore is called on an already initialized - # iterator resource. - - def _build_graph(start, stop): - dataset = dataset_ops.Dataset.range(start, stop) - iterator = dataset.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - save_op = self._save_op(iterator._iterator_resource) - restore_op = self._restore_op(iterator._iterator_resource) - return init_op, get_next, save_op, restore_op - - # Saving and restoring in different sessions. - start = 2 - stop = 10 - break_point = 5 - with ops.Graph().as_default() as g: - init_op, get_next, save_op, _ = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - for i in range(start, break_point): - self.assertEqual(i, sess.run(get_next)) - sess.run(save_op) - - with ops.Graph().as_default() as g: - init_op, get_next, _, restore_op = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(init_op) - sess.run(restore_op) - for i in range(break_point, stop): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testMultipleSaves(self): - - def _build_graph(start, stop): - iterator = dataset_ops.Dataset.range(start, - stop).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - save_op = self._save_op(iterator._iterator_resource) - restore_op = self._restore_op(iterator._iterator_resource) - return init_op, get_next, save_op, restore_op - - start = 2 - stop = 10 - break_point1 = 5 - break_point2 = 7 - - with ops.Graph().as_default() as g: - init_op, get_next, save_op, _ = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - for i in range(start, break_point1): - self.assertEqual(i, sess.run(get_next)) - sess.run(save_op) - - with ops.Graph().as_default() as g: - init_op, get_next, save_op, restore_op = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(restore_op) - for i in range(break_point1, break_point2): - self.assertEqual(i, sess.run(get_next)) - sess.run(save_op) - - break_point2 = 7 - with ops.Graph().as_default() as g: - init_op, get_next, save_op, restore_op = _build_graph(start, stop) - with self.test_session(graph=g) as sess: - sess.run(restore_op) - for i in range(break_point2, stop): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testSaveRestoreWithRepeat(self): - - def _build_graph(start, stop, num_epochs): - iterator = dataset_ops.Dataset.range( - start, stop).repeat(num_epochs).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - save_op = self._save_op(iterator._iterator_resource) - restore_op = self._restore_op(iterator._iterator_resource) - return init_op, get_next, save_op, restore_op - - start = 2 - stop = 10 - num_epochs = 5 - break_range = 5 - break_epoch = 3 - with ops.Graph().as_default() as g: - init_op, get_next, save_op, restore_op = _build_graph( - start, stop, num_epochs) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - # Note: There is no checkpoint saved currently so a NotFoundError is - # raised. - with self.assertRaises(errors.NotFoundError): - sess.run(restore_op) - for _ in range(break_epoch - 1): - for i in range(start, stop): - self.assertEqual(i, sess.run(get_next)) - for i in range(start, break_range): - self.assertEqual(i, sess.run(get_next)) - sess.run(save_op) - - with ops.Graph().as_default() as g: - init_op, get_next, _, restore_op = _build_graph(start, stop, num_epochs) - with self.test_session(graph=g) as sess: - sess.run(restore_op) - for i in range(break_range, stop): - self.assertEqual(i, sess.run(get_next)) - for _ in range(break_epoch, num_epochs): - for i in range(start, stop): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testSaveRestoreExhaustedIterator(self): - - def _build_graph(start, stop, num_epochs): - iterator = dataset_ops.Dataset.range( - start, stop).repeat(num_epochs).make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - save_op = self._save_op(iterator._iterator_resource) - restore_op = self._restore_op(iterator._iterator_resource) - return init_op, get_next, save_op, restore_op - - start = 2 - stop = 10 - num_epochs = 5 - with ops.Graph().as_default() as g: - init_op, get_next, save_op, restore_op = _build_graph( - start, stop, num_epochs) - with self.test_session(graph=g) as sess: - sess.run(variables.global_variables_initializer()) - sess.run(init_op) - # Note: There is no checkpoint saved currently so a NotFoundError is - # raised. - with self.assertRaises(errors.NotFoundError): - sess.run(restore_op) - for _ in range(num_epochs): - for i in range(start, stop): - self.assertEqual(i, sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - sess.run(save_op) - - with ops.Graph().as_default() as g: - init_op, get_next, _, restore_op = _build_graph(start, stop, num_epochs) - with self.test_session(graph=g) as sess: - sess.run(restore_op) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) + self.run_core_tests(lambda: self._build_range_dataset(start, stop), + lambda: self._build_range_dataset(start, stop_1), + stop - start) if __name__ == "__main__": diff --git a/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py b/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py index 1c42a3d855bc16c21e385d7108c3106884ae4f5e..6efe97444a375febc550ff3a3ea04bcd9330a3a5 100644 --- a/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py @@ -26,6 +26,7 @@ from tensorflow.contrib.data.python.ops import readers from tensorflow.core.example import example_pb2 from tensorflow.core.example import feature_pb2 from tensorflow.python.data.ops import iterator_ops +from tensorflow.python.data.ops import readers as core_readers from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -76,101 +77,12 @@ class TextLineDatasetTestBase(test.TestCase): return filenames -class TextLineDatasetTest(TextLineDatasetTestBase): - - def _testTextLineDataset(self, compression_type=None): - test_filenames = self._createFiles( - 2, 5, crlf=True, compression_type=compression_type) - filenames = array_ops.placeholder(dtypes.string, shape=[None]) - num_epochs = array_ops.placeholder(dtypes.int64, shape=[]) - batch_size = array_ops.placeholder(dtypes.int64, shape=[]) - - repeat_dataset = readers.TextLineDataset( - filenames, compression_type=compression_type).repeat(num_epochs) - batch_dataset = repeat_dataset.batch(batch_size) - - iterator = iterator_ops.Iterator.from_structure(batch_dataset.output_types) - init_op = iterator.make_initializer(repeat_dataset) - init_batch_op = iterator.make_initializer(batch_dataset) - get_next = iterator.get_next() - - with self.test_session() as sess: - # Basic test: read from file 0. - sess.run( - init_op, feed_dict={filenames: [test_filenames[0]], - num_epochs: 1}) - for i in range(5): - self.assertEqual(self._lineText(0, i), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Basic test: read from file 1. - sess.run( - init_op, feed_dict={filenames: [test_filenames[1]], - num_epochs: 1}) - for i in range(5): - self.assertEqual(self._lineText(1, i), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Basic test: read from both files. - sess.run(init_op, feed_dict={filenames: test_filenames, num_epochs: 1}) - for j in range(2): - for i in range(5): - self.assertEqual(self._lineText(j, i), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test repeated iteration through both files. - sess.run(init_op, feed_dict={filenames: test_filenames, num_epochs: 10}) - for _ in range(10): - for j in range(2): - for i in range(5): - self.assertEqual(self._lineText(j, i), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test batched and repeated iteration through both files. - sess.run( - init_batch_op, - feed_dict={filenames: test_filenames, - num_epochs: 10, - batch_size: 5}) - for _ in range(10): - self.assertAllEqual([self._lineText(0, i) for i in range(5)], - sess.run(get_next)) - self.assertAllEqual([self._lineText(1, i) for i in range(5)], - sess.run(get_next)) - - def testTextLineDatasetNoCompression(self): - self._testTextLineDataset() - - def testTextLineDatasetGzipCompression(self): - self._testTextLineDataset(compression_type="GZIP") - - def testTextLineDatasetZlibCompression(self): - self._testTextLineDataset(compression_type="ZLIB") - - def testTextLineDatasetBuffering(self): - test_filenames = self._createFiles(2, 5, crlf=True) - - repeat_dataset = readers.TextLineDataset(test_filenames, buffer_size=10) - iterator = repeat_dataset.make_one_shot_iterator() - - with self.test_session() as sess: - for j in range(2): - for i in range(5): - self.assertEqual(self._lineText(j, i), sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - class TextLineDatasetSerializationTest( TextLineDatasetTestBase, dataset_serialization_test_base.DatasetSerializationTestBase): def _build_iterator_graph(self, test_filenames, compression_type=None): - return readers.TextLineDataset( + return core_readers.TextLineDataset( test_filenames, compression_type=compression_type, buffer_size=10) def testTextLineCore(self): @@ -217,101 +129,13 @@ class FixedLengthRecordReaderTestBase(test.TestCase): return filenames -class FixedLengthRecordReaderTest(FixedLengthRecordReaderTestBase): - - def testFixedLengthRecordDataset(self): - test_filenames = self._createFiles() - filenames = array_ops.placeholder(dtypes.string, shape=[None]) - num_epochs = array_ops.placeholder(dtypes.int64, shape=[]) - batch_size = array_ops.placeholder(dtypes.int64, shape=[]) - - repeat_dataset = (readers.FixedLengthRecordDataset( - filenames, self._record_bytes, self._header_bytes, self._footer_bytes) - .repeat(num_epochs)) - batch_dataset = repeat_dataset.batch(batch_size) - - iterator = iterator_ops.Iterator.from_structure(batch_dataset.output_types) - init_op = iterator.make_initializer(repeat_dataset) - init_batch_op = iterator.make_initializer(batch_dataset) - get_next = iterator.get_next() - - with self.test_session() as sess: - # Basic test: read from file 0. - sess.run( - init_op, feed_dict={filenames: [test_filenames[0]], - num_epochs: 1}) - for i in range(self._num_records): - self.assertEqual(self._record(0, i), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Basic test: read from file 1. - sess.run( - init_op, feed_dict={filenames: [test_filenames[1]], - num_epochs: 1}) - for i in range(self._num_records): - self.assertEqual(self._record(1, i), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Basic test: read from both files. - sess.run(init_op, feed_dict={filenames: test_filenames, num_epochs: 1}) - for j in range(self._num_files): - for i in range(self._num_records): - self.assertEqual(self._record(j, i), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test repeated iteration through both files. - sess.run(init_op, feed_dict={filenames: test_filenames, num_epochs: 10}) - for _ in range(10): - for j in range(self._num_files): - for i in range(self._num_records): - self.assertEqual(self._record(j, i), sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test batched and repeated iteration through both files. - sess.run( - init_batch_op, - feed_dict={ - filenames: test_filenames, - num_epochs: 10, - batch_size: self._num_records - }) - for _ in range(10): - for j in range(self._num_files): - self.assertAllEqual( - [self._record(j, i) for i in range(self._num_records)], - sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testFixedLengthRecordDatasetBuffering(self): - test_filenames = self._createFiles() - dataset = readers.FixedLengthRecordDataset( - test_filenames, - self._record_bytes, - self._header_bytes, - self._footer_bytes, - buffer_size=10) - iterator = dataset.make_one_shot_iterator() - - with self.test_session() as sess: - for j in range(self._num_files): - for i in range(self._num_records): - self.assertEqual(self._record(j, i), sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - class FixedLengthRecordDatasetSerializationTest( FixedLengthRecordReaderTestBase, dataset_serialization_test_base.DatasetSerializationTestBase): def _build_iterator_graph(self, num_epochs, compression_type=None): filenames = self._createFiles() - return readers.FixedLengthRecordDataset( + return core_readers.FixedLengthRecordDataset( filenames, self._record_bytes, self._header_bytes, self._footer_bytes).repeat(num_epochs) @@ -338,9 +162,8 @@ class TFRecordDatasetTestBase(test.TestCase): self.compression_type = array_ops.placeholder_with_default("", shape=[]) self.batch_size = array_ops.placeholder(dtypes.int64, shape=[]) - repeat_dataset = readers.TFRecordDataset(self.filenames, - self.compression_type).repeat( - self.num_epochs) + repeat_dataset = core_readers.TFRecordDataset( + self.filenames, self.compression_type).repeat(self.num_epochs) batch_dataset = repeat_dataset.batch(self.batch_size) iterator = iterator_ops.Iterator.from_structure(batch_dataset.output_types) @@ -363,129 +186,6 @@ class TFRecordDatasetTestBase(test.TestCase): return filenames -class TFRecordDatasetTest(TFRecordDatasetTestBase): - - def testReadOneEpoch(self): - with self.test_session() as sess: - # Basic test: read from file 0. - sess.run( - self.init_op, - feed_dict={ - self.filenames: [self.test_filenames[0]], - self.num_epochs: 1 - }) - for i in range(self._num_records): - self.assertAllEqual(self._record(0, i), sess.run(self.get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(self.get_next) - - # Basic test: read from file 1. - sess.run( - self.init_op, - feed_dict={ - self.filenames: [self.test_filenames[1]], - self.num_epochs: 1 - }) - for i in range(self._num_records): - self.assertAllEqual(self._record(1, i), sess.run(self.get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(self.get_next) - - # Basic test: read from both files. - sess.run( - self.init_op, - feed_dict={self.filenames: self.test_filenames, - self.num_epochs: 1}) - for j in range(self._num_files): - for i in range(self._num_records): - self.assertAllEqual(self._record(j, i), sess.run(self.get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(self.get_next) - - def testReadTenEpochs(self): - with self.test_session() as sess: - sess.run( - self.init_op, - feed_dict={self.filenames: self.test_filenames, - self.num_epochs: 10}) - for _ in range(10): - for j in range(self._num_files): - for i in range(self._num_records): - self.assertAllEqual(self._record(j, i), sess.run(self.get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(self.get_next) - - def testReadTenEpochsOfBatches(self): - with self.test_session() as sess: - sess.run( - self.init_batch_op, - feed_dict={ - self.filenames: self.test_filenames, - self.num_epochs: 10, - self.batch_size: self._num_records - }) - for _ in range(10): - for j in range(self._num_files): - values = sess.run(self.get_next) - self.assertAllEqual( - [self._record(j, i) for i in range(self._num_records)], values) - with self.assertRaises(errors.OutOfRangeError): - sess.run(self.get_next) - - def testReadZlibFiles(self): - zlib_files = [] - for i, fn in enumerate(self.test_filenames): - with open(fn, "rb") as f: - cdata = zlib.compress(f.read()) - - zfn = os.path.join(self.get_temp_dir(), "tfrecord_%s.z" % i) - with open(zfn, "wb") as f: - f.write(cdata) - zlib_files.append(zfn) - - with self.test_session() as sess: - sess.run( - self.init_op, - feed_dict={self.filenames: zlib_files, - self.compression_type: "ZLIB"}) - for j in range(self._num_files): - for i in range(self._num_records): - self.assertAllEqual(self._record(j, i), sess.run(self.get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(self.get_next) - - def testReadGzipFiles(self): - gzip_files = [] - for i, fn in enumerate(self.test_filenames): - with open(fn, "rb") as f: - gzfn = os.path.join(self.get_temp_dir(), "tfrecord_%s.gz" % i) - with gzip.GzipFile(gzfn, "wb") as gzf: - gzf.write(f.read()) - gzip_files.append(gzfn) - - with self.test_session() as sess: - sess.run( - self.init_op, - feed_dict={self.filenames: gzip_files, - self.compression_type: "GZIP"}) - for j in range(self._num_files): - for i in range(self._num_records): - self.assertAllEqual(self._record(j, i), sess.run(self.get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(self.get_next) - - def testReadWithBuffer(self): - one_mebibyte = 2**20 - d = readers.TFRecordDataset(self.test_filenames, buffer_size=one_mebibyte) - iterator = d.make_one_shot_iterator() - with self.test_session() as sess: - for j in range(self._num_files): - for i in range(self._num_records): - self.assertAllEqual(self._record(j, i), sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - class TFRecordDatasetSerializationTest( TFRecordDatasetTestBase, dataset_serialization_test_base.DatasetSerializationTestBase): @@ -517,7 +217,7 @@ class TFRecordDatasetSerializationTest( gzip_files.append(gzfn) filenames = gzip_files - return readers.TFRecordDataset( + return core_readers.TFRecordDataset( filenames, compression_type, buffer_size=buffer_size).repeat(num_epochs).batch(batch_size) @@ -575,7 +275,7 @@ class ReadBatchFeaturesTest(test.TestCase): "record": parsing_ops.FixedLenFeature([], dtypes.int64), "keywords": parsing_ops.VarLenFeature(dtypes.string) }, - reader=readers.TFRecordDataset, + reader=core_readers.TFRecordDataset, randomize_input=False, num_epochs=self.num_epochs) @@ -714,12 +414,11 @@ class ReadBatchFeaturesTest(test.TestCase): self._next_actual_batch(sess) def testReadWithEquivalentDataset(self): - # TODO(mrry): Add support for tf.SparseTensor as a Dataset component. features = { "file": parsing_ops.FixedLenFeature([], dtypes.int64), "record": parsing_ops.FixedLenFeature([], dtypes.int64), } - dataset = (readers.TFRecordDataset(self.test_filenames) + dataset = (core_readers.TFRecordDataset(self.test_filenames) .map(lambda x: parsing_ops.parse_single_example(x, features)) .repeat(10).batch(2)) iterator = dataset.make_initializable_iterator() diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py index 0ac8d7359f7234d98167277724780bf31555e6fb..3c7b46629edb13459766b5ef3f392e8d00ad4db8 100644 --- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py @@ -19,8 +19,8 @@ from __future__ import print_function import numpy as np -from tensorflow.contrib.data.python.ops import dataset_ops from tensorflow.contrib.data.python.ops import resampling +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import errors from tensorflow.python.ops import string_ops from tensorflow.python.platform import test diff --git a/tensorflow/contrib/data/python/kernel_tests/sequence_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/sequence_dataset_op_test.py index 1a26da82e533ec01106ea10525c1cd96627c34fb..36ddf3004237ed042f21d691d83eafbaa20621e6 100644 --- a/tensorflow/contrib/data/python/kernel_tests/sequence_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/sequence_dataset_op_test.py @@ -20,194 +20,10 @@ from __future__ import print_function import numpy as np from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.ops import array_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.platform import test -class SequenceDatasetTest(test.TestCase): - - def testRepeatTensorDataset(self): - """Test a dataset that repeats its input multiple times.""" - components = (np.array(1), np.array([1, 2, 3]), np.array(37.0)) - # This placeholder can be fed when dataset-definition subgraph - # runs (i.e. `init_op` below) to configure the number of - # repetitions used in a particular iterator. - count_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) - - iterator = (dataset_ops.Dataset.from_tensors(components) - .repeat(count_placeholder).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - # Test a finite repetition. - sess.run(init_op, feed_dict={count_placeholder: 3}) - for _ in range(3): - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertAllEqual(component, result_component) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test a different finite repetition. - sess.run(init_op, feed_dict={count_placeholder: 7}) - for _ in range(7): - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertAllEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test an empty repetition. - sess.run(init_op, feed_dict={count_placeholder: 0}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Test an infinite repetition. - # NOTE(mrry): There's not a good way to test that the sequence - # actually is infinite. - sess.run(init_op, feed_dict={count_placeholder: -1}) - for _ in range(17): - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertAllEqual(component, result_component) - - def testTakeTensorDataset(self): - components = (np.arange(10),) - count_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) - - iterator = (dataset_ops.Dataset.from_tensor_slices(components) - .take(count_placeholder).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape[1:] for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - # Take fewer than input size - sess.run(init_op, feed_dict={count_placeholder: 4}) - for i in range(4): - results = sess.run(get_next) - self.assertAllEqual(results, components[0][i:i+1]) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Take more than input size - sess.run(init_op, feed_dict={count_placeholder: 25}) - for i in range(10): - results = sess.run(get_next) - self.assertAllEqual(results, components[0][i:i+1]) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Take all of input - sess.run(init_op, feed_dict={count_placeholder: -1}) - for i in range(10): - results = sess.run(get_next) - self.assertAllEqual(results, components[0][i:i+1]) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Take nothing - sess.run(init_op, feed_dict={count_placeholder: 0}) - - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testSkipTensorDataset(self): - components = (np.arange(10),) - count_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) - - iterator = (dataset_ops.Dataset.from_tensor_slices(components) - .skip(count_placeholder).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape[1:] for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - # Skip fewer than input size, we should skip - # the first 4 elements and then read the rest. - sess.run(init_op, feed_dict={count_placeholder: 4}) - for i in range(4, 10): - results = sess.run(get_next) - self.assertAllEqual(results, components[0][i:i+1]) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Skip more than input size: get nothing. - sess.run(init_op, feed_dict={count_placeholder: 25}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Skip exactly input size. - sess.run(init_op, feed_dict={count_placeholder: 10}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Set -1 for 'count': skip the entire dataset. - sess.run(init_op, feed_dict={count_placeholder: -1}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Skip nothing - sess.run(init_op, feed_dict={count_placeholder: 0}) - for i in range(0, 10): - results = sess.run(get_next) - self.assertAllEqual(results, components[0][i:i+1]) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testRepeatRepeatTensorDataset(self): - """Test the composition of repeat datasets.""" - components = (np.array(1), np.array([1, 2, 3]), np.array(37.0)) - inner_count = array_ops.placeholder(dtypes.int64, shape=[]) - outer_count = array_ops.placeholder(dtypes.int64, shape=[]) - - iterator = (dataset_ops.Dataset.from_tensors(components).repeat(inner_count) - .repeat(outer_count).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([c.shape for c in components], - [t.shape for t in get_next]) - - with self.test_session() as sess: - sess.run(init_op, feed_dict={inner_count: 7, outer_count: 14}) - for _ in range(7 * 14): - results = sess.run(get_next) - for component, result_component in zip(components, results): - self.assertAllEqual(component, result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testRepeatEmptyDataset(self): - """Test that repeating an empty dataset does not hang.""" - iterator = (dataset_ops.Dataset.from_tensors(0).repeat(10).skip(10) - .repeat(-1).make_initializable_iterator()) - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - sess.run(init_op) - with self.assertRaisesRegexp( - errors.OutOfRangeError, - "Attempted to repeat an empty dataset infinitely."): - sess.run(get_next) - - class SequenceDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization_integration_test.py b/tensorflow/contrib/data/python/kernel_tests/serialization_integration_test.py new file mode 100644 index 0000000000000000000000000000000000000000..0a6b74dc3eb80a6168117beed06935737198cecb --- /dev/null +++ b/tensorflow/contrib/data/python/kernel_tests/serialization_integration_test.py @@ -0,0 +1,85 @@ +# 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. +# ============================================================================== +"""Integration test for input pipeline serialization.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os + +from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import ops +from tensorflow.python.platform import test +from tensorflow.python.training import saver as saver_lib + + +class MultipleInputPipelinesTest(test.TestCase): + + def _build_input_pipeline(self, name, num_outputs): + with ops.name_scope(name): + ds = dataset_ops.Dataset.range(num_outputs).shuffle( + 10, reshuffle_each_iteration=False).prefetch(10) + iterator = ds.make_initializable_iterator() + saveable = contrib_iterator_ops.make_saveable_from_iterator(iterator) + ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) + return iterator.initializer, iterator.get_next() + + def _build_graph(self, num_pipelines, num_outputs): + init_ops = [] + get_next_ops = [] + for i in range(num_pipelines): + name = "input_pipeline_%d" % i + init_op, get_next_op = self._build_input_pipeline(name, num_outputs) + init_ops.append(init_op) + get_next_ops.append(get_next_op) + saver = saver_lib.Saver() + return init_ops, get_next_ops, saver + + def _ckpt_path(self): + return os.path.join(self.get_temp_dir(), "iterator") + + def testConcurrentSaves(self): + num_pipelines = 100 + num_outputs = 100 + break_point = 10 + all_outputs = [[] for _ in range(num_pipelines)] + with ops.Graph().as_default() as g: + init_ops, get_next_ops, saver = self._build_graph(num_pipelines, + num_outputs) + with self.test_session(graph=g) as sess: + sess.run(init_ops) + for _ in range(break_point): + output = sess.run(get_next_ops) + for i in range(num_pipelines): + all_outputs[i].append(output[i]) + saver.save(sess, self._ckpt_path()) + + with ops.Graph().as_default() as g: + init_ops, get_next_ops, saver = self._build_graph(num_pipelines, + num_outputs) + with self.test_session(graph=g) as sess: + saver.restore(sess, self._ckpt_path()) + for _ in range(num_outputs - break_point): + output = sess.run(get_next_ops) + for i in range(num_pipelines): + all_outputs[i].append(output[i]) + + for output in all_outputs: + self.assertSequenceEqual(sorted(output), range(num_outputs)) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/shard_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/shard_dataset_op_test.py deleted file mode 100644 index 0b3c32c06eb1d69244c9a02ca4ba571769f13f40..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/data/python/kernel_tests/shard_dataset_op_test.py +++ /dev/null @@ -1,111 +0,0 @@ -# Copyright 2017 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 the experimental input pipeline ops.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.python.framework import errors -from tensorflow.python.platform import test - - -class ShardDatasetOpTest(test.TestCase): - - def testSimpleCase(self): - dataset = dataset_ops.Dataset.range(10).shard(5, 2) - iterator = dataset.make_one_shot_iterator() - - with self.test_session() as sess: - self.assertEqual(2, sess.run(iterator.get_next())) - self.assertEqual(7, sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - def testNestedData(self): - dataset_a = dataset_ops.Dataset.range(10) - dataset_b = dataset_ops.Dataset.range(10, 0, -1) - dataset = dataset_ops.Dataset.zip((dataset_a, dataset_b)).shard(5, 2) - iterator = dataset.make_one_shot_iterator() - - with self.test_session() as sess: - self.assertEqual((2, 8), sess.run(iterator.get_next())) - self.assertEqual((7, 3), sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - def testOffsetZero(self): - dataset = dataset_ops.Dataset.range(10).shard(5, 0) - iterator = dataset.make_one_shot_iterator() - - with self.test_session() as sess: - self.assertEqual(0, sess.run(iterator.get_next())) - self.assertEqual(5, sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - def testOffsetGreaterNumShards(self): - with self.assertRaises(ValueError): - dataset_ops.Dataset.range(10).shard(5, 7) - - def testNegativeOffset(self): - with self.assertRaises(ValueError): - dataset_ops.Dataset.range(10).shard(5, -3) - - def testNegativeNumShards(self): - with self.assertRaises(ValueError): - dataset_ops.Dataset.range(10).shard(-3, 1) - - def testZeroNumShards(self): - with self.assertRaises(ValueError): - dataset_ops.Dataset.range(10).shard(0, 1) - - def testIteratorEndsBeforeFirstElem(self): - dataset = dataset_ops.Dataset.range(1).shard(5, 2) - iterator = dataset.make_one_shot_iterator() - - with self.test_session() as sess: - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - def testLargerWorkerPool(self): - dataset = dataset_ops.Dataset.range(10).shard(7, 5) - iterator = dataset.make_one_shot_iterator() - with self.test_session() as sess: - self.assertEqual(5, sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - def testIndexEqualsNumShards(self): - dataset = dataset_ops.Dataset.range(10).shard(5, 4) - iterator = dataset.make_one_shot_iterator() - with self.test_session() as sess: - self.assertEqual(4, sess.run(iterator.get_next())) - self.assertEqual(9, sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - def testIndexEqualsNumShards2(self): - dataset = dataset_ops.Dataset.range(10).shard(4, 3) - iterator = dataset.make_one_shot_iterator() - with self.test_session() as sess: - self.assertEqual(3, sess.run(iterator.get_next())) - self.assertEqual(7, sess.run(iterator.get_next())) - with self.assertRaises(errors.OutOfRangeError): - sess.run(iterator.get_next()) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/shuffle_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/shuffle_dataset_op_test.py index 72745ec7525ad0578934fb2051018f6531938088..bcc644c0971854d948025009dc7add2fea214048 100644 --- a/tensorflow/contrib/data/python/kernel_tests/shuffle_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/shuffle_dataset_op_test.py @@ -17,463 +17,59 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections -import os - import numpy as np from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base -from tensorflow.contrib.data.python.ops import dataset_ops as contrib_dataset_ops -from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops from tensorflow.contrib.data.python.ops import shuffle_ops from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.ops import iterator_ops -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.platform import gfile from tensorflow.python.platform import test -from tensorflow.python.training import saver as saver_lib - - -class ShuffleDatasetTest(test.TestCase): - - def testShuffleDataset(self): - components = ( - np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]), - np.array([9.0, 10.0, 11.0, 12.0]) - ) - count_placeholder = array_ops.placeholder_with_default( - constant_op.constant(5, dtypes.int64), shape=[]) - buffer_size_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) - seed_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) - - repeat_dataset = ( - contrib_dataset_ops.Dataset.from_tensor_slices(components) - .repeat(count_placeholder)) - - shuffle_dataset = repeat_dataset.shuffle(buffer_size_placeholder, - seed_placeholder) - - self.assertEqual(tuple([c.shape[1:] for c in components]), - shuffle_dataset.output_shapes) - - # Create initialization ops for iterators without and with - # shuffling, respectively. - iterator = iterator_ops.Iterator.from_structure( - shuffle_dataset.output_types, shuffle_dataset.output_shapes) - init_fifo_op = iterator.make_initializer(repeat_dataset) - init_shuffle_op = iterator.make_initializer(shuffle_dataset) - - get_next = iterator.get_next() - - with self.test_session() as sess: - # First run without shuffling to collect the "ground truth". - sess.run(init_fifo_op) - unshuffled_elements = [] - for _ in range(20): - unshuffled_elements.append(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - # Assert that the shuffled dataset has the same elements as the - # "ground truth". - sess.run( - init_shuffle_op, - feed_dict={buffer_size_placeholder: 100, - seed_placeholder: 37}) - shuffled_elements = [] - for _ in range(20): - shuffled_elements.append(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - self.assertAllEqual( - sorted(unshuffled_elements), sorted(shuffled_elements)) - - # Assert that shuffling twice with the same seeds gives the same sequence. - sess.run( - init_shuffle_op, - feed_dict={buffer_size_placeholder: 100, - seed_placeholder: 37}) - reshuffled_elements_same_seed = [] - for _ in range(20): - reshuffled_elements_same_seed.append(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - self.assertEqual(shuffled_elements, reshuffled_elements_same_seed) - - # Assert that shuffling twice with a different seed gives a different - # permutation of the same elements. - sess.run( - init_shuffle_op, - feed_dict={buffer_size_placeholder: 100, - seed_placeholder: 1037}) - reshuffled_elements_different_seed = [] - for _ in range(20): - reshuffled_elements_different_seed.append(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - self.assertNotEqual(shuffled_elements, reshuffled_elements_different_seed) - self.assertAllEqual( - sorted(shuffled_elements), sorted(reshuffled_elements_different_seed)) - - # Assert that the shuffled dataset has the same elements as the - # "ground truth" when the buffer size is smaller than the input - # dataset. - sess.run( - init_shuffle_op, - feed_dict={buffer_size_placeholder: 2, - seed_placeholder: 37}) - reshuffled_elements_small_buffer = [] - for _ in range(20): - reshuffled_elements_small_buffer.append(sess.run(get_next)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - self.assertAllEqual( - sorted(unshuffled_elements), sorted(reshuffled_elements_small_buffer)) - - # Test the case of shuffling an empty dataset. - sess.run(init_shuffle_op, feed_dict={buffer_size_placeholder: 2, - seed_placeholder: 37, - count_placeholder: 0}) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testDefaultArguments(self): - components = [0, 1, 2, 3, 4] - iterator = ( - contrib_dataset_ops.Dataset.from_tensor_slices(components).shuffle(5) - .repeat().make_one_shot_iterator()) - - get_next = iterator.get_next() - - with self.test_session() as sess: - counts = collections.defaultdict(lambda: 0) - for _ in range(10): - for _ in range(5): - counts[sess.run(get_next)] += 1 - for i in range(5): - self.assertEqual(10, counts[i]) - - -class ShuffleDatasetSerializationTest(test.TestCase): - def tearDown(self): - # Remove all checkpoint files. - prefix = self._ckpt_path() - pattern = prefix + "*" - files = gfile.Glob(pattern) - map(gfile.Remove, files) +class ShuffleDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): - def _build_graph(self, - range_limit=10, - num_repeats=5, - buffer_size=5, - seed=None, - reshuffle_each_iteration=None, - build_saveable=True): - iterator = dataset_ops.Dataset.range(range_limit).shuffle( + def _build_shuffle_dataset( + self, + range_limit=10, + num_repeats=5, + buffer_size=5, + seed=None, + reshuffle_each_iteration=None, + ): + return dataset_ops.Dataset.range(range_limit).shuffle( buffer_size, seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration).repeat( - num_repeats).make_initializable_iterator() - if build_saveable: - saveable = contrib_iterator_ops.make_saveable_from_iterator(iterator) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) - init_op = iterator.initializer - get_next = iterator.get_next() - ops.add_to_collection("iterator_ops", init_op) - ops.add_to_collection("iterator_ops", get_next) - saver = saver_lib.Saver(allow_empty=True) - return init_op, get_next, saver - - def _ckpt_path(self): - return os.path.join(self.get_temp_dir(), "iterator") + reshuffle_each_iteration=reshuffle_each_iteration).repeat(num_repeats) - def _latest_ckpt(self): - return saver_lib.latest_checkpoint(self.get_temp_dir()) + def testShuffleCore(self): - def _save(self, sess, saver): - saver.save(sess, self._ckpt_path()) - - def _restore(self, saver, sess): - saver.restore(sess, self._latest_ckpt()) - - def _import_meta_graph(self): - meta_file_path = self._ckpt_path() + ".meta" - return saver_lib.import_meta_graph(meta_file_path) - - def _testReadWithBreaks(self, break_points, init_before_restore=False): seed = 55 range_limit = 10 num_repeats = 5 num_outputs = range_limit * num_repeats buffer_sizes = [1, 3, 8, 10, 25, 50] reshuffle_each_iteration = False + # pylint: disable=cell-var-from-loop + # pylint: disable=g-long-lambda for buffer_size in buffer_sizes: - expected = [] - actual = [] - # Generate the ground truth. - with ops.Graph().as_default() as g: - g.seed = 10 - init_op, get_next_op, _ = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration) - with self.test_session(graph=g) as sess: - sess.run(init_op) - for _ in range(num_outputs): - expected.append(sess.run(get_next_op)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - - # Run and checkpoint after first break_point. - with ops.Graph().as_default() as g: - g.seed = 10 - init_op, get_next_op, saver = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration) - with self.test_session(graph=g) as sess: - sess.run(init_op) - for _ in range(break_points[0]): - actual.append(sess.run(get_next_op)) - self._save(sess, saver) - - # Load from checkpoint and continue running while stopping at each - # subsequent checkpoint. - for i in range(len(break_points)): - with ops.Graph().as_default() as g: - saver = self._import_meta_graph() - init_op, get_next_op = ops.get_collection("iterator_ops") - with self.test_session(graph=g) as sess: - if init_before_restore: - sess.run(init_op) - self._restore(saver, sess) - start = break_points[i] - end = break_points[ - i + 1] if i < len(break_points) - 1 else num_outputs - for _ in range(end - start): - actual.append(sess.run(get_next_op)) - self._save(sess, saver) - if end == num_outputs: - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - self.assertEqual(expected, actual) - - def testSaveRestore(self): - self._testReadWithBreaks([8]) # rng buffer_size: 0 - self._testReadWithBreaks([13]) # rng buffer_size: 1 - self._testReadWithBreaks([18]) # rng buffer_size: 2 - self._testReadWithBreaks([23]) # rng buffer_size: 3 - - def testSaveUnusedIterator(self): - self._testReadWithBreaks([0]) - - def testSaveFullyUsedIterator(self): - self._testReadWithBreaks([50]) - - def testMultipleBreaks(self): - self._testReadWithBreaks([0, 5, 9, 15, 25, 32]) - - def testIdempotence(self): - # Attempt to save iterator immediately after restoring. - self._testReadWithBreaks([1, 1, 5, 5, 5, 25, 32]) - - def testInitThenRestore(self): - self._testReadWithBreaks([0, 5, 9, 15, 25, 32], init_before_restore=True) - - def testRestoreExhaustedIterator(self): - seed = 55 - range_limit = 10 - num_repeats = 5 - num_outputs = range_limit * num_repeats - buffer_sizes = [1, 3, 8, 10, 25, 50] - reshuffle_each_iteration = False - for buffer_size in buffer_sizes: - with ops.Graph().as_default() as g: - g.seed = 10 - init_op, get_next_op, saver = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration) - with self.test_session(graph=g) as sess: - sess.run(init_op) - for _ in range(num_outputs): - sess.run(get_next_op) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - self._save(sess, saver) - - with ops.Graph().as_default() as g: - saver = self._import_meta_graph() - init_op, get_next_op = ops.get_collection("iterator_ops") - with self.test_session(graph=g) as sess: - self._restore(saver, sess) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - - def testResetRestoredIterator(self): - seed = 55 - range_limit = 10 - num_repeats = 5 - num_outputs = range_limit * num_repeats - buffer_sizes = [1, 3, 8, 10, 25, 50] - reshuffle_each_iteration = False - for buffer_size in buffer_sizes: - with ops.Graph().as_default() as g: - g.seed = 10 - init_op, get_next_op, saver = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration) - with self.test_session(graph=g) as sess: - sess.run(init_op) - for _ in range(num_outputs // 2): - sess.run(get_next_op) - self._save(sess, saver) - - outputs = [] - with ops.Graph().as_default() as g: - saver = self._import_meta_graph() - init_op, get_next_op = ops.get_collection("iterator_ops") - with self.test_session(graph=g) as sess: - self._restore(saver, sess) - sess.run(init_op) - for _ in range(num_outputs): - outputs.append(sess.run(get_next_op)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - expected_outputs_sorted = sorted( - np.array([range(range_limit) - for _ in range(num_repeats)]).flatten()) - self.assertEqual(expected_outputs_sorted, sorted(outputs)) - - def testRestoreInModifiedGraph(self): - seed = 55 - break_point = 25 - range_limit = 10 - num_repeats = 5 - num_outputs = range_limit * num_repeats - buffer_sizes = [3, 8, 10, 25, 50] - reshuffle_each_iteration = False - for buffer_size in buffer_sizes: - expected = [] - actual_without_restore = [] - actual = [] - with ops.Graph().as_default() as g: - g.seed = 10 - init_op, get_next_op, saver = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration) - with self.test_session(graph=g) as sess: - sess.run(init_op) - for _ in range(break_point): - expected.append(sess.run(get_next_op)) - actual.extend(expected) - self._save(sess, saver) - for _ in range(num_outputs - break_point): - expected.append(sess.run(get_next_op)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - - with ops.Graph().as_default() as g: - g.seed = 20 # Different seed than previous graph for shuffle rngs. - init_op, get_next_op, saver = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration) - with self.test_session(graph=g) as sess: - sess.run(init_op) - for _ in range(num_outputs): - actual_without_restore.append(sess.run(get_next_op)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - - with ops.Graph().as_default() as g: - g.seed = 20 # Different seed than previous graph for shuffle rngs. - init_op, get_next_op, saver = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration) - with self.test_session(graph=g) as sess: - self._restore(saver, sess) - for _ in range(num_outputs - break_point): - actual.append(sess.run(get_next_op)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - - # Since the modified graph has a different random seed it produces a - # different order of examples. - self.assertNotEqual(expected, actual_without_restore) - self.assertEqual(sorted(expected), sorted(actual_without_restore)) - self.assertEqual(expected, actual) - - def testDoNotBuildSaveable(self): - seed = 55 - break_point = 25 - range_limit = 10 - num_repeats = 5 - num_outputs = range_limit * num_repeats - buffer_sizes = [3, 8, 10, 25, 50] - reshuffle_each_iteration = False - for buffer_size in buffer_sizes: - actual = [] - with ops.Graph().as_default() as g: - g.seed = 10 - init_op, get_next_op, saver = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration) - with self.test_session(graph=g) as sess: - sess.run(init_op) - for _ in range(break_point): - sess.run(get_next_op) - self._save(sess, saver) - - with ops.Graph().as_default() as g: - g.seed = 20 # Different seed than previous graph for shuffle rngs. - init_op, get_next_op, saver = self._build_graph( - range_limit=range_limit, - num_repeats=num_repeats, - buffer_size=buffer_size, - seed=seed, - reshuffle_each_iteration=reshuffle_each_iteration, - build_saveable=False) - with self.test_session(graph=g) as sess: - # Since the SaveableObject was not added to Saver's list - # of saveables, iterator state is not restored by saver.restore(). - self._restore(saver, sess) - with self.assertRaises(errors.FailedPreconditionError): - sess.run(get_next_op) - sess.run(init_op) - for _ in range(num_outputs): - actual.append(sess.run(get_next_op)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next_op) - expected_outputs_sorted = sorted( - np.array([range(range_limit) for _ in range(num_repeats)]).flatten()) - self.assertEqual(expected_outputs_sorted, sorted(actual)) + self.run_core_tests( + lambda: self._build_shuffle_dataset( + range_limit=range_limit, + num_repeats=num_repeats, + buffer_size=buffer_size, + seed=seed, + reshuffle_each_iteration=reshuffle_each_iteration), + lambda: self._build_shuffle_dataset( + range_limit=range_limit, + num_repeats=num_repeats, + buffer_size=buffer_size, + seed=10, + reshuffle_each_iteration=reshuffle_each_iteration), + num_outputs) + # pylint: enable=cell-var-from-loop + # pylint: enable=g-long-lambda class ShuffleAndRepeatTest( @@ -541,6 +137,14 @@ class ShuffleAndRepeatTest( self.gen_outputs(lambda: self._build_ds(10, count=-1, num_elements=0), [], 100) + def testLargeBufferSize(self): + with ops.Graph().as_default() as g: + ds = dataset_ops.Dataset.range(20).apply( + shuffle_ops.shuffle_and_repeat(buffer_size=21)) + get_next_op = ds.make_one_shot_iterator().get_next() + with self.test_session(graph=g) as sess: + sess.run(get_next_op) + class ShuffleAndRepeatSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): diff --git a/tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test.py index efd864f866611bfd3bac1edcf98d84be852410fd..e26cef8ec522c7e69a0c19b2b30a969bbfc0ad78 100644 --- a/tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function import os + import sqlite3 from tensorflow.contrib.data.python.ops import readers diff --git a/tensorflow/contrib/data/python/kernel_tests/unique_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/unique_dataset_op_test.py index 55296d5710e7f66408bb7464cf790149d6df9fa1..3c436f7a0b45a13109960e87dd97ca56b10bb871 100644 --- a/tensorflow/contrib/data/python/kernel_tests/unique_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/unique_dataset_op_test.py @@ -18,8 +18,8 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base -from tensorflow.contrib.data.python.ops import dataset_ops from tensorflow.contrib.data.python.ops import unique +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.platform import test diff --git a/tensorflow/contrib/data/python/kernel_tests/zip_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/zip_dataset_op_test.py index 5d34b0024c472d0393544ff3dad8acea7964345f..e39fa957f0bbb9d3671274d5f58b993e8399814b 100644 --- a/tensorflow/contrib/data/python/kernel_tests/zip_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/zip_dataset_op_test.py @@ -20,97 +20,10 @@ from __future__ import print_function import numpy as np from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base -from tensorflow.contrib.data.python.ops import dataset_ops -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.ops import array_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.platform import test -class ZipDatasetTest(test.TestCase): - - def testZipDataset(self): - component_placeholders = [ - array_ops.placeholder(dtypes.int64), - array_ops.placeholder(dtypes.int64), - array_ops.placeholder(dtypes.float64) - ] - - datasets = tuple([ - dataset_ops.Dataset.from_tensor_slices(component_placeholder) - for component_placeholder in component_placeholders - ]) - zipped = dataset_ops.Dataset.zip(datasets) - - iterator = zipped.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - with self.test_session() as sess: - equal_length_components = [ - np.tile(np.array([[1], [2], [3], [4]]), 20), - np.tile(np.array([[12], [13], [14], [15]]), 22), - np.array([37.0, 38.0, 39.0, 40.0]) - ] - sess.run(init_op, feed_dict={ph: value for ph, value in zip( - component_placeholders, equal_length_components)}) - for i in range(4): - results = sess.run(get_next) - for component, result_component in zip( - equal_length_components, results): - self.assertAllEqual(component[i], result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - variable_length_components = [[1, 2, 3, 4], [1, 2, 3, 4, 5], [1.0, 2.0]] - sess.run(init_op, feed_dict={ph: value for ph, value in zip( - component_placeholders, variable_length_components)}) - for i in range(2): - results = sess.run(get_next) - for component, result_component in zip( - variable_length_components, results): - self.assertAllEqual(component[i], result_component) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - def testNestedZipDataset(self): - component_placeholders = [ - array_ops.placeholder(dtypes.int64, shape=[4, 20]), - array_ops.placeholder(dtypes.int64, shape=[4, 22]), - array_ops.placeholder(dtypes.float64, shape=[4]) - ] - - datasets = [ - dataset_ops.Dataset.from_tensor_slices(component_placeholder) - for component_placeholder in component_placeholders - ] - zipped = dataset_ops.Dataset.zip((datasets[0], (datasets[1], datasets[2]))) - - iterator = zipped.make_initializable_iterator() - init_op = iterator.initializer - get_next = iterator.get_next() - - self.assertEqual([20], get_next[0].shape) - self.assertEqual([22], get_next[1][0].shape) - self.assertEqual([], get_next[1][1].shape) - - with self.test_session() as sess: - equal_length_components = [ - np.tile(np.array([[1], [2], [3], [4]]), 20), - np.tile(np.array([[12], [13], [14], [15]]), 22), - np.array([37.0, 38.0, 39.0, 40.0]) - ] - sess.run(init_op, feed_dict={ph: value for ph, value in zip( - component_placeholders, equal_length_components)}) - for i in range(4): - result1, (result2, result3) = sess.run(get_next) - self.assertAllEqual(equal_length_components[0][i], result1) - self.assertAllEqual(equal_length_components[1][i], result2) - self.assertAllEqual(equal_length_components[2][i], result3) - with self.assertRaises(errors.OutOfRangeError): - sess.run(get_next) - - class ZipDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD index 4349085a10135b4dee842a29916aeb5febe9ddd4..b488357f226d0922bba3799cc1f4b5c75e2e8328 100644 --- a/tensorflow/contrib/data/python/ops/BUILD +++ b/tensorflow/contrib/data/python/ops/BUILD @@ -15,7 +15,7 @@ py_library( name = "dataset_ops", srcs = [ "counter.py", - "dataset_ops.py", + "get_single_element.py", ], srcs_version = "PY2AND3", deps = [ @@ -109,6 +109,8 @@ py_library( ], srcs_version = "PY2AND3", deps = [ + ":contrib_op_loader", + ":gen_dataset_ops", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", "//tensorflow/python:dataset_ops_gen", @@ -130,36 +132,44 @@ py_library( ) tf_gen_op_wrapper_py( - name = "prefetching_ops", - out = "gen_prefetching_ops.py", - deps = ["//tensorflow/contrib/data:prefetching_ops_op_lib"], + name = "gen_dataset_ops", + out = "gen_dataset_ops.py", + deps = ["//tensorflow/contrib/data:dataset_ops_op_lib"], ) tf_kernel_library( - name = "prefetching_ops_kernels", + name = "dataset_ops_kernels", deps = [ - "//tensorflow/contrib/data/kernels:prefetching_kernels", + "//tensorflow/contrib/data/kernels:dataset_kernels", "//tensorflow/core:framework", ], alwayslink = 1, ) tf_custom_op_py_library( - name = "prefetching_py", - srcs = ["prefetching_ops.py"], - dso = ["//tensorflow/contrib/data:_prefetching_ops.so"], + name = "contrib_op_loader", + srcs = ["contrib_op_loader.py"], + dso = ["//tensorflow/contrib/data:_dataset_ops.so"], kernels = [ - ":prefetching_ops_kernels", - "//tensorflow/contrib/data:prefetching_ops_op_lib", + ":dataset_ops_kernels", + "//tensorflow/contrib/data:dataset_ops_op_lib", ], srcs_version = "PY2AND3", deps = [ - ":prefetching_ops", + ":gen_dataset_ops", "//tensorflow/contrib/util:util_py", "//tensorflow/python:platform", ], ) +py_library( + name = "prefetching_ops", + srcs = ["prefetching_ops.py"], + deps = [ + ":contrib_op_loader", + ], +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/contrib/data/python/ops/batching.py b/tensorflow/contrib/data/python/ops/batching.py index e0860cbe8a712a2743c2bbe206afb7b1746b6fc9..6eb512dec67cb7b9c8c4518d03aee0b436205f9a 100644 --- a/tensorflow/contrib/data/python/ops/batching.py +++ b/tensorflow/contrib/data/python/ops/batching.py @@ -403,11 +403,11 @@ def map_and_batch(map_func, batch_size, num_parallel_batches=1): num_parallel_batches: A `tf.int64` scalar `tf.Tensor`, representing the number of batches to create in parallel. On one hand, higher values can help mitigate the effect of stragglers. On the other hand, higher values - can increasing contention if CPU is scarce. + can increase contention if CPU is scarce. Returns: A `Dataset` transformation function, which can be passed to - @{tf.contrib.data.Dataset.apply}. + @{tf.data.Dataset.apply}. """ def _apply_fn(dataset): diff --git a/tensorflow/contrib/data/python/ops/contrib_op_loader.py b/tensorflow/contrib/data/python/ops/contrib_op_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..8f495a9dc9c82311435e71d2ac9ed35fd9aea794 --- /dev/null +++ b/tensorflow/contrib/data/python/ops/contrib_op_loader.py @@ -0,0 +1,24 @@ +# 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. +# ============================================================================== +"""Python helper for loading contrib ops and kernels.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.util import loader +from tensorflow.python.platform import resource_loader + +_dataset_ops = loader.load_op_library( + resource_loader.get_path_to_datafile("../../_dataset_ops.so")) diff --git a/tensorflow/contrib/data/python/ops/dataset_ops.py b/tensorflow/contrib/data/python/ops/dataset_ops.py index 626a9e0edcea5928b1636c1a2a86e83657c966a5..ff15c4451ad987bcd77dbdd022a1c070056c47e1 100644 --- a/tensorflow/contrib/data/python/ops/dataset_ops.py +++ b/tensorflow/contrib/data/python/ops/dataset_ops.py @@ -364,7 +364,7 @@ class Dataset(dataset_ops.Dataset): When reading a single input file, you can skip elements as follows: ```python - d = tf.contrib.data.TFRecordDataset(FLAGS.input_file) + d = tf.data.TFRecordDataset(FLAGS.input_file) d = d.shard(FLAGS.num_workers, FLAGS.worker_index) d = d.repeat(FLAGS.num_epochs) d = d.shuffle(FLAGS.shuffle_buffer_size) @@ -382,12 +382,11 @@ class Dataset(dataset_ops.Dataset): sharding strategy within a complete pipeline: ```python - d = Dataset.list_files(FLAGS.pattern) + d = tf.data.Dataset.list_files(FLAGS.pattern) d = d.shard(FLAGS.num_workers, FLAGS.worker_index) d = d.repeat(FLAGS.num_epochs) d = d.shuffle(FLAGS.shuffle_buffer_size) - d = d.repeat() - d = d.interleave(tf.contrib.data.TFRecordDataset, + d = d.interleave(tf.data.TFRecordDataset, cycle_length=FLAGS.num_readers, block_length=1) d = d.map(parser_fn, num_parallel_calls=FLAGS.num_map_threads) ``` @@ -484,7 +483,7 @@ class Dataset(dataset_ops.Dataset): num_threads=None, output_buffer_size=None, num_parallel_calls=None): - """Maps `map_func` across this datset. + """Maps `map_func` across this dataset. Args: map_func: A function mapping a nested structure of tensors (having @@ -549,7 +548,7 @@ class Dataset(dataset_ops.Dataset): elements are produced. `cycle_length` controls the number of input elements that are processed concurrently. If you set `cycle_length` to 1, this transformation will handle one input element at a time, and will produce - identical results = to @{tf.contrib.data.Dataset.flat_map}. In general, + identical results = to @{tf.data.Dataset.flat_map}. In general, this transformation will apply `map_func` to `cycle_length` input elements, open iterators on the returned `Dataset` objects, and cycle through them producing `block_length` consecutive elements from each iterator, and diff --git a/tensorflow/contrib/data/python/ops/error_ops.py b/tensorflow/contrib/data/python/ops/error_ops.py index aa629cba479102ee4244884e7c546615b28cf4e5..6c21e489f7c35484ebacd465e3b46d6920df5933 100644 --- a/tensorflow/contrib/data/python/ops/error_ops.py +++ b/tensorflow/contrib/data/python/ops/error_ops.py @@ -17,10 +17,11 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import +from tensorflow.contrib.data.python.ops import gen_dataset_ops from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.util import nest from tensorflow.python.data.util import sparse -from tensorflow.python.ops import gen_dataset_ops def ignore_errors(): diff --git a/tensorflow/contrib/data/python/ops/get_single_element.py b/tensorflow/contrib/data/python/ops/get_single_element.py new file mode 100644 index 0000000000000000000000000000000000000000..a817b45b71b608810a9d7536ec123ab84f7cdc3b --- /dev/null +++ b/tensorflow/contrib/data/python/ops/get_single_element.py @@ -0,0 +1,67 @@ +# Copyright 2017 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. +# ============================================================================== +"""Python wrappers for Datasets and Iterators.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.data.util import nest +from tensorflow.python.ops import gen_dataset_ops + + +def get_single_element(dataset): + """Returns the single element in `dataset` as a nested structure of tensors. + + This function enables you to use a @{tf.data.Dataset} in a stateless + "tensor-in tensor-out" expression, without creating a @{tf.data.Iterator}. + This can be useful when your preprocessing transformations are expressed + as a `Dataset`, and you want to use the transformation at serving time. + For example: + + ```python + input_batch = tf.placeholder(tf.string, shape=[BATCH_SIZE]) + + def preprocessing_fn(input_str): + # ... + return image, label + + dataset = (tf.data.Dataset.from_tensor_slices(input_batch) + .map(preprocessing_fn, num_parallel_calls=BATCH_SIZE) + .batch(BATCH_SIZE)) + + image_batch, label_batch = tf.contrib.data.get_single_element(dataset) + ``` + + Args: + dataset: A @{tf.data.Dataset} object containing a single element. + + Returns: + A nested structure of @{tf.Tensor} objects, corresponding to the single + element of `dataset`. + + Raises: + TypeError: if `dataset` is not a `tf.data.Dataset` object. + InvalidArgumentError (at runtime): if `dataset` does not contain exactly + one element. + """ + if not isinstance(dataset, dataset_ops.Dataset): + raise TypeError("`dataset` must be a `tf.data.Dataset` object.") + return nest.pack_sequence_as( + dataset.output_types, + gen_dataset_ops.dataset_to_single_element( + dataset._as_variant_tensor(), # pylint: disable=protected-access + output_types=nest.flatten(dataset.output_types), + output_shapes=nest.flatten(dataset.output_shapes))) diff --git a/tensorflow/contrib/data/python/ops/grouping.py b/tensorflow/contrib/data/python/ops/grouping.py index ef91c56726e969053fdad667dda3e89430045652..67b085002aa7797d858837fea4646fb968ad5d97 100644 --- a/tensorflow/contrib/data/python/ops/grouping.py +++ b/tensorflow/contrib/data/python/ops/grouping.py @@ -45,7 +45,7 @@ def group_by_window(key_func, key_func: A function mapping a nested structure of tensors (having shapes and types defined by `self.output_shapes` and `self.output_types`) to a scalar `tf.int64` tensor. - reduce_func: A function mapping a key and a dataset of up to `batch_size` + reduce_func: A function mapping a key and a dataset of up to `window_size` consecutive elements matching that key to another dataset. window_size: A `tf.int64` scalar `tf.Tensor`, representing the number of consecutive elements matching the same key to combine in a single diff --git a/tensorflow/contrib/data/python/ops/prefetching_ops.py b/tensorflow/contrib/data/python/ops/prefetching_ops.py index cfe8012b5657995b78d701528ea35cbb3748adb9..96a9e9ed6649444dac5e56d7dd2fcdb62fc56459 100644 --- a/tensorflow/contrib/data/python/ops/prefetching_ops.py +++ b/tensorflow/contrib/data/python/ops/prefetching_ops.py @@ -17,12 +17,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.data.python.ops import gen_prefetching_ops -from tensorflow.contrib.util import loader -from tensorflow.python.platform import resource_loader - -_prefetching_ops = loader.load_op_library( - resource_loader.get_path_to_datafile("../../_prefetching_ops.so")) +from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import +from tensorflow.contrib.data.python.ops import gen_dataset_ops # TODO(rohanj): Add a python class that constructs resource in the __init__ @@ -35,7 +31,7 @@ def function_buffering_resource(string_arg, thread_pool_size=1, container="", name=None): - return gen_prefetching_ops.function_buffering_resource( + return gen_dataset_ops.function_buffering_resource( string_arg=string_arg, target_device=target_device, shared_name=shared_name, @@ -49,7 +45,7 @@ def function_buffering_resource(string_arg, def function_buffering_resource_get_next(function_buffer_resource, output_types, name=None): - return gen_prefetching_ops.function_buffering_resource_get_next( + return gen_dataset_ops.function_buffering_resource_get_next( function_buffer_resource=function_buffer_resource, output_types=output_types, name=name) diff --git a/tensorflow/contrib/data/python/ops/readers.py b/tensorflow/contrib/data/python/ops/readers.py index 347e5edc7b0d479dfa260e8cec500ffaaba375be..57f30102778f3bac47580f9bdf94e411dfe1b621 100644 --- a/tensorflow/contrib/data/python/ops/readers.py +++ b/tensorflow/contrib/data/python/ops/readers.py @@ -17,9 +17,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.data.python.ops import dataset_ops as contrib_dataset_ops from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.ops import readers from tensorflow.python.data.util import nest from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -27,74 +25,6 @@ from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import gen_dataset_ops from tensorflow.python.ops import parsing_ops from tensorflow.python.platform import gfile -from tensorflow.python.util import deprecation - - -class TextLineDataset(contrib_dataset_ops.Dataset): - """A `Dataset` comprising lines from one or more text files.""" - - @deprecation.deprecated(None, "Use `tf.data.TextLineDataset`.") - def __init__(self, filenames, compression_type=None, buffer_size=None): - """Creates a `TextLineDataset`. - - Args: - filenames: A `tf.string` tensor containing one or more filenames. - compression_type: (Optional.) A `tf.string` scalar evaluating to one of - `""` (no compression), `"ZLIB"`, or `"GZIP"`. - buffer_size: (Optional.) A `tf.int64` scalar denoting the number of bytes - to buffer. A value of 0 results in the default buffering values chosen - based on the compression type. - """ - dataset = readers.TextLineDataset(filenames, compression_type, - buffer_size) - super(TextLineDataset, self).__init__(dataset) - - -class TFRecordDataset(contrib_dataset_ops.Dataset): - """A `Dataset` comprising records from one or more TFRecord files.""" - - @deprecation.deprecated(None, "Use `tf.data.TFRecordDataset`.") - def __init__(self, filenames, compression_type=None, buffer_size=None): - """Creates a `TFRecordDataset`. - - Args: - filenames: A `tf.string` tensor containing one or more filenames. - compression_type: (Optional.) A `tf.string` scalar evaluating to one of - `""` (no compression), `"ZLIB"`, or `"GZIP"`. - buffer_size: (Optional.) A `tf.int64` scalar representing the number of - bytes in the read buffer. 0 means no buffering. - """ - dataset = readers.TFRecordDataset(filenames, compression_type, - buffer_size) - super(TFRecordDataset, self).__init__(dataset) - - -class FixedLengthRecordDataset(contrib_dataset_ops.Dataset): - """A `Dataset` of fixed-length records from one or more binary files.""" - - @deprecation.deprecated(None, "Use `tf.data.FixedLengthRecordDataset`.") - def __init__(self, - filenames, - record_bytes, - header_bytes=None, - footer_bytes=None, - buffer_size=None): - """Creates a `FixedLengthRecordDataset`. - - Args: - filenames: A `tf.string` tensor containing one or more filenames. - record_bytes: A `tf.int64` scalar representing the number of bytes in - each record. - header_bytes: (Optional.) A `tf.int64` scalar representing the number of - bytes to skip at the start of a file. - footer_bytes: (Optional.) A `tf.int64` scalar representing the number of - bytes to ignore at the end of a file. - buffer_size: (Optional.) A `tf.int64` scalar representing the number of - bytes to buffer when reading. - """ - dataset = readers.FixedLengthRecordDataset( - filenames, record_bytes, header_bytes, footer_bytes, buffer_size) - super(FixedLengthRecordDataset, self).__init__(dataset) def read_batch_features(file_pattern, @@ -216,14 +146,7 @@ def _get_file_names(file_pattern, randomize_input): return file_names -class SqlDataset(contrib_dataset_ops.Dataset): - - def __init__(self, driver_name, data_source_name, query, output_types): - dataset = _SqlDataset(driver_name, data_source_name, query, output_types) - super(SqlDataset, self).__init__(dataset) - - -class _SqlDataset(dataset_ops.Dataset): +class SqlDataset(dataset_ops.Dataset): """A `Dataset` consisting of the results from a SQL query.""" def __init__(self, driver_name, data_source_name, query, output_types): @@ -255,7 +178,7 @@ class _SqlDataset(dataset_ops.Dataset): output_types: A tuple of `tf.DType` objects representing the types of the columns returned by `query`. """ - super(_SqlDataset, self).__init__() + super(SqlDataset, self).__init__() self._driver_name = ops.convert_to_tensor( driver_name, dtype=dtypes.string, name="driver_name") self._data_source_name = ops.convert_to_tensor( diff --git a/tensorflow/contrib/data/python/ops/scan_ops.py b/tensorflow/contrib/data/python/ops/scan_ops.py index 2744786e9eec4c9268ba854df6ea761339bb0b4e..1c88366273f5d186509454188e02350d4ea9f66b 100644 --- a/tensorflow/contrib/data/python/ops/scan_ops.py +++ b/tensorflow/contrib/data/python/ops/scan_ops.py @@ -188,7 +188,7 @@ def scan(initial_state, scan_func): Returns: A `Dataset` transformation function, which can be passed to - @{tf.contrib.data.Dataset.apply}. + @{tf.data.Dataset.apply}. """ def _apply_fn(dataset): return _ScanDataset(dataset, initial_state, scan_func) diff --git a/tensorflow/contrib/data/python/ops/shuffle_ops.py b/tensorflow/contrib/data/python/ops/shuffle_ops.py index 410989fad4f2a3bb8c9051c094ce8ab7b2eee96c..99bb79bc06a421f811869ca9169aaa11deaca2f3 100644 --- a/tensorflow/contrib/data/python/ops/shuffle_ops.py +++ b/tensorflow/contrib/data/python/ops/shuffle_ops.py @@ -111,7 +111,7 @@ def shuffle_and_repeat(buffer_size, count=None, seed=None): Returns: A `Dataset` transformation function, which can be passed to - @{tf.contrib.data.Dataset.apply}. + @{tf.data.Dataset.apply}. """ def _apply_fn(dataset): # pylint: disable=missing-docstring diff --git a/tensorflow/contrib/data/python/ops/stats_ops.py b/tensorflow/contrib/data/python/ops/stats_ops.py index b8875bd533ddc9e2c195646619dccf3aab5225e4..9cd1701c397b5a0bf5cc47c1bcab033704794d80 100644 --- a/tensorflow/contrib/data/python/ops/stats_ops.py +++ b/tensorflow/contrib/data/python/ops/stats_ops.py @@ -20,6 +20,7 @@ from __future__ import print_function from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops from tensorflow.python.data.util import nest +from tensorflow.python.data.util import sparse from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import gen_dataset_ops @@ -117,7 +118,7 @@ def bytes_produced_stats(tag): Returns: A `Dataset` transformation function, which can be passed to - @{tf.contrib.data.Dataset.apply}. + @{tf.data.Dataset.apply}. """ def _apply_fn(dataset): @@ -139,7 +140,7 @@ def latency_stats(tag): Returns: A `Dataset` transformation function, which can be passed to - @{tf.contrib.data.Dataset.apply}. + @{tf.data.Dataset.apply}. """ def _apply_fn(dataset): @@ -161,8 +162,10 @@ class _StatsDataset(dataset_ops.Dataset): return self._op_function( self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access self._tag, - output_shapes=nest.flatten(self.output_shapes), - output_types=nest.flatten(self.output_types)) + output_types=nest.flatten( + sparse.as_dense_types(self.output_types, self.output_classes)), + output_shapes=nest.flatten( + sparse.as_dense_shapes(self.output_shapes, self.output_classes))) @property def output_shapes(self): diff --git a/tensorflow/contrib/decision_trees/proto/BUILD b/tensorflow/contrib/decision_trees/proto/BUILD index 87c80740a8f0c0721394b5d832bc96e548e3a313..f6de5998d73a4869d2444cd90c9b64d1a2c889ac 100644 --- a/tensorflow/contrib/decision_trees/proto/BUILD +++ b/tensorflow/contrib/decision_trees/proto/BUILD @@ -7,7 +7,11 @@ exports_files([ "generic_tree_model_proto.swig", ]) -load("//tensorflow/core:platform/default/build_config.bzl", "tf_proto_library") +load( + "//tensorflow/core:platform/default/build_config.bzl", + "tf_proto_library", + "tf_pyclif_proto_library", +) filegroup( name = "all_files", @@ -34,3 +38,10 @@ tf_proto_library( protodeps = [":generic_tree_model"], visibility = ["//visibility:public"], ) + +tf_pyclif_proto_library( + name = "generic_tree_model_pyclif", + proto_lib = ":generic_tree_model", + proto_srcfile = "generic_tree_model.proto", + visibility = ["//visibility:public"], +) diff --git a/tensorflow/contrib/distributions/BUILD b/tensorflow/contrib/distributions/BUILD index 95848af69950bdaa680c41daecd8cbd8f3174f8e..7f510c42215f48a9e795eb81bd9f66b0a2108335 100644 --- a/tensorflow/contrib/distributions/BUILD +++ b/tensorflow/contrib/distributions/BUILD @@ -128,6 +128,19 @@ cuda_py_test( tags = ["no_pip"], ) +cuda_py_test( + name = "autoregressive_test", + size = "small", + srcs = ["python/kernel_tests/autoregressive_test.py"], + additional_deps = [ + ":distributions_py", + "//third_party/py/numpy", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + cuda_py_test( name = "binomial_test", size = "small", @@ -918,6 +931,22 @@ cuda_py_test( ], ) +cuda_py_test( + name = "real_nvp_test", + size = "small", + srcs = ["python/kernel_tests/bijectors/real_nvp_test.py"], + additional_deps = [ + ":bijectors_py", + ":distributions_py", + "//third_party/py/numpy", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:platform_test", + ], +) + cuda_py_test( name = "permute_test", size = "small", diff --git a/tensorflow/contrib/distributions/__init__.py b/tensorflow/contrib/distributions/__init__.py index 7b401e178f35fe56e4eb461936565f5c630ec4cf..61c411271d0bb8d7b4cc3b14992b82ec1e5674ed 100644 --- a/tensorflow/contrib/distributions/__init__.py +++ b/tensorflow/contrib/distributions/__init__.py @@ -23,6 +23,7 @@ from __future__ import print_function # pylint: disable=unused-import,wildcard-import,line-too-long,g-importing-member from tensorflow.contrib.distributions.python.ops import bijectors +from tensorflow.contrib.distributions.python.ops.autoregressive import * from tensorflow.contrib.distributions.python.ops.binomial import * from tensorflow.contrib.distributions.python.ops.cauchy import * from tensorflow.contrib.distributions.python.ops.chi2 import * @@ -39,6 +40,7 @@ from tensorflow.contrib.distributions.python.ops.geometric import * from tensorflow.contrib.distributions.python.ops.half_normal import * from tensorflow.contrib.distributions.python.ops.independent import * from tensorflow.contrib.distributions.python.ops.inverse_gamma import * +from tensorflow.contrib.distributions.python.ops.kumaraswamy import * from tensorflow.contrib.distributions.python.ops.logistic import * from tensorflow.contrib.distributions.python.ops.mixture import * from tensorflow.contrib.distributions.python.ops.mixture_same_family import * @@ -84,6 +86,7 @@ from tensorflow.python.ops.distributions.uniform import * from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ + 'auto_correlation', 'bijectors', 'Cauchy', 'ConditionalDistribution', @@ -92,9 +95,9 @@ _allowed_symbols = [ 'NOT_REPARAMETERIZED', 'ReparameterizationType', 'Distribution', + 'Autoregressive', 'Binomial', 'Bernoulli', - 'BernoulliWithSigmoidProbs', 'Beta', 'BetaWithSoftplusConcentration', 'Categorical', @@ -112,6 +115,7 @@ _allowed_symbols = [ 'Independent', 'InverseGamma', 'InverseGammaWithSoftplusConcentrationRate', + 'Kumaraswamy', 'Laplace', 'LaplaceWithSoftplusScale', 'Logistic', diff --git a/tensorflow/contrib/distributions/python/kernel_tests/autoregressive_test.py b/tensorflow/contrib/distributions/python/kernel_tests/autoregressive_test.py new file mode 100644 index 0000000000000000000000000000000000000000..0928dc3f358ede693865a8d1ff9257a0ecbe9499 --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/autoregressive_test.py @@ -0,0 +1,94 @@ +# Copyright 2017 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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.distributions.python.ops import autoregressive as autoregressive_lib +from tensorflow.contrib.distributions.python.ops import independent as independent_lib +from tensorflow.contrib.distributions.python.ops import test_util +from tensorflow.contrib.distributions.python.ops.bijectors.affine import Affine +from tensorflow.contrib.distributions.python.ops.bijectors.masked_autoregressive import MaskedAutoregressiveFlow +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops.distributions import normal as normal_lib +from tensorflow.python.ops.distributions import transformed_distribution as transformed_distribution_lib +from tensorflow.python.ops.distributions import util as distribution_util +from tensorflow.python.platform import test + + +class AutogressiveTest(test_util.VectorDistributionTestHelpers, test.TestCase): + """Tests the Autoregressive distribution.""" + + def setUp(self): + self._rng = np.random.RandomState(42) + + def _random_scale_tril(self, event_size): + n = np.int32(event_size * (event_size + 1) // 2) + p = 2. * self._rng.random_sample(n).astype(np.float32) - 1. + return distribution_util.fill_triangular(0.25 * p) + + def _normal_fn(self, affine_bijector): + def _fn(samples): + scale = math_ops.exp(affine_bijector.forward(samples)) + return independent_lib.Independent( + normal_lib.Normal(loc=0., scale=scale, validate_args=True), + reinterpreted_batch_ndims=1) + return _fn + + def testSampleAndLogProbConsistency(self): + batch_shape = [] + event_size = 2 + with self.test_session() as sess: + batch_event_shape = np.concatenate([batch_shape, [event_size]], axis=0) + sample0 = array_ops.zeros(batch_event_shape) + affine = Affine(scale_tril=self._random_scale_tril(event_size)) + ar = autoregressive_lib.Autoregressive( + self._normal_fn(affine), sample0, validate_args=True) + self.run_test_sample_consistent_log_prob( + sess.run, ar, radius=1., center=0., rtol=0.01) + + def testCompareToBijector(self): + """Demonstrates equivalence between TD, Bijector approach and AR dist.""" + sample_shape = np.int32([4, 5]) + batch_shape = np.int32([]) + event_size = np.int32(2) + with self.test_session() as sess: + batch_event_shape = np.concatenate([batch_shape, [event_size]], axis=0) + sample0 = array_ops.zeros(batch_event_shape) + affine = Affine(scale_tril=self._random_scale_tril(event_size)) + ar = autoregressive_lib.Autoregressive( + self._normal_fn(affine), sample0, validate_args=True) + ar_flow = MaskedAutoregressiveFlow( + is_constant_jacobian=True, + shift_and_log_scale_fn=lambda x: [None, affine.forward(x)], + validate_args=True) + td = transformed_distribution_lib.TransformedDistribution( + distribution=normal_lib.Normal(loc=0., scale=1.), + bijector=ar_flow, + event_shape=[event_size], + batch_shape=batch_shape, + validate_args=True) + x_shape = np.concatenate( + [sample_shape, batch_shape, [event_size]], axis=0) + x = 2. * self._rng.random_sample(x_shape).astype(np.float32) - 1. + td_log_prob_, ar_log_prob_ = sess.run([td.log_prob(x), ar.log_prob(x)]) + self.assertAllClose(td_log_prob_, ar_log_prob_, atol=0., rtol=1e-6) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/masked_autoregressive_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/masked_autoregressive_test.py index 288d9d8dd6f17cd6348d3d72aea4408e26913ebd..dcfb0eb05185d36d96947905c2eb91b2201aece1 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/masked_autoregressive_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/masked_autoregressive_test.py @@ -149,5 +149,17 @@ class MaskedAutoregressiveFlowShiftOnlyTest(MaskedAutoregressiveFlowTest): } +class MaskedAutoregressiveFlowUnrollLoopTest(MaskedAutoregressiveFlowTest): + + @property + def _autoregressive_flow_kwargs(self): + return { + "shift_and_log_scale_fn": masked_autoregressive_default_template( + hidden_layers=[2], shift_only=False), + "is_constant_jacobian": False, + "unroll_loop": True, + } + + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/real_nvp_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/real_nvp_test.py new file mode 100644 index 0000000000000000000000000000000000000000..46fe7797419a9906ecdad60dd0dfe1e9d7c743ed --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/real_nvp_test.py @@ -0,0 +1,144 @@ +# Copyright 2017 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 MaskedAutoregressiveFlow.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +from tensorflow.contrib.distributions.python.ops import test_util +from tensorflow.contrib.distributions.python.ops.bijectors.invert import Invert +from tensorflow.contrib.distributions.python.ops.bijectors.real_nvp import real_nvp_default_template +from tensorflow.contrib.distributions.python.ops.bijectors.real_nvp import RealNVP +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import variables +from tensorflow.python.ops.distributions import normal as normal_lib +from tensorflow.python.ops.distributions import transformed_distribution as transformed_distribution_lib +from tensorflow.python.platform import test + + +class RealNVPTest(test_util.VectorDistributionTestHelpers, test.TestCase): + + @property + def _real_nvp_kwargs(self): + return { + "shift_and_log_scale_fn": real_nvp_default_template( + hidden_layers=[3], shift_only=False), + "is_constant_jacobian": False, + } + + def testBijector(self): + x_ = np.arange(3 * 4 * 2).astype(np.float32).reshape(3, 4 * 2) + with self.test_session() as sess: + nvp = RealNVP( + num_masked=4, + validate_args=True, + **self._real_nvp_kwargs) + x = constant_op.constant(x_) + forward_x = nvp.forward(x) + # Use identity to invalidate cache. + inverse_y = nvp.inverse(array_ops.identity(forward_x)) + fldj = nvp.forward_log_det_jacobian(x) + # Use identity to invalidate cache. + ildj = nvp.inverse_log_det_jacobian(array_ops.identity(forward_x)) + variables.global_variables_initializer().run() + [ + forward_x_, + inverse_y_, + ildj_, + fldj_, + ] = sess.run([ + forward_x, + inverse_y, + ildj, + fldj, + ]) + self.assertEqual("real_nvp", nvp.name) + self.assertAllClose(forward_x_, forward_x_, rtol=1e-6, atol=0.) + self.assertAllClose(x_, inverse_y_, rtol=1e-5, atol=0.) + self.assertAllClose(ildj_, -fldj_, rtol=1e-6, atol=0.) + + def testMutuallyConsistent(self): + dims = 4 + with self.test_session() as sess: + nvp = RealNVP( + num_masked=3, + validate_args=True, + **self._real_nvp_kwargs) + dist = transformed_distribution_lib.TransformedDistribution( + distribution=normal_lib.Normal(loc=0., scale=1.), + bijector=nvp, + event_shape=[dims], + validate_args=True) + self.run_test_sample_consistent_log_prob( + sess_run_fn=sess.run, + dist=dist, + num_samples=int(1e5), + radius=1., + center=0., + rtol=0.02) + + def testInvertMutuallyConsistent(self): + dims = 4 + with self.test_session() as sess: + nvp = Invert(RealNVP( + num_masked=3, + validate_args=True, + **self._real_nvp_kwargs)) + dist = transformed_distribution_lib.TransformedDistribution( + distribution=normal_lib.Normal(loc=0., scale=1.), + bijector=nvp, + event_shape=[dims], + validate_args=True) + self.run_test_sample_consistent_log_prob( + sess_run_fn=sess.run, + dist=dist, + num_samples=int(1e5), + radius=1., + center=0., + rtol=0.02) + + +class NICETest(RealNVPTest): + + @property + def _real_nvp_kwargs(self): + return { + "shift_and_log_scale_fn": real_nvp_default_template( + hidden_layers=[2], shift_only=True), + "is_constant_jacobian": True, + } + + +class RealNVPConstantShiftScaleTest(RealNVPTest): + + @property + def _real_nvp_kwargs(self): + + def constant_shift_log_scale_fn(x0, output_units): + del x0, output_units + shift = constant_op.constant([0.1]) + log_scale = constant_op.constant([0.5]) + return shift, log_scale + + return { + "shift_and_log_scale_fn": constant_shift_log_scale_fn, + "is_constant_jacobian": True, + } + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/reshape_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/reshape_test.py index 49451446b56d290f130c5db90c13b94974d92dc9..e216d88cb190dc16fc0056186f80817d6f2d7c67 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/reshape_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/reshape_test.py @@ -22,12 +22,15 @@ import numpy as np from tensorflow.contrib.distributions.python.ops.bijectors.reshape import Reshape from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops.distributions.bijector_test_util import assert_bijective_and_finite from tensorflow.python.platform import test +@test_util.with_c_api class _ReshapeBijectorTest(object): """Base class for testing the reshape transformation. @@ -136,7 +139,8 @@ class _ReshapeBijectorTest(object): sess.run(bijector.forward_event_shape_tensor(shape_in), feed_dict=feed_dict) - def testInvalidDimensionsOpError(self): + # pylint: disable=invalid-name + def _testInvalidDimensionsOpError(self, expected_error_message): with self.test_session() as sess: @@ -146,10 +150,10 @@ class _ReshapeBijectorTest(object): event_shape_in=shape_in, validate_args=True) - with self.assertRaisesError( - "elements must be either positive integers or `-1`."): + with self.assertRaisesError(expected_error_message): sess.run(bijector.forward_event_shape_tensor(shape_in), feed_dict=feed_dict) + # pylint: enable=invalid-name def testValidButNonMatchingInputOpError(self): x = np.random.randn(4, 3, 2) @@ -184,7 +188,8 @@ class _ReshapeBijectorTest(object): sess.run(bijector.forward(x), feed_dict=feed_dict) - def testInputOutputMismatchOpError(self): + # pylint: disable=invalid-name + def _testInputOutputMismatchOpError(self, expected_error_message): x1 = np.random.randn(4, 2, 3) x2 = np.random.randn(4, 1, 1, 5) @@ -196,13 +201,11 @@ class _ReshapeBijectorTest(object): event_shape_in=shape_in, validate_args=True) - # test that *all* methods check basic assertions - with self.assertRaisesError( - "Input to reshape is a tensor with"): + with self.assertRaisesError(expected_error_message): sess.run(bijector.forward(x1), feed_dict=fd_mismatched) - with self.assertRaisesError( - "Input to reshape is a tensor with"): + with self.assertRaisesError(expected_error_message): sess.run(bijector.inverse(x2), feed_dict=fd_mismatched) + # pylint: enable=invalid-name def testOneShapePartiallySpecified(self): expected_x = np.random.randn(4, 6) @@ -262,6 +265,7 @@ class _ReshapeBijectorTest(object): raise NotImplementedError("Subclass failed to implement `build_shapes`.") +@test_util.with_c_api class ReshapeBijectorTestStatic(test.TestCase, _ReshapeBijectorTest): def build_shapes(self, shape_in, shape_out): @@ -299,7 +303,22 @@ class ReshapeBijectorTestStatic(test.TestCase, _ReshapeBijectorTest): validate_args=True) assert_bijective_and_finite(bijector, x, y, rtol=1e-6, atol=0) + def testInvalidDimensionsOpError(self): + if ops._USE_C_API: + error_message = "Invalid value in tensor used for shape: -2" + else: + error_message = "elements must be either positive integers or `-1`." + self._testInvalidDimensionsOpError(error_message) + + def testInputOutputMismatchOpError(self): + if ops._USE_C_API: + error_message = "Cannot reshape a tensor with" + else: + error_message = "Input to reshape is a tensor with" + self._testInputOutputMismatchOpError(error_message) + +@test_util.with_c_api class ReshapeBijectorTestDynamic(test.TestCase, _ReshapeBijectorTest): def build_shapes(self, shape_in, shape_out): @@ -313,7 +332,15 @@ class ReshapeBijectorTestDynamic(test.TestCase, _ReshapeBijectorTest): def assertRaisesError(self, msg): return self.assertRaisesOpError(msg) + def testInvalidDimensionsOpError(self): + self._testInvalidDimensionsOpError( + "elements must be either positive integers or `-1`.") + + def testInputOutputMismatchOpError(self): + self._testInputOutputMismatchOpError("Input to reshape is a tensor with") + +@test_util.with_c_api class ReshapeBijectorTestDynamicNdims(test.TestCase, _ReshapeBijectorTest): def build_shapes(self, shape_in, shape_out): @@ -325,6 +352,13 @@ class ReshapeBijectorTestDynamicNdims(test.TestCase, _ReshapeBijectorTest): def assertRaisesError(self, msg): return self.assertRaisesOpError(msg) + def testInvalidDimensionsOpError(self): + self._testInvalidDimensionsOpError( + "elements must be either positive integers or `-1`.") + + def testInputOutputMismatchOpError(self): + self._testInputOutputMismatchOpError("Input to reshape is a tensor with") + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/distribution_util_test.py b/tensorflow/contrib/distributions/python/kernel_tests/distribution_util_test.py index a255d4fc890e67180532e342332a8e3f63a869cd..31d24aa9ea09007b8db40e4869371b1f62639ac7 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/distribution_util_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/distribution_util_test.py @@ -23,10 +23,15 @@ import itertools import numpy as np from tensorflow.contrib.distributions.python.ops import distribution_util +from tensorflow.contrib.distributions.python.ops import mixture +from tensorflow.contrib.distributions.python.ops import mixture_same_family +from tensorflow.contrib.distributions.python.ops import mvn_diag from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops +from tensorflow.python.ops.distributions import categorical +from tensorflow.python.ops.distributions import normal from tensorflow.python.ops.linalg import linear_operator_diag import tensorflow.python.ops.nn_grad # pylint: disable=unused-import from tensorflow.python.platform import test @@ -395,6 +400,41 @@ class MixtureStddevTest(test.TestCase): self.assertAllClose(actual_devs, expected_devs) +class PadMixtureDimensionsTest(test.TestCase): + + def test_pad_mixture_dimensions_mixture(self): + with self.test_session() as sess: + gm = mixture.Mixture( + cat=categorical.Categorical(probs=[[0.3, 0.7]]), + components=[ + normal.Normal(loc=[-1.0], scale=[1.0]), + normal.Normal(loc=[1.0], scale=[0.5]) + ]) + + x = array_ops.constant([[1.0, 2.0], [3.0, 4.0]]) + x_pad = distribution_util.pad_mixture_dimensions( + x, gm, gm.cat, gm.event_shape.ndims) + x_out, x_pad_out = sess.run([x, x_pad]) + + self.assertAllEqual(x_pad_out.shape, [2, 2]) + self.assertAllEqual(x_out.reshape([-1]), x_pad_out.reshape([-1])) + + def test_pad_mixture_dimensions_mixture_same_family(self): + with self.test_session() as sess: + gm = mixture_same_family.MixtureSameFamily( + mixture_distribution=categorical.Categorical(probs=[0.3, 0.7]), + components_distribution=mvn_diag.MultivariateNormalDiag( + loc=[[-1., 1], [1, -1]], scale_identity_multiplier=[1.0, 0.5])) + + x = array_ops.constant([[1.0, 2.0], [3.0, 4.0]]) + x_pad = distribution_util.pad_mixture_dimensions( + x, gm, gm.mixture_distribution, gm.event_shape.ndims) + x_out, x_pad_out = sess.run([x, x_pad]) + + self.assertAllEqual(x_pad_out.shape, [2, 2, 1]) + self.assertAllEqual(x_out.reshape([-1]), x_pad_out.reshape([-1])) + + class _PadTest(object): def testNegAxisCorrectness(self): diff --git a/tensorflow/contrib/distributions/python/kernel_tests/kumaraswamy_test.py b/tensorflow/contrib/distributions/python/kernel_tests/kumaraswamy_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ea3c86b5c0f42b64fc6e4e362cbcc162bccf74a2 --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/kumaraswamy_test.py @@ -0,0 +1,388 @@ +# Copyright 2016 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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import importlib + +import numpy as np + +from tensorflow.contrib.distributions.python.ops import kumaraswamy as kumaraswamy_lib +from tensorflow.python.client import session +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import random_seed +from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test +from tensorflow.python.platform import tf_logging + + +def try_import(name): # pylint: disable=invalid-name + module = None + try: + module = importlib.import_module(name) + except ImportError as e: + tf_logging.warning("Could not import %s: %s" % (name, str(e))) + return module + + +special = try_import("scipy.special") +stats = try_import("scipy.stats") + + +def _kumaraswamy_mode(a, b): + a = np.asarray(a) + b = np.asarray(b) + return ((a - 1) / (a * b - 1))**(1 / a) + + +def _kumaraswamy_moment(a, b, n): + a = np.asarray(a) + b = np.asarray(b) + return b * special.beta(1.0 + n / a, b) + + +def _harmonic_number(b): + b = np.asarray(b) + return special.psi(b + 1) - special.psi(1) + + +def _kumaraswamy_cdf(a, b, x): + a = np.asarray(a) + b = np.asarray(b) + x = np.asarray(x) + return 1 - (1 - x**a)**b + + +def _kumaraswamy_pdf(a, b, x): + a = np.asarray(a) + b = np.asarray(b) + x = np.asarray(x) + return a * b * x ** (a - 1) * (1 - x ** a) ** (b - 1) + + +class KumaraswamyTest(test.TestCase): + + def testSimpleShapes(self): + with self.test_session(): + a = np.random.rand(3) + b = np.random.rand(3) + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertAllEqual([], dist.event_shape_tensor().eval()) + self.assertAllEqual([3], dist.batch_shape_tensor().eval()) + self.assertEqual(tensor_shape.TensorShape([]), dist.event_shape) + self.assertEqual(tensor_shape.TensorShape([3]), dist.batch_shape) + + def testComplexShapes(self): + with self.test_session(): + a = np.random.rand(3, 2, 2) + b = np.random.rand(3, 2, 2) + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertAllEqual([], dist.event_shape_tensor().eval()) + self.assertAllEqual([3, 2, 2], dist.batch_shape_tensor().eval()) + self.assertEqual(tensor_shape.TensorShape([]), dist.event_shape) + self.assertEqual(tensor_shape.TensorShape([3, 2, 2]), dist.batch_shape) + + def testComplexShapesBroadcast(self): + with self.test_session(): + a = np.random.rand(3, 2, 2) + b = np.random.rand(2, 2) + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertAllEqual([], dist.event_shape_tensor().eval()) + self.assertAllEqual([3, 2, 2], dist.batch_shape_tensor().eval()) + self.assertEqual(tensor_shape.TensorShape([]), dist.event_shape) + self.assertEqual(tensor_shape.TensorShape([3, 2, 2]), dist.batch_shape) + + def testAProperty(self): + a = [[1., 2, 3]] + b = [[2., 4, 3]] + with self.test_session(): + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertEqual([1, 3], dist.concentration1.get_shape()) + self.assertAllClose(a, dist.concentration1.eval()) + + def testBProperty(self): + a = [[1., 2, 3]] + b = [[2., 4, 3]] + with self.test_session(): + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertEqual([1, 3], dist.concentration0.get_shape()) + self.assertAllClose(b, dist.concentration0.eval()) + + def testPdfXProper(self): + a = [[1., 2, 3]] + b = [[2., 4, 3]] + with self.test_session(): + dist = kumaraswamy_lib.Kumaraswamy(a, b, validate_args=True) + dist.prob([.1, .3, .6]).eval() + dist.prob([.2, .3, .5]).eval() + # Either condition can trigger. + with self.assertRaisesOpError("sample must be positive"): + dist.prob([-1., 0.1, 0.5]).eval() + with self.assertRaisesOpError("sample must be positive"): + dist.prob([0., 0.1, 0.5]).eval() + with self.assertRaisesOpError("sample must be no larger than `1`"): + dist.prob([.1, .2, 1.2]).eval() + + def testPdfTwoBatches(self): + with self.test_session(): + a = [1., 2] + b = [1., 2] + x = [.5, .5] + dist = kumaraswamy_lib.Kumaraswamy(a, b) + pdf = dist.prob(x) + expected_pdf = _kumaraswamy_pdf(a, b, x) + self.assertAllClose(expected_pdf, pdf.eval()) + self.assertEqual((2,), pdf.get_shape()) + + def testPdfTwoBatchesNontrivialX(self): + with self.test_session(): + a = [1., 2] + b = [1., 2] + x = [.3, .7] + dist = kumaraswamy_lib.Kumaraswamy(a, b) + pdf = dist.prob(x) + expected_pdf = _kumaraswamy_pdf(a, b, x) + self.assertAllClose(expected_pdf, pdf.eval()) + self.assertEqual((2,), pdf.get_shape()) + + def testPdfUniformZeroBatch(self): + with self.test_session(): + # This is equivalent to a uniform distribution + a = 1. + b = 1. + x = np.array([.1, .2, .3, .5, .8], dtype=np.float32) + dist = kumaraswamy_lib.Kumaraswamy(a, b) + pdf = dist.prob(x) + expected_pdf = _kumaraswamy_pdf(a, b, x) + self.assertAllClose(expected_pdf, pdf.eval()) + self.assertEqual((5,), pdf.get_shape()) + + def testPdfAStretchedInBroadcastWhenSameRank(self): + with self.test_session(): + a = [[1., 2]] + b = [[1., 2]] + x = [[.5, .5], [.3, .7]] + dist = kumaraswamy_lib.Kumaraswamy(a, b) + pdf = dist.prob(x) + expected_pdf = _kumaraswamy_pdf(a, b, x) + self.assertAllClose(expected_pdf, pdf.eval()) + self.assertEqual((2, 2), pdf.get_shape()) + + def testPdfAStretchedInBroadcastWhenLowerRank(self): + with self.test_session(): + a = [1., 2] + b = [1., 2] + x = [[.5, .5], [.2, .8]] + pdf = kumaraswamy_lib.Kumaraswamy(a, b).prob(x) + expected_pdf = _kumaraswamy_pdf(a, b, x) + self.assertAllClose(expected_pdf, pdf.eval()) + self.assertEqual((2, 2), pdf.get_shape()) + + def testPdfXStretchedInBroadcastWhenSameRank(self): + with self.test_session(): + a = [[1., 2], [2., 3]] + b = [[1., 2], [2., 3]] + x = [[.5, .5]] + pdf = kumaraswamy_lib.Kumaraswamy(a, b).prob(x) + expected_pdf = _kumaraswamy_pdf(a, b, x) + self.assertAllClose(expected_pdf, pdf.eval()) + self.assertEqual((2, 2), pdf.get_shape()) + + def testPdfXStretchedInBroadcastWhenLowerRank(self): + with self.test_session(): + a = [[1., 2], [2., 3]] + b = [[1., 2], [2., 3]] + x = [.5, .5] + pdf = kumaraswamy_lib.Kumaraswamy(a, b).prob(x) + expected_pdf = _kumaraswamy_pdf(a, b, x) + self.assertAllClose(expected_pdf, pdf.eval()) + self.assertEqual((2, 2), pdf.get_shape()) + + def testKumaraswamyMean(self): + with session.Session(): + a = [1., 2, 3] + b = [2., 4, 1.2] + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertEqual(dist.mean().get_shape(), (3,)) + if not stats: + return + expected_mean = _kumaraswamy_moment(a, b, 1) + self.assertAllClose(expected_mean, dist.mean().eval()) + + def testKumaraswamyVariance(self): + with session.Session(): + a = [1., 2, 3] + b = [2., 4, 1.2] + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertEqual(dist.variance().get_shape(), (3,)) + if not stats: + return + expected_variance = _kumaraswamy_moment(a, b, 2) - _kumaraswamy_moment( + a, b, 1)**2 + self.assertAllClose(expected_variance, dist.variance().eval()) + + def testKumaraswamyMode(self): + with session.Session(): + a = np.array([1.1, 2, 3]) + b = np.array([2., 4, 1.2]) + expected_mode = _kumaraswamy_mode(a, b) + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertEqual(dist.mode().get_shape(), (3,)) + self.assertAllClose(expected_mode, dist.mode().eval()) + + def testKumaraswamyModeInvalid(self): + with session.Session(): + a = np.array([1., 2, 3]) + b = np.array([2., 4, 1.2]) + dist = kumaraswamy_lib.Kumaraswamy(a, b, allow_nan_stats=False) + with self.assertRaisesOpError("Condition x < y.*"): + dist.mode().eval() + + a = np.array([2., 2, 3]) + b = np.array([1., 4, 1.2]) + dist = kumaraswamy_lib.Kumaraswamy(a, b, allow_nan_stats=False) + with self.assertRaisesOpError("Condition x < y.*"): + dist.mode().eval() + + def testKumaraswamyModeEnableAllowNanStats(self): + with session.Session(): + a = np.array([1., 2, 3]) + b = np.array([2., 4, 1.2]) + dist = kumaraswamy_lib.Kumaraswamy(a, b, allow_nan_stats=True) + + expected_mode = _kumaraswamy_mode(a, b) + expected_mode[0] = np.nan + self.assertEqual((3,), dist.mode().get_shape()) + self.assertAllClose(expected_mode, dist.mode().eval()) + + a = np.array([2., 2, 3]) + b = np.array([1., 4, 1.2]) + dist = kumaraswamy_lib.Kumaraswamy(a, b, allow_nan_stats=True) + + expected_mode = _kumaraswamy_mode(a, b) + expected_mode[0] = np.nan + self.assertEqual((3,), dist.mode().get_shape()) + self.assertAllClose(expected_mode, dist.mode().eval()) + + def testKumaraswamyEntropy(self): + with session.Session(): + a = np.array([1., 2, 3]) + b = np.array([2., 4, 1.2]) + dist = kumaraswamy_lib.Kumaraswamy(a, b) + self.assertEqual(dist.entropy().get_shape(), (3,)) + if not stats: + return + expected_entropy = (1 - 1. / a) + ( + 1 - 1. / b) * _harmonic_number(b) + np.log(a * b) + self.assertAllClose(expected_entropy, dist.entropy().eval()) + + def testKumaraswamySample(self): + with self.test_session(): + a = 1. + b = 2. + kumaraswamy = kumaraswamy_lib.Kumaraswamy(a, b) + n = constant_op.constant(100000) + samples = kumaraswamy.sample(n) + sample_values = samples.eval() + self.assertEqual(sample_values.shape, (100000,)) + self.assertFalse(np.any(sample_values < 0.0)) + if not stats: + return + self.assertLess( + stats.kstest( + # Kumaraswamy is a univariate distribution. + sample_values, + lambda x: _kumaraswamy_cdf(1., 2., x))[0], + 0.01) + # The standard error of the sample mean is 1 / (sqrt(18 * n)) + expected_mean = _kumaraswamy_moment(a, b, 1) + self.assertAllClose(sample_values.mean(axis=0), expected_mean, atol=1e-2) + expected_variance = _kumaraswamy_moment(a, b, 2) - _kumaraswamy_moment( + a, b, 1)**2 + self.assertAllClose( + np.cov(sample_values, rowvar=0), expected_variance, atol=1e-1) + + # Test that sampling with the same seed twice gives the same results. + def testKumaraswamySampleMultipleTimes(self): + with self.test_session(): + a_val = 1. + b_val = 2. + n_val = 100 + + random_seed.set_random_seed(654321) + kumaraswamy1 = kumaraswamy_lib.Kumaraswamy( + concentration1=a_val, concentration0=b_val, name="kumaraswamy1") + samples1 = kumaraswamy1.sample(n_val, seed=123456).eval() + + random_seed.set_random_seed(654321) + kumaraswamy2 = kumaraswamy_lib.Kumaraswamy( + concentration1=a_val, concentration0=b_val, name="kumaraswamy2") + samples2 = kumaraswamy2.sample(n_val, seed=123456).eval() + + self.assertAllClose(samples1, samples2) + + def testKumaraswamySampleMultidimensional(self): + with self.test_session(): + a = np.random.rand(3, 2, 2).astype(np.float32) + b = np.random.rand(3, 2, 2).astype(np.float32) + kumaraswamy = kumaraswamy_lib.Kumaraswamy(a, b) + n = constant_op.constant(100000) + samples = kumaraswamy.sample(n) + sample_values = samples.eval() + self.assertEqual(sample_values.shape, (100000, 3, 2, 2)) + self.assertFalse(np.any(sample_values < 0.0)) + if not stats: + return + self.assertAllClose( + sample_values[:, 1, :].mean(axis=0), + _kumaraswamy_moment(a, b, 1)[1, :], + atol=1e-1) + + def testKumaraswamyCdf(self): + with self.test_session(): + shape = (30, 40, 50) + for dt in (np.float32, np.float64): + a = 10. * np.random.random(shape).astype(dt) + b = 10. * np.random.random(shape).astype(dt) + x = np.random.random(shape).astype(dt) + actual = kumaraswamy_lib.Kumaraswamy(a, b).cdf(x).eval() + self.assertAllEqual(np.ones(shape, dtype=np.bool), 0. <= x) + self.assertAllEqual(np.ones(shape, dtype=np.bool), 1. >= x) + if not stats: + return + self.assertAllClose( + _kumaraswamy_cdf(a, b, x), actual, rtol=1e-4, atol=0) + + def testKumaraswamyLogCdf(self): + with self.test_session(): + shape = (30, 40, 50) + for dt in (np.float32, np.float64): + a = 10. * np.random.random(shape).astype(dt) + b = 10. * np.random.random(shape).astype(dt) + x = np.random.random(shape).astype(dt) + actual = math_ops.exp(kumaraswamy_lib.Kumaraswamy(a, + b).log_cdf(x)).eval() + self.assertAllEqual(np.ones(shape, dtype=np.bool), 0. <= x) + self.assertAllEqual(np.ones(shape, dtype=np.bool), 1. >= x) + if not stats: + return + self.assertAllClose( + _kumaraswamy_cdf(a, b, x), actual, rtol=1e-4, atol=0) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/mixture_test.py b/tensorflow/contrib/distributions/python/kernel_tests/mixture_test.py index 1e514fe0ff21cd53c8c235da417890773db50c37..02064891758a86c5108e11da6a3666f2d5c56c64 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/mixture_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/mixture_test.py @@ -107,7 +107,7 @@ def _test_capture_normal_sample_outputs(): ds.Normal._call_sample_n = true_normal_call_sample_n -def make_univariate_mixture(batch_shape, num_components): +def make_univariate_mixture(batch_shape, num_components, use_static_graph): batch_shape = ops.convert_to_tensor(batch_shape, dtypes.int32) logits = random_ops.random_uniform( array_ops.concat((batch_shape, [num_components]), axis=0), @@ -119,11 +119,11 @@ def make_univariate_mixture(batch_shape, num_components): for _ in range(num_components) ] cat = ds.Categorical(logits, dtype=dtypes.int32) - return ds.Mixture(cat, components) + return ds.Mixture(cat, components, use_static_graph=use_static_graph) def make_multivariate_mixture(batch_shape, num_components, event_shape, - batch_shape_tensor=None): + use_static_graph, batch_shape_tensor=None): if batch_shape_tensor is None: batch_shape_tensor = batch_shape batch_shape_tensor = ops.convert_to_tensor(batch_shape_tensor, dtypes.int32) @@ -145,15 +145,17 @@ def make_multivariate_mixture(batch_shape, num_components, event_shape, loc=loc, scale_diag=scale_diag) components = [create_component() for _ in range(num_components)] cat = ds.Categorical(logits, dtype=dtypes.int32) - return ds.Mixture(cat, components) + return ds.Mixture(cat, components, use_static_graph=use_static_graph) class MixtureTest(test.TestCase): + use_static_graph = False def testShapes(self): with self.test_session(): for batch_shape in ([], [1], [2, 3, 4]): - dist = make_univariate_mixture(batch_shape, num_components=10) + dist = make_univariate_mixture(batch_shape, num_components=10, + use_static_graph=self.use_static_graph) self.assertAllEqual(batch_shape, dist.batch_shape) self.assertAllEqual(batch_shape, dist.batch_shape_tensor().eval()) self.assertAllEqual([], dist.event_shape) @@ -161,7 +163,8 @@ class MixtureTest(test.TestCase): for event_shape in ([1], [2]): dist = make_multivariate_mixture( - batch_shape, num_components=10, event_shape=event_shape) + batch_shape, num_components=10, event_shape=event_shape, + use_static_graph=self.use_static_graph) self.assertAllEqual(batch_shape, dist.batch_shape) self.assertAllEqual(batch_shape, dist.batch_shape_tensor().eval()) self.assertAllEqual(event_shape, dist.event_shape) @@ -172,7 +175,8 @@ class MixtureTest(test.TestCase): r"cat.num_classes != len"): ds.Mixture( ds.Categorical([0.1, 0.5]), # 2 classes - [ds.Normal(loc=1.0, scale=2.0)]) + [ds.Normal(loc=1.0, scale=2.0)], + use_static_graph=self.use_static_graph) with self.assertRaisesWithPredicateMatch( ValueError, r"\(\) and \(2,\) are not compatible"): # The value error is raised because the batch shapes of the @@ -185,13 +189,15 @@ class MixtureTest(test.TestCase): loc=1.0, scale=2.0), # scalar dist ds.Normal( loc=[1.0, 1.0], scale=[2.0, 2.0]) - ]) + ], + use_static_graph=self.use_static_graph) with self.assertRaisesWithPredicateMatch(ValueError, r"Could not infer"): cat_logits = array_ops.placeholder(shape=[1, None], dtype=dtypes.float32) ds.Mixture( ds.Categorical(cat_logits), [ds.Normal( - loc=[1.0], scale=[2.0])]) + loc=[1.0], scale=[2.0])], + use_static_graph=self.use_static_graph) def testBrokenShapesDynamic(self): with self.test_session(): @@ -203,29 +209,37 @@ class MixtureTest(test.TestCase): loc=d0_param, scale=d0_param), ds.Normal( loc=d1_param, scale=d1_param) ], - validate_args=True) - with self.assertRaisesOpError(r"batch shape must match"): + validate_args=True, + use_static_graph=self.use_static_graph) + + if self.use_static_graph: + error_string = r"Shapes of all inputs must match" + else: + error_string = r"batch shape must match" + + with self.assertRaisesOpError(error_string): d.sample().eval(feed_dict={d0_param: [2.0, 3.0], d1_param: [1.0]}) - with self.assertRaisesOpError(r"batch shape must match"): + with self.assertRaisesOpError(error_string): d.sample().eval(feed_dict={d0_param: [2.0, 3.0], d1_param: 1.0}) def testBrokenTypes(self): with self.assertRaisesWithPredicateMatch(TypeError, "Categorical"): - ds.Mixture(None, []) + ds.Mixture(None, [], use_static_graph=self.use_static_graph) cat = ds.Categorical([0.3, 0.2]) # components must be a list of distributions with self.assertRaisesWithPredicateMatch( TypeError, "all .* must be Distribution instances"): - ds.Mixture(cat, [None]) + ds.Mixture(cat, [None], use_static_graph=self.use_static_graph) with self.assertRaisesWithPredicateMatch(TypeError, "same dtype"): ds.Mixture( cat, [ ds.Normal(loc=[1.0], scale=[2.0]), ds.Normal(loc=[np.float16(1.0)], scale=[np.float16(2.0)]), - ]) + ], use_static_graph=self.use_static_graph) with self.assertRaisesWithPredicateMatch(ValueError, "non-empty list"): - ds.Mixture(ds.Categorical([0.3, 0.2]), None) + ds.Mixture(ds.Categorical([0.3, 0.2]), None, + use_static_graph=self.use_static_graph) # TODO(ebrevdo): once distribution Domains have been added, add a # test to ensure that the domains of the distributions in a @@ -235,7 +249,8 @@ class MixtureTest(test.TestCase): with self.test_session() as sess: for batch_shape in ((), (2,), (2, 3)): dist = make_univariate_mixture( - batch_shape=batch_shape, num_components=2) + batch_shape=batch_shape, num_components=2, + use_static_graph=self.use_static_graph) mean = dist.mean() self.assertEqual(batch_shape, mean.get_shape()) @@ -256,7 +271,8 @@ class MixtureTest(test.TestCase): with self.test_session() as sess: for batch_shape in ((), (2,), (2, 3)): dist = make_multivariate_mixture( - batch_shape=batch_shape, num_components=2, event_shape=(4,)) + batch_shape=batch_shape, num_components=2, event_shape=(4,), + use_static_graph=self.use_static_graph) mean = dist.mean() self.assertEqual(batch_shape + (4,), mean.get_shape()) @@ -283,7 +299,8 @@ class MixtureTest(test.TestCase): with self.test_session() as sess: for batch_shape in ((), (2,), (2, 3)): dist = make_univariate_mixture( - batch_shape=batch_shape, num_components=num_components) + batch_shape=batch_shape, num_components=num_components, + use_static_graph=self.use_static_graph) dev = dist.stddev() self.assertEqual(batch_shape, dev.get_shape()) @@ -325,7 +342,8 @@ class MixtureTest(test.TestCase): dist = make_multivariate_mixture( batch_shape=batch_shape, num_components=num_components, - event_shape=(4,)) + event_shape=(4,), + use_static_graph=self.use_static_graph) dev = dist.stddev() self.assertEqual(batch_shape + (4,), dev.get_shape()) @@ -371,7 +389,8 @@ class MixtureTest(test.TestCase): scale=component_devs[0]), ds.Normal(loc=component_means[1], scale=component_devs[1]), - ]) + ], + use_static_graph=self.use_static_graph) mix_dev = mixture_dist.stddev() with self.test_session() as sess: actual_stddev = sess.run(mix_dev) @@ -379,7 +398,8 @@ class MixtureTest(test.TestCase): def testProbScalarUnivariate(self): with self.test_session() as sess: - dist = make_univariate_mixture(batch_shape=[], num_components=2) + dist = make_univariate_mixture(batch_shape=[], num_components=2, + use_static_graph=self.use_static_graph) for x in [ np.array( [1.0, 2.0], dtype=np.float32), np.array( @@ -405,7 +425,8 @@ class MixtureTest(test.TestCase): def testProbScalarMultivariate(self): with self.test_session() as sess: dist = make_multivariate_mixture( - batch_shape=[], num_components=2, event_shape=[3]) + batch_shape=[], num_components=2, event_shape=[3], + use_static_graph=self.use_static_graph) for x in [ np.array( [[-1.0, 0.0, 1.0], [0.5, 1.0, -0.3]], dtype=np.float32), np.array( @@ -432,7 +453,8 @@ class MixtureTest(test.TestCase): def testProbBatchUnivariate(self): with self.test_session() as sess: - dist = make_univariate_mixture(batch_shape=[2, 3], num_components=2) + dist = make_univariate_mixture(batch_shape=[2, 3], num_components=2, + use_static_graph=self.use_static_graph) for x in [ np.random.randn(2, 3).astype(np.float32), @@ -459,7 +481,8 @@ class MixtureTest(test.TestCase): def testProbBatchMultivariate(self): with self.test_session() as sess: dist = make_multivariate_mixture( - batch_shape=[2, 3], num_components=2, event_shape=[4]) + batch_shape=[2, 3], num_components=2, event_shape=[4], + use_static_graph=self.use_static_graph) for x in [ np.random.randn(2, 3, 4).astype(np.float32), @@ -487,7 +510,8 @@ class MixtureTest(test.TestCase): num_components = 3 batch_shape = [] dist = make_univariate_mixture( - batch_shape=batch_shape, num_components=num_components) + batch_shape=batch_shape, num_components=num_components, + use_static_graph=self.use_static_graph) n = 4 with _test_capture_normal_sample_outputs() as component_samples: samples = dist.sample(n, seed=123) @@ -502,7 +526,10 @@ class MixtureTest(test.TestCase): which_c = np.where(cat_sample_values == c)[0] size_c = which_c.size # Scalar Batch univariate case: batch_size == 1, rank 1 - which_dist_samples = dist_sample_values[c][:size_c] + if self.use_static_graph: + which_dist_samples = dist_sample_values[c][which_c] + else: + which_dist_samples = dist_sample_values[c][:size_c] self.assertAllClose(which_dist_samples, sample_values[which_c]) # Test that sampling with the same seed twice gives the same results. @@ -522,7 +549,8 @@ class MixtureTest(test.TestCase): ] cat = ds.Categorical( logits, dtype=dtypes.int32, name="cat1") - dist1 = ds.Mixture(cat, components, name="mixture1") + dist1 = ds.Mixture(cat, components, name="mixture1", + use_static_graph=self.use_static_graph) samples1 = dist1.sample(n, seed=123456).eval() random_seed.set_random_seed(654321) @@ -532,7 +560,8 @@ class MixtureTest(test.TestCase): ] cat2 = ds.Categorical( logits, dtype=dtypes.int32, name="cat2") - dist2 = ds.Mixture(cat2, components2, name="mixture2") + dist2 = ds.Mixture(cat2, components2, name="mixture2", + use_static_graph=self.use_static_graph) samples2 = dist2.sample(n, seed=123456).eval() self.assertAllClose(samples1, samples2) @@ -541,7 +570,8 @@ class MixtureTest(test.TestCase): with self.test_session() as sess: num_components = 3 dist = make_multivariate_mixture( - batch_shape=[], num_components=num_components, event_shape=[2]) + batch_shape=[], num_components=num_components, event_shape=[2], + use_static_graph=self.use_static_graph) n = 4 with _test_capture_mvndiag_sample_outputs() as component_samples: samples = dist.sample(n, seed=123) @@ -555,14 +585,18 @@ class MixtureTest(test.TestCase): which_c = np.where(cat_sample_values == c)[0] size_c = which_c.size # Scalar Batch multivariate case: batch_size == 1, rank 2 - which_dist_samples = dist_sample_values[c][:size_c, :] + if self.use_static_graph: + which_dist_samples = dist_sample_values[c][which_c, :] + else: + which_dist_samples = dist_sample_values[c][:size_c, :] self.assertAllClose(which_dist_samples, sample_values[which_c, :]) def testSampleBatchUnivariate(self): with self.test_session() as sess: num_components = 3 dist = make_univariate_mixture( - batch_shape=[2, 3], num_components=num_components) + batch_shape=[2, 3], num_components=num_components, + use_static_graph=self.use_static_graph) n = 4 with _test_capture_normal_sample_outputs() as component_samples: samples = dist.sample(n, seed=123) @@ -576,8 +610,12 @@ class MixtureTest(test.TestCase): which_c_s, which_c_b0, which_c_b1 = np.where(cat_sample_values == c) size_c = which_c_s.size # Batch univariate case: batch_size == [2, 3], rank 3 - which_dist_samples = dist_sample_values[c][range(size_c), which_c_b0, - which_c_b1] + if self.use_static_graph: + which_dist_samples = dist_sample_values[c][which_c_s, which_c_b0, + which_c_b1] + else: + which_dist_samples = dist_sample_values[c][range(size_c), which_c_b0, + which_c_b1] self.assertAllClose(which_dist_samples, sample_values[which_c_s, which_c_b0, which_c_b1]) @@ -594,7 +632,8 @@ class MixtureTest(test.TestCase): dist = make_multivariate_mixture( batch_shape=batch_shape, num_components=num_components, event_shape=[4], - batch_shape_tensor=batch_shape_tensor) + batch_shape_tensor=batch_shape_tensor, + use_static_graph=self.use_static_graph) n = 5 with _test_capture_mvndiag_sample_outputs() as component_samples: samples = dist.sample(n, seed=123) @@ -617,8 +656,12 @@ class MixtureTest(test.TestCase): which_c_s, which_c_b0, which_c_b1 = np.where(cat_sample_values == c) size_c = which_c_s.size # Batch univariate case: batch_size == [2, 3], rank 4 (multivariate) - which_dist_samples = dist_sample_values[c][range(size_c), which_c_b0, - which_c_b1, :] + if self.use_static_graph: + which_dist_samples = dist_sample_values[c][which_c_s, which_c_b0, + which_c_b1, :] + else: + which_dist_samples = dist_sample_values[c][range(size_c), which_c_b0, + which_c_b1, :] self.assertAllClose(which_dist_samples, sample_values[which_c_s, which_c_b0, which_c_b1, :]) @@ -632,7 +675,8 @@ class MixtureTest(test.TestCase): with self.test_session() as sess: for batch_shape in ((), (2,), (2, 3)): dist = make_multivariate_mixture( - batch_shape=batch_shape, num_components=2, event_shape=(4,)) + batch_shape=batch_shape, num_components=2, event_shape=(4,), + use_static_graph=self.use_static_graph) entropy_lower_bound = dist.entropy_lower_bound() self.assertEqual(batch_shape, entropy_lower_bound.get_shape()) @@ -673,7 +717,8 @@ class MixtureTest(test.TestCase): cat_tf = ds.Categorical(probs=mixture_weights) components_tf = [ds.Normal(loc=mu, scale=sigma) for (mu, sigma) in zip(means, sigmas)] - mixture_tf = ds.Mixture(cat=cat_tf, components=components_tf) + mixture_tf = ds.Mixture(cat=cat_tf, components=components_tf, + use_static_graph=self.use_static_graph) x_tensor = array_ops.placeholder(shape=(), dtype=dtypes.float32) @@ -721,7 +766,8 @@ class MixtureTest(test.TestCase): cat_tf = ds.Categorical(probs=mixture_weights) components_tf = [ds.Normal(loc=mu, scale=sigma) for (mu, sigma) in zip(means, sigmas)] - mixture_tf = ds.Mixture(cat=cat_tf, components=components_tf) + mixture_tf = ds.Mixture(cat=cat_tf, components=components_tf, + use_static_graph=self.use_static_graph) x_tensor = array_ops.placeholder(shape=psize, dtype=dtypes.float32) xs_to_check = [ @@ -760,12 +806,18 @@ class MixtureTest(test.TestCase): gm = ds.Mixture( cat=ds.Categorical(probs=[.3, .7]), components=[ds.Gamma(1., 2.), - ds.Gamma(2., 1.)]) + ds.Gamma(2., 1.)], + use_static_graph=self.use_static_graph) x_ = gm.sample().eval() self.assertAllEqual([], x_.shape) +class MixtureStaticSampleTest(MixtureTest): + use_static_graph = True + + class MixtureBenchmark(test.Benchmark): + use_static_graph = False def _runSamplingBenchmark(self, name, create_distribution, use_gpu, num_components, batch_size, num_features, @@ -811,7 +863,7 @@ class MixtureBenchmark(test.Benchmark): components = list( ds.MultivariateNormalDiag( loc=mu, scale_diag=sigma) for (mu, sigma) in zip(mus, sigmas)) - return ds.Mixture(cat, components) + return ds.Mixture(cat, components, use_static_graph=self.use_static_graph) for use_gpu in False, True: if use_gpu and not test.is_gpu_available(): @@ -853,7 +905,7 @@ class MixtureBenchmark(test.Benchmark): ds.MultivariateNormalTriL( loc=mu, scale_tril=linalg_ops.cholesky(sigma)) for (mu, sigma) in zip(mus, sigmas)) - return ds.Mixture(cat, components) + return ds.Mixture(cat, components, use_static_graph=self.use_static_graph) for use_gpu in False, True: if use_gpu and not test.is_gpu_available(): @@ -872,5 +924,9 @@ class MixtureBenchmark(test.Benchmark): sample_size=sample_size) +class MixtureStaticSampleBenchmark(MixtureBenchmark): + use_static_graph = True + + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/vector_diffeomixture_test.py b/tensorflow/contrib/distributions/python/kernel_tests/vector_diffeomixture_test.py index d292b04665e34196670ee4f1c1655f805e04e06a..04f047aa0c81b3f59b97f14554fb59cb1b3dd8af 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/vector_diffeomixture_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/vector_diffeomixture_test.py @@ -27,6 +27,8 @@ from tensorflow.python.ops.linalg import linear_operator_diag as linop_diag_lib from tensorflow.python.ops.linalg import linear_operator_identity as linop_identity_lib from tensorflow.python.platform import test +rng = np.random.RandomState(0) + class VectorDiffeomixtureTest( test_util.VectorDistributionTestHelpers, test.TestCase): @@ -37,7 +39,7 @@ class VectorDiffeomixtureTest( dims = 4 vdm = vdm_lib.VectorDiffeomixture( mix_loc=[[0.], [1.]], - mix_scale=[1.], + temperature=[1.], distribution=normal_lib.Normal(0., 1.), loc=[ None, @@ -66,7 +68,7 @@ class VectorDiffeomixtureTest( dims = 4 vdm = vdm_lib.VectorDiffeomixture( mix_loc=[[0.], [1.]], - mix_scale=[1.], + temperature=[1.], distribution=normal_lib.Normal(1., 1.5), loc=[ None, @@ -95,7 +97,7 @@ class VectorDiffeomixtureTest( dims = 4 vdm = vdm_lib.VectorDiffeomixture( mix_loc=[[0.], [1.]], - mix_scale=[1.], + temperature=[1.], distribution=normal_lib.Normal(0., 1.), loc=[ None, @@ -122,12 +124,39 @@ class VectorDiffeomixtureTest( self.run_test_sample_consistent_log_prob( sess.run, vdm, radius=4., center=2., rtol=0.01) + def testSampleProbConsistentBroadcastMixTwoBatchDims(self): + dims = 4 + loc_1 = rng.randn(2, 3, dims).astype(np.float32) + + with self.test_session() as sess: + vdm = vdm_lib.VectorDiffeomixture( + mix_loc=(rng.rand(2, 3, 1) - 0.5).astype(np.float32), + temperature=[1.], + distribution=normal_lib.Normal(0., 1.), + loc=[ + None, + loc_1, + ], + scale=[ + linop_identity_lib.LinearOperatorScaledIdentity( + num_rows=dims, + multiplier=[np.float32(1.1)], + is_positive_definite=True), + ] * 2, + validate_args=True) + # Ball centered at component0's mean. + self.run_test_sample_consistent_log_prob( + sess.run, vdm, radius=2., center=0., rtol=0.01) + # Larger ball centered at component1's mean. + self.run_test_sample_consistent_log_prob( + sess.run, vdm, radius=3., center=loc_1, rtol=0.02) + def testMeanCovarianceNoBatch(self): with self.test_session() as sess: dims = 3 vdm = vdm_lib.VectorDiffeomixture( mix_loc=[[0.], [4.]], - mix_scale=[10.], + temperature=[1 / 10.], distribution=normal_lib.Normal(0., 1.), loc=[ np.float32([-2.]), @@ -147,12 +176,94 @@ class VectorDiffeomixtureTest( self.run_test_sample_consistent_mean_covariance( sess.run, vdm, rtol=0.02, cov_rtol=0.08) + def testTemperatureControlsHowMuchThisLooksLikeDiscreteMixture(self): + # As temperature decreases, this should approach a mixture of normals, with + # components at -2, 2. + with self.test_session() as sess: + dims = 1 + vdm = vdm_lib.VectorDiffeomixture( + mix_loc=[0.], + temperature=[[2.], [1.], [0.2]], + distribution=normal_lib.Normal(0., 1.), + loc=[ + np.float32([-2.]), + np.float32([2.]), + ], + scale=[ + linop_identity_lib.LinearOperatorScaledIdentity( + num_rows=dims, + multiplier=np.float32(0.5), + is_positive_definite=True), + ] * 2, # Use the same scale for each component. + quadrature_size=8, + validate_args=True) + + samps = vdm.sample(10000) + self.assertAllEqual((10000, 3, 1), samps.shape) + samps_ = sess.run(samps).reshape(10000, 3) # Make scalar event shape. + + # One characteristic of a discrete mixture (as opposed to a "smear") is + # that more weight is put near the component centers at -2, 2, and thus + # less weight is put near the origin. + prob_of_being_near_origin = (np.abs(samps_) < 1).mean(axis=0) + self.assertGreater( + prob_of_being_near_origin[0], prob_of_being_near_origin[1]) + self.assertGreater( + prob_of_being_near_origin[1], prob_of_being_near_origin[2]) + + # Run this test as well, just because we can. + self.run_test_sample_consistent_mean_covariance( + sess.run, vdm, rtol=0.02, cov_rtol=0.08) + + def testConcentrationLocControlsHowMuchWeightIsOnEachComponent(self): + with self.test_session() as sess: + dims = 1 + vdm = vdm_lib.VectorDiffeomixture( + mix_loc=[[-1.], [0.], [1.]], + temperature=[0.5], + distribution=normal_lib.Normal(0., 1.), + loc=[ + np.float32([-2.]), + np.float32([2.]), + ], + scale=[ + linop_identity_lib.LinearOperatorScaledIdentity( + num_rows=dims, + multiplier=np.float32(0.5), + is_positive_definite=True), + ] * 2, # Use the same scale for each component. + quadrature_size=8, + validate_args=True) + + samps = vdm.sample(10000) + self.assertAllEqual((10000, 3, 1), samps.shape) + samps_ = sess.run(samps).reshape(10000, 3) # Make scalar event shape. + + # One characteristic of putting more weight on a component is that the + # mean is closer to that component's mean. + # Get the mean for each batch member, the names signify the value of + # concentration for that batch member. + mean_neg1, mean_0, mean_1 = samps_.mean(axis=0) + + # Since concentration is the concentration for component 0, + # concentration = -1 ==> more weight on component 1, which has mean = 2 + # concentration = 0 ==> equal weight + # concentration = 1 ==> more weight on component 0, which has mean = -2 + self.assertLess(-2, mean_1) + self.assertLess(mean_1, mean_0) + self.assertLess(mean_0, mean_neg1) + self.assertLess(mean_neg1, 2) + + # Run this test as well, just because we can. + self.run_test_sample_consistent_mean_covariance( + sess.run, vdm, rtol=0.02, cov_rtol=0.08) + def testMeanCovarianceNoBatchUncenteredNonStandardBase(self): with self.test_session() as sess: dims = 3 vdm = vdm_lib.VectorDiffeomixture( mix_loc=[[0.], [4.]], - mix_scale=[10.], + temperature=[0.1], distribution=normal_lib.Normal(-1., 1.5), loc=[ np.float32([-2.]), @@ -177,7 +288,7 @@ class VectorDiffeomixtureTest( dims = 3 vdm = vdm_lib.VectorDiffeomixture( mix_loc=[[0.], [4.]], - mix_scale=[10.], + temperature=[0.1], distribution=normal_lib.Normal(0., 1.), loc=[ np.float32([[-2.]]), @@ -205,7 +316,7 @@ class VectorDiffeomixtureTest( dims = 4 vdm = vdm_lib.VectorDiffeomixture( mix_loc=[0.], - mix_scale=[1.], + temperature=[0.1], distribution=normal_lib.Normal(0., 1.), loc=[ None, @@ -229,29 +340,6 @@ class VectorDiffeomixtureTest( self.run_test_sample_consistent_log_prob( sess.run, vdm, radius=4., center=2., rtol=0.005) - # TODO(jvdillon): We've tested that (i) .sample and .log_prob are consistent, - # (ii) .mean, .stddev etc... and .sample are consistent. However, we haven't - # tested that the quadrature approach well-approximates the integral. - # - # To that end, consider adding these tests: - # - # Test1: In the limit of high mix_scale, this approximates a discrete mixture, - # and there are many discrete mixtures where we can explicitly compute - # mean/var, etc... So test1 would choose one of those discrete mixtures and - # show our mean/var/etc... is close to that. - # - # Test2: In the limit of low mix_scale, the a diffeomixture of Normal(-5, 1), - # Normal(5, 1) should (I believe...must check) should look almost like - # Uniform(-5, 5), and thus (i) .prob(x) should be about 1/10 for x in (-5, 5), - # and (ii) the first few moments should approximately match that of - # Uniform(-5, 5) - # - # Test3: If mix_loc is symmetric, then for any mix_scale, our - # quadrature-based diffeomixture of Normal(-1, 1), Normal(1, 1) should have - # mean zero, exactly. - - # TODO(jvdillon): Add more tests which verify broadcasting. - if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distributions/python/ops/autoregressive.py b/tensorflow/contrib/distributions/python/ops/autoregressive.py new file mode 100644 index 0000000000000000000000000000000000000000..852298bf334666db003353d5fc8e172ffb738668 --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/autoregressive.py @@ -0,0 +1,208 @@ +# Copyright 2017 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. +# ============================================================================== +"""The Autoregressive distribution.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.python.framework import ops +from tensorflow.python.ops.distributions import distribution as distribution_lib +from tensorflow.python.ops.distributions import util as distribution_util + + +class Autoregressive(distribution_lib.Distribution): + """Autoregressive distributions. + + The Autoregressive distribution enables learning (often) richer multivariate + distributions by repeatedly applying a [diffeomorphic]( + https://en.wikipedia.org/wiki/Diffeomorphism) transformation (such as + implemented by `Bijector`s). Regarding terminology, + + "Autoregressive models decompose the joint density as a product of + conditionals, and model each conditional in turn. Normalizing flows + transform a base density (e.g. a standard Gaussian) into the target density + by an invertible transformation with tractable Jacobian." [1] + + In other words, the "autoregressive property" is equivalent to the + decomposition, `p(x) = prod{ p(x[i] | x[0:i]) : i=0, ..., d }`. The provided + `shift_and_log_scale_fn`, `masked_autoregressive_default_template`, achieves + this property by zeroing out weights in its `masked_dense` layers. + + Practically speaking the autoregressive property means that there exists a + permutation of the event coordinates such that each coordinate is a + diffeomorphic function of only preceding coordinates. [2] + + #### Mathematical Details + + The probability function is, + + ```none + prob(x; fn, n) = fn(x).prob(x) + ``` + + And a sample is generated by, + + ```none + x = fn(...fn(fn(x0).sample()).sample()).sample() + ``` + + where the ellipses (`...`) represent `n-2` composed calls to `fn`, `fn` + constructs a `tf.distributions.Distribution`-like instance, and `x0` is a + fixed initializing `Tensor`. + + #### Examples + + ```python + tfd = tf.contrib.distributions + + def normal_fn(self, event_size): + n = event_size * (event_size + 1) / 2 + p = tf.Variable(tfd.Normal(loc=0., scale=1.).sample(n)) + affine = tfd.bijectors.Affine( + scale_tril=tfd.fill_triangular(0.25 * p)) + def _fn(samples): + scale = math_ops.exp(affine.forward(samples)).eval() + return independent_lib.Independent( + normal_lib.Normal(loc=0., scale=scale, validate_args=True), + reinterpreted_batch_ndims=1) + return _fn + + batch_and_event_shape = [3, 2, 4] + sample0 = array_ops.zeros(batch_and_event_shape) + ar = autoregressive_lib.Autoregressive( + self._normal_fn(batch_and_event_shape[-1]), sample0) + x = ar.sample([6, 5]) + # ==> x.shape = [6, 5, 3, 2, 4] + prob_x = ar.prob(x) + # ==> x.shape = [6, 5, 3, 2] + + ``` + + [1]: "Masked Autoregressive Flow for Density Estimation." + George Papamakarios, Theo Pavlakou, Iain Murray. Arxiv. 2017. + https://arxiv.org/abs/1705.07057 + + [2]: "Conditional Image Generation with PixelCNN Decoders." + Aaron van den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex + Graves, Koray Kavukcuoglu. Arxiv, 2016. + https://arxiv.org/abs/1606.05328 + """ + + def __init__(self, + distribution_fn, + sample0=None, + num_steps=None, + validate_args=False, + allow_nan_stats=True, + name="Autoregressive"): + """Construct an `Autoregressive` distribution. + + Args: + distribution_fn: Python `callable` which constructs a + `tf.distributions.Distribution`-like instance from a `Tensor` (e.g., + `sample0`). The function must respect the "autoregressive property", + i.e., there exists a permutation of event such that each coordinate is a + diffeomorphic function of on preceding coordinates. + sample0: Initial input to `distribution_fn`; used to + build the distribution in `__init__` which in turn specifies this + distribution's properties, e.g., `event_shape`, `batch_shape`, `dtype`. + If unspecified, then `distribution_fn` should be default constructable. + num_steps: Number of times `distribution_fn` is composed from samples, + e.g., `num_steps=2` implies + `distribution_fn(distribution_fn(sample0).sample(n)).sample()`. + validate_args: Python `bool`. Whether to validate input with asserts. + If `validate_args` is `False`, and the inputs are invalid, + correct behavior is not guaranteed. + allow_nan_stats: Python `bool`, default `True`. When `True`, statistics + (e.g., mean, mode, variance) use the value "`NaN`" to indicate the + result is undefined. When `False`, an exception is raised if one or + more of the statistic's batch members are undefined. + name: Python `str` name prefixed to Ops created by this class. + Default value: "Autoregressive". + + Raises: + ValueError: if `num_steps` and + `distribution_fn(sample0).event_shape.num_elements()` are both `None`. + ValueError: if `num_steps < 1`. + """ + parameters = locals() + with ops.name_scope(name): + self._distribution_fn = distribution_fn + self._sample0 = sample0 + self._distribution0 = (distribution_fn() if sample0 is None + else distribution_fn(sample0)) + if num_steps is None: + num_steps = self._distribution0.event_shape.num_elements() + if num_steps is None: + raise ValueError("distribution_fn must generate a distribution " + "with fully known `event_shape`.") + if num_steps < 1: + raise ValueError("num_steps ({}) must be at least 1.".format(num_steps)) + self._num_steps = num_steps + super(Autoregressive, self).__init__( + dtype=self._distribution0.dtype, + reparameterization_type=self._distribution0.reparameterization_type, + validate_args=validate_args, + allow_nan_stats=allow_nan_stats, + parameters=parameters, + graph_parents=self._distribution0._graph_parents, # pylint: disable=protected-access + name=name) + + @property + def distribution_fn(self): + return self._distribution_fn + + @property + def sample0(self): + return self._sample0 + + @property + def num_steps(self): + return self._num_steps + + @property + def distribution0(self): + return self._distribution0 + + def _batch_shape(self): + return self.distribution0.batch_shape + + def _batch_shape_tensor(self): + return self.distribution0.batch_shape_tensor() + + def _event_shape(self): + return self.distribution0.event_shape + + def _event_shape_tensor(self): + return self.distribution0.event_shape_tensor() + + def _sample_n(self, n, seed=None): + if seed is None: + seed = distribution_util.gen_new_seed( + seed=np.random.randint(2**32 - 1), + salt="autoregressive") + samples = self.distribution0.sample(n, seed=seed) + for _ in range(self._num_steps): + samples = self.distribution_fn(samples).sample(seed=seed) + return samples + + def _log_prob(self, value): + return self.distribution_fn(value).log_prob(value) + + def _prob(self, value): + return self.distribution_fn(value).prob(value) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py b/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py index bc0ec7f195af009c87020ce8c4ea18f2e713759a..93923c3f083c7f5136b55e9021cbd6323684b976 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py @@ -29,6 +29,7 @@ @@MaskedAutoregressiveFlow @@Permute @@PowerTransform +@@RealNVP @@Reshape @@Sigmoid @@SigmoidCentered @@ -39,6 +40,7 @@ @@masked_autoregressive_default_template @@masked_dense +@@real_nvp_default_template """ from __future__ import absolute_import @@ -60,6 +62,7 @@ from tensorflow.contrib.distributions.python.ops.bijectors.invert import * from tensorflow.contrib.distributions.python.ops.bijectors.masked_autoregressive import * from tensorflow.contrib.distributions.python.ops.bijectors.permute import * from tensorflow.contrib.distributions.python.ops.bijectors.power_transform import * +from tensorflow.contrib.distributions.python.ops.bijectors.real_nvp import * from tensorflow.contrib.distributions.python.ops.bijectors.reshape import * from tensorflow.contrib.distributions.python.ops.bijectors.sigmoid import * from tensorflow.contrib.distributions.python.ops.bijectors.sigmoid_centered import * diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py b/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py index 06c7c61ec3dc3980e0d12a984739dca5a925ac9f..5251dbcb5748f75688aa43ce6e4e9dbd76be78bb 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py @@ -182,6 +182,7 @@ class MaskedAutoregressiveFlow(bijector_lib.Bijector): shift_and_log_scale_fn, is_constant_jacobian=False, validate_args=False, + unroll_loop=False, name=None): """Creates the MaskedAutoregressiveFlow bijector. @@ -201,17 +202,46 @@ class MaskedAutoregressiveFlow(bijector_lib.Bijector): inefficient.) validate_args: Python `bool` indicating whether arguments should be checked for correctness. + unroll_loop: Python `bool` indicating whether the `tf.while_loop` in + `_forward` should be replaced with a static for loop. Requires that + the final dimension of `x` be known at graph construction time. Defaults + to `False`. name: Python `str`, name given to ops managed by this object. """ name = name or "masked_autoregressive_flow" self._shift_and_log_scale_fn = shift_and_log_scale_fn + self._unroll_loop = unroll_loop super(MaskedAutoregressiveFlow, self).__init__( is_constant_jacobian=is_constant_jacobian, validate_args=validate_args, name=name) def _forward(self, x): + if self._unroll_loop: + event_size = x.shape.with_rank_at_least(1)[-1].value + if event_size is None: + raise ValueError( + "The final dimension of `x` must be known at graph construction " + "time if `unroll_loop=True`. `x.shape: %r`" % x.shape) + y = array_ops.zeros_like(x, name="y0") + + for _ in range(event_size): + shift, log_scale = self._shift_and_log_scale_fn(y) + # next_y = scale * x + shift + next_y = x + if log_scale is not None: + next_y *= math_ops.exp(log_scale) + if shift is not None: + next_y += shift + y = next_y + return y + event_size = array_ops.shape(x)[-1] + # If the event size is available at graph construction time, we can inform + # the graph compiler of the maximum number of steps. If not, + # static_event_size will be None, and the maximum_iterations argument will + # have no effect. + static_event_size = x.shape.with_rank_at_least(1)[-1].value y0 = array_ops.zeros_like(x, name="y0") # call the template once to ensure creation _ = self._shift_and_log_scale_fn(y0) @@ -233,7 +263,8 @@ class MaskedAutoregressiveFlow(bijector_lib.Bijector): _, y = control_flow_ops.while_loop( cond=lambda index, _: index < event_size, body=_loop_body, - loop_vars=[0, y0]) + loop_vars=(0, y0), + maximum_iterations=static_event_size) return y def _inverse(self, y): diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py b/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py new file mode 100644 index 0000000000000000000000000000000000000000..2840f52e742eac5e9e37a576bf7f6d6f05a07a35 --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py @@ -0,0 +1,282 @@ +# Copyright 2017 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. +# ============================================================================== +"""Real NVP bijector.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.layers import core as layers +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn_ops +from tensorflow.python.ops import template as template_ops +from tensorflow.python.ops.distributions import bijector as bijector_lib + + +__all__ = [ + "RealNVP", + "real_nvp_default_template" +] + + +class RealNVP(bijector_lib.Bijector): + """RealNVP "affine coupling layer" for vector-valued events. + + Real NVP models a normalizing flow on a `D`-dimensional distribution via a + single `D-d`-dimensional conditional distribution [1]: + + `y[d:D] = y[d:D] * math_ops.exp(log_scale_fn(y[d:D])) + shift_fn(y[d:D])` + `y[0:d] = x[0:d]` + + The last `D-d` units are scaled and shifted based on the first `d` units only, + while the first `d` units are 'masked' and left unchanged. Real NVP's + `shift_and_log_scale_fn` computes vector-valued quantities. For + scale-and-shift transforms that do not depend on any masked units, i.e. + `d=0`, use the `tfb.Affine` bijector with learned parameters instead. + + Masking is currently only supported for base distributions with + `event_ndims=1`. For more sophisticated masking schemes like checkerboard or + channel-wise masking [2], use the `tfb.Permute` bijector to re-order desired + masked units into the first `d` units. For base distributions with + `event_ndims > 1`, use the `tfb.Reshape` bijector to flatten the event shape. + + Recall that the MAF bijector [2] implements a normalizing flow via an + autoregressive transformation. MAF and IAF have opposite computational + tradeoffs - MAF can train all units in parallel but must sample units + sequentially, while IAF must train units sequentially but can sample in + parallel. In contrast, Real NVP can compute both forward and inverse + computations in parallel. However, the lack of an autoregressive + transformations makes it less expressive on a per-bijector basis. + + A "valid" `shift_and_log_scale_fn` must compute each `shift` (aka `loc` or + "mu" [2]) and `log(scale)` (aka "alpha" [2]) such that each are broadcastable + with the arguments to `forward` and `inverse`, i.e., such that the + calculations in `forward`, `inverse` [below] are possible. For convenience, + `real_nvp_default_nvp` is offered as a possible `shift_and_log_scale_fn` + function. + + NICE [3] is a special case of the Real NVP bijector which discards the scale + transformation, resulting in a constant-time inverse-log-determinant-Jacobian. + To use a NICE bijector instead of Real NVP, `shift_and_log_scale_fn` should + return `(shift, None)`, and `is_constant_jacobian` should be set to `True` in + the `RealNVP` constructor. Calling `real_nvp_default_template` with + `shift_only=True` returns one such NICE-compatible `shift_and_log_scale_fn`. + + Caching: the scalar input depth `D` of the base distribution is not known at + construction time. The first call to any of `forward(x)`, `inverse(x)`, + `inverse_log_det_jacobian(x)`, or `forward_log_det_jacobian(x)` memoizes + `D`, which is re-used in subsequent calls. This shape must be known prior to + graph execution (which is the case if using tf.layers). + + #### Example Use + + ```python + tfd = tf.contrib.distributions + tfb = tfd.bijectors + + # A common choice for a normalizing flow is to use a Gaussian for the base + # distribution. (However, any continuous distribution would work.) E.g., + nvp = tfd.TransformedDistribution( + distribution=tfd.MultivariateNormalDiag(loc=[0., 0., 0.])), + bijector=tfb.RealNVP( + num_masked=2, + shift_and_log_scale_fn=tfb.real_nvp_default_template( + hidden_layers=[512, 512]))) + + x = nvp.sample() + nvp.log_prob(x) + nvp.log_prob(0.) + ``` + + For more examples, see [4]. + + [1]: "Density Estimation using Real NVP." + Laurent Dinh, Jascha Sohl-Dickstein, Samy Bengio. ICLR. 2017. + https://arxiv.org/abs/1605.08803 + + [2]: "Masked Autoregressive Flow for Density Estimation." + George Papamakarios, Theo Pavlakou, Iain Murray. Arxiv. 2017. + https://arxiv.org/abs/1705.07057 + + [3]: "NICE: Non-linear Independent Components Estimation." + Laurent Dinh, David Krueger, Yoshua Bengio. ICLR. 2015. + https://arxiv.org/abs/1410.8516 + + [4]: "Normalizing Flows Tutorial, Part 2: Modern Normalizing Flows." + Eric Jang. Blog post. January 2018. + http://blog.evjang.com/2018/01/nf2.html + """ + + def __init__(self, + num_masked, + shift_and_log_scale_fn, + is_constant_jacobian=False, + validate_args=False, + name=None): + """Creates the Real NVP or NICE bijector. + + Args: + num_masked: Python `int` indicating that the first `d` units of the event + should be masked. Must be in the closed interval `[1, D-1]`, where `D` + is the event size of the base distribution. + shift_and_log_scale_fn: Python `callable` which computes `shift` and + `log_scale` from both the forward domain (`x`) and the inverse domain + (`y`). Calculation must respect the "autoregressive property" (see class + docstring). Suggested default + `masked_autoregressive_default_template(hidden_layers=...)`. + Typically the function contains `tf.Variables` and is wrapped using + `tf.make_template`. Returning `None` for either (both) `shift`, + `log_scale` is equivalent to (but more efficient than) returning zero. + is_constant_jacobian: Python `bool`. Default: `False`. When `True` the + implementation assumes `log_scale` does not depend on the forward domain + (`x`) or inverse domain (`y`) values. (No validation is made; + `is_constant_jacobian=False` is always safe but possibly computationally + inefficient.) + validate_args: Python `bool` indicating whether arguments should be + checked for correctness. + name: Python `str`, name given to ops managed by this object. + + Raises: + ValueError: If num_masked < 1. + """ + name = name or "real_nvp" + if num_masked <= 0: + raise ValueError("num_masked must be a positive integer.") + self._num_masked = num_masked + # At construction time, we don't know input_depth. + self._input_depth = None + self._shift_and_log_scale_fn = shift_and_log_scale_fn + super(RealNVP, self).__init__( + event_ndims=1, + is_constant_jacobian=is_constant_jacobian, + validate_args=validate_args, + name=name) + + def _cache_input_depth(self, x): + if self._input_depth is None: + self._input_depth = x.shape.with_rank_at_least(1)[-1].value + if self._input_depth is None: + raise NotImplementedError( + "Rightmost dimension must be known prior to graph execution.") + if self._num_masked >= self._input_depth: + raise ValueError( + "Number of masked units must be smaller than the event size.") + + def _forward(self, x): + self._cache_input_depth(x) + # Performs scale and shift. + x0, x1 = x[:, :self._num_masked], x[:, self._num_masked:] + shift, log_scale = self._shift_and_log_scale_fn( + x0, self._input_depth - self._num_masked) + y1 = x1 + if log_scale is not None: + y1 *= math_ops.exp(log_scale) + if shift is not None: + y1 += shift + y = array_ops.concat([x0, y1], axis=-1) + return y + + def _inverse(self, y): + self._cache_input_depth(y) + # Performs un-shift and un-scale. + y0, y1 = y[:, :self._num_masked], y[:, self._num_masked:] + shift, log_scale = self._shift_and_log_scale_fn( + y0, self._input_depth - self._num_masked) + x1 = y1 + if shift is not None: + x1 -= shift + if log_scale is not None: + x1 *= math_ops.exp(-log_scale) + x = array_ops.concat([y0, x1], axis=-1) + return x + + def _inverse_log_det_jacobian(self, y): + self._cache_input_depth(y) + y0 = y[:, :self._num_masked] + _, log_scale = self._shift_and_log_scale_fn( + y0, self._input_depth - self._num_masked) + if log_scale is None: + return constant_op.constant(0., dtype=y.dtype, name="ildj") + return -math_ops.reduce_sum(log_scale, axis=-1) + + def _forward_log_det_jacobian(self, x): + self._cache_input_depth(x) + x0 = x[:, :self._num_masked] + _, log_scale = self._shift_and_log_scale_fn( + x0, self._input_depth - self._num_masked) + if log_scale is None: + return constant_op.constant(0., dtype=x.dtype, name="ildj") + return math_ops.reduce_sum(log_scale, axis=-1) + + +def real_nvp_default_template( + hidden_layers, + shift_only=False, + activation=nn_ops.relu, + name=None, + *args, + **kwargs): + """Build a scale-and-shift function using a multi-layer neural network. + + This will be wrapped in a make_template to ensure the variables are only + created once. It takes the `d`-dimensional input x[0:d] and returns the `D-d` + dimensional outputs `loc` ("mu") and `log_scale` ("alpha"). + + Arguments: + hidden_layers: Python `list`-like of non-negative integer, scalars + indicating the number of units in each hidden layer. Default: `[512, 512]. + shift_only: Python `bool` indicating if only the `shift` term shall be + computed (i.e. NICE bijector). Default: `False`. + activation: Activation function (callable). Explicitly setting to `None` + implies a linear activation. + name: A name for ops managed by this function. Default: + "real_nvp_default_template". + *args: `tf.layers.dense` arguments. + **kwargs: `tf.layers.dense` keyword arguments. + + Returns: + shift: `Float`-like `Tensor` of shift terms (the "mu" in [2]). + log_scale: `Float`-like `Tensor` of log(scale) terms (the "alpha" in [2]). + + Raises: + NotImplementedError: if rightmost dimension of `inputs` is unknown prior to + graph execution. + """ + + with ops.name_scope(name, "real_nvp_default_template"): + def _fn(x, output_units): + """Fully connected MLP parameterized via `real_nvp_template`.""" + for units in hidden_layers: + x = layers.dense( + inputs=x, + units=units, + activation=activation, + *args, + **kwargs) + x = layers.dense( + inputs=x, + units=(1 if shift_only else 2) * output_units, + activation=None, + *args, + **kwargs) + if shift_only: + return x, None + shift, log_scale = array_ops.split(x, 2, axis=-1) + return shift, log_scale + return template_ops.make_template( + "real_nvp_default_template", _fn) diff --git a/tensorflow/contrib/distributions/python/ops/distribution_util.py b/tensorflow/contrib/distributions/python/ops/distribution_util.py index a4d249d41ec9733721a3583d3708e0da56db1733..289e1d50e1146a641c0cc433ece3465aed73b1c2 100644 --- a/tensorflow/contrib/distributions/python/ops/distribution_util.py +++ b/tensorflow/contrib/distributions/python/ops/distribution_util.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib import linalg +from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_ops @@ -442,6 +443,44 @@ def maybe_check_scalar_distribution( return assertions +def pad_mixture_dimensions(x, mixture_distribution, categorical_distribution, + event_ndims): + """Pad dimensions of event tensors for mixture distributions. + + See `Mixture._sample_n` and `MixtureSameFamily._sample_n` for usage examples. + + Args: + x: event tensor to pad. + mixture_distribution: Base distribution of the mixture. + categorical_distribution: `Categorical` distribution that mixes the base + distribution. + event_ndims: Integer specifying the number of event dimensions in the event + tensor. + + Returns: + A padded version of `x` that can broadcast with `categorical_distribution`. + """ + with ops.name_scope("pad_mix_dims", values=[x]): + def _get_ndims(d): + if d.batch_shape.ndims is not None: + return d.batch_shape.ndims + return array_ops.shape(d.batch_shape_tensor())[0] + dist_batch_ndims = _get_ndims(mixture_distribution) + cat_batch_ndims = _get_ndims(categorical_distribution) + pad_ndims = array_ops.where( + categorical_distribution.is_scalar_batch(), + dist_batch_ndims, + dist_batch_ndims - cat_batch_ndims) + s = array_ops.shape(x) + x = array_ops.reshape(x, shape=array_ops.concat([ + s[:-1], + array_ops.ones([pad_ndims], dtype=dtypes.int32), + s[-1:], + array_ops.ones([event_ndims], dtype=dtypes.int32), + ], axis=0)) + return x + + def static_value(x): """Returns the static value of a `Tensor` or `None`.""" return tensor_util.constant_value(ops.convert_to_tensor(x)) diff --git a/tensorflow/contrib/distributions/python/ops/kumaraswamy.py b/tensorflow/contrib/distributions/python/ops/kumaraswamy.py new file mode 100644 index 0000000000000000000000000000000000000000..74d5d8773cf3e69a52554c87d656fea2835c8354 --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/kumaraswamy.py @@ -0,0 +1,258 @@ +# Copyright 2016 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. +# ============================================================================== +"""The Kumaraswamy distribution class.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import check_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import random_ops +from tensorflow.python.ops import special_math_ops +from tensorflow.python.ops.distributions import beta +from tensorflow.python.ops.distributions import distribution +from tensorflow.python.ops.distributions import util as distribution_util +from tensorflow.python.util.tf_export import tf_export + +__all__ = [ + "Kumaraswamy", +] + +_kumaraswamy_sample_note = """Note: `x` must have dtype `self.dtype` and be in +`[0, 1].` It must have a shape compatible with `self.batch_shape()`.""" + + +def _harmonic_number(x): + """Compute the harmonic number from its analytic continuation. + + Derivation from [1] and Euler's constant [2]. + [1] - + https://en.wikipedia.org/wiki/Digamma_function#Relation_to_harmonic_numbers + [2] - https://en.wikipedia.org/wiki/Euler%E2%80%93Mascheroni_constant + + + Args: + x: input float. + + Returns: + z: The analytic continuation of the harmonic number for the input. + + """ + one = array_ops.ones([], dtype=x.dtype) + return math_ops.digamma(x + one) - math_ops.digamma(one) + + +@tf_export("distributions.Kumaraswamy") +class Kumaraswamy(beta.Beta): + """Kumaraswamy distribution. + + The Kumaraswamy distribution is defined over the `(0, 1)` interval using + parameters + `concentration1` (aka "alpha") and `concentration0` (aka "beta"). It has a + shape similar to the Beta distribution, but is reparameterizeable. + + #### Mathematical Details + + The probability density function (pdf) is, + + ```none + pdf(x; alpha, beta) = alpha * beta * x**(alpha - 1) * (1 - x**alpha)**(beta - + 1) + ``` + + where: + + * `concentration1 = alpha`, + * `concentration0 = beta`, + + Distribution parameters are automatically broadcast in all functions; see + examples for details. + + #### Examples + + ```python + # Create a batch of three Kumaraswamy distributions. + alpha = [1, 2, 3] + beta = [1, 2, 3] + dist = Kumaraswamy(alpha, beta) + + dist.sample([4, 5]) # Shape [4, 5, 3] + + # `x` has three batch entries, each with two samples. + x = [[.1, .4, .5], + [.2, .3, .5]] + # Calculate the probability of each pair of samples under the corresponding + # distribution in `dist`. + dist.prob(x) # Shape [2, 3] + ``` + + ```python + # Create batch_shape=[2, 3] via parameter broadcast: + alpha = [[1.], [2]] # Shape [2, 1] + beta = [3., 4, 5] # Shape [3] + dist = Kumaraswamy(alpha, beta) + + # alpha broadcast as: [[1., 1, 1,], + # [2, 2, 2]] + # beta broadcast as: [[3., 4, 5], + # [3, 4, 5]] + # batch_Shape [2, 3] + dist.sample([4, 5]) # Shape [4, 5, 2, 3] + + x = [.2, .3, .5] + # x will be broadcast as [[.2, .3, .5], + # [.2, .3, .5]], + # thus matching batch_shape [2, 3]. + dist.prob(x) # Shape [2, 3] + ``` + + """ + + def __init__(self, + concentration1=None, + concentration0=None, + validate_args=False, + allow_nan_stats=True, + name="Kumaraswamy"): + """Initialize a batch of Kumaraswamy distributions. + + Args: + concentration1: Positive floating-point `Tensor` indicating mean + number of successes; aka "alpha". Implies `self.dtype` and + `self.batch_shape`, i.e., + `concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`. + concentration0: Positive floating-point `Tensor` indicating mean + number of failures; aka "beta". Otherwise has same semantics as + `concentration1`. + validate_args: Python `bool`, default `False`. When `True` distribution + parameters are checked for validity despite possibly degrading runtime + performance. When `False` invalid inputs may silently render incorrect + outputs. + allow_nan_stats: Python `bool`, default `True`. When `True`, statistics + (e.g., mean, mode, variance) use the value "`NaN`" to indicate the + result is undefined. When `False`, an exception is raised if one or + more of the statistic's batch members are undefined. + name: Python `str` name prefixed to Ops created by this class. + """ + super(Kumaraswamy, self).__init__( + concentration1=concentration1, + concentration0=concentration0, + validate_args=validate_args, + allow_nan_stats=allow_nan_stats, + name=name) + self._reparameterization_type = distribution.FULLY_REPARAMETERIZED + + def _sample_n(self, n, seed=None): + expanded_concentration1 = array_ops.ones_like( + self.total_concentration, dtype=self.dtype) * self.concentration1 + expanded_concentration0 = array_ops.ones_like( + self.total_concentration, dtype=self.dtype) * self.concentration0 + shape = array_ops.concat([[n], self.batch_shape_tensor()], 0) + uniform_sample = random_ops.random_uniform( + shape=shape, minval=0.0, maxval=1.0, dtype=self.dtype, seed=seed) + + kumaraswamy_sample = (1 - uniform_sample**(1. / expanded_concentration0))**( + 1. / expanded_concentration1) + return kumaraswamy_sample + + @distribution_util.AppendDocstring(_kumaraswamy_sample_note) + def _log_cdf(self, x): + a = self.concentration1 + b = self.concentration0 + return math_ops.log1p(-(1 - x**a)**b) + + @distribution_util.AppendDocstring(_kumaraswamy_sample_note) + def _cdf(self, x): + a = self.concentration1 + b = self.concentration0 + return 1 - (1 - x**a)**b + + def _survival_function(self, x): + a = self.concentration1 + b = self.concentration0 + return (1 - x**a)**b + + def _log_survival_function(self, x): + a = self.concentration1 + b = self.concentration0 + return b * math_ops.log1p(-x**a) + + def _log_unnormalized_prob(self, x): + x = self._maybe_assert_valid_sample(x) + a = self.concentration1 + b = self.concentration0 + return (a - 1) * math_ops.log(x) + (b - 1) * math_ops.log1p(-x**a) + + def _log_normalization(self): + a = self.concentration1 + b = self.concentration0 + return -(math_ops.log(a) + math_ops.log(b)) + + def _entropy(self): + a = self.concentration1 + b = self.concentration0 + return (1 - 1. / a) + ( + 1 - 1. / b) * _harmonic_number(b) + math_ops.log(a) + math_ops.log(b) + + def _moment(self, n): + """Compute the n'th (uncentered) moment.""" + expanded_concentration1 = array_ops.ones_like( + self.total_concentration, dtype=self.dtype) * self.concentration1 + expanded_concentration0 = array_ops.ones_like( + self.total_concentration, dtype=self.dtype) * self.concentration0 + beta_arg0 = 1 + n / expanded_concentration1 + beta_arg = array_ops.stack([beta_arg0, expanded_concentration0], -1) + log_moment = math_ops.log(expanded_concentration0) + special_math_ops.lbeta( + beta_arg) + return math_ops.exp(log_moment) + + def _mean(self): + return self._moment(1) + + def _variance(self): + # TODO(b/72696533): Investigate a more numerically stable version. + return self._moment(2) - math_ops.square(self._moment(1)) + + @distribution_util.AppendDocstring( + """Note: The mode is undefined when `concentration1 <= 1` or + `concentration0 <= 1`. If `self.allow_nan_stats` is `True`, `NaN` + is used for undefined modes. If `self.allow_nan_stats` is `False` an + exception is raised when one or more modes are undefined.""") + def _mode(self): + a = self.concentration1 + b = self.concentration0 + mode = ((a - 1) / (a * b - 1))**(1. / a) + if self.allow_nan_stats: + nan = array_ops.fill( + self.batch_shape_tensor(), + np.array(np.nan, dtype=self.dtype.as_numpy_dtype), + name="nan") + is_defined = (self.concentration1 > 1.) & (self.concentration0 > 1.) + return array_ops.where(is_defined, mode, nan) + return control_flow_ops.with_dependencies([ + check_ops.assert_less( + array_ops.ones([], dtype=self.dtype), + self.concentration1, + message="Mode undefined for concentration1 <= 1."), + check_ops.assert_less( + array_ops.ones([], dtype=self.dtype), + self.concentration0, + message="Mode undefined for concentration0 <= 1.") + ], mode) diff --git a/tensorflow/contrib/distributions/python/ops/mixture.py b/tensorflow/contrib/distributions/python/ops/mixture.py index f2d492f5489a197157558ae727416b51db04793e..cef6a143fc615901315a3780bf4ed53b8c7cd177 100644 --- a/tensorflow/contrib/distributions/python/ops/mixture.py +++ b/tensorflow/contrib/distributions/python/ops/mixture.py @@ -71,6 +71,7 @@ class Mixture(distribution.Distribution): components, validate_args=False, allow_nan_stats=True, + use_static_graph=False, name="Mixture"): """Initialize a Mixture distribution. @@ -96,6 +97,11 @@ class Mixture(distribution.Distribution): exception if a statistic (e.g. mean/mode/etc...) is undefined for any batch member. If `True`, batch members with valid parameters leading to undefined statistics will return NaN for this statistic. + use_static_graph: Calls to `sample` will not rely on dynamic tensor + indexing, allowing for some static graph compilation optimizations, but + at the expense of sampling all underlying distributions in the mixture. + (Possibly useful when running on TPUs). + Default value: `False` (i.e., use dynamic indexing). name: A name for this distribution (optional). Raises: @@ -178,6 +184,10 @@ class Mixture(distribution.Distribution): self._static_event_shape = static_event_shape self._static_batch_shape = static_batch_shape + self._use_static_graph = use_static_graph + if use_static_graph and static_num_components is None: + raise ValueError("Number of categories must be known statically when " + "`static_sample=True`.") # We let the Mixture distribution access _graph_parents since its arguably # more like a baseclass. graph_parents = self._cat._graph_parents # pylint: disable=protected-access @@ -292,6 +302,31 @@ class Mixture(distribution.Distribution): return mixture_log_cdf def _sample_n(self, n, seed=None): + if self._use_static_graph: + # This sampling approach is almost the same as the approach used by + # `MixtureSameFamily`. The differences are due to having a list of + # `Distribution` objects rather than a single object, and maintaining + # random seed management that is consistent with the non-static code path. + samples = [] + cat_samples = self.cat.sample(n, seed=seed) + for c in range(self.num_components): + seed = distribution_util.gen_new_seed(seed, "mixture") + samples.append(self.components[c].sample(n, seed=seed)) + x = array_ops.stack( + samples, -self._static_event_shape.ndims - 1) # [n, B, k, E] + npdt = x.dtype.as_numpy_dtype + mask = array_ops.one_hot( + indices=cat_samples, # [n, B] + depth=self._num_components, # == k + on_value=np.ones([], dtype=npdt), + off_value=np.zeros([], dtype=npdt)) # [n, B, k] + mask = distribution_utils.pad_mixture_dimensions( + mask, self, self._cat, + self._static_event_shape.ndims) # [n, B, k, [1]*e] + return math_ops.reduce_sum( + x * mask, + axis=-1 - self._static_event_shape.ndims) # [n, B, E] + with ops.control_dependencies(self._assertions): n = ops.convert_to_tensor(n, name="n") static_n = tensor_util.constant_value(n) diff --git a/tensorflow/contrib/distributions/python/ops/mixture_same_family.py b/tensorflow/contrib/distributions/python/ops/mixture_same_family.py index 0ca236c3761f9d3a0fcc79ff9db792319108db0d..b93bdc5ab4010663baddda1410b302644853648b 100644 --- a/tensorflow/contrib/distributions/python/ops/mixture_same_family.py +++ b/tensorflow/contrib/distributions/python/ops/mixture_same_family.py @@ -20,7 +20,7 @@ from __future__ import print_function import numpy as np -from tensorflow.python.framework import dtypes +from tensorflow.contrib.distributions.python.ops import distribution_util as distribution_utils from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -239,7 +239,9 @@ class MixtureSameFamily(distribution.Distribution): depth=self._num_components, # == k on_value=np.ones([], dtype=npdt), off_value=np.zeros([], dtype=npdt)) # [n, B, k] - mask = self._pad_mix_dims(mask) # [n, B, k, [1]*e] + mask = distribution_utils.pad_mixture_dimensions( + mask, self, self.mixture_distribution, + self._event_shape().ndims) # [n, B, k, [1]*e] return math_ops.reduce_sum( x * mask, axis=-1 - self._event_ndims) # [n, B, E] @@ -248,14 +250,15 @@ class MixtureSameFamily(distribution.Distribution): x = self._pad_sample_dims(x) log_prob_x = self.components_distribution.log_prob(x) # [S, B, k] log_mix_prob = nn_ops.log_softmax( - self.mixture_distribution.logits, dim=-1) # [B, k] + self.mixture_distribution.logits, axis=-1) # [B, k] return math_ops.reduce_logsumexp( log_prob_x + log_mix_prob, axis=-1) # [S, B] def _mean(self): with ops.control_dependencies(self._runtime_assertions): - probs = self._pad_mix_dims( - self.mixture_distribution.probs) # [B, k, [1]*e] + probs = distribution_utils.pad_mixture_dimensions( + self.mixture_distribution.probs, self, self.mixture_distribution, + self._event_shape().ndims) # [B, k, [1]*e] return math_ops.reduce_sum( probs * self.components_distribution.mean(), axis=-1 - self._event_ndims) # [B, E] @@ -264,15 +267,16 @@ class MixtureSameFamily(distribution.Distribution): x = self._pad_sample_dims(x) log_cdf_x = self.components_distribution.log_cdf(x) # [S, B, k] log_mix_prob = nn_ops.log_softmax( - self.mixture_distribution.logits, dim=-1) # [B, k] + self.mixture_distribution.logits, axis=-1) # [B, k] return math_ops.reduce_logsumexp( log_cdf_x + log_mix_prob, axis=-1) # [S, B] def _variance(self): with ops.control_dependencies(self._runtime_assertions): # Law of total variance: Var(Y) = E[Var(Y|X)] + Var(E[Y|X]) - probs = self._pad_mix_dims( - self.mixture_distribution.probs) # [B, k, [1]*e] + probs = distribution_utils.pad_mixture_dimensions( + self.mixture_distribution.probs, self, self.mixture_distribution, + self._event_shape().ndims) # [B, k, [1]*e] mean_cond_var = math_ops.reduce_sum( probs * self.components_distribution.variance(), axis=-1 - self._event_ndims) # [B, E] @@ -291,8 +295,12 @@ class MixtureSameFamily(distribution.Distribution): with ops.control_dependencies(self._runtime_assertions): # Law of total variance: Var(Y) = E[Var(Y|X)] + Var(E[Y|X]) - probs = self._pad_mix_dims(self._pad_mix_dims( - self.mixture_distribution.probs)) # [B, k, 1, 1] + probs = distribution_utils.pad_mixture_dimensions( + distribution_utils.pad_mixture_dimensions( + self.mixture_distribution.probs, self, self.mixture_distribution, + self._event_shape().ndims), + self, self.mixture_distribution, + self._event_shape().ndims) # [B, k, 1, 1] mean_cond_var = math_ops.reduce_sum( probs * self.components_distribution.covariance(), axis=-3) # [B, e, e] @@ -312,27 +320,6 @@ class MixtureSameFamily(distribution.Distribution): shape[:d], [1], shape[d:]], axis=0)) return x - def _pad_mix_dims(self, x): - with ops.name_scope("pad_mix_dims", values=[x]): - def _get_ndims(d): - if d.batch_shape.ndims is not None: - return d.batch_shape.ndims - return array_ops.shape(d.batch_shape_tensor())[0] - dist_batch_ndims = _get_ndims(self) - cat_batch_ndims = _get_ndims(self.mixture_distribution) - pad_ndims = array_ops.where( - self.mixture_distribution.is_scalar_batch(), - dist_batch_ndims, - dist_batch_ndims - cat_batch_ndims) - s = array_ops.shape(x) - x = array_ops.reshape(x, shape=array_ops.concat([ - s[:-1], - array_ops.ones([pad_ndims], dtype=dtypes.int32), - s[-1:], - array_ops.ones([self._event_ndims], dtype=dtypes.int32), - ], axis=0)) - return x - def _outer_squared_difference(x, y): """Convenience function analogous to tf.squared_difference.""" diff --git a/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py b/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py index b6becfa9fc93f189a1a7bf7b2a7af8dc1f2e9720..2aa771a71efe52c8d86d459f090ea8ee137c4487 100644 --- a/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py +++ b/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py @@ -278,7 +278,7 @@ class ExpRelaxedOneHotCategorical(distribution.Distribution): * math_ops.log(self.temperature)) # compute the unnormalized density log_softmax = nn_ops.log_softmax(logits_2d - x_2d * self._temperature_2d) - log_unnorm_prob = math_ops.reduce_sum(log_softmax, [-1], keep_dims=False) + log_unnorm_prob = math_ops.reduce_sum(log_softmax, [-1], keepdims=False) # combine unnormalized density with normalization constant log_prob = log_norm_const + log_unnorm_prob # Reshapes log_prob to be consistent with shape of user-supplied logits diff --git a/tensorflow/contrib/distributions/python/ops/test_util.py b/tensorflow/contrib/distributions/python/ops/test_util.py index bfc727450f5e48ecbf98bf8ab0475ec67c9e7137..15b0820cbdf560e04a304c40a47e541006523b6d 100644 --- a/tensorflow/contrib/distributions/python/ops/test_util.py +++ b/tensorflow/contrib/distributions/python/ops/test_util.py @@ -327,7 +327,7 @@ class VectorDistributionTestHelpers(object): num_samples=int(1e5), seed=24, rtol=1e-2, - atol=0., + atol=0.1, cov_rtol=None, cov_atol=None): """Tests that sample/mean/covariance are consistent with each other. diff --git a/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py b/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py index 7ce8a83fd91e2dfaa0ccef633f803b3ae595e646..0c747f8e68529484ae6f695b8500cde74857bb11 100644 --- a/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py +++ b/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py @@ -50,20 +50,25 @@ __all__ = [ def quadrature_scheme_softmaxnormal_gauss_hermite( - loc, scale, quadrature_size, + normal_loc, normal_scale, quadrature_size, validate_args=False, name=None): """Use Gauss-Hermite quadrature to form quadrature on `K - 1` simplex. + A `SoftmaxNormal` random variable `Y` may be generated via + + ``` + Y = SoftmaxCentered(X), + X = Normal(normal_loc, normal_scale) + ``` + Note: for a given `quadrature_size`, this method is generally less accurate than `quadrature_scheme_softmaxnormal_quantiles`. Args: - loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0. - Represents the `location` parameter of the SoftmaxNormal used for - selecting one of the `K` affine transformations. - scale: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0. - Represents the `scale` parameter of the SoftmaxNormal used for - selecting one of the `K` affine transformations. + normal_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0. + The location parameter of the Normal used to construct the SoftmaxNormal. + normal_scale: `float`-like `Tensor`. Broadcastable with `normal_loc`. + The scale parameter of the Normal used to construct the SoftmaxNormal. quadrature_size: Python `int` scalar representing the number of quadrature points. validate_args: Python `bool`, default `False`. When `True` distribution @@ -80,24 +85,25 @@ def quadrature_scheme_softmaxnormal_gauss_hermite( associated with each grid point. """ with ops.name_scope(name, "quadrature_scheme_softmaxnormal_gauss_hermite", - [loc, scale]): - loc = ops.convert_to_tensor(loc, name="loc") - dt = loc.dtype.base_dtype - scale = ops.convert_to_tensor(scale, dtype=dt, name="scale") + [normal_loc, normal_scale]): + normal_loc = ops.convert_to_tensor(normal_loc, name="normal_loc") + dt = normal_loc.dtype.base_dtype + normal_scale = ops.convert_to_tensor( + normal_scale, dtype=dt, name="normal_scale") - loc = maybe_check_quadrature_param(loc, "loc", validate_args) - scale = maybe_check_quadrature_param(scale, "scale", validate_args) + normal_scale = maybe_check_quadrature_param( + normal_scale, "normal_scale", validate_args) grid, probs = np.polynomial.hermite.hermgauss(deg=quadrature_size) - grid = grid.astype(loc.dtype.as_numpy_dtype) - probs = probs.astype(loc.dtype.as_numpy_dtype) + grid = grid.astype(dt.dtype.as_numpy_dtype) + probs = probs.astype(dt.dtype.as_numpy_dtype) probs /= np.linalg.norm(probs, ord=1, keepdims=True) - probs = ops.convert_to_tensor(probs, name="probs", dtype=loc.dtype) + probs = ops.convert_to_tensor(probs, name="probs", dtype=dt) grid = softmax( -distribution_util.pad( - (loc[..., array_ops.newaxis] + - np.sqrt(2.) * scale[..., array_ops.newaxis] * grid), + (normal_loc[..., array_ops.newaxis] + + np.sqrt(2.) * normal_scale[..., array_ops.newaxis] * grid), axis=-2, front=True), axis=-2) # shape: [B, components, deg] @@ -106,18 +112,23 @@ def quadrature_scheme_softmaxnormal_gauss_hermite( def quadrature_scheme_softmaxnormal_quantiles( - loc, scale, quadrature_size, + normal_loc, normal_scale, quadrature_size, validate_args=False, name=None): """Use SoftmaxNormal quantiles to form quadrature on `K - 1` simplex. + A `SoftmaxNormal` random variable `Y` may be generated via + + ``` + Y = SoftmaxCentered(X), + X = Normal(normal_loc, normal_scale) + ``` + Args: - loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0. - Represents the `location` parameter of the SoftmaxNormal used for - selecting one of the `K` affine transformations. - scale: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0. - Represents the `scale` parameter of the SoftmaxNormal used for - selecting one of the `K` affine transformations. - quadrature_size: Python scalar `int` representing the number of quadrature + normal_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0. + The location parameter of the Normal used to construct the SoftmaxNormal. + normal_scale: `float`-like `Tensor`. Broadcastable with `normal_loc`. + The scale parameter of the Normal used to construct the SoftmaxNormal. + quadrature_size: Python `int` scalar representing the number of quadrature points. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime @@ -132,15 +143,17 @@ def quadrature_scheme_softmaxnormal_quantiles( probs: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the associated with each grid point. """ - with ops.name_scope(name, "softmax_normal_grid_and_probs", [loc, scale]): - loc = ops.convert_to_tensor(loc, name="loc") - dt = loc.dtype.base_dtype - scale = ops.convert_to_tensor(scale, dtype=dt, name="scale") + with ops.name_scope(name, "softmax_normal_grid_and_probs", + [normal_loc, normal_scale]): + normal_loc = ops.convert_to_tensor(normal_loc, name="normal_loc") + dt = normal_loc.dtype.base_dtype + normal_scale = ops.convert_to_tensor( + normal_scale, dtype=dt, name="normal_scale") - loc = maybe_check_quadrature_param(loc, "loc", validate_args) - scale = maybe_check_quadrature_param(scale, "scale", validate_args) + normal_scale = maybe_check_quadrature_param( + normal_scale, "normal_scale", validate_args) - dist = normal_lib.Normal(loc=loc, scale=scale) + dist = normal_lib.Normal(loc=normal_loc, scale=normal_scale) def _get_batch_ndims(): """Helper to get dist.batch_shape.ndims, statically if possible.""" @@ -195,114 +208,51 @@ def quadrature_scheme_softmaxnormal_quantiles( class VectorDiffeomixture(distribution_lib.Distribution): """VectorDiffeomixture distribution. - The VectorDiffeomixture is an approximation to a [compound distribution]( - https://en.wikipedia.org/wiki/Compound_probability_distribution), i.e., + A vector diffeomixture (VDM) is a distribution parameterized by a convex + combination of `K` component `loc` vectors, `loc[k], k = 0,...,K-1`, and `K` + `scale` matrices `scale[k], k = 0,..., K-1`. It approximates the following + [compound distribution] + (https://en.wikipedia.org/wiki/Compound_probability_distribution) ```none - p(x) = int_{X} q(x | v) p(v) dv - = lim_{Q->infty} sum{ prob[i] q(x | loc=sum_k^K lambda[k;i] loc[k], - scale=sum_k^K lambda[k;i] scale[k]) - : i=0, ..., Q-1 } + p(x) = int p(x | z) p(z) dz, + where z is in the K-simplex, and + p(x | z) := p(x | loc=sum_k z[k] loc[k], scale=sum_k z[k] scale[k]) ``` - where `q(x | v)` is a vector version of the `distribution` argument and `p(v)` - is a SoftmaxNormal parameterized by `mix_loc` and `mix_scale`. The - vector-ization of `distribution` entails an affine transformation of iid - samples from `distribution`. The `prob` term is from quadrature and - `lambda[k] = sigmoid(mix_loc[k] + sqrt(2) mix_scale[k] grid[k])` where the - `grid` points correspond to the `prob`s. - - In the non-approximation case, a draw from the mixture distribution (the - "prior") represents the convex weights for different affine transformations. - I.e., draw a mixing vector `v` (from the `K-1`-simplex) and let the final - sample be: `y = (sum_k^K v[k] scale[k]) @ x + (sum_k^K v[k] loc[k])` where `@` - denotes matrix multiplication. However, the non-approximate distribution does - not have an analytical probability density function (pdf). Therefore the - `VectorDiffeomixture` class implements an approximation based on - [numerical quadrature]( - https://en.wikipedia.org/wiki/Numerical_integration) (default: - [Gauss--Hermite quadrature]( - https://en.wikipedia.org/wiki/Gauss%E2%80%93Hermite_quadrature)). I.e., in - Note: although the `VectorDiffeomixture` is approximately the - `SoftmaxNormal-Distribution` compound distribution, it is itself a valid - distribution. It possesses a `sample`, `log_prob`, `mean`, `covariance` which - are all mutually consistent. - - #### Intended Use - - This distribution is noteworthy because it implements a mixture of - `Vector`-ized distributions yet has samples differentiable in the - distribution's parameters (aka "reparameterized"). It has an analytical - density function with `O(dKQ)` complexity. `d` is the vector dimensionality, - `K` is the number of components, and `Q` is the number of quadrature points. - These properties make it well-suited for Bayesian Variational Inference, i.e., - as a surrogate family for the posterior. - - For large values of `mix_scale`, the `VectorDistribution` behaves increasingly - like a discrete mixture. (In most cases this limit is only achievable by also - increasing the quadrature polynomial degree, `Q`.) - - The term `Vector` is consistent with similar named Tensorflow `Distribution`s. - For more details, see the "About `Vector` distributions in Tensorflow." - section. - - The term `Diffeomixture` is a portmanteau of - [diffeomorphism](https://en.wikipedia.org/wiki/Diffeomorphism) and [compound - mixture](https://en.wikipedia.org/wiki/Compound_probability_distribution). For - more details, see the "About `Diffeomixture`s and reparametrization.`" - section. - - #### Mathematical Details - - The `VectorDiffeomixture` approximates a SoftmaxNormal-mixed ("prior") - [compound distribution]( - https://en.wikipedia.org/wiki/Compound_probability_distribution). - Using variable-substitution and [numerical quadrature]( - https://en.wikipedia.org/wiki/Numerical_integration) (default: - [Gauss--Hermite quadrature]( - https://en.wikipedia.org/wiki/Gauss%E2%80%93Hermite_quadrature)) we can - redefine the distribution to be a parameter-less convex combination of `K` - different affine combinations of a `d` iid samples from `distribution`. - - That is, defined over `R**d` this distribution is parameterized by a - (batch of) length-`K` `mix_loc` and `mix_scale` vectors, a length-`K` list of - (a batch of) length-`d` `loc` vectors, and a length-`K` list of `scale` - `LinearOperator`s each operating on a (batch of) length-`d` vector space. - Finally, a `distribution` parameter specifies the underlying base distribution - which is "lifted" to become multivariate ("lifting" is the same concept as in - `TransformedDistribution`). - - The probability density function (pdf) is, + The integral `int p(x | z) p(z) dz` is approximated with a quadrature scheme + adapted to the mixture density `p(z)`. The `N` quadrature points `z_{N, n}` + and weights `w_{N, n}` (which are non-negative and sum to 1) are chosen + such that - ```none - pdf(y; mix_loc, mix_scale, loc, scale, phi) - = sum{ prob[i] phi(f_inverse(x; i)) / abs(det(interp_scale[i])) - : i=0, ..., Q-1 } - ``` + ```q_N(x) := sum_{n=1}^N w_{n, N} p(x | z_{N, n}) --> p(x)``` - where, `phi` is the base distribution pdf, and, + as `N --> infinity`. - ```none - f_inverse(x; i) = inv(interp_scale[i]) @ (x - interp_loc[i]) - interp_loc[i] = sum{ lambda[k; i] loc[k] : k=0, ..., K-1 } - interp_scale[i] = sum{ lambda[k; i] scale[k] : k=0, ..., K-1 } - ``` + Since `q_N(x)` is in fact a mixture (of `N` points), we may sample from + `q_N` exactly. It is important to note that the VDM is *defined* as `q_N` + above, and *not* `p(x)`. Therefore, sampling and pdf may be implemented as + exact (up to floating point error) methods. - and, + A common choice for the conditional `p(x | z)` is a multivariate Normal. - ```none - grid, weight = np.polynomial.hermite.hermgauss(quadrature_size) - prob[k] = weight[k] / sqrt(pi) - lambda[k; i] = sigmoid(mix_loc[k] + sqrt(2) mix_scale[k] grid[i]) + The implemented marginal `p(z)` is the `SoftmaxNormal`, which is a + `K-1` dimensional Normal transformed by a `SoftmaxCentered` bijector, making + it a density on the `K`-simplex. That is, + + ``` + Z = SoftmaxCentered(X), + X = Normal(mix_loc / temperature, 1 / temperature) ``` - The distribution corresponding to `phi` must be a scalar-batch, scalar-event - distribution. Typically it is reparameterized. If not, it must be a function - of non-trainable parameters. + The default quadrature scheme chooses `z_{N, n}` as `N` midpoints of + the quantiles of `p(z)` (generalized quantiles if `K > 2`). - WARNING: If you backprop through a VectorDiffeomixture sample and the "base" - distribution is both: not `FULLY_REPARAMETERIZED` and a function of trainable - variables, then the gradient is not guaranteed correct! + See [1] for more details. + + [1]. "Quadrature Compound: An approximating family of distributions" + Joshua Dillon, Ian Langmore, arXiv preprints + https://arxiv.org/abs/1801.03080 #### About `Vector` distributions in TensorFlow. @@ -310,12 +260,11 @@ class VectorDiffeomixture(distribution_lib.Distribution): particularly useful in [variational Bayesian methods](https://en.wikipedia.org/wiki/Variational_Bayesian_methods). - Conditioned on a draw from the SoftmaxNormal, `Y|v` is a vector whose + Conditioned on a draw from the SoftmaxNormal, `X|z` is a vector whose components are linear combinations of affine transformations, thus is itself - an affine transformation. Therefore `Y|v` lives in the vector space generated - by vectors of affine-transformed distributions. + an affine transformation. - Note: The marginals `Y_1|v, ..., Y_d|v` are *not* generally identical to some + Note: The marginals `X_1|v, ..., X_d|v` are *not* generally identical to some parameterization of `distribution`. This is due to the fact that the sum of draws from `distribution` are not generally itself the same `distribution`. @@ -331,12 +280,16 @@ class VectorDiffeomixture(distribution_lib.Distribution): optimize Monte-Carlo objectives. Such objectives are a finite-sample approximation of an expectation and arise throughout scientific computing. + WARNING: If you backprop through a VectorDiffeomixture sample and the "base" + distribution is both: not `FULLY_REPARAMETERIZED` and a function of trainable + variables, then the gradient is not guaranteed correct! + #### Examples ```python tfd = tf.contrib.distributions - # Create two batches of VectorDiffeomixtures, one with mix_loc=[0.] and + # Create two batches of VectorDiffeomixtures, one with mix_loc=[0.], # another with mix_loc=[1]. In both cases, `K=2` and the affine # transformations involve: # k=0: loc=zeros(dims) scale=LinearOperatorScaledIdentity @@ -344,7 +297,7 @@ class VectorDiffeomixture(distribution_lib.Distribution): dims = 5 vdm = tfd.VectorDiffeomixture( mix_loc=[[0.], [1]], - mix_scale=[1.], + temperature=[1.], distribution=tfd.Normal(loc=0., scale=1.), loc=[ None, # Equivalent to `np.zeros(dims, dtype=np.float32)`. @@ -364,7 +317,7 @@ class VectorDiffeomixture(distribution_lib.Distribution): def __init__(self, mix_loc, - mix_scale, + temperature, distribution, loc=None, scale=None, @@ -373,15 +326,24 @@ class VectorDiffeomixture(distribution_lib.Distribution): validate_args=False, allow_nan_stats=True, name="VectorDiffeomixture"): - """Constructs the VectorDiffeomixture on `R**d`. + """Constructs the VectorDiffeomixture on `R^d`. + + The vector diffeomixture (VDM) approximates the compound distribution + + ```none + p(x) = int p(x | z) p(z) dz, + where z is in the K-simplex, and + p(x | z) := p(x | loc=sum_k z[k] loc[k], scale=sum_k z[k] scale[k]) + ``` Args: - mix_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`. Represents - the `location` parameter of the SoftmaxNormal used for selecting one of - the `K` affine transformations. - mix_scale: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`. - Represents the `scale` parameter of the SoftmaxNormal used for selecting - one of the `K` affine transformations. + mix_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`. + In terms of samples, larger `mix_loc[..., k]` ==> + `Z` is more likely to put more weight on its `kth` component. + temperature: `float`-like `Tensor`. Broadcastable with `mix_loc`. + In terms of samples, smaller `temperature` means one component is more + likely to dominate. I.e., smaller `temperature` makes the VDM look more + like a standard mixture of `K` components. distribution: `tf.Distribution`-like instance. Distribution from which `d` iid samples are used as input to the selected affine transformation. Must be a scalar-batch, scalar-event distribution. Typically @@ -401,8 +363,9 @@ class VectorDiffeomixture(distribution_lib.Distribution): transformation. `LinearOperator`s must have shape `[B1, ..., Bb, d, d]`, `b >= 0`, i.e., characterizes `b`-batches of `d x d` matrices quadrature_size: Python `int` scalar representing number of - quadrature points. - quadrature_fn: Python callable taking `mix_loc`, `mix_scale`, + quadrature points. Larger `quadrature_size` means `q_N(x)` better + approximates `p(x)`. + quadrature_fn: Python callable taking `normal_loc`, `normal_scale`, `quadrature_size`, `validate_args` and returning `tuple(grid, probs)` representing the SoftmaxNormal grid and corresponding normalized weight. normalized) weight. @@ -430,7 +393,7 @@ class VectorDiffeomixture(distribution_lib.Distribution): ValueError: if `not distribution.is_scalar_event`. """ parameters = locals() - with ops.name_scope(name, values=[mix_loc, mix_scale]): + with ops.name_scope(name, values=[mix_loc, temperature]): if not scale or len(scale) < 2: raise ValueError("Must specify list (or list-like object) of scale " "LinearOperators, one for each component with " @@ -473,8 +436,15 @@ class VectorDiffeomixture(distribution_lib.Distribution): raise NotImplementedError("Currently only bimixtures are supported; " "len(scale)={} is not 2.".format(len(scale))) + mix_loc = ops.convert_to_tensor( + mix_loc, dtype=dtype, name="mix_loc") + temperature = ops.convert_to_tensor( + temperature, dtype=dtype, name="temperature") self._grid, probs = tuple(quadrature_fn( - mix_loc, mix_scale, quadrature_size, validate_args)) + mix_loc / temperature, + 1. / temperature, + quadrature_size, + validate_args)) # Note: by creating the logits as `log(prob)` we ensure that # `self.mixture_distribution.logits` is equivalent to @@ -618,7 +588,14 @@ class VectorDiffeomixture(distribution_lib.Distribution): weight = array_ops.gather( array_ops.reshape(self.grid, shape=[-1]), ids + offset) - weight = weight[..., array_ops.newaxis] + # At this point, weight flattened all batch dims into one. + # We also need to append a singleton to broadcast with event dims. + if self.batch_shape.is_fully_defined(): + new_shape = [-1] + self.batch_shape.as_list() + [1] + else: + new_shape = array_ops.concat( + ([-1], self.batch_shape_tensor(), [1]), axis=0) + weight = array_ops.reshape(weight, shape=new_shape) if len(x) != 2: # We actually should have already triggered this exception. However as a @@ -686,7 +663,7 @@ class VectorDiffeomixture(distribution_lib.Distribution): # To compute E[Cov(Z|V)], we'll add matrices within three categories: # scaled-identity, diagonal, and full. Then we'll combine these at the end. - scaled_identity = None + scale_identity_multiplier = None diag = None full = None @@ -694,10 +671,12 @@ class VectorDiffeomixture(distribution_lib.Distribution): s = aff.scale # Just in case aff.scale has side-effects, we'll call once. if (s is None or isinstance(s, linop_identity_lib.LinearOperatorIdentity)): - scaled_identity = add(scaled_identity, p[..., k, array_ops.newaxis]) + scale_identity_multiplier = add(scale_identity_multiplier, + p[..., k, array_ops.newaxis]) elif isinstance(s, linop_identity_lib.LinearOperatorScaledIdentity): - scaled_identity = add(scaled_identity, (p[..., k, array_ops.newaxis] * - math_ops.square(s.multiplier))) + scale_identity_multiplier = add( + scale_identity_multiplier, + (p[..., k, array_ops.newaxis] * math_ops.square(s.multiplier))) elif isinstance(s, linop_diag_lib.LinearOperatorDiag): diag = add(diag, (p[..., k, array_ops.newaxis] * math_ops.square(s.diag_part()))) @@ -709,12 +688,13 @@ class VectorDiffeomixture(distribution_lib.Distribution): full = add(full, x) # We must now account for the fact that the base distribution might have a - # non-unity variance. Recall that `Cov(SX+m) = S.T Cov(X) S = S.T S Var(X)`. + # non-unity variance. Recall that, since X ~ iid Law(X_0), + # `Cov(SX+m) = S Cov(X) S.T = S S.T Diag(Var(X_0))`. # We can scale by `Var(X)` (vs `Cov(X)`) since X corresponds to `d` iid # samples from a scalar-event distribution. v = self.distribution.variance() - if scaled_identity is not None: - scaled_identity *= v + if scale_identity_multiplier is not None: + scale_identity_multiplier *= v if diag is not None: diag *= v[..., array_ops.newaxis] if full is not None: @@ -723,10 +703,10 @@ class VectorDiffeomixture(distribution_lib.Distribution): if diag_only: # Apparently we don't need the full matrix, just the diagonal. r = add(diag, full) - if r is None and scaled_identity is not None: + if r is None and scale_identity_multiplier is not None: ones = array_ops.ones(self.event_shape_tensor(), dtype=self.dtype) - return scaled_identity * ones - return add(r, scaled_identity) + return scale_identity_multiplier[..., array_ops.newaxis] * ones + return add(r, scale_identity_multiplier) # `None` indicates we don't know if the result is positive-definite. is_positive_definite = (True if all(aff.scale.is_positive_definite @@ -742,10 +722,10 @@ class VectorDiffeomixture(distribution_lib.Distribution): to_add.append(linop_full_lib.LinearOperatorFullMatrix( matrix=full, is_positive_definite=is_positive_definite)) - if scaled_identity is not None: + if scale_identity_multiplier is not None: to_add.append(linop_identity_lib.LinearOperatorScaledIdentity( num_rows=self.event_shape_tensor()[0], - multiplier=scaled_identity, + multiplier=scale_identity_multiplier, is_positive_definite=is_positive_definite)) return (linop_add_lib.add_operators(to_add)[0].to_dense() diff --git a/tensorflow/contrib/eager/README.md b/tensorflow/contrib/eager/README.md index 09242ee47ddd044dfc99e22d5b7751a989c86485..9d2ca07c3a25fa7acb9b0f5806b763d9a57b51fa 100644 --- a/tensorflow/contrib/eager/README.md +++ b/tensorflow/contrib/eager/README.md @@ -41,28 +41,8 @@ support for distributed and multi-GPU training and CPU performance. ## Installation -Since eager execution is not yet part of a TensorFlow release, using it requires -either [building from source](https://www.tensorflow.org/install/install_sources) -or the latest nightly builds. The nightly builds are available as: - -- [`pip` packages](https://github.com/tensorflow/tensorflow/blob/master/README.md#installation) and - -- [docker](https://hub.docker.com/r/tensorflow/tensorflow/) images. - -For example, to run the latest nightly docker image: - -```sh -# If you have a GPU, use https://github.com/NVIDIA/nvidia-docker -nvidia-docker pull tensorflow/tensorflow:nightly-gpu -nvidia-docker run -it -p 8888:8888 tensorflow/tensorflow:nightly-gpu - -# If you do not have a GPU, use the CPU-only image -docker pull tensorflow/tensorflow:nightly -docker run -it -p 8888:8888 tensorflow/tensorflow:nightly -``` - -And then visit http://localhost:8888 in your browser for a Jupyter notebook -environment. Try out the notebooks below. +Eager execution is included in TensorFlow versions 1.5 and above. +Installation instructions at https://www.tensorflow.org/install/ ## Documentation diff --git a/tensorflow/contrib/eager/proto/checkpointable_object_graph.proto b/tensorflow/contrib/eager/proto/checkpointable_object_graph.proto index c962638aa11c06dcd5be6a794314e029ae84e572..024765acb28726fd102dfbf167f4e780072ce6e7 100644 --- a/tensorflow/contrib/eager/proto/checkpointable_object_graph.proto +++ b/tensorflow/contrib/eager/proto/checkpointable_object_graph.proto @@ -4,9 +4,9 @@ option cc_enable_arenas = true; package tensorflow.contrib.eager; -// Prototype for an addition to BundleHeaderProto which saves extra information -// about the objects which own variables, allowing for more robust checkpoint -// loading into modified programs. +// Prototype format which saves extra information about the objects which own +// variables, allowing for more robust checkpoint loading into modified +// programs. Currently stored in its own entry in a TensorBundle. message CheckpointableObjectGraph { message Object { @@ -14,40 +14,39 @@ message CheckpointableObjectGraph { // An index into `CheckpointableObjectGraph.nodes`, indicating the object // being referenced. int32 node_id = 1; - // A numeric identifier for this object within its parent. - int32 local_uid = 2; - // A user-provided name for the edge. May be blank/omitted, in which case - // there is no explicitly provided local name; fall back on local_uid. - string local_name = 3; + // A user-provided name for the edge. + string local_name = 2; } - message VariableReference { - // A name for the variable which is unique within the object which owns - // it. Does not include a name_scope or variable_scope prefix. - string local_name = 1; - // The full name of the variable. Used to allow name-based loading of - // checkpoints which were saved using an object-based API. + message SerializedTensor { + // A name for the Tensor. Simple variables have only one + // `SerializedTensor` named "VARIABLE_VALUE" by convention. This value may + // be restored on object creation as an optimization. + string name = 1; + // The full name of the variable/tensor, if applicable. Used to allow + // name-based loading of checkpoints which were saved using an + // object-based API. Should match the checkpoint key which would have been + // assigned by tf.train.Saver. string full_name = 2; + // The generated name of the Tensor in the checkpoint. + string checkpoint_key = 3; } message SlotVariableReference { - // An index into `CheckpointableObjectGraph.nodes`, indicating the object - // which created the variable that this variable is slotting for. + // An index into `CheckpointableObjectGraph.nodes`, indicating the + // variable object this slot was created for. int32 original_variable_node_id = 1; - // The local name of the variable being slotted for within the object that - // owns it. - string original_variable_local_name = 2; // The name of the slot (e.g. "m"/"v"). - string slot_name = 3; - // The full name of the slot variable. Used to allow name-based loading of - // checkpoints which were saved using an object-based API. - string full_name = 4; + string slot_name = 2; + // An index into `CheckpointableObjectGraph.nodes`, indicating the + // `Object` with the value of the slot variable. + int32 slot_variable_node_id = 3; } // Objects which this object depends on. repeated ObjectReference children = 1; - // Non-slot variables owned by this object. - repeated VariableReference variables = 2; + // Serialized data specific to this object. + repeated SerializedTensor attributes = 2; // Slot variables owned by this object. repeated SlotVariableReference slot_variables = 3; } diff --git a/tensorflow/contrib/eager/python/BUILD b/tensorflow/contrib/eager/python/BUILD index 086315464c99811371d836aed290b5068729adb0..ad40e55cb48aac08eca7022846a0bd07b8accb3f 100644 --- a/tensorflow/contrib/eager/python/BUILD +++ b/tensorflow/contrib/eager/python/BUILD @@ -20,6 +20,7 @@ py_library( "//tensorflow/python:numerics", "//tensorflow/python:resource_variable_ops", "//tensorflow/python:script_ops", + "//tensorflow/python:template", "//tensorflow/python:util", "//tensorflow/python:variable_scope", "//tensorflow/python/eager:backprop", @@ -51,7 +52,7 @@ py_library( srcs_version = "PY2AND3", visibility = ["//tensorflow:internal"], deps = [ - "//tensorflow/contrib/data/python/ops:prefetching_py", + "//tensorflow/contrib/data/python/ops:prefetching_ops", "//tensorflow/python:array_ops", "//tensorflow/python:dataset_ops_gen", "//tensorflow/python:errors", @@ -68,6 +69,7 @@ cuda_py_test( srcs = ["datasets_test.py"], additional_deps = [ ":datasets", + "//tensorflow/contrib/lookup:lookup_py", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", @@ -218,29 +220,42 @@ py_test( ) py_library( - name = "checkpointable", - srcs = ["checkpointable.py"], + name = "checkpointable_utils", + srcs = ["checkpointable_utils.py"], srcs_version = "PY2AND3", visibility = ["//tensorflow:internal"], deps = [ "//tensorflow/contrib/eager/proto:checkpointable_object_graph_proto_py", + "//tensorflow/python:constant_op", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:init_ops", + "//tensorflow/python:io_ops", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python:tensor_shape", "//tensorflow/python:training", "//tensorflow/python:variable_scope", + "//tensorflow/python/eager:context", ], ) py_test( - name = "checkpointable_test", - srcs = ["checkpointable_test.py"], + name = "checkpointable_utils_test", + srcs = ["checkpointable_utils_test.py"], srcs_version = "PY2AND3", deps = [ - ":checkpointable", + ":checkpointable_utils", ":network", "//tensorflow/python:constant_op", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", + "//tensorflow/python:init_ops", "//tensorflow/python:layers", + "//tensorflow/python:layers_base", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python:state_ops", "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", diff --git a/tensorflow/contrib/eager/python/checkpointable.py b/tensorflow/contrib/eager/python/checkpointable.py deleted file mode 100644 index b141ffb2bc03b8e38f8481bc044c3aae7e156c15..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/python/checkpointable.py +++ /dev/null @@ -1,392 +0,0 @@ -"""An object-local variable management scheme.""" -# Copyright 2017 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. -# ============================================================================== -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import collections -import re - -from tensorflow.contrib.eager.proto import checkpointable_object_graph_pb2 -from tensorflow.python.ops import variable_scope -from tensorflow.python.training import optimizer as optimizer_lib -from tensorflow.python.training import saver as saver_lib - -_CheckpointableReference = collections.namedtuple( - "_CheckpointableReference", - [ - "name", # The local name if explicitly specified, else None. - "local_uid", # 0 for the first dependency, 1 for the next, ... Used for - # routing checkpointed variables to their correct - # Checkpointables when "name" is not set (see docstring of - # `track_checkpointable`). - "ref" # The Checkpointable object being referenced. - ]) - -_OwnedVariable = collections.namedtuple( - "_OwnedVariable", - [ - "name", # The variable's (local) name. - "variable" # The owned variable object. - ]) - -# Validation regular expression for the local names of Checkpointable -# objects. In particular, disallows "/" in names, and reserves -# underscore-prefixed names. -_VALID_LOCAL_NAME = re.compile(r"^[A-Za-z0-9.][A-Za-z0-9_.-]*$") - -# Keyword for identifying that the next bit of a checkpoint variable name is a -# slot name. May not be the local name of a checkpointable. Checkpoint names for -# slot variables look like: -# -# /<_OPTIMIZER_SLOTS_NAME>// -# -# Where is a full path from the checkpoint root to the -# variable being slotted for. -_OPTIMIZER_SLOTS_NAME = "_OPTIMIZER_SLOT" - - -class Checkpointable(object): - """Manages variables and dependencies on other objects. - - To make reliable checkpoints, all `Checkpointable`s on which this object - depends must be registered in the constructor using `track_checkpointable` in - a deterministic order, and if possible they should be named. Variables may be - created using `add_variable` outside of the constructor and in any order, but - only these variables will be saved. - """ - - def __init__(self): - # Basically less useful OrderedDicts but without the reference cycles. - # TODO(allenl): Switch these to OrderedDict once TensorFlow supports only - # Python 3.6+. - self._checkpoint_dependencies = [] # A list of _CheckpointableReference - # objects. - self._dependency_names = set() - self._owned_variables = [] # A list of _OwnedVariable objects. - self._owned_variable_names = set() - - def add_variable(self, name, shape, dtype=None, initializer=None, **kwargs): - """Create a new variable object to be saved with this `Checkpointable`. - - If the user has requested that this object or another `Checkpointable` which - depends on this object be restored from a checkpoint (deferred loading - before variable object creation), `initializer` may be ignored and the value - from the checkpoint used instead. - - Args: - name: A name for the variable. Must be unique within this object. - shape: The shape of the variable. - dtype: The data type of the variable. - initializer: The initializer to use. Ignored if deferred loading has been - requested. - **kwargs: Passed to get_variable. - - Returns: - The new variable object. - - Raises: - ValueError: If the variable name is not unique. - """ - if name in self._owned_variable_names: - raise ValueError( - ("A variable named '%s' already exists in this Checkpointable, but " - "Checkpointable.add_variable called to create another with " - "that name. Variable names must be unique within a Checkpointable " - "object.") % (name,)) - if "getter" in kwargs: - # Allow the getter to be overridden, typically because there is a need for - # compatibility with some other variable creation mechanism. This should - # be relatively uncommon in user code. - getter = kwargs.pop("getter") - else: - getter = variable_scope.get_variable - # TODO(allenl): handle deferred loading - new_variable = getter( - name=name, shape=shape, dtype=dtype, initializer=initializer, **kwargs) - self._owned_variables.append( - _OwnedVariable(name=name, variable=new_variable)) - self._owned_variable_names.add(name) - return new_variable - - def track_checkpointable(self, checkpointable, name=None): - """Declare a dependency on another `Checkpointable` object. - - Indicates that checkpoints for this object should include variables from - `checkpointable`. - - Variables in a checkpoint are mapped to `Checkpointable`s based on names if - provided when the checkpoint was written, but otherwise use the order those - `Checkpointable`s were declared as dependencies. Both `name` arguments and - the dependency declaration order should be deterministic. - - There are two sufficient conditions to avoid breaking existing checkpoints - when modifying a class: (1) New dependencies must be declared after existing - dependencies, and (2) dependencies which were previously declared may never - be removed (a trivial placeholder with the same name may be used instead). - - Args: - checkpointable: A `Checkpointable` which this object depends on. - name: A local name for `checkpointable`, used for loading checkpoints into - the correct objects. If provided, it must be unique within this - `Checkpointable`. If None, dependency declaration order is used instead. - - Returns: - `checkpointable`, for convenience when declaring a dependency and - assigning to a member variable in one statement. - - Raises: - RuntimeError: If __init__ was not called. - TypeError: If `checkpointable` does not inherit from `Checkpointable`. - ValueError: For invalid names. - """ - if not hasattr(self, "_checkpoint_dependencies"): - raise RuntimeError("Need to call Checkpointable.__init__ before calling " - "Checkpointable.track_checkpointable().") - if not isinstance(checkpointable, Checkpointable): - raise TypeError( - ("Checkpointable.track_checkpointable() passed type %s, not a " - "Checkpointable.") % (type(checkpointable),)) - if name is not None: - if not _VALID_LOCAL_NAME.match(name): - raise ValueError( - ("Checkpointable names must match the regular expression '%s', but " - "got an invalid name '%s' instead.") % (_VALID_LOCAL_NAME.pattern, - name)) - if name in self._dependency_names: - raise ValueError( - ("Called Checkpointable.track_checkpointable() with name='%s', but " - "a Checkpointable with this name is already declared as a " - "dependency. If provided, names must be unique.") % (name,)) - self._dependency_names.add(name) - self._checkpoint_dependencies.append( - _CheckpointableReference( - name=name, - ref=checkpointable, - # TODO(allenl): Should this be exposed to allow users to stop - # depending on things and still load checkpoints when not using - # names? - local_uid=len(self._checkpoint_dependencies))) - return checkpointable - - @property - def checkpoint_dependencies(self): - """Other `Checkpointable` objects on which this object depends.""" - return self._checkpoint_dependencies - - -def _breadth_first_checkpointable_traversal(root_checkpointable): - """Find shortest paths to all variables owned by dependencies of root.""" - bfs_sorted = [] - root_checkpointable_reference = _CheckpointableReference( - name=None, local_uid=0, ref=root_checkpointable) - to_visit = collections.deque([root_checkpointable_reference]) - path_to_root = {root_checkpointable_reference: ()} - while to_visit: - current_checkpointable = to_visit.popleft() - bfs_sorted.append(current_checkpointable) - for child_checkpointable in ( - current_checkpointable.ref.checkpoint_dependencies): - if child_checkpointable not in path_to_root: - path_to_root[child_checkpointable] = ( - path_to_root[current_checkpointable] + (child_checkpointable,)) - to_visit.append(child_checkpointable) - return bfs_sorted, path_to_root - - -def _object_prefix_from_path(path_to_root): - return "/".join((checkpointable.name if checkpointable.name else "_%d" % ( - checkpointable.local_uid,)) for checkpointable in path_to_root) - - -def _escape_variable_name(variable_name): - # We need to support slashes in variable names for compatibility, since this - # naming scheme is being patched in to things like Layer.add_variable where - # slashes were previously accepted. We also want to use slashes to indicate - # edges traversed to reach the variable, so we escape forward slashes in - # variable names. - return variable_name.replace("_S_", "_S_.").replace(r"/", r"_S__") - - -def _variable_naming_for_object(path_to_root): - """Make a function for naming variables in an object.""" - # Name non-slot variables: - # - # / - # - # is not necessarily unique, but this is fine since we also - # save the graph of `Checkpointable`s with the checkpoint. Even if this path - # no longer exists because of a change in the Python program, we can look up - # the `Checkpointable` which owns the variable in the checkpoint's graph and - # use another path if one still exists. - - object_prefix = _object_prefix_from_path(path_to_root) - if object_prefix: - object_prefix += "/" - - def _name_single_variable(owned_variable): - """Names a variable within an object.""" - return object_prefix + _escape_variable_name(owned_variable.name) - - return _name_single_variable - - -def _slot_variable_naming_for_optimizer(optimizer, path_to_root): - """Make a function for naming slot variables in an optimizer.""" - # Name slot variables: - # - # /<_OPTIMIZER_SLOTS_NAME>// - # - # where is exactly the checkpoint name used for the original - # variable, including the path from the checkpoint root and the local name in - # the object which owns it. Note that we only save slot variables if the - # variable it's slotting for is also being saved. - - optimizer_identifier = "/%s/%s/" % (_OPTIMIZER_SLOTS_NAME, - _object_prefix_from_path(path_to_root)) - - def _name_slot_variable(variable_path, slot_name): - """With an optimizer specified, name a slot variable.""" - - if not _VALID_LOCAL_NAME.match(slot_name): - # Slot variable names include the name of the slot. We need to - # validate that part of the name to be sure that the checkpoint name - # is a valid name scope name. - raise ValueError( - ("Could not save slot variables for optimizer %s, because its " - "slot name has invalid characters (got '%s', was expecting it " - "to match the regular expression '%s').") % - (optimizer, slot_name, _VALID_LOCAL_NAME.pattern)) - - return variable_path + optimizer_identifier + slot_name - - return _name_slot_variable - - -def _serialize_non_slot_variables(checkpointable_objects, path_to_root, - object_graph_proto): - """Name non-slot variables and add them to `object_graph_proto`.""" - named_variables = {} - non_slot_variables = [] - checkpoint_node_ids = {} - - for checkpoint_id, checkpointable in enumerate(checkpointable_objects): - checkpoint_node_ids[checkpointable] = checkpoint_id - - for checkpoint_id, checkpointable in enumerate(checkpointable_objects): - naming_scheme = _variable_naming_for_object(path_to_root[checkpointable]) - object_proto = object_graph_proto.nodes.add() - for owned_variable in checkpointable.ref._owned_variables: # pylint: disable=protected-access - variable_name = naming_scheme(owned_variable) - named_variables[variable_name] = owned_variable.variable - non_slot_variables.append(( - variable_name, # The variable's full checkpoint name - owned_variable, # The variable's _OwnedVariable object - checkpoint_id)) # The checkpoint ID of the node which owns this - # variable. - variable_proto = object_proto.variables.add() - variable_proto.local_name = owned_variable.name - # Figure out the name-based Saver's name for this variable. - saver_dict = saver_lib.BaseSaverBuilder.OpListToDict( - [owned_variable.variable], convert_variable_to_tensor=False) - variable_full_name, = saver_dict.keys() - variable_proto.full_name = variable_full_name - - for child in checkpointable.ref.checkpoint_dependencies: - child_proto = object_proto.children.add() - child_proto.node_id = checkpoint_node_ids[child] - child_proto.local_uid = child.local_uid - if child.name is not None: - child_proto.local_name = child.name - return named_variables, non_slot_variables - - -def _serialize_slot_variables(checkpointable_objects, path_to_root, - non_slot_variables, object_graph_proto): - """Name slot variables and add them to `object_graph_proto`.""" - named_slot_variables = {} - for optimizer_checkpoint_id, checkpointable_ref in enumerate( - checkpointable_objects): - if isinstance(checkpointable_ref.ref, optimizer_lib.Optimizer): - optimizer_object_proto = object_graph_proto.nodes[optimizer_checkpoint_id] - naming_scheme = _slot_variable_naming_for_optimizer( - optimizer=checkpointable_ref.ref, - path_to_root=path_to_root[checkpointable_ref]) - slot_names = checkpointable_ref.ref.get_slot_names() - for (variable_path, owned_variable, - original_node_checkpoint_id) in non_slot_variables: - for slot_name in slot_names: - slot_variable = checkpointable_ref.ref.get_slot( - owned_variable.variable, slot_name) - if slot_variable is not None: - checkpoint_name = naming_scheme( - variable_path=variable_path, slot_name=slot_name) - named_slot_variables[checkpoint_name] = slot_variable - slot_variable_proto = optimizer_object_proto.slot_variables.add() - slot_variable_proto.slot_name = slot_name - # Figure out the name-based Saver's name for this variable. - saver_dict = saver_lib.BaseSaverBuilder.OpListToDict( - [slot_variable], convert_variable_to_tensor=False) - slot_variable_full_name, = saver_dict.keys() - slot_variable_proto.full_name = slot_variable_full_name - slot_variable_proto.original_variable_local_name = ( - owned_variable.name) - slot_variable_proto.original_variable_node_id = ( - original_node_checkpoint_id) - return named_slot_variables - - -# TODO(allenl): Convenience utility for saving multiple objects (i.e. construct -# a root Checkpointable if passed a list of Checkpointables). -def _serialize_object_graph(root_checkpointable): - """Determine checkpoint keys for variables and build a serialized graph. - - Non-slot variables are keyed based on a shortest path from the root saveable - to the object which owns the variable (i.e. the one which called - `Checkpointable.add_variable` to create it). - - Slot variables are keyed based on a shortest path to the variable being - slotted for, a shortest path to their optimizer, and the slot name. - - Args: - root_checkpointable: A `Checkpointable` object whose variables (including - the variables of dependencies, recursively) should be saved. - - Returns: - A tuple of (named_variables, object_graph_proto): - named_variables: A dictionary mapping names to variable objects. - object_graph_proto: A CheckpointableObjectGraph protocol buffer containing - the serialized object graph and variable references. - - Raises: - ValueError: If there are invalid characters in an optimizer's slot names. - """ - checkpointable_objects, path_to_root = ( - _breadth_first_checkpointable_traversal(root_checkpointable)) - object_graph_proto = ( - checkpointable_object_graph_pb2.CheckpointableObjectGraph()) - - # Gather non-slot variables. - named_variables, non_slot_variables = _serialize_non_slot_variables( - checkpointable_objects, path_to_root, object_graph_proto) - - # Gather slot variables which are associated with variables gathered above. - named_slot_variables = _serialize_slot_variables( - checkpointable_objects, path_to_root, non_slot_variables, - object_graph_proto) - - named_variables.update(named_slot_variables) - return named_variables, object_graph_proto diff --git a/tensorflow/contrib/eager/python/checkpointable_test.py b/tensorflow/contrib/eager/python/checkpointable_test.py deleted file mode 100644 index f820990bbe5fe6c9b4cdf890680aaad0847010c0..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/python/checkpointable_test.py +++ /dev/null @@ -1,277 +0,0 @@ -# Copyright 2017 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. -# ============================================================================== -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import functools -import six - -from tensorflow.contrib.eager.python import checkpointable -from tensorflow.contrib.eager.python import network as network_lib -from tensorflow.python.eager import context -from tensorflow.python.eager import test -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util -from tensorflow.python.layers import core -from tensorflow.python.ops import variable_scope -from tensorflow.python.ops import variables -from tensorflow.python.training import adam -from tensorflow.python.training import training_util - - -class CheckpointableDenseLayer(core.Dense, checkpointable.Checkpointable): - - def __init__(self, *args, **kwargs): - checkpointable.Checkpointable.__init__(self) - core.Dense.__init__(self, *args, **kwargs) - - def add_variable(self, name, shape, **kwargs): - # Calls both Checkpointable.add_variable and Layer.add_variable. Eventually - # Layer.add_variable should inherit from Checkpointable and simply call - # super and then do post-processing. - return checkpointable.Checkpointable.add_variable( - self, - name=name, - shape=shape, - getter=functools.partial(core.Dense.add_variable, self), - **kwargs) - - -# pylint: disable=not-callable -class CheckpointableNetwork(network_lib.Network, checkpointable.Checkpointable): - - def __init__(self): - network_lib.Network.__init__(self) - checkpointable.Checkpointable.__init__(self) - - def track_layer(self, layer, name=None): - self.track_checkpointable(layer, name=name) - return super(CheckpointableNetwork, self).track_layer(layer) - - -class CheckpointableAdam(adam.AdamOptimizer, checkpointable.Checkpointable): - - def __init__(self, *args, **kwargs): - checkpointable.Checkpointable.__init__(self) - adam.AdamOptimizer.__init__(self, *args, **kwargs) - - # NOTE: Copied from AdamOptimizer with modifications to use add_variable - # for non-slot variables. These contortions are necessary to maintain - # checkpoint compatibility with variable.name based saving. - def _create_slots(self, var_list): - # Create the beta1 and beta2 accumulators on the same device as the first - # variable. Sort the var_list to make sure this device is consistent across - # workers (these need to go on the same PS, otherwise some updates are - # silently ignored). - first_var = min(var_list, key=lambda x: x.name) - - create_new = self._beta1_power is None - if not create_new and context.in_graph_mode(): - create_new = (self._beta1_power.graph is not first_var.graph) - - if create_new: - with ops.colocate_with(first_var): - - def _variable_getter(name, shape, dtype, initializer): - del shape, dtype # not used, but there for compatibility - return variable_scope.variable( - name=name, initial_value=initializer, trainable=False) - - self._beta1_power = self.add_variable( - name="beta1_power", - shape=[], - initializer=self._beta1, - getter=_variable_getter) - self._beta2_power = self.add_variable( - name="beta2_power", - shape=[], - initializer=self._beta2, - getter=_variable_getter) - # Create slots for the first and second moments. - for v in var_list: - self._zeros_slot(v, "m", self._name) - self._zeros_slot(v, "v", self._name) - - # TODO(allenl): Override slot variable creation (_get_or_make_slot, - # _get_or_make_slot_with_initializer, _zeros_slot) to allow deferred - # loading. Likely no need to run this through add_variable, since gathering - # slot variables is special cased anyway. - - -class MyNetwork(CheckpointableNetwork): - """A concrete Network for testing.""" - - def __init__(self): - super(MyNetwork, self).__init__() - self._named = self.track_layer( - CheckpointableDenseLayer(1, use_bias=True), name="named_dense") - self._unnamed = self.track_layer( - CheckpointableDenseLayer(1, use_bias=False)) - - def call(self, values): - return self._unnamed(self._named(values)) - - -class Root(checkpointable.Checkpointable): - """A stand-in for a Trainer class.""" - - def __init__(self, optimizer, network): - super(Root, self).__init__() - self.track_checkpointable(optimizer, name="optimizer") - self.track_checkpointable(network, name="network") - self._global_step = None - - @property - def global_step(self): - if self._global_step is None: - # Get the default create_global_step utility to actually call - # self.add_variable, by setting a custom getter. - def _owned_variable_as_custom_getter(getter, *args, **kwargs): - return self.add_variable(*args, getter=getter, **kwargs) - - with variable_scope.variable_scope( - "", custom_getter=_owned_variable_as_custom_getter): - self._global_step = training_util.create_global_step() - return self._global_step - - -class CheckpointNamingTests(test.TestCase): - - @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) - def testNamingWithOptimizer(self): - input_value = constant_op.constant([[3.]]) - network = MyNetwork() - # A nuisance Network using the same optimizer. Its slot variables should not - # go in the checkpoint, since it is never depended on. - other_network = MyNetwork() - optimizer = CheckpointableAdam(0.001) - root_checkpointable = Root(optimizer=optimizer, network=network) - if context.in_eager_mode(): - optimizer.minimize( - lambda: network(input_value), - global_step=root_checkpointable.global_step) - optimizer.minimize( - lambda: other_network(input_value), - global_step=root_checkpointable.global_step) - else: - train_op = optimizer.minimize( - network(input_value), global_step=root_checkpointable.global_step) - optimizer.minimize( - other_network(input_value), - global_step=root_checkpointable.global_step) - self.evaluate(variables.global_variables_initializer()) - self.evaluate(train_op) - named_variables, serialized_graph = checkpointable._serialize_object_graph( - root_checkpointable) - expected_checkpoint_names = ( - # Created in the root node, so no prefix. - "global_step", - # No name provided to track_checkpointable(), so the position (1, after - # the named track_checkpointable() which is 0) is used instead. - "network/_1/kernel", - # track_checkpointable() with a name provided, so that's used - "network/named_dense/kernel", - "network/named_dense/bias", - # The optimizer creates two non-slot variables - "optimizer/beta1_power", - "optimizer/beta2_power", - # Slot variables - "network/_1/kernel/_OPTIMIZER_SLOT/optimizer/m", - "network/_1/kernel/_OPTIMIZER_SLOT/optimizer/v", - "network/named_dense/kernel/_OPTIMIZER_SLOT/optimizer/m", - "network/named_dense/kernel/_OPTIMIZER_SLOT/optimizer/v", - "network/named_dense/bias/_OPTIMIZER_SLOT/optimizer/m", - "network/named_dense/bias/_OPTIMIZER_SLOT/optimizer/v", - ) - six.assertCountEqual(self, expected_checkpoint_names, - named_variables.keys()) - # Check that we've mapped to the right variable objects (not exhaustive) - self.assertEqual("global_step:0", named_variables["global_step"].name) - self.assertEqual("my_network/checkpointable_dense_layer_1/kernel:0", - named_variables["network/_1/kernel"].name) - self.assertEqual("my_network/checkpointable_dense_layer/kernel:0", - named_variables["network/named_dense/kernel"].name) - self.assertEqual("beta1_power:0", - named_variables["optimizer/beta1_power"].name) - self.assertEqual("beta2_power:0", - named_variables["optimizer/beta2_power"].name) - # Spot check the generated protocol buffers. - self.assertEqual(0, serialized_graph.nodes[0].children[0].local_uid) - self.assertEqual("optimizer", - serialized_graph.nodes[0].children[0].local_name) - optimizer_node = serialized_graph.nodes[serialized_graph.nodes[0].children[ - 0].node_id] - self.assertEqual("beta1_power", optimizer_node.variables[0].local_name) - self.assertEqual("beta1_power", optimizer_node.variables[0].full_name) - self.assertEqual( - "kernel", optimizer_node.slot_variables[0].original_variable_local_name) - original_variable_owner = serialized_graph.nodes[ - optimizer_node.slot_variables[0].original_variable_node_id] - self.assertEqual("kernel", original_variable_owner.variables[0].local_name) - self.assertEqual("m", optimizer_node.slot_variables[0].slot_name) - # We strip off the :0 suffix, as variable.name-based saving does. - self.assertEqual("my_network/checkpointable_dense_layer/kernel/Adam", - optimizer_node.slot_variables[0].full_name) - self.assertEqual("my_network/checkpointable_dense_layer/kernel/Adam:0", - optimizer.get_slot( - var=named_variables["network/named_dense/kernel"], - name="m").name) - - def _get_checkpoint_name(self, name): - root = checkpointable.Checkpointable() - with variable_scope.variable_scope("get_checkpoint_name"): - # Create the variable in a variable scope so that we get more relaxed - # naming rules (variables outside a scope may not start with "_", "/" or - # "-"). Since we don't use the scope part of the name, these cases are - # somewhat annoying. - root.add_variable(name=name, shape=[1, 2], dtype=dtypes.float64) - named_variables, _ = checkpointable._serialize_object_graph(root) - checkpoint_name, = named_variables.keys() - with ops.name_scope("root/" + checkpoint_name): - pass # Make sure we can use this as an op name if we prefix it. - return checkpoint_name - - @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) - def testVariableNameEscaping(self): - self.assertEqual(r"a_S__b_S__c", self._get_checkpoint_name(r"a/b/c")) - self.assertEqual(r"", self._get_checkpoint_name(r"")) - self.assertEqual(r"_S__", self._get_checkpoint_name(r"/")) - self.assertEqual(r"_S___S_._", self._get_checkpoint_name(r"/_S__")) - - @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) - def testNumberedPath(self): - root = checkpointable.Checkpointable() - leaf = checkpointable.Checkpointable() - root.track_checkpointable(leaf) - leaf.add_variable(name="v", shape=[]) - named_variables, _ = checkpointable._serialize_object_graph(root) - variable_name, = named_variables.keys() - self.assertEqual(r"_0/v", variable_name) - - @test_util.run_in_graph_and_eager_modes() - def testLocalNameValidation(self): - root = checkpointable.Checkpointable() - leaf = checkpointable.Checkpointable() - with self.assertRaisesRegexp(ValueError, "invalid name"): - # Leading underscores are reserved, which avoids conflicts with - # un-named edges in paths and the optimizer slots identifier. - root.track_checkpointable(leaf, name="_12") - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/eager/python/checkpointable_utils.py b/tensorflow/contrib/eager/python/checkpointable_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d3c57bc606179a36d110eb5ac5b29a8f7e2469fb --- /dev/null +++ b/tensorflow/contrib/eager/python/checkpointable_utils.py @@ -0,0 +1,413 @@ +"""Utilities for working with Checkpointable objects.""" +# Copyright 2017 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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import collections + +from tensorflow.contrib.eager.proto import checkpointable_object_graph_pb2 +from tensorflow.python.eager import context +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import init_ops +from tensorflow.python.ops import io_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.training import checkpointable as core_checkpointable +from tensorflow.python.training import optimizer as optimizer_lib +from tensorflow.python.training import saver as saver_lib + + +_ESCAPE_CHAR = "." # For avoiding conflicts with user-specified names. + +# Keyword for identifying that the next bit of a checkpoint variable name is a +# slot name. Checkpoint names for slot variables look like: +# +# /<_OPTIMIZER_SLOTS_NAME>// +# +# Where is a full path from the checkpoint root to the +# variable being slotted for. +_OPTIMIZER_SLOTS_NAME = _ESCAPE_CHAR + "OPTIMIZER_SLOT" +# Keyword for separating the path to an object from the name of an +# attribute in checkpoint names. Used like: +# /<_OBJECT_ATTRIBUTES_NAME>/ +_OBJECT_ATTRIBUTES_NAME = _ESCAPE_CHAR + "ATTRIBUTES" +# Key where the object graph proto is saved in a TensorBundle +_OBJECT_GRAPH_PROTO_KEY = "_CHECKPOINTABLE_OBJECT_GRAPH" + + +# TODO(allenl): If this ends up in a public API, consider adding LINT.IfChange +# or consolidating the implementation with get_variable. +def _default_getter(name, shape, dtype, initializer=None, + partition_info=None, **kwargs): + """A pared-down version of get_variable which does not reuse variables.""" + dtype = dtypes.as_dtype(dtype) + shape_object = tensor_shape.as_shape(shape) + with ops.init_scope(): + if initializer is None: + initializer, initializing_from_value = ( + variable_scope._get_default_variable_store()._get_default_initializer( # pylint: disable=protected-access + name=name, shape=shape_object, dtype=dtype)) + else: + initializing_from_value = not callable(initializer) + # Same logic as get_variable + variable_dtype = dtype.base_dtype + if initializing_from_value: + if shape is not None: + raise ValueError("If initializer is a constant, do not specify shape.") + initial_value = initializer + else: + # Instantiate initializer if provided initializer is a type object. + if isinstance(initializer, type(init_ops.Initializer)): + initializer = initializer(dtype=dtype) + def initial_value(): + return initializer( + shape_object.as_list(), dtype=dtype, partition_info=partition_info) + return resource_variable_ops.ResourceVariable( + initial_value=initial_value, + name=name, + dtype=variable_dtype, + **kwargs + ) + + +def add_variable(checkpointable, name, shape=None, dtype=dtypes.float32, + initializer=None): + """Add a variable to a Checkpointable with no scope influence.""" + return checkpointable._add_variable_with_custom_getter( # pylint: disable=protected-access + name=name, shape=shape, dtype=dtype, + initializer=initializer, getter=_default_getter) + + +def _breadth_first_checkpointable_traversal(root_checkpointable): + """Find shortest paths to all variables owned by dependencies of root.""" + bfs_sorted = [] + to_visit = collections.deque([root_checkpointable]) + path_to_root = {root_checkpointable: ()} + while to_visit: + current_checkpointable = to_visit.popleft() + current_checkpointable._maybe_initialize_checkpointable() # pylint: disable=protected-access + bfs_sorted.append(current_checkpointable) + for child_checkpointable in ( + current_checkpointable._checkpoint_dependencies): # pylint: disable=protected-access + if child_checkpointable.ref not in path_to_root: + path_to_root[child_checkpointable.ref] = ( + path_to_root[current_checkpointable] + (child_checkpointable,)) + to_visit.append(child_checkpointable.ref) + return bfs_sorted, path_to_root + + +def _escape_local_name(name): + # We need to support slashes in local names for compatibility, since this + # naming scheme is being patched in to things like Layer.add_variable where + # slashes were previously accepted. We also want to use slashes to indicate + # edges traversed to reach the variable, so we escape forward slashes in + # names. + return (name.replace(_ESCAPE_CHAR, _ESCAPE_CHAR + _ESCAPE_CHAR) + .replace(r"/", _ESCAPE_CHAR + "S")) + + +def _object_prefix_from_path(path_to_root): + return "/".join( + (_escape_local_name(checkpointable.name) + for checkpointable in path_to_root)) + + +def _slot_variable_naming_for_optimizer(optimizer_path): + """Make a function for naming slot variables in an optimizer.""" + # Name slot variables: + # + # /<_OPTIMIZER_SLOTS_NAME>// + # + # where is exactly the checkpoint name used for the original + # variable, including the path from the checkpoint root and the local name in + # the object which owns it. Note that we only save slot variables if the + # variable it's slotting for is also being saved. + + optimizer_identifier = "/%s/%s/" % (_OPTIMIZER_SLOTS_NAME, optimizer_path) + + def _name_slot_variable(variable_path, slot_name): + """With an optimizer specified, name a slot variable.""" + return (variable_path + + optimizer_identifier + + _escape_local_name(slot_name)) + + return _name_slot_variable + + +def _serialize_slot_variables(checkpointable_objects, node_ids, object_names): + """Gather and name slot variables.""" + non_slot_objects = list(checkpointable_objects) + slot_variables = {} + for checkpointable in non_slot_objects: + if isinstance(checkpointable, optimizer_lib.Optimizer): + naming_scheme = _slot_variable_naming_for_optimizer( + optimizer_path=object_names[checkpointable]) + slot_names = checkpointable.get_slot_names() + for slot_name in slot_names: + for original_variable_node_id, original_variable in enumerate( + non_slot_objects): + try: + slot_variable = checkpointable.get_slot( + original_variable, slot_name) + except AttributeError: + slot_variable = None + if slot_variable is None: + continue + slot_variable._maybe_initialize_checkpointable() # pylint: disable=protected-access + if slot_variable._checkpoint_dependencies: # pylint: disable=protected-access + # TODO(allenl): Gather dependencies of slot variables. + raise NotImplementedError( + "Currently only variables with no dependencies can be saved as " + "slot variables. File a feature request if this limitation " + "bothers you.") + if slot_variable in node_ids: + raise NotImplementedError( + "A slot variable was re-used as a dependency of a " + "Checkpointable object. This is not currently allowed. File a " + "feature request if this limitation bothers you.") + checkpoint_name = naming_scheme( + variable_path=object_names[original_variable], + slot_name=slot_name) + object_names[slot_variable] = checkpoint_name + slot_variable_node_id = len(checkpointable_objects) + node_ids[slot_variable] = slot_variable_node_id + checkpointable_objects.append(slot_variable) + slot_variable_proto = ( + checkpointable_object_graph_pb2.CheckpointableObjectGraph + .Object.SlotVariableReference( + slot_name=slot_name, + original_variable_node_id=original_variable_node_id, + slot_variable_node_id=slot_variable_node_id)) + slot_variables.setdefault(checkpointable, []).append( + slot_variable_proto) + return slot_variables + + +def _serialize_checkpointables( + checkpointable_objects, node_ids, object_names, slot_variables): + """Name non-slot `Checkpointable`s and add them to `object_graph_proto`.""" + object_graph_proto = ( + checkpointable_object_graph_pb2.CheckpointableObjectGraph()) + named_saveables = {} + + for checkpoint_id, checkpointable in enumerate(checkpointable_objects): + assert node_ids[checkpointable] == checkpoint_id + object_proto = object_graph_proto.nodes.add() + object_proto.slot_variables.extend(slot_variables.get(checkpointable, ())) + object_name = object_names[checkpointable] + for name, saveable in ( + checkpointable._gather_tensors_for_checkpoint().items()): # pylint: disable=protected-access + attribute = object_proto.attributes.add() + attribute.name = name + attribute.checkpoint_key = "%s/%s/%s" % ( + object_name, _OBJECT_ATTRIBUTES_NAME, _escape_local_name(name)) + # Figure out the name-based Saver's name for this variable. + saver_dict = saver_lib.BaseSaverBuilder.OpListToDict( + [saveable], convert_variable_to_tensor=False) + attribute.full_name, = saver_dict.keys() + named_saveables[attribute.checkpoint_key] = saveable + + for child in checkpointable._checkpoint_dependencies: # pylint: disable=protected-access + child_proto = object_proto.children.add() + child_proto.node_id = node_ids[child.ref] + child_proto.local_name = child.name + + return named_saveables, object_graph_proto + + +def _serialize_object_graph(root_checkpointable): + """Determine checkpoint keys for variables and build a serialized graph. + + Non-slot variables are keyed based on a shortest path from the root saveable + to the object which owns the variable (i.e. the one which called + `Checkpointable._add_variable` to create it). + + Slot variables are keyed based on a shortest path to the variable being + slotted for, a shortest path to their optimizer, and the slot name. + + Args: + root_checkpointable: A `Checkpointable` object whose variables (including + the variables of dependencies, recursively) should be saved. + + Returns: + A tuple of (named_variables, object_graph_proto): + named_variables: A dictionary mapping names to variable objects. + object_graph_proto: A CheckpointableObjectGraph protocol buffer containing + the serialized object graph and variable references. + + Raises: + ValueError: If there are invalid characters in an optimizer's slot names. + """ + checkpointable_objects, path_to_root = ( + _breadth_first_checkpointable_traversal(root_checkpointable)) + object_names = { + obj: _object_prefix_from_path(path) + for obj, path in path_to_root.items()} + node_ids = {node: node_id for node_id, node + in enumerate(checkpointable_objects)} + slot_variables = _serialize_slot_variables( + checkpointable_objects=checkpointable_objects, + node_ids=node_ids, + object_names=object_names) + return _serialize_checkpointables( + checkpointable_objects=checkpointable_objects, + node_ids=node_ids, + object_names=object_names, + slot_variables=slot_variables) + + +class _NoRestoreSaveable(saver_lib.BaseSaverBuilder.SaveableObject): + + def __init__(self, tensor, name): + spec = saver_lib.BaseSaverBuilder.SaveSpec(tensor, "", name) + super(_NoRestoreSaveable, self).__init__(tensor, [spec], name) + + def restore(self, restored_tensors, restored_shapes): + return control_flow_ops.no_op() + + +def save(file_prefix, root_checkpointable, checkpoint_number=None, + session=None): + """Save a training checkpoint. + + Args: + file_prefix: A prefix to use for the checkpoint filenames + (/path/to/directory/and_a_prefix). Names are generated based on this + prefix and the global step, if provided. + root_checkpointable: A Checkpointable object to save. The checkpoint + includes variables created by this object and any Checkpointable objects + it depends on. + checkpoint_number: An integer variable or Tensor, used to number + checkpoints. Typically this value is saved along with other variables in + training checkpoints, which will happen automatically if it was created by + `root_checkpointable` or one of its dependencies (via + `Checkpointable._add_variable`). + session: The session to evaluate variables in. Ignored when executing + eagerly. If not provided when graph building, the default session is used. + + Returns: + The full path to the checkpoint. + """ + named_variables, serialized_graph = _serialize_object_graph( + root_checkpointable) + if context.in_graph_mode(): + if session is None: + session = ops.get_default_session() + else: + session = None + assert _OBJECT_GRAPH_PROTO_KEY not in named_variables + # TODO(allenl): Feed rather than embedding a constant. + named_variables[_OBJECT_GRAPH_PROTO_KEY] = _NoRestoreSaveable( + tensor=constant_op.constant( + serialized_graph.SerializeToString(), dtype=dtypes.string), + name=_OBJECT_GRAPH_PROTO_KEY) + with ops.device("/device:CPU:0"): + save_path = saver_lib.Saver(var_list=named_variables).save( + sess=session, + save_path=file_prefix, + write_meta_graph=False, + global_step=checkpoint_number) + return save_path + + +class CheckpointLoadStatus(object): + + def __init__(self, checkpoint): + self._checkpoint = checkpoint + + def assert_consumed(self): + """Asserts that all objects in the checkpoint have been created/matched.""" + for node_id, node in enumerate(self._checkpoint.object_graph_proto.nodes): + checkpointable = self._checkpoint.object_by_proto_id.get(node_id, None) + if checkpointable is None: + raise AssertionError("Unresolved object in checkpoint: %s" % (node)) + if checkpointable._update_uid < self._checkpoint.restore_uid: # pylint: disable=protected-access + raise AssertionError( + "Object not assigned a value from checkpoint: %s" % (node)) + return self + + +def restore(save_path, root_checkpointable, session=None): + """Restore a training checkpoint. + + Restores the values of variables created with `Checkpointable._add_variable` + in `root_checkpointable` and any objects that it tracks (transitive). Either + assigns values immediately if variables to restore have been created already, + or defers restoration until the variables are created. Dependencies added to + `root_checkpointable` after this call will be matched if they have a + corresponding object in the checkpoint. + + When building a graph, restorations are executed in the default session if + `session` is `None`. Variable initializers read checkpointed values. + + To disallow deferred loading, assert immediately that all checkpointed + variables have been matched to variable objects: + + ```python + restore(path, root).assert_consumed() + ``` + + An exception will be raised unless every object was matched and its variables + already exist. + + Args: + save_path: The path to the checkpoint, as returned by `save` or + `tf.train.latest_checkpoint`. If None (as when there is no latest + checkpoint for `tf.train.latest_checkpoint` to return), does nothing. + root_checkpointable: The root of the object graph to restore. Variables to + restore need not have been created yet, but all dependencies on other + Checkpointable objects should already be declared. Objects in the + dependency graph are matched to objects in the checkpointed graph, and + matching objects have their variables restored (or the checkpointed values + saved for eventual restoration when the variable is created). + session: The session to evaluate assignment ops in. Ignored when executing + eagerly. If not provided when graph building, the default session is used. + Returns: + A CheckpointLoadStatus object, which can be used to make assertions about + the status of checkpoint restoration. + """ + if save_path is None: + return + if context.in_graph_mode(): + if session is None: + session = ops.get_default_session() + else: + session = None + object_graph_string, = io_ops.restore_v2( + prefix=save_path, + tensor_names=[_OBJECT_GRAPH_PROTO_KEY], + shape_and_slices=[""], + dtypes=[dtypes.string], + name="object_graph_proto_read") + if session is not None: + object_graph_string = session.run(object_graph_string) + else: + object_graph_string = object_graph_string.numpy() + object_graph_proto = ( + checkpointable_object_graph_pb2.CheckpointableObjectGraph()) + object_graph_proto.ParseFromString(object_graph_string) + checkpoint = core_checkpointable._Checkpoint( # pylint: disable=protected-access + object_graph_proto=object_graph_proto, + save_path=save_path, + session=session) + core_checkpointable._CheckpointPosition( # pylint: disable=protected-access + checkpoint=checkpoint, proto_id=0).restore(root_checkpointable) + return CheckpointLoadStatus(checkpoint) diff --git a/tensorflow/contrib/eager/python/checkpointable_utils_test.py b/tensorflow/contrib/eager/python/checkpointable_utils_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1394f0cf0f8a1e2eefaa185c74e376b21f73b688 --- /dev/null +++ b/tensorflow/contrib/eager/python/checkpointable_utils_test.py @@ -0,0 +1,857 @@ +# Copyright 2017 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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +import os +import unittest + +import six + +from tensorflow.contrib.eager.python import checkpointable_utils +from tensorflow.contrib.eager.python import network as network_lib +from tensorflow.python.eager import context +from tensorflow.python.eager import test +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import test_util +from tensorflow.python.layers import base +from tensorflow.python.layers import core +from tensorflow.python.ops import init_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables +from tensorflow.python.training import adam +from tensorflow.python.training import checkpointable +from tensorflow.python.training import saver as core_saver +from tensorflow.python.training import training_util + + +class CheckpointableDenseLayer(core.Dense, checkpointable.Checkpointable): + + def __init__(self, *args, **kwargs): + checkpointable.Checkpointable.__init__(self) + core.Dense.__init__(self, *args, **kwargs) + + def add_variable(self, name, shape, **kwargs): + # Calls both Checkpointable._add_variable and Layer.add_variable. Eventually + # Layer.add_variable should inherit from Checkpointable and simply call + # super and then do post-processing. + return checkpointable.Checkpointable._add_variable_with_custom_getter( + self, + name=name, + shape=shape, + getter=functools.partial(core.Dense.add_variable, self), + **kwargs) + + +# pylint: disable=not-callable +class CheckpointableNetwork(network_lib.Network, checkpointable.Checkpointable): + + def __setattr__(self, name, value): + if isinstance(value, base.Layer): + self.track_layer(value, name=name) + # Checkpointable is next in the method resolution order, so this will catch + # Checkpointable objects which aren't Layers. + super(CheckpointableNetwork, self).__setattr__(name, value) + + def track_layer(self, layer, name): + self._track_checkpointable(layer, name=name) + return super(CheckpointableNetwork, self).track_layer(layer) + + +class CheckpointableAdam(adam.AdamOptimizer, checkpointable.Checkpointable): + + # NOTE: Copied from Optimizer with modifications to use add_variable + # for non-slot variables. These contortions are necessary to maintain + # checkpoint compatibility with variable.name based saving. + # TODO(allenl): Make this cleaner. + def _create_non_slot_variable(self, initial_value, name, colocate_with): + """Add an extra variable, not associated with a slot.""" + if context.in_graph_mode(): + graph = colocate_with.graph + else: + graph = None + + key = (name, graph) + v = self._non_slot_dict.get(key, None) + if v is None: + with ops.colocate_with(colocate_with): + def _variable_getter(name, shape, dtype, initializer): + del shape, dtype # not used, but there for compatibility + return variable_scope.variable( + name=name, initial_value=initializer, trainable=False) + + initial_value = ops.convert_to_tensor(initial_value) + v = self._add_variable_with_custom_getter( + name=name, + shape=initial_value.get_shape(), + initializer=initial_value, + getter=_variable_getter) + + self._non_slot_dict[key] = v + + return v + + +class NonLayerCheckpointable(checkpointable.Checkpointable): + + def __init__(self): + super(NonLayerCheckpointable, self).__init__() + self.a_variable = checkpointable_utils.add_variable( + self, name="a_variable", shape=[]) + + +class MyNetwork(CheckpointableNetwork): + """A concrete Network for testing.""" + + def __init__(self): + super(MyNetwork, self).__init__() + self._named_dense = CheckpointableDenseLayer(1, use_bias=True) + self._via_track_layer = self.track_layer( + CheckpointableDenseLayer(1, use_bias=False), name="via_track_layer") + # We can still track Checkpointables which aren't Layers. + self._non_layer = NonLayerCheckpointable() + + def call(self, values): + return self._via_track_layer(self._named_dense(values)) + + +class Checkpoint(checkpointable.Checkpointable): + """A utility class which groups `Checkpointable` objects.""" + + def __init__(self, **kwargs): + super(Checkpoint, self).__init__() + for k, v in sorted(kwargs.items(), key=lambda item: item[0]): + setattr(self, k, v) + self._save_counter = None + + @property + def save_counter(self): + """An integer variable which starts at zero and is incremented on save. + + Used to number checkpoints. + + Returns: + The save counter variable. + """ + if self._save_counter is None: + # Initialized to 0 and incremented before saving. + self._save_counter = checkpointable_utils.add_variable( + self, name="save_counter", initializer=0, dtype=dtypes.int64) + return self._save_counter + + def save(self, file_prefix, session=None): + assign_op = self.save_counter.assign_add(1) + if context.in_graph_mode(): + if session is None: + session = ops.get_default_session() + session.run(assign_op) + return checkpointable_utils.save( + file_prefix=file_prefix, + root_checkpointable=self, + checkpoint_number=self.save_counter, + session=session) + + def restore(self, save_path): + return checkpointable_utils.restore( + save_path=save_path, + root_checkpointable=self) + + +class InterfaceTests(test.TestCase): + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testAddVariable(self): + obj = NonLayerCheckpointable() + with self.assertRaisesRegexp(ValueError, "do not specify shape"): + checkpointable_utils.add_variable( + obj, name="shape_specified_twice", shape=[], initializer=1) + constant_initializer = checkpointable_utils.add_variable( + obj, name="constant_initializer", initializer=1) + with variable_scope.variable_scope("some_variable_scope"): + ones_initializer = checkpointable_utils.add_variable( + obj, + name="ones_initializer", + shape=[2], + initializer=init_ops.ones_initializer(dtype=dtypes.float32)) + bare_initializer = checkpointable_utils.add_variable( + obj, + name="bare_initializer", + shape=[2, 2], + dtype=dtypes.float64, + initializer=init_ops.zeros_initializer) + + # Even in graph mode, there are no naming conflicts between objects, only + # naming conflicts within an object. + other_duplicate = resource_variable_ops.ResourceVariable( + name="duplicate", initial_value=1.) + duplicate = checkpointable_utils.add_variable( + obj, name="duplicate", shape=[]) + with self.assertRaisesRegexp(ValueError, "'duplicate' already exists"): + checkpointable_utils.add_variable(obj, name="duplicate", shape=[]) + + if context.in_graph_mode(): + self.evaluate(variables.global_variables_initializer()) + self.assertEqual("constant_initializer:0", constant_initializer.name) + self.assertEqual(1, self.evaluate(constant_initializer)) + self.assertEqual("some_variable_scope/ones_initializer:0", + ones_initializer.name) + self.assertAllEqual([1, 1], self.evaluate(ones_initializer)) + self.assertAllEqual([[0., 0.], + [0., 0.]], self.evaluate(bare_initializer)) + self.assertEqual("a_variable:0", obj.a_variable.name) + self.assertEqual("duplicate:0", other_duplicate.name) + if context.in_graph_mode(): + # The .name attribute may be globally influenced, but the checkpoint name + # won't be (tested below). + self.assertEqual("duplicate_1:0", duplicate.name) + else: + # When executing eagerly, there's no uniquification of variable names. The + # checkpoint name will be the same. + self.assertEqual("duplicate:0", duplicate.name) + named_variables, _ = checkpointable_utils._serialize_object_graph(obj) + expected_checkpoint_names = ( + "a_variable/.ATTRIBUTES/VARIABLE_VALUE", + "bare_initializer/.ATTRIBUTES/VARIABLE_VALUE", + "constant_initializer/.ATTRIBUTES/VARIABLE_VALUE", + "duplicate/.ATTRIBUTES/VARIABLE_VALUE", + "ones_initializer/.ATTRIBUTES/VARIABLE_VALUE", + ) + six.assertCountEqual( + self, expected_checkpoint_names, named_variables.keys()) + + def testInitNotCalled(self): + + class NoInit(checkpointable.Checkpointable): + + def __init__(self): + pass + + # __init__ for Checkpointable will be called implicitly. + checkpointable_utils.add_variable(NoInit(), "var", shape=[]) + + def testShapeDtype(self): + root = checkpointable.Checkpointable() + v1 = checkpointable_utils.add_variable( + root, name="v1", initializer=3., dtype=dtypes.float64) + self.assertEqual(dtypes.float64, v1.dtype) + v2 = checkpointable_utils.add_variable( + root, + name="v2", + shape=[3], + initializer=init_ops.ones_initializer, + dtype=dtypes.float64) + self.assertEqual(dtypes.float64, v2.dtype) + self.assertAllEqual([1., 1., 1.], self.evaluate(v2)) + + +class CheckpointingTests(test.TestCase): + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testNamingWithOptimizer(self): + input_value = constant_op.constant([[3.]]) + network = MyNetwork() + # A nuisance Network using the same optimizer. Its slot variables should not + # go in the checkpoint, since it is never depended on. + other_network = MyNetwork() + optimizer = CheckpointableAdam(0.001) + optimizer_step = training_util.get_or_create_global_step() + root_checkpointable = Checkpoint( + optimizer=optimizer, network=network, optimizer_step=optimizer_step) + if context.in_eager_mode(): + optimizer.minimize( + lambda: network(input_value), + global_step=optimizer_step) + optimizer.minimize( + lambda: other_network(input_value), + global_step=optimizer_step) + else: + train_op = optimizer.minimize( + network(input_value), global_step=optimizer_step) + optimizer.minimize( + other_network(input_value), + global_step=optimizer_step) + self.evaluate(variables.global_variables_initializer()) + self.evaluate(train_op) + named_variables, serialized_graph = ( + checkpointable_utils._serialize_object_graph(root_checkpointable)) + expected_checkpoint_names = ( + # Created in the root node, so no prefix. + "optimizer_step", + # No name provided to track_checkpointable(), so the position is used + # instead (one-based). + "network/via_track_layer/kernel", + # track_checkpointable() with a name provided, so that's used + "network/_named_dense/kernel", + "network/_named_dense/bias", + # non-Layer dependency of the network + "network/_non_layer/a_variable", + # The optimizer creates two non-slot variables + "optimizer/beta1_power", + "optimizer/beta2_power", + # Slot variables + "network/via_track_layer/kernel/.OPTIMIZER_SLOT/optimizer/m", + "network/via_track_layer/kernel/.OPTIMIZER_SLOT/optimizer/v", + "network/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m", + "network/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/v", + "network/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/m", + "network/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/v", + ) + suffix = "/.ATTRIBUTES/VARIABLE_VALUE" + expected_checkpoint_names = [ + name + suffix for name in expected_checkpoint_names] + six.assertCountEqual(self, expected_checkpoint_names, + named_variables.keys()) + # Check that we've mapped to the right variable objects (not exhaustive) + self.assertEqual( + "global_step:0", + named_variables["optimizer_step" + suffix].name) + self.assertEqual( + "my_network/checkpointable_dense_layer_1/kernel:0", + named_variables["network/via_track_layer/kernel" + suffix].name) + self.assertEqual( + "my_network/checkpointable_dense_layer/kernel:0", + named_variables["network/_named_dense/kernel" + suffix].name) + self.assertEqual( + "beta1_power:0", + named_variables["optimizer/beta1_power" + suffix].name) + self.assertEqual( + "beta2_power:0", + named_variables["optimizer/beta2_power" + suffix].name) + # Spot check the generated protocol buffers. + self.assertEqual("optimizer", + serialized_graph.nodes[0].children[1].local_name) + optimizer_node = serialized_graph.nodes[serialized_graph.nodes[0].children[ + 1].node_id] + self.assertEqual("beta1_power", + optimizer_node.children[0].local_name) + self.assertEqual("beta1_power", + serialized_graph.nodes[optimizer_node.children[0].node_id] + .attributes[0].full_name) + self.assertEqual( + "my_network/checkpointable_dense_layer/kernel", + serialized_graph.nodes[optimizer_node.slot_variables[0] + .original_variable_node_id] + .attributes[0].full_name) + # We strip off the :0 suffix, as variable.name-based saving does. + self.assertEqual( + "my_network/checkpointable_dense_layer/kernel/Adam", + serialized_graph.nodes[optimizer_node.slot_variables[0] + .slot_variable_node_id] + .attributes[0].full_name) + self.assertEqual( + "my_network/checkpointable_dense_layer/kernel/Adam:0", + optimizer.get_slot( + var=named_variables["network/_named_dense/kernel" + suffix], + name="m").name) + self.assertEqual( + "network/_named_dense/kernel" + suffix, + serialized_graph.nodes[ + optimizer_node.slot_variables[0] + .original_variable_node_id].attributes[0].checkpoint_key) + self.assertEqual("m", optimizer_node.slot_variables[0].slot_name) + self.assertEqual( + "network/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m" + suffix, + serialized_graph.nodes[ + optimizer_node.slot_variables[0] + .slot_variable_node_id].attributes[0].checkpoint_key) + + @test_util.run_in_graph_and_eager_modes() + def testSaveRestore(self): + network = MyNetwork() + optimizer = CheckpointableAdam(0.001) + root_checkpointable = Checkpoint(optimizer=optimizer, network=network) + input_value = constant_op.constant([[3.]]) + if context.in_eager_mode(): + optimizer.minimize( + lambda: network(input_value)) + else: + train_op = optimizer.minimize(network(input_value)) + # TODO(allenl): Make initialization more pleasant when graph building. + root_checkpointable.save_counter # pylint: disable=pointless-statement + self.evaluate(variables.global_variables_initializer()) + self.evaluate(train_op) + prefix = os.path.join(self.get_temp_dir(), "ckpt") + self.evaluate(state_ops.assign(network._named_dense.variables[1], [42.])) + m_bias_slot = optimizer.get_slot(network._named_dense.variables[1], "m") + self.evaluate(state_ops.assign(m_bias_slot, [1.5])) + save_path = root_checkpointable.save(file_prefix=prefix) + self.evaluate(state_ops.assign(network._named_dense.variables[1], [43.])) + self.evaluate(state_ops.assign(root_checkpointable.save_counter, 3)) + optimizer_variables = self.evaluate(optimizer.variables()) + self.evaluate(state_ops.assign(m_bias_slot, [-2.])) + # Immediate restoration + root_checkpointable.restore(save_path=save_path).assert_consumed() + self.assertAllEqual([42.], self.evaluate(network._named_dense.variables[1])) + self.assertAllEqual(1, self.evaluate(root_checkpointable.save_counter)) + self.assertAllEqual([1.5], self.evaluate(m_bias_slot)) + with ops.Graph().as_default(): + on_create_network = MyNetwork() + on_create_optimizer = CheckpointableAdam(0.001) + on_create_root = Checkpoint( + optimizer=on_create_optimizer, network=on_create_network) + with self.test_session(graph=ops.get_default_graph()): + # Deferred restoration + status = on_create_root.restore(save_path=save_path) + on_create_network(constant_op.constant([[3.]])) # create variables + self.assertAllEqual(1, self.evaluate(on_create_root.save_counter)) + self.assertAllEqual([42.], + self.evaluate( + on_create_network._named_dense.variables[1])) + on_create_m_bias_slot = on_create_optimizer.get_slot( + on_create_network._named_dense.variables[1], "m") + # Optimizer slot variables are created when the original variable is + # restored. + self.assertAllEqual([1.5], self.evaluate(on_create_m_bias_slot)) + self.assertAllEqual(optimizer_variables[2:], + self.evaluate(on_create_optimizer.variables())) + on_create_optimizer._create_slots( + [resource_variable_ops.ResourceVariable([1.])]) + status.assert_consumed() + beta1_power, beta2_power = on_create_optimizer._get_beta_accumulators() + self.assertAllEqual(optimizer_variables[0], self.evaluate(beta1_power)) + self.assertAllEqual(optimizer_variables[1], self.evaluate(beta2_power)) + + def testDeferredRestorationUsageEager(self): + """An idiomatic eager execution example.""" + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + for training_continuation in range(3): + network = MyNetwork() + optimizer = CheckpointableAdam(0.001) + root = Checkpoint( + optimizer=optimizer, network=network, + optimizer_step=training_util.get_or_create_global_step()) + root.restore(core_saver.latest_checkpoint(checkpoint_directory)) + for _ in range(num_training_steps): + # TODO(allenl): Use a Dataset and serialize/checkpoint it. + input_value = constant_op.constant([[3.]]) + optimizer.minimize( + lambda: network(input_value), # pylint: disable=cell-var-from-loop + global_step=root.optimizer_step) + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + root.optimizer_step.numpy()) + + def testUsageGraph(self): + """Expected usage when graph building.""" + with context.graph_mode(): + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + for training_continuation in range(3): + with ops.Graph().as_default(): + network = MyNetwork() + optimizer = CheckpointableAdam(0.001) + root = Checkpoint( + optimizer=optimizer, network=network, + global_step=training_util.get_or_create_global_step()) + input_value = constant_op.constant([[3.]]) + train_op = optimizer.minimize( + network(input_value), + global_step=root.global_step) + root.save_counter # pylint: disable=pointless-statement + init_op = variables.global_variables_initializer() + checkpoint_path = core_saver.latest_checkpoint(checkpoint_directory) + with self.test_session(graph=ops.get_default_graph()) as session: + if checkpoint_path is None: + self.assertEqual(0, training_continuation) + session.run(init_op) + # Another alternative would be to run initializers automatically + # if no checkpoint is being loaded. This would make deferred + # loading a bit more useful with graph execution. + else: + checkpointable_utils.restore( + save_path=checkpoint_path, + root_checkpointable=root, + session=session) + for _ in range(num_training_steps): + session.run(train_op) + root.save(file_prefix=checkpoint_prefix, + session=session) + self.assertEqual((training_continuation + 1) * num_training_steps, + session.run(root.global_step)) + self.assertEqual(training_continuation + 1, + session.run(root.save_counter)) + + def _get_checkpoint_name(self, name): + root = checkpointable.Checkpointable() + checkpointable_utils.add_variable( + root, name=name, shape=[1, 2], dtype=dtypes.float64) + named_variables, _ = checkpointable_utils._serialize_object_graph(root) + checkpoint_name, = named_variables.keys() + with ops.name_scope("root/" + checkpoint_name): + pass # Make sure we can use this as an op name if we prefix it. + return checkpoint_name + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testVariableNameEscaping(self): + suffix = "/.ATTRIBUTES/VARIABLE_VALUE" + self.assertEqual(r"a.Sb.Sc" + suffix, self._get_checkpoint_name(r"a/b/c")) + self.assertEqual(r"b" + suffix, self._get_checkpoint_name(r"b")) + self.assertEqual(r"c.S" + suffix, self._get_checkpoint_name(r"c/")) + self.assertEqual(r"d.S..S" + suffix, self._get_checkpoint_name(r"d/.S")) + self.assertEqual(r"d.S..ATTRIBUTES.Sf" + suffix, + self._get_checkpoint_name(r"d/.ATTRIBUTES/f")) + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testNumberedPath(self): + root = checkpointable.Checkpointable() + leaf = checkpointable.Checkpointable() + root.leaf = leaf + checkpointable_utils.add_variable(leaf, name="v", shape=[]) + named_variables, _ = checkpointable_utils._serialize_object_graph(root) + variable_name, = named_variables.keys() + self.assertEqual(r"leaf/v/.ATTRIBUTES/VARIABLE_VALUE", variable_name) + + @test_util.run_in_graph_and_eager_modes() + def testLocalNameValidation(self): + root = checkpointable.Checkpointable() + leaf = checkpointable.Checkpointable() + # Dots are escaped, which avoids conflicts with reserved names. + root._track_checkpointable(leaf, name=".ATTRIBUTES") + checkpointable_utils.add_variable(checkpointable=leaf, name="a", shape=[]) + named_variables, _ = checkpointable_utils._serialize_object_graph(root) + name, = named_variables.keys() + self.assertEqual(name, "..ATTRIBUTES/a/.ATTRIBUTES/VARIABLE_VALUE") + + @test_util.run_in_graph_and_eager_modes() + def testLateDependencyTracking(self): + + class Dependency(checkpointable.Checkpointable): + + def build(self): + self.var = checkpointable_utils.add_variable( + self, "var", initializer=0.) + + class LateDependencies(checkpointable.Checkpointable): + + def add_dep(self): + self.dep = Dependency() + self.dep.build() + + original = LateDependencies() + original.add_dep() + self.evaluate(state_ops.assign(original.dep.var, 123.)) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + save_path = checkpointable_utils.save(checkpoint_prefix, original) + load_into = LateDependencies() + status = checkpointable_utils.restore(save_path, load_into) + with self.assertRaises(AssertionError): + status.assert_consumed() + load_into.add_dep() + status.assert_consumed() + self.assertEqual(123., self.evaluate(load_into.dep.var)) + + @test_util.run_in_graph_and_eager_modes() + def testDepAfterVar(self): + + class Dependency(checkpointable.Checkpointable): + + def build(self): + self.var = checkpointable_utils.add_variable( + self, "var", initializer=0.) + + class DepAfterVar(checkpointable.Checkpointable): + + def add_dep(self): + dep = Dependency() + dep.build() + self.dep = dep + + dep_after_var = DepAfterVar() + dep_after_var.add_dep() + self.evaluate(state_ops.assign(dep_after_var.dep.var, -14.)) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + save_path = checkpointable_utils.save( + checkpoint_prefix, dep_after_var) + + loaded_dep_after_var = DepAfterVar() + status = checkpointable_utils.restore( + save_path, loaded_dep_after_var) + loaded_dep_after_var.add_dep() + status.assert_consumed() + self.assertEqual(-14., self.evaluate(loaded_dep_after_var.dep.var)) + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testDeferredSlotRestoration(self): + checkpoint_directory = self.get_temp_dir() + + root = checkpointable.Checkpointable() + root.var = checkpointable_utils.add_variable( + root, name="var", initializer=0.) + optimizer = CheckpointableAdam(0.1) + if context.in_graph_mode(): + train_op = optimizer.minimize(root.var) + self.evaluate(variables.global_variables_initializer()) + self.evaluate(train_op) + else: + optimizer.minimize(root.var.read_value) + self.evaluate(state_ops.assign(root.var, 12.)) + no_slots_path = checkpointable_utils.save( + os.path.join(checkpoint_directory, "no_slots"), root) + root.optimizer = optimizer + self.evaluate(state_ops.assign(root.var, 13.)) + self.evaluate(state_ops.assign(optimizer.get_slot(name="m", var=root.var), + 14.)) + slots_path = checkpointable_utils.save( + os.path.join(checkpoint_directory, "with_slots"), root) + new_root = checkpointable.Checkpointable() + # Load the slot-containing checkpoint (deferred), then immediately overwrite + # the non-slot variable (also deferred). + slot_status = checkpointable_utils.restore( + slots_path, new_root) + no_slot_status = checkpointable_utils.restore( + no_slots_path, new_root) + with self.assertRaises(AssertionError): + no_slot_status.assert_consumed() + new_root.var = checkpointable_utils.add_variable( + new_root, name="var", shape=[]) + self.assertEqual(12., self.evaluate(new_root.var)) + no_slot_status.assert_consumed() + new_root.optimizer = CheckpointableAdam(0.1) + with self.assertRaisesRegexp(AssertionError, "beta1_power"): + slot_status.assert_consumed() + self.assertEqual(12., self.evaluate(new_root.var)) + self.assertEqual(14., self.evaluate( + new_root.optimizer.get_slot(name="m", var=new_root.var))) + if context.in_graph_mode(): + train_op = new_root.optimizer.minimize(new_root.var) + self.evaluate(train_op) + else: + new_root.optimizer.minimize(new_root.var.read_value) + slot_status.assert_consumed() + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testOverlappingRestores(self): + checkpoint_directory = self.get_temp_dir() + save_root = checkpointable.Checkpointable() + save_root.dep = checkpointable.Checkpointable() + save_root.dep.var = checkpointable_utils.add_variable( + save_root.dep, name="var", initializer=0.) + self.evaluate(state_ops.assign(save_root.dep.var, 12.)) + first_path = checkpointable_utils.save( + os.path.join(checkpoint_directory, "first"), save_root) + self.evaluate(state_ops.assign(save_root.dep.var, 13.)) + second_path = checkpointable_utils.save( + os.path.join(checkpoint_directory, "second"), save_root) + + first_root = checkpointable.Checkpointable() + second_root = checkpointable.Checkpointable() + first_status = checkpointable_utils.restore( + first_path, first_root) + second_status = checkpointable_utils.restore( + second_path, second_root) + load_dep = checkpointable.Checkpointable() + load_dep.var = checkpointable_utils.add_variable( + load_dep, name="var", shape=[]) + first_root.dep = load_dep + first_status.assert_consumed() + self.assertEqual(12., self.evaluate(load_dep.var)) + second_root.dep = load_dep + second_status.assert_consumed() + self.assertEqual(13., self.evaluate(load_dep.var)) + + # Try again with the order of the restore() reversed. The last restore + # determines the final value. + first_root = checkpointable.Checkpointable() + second_root = checkpointable.Checkpointable() + second_status = checkpointable_utils.restore( + second_path, second_root) + first_status = checkpointable_utils.restore( + first_path, first_root) + load_dep = checkpointable.Checkpointable() + load_dep.var = checkpointable_utils.add_variable( + load_dep, name="var", shape=[]) + first_root.dep = load_dep + first_status.assert_consumed() + self.assertEqual(12., self.evaluate(load_dep.var)) + second_root.dep = load_dep + second_status.assert_consumed() + self.assertEqual(12., self.evaluate(load_dep.var)) + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testAmbiguousLoad(self): + # Not OK to split one checkpoint object into two + checkpoint_directory = self.get_temp_dir() + save_root = checkpointable.Checkpointable() + save_root.dep_one = checkpointable.Checkpointable() + save_root.dep_two = checkpointable.Checkpointable() + dep_three = checkpointable.Checkpointable() + save_root.dep_one.dep_three = dep_three + save_root.dep_two.dep_three = dep_three + checkpointable_utils.add_variable(dep_three, name="var", initializer=0.) + self.evaluate(variables.global_variables_initializer()) + save_path = checkpointable_utils.save( + os.path.join(checkpoint_directory, "ckpt"), save_root) + load_root = checkpointable.Checkpointable() + checkpointable_utils.restore(save_path, load_root) + load_root.dep_one = checkpointable.Checkpointable() + load_root.dep_two = checkpointable.Checkpointable() + load_root.dep_one.dep_three = checkpointable.Checkpointable() + with self.assertRaisesRegexp(AssertionError, + "resolved to different objects"): + load_root.dep_two.dep_three = checkpointable.Checkpointable() + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testObjectsCombined(self): + # Currently fine to load two checkpoint objects into one Python object + checkpoint_directory = self.get_temp_dir() + save_root = checkpointable.Checkpointable() + save_root.dep_one = checkpointable.Checkpointable() + save_root.dep_two = checkpointable.Checkpointable() + checkpointable_utils.add_variable( + save_root.dep_one, name="var1", initializer=32., dtype=dtypes.float64) + checkpointable_utils.add_variable( + save_root.dep_two, name="var2", initializer=64., dtype=dtypes.float64) + self.evaluate(variables.global_variables_initializer()) + save_path = checkpointable_utils.save( + os.path.join(checkpoint_directory, "ckpt"), save_root) + load_root = checkpointable.Checkpointable() + load_root.dep_one = checkpointable.Checkpointable() + load_root.dep_two = load_root.dep_one + v1 = checkpointable_utils.add_variable( + load_root.dep_one, name="var1", shape=[], dtype=dtypes.float64) + v2 = checkpointable_utils.add_variable( + load_root.dep_one, name="var2", shape=[], dtype=dtypes.float64) + checkpointable_utils.restore(save_path, load_root).assert_consumed() + self.assertEqual(32., self.evaluate(v1)) + self.assertEqual(64., self.evaluate(v2)) + + @test_util.run_in_graph_and_eager_modes() + def testDependencyLoop(self): + # Note: this test creates garbage during eager execution because it + # purposefully creates a reference cycle. + first = checkpointable.Checkpointable() + second = checkpointable.Checkpointable() + first.second = second + second.first = first + first.v = checkpointable_utils.add_variable( + first, "v1", initializer=[3., 1., 4.]) + second.v = checkpointable_utils.add_variable( + second, "v2", initializer=[1., 1., 2., 3.]) + self.evaluate(variables.global_variables_initializer()) + checkpoint_directory = self.get_temp_dir() + save_path = checkpointable_utils.save( + os.path.join(checkpoint_directory, "ckpt"), first) + + # Test deferred loading + first_load = checkpointable.Checkpointable() + status = checkpointable_utils.restore(save_path, first_load) + second_load = checkpointable.Checkpointable() + first_load.second = second_load + second_load.first = first_load + with self.assertRaises(AssertionError): + status.assert_consumed() + first_load.v = checkpointable_utils.add_variable( + first_load, "v1", shape=[3]) + second_load.v = checkpointable_utils.add_variable( + second_load, "v2", shape=[4]) + status.assert_consumed() + self.assertAllEqual([3., 1., 4.], self.evaluate(first_load.v)) + self.assertAllEqual([1., 1., 2., 3.], self.evaluate(second_load.v)) + + # Test loading when variables have already been created + self.evaluate(first_load.v.assign([2., 7., 1.])) + self.assertAllEqual([2., 7., 1.], self.evaluate(first_load.v)) + self.evaluate(second_load.v.assign([2., 7., 1., 8.])) + self.assertAllEqual([2., 7., 1., 8.], self.evaluate(second_load.v)) + checkpointable_utils.restore( + save_path, first_load).assert_consumed() + self.assertAllEqual([3., 1., 4.], self.evaluate(first_load.v)) + self.assertAllEqual([1., 1., 2., 3.], self.evaluate(second_load.v)) + + @test_util.run_in_graph_and_eager_modes() + def testRestoreOnAssign(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + save_graph = ops.Graph() + with save_graph.as_default(), self.test_session(save_graph): + first = checkpointable.Checkpointable() + first.var1 = variable_scope.get_variable( + name="outside_var", initializer=0.) + first.var2 = variable_scope.get_variable( + name="blah", initializer=0.) + self.evaluate(first.var1.assign(4.)) + self.evaluate(first.var2.assign(8.)) + save_path = checkpointable_utils.save( + checkpoint_prefix, root_checkpointable=first) + restore_graph = ops.Graph() + with restore_graph.as_default(), self.test_session(restore_graph): + second = checkpointable.Checkpointable() + second.var2 = variable_scope.get_variable( + name="blah", initializer=0.) + checkpointable_utils.restore(save_path, root_checkpointable=second) + recreated_var1 = variable_scope.get_variable( + name="outside_var", initializer=0.) + self.assertEqual(8., self.evaluate(second.var2)) + self.evaluate(recreated_var1.assign(-2.)) + self.assertEqual(-2., self.evaluate(recreated_var1)) + second.var1 = recreated_var1 + self.assertEqual(4., self.evaluate(recreated_var1)) + + # TODO(allenl): Saver class that doesn't pollute the graph with constants. + @unittest.skip("todo") + def testManySavesGraph(self): + """Saves after the first should not modify the graph.""" + with context.graph_mode(): + graph = ops.Graph() + with graph.as_default(), self.test_session(graph): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + obj = checkpointable.Checkpointable() + obj.var = variable_scope.get_variable(name="v", initializer=0.) + obj.opt = CheckpointableAdam(0.1) + obj.opt.minimize(obj.var.read_value()) + self.evaluate(variables.global_variables_initializer()) + checkpointable_utils.save( + checkpoint_prefix, root_checkpointable=obj) + before_ops = graph.get_operations() + checkpointable_utils.save( + checkpoint_prefix, root_checkpointable=obj) + self.assertEqual(before_ops, graph.get_operations()) + + @unittest.skip("todo") + def testManyRestoresGraph(self): + """Restores after the first should not modify the graph.""" + with context.graph_mode(): + graph = ops.Graph() + with graph.as_default(), self.test_session(graph): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + obj = checkpointable.Checkpointable() + obj.var = variable_scope.get_variable(name="v", initializer=0.) + obj.opt = CheckpointableAdam(0.1) + obj.opt.minimize(obj.var.read_value()) + self.evaluate(variables.global_variables_initializer()) + save_path = checkpointable_utils.save( + checkpoint_prefix, root_checkpointable=obj) + checkpointable_utils.restore( + save_path, root_checkpointable=obj) + before_ops = graph.get_operations() + checkpointable_utils.restore( + save_path, root_checkpointable=obj) + self.assertEqual(before_ops, graph.get_operations()) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/eager/python/datasets.py b/tensorflow/contrib/eager/python/datasets.py index b559cce6b12a809d671ce7855680063f02a4ac22..d177bfeab2d1fdc05d7ced54df8723fae2c77fdb 100644 --- a/tensorflow/contrib/eager/python/datasets.py +++ b/tensorflow/contrib/eager/python/datasets.py @@ -23,6 +23,7 @@ import threading from tensorflow.contrib.data.python.ops import prefetching_ops from tensorflow.python.data.ops import iterator_ops from tensorflow.python.data.util import nest +from tensorflow.python.data.util import sparse from tensorflow.python.eager import context from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -41,7 +42,7 @@ def _generate_shared_name(prefix): global _uid_counter uid = _uid_counter _uid_counter += 1 - return "{}_{}".format(prefix, uid) + return "{}{}".format(prefix, uid) class Iterator(object): @@ -75,13 +76,16 @@ class Iterator(object): format(type(self))) with ops.device("/device:CPU:0"): ds_variant = dataset._as_variant_tensor() # pylint: disable=protected-access + self._output_classes = dataset.output_classes self._output_types = dataset.output_types self._output_shapes = dataset.output_shapes - self._flat_output_types = nest.flatten(dataset.output_types) - self._flat_output_shapes = nest.flatten(dataset.output_shapes) + self._flat_output_types = nest.flatten( + sparse.as_dense_types(self._output_types, self._output_classes)) + self._flat_output_shapes = nest.flatten( + sparse.as_dense_shapes(self._output_shapes, self._output_classes)) self._resource = gen_dataset_ops.iterator( - container="", - shared_name=_generate_shared_name("eager_iterator"), + shared_name="", + container=_generate_shared_name("eageriterator"), output_types=self._flat_output_types, output_shapes=self._flat_output_shapes) gen_dataset_ops.make_iterator(ds_variant, self._resource) @@ -108,7 +112,7 @@ class Iterator(object): remote_fn.add_to_graph(None) target = constant_op.constant("/device:CPU:0") with ops.device(self._device): - self._buffer_resource_handle = prefetching_ops.function_buffering_resource( + self._buffer_resource_handle = prefetching_ops.function_buffering_resource( # pylint: disable=line-too-long string_arg=iter_string_handle, f=remote_fn, target_device=target, @@ -116,8 +120,9 @@ class Iterator(object): thread_pool_size=1, container="", shared_name=_generate_shared_name("function_buffer_resource")) - self._buffer_resource_deleter = resource_variable_ops.EagerResourceDeleter( - handle=self._buffer_resource_handle, handle_device=self._device) + self._buffer_resource_deleter = resource_variable_ops.EagerResourceDeleter( # pylint: disable=line-too-long + handle=self._buffer_resource_handle, + handle_device=self._device) def __iter__(self): return self @@ -125,22 +130,83 @@ class Iterator(object): def __next__(self): # For Python 3 compatibility return self.next() - def next(self): - """Return the next tf.Tensor from the dataset.""" + def _next_internal(self): + """Returns a nested structure of `tf.Tensor`s containing the next element. + """ with ops.device(self._device): - try: - if self._buffer_resource_handle is not None: - ret = prefetching_ops.function_buffering_resource_get_next( - function_buffer_resource=self._buffer_resource_handle, - output_types=self._flat_output_types) - else: - # TODO(ashankar): Consider removing this ops.device() contextmanager - # and instead mimic ops placement in graphs: Operations on resource - # handles execute on the same device as where the resource is placed. - ret = gen_dataset_ops.iterator_get_next( - self._resource, - output_types=self._flat_output_types, - output_shapes=self._flat_output_shapes) - except errors.OutOfRangeError: - raise StopIteration - return nest.pack_sequence_as(self._output_types, ret) + if self._buffer_resource_handle is not None: + ret = prefetching_ops.function_buffering_resource_get_next( + function_buffer_resource=self._buffer_resource_handle, + output_types=self._flat_output_types) + else: + # TODO(ashankar): Consider removing this ops.device() contextmanager + # and instead mimic ops placement in graphs: Operations on resource + # handles execute on the same device as where the resource is placed. + # NOTE(mrry): Here we use the "_sync" variant of `iterator_get_next` + # because in eager mode this code will run synchronously on the calling + # thread. Therefore we do not need to make a defensive context switch + # to a background thread, and can achieve a small constant performance + # boost by invoking the iterator synchronously. + ret = gen_dataset_ops.iterator_get_next_sync( + self._resource, + output_types=self._flat_output_types, + output_shapes=self._flat_output_shapes) + + return sparse.deserialize_sparse_tensors( + nest.pack_sequence_as(self._output_types, ret), self._output_types, + self._output_shapes, self._output_classes) + + def next(self): + """Returns a nested structure of `tf.Tensor`s containing the next element. + """ + try: + return self._next_internal() + except errors.OutOfRangeError: + raise StopIteration + + @property + def output_classes(self): + """Returns the class of each component of an element of this iterator. + + The expected values are `tf.Tensor` and `tf.SparseTensor`. + + Returns: + A nested structure of Python `type` objects corresponding to each + component of an element of this dataset. + """ + return self._output_classes + + @property + def output_shapes(self): + """Returns the shape of each component of an element of this iterator. + + Returns: + A nested structure of `tf.TensorShape` objects corresponding to each + component of an element of this dataset. + """ + return self._output_shapes + + @property + def output_types(self): + """Returns the type of each component of an element of this iterator. + + Returns: + A nested structure of `tf.DType` objects corresponding to each component + of an element of this dataset. + """ + return self._output_types + + def get_next(self, name=None): + """Returns a nested structure of `tf.Tensor`s containing the next element. + + Args: + name: (Optional.) A name for the created operation. Currently unused. + + Returns: + A nested structure of `tf.Tensor` objects. + + Raises: + `tf.errors.OutOfRangeError`: If the end of the dataset has been reached. + """ + del name + return self._next_internal() diff --git a/tensorflow/contrib/eager/python/datasets_test.py b/tensorflow/contrib/eager/python/datasets_test.py index c924d81c9d85e638e4f35f260664c0ee7d03257e..a1611e92b113839c2dd2a3b2560b0ba90c0a7ef0 100644 --- a/tensorflow/contrib/eager/python/datasets_test.py +++ b/tensorflow/contrib/eager/python/datasets_test.py @@ -16,11 +16,19 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import time + +import numpy as np + +from tensorflow.contrib import lookup from tensorflow.contrib.eager.python import datasets from tensorflow.python.data import Dataset from tensorflow.python.eager import test +from tensorflow.python.framework import constant_op 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.ops import math_ops from tensorflow.python.ops import script_ops @@ -33,6 +41,15 @@ class IteratorTest(test.TestCase): got.append(t.numpy()) self.assertAllEqual([0, 1, 2, 3], got) + def testGetNext(self): + iterator = datasets.Iterator(Dataset.range(4)) + self.assertEqual(0, iterator.get_next().numpy()) + self.assertEqual(1, iterator.get_next().numpy()) + self.assertEqual(2, iterator.get_next().numpy()) + self.assertEqual(3, iterator.get_next().numpy()) + with self.assertRaises(errors.OutOfRangeError): + iterator.get_next() + def testMultipleIteratorsOnTheSameDataset(self): ds = Dataset.range(4) it1 = datasets.Iterator(ds) @@ -64,6 +81,18 @@ class IteratorTest(test.TestCase): got = [x.numpy() for x in it] self.assertAllEqual([0, 4, 16, 36], got) + def testMapCaptureLookupTable(self): + default_val = -1 + keys = constant_op.constant(['brain', 'salad', 'surgery']) + values = constant_op.constant([0, 1, 2], dtypes.int64) + table = lookup.HashTable( + lookup.KeyValueTensorInitializer(keys, values), default_val) + dataset = Dataset.from_tensor_slices(['brain', 'salad', 'surgery']) + dataset = dataset.map(table.lookup) + it = datasets.Iterator(dataset) + got = [x.numpy() for x in it] + self.assertAllEqual([0, 1, 2], got) + def testMultipleIteratorsOnADatasetThatUsesFunctions(self): ds = Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6]).map(math_ops.square) @@ -72,6 +101,53 @@ class IteratorTest(test.TestCase): got2 = [x.numpy() for x in datasets.Iterator(ds)] self.assertAllEqual(got1, got2) + def assertSparseValuesEqual(self, a, b): + self.assertAllEqual(a.indices, b.indices) + self.assertAllEqual(a.values, b.values) + self.assertAllEqual(a.dense_shape, b.dense_shape) + + def testSparseTensorElements(self): + components = (sparse_tensor.SparseTensorValue( + indices=np.array([[0, 0], [1, 0], [2, 0]]), + values=np.array([0, 0, 0]), + dense_shape=np.array([3, 1])), + sparse_tensor.SparseTensorValue( + indices=np.array([[0, 0], [1, 1], [2, 2]]), + values=np.array([1, 2, 3]), + dense_shape=np.array([3, 3]))) + + expected = [ + (sparse_tensor.SparseTensorValue( + indices=np.array([[0]]), + values=np.array([0]), + dense_shape=np.array([1])), + sparse_tensor.SparseTensorValue( + indices=np.array([[0]]), + values=np.array([1]), + dense_shape=np.array([3]))), + (sparse_tensor.SparseTensorValue( + indices=np.array([[0]]), + values=np.array([0]), + dense_shape=np.array([1])), + sparse_tensor.SparseTensorValue( + indices=np.array([[1]]), + values=np.array([2]), + dense_shape=np.array([3]))), + (sparse_tensor.SparseTensorValue( + indices=np.array([[0]]), + values=np.array([0]), + dense_shape=np.array([1])), + sparse_tensor.SparseTensorValue( + indices=np.array([[2]]), + values=np.array([3]), + dense_shape=np.array([3]))), + ] + + for i, result in enumerate( + datasets.Iterator(Dataset.from_tensor_slices(components))): + self.assertSparseValuesEqual(expected[i][0], result[0]) + self.assertSparseValuesEqual(expected[i][1], result[1]) + def testPyFunc(self): def my_map(inp): @@ -90,5 +166,64 @@ class IteratorTest(test.TestCase): self.assertAllEqual([0., 2.], x.numpy()) +class DatasetConstructorBenchmark(test.Benchmark): + + def benchmarkSliceRepeatBatchEager(self): + input_size = 10000 + batch_size = 100 + num_epochs = 100 + + input_data = np.random.randn(input_size) + + dataset = ( + Dataset.from_tensor_slices(input_data).repeat(num_epochs) + .batch(batch_size)) + iterator = datasets.Iterator(dataset) + + ends = [time.time()] + for _ in iterator: + ends.append(time.time()) + + deltas = np.ediff1d(ends) + median_wall_time = np.median(deltas) + print( + 'Slice/repeat/batch eager input size: %d batch size: %d Median wall ' + 'time per element: %f' + % (input_size, batch_size, median_wall_time)) + self.report_benchmark( + iters=len(deltas), + wall_time=median_wall_time, + name='benchmark_slice_repeat_batch_eager_input_%d_batch_%d' % + (input_size, batch_size)) + + def benchmarkSliceBatchCacheRepeatCallable(self): + input_size = 10000 + batch_size = 100 + num_epochs = 100 + + input_data = np.random.randn(input_size) + + dataset = ( + Dataset.from_tensor_slices(input_data).batch(batch_size).cache() + .repeat(num_epochs)) + iterator = datasets.Iterator(dataset) + + ends = [time.time()] + for _ in iterator: + ends.append(time.time()) + + deltas = np.ediff1d(ends) + median_wall_time = np.median(deltas) + print( + 'Slice/batch/cache/repeat eager input size: %d batch size: %d Median ' + 'wall time per element: %f' + % (input_size, batch_size, median_wall_time)) + self.report_benchmark( + iters=len(deltas), + wall_time=median_wall_time, + name='benchmark_slice_batch_cache_repeat_eager_input_%d_batch_%d' % + (input_size, batch_size)) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/eager/python/evaluator.py b/tensorflow/contrib/eager/python/evaluator.py index 3faaeef5903615ea122800a6690117dde682e830..68e7b5421fec7f73f10e381ca45f9d900de299d7 100644 --- a/tensorflow/contrib/eager/python/evaluator.py +++ b/tensorflow/contrib/eager/python/evaluator.py @@ -178,7 +178,7 @@ class Evaluator(object): call_op: An op that updates evaluation state on a mini-batch of examples. Must generate an tf.errors.OutOfRangeError when done. results_op: A dictionary of tensors that compute the final evaluation - results from the evaulation state. + results from the evaluation state. sess: The Session to run the evaluation in. Defaults to the default Session. diff --git a/tensorflow/contrib/eager/python/examples/BUILD b/tensorflow/contrib/eager/python/examples/BUILD index 6aef010a2139c4cd2ae19c008aa21d4e3592ca98..15a21885f66eface291a39fa0ee1ff28bc297548 100644 --- a/tensorflow/contrib/eager/python/examples/BUILD +++ b/tensorflow/contrib/eager/python/examples/BUILD @@ -6,6 +6,7 @@ package(default_visibility = ["//tensorflow:internal"]) py_library( name = "examples_pip", deps = [ + "//tensorflow/contrib/eager/python/examples/gan:mnist", "//tensorflow/contrib/eager/python/examples/linear_regression", "//tensorflow/contrib/eager/python/examples/mnist", "//tensorflow/contrib/eager/python/examples/resnet50", diff --git a/tensorflow/contrib/eager/python/examples/gan/BUILD b/tensorflow/contrib/eager/python/examples/gan/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..c61ec2dbae60a782c0e6589701554b045dcb92ae --- /dev/null +++ b/tensorflow/contrib/eager/python/examples/gan/BUILD @@ -0,0 +1,36 @@ +licenses(["notice"]) # Apache 2.0 + +package(default_visibility = ["//tensorflow:internal"]) + +load("//tensorflow:tensorflow.bzl", "cuda_py_test") + +py_binary( + name = "mnist", + srcs = ["mnist.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow:tensorflow_py", + "//tensorflow/contrib/eager/python:tfe", + "//tensorflow/examples/tutorials/mnist:input_data", + ], +) + +cuda_py_test( + name = "mnist_test", + srcs = ["mnist_test.py"], + additional_deps = [ + ":mnist", + "//tensorflow/contrib/eager/python:tfe", + "//tensorflow:tensorflow_py", + ], +) + +cuda_py_test( + name = "mnist_graph_test", + srcs = ["mnist_graph_test.py"], + additional_deps = [ + ":mnist", + "//third_party/py/numpy", + "//tensorflow:tensorflow_py", + ], +) diff --git a/tensorflow/contrib/eager/python/examples/gan/README.md b/tensorflow/contrib/eager/python/examples/gan/README.md new file mode 100644 index 0000000000000000000000000000000000000000..208a64b05d47eea10b49a1bf967a5453677bfd21 --- /dev/null +++ b/tensorflow/contrib/eager/python/examples/gan/README.md @@ -0,0 +1,38 @@ +# GAN with TensorFlow eager execution + +A simple Generative Adversarial Network (GAN) example using eager execution. +The discriminator and generator networks each contain a few convolution and +fully connected layers. + +Other eager execution examples can be found under the parent directory. + +## Content + +- `mnist.py`: Model definitions and training routines. +- `mnist_test.py`: Benchmarks for training and using the models using eager +execution. +- `mnist_graph_test.py`: Benchmarks for training and using the models using +graph execution. The same model definitions and loss functions are used in +all benchmarks. + + +## To run + +- Make sure you have installed TensorFlow 1.5+ or the latest `tf-nightly` +or `tf-nightly-gpu` pip package in order to access the eager execution feature. + +- Train model. E.g., + + ```bash + python mnist.py + ``` + + Use `--output_dir=` to direct the script to save TensorBoard summaries + during training. Disabled by default. + + Use `--checkpoint_dir=` to direct the script to save checkpoints to + `` during training. DIR defaults to /tmp/tensorflow/mnist/checkpoints/. + The script will load the latest saved checkpoint from this directory if + one exists. + + Use `-h` for other options. diff --git a/tensorflow/contrib/eager/python/examples/gan/mnist.py b/tensorflow/contrib/eager/python/examples/gan/mnist.py new file mode 100644 index 0000000000000000000000000000000000000000..b9ac79f46c83bb709918e3b72830b90ddcfd71b4 --- /dev/null +++ b/tensorflow/contrib/eager/python/examples/gan/mnist.py @@ -0,0 +1,368 @@ +# Copyright 2017 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. +# ============================================================================== +"""A deep MNIST classifier using convolutional layers. + +Sample usage: + python mnist.py --help +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import argparse +import os +import sys +import time + +import tensorflow as tf + +import tensorflow.contrib.eager as tfe +from tensorflow.examples.tutorials.mnist import input_data + +FLAGS = None + + +class Discriminator(tfe.Network): + """GAN Discriminator. + + A network to differentiate between generated and real handwritten digits. + """ + + def __init__(self, data_format): + """Creates a model for discriminating between real and generated digits. + + Args: + data_format: Either 'channels_first' or 'channels_last'. + 'channels_first' is typically faster on GPUs while 'channels_last' is + typically faster on CPUs. See + https://www.tensorflow.org/performance/performance_guide#data_formats + """ + super(Discriminator, self).__init__(name='') + if data_format == 'channels_first': + self._input_shape = [-1, 1, 28, 28] + else: + assert data_format == 'channels_last' + self._input_shape = [-1, 28, 28, 1] + self.conv1 = self.track_layer(tf.layers.Conv2D(64, 5, padding='SAME', + data_format=data_format, + activation=tf.tanh)) + self.pool1 = self.track_layer( + tf.layers.AveragePooling2D(2, 2, data_format=data_format)) + self.conv2 = self.track_layer(tf.layers.Conv2D(128, 5, + data_format=data_format, + activation=tf.tanh)) + self.pool2 = self.track_layer( + tf.layers.AveragePooling2D(2, 2, data_format=data_format)) + self.flatten = self.track_layer(tf.layers.Flatten()) + self.fc1 = self.track_layer(tf.layers.Dense(1024, activation=tf.tanh)) + self.fc2 = self.track_layer(tf.layers.Dense(1, activation=None)) + + def call(self, inputs): + """Return two logits per image estimating input authenticity. + + Users should invoke __call__ to run the network, which delegates to this + method (and not call this method directly). + + Args: + inputs: A batch of images as a Tensor with shape [batch_size, 28, 28, 1] + or [batch_size, 1, 28, 28] + + Returns: + A Tensor with shape [batch_size] containing logits estimating + the probability that corresponding digit is real. + """ + x = tf.reshape(inputs, self._input_shape) + x = self.conv1(x) + x = self.pool1(x) + x = self.conv2(x) + x = self.pool2(x) + x = self.flatten(x) + x = self.fc1(x) + x = self.fc2(x) + return x + + +class Generator(tfe.Network): + """Generator of handwritten digits similar to the ones in the MNIST dataset. + """ + + def __init__(self, data_format): + """Creates a model for discriminating between real and generated digits. + + Args: + data_format: Either 'channels_first' or 'channels_last'. + 'channels_first' is typically faster on GPUs while 'channels_last' is + typically faster on CPUs. See + https://www.tensorflow.org/performance/performance_guide#data_formats + """ + super(Generator, self).__init__(name='') + self.data_format = data_format + # We are using 128 6x6 channels as input to the first deconvolution layer + if data_format == 'channels_first': + self._pre_conv_shape = [-1, 128, 6, 6] + else: + assert data_format == 'channels_last' + self._pre_conv_shape = [-1, 6, 6, 128] + self.fc1 = self.track_layer(tf.layers.Dense(6 * 6 * 128, + activation=tf.tanh)) + + # In call(), we reshape the output of fc1 to _pre_conv_shape + + # Deconvolution layer. Resulting image shape: (batch, 14, 14, 64) + self.conv1 = self.track_layer(tf.layers.Conv2DTranspose( + 64, 4, strides=2, activation=None, data_format=data_format)) + + # Deconvolution layer. Resulting image shape: (batch, 28, 28, 1) + self.conv2 = self.track_layer(tf.layers.Conv2DTranspose( + 1, 2, strides=2, activation=tf.nn.sigmoid, data_format=data_format)) + + def call(self, inputs): + """Return a batch of generated images. + + Users should invoke __call__ to run the network, which delegates to this + method (and not call this method directly). + + Args: + inputs: A batch of noise vectors as a Tensor with shape + [batch_size, length of noise vectors]. + + Returns: + A Tensor containing generated images. If data_format is 'channels_last', + the shape of returned images is [batch_size, 28, 28, 1], else + [batch_size, 1, 28, 28] + """ + + x = self.fc1(inputs) + x = tf.reshape(x, shape=self._pre_conv_shape) + x = self.conv1(x) + x = self.conv2(x) + return x + + +def discriminator_loss(discriminator_real_outputs, discriminator_gen_outputs): + """Original discriminator loss for GANs, with label smoothing. + + See `Generative Adversarial Nets` (https://arxiv.org/abs/1406.2661) for more + details. + + Args: + discriminator_real_outputs: Discriminator output on real data. + discriminator_gen_outputs: Discriminator output on generated data. Expected + to be in the range of (-inf, inf). + + Returns: + A scalar loss Tensor. + """ + + loss_on_real = tf.losses.sigmoid_cross_entropy( + tf.ones_like(discriminator_real_outputs), discriminator_real_outputs, + label_smoothing=0.25) + loss_on_generated = tf.losses.sigmoid_cross_entropy( + tf.zeros_like(discriminator_gen_outputs), discriminator_gen_outputs) + loss = loss_on_real + loss_on_generated + tf.contrib.summary.scalar('discriminator_loss', loss) + return loss + + +def generator_loss(discriminator_gen_outputs): + """Original generator loss for GANs. + + L = -log(sigmoid(D(G(z)))) + + See `Generative Adversarial Nets` (https://arxiv.org/abs/1406.2661) + for more details. + + Args: + discriminator_gen_outputs: Discriminator output on generated data. Expected + to be in the range of (-inf, inf). + + Returns: + A scalar loss Tensor. + """ + loss = tf.losses.sigmoid_cross_entropy( + tf.ones_like(discriminator_gen_outputs), discriminator_gen_outputs) + tf.contrib.summary.scalar('generator_loss', loss) + return loss + + +def train_one_epoch(generator, discriminator, + generator_optimizer, discriminator_optimizer, + dataset, log_interval, noise_dim): + """Trains `generator` and `discriminator` models on `dataset`. + + Args: + generator: Generator model. + discriminator: Discriminator model. + generator_optimizer: Optimizer to use for generator. + discriminator_optimizer: Optimizer to use for discriminator. + dataset: Dataset of images to train on. + log_interval: How many global steps to wait between logging and collecting + summaries. + noise_dim: Dimension of noise vector to use. + """ + + total_generator_loss = 0.0 + total_discriminator_loss = 0.0 + for (batch_index, images) in enumerate(tfe.Iterator(dataset)): + with tf.device('/cpu:0'): + tf.assign_add(tf.train.get_global_step(), 1) + + with tf.contrib.summary.record_summaries_every_n_global_steps(log_interval): + current_batch_size = images.shape[0] + noise = tf.random_uniform(shape=[current_batch_size, noise_dim], + minval=-1., maxval=1., seed=batch_index) + + with tfe.GradientTape(persistent=True) as g: + generated_images = generator(noise) + tf.contrib.summary.image('generated_images', + tf.reshape(generated_images, [-1, 28, 28, 1]), + max_images=10) + + discriminator_gen_outputs = discriminator(generated_images) + discriminator_real_outputs = discriminator(images) + discriminator_loss_val = discriminator_loss(discriminator_real_outputs, + discriminator_gen_outputs) + total_discriminator_loss += discriminator_loss_val + + generator_loss_val = generator_loss(discriminator_gen_outputs) + total_generator_loss += generator_loss_val + + generator_grad = g.gradient(generator_loss_val, generator.variables) + discriminator_grad = g.gradient(discriminator_loss_val, + discriminator.variables) + + with tf.variable_scope('generator'): + generator_optimizer.apply_gradients(zip(generator_grad, + generator.variables)) + with tf.variable_scope('discriminator'): + discriminator_optimizer.apply_gradients(zip(discriminator_grad, + discriminator.variables)) + + if log_interval and batch_index > 0 and batch_index % log_interval == 0: + print('Batch #%d\tAverage Generator Loss: %.6f\t' + 'Average Discriminator Loss: %.6f' % ( + batch_index, total_generator_loss/batch_index, + total_discriminator_loss/batch_index)) + + +def main(_): + (device, data_format) = ('/gpu:0', 'channels_first') + if FLAGS.no_gpu or tfe.num_gpus() <= 0: + (device, data_format) = ('/cpu:0', 'channels_last') + print('Using device %s, and data format %s.' % (device, data_format)) + + # Load the datasets + data = input_data.read_data_sets(FLAGS.data_dir) + dataset = (tf.data.Dataset + .from_tensor_slices(data.train.images) + .shuffle(60000) + .batch(FLAGS.batch_size)) + + # Create the models and optimizers + generator = Generator(data_format) + discriminator = Discriminator(data_format) + with tf.variable_scope('generator'): + generator_optimizer = tf.train.AdamOptimizer(FLAGS.lr) + with tf.variable_scope('discriminator'): + discriminator_optimizer = tf.train.AdamOptimizer(FLAGS.lr) + + # Prepare summary writer and checkpoint info + summary_writer = tf.contrib.summary.create_summary_file_writer( + FLAGS.output_dir, flush_millis=1000) + checkpoint_prefix = os.path.join(FLAGS.checkpoint_dir, 'ckpt') + latest_cpkt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir) + if latest_cpkt: + print('Using latest checkpoint at ' + latest_cpkt) + + with tf.device(device): + for epoch in range(1, 101): + with tfe.restore_variables_on_create(latest_cpkt): + global_step = tf.train.get_or_create_global_step() + start = time.time() + with summary_writer.as_default(): + train_one_epoch(generator, discriminator, generator_optimizer, + discriminator_optimizer, + dataset, FLAGS.log_interval, FLAGS.noise) + end = time.time() + print('\nTrain time for epoch #%d (global step %d): %f' % ( + epoch, global_step.numpy(), end - start)) + + all_variables = ( + generator.variables + + discriminator.variables + + generator_optimizer.variables() + + discriminator_optimizer.variables() + + [global_step]) + tfe.Saver(all_variables).save( + checkpoint_prefix, global_step=global_step) + + +if __name__ == '__main__': + tfe.enable_eager_execution() + + parser = argparse.ArgumentParser() + parser.add_argument( + '--data-dir', + type=str, + default='/tmp/tensorflow/mnist/input_data', + help=('Directory for storing input data (default ' + '/tmp/tensorflow/mnist/input_data)')) + parser.add_argument( + '--batch-size', + type=int, + default=128, + metavar='N', + help='input batch size for training (default: 128)') + parser.add_argument( + '--log-interval', + type=int, + default=100, + metavar='N', + help=('number of batches between logging and writing summaries ' + '(default: 100)')) + parser.add_argument( + '--output_dir', + type=str, + default=None, + metavar='DIR', + help='Directory to write TensorBoard summaries (defaults to none)') + parser.add_argument( + '--checkpoint_dir', + type=str, + default='/tmp/tensorflow/mnist/checkpoints/', + metavar='DIR', + help=('Directory to save checkpoints in (once per epoch) (default ' + '/tmp/tensorflow/mnist/checkpoints/)')) + parser.add_argument( + '--lr', + type=float, + default=0.001, + metavar='LR', + help='learning rate (default: 0.001)') + parser.add_argument( + '--noise', + type=int, + default=100, + metavar='N', + help='Length of noise vector for generator input (default: 100)') + parser.add_argument( + '--no-gpu', + action='store_true', + default=False, + help='disables GPU usage even if a GPU is available') + + FLAGS, unparsed = parser.parse_known_args() + tf.app.run(main=main, argv=[sys.argv[0]] + unparsed) diff --git a/tensorflow/contrib/eager/python/examples/gan/mnist_graph_test.py b/tensorflow/contrib/eager/python/examples/gan/mnist_graph_test.py new file mode 100644 index 0000000000000000000000000000000000000000..12b39b0cde49d4c017acfa74572c725036c54eff --- /dev/null +++ b/tensorflow/contrib/eager/python/examples/gan/mnist_graph_test.py @@ -0,0 +1,151 @@ +# Copyright 2017 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. +# ============================================================================== + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import tempfile +import time + +import numpy as np +import tensorflow as tf + +from tensorflow.contrib.eager.python.examples.gan import mnist + +NOISE_DIM = 100 +# Big enough so that summaries are never recorded. +# Lower this value if would like to benchmark with some summaries. +SUMMARY_INTERVAL = 10000 +SUMMARY_FLUSH_MS = 100 # Flush summaries every 100ms + + +def data_format(): + return 'channels_first' if tf.test.is_gpu_available() else 'channels_last' + + +class MnistGraphGanBenchmark(tf.test.Benchmark): + + def _create_graph(self, batch_size): + # Generate some random data. + images_data = np.random.randn(batch_size, 784).astype(np.float32) + dataset = tf.data.Dataset.from_tensors(images_data) + images = dataset.repeat().make_one_shot_iterator().get_next() + + # Create the models and optimizers + generator = mnist.Generator(data_format()) + discriminator = mnist.Discriminator(data_format()) + with tf.variable_scope('generator'): + generator_optimizer = tf.train.AdamOptimizer(0.001) + with tf.variable_scope('discriminator'): + discriminator_optimizer = tf.train.AdamOptimizer(0.001) + + # Run models and compute loss + noise_placeholder = tf.placeholder(tf.float32, + shape=[batch_size, NOISE_DIM]) + generated_images = generator(noise_placeholder) + tf.contrib.summary.image('generated_images', + tf.reshape(generated_images, [-1, 28, 28, 1]), + max_images=10) + discriminator_gen_outputs = discriminator(generated_images) + discriminator_real_outputs = discriminator(images) + generator_loss = mnist.generator_loss(discriminator_gen_outputs) + discriminator_loss = mnist.discriminator_loss(discriminator_real_outputs, + discriminator_gen_outputs) + # Get train ops + with tf.variable_scope('generator'): + generator_train = generator_optimizer.minimize( + generator_loss, var_list=generator.variables) + with tf.variable_scope('discriminator'): + discriminator_train = discriminator_optimizer.minimize( + discriminator_loss, var_list=discriminator.variables) + + return (generator_train, discriminator_train, noise_placeholder) + + def _report(self, test_name, start, num_iters, batch_size): + avg_time = (time.time() - start) / num_iters + dev = 'gpu' if tf.test.is_gpu_available() else 'cpu' + name = 'graph_%s_%s_batch_%d_%s' % (test_name, dev, batch_size, + data_format()) + extras = {'examples_per_sec': batch_size / avg_time} + self.report_benchmark( + iters=num_iters, wall_time=avg_time, name=name, extras=extras) + + def benchmark_train(self): + for batch_size in [64, 128, 256]: + with tf.Graph().as_default(): + global_step = tf.train.get_or_create_global_step() + increment_global_step = tf.assign_add(global_step, 1) + with tf.contrib.summary.create_file_writer( + tempfile.mkdtemp(), flush_millis=SUMMARY_FLUSH_MS).as_default(), ( + tf.contrib.summary.record_summaries_every_n_global_steps( + SUMMARY_INTERVAL)): + (generator_train, discriminator_train, noise_placeholder + ) = self._create_graph(batch_size) + + with tf.Session() as sess: + tf.contrib.summary.initialize(graph=tf.get_default_graph(), + session=sess) + + sess.run(tf.global_variables_initializer()) + + num_burn, num_iters = (3, 100) + for _ in range(num_burn): + noise = np.random.uniform(-1.0, 1.0, size=[batch_size, NOISE_DIM]) + # Increment global step before evaluating summary ops to avoid + # race condition. + sess.run(increment_global_step) + sess.run([generator_train, discriminator_train, + tf.contrib.summary.all_summary_ops()], + feed_dict={noise_placeholder: noise}) + + # Run and benchmark 2 epochs + start = time.time() + for _ in range(num_iters): + noise = np.random.uniform(-1.0, 1.0, size=[batch_size, NOISE_DIM]) + sess.run(increment_global_step) + sess.run([generator_train, discriminator_train, + tf.contrib.summary.all_summary_ops()], + feed_dict={noise_placeholder: noise}) + self._report('train', start, num_iters, batch_size) + + def benchmark_generate(self): + for batch_size in [64, 128, 256]: + with tf.Graph().as_default(): + # Using random weights. This will generate garbage. + generator = mnist.Generator(data_format()) + noise_placeholder = tf.placeholder(tf.float32, + shape=[batch_size, NOISE_DIM]) + generated_images = generator(noise_placeholder) + + init = tf.global_variables_initializer() + with tf.Session() as sess: + sess.run(init) + noise = np.random.uniform(-1.0, 1.0, size=[batch_size, NOISE_DIM]) + num_burn, num_iters = (30, 1000) + for _ in range(num_burn): + sess.run(generated_images, feed_dict={noise_placeholder: noise}) + + start = time.time() + for _ in range(num_iters): + # Comparison with the eager execution benchmark in mnist_test.py + # isn't entirely fair as the time here includes the cost of copying + # the feeds from CPU memory to GPU. + sess.run(generated_images, feed_dict={noise_placeholder: noise}) + self._report('generate', start, num_iters, batch_size) + + +if __name__ == '__main__': + tf.test.main() diff --git a/tensorflow/contrib/eager/python/examples/gan/mnist_test.py b/tensorflow/contrib/eager/python/examples/gan/mnist_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4a3ca8d82bc2619b05a734f6d2e58431c1a45995 --- /dev/null +++ b/tensorflow/contrib/eager/python/examples/gan/mnist_test.py @@ -0,0 +1,113 @@ +# Copyright 2017 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. +# ============================================================================== + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import tempfile +import time + +import tensorflow as tf + +import tensorflow.contrib.eager as tfe +from tensorflow.contrib.eager.python.examples.gan import mnist + +NOISE_DIM = 100 +# Big enough so that summaries are never recorded. +# Lower this value if would like to benchmark with some summaries. +SUMMARY_INTERVAL = 10000 +SUMMARY_FLUSH_MS = 100 # Flush summaries every 100ms + + +def data_format(): + return 'channels_first' if tf.test.is_gpu_available() else 'channels_last' + + +def device(): + return '/gpu:0' if tfe.num_gpus() else '/cpu:0' + + +class MnistEagerGanBenchmark(tf.test.Benchmark): + + def _report(self, test_name, start, num_iters, batch_size): + avg_time = (time.time() - start) / num_iters + dev = 'gpu' if tfe.num_gpus() else 'cpu' + name = 'eager_%s_%s_batch_%d_%s' % (test_name, dev, batch_size, + data_format()) + extras = {'examples_per_sec': batch_size / avg_time} + self.report_benchmark( + iters=num_iters, wall_time=avg_time, name=name, extras=extras) + + def benchmark_train(self): + for batch_size in [64, 128, 256]: + # Generate some random data. + burn_batches, measure_batches = (3, 100) + burn_images = [tf.random_normal([batch_size, 784]) + for _ in range(burn_batches)] + burn_dataset = tf.data.Dataset.from_tensor_slices(burn_images) + measure_images = [tf.random_normal([batch_size, 784]) + for _ in range(measure_batches)] + measure_dataset = tf.data.Dataset.from_tensor_slices(measure_images) + + tf.train.get_or_create_global_step() + with tf.device(device()): + # Create the models and optimizers + generator = mnist.Generator(data_format()) + discriminator = mnist.Discriminator(data_format()) + with tf.variable_scope('generator'): + generator_optimizer = tf.train.AdamOptimizer(0.001) + with tf.variable_scope('discriminator'): + discriminator_optimizer = tf.train.AdamOptimizer(0.001) + + with tf.contrib.summary.create_file_writer( + tempfile.mkdtemp(), flush_millis=SUMMARY_FLUSH_MS).as_default(): + + # warm up + mnist.train_one_epoch(generator, discriminator, generator_optimizer, + discriminator_optimizer, + burn_dataset, log_interval=SUMMARY_INTERVAL, + noise_dim=NOISE_DIM) + # measure + start = time.time() + mnist.train_one_epoch(generator, discriminator, generator_optimizer, + discriminator_optimizer, + measure_dataset, log_interval=SUMMARY_INTERVAL, + noise_dim=NOISE_DIM) + self._report('train', start, measure_batches, batch_size) + + def benchmark_generate(self): + for batch_size in [64, 128, 256]: + with tf.device(device()): + # Using random weights. This will generate garbage. + generator = mnist.Generator(data_format()) + + num_burn, num_iters = (30, 1000) + for _ in range(num_burn): + noise = tf.random_uniform(shape=[batch_size, NOISE_DIM], + minval=-1., maxval=1.) + generator(noise) + + start = time.time() + for _ in range(num_iters): + noise = tf.random_uniform(shape=[batch_size, NOISE_DIM], + minval=-1., maxval=1.) + generator(noise) + self._report('generate', start, num_iters, batch_size) + + +if __name__ == '__main__': + tfe.enable_eager_execution() + tf.test.main() diff --git a/tensorflow/contrib/eager/python/examples/linear_regression/BUILD b/tensorflow/contrib/eager/python/examples/linear_regression/BUILD index bab7ad0c701b2110fda9a8d27792fd361a5fc1c0..f86331af6f7928f0f86c888e22706c6e0a5978b2 100644 --- a/tensorflow/contrib/eager/python/examples/linear_regression/BUILD +++ b/tensorflow/contrib/eager/python/examples/linear_regression/BUILD @@ -23,3 +23,13 @@ cuda_py_test( "//tensorflow:tensorflow_py", ], ) + +cuda_py_test( + name = "linear_regression_graph_test", + size = "small", + srcs = ["linear_regression_graph_test.py"], + additional_deps = [ + ":linear_regression", + "//tensorflow:tensorflow_py", + ], +) diff --git a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py index 7bc5007c5655bed81b5600ee283c35bd332a1ebe..6ce4de6ee0bf50400eff339ac04e132252a2b53e 100644 --- a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py +++ b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py @@ -41,7 +41,7 @@ class LinearModel(tfe.Network): For those familiar with TensorFlow graphs, notice the absence of `tf.Session`. The `forward()` method here immediately executes and returns output values. The `loss()` method immediately compares the - output of `forward()` with the target adn returns the MSE loss value. + output of `forward()` with the target and returns the MSE loss value. The `fit()` performs gradient-descent training on the model's weights and bias. """ @@ -63,6 +63,10 @@ class LinearModel(tfe.Network): return self._hidden_layer(xs) +def mean_square_loss(model, xs, ys): + return tf.reduce_mean(tf.square(model(xs) - ys)) + + def fit(model, dataset, optimizer, verbose=False, logdir=None): """Fit the linear-regression model. @@ -76,10 +80,8 @@ def fit(model, dataset, optimizer, verbose=False, logdir=None): """ # The loss function to optimize. - def mean_square_loss(xs, ys): - return tf.reduce_mean(tf.square(model(xs) - ys)) - - loss_and_grads = tfe.implicit_value_and_gradients(mean_square_loss) + mse = lambda xs, ys: mean_square_loss(model, xs, ys) + loss_and_grads = tfe.implicit_value_and_gradients(mse) tf.train.get_or_create_global_step() if logdir: @@ -103,14 +105,20 @@ def fit(model, dataset, optimizer, verbose=False, logdir=None): def synthetic_dataset(w, b, noise_level, batch_size, num_batches): """tf.data.Dataset that yields synthetic data for linear regression.""" + return synthetic_dataset_helper(w, b, + tf.shape(w)[0], noise_level, batch_size, + num_batches) + +def synthetic_dataset_helper(w, b, num_features, noise_level, batch_size, + num_batches): # w is a matrix with shape [N, M] # b is a vector with shape [M] # So: # - Generate x's as vectors with shape [batch_size N] # - y = tf.matmul(x, W) + b + noise def batch(_): - x = tf.random_normal([batch_size, tf.shape(w)[0]]) + x = tf.random_normal([batch_size, num_features]) y = tf.matmul(x, w) + b + noise_level * tf.random_normal([]) return x, y diff --git a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_graph_test.py b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_graph_test.py new file mode 100644 index 0000000000000000000000000000000000000000..557ad42752144243ae3da61b955b31398cba846e --- /dev/null +++ b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_graph_test.py @@ -0,0 +1,85 @@ +# Copyright 2017 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. +"""Graph benchmark for linear regression, to contrast with eager execution.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import time + +import tensorflow as tf +from tensorflow.contrib.eager.python.examples.linear_regression import linear_regression + + +class GraphLinearRegressionBenchmark(tf.test.Benchmark): + + def benchmarkGraphLinearRegression(self): + num_epochs = 10 + num_batches = 200 + batch_size = 64 + dataset = linear_regression.synthetic_dataset_helper( + w=tf.random_uniform([3, 1]), + b=tf.random_uniform([1]), + num_features=3, + noise_level=0.01, + batch_size=batch_size, + num_batches=num_batches) + iterator = dataset.make_initializable_iterator() + x, y = iterator.get_next() + + model = linear_regression.LinearModel() + + if tf.test.is_gpu_available(): + use_gpu = True + device = "/device:GPU:0" + else: + use_gpu = False + device = "/device:CPU:0" + + with tf.device(device): + loss = linear_regression.mean_square_loss(model, x, y) + optimization_step = tf.train.GradientDescentOptimizer( + learning_rate=0.1).minimize(loss) + + with tf.Session() as sess: + sess.run(tf.global_variables_initializer()) + + def train(num_epochs): + for _ in range(num_epochs): + sess.run(iterator.initializer) + try: + while True: + _, _ = sess.run([optimization_step, loss]) + except tf.errors.OutOfRangeError: + pass + + # Warmup: a single epoch. + train(1) + + start_time = time.time() + train(num_epochs) + wall_time = time.time() - start_time + + examples_per_sec = num_epochs * num_batches * batch_size / wall_time + self.report_benchmark( + name="graph_train_%s" % + ("gpu" if use_gpu else "cpu"), + iters=num_epochs * num_batches, + extras={"examples_per_sec": examples_per_sec}, + wall_time=wall_time) + + +if __name__ == "__main__": + tf.test.main() diff --git a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_test.py b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_test.py index 39e7aabd7be04ba36a786a4c08d0df6c2ce916d0..e53234b51a7dccc11e548ac81a7ef070c628aa52 100644 --- a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_test.py +++ b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_test.py @@ -83,6 +83,7 @@ class LinearRegressionTest(tf.test.TestCase): class EagerLinearRegressionBenchmark(tf.test.Benchmark): def benchmarkEagerLinearRegression(self): + num_epochs = 10 num_batches = 200 batch_size = 64 dataset = linear_regression.synthetic_dataset( @@ -102,14 +103,15 @@ class EagerLinearRegressionBenchmark(tf.test.Benchmark): linear_regression.fit(model, burn_in_dataset, optimizer) start_time = time.time() - linear_regression.fit(model, dataset, optimizer) + for _ in range(num_epochs): + linear_regression.fit(model, dataset, optimizer) wall_time = time.time() - start_time - examples_per_sec = num_batches * batch_size / wall_time + examples_per_sec = num_epochs * num_batches * batch_size / wall_time self.report_benchmark( name="eager_train_%s" % ("gpu" if tfe.num_gpus() > 0 else "cpu"), - iters=num_batches, + iters=num_epochs * num_batches, extras={"examples_per_sec": examples_per_sec}, wall_time=wall_time) diff --git a/tensorflow/contrib/eager/python/examples/mnist/mnist.py b/tensorflow/contrib/eager/python/examples/mnist/mnist.py index bb121c7704b4772dde520ddc928a13c50ec8bb18..58b1e89d15895cf38331e6f7bd5a311a2f5f6467 100644 --- a/tensorflow/contrib/eager/python/examples/mnist/mnist.py +++ b/tensorflow/contrib/eager/python/examples/mnist/mnist.py @@ -23,7 +23,6 @@ from __future__ import division from __future__ import print_function import argparse -import functools import os import sys import time @@ -36,11 +35,11 @@ from tensorflow.examples.tutorials.mnist import input_data FLAGS = None -class MNISTModel(tfe.Network): +class MNISTModel(tf.keras.Model): """MNIST Network. Network structure is equivalent to: - https://github.com/tensorflow/tensorflow/blob/r1.4/tensorflow/examples/tutorials/mnist/mnist_deep.py + https://github.com/tensorflow/tensorflow/blob/r1.6/tensorflow/examples/tutorials/mnist/mnist_deep.py and https://github.com/tensorflow/models/blob/master/tutorials/image/mnist/convolutional.py @@ -62,18 +61,17 @@ class MNISTModel(tfe.Network): else: assert data_format == 'channels_last' self._input_shape = [-1, 28, 28, 1] - self.conv1 = self.track_layer( - tf.layers.Conv2D(32, 5, data_format=data_format, activation=tf.nn.relu)) - self.conv2 = self.track_layer( - tf.layers.Conv2D(64, 5, data_format=data_format, activation=tf.nn.relu)) - self.fc1 = self.track_layer(tf.layers.Dense(1024, activation=tf.nn.relu)) - self.fc2 = self.track_layer(tf.layers.Dense(10)) - self.dropout = self.track_layer(tf.layers.Dropout(0.5)) - self.max_pool2d = self.track_layer( - tf.layers.MaxPooling2D( - (2, 2), (2, 2), padding='SAME', data_format=data_format)) - - def call(self, inputs, training): + self.conv1 = tf.layers.Conv2D( + 32, 5, data_format=data_format, activation=tf.nn.relu) + self.conv2 = tf.layers.Conv2D( + 64, 5, data_format=data_format, activation=tf.nn.relu) + self.fc1 = tf.layers.Dense(1024, activation=tf.nn.relu) + self.fc2 = tf.layers.Dense(10) + self.dropout = tf.layers.Dropout(0.5) + self.max_pool2d = tf.layers.MaxPooling2D( + (2, 2), (2, 2), padding='SAME', data_format=data_format) + + def call(self, inputs, training=False): """Computes labels from inputs. Users should invoke __call__ to run the network, which delegates to this @@ -96,8 +94,7 @@ class MNISTModel(tfe.Network): x = self.max_pool2d(x) x = tf.layers.flatten(x) x = self.fc1(x) - if training: - x = self.dropout(x) + x = self.dropout(x, training=training) x = self.fc2(x) return x @@ -124,21 +121,18 @@ def train_one_epoch(model, optimizer, dataset, log_interval=None): tf.train.get_or_create_global_step() - def model_loss(labels, images): - prediction = model(images, training=True) - loss_value = loss(prediction, labels) - tf.contrib.summary.scalar('loss', loss_value) - tf.contrib.summary.scalar('accuracy', - compute_accuracy(prediction, labels)) - return loss_value - for (batch, (images, labels)) in enumerate(tfe.Iterator(dataset)): with tf.contrib.summary.record_summaries_every_n_global_steps(10): - batch_model_loss = functools.partial(model_loss, labels, images) - optimizer.minimize( - batch_model_loss, global_step=tf.train.get_global_step()) + with tfe.GradientTape() as tape: + prediction = model(images, training=True) + loss_value = loss(prediction, labels) + tf.contrib.summary.scalar('loss', loss_value) + tf.contrib.summary.scalar('accuracy', + compute_accuracy(prediction, labels)) + grads = tape.gradient(loss_value, model.variables) + optimizer.apply_gradients(zip(grads, model.variables)) if log_interval and batch % log_interval == 0: - print('Batch #%d\tLoss: %.6f' % (batch, batch_model_loss())) + print('Batch #%d\tLoss: %.6f' % (batch, loss_value)) def test(model, dataset): diff --git a/tensorflow/contrib/eager/python/examples/mnist/mnist_test.py b/tensorflow/contrib/eager/python/examples/mnist/mnist_test.py index 205709fe2edd3c260c30a84b624e322e120edf8e..136085eba21284a42282395e54f32c33bf63b5c3 100644 --- a/tensorflow/contrib/eager/python/examples/mnist/mnist_test.py +++ b/tensorflow/contrib/eager/python/examples/mnist/mnist_test.py @@ -39,22 +39,40 @@ def random_dataset(): return tf.data.Dataset.from_tensors((images, labels)) +def train_one_epoch(defun=False): + model = mnist.MNISTModel(data_format()) + if defun: + model.call = tfe.defun(model.call) + optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01) + dataset = random_dataset() + with tf.device(device()): + tf.train.get_or_create_global_step() + mnist.train_one_epoch(model, optimizer, dataset) + + +def evaluate(defun=False): + model = mnist.MNISTModel(data_format()) + dataset = random_dataset() + if defun: + model.call = tfe.defun(model.call) + with tf.device(device()): + tf.train.get_or_create_global_step() + mnist.test(model, dataset) + + class MNISTTest(tf.test.TestCase): def testTrainOneEpoch(self): - model = mnist.MNISTModel(data_format()) - optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01) - dataset = random_dataset() - with tf.device(device()): - tf.train.get_or_create_global_step() - mnist.train_one_epoch(model, optimizer, dataset) + train_one_epoch(defun=False) def testTest(self): - model = mnist.MNISTModel(data_format()) - dataset = random_dataset() - with tf.device(device()): - tf.train.get_or_create_global_step() - mnist.test(model, dataset) + evaluate(defun=False) + + def testTrainOneEpochWithDefunCall(self): + train_one_epoch(defun=True) + + def testTestWithDefunCall(self): + evaluate(defun=True) if __name__ == "__main__": diff --git a/tensorflow/contrib/eager/python/examples/resnet50/README.md b/tensorflow/contrib/eager/python/examples/resnet50/README.md index db023e6c976c8eda09ef0dee7eecb144678773c4..79e460052945718eac194653015d60d900998e2d 100644 --- a/tensorflow/contrib/eager/python/examples/resnet50/README.md +++ b/tensorflow/contrib/eager/python/examples/resnet50/README.md @@ -34,7 +34,7 @@ bazel run -c opt --config=cuda :resnet50_graph_test -- --benchmarks=. (Or remove the `--config=cuda` flag for running on CPU instead of GPU). -On October 31, 2017, the benchmarks demostrated comparable performance +On October 31, 2017, the benchmarks demonstrated comparable performance for eager and graph execution of this particular model when using a single NVIDIA Titan X (Pascal) GPU on a host with an Intel Xeon E5-1650 CPU @ 3.50GHz and a batch size of 32. diff --git a/tensorflow/contrib/eager/python/examples/resnet50/resnet50.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50.py index b302a87e0e8a61d2456db1eba847f31bd70f552e..9982fdb07eefa665379e7be095f4f8017d92cf97 100644 --- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50.py +++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50.py @@ -97,7 +97,7 @@ class _ConvBlock(tfe.Network): Args: kernel_size: the kernel size of middle conv layer at main path - filters: list of integers, the filterss of 3 conv layer at main path + filters: list of integers, the filters of 3 conv layer at main path stage: integer, current stage label, used for generating layer names block: 'a','b'..., current block label, used for generating layer names data_format: data_format for the input ('channels_first' or diff --git a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py index d8d8644dde10498e5fd480f92b69656fca1558dd..0ff8746884c288f824f5f22ab4c550370d0e0302 100644 --- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py +++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py @@ -22,6 +22,7 @@ import gc import tempfile import time +from six.moves import xrange # pylint: disable=redefined-builtin import tensorflow as tf import tensorflow.contrib.eager as tfe @@ -52,26 +53,33 @@ def random_batch(batch_size): def train_one_step(model, images, labels, optimizer): - def model_loss(): + with tfe.GradientTape() as tape: logits = model(images, training=True) loss = tf.losses.softmax_cross_entropy( logits=logits, onehot_labels=labels) tf.contrib.summary.scalar(name='loss', tensor=loss) - return loss - - optimizer.minimize(model_loss) + grads = tape.gradient(loss, model.variables) + optimizer.apply_gradients(zip(grads, model.variables)) class ResNet50Test(tf.test.TestCase): - def test_apply(self): + def _apply(self, defun=False): device, data_format = device_and_data_format() model = resnet50.ResNet50(data_format) + if defun: + model.call = tfe.defun(model.call) with tf.device(device): images, _ = random_batch(2) output = model(images) self.assertEqual((2, 1000), output.shape) + def test_apply(self): + self._apply(defun=False) + + def test_apply_with_defun(self): + self._apply(defun=True) + def test_apply_no_top(self): device, data_format = device_and_data_format() model = resnet50.ResNet50(data_format, include_top=False) @@ -175,9 +183,11 @@ class ResNet50Benchmarks(tf.test.Benchmark): # a sync. This is a roundabout way, yes. tf.constant(1.).cpu() - def benchmark_eager_apply(self): + def _benchmark_eager_apply(self, label, defun=False): device, data_format = device_and_data_format() model = resnet50.ResNet50(data_format) + if defun: + model.call = tfe.defun(model.call) batch_size = 64 num_burn = 5 num_iters = 30 @@ -189,16 +199,23 @@ class ResNet50Benchmarks(tf.test.Benchmark): start = time.time() for _ in xrange(num_iters): model(images).cpu() - self._report('eager_apply', start, num_iters, device, batch_size, - data_format) + self._report(label, start, num_iters, device, batch_size, data_format) + + def benchmark_eager_apply(self): + self._benchmark_eager_apply('eager_apply', defun=False) + + def benchmark_eager_apply_with_defun(self): + self._benchmark_eager_apply('eager_apply_with_defun', defun=True) - def _benchmark_eager_train(self, label, make_iterator): + def _benchmark_eager_train(self, label, make_iterator, defun=False): device, data_format = device_and_data_format() for batch_size in self._train_batch_sizes(): (images, labels) = random_batch(batch_size) num_burn = 3 num_iters = 10 model = resnet50.ResNet50(data_format) + if defun: + model.call = tfe.defun(model.call) optimizer = tf.train.GradientDescentOptimizer(0.1) with tf.device(device): @@ -217,7 +234,11 @@ class ResNet50Benchmarks(tf.test.Benchmark): self._report(label, start, num_iters, device, batch_size, data_format) def benchmark_eager_train(self): - self._benchmark_eager_train('eager_train', MockIterator) + self._benchmark_eager_train('eager_train', MockIterator, defun=False) + + def benchmark_eager_train_with_defun(self): + self._benchmark_eager_train( + 'eager_train_with_defun', MockIterator, defun=True) def benchmark_eager_train_datasets(self): @@ -226,7 +247,18 @@ class ResNet50Benchmarks(tf.test.Benchmark): ds = tf.data.Dataset.from_tensors(tensors).repeat() return tfe.Iterator(ds) - self._benchmark_eager_train('eager_train_dataset', make_iterator) + self._benchmark_eager_train( + 'eager_train_dataset', make_iterator, defun=False) + + def benchmark_eager_train_datasets_with_defun(self): + + def make_iterator(tensors): + with tf.device('/device:CPU:0'): + ds = tf.data.Dataset.from_tensors(tensors).repeat() + return tfe.Iterator(ds) + + self._benchmark_eager_train( + 'eager_train_dataset_with_defun', make_iterator, defun=True) if __name__ == '__main__': diff --git a/tensorflow/contrib/eager/python/examples/rnn_colorbot/rnn_colorbot.py b/tensorflow/contrib/eager/python/examples/rnn_colorbot/rnn_colorbot.py index 40919f2d4cf511eb35fac954719286366aef6c7c..aa87b94e7b0876e65405f6bcb2d6aabde36582bf 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_colorbot/rnn_colorbot.py +++ b/tensorflow/contrib/eager/python/examples/rnn_colorbot/rnn_colorbot.py @@ -65,7 +65,6 @@ import six import tensorflow as tf from tensorflow.contrib.eager.python import tfe -from tensorflow.python.eager import context try: import matplotlib.pyplot as plt # pylint: disable=g-import-not-at-top diff --git a/tensorflow/contrib/eager/python/examples/rnn_ptb/README.md b/tensorflow/contrib/eager/python/examples/rnn_ptb/README.md index 743ebb68ee5bba5635899267cc4839828f7e4e2f..966177e91c212c1aa132fe3af6f7dc9a50fb984e 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_ptb/README.md +++ b/tensorflow/contrib/eager/python/examples/rnn_ptb/README.md @@ -40,7 +40,7 @@ bazel run -c opt --config=cuda :rnn_ptb_graph_test -- --benchmarks=. (Or remove the `--config=cuda` flag for running on CPU instead of GPU). -On October 31, 2017, the benchmarks demostrated slightly better performance +On October 31, 2017, the benchmarks demonstrated slightly better performance (3-6%) for graph execution over eager execution for this particular model when using a single NVIDIA Titan X (Pascal) GPU on a host with an Intel Xeon E5-1650 CPU @ 3.50GHz and a batch size of 32. diff --git a/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py b/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py index b8388f93a4d8b985789ae9aef26a6ad07ac7d857..5c5c59c87744f4ffa6db90e5d8d3aa3bc8132756 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py +++ b/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py @@ -88,7 +88,7 @@ class Embedding(tf.layers.Layer): class PTBModel(tfe.Network): - """LSTM for word language modelling. + """LSTM for word language modeling. Model described in: (Zaremba, et. al.) Recurrent Neural Network Regularization @@ -214,7 +214,7 @@ class Datasets(object): """Load the Penn Treebank dataset. Args: - path: Path to the data/ directory of the dataset from from Tomas Mikolov's + path: Path to the data/ directory of the dataset from Tomas Mikolov's webpage - http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz """ @@ -339,8 +339,7 @@ if __name__ == "__main__": "http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz") parser.add_argument( "--logdir", type=str, default="", help="Directory for checkpoint.") - parser.add_argument( - "--epoch", type=int, default=20, help="Number of epoches.") + parser.add_argument("--epoch", type=int, default=20, help="Number of epochs.") parser.add_argument("--batch-size", type=int, default=20, help="Batch size.") parser.add_argument( "--seq-len", type=int, default=35, help="Sequence length.") diff --git a/tensorflow/contrib/eager/python/examples/spinn/data.py b/tensorflow/contrib/eager/python/examples/spinn/data.py index a6e046320f78541bef4e091e97f08fd51857af83..3bc3bb49bcbbc26f7a3134a8bfc385ec080dde1e 100644 --- a/tensorflow/contrib/eager/python/examples/spinn/data.py +++ b/tensorflow/contrib/eager/python/examples/spinn/data.py @@ -51,11 +51,11 @@ def get_non_parenthesis_words(items): """Get the non-parenthesis items from a SNLI parsed sentence. Args: - items: Data items from a parsed SNLI setence, with parentheses. E.g., + items: Data items from a parsed SNLI sentence, with parentheses. E.g., ["(", "Man", "(", "(", "(", "(", "(", "wearing", "pass", ")", ... Returns: - A list of non-parenthis word items, all converted to lower case. E.g., + A list of non-parentheses word items, all converted to lower case. E.g., ["man", "wearing", "pass", ... """ return [x.lower() for x in items if x not in PARENTHESES and x] @@ -201,7 +201,7 @@ def load_word_vectors(data_root, vocab): def calculate_bins(length2count, min_bin_size): - """Cacluate bin boundaries given a histogram of lengths and mininum bin size. + """Calculate bin boundaries given a histogram of lengths and minimum bin size. Args: length2count: A `dict` mapping length to sentence count. @@ -227,6 +227,29 @@ def calculate_bins(length2count, min_bin_size): return bounds +def encode_sentence(sentence, word2index): + """Encode a single sentence as word indices and shift-reduce code. + + Args: + sentence: The sentence with added binary parse information, represented as + a string, with all the word items and parentheses separated by spaces. + E.g., '( ( The dog ) ( ( is ( playing toys ) ) . ) )'. + word2index: A `dict` mapping words to their word indices. + + Returns: + 1. Word indices as a numpy array, with shape `(sequence_len, 1)`. + 2. Shift-reduce sequence as a numpy array, with shape + `(sequence_len * 2 - 3, 1)`. + """ + items = [w for w in sentence.split(" ") if w] + words = get_non_parenthesis_words(items) + shift_reduce = get_shift_reduce(items) + word_indices = pad_and_reverse_word_ids( + [[word2index.get(word, UNK_CODE) for word in words]]).T + return (word_indices, + np.expand_dims(np.array(shift_reduce, dtype=np.int64), -1)) + + class SnliData(object): """A split of SNLI data.""" @@ -335,9 +358,9 @@ class SnliData(object): # The sorting above and the batching here makes sure that sentences of # similar max lengths are batched together, minimizing the inefficiency # due to uneven max lengths. The sentences are batched differently in - # each call to get_generator() due to the shuffling before sotring + # each call to get_generator() due to the shuffling before sorting # above. The pad_and_reverse_word_ids() and pad_transitions() functions - # take care of any remaning unevenness of the max sentence lengths. + # take care of any remaining unevenness of the max sentence lengths. end = min(begin + batch_size, len(labels)) # Transpose, because the SPINN model requires time-major, instead of # batch-major. diff --git a/tensorflow/contrib/eager/python/examples/spinn/data_test.py b/tensorflow/contrib/eager/python/examples/spinn/data_test.py index e4f0b37c5099e45b7e3b258b258c0a203c36b3b7..54fef2c3fe4111cd2d93ac109a5b8fffad0c2fad 100644 --- a/tensorflow/contrib/eager/python/examples/spinn/data_test.py +++ b/tensorflow/contrib/eager/python/examples/spinn/data_test.py @@ -22,6 +22,7 @@ import os import shutil import tempfile +import numpy as np import tensorflow as tf from tensorflow.contrib.eager.python.examples.spinn import data @@ -173,14 +174,9 @@ class DataTest(tf.test.TestCase): ValueError, "Cannot find GloVe embedding file at"): data.load_word_vectors(self._temp_data_dir, vocab) - def testSnliData(self): - """Unit test for SnliData objects.""" - snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0") - fake_train_file = os.path.join(snli_1_0_dir, "snli_1.0_train.txt") - os.makedirs(snli_1_0_dir) - + def _createFakeSnliData(self, fake_snli_file): # Four sentences in total. - with open(fake_train_file, "wt") as f: + with open(fake_snli_file, "wt") as f: f.write("gold_label\tsentence1_binary_parse\tsentence2_binary_parse\t" "sentence1_parse\tsentence2_parse\tsentence1\tsentence2\t" "captionID\tpairID\tlabel1\tlabel2\tlabel3\tlabel4\tlabel5\n") @@ -205,10 +201,7 @@ class DataTest(tf.test.TestCase): "4705552913.jpg#2\t4705552913.jpg#2r1n\t" "neutral\tentailment\tneutral\tneutral\tneutral\n") - glove_dir = os.path.join(self._temp_data_dir, "glove") - os.makedirs(glove_dir) - glove_file = os.path.join(glove_dir, "glove.42B.300d.txt") - + def _createFakeGloveData(self, glove_file): words = [".", "foo", "bar", "baz", "quux", "quuz", "grault", "garply"] with open(glove_file, "wt") as f: for i, word in enumerate(words): @@ -220,6 +213,40 @@ class DataTest(tf.test.TestCase): else: f.write("\n") + def testEncodeSingleSentence(self): + snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0") + fake_train_file = os.path.join(snli_1_0_dir, "snli_1.0_train.txt") + os.makedirs(snli_1_0_dir) + self._createFakeSnliData(fake_train_file) + vocab = data.load_vocabulary(self._temp_data_dir) + glove_dir = os.path.join(self._temp_data_dir, "glove") + os.makedirs(glove_dir) + glove_file = os.path.join(glove_dir, "glove.42B.300d.txt") + self._createFakeGloveData(glove_file) + word2index, _ = data.load_word_vectors(self._temp_data_dir, vocab) + + sentence_variants = [ + "( Foo ( ( bar baz ) . ) )", + " ( Foo ( ( bar baz ) . ) ) ", + "( Foo ( ( bar baz ) . ) )"] + for sentence in sentence_variants: + word_indices, shift_reduce = data.encode_sentence(sentence, word2index) + self.assertEqual(np.int64, word_indices.dtype) + self.assertEqual((5, 1), word_indices.shape) + self.assertAllClose( + np.array([[3, 3, 3, 2, 3, 2, 2]], dtype=np.int64).T, shift_reduce) + + def testSnliData(self): + snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0") + fake_train_file = os.path.join(snli_1_0_dir, "snli_1.0_train.txt") + os.makedirs(snli_1_0_dir) + self._createFakeSnliData(fake_train_file) + + glove_dir = os.path.join(self._temp_data_dir, "glove") + os.makedirs(glove_dir) + glove_file = os.path.join(glove_dir, "glove.42B.300d.txt") + self._createFakeGloveData(glove_file) + vocab = data.load_vocabulary(self._temp_data_dir) word2index, _ = data.load_word_vectors(self._temp_data_dir, vocab) @@ -230,7 +257,7 @@ class DataTest(tf.test.TestCase): self.assertEqual(1, train_data.num_batches(4)) generator = train_data.get_generator(2)() - for i in range(2): + for _ in range(2): label, prem, prem_trans, hypo, hypo_trans = next(generator) self.assertEqual(2, len(label)) self.assertEqual((4, 2), prem.shape) diff --git a/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py b/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py index 84e25cf81a2223800c47994b26d000caddee6b01..081b0af14fcc983a3f85d2a50e2bb04d2f2493b3 100644 --- a/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py +++ b/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py @@ -26,6 +26,7 @@ import tempfile import time import numpy as np +from six.moves import xrange # pylint: disable=redefined-builtin import tensorflow as tf # pylint: disable=g-bad-import-order @@ -35,6 +36,7 @@ from third_party.examples.eager.spinn import spinn from tensorflow.contrib.summary import summary_test_util from tensorflow.python.eager import test from tensorflow.python.framework import test_util +from tensorflow.python.training import checkpoint_utils # pylint: enable=g-bad-import-order @@ -65,13 +67,30 @@ def _generate_synthetic_snli_data_batch(sequence_length, return labels, prem, prem_trans, hypo, hypo_trans -def _test_spinn_config(d_embed, d_out, logdir=None): +def _test_spinn_config(d_embed, d_out, logdir=None, inference_sentences=None): + """Generate a config tuple for testing. + + Args: + d_embed: Embedding dimensions. + d_out: Model output dimensions. + logdir: Optional logdir. + inference_sentences: A 2-tuple of strings representing the sentences (with + binary parsing result), e.g., + ("( ( The dog ) ( ( is running ) . ) )", "( ( The dog ) ( moves . ) )"). + + Returns: + A config tuple. + """ config_tuple = collections.namedtuple( "Config", ["d_hidden", "d_proj", "d_tracker", "predict", "embed_dropout", "mlp_dropout", "n_mlp_layers", "d_mlp", "d_out", "projection", "lr", "batch_size", "epochs", "force_cpu", "logdir", "log_every", "dev_every", "save_every", - "lr_decay_every", "lr_decay_by"]) + "lr_decay_every", "lr_decay_by", "inference_premise", + "inference_hypothesis"]) + + inference_premise = inference_sentences[0] if inference_sentences else None + inference_hypothesis = inference_sentences[1] if inference_sentences else None return config_tuple( d_hidden=d_embed, d_proj=d_embed * 2, @@ -85,14 +104,16 @@ def _test_spinn_config(d_embed, d_out, logdir=None): projection=True, lr=2e-2, batch_size=2, - epochs=10, + epochs=20, force_cpu=False, logdir=logdir, log_every=1, dev_every=2, save_every=2, lr_decay_every=1, - lr_decay_by=0.75) + lr_decay_by=0.75, + inference_premise=inference_premise, + inference_hypothesis=inference_hypothesis) class SpinnTest(test_util.TensorFlowTestCase): @@ -287,11 +308,7 @@ class SpinnTest(test_util.TensorFlowTestCase): # Training on the batch should have led to a change in the loss value. self.assertNotEqual(loss1.numpy(), loss2.numpy()) - def testTrainSpinn(self): - """Test with fake toy SNLI data and GloVe vectors.""" - - # 1. Create and load a fake SNLI data file and a fake GloVe embedding file. - snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0") + def _create_test_data(self, snli_1_0_dir): fake_train_file = os.path.join(snli_1_0_dir, "snli_1.0_train.txt") os.makedirs(snli_1_0_dir) @@ -336,13 +353,52 @@ class SpinnTest(test_util.TensorFlowTestCase): else: f.write("\n") + return fake_train_file + + def testInferSpinnWorks(self): + """Test inference with the spinn model.""" + snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0") + self._create_test_data(snli_1_0_dir) + + vocab = data.load_vocabulary(self._temp_data_dir) + word2index, embed = data.load_word_vectors(self._temp_data_dir, vocab) + + config = _test_spinn_config( + data.WORD_VECTOR_LEN, 4, + logdir=os.path.join(self._temp_data_dir, "logdir"), + inference_sentences=("( foo ( bar . ) )", "( bar ( foo . ) )")) + logits = spinn.train_or_infer_spinn( + embed, word2index, None, None, None, config) + self.assertEqual(tf.float32, logits.dtype) + self.assertEqual((3,), logits.shape) + + def testInferSpinnThrowsErrorIfOnlyOneSentenceIsSpecified(self): + snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0") + self._create_test_data(snli_1_0_dir) + + vocab = data.load_vocabulary(self._temp_data_dir) + word2index, embed = data.load_word_vectors(self._temp_data_dir, vocab) + + config = _test_spinn_config( + data.WORD_VECTOR_LEN, 4, + logdir=os.path.join(self._temp_data_dir, "logdir"), + inference_sentences=("( foo ( bar . ) )", None)) + with self.assertRaises(ValueError): + spinn.train_or_infer_spinn(embed, word2index, None, None, None, config) + + def testTrainSpinn(self): + """Test with fake toy SNLI data and GloVe vectors.""" + + # 1. Create and load a fake SNLI data file and a fake GloVe embedding file. + snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0") + fake_train_file = self._create_test_data(snli_1_0_dir) + vocab = data.load_vocabulary(self._temp_data_dir) word2index, embed = data.load_word_vectors(self._temp_data_dir, vocab) train_data = data.SnliData(fake_train_file, word2index) dev_data = data.SnliData(fake_train_file, word2index) test_data = data.SnliData(fake_train_file, word2index) - print(embed) # 2. Create a fake config. config = _test_spinn_config( @@ -350,7 +406,8 @@ class SpinnTest(test_util.TensorFlowTestCase): logdir=os.path.join(self._temp_data_dir, "logdir")) # 3. Test training of a SPINN model. - spinn.train_spinn(embed, train_data, dev_data, test_data, config) + trainer = spinn.train_or_infer_spinn( + embed, word2index, train_data, dev_data, test_data, config) # 4. Load train loss values from the summary files and verify that they # decrease with training. @@ -362,6 +419,15 @@ class SpinnTest(test_util.TensorFlowTestCase): self.assertEqual(config.epochs, len(train_losses)) self.assertLess(train_losses[-1], train_losses[0]) + # 5. Verify that checkpoints exist and contains all the expected variables. + self.assertTrue(glob.glob(os.path.join(config.logdir, "ckpt*"))) + ckpt_variable_names = [ + item[0] for item in checkpoint_utils.list_variables(config.logdir)] + self.assertIn("global_step", ckpt_variable_names) + for v in trainer.variables: + variable_name = v.name[:v.name.index(":")] if ":" in v.name else v.name + self.assertIn(variable_name, ckpt_variable_names) + class EagerSpinnSNLIClassifierBenchmark(test.Benchmark): diff --git a/tensorflow/contrib/eager/python/g3doc/guide.md b/tensorflow/contrib/eager/python/g3doc/guide.md index 0095ffa0db99d46d25654d73504d0d7d41c18b6f..d97ff6b74cf033617154f7cbbd00cb6492a1d2f4 100644 --- a/tensorflow/contrib/eager/python/g3doc/guide.md +++ b/tensorflow/contrib/eager/python/g3doc/guide.md @@ -19,29 +19,34 @@ to models defined without using eager execution. ## Installation -Eager execution is **not** included in the latest release (version 1.4) of -TensorFlow. To use it, you will need to [build TensorFlow from -source](https://www.tensorflow.org/install/install_sources) or install the -nightly builds. +Eager execution is included in TensorFlow versions 1.5 and above. +Installation instructions at https://www.tensorflow.org/install/ -For example, the nightly builds can be installed using `pip`: +The contents of this guide are compatible with TensorFlow 1.5. +However, if you run into bugs that are fixed in source but not the +release, you may want to either [build from +source](https://www.tensorflow.org/install/install_sources) +or try a nightly build. The nightly builds are available as: -- `pip install tf-nightly` (for CPU-only TensorFlow) -- `pip install tf-nightly-gpu` (for GPU-enabled TensorFlow) +- [`pip` packages](https://github.com/tensorflow/tensorflow/blob/master/README.md#installation) and -Or using `docker`, with [Jupyter Notebook](http://jupyter.org/) support: +- [docker](https://hub.docker.com/r/tensorflow/tensorflow/) images. + +For example, to run the latest nightly docker image: ```sh -# For CPU-only TensorFlow +# If you have a GPU, use https://github.com/NVIDIA/nvidia-docker +docker pull tensorflow/tensorflow:nightly-gpu +docker run --runtime=nvidia -it -p 8888:8888 tensorflow/tensorflow:nightly-gpu + +# If you do not have a GPU, use the CPU-only image docker pull tensorflow/tensorflow:nightly docker run -it -p 8888:8888 tensorflow/tensorflow:nightly - -# For GPU-enabled TensorFlow: -# (Requires https://github.com/NVIDIA/nvidia-docker) -nvidia-docker pull tensorflow/tensorflow:nightly-gpu -nvidia-docker run -it -p 8888:8888 tensorflow/tensorflow:nightly-gpu ``` +And then visit http://localhost:8888 in your browser for a Jupyter notebook +environment. + ## Getting Started With TensorFlow installed, eager execution is enabled via a single call: @@ -292,7 +297,7 @@ def loss(weight, bias): error = prediction(training_inputs, weight, bias) - training_outputs return tf.reduce_mean(tf.square(error)) -# Function that returns the the derivative of loss with respect to +# Function that returns the derivative of loss with respect to # weight and bias grad = tfe.gradients_function(loss) diff --git a/tensorflow/contrib/eager/python/metrics_impl.py b/tensorflow/contrib/eager/python/metrics_impl.py index bf029ca5f9dddb152274da6a1cc96bea7981d8fd..ea8dbf2b46ea4bd0e33645ae3c590c4dd13f7a52 100644 --- a/tensorflow/contrib/eager/python/metrics_impl.py +++ b/tensorflow/contrib/eager/python/metrics_impl.py @@ -291,6 +291,9 @@ class Mean(Metric): Args: values: Tensor with the per-example value. weights: Optional weighting of each example. Defaults to 1. + + Returns: + The arguments, for easy chaining. """ if weights is None: self.denom.assign_add( @@ -302,6 +305,9 @@ class Mean(Metric): self.denom.assign_add(math_ops.reduce_sum(weights)) values = math_ops.cast(values, self.dtype) * weights self.numer.assign_add(math_ops.reduce_sum(values)) + if weights is None: + return values + return values, weights def result(self): t = self.numer / self.denom @@ -329,7 +335,13 @@ class Accuracy(Mean): per element of the Tensor. predictions: Tensor with the predicted label for each example. weights: Optional weighting of each example. Defaults to 1. + + Returns: + The arguments, for easy chaining. """ matches = math_ops.equal(labels, predictions) matches = math_ops.cast(matches, dtypes.float64) super(Accuracy, self).call(matches, weights=weights) + if weights is None: + return labels, predictions + return labels, predictions, weights diff --git a/tensorflow/contrib/eager/python/metrics_test.py b/tensorflow/contrib/eager/python/metrics_test.py index 9cf34fd9b2dcf1b123cacc6863af817419eda007..a9ecaa3f8bced3043ea0eb0ac3aa8bfa65e9e1ff 100644 --- a/tensorflow/contrib/eager/python/metrics_test.py +++ b/tensorflow/contrib/eager/python/metrics_test.py @@ -180,6 +180,19 @@ class MetricsTest(test.TestCase): m2 = metrics.Mean() m2(2) + def testMetricsChain(self): + with context.graph_mode(), self.test_session(): + m1 = metrics.Mean() + m2 = metrics.Mean(name="m2") + update_m2 = m2(3.0) + update_m2_2 = m2(m1(1.0)) + m1.init_variables().run() + m2.init_variables().run() + update_m2.eval() + update_m2_2.eval() + self.assertAllEqual(m2.result().eval(), 2.0) + self.assertAllEqual(m1.result().eval(), 1.0) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/eager/python/network_test.py b/tensorflow/contrib/eager/python/network_test.py index 3eb4f5f8b3954a7ed04d2ef1d4f119ad137e1e65..3329fc6c513265deff41a368f5688dd605209c14 100644 --- a/tensorflow/contrib/eager/python/network_test.py +++ b/tensorflow/contrib/eager/python/network_test.py @@ -105,15 +105,13 @@ class NetworkTest(test.TestCase): result = net(constant_op.constant([[2.0]])) self.assertEqual(34.0, self.evaluate(result)) - # TODO(akshayka): This test should be changed once an API for compiling - # `call` into a defun is implemented. def testReplacingNetworkCallWithDefun(self): net = MyNetwork(name="abcd") + net.call = function.defun(net.call) x = constant_op.constant([[2.0]]) net(x) # Force variables to be created. self.evaluate(net.trainable_variables[0].assign([[17.0]])) - net.call = function.defun(net.call) result = net(x) # Build and execute the TensorFlow function self.assertEqual(34.0, self.evaluate(result)) @@ -541,7 +539,7 @@ class NetworkTest(test.TestCase): # No issue here since the name is unique within its scope. name_conflict3 = MyNetwork(name="name_conflict") net2 = MyNetwork() # name=outside_scope/my_network_2 to avoid the - # variable_scope my_network_1 below. + # variable_scope my_network_1 below. vs_name_conflict = MyNetwork(name="vs_name_conflict") # conflict below with variable_scope.variable_scope("intervening_scope"): with variable_scope.variable_scope(captured_scope): @@ -690,7 +688,7 @@ class NetworkTest(test.TestCase): net2(one) # Layer names typically are globally unique rather than being unique within # the scope of their first use. However, within a Network they must be named - # locally so that previous Layer consutrciton does not interfere with + # locally so that previous Layer construction does not interfere with # variable naming (e.g. add a Layer construction before the Network, # suddenly your previously saved checkpoint is incompatible). self.assertEqual("dense", net1.l1.name) diff --git a/tensorflow/contrib/eager/python/saver.py b/tensorflow/contrib/eager/python/saver.py index 57b070ec6eeac00c77f199a846639d64c4957cd8..62421849c766a1124c726812428985c913c653a3 100644 --- a/tensorflow/contrib/eager/python/saver.py +++ b/tensorflow/contrib/eager/python/saver.py @@ -82,7 +82,7 @@ def restore_variables_on_create(save_path, map_func=None): map_func_wrapper = lambda self, x: x else: if not callable(map_func): - raise ValueError("map_func must be callaled.") + raise ValueError("map_func must be callable.") map_func_wrapper = lambda self, x: map_func(x) ckpt_var_cache = dict() diff --git a/tensorflow/contrib/eager/python/saver_test.py b/tensorflow/contrib/eager/python/saver_test.py index abc7e3690c76c4446bce6b945325f1ca15ef1c8b..1a7f7b85e688e80e3cf482f2754462888187d311 100644 --- a/tensorflow/contrib/eager/python/saver_test.py +++ b/tensorflow/contrib/eager/python/saver_test.py @@ -73,16 +73,6 @@ class SaverTest(test.TestCase): with self.assertRaisesRegexp(ValueError, 'v1'): saver.save(ckpt_prefix) - def testDifferentGraphError(self): - with ops.device(self._dev()): - with ops.Graph().as_default(): - v1 = resource_variable_ops.ResourceVariable(1.0, name='v1') - with ops.Graph().as_default(): - saver = _saver.Saver([v1]) - ckpt_prefix = os.path.join(test.get_temp_dir(), 'ckpt') - with self.assertRaisesRegexp(ValueError, 'Graph'): - saver.save(ckpt_prefix) - def testSameObjectOK(self): with ops.device(self._dev()): v1 = resource_variable_ops.ResourceVariable(1.0, name='v1') diff --git a/tensorflow/contrib/eager/python/tfe.py b/tensorflow/contrib/eager/python/tfe.py index 770a7e3e7a01f3351c229b7fb53383240dd1f1c8..d32bebf90c1e768d1efec26b3b78bf1a522a8f00 100644 --- a/tensorflow/contrib/eager/python/tfe.py +++ b/tensorflow/contrib/eager/python/tfe.py @@ -25,6 +25,7 @@ To use, at program startup, call `tfe.enable_eager_execution()`. @@py_func @@defun +@@make_template @@implicit_gradients @@implicit_value_and_gradients @@gradients_function @@ -51,13 +52,13 @@ To use, at program startup, call `tfe.enable_eager_execution()`. @@EagerVariableStore @@Network +@@Sequential @@save_network_checkpoint @@restore_network_checkpoint @@in_eager_mode @@in_graph_mode -@@IsolateTest @@run_test_in_graph_and_eager_modes @@DEVICE_PLACEMENT_EXPLICIT @@ -75,6 +76,7 @@ from __future__ import print_function from tensorflow.contrib.eager.python import metrics from tensorflow.contrib.eager.python.datasets import Iterator from tensorflow.contrib.eager.python.network import Network +from tensorflow.contrib.eager.python.network import Sequential from tensorflow.contrib.eager.python.network import save_network_checkpoint from tensorflow.contrib.eager.python.network import restore_network_checkpoint from tensorflow.contrib.eager.python.saver import get_optimizer_variables @@ -98,15 +100,16 @@ from tensorflow.python.eager.execution_callbacks import nan_callback from tensorflow.python.eager.execution_callbacks import seterr from tensorflow.python.framework.ops import enable_eager_execution from tensorflow.python.framework.ops import eager_run as run -from tensorflow.python.framework.test_util import IsolateTest from tensorflow.python.framework.test_util import run_in_graph_and_eager_modes as run_test_in_graph_and_eager_modes from tensorflow.python.ops.resource_variable_ops import ResourceVariable as Variable from tensorflow.python.ops.variable_scope import EagerVariableStore from tensorflow.python.ops import script_ops +from tensorflow.python.ops import template from tensorflow.python.util.all_util import remove_undocumented py_func = script_ops.eager_py_func defun = function.defun +make_template = template.make_template_internal implicit_gradients = backprop.implicit_grad implicit_value_and_gradients = backprop.implicit_val_and_grad gradients_function = backprop.gradients_function diff --git a/tensorflow/contrib/eager/python/tfe_test.py b/tensorflow/contrib/eager/python/tfe_test.py index 0dedb2fd7c0905801cd87c239ff2ee09eecb6080..b6659c2a1797feab261d756e78b45231dbea5a02 100644 --- a/tensorflow/contrib/eager/python/tfe_test.py +++ b/tensorflow/contrib/eager/python/tfe_test.py @@ -102,10 +102,6 @@ class TFETest(test_util.TensorFlowTestCase): # Expect at least one device. self.assertTrue(tfe.list_devices()) - def testNumGPUs(self): - devices = tfe.list_devices() - self.assertEqual(len(devices) - 1, tfe.num_gpus()) - def testAddCheckNumericsOpsRaisesError(self): with self.assertRaisesRegexp( RuntimeError, diff --git a/tensorflow/contrib/estimator/BUILD b/tensorflow/contrib/estimator/BUILD index bd65ece85d2bfc6b38ba3507d3e702241eaf6067..6cdbed5b896577f5622b1bd0123c289c798bc0a5 100644 --- a/tensorflow/contrib/estimator/BUILD +++ b/tensorflow/contrib/estimator/BUILD @@ -163,7 +163,7 @@ py_library( srcs_version = "PY2AND3", deps = [ "//tensorflow/python:array_ops", - "//tensorflow/python:control_flow_ops", + "//tensorflow/python:check_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:lookup_ops", @@ -177,7 +177,6 @@ py_library( "//tensorflow/python/estimator:metric_keys", "//tensorflow/python/estimator:model_fn", "//tensorflow/python/estimator:prediction_keys", - "//tensorflow/python/estimator:util", "//tensorflow/python/ops/losses", "//tensorflow/python/saved_model:signature_constants", ], @@ -205,6 +204,7 @@ py_test( "//tensorflow/python/estimator:metric_keys", "//tensorflow/python/estimator:model_fn", "//tensorflow/python/estimator:prediction_keys", + "//tensorflow/python/ops/losses", "//tensorflow/python/saved_model:signature_constants", "//third_party/py/numpy", "@six_archive//:six", @@ -376,5 +376,9 @@ cuda_py_test( "//tensorflow/python:variables", ":replicate_model_fn", ], - tags = ["multi_gpu"], + tags = [ + "manual", + "multi_gpu", + "notap", + ], ) diff --git a/tensorflow/contrib/estimator/__init__.py b/tensorflow/contrib/estimator/__init__.py index 7533943d4834fd1211c0907b322854338e6d08f3..0f75b77050b0ba4c752a6a74fdc7024170b6f318 100644 --- a/tensorflow/contrib/estimator/__init__.py +++ b/tensorflow/contrib/estimator/__init__.py @@ -47,7 +47,7 @@ _allowed_symbols = [ 'dnn_logit_fn_builder', 'linear_logit_fn_builder', 'replicate_model_fn', - 'GatheringOptimizer', + 'TowerOptimizer', ] remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/estimator/python/estimator/extenders.py b/tensorflow/contrib/estimator/python/estimator/extenders.py index 29c3c7358534f6e8ebbd31cbfcd7e34086d9b506..c99bf8badb35e6fffb7cae8761db9d402b8b3a8f 100644 --- a/tensorflow/contrib/estimator/python/estimator/extenders.py +++ b/tensorflow/contrib/estimator/python/estimator/extenders.py @@ -100,7 +100,7 @@ def add_metrics(estimator, metric_fn): def clip_gradients_by_norm(optimizer, clip_norm): - """Returns an optimizer which clips gradients before appliying them. + """Returns an optimizer which clips gradients before applying them. Example: diff --git a/tensorflow/contrib/estimator/python/estimator/extenders_test.py b/tensorflow/contrib/estimator/python/estimator/extenders_test.py index 5f4a3cc902c9cc07c0688ad41dab7391a641c133..ad1a8ef152b07ecbab33d9eb3184a2ae89def27d 100644 --- a/tensorflow/contrib/estimator/python/estimator/extenders_test.py +++ b/tensorflow/contrib/estimator/python/estimator/extenders_test.py @@ -20,8 +20,8 @@ from __future__ import print_function import numpy as np -from tensorflow.contrib.data.python.ops import dataset_ops from tensorflow.contrib.estimator.python.estimator import extenders +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.estimator import estimator_lib from tensorflow.python.estimator.canned import linear from tensorflow.python.feature_column import feature_column as fc diff --git a/tensorflow/contrib/estimator/python/estimator/head.py b/tensorflow/contrib/estimator/python/estimator/head.py index a9311a20f127d92f02a95b8b48082fc90850635a..238cf287b768eee28b20202084eb244c085c8b75 100644 --- a/tensorflow/contrib/estimator/python/estimator/head.py +++ b/tensorflow/contrib/estimator/python/estimator/head.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function from tensorflow.python.estimator import model_fn -from tensorflow.python.estimator import util from tensorflow.python.estimator.canned import head as head_lib from tensorflow.python.estimator.canned import metric_keys from tensorflow.python.estimator.canned import prediction_keys @@ -29,7 +28,6 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops -from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import lookup_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics as metrics_lib @@ -44,6 +42,8 @@ _DEFAULT_SERVING_KEY = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY def multi_class_head(n_classes, weight_column=None, label_vocabulary=None, + loss_reduction=losses.Reduction.SUM, + loss_fn=None, name=None): """Creates a `_Head` for multi class classification. @@ -64,6 +64,12 @@ def multi_class_head(n_classes, labels have shape `[batch_size, 1]`, the loss is the weighted sum over `batch_size`. + Also supports custom `loss_fn`. `loss_fn` takes `(labels, logits)` or + `(labels, logits, features)` as arguments and returns unreduced loss with + shape `[D0, D1, ... DN, 1]`. `loss_fn` must support integer `labels` with + shape `[D0, D1, ... DN, 1]`. Namely, the head applies `label_vocabulary` to + the input labels before passing them to `loss_fn`. + Args: n_classes: Number of classes, must be greater than 2 (for 2 classes, use `binary_classification_head`). @@ -76,6 +82,9 @@ def multi_class_head(n_classes, integer within [0, n_classes). If given, labels must be of string type and have any value in `label_vocabulary`. Note that errors will be raised if `label_vocabulary` is not provided but labels are strings. + loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how to + reduce training loss over batch. Defaults to `SUM`. + loss_fn: Optional loss function. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. @@ -83,17 +92,25 @@ def multi_class_head(n_classes, An instance of `_Head` for multi class classification. Raises: - ValueError: if `n_classes`, `metric_class_ids` or `label_keys` is invalid. + ValueError: if `n_classes`, `label_vocabulary` or `loss_reduction` is + invalid. """ return head_lib._multi_class_head_with_softmax_cross_entropy_loss( # pylint:disable=protected-access n_classes=n_classes, weight_column=weight_column, label_vocabulary=label_vocabulary, + loss_reduction=loss_reduction, + loss_fn=loss_fn, name=name) def binary_classification_head( - weight_column=None, thresholds=None, label_vocabulary=None, name=None): + weight_column=None, + thresholds=None, + label_vocabulary=None, + loss_reduction=losses.Reduction.SUM, + loss_fn=None, + name=None): """Creates a `_Head` for single label binary classification. This head uses `sigmoid_cross_entropy_with_logits` loss. @@ -113,6 +130,12 @@ def binary_classification_head( labels have shape `[batch_size, 1]`, the loss is the weighted sum over `batch_size`. + Also supports custom `loss_fn`. `loss_fn` takes `(labels, logits)` or + `(labels, logits, features)` as arguments and returns unreduced loss with + shape `[D0, D1, ... DN, 1]`. `loss_fn` must support float `labels` with + shape `[D0, D1, ... DN, 1]`. Namely, the head applies `label_vocabulary` to + the input labels before passing them to `loss_fn`. + Args: weight_column: A string or a `_NumericColumn` created by `tf.feature_column.numeric_column` defining feature column representing @@ -128,6 +151,9 @@ def binary_classification_head( [0, 1]. If given, labels must be string type and have any value in `label_vocabulary`. Note that errors will be raised if `label_vocabulary` is not provided but labels are strings. + loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how to + reduce training loss over batch. Defaults to `SUM`. + loss_fn: Optional loss function. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. @@ -135,17 +161,22 @@ def binary_classification_head( An instance of `_Head` for binary classification. Raises: - ValueError: if `thresholds` contains a value outside of `(0, 1)`. + ValueError: If `thresholds` contains a value outside of `(0, 1)`. + ValueError: If `loss_reduction` is invalid. """ return head_lib._binary_logistic_head_with_sigmoid_cross_entropy_loss( # pylint:disable=protected-access weight_column=weight_column, thresholds=thresholds, label_vocabulary=label_vocabulary, + loss_reduction=loss_reduction, + loss_fn=loss_fn, name=name) def regression_head(weight_column=None, label_dimension=1, + loss_reduction=losses.Reduction.SUM, + loss_fn=None, name=None): """Creates a `_Head` for regression using the `mean_squared_error` loss. @@ -164,6 +195,10 @@ def regression_head(weight_column=None, `[D0, D1, ... DN]`, `[D0, D1, ... DN, 1]` or `[D0, D1, ... DN, label_dimension]`. + Also supports custom `loss_fn`. `loss_fn` takes `(labels, logits)` or + `(labels, logits, features)` as arguments and returns unreduced loss with + shape `[D0, D1, ... DN, label_dimension]`. + Args: weight_column: A string or a `_NumericColumn` created by `tf.feature_column.numeric_column` defining feature column representing @@ -172,15 +207,23 @@ def regression_head(weight_column=None, label_dimension: Number of regression labels per example. This is the size of the last dimension of the labels `Tensor` (typically, this has shape `[batch_size, label_dimension]`). + loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how to + reduce training loss over batch. Defaults to `SUM`. + loss_fn: Optional loss function. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. Returns: An instance of `_Head` for linear regression. + + Raises: + ValueError: If `label_dimension` or `loss_reduction` is invalid. """ return head_lib._regression_head_with_mean_squared_error_loss( # pylint:disable=protected-access weight_column=weight_column, label_dimension=label_dimension, + loss_reduction=loss_reduction, + loss_fn=loss_fn, name=name) @@ -188,6 +231,7 @@ def multi_label_head(n_classes, weight_column=None, thresholds=None, label_vocabulary=None, + loss_reduction=losses.Reduction.SUM, loss_fn=None, name=None): """Creates a `_Head` for multi-label classification. @@ -202,7 +246,7 @@ def multi_label_head(n_classes, `batch_size`. The head expects `logits` with shape `[D0, D1, ... DN, n_classes]`. In many - applications, the shape is `[batch_size, label_n_classes]`. + applications, the shape is `[batch_size, n_classes]`. Labels can be: * A multi-hot tensor of shape `[D0, D1, ... DN, n_classes]` @@ -237,6 +281,8 @@ def multi_label_head(n_classes, [0, n_classes) or multi-hot Tensor. If given, labels must be SparseTensor string type and have any value in `label_vocabulary`. Also there will be errors if vocabulary is not provided and labels are string. + loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how to + reduce training loss over batch. Defaults to `SUM`. loss_fn: Optional loss function. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. @@ -245,7 +291,8 @@ def multi_label_head(n_classes, An instance of `_Head` for multi-label classification. Raises: - ValueError: if `n_classes`, `thresholds`, or `loss_fn` is invalid. + ValueError: if `n_classes`, `thresholds`, `loss_reduction` or `loss_fn` is + invalid. """ thresholds = tuple(thresholds) if thresholds else tuple() if n_classes is None or n_classes < 2: @@ -266,10 +313,14 @@ def multi_label_head(n_classes, 'Length of label_vocabulary must be n_classes ({}). ' 'Given: {}'.format(n_classes, len(label_vocabulary))) if loss_fn: - _validate_loss_fn_args(loss_fn) + head_lib._validate_loss_fn_args(loss_fn) # pylint:disable=protected-access + if (loss_reduction not in losses.Reduction.all() or + loss_reduction == losses.Reduction.NONE): + raise ValueError('Invalid loss_reduction: {}'.format(loss_reduction)) return _MultiLabelHead( n_classes=n_classes, weight_column=weight_column, thresholds=thresholds, - label_vocabulary=label_vocabulary, loss_fn=loss_fn, name=name) + label_vocabulary=label_vocabulary, loss_reduction=loss_reduction, + loss_fn=loss_fn, name=name) class _MultiLabelHead(head_lib._Head): # pylint:disable=protected-access @@ -280,12 +331,14 @@ class _MultiLabelHead(head_lib._Head): # pylint:disable=protected-access weight_column=None, thresholds=None, label_vocabulary=None, + loss_reduction=losses.Reduction.SUM, loss_fn=None, name=None): self._n_classes = n_classes self._weight_column = weight_column self._thresholds = thresholds self._label_vocabulary = label_vocabulary + self._loss_reduction = loss_reduction self._loss_fn = loss_fn self._name = name @@ -344,9 +397,9 @@ class _MultiLabelHead(head_lib._Head): # pylint:disable=protected-access labels=processed_labels, logits=logits, expected_labels_dimension=self.logits_dimension) if self._loss_fn: - unweighted_loss = _call_loss_fn( + unweighted_loss = head_lib._call_loss_fn( # pylint:disable=protected-access loss_fn=self._loss_fn, labels=processed_labels, logits=logits, - features=features) + features=features, expected_loss_dim=1) else: unweighted_loss = losses.sigmoid_cross_entropy( multi_class_labels=processed_labels, logits=logits, @@ -356,19 +409,41 @@ class _MultiLabelHead(head_lib._Head): # pylint:disable=protected-access unweighted_loss, axis=-1, keep_dims=True) weights = head_lib._get_weights_and_check_match_logits( # pylint:disable=protected-access, features=features, weight_column=self._weight_column, logits=logits) - weighted_sum_loss = losses.compute_weighted_loss( - unweighted_loss, weights=weights, reduction=losses.Reduction.SUM) - # _weights() can return 1. - example_weight_sum = math_ops.reduce_sum( - weights * array_ops.ones_like(unweighted_loss)) + training_loss = losses.compute_weighted_loss( + unweighted_loss, weights=weights, reduction=self._loss_reduction) return head_lib.LossSpec( - weighted_sum_loss=weighted_sum_loss, - example_weight_sum=example_weight_sum, + training_loss=training_loss, + unreduced_loss=unweighted_loss, + weights=weights, processed_labels=processed_labels) def create_estimator_spec( - self, features, mode, logits, labels=None, train_op_fn=None): - """See `Head`.""" + self, features, mode, logits, labels=None, train_op_fn=None, + regularization_losses=None): + """Returns an `EstimatorSpec`. + + Args: + features: Input `dict` of `Tensor` or `SparseTensor` objects. + mode: Estimator's `ModeKeys`. + logits: logits `Tensor` with shape `[D0, D1, ... DN, n_classes]`. + For many applications, the shape is `[batch_size, n_classes]`. + labels: Labels with shape matching `logits`. Can be multi-hot `Tensor` + with shape `[D0, D1, ... DN, n_classes]` or `SparseTensor` with + `dense_shape` `[D0, D1, ... DN, ?]`. `labels` is required argument when + `mode` equals `TRAIN` or `EVAL`. + train_op_fn: Function that takes a scalar loss `Tensor` and returns + `train_op`. Required in TRAIN mode. + regularization_losses: A list of additional scalar losses to be added to + the training loss, such as regularization losses. These losses are + usually expressed as a batch average, so for best results users need to + set `loss_reduction=SUM_OVER_BATCH_SIZE` or + `loss_reduction=SUM_OVER_NONZERO_WEIGHTS` when creating the head to + avoid scaling errors. + Returns: + `EstimatorSpec`. + Raises: + ValueError: If `train_op_fn` is `None` in TRAIN mode. + """ with ops.name_scope(self._name, 'head'): logits = head_lib._check_logits_final_dim(logits, self.logits_dimension) # pylint:disable=protected-access @@ -394,60 +469,74 @@ class _MultiLabelHead(head_lib._Head): # pylint:disable=protected-access export_output.PredictOutput(predictions)) }) - (weighted_sum_loss, example_weight_sum, + (training_loss, unreduced_loss, weights, processed_labels) = self.create_loss( features=features, mode=mode, logits=logits, labels=labels) + if regularization_losses: + regularization_loss = math_ops.add_n(regularization_losses) + regularized_training_loss = math_ops.add_n( + [training_loss, regularization_loss]) + else: + regularization_loss = None + regularized_training_loss = training_loss # Eval. if mode == model_fn.ModeKeys.EVAL: - weights = head_lib._get_weights_and_check_match_logits( # pylint:disable=protected-access, - features=features, weight_column=self._weight_column, logits=logits) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.EVAL, predictions=predictions, - loss=weighted_sum_loss, + loss=regularized_training_loss, eval_metric_ops=self._eval_metric_ops( labels=processed_labels, probabilities=probabilities, weights=weights, - weighted_sum_loss=weighted_sum_loss, - example_weight_sum=example_weight_sum)) + unreduced_loss=unreduced_loss, + regularization_loss=regularization_loss)) # Train. if train_op_fn is None: raise ValueError('train_op_fn can not be None.') + # Only summarize mean_loss for SUM reduction to preserve backwards + # compatibility. Otherwise skip it to avoid unnecessary computation. + if self._loss_reduction == losses.Reduction.SUM: + example_weight_sum = math_ops.reduce_sum( + weights * array_ops.ones_like(unreduced_loss)) + mean_loss = training_loss / example_weight_sum + else: + mean_loss = None with ops.name_scope(''): + keys = metric_keys.MetricKeys summary.scalar( - head_lib._summary_key(self._name, metric_keys.MetricKeys.LOSS), # pylint:disable=protected-access - weighted_sum_loss) - summary.scalar( - head_lib._summary_key( # pylint:disable=protected-access - self._name, metric_keys.MetricKeys.LOSS_MEAN), - weighted_sum_loss / example_weight_sum) + head_lib._summary_key(self._name, keys.LOSS), # pylint:disable=protected-access + regularized_training_loss) + if mean_loss is not None: + summary.scalar( + head_lib._summary_key(self._name, keys.LOSS_MEAN), # pylint:disable=protected-access + mean_loss) + if regularization_loss is not None: + summary.scalar( + head_lib._summary_key(self._name, keys.LOSS_REGULARIZATION), # pylint:disable=protected-access + regularization_loss) return model_fn.EstimatorSpec( mode=model_fn.ModeKeys.TRAIN, predictions=predictions, - loss=weighted_sum_loss, - train_op=train_op_fn(weighted_sum_loss)) + loss=regularized_training_loss, + train_op=train_op_fn(regularized_training_loss)) - def _eval_metric_ops(self, labels, probabilities, weights, weighted_sum_loss, - example_weight_sum): + def _eval_metric_ops( + self, labels, probabilities, weights, unreduced_loss, + regularization_loss): """Returns a dict of metrics for eval_metric_ops.""" with ops.name_scope( None, 'metrics', - [labels, probabilities, weights, weighted_sum_loss, example_weight_sum - ]): + [labels, probabilities, weights, unreduced_loss, regularization_loss]): keys = metric_keys.MetricKeys metric_ops = { # Estimator already adds a metric for loss. head_lib._summary_key(self._name, keys.LOSS_MEAN): # pylint:disable=protected-access metrics_lib.mean( - # Both values and weights here are reduced, scalar Tensors. - # values is the actual mean we want, but we pass the scalar - # example_weight_sum in order to return the correct update_op - # alongside the value_op for streaming metrics. - values=(weighted_sum_loss / example_weight_sum), - weights=example_weight_sum, + values=unreduced_loss, + weights=weights, name=keys.LOSS_MEAN), head_lib._summary_key(self._name, keys.AUC): # pylint:disable=protected-access metrics_lib.auc(labels=labels, predictions=probabilities, @@ -457,6 +546,13 @@ class _MultiLabelHead(head_lib._Head): # pylint:disable=protected-access weights=weights, curve='PR', name=keys.AUC_PR), } + if regularization_loss is not None: + loss_regularization_key = head_lib._summary_key( # pylint:disable=protected-access + self._name, keys.LOSS_REGULARIZATION) + metric_ops[loss_regularization_key] = ( + metrics_lib.mean( + values=regularization_loss, + name=keys.LOSS_REGULARIZATION)) for threshold in self._thresholds: accuracy_key = keys.ACCURACY_AT_THRESHOLD % threshold metric_ops[head_lib._summary_key(self._name, accuracy_key)] = ( # pylint:disable=protected-access @@ -485,52 +581,3 @@ class _MultiLabelHead(head_lib._Head): # pylint:disable=protected-access threshold=threshold, name=recall_key)) return metric_ops - - -def _validate_loss_fn_args(loss_fn): - """Validates loss_fn arguments. - - Required arguments: labels, logits. - Optional arguments: features. - - Args: - loss_fn: The loss function. - Raises: - ValueError: If the signature is unexpected. - """ - loss_fn_args = util.fn_args(loss_fn) - for required_arg in ['labels', 'logits']: - if required_arg not in loss_fn_args: - raise ValueError( - 'loss_fn must contain argument: {}. ' - 'Given arguments: {}'.format(required_arg, loss_fn_args)) - invalid_args = list(set(loss_fn_args) - set(['labels', 'logits', 'features'])) - if invalid_args: - raise ValueError('loss_fn has unexpected args: {}'.format(invalid_args)) - - -def _call_loss_fn(loss_fn, labels, logits, features): - """Calls loss_fn and checks the returned shape. - - Args: - loss_fn: The loss function. - labels: Processed labels Tensor. - logits: Logits Tensor of shape [batch_size, logits_dimension]. - features: Features dict. - Returns: - Loss Tensor with shape [batch_size, 1]. - """ - loss_fn_args = util.fn_args(loss_fn) - kwargs = {} - if 'features' in loss_fn_args: - kwargs['features'] = features - unweighted_loss = loss_fn(labels=labels, logits=logits, **kwargs) - batch_size = array_ops.shape(logits)[0] - loss_shape = array_ops.shape(unweighted_loss) - check_shape_op = control_flow_ops.Assert( - math_ops.reduce_all(math_ops.equal(loss_shape, [batch_size, 1])), - data=[ - 'loss_fn must return Tensor of shape [batch_size, 1]. Given: ', - loss_shape]) - with ops.control_dependencies([check_shape_op]): - return array_ops.identity(unweighted_loss) diff --git a/tensorflow/contrib/estimator/python/estimator/head_test.py b/tensorflow/contrib/estimator/python/estimator/head_test.py index d1cf9090048470181818c573647923c9f5824dfa..1411635228457218578c0297d4d901e9c86ca91a 100644 --- a/tensorflow/contrib/estimator/python/estimator/head_test.py +++ b/tensorflow/contrib/estimator/python/estimator/head_test.py @@ -35,6 +35,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import string_ops +from tensorflow.python.ops.losses import losses from tensorflow.python.platform import test from tensorflow.python.saved_model import signature_constants from tensorflow.python.training import monitored_session @@ -132,6 +133,16 @@ class MultiLabelHead(test.TestCase): r'Length of label_vocabulary must be n_classes \(3\). Given: 2'): head_lib.multi_label_head(n_classes=3, label_vocabulary=['foo', 'bar']) + def test_invalid_loss_reduction(self): + with self.assertRaisesRegexp( + ValueError, r'Invalid loss_reduction: invalid_loss_reduction'): + head_lib.multi_label_head( + n_classes=3, loss_reduction='invalid_loss_reduction') + with self.assertRaisesRegexp( + ValueError, r'Invalid loss_reduction: none'): + head_lib.multi_label_head( + n_classes=3, loss_reduction=losses.Reduction.NONE) + def test_loss_fn_arg_labels_missing(self): def _loss_fn(logits): del logits # Unused @@ -262,17 +273,17 @@ class MultiLabelHead(test.TestCase): labels = np.array([[1, 0], [1, 1]], dtype=np.int64) # loss = labels * -log(sigmoid(logits)) + # (1 - labels) * -log(1 - sigmoid(logits)) - expected_weighted_sum_loss = np.sum( + expected_training_loss = np.sum( _sigmoid_cross_entropy(labels=labels, logits=logits)) - actual_weighted_sum_loss = head.create_loss( + actual_training_loss = head.create_loss( features={'x': np.array(((42,),), dtype=np.int32)}, mode=model_fn.ModeKeys.EVAL, logits=logits, labels=labels)[0] with self.test_session(): _initialize_variables(self, monitored_session.Scaffold()) - self.assertAllClose(expected_weighted_sum_loss, - actual_weighted_sum_loss.eval()) + self.assertAllClose(expected_training_loss, + actual_training_loss.eval()) def test_eval_create_loss_large_logits(self): """Tests head.create_loss for eval mode and large logits.""" @@ -286,9 +297,9 @@ class MultiLabelHead(test.TestCase): # For large logits, this is approximated as: # loss = labels * (logits < 0) * (-logits) + # (1 - labels) * (logits > 0) * logits - expected_weighted_sum_loss = np.sum( + expected_training_loss = np.sum( np.array([[(10. + 10.) / 2.], [(15. + 0.) / 2.]], dtype=np.float32)) - actual_weighted_sum_loss = head.create_loss( + actual_training_loss = head.create_loss( features={'x': np.array(((42,),), dtype=np.int32)}, mode=model_fn.ModeKeys.EVAL, logits=logits, @@ -296,9 +307,7 @@ class MultiLabelHead(test.TestCase): with self.test_session(): _initialize_variables(self, monitored_session.Scaffold()) self.assertAllClose( - expected_weighted_sum_loss, - actual_weighted_sum_loss.eval(), - atol=1e-4) + expected_training_loss, actual_training_loss.eval(), atol=1e-4) def test_eval_create_loss_labels_wrong_shape(self): """Tests head.create_loss for eval mode when labels has the wrong shape.""" @@ -307,7 +316,7 @@ class MultiLabelHead(test.TestCase): logits = np.array([[-1., 1.], [-1.5, 1.]], dtype=np.float32) labels_placeholder = array_ops.placeholder(dtype=dtypes.int64) - actual_weighted_sum_loss = head.create_loss( + actual_training_loss = head.create_loss( features={'x': np.array(((42,),), dtype=np.int32)}, mode=model_fn.ModeKeys.EVAL, logits=logits, @@ -317,14 +326,14 @@ class MultiLabelHead(test.TestCase): with self.assertRaisesRegexp( errors.InvalidArgumentError, r'\[expected_labels_shape: \] \[2 2\] \[labels_shape: \] \[2 1\]'): - actual_weighted_sum_loss.eval({ + actual_training_loss.eval({ labels_placeholder: np.array([[1], [1]], dtype=np.int64) }) with self.assertRaisesRegexp( errors.InvalidArgumentError, r'labels shape must be \[D0, D1, ... DN, 2\]\..*' r'\[Received shape: \] \[2\]'): - actual_weighted_sum_loss.eval({ + actual_training_loss.eval({ labels_placeholder: np.array([1, 1], dtype=np.int64) }) @@ -344,14 +353,14 @@ class MultiLabelHead(test.TestCase): return constant_op.constant(loss) head = head_lib.multi_label_head(n_classes=2, loss_fn=_loss_fn) - actual_weighted_sum_loss = head.create_loss( + actual_training_loss = head.create_loss( features={'x': np.array(((42,),), dtype=np.int32)}, mode=model_fn.ModeKeys.EVAL, logits=logits_input, labels=labels_input)[0] with self.test_session(): _initialize_variables(self, monitored_session.Scaffold()) - self.assertAllClose(np.sum(loss), actual_weighted_sum_loss.eval()) + self.assertAllClose(np.sum(loss), actual_training_loss.eval()) def test_eval_create_loss_loss_fn_wrong_shape(self): """Tests custom loss_fn that returns Tensor of unexpected shape.""" @@ -363,7 +372,7 @@ class MultiLabelHead(test.TestCase): logits = np.array([[-10., 10.], [-15., 10.]], dtype=np.float32) labels = np.array([[1, 0], [1, 1]], dtype=np.int64) - actual_weighted_sum_loss = head.create_loss( + actual_training_loss = head.create_loss( features={'x': np.array(((42,),), dtype=np.int32)}, mode=model_fn.ModeKeys.EVAL, logits=logits, @@ -372,9 +381,9 @@ class MultiLabelHead(test.TestCase): _initialize_variables(self, monitored_session.Scaffold()) with self.assertRaisesRegexp( errors.InvalidArgumentError, - r'loss_fn must return Tensor of shape \[batch_size, 1\]\. ' - r'Given: \] \[2\]'): - actual_weighted_sum_loss.eval() + r'\[loss_fn must return Tensor of shape \[D0, D1, ... DN, 1\]\. \] ' + r'\[logits_shape: \] \[2 2\] \[loss_shape: \] \[2\]'): + actual_training_loss.eval() def test_eval_labels_none(self): """Tests that error is raised when labels is None.""" @@ -390,12 +399,13 @@ class MultiLabelHead(test.TestCase): def _test_eval( self, head, logits, labels, expected_loss, expected_metrics, - features=None): + features=None, regularization_losses=None): spec = head.create_estimator_spec( features=features or {}, mode=model_fn.ModeKeys.EVAL, logits=logits, - labels=labels) + labels=labels, + regularization_losses=regularization_losses) # Assert spec contains expected tensors. self.assertIsNotNone(spec.loss) @@ -436,7 +446,7 @@ class MultiLabelHead(test.TestCase): # auc and auc_pr cannot be reliably calculated for only 4 samples, but # this assert tests that the algorithm remains consistent. keys.AUC: 0.3333, - keys.AUC_PR: 0.7639, + keys.AUC_PR: 0.5972, } self._test_eval( head=head, @@ -468,7 +478,7 @@ class MultiLabelHead(test.TestCase): # auc and auc_pr cannot be reliably calculated for only 4 samples, but # this assert tests that the algorithm remains consistent. keys.AUC: 0.3333, - keys.AUC_PR: 0.7639, + keys.AUC_PR: 0.5972, } self._test_eval( head=head, @@ -477,6 +487,38 @@ class MultiLabelHead(test.TestCase): expected_loss=expected_loss, expected_metrics=expected_metrics) + def test_eval_with_regularization_losses(self): + n_classes = 2 + head = head_lib.multi_label_head( + n_classes, loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE) + logits = np.array([[-1., 1.], [-1.5, 1.5]], dtype=np.float32) + labels = np.array([[1, 0], [1, 1]], dtype=np.int64) + regularization_losses = [1.5, 0.5] + expected_regularization_loss = 2. + # unregularized_loss = sum( + # labels * -log(sigmoid(logits)) + + # (1 - labels) * -log(1 - sigmoid(logits))) / batch_size + expected_unregularized_loss = np.sum( + _sigmoid_cross_entropy(labels=labels, logits=logits)) / 2. + expected_regularized_loss = ( + expected_unregularized_loss + expected_regularization_loss) + keys = metric_keys.MetricKeys + expected_metrics = { + keys.LOSS_MEAN: expected_unregularized_loss, + keys.LOSS_REGULARIZATION: expected_regularization_loss, + # auc and auc_pr cannot be reliably calculated for only 4 samples, but + # this assert tests that the algorithm remains consistent. + keys.AUC: 0.3333, + keys.AUC_PR: 0.5972, + } + self._test_eval( + head=head, + logits=logits, + labels=labels, + expected_loss=expected_regularized_loss, + expected_metrics=expected_metrics, + regularization_losses=regularization_losses) + def test_eval_with_label_vocabulary(self): n_classes = 2 head = head_lib.multi_label_head( @@ -501,7 +543,7 @@ class MultiLabelHead(test.TestCase): # auc and auc_pr cannot be reliably calculated for only 4 samples, but # this assert tests that the algorithm remains consistent. keys.AUC: 0.3333, - keys.AUC_PR: 0.7639, + keys.AUC_PR: 0.5972, } self._test_eval( head=head, @@ -531,7 +573,7 @@ class MultiLabelHead(test.TestCase): # auc and auc_pr cannot be reliably calculated for only 4 samples, but # this assert tests that the algorithm remains consistent. keys.AUC: 0.3333, - keys.AUC_PR: 0.7639, + keys.AUC_PR: 0.5972, keys.ACCURACY_AT_THRESHOLD % thresholds[0]: 2. / 4., keys.PRECISION_AT_THRESHOLD % thresholds[0]: 2. / 3., keys.RECALL_AT_THRESHOLD % thresholds[0]: 2. / 3., @@ -579,7 +621,7 @@ class MultiLabelHead(test.TestCase): # auc and auc_pr cannot be reliably calculated for only 4 samples, but # this assert tests that the algorithm remains consistent. keys.AUC: 0.2000, - keys.AUC_PR: 0.7833, + keys.AUC_PR: 0.5833, } # Assert spec contains expected tensors. @@ -618,12 +660,10 @@ class MultiLabelHead(test.TestCase): # For large logits, this is approximated as: # loss = labels * (logits < 0) * (-logits) + # (1 - labels) * (logits > 0) * logits - expected_weighted_sum_loss = np.sum( - np.array( - [[1. * (10. + 10.) / 2.], [2. * (15. + 0.) / 2.]], - dtype=np.float32)) - expected_example_weight_sum = 1. + 2. - actual_weighted_sum_loss, actual_example_weight_sum, _ = head.create_loss( + expected_unreduced_loss = [[(10. + 10.) / 2.], [(15. + 0.) / 2.]] + expected_weights = [[1.], [2.]] + expected_training_loss = 1. * (10. + 10.) / 2. + 2. * (15. + 0.) / 2. + training_loss, unreduced_loss, actual_weights, _ = head.create_loss( features={ 'x': np.array(((42,),), dtype=np.int32), 'example_weights': weights @@ -634,13 +674,44 @@ class MultiLabelHead(test.TestCase): with self.test_session(): _initialize_variables(self, monitored_session.Scaffold()) self.assertAllClose( - expected_weighted_sum_loss, - actual_weighted_sum_loss.eval(), - atol=1e-4) + expected_training_loss, training_loss.eval(), atol=1e-4) self.assertAllClose( - expected_example_weight_sum, - actual_example_weight_sum.eval(), - atol=1e-4) + expected_unreduced_loss, unreduced_loss.eval(), atol=1e-4) + self.assertAllClose(expected_weights, actual_weights.eval()) + + def test_train_create_loss_loss_reduction(self): + """Tests head.create_loss with loss_reduction.""" + n_classes = 2 + head = head_lib.multi_label_head( + n_classes, weight_column='example_weights', + loss_reduction=losses.Reduction.SUM_BY_NONZERO_WEIGHTS) + + logits = np.array([[-10., 10.], [-15., 10.]], dtype=np.float32) + labels = np.array([[1, 0], [1, 1]], dtype=np.int64) + weights = np.array([[1.], [2.]], dtype=np.float32) + # loss = labels * -log(sigmoid(logits)) + + # (1 - labels) * -log(1 - sigmoid(logits)) + # For large logits, this is approximated as: + # loss = labels * (logits < 0) * (-logits) + + # (1 - labels) * (logits > 0) * logits + expected_unreduced_loss = [[(10. + 10.) / 2.], [(15. + 0.) / 2.]] + expected_weights = [[1.], [2.]] + expected_training_loss = (1. * (10. + 10.) / 2. + 2. * (15. + 0.) / 2.) / 2. + training_loss, unreduced_loss, actual_weights, _ = head.create_loss( + features={ + 'x': np.array(((42,),), dtype=np.int32), + 'example_weights': weights + }, + mode=model_fn.ModeKeys.TRAIN, + logits=logits, + labels=labels) + with self.test_session(): + _initialize_variables(self, monitored_session.Scaffold()) + self.assertAllClose( + expected_training_loss, training_loss.eval(), atol=1e-4) + self.assertAllClose( + expected_unreduced_loss, unreduced_loss.eval(), atol=1e-4) + self.assertAllClose(expected_weights, actual_weights.eval()) def test_train_labels_none(self): """Tests that error is raised when labels is None.""" @@ -791,6 +862,49 @@ class MultiLabelHead(test.TestCase): self._test_train( head=head, logits=logits, labels=labels, expected_loss=expected_loss) + def test_train_with_regularization_losses(self): + head = head_lib.multi_label_head( + n_classes=2, loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE) + logits = np.array([[-10., 10.], [-15., 10.]], dtype=np.float32) + labels = np.array([[1, 0], [1, 1]], dtype=np.int64) + regularization_losses = [1.5, 0.5] + # For large logits, sigmoid cross entropy loss is approximated as: + # loss = labels * (logits < 0) * (-logits) + + # (1 - labels) * (logits > 0) * logits => + # expected_unweighted_loss = [[10., 10.], [15., 0.]] + # Average over classes and over batch and add regularization loss. + expected_loss = 35. / 4. + 2. + expected_summaries = { + metric_keys.MetricKeys.LOSS: expected_loss, + metric_keys.MetricKeys.LOSS_REGULARIZATION: 2., + } + expected_train_result = 'my_train_op' + def _train_op_fn(loss): + return string_ops.string_join( + [constant_op.constant(expected_train_result), + string_ops.as_string(loss, precision=3)]) + + spec = head.create_estimator_spec( + features={'x': np.array(((42,),), dtype=np.int32)}, + mode=model_fn.ModeKeys.TRAIN, + logits=logits, + labels=labels, + train_op_fn=_train_op_fn, + regularization_losses=regularization_losses) + + # Assert predictions, loss, train_op, and summaries. + tol = 1e-3 + with self.test_session() as sess: + _initialize_variables(self, spec.scaffold) + self.assertIsNotNone(spec.scaffold.summary_op) + loss, train_result, summary_str = sess.run((spec.loss, spec.train_op, + spec.scaffold.summary_op)) + self.assertAllClose(expected_loss, loss, rtol=tol, atol=tol) + self.assertEqual( + six.b('{0:s}{1:.3f}'.format(expected_train_result, expected_loss)), + train_result) + _assert_simple_summaries(self, expected_summaries, summary_str, tol) + def test_train_with_weights(self): n_classes = 2 head = head_lib.multi_label_head(n_classes, weight_column='example_weights') @@ -851,12 +965,15 @@ class MultiLabelHead(test.TestCase): labels = np.array([[[1, 0, 0], [1, 0, 0]], [[0, 1, 1], [0, 1, 1]]], dtype=np.int64) weights = np.array([[1., 1.5], [2., 2.5]], dtype=np.float32) - # loss = [[10 + 10 + 0, 0 + 0 + 10], [0 + 0 + 12, 12 + 12 + 0]] / 3 - # = [[20/3, 10/3], [4, 8]] + # unreduced_loss = + # [[10 + 10 + 0, 0 + 0 + 10], [0 + 0 + 12, 12 + 12 + 0]] / 3 + # = [[20/3, 10/3], [4, 8]] + expected_unreduced_loss = [[[20./3.], [10./3.]], [[4.], [8.]]] + # weights are reshaped to [2, 2, 1] to match logits. + expected_weights = [[[1.], [1.5]], [[2.], [2.5]]] # weighted_sum_loss = 1*20/3 + 1.5*10/3 + 2*4 + 2.5*8 = 39.6667 - expected_weighted_sum_loss = 39.6667 - expected_example_weight_sum = np.sum(weights) - actual_weighted_sum_loss, actual_example_weight_sum, _ = head.create_loss( + expected_training_loss = 39.6667 + training_loss, unreduced_loss, actual_weights, _ = head.create_loss( features={'weights': weights}, mode=model_fn.ModeKeys.TRAIN, logits=logits, @@ -865,11 +982,10 @@ class MultiLabelHead(test.TestCase): with self.test_session(): _initialize_variables(self, monitored_session.Scaffold()) self.assertAllClose( - expected_weighted_sum_loss, actual_weighted_sum_loss.eval(), - atol=atol) + expected_training_loss, training_loss.eval(), atol=atol) self.assertAllClose( - expected_example_weight_sum, actual_example_weight_sum.eval(), - atol=atol) + expected_unreduced_loss, unreduced_loss.eval(), atol=atol) + self.assertAllClose(expected_weights, actual_weights.eval()) def test_multi_dim_weighted_train(self): """Logits and labels of shape [2, 2, 3], weights [2, 2].""" @@ -979,7 +1095,7 @@ class MultiLabelHead(test.TestCase): # auc and auc_pr cannot be reliably calculated for only 4 samples, but # this assert tests that the algorithm remains consistent. keys.AUC: 0.4977, - keys.AUC_PR: 0.6645, + keys.AUC_PR: 0.4037, } self._test_eval( head=head, diff --git a/tensorflow/contrib/estimator/python/estimator/multi_head.py b/tensorflow/contrib/estimator/python/estimator/multi_head.py index f2a6eae03ec021e5c28d48b3887870d8a057e077..0346ddc24bffd61068177f4622bd03be4acd53d9 100644 --- a/tensorflow/contrib/estimator/python/estimator/multi_head.py +++ b/tensorflow/contrib/estimator/python/estimator/multi_head.py @@ -186,40 +186,44 @@ class _MultiHead(head_lib._Head): # pylint:disable=protected-access logits_dict = logits else: logits_dict = self._split_logits(logits) - weighted_sum_losses = [] - example_weight_sums = [] + training_losses = [] labels_by_head = {} - for head in self._heads: - (weighted_sum_loss, - example_weight_sum, processed_labels) = head.create_loss( + unreduced_losses_by_head = {} + example_weights_by_head = {} + for i, head in enumerate(self._heads): + (training_loss, unreduced_loss, + weights, processed_labels) = head.create_loss( features, mode, logits_dict[head.name], labels[head.name]) - weighted_sum_losses.append(weighted_sum_loss) - example_weight_sums.append(example_weight_sum) + training_losses.append(training_loss) labels_by_head[head.name] = processed_labels + if self._head_weights: + head_weight = self._head_weights[i] + unreduced_losses_by_head[head.name] = math_ops.multiply( + unreduced_loss, head_weight) + example_weights_by_head[head.name] = math_ops.multiply( + weights, head_weight) + else: + unreduced_losses_by_head[head.name] = unreduced_loss + example_weights_by_head[head.name] = weights - weighted_sum_losses = tuple(weighted_sum_losses) - with ops.name_scope('merge_losses', - values=weighted_sum_losses + (self._head_weights or - tuple())): + training_losses = tuple(training_losses) + with ops.name_scope( + 'merge_losses', + values=training_losses + (self._head_weights or tuple())): if self._head_weights: - head_weighted_losses = [] - head_weighted_example_weight_sums = [] - for loss, example_weight_sum, weight in zip(weighted_sum_losses, - example_weight_sums, - self._head_weights): - head_weighted_losses.append(math_ops.multiply(loss, weight)) - head_weighted_example_weight_sums.append(math_ops.multiply( - example_weight_sum, weight)) - merged_weighted_sum_loss = math_ops.add_n(head_weighted_losses) - merged_example_weight_sum = math_ops.add_n( - head_weighted_example_weight_sums) + head_weighted_training_losses = [] + for training_loss, head_weight in zip( + training_losses, self._head_weights): + head_weighted_training_losses.append( + math_ops.multiply(training_loss, head_weight)) + merged_training_loss = math_ops.add_n(head_weighted_training_losses) else: - merged_weighted_sum_loss = math_ops.add_n(weighted_sum_losses) - merged_example_weight_sum = math_ops.add_n(example_weight_sums) + merged_training_loss = math_ops.add_n(training_losses) return head_lib.LossSpec( - weighted_sum_loss=merged_weighted_sum_loss, - example_weight_sum=merged_example_weight_sum, + training_loss=merged_training_loss, + unreduced_loss=unreduced_losses_by_head, + weights=example_weights_by_head, processed_labels=labels_by_head) def create_estimator_spec( diff --git a/tensorflow/contrib/estimator/python/estimator/multi_head_test.py b/tensorflow/contrib/estimator/python/estimator/multi_head_test.py index 68f2d5d1cd53456f7dd82222e171b3619052321a..e47a6788f3b5440c4906b9f0430c802cf73237e3 100644 --- a/tensorflow/contrib/estimator/python/estimator/multi_head_test.py +++ b/tensorflow/contrib/estimator/python/estimator/multi_head_test.py @@ -306,8 +306,8 @@ class MultiHeadTest(test.TestCase): # this assert tests that the algorithm remains consistent. keys.AUC + '/head1': 0.1667, keys.AUC + '/head2': 0.3333, - keys.AUC_PR + '/head1': 0.6667, - keys.AUC_PR + '/head2': 0.5000, + keys.AUC_PR + '/head1': 0.49999964, + keys.AUC_PR + '/head2': 0.33333313, } # Assert spec contains expected tensors. @@ -370,7 +370,7 @@ class MultiHeadTest(test.TestCase): 'head1': np.array([[1, 0], [1, 1]], dtype=np.int64), 'head2': np.array([[0, 1, 0], [1, 1, 0]], dtype=np.int64), } - weighted_sum_loss, example_weight_sum, _ = multi_head.create_loss( + training_loss, unreduced_losses, weights, _ = multi_head.create_loss( features={ 'x': np.array(((42,),), dtype=np.int32), 'weights1': weights1, @@ -383,14 +383,23 @@ class MultiHeadTest(test.TestCase): with self.test_session(): # loss of the first head is [[(10 + 10) / 2], [(15 + 0) / 2]] # = [10, 7.5] - # weighted_sum_loss = 1 * 10 + 2 * 7.5 = 25 + # training_loss = 1 * 10 + 2 * 7.5 = 25 + # head-weighted unreduced_loss = 1 * [10, 7.5] + self.assertAllClose( + [[10.], [7.5]], unreduced_losses['head1'].eval(), rtol=tol, atol=tol) # loss of the second head is [[(20 + 20 + 20) / 3], [(30 + 0 + 0) / 3]] # = [20, 10] - # weighted_sum_loss = 2 * 20 + 3 * 10 = 70 - # head-weighted merge = 1 * 25 + 2 * 70 = 165 - self.assertAllClose(165, weighted_sum_loss.eval(), rtol=tol, atol=tol) - # example_weight_sum = 1 * (1 + 2) + 2 * (2 + 3) = 13 - self.assertAllClose(13., example_weight_sum.eval(), rtol=tol, atol=tol) + # training_loss = 2 * 20 + 3 * 10 = 70 + # head-weighted unreduced_loss = 2 * [20, 10] + self.assertAllClose( + [[40.], [20.]], unreduced_losses['head2'].eval(), rtol=tol, atol=tol) + # head-weighted training_loss = 1 * 25 + 2 * 70 = 165 + self.assertAllClose(165, training_loss.eval(), rtol=tol, atol=tol) + # head-weighted example weights + self.assertAllClose( + [[1.], [2.]], weights['head1'].eval(), rtol=tol, atol=tol) + self.assertAllClose( + [[4.], [6.]], weights['head2'].eval(), rtol=tol, atol=tol) def test_train_create_loss_logits_tensor(self): """Tests create_loss with logits Tensor.""" @@ -409,7 +418,7 @@ class MultiHeadTest(test.TestCase): 'head1': np.array([[1, 0], [1, 1]], dtype=np.int64), 'head2': np.array([[0, 1, 0], [1, 1, 0]], dtype=np.int64), } - weighted_sum_loss, example_weight_sum, _ = multi_head.create_loss( + training_loss, unreduced_losses, weights, _ = multi_head.create_loss( features={ 'x': np.array(((42,),), dtype=np.int32), 'weights1': weights1, @@ -422,14 +431,23 @@ class MultiHeadTest(test.TestCase): with self.test_session(): # loss of the first head is [[(10 + 10) / 2], [(15 + 0) / 2]] # = [10, 7.5] - # weighted_sum_loss = 1 * 10 + 2 * 7.5 = 25 + # training_loss = 1 * 10 + 2 * 7.5 = 25 + # head-weighted unreduced_loss = 1 * [10, 7.5] + self.assertAllClose( + [[10.], [7.5]], unreduced_losses['head1'].eval(), rtol=tol, atol=tol) # loss of the second head is [[(20 + 20 + 20) / 3], [(30 + 0 + 0) / 3]] # = [20, 10] - # weighted_sum_loss = 2 * 20 + 3 * 10 = 70 - # head-weighted merge = 1 * 25 + 2 * 70 = 165 - self.assertAllClose(165, weighted_sum_loss.eval(), rtol=tol, atol=tol) - # example_weight_sum = 1 * (1 + 2) + 2 * (2 + 3) = 13 - self.assertAllClose(13., example_weight_sum.eval(), rtol=tol, atol=tol) + # training_loss = 2 * 20 + 3 * 10 = 70 + # head-weighted unreduced_loss = 2 * [20, 10] + self.assertAllClose( + [[40.], [20.]], unreduced_losses['head2'].eval(), rtol=tol, atol=tol) + # head-weighted training_loss = 1 * 25 + 2 * 70 = 165 + self.assertAllClose(165, training_loss.eval(), rtol=tol, atol=tol) + # head-weighted example weights + self.assertAllClose( + [[1.], [2.]], weights['head1'].eval(), rtol=tol, atol=tol) + self.assertAllClose( + [[4.], [6.]], weights['head2'].eval(), rtol=tol, atol=tol) def test_train_create_loss_logits_tensor_multi_dim(self): """Tests create_loss with multi-dimensional logits of shape [2, 2, 5].""" @@ -455,20 +473,17 @@ class MultiHeadTest(test.TestCase): # loss2 = (0-2)^2 + (1+2)^2 + (0-2)^2 + (0-2)^2 + (1+2)^2 + (0-2)^2 + # (2+2)^2 + (2-2)^2 + (0+2)^2 + (2+2)^2 + (2-2)^2 + (0+2)^2 # = 74 - expected_weighted_sum_loss = 28. + 74. + expected_training_loss = 28. + 74. - weighted_sum_loss, example_weight_sum, _ = multi_head.create_loss( + training_loss = multi_head.create_loss( features={}, mode=model_fn.ModeKeys.TRAIN, logits=logits, - labels=labels) + labels=labels)[0] tol = 1e-3 with self.test_session(): self.assertAllClose( - expected_weighted_sum_loss, weighted_sum_loss.eval(), - rtol=tol, atol=tol) - self.assertAllClose( - 2. * 2. * 5., example_weight_sum.eval(), rtol=tol, atol=tol) + expected_training_loss, training_loss.eval(), rtol=tol, atol=tol) def test_train_one_head(self): head1 = head_lib.multi_label_head(n_classes=2, name='head1') diff --git a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py index f4caa60460dbd4a9149afb45ab4d4e0148f76bca..7134cd3f5a457a322f51066eb791133c3181d3fb 100644 --- a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py +++ b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py @@ -23,6 +23,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from collections import defaultdict from contextlib import contextmanager import copy @@ -49,9 +50,9 @@ from tensorflow.python.training import optimizer as optimizer_lib def replicate_model_fn(model_fn, - loss_reduction=losses.Reduction.SUM, + loss_reduction=losses.Reduction.SUM_BY_NONZERO_WEIGHTS, devices=None): - """Replicate `Estimator.model_fn` over GPUs within a single host. + """Replicate `Estimator.model_fn` over GPUs. The given `model_fn` specifies a single forward pass of a model. To replicate such a model over GPUs, each GPU gets its own instance of the forward pass @@ -68,7 +69,7 @@ def replicate_model_fn(model_fn, Two modes of local replication over available GPUs are supported: 1) If exactly 1 GPU is detected, then variables and operations are placed - onto GPU. + onto the GPU. 2) If more than 1 GPU is detected, then variables are going to be placed on the CPU. Replicas of operations are placed on each individual GPU. @@ -78,7 +79,7 @@ def replicate_model_fn(model_fn, def model_fn(...): # See `model_fn` in `Estimator`. loss = ... optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001) - optimizer = tf.contrib.estimator.GatheringOptimizer(optimizer) + optimizer = tf.contrib.estimator.TowerOptimizer(optimizer) if mode == tf.estimator.ModeKeys.TRAIN: # See the section below on `EstimatorSpec.train_op`. return EstimatorSpec(mode=mode, loss=loss, @@ -91,12 +92,15 @@ def replicate_model_fn(model_fn, model_fn=tf.contrib.estimator.replicate_model_fn(model_fn)) ``` + Please see `DNNClassifierIntegrationTest` for an example with a canned + Estimator. + On `EstimatorSpec.train_op`: `model_fn` returns `EstimatorSpec.train_op` for `tf.estimator.GraphKeys.TRAIN`. It is typically derived using an optimizer. Towers are expected to populate it in the same way. Gradients from all towers - are reduced and applied in the last tower. To achieve that, - `GatheringOptimizer` needs to be used. See `GatheringOptimizer`. + are reduced and applied in the last tower. To achieve that in the case of + multiple towers, `TowerOptimizer` needs to be used. See `TowerOptimizer`. On sharding input features and labels: Input features and labels are split for consumption by each tower. They are @@ -117,9 +121,14 @@ def replicate_model_fn(model_fn, Variables are not duplicated between towers. Instead, they are placed on a single device as defined above and shared across towers. - Other current limitations: - - `predictions` are not supported for `ModeKeys.EVAL`. That is required for - `tf.contrib.estimator.add_metrics`. + On overhead: + If only one device is specified, then aggregation of loss and gradients + doesn't happen. Replication consists of placing `model_fn` onto the + specified device. + + On current limitations: + - `predictions` are not supported for `ModeKeys.EVAL`. They are required + for `tf.contrib.estimator.add_metrics`. Args: model_fn: `model_fn` as defined in `Estimator`. See the section above about @@ -130,6 +139,10 @@ def replicate_model_fn(model_fn, If `None`, then all available GPUs are going to be used for replication. If no GPUs are available, then the model is going to be placed on the CPU. + Raises: + ValueError: if there is no `loss_reduction` or if TowerOptimizer is + mis-used. + Returns: A replicated version of the supplied `model_fn`. Returned function that conforms to the requirements of `Estimator`'s `model_fn` and can be used @@ -172,7 +185,7 @@ class _VariableDistributionMode(object): def _replicate_model_fn_with_mode( model_fn, - loss_reduction=losses.Reduction.SUM, + loss_reduction, devices=None, mode=_VariableDistributionMode.SHARED_LOCAL_PARAMETER_SERVER): """A version of `replicate_model_fn` that allows to specify a `mode`.""" @@ -182,9 +195,8 @@ def _replicate_model_fn_with_mode( if not devices: devices = _get_local_devices('GPU') or _get_local_devices('CPU') - is_a_single_gpu_case = len(devices) == 1 and 'GPU' in devices[0] - consolidation_device = '/{}:0'.format('GPU' - if is_a_single_gpu_case else 'CPU') + is_a_single_gpu_case = len(devices) == 1 and 'GPU' in devices[0].upper() + consolidation_device = devices[0] if is_a_single_gpu_case else '/CPU:0' ps_devices = [consolidation_device] if mode == _VariableDistributionMode.SHARED_ROUND_ROBIN: @@ -194,6 +206,19 @@ def _replicate_model_fn_with_mode( 'to be placed on {}. Consolidation device is going to be {}.' .format(devices, ps_devices, consolidation_device)) + def single_device_model_fn(features, labels, mode, params=None, config=None): + """`model_fn` on a single device without reduction overhead.""" + return _get_loss_towers( + model_fn=model_fn, + mode=mode, + features=[features], + labels=[labels], + params=params, + loss_reduction=loss_reduction, + config=config, + devices=devices, + local_ps_devices=ps_devices)[0] # One device, so one spec is out. + def replicated_model_fn(features, labels, mode, params=None, config=None): """Replicated version of `model_fn` to be used instead.""" feature_shards, label_shards = _split_batch( @@ -218,10 +243,13 @@ def _replicate_model_fn_with_mode( elif mode == model_fn_lib.ModeKeys.PREDICT: return _predict_spec(tower_specs, aggregation_device=consolidation_device) - return replicated_model_fn + if len(devices) == 1: + return single_device_model_fn + else: + return replicated_model_fn -class GatheringOptimizer(optimizer_lib.Optimizer): +class TowerOptimizer(optimizer_lib.Optimizer): """Gathers gradients from all towers and reduces them in the last one.""" COLLECTION_FOR_GRAPH_STATES = 'replicate_model_fn_graph_states' @@ -229,11 +257,18 @@ class GatheringOptimizer(optimizer_lib.Optimizer): def __init__(self, optimizer_or_optimizer_fn): """Wrap an existing optimizer for gathering gradients across towers. - Each invocation of model_fn has to call optimizers in the same order. + Each invocation of model_fn has to call the same optimizers in the same + order. Multiple optimizers that use the same or different losses are supported. - Optimizers, however, need to be of different type as per `__class__` - in order to increment the global_step correctly. + + If TowerOptimizer is used but `replicate_model_fn` isn't, then no + aggregation will happen. All calls will simply be forwarded to the + underlying optimizer. The behavior is similar if there is only one tower. + + If TowerOptimizer is used together with SyncReplicasOptimizer that wraps + the user's optimizer, then it's the SyncReplicasOptimizer that needs to be + wrapped with TowerOptimizer. Args: optimizer_or_optimizer_fn: an instance of optimizer to wrap. That @@ -244,7 +279,7 @@ class GatheringOptimizer(optimizer_lib.Optimizer): @staticmethod def has_been_used(): - return GatheringOptimizer._graph_state().has_gathering_optimizer_been_used + return TowerOptimizer._graph_state().has_tower_optimizer_been_used def get_slot(self, *args, **kwargs): return self._get_optimizer().get_slot(*args, **kwargs) @@ -267,11 +302,20 @@ class GatheringOptimizer(optimizer_lib.Optimizer): def apply_gradients(self, grads_and_vars, global_step=None, **kwargs): """Collect gradients updates to apply them with the last tower.""" - self._graph_state().collect_gradients(grads_and_vars, - self._get_optimizer()) + if self._graph_state().number_of_towers == 1: + # Avoid the overhead of reduction if there's only one tower. + # + # There assumed to be only one tower if aggregation-related methods were + # not called by `_get_loss_towers`, for example if the model_fn uses + # TowerEstimator, but `replicate_model_fn` isn't used. + return self._get_optimizer().apply_gradients(grads_and_vars, global_step, + **kwargs) + + self._graph_state().collect_gradients(grads_and_vars) if not self._graph_state().is_the_last_tower: - return self._construct_no_op_train_op() + with ops_lib.control_dependencies(_extract_tensors(grads_and_vars)): + return self._construct_no_op_train_op() else: # Gradients need to be gathered and applied in the scope of the first # tower, so that the tensors are accessible via names without prefixes. @@ -283,11 +327,9 @@ class GatheringOptimizer(optimizer_lib.Optimizer): def _apply_gathered_gradients(self, global_step, **kwargs): graph_state = self._graph_state() optimizer = self._get_optimizer() - train_ops = [] grad_lists = {} - # Only aggregate gradients for `optimizer.__class__` type of Optimizer. - for grad, var in graph_state.get_grad_and_vars_for_optimizer(optimizer): + for grad, var in graph_state.get_latest_gradients_from_all_towers(): if grad is not None: grad_lists.setdefault(var, []).append(grad) @@ -296,24 +338,16 @@ class GatheringOptimizer(optimizer_lib.Optimizer): for var, grads in six.iteritems(grad_lists): grad = _compute_sum_on_device(grads, var.device) aggregated_grads.append((grad, var)) - train_ops.append(optimizer.apply_gradients(aggregated_grads)) - - # A model might use multiple optimizers. We only want to increment global - # step after apply_gradients of the last optimizer inside the tower. - if global_step and graph_state.is_the_last_optimizer_within_a_tower( - optimizer): - with ops_lib.control_dependencies(train_ops): - with ops_lib.colocate_with(global_step): - return state_ops.assign_add(global_step, 1) - else: - return control_flow_ops.group(train_ops) + return optimizer.apply_gradients( + aggregated_grads, global_step=global_step, **kwargs) def _get_optimizer(self): - if not isinstance(self._optimizer_or_optimizer_fn, optimizer_lib.Optimizer): + if callable(self._optimizer_or_optimizer_fn): # If optimizer is given as a function then we need to wait till we are - # under the right graph context before constructing it. + # under the right graph context before constructing it. That's why the + # optimizer is constructed in _get_optimizer() rather than __init__(). self._optimizer_or_optimizer_fn = self._optimizer_or_optimizer_fn() - self._graph_state().has_gathering_optimizer_been_used = True + self._graph_state().has_tower_optimizer_been_used = True return self._optimizer_or_optimizer_fn def _construct_no_op_train_op(self): @@ -322,44 +356,50 @@ class GatheringOptimizer(optimizer_lib.Optimizer): @staticmethod def _graph_state(): graph_states = ops_lib.get_default_graph().get_collection_ref( - GatheringOptimizer.COLLECTION_FOR_GRAPH_STATES) + TowerOptimizer.COLLECTION_FOR_GRAPH_STATES) if not graph_states: - graph_states.append(GatheringOptimizer._PerGraphState()) + graph_states.append(TowerOptimizer._PerGraphState()) return graph_states[-1] + @staticmethod + def _did_towers_have_same_optimizer_calls(): + graph_state = TowerOptimizer._graph_state() + return graph_state.did_towers_have_same_optimizer_calls() + @staticmethod def _clear_graph_state(): - # Clearing a collection in Graph will prevent _PerGraphState from being + # Clearing the Graph collection will prevent _PerGraphState from being # serialized. ops_lib.get_default_graph().clear_collection( - GatheringOptimizer.COLLECTION_FOR_GRAPH_STATES) + TowerOptimizer.COLLECTION_FOR_GRAPH_STATES) class _PerGraphState(object): """Gradient reduction related state of a Tensorflow graph.""" def __init__(self): - # For every type of optimizer, collect all gradients and variables. - self._optimizer_grads_and_vars = {} - # In what order were optimizers invoked within each tower? - self._ordered_optimizer_types = [] - self._number_of_towers = 0 - self._is_the_last_tower = False + self._collected_grads_and_vars = defaultdict(list) + self._current_tower_index = 0 + self._number_of_towers = 1 self._loss_reduction = None # Scopes of the first tower that don't have a prefix: self._variable_scope = None self._name_scope = None - # If needed, alert that GatheringOptimizer needs to be used with model_fn. - self._has_gathering_optimizer_been_used = False - - def collect_gradients(self, grads_and_vars, optimizer): - if optimizer.__class__ not in self._ordered_optimizer_types: - self._ordered_optimizer_types.append(optimizer.__class__) - - self._optimizer_grads_and_vars.setdefault(optimizer.__class__, - []).extend(grads_and_vars) - - def get_grad_and_vars_for_optimizer(self, optimizer): - return self._optimizer_grads_and_vars[optimizer.__class__] + # If needed, alert that TowerOptimizer needs to be used with model_fn. + self._has_tower_optimizer_been_used = False + + def collect_gradients(self, grads_and_vars): + self._collected_grads_and_vars[self._current_tower_index].append( + grads_and_vars) + + def get_latest_gradients_from_all_towers(self): + """Get gradients across towers for the last called optimizer.""" + grads_and_vars = [] + index_of_last_gradients = len( + self._collected_grads_and_vars[self._current_tower_index]) - 1 + for tower_id in range(self._current_tower_index + 1): + grads_and_vars.extend( + self._collected_grads_and_vars[tower_id][index_of_last_gradients]) + return grads_and_vars def set_reduction_across_towers(self, loss_reduction, number_of_towers): self._loss_reduction = loss_reduction @@ -370,10 +410,8 @@ class GatheringOptimizer(optimizer_lib.Optimizer): if tower_id == 0: self._variable_scope = var_scope self._name_scope = name_scope - if tower_id == (self._number_of_towers - 1): - self._is_the_last_tower = True + self._current_tower_index = tower_id yield - self._is_the_last_tower = False @property def scopes_of_the_first_tower(self): @@ -381,10 +419,7 @@ class GatheringOptimizer(optimizer_lib.Optimizer): @property def is_the_last_tower(self): - return self._is_the_last_tower - - def is_the_last_optimizer_within_a_tower(self, optimizer): - return optimizer.__class__ == self._ordered_optimizer_types[-1] + return self._current_tower_index == (self._number_of_towers - 1) @property def number_of_towers(self): @@ -395,12 +430,19 @@ class GatheringOptimizer(optimizer_lib.Optimizer): return self._loss_reduction @property - def has_gathering_optimizer_been_used(self): - return self._has_gathering_optimizer_been_used + def has_tower_optimizer_been_used(self): + return self._has_tower_optimizer_been_used - @has_gathering_optimizer_been_used.setter - def has_gathering_optimizer_been_used(self, value): - self._has_gathering_optimizer_been_used = value + @has_tower_optimizer_been_used.setter + def has_tower_optimizer_been_used(self, value): + self._has_tower_optimizer_been_used = value + + def did_towers_have_same_optimizer_calls(self): + total_number_of_grads = sum([ + len(grads) + for _, grads in six.iteritems(self._collected_grads_and_vars) + ]) + return total_number_of_grads % self._number_of_towers == 0 def _get_local_devices(device_type): @@ -415,6 +457,13 @@ def _get_local_devices(device_type): def _split_batch(features, labels, number_of_shards, device): """Split input features and labes into batches.""" + def ensure_divisible_by_shards(sequence): + batch_size = ops_lib.convert_to_tensor(sequence).get_shape()[0] + if batch_size % number_of_shards != 0: + raise ValueError( + 'Batch size {} needs to be divisible by the number of GPUs, which ' + 'is {}.'.format(batch_size, number_of_shards)) + def split_dictionary(dictionary): """Split a dictionary into shards.""" shards = [{} for _ in range(number_of_shards)] @@ -425,6 +474,7 @@ def _split_batch(features, labels, number_of_shards, device): sp_input=tensor, num_split=number_of_shards, axis=0)): shards[i][name] = shard else: + ensure_divisible_by_shards(tensor) for i, shard in enumerate(array_ops.split(tensor, number_of_shards)): shards[i][name] = shard return shards @@ -434,6 +484,7 @@ def _split_batch(features, labels, number_of_shards, device): if isinstance(features, dict): feature_shards = split_dictionary(features) else: + ensure_divisible_by_shards(features) feature_shards = array_ops.split(features, number_of_shards) if labels is None: @@ -441,6 +492,7 @@ def _split_batch(features, labels, number_of_shards, device): elif isinstance(labels, dict): label_shards = split_dictionary(labels) else: + ensure_divisible_by_shards(labels) label_shards = array_ops.split(labels, number_of_shards) return feature_shards, label_shards @@ -456,7 +508,7 @@ def _get_loss_towers(model_fn, config, devices, local_ps_devices, - loss_reduction=losses.Reduction.SUM, + loss_reduction, name_scope_pattern=_DEFAULT_NAME_SCOPE_PATTERN): """Replicate the loss computation across devices.""" tower_specs = [] @@ -471,7 +523,7 @@ def _get_loss_towers(model_fn, # pylint: disable=protected-access round_robin_strategy = device_setter_lib._RoundRobinStrategy( num_tasks=len(local_ps_devices)) - GatheringOptimizer._graph_state().set_reduction_across_towers( + TowerOptimizer._graph_state().set_reduction_across_towers( loss_reduction, len(devices)) for i, device in enumerate(devices): @@ -492,7 +544,7 @@ def _get_loss_towers(model_fn, with variable_scope.variable_scope( '', reuse=not is_the_first_tower) as var_scope: with ops_lib.name_scope(name_scope.format(i)) as name_scope: - with GatheringOptimizer._graph_state().tower( + with TowerOptimizer._graph_state().tower( tower_id=i, var_scope=var_scope, name_scope=name_scope): with ops_lib.device(device_setter): labels_shard = None @@ -505,18 +557,22 @@ def _get_loss_towers(model_fn, labels=labels_shard, **optional_params) - if (tower_spec.train_op is not None and - not GatheringOptimizer.has_been_used()): - raise ValueError('Please wrap optimizers with GatheringOptimizer' - ' in order to use replicate_model_fn.') + if (tower_spec.train_op is not None and len(devices) > 1 and + not TowerOptimizer.has_been_used()): + raise ValueError('Please wrap optimizers with TowerOptimizer' + ' in order to use replicate_model_fn with' + ' multiple `devices`.') # Scaling the loss here doesn't actually affect gradients. Another - # instance of scaling happens inside the GatheringOptimizer. + # instance of scaling happens inside the TowerOptimizer. tower_spec = _scale_tower_loss( tower_spec, loss_reduction, number_of_towers=len(devices)) tower_specs.append(tower_spec) - GatheringOptimizer._clear_graph_state() + if not TowerOptimizer._did_towers_have_same_optimizer_calls(): + raise ValueError('Each invocation of model_fn was supposed to make the same' + ' optimizer calls.') + TowerOptimizer._clear_graph_state() # pylint: enable=protected-access return tower_specs @@ -559,6 +615,8 @@ def _scale_loss(loss, loss_reduction, number_of_towers): """If needed, scale down the loss for averaging loss by summing.""" if loss is None: return None + if number_of_towers == 1: + return loss if loss_reduction != losses.Reduction.SUM: return math_ops.div(loss, 1.0 * number_of_towers, name='averaged_loss') @@ -591,7 +649,12 @@ def _train_spec(tower_specs, aggregation_device, aggregated_loss_name='loss'): """Populate replicated EstimatorSpec for `GraphKeys.TRAIN`.""" - estimator_spec = _asdict(tower_specs[0]) + # Spec of the last tower is used as the template for the final spec, because + # some `EstimatorSpec.training_hooks` rely on calls made in model_fn. For + # example, `SyncReplicasOptimizerHook` validates the + # `SyncReplicasOptimizer.apply_gradients` call. `TowerEstimator` makes that + # call only in the last tower. + estimator_spec = _asdict(tower_specs[-1]) estimator_spec['mode'] = model_fn_lib.ModeKeys.TRAIN estimator_spec['train_op'] = train_op estimator_spec['loss'] = _compute_sum_on_device( @@ -721,6 +784,17 @@ def _concat_tensor_dicts(*tensor_dicts): } +def _extract_tensors(tensors_and_vars): + tensors = [] + for tensor_and_var in tensors_and_vars: + tensor, _ = tensor_and_var + if isinstance(tensor, ops_lib.IndexedSlices): + tensors.append(tensor.values) + elif tensor is not None: + tensors.append(tensor) + return tensors + + def _dict_concat(*dicts): list_dict = {} for d in dicts: diff --git a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py index 0b6c0e957b75d641ca93b946267a5797083d16ab..d46a18aacfcd911c56a9f22dc9581060c7b458a6 100644 --- a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py +++ b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn_test.py @@ -37,6 +37,7 @@ from tensorflow.python.feature_column import feature_column from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops as ops_lib +from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -53,6 +54,7 @@ from tensorflow.python.summary.writer import writer_cache from tensorflow.python.training import adam from tensorflow.python.training import device_setter from tensorflow.python.training import gradient_descent +from tensorflow.python.training import training # TODO(isaprykin): Parametrize all the tests on @@ -119,8 +121,8 @@ class DNNClassifierIntegrationTest(test_util.TensorFlowTestCase): estimator = dnn.DNNClassifier( hidden_units=(2, 2), # Adagrad is configured with `get_optimizer_instance`, so the function - # form of `GatheringOptimizer.__init__` is used. - optimizer=replicate_model_fn.GatheringOptimizer(optimizer_fn), + # form of `TowerOptimizer.__init__` is used. + optimizer=replicate_model_fn.TowerOptimizer(optimizer_fn), feature_columns=feature_columns, n_classes=n_classes, model_dir=self._model_dir) @@ -130,7 +132,10 @@ class DNNClassifierIntegrationTest(test_util.TensorFlowTestCase): estimator.model_fn, devices=['/gpu:0', '/gpu:1', '/gpu:2']) else: model_fn = replicate_model_fn._replicate_model_fn_with_mode( - estimator.model_fn, devices=['/gpu:0', '/gpu:1', '/gpu:2'], mode=mode) + estimator.model_fn, + devices=['/gpu:0', '/gpu:1', '/gpu:2'], + loss_reduction=losses.Reduction.SUM, + mode=mode) estimator = estimator_lib.Estimator( model_fn=model_fn, @@ -160,7 +165,7 @@ class DNNClassifierIntegrationTest(test_util.TensorFlowTestCase): # Nothing should be left in the graph so that it doesn't get serialized. self.assertFalse(ops_lib.get_default_graph().get_collection_ref( - replicate_model_fn.GatheringOptimizer.COLLECTION_FOR_GRAPH_STATES)) + replicate_model_fn.TowerOptimizer.COLLECTION_FOR_GRAPH_STATES)) def _as_label(self, data_in_float): return np.rint(data_in_float).astype(np.int64) @@ -190,7 +195,7 @@ class ReplicateModelTest(test_util.TensorFlowTestCase): 'auc': metrics_lib.auc(labels, predictions) } - optimizer = replicate_model_fn.GatheringOptimizer( + optimizer = replicate_model_fn.TowerOptimizer( gradient_descent.GradientDescentOptimizer(params['learning_rate'])) return model_fn_lib.EstimatorSpec( @@ -212,7 +217,9 @@ class ReplicateModelTest(test_util.TensorFlowTestCase): with self.test_session() as session: replicated_model_fn = replicate_model_fn.replicate_model_fn( - self.model_fn, devices=['/gpu:0', '/gpu:1']) + self.model_fn, + loss_reduction=losses.Reduction.SUM, + devices=['/gpu:0', '/gpu:1']) estimator_spec = replicated_model_fn( features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) session.run(variables.global_variables_initializer()) @@ -233,6 +240,13 @@ class ReplicateModelTest(test_util.TensorFlowTestCase): labels = np.array([[1.0], [2.0]]) with self.test_session() as session: + # Add another trainable variable that doesn't produce a gradient to + # verify that None gradients are supported. + _ = variable_scope.get_variable( + 'another_variable', + initializer=constant_op.constant(1, dtype=dtypes.float64), + dtype=dtypes.float64) + replicated_model_fn = replicate_model_fn.replicate_model_fn( self.model_fn, losses.Reduction.MEAN, devices=['/gpu:0', '/gpu:1']) estimator_spec = replicated_model_fn( @@ -270,7 +284,9 @@ class ReplicateModelTest(test_util.TensorFlowTestCase): with self.test_session() as session, variable_scope.variable_scope( '', reuse=variable_scope.AUTO_REUSE): replicated_model_fn = replicate_model_fn.replicate_model_fn( - self.model_fn, devices=['/gpu:0', '/gpu:1']) + self.model_fn, + loss_reduction=losses.Reduction.SUM, + devices=['/gpu:0', '/gpu:1']) estimator_spec = replicated_model_fn( features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) session.run(variables.global_variables_initializer()) @@ -291,7 +307,9 @@ class ReplicateModelTest(test_util.TensorFlowTestCase): with self.test_session() as session: replicated_model_fn = replicate_model_fn.replicate_model_fn( - self.model_fn, devices=['/gpu:0', '/gpu:1']) + self.model_fn, + loss_reduction=losses.Reduction.SUM, + devices=['/gpu:0', '/gpu:1']) estimator_spec = replicated_model_fn( features, labels, model_fn_lib.ModeKeys.EVAL, self.params) session.run(variables.local_variables_initializer()) @@ -423,12 +441,224 @@ class ReplicateModelTest(test_util.TensorFlowTestCase): 'probabilities': np.array([[0.1], [0.02]]) }, session.run(estimator_spec.predictions)) + def test_batch_size_that_is_not_divisible_by_the_number_of_gpus(self): + features = np.array([[1.0], [2.0], [3.0]]) + labels = np.array([[1.0], [2.0], [3.0]]) + + with self.assertRaisesRegexp( + ValueError, '.*Batch.+size.+needs.+to.+be.+divisible.+by.+GPUs.+'): + replicated_model_fn = replicate_model_fn.replicate_model_fn( + self.model_fn, devices=['/gpu:0', '/gpu:1']) + _ = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) + def test_unsupported_loss_reduction(self): with self.assertRaisesRegexp(ValueError, '.+none.+reduction.+is.+specified.+'): _ = replicate_model_fn.replicate_model_fn(self.model_fn, losses.Reduction.NONE) + def test_places_on_gpu_with_upper_case_spelling(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session(): + replicated_model_fn = replicate_model_fn.replicate_model_fn( + self.model_fn, devices=['/GPU:0']) + _ = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) + + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual('/device:GPU:0', c.device) + + def test_places_on_gpu_with_lower_case_spelling(self): + features = np.array([[0.01], [0.002]]) + labels = np.array([[0.01], [0.02]]) + + with self.test_session(): + replicated_model_fn = replicate_model_fn.replicate_model_fn( + self.model_fn, devices=['/gpu:0']) + _ = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) + + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual('/device:GPU:0', c.device) + + +class ReplicateAcrossASingleDeviceWithoutTowerOptimizer( + test_util.TensorFlowTestCase): + + def model_fn(self, mode, features, labels, params): + c = variable_scope.get_variable( + 'c', + initializer=constant_op.constant(10, dtype=dtypes.float64), + dtype=dtypes.float64) + + predictions = math_ops.multiply(features, c) + + loss = losses.absolute_difference( + labels=labels, predictions=predictions, reduction=losses.Reduction.SUM) + loss = math_ops.reduce_sum(loss) + + metrics = { + 'accuracy': metrics_lib.accuracy(labels, predictions), + 'auc': metrics_lib.auc(labels, predictions) + } + + optimizer = gradient_descent.GradientDescentOptimizer( + params['learning_rate']) + + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=loss, + eval_metric_ops=metrics, + predictions={'probabilities': predictions}, + train_op=optimizer.minimize(loss)) + + @property + def params(self): + params = {} + params['learning_rate'] = 1.0 + return params + + def test_train_single_tower(self): + features = np.array([[1.0], [2.0]]) + labels = np.array([[1.0], [2.0]]) + + with self.test_session() as session: + replicated_model_fn = replicate_model_fn.replicate_model_fn( + self.model_fn, devices=['/gpu:0']) + estimator_spec = replicated_model_fn( + features, labels, model_fn_lib.ModeKeys.TRAIN, self.params) + session.run(variables.global_variables_initializer()) + + # loss = feature * c - label + total_loss = (1.0 * 10 - 1.0) + (2.0 * 10 - 2.0) + self.assertEqual(total_loss, session.run(estimator_spec.loss)) + + # loss' of c is 3. + # new value of c = 10 - learning rate * 3 = 7.0. + session.run(estimator_spec.train_op) + with variable_scope.variable_scope('', reuse=True): + c = variable_scope.get_variable('c', dtype=dtypes.float64) + self.assertEqual(7.0, session.run(c)) + + +class UseTowerEstimatorWithoutReplication(test_util.TensorFlowTestCase): + + def model_fn(self, mode, features, labels, params): + c = variable_scope.get_variable( + 'c', + initializer=constant_op.constant(10, dtype=dtypes.float64), + dtype=dtypes.float64) + + features = features['features'] + predictions = math_ops.multiply(features, c) + + loss = losses.absolute_difference( + labels=labels, predictions=predictions, reduction=losses.Reduction.SUM) + loss = math_ops.reduce_sum(loss) + + metrics = { + 'accuracy': metrics_lib.accuracy(labels, predictions), + 'auc': metrics_lib.auc(labels, predictions) + } + + optimizer = replicate_model_fn.TowerOptimizer( + gradient_descent.GradientDescentOptimizer(params['learning_rate'])) + + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=loss, + eval_metric_ops=metrics, + predictions={'probabilities': predictions}, + train_op=optimizer.minimize(loss)) + + @property + def params(self): + params = {} + params['learning_rate'] = 1.0 + return params + + def test_train_single_tower(self): + features = np.array([[1.0], [2.0]]) + labels = np.array([[1.0], [2.0]]) + + train_input_fn = numpy_io.numpy_input_fn( + x={'features': features}, y=labels, batch_size=2, shuffle=False) + + with self.test_session(): + estimator = estimator_lib.Estimator( + model_fn=self.model_fn, + model_dir=tempfile.mkdtemp(), + params=self.params) + estimator.train(train_input_fn, steps=1) + + self.assertEqual(7.0, estimator.get_variable_value('c')) + + +class MakeSureSyncReplicasOptimizerWorks(test_util.TensorFlowTestCase): + + def model_fn(self, mode, features, labels, params): + c = variable_scope.get_variable( + 'c', + initializer=constant_op.constant(10, dtype=dtypes.float64), + dtype=dtypes.float64) + + features = features['features'] + predictions = math_ops.multiply(features, c) + + loss = losses.absolute_difference( + labels=labels, predictions=predictions, reduction=losses.Reduction.SUM) + loss = math_ops.reduce_sum(loss) + + metrics = { + 'accuracy': metrics_lib.accuracy(labels, predictions), + 'auc': metrics_lib.auc(labels, predictions) + } + + optimizer = gradient_descent.GradientDescentOptimizer( + params['learning_rate']) + optimizer = training.SyncReplicasOptimizer( + optimizer, replicas_to_aggregate=1) + sync_hook = optimizer.make_session_run_hook(True) + optimizer = replicate_model_fn.TowerOptimizer(optimizer) + + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=loss, + eval_metric_ops=metrics, + training_hooks=[sync_hook], + predictions={'probabilities': predictions}, + train_op=optimizer.minimize( + loss, global_step=training.get_global_step())) + + @property + def params(self): + params = {} + params['learning_rate'] = 1.0 + return params + + def test_train_multiple_towers(self): + features = np.array([[1.0], [2.0]]) + labels = np.array([[1.0], [2.0]]) + + train_input_fn = numpy_io.numpy_input_fn( + x={'features': features}, y=labels, batch_size=2, shuffle=False) + + model_fn = replicate_model_fn.replicate_model_fn( + self.model_fn, + loss_reduction=losses.Reduction.SUM, + devices=['/gpu:0', '/gpu:1']) + + estimator = estimator_lib.Estimator( + model_fn=model_fn, model_dir=tempfile.mkdtemp(), params=self.params) + estimator.train(train_input_fn, steps=1) + + self.assertEqual(7.0, estimator.get_variable_value('c')) + class ReplicateWithTwoOptimizersTest(test_util.TensorFlowTestCase): @@ -438,6 +668,12 @@ class ReplicateWithTwoOptimizersTest(test_util.TensorFlowTestCase): initializer=constant_op.constant(10, dtype=dtypes.float64), dtype=dtypes.float64) + side_effects = variable_scope.get_variable( + 'side_effects', + initializer=constant_op.constant(0, dtype=dtypes.float64), + dtype=dtypes.float64, + trainable=False) + predictions = math_ops.multiply(features, c) loss = losses.absolute_difference( @@ -449,13 +685,16 @@ class ReplicateWithTwoOptimizersTest(test_util.TensorFlowTestCase): 'auc': metrics_lib.auc(labels, predictions) } - first_optimizer = replicate_model_fn.GatheringOptimizer( + first_optimizer = replicate_model_fn.TowerOptimizer( gradient_descent.GradientDescentOptimizer(1.0)) - second_optimizer = replicate_model_fn.GatheringOptimizer( + second_optimizer = replicate_model_fn.TowerOptimizer( adam.AdamOptimizer(1.0)) + with ops_lib.control_dependencies([side_effects.assign_add(1.0)]): + first_grads_and_vars = first_optimizer.compute_gradients(loss) + train_op = control_flow_ops.group( - [first_optimizer.minimize(loss), + [first_optimizer.apply_gradients(first_grads_and_vars), second_optimizer.minimize(loss)]) return model_fn_lib.EstimatorSpec( @@ -471,7 +710,9 @@ class ReplicateWithTwoOptimizersTest(test_util.TensorFlowTestCase): with self.test_session() as session: replicated_model_fn = replicate_model_fn.replicate_model_fn( - self.model_fn, devices=['/gpu:0', '/gpu:1']) + self.model_fn, + loss_reduction=losses.Reduction.SUM, + devices=['/gpu:0', '/gpu:1']) estimator_spec = replicated_model_fn(features, labels, model_fn_lib.ModeKeys.TRAIN, {}) session.run(variables.global_variables_initializer()) @@ -488,11 +729,29 @@ class ReplicateWithTwoOptimizersTest(test_util.TensorFlowTestCase): c = variable_scope.get_variable('c', dtype=dtypes.float64) self.assertNear(6.0, session.run(c), 0.000001) + side_effects = variable_scope.get_variable( + 'side_effects', dtype=dtypes.float64) + self.assertNear(2.0, session.run(side_effects), 0.000001) + class ReplicateWithTwoLossesAndOneOptimizer(test_util.TensorFlowTestCase): - class AnotherOptimizer(gradient_descent.GradientDescentOptimizer): - pass + def setUp(self): + self._should_skip_optimizer = False + self._towers_left_before_skipping_optimizer = -1 + + def incorrectly_skip_optimizer_for_tower(self, tower_number): + self._should_skip_optimizer = True + self._towers_left_before_skipping_optimizer = tower_number + + def should_skip_optimizer(self): + if not self._should_skip_optimizer: + return False + if self._towers_left_before_skipping_optimizer == 0: + return True + else: + self._towers_left_before_skipping_optimizer -= 1 + return False def model_fn(self, mode, features, labels, params): c = variable_scope.get_variable( @@ -524,16 +783,17 @@ class ReplicateWithTwoLossesAndOneOptimizer(test_util.TensorFlowTestCase): 'auc': metrics_lib.auc(labels, predictions) } - optimizer = replicate_model_fn.GatheringOptimizer( - gradient_descent.GradientDescentOptimizer(1.0)) - another_optimizer = replicate_model_fn.GatheringOptimizer( - self.AnotherOptimizer(1.0)) + train_ops = [] - train_op = control_flow_ops.group([ - optimizer.minimize(loss, var_list=[c]), - another_optimizer.minimize(another_loss, var_list=[d]) - ]) + optimizer = replicate_model_fn.TowerOptimizer( + gradient_descent.GradientDescentOptimizer(1.0)) + train_ops.append(optimizer.minimize(loss, var_list=[c])) + if not self.should_skip_optimizer(): + another_optimizer = replicate_model_fn.TowerOptimizer( + gradient_descent.GradientDescentOptimizer(1.0)) + train_ops.append(another_optimizer.minimize(another_loss, var_list=[d])) + train_op = control_flow_ops.group(train_ops) return model_fn_lib.EstimatorSpec( mode=mode, loss=total_loss, @@ -547,7 +807,9 @@ class ReplicateWithTwoLossesAndOneOptimizer(test_util.TensorFlowTestCase): with self.test_session() as session: replicated_model_fn = replicate_model_fn.replicate_model_fn( - self.model_fn, devices=['/gpu:0', '/gpu:1']) + self.model_fn, + loss_reduction=losses.Reduction.SUM, + devices=['/gpu:0', '/gpu:1']) estimator_spec = replicated_model_fn(features, labels, model_fn_lib.ModeKeys.TRAIN, {}) session.run(variables.global_variables_initializer()) @@ -568,6 +830,20 @@ class ReplicateWithTwoLossesAndOneOptimizer(test_util.TensorFlowTestCase): d = variable_scope.get_variable('d', dtype=dtypes.float64) self.assertNear(-1.0, session.run(d), 0.000001) + def test_different_optimizer_calls_within_towers(self): + self.incorrectly_skip_optimizer_for_tower(1) + + features = np.array([[1.0], [2.0]]) + labels = np.array([[1.0], [2.0]]) + + with self.test_session(), ops_lib.Graph().as_default(): + with self.assertRaisesRegexp( + ValueError, '.+was.+supposed.+to.+make.+same.+optimizer.+calls.+'): + replicated_model_fn = replicate_model_fn.replicate_model_fn( + self.model_fn, devices=['/gpu:0', '/gpu:1']) + _ = replicated_model_fn(features, labels, model_fn_lib.ModeKeys.TRAIN, + {}) + class FailToWrapOptimizerInTheModelFn(test_util.TensorFlowTestCase): @@ -604,7 +880,7 @@ class FailToWrapOptimizerInTheModelFn(test_util.TensorFlowTestCase): with self.test_session(): with self.assertRaisesRegexp(ValueError, - 'Please.+wrap.+with.+GatheringOptimizer'): + 'Please.+wrap.+with.+TowerOptimizer'): replicated_model_fn = replicate_model_fn.replicate_model_fn( self.model_fn, devices=['/gpu:0', '/gpu:1']) _ = replicated_model_fn(features, labels, model_fn_lib.ModeKeys.TRAIN, @@ -636,6 +912,7 @@ class GetLossTowersTest(test_util.TensorFlowTestCase): labels=[[0.6], [0.6]], params=None, config=None, + loss_reduction=losses.Reduction.SUM, devices=['/gpu:0', '/gpu:1'], local_ps_devices=['/gpu:0'], name_scope_pattern='test_tower_{}') @@ -721,6 +998,7 @@ class GetLossTowersTest(test_util.TensorFlowTestCase): features=[[0.6], [1.6], [2.6]], labels=[[0.6], [0.6], [2.6]], params=None, + loss_reduction=losses.Reduction.SUM, config=None, devices=['/gpu:0', '/gpu:1', '/gpu:3'], local_ps_devices=['/gpu:0', '/gpu:1', '/gpu:3'], @@ -750,8 +1028,13 @@ class SplitBatchTest(test_util.TensorFlowTestCase): return list(map(evaluate_items, first_list)), list( map(evaluate_items, second_list)) + def assertSparseValuesEqual(self, a, b): + self.assertAllEqual(a.indices, b.indices) + self.assertAllEqual(a.values, b.values) + self.assertAllEqual(a.dense_shape, b.dense_shape) + def test_simple_half_split(self): - with self.test_session() as session: # pylint: disable=unused-variable + with self.test_session(): features = [0.0, 1.0, 2.0, 3.0] labels = [10.0, 11.0, 12.0, 13.0] feature_shards, label_shards = replicate_model_fn._split_batch( @@ -764,7 +1047,7 @@ class SplitBatchTest(test_util.TensorFlowTestCase): self.assertAllEqual([[10.0, 11.0], [12.0, 13.0]], label_shards) def test_to_each_their_own(self): - with self.test_session() as session: # pylint: disable=unused-variable + with self.test_session(): features = [0.0, 1.0, 2.0, 3.0] labels = [10.0, 11.0, 12.0, 13.0] feature_shards, label_shards = replicate_model_fn._split_batch( @@ -777,7 +1060,7 @@ class SplitBatchTest(test_util.TensorFlowTestCase): self.assertAllEqual([[10.0], [11.0], [12.0], [13.0]], label_shards) def test_one_batch(self): - with self.test_session() as session: # pylint: disable=unused-variable + with self.test_session(): features = [0.0, 1.0, 2.0, 3.0] labels = [10.0, 11.0, 12.0, 13.0] feature_shards, label_shards = replicate_model_fn._split_batch( @@ -790,7 +1073,7 @@ class SplitBatchTest(test_util.TensorFlowTestCase): self.assertAllEqual([[10.0, 11.0, 12.0, 13.0]], label_shards) def test_half_split_in_dictionary(self): - with self.test_session() as session: # pylint: disable=unused-variable + with self.test_session(): features = {'first': [0.0, 1.0, 2.0, 3.0], 'second': [4.0, 5.0, 6.0, 7.0]} labels = [10.0, 11.0, 12.0, 13.0] @@ -804,6 +1087,58 @@ class SplitBatchTest(test_util.TensorFlowTestCase): self.assertAllEqual([10.0, 11.0], label_shards[0].eval()) self.assertAllEqual([12.0, 13.0], label_shards[1].eval()) + def test_sparse_tensor_can_be_split_unevenly(self): + with self.test_session(): + features = { + 'x': + sparse_tensor.SparseTensor( + indices=[[0, 0], [1, 2], [2, 2]], + values=[1.0, 2.0, 3.0], + dense_shape=[3, 4]) + } + labels = np.array([[1.0], [2.0]]) + + feature_shards, label_shards = replicate_model_fn._split_batch( + features, labels, 2, device='/gpu:0') + + self.assertSparseValuesEqual( + sparse_tensor.SparseTensorValue( + indices=[[0, 0], [1, 2]], values=[1., 2.], dense_shape=[2, 4]), + feature_shards[0]['x'].eval()) + self.assertSparseValuesEqual( + sparse_tensor.SparseTensorValue( + indices=[[0, 2]], values=[3.], dense_shape=[1, 4]), + feature_shards[1]['x'].eval()) + self.assertAllEqual([[1.0]], label_shards[0].eval()) + self.assertAllEqual([[2.0]], label_shards[1].eval()) + + def test_sparse_tensor_can_be_split_unevenly_repeated_row(self): + with self.test_session(): + features = { + 'x': + sparse_tensor.SparseTensor( + indices=[[0, 0], [1, 0], [1, 1]], + values=[1.0, 2.0, 3.0], + dense_shape=[3, 4]) + } + labels = np.array([[1.0], [2.0]]) + + feature_shards, label_shards = replicate_model_fn._split_batch( + features, labels, 2, device='/gpu:0') + + self.assertSparseValuesEqual( + sparse_tensor.SparseTensorValue( + indices=[[0, 0], [1, 0], [1, 1]], + values=[1., 2., 3.], + dense_shape=[2, 4]), feature_shards[0]['x'].eval()) + + second_batch = feature_shards[1]['x'].eval() + self.assertFalse(len(second_batch.indices)) + self.assertFalse(len(second_batch.values)) + self.assertAllEqual([1, 4], second_batch.dense_shape) + self.assertAllEqual([[1.0]], label_shards[0].eval()) + self.assertAllEqual([[2.0]], label_shards[1].eval()) + def test_one_batch_in_dictionary(self): with self.test_session() as session: # pylint: disable=unused-variable features = {'first': [0.0, 1.0, 2.0, 3.0], 'second': [4.0, 5.0, 6.0, 7.0]} @@ -960,6 +1295,7 @@ class PredictSpecTest(test_util.TensorFlowTestCase): self.model_fn, mode=None, features=[[0.1], [0.2]], + loss_reduction=losses.Reduction.SUM, labels=[[], []], params=None, config=None, diff --git a/tensorflow/contrib/factorization/BUILD b/tensorflow/contrib/factorization/BUILD index fe86a20ab1f69a0eaf9d7486142451dac6337274..180f1b68f3b56113dfbbfc100bd04efc3bb8b31f 100644 --- a/tensorflow/contrib/factorization/BUILD +++ b/tensorflow/contrib/factorization/BUILD @@ -221,6 +221,7 @@ py_test( name = "kmeans_test", size = "medium", srcs = ["python/ops/kmeans_test.py"], + shard_count = 4, srcs_version = "PY2AND3", tags = ["notsan"], # b/67512932 deps = [ diff --git a/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc b/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc index 31d08bfb65ea49e1378ffba480771d38ce16abec..a8c5d0763c28ba2b54f217405f0da65533f26b91 100644 --- a/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc +++ b/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc @@ -57,11 +57,11 @@ typedef Eigen::Map< class MaskedMatmulOp : public OpKernel { public: - explicit MaskedMatmulOp(OpKernelConstruction* context) - : OpKernel(context) { - OP_REQUIRES_OK(context, context->MatchSignature( - {DT_FLOAT, DT_FLOAT, DT_INT64, DT_BOOL, DT_BOOL}, - {DT_FLOAT})); + explicit MaskedMatmulOp(OpKernelConstruction* context) : OpKernel(context) { + OP_REQUIRES_OK( + context, + context->MatchSignature( + {DT_FLOAT, DT_FLOAT, DT_INT64, DT_BOOL, DT_BOOL}, {DT_FLOAT})); } void Compute(OpKernelContext* context) override { @@ -110,12 +110,11 @@ class MaskedMatmulOp : public OpKernel { num_nonzero_elements, 2); Tensor* prod_values_tensor; - OP_REQUIRES_OK(context, - context->allocate_output( - 0, TensorShape({num_nonzero_elements}), - &prod_values_tensor)); - EigenMatFloatMap prod_values(prod_values_tensor->vec().data(), - 1, num_nonzero_elements); + OP_REQUIRES_OK(context, context->allocate_output( + 0, TensorShape({num_nonzero_elements}), + &prod_values_tensor)); + EigenMatFloatMap prod_values(prod_values_tensor->vec().data(), 1, + num_nonzero_elements); auto get_a_index = [&indices_mat, &a_dim_0](int64 i) { int64 a_index = internal::SubtleMustCopy(indices_mat(i, 0)); @@ -182,8 +181,8 @@ class MaskedMatmulOp : public OpKernel { } }; // Shard the work. - worker_threads.workers->ParallelFor( - num_nonzero_elements, cost_per_unit, work); + worker_threads.workers->ParallelFor(num_nonzero_elements, cost_per_unit, + work); } }; REGISTER_KERNEL_BUILDER(Name("MaskedMatmul").Device(DEVICE_CPU), diff --git a/tensorflow/contrib/factorization/python/ops/clustering_ops.py b/tensorflow/contrib/factorization/python/ops/clustering_ops.py index 6d3acb2750743318aad83991bc1e89d64c329423..23137e0a973c0bdd2cdbd97159f7fd310178bf54 100644 --- a/tensorflow/contrib/factorization/python/ops/clustering_ops.py +++ b/tensorflow/contrib/factorization/python/ops/clustering_ops.py @@ -192,11 +192,11 @@ class KMeans(object): # Computes Euclidean distance. Note the first and third terms are # broadcast additions. squared_distance = ( - math_ops.reduce_sum(math_ops.square(inp), 1, keep_dims=True) - + math_ops.reduce_sum(math_ops.square(inp), 1, keepdims=True) - 2 * math_ops.matmul(inp, clusters, transpose_b=True) + array_ops.transpose( math_ops.reduce_sum( - math_ops.square(clusters), 1, keep_dims=True))) + math_ops.square(clusters), 1, keepdims=True))) output.append(squared_distance) return output diff --git a/tensorflow/contrib/factorization/python/ops/gmm.py b/tensorflow/contrib/factorization/python/ops/gmm.py index f72280c4ecf19e33278ffe74061f44bbb7b21709..b2dfe48b2dbe0ec0975f865bba95a7ceba0f590c 100644 --- a/tensorflow/contrib/factorization/python/ops/gmm.py +++ b/tensorflow/contrib/factorization/python/ops/gmm.py @@ -24,17 +24,16 @@ import numpy as np from tensorflow.contrib import framework from tensorflow.contrib.factorization.python.ops import gmm_ops from tensorflow.contrib.framework.python.framework import checkpoint_utils -from tensorflow.python.training import training_util from tensorflow.contrib.learn.python.learn.estimators import estimator from tensorflow.contrib.learn.python.learn.estimators import model_fn as model_fn_lib from tensorflow.python.framework import constant_op from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import logging_ops as logging -from tensorflow.python.ops import math_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops.control_flow_ops import with_dependencies from tensorflow.python.training import session_run_hook +from tensorflow.python.training import training_util def _streaming_sum(scalar_tensor): @@ -70,8 +69,8 @@ class _InitializeClustersHook(session_run_hook.SessionRunHook): class GMM(estimator.Estimator): """An estimator for GMM clustering.""" SCORES = 'scores' + LOG_LIKELIHOOD = 'loss' ASSIGNMENTS = 'assignments' - ALL_SCORES = 'all_scores' def __init__(self, num_clusters, @@ -113,10 +112,7 @@ class GMM(estimator.Estimator): yield result[GMM.ASSIGNMENTS] def score(self, input_fn=None, batch_size=None, steps=None): - """Predict total sum of distances to nearest clusters. - - Note that this function is different from the corresponding one in sklearn - which returns the negative of the sum of distances. + """Predict total log-likelihood. Args: input_fn: see predict. @@ -124,11 +120,11 @@ class GMM(estimator.Estimator): steps: see predict. Returns: - Total sum of distances to nearest clusters. + Total log-likelihood. """ results = self.evaluate(input_fn=input_fn, batch_size=batch_size, steps=steps) - return np.sum(results[GMM.SCORES]) + return np.log(np.sum(np.exp(results[GMM.SCORES]))) def weights(self): """Returns the cluster weights.""" @@ -158,9 +154,10 @@ class GMM(estimator.Estimator): def _model_fn(features, labels, mode, config): """Model function.""" assert labels is None, labels - (all_scores, + (loss, + scores, model_predictions, - losses, training_op, + training_op, init_op, is_initialized) = gmm_ops.gmm(self._parse_tensor_or_dict(features), self._training_initial_clusters, @@ -168,16 +165,15 @@ class GMM(estimator.Estimator): self._covariance_type, self._params) incr_step = state_ops.assign_add(training_util.get_global_step(), 1) - loss = math_ops.reduce_sum(losses) training_op = with_dependencies([training_op, incr_step], loss) training_hooks = [_InitializeClustersHook( init_op, is_initialized, config.is_chief)] predictions = { - GMM.ALL_SCORES: all_scores[0], GMM.ASSIGNMENTS: model_predictions[0][0], } eval_metric_ops = { - GMM.SCORES: _streaming_sum(loss), + GMM.SCORES: scores, + GMM.LOG_LIKELIHOOD: _streaming_sum(loss), } return model_fn_lib.ModelFnOps(mode=mode, predictions=predictions, eval_metric_ops=eval_metric_ops, diff --git a/tensorflow/contrib/factorization/python/ops/gmm_ops.py b/tensorflow/contrib/factorization/python/ops/gmm_ops.py index a61681c7f5a69a0fff1089404fc80b95c1c3106e..98d6434f4752b224201e38bed05ccd14428a758b 100644 --- a/tensorflow/contrib/factorization/python/ops/gmm_ops.py +++ b/tensorflow/contrib/factorization/python/ops/gmm_ops.py @@ -21,7 +21,6 @@ from __future__ import division from __future__ import print_function import numpy as np -from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -36,7 +35,6 @@ from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.ops.embedding_ops import embedding_lookup -from tensorflow.python.summary import summary # Machine epsilon. MEPS = np.finfo(float).eps @@ -253,14 +251,16 @@ class GmmAlgorithm(object): return ret def scores(self): - """Returns the distances to each class. + """Returns the per-sample likelihood fo the data. Returns: - A tuple with two Tensors. The first contains the distance to - each class. The second contains the distance to the assigned - class. + Log probabilities of each data point. """ - return (self._all_scores, self._scores) + return self._scores + + def log_likelihood_op(self): + """Returns the log-likelihood operation.""" + return self._log_likelihood_op def _define_graph(self, data): """Define graph for a single iteration. @@ -276,7 +276,8 @@ class GmmAlgorithm(object): self._define_expectation_operation(shard_id) self._define_partial_maximization_operation(shard_id, shard) self._define_maximization_operation(len(data)) - self._define_distance_to_clusters(data) + self._define_loglikelihood_operation() + self._define_score_samples() def _define_full_covariance_probs(self, shard_id, shard): """Defines the full covariance probabilties per example in a class. @@ -440,50 +441,20 @@ class GmmAlgorithm(object): state_ops.assign( self._covs, new_covs, validate_shape=False)) - def _define_distance_to_clusters(self, data): - """Defines the Mahalanobis distance to the assigned Gaussian.""" - # TODO(xavigonzalvo): reuse (input - mean) * cov^-1 * (input - - # mean) from log probability function. - self._all_scores = [] - for shard in data: - all_scores = [] - shard = array_ops.expand_dims(shard, 0) - for c in xrange(self._num_classes): - if self._covariance_type == FULL_COVARIANCE: - cov = self._covs[c, :, :] - elif self._covariance_type == DIAG_COVARIANCE: - cov = array_ops.diag(self._covs[c, :]) - inverse = linalg_ops.matrix_inverse(cov + self._min_var) - inv_cov = array_ops.tile( - array_ops.expand_dims(inverse, 0), - array_ops.stack([self._num_examples, 1, 1])) - diff = array_ops.transpose(shard - self._means[c, :, :], perm=[1, 0, 2]) - m_left = math_ops.matmul(diff, inv_cov) - all_scores.append( - math_ops.sqrt( - math_ops.matmul( - m_left, array_ops.transpose( - diff, perm=[0, 2, 1])))) - self._all_scores.append( - array_ops.reshape( - array_ops.concat(all_scores, 1), - array_ops.stack([self._num_examples, self._num_classes]))) - - # Distance to the associated class. - self._all_scores = array_ops.concat(self._all_scores, 0) - assignments = array_ops.concat(self.assignments(), 0) - rows = math_ops.to_int64(math_ops.range(0, self._num_examples)) - indices = array_ops.concat( - [array_ops.expand_dims(rows, 1), array_ops.expand_dims(assignments, 1)], - 1) - self._scores = array_ops.gather_nd(self._all_scores, indices) - def _define_loglikelihood_operation(self): """Defines the total log-likelihood of current iteration.""" - self._ll_op = [] + op = [] for prior_probs in self._prior_probs: - self._ll_op.append(math_ops.reduce_sum(math_ops.log(prior_probs))) - summary.scalar('ll', math_ops.reduce_sum(self._ll_op)) + op.append(math_ops.reduce_logsumexp(prior_probs)) + self._log_likelihood_op = math_ops.reduce_logsumexp(op) + + def _define_score_samples(self): + """Defines the likelihood of each data sample.""" + op = [] + for shard_id, prior_probs in enumerate(self._prior_probs): + op.append(prior_probs + math_ops.log(self._w[shard_id])) + self._scores = array_ops.squeeze( + math_ops.reduce_logsumexp(op, axis=2, keep_dims=True), axis=0) def gmm(inp, @@ -511,14 +482,9 @@ def gmm(inp, Returns: Note: tuple of lists returned to be consistent with skflow A tuple consisting of: - all_scores: A matrix (or list of matrices) of dimensions (num_input, - num_clusters) where the value is the distance of an input vector and a - cluster center. assignments: A vector (or list of vectors). Each element in the vector corresponds to an input row in 'inp' and specifies the cluster id corresponding to the input. - scores: Similar to assignments but specifies the distance to the - assigned cluster instead. training_op: an op that runs an iteration of training. init_op: an op that runs the initialization. """ @@ -532,6 +498,7 @@ def gmm(inp, gmm_tool = GmmAlgorithm(inp, num_clusters, initial_means, params, covariance_type, random_seed) assignments = gmm_tool.assignments() - all_scores, scores = gmm_tool.scores() - return ([all_scores], [assignments], [scores], gmm_tool.training_ops(), + scores = gmm_tool.scores() + loss = gmm_tool.log_likelihood_op() + return (loss, scores, [assignments], gmm_tool.training_ops(), gmm_tool.init_ops(), gmm_tool.is_initialized()) diff --git a/tensorflow/contrib/factorization/python/ops/gmm_ops_test.py b/tensorflow/contrib/factorization/python/ops/gmm_ops_test.py index c50e82db8a230012ba13c1d7ad7e28c23bd27355..888c3c238c2654ea11ea3bf8270d6c3fcd951a03 100644 --- a/tensorflow/contrib/factorization/python/ops/gmm_ops_test.py +++ b/tensorflow/contrib/factorization/python/ops/gmm_ops_test.py @@ -122,17 +122,23 @@ class GmmOpsTest(test.TestCase): g.seed = 5 with self.test_session() as sess: data = constant_op.constant(self.data, dtype=dtypes.float32) - _, assignments, _, training_op, init_op, _ = gmm_ops.gmm( + loss_op, scores, assignments, training_op, init_op, _ = gmm_ops.gmm( data, 'random', num_classes, random_seed=self.seed) variables.global_variables_initializer().run() sess.run(init_op) + first_loss = sess.run(loss_op) for _ in xrange(self.iterations): sess.run(training_op) assignments = sess.run(assignments) + end_loss = sess.run(loss_op) + scores = sess.run(scores) + self.assertEqual((self.num_examples, 1), scores.shape) accuracy = np.mean( np.asarray(self.true_assignments) == np.squeeze(assignments)) logging.info('Accuracy: %f', accuracy) + logging.info('First loss: %f, end loss: %f', first_loss, end_loss) + self.assertGreater(end_loss, first_loss) self.assertGreater(accuracy, 0.98) def testParams(self): diff --git a/tensorflow/contrib/factorization/python/ops/gmm_test.py b/tensorflow/contrib/factorization/python/ops/gmm_test.py index 7717b47daefce9ff65b1f1e84f671a463cf2e826..00a4734eb6d89cd02484f1c5161366377cc71208 100644 --- a/tensorflow/contrib/factorization/python/ops/gmm_test.py +++ b/tensorflow/contrib/factorization/python/ops/gmm_test.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function import numpy as np -from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.contrib.factorization.python.ops import gmm as gmm_lib from tensorflow.contrib.learn.python.learn.estimators import kmeans @@ -30,12 +29,9 @@ from tensorflow.python.framework import random_seed as random_seed_lib from tensorflow.python.ops import array_ops from tensorflow.python.ops import data_flow_ops from tensorflow.python.ops import random_ops -from tensorflow.python.platform import flags from tensorflow.python.platform import test from tensorflow.python.training import queue_runner -FLAGS = flags.FLAGS - class GMMTest(test.TestCase): @@ -64,9 +60,8 @@ class GMMTest(test.TestCase): self.batch_size = self.num_points self.true_centers = self.make_random_centers(self.num_centers, self.num_dims) - self.points, self.assignments, self.scores = self.make_random_points( + self.points, self.assignments = self.make_random_points( self.true_centers, self.num_points) - self.true_score = np.add.reduce(self.scores) # Use initial means from kmeans (just like scikit-learn does). clusterer = kmeans.KMeansClustering(num_clusters=self.num_centers) @@ -86,24 +81,7 @@ class GMMTest(test.TestCase): offsets = np.round( np.random.randn(num_points, num_dims).astype(np.float32) * 20) points = centers[assignments] + offsets - means = [ - np.mean( - points[assignments == center], axis=0) - for center in xrange(num_centers) - ] - covs = [ - np.cov(points[assignments == center].T) - for center in xrange(num_centers) - ] - scores = [] - for r in xrange(num_points): - scores.append( - np.sqrt( - np.dot( - np.dot(points[r, :] - means[assignments[r]], - np.linalg.inv(covs[assignments[r]])), points[r, :] - - means[assignments[r]]))) - return (points, assignments, scores) + return (points, assignments) def test_weights(self): """Tests the shape of the weights.""" @@ -136,8 +114,7 @@ class GMMTest(test.TestCase): gmm.fit(input_fn=self.input_fn(), steps=10) score2 = gmm.score(input_fn=self.input_fn(batch_size=self.num_points), steps=1) - self.assertGreater(score1, score2) - self.assertNear(self.true_score, score2, self.true_score * 0.15) + self.assertLess(score1, score2) def test_infer(self): gmm = gmm_lib.GMM(self.num_centers, @@ -149,8 +126,7 @@ class GMMTest(test.TestCase): # Make a small test set num_points = 40 - points, true_assignments, true_offsets = ( - self.make_random_points(clusters, num_points)) + points, true_assignments = self.make_random_points(clusters, num_points) assignments = [] for item in gmm.predict_assignments( @@ -159,11 +135,6 @@ class GMMTest(test.TestCase): assignments = np.ravel(assignments) self.assertAllEqual(true_assignments, assignments) - # Test score - score = gmm.score(input_fn=self.input_fn(points=points, - batch_size=num_points), steps=1) - self.assertNear(score, np.sum(true_offsets), 4.05) - def _compare_with_sklearn(self, cov_type): # sklearn version. iterations = 40 diff --git a/tensorflow/contrib/factorization/python/ops/kmeans.py b/tensorflow/contrib/factorization/python/ops/kmeans.py index 9a5413fc3f2642443621b33d325e3d8c893fd6ac..c861cfff544a78617aa1ace730b50c094cf16330 100644 --- a/tensorflow/contrib/factorization/python/ops/kmeans.py +++ b/tensorflow/contrib/factorization/python/ops/kmeans.py @@ -25,6 +25,7 @@ import time from tensorflow.contrib.factorization.python.ops import clustering_ops from tensorflow.python.estimator import estimator from tensorflow.python.estimator import model_fn as model_fn_lib +from tensorflow.python.estimator.export import export_output from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -32,6 +33,7 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics from tensorflow.python.ops import state_ops from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.saved_model import signature_constants from tensorflow.python.summary import summary from tensorflow.python.training import session_run_hook from tensorflow.python.training import training_util @@ -141,7 +143,7 @@ class _ModelFn(object): def model_fn(self, features, mode, config): """Model function for the estimator. - Note that this does not take a `1abels` arg. This works, but `input_fn` must + Note that this does not take a `labels` arg. This works, but `input_fn` must return either `features` or, equivalently, `(features, None)`. Args: @@ -207,6 +209,15 @@ class _ModelFn(object): training_hooks.append( _LossRelativeChangeHook(loss, self._relative_tolerance)) + export_outputs = { + KMeansClustering.ALL_DISTANCES: + export_output.PredictOutput(all_distances[0]), + KMeansClustering.CLUSTER_INDEX: + export_output.PredictOutput(model_predictions[0]), + signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: + export_output.PredictOutput(model_predictions[0]) + } + return model_fn_lib.EstimatorSpec( mode=mode, predictions={ @@ -216,7 +227,8 @@ class _ModelFn(object): loss=loss, train_op=training_op, eval_metric_ops={KMeansClustering.SCORE: metrics.mean(loss)}, - training_hooks=training_hooks) + training_hooks=training_hooks, + export_outputs=export_outputs) # TODO(agarwal,ands): support sharded input. diff --git a/tensorflow/contrib/feature_column/BUILD b/tensorflow/contrib/feature_column/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..6fc053759c58d30c24657dd22e7d12be46fc7a7e --- /dev/null +++ b/tensorflow/contrib/feature_column/BUILD @@ -0,0 +1,37 @@ +package( + default_visibility = [ + "//tensorflow:internal", + ], +) + +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "py_test") + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +py_library( + name = "feature_column_py", + srcs = ["__init__.py"], + srcs_version = "PY2AND3", + deps = [ + ":sequential_feature_column", + ], +) + +py_library( + name = "sequential_feature_column", + srcs = ["python/feature_column/sequential_feature_column.py"], + srcs_version = "PY2AND3", + deps = [], +) diff --git a/tensorflow/contrib/feature_column/__init__.py b/tensorflow/contrib/feature_column/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6da7b126931effae9cc97091a27070d7013450d4 --- /dev/null +++ b/tensorflow/contrib/feature_column/__init__.py @@ -0,0 +1,30 @@ +# 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 utilities for tf.feature_column.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# pylint: disable=unused-import,line-too-long,wildcard-import +from tensorflow.contrib.feature_column.python.feature_column.sequential_feature_column import * + +from tensorflow.python.util.all_util import remove_undocumented +# pylint: enable=unused-import,line-too-long,wildcard-import + +_allowed_symbols = [ +] + +remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequential_feature_column.py b/tensorflow/contrib/feature_column/python/feature_column/sequential_feature_column.py new file mode 100644 index 0000000000000000000000000000000000000000..690a44ff4368663306733300a1ea70397fb93e1e --- /dev/null +++ b/tensorflow/contrib/feature_column/python/feature_column/sequential_feature_column.py @@ -0,0 +1,19 @@ +# 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 sequential input.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function diff --git a/tensorflow/contrib/ffmpeg/decode_audio_op.cc b/tensorflow/contrib/ffmpeg/decode_audio_op.cc index 92fad70b1f9cc55e0690a3fbb35abcf56aa68f16..5ab57ca4cd413bd92f1576278b22d2602c905309 100644 --- a/tensorflow/contrib/ffmpeg/decode_audio_op.cc +++ b/tensorflow/contrib/ffmpeg/decode_audio_op.cc @@ -44,7 +44,7 @@ const char* kValidFileFormats[] = {"mp3", "mp4", "ogg", "wav"}; void Decode(OpKernelContext* context, const tensorflow::StringPiece& file_contents, const string& file_format, const int32 samples_per_second, - const int32 channel_count) { + const int32 channel_count, const string& stream) { // Write the input data to a temp file. const string temp_filename = io::GetTempFilename(file_format); OP_REQUIRES_OK(context, WriteFile(temp_filename, file_contents)); @@ -54,7 +54,7 @@ void Decode(OpKernelContext* context, std::vector output_samples; Status result = ffmpeg::ReadAudioFile(temp_filename, file_format, samples_per_second, - channel_count, &output_samples); + channel_count, stream, &output_samples); if (result.code() == error::Code::NOT_FOUND) { OP_REQUIRES( context, result.ok(), @@ -99,7 +99,12 @@ void Decode(OpKernelContext* context, */ class DecodeAudioOpV2 : public OpKernel { public: - explicit DecodeAudioOpV2(OpKernelConstruction* context) : OpKernel(context) {} + explicit DecodeAudioOpV2(OpKernelConstruction* context) : OpKernel(context) { + string stream; + if (context->GetAttr("stream", &stream).ok()) { + stream_ = stream; + } + } void Compute(OpKernelContext* context) override { OP_REQUIRES( @@ -153,8 +158,12 @@ class DecodeAudioOpV2 : public OpKernel { errors::InvalidArgument("channel_count must be positive, but got: ", channel_count)); - Decode(context, contents, file_format, samples_per_second, channel_count); + Decode(context, contents, file_format, samples_per_second, channel_count, + stream_); } + + private: + string stream_; }; REGISTER_KERNEL_BUILDER(Name("DecodeAudioV2").Device(DEVICE_CPU), @@ -166,6 +175,7 @@ REGISTER_OP("DecodeAudioV2") .Input("samples_per_second: int32") .Input("channel_count: int32") .Output("sampled_audio: float") + .Attr("stream: string = ''") .SetShapeFn([](shape_inference::InferenceContext* c) { const Tensor* channels_tensor = c->input_tensor(3); if (channels_tensor == nullptr) { @@ -237,7 +247,7 @@ class DecodeAudioOp : public OpKernel { const tensorflow::StringPiece file_contents = contents.scalar()(); Decode(context, file_contents, file_format_, samples_per_second_, - channel_count_); + channel_count_, ""); } private: diff --git a/tensorflow/contrib/ffmpeg/decode_audio_op_test.py b/tensorflow/contrib/ffmpeg/decode_audio_op_test.py index 0d7c9cb99e8a5fad4a7ccf86d7253170ace91fd7..3dc663bb6f589d09ed067eae09d7d7dd0c40ec95 100644 --- a/tensorflow/contrib/ffmpeg/decode_audio_op_test.py +++ b/tensorflow/contrib/ffmpeg/decode_audio_op_test.py @@ -33,7 +33,8 @@ class DecodeAudioOpTest(test.TestCase): def _loadFileAndTest(self, filename, file_format, duration_sec, samples_per_second, channel_count, - samples_per_second_tensor=None, feed_dict=None): + samples_per_second_tensor=None, feed_dict=None, + stream=None): """Loads an audio file and validates the output tensor. Args: @@ -49,6 +50,9 @@ class DecodeAudioOpTest(test.TestCase): feed_dict: Used when evaluating the `decode_audio` op. If not provided, will be empty. Useful when providing a placeholder for `samples_per_second_tensor`. + stream: A string specifying which stream from the content file + should be decoded. The default value is '' which leaves the + decision to ffmpeg. """ if samples_per_second_tensor is None: samples_per_second_tensor = samples_per_second @@ -62,7 +66,7 @@ class DecodeAudioOpTest(test.TestCase): contents, file_format=file_format, samples_per_second=samples_per_second_tensor, - channel_count=channel_count) + channel_count=channel_count, stream=stream) audio = audio_op.eval(feed_dict=feed_dict or {}) self.assertEqual(len(audio.shape), 2) self.assertNear( @@ -72,6 +76,17 @@ class DecodeAudioOpTest(test.TestCase): 0.1 * audio.shape[0]) self.assertEqual(audio.shape[1], channel_count) + def testStreamIdentifier(self): + # mono_16khz_mp3_32khz_aac.mp4 was generated from: + # ffmpeg -i tensorflow/contrib/ffmpeg/testdata/mono_16khz_mp3.mp4 \ + # -i tensorflow/contrib/ffmpeg/testdata/mono_32khz_aac.mp4 \ + # -strict -2 -map 0:a -map 1:a \ + # tensorflow/contrib/ffmpeg/testdata/mono_16khz_mp3_32khz_aac.mp4 + self._loadFileAndTest('mono_16khz_mp3_32khz_aac.mp4', 'mp4', 2.77, 20000, + 1, stream='0') + self._loadFileAndTest('mono_16khz_mp3_32khz_aac.mp4', 'mp4', 2.77, 20000, + 1, stream='1') + def testMonoMp3(self): self._loadFileAndTest('mono_16khz.mp3', 'mp3', 0.57, 20000, 1) self._loadFileAndTest('mono_16khz.mp3', 'mp3', 0.57, 20000, 2) diff --git a/tensorflow/contrib/ffmpeg/decode_video_op.cc b/tensorflow/contrib/ffmpeg/decode_video_op.cc index d44032968d559bec14722902a4d47d22c46ea4aa..6f8ad486d10a825a277749157d68fa671b9f8d3a 100644 --- a/tensorflow/contrib/ffmpeg/decode_video_op.cc +++ b/tensorflow/contrib/ffmpeg/decode_video_op.cc @@ -102,16 +102,12 @@ REGISTER_OP("DecodeVideo") return Status::OK(); }) .Doc(R"doc( -Processes the contents of an audio file into a tensor using FFmpeg to decode +Processes the contents of an video file into a tensor using FFmpeg to decode the file. -One row of the tensor is created for each channel in the audio file. Each -channel contains audio samples starting at the beginning of the audio and -having `1/samples_per_second` time between them. If the `channel_count` is -different from the contents of the file, channels will be merged or created. - -contents: The binary audio file contents, as a string or rank-0 string - tensor. +contents: The binary contents of the video file to decode. This is a + scalar. +output: A rank-4 `Tensor` that has `[frames, height, width, 3]` RGB as output. )doc"); } // namespace ffmpeg diff --git a/tensorflow/contrib/ffmpeg/default/BUILD b/tensorflow/contrib/ffmpeg/default/BUILD index 949ae9ad9e4b045ee1b5cc82d49c0e7468c2005d..6b455567d766dbe6d380a498bd7f521db27e077b 100644 --- a/tensorflow/contrib/ffmpeg/default/BUILD +++ b/tensorflow/contrib/ffmpeg/default/BUILD @@ -19,6 +19,7 @@ cc_library( ], deps = [ "//tensorflow/core:framework_headers_lib", + "//third_party/eigen3", "@protobuf_archive//:protobuf_headers", ], ) diff --git a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc index 1e8af1458cea13b2ddb89b7d93a4ffb8b974ecd2..e61221a6b0d34373279a379f356c99c379488182 100644 --- a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc +++ b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc @@ -44,41 +44,43 @@ std::vector FfmpegAudioCommandLine(const string& input_filename, const string& output_filename, const string& input_format_id, int32 samples_per_second, - int32 channel_count) { - return {"-nostats", // No additional progress display. - "-nostdin", // No interactive commands accepted. - "-f", input_format_id, // eg: "mp3" - "-probesize", StrCat(kDefaultProbeSize), "-i", input_filename, - "-loglevel", "error", // Print errors only. - "-hide_banner", // Skip printing build options, version, etc. - "-map_metadata", "-1", // Copy global metadata from input to output. - "-vn", // No video recording. - "-ac:a:0", StrCat(channel_count), "-ar:a:0", - StrCat(samples_per_second), - // Output set (in several ways) to signed 16-bit little-endian ints. - "-codec:a:0", "pcm_s16le", "-sample_fmt", "s16", "-f", "s16le", - "-sn", // No subtitle recording. - "-y", // Overwrite output file. - StrCat(output_filename)}; + int32 channel_count, + const string& stream) { + std::vector command({ + "-nostats", // No additional progress display. + "-nostdin", // No interactive commands accepted. + "-f", input_format_id, // eg: "mp3" + "-probesize", StrCat(kDefaultProbeSize), "-i", input_filename, + "-loglevel", "error", // Print errors only. + "-hide_banner", // Skip printing build options, version, etc. + "-map_metadata", "-1", // Copy global metadata from input to output. + "-vn", // No video recording. + "-ac:a:0", StrCat(channel_count), "-ar:a:0", StrCat(samples_per_second), + // Output set (in several ways) to signed 16-bit little-endian ints. + "-codec:a:0", "pcm_s16le", "-sample_fmt", "s16", "-f", "s16le", + "-sn", // No subtitle recording. + "-y" // Overwrite output file. + }); + if (!stream.empty()) { + command.emplace_back("-map"); + command.emplace_back(StrCat("0:", stream)); + } + command.emplace_back(StrCat(output_filename)); + + return command; } std::vector FfmpegVideoCommandLine(const string& input_filename, const string& output_filename) { return {"-nostats", // No additional progress display. "-nostdin", // No interactive commands accepted. - "-i", - input_filename, - "-f", - "image2pipe", - "-probesize", - StrCat(kDefaultProbeSize), - "-loglevel", - "error", // Print errors only. + "-i", input_filename, "-f", "image2pipe", "-probesize", + StrCat(kDefaultProbeSize), "-loglevel", + // Info is needed to get the information about stream, etc. + // It is generated to a separate file, not stdout/stderr. + "info", "-hide_banner", // Skip printing build options, version, etc. - "-vcodec", - "rawvideo", - "-pix_fmt", - "rgb24", + "-vcodec", "rawvideo", "-pix_fmt", "rgb24", "-y", // Overwrite output file. StrCat(output_filename)}; } @@ -123,7 +125,6 @@ bool IsBinaryInstalled(const string& binary_name) { std::transform(args.begin(), args.end(), std::back_inserter(args_chars), [](const string& s) { return const_cast(s.c_str()); }); args_chars.push_back(nullptr); - ::execvp(kFfmpegExecutable, args_chars.data()); // exec only returns on error. const int error = errno; @@ -308,13 +309,12 @@ Status WriteFile(const string& filename, StringPiece contents) { Status ReadAudioFile(const string& filename, const string& audio_format_id, int32 samples_per_second, int32 channel_count, - std::vector* output_samples) { + const string& stream, std::vector* output_samples) { // Create an argument list. string output_filename = io::GetTempFilename("raw"); const std::vector args = FfmpegAudioCommandLine(filename, output_filename, audio_format_id, - samples_per_second, channel_count); - + samples_per_second, channel_count, stream); // Unfortunately, it's impossible to differentiate an exec failure due to the // binary being missing and an error from the binary's execution. Therefore, // check to see if the binary *should* be available. If not, return an error @@ -368,7 +368,6 @@ Status ReadVideoFile(const string& filename, std::vector* output_data, // Create an argument list. const std::vector args = FfmpegVideoCommandLine(filename, output_filename); - // Execute ffmpeg and report errors. pid_t child_pid = ::fork(); if (child_pid < 0) { diff --git a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib_test.cc b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib_test.cc index 85b61b26163d87a10d4e316720b4f633e038bbec..05728b3d37570d06f2f8af67e3b0612d21d07601 100644 --- a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib_test.cc +++ b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib_test.cc @@ -32,10 +32,8 @@ namespace tensorflow { namespace ffmpeg { namespace { -const char kTestWavFilename[] = - "contrib/ffmpeg/testdata/mono_10khz.wav"; -const char kTestMp3Filename[] = - "contrib/ffmpeg/testdata/test_sound1.mp3"; +const char kTestWavFilename[] = "contrib/ffmpeg/testdata/mono_10khz.wav"; +const char kTestMp3Filename[] = "contrib/ffmpeg/testdata/test_sound1.mp3"; // Set to true via a command line flag iff the test is expected to have FFmpeg // installed. @@ -139,7 +137,7 @@ TEST(FfmpegLibTest, TestRoundTripWav) { } // namespace ffmpeg } // namespace tensorflow -int main(int argc, char **argv) { +int main(int argc, char** argv) { tensorflow::string usage = tensorflow::ffmpeg::ParseTestFlags(&argc, argv); testing::InitGoogleTest(&argc, argv); if (argc != 1) { diff --git a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib_utility_test.cc b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib_utility_test.cc index 36fc71794b06e0f3cb86c40b325ce50e8999c667..d6c885a32424334bfc28c830e3701f219aa244ee 100644 --- a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib_utility_test.cc +++ b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib_utility_test.cc @@ -20,8 +20,6 @@ #include #include - -#include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/lib/core/threadpool.h" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/platform/env.h" diff --git a/tensorflow/contrib/ffmpeg/ffmpeg_lib.h b/tensorflow/contrib/ffmpeg/ffmpeg_lib.h index c5ea1432bf8b61c87615074a93a45325371c4c87..a8d5a0dd83fb504b5e6671c3e82dc7d2dd3e6a9b 100644 --- a/tensorflow/contrib/ffmpeg/ffmpeg_lib.h +++ b/tensorflow/contrib/ffmpeg/ffmpeg_lib.h @@ -13,8 +13,8 @@ // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_FFMPEG_FFMPEG_LIB_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_FFMPEG_FFMPEG_LIB_H_ +#ifndef TENSORFLOW_CONTRIB_FFMPEG_FFMPEG_LIB_H_ +#define TENSORFLOW_CONTRIB_FFMPEG_FFMPEG_LIB_H_ #include #include @@ -42,7 +42,7 @@ Status WriteFile(const string& filename, tensorflow::StringPiece contents); // contain a separate sample for each channel. Frames are ordered by time. Status ReadAudioFile(const string& filename, const string& audio_format_id, int32 samples_per_second, int32 channel_count, - std::vector* output_samples); + const string& stream, std::vector* output_samples); // Creates an audio file using ffmpeg in a specific format. The samples are in // [-1.0, 1.0]. If there are multiple channels in the audio then each frame will @@ -61,4 +61,4 @@ Status ReadVideoFile(const string& filename, std::vector* output_data, } // namespace ffmpeg } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_FFMPEG_DEFAULT_FFMPEG_LIB_H_ +#endif // TENSORFLOW_CONTRIB_FFMPEG_DEFAULT_FFMPEG_LIB_H_ diff --git a/tensorflow/contrib/ffmpeg/ffmpeg_ops.py b/tensorflow/contrib/ffmpeg/ffmpeg_ops.py index 08b5a6ea48c2d4959af68a2ee9d27d21c6245457..020b5c99c61019254bef0b1dff6bc5901c92758a 100644 --- a/tensorflow/contrib/ffmpeg/ffmpeg_ops.py +++ b/tensorflow/contrib/ffmpeg/ffmpeg_ops.py @@ -31,7 +31,7 @@ _ffmpeg_so = loader.load_op_library( def decode_audio(contents, file_format=None, samples_per_second=None, - channel_count=None): + channel_count=None, stream=None): """Create an op that decodes the contents of an audio file. Note that ffmpeg is free to select the "best" audio track from an mp4. @@ -51,6 +51,9 @@ def decode_audio(contents, file_format=None, samples_per_second=None, `contents` have more than this number, then some channels will be merged or dropped. If `contents` has fewer than this, then additional channels will be created from the existing ones. + stream: A string specifying which stream from the content file + should be decoded, e.g., '0' means the 0-th stream. + The default value is '' which leaves the decision to ffmpeg. Returns: A rank-2 tensor that has time along dimension 0 and channels along @@ -61,7 +64,7 @@ def decode_audio(contents, file_format=None, samples_per_second=None, """ return gen_decode_audio_op_py.decode_audio_v2( contents, file_format=file_format, samples_per_second=samples_per_second, - channel_count=channel_count) + channel_count=channel_count, stream=stream) ops.NotDifferentiable('DecodeAudio') diff --git a/tensorflow/contrib/ffmpeg/testdata/mono_16khz_mp3_32khz_aac.mp4 b/tensorflow/contrib/ffmpeg/testdata/mono_16khz_mp3_32khz_aac.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..2485da86d60837800fbb0b390c440e674de25993 Binary files /dev/null and b/tensorflow/contrib/ffmpeg/testdata/mono_16khz_mp3_32khz_aac.mp4 differ diff --git a/tensorflow/contrib/framework/BUILD b/tensorflow/contrib/framework/BUILD index 5b659ddaa1386736eb8cc05a203ed1827ccd160e..9e5f54f0973eae899ca65e4098358107053cb7d4 100644 --- a/tensorflow/contrib/framework/BUILD +++ b/tensorflow/contrib/framework/BUILD @@ -11,11 +11,12 @@ package(default_visibility = [ ]) load("//tensorflow:tensorflow.bzl", "py_test") -load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") load("//tensorflow:tensorflow.bzl", "tf_custom_op_library") load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py") load("//tensorflow:tensorflow.bzl", "tf_gen_op_libs") load("//tensorflow:tensorflow.bzl", "tf_kernel_library") +load("//tensorflow:tensorflow.bzl", "cuda_py_test") +load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") tf_custom_op_py_library( name = "framework_py", @@ -31,8 +32,10 @@ tf_custom_op_py_library( "python/ops/arg_scope.py", "python/ops/audio_ops.py", "python/ops/checkpoint_ops.py", + "python/ops/critical_section_ops.py", "python/ops/ops.py", "python/ops/prettyprint_ops.py", + "python/ops/script_ops.py", "python/ops/sort_ops.py", "python/ops/variables.py", ], @@ -60,6 +63,7 @@ tf_custom_op_py_library( "//tensorflow/python:math_ops", "//tensorflow/python:platform", "//tensorflow/python:pywrap_tensorflow", + "//tensorflow/python:script_ops", "//tensorflow/python:sparse_tensor", "//tensorflow/python:state_ops", "//tensorflow/python:state_ops_gen", @@ -70,6 +74,7 @@ tf_custom_op_py_library( "//tensorflow/python:variable_scope", "//tensorflow/python:variables", "//tensorflow/python/eager:context", + "//tensorflow/python/eager:function", "//third_party/py/numpy", "@six_archive//:six", ], @@ -173,6 +178,21 @@ py_test( ], ) +cuda_py_test( + name = "critical_section_test", + size = "medium", + srcs = ["python/ops/critical_section_test.py"], + additional_deps = [ + "//tensorflow/python:client_testlib", + ":framework_py", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:gradients", + "//tensorflow/python:platform_test", + "//tensorflow/python:resource_variable_ops", + ], +) + py_test( name = "accumulate_n_v2_eager_test", size = "small", diff --git a/tensorflow/contrib/framework/__init__.py b/tensorflow/contrib/framework/__init__.py index 4edc77f86ba786ca547b8d3842e2cf02833fbbac..a49d42cd525434d4ffd4a6bb0d8854dc707b9280 100644 --- a/tensorflow/contrib/framework/__init__.py +++ b/tensorflow/contrib/framework/__init__.py @@ -53,6 +53,7 @@ See the @{$python/contrib.framework} guide. @@assign_from_values_fn @@create_global_step @@filter_variables +@@fuse_op @@get_global_step @@get_or_create_global_step @@get_local_variables @@ -81,7 +82,15 @@ See the @{$python/contrib.framework} guide. @@load_linear_multiclass_bias_initializer @@load_variable_slot_initializer +@@py_func @@sort + +@@get_placeholders + +@@CriticalSection + +@@BoundedTensorSpec +@@TensorSpec """ from __future__ import absolute_import @@ -96,6 +105,9 @@ from tensorflow.contrib.framework.python.ops import * from tensorflow.python.framework.ops import prepend_name_scope from tensorflow.python.framework.ops import strip_name_scope +from tensorflow.python.framework.tensor_spec import BoundedTensorSpec +from tensorflow.python.framework.tensor_spec import TensorSpec + from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = ['nest'] diff --git a/tensorflow/contrib/framework/kernels/zero_initializer_op.cc b/tensorflow/contrib/framework/kernels/zero_initializer_op.cc index 6677dca752f84fc1ba7548b7739df04b7aaf14f7..5bf6b67529579e71a615c27e035111a58d5c02e0 100644 --- a/tensorflow/contrib/framework/kernels/zero_initializer_op.cc +++ b/tensorflow/contrib/framework/kernels/zero_initializer_op.cc @@ -21,8 +21,8 @@ limitations under the License. #include "tensorflow/contrib/framework/kernels/zero_initializer_op.h" -#include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/register_types.h" namespace tensorflow { @@ -81,8 +81,8 @@ TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPEC); #define REGISTER_GPU_KERNELS(T) REGISTER_KERNELS(GPU, T); TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS); #undef REGISTER_GPU_KERNELS -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA #undef REGISTER_KERNELS -} // namespace tensorflow +} // namespace tensorflow diff --git a/tensorflow/contrib/framework/kernels/zero_initializer_op.h b/tensorflow/contrib/framework/kernels/zero_initializer_op.h index 14c9268efa869ffd48b01dd2add44990ef7a43f8..99389a5ab6aa73c2ab0e522dd0f9fbc7093c8f4a 100644 --- a/tensorflow/contrib/framework/kernels/zero_initializer_op.h +++ b/tensorflow/contrib/framework/kernels/zero_initializer_op.h @@ -29,5 +29,5 @@ struct TensorSetZero { }; } // namespace functor -} // end namespace tensorflow -#endif // TENSORFLOW_CONTRIB_FRAMEWORK_KERNELS_ZERO_INITIALIZER_OP_H_ +} // end namespace tensorflow +#endif // TENSORFLOW_CONTRIB_FRAMEWORK_KERNELS_ZERO_INITIALIZER_OP_H_ diff --git a/tensorflow/contrib/framework/ops/variable_ops.cc b/tensorflow/contrib/framework/ops/variable_ops.cc index 1ee8e1498cf07559fe3db78ef832e2cdf26bea1c..706134ba9a51de6253ba7463b17ff662ea740ed0 100644 --- a/tensorflow/contrib/framework/ops/variable_ops.cc +++ b/tensorflow/contrib/framework/ops/variable_ops.cc @@ -26,8 +26,8 @@ REGISTER_OP("ZeroInitializer") .Attr("T: realnumbertype") .SetAllowsUninitializedInput() .SetShapeFn([](InferenceContext* c) { - c->set_output(0, c->input(0)); - return Status::OK(); + c->set_output(0, c->input(0)); + return Status::OK(); }) .Doc(R"doc( Initialize 'ref' with all zeros. This op requires that the tensor is not diff --git a/tensorflow/contrib/framework/python/framework/graph_util.py b/tensorflow/contrib/framework/python/framework/graph_util.py index a18ff2320d99726bb355ff6179fc97a070c2fec7..49eec3a3f1a0f357ea3adfade51e71cb0f89942d 100644 --- a/tensorflow/contrib/framework/python/framework/graph_util.py +++ b/tensorflow/contrib/framework/python/framework/graph_util.py @@ -133,6 +133,18 @@ def fuse_op(graph_def, input_nodes, output_nodes, output_dtypes, def get_placeholders(graph): """Get placeholders of a graph. + For example: + + ```python + a = tf.placeholder(dtype=tf.float32, shape=[2, 2], name='a') + a = tf.placeholder(dtype=tf.int32, shape=[3, 2], name='b') + + tf.contrib.framework.get_placeholders(tf.get_default_graph()) + # Returns: + # [, + # ] + ``` + Args: graph: A tf.Graph. Returns: diff --git a/tensorflow/contrib/framework/python/framework/tensor_util_test.py b/tensorflow/contrib/framework/python/framework/tensor_util_test.py index 2effe8eb26e98caa2707315d5f2e0e530ead31d3..8cdb340f2ddd9b3a7f55c1937ef045f4627e99be 100644 --- a/tensorflow/contrib/framework/python/framework/tensor_util_test.py +++ b/tensorflow/contrib/framework/python/framework/tensor_util_test.py @@ -30,6 +30,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors_impl from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor +from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import variables as variables_lib from tensorflow.python.platform import test @@ -77,6 +78,7 @@ class AssertScalarIntTest(test.TestCase): [3, 4], dtype=dtypes.int32)) +@test_util.with_c_api class WithShapeTest(test.TestCase): def _assert_with_shape(self, tensor, expected_value, expected_shape, @@ -213,16 +215,25 @@ class WithShapeTest(test.TestCase): tensor_partial_shape.set_shape([None, 2]) for incompatible_shape in [[0], [1]]: + if ops._USE_C_API: + error_message = "Shapes must be equal rank, but are 2 and 1" + else: + error_message = r"Shapes \(\?, 2\) and \([01],\) are not compatible" self.assertRaisesRegexp( - ValueError, r"Shapes \(\?, 2\) and \([01],\) are not compatible", + ValueError, error_message, tensor_util.with_shape, incompatible_shape, tensor_partial_shape) for incompatible_shape in [[1, 2, 1]]: self.assertRaisesRegexp(ValueError, "Dimensions must be equal", tensor_util.with_shape, incompatible_shape, tensor_partial_shape) for incompatible_shape in [[2, 1]]: + if ops._USE_C_API: + error_message = (r"Dimension 1 in both shapes must be equal, but are " + r"2 and 1. Shapes are \[\?,2\] and \[2,1\].") + else: + error_message = r"Shapes \(\?, 2\) and \(2, 1\) are not compatible" self.assertRaisesRegexp( - ValueError, r"Shapes \(\?, 2\) and \(2, 1\) are not compatible", + ValueError, error_message, tensor_util.with_shape, incompatible_shape, tensor_partial_shape) compatible_shape = [2, 2] diff --git a/tensorflow/contrib/framework/python/ops/__init__.py b/tensorflow/contrib/framework/python/ops/__init__.py index 685bb94779762ce46ee342e7e0a182c54be64743..c4976497f5fa95d82e492153b117681f693eaa13 100644 --- a/tensorflow/contrib/framework/python/ops/__init__.py +++ b/tensorflow/contrib/framework/python/ops/__init__.py @@ -22,8 +22,10 @@ from __future__ import print_function # pylint: disable=wildcard-import from tensorflow.contrib.framework.python.ops.arg_scope import * from tensorflow.contrib.framework.python.ops.checkpoint_ops import * +from tensorflow.contrib.framework.python.ops.critical_section_ops import * from tensorflow.contrib.framework.python.ops.ops import * from tensorflow.contrib.framework.python.ops.prettyprint_ops import * +from tensorflow.contrib.framework.python.ops.script_ops import * from tensorflow.contrib.framework.python.ops.sort_ops import * from tensorflow.contrib.framework.python.ops.variables import * # pylint: enable=wildcard-import diff --git a/tensorflow/contrib/framework/python/ops/accumulate_n_v2.py b/tensorflow/contrib/framework/python/ops/accumulate_n_v2.py index 2375ee4f550616ff60d20b87b5773704d8fbbe1e..476528b0dd3df05239d5dc402b466e06dd789985 100644 --- a/tensorflow/contrib/framework/python/ops/accumulate_n_v2.py +++ b/tensorflow/contrib/framework/python/ops/accumulate_n_v2.py @@ -22,6 +22,7 @@ from __future__ import print_function from tensorflow.python.eager import context 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 gen_math_ops from tensorflow.python.ops import math_ops @@ -108,4 +109,3 @@ def _AddNGrad(op, grad): """Same as gradient for AddN. Copies the gradient to all inputs.""" # Not broadcasting. return [grad] * len(op.inputs) - diff --git a/tensorflow/contrib/framework/python/ops/accumulate_n_v2_eager_test.py b/tensorflow/contrib/framework/python/ops/accumulate_n_v2_eager_test.py index 8f44698da851b48abf831e957c80fa1643a58bda..35974b9e21d2d7423777a95a99f51c9cb4b453b2 100644 --- a/tensorflow/contrib/framework/python/ops/accumulate_n_v2_eager_test.py +++ b/tensorflow/contrib/framework/python/ops/accumulate_n_v2_eager_test.py @@ -27,16 +27,11 @@ import numpy as np from tensorflow.contrib.framework.python.ops import accumulate_n_v2 as av2 from tensorflow.python.eager import backprop -from tensorflow.python.eager import context as eager_context -from tensorflow.python.eager import tape from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes as dtypes_lib from tensorflow.python.framework import ops from tensorflow.python.framework import test_util -from tensorflow.python.ops import gradients -from tensorflow.python.ops import math_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.platform import test diff --git a/tensorflow/contrib/framework/python/ops/accumulate_n_v2_test.py b/tensorflow/contrib/framework/python/ops/accumulate_n_v2_test.py index b5e9f8df79262635bf579a6bf2260bc40c140c6f..45962098e93acfac414396ddbeaa847701ff2b4b 100644 --- a/tensorflow/contrib/framework/python/ops/accumulate_n_v2_test.py +++ b/tensorflow/contrib/framework/python/ops/accumulate_n_v2_test.py @@ -22,7 +22,6 @@ import numpy as np from tensorflow.contrib.framework.python.ops import accumulate_n_v2 as av2 -from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes as dtypes_lib from tensorflow.python.framework import ops from tensorflow.python.framework import test_util @@ -31,7 +30,6 @@ from tensorflow.python.ops import variables from tensorflow.python.platform import googletest - class AccumulateNV2Test(test_util.TensorFlowTestCase): """Tests of the new, differentiable version of accumulate_n""" @@ -62,8 +60,9 @@ class AccumulateNV2Test(test_util.TensorFlowTestCase): accum_n = av2.accumulate_n_v2(input_vars) sess.run(variables.global_variables_initializer()) accum_n_grad = gradients.gradients(accum_n, input_vars) - self.assertAllEqual(np.repeat(1.0, num_inputs), # d/dx (x + y + ...) = 1 - [g.eval() for g in accum_n_grad]) + self.assertAllEqual( + np.repeat(1.0, num_inputs), # d/dx (x + y + ...) = 1 + [g.eval() for g in accum_n_grad]) # The tests below used to be in a separate class under cwise_ops_test.py, # which did not run in the default test target. @@ -75,8 +74,8 @@ class AccumulateNV2Test(test_util.TensorFlowTestCase): np.random.rand(16, 16, 16, 16).astype(np.float32) for _ in range(20) ] random_tensors = [ - ops.convert_to_tensor( - x, dtype=dtypes_lib.float32) for x in random_arrays + ops.convert_to_tensor(x, dtype=dtypes_lib.float32) + for x in random_arrays ] tf_val = av2.accumulate_n_v2(random_tensors) np_val = random_arrays[0] @@ -95,21 +94,21 @@ class AccumulateNV2Test(test_util.TensorFlowTestCase): with self.assertRaises(ValueError): a = variables.Variable(0.2) b = variables.Variable(0.1) - tf_val = av2.accumulate_n_v2([a,b], shape=[2,2]) # Should be shape=[] + tf_val = av2.accumulate_n_v2([a, b], shape=[2, 2]) # Should be shape=[] def testIncompatibleShapes(self): with self.test_session(): with self.assertRaises(ValueError): - a = variables.Variable(np.array([0.1,0.2])) - b = variables.Variable(np.array([[0.3],[0.4]])) - tf_val = av2.accumulate_n_v2([a,b]) + a = variables.Variable(np.array([0.1, 0.2])) + b = variables.Variable(np.array([[0.3], [0.4]])) + tf_val = av2.accumulate_n_v2([a, b]) def testWrongType(self): with self.test_session(): with self.assertRaises(TypeError): a = variables.Variable(0.2, dtype=np.float32) b = variables.Variable(0.1, dtype=np.float32) - tf_val = av2.accumulate_n_v2([a,b], tensor_dtype=np.int32) + tf_val = av2.accumulate_n_v2([a, b], tensor_dtype=np.int32) def testWrongTypeOneInput(self): # Scenario that used to trigger a bug, even when testWrongType() worked diff --git a/tensorflow/contrib/framework/python/ops/arg_scope.py b/tensorflow/contrib/framework/python/ops/arg_scope.py index 2bce00fde2459878a12027bb4d98bd3818bc92a2..409657fe1da0e5540cd2ad6070d86737c039e91f 100644 --- a/tensorflow/contrib/framework/python/ops/arg_scope.py +++ b/tensorflow/contrib/framework/python/ops/arg_scope.py @@ -53,7 +53,8 @@ net = layers.conv2d(net, 256, [5, 5], scope='conv2') ``` - Example of how to use tf.contrib.framework.add_arg_scope to enable your function to be called within an arg_scope later: + Example of how to use tf.contrib.framework.add_arg_scope to enable your + function to be called within an arg_scope later: @tf.contrib.framework.add_arg_scope def conv2d(*args, **kwargs) @@ -65,11 +66,10 @@ from __future__ import print_function from tensorflow.python.util import tf_contextlib from tensorflow.python.util import tf_decorator -__all__ = ['arg_scope', - 'add_arg_scope', - 'current_arg_scope', - 'has_arg_scope', - 'arg_scoped_arguments'] +__all__ = [ + 'arg_scope', 'add_arg_scope', 'current_arg_scope', 'has_arg_scope', + 'arg_scoped_arguments' +] _ARGSTACK = [{}] @@ -172,6 +172,7 @@ def add_arg_scope(func): Returns: A tuple with the decorated function func_with_args(). """ + def func_with_args(*args, **kwargs): current_scope = current_arg_scope() current_args = kwargs @@ -180,6 +181,7 @@ def add_arg_scope(func): current_args = current_scope[key_func].copy() current_args.update(kwargs) return func(*args, **current_args) + _add_op(func) setattr(func_with_args, '_key_op', _key_op(func)) return tf_decorator.make_decorator(func, func_with_args) diff --git a/tensorflow/contrib/framework/python/ops/critical_section_ops.py b/tensorflow/contrib/framework/python/ops/critical_section_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..182fec924febb74a23b82b1664d137f033f3b1b4 --- /dev/null +++ b/tensorflow/contrib/framework/python/ops/critical_section_ops.py @@ -0,0 +1,324 @@ +# Copyright 2017 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. +# ============================================================================== +"""Critical Section object and execution logic.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import collections + +# TODO(ebrevdo): Re-enable once CriticalSection is in core. +# from tensorflow.core.protobuf import critical_section_pb2 + +from tensorflow.python.eager import context +from tensorflow.python.eager import function +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.ops import gen_resource_variable_ops +from tensorflow.python.util import nest + + +# Graph Keys +CRITICAL_SECTIONS = "critical_sections" +CRITICAL_SECTION_EXECUTIONS = "critical_section_executions" + + +class _ExecutionSignature( + collections.namedtuple("_ExecutionSignature", + ("op", "exclusive_resource_access"))): + """A class storing an `ExecuteInCriticalResource` op and associated attrs.""" + pass + + +class CriticalSection(object): + """Critical section. + + A `CriticalSection` object is a resource in the graph which executes subgraphs + in **serial** order. A common example of a subgraph one may wish to run + exclusively is the one given by the following function: + + ```python + v = resource_variable_ops.ResourceVariable(0.0, name="v") + + def count(): + value = v.read_value() + with tf.control_dependencies([value]): + with tf.control_dependencies([v.assign_add(1)]): + return tf.identity(value) + ``` + + Here, a snapshot of `v` is captured in `value`; and then `v` is updated. + The snapshot value is returned. + + If multiple workers or threads all execute `count` in parallel, there is no + guarantee that access to the variable `v` is atomic at any point within + any thread's calculation of `count`. In fact, even implementing an atomic + counter that guarantees that the user will see each value `0, 1, ...,` is + currently impossible. + + The solution is to ensure any access to the underlying resource `v` is + only processed through a critical section: + + ```python + cs = CriticalSection() + f1 = cs.execute(count) + f2 = cs.execute(count) + output = f1 + f2 + session.run(output) + ``` + The functions `f1` and `f2` will be executed serially, and updates to `v` + will be atomic. + + **NOTES** + + All resource objects, including the critical section and any captured + variables of functions executed on that critical section, will be + colocated to the same device (host and cpu/gpu). + + When using multiple critical sections on the same resources, there is no + guarantee of exclusive access to those resources. This behavior is disallowed + by default (but see the kwarg `exclusive_resource_access`). + + For example, running the same function in two separate critical sections + will not ensure serial execution: + + ```python + v = tf.get_variable("v", initializer=0.0, use_resource=True) + def accumulate(up): + x = v.read_value() + with tf.control_dependencies([x]): + with tf.control_dependencies([v.assign_add(up)]): + return tf.identity(x) + ex1 = CriticalSection().execute( + accumulate, 1.0, exclusive_resource_access=False) + ex2 = CriticalSection().execute( + accumulate, 1.0, exclusive_resource_access=False) + bad_sum = ex1 + ex2 + sess.run(v.initializer) + sess.run(bad_sum) # May return 0.0 + ``` + """ + + def __init__(self, name=None, critical_section_def=None, import_scope=None): + """Creates a critical section.""" + if critical_section_def and name is not None: + raise ValueError("critical_section_def and name are mutually exclusive.") + if critical_section_def: + self._init_from_proto(critical_section_def, import_scope=import_scope) + else: + self._init_from_args(name) + + def _init_from_proto(self, critical_section_def, import_scope): + raise NotImplementedError("Not yet implemented") + # TODO(ebrevdo): Re-enable once CriticalSection is in core. + # assert isinstance( + # critical_section_def, critical_section_pb2.CriticalSectionDef) + # # Create from critical_section_def. + # g = ops.get_default_graph() + # self._handle = g.as_graph_element( + # ops.prepend_name_scope( + # critical_section_def.critical_section_name, + # import_scope=import_scope)) + + def _init_from_args(self, name): + """Initialize the CriticalSection from constructor arguments.""" + with ops.name_scope(name, "CriticalSection", []) as name: + with ops.control_dependencies(None): + # pylint: disable=protected-access + handle_name = ops._name_from_scope_name(name) + container = ops.get_default_graph()._container + # pylint: enable=protected-access + if container is None: + container = "" + self._handle = gen_resource_variable_ops.critical_section_op( + shared_name=handle_name, name=name) + if context.in_graph_mode(): + ops.add_to_collections(CRITICAL_SECTIONS, self) + + @property + def name(self): + return self._handle.op.name + + def execute(self, fn, *args, **kwargs): + """Execute function `fn(*args, **kwargs)` inside the CriticalSection. + + Args: + fn: The function to execute. Must return at least one tensor. + *args: Additional positional arguments to `fn`. + **kwargs: Additional keyword arguments to `fn`. + Several keywords are reserved for `execute`. These are: + + - name; The name to use when creating the execute operation. + - exclusive_resource_access; Whether the resources required by + `fn` should be exclusive to this `CriticalSection`. Default: `True`. + You may want to set this to `False` if you will be accessing a + resource in read-only mode in two different CriticalSections. + + Returns: + The tensors returned from `fn(*args, **kwargs)`. + + Raises: + ValueError: If `fn` attempts to use this `CriticalSection` in any nested + way. + ValueError: If `exclusive_resource_access` is not provided (is `True`) and + another `CriticalSection` has an execution requesting the same + resources as in `*args`, `**kwargs`, and any additionaly captured + inputs in `fn`. Note, even if `exclusive_resource_access` is `True`, + if another execution in another `CriticalSection` was created without + `exclusive_resource_access=True`, a `ValueError` will be raised. + """ + name = kwargs.pop("name", None) + exclusive_resource_access = kwargs.pop("exclusive_resource_access", True) + + args = nest.map_structure(ops.convert_to_tensor, args) + with ops.name_scope(name, "critical_section_execute", []): + fn_op = function.make_defun_op(fn, *args, **kwargs) + flat_dtypes = nest.flatten(fn_op.output_dtypes) + flat_shapes = nest.flatten(fn_op.output_shapes) + all_inputs = nest.flatten(args) + fn_op.captured_inputs + if self._handle in all_inputs: + raise ValueError("The function fn attempts to access the " + "CriticalSection in which it would be running. This " + "is illegal and would cause deadlocks. " + "CriticalSection: %s." % self._handle) + + if context.in_graph_mode(): + # Collections and op introspection does not work in eager + # mode. This is generally ok; since eager mode (as of + # writing) executes sequentially anyway. + all_input_resources = [ + x for x in all_inputs if x.dtype == dtypes.resource] + for sg in ops.get_collection(CRITICAL_SECTION_EXECUTIONS): + if sg.op.inputs[0].name == self._handle.name: + # Other executions in the same critical section are allowed. + continue + if not (exclusive_resource_access or sg.exclusive_resource_access): + # Neither execution requested exclusive access. + continue + sg_input_names = [y.name for y in sg.op.inputs[1:]] + for res in all_input_resources: + if res.name in sg_input_names: + raise ValueError( + "This execution would access resource %s; but either this " + "execution (CriticalSection: %s) or Execution '%s' " + "(CriticalSection: %s) requested exclusive resource access " + "of this resource for their critical section. Did you mean " + "to call execute with keyword argument " + "exclusive_resource_access=False?" + % (res.name, + self.name, + sg.op.name, + sg.op.inputs[0].op.name)) + + flat_outputs = gen_resource_variable_ops.execute_in_critical_section( + critical_section=self._handle, + arguments=all_inputs, + f=fn_op, + output_types=flat_dtypes, + output_shapes=flat_shapes) + + if context.in_graph_mode(): + if isinstance(flat_outputs, ops.Operation): + flat_outputs = [flat_outputs] + op = (flat_outputs[0].op if isinstance(flat_outputs[0], ops.Tensor) + else flat_outputs[0]) + signature = _ExecutionSignature( + op=op, + exclusive_resource_access=exclusive_resource_access) + ops.add_to_collections( + CRITICAL_SECTION_EXECUTIONS, signature) + + return (flat_outputs[0] + if (len(flat_outputs) == 1 + and isinstance(flat_outputs[0], ops.Operation)) + else nest.pack_sequence_as(fn_op.output_dtypes, flat_outputs)) + + # TODO(ebrevdo): Re-enable once CriticalSection is in core. + + # def to_proto(self, export_scope=None): + # """Converts a `CriticalSection` to a `CriticalSectoinDef` protocol buffer. + + # Args: + # export_scope: Optional `string`. Name scope to remove. + + # Returns: + # A `CriticalSectionDef` protocol buffer, or `None` if the + # `CriticalSection` is not in the specified name scope. + # """ + # if export_scope is None or self.handle.name.startswith(export_scope): + # cs_def = critical_section_pb2.CriticalSectionDef() + # cs_def.critical_section_name = ops.strip_name_scope( + # self._handle.name, export_scope) + # return cs_def + # else: + # return None + + # @staticmethod + # def from_proto(critical_section_def, import_scope=None): + # return CriticalSection( + # critical_section_def=critical_section_def, import_scope=import_scope) + + +# TODO(ebrevdo): Re-enable once CriticalSection is in core. + +# def _execution_to_proto_fn(execution_signature, export_scope=None): +# """Converts `_ExecutionSignature` to a `CriticalSectionExecutionDef`. + +# Args: +# execution_signature: Instance of `_ExecutionSignature`. +# export_scope: The export scope, if any. + +# Returns: +# An instance of `CriticalSectionExecutionDef`. +# """ +# if (export_scope is None +# or execution_signature.op.name.startswith(export_scope)): +# op_def = critical_section_pb2.CriticalSectionExecutionDef() +# op_def.execute_in_critical_section_name = ops.strip_name_scope( +# execution_signature.op.name, export_scope) +# op_def.exclusive_resource_access = ( +# execution_signature.exclusive_resource_access) +# return op_def +# else: +# return None + + +# def _execution_from_proto_fn(op_def, import_scope=None): +# """Converts a `CriticalSectionExecutionDef` to a `_ExecutionSignature`.""" +# assert isinstance( +# op_def, critical_section_pb2.CriticalSectionExecutionDef) + +# # Create from op_def. +# g = ops.get_default_graph() +# execution_op = g.as_graph_element( +# ops.prepend_name_scope( +# op_def.execute_in_critical_section_name, +# import_scope=import_scope)) +# return _ExecutionSignature( +# op=execution_op, +# exclusive_resource_access=op_def.exclusive_resource_access) + +# ops.register_proto_function( +# CRITICAL_SECTIONS, +# proto_type=critical_section_pb2.CriticalSectionDef, +# to_proto=CriticalSection.to_proto, +# from_proto=CriticalSection.from_proto) + +# ops.register_proto_function( +# CRITICAL_SECTION_EXECUTIONS, +# proto_type=critical_section_pb2.CriticalSectionExecutionDef, +# to_proto=_execution_to_proto_fn, +# from_proto=_execution_from_proto_fn) diff --git a/tensorflow/contrib/framework/python/ops/critical_section_test.py b/tensorflow/contrib/framework/python/ops/critical_section_test.py new file mode 100644 index 0000000000000000000000000000000000000000..a416724d3ba1719471d70667e140f9cd2daf86c7 --- /dev/null +++ b/tensorflow/contrib/framework/python/ops/critical_section_test.py @@ -0,0 +1,178 @@ +# Copyright 2017 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. +# ============================================================================== +"""critical section tests.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.framework.python.ops import critical_section_ops +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import function +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.platform import test +# TODO(ebrevdo): Re-enable once CriticalSection is in core. +# from tensorflow.python.training import saver as saver_lib + + +class CriticalSectionTest(test.TestCase): + + @test_util.run_in_graph_and_eager_modes() + def testCreateCriticalSection(self): + cs = critical_section_ops.CriticalSection(name="cs") + v = resource_variable_ops.ResourceVariable(0.0, name="v") + + def fn(a, b): + c = v.read_value() + with ops.control_dependencies([c]): + nv = v.assign_add(a * b) + with ops.control_dependencies([nv]): + return array_ops.identity(c) + + num_concurrent = 1000 + r = [cs.execute(fn, 1.0, 2.0) for _ in range(num_concurrent)] + self.evaluate(v.initializer) + r_value = self.evaluate(r) + self.assertAllClose([2.0 * i for i in range(num_concurrent)], + sorted(r_value)) + + @test_util.run_in_graph_and_eager_modes() + def testCreateCriticalSectionFnReturnsOp(self): + cs = critical_section_ops.CriticalSection(name="cs") + v = resource_variable_ops.ResourceVariable(0.0, name="v") + + def fn_return_op(a, b): + c = v.read_value() + with ops.control_dependencies([c]): + nv = v.assign_add(a * b) + with ops.control_dependencies([nv]): + return () + + num_concurrent = 100 + r = [cs.execute(fn_return_op, 1.0, 2.0) for _ in range(num_concurrent)] + self.evaluate(v.initializer) + self.evaluate(r) + final_v = self.evaluate(v) + self.assertAllClose(2.0 * num_concurrent, final_v) + + def testCreateCriticalSectionRaw(self): + cs = critical_section_ops.CriticalSection(name="cs") + v = resource_variable_ops.ResourceVariable(0.0, name="v") + + @function.Defun(dtypes.float32, dtypes.float32) + def fn(a, b): + c = v.read_value() + with ops.control_dependencies([c]): + nv = v.assign_add(a * b) + with ops.control_dependencies([nv]): + return array_ops.identity(c) + + def execute(fn, *args): + output_args = fn.definition.signature.output_arg + return resource_variable_ops.execute_in_critical_section( + critical_section=cs._handle, + arguments=list(args) + fn.captured_inputs, + f=fn, + output_types=[out.type for out in output_args], + output_shapes=[tensor_shape.TensorShape(None) for _ in output_args]) + + num_concurrent = 1000 + r = [execute(fn, 1.0, 2.0)[0] for _ in range(num_concurrent)] + self.evaluate(v.initializer) + r_value = self.evaluate(r) + self.assertAllClose([2.0 * i for i in range(num_concurrent)], + sorted(r_value)) + + def testCollection(self): + cs = critical_section_ops.CriticalSection(name="cs") + self.assertIn( + cs, ops.get_collection(critical_section_ops.CRITICAL_SECTIONS)) + execute_op = cs.execute(lambda x: x + 1, 1.0).op + self.assertIn( + execute_op, + [signature.op for signature in + ops.get_collection(critical_section_ops.CRITICAL_SECTION_EXECUTIONS)]) + + @test_util.run_in_graph_and_eager_modes() + def testRecursiveCriticalSectionAccessIsIllegal(self): + cs = critical_section_ops.CriticalSection(name="cs") + def fn(x): + return cs.execute(lambda x: x+1, x) + with self.assertRaisesRegexp( + ValueError, + r"attempts to access the CriticalSection in which it would be running"): + cs.execute(fn, 1.0) + + def testMultipleCSExecutionsRequestSameResource(self): + cs0 = critical_section_ops.CriticalSection() + cs1 = critical_section_ops.CriticalSection() + v = resource_variable_ops.ResourceVariable(0.0, name="v") + cs0.execute(lambda: v + 1) + # It's OK for the same CriticalSection to access this resource. + cs0.execute(lambda: v - 1) + # It's *not* OK for a different CriticalSection to access it by + # default. + with self.assertRaisesRegexp( + ValueError, "requested exclusive resource access"): + cs1.execute(lambda: v + 1) + # It's not even OK if the second call doesn't request exclusive access. + with self.assertRaisesRegexp( + ValueError, "requested exclusive resource access"): + cs1.execute(lambda: v + 1, exclusive_resource_access=False) + + v2 = resource_variable_ops.ResourceVariable(0.0, name="v2") + cs0.execute(lambda: v2 + 1, exclusive_resource_access=False) + # It's OK if neither requests exclusive resource access. + cs1.execute(lambda: v2 + 1, exclusive_resource_access=False) + + # It's not OK if the second request requires exlusive resource + # access. + with self.assertRaisesRegexp( + ValueError, "requested exclusive resource access"): + cs1.execute(lambda: v2 + 1) + + # TODO(ebrevdo): Re-enable once CriticalSection is in core. + # + # def testCriticalSectionAndExecuteOpSaverRoundTrip(self): + # cs = critical_section_ops.CriticalSection() + # r = cs.execute(lambda x: x + 1, 1.0) + # graph = ops.get_default_graph() + # meta_graph = saver_lib.export_meta_graph( + # graph=graph, collection_list=graph.get_all_collection_keys()) + # graph_copy = ops.Graph() + # with graph_copy.as_default(): + # _ = saver_lib.import_meta_graph(meta_graph, import_scope="imported") + # restored_cs = ops.get_collection(critical_section_ops.CRITICAL_SECTIONS) + # restored_exec = ops.get_collection( + # critical_section_ops.CRITICAL_SECTION_EXECUTIONS) + # self.assertEqual(1, len(restored_cs)) + # self.assertEqual(1, len(restored_exec)) + # self.assertEqual(restored_cs[0].name, "imported/%s" % cs.name) + # self.assertEqual(restored_exec[0].op.name, "imported/%s" % r.op.name) + + # def testToProto(self): + # cs = critical_section_ops.CriticalSection(name="cs") + # proto = cs.to_proto() + # self.assertEqual(proto.critical_section_name, cs._handle.name) + # cs_copy = critical_section_ops.CriticalSection.from_proto(proto) + # self.assertEqual(cs_copy._handle, cs._handle) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/framework/python/ops/script_ops.py b/tensorflow/contrib/framework/python/ops/script_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5d269fefdcfae7902b35e0f29f8cd12fcc58b882 --- /dev/null +++ b/tensorflow/contrib/framework/python/ops/script_ops.py @@ -0,0 +1,143 @@ +# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Script Language Operators. See the @{$python/script_ops} guide. + +@@py_func +""" + +# pylint: disable=g-bad-name +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops.script_ops import py_func as _py_func +from tensorflow.python.util import nest + +__all__ = ['py_func'] + + +def py_func(func, + args=(), + kwargs=None, + output_types=None, + output_shapes=None, + stateful=True, + name=None): + """Wraps a python function and uses it as a TensorFlow op. + + This function is a wrapper around `tf.py_func` and improve it with kwargs + and output_shapes. Further it changed some argument names. + + Given a python function `func`, which takes numpy arrays as its + inputs and returns numpy arrays as its outputs, wrap this function as an + operation in a TensorFlow graph. The following snippet constructs a simple + TensorFlow graph that invokes the `np.sinh()` NumPy function as a operation + in the graph: + + ```python + def my_func(x): + # x will be a numpy array with the contents of the placeholder below + return np.sinh(x) + inp = tf.placeholder(tf.float32) + y = tf.py_func(my_func, [inp], tf.float32) + ``` + + + **N.B.** The `tf.py_func()` operation has the following known limitations: + + * The body of the function (i.e. `func`) will not be serialized in a + `GraphDef`. Therefore, you should not use this function if you need to + serialize your model and restore it in a different environment. + + * The operation must run in the same address space as the Python program + that calls `tf.py_func()`. If you are using distributed TensorFlow, you + must run a `tf.train.Server` in the same process as the program that calls + `tf.py_func()` and you must pin the created operation to a device in that + server (e.g. using `with tf.device():`). + + Args: + func: A Python function, which accepts a list of NumPy `ndarray` objects + having element types that match the corresponding `tf.Tensor` objects + in `inp`, and returns a list of `ndarray` objects (or a single `ndarray`) + having element types that match the corresponding values in `Tout`. + args: A list of `Tensor` objects. + kwargs: A dict with `Tensor` objects as values. + output_types: A nested structure of tensorflow data types or a single + tensorflow data type if there is only one, indicating what `func` returns. + output_shapes: Same as output_types, except the types are replaces with + shapes (optional). + stateful: (Boolean.) If True, the function should be considered stateful. + If a function is stateless, when given the same input it will return the + same output and have no observable side effects. Optimizations such as + common subexpression elimination are only performed on stateless + operations. + name: A name for the operation (optional). + + Returns: + Tensorflow op that wraps the input python function. + """ + + if kwargs is None: + kwargs = {} + + if not isinstance(args, (list, tuple)): + raise TypeError('args must be list and not {}. args: {}'.format( + type(args), args)) + + if not isinstance(kwargs, dict): + raise TypeError('kwargs must be dict and not {}. args: {}'.format( + type(kwargs), kwargs)) + + # For dynamic type inference use callable output_types and output_shapes + if callable(output_types): + # If callable assume same signature and call with tensors and get the types + output_types = output_types(*args, **kwargs) + if callable(output_shapes): + # If callable assume same signature and call with tensors and get the shapes + output_shapes = output_shapes(*args, **kwargs) + + flat_output_types = nest.flatten(output_types) + args = (args, kwargs) + flat_args = nest.flatten(args) + + def python_function_wrapper(*py_args): + py_args, py_kwargs = nest.pack_sequence_as(args, py_args) + + ret = func(*py_args, **py_kwargs) + # TODO(alextp): Catch Exceptions and improve msg, because tensorflow + # ist not able to preserve the traceback, i.e. the Exceptions does not + # contain any information where the Exception was raised. + nest.assert_shallow_structure(output_types, ret) + return nest.flatten(ret) + + flat_values = _py_func( + python_function_wrapper, + flat_args, + flat_output_types, + stateful=stateful, + name=name) + + if output_shapes is not None: + # I am not sure if this is nessesary + output_shapes = nest.map_structure_up_to( + output_types, tensor_shape.as_shape, output_shapes) + + flattened_shapes = nest.flatten(output_shapes) + for ret_t, shape in zip(flat_values, flattened_shapes): + ret_t.set_shape(shape) + + return nest.pack_sequence_as(output_types, flat_values) diff --git a/tensorflow/contrib/framework/python/ops/variables.py b/tensorflow/contrib/framework/python/ops/variables.py index 3f1ece4510578b5ac39849c577fffbb2a3be45a7..0754c3e0e30a340910a43a3ce86f6ca10afe848e 100644 --- a/tensorflow/contrib/framework/python/ops/variables.py +++ b/tensorflow/contrib/framework/python/ops/variables.py @@ -25,6 +25,7 @@ import re from tensorflow.contrib.framework.python.ops import add_arg_scope as contrib_add_arg_scope from tensorflow.contrib.framework.python.ops import gen_variable_ops from tensorflow.contrib.util import loader +from tensorflow.core.protobuf import saver_pb2 from tensorflow.python import pywrap_tensorflow from tensorflow.python.framework import device as tf_device from tensorflow.python.framework import dtypes @@ -32,9 +33,8 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import variable_scope -from tensorflow.python.ops import gen_state_ops -from tensorflow.python.platform import tf_logging as logging from tensorflow.python.platform import resource_loader +from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import saver as tf_saver from tensorflow.python.training import training_util from tensorflow.python.util.deprecation import deprecated @@ -685,7 +685,8 @@ def assign_from_checkpoint_fn(model_path, var_list, ignore_missing_vars=False, 'Variable %s missing in checkpoint %s', var, model_path) var_list = available_vars if var_list: - saver = tf_saver.Saver(var_list, reshape=reshape_variables) + saver = tf_saver.Saver(var_list, reshape=reshape_variables, + write_version=saver_pb2.SaverDef.V1) def callback(session): saver.restore(session, model_path) return callback diff --git a/tensorflow/contrib/fused_conv/kernels/fused_conv_ops_gpu.h b/tensorflow/contrib/fused_conv/kernels/fused_conv_ops_gpu.h index fa7a3c03aa35c756252b22a004be91fa24c10e41..ba52697679dafc239b1dac5562573b3589877a8c 100644 --- a/tensorflow/contrib/fused_conv/kernels/fused_conv_ops_gpu.h +++ b/tensorflow/contrib/fused_conv/kernels/fused_conv_ops_gpu.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_FUSED_CONV_KERNELS_FUSED_CONV_OPS_GPU_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_FUSED_CONV_KERNELS_FUSED_CONV_OPS_GPU_H_ +#ifndef TENSORFLOW_CONTRIB_FUSED_CONV_KERNELS_FUSED_CONV_OPS_GPU_H_ +#define TENSORFLOW_CONTRIB_FUSED_CONV_KERNELS_FUSED_CONV_OPS_GPU_H_ #if GOOGLE_CUDA @@ -72,4 +72,4 @@ class FusedConvParameters : public ConvParameters { #endif // GOOGLE_CUDA -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_FUSED_CONV_KERNELS_FUSED_CONV_OPS_GPU_H_ +#endif // TENSORFLOW_CONTRIB_FUSED_CONV_KERNELS_FUSED_CONV_OPS_GPU_H_ diff --git a/tensorflow/contrib/fused_conv/ops/fused_conv2d_bias_activation_op.cc b/tensorflow/contrib/fused_conv/ops/fused_conv2d_bias_activation_op.cc index 6a56237f67c844a3daa546eb02d64c9e2658f639..bafd1d59418f0ba47ebbdaabbf06f8e5471fc1a1 100644 --- a/tensorflow/contrib/fused_conv/ops/fused_conv2d_bias_activation_op.cc +++ b/tensorflow/contrib/fused_conv/ops/fused_conv2d_bias_activation_op.cc @@ -25,13 +25,6 @@ limitations under the License. namespace tensorflow { -namespace { -// Return the string containing the list of valid activation modes, that can be -// used as an Attr() in REGISTER_OP. -string GetAllActivationModeAttrString() { return "activation_mode: {'Relu'}"; } - -} // namespace - // -------------------------------------------------------------------------- // TODO(pauldonnelly): Add support for double inputs and scales to this Op, diff --git a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_benchmark.py b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_benchmark.py index a65d4bc50ff796977e8ea7f652b7cbe3fe37f673..96cdd8b1ca4d56d12d38ea961ae73f3a3aa28968 100644 --- a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_benchmark.py +++ b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_benchmark.py @@ -116,7 +116,7 @@ def build_fused_conv_bias_relu_graph(device, input_shape, filter_shape, strides, for _ in range(1, num_iters): with ops.control_dependencies([fused_out]): # pylint: disable=g-line-too-long - fused_out = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation( + fused_out = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation( # pylint: disable=line-too-long inp, filt, bias, @@ -166,10 +166,10 @@ class FusedConv2DBiasActivationBenchmark(test.Benchmark): duration = (time.time() - start_time) / num_iters print("%s inputshape:%s filtershape:%s strides:%s padding:%s " - "%d iters: %.8f sec" % - (device, str(input_shape).replace(" ", ""), - str(filter_shape).replace(" ", ""), - str(strides).replace(" ", ""), padding, num_iters, duration)) + "%d iters: %.8f sec" % (device, str(input_shape).replace(" ", ""), + str(filter_shape).replace(" ", ""), + str(strides).replace(" ", ""), padding, + num_iters, duration)) name_template = ( "conv2d_{device}_input_shape_{inputshape}_filter_shape_{filtershape}_" "strides_{strides}_padding_{padding}") diff --git a/tensorflow/contrib/gan/BUILD b/tensorflow/contrib/gan/BUILD index b355a79b1a5d967eb82a30d41c073bbb52e0364c..0eb0e3cbe20f5804db5476c08167d4e1c9080cfa 100644 --- a/tensorflow/contrib/gan/BUILD +++ b/tensorflow/contrib/gan/BUILD @@ -55,6 +55,7 @@ py_test( name = "train_test", srcs = ["python/train_test.py"], srcs_version = "PY2AND3", + tags = ["notsan"], deps = [ ":features", ":namedtuples", @@ -177,6 +178,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":losses_impl", + ":namedtuples", "//tensorflow/python:util", ], ) @@ -188,6 +190,9 @@ py_test( deps = [ ":tuple_losses", "//tensorflow/python:client_testlib", + "//tensorflow/python:constant_op", + "//tensorflow/python:dtypes", + "//tensorflow/python:variables", "//third_party/py/numpy", ], ) @@ -395,6 +400,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":eval_utils", + ":namedtuples", "//tensorflow/python:array_ops", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", diff --git a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py index d3dca3d9e75fe1ef3be67143e18c0b51e84ad24c..082c42eba180917e732bb7890129dfa94bf00fec 100644 --- a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py +++ b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py @@ -59,7 +59,11 @@ _summary_type_map = { class GANEstimator(estimator.Estimator): """An estimator for Generative Adversarial Networks (GANs). - This Estimator is backed by TFGAN. + This Estimator is backed by TFGAN. The network functions follow the TFGAN API + except for one exception: if either `generator_fn` or `discriminator_fn` have + an argument called `mode`, then the tf.Estimator mode is passed in for that + argument. This helps with operations like batch normalization, which have + different train and evaluation behavior. Example: @@ -96,7 +100,7 @@ class GANEstimator(estimator.Estimator): # Generate samples from generator. predictions = np.array([ x for x in gan_estimator.predict(predict_input_fn)]) - ``` + ``` """ def __init__(self, @@ -233,9 +237,11 @@ def _gan_model_fn( def _make_gan_model(generator_fn, discriminator_fn, real_data, generator_inputs, generator_scope, add_summaries, mode): """Make a `GANModel`, and optionally pass in `mode`.""" - # If `generator_fn` has an argument `mode`, pass mode to it. + # If network functions have an argument `mode`, pass mode to it. if 'mode' in inspect.getargspec(generator_fn).args: generator_fn = functools.partial(generator_fn, mode=mode) + if 'mode' in inspect.getargspec(discriminator_fn).args: + discriminator_fn = functools.partial(discriminator_fn, mode=mode) gan_model = tfgan_train.gan_model( generator_fn, discriminator_fn, diff --git a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py index e752f0bcccda418b79d4fdabb27807394cbbb425..387a62bd741bd42c03dc1bf70592060c29ccd7a8 100644 --- a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py +++ b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py @@ -54,7 +54,8 @@ def generator_fn(noise_dict, mode): return layers.fully_connected(noise, noise.shape[1].value) -def discriminator_fn(data, _): +def discriminator_fn(data, unused_conditioning, mode): + del unused_conditioning, mode return layers.fully_connected(data, 1) @@ -99,7 +100,6 @@ def mock_head(testcase, expected_generator_inputs, expected_real_data, else: testcase.assertEqual(discriminator_scope_name, gan_model.discriminator_scope.name) - testcase.assertEqual(_or_none(discriminator_fn), gan_model.discriminator_fn) with ops.control_dependencies(assertions): if mode == model_fn_lib.ModeKeys.TRAIN: diff --git a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py index 82293b575aefa198a618ae7286ca24ebabd6987d..fdfabd07c13f689d075ecbb8786d725fa8a62d01 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py @@ -28,6 +28,7 @@ from __future__ import division from __future__ import print_function import functools +import os import sys import tarfile @@ -189,20 +190,34 @@ def get_graph_def_from_resource(filename): return graph_pb2.GraphDef.FromString(resource_loader.load_resource(filename)) -def get_graph_def_from_url_tarball(url, filename): - """Get a GraphDef proto from a tarball on the web.""" - def _progress(count, block_size, total_size): - sys.stdout.write('\r>> Downloading %s %.1f%%' % ( - url, float(count * block_size) / float(total_size) * 100.0)) - sys.stdout.flush() - tar_filename, _ = urllib.request.urlretrieve(url, reporthook=_progress) +def get_graph_def_from_url_tarball(url, filename, tar_filename=None): + """Get a GraphDef proto from a tarball on the web. + + Args: + url: Web address of tarball + filename: Filename of graph definition within tarball + tar_filename: Temporary download filename (None = always download) + + Returns: + A GraphDef loaded from a file in the downloaded tarball. + """ + if not (tar_filename and os.path.exists(tar_filename)): + + def _progress(count, block_size, total_size): + sys.stdout.write('\r>> Downloading %s %.1f%%' % + (url, + float(count * block_size) / float(total_size) * 100.0)) + sys.stdout.flush() + + tar_filename, _ = urllib.request.urlretrieve(url, tar_filename, _progress) with tarfile.open(tar_filename, 'r:gz') as tar: proto_str = tar.extractfile(filename).read() return graph_pb2.GraphDef.FromString(proto_str) def _default_graph_def_fn(): - return get_graph_def_from_url_tarball(INCEPTION_URL, INCEPTION_FROZEN_GRAPH) + return get_graph_def_from_url_tarball(INCEPTION_URL, INCEPTION_FROZEN_GRAPH, + os.path.basename(INCEPTION_URL)) def run_inception(images, @@ -300,6 +315,11 @@ def classifier_score(images, classifier_fn, num_batches=1): which captures how different the network's classification prediction is from the prior distribution over classes. + NOTE: This function consumes images, computes their logits, and then + computes the classifier score. If you would like to precompute many logits for + large batches, use clasifier_score_from_logits(), which this method also + uses. + Args: images: Images to calculate the classifier score for. classifier_fn: A function that takes images and produces logits based on a @@ -328,12 +348,15 @@ def classifier_score(images, classifier_fn, num_batches=1): def classifier_score_from_logits(logits): - """Classifier score for evaluating a conditional generative model. + """Classifier score for evaluating a generative model from logits. - This is based on the Inception Score, but for an arbitrary classifier. + This method computes the classifier score for a set of logits. This can be + used independently of the classifier_score() method, especially in the case + of using large batches during evaluation where we would like precompute all + of the logits before computing the classifier score. This technique is described in detail in https://arxiv.org/abs/1606.03498. In - summary, this function calculates + summary, this function calculates: exp( E[ KL(p(y|x) || p(y)) ] ) @@ -341,7 +364,8 @@ def classifier_score_from_logits(logits): the prior distribution over classes. Args: - logits: A 2D Tensor of logits. + logits: Precomputed 2D tensor of logits that will be used to + compute the classifier score. Returns: The classifier score. A floating-point scalar of the same type as the output @@ -363,6 +387,7 @@ def classifier_score_from_logits(logits): if logits_dtype != dtypes.float64: final_score = math_ops.cast(final_score, logits_dtype) + return final_score @@ -441,6 +466,11 @@ def frechet_classifier_distance(real_images, sample size to compute frechet classifier distance when comparing two generative models. + NOTE: This function consumes images, computes their activations, and then + computes the classifier score. If you would like to precompute many + activations for real and generated images for large batches, please use + frechet_clasifier_distance_from_activations(), which this method also uses. + Args: real_images: Real images to use to compute Frechet Inception distance. generated_images: Generated images to use to compute Frechet Inception @@ -452,7 +482,7 @@ def frechet_classifier_distance(real_images, Returns: The Frechet Inception distance. A floating-point scalar of the same type - as the output of `classifier_fn` + as the output of `classifier_fn`. """ real_images_list = array_ops.split( @@ -483,10 +513,13 @@ def frechet_classifier_distance(real_images, def frechet_classifier_distance_from_activations( real_activations, generated_activations): - """Classifier distance for evaluating a generative model. + """Classifier distance for evaluating a generative model from activations. - This is based on the Frechet Inception distance, but for an arbitrary - classifier. + This methods computes the Frechet classifier distance from activations of + real images and generated images. This can be used independently of the + frechet_classifier_distance() method, especially in the case of using large + batches during evaluation where we would like precompute all of the + activations before computing the classifier distance. This technique is described in detail in https://arxiv.org/abs/1706.08500. Given two Gaussian distribution with means m and m_w and covariance matrices @@ -499,21 +532,16 @@ def frechet_classifier_distance_from_activations( Inception score, this is a true distance and utilizes information about real world images. - Note that when computed using sample means and sample covariance matrices, - Frechet distance is biased. It is more biased for small sample sizes. (e.g. - even if the two distributions are the same, for a small sample size, the - expected Frechet distance is large). It is important to use the same - sample size to compute frechet classifier distance when comparing two - generative models. - Args: - real_activations: Real images to use to compute Frechet Inception distance. - generated_activations: Generated images to use to compute Frechet Inception - distance. + real_activations: 2D Tensor containing activations of real data. Shape is + [batch_size, activation_size]. + generated_activations: 2D Tensor containing activations of generated data. + Shape is [batch_size, activation_size]. Returns: - The Frechet Inception distance. A floating-point scalar of the same type - as the output of the activations. + The Frechet Inception distance. A floating-point scalar of the same type + as the output of the activations. + """ real_activations.shape.assert_has_rank(2) generated_activations.shape.assert_has_rank(2) diff --git a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py index 1e18c699ba93b5f524341c65d0a2db84556b65a2..61dc8646ddc10605561ae6b19e90f4739c346608 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py @@ -181,7 +181,8 @@ class ClassifierMetricsTest(test.TestCase): batch_size = 3 img = array_ops.ones([batch_size, 299, 299, 3]) pool = _run_with_mock( - classifier_metrics.run_inception, img, + classifier_metrics.run_inception, + img, output_tensor=classifier_metrics.INCEPTION_FINAL_POOL) self.assertTrue(isinstance(pool, ops.Tensor)) @@ -195,9 +196,12 @@ class ClassifierMetricsTest(test.TestCase): batch_size = 3 img = array_ops.ones([batch_size, 299, 299, 3]) logits, pool = _run_with_mock( - classifier_metrics.run_inception, img, - output_tensor=[classifier_metrics.INCEPTION_OUTPUT, - classifier_metrics.INCEPTION_FINAL_POOL]) + classifier_metrics.run_inception, + img, + output_tensor=[ + classifier_metrics.INCEPTION_OUTPUT, + classifier_metrics.INCEPTION_FINAL_POOL + ]) self.assertTrue(isinstance(logits, ops.Tensor)) self.assertTrue(isinstance(pool, ops.Tensor)) @@ -209,8 +213,10 @@ class ClassifierMetricsTest(test.TestCase): def test_inception_score_graph(self): """Test `inception_score` graph construction.""" - score = _run_with_mock(classifier_metrics.inception_score, - array_ops.zeros([6, 299, 299, 3]), num_batches=3) + score = _run_with_mock( + classifier_metrics.inception_score, + array_ops.zeros([6, 299, 299, 3]), + num_batches=3) self.assertTrue(isinstance(score, ops.Tensor)) score.shape.assert_has_rank(0) @@ -248,12 +254,14 @@ class ClassifierMetricsTest(test.TestCase): array_ops.zeros([8, 10], dtype=dtypes.int32), p_logits, q) with self.assertRaisesRegexp(ValueError, 'must be floating type'): - classifier_metrics._kl_divergence( - p, array_ops.zeros([8, 10], dtype=dtypes.int32), q) + classifier_metrics._kl_divergence(p, + array_ops.zeros( + [8, 10], dtype=dtypes.int32), q) with self.assertRaisesRegexp(ValueError, 'must be floating type'): - classifier_metrics._kl_divergence( - p, p_logits, array_ops.zeros([10], dtype=dtypes.int32)) + classifier_metrics._kl_divergence(p, p_logits, + array_ops.zeros( + [10], dtype=dtypes.int32)) with self.assertRaisesRegexp(ValueError, 'must have rank 2'): classifier_metrics._kl_divergence(array_ops.zeros([8]), p_logits, q) @@ -266,8 +274,9 @@ class ClassifierMetricsTest(test.TestCase): def test_inception_score_value(self): """Test that `inception_score` gives the correct value.""" - logits = np.array([np.array([1, 2] * 500 + [4]), - np.array([4, 5] * 500 + [6])]) + logits = np.array( + [np.array([1, 2] * 500 + [4]), + np.array([4, 5] * 500 + [6])]) unused_image = array_ops.zeros([2, 299, 299, 3]) incscore = _run_with_mock(classifier_metrics.inception_score, unused_image) @@ -285,9 +294,11 @@ class ClassifierMetricsTest(test.TestCase): test_pool_real_a = np.float32(np.random.randn(512, 256)) test_pool_gen_a = np.float32(np.random.randn(512, 256)) - fid_op = _run_with_mock(classifier_metrics.frechet_classifier_distance, - test_pool_real_a, test_pool_gen_a, - classifier_fn=lambda x: x) + fid_op = _run_with_mock( + classifier_metrics.frechet_classifier_distance, + test_pool_real_a, + test_pool_gen_a, + classifier_fn=lambda x: x) with self.test_session() as sess: actual_fid = sess.run(fid_op) @@ -296,6 +307,33 @@ class ClassifierMetricsTest(test.TestCase): self.assertAllClose(expected_fid, actual_fid, 0.0001) + def test_frechet_classifier_distance_covariance(self): + """Test that `frechet_classifier_distance` takes covariance into account.""" + np.random.seed(0) + + # Make num_examples > num_features to ensure scipy's sqrtm function + # doesn't return a complex matrix. + test_pool_reals, test_pool_gens = [], [] + for i in range(1, 11, 2): + test_pool_reals.append(np.float32(np.random.randn(2048, 256) * i)) + test_pool_gens.append(np.float32(np.random.randn(2048, 256) * i)) + + fid_ops = [] + for i in range(len(test_pool_reals)): + fid_ops.append(_run_with_mock( + classifier_metrics.frechet_classifier_distance, + test_pool_reals[i], + test_pool_gens[i], + classifier_fn=lambda x: x)) + + fids = [] + with self.test_session() as sess: + for fid_op in fid_ops: + fids.append(sess.run(fid_op)) + + # Check that the FIDs increase monotonically. + self.assertTrue(all(fid_a < fid_b for fid_a, fid_b in zip(fids, fids[1:]))) + def test_trace_sqrt_product_value(self): """Test that `trace_sqrt_product` gives the correct value.""" np.random.seed(0) diff --git a/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_test.py b/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_test.py index b960af28eaa969079b72c7aabcde2ad6cd1f5c68..871f1ad54e2559f5df28efa78f99997a866f7087 100644 --- a/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_test.py +++ b/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_test.py @@ -84,11 +84,11 @@ class ClassifierMetricsTest(test.TestCase): self.assertAllClose( np.array([0.014, 0.014], 'f'), np.array([x[0] for x in wscores], 'f'), - rtol=0.1) + rtol=0.15) self.assertAllClose( np.array([0.014, 0.020], 'f'), np.array([x[1] for x in wscores], 'f'), - rtol=0.1) + rtol=0.15) def test_sliced_wasserstein_distance_svd(self): """Test the distance.""" diff --git a/tensorflow/contrib/gan/python/eval/python/summaries_impl.py b/tensorflow/contrib/gan/python/eval/python/summaries_impl.py index 508b4d20d8767f42246a0d0c87f911b7ac612f45..0d1afad72da8a8e087239868e25ddebe23490d1e 100644 --- a/tensorflow/contrib/gan/python/eval/python/summaries_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/summaries_impl.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib.gan.python import namedtuples from tensorflow.contrib.gan.python.eval.python import eval_utils from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops @@ -38,16 +39,26 @@ def _assert_is_image(data): data.shape[1:].assert_is_fully_defined() -def add_gan_model_image_summaries(gan_model, grid_size=4): +def add_gan_model_image_summaries(gan_model, grid_size=4, model_summaries=True): """Adds image summaries for real and fake images. Args: gan_model: A GANModel tuple. grid_size: The size of an image grid. + model_summaries: Also add summaries of the model. Raises: ValueError: If real and generated data aren't images. """ + if isinstance(gan_model, namedtuples.CycleGANModel): + saved_params = locals() + saved_params.pop('gan_model', None) + with ops.name_scope('cyclegan_x2y_image_summaries'): + add_gan_model_image_summaries(gan_model.model_x2y, **saved_params) + with ops.name_scope('cyclegan_y2x_image_summaries'): + add_gan_model_image_summaries(gan_model.model_y2x, **saved_params) + return + _assert_is_image(gan_model.real_data) _assert_is_image(gan_model.generated_data) @@ -73,7 +84,9 @@ def add_gan_model_image_summaries(gan_model, grid_size=4): image_shape=generated_image_shape, num_channels=generated_channels), max_outputs=1) - add_gan_model_summaries(gan_model) + + if model_summaries: + add_gan_model_summaries(gan_model) def add_image_comparison_summaries(gan_model, num_comparisons=2, @@ -96,6 +109,15 @@ def add_image_comparison_summaries(gan_model, num_comparisons=2, ValueError: If the generator input, real, and generated data aren't all the same size. """ + if isinstance(gan_model, namedtuples.CycleGANModel): + saved_params = locals() + saved_params.pop('gan_model', None) + with ops.name_scope('cyclegan_x2y_image_comparison_summaries'): + add_image_comparison_summaries(gan_model.model_x2y, **saved_params) + with ops.name_scope('cyclegan_y2x_image_comparison_summaries'): + add_image_comparison_summaries(gan_model.model_y2x, **saved_params) + return + _assert_is_image(gan_model.generator_inputs) _assert_is_image(gan_model.generated_data) _assert_is_image(gan_model.real_data) @@ -133,6 +155,13 @@ def add_gan_model_summaries(gan_model): Args: gan_model: A GANModel tuple. """ + if isinstance(gan_model, namedtuples.CycleGANModel): + with ops.name_scope('cyclegan_x2y_summaries'): + add_gan_model_summaries(gan_model.model_x2y) + with ops.name_scope('cyclegan_y2x_summaries'): + add_gan_model_summaries(gan_model.model_y2x) + return + with ops.name_scope('generator_variables'): for var in gan_model.generator_variables: summary.histogram(var.name, var) @@ -147,6 +176,13 @@ def add_regularization_loss_summaries(gan_model): Args: gan_model: A GANModel tuple. """ + if isinstance(gan_model, namedtuples.CycleGANModel): + with ops.name_scope('cyclegan_x2y_regularization_loss_summaries'): + add_regularization_loss_summaries(gan_model.model_x2y) + with ops.name_scope('cyclegan_y2x_regularization_loss_summaries'): + add_regularization_loss_summaries(gan_model.model_y2x) + return + if gan_model.generator_scope: summary.scalar( 'generator_regularization_loss', diff --git a/tensorflow/contrib/gan/python/eval/python/summaries_test.py b/tensorflow/contrib/gan/python/eval/python/summaries_test.py index a3b02bcefc6cbaa6e24131b336b5c9c072bde52c..7956db43348c0cc0f3d372e92a2e343f5aa62013 100644 --- a/tensorflow/contrib/gan/python/eval/python/summaries_test.py +++ b/tensorflow/contrib/gan/python/eval/python/summaries_test.py @@ -57,40 +57,89 @@ def get_gan_model(): discriminator_fn=discriminator_model) +def get_cyclegan_model(): + with variable_scope.variable_scope('x2y'): + model_x2y = get_gan_model() + with variable_scope.variable_scope('y2x'): + model_y2x = get_gan_model() + return namedtuples.CycleGANModel( + model_x2y=model_x2y, + model_y2x=model_y2x, + reconstructed_x=array_ops.zeros([3, 30, 35, 6]), + reconstructed_y=array_ops.zeros([3, 30, 35, 6])) + + class SummariesTest(test.TestCase): - def testAddGanModelImageSummaries(self): - summaries.add_gan_model_image_summaries(get_gan_model(), grid_size=2) + 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(5, len(ops.get_collection(ops.GraphKeys.SUMMARIES))) + self.assertEquals(expected_num_summary_ops, + len(ops.get_collection(ops.GraphKeys.SUMMARIES))) with self.test_session(use_gpu=True): variables.global_variables_initializer().run() summary.merge_all().eval() - def testAddGanModelSummaries(self): - summaries.add_gan_model_summaries(get_gan_model()) + def test_add_gan_model_image_summaries(self): + self._test_add_gan_model_image_summaries_impl(get_gan_model, 5, True) + + def test_add_gan_model_image_summaries_no_model(self): + self._test_add_gan_model_image_summaries_impl(get_gan_model, 2, False) - self.assertEquals(3, len(ops.get_collection(ops.GraphKeys.SUMMARIES))) + def test_add_gan_model_image_summaries_for_cyclegan(self): + self._test_add_gan_model_image_summaries_impl(get_cyclegan_model, 10, + True) + + def _test_add_gan_model_summaries_impl(self, get_model_fn, + expected_num_summary_ops): + summaries.add_gan_model_summaries(get_model_fn()) + + self.assertEquals(expected_num_summary_ops, + len(ops.get_collection(ops.GraphKeys.SUMMARIES))) with self.test_session(use_gpu=True): variables.global_variables_initializer().run() summary.merge_all().eval() - def testAddRegularizationLossSummaries(self): - summaries.add_regularization_loss_summaries(get_gan_model()) + def test_add_gan_model_summaries(self): + self._test_add_gan_model_summaries_impl(get_gan_model, 3) + + def test_add_gan_model_summaries_for_cyclegan(self): + self._test_add_gan_model_summaries_impl(get_cyclegan_model, 6) - self.assertEquals(2, len(ops.get_collection(ops.GraphKeys.SUMMARIES))) + def _test_add_regularization_loss_summaries_impl(self, get_model_fn, + expected_num_summary_ops): + summaries.add_regularization_loss_summaries(get_model_fn()) + + self.assertEquals(expected_num_summary_ops, + len(ops.get_collection(ops.GraphKeys.SUMMARIES))) with self.test_session(use_gpu=True): summary.merge_all().eval() + def test_add_regularization_loss_summaries(self): + self._test_add_regularization_loss_summaries_impl(get_gan_model, 2) + + def test_add_regularization_loss_summaries_for_cyclegan(self): + self._test_add_regularization_loss_summaries_impl(get_cyclegan_model, 4) + # TODO(joelshor): Add correctness test. - def testAddImageComparisonSummaries(self): - summaries.add_image_comparison_summaries( - get_gan_model(), display_diffs=True) + def _test_add_image_comparison_summaries_impl(self, get_model_fn, + expected_num_summary_ops): + summaries.add_image_comparison_summaries(get_model_fn(), display_diffs=True) - self.assertEquals(1, len(ops.get_collection(ops.GraphKeys.SUMMARIES))) + self.assertEquals(expected_num_summary_ops, + len(ops.get_collection(ops.GraphKeys.SUMMARIES))) with self.test_session(use_gpu=True): summary.merge_all().eval() + def test_add_image_comparison_summaries(self): + self._test_add_image_comparison_summaries_impl(get_gan_model, 1) + + def test_add_image_comparison_summaries_for_cyclegan(self): + self._test_add_image_comparison_summaries_impl(get_cyclegan_model, 2) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/gan/python/features/python/random_tensor_pool_impl.py b/tensorflow/contrib/gan/python/features/python/random_tensor_pool_impl.py index 9d10db0f5a3d09dc4dd7d8b1c97c16c29808547c..4cfae0de4451880cf8229903b0eb74b1c6e2e04d 100644 --- a/tensorflow/contrib/gan/python/features/python/random_tensor_pool_impl.py +++ b/tensorflow/contrib/gan/python/features/python/random_tensor_pool_impl.py @@ -128,6 +128,10 @@ def tensor_pool(input_values, pool_queue.size() < pool_size, _get_input_value_pooled, _get_random_pool_value_and_enqueue_input)) + # Make sure that the shape of `output_value` is set. + for input_value, output_value in zip(input_values, output_values): + output_value.set_shape(input_value.shape) + if isinstance(original_input_values, list): return list(output_values) elif isinstance(original_input_values, tuple): diff --git a/tensorflow/contrib/gan/python/features/python/random_tensor_pool_test.py b/tensorflow/contrib/gan/python/features/python/random_tensor_pool_test.py index cef3a87ab34f9754099073eefcb3f1b1c97a3762..d8cf549cf71838178c9da01df462d41d81595fe5 100644 --- a/tensorflow/contrib/gan/python/features/python/random_tensor_pool_test.py +++ b/tensorflow/contrib/gan/python/features/python/random_tensor_pool_test.py @@ -33,6 +33,7 @@ class TensorPoolTest(test.TestCase): input_value = array_ops.placeholder( dtype=dtypes.int32, shape=[None, None, 3]) output_value = tensor_pool(input_value, pool_size=10) + self.assertEqual(output_value.shape.as_list(), [None, None, 3]) with self.test_session(use_gpu=True) as session: for i in range(10): @@ -44,6 +45,7 @@ class TensorPoolTest(test.TestCase): """Checks that values are pooled and returned maximally twice.""" input_value = array_ops.placeholder(dtype=dtypes.int32, shape=[]) output_value = tensor_pool(input_value, pool_size=10) + self.assertEqual(output_value.shape.as_list(), []) with self.test_session(use_gpu=True) as session: outs = [] @@ -61,6 +63,7 @@ class TensorPoolTest(test.TestCase): input_value = array_ops.placeholder(dtype=dtypes.int32, shape=[]) output_value = tensor_pool( input_value, pool_size=10, pooling_probability=0.0) + self.assertEqual(output_value.shape.as_list(), []) with self.test_session(use_gpu=True) as session: for i in range(50): @@ -76,6 +79,7 @@ class TensorPoolTest(test.TestCase): input_value, pool_size=pool_size, pooling_probability=pooling_probability) + self.assertEqual(output_value.shape.as_list(), []) with self.test_session(use_gpu=True) as session: not_pooled = 0 @@ -95,6 +99,8 @@ class TensorPoolTest(test.TestCase): array_ops.placeholder(dtype=dtypes.int32, shape=[])) output_values = tensor_pool(input_values, pool_size=3) self.assertEqual(len(output_values), len(input_values)) + for output_value in output_values: + self.assertEqual(output_value.shape.as_list(), []) with self.test_session(use_gpu=True) as session: for i in range(10): diff --git a/tensorflow/contrib/gan/python/losses/python/losses_impl.py b/tensorflow/contrib/gan/python/losses/python/losses_impl.py index 940762cf2aa0f473cd41d9d543e2773b565a5248..39588b7219ebac1cc4855532be3fcc38e6381134 100644 --- a/tensorflow/contrib/gan/python/losses/python/losses_impl.py +++ b/tensorflow/contrib/gan/python/losses/python/losses_impl.py @@ -67,6 +67,7 @@ __all__ = [ 'wasserstein_gradient_penalty', 'mutual_information_penalty', 'combine_adversarial_loss', + 'cycle_consistency_loss', ] @@ -304,6 +305,7 @@ def wasserstein_gradient_penalty( discriminator_fn, discriminator_scope, epsilon=1e-10, + target=1.0, weights=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES, @@ -323,6 +325,8 @@ def wasserstein_gradient_penalty( discriminator_scope: If not `None`, reuse discriminators from this scope. epsilon: A small positive number added for numerical stability when computing the gradient norm. + target: Optional Python number or `Tensor` indicating the target value of + gradient norm. Defaults to 1.0. weights: Optional `Tensor` whose rank is either 0, or the same rank as `real_data` and `generated_data`, and must be broadcastable to them (i.e., all dimensions must be either `1`, or the same as the @@ -373,7 +377,7 @@ def wasserstein_gradient_penalty( # For numerical stability, add epsilon to the sum before taking the square # root. Note tf.norm does not add epsilon. slopes = math_ops.sqrt(gradient_squares + epsilon) - penalties = math_ops.square(slopes - 1.0) + penalties = math_ops.square(slopes / target - 1.0) penalty = losses.compute_weighted_loss( penalties, weights, scope=scope, loss_collection=loss_collection, reduction=reduction) @@ -915,3 +919,63 @@ def combine_adversarial_loss(main_loss, array_ops.stop_gradient(adv_coeff) * adversarial_loss) return final_loss + + +def cycle_consistency_loss(data_x, + reconstructed_data_x, + data_y, + reconstructed_data_y, + scope=None, + add_summaries=False): + """Defines the cycle consistency loss. + + The cyclegan model has two partial models where `model_x2y` generator F maps + data set X to Y, `model_y2x` generator G maps data set Y to X. For a `data_x` + in data set X, we could reconstruct it by + * reconstructed_data_x = G(F(data_x)) + Similarly + * reconstructed_data_y = F(G(data_y)) + + The cycle consistency loss is about the difference between data and + reconstructed data, namely + * loss_x2x = |data_x - G(F(data_x))| (L1-norm) + * loss_y2y = |data_y - F(G(data_y))| (L1-norm) + * loss = (loss_x2x + loss_y2y) / 2 + where `loss` is the final result. + + See https://arxiv.org/abs/1703.10593 for more details. + + Args: + data_x: A `Tensor` of data X. + reconstructed_data_x: A `Tensor` of reconstructed data X. + data_y: A `Tensor` of data Y. + reconstructed_data_y: A `Tensor` of reconstructed data Y. + scope: The scope for the operations performed in computing the loss. + Defaults to None. + add_summaries: Whether or not to add detailed summaries for the loss. + Defaults to False. + + Returns: + A scalar `Tensor` of cycle consistency loss. + """ + + def _partial_cycle_consistency_loss(data, reconstructed_data): + # Following the original implementation + # https://github.com/junyanz/CycleGAN/blob/master/models/cycle_gan_model.lua + # use L1-norm of pixel-wise error normalized by data size so that + # `cycle_loss_weight` can be specified independent of image size. + return math_ops.reduce_mean(math_ops.abs(data - reconstructed_data)) + + with ops.name_scope( + scope, + 'cycle_consistency_loss', + values=[data_x, reconstructed_data_x, data_y, reconstructed_data_y]): + loss_x2x = _partial_cycle_consistency_loss(data_x, reconstructed_data_x) + loss_y2y = _partial_cycle_consistency_loss(data_y, reconstructed_data_y) + loss = (loss_x2x + loss_y2y) / 2.0 + if add_summaries: + summary.scalar('cycle_consistency_loss_x2x', loss_x2x) + summary.scalar('cycle_consistency_loss_y2y', loss_y2y) + summary.scalar('cycle_consistency_loss', loss) + + return loss diff --git a/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py b/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py index b5cd8c92ba180e981e0faf877021cb6d69dc34b4..dbaa624ae9d6a5a5949db692e52c0c1deb18b8df 100644 --- a/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py +++ b/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py @@ -481,6 +481,29 @@ class GradientPenaltyTest(test.TestCase, _PenaltyTest): }) self.assertAlmostEqual(self._expected_loss, loss, 5) + def test_loss_with_gradient_norm_target(self): + """Test loss value with non default gradient norm target.""" + generated_data = array_ops.placeholder(dtypes.float32, shape=(None, None)) + real_data = array_ops.placeholder(dtypes.float32, shape=(None, None)) + + loss = tfgan_losses.wasserstein_gradient_penalty( + generated_data, + real_data, + self._kwargs['generator_inputs'], + self._kwargs['discriminator_fn'], + self._kwargs['discriminator_scope'], + target=2.0) + + with self.test_session() as sess: + variables.global_variables_initializer().run() + loss = sess.run( + loss, + feed_dict={ + generated_data: self._generated_data_np, + real_data: self._real_data_np, + }) + self.assertAlmostEqual(1.0, loss, 5) + def test_reuses_scope(self): """Test that gradient penalty reuses discriminator scope.""" num_vars = len(ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)) @@ -620,7 +643,34 @@ class CombineAdversarialLossTest(test.TestCase): with self.test_session(use_gpu=True) as sess: for _ in range(10): # spot check closeness on more than one sample. gnorm_np, precond_gnorm_np = sess.run([gnorm, precond_gnorm]) - self.assertNear(gnorm_np, precond_gnorm_np, 1e-5) + self.assertNear(gnorm_np, precond_gnorm_np, 1e-4) + + +class CycleConsistencyLossTest(test.TestCase): + """Tests for cycle_consistency_loss.""" + + def setUp(self): + super(CycleConsistencyLossTest, self).setUp() + + self._data_x_np = [[1.0, 2, 3], [4, 5, 6]] + self._reconstructed_data_x_np = [[7.0, 8, 9], [10, 11, 12]] + self._data_y_np = [1.0, 9] + self._reconstructed_data_y_np = [-2.0, 3] + + self._data_x = constant_op.constant(self._data_x_np, dtype=dtypes.float32) + self._reconstructed_data_x = constant_op.constant( + self._reconstructed_data_x_np, dtype=dtypes.float32) + self._data_y = constant_op.constant(self._data_y_np, dtype=dtypes.float32) + self._reconstructed_data_y = constant_op.constant( + self._reconstructed_data_y_np, dtype=dtypes.float32) + + def test_correct_loss(self): + loss = tfgan_losses.cycle_consistency_loss( + self._data_x, self._reconstructed_data_x, self._data_y, + self._reconstructed_data_y) + with self.test_session(use_gpu=True): + variables.global_variables_initializer().run() + self.assertNear(5.25, loss.eval(), 1e-5) if __name__ == '__main__': diff --git a/tensorflow/contrib/gan/python/losses/python/tuple_losses_impl.py b/tensorflow/contrib/gan/python/losses/python/tuple_losses_impl.py index b341f03a0ddaacca8b036189516c71908bee50eb..dcc3f94c2d6b9e5e44036e7cc1a9d1bb39104fb5 100644 --- a/tensorflow/contrib/gan/python/losses/python/tuple_losses_impl.py +++ b/tensorflow/contrib/gan/python/losses/python/tuple_losses_impl.py @@ -60,6 +60,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib.gan.python import namedtuples from tensorflow.contrib.gan.python.losses.python import losses_impl from tensorflow.python.util import tf_inspect @@ -78,6 +79,7 @@ __all__ = [ 'wasserstein_gradient_penalty', 'mutual_information_penalty', 'combine_adversarial_loss', + 'cycle_consistency_loss', ] @@ -246,3 +248,32 @@ def combine_adversarial_loss(gan_loss, scalar_summaries, gradient_summaries) return gan_loss._replace(generator_loss=combined_loss) + + +def cycle_consistency_loss(cyclegan_model, scope=None, add_summaries=False): + """Defines the cycle consistency loss. + + Uses `cycle_consistency_loss` to compute the cycle consistency loss for a + `cyclegan_model`. + + Args: + cyclegan_model: A `CycleGANModel` namedtuple. + scope: The scope for the operations performed in computing the loss. + Defaults to None. + add_summaries: Whether or not to add detailed summaries for the loss. + Defaults to False. + + Returns: + A scalar `Tensor` of cycle consistency loss. + + Raises: + ValueError: If `cyclegan_model` is not a `CycleGANModel` namedtuple. + """ + if not isinstance(cyclegan_model, namedtuples.CycleGANModel): + raise ValueError( + '`cyclegan_model` must be a `CycleGANModel`. Instead, was %s.' % + type(cyclegan_model)) + return losses_impl.cycle_consistency_loss( + cyclegan_model.model_x2y.generator_inputs, cyclegan_model.reconstructed_x, + cyclegan_model.model_y2x.generator_inputs, cyclegan_model.reconstructed_y, + scope, add_summaries) diff --git a/tensorflow/contrib/gan/python/losses/python/tuple_losses_test.py b/tensorflow/contrib/gan/python/losses/python/tuple_losses_test.py index 215b15ef6915d0b8113def35987ed6ab85617bcc..aa1ef11172dee6799994b87f70a3883cd67fd15b 100644 --- a/tensorflow/contrib/gan/python/losses/python/tuple_losses_test.py +++ b/tensorflow/contrib/gan/python/losses/python/tuple_losses_test.py @@ -22,8 +22,11 @@ import collections import numpy as np +from tensorflow.contrib.gan.python import namedtuples from tensorflow.contrib.gan.python.losses.python import tuple_losses_impl as tfgan_losses - +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import variables from tensorflow.python.platform import test @@ -125,6 +128,7 @@ manual_tests = [ 'combine_adversarial_loss', 'mutual_information_penalty', 'wasserstein_gradient_penalty', + 'cycle_consistency_loss', ] discriminator_keyword_args = { @@ -139,6 +143,38 @@ generator_keyword_args = { } +class CycleConsistencyLossTest(test.TestCase): + + def setUp(self): + super(CycleConsistencyLossTest, self).setUp() + + def _partial_model(generator_inputs_np): + model = namedtuples.GANModel(*[None] * 11) + return model._replace( + generator_inputs=constant_op.constant( + generator_inputs_np, dtype=dtypes.float32)) + + self._model_x2y = _partial_model([1, 2]) + self._model_y2x = _partial_model([5, 6]) + + def test_model_type(self): + """Test the input model type for `cycle_consistency_loss`.""" + with self.assertRaises(ValueError): + tfgan_losses.cycle_consistency_loss(self._model_x2y) + + def test_correct_loss(self): + """Test the output of `cycle_consistency_loss`.""" + loss = tfgan_losses.cycle_consistency_loss( + namedtuples.CycleGANModel( + model_x2y=self._model_x2y, + model_y2x=self._model_y2x, + reconstructed_x=constant_op.constant([9, 8], dtype=dtypes.float32), + reconstructed_y=constant_op.constant([7, 2], dtype=dtypes.float32))) + with self.test_session(use_gpu=True): + variables.global_variables_initializer().run() + self.assertNear(5.0, loss.eval(), 1e-5) + + if __name__ == '__main__': for loss_name in tfgan_losses.__all__: if loss_name in manual_tests: continue diff --git a/tensorflow/contrib/gan/python/namedtuples.py b/tensorflow/contrib/gan/python/namedtuples.py index 3d4e315ebd0bd52b3b5e3e4a8655df8bfe9cebe8..25cfeafeec9000b0dc3849ebe646e59c1b4d1cc3 100644 --- a/tensorflow/contrib/gan/python/namedtuples.py +++ b/tensorflow/contrib/gan/python/namedtuples.py @@ -30,7 +30,9 @@ __all__ = [ 'GANModel', 'InfoGANModel', 'ACGANModel', + 'CycleGANModel', 'GANLoss', + 'CycleGANLoss', 'GANTrainOps', 'GANTrainSteps', ] @@ -115,6 +117,25 @@ class ACGANModel( """ +class CycleGANModel( + collections.namedtuple( + 'CycleGANModel', + ('model_x2y', 'model_y2x', 'reconstructed_x', 'reconstructed_y'))): + """An CycleGANModel contains all the pieces needed for CycleGAN training. + + The model `model_x2y` generator F maps data set X to Y, while the model + `model_y2x` generator G maps data set Y to X. + + See https://arxiv.org/abs/1703.10593 for more details. + + Args: + model_x2y: A `GANModel` namedtuple whose generator maps data set X to Y. + model_y2x: A `GANModel` namedtuple whose generator maps data set Y to X. + reconstructed_x: A `Tensor` of reconstructed data X which is G(F(X)). + reconstructed_y: A `Tensor` of reconstructed data Y which is F(G(Y)). + """ + + class GANLoss( collections.namedtuple('GANLoss', ( 'generator_loss', @@ -128,6 +149,18 @@ class GANLoss( """ +class CycleGANLoss( + collections.namedtuple('CycleGANLoss', ('loss_x2y', 'loss_y2x'))): + """CycleGANLoss contains the losses for `CycleGANModel`. + + See https://arxiv.org/abs/1703.10593 for more details. + + Args: + loss_x2y: A `GANLoss` namedtuple representing the loss of `model_x2y`. + loss_y2x: A `GANLoss` namedtuple representing the loss of `model_y2x`. + """ + + class GANTrainOps( collections.namedtuple('GANTrainOps', ( 'generator_train_op', diff --git a/tensorflow/contrib/gan/python/train.py b/tensorflow/contrib/gan/python/train.py index edd0113977ff4ddc672b0ec134be1a48c621b579..776eb11ecb1624544d24611d8fe6ca19768b8313 100644 --- a/tensorflow/contrib/gan/python/train.py +++ b/tensorflow/contrib/gan/python/train.py @@ -52,7 +52,9 @@ __all__ = [ 'gan_model', 'infogan_model', 'acgan_model', + 'cyclegan_model', 'gan_loss', + 'cyclegan_loss', 'gan_train_ops', 'gan_train', 'get_sequential_train_hooks', @@ -277,14 +279,16 @@ def acgan_model( generator_inputs = _convert_tensor_or_l_or_d(generator_inputs) generated_data = generator_fn(generator_inputs) with variable_scope.variable_scope(discriminator_scope) as dis_scope: - (discriminator_gen_outputs, discriminator_gen_classification_logits - ) = _validate_acgan_discriminator_outputs( - discriminator_fn(generated_data, generator_inputs)) + with ops.name_scope(dis_scope.name+'/generated/'): + (discriminator_gen_outputs, discriminator_gen_classification_logits + ) = _validate_acgan_discriminator_outputs( + discriminator_fn(generated_data, generator_inputs)) with variable_scope.variable_scope(dis_scope, reuse=True): - real_data = ops.convert_to_tensor(real_data) - (discriminator_real_outputs, discriminator_real_classification_logits - ) = _validate_acgan_discriminator_outputs( - discriminator_fn(real_data, generator_inputs)) + with ops.name_scope(dis_scope.name+'/real/'): + real_data = ops.convert_to_tensor(real_data) + (discriminator_real_outputs, discriminator_real_classification_logits + ) = _validate_acgan_discriminator_outputs( + discriminator_fn(real_data, generator_inputs)) if check_shapes: if not generated_data.shape.is_compatible_with(real_data.shape): raise ValueError( @@ -305,6 +309,76 @@ def acgan_model( discriminator_gen_classification_logits) +def cyclegan_model( + # Lambdas defining models. + generator_fn, + discriminator_fn, + # data X and Y. + data_x, + data_y, + # Optional scopes. + generator_scope='Generator', + discriminator_scope='Discriminator', + model_x2y_scope='ModelX2Y', + model_y2x_scope='ModelY2X', + # Options. + check_shapes=True): + """Returns a CycleGAN model outputs and variables. + + See https://arxiv.org/abs/1703.10593 for more details. + + Args: + generator_fn: A python lambda that takes `data_x` or `data_y` as inputs and + returns the outputs of the GAN generator. + discriminator_fn: A python lambda that takes `real_data`/`generated data` + and `generator_inputs`. Outputs a Tensor in the range [-inf, inf]. + data_x: A `Tensor` of dataset X. Must be the same shape as `data_y`. + data_y: A `Tensor` of dataset Y. Must be the same shape as `data_x`. + generator_scope: Optional generator variable scope. Useful if you want to + reuse a subgraph that has already been created. Defaults to 'Generator'. + discriminator_scope: Optional discriminator variable scope. Useful if you + want to reuse a subgraph that has already been created. Defaults to + 'Discriminator'. + model_x2y_scope: Optional variable scope for model x2y variables. Defaults + to 'ModelX2Y'. + model_y2x_scope: Optional variable scope for model y2x variables. Defaults + to 'ModelY2X'. + check_shapes: If `True`, check that generator produces Tensors that are the + same shape as `data_x` (`data_y`). Otherwise, skip this check. + + Returns: + A `CycleGANModel` namedtuple. + + Raises: + ValueError: If `check_shapes` is True and `data_x` or the generator output + does not have the same shape as `data_y`. + """ + + # Create models. + def _define_partial_model(input_data, output_data): + return gan_model( + generator_fn=generator_fn, + discriminator_fn=discriminator_fn, + real_data=output_data, + generator_inputs=input_data, + generator_scope=generator_scope, + discriminator_scope=discriminator_scope, + check_shapes=check_shapes) + + with variable_scope.variable_scope(model_x2y_scope): + model_x2y = _define_partial_model(data_x, data_y) + with variable_scope.variable_scope(model_y2x_scope): + model_y2x = _define_partial_model(data_y, data_x) + + with variable_scope.variable_scope(model_y2x.generator_scope, reuse=True): + reconstructed_x = model_y2x.generator_fn(model_x2y.generated_data) + with variable_scope.variable_scope(model_x2y.generator_scope, reuse=True): + reconstructed_y = model_x2y.generator_fn(model_y2x.generated_data) + + return namedtuples.CycleGANModel(model_x2y, model_y2x, reconstructed_x, + reconstructed_y) + + def _validate_aux_loss_weight(aux_loss_weight, name='aux_loss_weight'): if isinstance(aux_loss_weight, ops.Tensor): aux_loss_weight.shape.assert_is_compatible_with([]) @@ -343,7 +417,7 @@ def _tensor_pool_adjusted_model(model, tensor_pool_fn): `tensor_pool_fn` is None. Raises: - ValueError: If tensor pool does not suport the `model`. + ValueError: If tensor pool does not support the `model`. """ if tensor_pool_fn is None: return model @@ -352,21 +426,24 @@ def _tensor_pool_adjusted_model(model, tensor_pool_fn): (model.generated_data, model.generator_inputs)) if isinstance(model, namedtuples.GANModel): - dis_gen_outputs = model.discriminator_fn(pooled_generated_data, - pooled_generator_inputs) + with variable_scope.variable_scope(model.discriminator_scope, reuse=True): + dis_gen_outputs = model.discriminator_fn(pooled_generated_data, + pooled_generator_inputs) return model._replace(discriminator_gen_outputs=dis_gen_outputs) elif isinstance(model, namedtuples.ACGANModel): - (dis_pooled_gen_outputs, - dis_pooled_gen_classification_logits) = model.discriminator_fn( - pooled_generated_data, pooled_generator_inputs) + with variable_scope.variable_scope(model.discriminator_scope, reuse=True): + (dis_pooled_gen_outputs, + dis_pooled_gen_classification_logits) = model.discriminator_fn( + pooled_generated_data, pooled_generator_inputs) return model._replace( discriminator_gen_outputs=dis_pooled_gen_outputs, discriminator_gen_classification_logits= dis_pooled_gen_classification_logits) elif isinstance(model, namedtuples.InfoGANModel): - (dis_pooled_gen_outputs, - pooled_predicted_distributions) = model.discriminator_and_aux_fn( - pooled_generated_data, pooled_generator_inputs) + with variable_scope.variable_scope(model.discriminator_scope, reuse=True): + (dis_pooled_gen_outputs, + pooled_predicted_distributions) = model.discriminator_and_aux_fn( + pooled_generated_data, pooled_generator_inputs) return model._replace( discriminator_gen_outputs=dis_pooled_gen_outputs, predicted_distributions=pooled_predicted_distributions) @@ -383,6 +460,7 @@ def gan_loss( # Auxiliary losses. gradient_penalty_weight=None, gradient_penalty_epsilon=1e-10, + gradient_penalty_target=1.0, mutual_information_penalty_weight=None, aux_cond_generator_weight=None, aux_cond_discriminator_weight=None, @@ -404,6 +482,9 @@ def gan_loss( small positive value used by the gradient penalty function for numerical stability. Note some applications will need to increase this value to avoid NaNs. + gradient_penalty_target: If `gradient_penalty_weight` is not None, a Python + number or `Tensor` indicating the target value of gradient norm. See the + CIFAR10 section of https://arxiv.org/abs/1710.10196. Defaults to 1.0. mutual_information_penalty_weight: If not `None`, must be a non-negative Python number or Tensor indicating how much to weight the mutual information penalty. See https://arxiv.org/abs/1606.03657 for more @@ -462,7 +543,10 @@ def gan_loss( # Add optional extra losses. if _use_aux_loss(gradient_penalty_weight): gp_loss = tfgan_losses.wasserstein_gradient_penalty( - model, epsilon=gradient_penalty_epsilon, add_summaries=add_summaries) + model, + epsilon=gradient_penalty_epsilon, + target=gradient_penalty_target, + add_summaries=add_summaries) dis_loss += gradient_penalty_weight * gp_loss if _use_aux_loss(mutual_information_penalty_weight): info_loss = tfgan_losses.mutual_information_penalty( @@ -491,6 +575,69 @@ def gan_loss( return namedtuples.GANLoss(gen_loss + gen_reg_loss, dis_loss + dis_reg_loss) +def cyclegan_loss( + model, + # Loss functions. + generator_loss_fn=tfgan_losses.least_squares_generator_loss, + discriminator_loss_fn=tfgan_losses.least_squares_discriminator_loss, + # Auxiliary losses. + cycle_consistency_loss_fn=tfgan_losses.cycle_consistency_loss, + cycle_consistency_loss_weight=10.0, + # Options + **kwargs): + """Returns the losses for a `CycleGANModel`. + + See https://arxiv.org/abs/1703.10593 for more details. + + Args: + model: A `CycleGANModel` namedtuple. + generator_loss_fn: The loss function on the generator. Takes a `GANModel` + named tuple. + discriminator_loss_fn: The loss function on the discriminator. Takes a + `GANModel` namedtuple. + cycle_consistency_loss_fn: The cycle consistency loss function. Takes a + `CycleGANModel` namedtuple. + cycle_consistency_loss_weight: A non-negative Python number or a scalar + `Tensor` indicating how much to weigh the cycle consistency loss. + **kwargs: Keyword args to pass directly to `gan_loss` to construct the loss + for each partial model of `model`. + + Returns: + A `CycleGANLoss` namedtuple. + + Raises: + ValueError: If `model` is not a `CycleGANModel` namedtuple. + """ + # Sanity checks. + if not isinstance(model, namedtuples.CycleGANModel): + raise ValueError( + '`model` must be a `CycleGANModel`. Instead, was %s.' % type(model)) + + # Defines cycle consistency loss. + cycle_consistency_loss = cycle_consistency_loss_fn( + model, add_summaries=kwargs.get('add_summaries', True)) + cycle_consistency_loss_weight = _validate_aux_loss_weight( + cycle_consistency_loss_weight, 'cycle_consistency_loss_weight') + aux_loss = cycle_consistency_loss_weight * cycle_consistency_loss + + # Defines losses for each partial model. + def _partial_loss(partial_model): + partial_loss = gan_loss( + partial_model, + generator_loss_fn=generator_loss_fn, + discriminator_loss_fn=discriminator_loss_fn, + **kwargs) + return partial_loss._replace( + generator_loss=partial_loss.generator_loss + aux_loss) + + with ops.name_scope('cyclegan_loss_x2y'): + loss_x2y = _partial_loss(model.model_x2y) + with ops.name_scope('cyclegan_loss_y2x'): + loss_y2x = _partial_loss(model.model_y2x) + + return namedtuples.CycleGANLoss(loss_x2y, loss_y2x) + + def _get_update_ops(kwargs, gen_scope, dis_scope, check_for_unused_ops=True): """Gets generator and discriminator update ops. @@ -558,6 +705,24 @@ def gan_train_ops( A GANTrainOps tuple of (generator_train_op, discriminator_train_op) that can be used to train a generator/discriminator pair. """ + if isinstance(model, namedtuples.CycleGANModel): + saved_params = locals() + saved_params.pop('model', None) + saved_params.pop('loss', None) + kwargs = saved_params.pop('kwargs', {}) + saved_params.update(kwargs) + with ops.name_scope('cyclegan_x2y_train'): + train_ops_x2y = gan_train_ops(model.model_x2y, loss.loss_x2y, + **saved_params) + with ops.name_scope('cyclegan_y2x_train'): + train_ops_y2x = gan_train_ops(model.model_y2x, loss.loss_y2x, + **saved_params) + return namedtuples.GANTrainOps( + (train_ops_x2y.generator_train_op, train_ops_y2x.generator_train_op), + (train_ops_x2y.discriminator_train_op, + train_ops_y2x.discriminator_train_op), + training_util.get_or_create_global_step().assign_add(1)) + # Create global step increment op. global_step = training_util.get_or_create_global_step() global_step_inc = global_step.assign_add(1) diff --git a/tensorflow/contrib/gan/python/train_test.py b/tensorflow/contrib/gan/python/train_test.py index 519d101e07f4f28d684017b86102fce8fa7677ef..f9bdaa74c948ecee11d5cfd89f06087924f8dace 100644 --- a/tensorflow/contrib/gan/python/train_test.py +++ b/tensorflow/contrib/gan/python/train_test.py @@ -210,12 +210,63 @@ def create_callable_acgan_model(): one_hot_labels=array_ops.one_hot([0, 1, 2], 10)) +def get_cyclegan_model(): + return namedtuples.CycleGANModel( + model_x2y=get_gan_model(), + model_y2x=get_gan_model(), + reconstructed_x=array_ops.ones([1, 2, 3]), + reconstructed_y=array_ops.zeros([1, 2, 3])) + + +def get_callable_cyclegan_model(): + return namedtuples.CycleGANModel( + model_x2y=get_callable_gan_model(), + model_y2x=get_callable_gan_model(), + reconstructed_x=array_ops.ones([1, 2, 3]), + reconstructed_y=array_ops.zeros([1, 2, 3])) + + +def create_cyclegan_model(): + return train.cyclegan_model( + generator_model, + discriminator_model, + data_x=array_ops.zeros([1, 2]), + data_y=array_ops.ones([1, 2])) + + +def create_callable_cyclegan_model(): + return train.cyclegan_model( + Generator(), + Discriminator(), + data_x=array_ops.zeros([1, 2]), + data_y=array_ops.ones([1, 2])) + + def get_sync_optimizer(): return sync_replicas_optimizer.SyncReplicasOptimizer( gradient_descent.GradientDescentOptimizer(learning_rate=1.0), replicas_to_aggregate=1) +def get_tensor_pool_fn(pool_size): + + def tensor_pool_fn_impl(input_values): + return random_tensor_pool.tensor_pool(input_values, pool_size=pool_size) + + return tensor_pool_fn_impl + + +def get_tensor_pool_fn_for_infogan(pool_size): + + def tensor_pool_fn_impl(input_values): + generated_data, generator_inputs = input_values + output_values = random_tensor_pool.tensor_pool( + [generated_data] + generator_inputs, pool_size=pool_size) + return output_values[0], output_values[1:] + + return tensor_pool_fn_impl + + class GANModelTest(test.TestCase): """Tests for `gan_model`.""" @@ -242,6 +293,13 @@ class GANModelTest(test.TestCase): self._test_output_type_helper( get_callable_acgan_model, namedtuples.ACGANModel) + def test_output_type_cyclegan(self): + self._test_output_type_helper(get_cyclegan_model, namedtuples.CycleGANModel) + + def test_output_type_callable_cyclegan(self): + self._test_output_type_helper(get_callable_cyclegan_model, + namedtuples.CycleGANModel) + def test_no_shape_check(self): def dummy_generator_model(_): return (None, None) @@ -289,6 +347,17 @@ class GANLossTest(test.TestCase): def test_output_type_callable_acgan(self): self._test_output_type_helper(get_callable_acgan_model) + def test_output_type_cyclegan(self): + loss = train.cyclegan_loss(create_cyclegan_model(), add_summaries=True) + self.assertIsInstance(loss, namedtuples.CycleGANLoss) + self.assertGreater(len(ops.get_collection(ops.GraphKeys.SUMMARIES)), 0) + + def test_output_type_callable_cyclegan(self): + loss = train.cyclegan_loss( + create_callable_cyclegan_model(), add_summaries=True) + self.assertIsInstance(loss, namedtuples.CycleGANLoss) + self.assertGreater(len(ops.get_collection(ops.GraphKeys.SUMMARIES)), 0) + # Test gradient penalty option. def _test_grad_penalty_helper(self, create_gan_model_fn): model = create_gan_model_fn() @@ -412,24 +481,115 @@ class GANLossTest(test.TestCase): def test_callable_acgan(self): self._test_acgan_helper(create_callable_acgan_model) + # Test that CycleGan models work. + def _test_cyclegan_helper(self, create_gan_model_fn): + model = create_gan_model_fn() + loss = train.cyclegan_loss(model) + self.assertIsInstance(loss, namedtuples.CycleGANLoss) + + # Check values. + with self.test_session(use_gpu=True) as sess: + variables.global_variables_initializer().run() + (loss_x2y_gen_np, loss_x2y_dis_np, loss_y2x_gen_np, + loss_y2x_dis_np) = sess.run([ + loss.loss_x2y.generator_loss, loss.loss_x2y.discriminator_loss, + loss.loss_y2x.generator_loss, loss.loss_y2x.discriminator_loss + ]) + + self.assertGreater(loss_x2y_gen_np, loss_x2y_dis_np) + self.assertGreater(loss_y2x_gen_np, loss_y2x_dis_np) + self.assertTrue(np.isscalar(loss_x2y_gen_np)) + self.assertTrue(np.isscalar(loss_x2y_dis_np)) + self.assertTrue(np.isscalar(loss_y2x_gen_np)) + self.assertTrue(np.isscalar(loss_y2x_dis_np)) + + def test_cyclegan(self): + self._test_cyclegan_helper(create_cyclegan_model) + + def test_callable_cyclegan(self): + self._test_cyclegan_helper(create_callable_cyclegan_model) + + def _check_tensor_pool_adjusted_model_outputs(self, tensor1, tensor2, + pool_size): + history_values = [] + with self.test_session(use_gpu=True) as sess: + variables.global_variables_initializer().run() + for i in range(2 * pool_size): + t1, t2 = sess.run([tensor1, tensor2]) + history_values.append(t1) + if i < pool_size: + # For [0, pool_size), the pool is not full, tensor1 should be equal + # to tensor2 as the pool. + self.assertAllEqual(t1, t2) + else: + # For [pool_size, ?), the pool is full, tensor2 must be equal to some + # historical values of tensor1 (which is previously stored in the + # pool). + self.assertTrue(any([(v == t2).all() for v in history_values])) + + # Test `_tensor_pool_adjusted_model` for gan model. + def test_tensor_pool_adjusted_model_gan(self): + model = create_gan_model() + + new_model = train._tensor_pool_adjusted_model(model, None) + # 'Generator/dummy_g:0' and 'Discriminator/dummy_d:0' + self.assertEqual(2, len(ops.get_collection(ops.GraphKeys.VARIABLES))) + self.assertIs(new_model.discriminator_gen_outputs, + model.discriminator_gen_outputs) + + pool_size = 5 + new_model = train._tensor_pool_adjusted_model( + model, get_tensor_pool_fn(pool_size=pool_size)) + self.assertIsNot(new_model.discriminator_gen_outputs, + model.discriminator_gen_outputs) + # Check values. + self._check_tensor_pool_adjusted_model_outputs( + model.discriminator_gen_outputs, new_model.discriminator_gen_outputs, + pool_size) + + # Test _tensor_pool_adjusted_model for infogan model. + def test_tensor_pool_adjusted_model_infogan(self): + model = create_infogan_model() + + pool_size = 5 + new_model = train._tensor_pool_adjusted_model( + model, get_tensor_pool_fn_for_infogan(pool_size=pool_size)) + # 'Generator/dummy_g:0' and 'Discriminator/dummy_d:0' + self.assertEqual(2, len(ops.get_collection(ops.GraphKeys.VARIABLES))) + self.assertIsNot(new_model.discriminator_gen_outputs, + model.discriminator_gen_outputs) + self.assertIsNot(new_model.predicted_distributions, + model.predicted_distributions) + # Check values. + self._check_tensor_pool_adjusted_model_outputs( + model.discriminator_gen_outputs, new_model.discriminator_gen_outputs, + pool_size) + + # Test _tensor_pool_adjusted_model for acgan model. + def test_tensor_pool_adjusted_model_acgan(self): + model = create_acgan_model() + + pool_size = 5 + new_model = train._tensor_pool_adjusted_model( + model, get_tensor_pool_fn(pool_size=pool_size)) + # 'Generator/dummy_g:0' and 'Discriminator/dummy_d:0' + self.assertEqual(2, len(ops.get_collection(ops.GraphKeys.VARIABLES))) + self.assertIsNot(new_model.discriminator_gen_outputs, + model.discriminator_gen_outputs) + self.assertIsNot(new_model.discriminator_gen_classification_logits, + model.discriminator_gen_classification_logits) + # Check values. + self._check_tensor_pool_adjusted_model_outputs( + model.discriminator_gen_outputs, new_model.discriminator_gen_outputs, + pool_size) + # Test tensor pool. def _test_tensor_pool_helper(self, create_gan_model_fn): model = create_gan_model_fn() if isinstance(model, namedtuples.InfoGANModel): - - def tensor_pool_fn_impl(input_values): - generated_data, generator_inputs = input_values - output_values = random_tensor_pool.tensor_pool( - [generated_data] + generator_inputs, pool_size=5) - return output_values[0], output_values[1:] - - tensor_pool_fn = tensor_pool_fn_impl + tensor_pool_fn = get_tensor_pool_fn_for_infogan(pool_size=5) else: - - def tensor_pool_fn_impl(input_values): - return random_tensor_pool.tensor_pool(input_values, pool_size=5) - - tensor_pool_fn = tensor_pool_fn_impl + tensor_pool_fn = get_tensor_pool_fn(pool_size=5) loss = train.gan_loss(model, tensor_pool_fn=tensor_pool_fn) self.assertTrue(isinstance(loss, namedtuples.GANLoss)) diff --git a/tensorflow/contrib/gdr/BUILD b/tensorflow/contrib/gdr/BUILD index bdbe6f0a72621e59562fe113da101ff5a2b8c06d..707ae25d485c64f15694ee0e357f32b619d3cd33 100644 --- a/tensorflow/contrib/gdr/BUILD +++ b/tensorflow/contrib/gdr/BUILD @@ -82,6 +82,7 @@ tf_cuda_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core/distributed_runtime:graph_mgr", + "//tensorflow/core/distributed_runtime:recent_request_ids", "//tensorflow/core/distributed_runtime:rendezvous_mgr_interface", "//tensorflow/core/distributed_runtime:worker", "//tensorflow/core/distributed_runtime:worker_cache", @@ -103,6 +104,7 @@ cc_library( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core/distributed_runtime:base_rendezvous_mgr", + "//tensorflow/core/distributed_runtime:request_id", "//tensorflow/core/distributed_runtime:tensor_coding", "//tensorflow/core/distributed_runtime:worker_cache", "//tensorflow/core/distributed_runtime:worker_env", diff --git a/tensorflow/contrib/gdr/README.md b/tensorflow/contrib/gdr/README.md index 34ce60b360822888aa6223c89362ae1b0d9d991f..8242d93f129904828a11b61d48f2df8fb0f88bc3 100644 --- a/tensorflow/contrib/gdr/README.md +++ b/tensorflow/contrib/gdr/README.md @@ -119,4 +119,4 @@ In the original design (as in the reference), tensor buffers are only registered Reference === -Bairen Yi, Jiacheng Xia, Li Chen, and Kai Chen. 2017. Towards Zero Copy Dataflows using RDMA. In Proceedings of SIGCOMM Posters and Demos'17, Los Angeles, CA, USA, August 22-24, 2017, 3 pages. https://doi.org/10.1145/3123878.3123907 +Bairen Yi, Jiacheng Xia, Li Chen, and Kai Chen. 2017. Towards Zero Copy Dataflows using RDMA. In Proceedings of SIGCOMM Posters and Demos'17, Los Angeles, CA, USA, August 22-24, 2017, 3 pages. https://doi.org/10.1145/3123878.3131975 diff --git a/tensorflow/contrib/gdr/gdr_memory_manager.cc b/tensorflow/contrib/gdr/gdr_memory_manager.cc index 5c7ac744289ab7729b4cc43ab9bedc9342284e65..81e70ae30a4c72dbcedd1aabfe758ecca4c8b366 100644 --- a/tensorflow/contrib/gdr/gdr_memory_manager.cc +++ b/tensorflow/contrib/gdr/gdr_memory_manager.cc @@ -86,8 +86,9 @@ int TryToReadNumaNode(ibv_device* device) { if (strings::safe_strto32(content, &value)) { if (value < 0) { LOG(INFO) << "Successful NUMA node read from SysFS had negative value (" - << value << "), but there must be at least one NUMA node" - ", so returning NUMA node zero"; + << value + << "), but there must be at least one NUMA node" + ", so returning NUMA node zero"; return 0; } LOG(INFO) << "NUMA node for device: " << device->name << " is " << value; @@ -290,8 +291,8 @@ Status GdrMemoryManager::Init() { // Host memory allocators for (Allocator* allocator : allocators) { auto* visitable_allocator = dynamic_cast(allocator); - CHECK(visitable_allocator) << "is not visitable for instrumentation" - << allocator->Name(); + CHECK(visitable_allocator) + << "is not visitable for instrumentation" << allocator->Name(); // Make sure we don't instrument the same allocator twice if (instrumented_.find(allocator) == std::end(instrumented_)) { visitable_allocator->AddAllocVisitor(alloc_visitor); @@ -635,8 +636,8 @@ void GdrMemoryManager::TensorFromTransportOptions( } else { checksum = GPUUtil::Checksum(*tensor); } - CHECK(checksum == remote_mr.checksum()) << "Checksum mismatch: " << checksum - << "!=" << remote_mr.checksum(); + CHECK(checksum == remote_mr.checksum()) + << "Checksum mismatch: " << checksum << "!=" << remote_mr.checksum(); #endif } done(Status::OK()); diff --git a/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc b/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc index adef2aac33e3e0839a268eabe2496e58861535c5..28f68cec8cce126f1b177a73e197ccd7ab749f4a 100644 --- a/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc +++ b/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/process_util.h" +#include "tensorflow/core/distributed_runtime/request_id.h" #include "tensorflow/core/distributed_runtime/tensor_coding.h" #include "tensorflow/core/distributed_runtime/worker_cache.h" #include "tensorflow/core/distributed_runtime/worker_interface.h" @@ -47,6 +48,7 @@ class GdrRecvTensorCall : public BaseRecvTensorCall { recv_args_(recv_args) { req_.set_step_id(step_id); req_.set_rendezvous_key(key.data(), key.size()); + req_.set_request_id(GetUniqueRequestId()); } ~GdrRecvTensorCall() override {} diff --git a/tensorflow/contrib/gdr/gdr_worker.cc b/tensorflow/contrib/gdr/gdr_worker.cc index 568641234731a458a05886d12066ee9f55fa58aa..ce1d8d2d73000559f03046aceacb169890ecc1b6 100644 --- a/tensorflow/contrib/gdr/gdr_worker.cc +++ b/tensorflow/contrib/gdr/gdr_worker.cc @@ -41,17 +41,26 @@ namespace tensorflow { GdrWorker::GdrWorker(WorkerEnv* worker_env, RemoteMemoryManager* remote_memory_manager) - : GrpcWorker(worker_env), remote_memory_manager_(remote_memory_manager) {} + : GrpcWorker(worker_env), + remote_memory_manager_(remote_memory_manager), + recv_tensor_recent_request_ids_(100000) {} void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, const RecvTensorRequest* request, ::grpc::ByteBuffer* response, StatusCallback done) { + Status s = recv_tensor_recent_request_ids_.TrackUnique( + request->request_id(), "RecvTensor (GdrWorker)", *request); + if (!s.ok()) { + done(s); + return; + } + const int64 step_id = request->step_id(); const string& key = request->rendezvous_key(); TRACEPRINTF("RecvTensor: %lld %s", step_id, key.c_str()); Rendezvous::ParsedKey parsed; - Status s = Rendezvous::ParseKey(key, &parsed); + s = Rendezvous::ParseKey(key, &parsed); Device* src_dev = nullptr; if (s.ok()) { s = PrepareRecvTensor(parsed, &src_dev); diff --git a/tensorflow/contrib/gdr/gdr_worker.h b/tensorflow/contrib/gdr/gdr_worker.h index a30b7baaedcbc80d93d7f37756732c37d2435935..54081f655ec087d78ac07974656257dcf478bcef 100644 --- a/tensorflow/contrib/gdr/gdr_worker.h +++ b/tensorflow/contrib/gdr/gdr_worker.h @@ -18,6 +18,7 @@ limitations under the License. #include "tensorflow/contrib/gdr/gdr_memory_manager.h" +#include "tensorflow/core/distributed_runtime/recent_request_ids.h" #include "tensorflow/core/distributed_runtime/rpc/grpc_worker_service.h" namespace tensorflow { @@ -38,6 +39,7 @@ class GdrWorker : public GrpcWorker { private: RemoteMemoryManager* remote_memory_manager_; // Not owned + RecentRequestIds recv_tensor_recent_request_ids_; }; } // namespace tensorflow diff --git a/tensorflow/contrib/hvx/README.md b/tensorflow/contrib/hvx/README.md index 5a6f2f3086d708e5264b0483c211902ac8dce5f6..163993a3f6bb1bedcdffb32944a98c7cc846878e 100644 --- a/tensorflow/contrib/hvx/README.md +++ b/tensorflow/contrib/hvx/README.md @@ -1,60 +1,67 @@ # TensorFlow Runtime with HVX Acceleration -## Description +This README explain how to build and use the TensorFlow runtime with HVX Acceleration. HVX is an extension of Hexagon, a DSP provided by Qualcomm, which can compute vector calculations faster using less energy than ARM processors. -This README explain how to build and use the TensorFlow Runtime with HVX Acceleration. HVX is an extension of Hexagon which is a DSP provided by qualcomm which can compute vector calculations faster using lower energy than ARM processors. +## Dependencies + +* [Android SDK](https://developer.android.com/studio/index.html). +* [Android NDK](https://developer.android.com/ndk/index.html). Save the path in `${NDK_ROOT}`. +* A rooted Qualcomm-based Android device connected to the computer (preferably, a [Snapdragon Development Board](https://developer.qualcomm.com/hardware/additional-snapdragon), but it could be a rooted phone with a Qualcomm SoC, albeit this guide may not work with it). The device needs to be rooted for development and testing purposes, and shouldn't be needed in production. See [Behold, The Snapdragon MDP](https://developer.qualcomm.com/blog/behold-snapdragon-mdp) for more information. +* [Hexagon SDK v3.0](https://developer.qualcomm.com/software/hexagon-dsp-sdk/tools). Save the path in `${QUALCOMM_SDK}`. +* The current directory should be TensorFlow source code (`git clone https://github.com/tensorflow/tensorflow.git && cd tensorflow`), and saved into `${TF_ROOT_DIR}`. + +You may also need to add a test signature in the device to run HVX-based binaries. Follow the instructions in `${QUALCOMM_SDK}/docs/Tools_Signing.html`, using Python 2. + +Note that if the device is not rooted, you may not be able to get the serial number, push the test signature and/or run binary files that call HVX libraries. ## Quick Start Guide -We provides several tools to build and run inference with this runtime quickly. +We provide several tools to build and run inference with this runtime quickly. -#### All-in-one script to run inception model with prebuild hexagon library -If you don’t need to build your own implementation of hexagon HVX, we provide a shortcut to execute graphs by using pre-compiled binaries. +### Run inception model with a prebuilt Hexagon library +If you don’t need to build your own implementation of Hexagon HVX, we provide a shortcut to execute graphs by using pre-compiled binaries. + +```shell +./tensorflow/contrib/makefile/samples/build_and_run_inception_hexagon.sh -p ``` -git clone https://github.com/tensorflow/tensorflow.git -cd tensorflow -NDK_ROOT="/path/to/ndk" ./tensorflow/contrib/makefile/build_all_android.sh -X -``` -(-X downloads dependencies to hexagon HVX and graphs, and copy all dependencies to android and execute a test) -#### All-in-one script to run inception model by building entire libraries from source code - If you want to build your own implementation of hexagon HVX, we provide a sample all-in-one script to execute graphs which downloads source and build everything for hexagon. +The `-p` option makes the script download dependencies (i.e., Hexagon HVX binaries and graphs models), copy them to the Android device and execute a test. -``` -git clone https://github.com/tensorflow/tensorflow.git -cd tensorflow -QUALCOMM_SDK="/path/to/qualcomm/sdk" NDK_ROOT="/path/to/ndk" ./tensorflow/contrib/makefile/samples/build_and_run_inception_hexagon.sh +### Run inception model by building all from the source code + +If you want to build your own implementation of Hexagon HVX, we provide a sample all-in-one script to execute graphs which downloads the source and builds everything that's necessary. + +```shell +./tensorflow/contrib/makefile/samples/build_and_run_inception_hexagon.sh ``` ## Building libraries If you've finished walking through the quick start guide, you may want to try building each binary manually. -#### Build libhexagon_nn_skel.so -Download hexagon nn library from codeaurora.org and build it. +### Build libhexagon\_nn\_skel.so -``` +Download Hexagon NN library from codeaurora.org and build it. + +```shell git clone https://source.codeaurora.org/quic/hexagon_nn/nnlib cd nnlib ``` -(Just follow instructions in README.HOW_TO_BUILD. You can find libhexagon_nn_skel.so in hexagon_Release_dynamic_toolv72_v60/ship) -Then copy the generated binary to GEN_LIBS_DIR +Just follow the instructions in `README.HOW_TO_BUILD`. You can find the file `libhexagon_nn_skel.so` in `hexagon_Release_dynamic_toolv72_v60/ship`. +Then copy the generated binary to `${GEN_LIBS_DIR}`. -``` +```shell GEN_LIBS_DIR="/path/to/a/dir/to/store/hexagon/libraries" cp -v "hexagon_Release_dynamic_toolv72_v60/ship/libhexagon_nn_skel.so" "${GEN_LIBS_DIR}" ``` -#### Build libhexagon_controller.so +### Build libhexagon\_controller.so + Download tensorflow and build hexagon controller. -``` -git clone https://github.com/tensorflow/tensorflow.git -cd tensorflow -TF_ROOT_DIR="$(pwd)" -QUALCOMM_SDK="/path/to/qualcomm/sdk" +```shell GENERATED_NNLIB_DIRECTORY="/path/to/nnlib" GENERATED_HEXAGON_CONTROLLER_DIRECTORY="${QUALCOMM_SDK}/examples/common/generated_hexagon_controller" rm -rf "${GENERATED_HEXAGON_CONTROLLER_DIRECTORY}" @@ -70,12 +77,12 @@ make tree VERBOSE=1 V=android_Release cp -v "${GENERATED_HEXAGON_CONTROLLER_DIRECTORY}/android_Release/ship/libhexagon_controller.so" "${GEN_LIBS_DIR}" ``` -#### Build tensorflow linking hexagon library -Build tensorflow with the build_all_android.sh with specifying -x option. +### Build TensorFlow linking Hexagon library -``` +Build TensorFlow with `build_all_android.sh` specifying the `-x` option. + +```shell BUILD_ALL_ANDROID_PATH="${TF_ROOT_DIR}/tensorflow/contrib/makefile/build_all_android.sh" -NDK_ROOT="/path/to/ndk/root" CC_PREFIX=${CC_PREFIX} NDK_ROOT=${NDK_ROOT} "${BUILD_ALL_ANDROID_PATH}" \ -x "${GEN_LIBS_DIR}" \ @@ -83,11 +90,11 @@ CC_PREFIX=${CC_PREFIX} NDK_ROOT=${NDK_ROOT} "${BUILD_ALL_ANDROID_PATH}" \ -t hexagon_graph_execution ``` -#### Push binaries to your Android device +### Push binaries to your Android device Before running tests on your Android device, you need to push several binaries to it. -``` +```shell adb push "${GEN_LIBS_DIR}/libhexagon_controller.so" "/data/local/tmp" adb push "${GEN_LIBS_DIR}/libhexagon_nn_skel.so" "/vendor/lib/rfsa/adsp" adb push -p \ @@ -100,40 +107,54 @@ adb shell chmod "${ANDROID_EXEC_FILE_MODE}" \ adb wait-for-device ``` -#### Run tests on the device +### Run tests on the device Finally, you can run the inference tests on your device. -``` +```shell adb shell 'LD_LIBRARY_PATH=/data/local/tmp:$LD_LIBRARY_PATH' \ "/data/local/tmp/hexagon_graph_execution" ``` -#### Troubleshooting -If you're using the Open-Q 820 Snapdragon development kit, you may run into an issue with running the executable due to a missing testsig library. From the Hexagon SDK documentation: *Dynamic shared objects are required to be digitally signed and then authenticated at runtime before they are allowed to be loaded and executed.* Generating a testsig library is necessary to run the unsigned sample library built from this project. +### Troubleshooting + +#### Testsig issue + +If you're using the Open-Q 820 Snapdragon Development Kit, you may run into an issue with running the executable due to a missing `testsig` library. From the Hexagon SDK documentation: *Dynamic shared objects are required to be digitally signed and then authenticated at runtime before they are allowed to be loaded and executed.* Generating a testsig library is necessary to run the unsigned sample library built from this project. -If the lack of a testsig library is your problem, you will see errors of the type: +If the lack of a `testsig` library is your problem, you will see errors of the type: `vendor/qcom/proprietary/adsprpc/src/fastrpc_apps_user.c:169::error: -1: 0 == (nErr = remotectl_open(name, (int*)ph, dlerrstr, sizeof(dlerrstr), &dlerr))` -appearing in adb logcat. - -There are several ways to create the testsig library, the only prerequisite is Python and the correct version of the Hexagon-SDK. The following steps is one way to create this library: -1. Run adb as root: `adb root` -2. Run the command `adb shell cat /sys/devices/soc0/serial_number` -3. Convert the decimal number you get as output to hex -4. Run the python script: `python ${QUALCOMM_SDK}/tools/elfsigner/elfsigner.py -t $(SERIAL_NUMBER_HEX_VALUE)` -5. The output of the python script is a shared library stored in ${QUALCOMM_SDK}/tools/elfsigner/output/testsig-$(SERIAL_NUMBER_HEX_VALUE).so -6. Push the shared library to your device: +appearing in `adb logcat` or ["Expected: (version) >= (1), actual: 0 vs 1" while running a binary from adb](https://github.com/tensorflow/tensorflow/issues/11210). + +You need to add a test signature, as described at the beginning of this README. After rebooting your device, you should be able to run the sample application. + +#### Qualcomm SDK Linux installation fails with "Malformed \uxxxx encoding" + +The installation file is based on LaunchAnywhere, which fails in Linux if the `PS1` env variable contains non-common Unicode chars: + ``` -adb root -adb wait-for-device -adb remount -adb wait-for-device -adb shell mkdir /system/lib/rfsa -adb shell mkdir /system/lib/rfsa/adsp -adb push ${QUALCOMM_SDK}/tools/elfsigner/output/testsig-$(SERIAL_NUMBER_HEX_VALUE).so /system/lib/rfsa/adsp/ +Preparing to install... +Extracting the JRE from the installer archive... +Unpacking the JRE... +Extracting the installation resources from the installer archive... +Configuring the installer for this system's environment... + +Launching installer... + +An internal LaunchAnywhere application error has occurred and this application cannot proceed. (LAX) + +Stack Trace: +java.lang.IllegalArgumentException: Malformed \uxxxx encoding. + at java.util.Properties.loadConvert(Properties.java:574) + at java.util.Properties.load0(Properties.java:391) + at java.util.Properties.load(Properties.java:317) + at com.zerog.common.java.util.PropertiesUtil.loadProperties(Unknown Source) + at com.zerog.lax.LAX.(Unknown Source) + at com.zerog.lax.LAX.main(Unknown Source) ``` -After rebooting your device, you should be able to run the sample application. +It can be solved by temporarily assigning the `PS1` environment variable to something simple, such as '$'. + +## Maintainers -Maintainers: -- Satoshi Kataoka (satok@google.com, github.com/satok16) +* Satoshi Kataoka (satok@google.com, github.com/satok16) diff --git a/tensorflow/contrib/image/BUILD b/tensorflow/contrib/image/BUILD index 54502cfc6eecb9d064ffde9773e97d893a24133a..3ff02e085ee63fabf42b3cc4389f4605455f3800 100755 --- a/tensorflow/contrib/image/BUILD +++ b/tensorflow/contrib/image/BUILD @@ -14,6 +14,7 @@ load( "tf_gen_op_libs", "tf_gen_op_wrapper_py", "tf_kernel_library", + "tf_py_test", ) load("//tensorflow:tensorflow.bzl", "cuda_py_test") load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") @@ -24,6 +25,8 @@ tf_custom_op_library( "kernels/bipartite_match_op.cc", "kernels/image_ops.cc", "kernels/image_ops.h", + "kernels/segmentation_ops.cc", + "kernels/segmentation_ops.h", "ops/image_ops.cc", ], gpu_srcs = [ @@ -38,6 +41,8 @@ tf_kernel_library( "kernels/bipartite_match_op.cc", "kernels/image_ops.cc", "kernels/image_ops.h", + "kernels/segmentation_ops.cc", + "kernels/segmentation_ops.h", ], gpu_srcs = [ "kernels/image_ops_gpu.cu.cc", @@ -78,6 +83,7 @@ tf_custom_op_py_library( "//tensorflow/python:array_ops", "//tensorflow/python:common_shapes", "//tensorflow/python:constant_op", + "//tensorflow/python:control_flow_ops", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:linalg_ops", "//tensorflow/python:math_ops", @@ -188,6 +194,21 @@ cuda_py_test( ], ) +tf_py_test( + name = "segmentation_test", + size = "medium", + srcs = ["python/kernel_tests/segmentation_test.py"], + additional_deps = [ + ":distort_image_py", + ":image_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + ], +) + tf_custom_op_library( name = "python/ops/_single_image_random_dot_stereograms.so", srcs = [ @@ -233,6 +254,23 @@ py_library( ], ) +cuda_py_test( + name = "single_image_random_dot_stereograms_ops_test", + size = "medium", + srcs = ["python/kernel_tests/single_image_random_dot_stereograms_ops_test.py"], + additional_deps = [ + ":distort_image_py", + ":image_py", + ":single_image_random_dot_stereograms_py", + "//third_party/py/numpy", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + ], +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/contrib/image/__init__.py b/tensorflow/contrib/image/__init__.py index d030dffadeb9d67f7ffcbc197a2a3feb9b3b122d..cc8ed117ba2edcc7a53e609381166f17a2fbb45e 100755 --- a/tensorflow/contrib/image/__init__.py +++ b/tensorflow/contrib/image/__init__.py @@ -20,6 +20,8 @@ This module provides functions for image manipulation; currently, chrominance transformas (including changing saturation and hue) in YIQ space and projective transforms (including rotation) are supported. +## Image Transformation `Ops` + @@angles_to_projective_transforms @@compose_transforms @@adjust_yiq_hsv @@ -28,19 +30,29 @@ projective transforms (including rotation) are supported. @@transform @@translate @@translations_to_projective_transforms + +## Image Segmentation `Ops` + +@@connected_components + +## Matching `Ops` + @@bipartite_match + +## Random Dot Stereogram `Ops` + @@single_image_random_dot_stereograms """ from __future__ import absolute_import from __future__ import division from __future__ import print_function -# pylint: disable=line-too-long from tensorflow.contrib.image.python.ops.distort_image_ops import adjust_hsv_in_yiq from tensorflow.contrib.image.python.ops.distort_image_ops import random_hsv_in_yiq from tensorflow.contrib.image.python.ops.image_ops import angles_to_projective_transforms from tensorflow.contrib.image.python.ops.image_ops import compose_transforms +from tensorflow.contrib.image.python.ops.image_ops import connected_components from tensorflow.contrib.image.python.ops.image_ops import rotate from tensorflow.contrib.image.python.ops.image_ops import transform from tensorflow.contrib.image.python.ops.image_ops import translate diff --git a/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.h b/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.h index 194ae2ba47456cac66c01989a78ab4ce607d1295..8968da6d8241ca7cd548910a024a618913c3ed70 100644 --- a/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.h +++ b/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.h @@ -11,8 +11,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_IMAGE_KERNELS_ADJUST_HSV_IN_YIQ_OP_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_IMAGE_KERNELS_ADJUST_HSV_IN_YIQ_OP_H_ +#ifndef TENSORFLOW_CONTRIB_IMAGE_KERNELS_ADJUST_HSV_IN_YIQ_OP_H_ +#define TENSORFLOW_CONTRIB_IMAGE_KERNELS_ADJUST_HSV_IN_YIQ_OP_H_ #if GOOGLE_CUDA #define EIGEN_USE_GPU @@ -84,4 +84,4 @@ struct AdjustHsvInYiqGPU { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_IMAGE_KERNELS_ADJUST_HSV_IN_YIQ_OP_H_ +#endif // TENSORFLOW_CONTRIB_IMAGE_KERNELS_ADJUST_HSV_IN_YIQ_OP_H_ diff --git a/tensorflow/contrib/image/kernels/bipartite_match_op.cc b/tensorflow/contrib/image/kernels/bipartite_match_op.cc index 7d207c388b159c4ad0f25032811e97b153fd50d6..726adb07775e3243fdc96a7f1a00dbb0304d3dd9 100644 --- a/tensorflow/contrib/image/kernels/bipartite_match_op.cc +++ b/tensorflow/contrib/image/kernels/bipartite_match_op.cc @@ -85,7 +85,7 @@ class BipartiteMatchOp : public OpKernel { context->allocate_output(1, TensorShape({num_input_columns}), &column_to_row_match_indices)); - typename TTypes::ConstTensor distance_mat = + TTypes::ConstTensor distance_mat = input_distance_mat.shaped( {num_input_rows, num_input_columns}); diff --git a/tensorflow/contrib/image/kernels/image_ops.cc b/tensorflow/contrib/image/kernels/image_ops.cc index 6adf837ca0ab506bd18f5e2e1fc1847e31d782bf..c2e32da133b32c8fe169302668031af8bace2c22 100644 --- a/tensorflow/contrib/image/kernels/image_ops.cc +++ b/tensorflow/contrib/image/kernels/image_ops.cc @@ -43,9 +43,9 @@ template struct FillProjectiveTransform; typedef Eigen::ThreadPoolDevice CPUDevice; using functor::FillProjectiveTransform; +using generator::Interpolation; using generator::INTERPOLATION_BILINEAR; using generator::INTERPOLATION_NEAREST; -using generator::Interpolation; using generator::ProjectiveGenerator; template @@ -72,11 +72,12 @@ class ImageProjectiveTransform : public OpKernel { const Tensor& transform_t = ctx->input(1); OP_REQUIRES(ctx, images_t.shape().dims() == 4, errors::InvalidArgument("Input images must have rank 4")); - OP_REQUIRES(ctx, (TensorShapeUtils::IsMatrix(transform_t.shape()) && - (transform_t.dim_size(0) == images_t.dim_size(0) || - transform_t.dim_size(0) == 1) && - transform_t.dim_size(1) == - ProjectiveGenerator::kNumParameters), + OP_REQUIRES(ctx, + (TensorShapeUtils::IsMatrix(transform_t.shape()) && + (transform_t.dim_size(0) == images_t.dim_size(0) || + transform_t.dim_size(0) == 1) && + transform_t.dim_size(1) == + ProjectiveGenerator::kNumParameters), errors::InvalidArgument( "Input transform should be num_images x 8 or 1 x 8")); auto images = images_t.tensor(); diff --git a/tensorflow/contrib/image/kernels/segmentation_ops.cc b/tensorflow/contrib/image/kernels/segmentation_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..fe8bf6e21c7b7310527668324571774e8bc50893 --- /dev/null +++ b/tensorflow/contrib/image/kernels/segmentation_ops.cc @@ -0,0 +1,139 @@ +/* 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. +==============================================================================*/ + +// See docs for ImageConnectedComponents in ../ops/image_ops.cc, and description +// of the algorithm in segmentation_ops.h. + +#define EIGEN_USE_THREADS + +#include "tensorflow/contrib/image/kernels/segmentation_ops.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/register_types.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { + +using tensorflow::functor::BlockedImageUnionFindFunctor; +using tensorflow::functor::FindRootFunctor; +using tensorflow::functor::ImageConnectedComponentsFunctor; +using tensorflow::functor::TensorRangeFunctor; + +using OutputType = typename BlockedImageUnionFindFunctor::OutputType; + +// Computes connected components on batches of 2D images. +template +class ImageConnectedComponents : public OpKernel { + public: + explicit ImageConnectedComponents(OpKernelConstruction* ctx) + : OpKernel(ctx) {} + + void Compute(OpKernelContext* ctx) override { + const Tensor& images_t = ctx->input(0); + OP_REQUIRES(ctx, images_t.shape().dims() == 3, + errors::InvalidArgument("Input images must have rank 3")); + Tensor forest_t, rank_t; + OP_REQUIRES_OK(ctx, ctx->allocate_temp(tensorflow::DT_INT64, + images_t.shape(), &forest_t)); + OP_REQUIRES_OK(ctx, ctx->allocate_temp(tensorflow::DT_INT64, + images_t.shape(), &rank_t)); + Tensor* output_t; + OP_REQUIRES_OK(ctx, ctx->allocate_output(0, images_t.shape(), &output_t)); + + // Fill forest with values from 0 to n - 1, so that each node points to + // itself. + TensorRangeFunctor()(ctx->eigen_device(), + forest_t.flat()); + auto rank = rank_t.tensor(); + rank.device(ctx->eigen_device()) = rank.constant(OutputType(0)); + + const auto images = images_t.tensor(); + auto forest = forest_t.tensor(); + ImageConnectedComponentsFunctor()( + ctx, output_t->flat(), images, forest, rank); + } +}; + +using CPUDevice = Eigen::ThreadPoolDevice; + +namespace functor { + +// Connected components CPU implementation. See `segmentation_ops.h` for a +// description of the algorithm. +template +struct ImageConnectedComponentsFunctor { + void operator()(OpKernelContext* ctx, + typename TTypes::Flat output, + typename TTypes::ConstTensor images, + typename TTypes::Tensor forest, + typename TTypes::Tensor rank) { + const int64 num_images = images.dimension(0), + num_rows = images.dimension(1), num_cols = images.dimension(2), + num_elements = images.size(); + // Bail out early for an empty image--no work to do. + if (num_elements == 0) { + return; + } + auto worker_threads = ctx->device()->tensorflow_cpu_worker_threads(); + BlockedImageUnionFindFunctor union_find( + images.data(), num_rows, num_cols, forest.data(), rank.data()); + while (union_find.can_merge()) { + union_find.merge_blocks(); + int64 num_blocks_vertically = union_find.num_blocks_vertically(); + int64 num_blocks_horizontally = union_find.num_blocks_horizontally(); + // Merging each block calls union_down for each pixel in a row of the + // block, and union_right for each pixel in a column of the block. Assume + // 20 instructions for each call to union_down or union_right. find() may + // loop more while searching for the root, but this should not be very + // significant. + int cost = (union_find.block_height() + union_find.block_width()) * 20; + Shard(worker_threads->num_threads, worker_threads->workers, + num_images * num_blocks_vertically * num_blocks_horizontally, cost, + [&union_find, num_images, num_blocks_vertically, + num_blocks_horizontally](int64 start_block, int64 limit_block) { + for (int64 i = start_block; i < limit_block; i++) { + int64 block_x = i % num_blocks_horizontally; + int64 block_y = + (i / num_blocks_horizontally) % num_blocks_vertically; + int64 image = + i / (num_blocks_horizontally * num_blocks_vertically); + union_find.merge_internal_block_edges(image, block_y, block_x); + } + }); + } + FindRootFunctor()(ctx->eigen_device(), output, + images.data(), union_find); + } +}; + +} // end namespace functor + +#define REGISTER_IMAGE_CONNECTED_COMPONENTS(TYPE) \ + REGISTER_KERNEL_BUILDER(Name("ImageConnectedComponents") \ + .Device(DEVICE_CPU) \ + .TypeConstraint("dtype"), \ + ImageConnectedComponents) +// Connected components (arguably) make sense for number, bool, and string types +TF_CALL_NUMBER_TYPES(REGISTER_IMAGE_CONNECTED_COMPONENTS); +TF_CALL_bool(REGISTER_IMAGE_CONNECTED_COMPONENTS); +TF_CALL_string(REGISTER_IMAGE_CONNECTED_COMPONENTS); +#undef REGISTER_IMAGE_CONNECTED_COMPONENTS + +// TODO(ringwalt): Implement on GPU. We probably want to stick to the original +// algorithm by Stava and Benes there for efficiency (computing small blocks in +// shared memory in CUDA thread blocks, instead of starting with single-pixel +// blocks). + +} // end namespace tensorflow diff --git a/tensorflow/contrib/image/kernels/segmentation_ops.h b/tensorflow/contrib/image/kernels/segmentation_ops.h new file mode 100644 index 0000000000000000000000000000000000000000..0957d5fd10f02daad3d8d51aadec9ce9da2660b5 --- /dev/null +++ b/tensorflow/contrib/image/kernels/segmentation_ops.h @@ -0,0 +1,303 @@ +/* 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_IMAGE_KERNELS_SEGMENTATION_OPS_H_ +#define TENSORFLOW_CONTRIB_IMAGE_KERNELS_SEGMENTATION_OPS_H_ + +// Connected component analysis. The op is described in ../ops/image_ops.cc. A +// description of the algorithm appears below. + +#define EIGEN_USE_THREADS + +#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor_types.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/work_sharder.h" + +namespace tensorflow { + +namespace functor { + +template +bool is_nonzero(T value) { + return value != T(0); +} + +template <> +bool is_nonzero(string value) { + return value.size() != 0; +} + +// Processes each pixel of an image for union-find, in parallel blocks. This is +// loosely based on the algorithm in "GPU Computing Gems" by Ondrej Stava and +// Bedrich Benes, available here: +// http://hpcg.purdue.edu/bbenes/papers/Stava2011CCL.pdf +// The bulk of the process uses blocks of each image, which have each been +// processed separately. As long as there are multiple blocks in the image, we +// double the height and width of the blocks, creating new blocks which each +// consist of 2x2 previous sub-blocks. On each new block, we process adjacent +// pixels from the previous sub-blocks serially. However, the new blocks are not +// connected, so we can process each block in parallel. +// The GPU algorithm first processes blocks of a fixed size in GPU shared +// memory, with one image block per CUDA thread block. On the CPU, we just start +// with a block size of a single pixel, and borrow the rest of the algorithm +// unchanged. +template +class BlockedImageUnionFindFunctor { + public: + using OutputType = int64; + + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE BlockedImageUnionFindFunctor( + const T* images, const int64 num_rows, const int64 num_cols, + OutputType* forest, OutputType* rank) + : images_(images), + num_rows_(num_rows), + num_cols_(num_cols), + block_height_(1), + block_width_(1), + forest_(forest), + rank_(rank) {} + + // Returns the root of the tree that the pixel at the given index belongs to. + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE OutputType + find(OutputType index) const { + while (forest_[index] != index) { + index = forest_[index]; + } + return index; + } + + // Returns the number of blocks along the y axis. + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE int64 num_blocks_vertically() const { + return (num_rows_ + block_height_ - 1) / block_height_; + } + + // Returns the number of blocks along the x axis. + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE int64 num_blocks_horizontally() const { + return (num_cols_ + block_width_ - 1) / block_width_; + } + + // Returns the total number of blocks in each image. + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE int64 num_blocks() const { + return num_blocks_vertically() * num_blocks_horizontally(); + } + + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE int64 block_height() const { + return block_height_; + } + + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE int64 block_width() const { + return block_width_; + } + + // Returns whether we may merge again (the image contains more than one + // block). + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE bool can_merge() const { + return block_height_ < num_rows_ || block_width_ < num_cols_; + } + + // Doubles the block size. After this method, you must call + // `merge_internal_block_edges` for each image and each *new* block's xy + // coordinates (typically in parallel). + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void merge_blocks() { + block_height_ *= 2; + block_width_ *= 2; + } + + // Processes pairs of pixels within the block which were adjacent in the four + // sub-blocks. This must be done at each stage so that the connected + // components in each block are joined correctly. + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void merge_internal_block_edges( + int64 image_index, int64 block_vertical_index, + int64 block_horizontal_index) const { + int64 block_start_y = block_vertical_index * block_height_; + int64 block_start_x = block_horizontal_index * block_width_; + // Merge the 4 sub-blocks horizontally (fixing the vertical seam). + int64 block_center_x = block_start_x + block_width_ / 2 - 1; + if (0 <= block_center_x && block_center_x + 1 < num_cols_) { + int64 merge_blocks_limit_y = + std::min(num_rows_, block_start_y + block_height_); + for (int64 y = block_start_y; y < merge_blocks_limit_y; y++) { + union_right(image_index, y, block_center_x); + } + } + // Merge the 4 sub-blocks vertically (fixing the horizontal seam). + int64 block_center_y = block_start_y + block_height_ / 2 - 1; + if (0 <= block_center_y && block_center_y + 1 < num_rows_) { + int64 merge_blocks_limit_x = + std::min(num_cols_, block_start_x + block_width_); + for (int64 x = block_start_x; x < merge_blocks_limit_x; x++) { + union_down(image_index, block_center_y, x); + } + } + } + + private: + // The input image(s). + const T* const images_; + const int64 num_rows_; + const int64 num_cols_; + // Current height of each sub-block of the image. + int64 block_height_; + // Current width of each sub-block of the image. + int64 block_width_; + // Union-find forest. This has the same size as `images_`, and each entry + // holds the index of its parent in `images_` (roots hold their own index). + // Cycles should not occur. + OutputType* const forest_; + // Union-find rank of each pixel. + OutputType* const rank_; + + // Unions the pixel with the pixel below it if applicable (both pixels are + // true, and the pixel is not in the last row). + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void union_down(OutputType batch, + OutputType row, + OutputType col) const { + T pixel = read_pixel(batch, row, col); + if (is_nonzero(pixel)) { + const int64 index_a = col + num_cols_ * (row + num_rows_ * batch); + if (row + 1 < num_rows_ && read_pixel(batch, row + 1, col) == pixel) { + const int64 index_b = col + num_cols_ * (row + 1 + num_rows_ * batch); + do_union(index_a, index_b); + } + } + } + + // Unions the pixel with the pixel to the right of it if applicable. + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void union_right(OutputType batch, + OutputType row, + OutputType col) const { + T pixel = read_pixel(batch, row, col); + if (is_nonzero(pixel)) { + const int64 index_a = col + num_cols_ * (row + num_rows_ * batch); + if (col + 1 < num_cols_ && read_pixel(batch, row, col + 1) == pixel) { + const int64 index_b = col + 1 + num_cols_ * (row + num_rows_ * batch); + do_union(index_a, index_b); + } + } + } + + // Reads a pixel value in the images. + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE T + read_pixel(const OutputType batch, const OutputType row, + const OutputType col) const { + return images_[col + num_cols_ * (row + num_rows_ * batch)]; + } + + // Unions the trees that the two pixels belong to, using their index in the + // `images_` array. + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void do_union( + OutputType index_a, OutputType index_b) const { + // Find the roots of index_a and index_b in the forest, and make one the + // child of the other. + index_a = find(index_a); + index_b = find(index_b); + const OutputType rank_a = rank_[index_a]; + const OutputType rank_b = rank_[index_b]; + OutputType parent, child; + if (index_a == index_b) { + return; + } else if (rank_a < rank_b) { + parent = index_a; + child = index_b; + } else { + parent = index_b; + child = index_a; + rank_[parent]++; + } + forest_[child] = parent; + } +}; + +// Runs the ImageUnionFindFunctor on all pixels. Will require different CPU and +// GPU implementations. +template +class ImageConnectedComponentsFunctor { + public: + using OutputType = typename BlockedImageUnionFindFunctor::OutputType; + + void operator()(OpKernelContext* ctx, + typename TTypes::ConstTensor images, + typename TTypes::Tensor forest, + typename TTypes::Tensor rank); +}; + +// Fills a flat Tensor with indices from 0 to n - 1. +template +class TensorRangeFunctor { + public: + using OutputType = typename BlockedImageUnionFindFunctor::OutputType; + + void operator()(const Device& device, + typename TTypes::Flat tensor) { + tensor.device(device) = tensor.generate(TensorRangeGenerator()); + } + + private: + class TensorRangeGenerator { + public: + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE OutputType + operator()(const Eigen::array& coords) const { + return coords[0]; + } + }; +}; + +// Given the union-find forest, generates the root index for each node. This +// gives us arbitrary, usually non-consecutive ids for each connected component. +// The ids are massaged in Python to get deterministic, consecutive ids. +template +class FindRootFunctor { + public: + using OutputType = typename BlockedImageUnionFindFunctor::OutputType; + + void operator()(const Device& device, + typename TTypes::Flat component_ids, + const T* images, + const BlockedImageUnionFindFunctor& union_find) { + component_ids.device(device) = + component_ids.generate(FindRootGenerator(images, union_find)); + } + + private: + class FindRootGenerator { + const T* const images_; + const BlockedImageUnionFindFunctor union_find_; + + public: + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE FindRootGenerator( + const T* images, BlockedImageUnionFindFunctor union_find) + : images_(images), union_find_(union_find) {} + + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE OutputType + operator()(const Eigen::array& coords) const { + if (is_nonzero(images_[coords[0]])) { + // True pixels have an arbitrary segment id > 0. The segment ids will be + // made contiguous later. + return union_find_.find(coords[0]) + 1; + } else { + // False pixels have a segment of 0. + return 0; + } + } + }; +}; + +} // end namespace functor + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IMAGE_KERNELS_SEGMENTATION_OPS_H_ diff --git a/tensorflow/contrib/image/kernels/single_image_random_dot_stereograms_ops.cc b/tensorflow/contrib/image/kernels/single_image_random_dot_stereograms_ops.cc index 9f0bf37aed3fc9aeefb7602ef3fda4cfd76f1917..8f9a5c28039b74a874028826ca8a6d5a36ab7cf4 100755 --- a/tensorflow/contrib/image/kernels/single_image_random_dot_stereograms_ops.cc +++ b/tensorflow/contrib/image/kernels/single_image_random_dot_stereograms_ops.cc @@ -143,8 +143,8 @@ class SingleImageRandomDotStereogramsOp : public OpKernel { } data_box_left = deltaX_border_image / 2; // Center DATA in X dimension - data_box_width = data_Xwindow; // width of scan line - data_box_height = data_Ywindow; // hight of image + data_box_width = data_Xwindow; // width of scan line + data_box_height = data_Ywindow; // hight of image const T* inputZ = input_tensor.flat().data(); // Flatten input Z buffer diff --git a/tensorflow/contrib/image/ops/image_ops.cc b/tensorflow/contrib/image/ops/image_ops.cc index 4527fdd87a8be3390fb0840410218ab74a27f0d2..68771b3d054a64ba94141c092e20df1ed6b2339b 100644 --- a/tensorflow/contrib/image/ops/image_ops.cc +++ b/tensorflow/contrib/image/ops/image_ops.cc @@ -98,4 +98,34 @@ col_to_row_match_indices: A vector of length num_columns, which is the number `col_to_row_match_indices[j]`. )doc"); +REGISTER_OP("ImageConnectedComponents") + .Input("image: dtype") + .Output("components: int64") + .Attr( + "dtype: {int64, int32, uint16, int16, uint8, int8, half, float, " + "double, bool, string}") + .SetShapeFn([](InferenceContext* c) { + return shape_inference::UnchangedShape(c); + }) + .Doc(R"doc( +Find the connected components of image(s). + +For each image (along the 0th axis), all connected components of adjacent pixels +with the same non-zero value are detected and given unique ids. + +The returned `components` tensor has 0s for the zero pixels of `images`, and +arbitrary nonzero ids for the connected components of nonzero values. Ids are +unique across all of the images, and are in row-major order by the first pixel +in the component. + +Uses union-find with union by rank but not path compression, giving a runtime of +`O(n log n)`. See: + https://en.wikipedia.org/wiki/Disjoint-set_data_structure#Time_Complexity + +image: Image(s) with shape (N, H, W). +components: Component ids for each pixel in "image". Same shape as "image". Zero + pixels all have an output of 0, and all components of adjacent pixels with + the same value are given consecutive ids, starting from 1. +)doc"); + } // namespace tensorflow diff --git a/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc b/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc index f8b56ab1c5400694b3aa8d4a0c19c7769aa8cbce..8139d4272d6950815bd39a64e86e0f7422e6f799 100755 --- a/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc +++ b/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc @@ -19,6 +19,10 @@ limitations under the License. namespace tensorflow { +using shape_inference::DimensionHandle; +using shape_inference::InferenceContext; +using shape_inference::ShapeHandle; + REGISTER_OP("SingleImageRandomDotStereograms") .Attr("T: {double,float,int64,int32}") .Input("depth_values: T") @@ -37,6 +41,28 @@ REGISTER_OP("SingleImageRandomDotStereograms") "output_image_shape: shape = { dim {size:1024} dim {size: 768} dim " "{size: 1}}") .Attr("output_data_window: shape = { dim {size:1022} dim {size: 757}}") + .SetShapeFn([](InferenceContext* c) { + // Validate that the output_image_shape attr is correct. + // NOTE: The output_image_shape is [X, Y, C] + // while the output data is [Y, X, C] (or [H, W, C]). + // As a result, by default the output_image_shape has the value + // of [1024, 768, 1] but the output data will be [768, 1024, 1]. + PartialTensorShape shape; + TF_RETURN_IF_ERROR(c->GetAttr("output_image_shape", &shape)); + ShapeHandle output_image_shape; + TF_RETURN_IF_ERROR( + c->MakeShapeFromPartialTensorShape(shape, &output_image_shape)); + DimensionHandle x_dim = c->Dim(output_image_shape, 0); + DimensionHandle y_dim = c->Dim(output_image_shape, 1); + + int colors; + TF_RETURN_IF_ERROR(c->GetAttr("number_colors", &colors)); + + c->set_output( + 0, c->MakeShape( + {y_dim, x_dim, colors > 256 ? c->MakeDim(3) : c->MakeDim(1)})); + return Status::OK(); + }) .Doc(R"doc( Outputs a single image random dot stereogram for export via encode_PNG/JPG OP. diff --git a/tensorflow/contrib/image/python/kernel_tests/segmentation_test.py b/tensorflow/contrib/image/python/kernel_tests/segmentation_test.py new file mode 100644 index 0000000000000000000000000000000000000000..48066cbacefe6b229a1f485486f11e8b8af7704f --- /dev/null +++ b/tensorflow/contrib/image/python/kernel_tests/segmentation_test.py @@ -0,0 +1,189 @@ +# 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 connected component analysis.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import logging + +import numpy as np + +from tensorflow.contrib.image.python.ops import image_ops +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import googletest + +# Image for testing connected_components, with a single, winding component. +SNAKE = np.asarray( + [[0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 1, 1, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 0, 0, 0, 0, 0, 0, 0], + [0, 1, 0, 1, 1, 1, 1, 1, 0], + [0, 1, 0, 0, 0, 0, 0, 1, 0], + [0, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0]]) # pyformat: disable + + +class SegmentationTest(test_util.TensorFlowTestCase): + + def testDisconnected(self): + arr = math_ops.cast( + [[1, 0, 0, 1, 0, 0, 0, 0, 1], + [0, 1, 0, 0, 0, 1, 0, 1, 0], + [1, 0, 1, 0, 0, 0, 1, 0, 0], + [0, 0, 0, 0, 1, 0, 0, 0, 0], + [0, 0, 1, 0, 0, 0, 0, 0, 0]], + dtypes.bool) # pyformat: disable + expected = ( + [[1, 0, 0, 2, 0, 0, 0, 0, 3], + [0, 4, 0, 0, 0, 5, 0, 6, 0], + [7, 0, 8, 0, 0, 0, 9, 0, 0], + [0, 0, 0, 0, 10, 0, 0, 0, 0], + [0, 0, 11, 0, 0, 0, 0, 0, 0]]) # pyformat: disable + with self.test_session(): + self.assertAllEqual(image_ops.connected_components(arr).eval(), expected) + + def testSimple(self): + arr = [[0, 1, 0], [1, 1, 1], [0, 1, 0]] + with self.test_session(): + # Single component with id 1. + self.assertAllEqual( + image_ops.connected_components(math_ops.cast( + arr, dtypes.bool)).eval(), arr) + + def testSnake(self): + with self.test_session(): + # Single component with id 1. + self.assertAllEqual( + image_ops.connected_components(math_ops.cast( + SNAKE, dtypes.bool)).eval(), SNAKE) + + def testSnake_disconnected(self): + for i in range(SNAKE.shape[0]): + for j in range(SNAKE.shape[1]): + with self.test_session(): + # If we disconnect any part of the snake except for the endpoints, + # there will be 2 components. + if SNAKE[i, j] and (i, j) not in [(1, 1), (6, 3)]: + disconnected_snake = SNAKE.copy() + disconnected_snake[i, j] = 0 + components = image_ops.connected_components( + math_ops.cast(disconnected_snake, dtypes.bool)).eval() + self.assertEqual(components.max(), 2, 'disconnect (%d, %d)' % (i, + j)) + bins = np.bincount(components.ravel()) + # Nonzero number of pixels labeled 0, 1, or 2. + self.assertGreater(bins[0], 0) + self.assertGreater(bins[1], 0) + self.assertGreater(bins[2], 0) + + def testMultipleImages(self): + images = [[[1, 1, 1, 1], + [1, 0, 0, 1], + [1, 0, 0, 1], + [1, 1, 1, 1]], + [[1, 0, 0, 1], + [0, 0, 0, 0], + [0, 0, 0, 0], + [1, 0, 0, 1]], + [[1, 1, 0, 1], + [0, 1, 1, 0], + [1, 0, 1, 0], + [0, 0, 1, 1]]] # pyformat: disable + expected = [[[1, 1, 1, 1], + [1, 0, 0, 1], + [1, 0, 0, 1], + [1, 1, 1, 1]], + [[2, 0, 0, 3], + [0, 0, 0, 0], + [0, 0, 0, 0], + [4, 0, 0, 5]], + [[6, 6, 0, 7], + [0, 6, 6, 0], + [8, 0, 6, 0], + [0, 0, 6, 6]]] # pyformat: disable + with self.test_session(): + self.assertAllEqual( + image_ops.connected_components(math_ops.cast( + images, dtypes.bool)).eval(), expected) + + def testZeros(self): + with self.test_session(): + self.assertAllEqual( + image_ops.connected_components( + array_ops.zeros((100, 20, 50), dtypes.bool)).eval(), + np.zeros((100, 20, 50))) + + def testOnes(self): + with self.test_session(): + self.assertAllEqual( + image_ops.connected_components( + array_ops.ones((100, 20, 50), dtypes.bool)).eval(), + np.tile(np.arange(100)[:, None, None] + 1, [1, 20, 50])) + + def testOnes_small(self): + with self.test_session(): + self.assertAllEqual( + image_ops.connected_components(array_ops.ones((3, 5), + dtypes.bool)).eval(), + np.ones((3, 5))) + + def testRandom_scipy(self): + np.random.seed(42) + images = np.random.randint(0, 2, size=(10, 100, 200)).astype(np.bool) + expected = connected_components_reference_implementation(images) + if expected is None: + return + with self.test_session(): + self.assertAllEqual( + image_ops.connected_components(images).eval(), expected) + + +def connected_components_reference_implementation(images): + try: + # pylint: disable=g-import-not-at-top + from scipy.ndimage import measurements + except ImportError: + logging.exception('Skipping test method because scipy could not be loaded') + return + image_or_images = np.asarray(images) + if len(image_or_images.shape) == 2: + images = image_or_images[None, :, :] + elif len(image_or_images.shape) == 3: + images = image_or_images + components = np.asarray([measurements.label(image)[0] for image in images]) + # Get the count of nonzero ids for each image, and offset each image's nonzero + # ids using the cumulative sum. + num_ids_per_image = components.reshape( + [-1, components.shape[1] * components.shape[2]]).max(axis=-1) + positive_id_start_per_image = np.cumsum(num_ids_per_image) + for i in range(components.shape[0]): + new_id_start = positive_id_start_per_image[i - 1] if i > 0 else 0 + components[i, components[i] > 0] += new_id_start + if len(image_or_images.shape) == 2: + return components[0, :, :] + else: + return components + + +if __name__ == '__main__': + googletest.main() diff --git a/tensorflow/contrib/image/python/kernel_tests/single_image_random_dot_stereograms_ops_test.py b/tensorflow/contrib/image/python/kernel_tests/single_image_random_dot_stereograms_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3f4029e558d92a2b6539456bf9cf49ec2d21c9f3 --- /dev/null +++ b/tensorflow/contrib/image/python/kernel_tests/single_image_random_dot_stereograms_ops_test.py @@ -0,0 +1,82 @@ +# Copyright 2017 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 python single_image_random_dot_stereograms_ops.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.image.python.ops.single_image_random_dot_stereograms \ + import single_image_random_dot_stereograms +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import test_util +from tensorflow.python.platform import googletest + +class SingleImageRandomDotStereogramsTest(test_util.TensorFlowTestCase): + + def test_shape_function_default(self): + """ + NOTE: The output_image_shape is [X, Y, C] + while the output data is [Y, X, C] (or [H, W, C]). + As a result, by default the output_image_shape has the value + of [1024, 768, 1], but the output data will be [768, 1024, 1]. + """ + x_np = [[1, 2, 3, 3, 2, 1], + [1, 2, 3, 4, 5, 2], + [1, 2, 3, 4, 5, 3], + [1, 2, 3, 4, 5, 4], + [6, 5, 4, 4, 5, 5]] + x_tf = constant_op.constant(x_np) + # By default [1024, 768, 1] => [768, 1024, 1]. + sirds_1 = single_image_random_dot_stereograms( + x_tf, + convergence_dots_size=8, + number_colors=256, + normalize=True) + shape_1 = sirds_1.get_shape().as_list() + self.assertEqual(shape_1, [768, 1024, 1]) + with self.test_session(): + r_tf_1 = sirds_1.eval() + self.assertAllEqual(shape_1, r_tf_1.shape) + + # If color > 256 then [1024, 768, 3] => [768, 1024, 3]. + sirds_2 = single_image_random_dot_stereograms( + x_tf, + convergence_dots_size=8, + number_colors=512, + normalize=True) + shape_2 = sirds_2.get_shape().as_list() + self.assertEqual(shape_2, [768, 1024, 3]) + with self.test_session(): + r_tf_2 = sirds_2.eval() + self.assertAllEqual(shape_2, r_tf_2.shape) + + # If explicitly set output_image_shape to [1200, 800, 1], + # then the output data should be [800, 1200, 1]. + sirds_3 = single_image_random_dot_stereograms( + x_tf, + convergence_dots_size=8, + number_colors=256, + normalize=True, + output_image_shape=[1200, 800, 1]) + shape_3 = sirds_3.get_shape().as_list() + self.assertEqual(shape_3, [800, 1200, 1]) + with self.test_session(): + r_tf_3 = sirds_3.eval() + self.assertAllEqual(shape_3, r_tf_3.shape) + + +if __name__ == '__main__': + googletest.main() diff --git a/tensorflow/contrib/image/python/ops/image_ops.py b/tensorflow/contrib/image/python/ops/image_ops.py index faedee6f87772016561671bacd87f88657eafffb..c139ae89d8d682d6b87813c3a21703ffa762f28e 100644 --- a/tensorflow/contrib/image/python/ops/image_ops.py +++ b/tensorflow/contrib/image/python/ops/image_ops.py @@ -24,6 +24,7 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import resource_loader @@ -34,11 +35,12 @@ _image_ops_so = loader.load_op_library( _IMAGE_DTYPES = set( [dtypes.uint8, dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64]) +ops.RegisterShape("ImageConnectedComponents")(common_shapes.call_cpp_shape_fn) ops.RegisterShape("ImageProjectiveTransform")(common_shapes.call_cpp_shape_fn) def rotate(images, angles, interpolation="NEAREST", name=None): - """Rotate image(s) by the passed angle(s) in radians. + """Rotate image(s) counterclockwise by the passed angle(s) in radians. Args: images: A tensor of shape (num_images, num_rows, num_columns, num_channels) @@ -288,31 +290,76 @@ def compose_transforms(*transforms): """ assert transforms, "transforms cannot be empty" with ops.name_scope("compose_transforms"): - composed = _flat_transforms_to_matrices(transforms[0]) + composed = flat_transforms_to_matrices(transforms[0]) for tr in transforms[1:]: # Multiply batches of matrices. - composed = math_ops.matmul(composed, _flat_transforms_to_matrices(tr)) - return _transform_matrices_to_flat(composed) + composed = math_ops.matmul(composed, flat_transforms_to_matrices(tr)) + return matrices_to_flat_transforms(composed) -def _flat_transforms_to_matrices(transforms): - # Make the transform(s) 2D in case the input is a single transform. - transforms = array_ops.reshape(transforms, constant_op.constant([-1, 8])) - num_transforms = array_ops.shape(transforms)[0] - # Add a column of ones for the implicit last entry in the matrix. - return array_ops.reshape( - array_ops.concat( - [transforms, array_ops.ones([num_transforms, 1])], axis=1), - constant_op.constant([-1, 3, 3])) +def flat_transforms_to_matrices(transforms): + """Converts `tf.contrib.image` projective transforms to affine matrices. + Note that the output matrices map output coordinates to input coordinates. For + the forward transformation matrix, call `tf.linalg.inv` on the result. -def _transform_matrices_to_flat(transform_matrices): - # Flatten each matrix. - transforms = array_ops.reshape(transform_matrices, - constant_op.constant([-1, 9])) - # Divide each matrix by the last entry (normally 1). - transforms /= transforms[:, 8:9] - return transforms[:, :8] + Args: + transforms: Vector of length 8, or batches of transforms with shape + `(N, 8)`. + + Returns: + 3D tensor of matrices with shape `(N, 3, 3)`. The output matrices map the + *output coordinates* (in homogeneous coordinates) of each transform to the + corresponding *input coordinates*. + + Raises: + ValueError: If `transforms` have an invalid shape. + """ + with ops.name_scope("flat_transforms_to_matrices"): + transforms = ops.convert_to_tensor(transforms, name="transforms") + if transforms.shape.ndims not in (1, 2): + raise ValueError("Transforms should be 1D or 2D, got: %s" % transforms) + # Make the transform(s) 2D in case the input is a single transform. + transforms = array_ops.reshape(transforms, constant_op.constant([-1, 8])) + num_transforms = array_ops.shape(transforms)[0] + # Add a column of ones for the implicit last entry in the matrix. + return array_ops.reshape( + array_ops.concat( + [transforms, array_ops.ones([num_transforms, 1])], axis=1), + constant_op.constant([-1, 3, 3])) + + +def matrices_to_flat_transforms(transform_matrices): + """Converts affine matrices to `tf.contrib.image` projective transforms. + + Note that we expect matrices that map output coordinates to input coordinates. + To convert forward transformation matrices, call `tf.linalg.inv` on the + matrices and use the result here. + + Args: + transform_matrices: One or more affine transformation matrices, for the + reverse transformation in homogeneous coordinates. Shape `(3, 3)` or + `(N, 3, 3)`. + + Returns: + 2D tensor of flat transforms with shape `(N, 8)`, which may be passed into + `tf.contrib.image.transform`. + + Raises: + ValueError: If `transform_matrices` have an invalid shape. + """ + with ops.name_scope("matrices_to_flat_transforms"): + transform_matrices = ops.convert_to_tensor( + transform_matrices, name="transform_matrices") + if transform_matrices.shape.ndims not in (2, 3): + raise ValueError( + "Matrices should be 2D or 3D, got: %s" % transform_matrices) + # Flatten each matrix. + transforms = array_ops.reshape(transform_matrices, + constant_op.constant([-1, 9])) + # Divide each matrix by the last entry (normally 1). + transforms /= transforms[:, 8:9] + return transforms[:, :8] @ops.RegisterGradient("ImageProjectiveTransform") @@ -344,9 +391,9 @@ def _image_projective_transform_grad(op, grad): raise TypeError("Transforms should have rank 1 or 2.") # Invert transformations - transforms = _flat_transforms_to_matrices(transforms=transforms) + transforms = flat_transforms_to_matrices(transforms=transforms) inverse = linalg_ops.matrix_inverse(transforms) - transforms = _transform_matrices_to_flat(inverse) + transforms = matrices_to_flat_transforms(inverse) output = gen_image_ops.image_projective_transform( grad, transforms, interpolation=interpolation) if len(image_or_images.get_shape()) == 2: @@ -395,4 +442,72 @@ def bipartite_match(distance_mat, return result +def connected_components(images): + """Labels the connected components in a batch of images. + + A component is a set of pixels in a single input image, which are all adjacent + and all have the same non-zero value. The components using a squared + connectivity of one (all True entries are joined with their neighbors above, + below, left, and right). Components across all images have consecutive ids 1 + through n. Components are labeled according to the first pixel of the + component appearing in row-major order (lexicographic order by + image_index_in_batch, row, col). Zero entries all have an output id of 0. + + This op is equivalent with `scipy.ndimage.measurements.label` on a 2D array + with the default structuring element (which is the connectivity used here). + + Args: + images: A 2D (H, W) or 3D (N, H, W) Tensor of boolean image(s). + + Returns: + Components with the same shape as `images`. False entries in `images` have + value 0, and all True entries map to a component id > 0. + + Raises: + TypeError: if `images` is not 2D or 3D. + """ + with ops.name_scope("connected_components"): + image_or_images = ops.convert_to_tensor(images, name="images") + if len(image_or_images.get_shape()) == 2: + images = image_or_images[None, :, :] + elif len(image_or_images.get_shape()) == 3: + images = image_or_images + else: + raise TypeError( + "images should have rank 2 (HW) or 3 (NHW). Static shape is %s" % + image_or_images.get_shape()) + components = gen_image_ops.image_connected_components(images) + + # TODO(ringwalt): Component id renaming should be done in the op, to avoid + # constructing multiple additional large tensors. + components_flat = array_ops.reshape(components, [-1]) + unique_ids, id_index = array_ops.unique(components_flat) + id_is_zero = array_ops.where(math_ops.equal(unique_ids, 0))[:, 0] + # Map each nonzero id to consecutive values. + nonzero_consecutive_ids = math_ops.range( + array_ops.shape(unique_ids)[0] - array_ops.shape(id_is_zero)[0]) + 1 + + def no_zero(): + # No need to insert a zero into the ids. + return nonzero_consecutive_ids + + def has_zero(): + # Insert a zero in the consecutive ids where zero appears in unique_ids. + # id_is_zero has length 1. + zero_id_ind = math_ops.to_int32(id_is_zero[0]) + ids_before = nonzero_consecutive_ids[:zero_id_ind] + ids_after = nonzero_consecutive_ids[zero_id_ind:] + return array_ops.concat([ids_before, [0], ids_after], axis=0) + + new_ids = control_flow_ops.cond( + math_ops.equal(array_ops.shape(id_is_zero)[0], 0), no_zero, has_zero) + components = array_ops.reshape( + array_ops.gather(new_ids, id_index), array_ops.shape(components)) + if len(image_or_images.get_shape()) == 2: + return components[0, :, :] + else: + return components + + ops.NotDifferentiable("BipartiteMatch") +ops.NotDifferentiable("ImageConnectedComponents") diff --git a/tensorflow/contrib/image/python/ops/single_image_random_dot_stereograms.py b/tensorflow/contrib/image/python/ops/single_image_random_dot_stereograms.py index bb766e59d2cee648042cc08be466796d9233ad66..d4a6a5bcbb52511d4093587814100b2a0e8b2420 100755 --- a/tensorflow/contrib/image/python/ops/single_image_random_dot_stereograms.py +++ b/tensorflow/contrib/image/python/ops/single_image_random_dot_stereograms.py @@ -26,18 +26,20 @@ _sirds_ops = loader.load_op_library( resource_loader.get_path_to_datafile( "_single_image_random_dot_stereograms.so")) -def single_image_random_dot_stereograms( - depth_values, - hidden_surface_removal=None, - convergence_dots_size=None, - dots_per_inch=None, - eye_separation=None, mu=None, - normalize=None, normalize_max=None, - normalize_min=None, - border_level=None, - number_colors=None, - output_image_shape=None, - output_data_window=None): + +def single_image_random_dot_stereograms(depth_values, + hidden_surface_removal=None, + convergence_dots_size=None, + dots_per_inch=None, + eye_separation=None, + mu=None, + normalize=None, + normalize_max=None, + normalize_min=None, + border_level=None, + number_colors=None, + output_image_shape=None, + output_data_window=None): """Output a RandomDotStereogram Tensor for export via encode_PNG/JPG OP. Given the 2-D tensor 'depth_values' with encoded Z values, this operation @@ -45,7 +47,8 @@ def single_image_random_dot_stereograms( for the encode_PNG/JPG ops. Be careful with image compression as this may corrupt the encode 3-D data witin the image. - Based upon [this paper](http://www.learningace.com/doc/4331582/b6ab058d1e206d68ab60e4e1ead2fe6e/sirds-paper). + Based upon [this + paper](http://www.learningace.com/doc/4331582/b6ab058d1e206d68ab60e4e1ead2fe6e/sirds-paper). This outputs a SIRDS image as picture_out.png: @@ -113,7 +116,8 @@ def single_image_random_dot_stereograms( hidden_surface_removal=hidden_surface_removal, convergence_dots_size=convergence_dots_size, dots_per_inch=dots_per_inch, - eye_separation=eye_separation, mu=mu, + eye_separation=eye_separation, + mu=mu, normalize=normalize, normalize_max=normalize_max, normalize_min=normalize_min, @@ -123,4 +127,5 @@ def single_image_random_dot_stereograms( output_data_window=output_data_window) return result + ops.NotDifferentiable("SingleImageRandomDotStereograms") diff --git a/tensorflow/contrib/input_pipeline/kernels/input_pipeline_kernels.cc b/tensorflow/contrib/input_pipeline/kernels/input_pipeline_kernels.cc index ca288c1f737d25faac678f5c199d5c1e49f721cb..886f6798150c57d8066546b0919481d3878882fc 100644 --- a/tensorflow/contrib/input_pipeline/kernels/input_pipeline_kernels.cc +++ b/tensorflow/contrib/input_pipeline/kernels/input_pipeline_kernels.cc @@ -34,9 +34,8 @@ class ObtainNextOp : public OpKernel { // Allocate output. Tensor* output_tensor = nullptr; - OP_REQUIRES_OK( - ctx, - ctx->allocate_output("out_element", TensorShape({}), &output_tensor)); + OP_REQUIRES_OK(ctx, ctx->allocate_output("out_element", TensorShape({}), + &output_tensor)); // Obtain mutex for the "counter" tensor. mutex* mu; diff --git a/tensorflow/contrib/kafka/BUILD b/tensorflow/contrib/kafka/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..efb403462a6e5df5b69ac0735ffc03f40d4a252c --- /dev/null +++ b/tensorflow/contrib/kafka/BUILD @@ -0,0 +1,105 @@ +package( + default_visibility = ["//visibility:private"], +) + +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + +load("//tensorflow:tensorflow.bzl", "tf_gen_op_libs") +load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py") +load("//tensorflow:tensorflow.bzl", "tf_kernel_library") +load("//tensorflow:tensorflow.bzl", "tf_py_test") + +tf_kernel_library( + name = "kafka_kernels", + srcs = ["kernels/kafka_dataset_ops.cc"], + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core/kernels:bounds_check_lib", + "//tensorflow/core/kernels:dataset", + "//third_party/eigen3", + "@kafka", + ], +) + +tf_gen_op_libs( + op_lib_names = ["kafka_ops"], + deps = [ + "//tensorflow/core:lib", + ], +) + +tf_gen_op_wrapper_py( + name = "gen_kafka_ops", + out = "python/ops/gen_kafka_ops.py", + require_shape_functions = True, + deps = [":kafka_ops_op_lib"], +) + +py_library( + name = "kafka", + srcs = [ + "__init__.py", + "python/ops/kafka_dataset_ops.py", + ], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + ":gen_kafka_ops", + "//tensorflow/contrib/util:util_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:framework", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform", + "//tensorflow/python:state_ops", + "//tensorflow/python:training", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/ops:iterator_ops", + "//tensorflow/python/data/ops:readers", + ], +) + +# The Kafka server has to be setup before running the test. +# The Kafka server is setup through Docker so the Docker engine +# has to be installed. +# +# Once the Docker engine is ready: +# To setup the Kafka server: +# $ bash tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh start kafka +# +# After the test is complete: +# To team down the Kafka server: +# $ bash tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh stop kafka +tf_py_test( + name = "kafka_test", + srcs = ["python/kernel_tests/kafka_test.py"], + additional_deps = [ + ":kafka", + "//third_party/py/numpy", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:platform_test", + ], + tags = [ + "manual", + "notap", + ], +) + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/contrib/ndlstm/python/__init__.py b/tensorflow/contrib/kafka/__init__.py similarity index 68% rename from tensorflow/contrib/ndlstm/python/__init__.py rename to tensorflow/contrib/kafka/__init__.py index 1aa51a6ec40c042ca3c26c6b08e5bdb8a42a12bd..4d755c40568dfa2f7f6f617cf3180268837a5ca0 100644 --- a/tensorflow/contrib/ndlstm/python/__init__.py +++ b/tensorflow/contrib/kafka/__init__.py @@ -4,7 +4,7 @@ # 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 +# 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, @@ -12,14 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Init file, giving convenient access to all ndlstm ops.""" +"""Kafka Dataset. + +@@KafkaDataset +""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -# pylint: disable=wildcard-import,g-importing-member -from tensorflow.contrib.ndlstm.python.lstm1d import * -from tensorflow.contrib.ndlstm.python.lstm2d import * -from tensorflow.contrib.ndlstm.python.misc import * -# pylint: enable=wildcard-import +from tensorflow.contrib.kafka.python.ops.kafka_dataset_ops import KafkaDataset + +from tensorflow.python.util.all_util import remove_undocumented + +_allowed_symbols = [ + "KafkaDataset", +] + +remove_undocumented(__name__) diff --git a/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc b/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..88ef5f357113372b0a2d0cb13382ac980a61252d --- /dev/null +++ b/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc @@ -0,0 +1,321 @@ +/* Copyright 2017 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/kernels/dataset.h" + +#include "tensorflow/core/framework/tensor.h" + +#include "src-cpp/rdkafkacpp.h" + +namespace tensorflow { + +class KafkaDatasetOp : public DatasetOpKernel { + public: + using DatasetOpKernel::DatasetOpKernel; + + void MakeDataset(OpKernelContext* ctx, DatasetBase** output) override { + const Tensor* topics_tensor; + OP_REQUIRES_OK(ctx, ctx->input("topics", &topics_tensor)); + OP_REQUIRES( + ctx, topics_tensor->dims() <= 1, + errors::InvalidArgument("`topics` must be a scalar or a vector.")); + + std::vector topics; + topics.reserve(topics_tensor->NumElements()); + for (int i = 0; i < topics_tensor->NumElements(); ++i) { + topics.push_back(topics_tensor->flat()(i)); + } + + std::string servers = ""; + OP_REQUIRES_OK(ctx, + ParseScalarArgument(ctx, "servers", &servers)); + std::string group = ""; + OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "group", &group)); + bool eof = false; + OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "eof", &eof)); + int64 timeout = -1; + OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "timeout", &timeout)); + OP_REQUIRES(ctx, (timeout > 0), + errors::InvalidArgument( + "Timeout value should be large than 0, got ", timeout)); + *output = new Dataset(ctx, std::move(topics), servers, group, eof, timeout); + } + + private: + class Dataset : public GraphDatasetBase { + public: + Dataset(OpKernelContext* ctx, std::vector topics, + const string& servers, const string& group, const bool eof, + const int64 timeout) + : GraphDatasetBase(ctx), + topics_(std::move(topics)), + servers_(servers), + group_(group), + eof_(eof), + timeout_(timeout) {} + + std::unique_ptr MakeIterator( + const string& prefix) const override { + return std::unique_ptr( + new Iterator({this, strings::StrCat(prefix, "::Kafka")})); + } + + const DataTypeVector& output_dtypes() const override { + static DataTypeVector* dtypes = new DataTypeVector({DT_STRING}); + return *dtypes; + } + + const std::vector& output_shapes() const override { + static std::vector* shapes = + new std::vector({{}}); + return *shapes; + } + + string DebugString() override { return "KafkaDatasetOp::Dataset"; } + + protected: + Status AsGraphDefInternal(DatasetGraphDefBuilder* b, + Node** output) const override { + Node* topics = nullptr; + TF_RETURN_IF_ERROR(b->AddVector(topics_, &topics)); + Node* servers = nullptr; + TF_RETURN_IF_ERROR(b->AddScalar(servers_, &servers)); + Node* group = nullptr; + TF_RETURN_IF_ERROR(b->AddScalar(group_, &group)); + Node* eof = nullptr; + TF_RETURN_IF_ERROR(b->AddScalar(eof_, &eof)); + Node* timeout = nullptr; + TF_RETURN_IF_ERROR(b->AddScalar(timeout_, &timeout)); + TF_RETURN_IF_ERROR( + b->AddDataset(this, {topics, servers, group, eof, timeout}, output)); + return Status::OK(); + } + + private: + class Iterator : public DatasetIterator { + public: + explicit Iterator(const Params& params) + : DatasetIterator(params) {} + + Status GetNextInternal(IteratorContext* ctx, + std::vector* out_tensors, + bool* end_of_sequence) override { + mutex_lock l(mu_); + do { + // We are currently processing a topic, so try to read the next line. + if (consumer_.get()) { + while (true) { + if (limit_ >= 0 && + (topic_partition_->offset() >= limit_ || offset_ >= limit_)) { + // EOF current topic + break; + } + std::unique_ptr message( + consumer_->consume(dataset()->timeout_)); + if (message->err() == RdKafka::ERR_NO_ERROR) { + // Produce the line as output. + Tensor line_tensor(cpu_allocator(), DT_STRING, {}); + line_tensor.scalar()() = + std::string(static_cast(message->payload()), + message->len()); + out_tensors->emplace_back(std::move(line_tensor)); + *end_of_sequence = false; + // Sync offset + offset_ = message->offset(); + return Status::OK(); + } + + if (message->err() == RdKafka::ERR__PARTITION_EOF && + dataset()->eof_) { + // EOF current topic + break; + } + if (message->err() != RdKafka::ERR__TIMED_OUT) { + return errors::Internal("Failed to consume:", + message->errstr()); + } + message.reset(nullptr); + consumer_->poll(0); + } + + // We have reached the end of the current topic, so maybe + // move on to next topic. + ResetStreamsLocked(); + ++current_topic_index_; + } + + // Iteration ends when there are no more topic to process. + if (current_topic_index_ == dataset()->topics_.size()) { + *end_of_sequence = true; + return Status::OK(); + } + + TF_RETURN_IF_ERROR(SetupStreamsLocked(ctx->env())); + } while (true); + } + + protected: + Status SaveInternal(IteratorStateWriter* writer) override { + mutex_lock l(mu_); + TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("current_topic_index"), + current_topic_index_)); + + // `consumer_` is empty if + // 1. GetNext has not been called even once. + // 2. All topics have been read and iterator has been exhausted. + if (consumer_.get()) { + TF_RETURN_IF_ERROR( + writer->WriteScalar(full_name("current_pos"), offset_)); + } + return Status::OK(); + } + + Status RestoreInternal(IteratorContext* ctx, + IteratorStateReader* reader) override { + mutex_lock l(mu_); + ResetStreamsLocked(); + int64 current_topic_index; + TF_RETURN_IF_ERROR(reader->ReadScalar(full_name("current_topic_index"), + ¤t_topic_index)); + current_topic_index_ = size_t(current_topic_index); + // The key "current_pos" is written only if the iterator was saved + // with an open topic. + if (reader->Contains(full_name("current_pos"))) { + int64 current_pos; + TF_RETURN_IF_ERROR( + reader->ReadScalar(full_name("current_pos"), ¤t_pos)); + + TF_RETURN_IF_ERROR(SetupStreamsLocked(ctx->env())); + topic_partition_->set_offset(current_pos); + if (topic_partition_->offset() != current_pos) { + return errors::Internal("Failed to restore to offset ", + current_pos); + } + offset_ = current_pos; + } + return Status::OK(); + } + + private: + // Sets up Kafka streams to read from the topic at + // `current_topic_index_`. + Status SetupStreamsLocked(Env* env) EXCLUSIVE_LOCKS_REQUIRED(mu_) { + if (current_topic_index_ >= dataset()->topics_.size()) { + return errors::InvalidArgument( + "current_topic_index_:", current_topic_index_, + " >= topics_.size():", dataset()->topics_.size()); + } + + // Actually move on to next topic. + string entry = dataset()->topics_[current_topic_index_]; + + std::vector parts = str_util::Split(entry, ":"); + if (parts.size() < 1) { + return errors::InvalidArgument("Invalid parameters: ", entry); + } + string topic = parts[0]; + int32 partition = 0; + if (parts.size() > 1) { + if (!strings::safe_strto32(parts[1], &partition)) { + return errors::InvalidArgument("Invalid parameters: ", entry); + } + } + int64 offset = 0; + if (parts.size() > 2) { + if (!strings::safe_strto64(parts[2], &offset)) { + return errors::InvalidArgument("Invalid parameters: ", entry); + } + } + + topic_partition_.reset( + RdKafka::TopicPartition::create(topic, partition, offset)); + + offset_ = topic_partition_->offset(); + limit_ = -1; + if (parts.size() > 3) { + if (!strings::safe_strto64(parts[3], &limit_)) { + return errors::InvalidArgument("Invalid parameters: ", entry); + } + } + + std::unique_ptr conf( + RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL)); + std::unique_ptr topic_conf( + RdKafka::Conf::create(RdKafka::Conf::CONF_TOPIC)); + + std::string errstr; + + RdKafka::Conf::ConfResult result = + conf->set("default_topic_conf", topic_conf.get(), errstr); + if (result != RdKafka::Conf::CONF_OK) { + return errors::Internal("Failed to set default_topic_conf:", errstr); + } + + result = conf->set("bootstrap.servers", dataset()->servers_, errstr); + if (result != RdKafka::Conf::CONF_OK) { + return errors::Internal("Failed to set bootstrap.servers ", + dataset()->servers_, ":", errstr); + } + result = conf->set("group.id", dataset()->group_, errstr); + if (result != RdKafka::Conf::CONF_OK) { + return errors::Internal("Failed to set group.id ", dataset()->group_, + ":", errstr); + } + + consumer_.reset(RdKafka::KafkaConsumer::create(conf.get(), errstr)); + if (!consumer_.get()) { + return errors::Internal("Failed to create consumer:", errstr); + } + + std::vector partitions; + partitions.emplace_back(topic_partition_.get()); + RdKafka::ErrorCode err = consumer_->assign(partitions); + if (err != RdKafka::ERR_NO_ERROR) { + return errors::Internal( + "Failed to assign partition [", topic_partition_->topic(), ", ", + topic_partition_->partition(), ", ", topic_partition_->offset(), + "]:", RdKafka::err2str(err)); + } + + return Status::OK(); + } + + // Resets all Kafka streams. + void ResetStreamsLocked() EXCLUSIVE_LOCKS_REQUIRED(mu_) { + consumer_->unassign(); + consumer_->close(); + consumer_.reset(nullptr); + } + + mutex mu_; + size_t current_topic_index_ GUARDED_BY(mu_) = 0; + int64 offset_ GUARDED_BY(mu_) = 0; + int64 limit_ GUARDED_BY(mu_) = -1; + std::unique_ptr topic_partition_ GUARDED_BY(mu_); + std::unique_ptr consumer_ GUARDED_BY(mu_); + }; + + const std::vector topics_; + const std::string servers_; + const std::string group_; + const bool eof_; + const int64 timeout_; + }; +}; + +REGISTER_KERNEL_BUILDER(Name("KafkaDataset").Device(DEVICE_CPU), + KafkaDatasetOp); + +} // namespace tensorflow diff --git a/tensorflow/contrib/kafka/ops/kafka_ops.cc b/tensorflow/contrib/kafka/ops/kafka_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..8cdf16103bab2b22d51c144d21a589e1e39f2f0b --- /dev/null +++ b/tensorflow/contrib/kafka/ops/kafka_ops.cc @@ -0,0 +1,44 @@ +/* Copyright 2017 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/common_shape_fns.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" + +namespace tensorflow { + +REGISTER_OP("KafkaDataset") + .Input("topics: string") + .Input("servers: string") + .Input("group: string") + .Input("eof: bool") + .Input("timeout: int64") + .Output("handle: variant") + .SetIsStateful() + .SetShapeFn(shape_inference::ScalarShape) + .Doc(R"doc( +Creates a dataset that emits the messages of one or more Kafka topics. + +topics: A `tf.string` tensor containing one or more subscriptions, + in the format of [topic:partition:offset:length], + by default length is -1 for unlimited. +servers: A list of bootstrap servers. +group: The consumer group id. +eof: If True, the kafka reader will stop on EOF. +timeout: The timeout value for the Kafka Consumer to wait + (in millisecond). +)doc"); + +} // namespace tensorflow diff --git a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py new file mode 100644 index 0000000000000000000000000000000000000000..621911876fc502ece76b08eb6c28697b3c12c863 --- /dev/null +++ b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py @@ -0,0 +1,115 @@ +# Copyright 2017 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 KafkaDataset.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.kafka.python.ops import kafka_dataset_ops +from tensorflow.python.data.ops import iterator_ops +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class KafkaDatasetTest(test.TestCase): + + def setUp(self): + # The Kafka server has to be setup before the test + # and tear down after the test manually. + # The docker engine has to be installed. + # + # To setup the Kafka server: + # $ bash kafka_test.sh start kafka + # + # To team down the Kafka server: + # $ bash kafka_test.sh stop kafka + pass + + def testKafkaDataset(self): + topics = array_ops.placeholder(dtypes.string, shape=[None]) + num_epochs = array_ops.placeholder(dtypes.int64, shape=[]) + batch_size = array_ops.placeholder(dtypes.int64, shape=[]) + + repeat_dataset = kafka_dataset_ops.KafkaDataset( + topics, group="test", eof=True).repeat(num_epochs) + batch_dataset = repeat_dataset.batch(batch_size) + + iterator = iterator_ops.Iterator.from_structure(batch_dataset.output_types) + init_op = iterator.make_initializer(repeat_dataset) + init_batch_op = iterator.make_initializer(batch_dataset) + get_next = iterator.get_next() + + with self.test_session() as sess: + # Basic test: read from topic 0. + sess.run(init_op, feed_dict={topics: ["test:0:0:4"], num_epochs: 1}) + for i in range(5): + self.assertEqual("D" + str(i), sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + # Basic test: read from topic 1. + sess.run(init_op, feed_dict={topics: ["test:0:5:-1"], num_epochs: 1}) + for i in range(5): + self.assertEqual("D" + str(i + 5), sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + # Basic test: read from both topics. + sess.run( + init_op, + feed_dict={ + topics: ["test:0:0:4", "test:0:5:-1"], + num_epochs: 1 + }) + for j in range(2): + for i in range(5): + self.assertEqual("D" + str(i + j * 5), sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + # Test repeated iteration through both files. + sess.run( + init_op, + feed_dict={ + topics: ["test:0:0:4", "test:0:5:-1"], + num_epochs: 10 + }) + for _ in range(10): + for j in range(2): + for i in range(5): + self.assertEqual("D" + str(i + j * 5), sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + + # Test batched and repeated iteration through both files. + sess.run( + init_batch_op, + feed_dict={ + topics: ["test:0:0:4", "test:0:5:-1"], + num_epochs: 10, + batch_size: 5 + }) + for _ in range(10): + self.assertAllEqual(["D" + str(i) for i in range(5)], + sess.run(get_next)) + self.assertAllEqual(["D" + str(i + 5) for i in range(5)], + sess.run(get_next)) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh new file mode 100644 index 0000000000000000000000000000000000000000..adf027b8e714124cde2b4618546e20c6b7162e1f --- /dev/null +++ b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +# 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. +# ============================================================================== + +set -e +set -o pipefail + +if [ "$#" -ne 2 ]; then + echo "Usage: $0 start|stop " >&2 + exit 1 +fi + +container=$2 +if [ "$1" == "start" ]; then + docker run -d --rm --net=host --name=$container spotify/kafka + echo Wait 5 secs until kafka is up and running + sleep 5 + echo Create test topic + docker exec $container bash -c '/opt/kafka_2.11-0.10.1.0/bin/kafka-topics.sh --create --zookeeper localhost:2181 --replication-factor 1 --partitions 1 --topic test' + echo Create test message + docker exec $container bash -c 'echo -e "D0\nD1\nD2\nD3\nD4\nD5\nD6\nD7\nD8\nD9" > /test' + echo Produce test message + docker exec $container bash -c '/opt/kafka_2.11-0.10.1.0/bin/kafka-console-producer.sh --topic test --broker-list 127.0.0.1:9092 < /test' + + echo Container $container started successfully +elif [ "$1" == "stop" ]; then + docker rm -f $container + + echo Container $container stopped successfully +else + echo "Usage: $0 start|stop " >&2 + exit 1 +fi + + + diff --git a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..8e51d27a342359881de072c3979a2b5a7fc034ea --- /dev/null +++ b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py @@ -0,0 +1,74 @@ +# Copyright 2017 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. +# ============================================================================== +"""Kafka Dataset.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.kafka.python.ops import gen_kafka_ops +from tensorflow.python.data.ops.readers import Dataset +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape + + +class KafkaDataset(Dataset): + """A Kafka Dataset that consumes the message. + """ + + def __init__(self, + topics, + servers="localhost", + group="", + eof=False, + timeout=1000): + """Create a KafkaReader. + + Args: + topics: A `tf.string` tensor containing one or more subscriptions, + in the format of [topic:partition:offset:length], + by default length is -1 for unlimited. + servers: A list of bootstrap servers. + group: The consumer group id. + eof: If True, the kafka reader will stop on EOF. + timeout: The timeout value for the Kafka Consumer to wait + (in millisecond). + """ + super(KafkaDataset, self).__init__() + self._topics = ops.convert_to_tensor( + topics, dtype=dtypes.string, name="topics") + self._servers = ops.convert_to_tensor( + servers, dtype=dtypes.string, name="servers") + self._group = ops.convert_to_tensor( + group, dtype=dtypes.string, name="group") + self._eof = ops.convert_to_tensor(eof, dtype=dtypes.bool, name="eof") + self._timeout = ops.convert_to_tensor( + timeout, dtype=dtypes.int64, name="timeout") + + def _as_variant_tensor(self): + return gen_kafka_ops.kafka_dataset(self._topics, self._servers, self._group, + self._eof, self._timeout) + + @property + def output_classes(self): + return ops.Tensor + + @property + def output_shapes(self): + return tensor_shape.scalar() + + @property + def output_types(self): + return dtypes.string diff --git a/tensorflow/contrib/kernel_methods/python/losses_test.py b/tensorflow/contrib/kernel_methods/python/losses_test.py index d38d8041ce1216dfb5af6e93984b35e71008610a..72507539f813d14064bc58f03b6db4781abc9438 100644 --- a/tensorflow/contrib/kernel_methods/python/losses_test.py +++ b/tensorflow/contrib/kernel_methods/python/losses_test.py @@ -119,19 +119,20 @@ class SparseMulticlassHingeLossTest(test.TestCase): def testUnknownShape(self): """Result keeps same with `testZeroLossInt32Labels`""" - logits_np = np.array([[1.2, -1.4, -1.0], - [1.4, 1.8, 4.0], - [0.5, 1.8, -1.0]]) + logits_np = np.array([[1.2, -1.4, -1.0], [1.4, 1.8, 4.0], [0.5, 1.8, -1.0]]) labels_np = np.array([0, 2, 1], dtype=np.int32) - logits_shapes = [[3, 3], # batch_size, num_classes - [None, 3], - [3, None], - [None, None]] + logits_shapes = [ + [3, 3], # batch_size, num_classes + [None, 3], + [3, None], + [None, None] + ] for batch_size, num_classes in logits_shapes: with self.test_session(): - logits = array_ops.placeholder(dtypes.float32, shape=(batch_size, num_classes)) + logits = array_ops.placeholder( + dtypes.float32, shape=(batch_size, num_classes)) labels = array_ops.placeholder(dtypes.int32, shape=(batch_size,)) loss = losses.sparse_multiclass_hinge_loss(labels, logits) result = loss.eval(feed_dict={logits: logits_np, labels: labels_np}) diff --git a/tensorflow/contrib/kfac/examples/convnet.py b/tensorflow/contrib/kfac/examples/convnet.py index 558bc294bc8ac129b3055ed46623c78a0d5a33e3..39d80addaac1fe855a37255b32bf4412b99df46a 100644 --- a/tensorflow/contrib/kfac/examples/convnet.py +++ b/tensorflow/contrib/kfac/examples/convnet.py @@ -286,7 +286,7 @@ def minimize_loss_distributed(task_id, num_worker_tasks, num_ps_tasks, master, damping=0.001, layer_collection=layer_collection, momentum=0.9) - inv_update_queue = oq.OpQueue(optimizer.inv_updates_dict.values()) + inv_update_queue = oq.OpQueue(optimizer.inv_update_ops) sync_optimizer = tf.train.SyncReplicasOptimizer( opt=optimizer, replicas_to_aggregate=_num_gradient_tasks(num_worker_tasks)) diff --git a/tensorflow/contrib/kfac/examples/mlp.py b/tensorflow/contrib/kfac/examples/mlp.py index 4275ceadc210ff471109b596e1c9aa260ce31ab5..87eed03888c894a04c0521d1ce5ee8975b60776b 100644 --- a/tensorflow/contrib/kfac/examples/mlp.py +++ b/tensorflow/contrib/kfac/examples/mlp.py @@ -239,3 +239,88 @@ def train_mnist_multitower(data_dir, }) return minimize( loss, accuracy, layer_collection, session_config=session_config) + + +def train_mnist_estimator(data_dir, num_epochs, use_fake_data=False): + """Train an MLP on MNIST using tf.estimator. + + Args: + data_dir: string. Directory to read MNIST examples from. + num_epochs: int. Number of passes to make over the training set. + use_fake_data: bool. If True, generate a synthetic dataset. + + Returns: + accuracy of model on the final minibatch of training data. + """ + + # Load a dataset. + def input_fn(): + tf.logging.info("Loading MNIST into memory.") + return mnist.load_mnist( + data_dir, + num_epochs=num_epochs, + batch_size=64, + flatten_images=True, + use_fake_data=use_fake_data) + + def model_fn(features, labels, mode, params): + """Model function for MLP trained with K-FAC. + + Args: + features: Tensor of shape [batch_size, input_size]. Input features. + labels: Tensor of shape [batch_size]. Target labels for training. + mode: tf.estimator.ModeKey. Must be TRAIN. + params: ignored. + + Returns: + EstimatorSpec for training. + + Raises: + ValueError: If 'mode' is anything other than TRAIN. + """ + del params + + if mode != tf.estimator.ModeKeys.TRAIN: + raise ValueError("Only training is supposed with this API.") + + # Build a ConvNet. + layer_collection = lc.LayerCollection() + loss, accuracy = build_model( + features, labels, num_labels=10, layer_collection=layer_collection) + + # Train with K-FAC. + global_step = tf.train.get_or_create_global_step() + optimizer = opt.KfacOptimizer( + learning_rate=tf.train.exponential_decay( + 0.00002, global_step, 10000, 0.5, staircase=True), + cov_ema_decay=0.95, + damping=0.0001, + layer_collection=layer_collection, + momentum=0.99) + + # Run cov_update_op every step. Run 1 inv_update_ops per step. + cov_update_op = optimizer.cov_update_op + inv_update_op = tf.group( + tf.contrib.kfac.utils.batch_execute( + global_step, optimizer.inv_update_thunks, batch_size=1)) + with tf.control_dependencies([cov_update_op, inv_update_op]): + train_op = optimizer.minimize(loss, global_step=global_step) + + # Print metrics every 5 sec. + hooks = [ + tf.train.LoggingTensorHook( + { + "loss": loss, + "accuracy": accuracy + }, every_n_secs=5), + ] + return tf.estimator.EstimatorSpec( + mode=mode, loss=loss, train_op=train_op, training_hooks=hooks) + + run_config = tf.estimator.RunConfig( + model_dir="/tmp/mnist", save_checkpoints_steps=1, keep_checkpoint_max=100) + + # Train until input_fn() is empty with Estimator. This is a prerequisite for + # TPU compatibility. + estimator = tf.estimator.Estimator(model_fn=model_fn, config=run_config) + estimator.train(input_fn=input_fn) diff --git a/tensorflow/contrib/kfac/examples/mlp_mnist_main.py b/tensorflow/contrib/kfac/examples/mlp_mnist_main.py index b318c71a568be2d717745579df24134ceb3b6a0b..9c34ade1d2018135b3636fddb9dcc65839cd59de 100644 --- a/tensorflow/contrib/kfac/examples/mlp_mnist_main.py +++ b/tensorflow/contrib/kfac/examples/mlp_mnist_main.py @@ -33,7 +33,11 @@ FLAGS = None def main(argv): _ = argv - if FLAGS.num_towers > 1: + if FLAGS.use_estimator: + if FLAGS.num_towers != 1: + raise ValueError("Only 1 device supported in tf.estimator example.") + mlp.train_mnist_estimator(FLAGS.data_dir, num_epochs=200) + elif FLAGS.num_towers > 1: mlp.train_mnist_multitower( FLAGS.data_dir, num_epochs=200, num_towers=FLAGS.num_towers) else: @@ -52,5 +56,9 @@ if __name__ == "__main__": type=int, default=1, help="Number of CPUs to split minibatch across.") + parser.add_argument( + "--use_estimator", + action="store_true", + help="Use tf.estimator API to train.") FLAGS, unparsed = parser.parse_known_args() tf.app.run(main=main, argv=[sys.argv[0]] + unparsed) diff --git a/tensorflow/contrib/kfac/examples/mnist.py b/tensorflow/contrib/kfac/examples/mnist.py index cf92c909f4b5201bc0ffda5703136f46c7058ec6..547c4ab25d589192f2a5b65987be3b05128fe298 100644 --- a/tensorflow/contrib/kfac/examples/mnist.py +++ b/tensorflow/contrib/kfac/examples/mnist.py @@ -63,7 +63,7 @@ def load_mnist(data_dir, images = mnist_data.train.images labels = mnist_data.train.labels - dataset = tf.contrib.data.Dataset.from_tensor_slices((np.asarray( + dataset = tf.data.Dataset.from_tensor_slices((np.asarray( images, dtype=np.float32), np.asarray(labels, dtype=np.int64))) return (dataset.repeat(num_epochs).shuffle(num_examples).batch(batch_size) .make_one_shot_iterator().get_next()) diff --git a/tensorflow/contrib/kfac/examples/tests/convnet_test.py b/tensorflow/contrib/kfac/examples/tests/convnet_test.py index 3c98c54ef6cbd527aa0035e0b6f40be961c6308d..8d86c2bb5150cd4bc8a2b21ba050e904929e0fe9 100644 --- a/tensorflow/contrib/kfac/examples/tests/convnet_test.py +++ b/tensorflow/contrib/kfac/examples/tests/convnet_test.py @@ -96,7 +96,7 @@ class ConvNetTest(tf.test.TestCase): """ x = np.asarray([[1.], [2.]]).astype(np.float32) y = np.asarray([1., 2.]).astype(np.float32) - x, y = (tf.contrib.data.Dataset.from_tensor_slices((x, y)) + x, y = (tf.data.Dataset.from_tensor_slices((x, y)) .repeat(100).batch(2).make_one_shot_iterator().get_next()) w = tf.get_variable("w", shape=[1, 1], initializer=tf.zeros_initializer()) y_hat = tf.matmul(x, w) diff --git a/tensorflow/contrib/kfac/examples/tests/mlp_test.py b/tensorflow/contrib/kfac/examples/tests/mlp_test.py index 34a942d27f64e2583c686c2ba3240bc636ed918b..22da6c29f1b364d94432315988d844db9b95ec28 100644 --- a/tensorflow/contrib/kfac/examples/tests/mlp_test.py +++ b/tensorflow/contrib/kfac/examples/tests/mlp_test.py @@ -53,6 +53,11 @@ class MlpTest(tf.test.TestCase): mlp.train_mnist_multitower( data_dir=None, num_epochs=1, num_towers=2, use_fake_data=True) + def testTrainMnistEstimator(self): + with tf.Graph().as_default(): + # Ensure model training doesn't crash. + mlp.train_mnist_estimator(data_dir=None, num_epochs=1, use_fake_data=True) + if __name__ == "__main__": tf.test.main() diff --git a/tensorflow/contrib/kfac/python/kernel_tests/BUILD b/tensorflow/contrib/kfac/python/kernel_tests/BUILD index 4928bf2c10e5063cec6ce8d7374748239079998c..f4ed978174a9ddd8b54a88e60bfb48a67a2e76d2 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/BUILD +++ b/tensorflow/contrib/kfac/python/kernel_tests/BUILD @@ -17,12 +17,17 @@ py_test( "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", + "//tensorflow/python:control_flow_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:init_ops", + "//tensorflow/python:linalg_ops", "//tensorflow/python:math_ops", "//tensorflow/python:random_ops", + "//tensorflow/python:training", "//tensorflow/python:variable_scope", + "//tensorflow/python:variables", + "//third_party/py/numpy", ], ) @@ -117,6 +122,8 @@ py_test( "//tensorflow/python:framework_ops", "//tensorflow/python:linalg_ops", "//tensorflow/python:random_seed", + "//tensorflow/python:variable_scope", + "//tensorflow/python:variables", "//third_party/py/numpy", ], ) diff --git a/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py b/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py index 9b28c45c7263208d21b1514ae5f05b7e81e315a3..bfdb69ad02caaa57827e0ae6b3c9fc0d0ed03754 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/estimator_test.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numpy as np + from tensorflow.contrib.kfac.python.ops import estimator from tensorflow.contrib.kfac.python.ops import layer_collection as lc from tensorflow.contrib.kfac.python.ops import utils @@ -25,11 +27,15 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops +from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables from tensorflow.python.platform import test +from tensorflow.python.training import training_util _ALL_ESTIMATION_MODES = ["gradients", "empirical", "curvature_prop", "exact"] @@ -119,6 +125,114 @@ class EstimatorTest(test.TestCase): estimator.FisherEstimator([self.weights], 0.1, 0.2, self.layer_collection, mode) + def test_cov_update_thunks(self): + """Ensures covariance update ops run once per global_step.""" + with self._graph.as_default(), self.test_session() as sess: + fisher_estimator = estimator.FisherEstimator( + variables=[self.weights], + layer_collection=self.layer_collection, + cov_ema_decay=0.0, + damping=0.0) + + # Construct an op that executes one covariance update per step. + global_step = training_util.get_or_create_global_step() + cov_matrices = [ + fisher_factor.get_cov() + for fisher_factor in self.layer_collection.get_factors() + ] + cov_update_op_thunks = fisher_estimator.cov_update_thunks + cov_update_op = control_flow_ops.case( + [(math_ops.equal(global_step, i), thunk) + for i, thunk in enumerate(cov_update_op_thunks)]) + increment_global_step = global_step.assign_add(1) + + sess.run(variables.global_variables_initializer()) + initial_cov_values = sess.run(cov_matrices) + + # Ensure there's one update per covariance matrix. + self.assertEqual(len(cov_matrices), len(cov_update_op_thunks)) + + # Test is no-op if only 1 covariance matrix. + assert len(cov_matrices) > 1 + + for i in range(len(cov_matrices)): + # Compare new and old covariance values + new_cov_values = sess.run(cov_matrices) + is_cov_equal = [ + np.allclose(initial_cov_value, new_cov_value) + for (initial_cov_value, + new_cov_value) in zip(initial_cov_values, new_cov_values) + ] + num_cov_equal = sum(is_cov_equal) + + # Ensure exactly one covariance matrix changes per step. + self.assertEqual(num_cov_equal, len(cov_matrices) - i) + + # Run all covariance update ops. + sess.run(cov_update_op) + sess.run(increment_global_step) + + def test_inv_update_thunks(self): + """Ensures inverse update ops run once per global_step.""" + with self._graph.as_default(), self.test_session() as sess: + fisher_estimator = estimator.FisherEstimator( + variables=[self.weights], + layer_collection=self.layer_collection, + cov_ema_decay=0.0, + damping=0.0) + + # Construct op that updates one inverse per global step. + global_step = training_util.get_or_create_global_step() + inv_matrices = [ + matrix + for fisher_factor in self.layer_collection.get_factors() + for matrix in fisher_factor._inverses_by_damping.values() + ] + inv_update_op_thunks = fisher_estimator.inv_update_thunks + inv_update_op = control_flow_ops.case( + [(math_ops.equal(global_step, i), thunk) + for i, thunk in enumerate(inv_update_op_thunks)]) + increment_global_step = global_step.assign_add(1) + + sess.run(variables.global_variables_initializer()) + initial_inv_values = sess.run(inv_matrices) + + # Ensure there's one update per inverse matrix. This is true as long as + # there's no fan-in/fan-out or parameter re-use. + self.assertEqual(len(inv_matrices), len(inv_update_op_thunks)) + + # Test is no-op if only 1 invariance matrix. + assert len(inv_matrices) > 1 + + # Assign each covariance matrix a value other than the identity. This + # ensures that the inverse matrices are updated to something different as + # well. + cov_matrices = [ + fisher_factor.get_cov() + for fisher_factor in self.layer_collection.get_factors() + ] + sess.run([ + cov_matrix.assign(2 * linalg_ops.eye(int(cov_matrix.shape[0]))) + for cov_matrix in cov_matrices + ]) + + for i in range(len(inv_matrices)): + # Compare new and old inverse values + new_inv_values = sess.run(inv_matrices) + is_inv_equal = [ + np.allclose(initial_inv_value, new_inv_value) + for (initial_inv_value, + new_inv_value) in zip(initial_inv_values, new_inv_values) + ] + num_inv_equal = sum(is_inv_equal) + + # Ensure exactly one inverse matrix changes per step. + self.assertEqual(num_inv_equal, len(inv_matrices) - i) + + # Run all inverse update ops. + sess.run(inv_update_op) + sess.run(increment_global_step) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/kfac/python/kernel_tests/fisher_blocks_test.py b/tensorflow/contrib/kfac/python/kernel_tests/fisher_blocks_test.py index 2d9b28185ce0db32d5cd7d84737fdf96e2c98851..fb4b3a241c1e9fd82e7bf630fd57295917048fbd 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/fisher_blocks_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/fisher_blocks_test.py @@ -26,6 +26,7 @@ from tensorflow.contrib.kfac.python.ops import utils from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed from tensorflow.python.ops import array_ops +from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import state_ops @@ -49,7 +50,7 @@ class UtilsTest(test.TestCase): right_factor = array_ops.ones([2., 2.]) # pi is the sqrt of the left trace norm divided by the right trace norm - pi = fb._compute_pi_tracenorm(left_factor, right_factor) + pi = fb.compute_pi_tracenorm(left_factor, right_factor) pi_val = sess.run(pi) self.assertEqual(1., pi_val) @@ -236,10 +237,10 @@ class NaiveDiagonalFBTest(test.TestCase): self.assertAllClose(output_flat, explicit) -class FullyConnectedDiagonalFB(test.TestCase): +class FullyConnectedDiagonalFBTest(test.TestCase): def setUp(self): - super(FullyConnectedDiagonalFB, self).setUp() + super(FullyConnectedDiagonalFBTest, self).setUp() self.batch_size = 4 self.input_size = 6 @@ -375,6 +376,65 @@ class FullyConnectedDiagonalFB(test.TestCase): return multiply_result, multiply_inverse_result +class EmbeddingKFACFBTest(test.TestCase): + + def testInstantiateFactors(self): + with ops.Graph().as_default(): + random_seed.set_random_seed(200) + + # Create a Fisher Block. + vocab_size = 5 + block = fb.EmbeddingKFACFB(lc.LayerCollection(), vocab_size) + + # Add some examples. + inputs = array_ops.constant([[0, 1], [1, 2], [2, 3]]) + outputs = array_ops.constant([[0.], [1.], [2.]]) + block.register_additional_minibatch(inputs, outputs) + + # Instantiate factor's variables. Ensure it doesn't fail. + grads = outputs**2. + damping = array_ops.constant(0.) + block.instantiate_factors(([grads],), damping) + + def testMultiplyInverse(self): + with ops.Graph().as_default(), self.test_session() as sess: + random_seed.set_random_seed(200) + + # Create a Fisher Block. + vocab_size = 5 + block = fb.EmbeddingKFACFB(lc.LayerCollection(), vocab_size) + + # Add some examples. + inputs = array_ops.constant([[0, 1], [1, 2], [2, 3]]) + outputs = array_ops.constant([[0.], [1.], [2.]]) + block.register_additional_minibatch(inputs, outputs) + + # Instantiate factor's variables. Ensure it doesn't fail. + grads = outputs**2. + damping = array_ops.constant(0.) + block.instantiate_factors(([grads],), damping) + + # Create a sparse update. + indices = array_ops.constant([1, 3, 4]) + values = array_ops.constant([[1.], [1.], [1.]]) + sparse_vector = ops.IndexedSlices( + values, indices, dense_shape=[vocab_size, 1]) + dense_vector = array_ops.reshape([0., 1., 0., 1., 1.], [vocab_size, 1]) + + # Compare Fisher-vector product against explicit result. + result = block.multiply_inverse(sparse_vector) + expected_result = linalg_ops.matrix_solve(block.full_fisher_block(), + dense_vector) + + sess.run(tf_variables.global_variables_initializer()) + self.assertAlmostEqual( + sess.run(expected_result[1]), sess.run(result.values[0])) + self.assertAlmostEqual( + sess.run(expected_result[3]), sess.run(result.values[1])) + self.assertAlmostEqual( + sess.run(expected_result[4]), sess.run(result.values[2])) + + class FullyConnectedKFACBasicFBTest(test.TestCase): def testFullyConnectedKFACBasicFBInit(self): diff --git a/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py b/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py index 70e56db055078bd4399b03e4d4a877e34249cc5e..66e18974abfadaad5d7a20b40d0b1352bfda67ee 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/fisher_factors_test.py @@ -35,18 +35,27 @@ from tensorflow.python.platform import test class MaybeColocateTest(test.TestCase): + def setUp(self): + self._colocate_cov_ops_with_inputs = ff.COLOCATE_COV_OPS_WITH_INPUTS + + def tearDown(self): + ff.set_global_constants( + colocate_cov_ops_with_inputs=self._colocate_cov_ops_with_inputs) + def testFalse(self): + ff.set_global_constants(colocate_cov_ops_with_inputs=False) with tf_ops.Graph().as_default(): a = constant_op.constant([2.0], name='a') - with ff._maybe_colocate_with(a, False): + with ff.maybe_colocate_with(a): b = constant_op.constant(3.0, name='b') self.assertEqual([b'loc:@a'], a.op.colocation_groups()) self.assertEqual([b'loc:@b'], b.op.colocation_groups()) def testTrue(self): + ff.set_global_constants(colocate_cov_ops_with_inputs=True) with tf_ops.Graph().as_default(): a = constant_op.constant([2.0], name='a') - with ff._maybe_colocate_with(a, True): + with ff.maybe_colocate_with(a): b = constant_op.constant(3.0, name='b') self.assertEqual([b'loc:@a'], a.op.colocation_groups()) self.assertEqual([b'loc:@a'], b.op.colocation_groups()) @@ -80,6 +89,21 @@ class FisherFactorTestingDummy(ff.FisherFactor): def make_inverse_update_ops(self): return [] + def get_cov(self): + return NotImplementedError + + def left_multiply(self, x, damping): + return NotImplementedError + + def right_multiply(self, x, damping): + return NotImplementedError + + def left_multiply_inverse(self, x, damping): + return NotImplementedError + + def right_multiply_inverse(self, x, damping): + return NotImplementedError + class InverseProvidingFactorTestingDummy(ff.InverseProvidingFactor): """Dummy class to test the non-abstract methods on ff.InverseProvidingFactor. @@ -120,7 +144,7 @@ class NumericalUtilsTest(test.TestCase): random_seed.set_random_seed(200) x = npr.randn(100, 3) - cov = ff._compute_cov(array_ops.constant(x)) + cov = ff.compute_cov(array_ops.constant(x)) np_cov = np.dot(x.T, x) / x.shape[0] self.assertAllClose(sess.run(cov), np_cov) @@ -132,7 +156,7 @@ class NumericalUtilsTest(test.TestCase): normalizer = 10. x = npr.randn(100, 3) - cov = ff._compute_cov(array_ops.constant(x), normalizer=normalizer) + cov = ff.compute_cov(array_ops.constant(x), normalizer=normalizer) np_cov = np.dot(x.T, x) / normalizer self.assertAllClose(sess.run(cov), np_cov) @@ -143,7 +167,7 @@ class NumericalUtilsTest(test.TestCase): m, n = 3, 4 a = npr.randn(m, n) - a_homog = ff._append_homog(array_ops.constant(a)) + a_homog = ff.append_homog(array_ops.constant(a)) np_result = np.hstack([a, np.ones((m, 1))]) self.assertAllClose(sess.run(a_homog), np_result) @@ -370,7 +394,7 @@ class NaiveDiagonalFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.ones((2, 3), name='a/b/c') factor = ff.NaiveDiagonalFactor((tensor,), 32) - self.assertEqual([6, 1], factor.get_cov().get_shape().as_list()) + self.assertEqual([6, 1], factor.get_cov_var().get_shape().as_list()) def testNaiveDiagonalFactorInitFloat64(self): with tf_ops.Graph().as_default(): @@ -378,7 +402,7 @@ class NaiveDiagonalFactorTest(test.TestCase): random_seed.set_random_seed(200) tensor = array_ops.ones((2, 3), dtype=dtype, name='a/b/c') factor = ff.NaiveDiagonalFactor((tensor,), 32) - cov = factor.get_cov() + cov = factor.get_cov_var() self.assertEqual(cov.dtype, dtype) self.assertEqual([6, 1], cov.get_shape().as_list()) @@ -393,6 +417,29 @@ class NaiveDiagonalFactorTest(test.TestCase): self.assertAllClose([[0.75], [1.5]], new_cov) +class EmbeddingInputKroneckerFactorTest(test.TestCase): + + def testInitialization(self): + with tf_ops.Graph().as_default(): + input_ids = array_ops.constant([[0], [1], [4]]) + vocab_size = 5 + factor = ff.EmbeddingInputKroneckerFactor((input_ids,), vocab_size) + cov = factor.get_cov_var() + self.assertEqual(cov.shape.as_list(), [vocab_size]) + + def testCovarianceUpdateOp(self): + with tf_ops.Graph().as_default(): + input_ids = array_ops.constant([[0], [1], [4]]) + vocab_size = 5 + factor = ff.EmbeddingInputKroneckerFactor((input_ids,), vocab_size) + cov_update_op = factor.make_covariance_update_op(0.0) + + with self.test_session() as sess: + sess.run(tf_variables.global_variables_initializer()) + new_cov = sess.run(cov_update_op) + self.assertAllClose(np.array([1., 1., 0., 0., 1.]) / 3., new_cov) + + class FullyConnectedKroneckerFactorTest(test.TestCase): def _testFullyConnectedKroneckerFactorInit(self, diff --git a/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py b/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py index 39ce3e9337157c8206107bc40c489e44019743ab..ae787b6f1ac90218f2ac73d37fb270df0b822de2 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/loss_functions_test.py @@ -113,6 +113,113 @@ class CategoricalLogitsNegativeLogProbLossTest(test.TestCase): self.assertListEqual(loss.input_minibatches, tower_logits) self.assertEqual(loss.num_registered_minibatches, num_towers) + def testMultiplyFisherSingleVector(self): + with ops.Graph().as_default(), self.test_session() as sess: + logits = np.array([1., 2., 3.]) + loss = loss_functions.CategoricalLogitsNegativeLogProbLoss(logits) + + # the LossFunction.multiply_fisher docstring only says it supports the + # case where the vector is the same shape as the input natural parameters + # (i.e. the logits here), but here we also test leading dimensions + vector = np.array([1., 2., 3.]) + vectors = [vector, vector.reshape(1, -1), np.stack([vector] * 4)] + + probs = np.exp(logits - np.logaddexp.reduce(logits)) + fisher = np.diag(probs) - np.outer(probs, probs) + + for vector in vectors: + result = loss.multiply_fisher(vector) + expected_result = np.dot(vector, fisher) + self.assertAllClose(expected_result, sess.run(result)) + + def testMultiplyFisherBatch(self): + with ops.Graph().as_default(), self.test_session() as sess: + logits = np.array([[1., 2., 3.], [4., 6., 8.]]) + loss = loss_functions.CategoricalLogitsNegativeLogProbLoss(logits) + + vector = np.array([[1., 2., 3.], [5., 3., 1.]]) + + na = np.newaxis + probs = np.exp(logits - np.logaddexp.reduce(logits, axis=-1, + keepdims=True)) + fishers = probs[..., na] * np.eye(3) - probs[..., na] * probs[..., na, :] + + result = loss.multiply_fisher(vector) + expected_result = np.matmul(vector[..., na, :], fishers)[..., 0, :] + self.assertEqual(sess.run(result).shape, logits.shape) + self.assertAllClose(expected_result, sess.run(result)) + + +class OnehotCategoricalLogitsNegativeLogProbLossTest(test.TestCase): + + def testSample(self): + """Ensure samples can be drawn.""" + with ops.Graph().as_default(), self.test_session() as sess: + logits = np.asarray([ + [0., 0., 0.], # + [1., -1., 0.] + ]).astype(np.float32) + loss = loss_functions.OnehotCategoricalLogitsNegativeLogProbLoss( + array_ops.constant(logits)) + sample = loss.sample(42) + sample = sess.run(sample) + self.assertEqual(sample.shape, (2, 3)) + + def testEvaluateOnTargets(self): + """Ensure log probability can be evaluated correctly.""" + with ops.Graph().as_default(), self.test_session() as sess: + logits = np.asarray([ + [0., 0., 0.], # + [1., -1., 0.] + ]).astype(np.float32) + targets = np.asarray([2, 1]).astype(np.int32) + loss = loss_functions.OnehotCategoricalLogitsNegativeLogProbLoss( + array_ops.constant(logits), targets=array_ops.one_hot(targets, 3)) + neg_log_prob = loss.evaluate() + neg_log_prob = sess.run(neg_log_prob) + + # Calculate explicit log probability of targets. + probs = np.exp(logits) / np.sum(np.exp(logits), axis=1, keepdims=True) + log_probs = np.log([ + probs[0, targets[0]], # + probs[1, targets[1]] + ]) + expected_log_prob = np.sum(log_probs) + + self.assertAllClose(neg_log_prob, -expected_log_prob) + + def testEvaluateOnSample(self): + """Ensure log probability of a sample can be drawn.""" + with ops.Graph().as_default(), self.test_session() as sess: + logits = np.asarray([ + [0., 0., 0.], # + [1., -1., 0.] + ]).astype(np.float32) + loss = loss_functions.OnehotCategoricalLogitsNegativeLogProbLoss( + array_ops.constant(logits)) + neg_log_prob = loss.evaluate_on_sample(42) + + # Simply ensure this doesn't crash. As the output is random, it's + # difficult to say if the output is correct or not... + neg_log_prob = sess.run(neg_log_prob) + + def testMultiMinibatchRegistration(self): + """Ensure this loss function supports registering multiple minibatches.""" + with ops.Graph().as_default(): + tower_logits = [] + loss = None + num_towers = 5 + for _ in range(num_towers): + logits = random_ops.random_uniform(shape=[2, 3]) + tower_logits.append(logits) + if loss is None: + loss = loss_functions.OnehotCategoricalLogitsNegativeLogProbLoss( + logits) + else: + loss.register_additional_minibatch(logits) + self.assertListEqual(loss.input_minibatches, tower_logits) + self.assertEqual(loss.num_registered_minibatches, num_towers) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py b/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py index c8631ed89ba58c5176b4b62344663e5d9e330926..97a97adbf5577cd2694d3055acaa59258ad27964 100644 --- a/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py +++ b/tensorflow/contrib/kfac/python/kernel_tests/utils_test.py @@ -30,6 +30,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables from tensorflow.python.platform import test @@ -287,6 +288,43 @@ class UtilsTest(test.TestCase): tensor = array_ops.zeros([], dtype=dtypes.float32) mean = utils.cross_replica_mean(tensor) + def testBatchExecute(self): + """Ensure batch_execute runs in a round-robin fashion.""" + + def increment_var(var): + return lambda: var.assign_add(1) + + with ops.Graph().as_default(), self.test_session() as sess: + i = variable_scope.get_variable('i', initializer=0) + accumulators = [ + variable_scope.get_variable('var%d' % j, initializer=0) + for j in range(3) + ] + thunks = [increment_var(var) for var in accumulators] + increment_accumulators = utils.batch_execute(i, thunks, 2) + increment_i = i.assign_add(1) + + sess.run(variables.global_variables_initializer()) + + # Ensure one op per thunk. + self.assertEqual(3, len(increment_accumulators)) + + # Ensure round-robin execution. + values = [] + for _ in range(5): + sess.run(increment_accumulators) + sess.run(increment_i) + values.append(sess.run(accumulators)) + self.assertAllClose( + [ + [1, 1, 0], # + [2, 1, 1], # + [2, 2, 2], # + [3, 3, 2], # + [4, 3, 3] + ], + values) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/kfac/python/ops/BUILD b/tensorflow/contrib/kfac/python/ops/BUILD index 9be3d60dc06e0282370fed3882ad994f7e4bc64c..ee6549b109399766579b6ea18a987ae2c8275983 100644 --- a/tensorflow/contrib/kfac/python/ops/BUILD +++ b/tensorflow/contrib/kfac/python/ops/BUILD @@ -65,6 +65,7 @@ py_library( srcs = ["loss_functions.py"], srcs_version = "PY2AND3", deps = [ + "//tensorflow/contrib/distributions:distributions_py", "//tensorflow/python:array_ops", "//tensorflow/python:math_ops", "//tensorflow/python:tensor_shape", @@ -198,6 +199,7 @@ py_library( deps = [ "//tensorflow/contrib/tpu", "//tensorflow/python:array_ops", + "//tensorflow/python:control_flow_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:gradients", diff --git a/tensorflow/contrib/kfac/python/ops/estimator.py b/tensorflow/contrib/kfac/python/ops/estimator.py index 5e1680967c184bf19f2a2578219db07a48264dc9..a7b1f9d35c931fc44408be804479e758f28f7110 100644 --- a/tensorflow/contrib/kfac/python/ops/estimator.py +++ b/tensorflow/contrib/kfac/python/ops/estimator.py @@ -66,7 +66,21 @@ class _DeviceContextGenerator(object): class FisherEstimator(object): - """Fisher estimator class supporting various approximations of the Fisher.""" + """Fisher estimator class supporting various approximations of the Fisher. + + Attributes: + cov_update_thunks: list of no-arg functions. Executing a function adds + covariance update ops for a single FisherFactor to the graph. + cov_update_ops: List of Ops. Running an op updates covariance matrices for a + single FisherFactor. + cov_update_op: Op. Running updates covariance matrices for all + FisherFactors. + inv_update_thunks: list of no-arg functions. Executing a function adds + inverse update ops for a single FisherFactor to the graph. + inv_update_ops: List of Ops. Running an op updates inverse matrices for a + single FisherFactor. + inv_update_op: Op. Running updates inverse matrices for all FisherFactors. + """ def __init__(self, variables, @@ -74,7 +88,7 @@ class FisherEstimator(object): damping, layer_collection, estimation_mode="gradients", - colocate_gradients_with_ops=False, + colocate_gradients_with_ops=True, cov_devices=None, inv_devices=None): """Create a FisherEstimator object. @@ -110,7 +124,7 @@ class FisherEstimator(object): is more expensive to compute than the other three options by a factor equal to the output dimension, roughly speaking. colocate_gradients_with_ops: Whether we should request gradients be - colocated with their respective ops. + colocated with their respective ops. (Default: True) cov_devices: Iterable of device strings (e.g. '/gpu:0'). Covariance computations will be placed on these devices in a round-robin fashion. Can be None, which means that no devices are specified. @@ -122,6 +136,7 @@ class FisherEstimator(object): ValueError: If no losses have been registered with layer_collection. """ + self._cov_ema_decay = cov_ema_decay self._variables = variables self._damping = damping self._estimation_mode = estimation_mode @@ -135,13 +150,31 @@ class FisherEstimator(object): "exact": self._get_grads_lists_exact } self._colocate_gradients_with_ops = colocate_gradients_with_ops + + # TODO(b/70674513): Factor device placement outside of this class. self._cov_device_context_generator = _DeviceContextGenerator(cov_devices) if inv_devices == cov_devices: self._inv_device_context_generator = self._cov_device_context_generator else: self._inv_device_context_generator = _DeviceContextGenerator(inv_devices) - setup = self._setup(cov_ema_decay) - self.cov_update_op, self.inv_update_op, self.inv_updates_dict = setup + + self._instantiate_factors() + + self.cov_update_thunks = [ + self._create_cov_update_thunk(factor) + for factor in self._layers.get_factors() + ] + self.cov_update_ops = [thunk() for thunk in self.cov_update_thunks] + self.cov_update_op = control_flow_ops.group( + self.cov_update_ops, name="cov_update_op") + + self.inv_update_thunks = [ + self._create_inv_update_thunk(factor) + for factor in self._layers.get_factors() + ] + self.inv_update_ops = [thunk() for thunk in self.inv_update_thunks] + self.inv_update_op = control_flow_ops.group( + self.inv_update_ops, name="inv_update_op") @property def variables(self): @@ -202,18 +235,8 @@ class FisherEstimator(object): return self._apply_transformation(vecs_and_vars, lambda fb, vec: fb.multiply(vec)) - def _setup(self, cov_ema_decay): - """Sets up the various operations. - - Args: - cov_ema_decay: The decay factor used when calculating the covariance - estimate moving averages. - - Returns: - A triple (covs_update_op, invs_update_op, inv_updates_dict), where - covs_update_op is the grouped Op to update all the covariance estimates, - invs_update_op is the grouped Op to update all the inverses, and - inv_updates_dict is a dict mapping Op names to individual inverse updates. + def _instantiate_factors(self): + """Instantiates FisherFactors' variables. Raises: ValueError: If estimation_mode was improperly specified at construction. @@ -238,20 +261,25 @@ class FisherEstimator(object): with self._cov_device_context_generator(): fb.instantiate_factors(grads_list, self.damping) - cov_updates = [ - factor.make_covariance_update_op(cov_ema_decay) - for factor in self._layers.get_factors() - ] - inv_updates = {op.name: op for op in self._get_all_inverse_update_ops()} + def _create_cov_update_thunk(self, factor): + """Constructs a covariance update thunk for a single FisherFactor.""" + + def thunk(): + with tf_ops.name_scope( + "create_cov_update_thunk", values=[self._cov_ema_decay]): + return factor.make_covariance_update_op(self._cov_ema_decay) + + return thunk + + def _create_inv_update_thunk(self, factor): + """Constructs an inverse update thunk for a single FisherFactor.""" - return control_flow_ops.group(*cov_updates), control_flow_ops.group( - *inv_updates.values()), inv_updates + def thunk(): + with tf_ops.name_scope("create_inv_update_thunk"): + with self._inv_device_context_generator(): + return control_flow_ops.group(factor.make_inverse_update_ops()) - def _get_all_inverse_update_ops(self): - for factor in self._layers.get_factors(): - with self._inv_device_context_generator(): - for op in factor.make_inverse_update_ops(): - yield op + return thunk def _get_grads_lists_gradients(self, tensors): grads_flat = gradients_impl.gradients( diff --git a/tensorflow/contrib/kfac/python/ops/fisher_blocks.py b/tensorflow/contrib/kfac/python/ops/fisher_blocks.py index 1ccb9e040f2bb6bcfd217886918abd40e3cc1cfb..cf38d28b43836dced8babe2ffa7853b1c4b1b369 100644 --- a/tensorflow/contrib/kfac/python/ops/fisher_blocks.py +++ b/tensorflow/contrib/kfac/python/ops/fisher_blocks.py @@ -54,7 +54,7 @@ from tensorflow.python.ops import math_ops NORMALIZE_DAMPING_POWER = 1.0 # Methods for adjusting damping for FisherBlocks. See -# _compute_pi_adjusted_damping() for details. +# compute_pi_adjusted_damping() for details. PI_OFF_NAME = "off" PI_TRACENORM_NAME = "tracenorm" PI_TYPE = PI_TRACENORM_NAME @@ -72,7 +72,14 @@ def set_global_constants(normalize_damping_power=None, pi_type=None): PI_TYPE = pi_type -def _compute_pi_tracenorm(left_cov, right_cov): +def normalize_damping(damping, num_replications): + """Normalize damping after adjusting scale by NORMALIZE_DAMPING_POWER.""" + if NORMALIZE_DAMPING_POWER: + return damping / (num_replications ** NORMALIZE_DAMPING_POWER) + return damping + + +def compute_pi_tracenorm(left_cov, right_cov): """Computes the scalar constant pi for Tikhonov regularization/damping. pi = sqrt( (trace(A) / dim(A)) / (trace(B) / dim(B)) ) @@ -85,17 +92,29 @@ def _compute_pi_tracenorm(left_cov, right_cov): Returns: The computed scalar constant pi for these Kronecker Factors (as a Tensor). """ + + def _trace(cov): + if len(cov.shape) == 1: + # Diagonal matrix. + return math_ops.reduce_sum(cov) + elif len(cov.shape) == 2: + # Full matrix. + return math_ops.trace(cov) + else: + raise ValueError( + "What's the trace of a Tensor of rank %d?" % len(cov.shape)) + # Instead of dividing by the dim of the norm, we multiply by the dim of the # other norm. This works out the same in the ratio. - left_norm = math_ops.trace(left_cov) * right_cov.shape.as_list()[0] - right_norm = math_ops.trace(right_cov) * left_cov.shape.as_list()[0] + left_norm = _trace(left_cov) * right_cov.shape.as_list()[0] + right_norm = _trace(right_cov) * left_cov.shape.as_list()[0] return math_ops.sqrt(left_norm / right_norm) -def _compute_pi_adjusted_damping(left_cov, right_cov, damping): +def compute_pi_adjusted_damping(left_cov, right_cov, damping): if PI_TYPE == PI_TRACENORM_NAME: - pi = _compute_pi_tracenorm(left_cov, right_cov) + pi = compute_pi_tracenorm(left_cov, right_cov) return (damping * pi, damping / pi) elif PI_TYPE == PI_OFF_NAME: @@ -194,15 +213,15 @@ class FullFB(FisherBlock): self._factor.register_damped_inverse(damping) def multiply_inverse(self, vector): - inverse = self._factor.get_damped_inverse(self._damping) - out_flat = math_ops.matmul(inverse, utils.tensors_to_column(vector)) + vector_flat = utils.tensors_to_column(vector) + out_flat = self._factor.left_multiply_inverse( + vector_flat, self._damping) return utils.column_to_tensors(vector, out_flat) def multiply(self, vector): vector_flat = utils.tensors_to_column(vector) - out_flat = ( - math_ops.matmul(self._factor.get_cov(), vector_flat) + - self._damping * vector_flat) + out_flat = self._factor.left_multiply( + vector_flat, self._damping) return utils.column_to_tensors(vector, out_flat) def full_fisher_block(self): @@ -258,16 +277,20 @@ class NaiveDiagonalFB(FisherBlock): def multiply_inverse(self, vector): vector_flat = utils.tensors_to_column(vector) - out_flat = vector_flat / (self._factor.get_cov() + self._damping) + print("vector_flat: %s" % vector_flat) + out_flat = self._factor.left_multiply_inverse( + vector_flat, self._damping) + print("out_flat: %s" % out_flat) return utils.column_to_tensors(vector, out_flat) def multiply(self, vector): vector_flat = utils.tensors_to_column(vector) - out_flat = vector_flat * (self._factor.get_cov() + self._damping) + out_flat = self._factor.left_multiply( + vector_flat, self._damping) return utils.column_to_tensors(vector, out_flat) def full_fisher_block(self): - return array_ops.diag(array_ops.reshape(self._factor.get_cov(), (-1,))) + return self._factor.get_cov() def tensors_to_compute_grads(self): return self._params @@ -349,8 +372,9 @@ class FullyConnectedDiagonalFB(FisherBlock): Tensor of the same shape, corresponding to the inverse Fisher-vector product. """ - reshaped_vect = utils.layer_params_to_mat2d(vector) - reshaped_out = reshaped_vect / (self._factor.get_cov() + self._damping) + reshaped_vec = utils.layer_params_to_mat2d(vector) + reshaped_out = self._factor.left_multiply_inverse( + reshaped_vec, self._damping) return utils.mat2d_to_layer_params(vector, reshaped_out) def multiply(self, vector): @@ -365,8 +389,9 @@ class FullyConnectedDiagonalFB(FisherBlock): Returns: Tensor of the same shape, corresponding to the Fisher-vector product. """ - reshaped_vect = utils.layer_params_to_mat2d(vector) - reshaped_out = reshaped_vect * (self._factor.get_cov() + self._damping) + reshaped_vec = utils.layer_params_to_mat2d(vector) + reshaped_out = self._factor.left_multiply( + reshaped_vec, self._damping) return utils.mat2d_to_layer_params(vector, reshaped_out) def tensors_to_compute_grads(self): @@ -451,9 +476,8 @@ class ConvDiagonalFB(FisherBlock): inputs_shape[1] * inputs_shape[2] // (self._strides[1] * self._strides[2])) - if NORMALIZE_DAMPING_POWER: - damping /= self._num_locations**NORMALIZE_DAMPING_POWER - self._damping = damping + self._damping = (self._num_locations + * normalize_damping(damping, self._num_locations)) self._factor = self._layer_collection.make_or_get_factor( fisher_factors.ConvDiagonalFactor, @@ -462,12 +486,14 @@ class ConvDiagonalFB(FisherBlock): def multiply_inverse(self, vector): reshaped_vect = utils.layer_params_to_mat2d(vector) - reshaped_out = reshaped_vect / (self._factor.get_cov() + self._damping) + reshaped_out = self._factor.left_multiply_inverse( + reshaped_vect, self._damping) return utils.mat2d_to_layer_params(vector, reshaped_out) def multiply(self, vector): reshaped_vect = utils.layer_params_to_mat2d(vector) - reshaped_out = reshaped_vect * (self._factor.get_cov() + self._damping) + reshaped_out = self._factor.left_multiply( + reshaped_vect, self._damping) return utils.mat2d_to_layer_params(vector, reshaped_out) def tensors_to_compute_grads(self): @@ -506,7 +532,7 @@ class KroneckerProductFB(FisherBlock): Args: damping: The base damping factor (float or Tensor) for the damped inverse. """ - self._input_damping, self._output_damping = _compute_pi_adjusted_damping( + self._input_damping, self._output_damping = compute_pi_adjusted_damping( self._input_factor.get_cov(), self._output_factor.get_cov(), damping**0.5) @@ -527,28 +553,24 @@ class KroneckerProductFB(FisherBlock): return 1.0 def multiply_inverse(self, vector): - left_factor_inv = self._input_factor.get_damped_inverse(self._input_damping) - right_factor_inv = self._output_factor.get_damped_inverse( - self._output_damping) reshaped_vector = utils.layer_params_to_mat2d(vector) - reshaped_out = math_ops.matmul(left_factor_inv, - math_ops.matmul(reshaped_vector, - right_factor_inv)) + reshaped_out = self._output_factor.right_multiply_inverse( + reshaped_vector, + self._output_damping) + reshaped_out = self._input_factor.left_multiply_inverse( + reshaped_out, self._input_damping) if self._renorm_coeff != 1.0: reshaped_out /= math_ops.cast( self._renorm_coeff, dtype=reshaped_out.dtype) return utils.mat2d_to_layer_params(vector, reshaped_out) def multiply(self, vector): - left_factor = self._input_factor.get_cov() - right_factor = self._output_factor.get_cov() reshaped_vector = utils.layer_params_to_mat2d(vector) - reshaped_out = ( - math_ops.matmul(reshaped_vector, right_factor) + - self._output_damping * reshaped_vector) - reshaped_out = ( - math_ops.matmul(left_factor, reshaped_out) + - self._input_damping * reshaped_out) + reshaped_out = self._output_factor.right_multiply( + reshaped_vector, + self._output_damping) + reshaped_out = self._input_factor.left_multiply( + reshaped_out, self._input_damping) if self._renorm_coeff != 1.0: reshaped_out *= math_ops.cast( self._renorm_coeff, dtype=reshaped_out.dtype) @@ -568,6 +590,74 @@ class KroneckerProductFB(FisherBlock): right_factor) +class EmbeddingKFACFB(KroneckerProductFB): + """K-FAC FisherBlock for embedding layers. + + This FisherBlock is similar to EmbeddingKFACFB, except that its + input factor is approximated by a diagonal matrix. In the case that each + example references exactly one embedding, this approximation is exact. + + Does not support bias parameters. + """ + + def __init__(self, layer_collection, vocab_size): + """Creates a EmbeddingKFACFB block. + + Args: + layer_collection: The collection of all layers in the K-FAC approximate + Fisher information matrix to which this FisherBlock belongs. + vocab_size: int. Size of vocabulary for this embedding layer. + """ + self._inputs = [] + self._outputs = [] + self._vocab_size = vocab_size + + super(EmbeddingKFACFB, self).__init__(layer_collection) + + def instantiate_factors(self, grads_list, damping): + """Instantiate Kronecker Factors for this FisherBlock. + + Args: + grads_list: List of list of Tensors. grads_list[i][j] is the + gradient of the loss with respect to 'outputs' from source 'i' and + tower 'j'. Each Tensor has shape [tower_minibatch_size, output_size]. + damping: 0-D Tensor or float. 'damping' * identity is approximately added + to this FisherBlock's Fisher approximation. + """ + # TODO(b/68033310): Validate which of, + # (1) summing on a single device (as below), or + # (2) on each device in isolation and aggregating + # is faster. + inputs = _concat_along_batch_dim(self._inputs) + grads_list = tuple(_concat_along_batch_dim(grads) for grads in grads_list) + + self._input_factor = self._layer_collection.make_or_get_factor( # + fisher_factors.EmbeddingInputKroneckerFactor, # + ((inputs,), self._vocab_size)) + self._output_factor = self._layer_collection.make_or_get_factor( # + fisher_factors.FullyConnectedKroneckerFactor, # + (grads_list,)) + self._register_damped_input_and_output_inverses(damping) + + def tensors_to_compute_grads(self): + return self._outputs + + def register_additional_minibatch(self, inputs, outputs): + """Registers an additional minibatch to the FisherBlock. + + Args: + inputs: Tensor of shape [batch_size, input_size]. Inputs to the + matrix-multiply. + outputs: Tensor of shape [batch_size, output_size]. Layer preactivations. + """ + self._inputs.append(inputs) + self._outputs.append(outputs) + + @property + def num_registered_minibatches(self): + return len(self._inputs) + + class FullyConnectedKFACBasicFB(KroneckerProductFB): """K-FAC FisherBlock for fully-connected (dense) layers. @@ -691,8 +781,8 @@ class ConvKFCBasicFB(KroneckerProductFB): grads_list = tuple(_concat_along_batch_dim(grads) for grads in grads_list) # Infer number of locations upon which convolution is applied. - self._num_locations = _num_conv_locations(inputs.shape.as_list(), - self._strides) + self._num_locations = num_conv_locations(inputs.shape.as_list(), + self._strides) self._input_factor = self._layer_collection.make_or_get_factor( fisher_factors.ConvInputKroneckerFactor, @@ -701,11 +791,9 @@ class ConvKFCBasicFB(KroneckerProductFB): self._output_factor = self._layer_collection.make_or_get_factor( fisher_factors.ConvOutputKroneckerFactor, (grads_list,)) - if NORMALIZE_DAMPING_POWER: - damping /= self._num_locations**NORMALIZE_DAMPING_POWER - self._damping = damping - + damping = normalize_damping(damping, self._num_locations) self._register_damped_input_and_output_inverses(damping) + self._damping = damping @property def _renorm_coeff(self): @@ -758,8 +846,16 @@ def _concat_along_batch_dim(tensor_list): return array_ops.concat(tensor_list, axis=0) -def _num_conv_locations(input_shape, strides): - """Returns the number of locations a Conv kernel is applied to.""" +def num_conv_locations(input_shape, strides): + """Returns the number of spatial locations a 2D Conv kernel is applied to. + + Args: + input_shape: list representing shape of inputs to the Conv layer. + strides: list representing strides for the Conv kernel. + + Returns: + A scalar |T| denoting the number of spatial locations for the Conv layer. + """ return input_shape[1] * input_shape[2] // (strides[1] * strides[2]) @@ -804,9 +900,7 @@ class FullyConnectedMultiIndepFB(KroneckerProductFB): self._output_factor = self._layer_collection.make_or_get_factor( fisher_factors.FullyConnectedMultiKF, (grads_list,)) - if NORMALIZE_DAMPING_POWER: - damping /= self._num_uses**NORMALIZE_DAMPING_POWER - + damping = normalize_damping(damping, self._num_uses) self._register_damped_input_and_output_inverses(damping) @property @@ -885,10 +979,8 @@ class FullyConnectedSeriesFB(FisherBlock): self._output_factor = self._layer_collection.make_or_get_factor( fisher_factors.FullyConnectedMultiKF, (grads_list,)) - if NORMALIZE_DAMPING_POWER: - damping /= self._num_timesteps**NORMALIZE_DAMPING_POWER - - self._damping_input, self._damping_output = _compute_pi_adjusted_damping( + damping = normalize_damping(damping, self._num_timesteps) + self._damping_input, self._damping_output = compute_pi_adjusted_damping( self._input_factor.get_cov(), self._output_factor.get_cov(), damping**0.5) diff --git a/tensorflow/contrib/kfac/python/ops/fisher_blocks_lib.py b/tensorflow/contrib/kfac/python/ops/fisher_blocks_lib.py index 59389f8d385c18f50914d690cfaa2825ef807ed3..c04cf727fa958160d61c7a3638ec65f6c93c2f24 100644 --- a/tensorflow/contrib/kfac/python/ops/fisher_blocks_lib.py +++ b/tensorflow/contrib/kfac/python/ops/fisher_blocks_lib.py @@ -29,10 +29,17 @@ _allowed_symbols = [ 'NaiveDiagonalFB', 'FullyConnectedDiagonalFB', 'KroneckerProductFB', + 'EmbeddingKFACFB', 'FullyConnectedKFACBasicFB', 'ConvKFCBasicFB', 'ConvDiagonalFB', 'set_global_constants', + 'compute_pi_tracenorm', + 'compute_pi_adjusted_damping', + 'num_conv_locations', + 'normalize_damping', + 'LEFT_MULTIPLY', + 'RIGHT_MULTIPLY', ] remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/kfac/python/ops/fisher_factors.py b/tensorflow/contrib/kfac/python/ops/fisher_factors.py index 2c00fc14d9a772a91075facdfdcc800bcc62eedd..603d8b8b210279ee6d8f1de0ce10869fde23f4d9 100644 --- a/tensorflow/contrib/kfac/python/ops/fisher_factors.py +++ b/tensorflow/contrib/kfac/python/ops/fisher_factors.py @@ -25,6 +25,7 @@ import numpy as np import six from tensorflow.contrib.kfac.python.ops import utils +from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops as tf_ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -52,11 +53,15 @@ EIGENVALUE_DECOMPOSITION_THRESHOLD = 2 # matrix powers. Must be nonnegative. EIGENVALUE_CLIPPING_THRESHOLD = 0.0 +# Colocate the covariance ops and variables with the input tensors for each +# factor. +COLOCATE_COV_OPS_WITH_INPUTS = True + @contextlib.contextmanager -def _maybe_colocate_with(op, colocate_cov_ops_with_inputs): - """Context to colocate with `op` if `colocate_cov_ops_with_inputs`.""" - if colocate_cov_ops_with_inputs: +def maybe_colocate_with(op): + """Context to colocate with `op` if `COLOCATE_COV_OPS_WITH_INPUTS`.""" + if COLOCATE_COV_OPS_WITH_INPUTS: if isinstance(op, (list, tuple)): with tf_ops.colocate_with(op[0]): yield @@ -70,12 +75,14 @@ def _maybe_colocate_with(op, colocate_cov_ops_with_inputs): def set_global_constants(init_covariances_at_zero=None, zero_debias=None, eigenvalue_decomposition_threshold=None, - eigenvalue_clipping_threshold=None): + eigenvalue_clipping_threshold=None, + colocate_cov_ops_with_inputs=None): """Sets various global constants used by the classes in this module.""" global INIT_COVARIANCES_AT_ZERO global ZERO_DEBIAS global EIGENVALUE_DECOMPOSITION_THRESHOLD global EIGENVALUE_CLIPPING_THRESHOLD + global COLOCATE_COV_OPS_WITH_INPUTS if init_covariances_at_zero is not None: INIT_COVARIANCES_AT_ZERO = init_covariances_at_zero @@ -85,6 +92,8 @@ def set_global_constants(init_covariances_at_zero=None, EIGENVALUE_DECOMPOSITION_THRESHOLD = eigenvalue_decomposition_threshold if eigenvalue_clipping_threshold is not None: EIGENVALUE_CLIPPING_THRESHOLD = eigenvalue_clipping_threshold + if colocate_cov_ops_with_inputs is not None: + COLOCATE_COV_OPS_WITH_INPUTS = colocate_cov_ops_with_inputs def inverse_initializer(shape, dtype, partition_info=None): # pylint: disable=unused-argument @@ -103,7 +112,7 @@ def diagonal_covariance_initializer(shape, dtype, partition_info): # pylint: di return array_ops.ones(shape, dtype) -def _compute_cov(tensor, tensor_right=None, normalizer=None): +def compute_cov(tensor, tensor_right=None, normalizer=None): """Compute the empirical second moment of the rows of a 2D Tensor. This function is meant to be applied to random matrices for which the true row @@ -131,7 +140,7 @@ def _compute_cov(tensor, tensor_right=None, normalizer=None): math_ops.cast(normalizer, tensor.dtype)) -def _append_homog(tensor): +def append_homog(tensor): """Appends a homogeneous coordinate to the last dimension of a Tensor. Args: @@ -202,12 +211,21 @@ def scalar_or_tensor_to_string(val): class FisherFactor(object): """Base class for objects modeling factors of approximate Fisher blocks. - Note that for blocks that aren't based on approximations, a 'factor' can - be the entire block itself, as is the case for the diagonal and full - representations. + A FisherFactor represents part of an approximate Fisher Information matrix. + For example, one approximation to the Fisher uses the Kronecker product of two + FisherFactors A and B, F = kron(A, B). FisherFactors are composed with + FisherBlocks to construct a block-diagonal approximation to the full Fisher. + + FisherFactors are backed by a single, non-trainable variable that is updated + by running FisherFactor.make_covariance_update_op(). The shape and type of + this variable is implementation specific. - Subclasses must implement the _compute_new_cov method, and the _var_scope - and _cov_shape properties. + Note that for blocks that aren't based on approximations, a 'factor' can + be the entire block itself, as is the case for the diagonal and full + representations. + + Subclasses must implement the _compute_new_cov() method, and the _var_scope + and _cov_shape properties. """ def __init__(self): @@ -215,16 +233,21 @@ class FisherFactor(object): @abc.abstractproperty def _var_scope(self): + """Variable scope for this FisherFactor instance. + + Returns: + string that unique identifies this FisherFactor instance. + """ pass @abc.abstractproperty def _cov_shape(self): - """The shape of the cov matrix.""" + """The shape of the variable backing this FisherFactor.""" pass @abc.abstractproperty def _num_sources(self): - """The number of things to sum over when computing cov. + """The number of things to sum over when updating covariance variable. The default make_covariance_update_op function will call _compute_new_cov with indices ranging from 0 to _num_sources-1. The typical situation is @@ -236,10 +259,12 @@ class FisherFactor(object): @abc.abstractproperty def _dtype(self): + """dtype for variable backing this factor.""" pass @property def _cov_initializer(self): + """Function for initializing covariance variable.""" return covariance_initializer def instantiate_covariance(self): @@ -254,6 +279,15 @@ class FisherFactor(object): @abc.abstractmethod def _compute_new_cov(self, idx=0): + """Computes minibatch-estimated covariance for a single source. + + Args: + idx: int in [0, self._num_sources). Which source to use when estimating + covariance. + + Returns: + Tensor of same shape as self.get_cov_var(). + """ pass def make_covariance_update_op(self, ema_decay): @@ -264,29 +298,123 @@ class FisherFactor(object): Returns: An Op for updating the covariance Variable referenced by _cov. """ - new_cov = math_ops.add_n( - tuple(self._compute_new_cov(idx) for idx in range(self._num_sources))) - - # Synchronize value across all TPU cores. - if utils.on_tpu(): - new_cov = utils.cross_replica_mean(new_cov) - - return moving_averages.assign_moving_average( - self._cov, new_cov, ema_decay, zero_debias=ZERO_DEBIAS) + new_cov_contribs = tuple(self._compute_new_cov(idx) + for idx in range(self._num_sources)) + # This gets the job done but we might want a better solution in the future. + # In particular, we could have a separate way of specifying where the + # the cov variables finally end up, independent of where their various + # contributions are computed. Right now these are the same thing, but in + # the future we might want to perform the cov computations on each tower, + # so that each tower will be considered a "source" (allowing us to reuse + # the existing "source" code for this). + with maybe_colocate_with(new_cov_contribs[0]): + new_cov = math_ops.add_n(new_cov_contribs) + # Synchronize value across all TPU cores. + if utils.on_tpu(): + new_cov = utils.cross_replica_mean(new_cov) + return moving_averages.assign_moving_average( + self._cov, new_cov, ema_decay, zero_debias=ZERO_DEBIAS) @abc.abstractmethod def make_inverse_update_ops(self): """Create and return update ops corresponding to registered computations.""" pass + @abc.abstractmethod def get_cov(self): + """Get full covariance matrix. + + Returns: + Tensor of shape [n, n]. Represents all parameter-parameter correlations + captured by this FisherFactor. + """ + pass + + def get_cov_var(self): + """Get variable backing this FisherFactor. + + May or may not be the same as self.get_cov() + + Returns: + Variable of shape self._cov_shape. + """ return self._cov + @abc.abstractmethod + def left_multiply(self, x, damping): + """Multiplies 'x' by the damped covariance of this factor. + + Let C be the covariance matrix this factor represents, and + D = C + damping * I be its damped variant. This method calculates + matmul(D, vec(x)). + + Args: + x: Tensor. Represents a single vector. Shape depends on implementation. + damping: 0-D Tensor. Damping to add to C's diagonal. + + Returns: + Tensor of same shape as 'x'. + """ + pass + + @abc.abstractmethod + def right_multiply(self, x, damping): + """Multiplies 'x' by the damped covariance of this factor. + + Let C be the covariance matrix this factor represents, and + D = C + damping * I be its damped variant. This method calculates + matmul(vec(x), D). + + Args: + x: Tensor. Represents a single vector. Shape depends on implementation. + damping: 0-D Tensor. Damping to add to C's diagonal. + + Returns: + Tensor of same shape as 'x'. + """ + pass + + @abc.abstractmethod + def left_multiply_inverse(self, x, damping): + """Multiplies 'x' by damped inverse of this factor. + + Let C be the covariance matrix this factor represents and + E = inv(C + damping * I) be its damped inverse. This method calculates + matmul(E, vec(x)). + + Args: + x: Tensor. Represents a single vector. Shape depends on implementation. + damping: 0-D Tensor. Damping to add to C's diagonal. + + Returns: + Tensor of same shape as 'x'. + """ + pass + + @abc.abstractmethod + def right_multiply_inverse(self, x, damping): + """Multiplies 'x' by damped inverse of this factor. + + Let C be the covariance matrix this factor represents and + E = inv(C + damping * I) be its damped inverse. This method calculates + matmul(vec(x), E). + + Args: + x: Tensor. Represents a single vector. Shape depends on implementation. + damping: 0-D Tensor. Damping to add to C's diagonal. + + Returns: + Tensor of same shape as 'x'. + """ + pass + class InverseProvidingFactor(FisherFactor): - """Base class for FisherFactors that maintain inverses, powers, etc of _cov. + """Base class for FisherFactors that maintain inverses explicitly. - Assumes that the _cov property is a square PSD matrix. + This class explicitly calculates and stores inverses of covariance matrices + provided by the underlying FisherFactor implementation. It is assumed that + vectors can be represented as 2-D matrices. Subclasses must implement the _compute_new_cov method, and the _var_scope and _cov_shape properties. @@ -350,28 +478,16 @@ class InverseProvidingFactor(FisherFactor): dtype=self._dtype) self._matpower_by_exp_and_damping[(exp, damping)] = matpower - def register_eigendecomp(self): - """Registers an eigendecomposition. - - Unlike register_damp_inverse and register_matpower this doesn't create - any variables or inverse ops. Instead it merely makes tensors containing - the eigendecomposition available to anyone that wants them. They will be - recomputed (once) for each session.run() call (when they needed by some op). - """ - if not self._eigendecomp: - eigenvalues, eigenvectors = linalg_ops.self_adjoint_eig(self._cov) - - # The matrix self._cov is positive semidefinite by construction, but the - # numerical eigenvalues could be negative due to numerical errors, so here - # we clip them to be at least FLAGS.eigenvalue_clipping_threshold - clipped_eigenvalues = math_ops.maximum(eigenvalues, - EIGENVALUE_CLIPPING_THRESHOLD) - self._eigendecomp = (clipped_eigenvalues, eigenvectors) - def make_inverse_update_ops(self): """Create and return update ops corresponding to registered computations.""" ops = [] + # We do this to ensure that we don't reuse the eigendecomp from old calls + # to make_inverse_update_ops that may be placed on different devices. This + # can happen is the user has both a permanent and lazily constructed + # version of the inverse ops (and only uses one of them). + self.reset_eigendecomp() + num_inverses = len(self._inverses_by_damping) matrix_power_registered = bool(self._matpower_by_exp_and_damping) use_eig = ( @@ -379,8 +495,7 @@ class InverseProvidingFactor(FisherFactor): num_inverses >= EIGENVALUE_DECOMPOSITION_THRESHOLD) if use_eig: - self.register_eigendecomp() # ensures self._eigendecomp is set - eigenvalues, eigenvectors = self._eigendecomp # pylint: disable=unpacking-non-sequence + eigenvalues, eigenvectors = self.get_eigendecomp() # pylint: disable=unpacking-non-sequence for damping, inv in self._inverses_by_damping.items(): ops.append( @@ -415,11 +530,80 @@ class InverseProvidingFactor(FisherFactor): return self._matpower_by_exp_and_damping[(exp, damping)] def get_eigendecomp(self): + """Creates or retrieves eigendecomposition of self._cov.""" # Unlike get_inverse and get_matpower this doesn't retrieve a stored # variable, but instead always computes a fresh version from the current # value of get_cov(). + if not self._eigendecomp: + eigenvalues, eigenvectors = linalg_ops.self_adjoint_eig(self._cov) + + # The matrix self._cov is positive semidefinite by construction, but the + # numerical eigenvalues could be negative due to numerical errors, so here + # we clip them to be at least FLAGS.eigenvalue_clipping_threshold + clipped_eigenvalues = math_ops.maximum(eigenvalues, + EIGENVALUE_CLIPPING_THRESHOLD) + self._eigendecomp = (clipped_eigenvalues, eigenvectors) + return self._eigendecomp + def reset_eigendecomp(self): + self._eigendecomp = None + + def get_cov(self): + # Variable contains full covariance matrix. + return self.get_cov_var() + + def left_multiply(self, x, damping): + n = self.get_cov().shape[0] + damped_cov = self.get_cov() + damping * array_ops.eye(n) + + if isinstance(x, tf_ops.IndexedSlices): + raise NotImplementedError( + "Left-multiply not yet supported for IndexedSlices.") + + if len(x.shape) != 2: + raise ValueError( + "InverseProvidingFactors apply to matrix-shaped vectors. Found: %s." + % (x,)) + + return math_ops.matmul(damped_cov, x) + + def right_multiply(self, x, damping): + n = self.get_cov().shape[0] + damped_cov = self.get_cov() + damping * array_ops.eye(n) + + if isinstance(x, tf_ops.IndexedSlices): + return utils.matmul_sparse_dense(x, damped_cov) + + if len(x.shape) != 2: + raise ValueError( + "InverseProvidingFactors apply to matrix-shaped vectors. Found: %s." + % (x,)) + + return math_ops.matmul(x, damped_cov) + + def left_multiply_inverse(self, x, damping): + if isinstance(x, tf_ops.IndexedSlices): + raise ValueError("Left-multiply not yet supported for IndexedSlices.") + + if x.shape.ndims != 2: + raise ValueError( + "InverseProvidingFactors apply to matrix-shaped vectors. Found: %s." + % (x,)) + + return math_ops.matmul(self.get_damped_inverse(damping), x) + + def right_multiply_inverse(self, x, damping): + if isinstance(x, tf_ops.IndexedSlices): + return utils.matmul_sparse_dense(x, self.get_damped_inverse(damping)) + + if x.shape.ndims != 2: + raise ValueError( + "InverseProvidingFactors apply to matrix-shaped vectors. Found: %s." + % (x,)) + + return math_ops.matmul(x, self.get_damped_inverse(damping)) + class FullFactor(InverseProvidingFactor): """FisherFactor for a full matrix representation of the Fisher of a parameter. @@ -430,49 +614,46 @@ class FullFactor(InverseProvidingFactor): def __init__(self, params_grads, - batch_size, - colocate_cov_ops_with_inputs=False): + batch_size): self._batch_size = batch_size - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs - self._orig_params_grads_name = scope_string_from_params( - [params_grads, self._batch_size]) - params_grads_flat = [] - for params_grad in params_grads: - with _maybe_colocate_with(params_grad, - self._colocate_cov_ops_with_inputs): - col = utils.tensors_to_column(params_grad) - params_grads_flat.append(col) - self._params_grads_flat = tuple(params_grads_flat) + self._params_grads = tuple(utils.ensure_sequence(params_grad) + for params_grad in params_grads) super(FullFactor, self).__init__() @property def _var_scope(self): - return "ff_full/" + self._orig_params_grads_name + return "ff_full/" + scope_string_from_params( + [self._params_grads, self._batch_size]) @property def _cov_shape(self): - size = self._params_grads_flat[0].shape[0] - return [size, size] + size = sum(param_grad.shape.num_elements() + for param_grad in self._params_grads[0]) + return (size, size) @property def _num_sources(self): - return len(self._params_grads_flat) + return len(self._params_grads) @property def _dtype(self): - return self._params_grads_flat[0].dtype + return self._params_grads[0][0].dtype def _compute_new_cov(self, idx=0): # This will be a very basic rank 1 estimate - with _maybe_colocate_with(self._params_grads_flat[idx], - self._colocate_cov_ops_with_inputs): - return ((self._params_grads_flat[idx] * array_ops.transpose( - self._params_grads_flat[idx])) / math_ops.cast( - self._batch_size, self._params_grads_flat[idx].dtype)) + with maybe_colocate_with(self._params_grads[idx]): + params_grads_flat = utils.tensors_to_column(self._params_grads[idx]) + return ((params_grads_flat * array_ops.transpose( + params_grads_flat)) / math_ops.cast(self._batch_size, + params_grads_flat.dtype)) class DiagonalFactor(FisherFactor): - """A base class for FisherFactors that use diagonal approximations.""" + """A base class for FisherFactors that use diagonal approximations. + + A DiagonalFactor's covariance variable can be of any shape, but must contain + exactly one entry per parameter. + """ def __init__(self): super(DiagonalFactor, self).__init__() @@ -484,6 +665,45 @@ class DiagonalFactor(FisherFactor): def make_inverse_update_ops(self): return [] + def get_cov(self): + # self.get_cov() could be any shape, but it must have one entry per + # parameter. Flatten it into a vector. + cov_diag_vec = array_ops.reshape(self.get_cov_var(), [-1]) + return array_ops.diag(cov_diag_vec) + + def left_multiply(self, x, damping): + damped_cov = self.get_cov_var() + damping + if isinstance(x, tf_ops.IndexedSlices): + return utils.matmul_diag_sparse(array_ops.reshape(damped_cov, [-1]), x) + + if x.shape != damped_cov.shape: + raise ValueError("x (%s) and cov (%s) must have same shape." % + (x, damped_cov)) + + return damped_cov * x + + def right_multiply(self, x, damping): + raise NotImplementedError("Only left-multiply is currently supported.") + + def left_multiply_inverse(self, x, damping): + inverse = 1. / (self.get_cov_var() + damping) + + if isinstance(x, tf_ops.IndexedSlices): + return utils.matmul_diag_sparse(array_ops.reshape(inverse, [-1]), x) + + if x.shape != inverse.shape: + raise ValueError("x (%s) and cov (%s) must have same shape." % + (x, inverse)) + + return inverse * x + + def right_multiply_inverse(self, x, damping): + raise NotImplementedError("Only left-multiply is currently supported.") + + def register_damped_inverse(self, damping): + # DiagonalFactors don't keep explicit inverses. + pass + class NaiveDiagonalFactor(DiagonalFactor): """FisherFactor for a diagonal approximation of any type of param's Fisher. @@ -494,28 +714,30 @@ class NaiveDiagonalFactor(DiagonalFactor): def __init__(self, params_grads, - batch_size, - colocate_cov_ops_with_inputs=False): + batch_size): + """Initializes NaiveDiagonalFactor instance. + + Args: + params_grads: Sequence of Tensors, each with same shape as parameters this + FisherFactor corresponds to. For example, the gradient of the loss with + respect to parameters. + batch_size: int or 0-D Tensor. Size + """ + self._params_grads = tuple(utils.ensure_sequence(params_grad) + for params_grad in params_grads) self._batch_size = batch_size - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs - params_grads_flat = [] - for params_grad in params_grads: - with _maybe_colocate_with(params_grad, - self._colocate_cov_ops_with_inputs): - col = utils.tensors_to_column(params_grad) - params_grads_flat.append(col) - self._params_grads = tuple(params_grads_flat) - self._orig_params_grads_name = scope_string_from_params( - [self._params_grads, self._batch_size]) super(NaiveDiagonalFactor, self).__init__() @property def _var_scope(self): - return "ff_naivediag/" + self._orig_params_grads_name + return "ff_naivediag/" + scope_string_from_params( + [self._params_grads, self._batch_size]) @property def _cov_shape(self): - return self._params_grads[0].shape + size = sum(param_grad.shape.num_elements() + for param_grad in self._params_grads[0]) + return [size, 1] @property def _num_sources(self): @@ -523,13 +745,91 @@ class NaiveDiagonalFactor(DiagonalFactor): @property def _dtype(self): - return self._params_grads[0].dtype + return self._params_grads[0][0].dtype def _compute_new_cov(self, idx=0): - with _maybe_colocate_with(self._params_grads[idx], - self._colocate_cov_ops_with_inputs): - return (math_ops.square(self._params_grads[idx]) / math_ops.cast( - self._batch_size, self._params_grads[idx].dtype)) + with maybe_colocate_with(self._params_grads[idx]): + params_grads_flat = utils.tensors_to_column(self._params_grads[idx]) + return (math_ops.square(params_grads_flat) / math_ops.cast( + self._batch_size, params_grads_flat.dtype)) + + +class EmbeddingInputKroneckerFactor(DiagonalFactor): + r"""FisherFactor for input to an embedding layer. + + Given input_ids = [batch_size, input_size] representing indices into an + [vocab_size, embedding_size] embedding matrix, approximate input covariance by + a diagonal matrix, + + Cov(input_ids, input_ids) = + (1/batch_size) sum_{i} diag(n_hot(input[i]) ** 2). + + where n_hot() constructs an n-hot binary vector and diag() constructs a + diagonal matrix of size [vocab_size, vocab_size]. + """ + + def __init__(self, input_ids, vocab_size, dtype=None): + """Instantiate EmbeddingInputKroneckerFactor. + + Args: + input_ids: Tuple of Tensors of shape [batch_size, input_size] and dtype + int32. Indices into embedding matrix. + vocab_size: int or 0-D Tensor. Maximum value for entries in 'input_ids'. + dtype: dtype for covariance statistics. Must be a floating point type. + Defaults to float32. + """ + self._input_ids = input_ids + self._vocab_size = vocab_size + self._cov_dtype = dtype or dtypes.float32 + + super(EmbeddingInputKroneckerFactor, self).__init__() + + @property + def _var_scope(self): + return "ff_diag_embedding/" + scope_string_from_params(self._input_ids) + + @property + def _cov_shape(self): + return [self._vocab_size] + + @property + def _num_sources(self): + return len(self._input_ids) + + @property + def _dtype(self): + return self._cov_dtype + + def _compute_new_cov(self, idx=0): + with maybe_colocate_with(self._input_ids): + input_ids = self._input_ids[idx] + if len(input_ids.shape) > 2: + raise ValueError( + "Input to embeddings must have rank <= 2. Found rank %d." % len( + input_ids.shape)) + + batch_size = array_ops.shape(input_ids)[0] + + # Transform indices into one-hot vectors. + # + # TODO(b/72714822): There must be a faster way to construct the diagonal + # covariance matrix! This operation is O(batch_size * vocab_size), where + # it should be O(batch_size * input_size). + flat_input_ids = array_ops.reshape(input_ids, [-1]) + one_hots = array_ops.one_hot(flat_input_ids, + self._vocab_size) # [?, vocab_size] + + # Take average across examples. Note that, because all entries have + # magnitude zero or one, there's no need to square the entries. + # + # TODO(b/72714822): Support for SparseTensor, other kinds of aggregation + # within an example such as average. + # + # TODO(b/72714822): Support for partitioned embeddings. + new_cov = math_ops.reduce_sum(one_hots, axis=0) # [vocab_size] + new_cov /= math_ops.cast(batch_size, new_cov.dtype) + + return new_cov class FullyConnectedDiagonalFactor(DiagonalFactor): @@ -543,13 +843,10 @@ class FullyConnectedDiagonalFactor(DiagonalFactor): where the square is taken element-wise. """ - # TODO(jamesmartens): add units tests for this class - def __init__(self, inputs, outputs_grads, - has_bias=False, - colocate_cov_ops_with_inputs=False): + has_bias=False): """Instantiate FullyConnectedDiagonalFactor. Args: @@ -558,32 +855,25 @@ class FullyConnectedDiagonalFactor(DiagonalFactor): outputs_grads: List of Tensors of shape [batch_size, output_size]. Gradient of loss with respect to layer's preactivations. has_bias: bool. If True, append '1' to each input. - colocate_cov_ops_with_inputs: Whether to colocate cov_update ops with - their inputs. """ + self._inputs = inputs + self._has_bias = has_bias self._outputs_grads = outputs_grads - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs self._batch_size = array_ops.shape(inputs)[0] - self._orig_tensors_name = scope_string_from_params( - (inputs,) + tuple(outputs_grads)) - - # Note that we precompute the required operations on the inputs since the - # inputs don't change with the 'idx' argument to _compute_new_cov. (Only - # the target entry of _outputs_grads changes with idx.) - with _maybe_colocate_with(inputs, self._colocate_cov_ops_with_inputs): - if has_bias: - inputs = _append_homog(inputs) - self._squared_inputs = math_ops.square(inputs) + self._squared_inputs = None super(FullyConnectedDiagonalFactor, self).__init__() @property def _var_scope(self): - return "ff_diagfc/" + self._orig_tensors_name + return "ff_diagfc/" + scope_string_from_params( + (self._inputs,) + tuple(self._outputs_grads)) @property def _cov_shape(self): - return [self._squared_inputs.shape[1], self._outputs_grads[0].shape[1]] + input_size = self._inputs.shape[1] + self._has_bias + output_size = self._outputs_grads[0].shape[1] + return [input_size, output_size] @property def _num_sources(self): @@ -598,8 +888,14 @@ class FullyConnectedDiagonalFactor(DiagonalFactor): # square of an outer product is the outer-product of the entry-wise squares. # The gradient is the outer product of the input and the output gradients, # so we just square both and then take their outer-product. - with _maybe_colocate_with(self._squared_inputs, - self._colocate_cov_ops_with_inputs): + with maybe_colocate_with(self._outputs_grads[idx]): + # We only need to compute squared_inputs once + if self._squared_inputs is None: + inputs = self._inputs + if self._has_bias: + inputs = append_homog(self._inputs) + self._squared_inputs = math_ops.square(inputs) + new_cov = math_ops.matmul( self._squared_inputs, math_ops.square(self._outputs_grads[idx]), @@ -611,16 +907,13 @@ class FullyConnectedDiagonalFactor(DiagonalFactor): class ConvDiagonalFactor(DiagonalFactor): """FisherFactor for a diagonal approx of a convolutional layer's Fisher.""" - # TODO(jamesmartens): add units tests for this class - def __init__(self, inputs, outputs_grads, filter_shape, strides, padding, - has_bias=False, - colocate_cov_ops_with_inputs=False): + has_bias=False): """Creates a ConvDiagonalFactor object. Args: @@ -635,42 +928,21 @@ class ConvDiagonalFactor(DiagonalFactor): padding: The padding in this layer (1-D of Tensor length 4). has_bias: Python bool. If True, the layer is assumed to have a bias parameter in addition to its filter parameter. - colocate_cov_ops_with_inputs: Whether to colocate cov_update ops with - their inputs. """ + self._inputs = inputs self._filter_shape = filter_shape + self._strides = strides + self._padding = padding self._has_bias = has_bias self._outputs_grads = outputs_grads - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs - - self._orig_tensors_name = scope_string_from_name( - (inputs,) + tuple(outputs_grads)) - - # Note that we precompute the required operations on the inputs since the - # inputs don't change with the 'idx' argument to _compute_new_cov. (Only - # the target entry of _outputs_grads changes with idx.) - with _maybe_colocate_with(inputs, self._colocate_cov_ops_with_inputs): - filter_height, filter_width, _, _ = self._filter_shape - - # TODO(b/64144716): there is potential here for a big savings in terms of - # memory use. - patches = array_ops.extract_image_patches( - inputs, - ksizes=[1, filter_height, filter_width, 1], - strides=strides, - rates=[1, 1, 1, 1], - padding=padding) - - if has_bias: - patches = _append_homog(patches) - - self._patches = patches + self._patches = None super(ConvDiagonalFactor, self).__init__() @property def _var_scope(self): - return "ff_convdiag/" + self._orig_tensors_name + return "ff_convdiag/" + scope_string_from_name( + (self._inputs,) + tuple(self._outputs_grads)) @property def _cov_shape(self): @@ -688,9 +960,32 @@ class ConvDiagonalFactor(DiagonalFactor): def _dtype(self): return self._outputs_grads[0].dtype + def make_covariance_update_op(self, ema_decay): + with maybe_colocate_with(self._inputs): + filter_height, filter_width, _, _ = self._filter_shape + + # TODO(b/64144716): there is potential here for a big savings in terms + # of memory use. + patches = array_ops.extract_image_patches( + self._inputs, + ksizes=[1, filter_height, filter_width, 1], + strides=self._strides, + rates=[1, 1, 1, 1], + padding=self._padding) + + if self._has_bias: + patches = append_homog(patches) + + self._patches = patches + + op = super(ConvDiagonalFactor, self).make_covariance_update_op(ema_decay) + + self._patches = None + + return op + def _compute_new_cov(self, idx=0): - with _maybe_colocate_with(self._outputs_grads[idx], - self._colocate_cov_ops_with_inputs): + with maybe_colocate_with(self._outputs_grads[idx]): outputs_grad = self._outputs_grads[idx] batch_size = array_ops.shape(self._patches)[0] @@ -714,22 +1009,18 @@ class FullyConnectedKroneckerFactor(InverseProvidingFactor): def __init__(self, tensors, - has_bias=False, - colocate_cov_ops_with_inputs=False): + has_bias=False): """Instantiate FullyConnectedKroneckerFactor. Args: tensors: List of Tensors of shape [batch_size, n]. Represents either a layer's inputs or its output's gradients. has_bias: bool. If True, append '1' to each row. - colocate_cov_ops_with_inputs: Whether to colocate cov_update ops with - their inputs. """ # The tensor argument is either a tensor of input activations or a tensor of # output pre-activation gradients. self._has_bias = has_bias self._tensors = tensors - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs super(FullyConnectedKroneckerFactor, self).__init__() @property @@ -751,12 +1042,11 @@ class FullyConnectedKroneckerFactor(InverseProvidingFactor): return self._tensors[0].dtype def _compute_new_cov(self, idx=0): - with _maybe_colocate_with(self._tensors[idx], - self._colocate_cov_ops_with_inputs): + with maybe_colocate_with(self._tensors[idx]): tensor = self._tensors[idx] if self._has_bias: - tensor = _append_homog(tensor) - return _compute_cov(tensor) + tensor = append_homog(tensor) + return compute_cov(tensor) class ConvInputKroneckerFactor(InverseProvidingFactor): @@ -774,8 +1064,7 @@ class ConvInputKroneckerFactor(InverseProvidingFactor): filter_shape, strides, padding, - has_bias=False, - colocate_cov_ops_with_inputs=False): + has_bias=False): """Initializes ConvInputKroneckerFactor. Args: @@ -787,15 +1076,12 @@ class ConvInputKroneckerFactor(InverseProvidingFactor): width_stride, in_channel_stride]. padding: str. Padding method for layer. "SAME" or "VALID". has_bias: bool. If True, append 1 to in_channel. - colocate_cov_ops_with_inputs: Whether to colocate cov_update ops with - their inputs. """ self._filter_shape = filter_shape self._strides = strides self._padding = padding self._has_bias = has_bias self._inputs = inputs - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs super(ConvInputKroneckerFactor, self).__init__() @property @@ -823,8 +1109,7 @@ class ConvInputKroneckerFactor(InverseProvidingFactor): if idx != 0: raise ValueError("ConvInputKroneckerFactor only supports idx = 0") - # TODO(jamesmartens): factor this patches stuff out into a utility function - with _maybe_colocate_with(self._inputs, self._colocate_cov_ops_with_inputs): + with maybe_colocate_with(self._inputs): filter_height, filter_width, in_channels, _ = self._filter_shape # TODO(b/64144716): there is potential here for a big savings in terms of @@ -846,15 +1131,15 @@ class ConvInputKroneckerFactor(InverseProvidingFactor): # We append a homogenous coordinate to patches_flat if the layer has # bias parameters. This gives us [[A_l]]_H from the paper. if self._has_bias: - patches_flat = _append_homog(patches_flat) - # We call _compute_cov without passing in a normalizer. _compute_cov uses + patches_flat = append_homog(patches_flat) + # We call compute_cov without passing in a normalizer. compute_cov uses # the first dimension of patches_flat i.e. M|T| as the normalizer by # default. Hence we end up computing 1/M|T| * [[A_l]]^T [[A_l]], with # shape J|Delta| x J|Delta|. This is related to hat{Omega}_l from # the paper but has a different scale here for consistency with # ConvOutputKroneckerFactor. # (Tilde omitted over A for clarity.) - return _compute_cov(patches_flat) + return compute_cov(patches_flat) class ConvOutputKroneckerFactor(InverseProvidingFactor): @@ -868,18 +1153,15 @@ class ConvOutputKroneckerFactor(InverseProvidingFactor): Section 3.1 Estimating the factors. """ - def __init__(self, outputs_grads, colocate_cov_ops_with_inputs=False): + def __init__(self, outputs_grads): """Initializes ConvOutputKroneckerFactor. Args: outputs_grads: list of Tensors. Each Tensor is of shape [batch_size, height, width, out_channels]. - colocate_cov_ops_with_inputs: Whether to colocate cov_update ops with - their inputs. """ self._out_channels = outputs_grads[0].shape.as_list()[3] self._outputs_grads = outputs_grads - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs super(ConvOutputKroneckerFactor, self).__init__() @property @@ -900,8 +1182,7 @@ class ConvOutputKroneckerFactor(InverseProvidingFactor): return self._outputs_grads[0].dtype def _compute_new_cov(self, idx=0): - with _maybe_colocate_with(self._outputs_grads[idx], - self._colocate_cov_ops_with_inputs): + with maybe_colocate_with(self._outputs_grads[idx]): # reshaped_tensor below is the matrix DS_l defined in the KFC paper # (tilde omitted over S for clarity). It has shape M|T| x I, where # M = minibatch size, |T| = number of spatial locations, and @@ -909,10 +1190,10 @@ class ConvOutputKroneckerFactor(InverseProvidingFactor): reshaped_tensor = array_ops.reshape(self._outputs_grads[idx], [-1, self._out_channels]) # Following the reasoning in ConvInputKroneckerFactor._compute_new_cov, - # _compute_cov here returns 1/M|T| * DS_l^T DS_l = hat{Gamma}_l + # compute_cov here returns 1/M|T| * DS_l^T DS_l = hat{Gamma}_l # as defined in the paper, with shape I x I. # (Tilde omitted over S for clarity.) - return _compute_cov(reshaped_tensor) + return compute_cov(reshaped_tensor) class FullyConnectedMultiKF(InverseProvidingFactor): @@ -920,57 +1201,49 @@ class FullyConnectedMultiKF(InverseProvidingFactor): def __init__(self, tensor_lists, - has_bias=False, - colocate_cov_ops_with_inputs=False): + has_bias=False): """Constructs a new `FullyConnectedMultiKF`. Args: tensor_lists: List of lists of Tensors of shape [batch_size, n]. has_bias: bool. If True, '1' is appended to each row. - colocate_cov_ops_with_inputs: Whether to colocate cov_update ops with - their inputs. """ - self._orig_tensors_name = scope_string_from_params(tensor_lists) + self._tensor_lists = tensor_lists + self._has_bias = has_bias self._batch_size = array_ops.shape(tensor_lists[0][0])[0] self._num_timesteps = len(tensor_lists[0]) - - tensors = tuple( - array_ops.concat(tensor_list, 0) for tensor_list in tensor_lists) - if has_bias: - tensors = tuple(_append_homog(tensor) for tensor in tensors) - self._tensors = tensors + self._tensors = [None] * len(tensor_lists) self._cov_dt1 = None self._option1quants_by_damping = {} self._option2quants_by_damping = {} - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs super(FullyConnectedMultiKF, self).__init__() @property def _var_scope(self): - return "ff_fc_multi/" + self._orig_tensors_name + return "ff_fc_multi/" + scope_string_from_params(self._tensor_lists) @property def _num_sources(self): - return len(self._tensors) + return len(self._tensor_lists) @property def _dtype(self): - return self._tensors[0].dtype + return self._tensor_lists[0][0].dtype def make_covariance_update_op(self, ema_decay): - with _maybe_colocate_with(self._tensors, - self._colocate_cov_ops_with_inputs): - op = super(FullyConnectedMultiKF, - self).make_covariance_update_op(ema_decay) - - if self._cov_dt1 is not None: - new_cov_dt1 = math_ops.add_n( - tuple( - self._compute_new_cov_dt1(idx) - for idx in range(self._num_sources))) + + op = super(FullyConnectedMultiKF, self).make_covariance_update_op(ema_decay) + + if self._cov_dt1 is not None: + new_cov_dt1_contribs = tuple(self._compute_new_cov_dt1(idx) + for idx in range(self._num_sources)) + + with maybe_colocate_with(new_cov_dt1_contribs[0]): + new_cov_dt1 = math_ops.add_n(new_cov_dt1_contribs) + op2 = moving_averages.assign_moving_average( self._cov_dt1, new_cov_dt1, ema_decay, zero_debias=ZERO_DEBIAS) @@ -984,26 +1257,35 @@ class FullyConnectedMultiKF(InverseProvidingFactor): return op def _compute_new_cov(self, idx=0): - tensor = self._tensors[idx] - normalizer = self._num_timesteps * self._batch_size - return _compute_cov(tensor, normalizer=normalizer) + with maybe_colocate_with(self._tensor_lists[idx]): + tensor = array_ops.concat(self._tensor_lists[idx], 0) + if self._has_bias: + tensor = append_homog(tensor) + # We save these so they can be used by _compute_new_cov_dt1 + self._tensors[idx] = tensor + return compute_cov(tensor) def _compute_new_cov_dt1(self, idx=0): tensor = self._tensors[idx] - normalizer = self._num_timesteps * self._batch_size - tensor_present = tensor[:-self._batch_size, :] - tensor_future = tensor[self._batch_size:, :] - return _compute_cov( - tensor_future, tensor_right=tensor_present, normalizer=normalizer) + with maybe_colocate_with(tensor): + # Is there a more elegant way to do this computation? + tensor_present = tensor[:-self._batch_size, :] + tensor_future = tensor[self._batch_size:, :] + # We specify a normalizer for this computation to ensure a PSD Fisher + # block estimate. This is equivalent to padding with zeros, as was done + # in Section B.2 of the appendix. + normalizer = self._num_timesteps * self._batch_size + return compute_cov( + tensor_future, tensor_right=tensor_present, normalizer=normalizer) @property def _cov_shape(self): - size = self._tensors[0].shape[1] + size = self._tensor_lists[0][0].shape[1] + self._has_bias return [size, size] @property def _vec_shape(self): - size = self._tensors[0].shape[1] + size = self._tensor_lists[0][0].shape[1] + self._has_bias return [size] def get_option1quants(self, damping): @@ -1033,7 +1315,6 @@ class FullyConnectedMultiKF(InverseProvidingFactor): def register_option1quants(self, damping): - self.register_eigendecomp() self.register_cov_dt1() if damping not in self._option1quants_by_damping: @@ -1059,7 +1340,6 @@ class FullyConnectedMultiKF(InverseProvidingFactor): def register_option2quants(self, damping): - self.register_eigendecomp() self.register_cov_dt1() if damping not in self._option2quants_by_damping: diff --git a/tensorflow/contrib/kfac/python/ops/fisher_factors_lib.py b/tensorflow/contrib/kfac/python/ops/fisher_factors_lib.py index 23ee93cd405bbf719939df89d525c812ee061f8b..2d8e378a932c16d48360bc4b15ff4f3239c0ed1f 100644 --- a/tensorflow/contrib/kfac/python/ops/fisher_factors_lib.py +++ b/tensorflow/contrib/kfac/python/ops/fisher_factors_lib.py @@ -24,23 +24,15 @@ from tensorflow.python.util.all_util import remove_undocumented # pylint: enable=unused-import,line-too-long,wildcard-import _allowed_symbols = [ - "inverse_initializer", - "covariance_initializer", - "diagonal_covariance_initializer", - "scope_string_from_params", - "scope_string_from_name", - "scalar_or_tensor_to_string", - "FisherFactor", - "InverseProvidingFactor", - "FullFactor", - "DiagonalFactor", - "NaiveDiagonalFactor", - "FullyConnectedDiagonalFactor", - "FullyConnectedKroneckerFactor", - "ConvInputKroneckerFactor", - "ConvOutputKroneckerFactor", - "ConvDiagonalFactor", - "set_global_constants", + "inverse_initializer", "covariance_initializer", + "diagonal_covariance_initializer", "scope_string_from_params", + "scope_string_from_name", "scalar_or_tensor_to_string", "FisherFactor", + "InverseProvidingFactor", "FullFactor", "DiagonalFactor", + "NaiveDiagonalFactor", "EmbeddingInputKroneckerFactor", + "FullyConnectedDiagonalFactor", "FullyConnectedKroneckerFactor", + "ConvInputKroneckerFactor", "ConvOutputKroneckerFactor", + "ConvDiagonalFactor", "set_global_constants", "maybe_colocate_with", + "compute_cov", "append_homog" ] remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/kfac/python/ops/layer_collection.py b/tensorflow/contrib/kfac/python/ops/layer_collection.py index ca42afe6fb2f5c7d7de8b5b087dc11be30a75d5e..ce9005b9ce99a4efa5f2821c56e199dd2086482e 100644 --- a/tensorflow/contrib/kfac/python/ops/layer_collection.py +++ b/tensorflow/contrib/kfac/python/ops/layer_collection.py @@ -76,14 +76,6 @@ _FULLY_CONNECTED_MULTI_APPROX_TO_BLOCK_TYPES = { VARIABLE_SCOPE = "VARIABLE_SCOPE" -def ensure_sequence(obj): - """If `obj` isn't a tuple or list, return a tuple containing `obj`.""" - if isinstance(obj, (tuple, list)): - return obj - else: - return (obj,) - - class LayerParametersDict(OrderedDict): """An OrderedDict where keys are Tensors or tuples of Tensors. @@ -142,7 +134,6 @@ class LayerCollection(object): def __init__(self, graph=None, - colocate_cov_ops_with_inputs=False, name="LayerCollection"): self.fisher_blocks = LayerParametersDict() self.fisher_factors = OrderedDict() @@ -152,11 +143,11 @@ class LayerCollection(object): self._loss_dict = {} # {str: LossFunction} self._subgraph = None self._default_generic_approximation = APPROX_FULL_NAME + self._default_embedding_approximation = APPROX_KRONECKER_NAME self._default_fully_connected_approximation = APPROX_KRONECKER_NAME self._default_convolution_2d_approximation = APPROX_KRONECKER_NAME self._default_fully_connected_multi_approximation = ( APPROX_KRONECKER_SERIES_2_NAME) - self._colocate_cov_ops_with_inputs = colocate_cov_ops_with_inputs with variable_scope.variable_scope(None, default_name=name) as scope: self._var_scope = scope.name @@ -169,7 +160,7 @@ class LayerCollection(object): @property def registered_variables(self): """A tuple of all of the variables currently registered.""" - tuple_of_tuples = (ensure_sequence(key) for key, block + tuple_of_tuples = (utils.ensure_sequence(key) for key, block in six.iteritems(self.fisher_blocks)) flat_tuple = tuple(item for tuple_ in tuple_of_tuples for item in tuple_) return flat_tuple @@ -188,6 +179,17 @@ class LayerCollection(object): """ return self._linked_parameters + @property + def default_embedding_approximation(self): + return self._default_embedding_approximation + + def set_default_embedding_approximation(self, value): + if value != APPROX_KRONECKER_NAME: + raise ValueError( + "{} is not a valid approximation for embedding variables.".format( + value)) + self._default_embedding_approximation = value + @property def default_generic_approximation(self): return self._default_generic_approximation @@ -276,9 +278,9 @@ class LayerCollection(object): variable_to_block = { var: (params, block) for (params, block) in self.fisher_blocks.items() - for var in ensure_sequence(params) + for var in utils.ensure_sequence(params) } - for variable in ensure_sequence(layer_key): + for variable in utils.ensure_sequence(layer_key): if variable in variable_to_block: prev_key, prev_block = variable_to_block[variable] raise ValueError( @@ -301,7 +303,7 @@ class LayerCollection(object): block.num_inputs()*block.num_registered_minibatches if isinstance( block, (fb.FullyConnectedSeriesFB, fb.FullyConnectedMultiIndepFB)) else block.num_registered_minibatches) - key = ensure_sequence(key) + key = utils.ensure_sequence(key) for k in key: vars_to_uses[k] += n return vars_to_uses @@ -382,12 +384,12 @@ class LayerCollection(object): ValueError: If the parameters were already registered in a layer or identified as part of an incompatible group. """ - params = frozenset(ensure_sequence(params)) + params = frozenset(utils.ensure_sequence(params)) # Check if any of the variables in 'params' is already in # 'self.fisher_blocks.keys()'. for registered_params, fisher_block in self.fisher_blocks.items(): - registered_params_set = set(ensure_sequence(registered_params)) + registered_params_set = set(utils.ensure_sequence(registered_params)) for variable in params: if (variable in registered_params_set and params != registered_params_set): @@ -421,12 +423,52 @@ class LayerCollection(object): def _get_linked_approx(self, params): """If params were linked, return their specified approximation.""" - params_set = frozenset(ensure_sequence(params)) + params_set = frozenset(utils.ensure_sequence(params)) if params_set in self.linked_parameters: return self.linked_parameters[params_set] else: return None + def register_embedding(self, + params, + inputs, + outputs, + approx=None, + reuse=VARIABLE_SCOPE): + """Registers a fully connnected layer. + + Args: + params: Embedding matrix of shape [vocab_size, embedding_size]. + inputs: Tensor of shape [batch_size, input_size] and dtype int32. Indices + into embedding matrix. + outputs: Tensor of shape [batch_size, output_size]. Outputs + produced by layer. + approx: str. Must be "kron". + reuse: bool or str. If True, reuse an existing FisherBlock. If False, + create a new FisherBlock. If "VARIABLE_SCOPE", use + tf.get_variable_scope().reuse. + + Raises: + ValueError: For improper value to 'approx'. + KeyError: If reuse == True but no FisherBlock found for 'params'. + ValueError: If reuse == True and FisherBlock found but of the wrong type. + """ + if approx is None: + approx = self._get_linked_approx(params) + if approx is None: + approx = self.default_embedding_approximation + + if approx != APPROX_KRONECKER_NAME: + raise ValueError("Bad value {} for approx.".format(approx)) + + if isinstance(params, (tuple, list)): + raise ValueError("Bias not supported.") + + vocab_size = int(params.shape[0]) + block = self.register_block( + params, fb.EmbeddingKFACFB(self, vocab_size), reuse=reuse) + block.register_additional_minibatch(inputs, outputs) + def register_fully_connected(self, params, inputs, @@ -727,7 +769,6 @@ class LayerCollection(object): key = cls, args if key not in self.fisher_factors: - colo = self._colocate_cov_ops_with_inputs with variable_scope.variable_scope(self._var_scope): - self.fisher_factors[key] = cls(*args, colocate_cov_ops_with_inputs=colo) + self.fisher_factors[key] = cls(*args) return self.fisher_factors[key] diff --git a/tensorflow/contrib/kfac/python/ops/layer_collection_lib.py b/tensorflow/contrib/kfac/python/ops/layer_collection_lib.py index d6bf61a210203dd74d4e93b65005f660b1fab4ff..f8aa230d9ca1f542950f56b1e6cf1ab7ccd3d05f 100644 --- a/tensorflow/contrib/kfac/python/ops/layer_collection_lib.py +++ b/tensorflow/contrib/kfac/python/ops/layer_collection_lib.py @@ -36,6 +36,9 @@ _allowed_symbols = [ "APPROX_DIAGONAL_NAME", "APPROX_FULL_NAME", "VARIABLE_SCOPE", + "APPROX_KRONECKER_INDEP_NAME", + "APPROX_KRONECKER_SERIES_1_NAME", + "APPROX_KRONECKER_SERIES_2_NAME" ] remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/kfac/python/ops/loss_functions.py b/tensorflow/contrib/kfac/python/ops/loss_functions.py index d449abcfa78b361b9d4774ca5c2e936f14f65433..cb3e698b9ceab920785adf735f88bd8e535a628f 100644 --- a/tensorflow/contrib/kfac/python/ops/loss_functions.py +++ b/tensorflow/contrib/kfac/python/ops/loss_functions.py @@ -22,6 +22,7 @@ import abc import six +from tensorflow.contrib.distributions.python.ops import onehot_categorical from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops @@ -659,19 +660,20 @@ class CategoricalLogitsNegativeLogProbLoss(DistributionNegativeLogProbLoss, def multiply_fisher(self, vector): probs = self._probs - return vector * probs - math_ops.reduce_sum(vector * probs, axis=1) * probs + return vector * probs - probs * math_ops.reduce_sum( + vector * probs, axis=-1, keep_dims=True) def multiply_fisher_factor(self, vector): probs = self._probs sqrt_probs = self._sqrt_probs return sqrt_probs * vector - probs * math_ops.reduce_sum( - sqrt_probs * vector, axis=1, keep_dims=True) + sqrt_probs * vector, axis=-1, keep_dims=True) def multiply_fisher_factor_transpose(self, vector): probs = self._probs sqrt_probs = self._sqrt_probs return sqrt_probs * vector - sqrt_probs * math_ops.reduce_sum( - probs * vector, axis=1, keep_dims=True) + probs * vector, axis=-1, keep_dims=True) def multiply_fisher_factor_replicated_one_hot(self, index): assert len(index) == 1, "Length of index was {}".format(len(index)) @@ -785,3 +787,16 @@ def insert_slice_in_zeros(slice_to_insert, dim, dim_size, position): after[dim] = dim_size - position - 1 return array_ops.pad(slice_to_insert, list(zip(before, after))) + + +class OnehotCategoricalLogitsNegativeLogProbLoss( + CategoricalLogitsNegativeLogProbLoss): + """Neg log prob loss for a categorical distribution with onehot targets. + + Identical to CategoricalLogitsNegativeLogProbLoss except that the underlying + distribution is OneHotCategorical as opposed to Categorical. + """ + + @property + def dist(self): + return onehot_categorical.OneHotCategorical(logits=self._logits) diff --git a/tensorflow/contrib/kfac/python/ops/loss_functions_lib.py b/tensorflow/contrib/kfac/python/ops/loss_functions_lib.py index e9bb4f14e9e24128382832fcdaccdc9b24017046..705a871d482565897e7ac850327729a6186f1746 100644 --- a/tensorflow/contrib/kfac/python/ops/loss_functions_lib.py +++ b/tensorflow/contrib/kfac/python/ops/loss_functions_lib.py @@ -31,6 +31,7 @@ _allowed_symbols = [ "NormalMeanNegativeLogProbLoss", "NormalMeanVarianceNegativeLogProbLoss", "CategoricalLogitsNegativeLogProbLoss", + "OnehotCategoricalLogitsNegativeLogProbLoss", "MultiBernoulliNegativeLogProbLoss", "MultiBernoulliNegativeLogProbLoss", "insert_slice_in_zeros", diff --git a/tensorflow/contrib/kfac/python/ops/op_queue.py b/tensorflow/contrib/kfac/python/ops/op_queue.py index 831870fca451c585cb1a1dc6b24aad757e2bbaa8..b6d9d37a31a949b154b79e6f3677289a0d167373 100644 --- a/tensorflow/contrib/kfac/python/ops/op_queue.py +++ b/tensorflow/contrib/kfac/python/ops/op_queue.py @@ -18,7 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.data.python.ops import dataset_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import ops as tf_ops diff --git a/tensorflow/contrib/kfac/python/ops/optimizer.py b/tensorflow/contrib/kfac/python/ops/optimizer.py index ecf7f3e4e5ab7d9c151f760fdab733bc3830e37b..1974b07acfc879dc4bc844db9af88fd1043d6698 100644 --- a/tensorflow/contrib/kfac/python/ops/optimizer.py +++ b/tensorflow/contrib/kfac/python/ops/optimizer.py @@ -41,12 +41,12 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): damping, layer_collection, var_list=None, - momentum=0., + momentum=0.9, momentum_type="regular", norm_constraint=None, name="KFAC", estimation_mode="gradients", - colocate_gradients_with_ops=False, + colocate_gradients_with_ops=True, cov_devices=None, inv_devices=None): """Initializes the KFAC optimizer with the given settings. @@ -70,8 +70,8 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): var_list: Optional list or tuple of variables to train. Defaults to the list of variables collected in the graph under the key `GraphKeys.TRAINABLE_VARIABLES`. - momentum: The momentum value for this optimizer. Only applies when - momentum_type is 'regular' or 'adam'. (Default: 0) + momentum: The momentum decay constant to use. Only applies when + momentum_type is 'regular' or 'adam'. (Default: 0.9) momentum_type: The type of momentum to use in this optimizer, one of 'regular', 'adam', or 'qmodel'. (Default: 'regular') norm_constraint: float or Tensor. If specified, the update is scaled down @@ -85,6 +85,7 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): more a more detailed description of these options. colocate_gradients_with_ops: Whether we should request gradients we compute in the estimator be colocated with their respective ops. + (Default: True) cov_devices: Iterable of device strings (e.g. '/gpu:0'). Covariance computations will be placed on these devices in a round-robin fashion. Can be None, which means that no devices are specified. @@ -136,12 +137,32 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer): self._batch_size = array_ops.shape(layer_collection.losses[0].inputs)[0] self._losses = layer_collection.losses - self.cov_update_op = self._fisher_est.cov_update_op - self.inv_update_op = self._fisher_est.inv_update_op - self.inv_updates_dict = self._fisher_est.inv_updates_dict - super(KfacOptimizer, self).__init__(learning_rate, name=name) + @property + def cov_update_thunks(self): + return self._fisher_est.cov_update_thunks + + @property + def cov_update_ops(self): + return self._fisher_est.cov_update_ops + + @property + def cov_update_op(self): + return self._fisher_est.cov_update_op + + @property + def inv_update_thunks(self): + return self._fisher_est.inv_update_thunks + + @property + def inv_update_ops(self): + return self._fisher_est.inv_update_ops + + @property + def inv_update_op(self): + return self._fisher_est.inv_update_op + @property def variables(self): return self._fisher_est.variables diff --git a/tensorflow/contrib/kfac/python/ops/utils.py b/tensorflow/contrib/kfac/python/ops/utils.py index 48b191ef501f60118dd4583483d9d27dd5265621..f5bd97cb4e7d547394050e944f75b43a40887f34 100644 --- a/tensorflow/contrib/kfac/python/ops/utils.py +++ b/tensorflow/contrib/kfac/python/ops/utils.py @@ -25,6 +25,7 @@ from tensorflow.contrib.tpu.python.tpu import tpu_function from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops @@ -143,7 +144,9 @@ def layer_params_to_mat2d(vector): [-1, w_part.shape.as_list()[-1]]) return array_ops.concat( (w_part_reshaped, array_ops.reshape(b_part, [1, -1])), axis=0) - else: + elif isinstance(vector, ops.IndexedSlices): + return vector + else: # Tensor or Tensor-like. return array_ops.reshape(vector, [-1, vector.shape.as_list()[-1]]) @@ -162,6 +165,11 @@ def mat2d_to_layer_params(vector_template, mat2d): if isinstance(vector_template, (tuple, list)): w_part, b_part = mat2d[:-1], mat2d[-1] return array_ops.reshape(w_part, vector_template[0].shape), b_part + elif isinstance(vector_template, ops.IndexedSlices): + if not isinstance(mat2d, ops.IndexedSlices): + raise TypeError( + "If vector_template is an IndexedSlices, so should mat2d.") + return mat2d else: return array_ops.reshape(mat2d, vector_template.shape) @@ -344,5 +352,132 @@ def cross_replica_mean(tensor, name=None): return tpu_ops.cross_replica_sum(tensor / num_shards) +def ensure_sequence(obj): + """If `obj` isn't a tuple or list, return a tuple containing `obj`.""" + if isinstance(obj, (tuple, list)): + return obj + else: + return (obj,) + + +def batch_execute(global_step, thunks, batch_size, name=None): + """Executes a subset of ops per global step. + + Given a list of thunks, each of which produces a single stateful op, + ensures that exactly 'batch_size' ops are run per global step. Ops are + scheduled in a round-robin fashion. For example, with 3 ops + + global_step | op0 | op1 | op2 + ------------+-----+-----+----- + 0 | x | x | + ------------+-----+-----+----- + 1 | x | | x + ------------+-----+-----+----- + 2 | | x | x + ------------+-----+-----+----- + 3 | x | x | + ------------+-----+-----+----- + 4 | x | | x + + Does not guarantee order of op execution within a single global step. + + Args: + global_step: Tensor indicating time. Determines which ops run. + thunks: List of thunks. Each thunk encapsulates one op. Return values are + ignored. + batch_size: int. Number of ops to execute per global_step. + name: string or None. Name scope for newly added ops. + + Returns: + List of ops. Exactly 'batch_size' ops are guaranteed to have an effect + every global step. + """ + + def true_fn(thunk): + """Ensures thunk is executed and returns an Op (not a Tensor).""" + + def result(): + with ops.control_dependencies([thunk()]): + return control_flow_ops.no_op() + + return result + + def false_fn(_): + """Executes a no-op.""" + + def result(): + return control_flow_ops.no_op() + + return result + + with ops.name_scope(name, "batch_execute"): + true_fns = [true_fn(thunk) for thunk in thunks] + false_fns = [false_fn(thunk) for thunk in thunks] + num_thunks = len(thunks) + conditions = [ + math_ops.less( + math_ops.mod(batch_size - 1 + global_step * batch_size - j, + num_thunks), batch_size) for j in range(num_thunks) + ] + result = [ + control_flow_ops.cond(condition, true_fn, false_fn) + for (condition, true_fn, + false_fn) in zip(conditions, true_fns, false_fns) + ] + return result + + +def matmul_sparse_dense(A, B, name=None): # pylint: disable=invalid-name + """Computes matmul(A, B) where A is sparse, B is dense. + + Args: + A: tf.IndexedSlices with dense shape [m, n]. + B: tf.Tensor with shape [n, k]. + name: str. Name of op. + + Returns: + tf.IndexedSlices resulting from matmul(A, B). + + Raises: + ValueError: If A doesn't represent a matrix. + ValueError: If B is not rank-2. + """ + with ops.name_scope(name, "matmul_sparse_dense", [A, B]): + if A.indices.shape.ndims != 1 or A.values.shape.ndims != 2: + raise ValueError("A must represent a matrix. Found: %s." % A) + if B.shape.ndims != 2: + raise ValueError("B must be a matrix.") + new_values = math_ops.matmul(A.values, B) + return ops.IndexedSlices( + new_values, + A.indices, + dense_shape=array_ops.stack([A.dense_shape[0], new_values.shape[1]])) + + +def matmul_diag_sparse(A_diag, B, name=None): # pylint: disable=invalid-name + """Computes matmul(A, B) where A is a diagonal matrix, B is sparse. + + Args: + A_diag: diagonal entries of matrix A of shape [m, m]. + B: tf.IndexedSlices. Represents matrix of shape [m, n]. + name: str. Name of op. + + Returns: + tf.IndexedSlices resulting from matmul(A, B). + + Raises: + ValueError: If A_diag is not rank-1. + ValueError: If B doesn't represent a matrix. + """ + with ops.name_scope(name, "matmul_diag_sparse", [A_diag, B]): + A_diag = ops.convert_to_tensor(A_diag) + if A_diag.shape.ndims != 1: + raise ValueError("A_diag must be a rank-1 Tensor.") + if B.indices.shape.ndims != 1 or B.values.shape.ndims != 2: + raise ValueError("B must represent a matrix. Found: %s." % B) + a = array_ops.gather(A_diag, B.indices) + a = array_ops.reshape(a, list(a.shape) + [1] * (B.values.shape.ndims - 1)) + return ops.IndexedSlices(a * B.values, B.indices, dense_shape=B.dense_shape) + # TODO(b/69623235): Add a function for finding tensors that share gradients # to eliminate redundant fisher factor computations. diff --git a/tensorflow/contrib/kfac/python/ops/utils_lib.py b/tensorflow/contrib/kfac/python/ops/utils_lib.py index 8903c90fbce6a890aa419d89b3b79d75f69509fc..8e424a794691484fdea7d8481677aa641c433d4c 100644 --- a/tensorflow/contrib/kfac/python/ops/utils_lib.py +++ b/tensorflow/contrib/kfac/python/ops/utils_lib.py @@ -24,6 +24,7 @@ from tensorflow.python.util.all_util import remove_undocumented # pylint: enable=unused-import,line-too-long,wildcard-import _allowed_symbols = [ + "set_global_constants", "SequenceDict", "tensors_to_column", "column_to_tensors", @@ -37,6 +38,10 @@ _allowed_symbols = [ "SubGraph", "generate_random_signs", "fwd_gradients", + "ensure_sequence", + "batch_execute", + "matmul_sparse_dense", + "matmul_diag_sparse", ] remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/labeled_tensor/python/ops/core_test.py b/tensorflow/contrib/labeled_tensor/python/ops/core_test.py index 1f4a3ef568efc459d4a36fcb0d5de7e0bce8335c..e70b4923749d89aba1bd0187857d762305daeb07 100644 --- a/tensorflow/contrib/labeled_tensor/python/ops/core_test.py +++ b/tensorflow/contrib/labeled_tensor/python/ops/core_test.py @@ -225,7 +225,7 @@ class LabeledTensorTest(test_util.Base): tensor = array_ops.placeholder(dtypes.string, [None]) actual = core.LabeledTensor(tensor, ['x']) self.assertIsNone(actual.axes['x'].size) - self.assertIs(actual.axes['x'].value, tensor.get_shape()[0]) + self.assertIsNone(actual.axes['x'].value.value) def test_eq(self): self.assertEqual(self.lt, self.lt) diff --git a/tensorflow/contrib/layers/__init__.py b/tensorflow/contrib/layers/__init__.py index 6c624929f20503054e0258aad8a843f4a201be64..337c9e06b870b2cca53fcdbf3d94225660e193c4 100644 --- a/tensorflow/contrib/layers/__init__.py +++ b/tensorflow/contrib/layers/__init__.py @@ -27,6 +27,7 @@ See the @{$python/contrib.layers} guide. @@convolution2d_transpose @@conv3d_transpose @@convolution3d_transpose +@@dense_to_sparse @@dropout @@elu @@embedding_lookup_unique @@ -34,6 +35,7 @@ See the @{$python/contrib.layers} guide. @@fully_connected @@GDN @@gdn +@@images_to_sequence @@layer_norm @@linear @@max_pool2d @@ -49,6 +51,7 @@ See the @{$python/contrib.layers} guide. @@scale_gradient @@separable_conv2d @@separable_convolution2d +@@sequence_to_images @@softmax @@spatial_softmax @@stack diff --git a/tensorflow/contrib/layers/kernels/sparse_feature_cross_kernel.cc b/tensorflow/contrib/layers/kernels/sparse_feature_cross_kernel.cc index 932c5ab99249feda1e3a7f2d707ce4237fe7177f..01893d60615a9b4ded2afc88c6de0168d4be0921 100644 --- a/tensorflow/contrib/layers/kernels/sparse_feature_cross_kernel.cc +++ b/tensorflow/contrib/layers/kernels/sparse_feature_cross_kernel.cc @@ -423,8 +423,9 @@ class SparseFeatureCrossOp : public OpKernel { "Input values should be a std::vector but received shape ", values_list_in[i].shape().DebugString(), " at position ", i)); OP_REQUIRES( - context, indices_list_in[i].shape().dim_size(0) == - values_list_in[i].shape().dim_size(0), + context, + indices_list_in[i].shape().dim_size(0) == + values_list_in[i].shape().dim_size(0), errors::InvalidArgument( "Expected size of values to be ", indices_list_in[i].shape().dim_size(0), " got ", diff --git a/tensorflow/contrib/layers/python/layers/feature_column.py b/tensorflow/contrib/layers/python/layers/feature_column.py index 8d2931b4867938024a494459c77976e1e714de5a..b7d34d6435789e54403926a342481971e854b449 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column.py +++ b/tensorflow/contrib/layers/python/layers/feature_column.py @@ -752,6 +752,10 @@ class _WeightedSparseColumn( {self.weight_column_name: parsing_ops.VarLenFeature(self.dtype)}) return config + @property + def lookup_config(self): + return self.sparse_id_column.lookup_config + @property def key(self): """Returns a string which will be used as a key when we do sorting.""" diff --git a/tensorflow/contrib/layers/python/layers/feature_column_test.py b/tensorflow/contrib/layers/python/layers/feature_column_test.py index 5ae885b7202357326bd8494d382adb57fa636d20..fc8f153fe3abdc83aca5abfa9a4bb5f5d5531480 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column_test.py +++ b/tensorflow/contrib/layers/python/layers/feature_column_test.py @@ -102,6 +102,16 @@ class FeatureColumnTest(test.TestCase): weighted_ids = fc.weighted_sparse_column(ids, "weights") self.assertEqual(weighted_ids.name, "ids_weighted_by_weights") + def testWeightedSparseColumnWithVocabularyFile(self): + ids = fc.sparse_column_with_vocabulary_file( + "ids", "a_file", num_oov_buckets=7, vocab_size=3) + weighted_ids = fc.weighted_sparse_column(ids, "weights") + self.assertEqual(weighted_ids.name, "ids_weighted_by_weights") + self.assertEqual(weighted_ids.lookup_config, ids.lookup_config) + self.assertEqual(weighted_ids.lookup_config.vocab_size, 3) + self.assertEqual(weighted_ids.lookup_config.num_oov_buckets, 7) + self.assertEqual(weighted_ids.lookup_config.vocabulary_file, "a_file") + def testWeightedSparseColumnDeepCopy(self): ids = fc.sparse_column_with_keys("ids", ["marlo", "omar", "stringer"]) weighted = fc.weighted_sparse_column(ids, "weights") @@ -211,8 +221,8 @@ class FeatureColumnTest(test.TestCase): weighted_sparse_col = fc.weighted_sparse_column(ids, "weights") self.assertEqual(weighted_sparse_col.name, "ids_weighted_by_weights") - b = fc.shared_embedding_columns([sparse_col, weighted_sparse_col], - dimension=4, combiner="mean") + b = fc.shared_embedding_columns( + [sparse_col, weighted_sparse_col], dimension=4, combiner="mean") self.assertEqual(len(b), 2) self.assertEqual(b[0].shared_embedding_name, "a1_ids_weighted_by_weights_shared_embedding") @@ -220,8 +230,8 @@ class FeatureColumnTest(test.TestCase): "a1_ids_weighted_by_weights_shared_embedding") # Tries reversing order to check compatibility condition. - b = fc.shared_embedding_columns([weighted_sparse_col, sparse_col], - dimension=4, combiner="mean") + b = fc.shared_embedding_columns( + [weighted_sparse_col, sparse_col], dimension=4, combiner="mean") self.assertEqual(len(b), 2) self.assertEqual(b[0].shared_embedding_name, "a1_ids_weighted_by_weights_shared_embedding") @@ -230,18 +240,17 @@ class FeatureColumnTest(test.TestCase): # Tries adding two weighted columns to check compatibility between them. weighted_sparse_col_2 = fc.weighted_sparse_column(ids, "weights_2") - b = fc.shared_embedding_columns([weighted_sparse_col, - weighted_sparse_col_2], - dimension=4, combiner="mean") + b = fc.shared_embedding_columns( + [weighted_sparse_col, weighted_sparse_col_2], + dimension=4, + combiner="mean") self.assertEqual(len(b), 2) self.assertEqual( b[0].shared_embedding_name, - "ids_weighted_by_weights_ids_weighted_by_weights_2_shared_embedding" - ) + "ids_weighted_by_weights_ids_weighted_by_weights_2_shared_embedding") self.assertEqual( b[1].shared_embedding_name, - "ids_weighted_by_weights_ids_weighted_by_weights_2_shared_embedding" - ) + "ids_weighted_by_weights_ids_weighted_by_weights_2_shared_embedding") def testSharedEmbeddingColumnDeterminism(self): # Tests determinism in auto-generated shared_embedding_name. @@ -276,10 +285,10 @@ class FeatureColumnTest(test.TestCase): columns = fc.shared_embedding_columns( [a1, a2], dimension=4, combiner="mean") columns_copy = copy.deepcopy(columns) - self.assertEqual( - columns_copy[0].shared_embedding_name, "a1_a2_shared_embedding") - self.assertEqual( - columns_copy[1].shared_embedding_name, "a1_a2_shared_embedding") + self.assertEqual(columns_copy[0].shared_embedding_name, + "a1_a2_shared_embedding") + self.assertEqual(columns_copy[1].shared_embedding_name, + "a1_a2_shared_embedding") def testOneHotColumn(self): a = fc.sparse_column_with_keys("a", ["a", "b", "c", "d"]) @@ -326,11 +335,11 @@ class FeatureColumnTest(test.TestCase): weighted_ids = fc.weighted_sparse_column(ids, "weights") one_hot = fc.one_hot_column(weighted_ids) features = { - 'ids': constant_op.constant([['marlo', 'unknown', 'omar']]), - 'weights': constant_op.constant([[2., 4., 6.]]) + "ids": constant_op.constant([["marlo", "unknown", "omar"]]), + "weights": constant_op.constant([[2., 4., 6.]]) } one_hot_tensor = feature_column_ops.input_from_feature_columns( - features, [one_hot]) + features, [one_hot]) with self.test_session() as sess: sess.run(variables.global_variables_initializer()) sess.run(lookup_ops.tables_initializer()) @@ -339,11 +348,9 @@ class FeatureColumnTest(test.TestCase): def testMissingValueInOneHotColumnForSparseColumnWithKeys(self): ids = fc.sparse_column_with_keys("ids", ["marlo", "omar", "stringer"]) one_hot = fc.one_hot_column(ids) - features = { - 'ids': constant_op.constant([['marlo', 'unknown', 'omar']]) - } + features = {"ids": constant_op.constant([["marlo", "unknown", "omar"]])} one_hot_tensor = feature_column_ops.input_from_feature_columns( - features, [one_hot]) + features, [one_hot]) with self.test_session() as sess: sess.run(variables.global_variables_initializer()) sess.run(lookup_ops.tables_initializer()) @@ -369,8 +376,7 @@ class FeatureColumnTest(test.TestCase): self.assertEqual(d4.default_value, None) self.assertEqual(d4.is_sparse, True) # Default value is a list but dimension is None. - with self.assertRaisesRegexp(ValueError, - "Only scalar default value.*"): + with self.assertRaisesRegexp(ValueError, "Only scalar default value.*"): fc._real_valued_var_len_column("g5", default_value=[2., 3.]) def testRealValuedVarLenColumnDtypes(self): @@ -380,18 +386,19 @@ class FeatureColumnTest(test.TestCase): "rvc": parsing_ops.VarLenFeature(dtype=dtypes.float32) }, rvc.config) - rvc = fc._real_valued_var_len_column("rvc", default_value=0, - is_sparse=False) - self.assertDictEqual( - { - "rvc": parsing_ops.FixedLenSequenceFeature(shape=[], - dtype=dtypes.float32, - allow_missing=True, - default_value=0.0) - }, rvc.config) - - rvc = fc._real_valued_var_len_column("rvc", dtype=dtypes.int32, - default_value=0, is_sparse=True) + rvc = fc._real_valued_var_len_column( + "rvc", default_value=0, is_sparse=False) + self.assertDictEqual({ + "rvc": + parsing_ops.FixedLenSequenceFeature( + shape=[], + dtype=dtypes.float32, + allow_missing=True, + default_value=0.0) + }, rvc.config) + + rvc = fc._real_valued_var_len_column( + "rvc", dtype=dtypes.int32, default_value=0, is_sparse=True) self.assertDictEqual( { "rvc": parsing_ops.VarLenFeature(dtype=dtypes.int32) @@ -399,8 +406,8 @@ class FeatureColumnTest(test.TestCase): with self.assertRaisesRegexp(TypeError, "dtype must be convertible to float"): - fc._real_valued_var_len_column("rvc", dtype=dtypes.string, - default_value="", is_sparse=True) + fc._real_valued_var_len_column( + "rvc", dtype=dtypes.string, default_value="", is_sparse=True) def testRealValuedColumn(self): a = fc.real_valued_column("aaa") @@ -494,13 +501,13 @@ class FeatureColumnTest(test.TestCase): for output_rank in range(1, 3 + len(dimensions)): with variable_scope.variable_scope("output_rank_{}".format(output_rank)): real_valued_output = real_valued_column._to_dnn_input_layer( - constant_op.constant( - real_valued_input, dtype=dtypes.float32), + constant_op.constant(real_valued_input, dtype=dtypes.float32), output_rank=output_rank) with self.test_session() as sess: real_valued_eval = sess.run(real_valued_output) - expected_shape = (input_shape[:output_rank - 1] + - [np.prod(input_shape[output_rank - 1:])]) + expected_shape = ( + input_shape[:output_rank - 1] + + [np.prod(input_shape[output_rank - 1:])]) self.assertEquals(expected_shape, list(real_valued_eval.shape)) def testRealValuedColumnDensification(self): @@ -510,8 +517,7 @@ class FeatureColumnTest(test.TestCase): "sparse_real_valued1", is_sparse=True) sparse_tensor = sparse_tensor_lib.SparseTensor( values=[2.0, 5.0], indices=[[0, 0], [2, 0]], dense_shape=[3, 1]) - with self.assertRaisesRegexp( - ValueError, "Set is_sparse to False"): + with self.assertRaisesRegexp(ValueError, "Set is_sparse to False"): real_valued_column._to_dnn_input_layer(sparse_tensor) def testRealValuedColumnDeepCopy(self): @@ -539,9 +545,8 @@ class FeatureColumnTest(test.TestCase): def testBucketizedColumnRequiresRealValuedColumnDimension(self): with self.assertRaisesRegexp( TypeError, "source_column must be an instance of _RealValuedColumn.*"): - fc.bucketized_column(fc._real_valued_var_len_column("bbb", - is_sparse=True), - [0]) + fc.bucketized_column( + fc._real_valued_var_len_column("bbb", is_sparse=True), [0]) def testBucketizedColumnRequiresSortedBuckets(self): with self.assertRaisesRegexp(ValueError, @@ -644,20 +649,14 @@ class FeatureColumnTest(test.TestCase): def testRealValuedColumnDtypes(self): rvc = fc.real_valued_column("rvc") - self.assertDictEqual( - { - "rvc": parsing_ops.FixedLenFeature( - [1], dtype=dtypes.float32) - }, - rvc.config) + self.assertDictEqual({ + "rvc": parsing_ops.FixedLenFeature([1], dtype=dtypes.float32) + }, rvc.config) rvc = fc.real_valued_column("rvc", dtype=dtypes.int32) - self.assertDictEqual( - { - "rvc": parsing_ops.FixedLenFeature( - [1], dtype=dtypes.int32) - }, - rvc.config) + self.assertDictEqual({ + "rvc": parsing_ops.FixedLenFeature([1], dtype=dtypes.int32) + }, rvc.config) with self.assertRaisesRegexp(ValueError, "dtype must be convertible to float"): @@ -692,8 +691,9 @@ class FeatureColumnTest(test.TestCase): batch_size = 4 dense_scalar_input = [1, 2, 3, 4] sparse_column = fc.sparse_column_with_integerized_feature("values", 10) - features = {"values": - constant_op.constant(dense_scalar_input, dtype=dtypes.int64)} + features = { + "values": constant_op.constant(dense_scalar_input, dtype=dtypes.int64) + } sparse_column.insert_transformed_feature(features) sparse_output = features[sparse_column] expected_shape = [batch_size, 1] @@ -721,8 +721,7 @@ class FeatureColumnTest(test.TestCase): def testSparseColumnKeysDeepCopy(self): """Tests deepcopy of sparse_column_with_keys.""" - column = fc.sparse_column_with_keys( - "a", keys=["key0", "key1", "key2"]) + column = fc.sparse_column_with_keys("a", keys=["key0", "key1", "key2"]) self.assertEqual("a", column.name) column_copy = copy.deepcopy(column) self.assertEqual("a", column_copy.name) @@ -775,8 +774,9 @@ class FeatureColumnTest(test.TestCase): a = fc.sparse_column_with_hash_bucket("cross_aaa", hash_bucket_size=100) b = fc.sparse_column_with_hash_bucket("cross_bbb", hash_bucket_size=100) cross_col = fc.crossed_column(set([a, b]), hash_bucket_size=10000) - one_hot_col = fc.one_hot_column(fc.sparse_column_with_hash_bucket( - "sparse_column_for_one_hot", hash_bucket_size=100)) + one_hot_col = fc.one_hot_column( + fc.sparse_column_with_hash_bucket( + "sparse_column_for_one_hot", hash_bucket_size=100)) scattered_embedding_col = fc.scattered_embedding_column( "scattered_embedding_column", size=100, dimension=10, hash_key=1) feature_columns = set([ @@ -799,17 +799,13 @@ class FeatureColumnTest(test.TestCase): "str_id_weights_column": parsing_ops.VarLenFeature(dtypes.float32), "real_valued_column1": - parsing_ops.FixedLenFeature( - [1], dtype=dtypes.float32), + parsing_ops.FixedLenFeature([1], dtype=dtypes.float32), "real_valued_column2": - parsing_ops.FixedLenFeature( - [5], dtype=dtypes.float32), + parsing_ops.FixedLenFeature([5], dtype=dtypes.float32), "real_valued_column_for_bucketization1": - parsing_ops.FixedLenFeature( - [1], dtype=dtypes.float32), + parsing_ops.FixedLenFeature([1], dtype=dtypes.float32), "real_valued_column_for_bucketization2": - parsing_ops.FixedLenFeature( - [4], dtype=dtypes.float32), + parsing_ops.FixedLenFeature([4], dtype=dtypes.float32), "cross_aaa": parsing_ops.VarLenFeature(dtypes.string), "cross_bbb": @@ -839,11 +835,14 @@ class FeatureColumnTest(test.TestCase): real_valued_col0 = fc._real_valued_var_len_column( "real_valued_column0", is_sparse=True) real_valued_col1 = fc._real_valued_var_len_column( - "real_valued_column1", dtype=dtypes.int64, default_value=0, + "real_valued_column1", + dtype=dtypes.int64, + default_value=0, is_sparse=False) feature_columns = set([real_valued_col0, real_valued_col1]) expected_config = { - "real_valued_column0": parsing_ops.VarLenFeature(dtype=dtypes.float32), + "real_valued_column0": + parsing_ops.VarLenFeature(dtype=dtypes.float32), "real_valued_column1": parsing_ops.FixedLenSequenceFeature( [], dtype=dtypes.int64, allow_missing=True, default_value=0), @@ -864,7 +863,9 @@ class FeatureColumnTest(test.TestCase): real_valued_col5 = fc._real_valued_var_len_column( "real_valued_column5", default_value=2, is_sparse=True) real_valued_col6 = fc._real_valued_var_len_column( - "real_valued_column6", dtype=dtypes.int64, default_value=1, + "real_valued_column6", + dtype=dtypes.int64, + default_value=1, is_sparse=False) feature_columns = [ real_valued_col1, real_valued_col2, real_valued_col3, real_valued_col4, @@ -892,8 +893,7 @@ class FeatureColumnTest(test.TestCase): parsing_ops.VarLenFeature(dtype=dtypes.float32), "real_valued_column6": parsing_ops.FixedLenSequenceFeature( - [], dtype=dtypes.int64, allow_missing=True, - default_value=1) + [], dtype=dtypes.int64, allow_missing=True, default_value=1) }, config) @@ -1094,8 +1094,8 @@ class FeatureColumnTest(test.TestCase): # This will initialize the crossed column weights from provided checkpoint # and return a [4, 1] tensor which is same as weights variable. Since we # won't modify weights, this should be same as 'saved_col_weights'. - _, col_weights, _ = (feature_column_ops.weighted_sum_from_feature_columns( - { + _, col_weights, _ = ( + feature_column_ops.weighted_sum_from_feature_columns({ sparse_col_1.name: input_tensor, sparse_col_2.name: input_tensor }, [crossed_col_initialized], 1)) diff --git a/tensorflow/contrib/layers/python/layers/layers.py b/tensorflow/contrib/layers/python/layers/layers.py index 0d25a09852544a7eb1ed5eb9c2f3402d9064d91a..5c1ff9ec267f1bccd9bee44a4b19e7ed3ec24cf0 100644 --- a/tensorflow/contrib/layers/python/layers/layers.py +++ b/tensorflow/contrib/layers/python/layers/layers.py @@ -29,6 +29,7 @@ from tensorflow.contrib.framework.python.ops import variables from tensorflow.contrib.layers.python.layers import initializers from tensorflow.contrib.layers.python.layers import utils from tensorflow.python.eager import context +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import function from tensorflow.python.framework import ops @@ -54,47 +55,18 @@ from tensorflow.python.layers.maxout import maxout # TODO(b/28426988): Replace legacy_* fns migrated from slim. # TODO(b/28426988): Remove legacy_* when all uses have migrated to new API. -__all__ = ['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', - 'dropout', - 'elu', - 'flatten', - 'fully_connected', - 'GDN', - 'gdn', - 'layer_norm', - 'linear', - 'pool', - 'max_pool2d', - 'max_pool3d', - 'one_hot_encoding', - 'relu', - 'relu6', - 'repeat', - 'scale_gradient', - 'separable_conv2d', - 'separable_convolution2d', - 'softmax', - 'spatial_softmax', - 'stack', - 'unit_norm', - 'legacy_fully_connected', - 'legacy_linear', - 'legacy_relu', - 'maxout'] +__all__ = [ + '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', + 'legacy_fully_connected', 'legacy_linear', 'legacy_relu', 'maxout' +] DATA_FORMAT_NCHW = 'NCHW' DATA_FORMAT_NHWC = 'NHWC' @@ -139,13 +111,14 @@ def avg_pool2d(inputs, raise ValueError('data_format has to be either NCHW or NHWC.') with ops.name_scope(scope, 'AvgPool2D', [inputs]) as sc: inputs = ops.convert_to_tensor(inputs) - df = ('channels_first' if data_format and data_format.startswith('NC') - else 'channels_last') - layer = pooling_layers.AveragePooling2D(pool_size=kernel_size, - strides=stride, - padding=padding, - data_format=df, - _scope=sc) + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') + layer = pooling_layers.AveragePooling2D( + pool_size=kernel_size, + strides=stride, + padding=padding, + data_format=df, + _scope=sc) outputs = layer.apply(inputs) return utils.collect_named_outputs(outputs_collections, sc, outputs) @@ -187,13 +160,14 @@ def avg_pool3d(inputs, raise ValueError('data_format has to be either NCDHW or NDHWC.') with ops.name_scope(scope, 'AvgPool3D', [inputs]) as sc: inputs = ops.convert_to_tensor(inputs) - df = ('channels_first' if data_format and data_format.startswith('NC') - else 'channels_last') - layer = pooling_layers.AveragePooling3D(pool_size=kernel_size, - strides=stride, - padding=padding, - data_format=df, - _scope=sc) + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') + layer = pooling_layers.AveragePooling3D( + pool_size=kernel_size, + strides=stride, + padding=padding, + data_format=df, + _scope=sc) outputs = layer.apply(inputs) return utils.collect_named_outputs(outputs_collections, sc, outputs) @@ -298,8 +272,8 @@ def _fused_batch_norm(inputs, raise ValueError('Inputs %s has undefined rank' % inputs.name) elif original_rank not in [2, 4]: raise ValueError('Inputs %s has unsupported rank.' - ' Expected 2 or 4 but got %d' % ( - inputs.name, original_rank)) + ' Expected 2 or 4 but got %d' % (inputs.name, + original_rank)) if original_rank == 2: channels = inputs.get_shape()[-1].value if channels is None: @@ -393,6 +367,7 @@ def _fused_batch_norm(inputs, def _fused_batch_norm_training(): return nn.fused_batch_norm( inputs, gamma, beta, epsilon=epsilon, data_format=data_format) + def _fused_batch_norm_inference(): return nn.fused_batch_norm( inputs, @@ -403,9 +378,9 @@ def _fused_batch_norm(inputs, epsilon=epsilon, is_training=False, data_format=data_format) - outputs, mean, variance = utils.smart_cond(is_training, - _fused_batch_norm_training, - _fused_batch_norm_inference) + + outputs, mean, variance = utils.smart_cond( + is_training, _fused_batch_norm_training, _fused_batch_norm_inference) # If `is_training` doesn't have a constant value, because it is a `Tensor`, # a `Variable` or `Placeholder` then is_training_value will be None and @@ -415,6 +390,7 @@ def _fused_batch_norm(inputs, if need_updates: if updates_collections is None: no_updates = lambda: outputs + def _force_updates(): """Internal function forces updates moving_vars if is_training.""" update_moving_mean = moving_averages.assign_moving_average( @@ -424,9 +400,11 @@ def _fused_batch_norm(inputs, with ops.control_dependencies( [update_moving_mean, update_moving_variance]): return array_ops.identity(outputs) + outputs = utils.smart_cond(is_training, _force_updates, no_updates) else: moving_vars_fn = lambda: (moving_mean, moving_variance) + def _delay_updates(): """Internal function that delay updates moving_vars if is_training.""" update_moving_mean = moving_averages.assign_moving_average( @@ -434,9 +412,9 @@ def _fused_batch_norm(inputs, update_moving_variance = moving_averages.assign_moving_average( moving_variance, variance, decay, zero_debias=False) return update_moving_mean, update_moving_variance - update_mean, update_variance = utils.smart_cond(is_training, - _delay_updates, - moving_vars_fn) + + update_mean, update_variance = utils.smart_cond( + is_training, _delay_updates, moving_vars_fn) ops.add_to_collections(updates_collections, update_mean) ops.add_to_collections(updates_collections, update_variance) @@ -479,7 +457,12 @@ def batch_norm(inputs, Sergey Ioffe, Christian Szegedy - Can be used as a normalizer function for conv2d and fully_connected. + Can be used as a normalizer function for conv2d and fully_connected. The + normalization is over all but the last dimension if `data_format` is `NHWC` + and all but the second dimension if `data_format` is `NCHW`. In case of a 2D + tensor this corresponds to the batch dimension, while in case of a 4D tensor + this + corresponds to the batch and space dimensions. Note: when training, the moving_mean and moving_variance need to be updated. By default the update ops are placed in `tf.GraphKeys.UPDATE_OPS`, so they @@ -535,8 +518,8 @@ def batch_norm(inputs, then the batch normalization uses weighted mean and variance. (This can be used to correct for bias in training example selection.) - fused: if `True`, use a faster, fused implementation if possible. - If `None`, use the system recommended implementation. + fused: if `None` or `True`, use a faster, fused implementation if possible. + If `False`, use the system recommended implementation. data_format: A string. `NHWC` (default) and `NCHW` are supported. zero_debias_moving_mean: Use zero_debias for moving_mean. It creates a new pair of variables 'moving_mean/biased' and 'moving_mean/local_step'. @@ -588,10 +571,9 @@ def batch_norm(inputs, # implementation in normalization_layers.BatchNormalization. inputs = ops.convert_to_tensor(inputs) rank = inputs.get_shape().ndims - possible_to_fuse = (batch_weights is None and - not renorm and - rank in [2, 4] and - adjustment is None) + possible_to_fuse = ( + batch_weights is None and not renorm and rank in [2, 4] and + adjustment is None) if fused and possible_to_fuse and ( zero_debias_moving_mean or rank == 2 or updates_collections is not ops.GraphKeys.UPDATE_OPS): @@ -619,7 +601,9 @@ def batch_norm(inputs, layer_variable_getter = _build_variable_getter() with variable_scope.variable_scope( - scope, 'BatchNorm', [inputs], reuse=reuse, + scope, + 'BatchNorm', [inputs], + reuse=reuse, custom_getter=layer_variable_getter) as sc: inputs = ops.convert_to_tensor(inputs) @@ -667,15 +651,15 @@ def batch_norm(inputs, outputs = layer.apply(inputs, training=is_training) # Add variables to collections. - _add_variable_to_collections( - layer.moving_mean, variables_collections, 'moving_mean') - _add_variable_to_collections( - layer.moving_variance, variables_collections, 'moving_variance') + _add_variable_to_collections(layer.moving_mean, variables_collections, + 'moving_mean') + _add_variable_to_collections(layer.moving_variance, variables_collections, + 'moving_variance') if layer.beta is not None: _add_variable_to_collections(layer.beta, variables_collections, 'beta') if layer.gamma is not None: - _add_variable_to_collections( - layer.gamma, variables_collections, 'gamma') + _add_variable_to_collections(layer.gamma, variables_collections, + 'gamma') if activation_fn is not None: outputs = activation_fn(outputs) @@ -715,8 +699,8 @@ def batch_norm(inputs, params_shape = inputs_shape[-1:] params_shape_broadcast = None if not params_shape.is_fully_defined(): - raise ValueError('Inputs %s has undefined channels dimension %s.' % ( - inputs.name, params_shape)) + raise ValueError('Inputs %s has undefined channels dimension %s.' % + (inputs.name, params_shape)) # Allocate parameters for the beta and gamma of the normalization. beta, gamma = None, None @@ -727,23 +711,25 @@ def batch_norm(inputs, 'beta') beta_initializer = param_initializers.get('beta', init_ops.zeros_initializer()) - beta = variables.model_variable('beta', - shape=params_shape, - dtype=dtype, - initializer=beta_initializer, - collections=beta_collections, - trainable=trainable) + beta = variables.model_variable( + 'beta', + shape=params_shape, + dtype=dtype, + initializer=beta_initializer, + collections=beta_collections, + trainable=trainable) if scale: - gamma_collections = utils.get_variable_collections(variables_collections, - 'gamma') + gamma_collections = utils.get_variable_collections( + variables_collections, 'gamma') gamma_initializer = param_initializers.get('gamma', init_ops.ones_initializer()) - gamma = variables.model_variable('gamma', - shape=params_shape, - dtype=dtype, - initializer=gamma_initializer, - collections=gamma_collections, - trainable=trainable) + gamma = variables.model_variable( + 'gamma', + shape=params_shape, + dtype=dtype, + initializer=gamma_initializer, + collections=gamma_collections, + trainable=trainable) # Create moving_mean and moving_variance variables and add them to the # appropriate collections. We disable variable partitioning while creating @@ -792,8 +778,8 @@ def batch_norm(inputs, mean, variance = nn.moments(inputs, moments_axes) else: if data_format == DATA_FORMAT_NCHW: - mean, variance = nn.weighted_moments(inputs, moments_axes, - batch_weights, keep_dims=True) + mean, variance = nn.weighted_moments( + inputs, moments_axes, batch_weights, keepdims=True) mean = array_ops.reshape(mean, [-1]) variance = array_ops.reshape(variance, [-1]) else: @@ -802,19 +788,21 @@ def batch_norm(inputs, moving_vars_fn = lambda: (moving_mean, moving_variance) if updates_collections is None: + def _force_updates(): """Internal function forces updates moving_vars if is_training.""" update_moving_mean = moving_averages.assign_moving_average( moving_mean, mean, decay, zero_debias=zero_debias_moving_mean) update_moving_variance = moving_averages.assign_moving_average( moving_variance, variance, decay, zero_debias=False) - with ops.control_dependencies([update_moving_mean, - update_moving_variance]): + with ops.control_dependencies( + [update_moving_mean, update_moving_variance]): return array_ops.identity(mean), array_ops.identity(variance) - mean, variance = utils.smart_cond(is_training, - _force_updates, + + mean, variance = utils.smart_cond(is_training, _force_updates, moving_vars_fn) else: + def _delay_updates(): """Internal function that delay updates moving_vars if is_training.""" update_moving_mean = moving_averages.assign_moving_average( @@ -823,9 +811,8 @@ def batch_norm(inputs, moving_variance, variance, decay, zero_debias=False) return update_moving_mean, update_moving_variance - update_mean, update_variance = utils.smart_cond(is_training, - _delay_updates, - moving_vars_fn) + update_mean, update_variance = utils.smart_cond( + is_training, _delay_updates, moving_vars_fn) ops.add_to_collections(updates_collections, update_mean) ops.add_to_collections(updates_collections, update_variance) # Use computed moments during training and moving_vars otherwise. @@ -893,8 +880,8 @@ def bias_add(inputs, """ if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC): raise ValueError('data_format has to be either NCHW or NHWC.') - with variable_scope.variable_scope(scope, 'BiasAdd', [inputs], - reuse=reuse) as sc: + with variable_scope.variable_scope( + scope, 'BiasAdd', [inputs], reuse=reuse) as sc: inputs = ops.convert_to_tensor(inputs) dtype = inputs.dtype.base_dtype inputs_shape = inputs.get_shape() @@ -909,13 +896,16 @@ def bias_add(inputs, raise ValueError('`C` dimension must be known but is None') biases_collections = utils.get_variable_collections(variables_collections, 'biases') - biases = variables.model_variable('biases', - shape=[num_features,], - dtype=dtype, - initializer=initializer, - regularizer=regularizer, - collections=biases_collections, - trainable=trainable) + biases = variables.model_variable( + 'biases', + shape=[ + num_features, + ], + dtype=dtype, + initializer=initializer, + regularizer=regularizer, + collections=biases_collections, + trainable=trainable) outputs = nn.bias_add(inputs, biases, data_format=data_format) if activation_fn is not None: outputs = activation_fn(outputs) @@ -1015,8 +1005,10 @@ def convolution(inputs, if data_format not in [None, 'NWC', 'NCW', 'NHWC', 'NCHW', 'NDHWC', 'NCDHW']: raise ValueError('Invalid data_format: %r' % (data_format,)) - layer_variable_getter = _build_variable_getter( - {'bias': 'biases', 'kernel': 'weights'}) + layer_variable_getter = _build_variable_getter({ + 'bias': 'biases', + 'kernel': 'weights' + }) with variable_scope.variable_scope( scope, 'Conv', [inputs], reuse=reuse, @@ -1034,26 +1026,27 @@ def convolution(inputs, raise ValueError('Convolution not supported for input with rank', input_rank) - df = ('channels_first' if data_format and data_format.startswith('NC') - else 'channels_last') - layer = layer_class(filters=num_outputs, - kernel_size=kernel_size, - strides=stride, - padding=padding, - data_format=df, - dilation_rate=rate, - activation=None, - use_bias=not normalizer_fn and biases_initializer, - kernel_initializer=weights_initializer, - bias_initializer=biases_initializer, - kernel_regularizer=weights_regularizer, - bias_regularizer=biases_regularizer, - activity_regularizer=None, - trainable=trainable, - name=sc.name, - dtype=inputs.dtype.base_dtype, - _scope=sc, - _reuse=reuse) + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') + layer = layer_class( + filters=num_outputs, + kernel_size=kernel_size, + strides=stride, + padding=padding, + data_format=df, + dilation_rate=rate, + activation=None, + use_bias=not normalizer_fn and biases_initializer, + kernel_initializer=weights_initializer, + bias_initializer=biases_initializer, + kernel_regularizer=weights_regularizer, + bias_regularizer=biases_regularizer, + activity_regularizer=None, + trainable=trainable, + name=sc.name, + dtype=inputs.dtype.base_dtype, + _scope=sc, + _reuse=reuse) outputs = layer.apply(inputs) # Add variables to collections. @@ -1069,6 +1062,7 @@ def convolution(inputs, outputs = activation_fn(outputs) return utils.collect_named_outputs(outputs_collections, sc.name, outputs) + convolution2d = convolution convolution3d = convolution @@ -1144,13 +1138,14 @@ def convolution2d_in_plane( weights_shape = [kernel_h, kernel_w, 1, 1] weights_collections = utils.get_variable_collections( variables_collections, 'weights') - weights = variables.model_variable('weights', - shape=weights_shape, - dtype=dtype, - initializer=weights_initializer, - regularizer=weights_regularizer, - collections=weights_collections, - trainable=trainable) + weights = variables.model_variable( + 'weights', + shape=weights_shape, + dtype=dtype, + initializer=weights_initializer, + regularizer=weights_regularizer, + collections=weights_collections, + trainable=trainable) depthwise_weights = array_ops.tile(weights, [1, 1, num_filters_in, 1]) outputs = nn.depthwise_conv2d(inputs, depthwise_weights, [1, stride_h, stride_w, 1], padding) @@ -1161,13 +1156,16 @@ def convolution2d_in_plane( if biases_initializer is not None: biases_collections = utils.get_variable_collections( variables_collections, 'biases') - biases = variables.model_variable('biases', - shape=[num_filters_in,], - dtype=dtype, - initializer=biases_initializer, - regularizer=biases_regularizer, - collections=biases_collections, - trainable=trainable) + biases = variables.model_variable( + 'biases', + shape=[ + num_filters_in, + ], + dtype=dtype, + initializer=biases_initializer, + regularizer=biases_regularizer, + collections=biases_collections, + trainable=trainable) outputs = nn.bias_add(outputs, biases) if activation_fn is not None: @@ -1240,19 +1238,23 @@ def convolution2d_transpose( ValueError: If `data_format` is neither `NHWC` nor `NCHW`. ValueError: If `C` dimension of `inputs` is None. """ - layer_variable_getter = _build_variable_getter( - {'bias': 'biases', 'kernel': 'weights'}) + layer_variable_getter = _build_variable_getter({ + 'bias': 'biases', + 'kernel': 'weights' + }) with variable_scope.variable_scope( - scope, 'Conv2d_transpose', [inputs], reuse=reuse, + scope, + 'Conv2d_transpose', [inputs], + reuse=reuse, custom_getter=layer_variable_getter) as sc: if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC): raise ValueError('data_format has to be either NCHW or NHWC.') inputs = ops.convert_to_tensor(inputs) - df = ('channels_first' if data_format and data_format.startswith('NC') - else 'channels_last') + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') layer = convolutional_layers.Convolution2DTranspose( filters=num_outputs, kernel_size=kernel_size, @@ -1349,19 +1351,23 @@ def convolution3d_transpose( ValueError: If `data_format` is neither `NDHWC` nor `NCDHW`. ValueError: If `C` dimension of `inputs` is None. """ - layer_variable_getter = _build_variable_getter( - {'bias': 'biases', 'kernel': 'weights'}) + layer_variable_getter = _build_variable_getter({ + 'bias': 'biases', + 'kernel': 'weights' + }) with variable_scope.variable_scope( - scope, 'Conv3d_transpose', [inputs], reuse=reuse, + scope, + 'Conv3d_transpose', [inputs], + reuse=reuse, custom_getter=layer_variable_getter) as sc: if data_format not in (DATA_FORMAT_NCDHW, DATA_FORMAT_NDHWC): raise ValueError('data_format has to be either NCDHW or NDHWC.') inputs = ops.convert_to_tensor(inputs) - df = ('channels_first' if data_format and data_format.startswith('NC') - else 'channels_last') + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') layer = convolutional_layers.Convolution3DTranspose( filters=num_outputs, kernel_size=kernel_size, @@ -1396,6 +1402,30 @@ def convolution3d_transpose( return utils.collect_named_outputs(outputs_collections, sc.name, outputs) +@add_arg_scope +def dense_to_sparse(tensor, eos_token=0, outputs_collections=None, scope=None): + """Converts a dense tensor into a sparse tensor. + An example use would be to convert dense labels to sparse ones + so that they can be fed to the ctc_loss. + + Args: + tensor: An `int` `Tensor` to be converted to a `Sparse`. + eos_token: An integer. + It is part of the target label that signfies the end of a sentence. + outputs_collections: Collection to add the outputs. + scope: Optional scope for name_scope. + """ + with variable_scope.variable_scope(scope, 'dense_to_sparse', [tensor]) as sc: + tensor = ops.convert_to_tensor(tensor) + indices = array_ops.where( + math_ops.not_equal(tensor, constant_op.constant(eos_token, + tensor.dtype))) + values = array_ops.gather_nd(tensor, indices) + shape = array_ops.shape(tensor, out_type=dtypes.int64) + outputs = sparse_tensor.SparseTensor(indices, values, shape) + return utils.collect_named_outputs(outputs_collections, sc.name, outputs) + + @add_arg_scope def dropout(inputs, keep_prob=0.5, @@ -1430,19 +1460,18 @@ def dropout(inputs, with variable_scope.variable_scope( scope, 'Dropout', [inputs], custom_getter=_model_variable_getter) as sc: inputs = ops.convert_to_tensor(inputs) - layer = core_layers.Dropout(rate=1 - keep_prob, - noise_shape=noise_shape, - seed=seed, - name=sc.name, - _scope=sc) + layer = core_layers.Dropout( + rate=1 - keep_prob, + noise_shape=noise_shape, + seed=seed, + name=sc.name, + _scope=sc) outputs = layer.apply(inputs, training=is_training) return utils.collect_named_outputs(outputs_collections, sc.name, outputs) @add_arg_scope -def flatten(inputs, - outputs_collections=None, - scope=None): +def flatten(inputs, outputs_collections=None, scope=None): """Flattens the input while maintaining the batch_size. Assumes that the first dimension represents the batch. @@ -1474,8 +1503,8 @@ def _sparse_inner_flatten(inputs, new_rank): outer_dimensions = inputs.dense_shape[:new_rank - 1] inner_dimensions = inputs.dense_shape[new_rank - 1:] - new_shape = array_ops.concat((outer_dimensions, - [math_ops.reduce_prod(inner_dimensions)]), 0) + new_shape = array_ops.concat( + (outer_dimensions, [math_ops.reduce_prod(inner_dimensions)]), 0) flattened = sparse_ops.sparse_reshape(inputs, new_shape) return flattened @@ -1541,10 +1570,18 @@ def _inner_flatten(inputs, new_rank, output_collections=None, scope=None): return utils.collect_named_outputs(output_collections, sc, flattened) -def _model_variable_getter(getter, name, shape=None, dtype=None, - initializer=None, regularizer=None, trainable=True, - collections=None, caching_device=None, - partitioner=None, rename=None, use_resource=None, +def _model_variable_getter(getter, + name, + shape=None, + dtype=None, + initializer=None, + regularizer=None, + trainable=True, + collections=None, + caching_device=None, + partitioner=None, + rename=None, + use_resource=None, **_): """Getter that uses model_variable for compatibility with core layers.""" short_name = name.split('/')[-1] @@ -1553,25 +1590,34 @@ def _model_variable_getter(getter, name, shape=None, dtype=None, name_components[-1] = rename[short_name] name = '/'.join(name_components) return variables.model_variable( - name, shape=shape, dtype=dtype, initializer=initializer, - regularizer=regularizer, collections=collections, trainable=trainable, - caching_device=caching_device, partitioner=partitioner, - custom_getter=getter, use_resource=use_resource) + name, + shape=shape, + dtype=dtype, + initializer=initializer, + regularizer=regularizer, + collections=collections, + trainable=trainable, + caching_device=caching_device, + partitioner=partitioner, + custom_getter=getter, + use_resource=use_resource) def _build_variable_getter(rename=None): """Build a model variable getter that respects scope getter and renames.""" + # VariableScope will nest the getters def layer_variable_getter(getter, *args, **kwargs): kwargs['rename'] = rename return _model_variable_getter(getter, *args, **kwargs) + return layer_variable_getter def _add_variable_to_collections(variable, collections_set, collections_name): """Adds variable (or all its parts) to all collections with that name.""" - collections = utils.get_variable_collections( - collections_set, collections_name) or [] + collections = utils.get_variable_collections(collections_set, + collections_name) or [] variables_list = [variable] if isinstance(variable, tf_variables.PartitionedVariable): variables_list = [v for v in variable] @@ -1640,15 +1686,19 @@ def fully_connected(inputs, ValueError: If x has rank less than 2 or if its last dimension is not set. """ if not isinstance(num_outputs, six.integer_types): - raise ValueError( - 'num_outputs should be int or long, got %s.' % (num_outputs,)) + raise ValueError('num_outputs should be int or long, got %s.' % + (num_outputs,)) - layer_variable_getter = _build_variable_getter({'bias': 'biases', - 'kernel': 'weights'}) + layer_variable_getter = _build_variable_getter({ + 'bias': 'biases', + 'kernel': 'weights' + }) with variable_scope.variable_scope( - scope, 'fully_connected', [inputs], - reuse=reuse, custom_getter=layer_variable_getter) as sc: + scope, + 'fully_connected', [inputs], + reuse=reuse, + custom_getter=layer_variable_getter) as sc: inputs = ops.convert_to_tensor(inputs) layer = core_layers.Dense( units=num_outputs, @@ -1754,15 +1804,17 @@ class GDN(base.Layer): inverse=False, beta_min=1e-6, gamma_init=.1, - reparam_offset=2 ** -18, + reparam_offset=2**-18, data_format='channels_last', activity_regularizer=None, trainable=True, name=None, **kwargs): - super(GDN, self).__init__(trainable=trainable, name=name, - activity_regularizer=activity_regularizer, - **kwargs) + super(GDN, self).__init__( + trainable=trainable, + name=name, + activity_regularizer=activity_regularizer, + **kwargs) self.inverse = inverse self._beta_min = beta_min self._gamma_init = gamma_init @@ -1797,8 +1849,9 @@ class GDN(base.Layer): with ops.name_scope(name, 'GDNLowerBound', [inputs, bound]) as scope: inputs = ops.convert_to_tensor(inputs, name='inputs') bound = ops.convert_to_tensor(bound, name='bound') - with ops.get_default_graph().gradient_override_map( - {'Maximum': 'GDNLowerBound'}): + with ops.get_default_graph().gradient_override_map({ + 'Maximum': 'GDNLowerBound' + }): return math_ops.maximum(inputs, bound, name=scope) @staticmethod @@ -1825,12 +1878,14 @@ class GDN(base.Layer): raise ValueError('The channel dimension of the inputs to `GDN` ' 'must be defined.') self._input_rank = input_shape.ndims - self.input_spec = base.InputSpec(ndim=input_shape.ndims, - axes={channel_axis: num_channels}) + self.input_spec = base.InputSpec( + ndim=input_shape.ndims, axes={ + channel_axis: num_channels + }) - pedestal = array_ops.constant(self._reparam_offset ** 2, dtype=self.dtype) + pedestal = array_ops.constant(self._reparam_offset**2, dtype=self.dtype) beta_bound = array_ops.constant( - (self._beta_min + self._reparam_offset ** 2) ** .5, dtype=self.dtype) + (self._beta_min + self._reparam_offset**2)**.5, dtype=self.dtype) gamma_bound = array_ops.constant(self._reparam_offset, dtype=self.dtype) def beta_initializer(shape, dtype=None, partition_info=None): @@ -1844,19 +1899,21 @@ class GDN(base.Layer): eye = linalg_ops.eye(shape[0], dtype=dtype) return math_ops.sqrt(self._gamma_init * eye + pedestal) - beta = self.add_variable('reparam_beta', - shape=[num_channels], - initializer=beta_initializer, - dtype=self.dtype, - trainable=True) + beta = self.add_variable( + 'reparam_beta', + shape=[num_channels], + initializer=beta_initializer, + dtype=self.dtype, + trainable=True) beta = self._lower_bound(beta, beta_bound) self.beta = math_ops.square(beta) - pedestal - gamma = self.add_variable('reparam_gamma', - shape=[num_channels, num_channels], - initializer=gamma_initializer, - dtype=self.dtype, - trainable=True) + gamma = self.add_variable( + 'reparam_gamma', + shape=[num_channels, num_channels], + initializer=gamma_initializer, + dtype=self.dtype, + trainable=True) gamma = self._lower_bound(gamma, gamma_bound) self.gamma = math_ops.square(gamma) - pedestal @@ -1871,8 +1928,11 @@ class GDN(base.Layer): # Compute normalization pool. if self.data_format == 'channels_first': - norm_pool = nn.convolution(math_ops.square(inputs), gamma, 'VALID', - data_format='NC' + 'DHW'[-(ndim - 2):]) + norm_pool = nn.convolution( + math_ops.square(inputs), + gamma, + 'VALID', + data_format='NC' + 'DHW' [-(ndim - 2):]) if ndim == 3: norm_pool = array_ops.expand_dims(norm_pool, 2) norm_pool = nn.bias_add(norm_pool, self.beta, data_format='NCHW') @@ -1896,7 +1956,7 @@ class GDN(base.Layer): outputs.set_shape(inputs.get_shape()) return outputs - def _compute_output_shape(self, input_shape): + def compute_output_shape(self, input_shape): channel_axis = self._channel_axis() input_shape = tensor_shape.TensorShape(input_shape) if not 3 <= input_shape.ndim <= 5: @@ -1914,7 +1974,7 @@ def gdn(inputs, inverse=False, beta_min=1e-6, gamma_init=.1, - reparam_offset=2 ** -18, + reparam_offset=2**-18, data_format='channels_last', activity_regularizer=None, trainable=True, @@ -1980,17 +2040,18 @@ def gdn(inputs, Returns: Output tensor. """ - layer = GDN(inverse=inverse, - beta_min=beta_min, - gamma_init=gamma_init, - reparam_offset=reparam_offset, - data_format=data_format, - activity_regularizer=activity_regularizer, - trainable=trainable, - name=name, - dtype=inputs.dtype.base_dtype, - _scope=name, - _reuse=reuse) + layer = GDN( + inverse=inverse, + beta_min=beta_min, + gamma_init=gamma_init, + reparam_offset=reparam_offset, + data_format=data_format, + activity_regularizer=activity_regularizer, + trainable=trainable, + name=name, + dtype=inputs.dtype.base_dtype, + _scope=name, + _reuse=reuse) return layer.apply(inputs) @@ -2066,8 +2127,8 @@ def layer_norm(inputs, or if `inputs.shape[begin_params_axis:]` is not fully defined at graph build time. """ - with variable_scope.variable_scope(scope, 'LayerNorm', [inputs], - reuse=reuse) as sc: + with variable_scope.variable_scope( + scope, 'LayerNorm', [inputs], reuse=reuse) as sc: inputs = ops.convert_to_tensor(inputs) inputs_shape = inputs.shape inputs_rank = inputs_shape.ndims @@ -2077,15 +2138,14 @@ def layer_norm(inputs, if begin_norm_axis < 0: begin_norm_axis = inputs_rank + begin_norm_axis if begin_params_axis >= inputs_rank or begin_norm_axis >= inputs_rank: - raise ValueError( - 'begin_params_axis (%d) and begin_norm_axis (%d) ' - 'must be < rank(inputs) (%d)' - % (begin_params_axis, begin_norm_axis, inputs_rank)) + raise ValueError('begin_params_axis (%d) and begin_norm_axis (%d) ' + 'must be < rank(inputs) (%d)' % + (begin_params_axis, begin_norm_axis, inputs_rank)) params_shape = inputs_shape[begin_params_axis:] if not params_shape.is_fully_defined(): raise ValueError( - 'Inputs %s: shape(inputs)[%s:] is not fully defined: %s' % ( - inputs.name, begin_params_axis, inputs_shape)) + 'Inputs %s: shape(inputs)[%s:] is not fully defined: %s' % + (inputs.name, begin_params_axis, inputs_shape)) # Allocate parameters for the beta and gamma of the normalization. beta, gamma = None, None if center: @@ -2099,8 +2159,8 @@ def layer_norm(inputs, collections=beta_collections, trainable=trainable) if scale: - gamma_collections = utils.get_variable_collections(variables_collections, - 'gamma') + gamma_collections = utils.get_variable_collections( + variables_collections, 'gamma') gamma = variables.model_variable( 'gamma', shape=params_shape, @@ -2114,7 +2174,11 @@ def layer_norm(inputs, # Compute layer normalization using the batch_normalization function. variance_epsilon = 1e-12 outputs = nn.batch_normalization( - inputs, mean, variance, offset=beta, scale=gamma, + inputs, + mean, + variance, + offset=beta, + scale=gamma, variance_epsilon=variance_epsilon) outputs.set_shape(inputs_shape) if activation_fn is not None: @@ -2122,6 +2186,34 @@ def layer_norm(inputs, return utils.collect_named_outputs(outputs_collections, sc.name, outputs) +@add_arg_scope +def images_to_sequence(inputs, data_format=DATA_FORMAT_NHWC, + outputs_collections=None, scope=None): + """Convert a batch of images into a batch of sequences. + Args: + inputs: a (num_images, height, width, depth) tensor + data_format: A string. `NHWC` (default) and `NCHW` are supported. + outputs_collections: The collections to which the outputs are added. + scope: Optional scope for name_scope. + Returns: + (width, num_images*height, depth) sequence tensor + """ + if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC): + raise ValueError('data_format has to be either NCHW or NHWC.') + with ops.name_scope(scope, 'ImagesToSequence', [inputs]) as sc: + inputs = ops.convert_to_tensor(inputs) + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') + if df == 'channels_first': + inputs = array_ops.transpose(inputs, [0, 2, 3, 1]) + _, _, width, depth = inputs.get_shape().as_list() + s = array_ops.shape(inputs) + batch_size, height = s[0], s[1] + transposed = array_ops.transpose(inputs, [2, 0, 1, 3]) + outputs = array_ops.reshape(transposed, [width, batch_size * height, depth]) + return utils.collect_named_outputs(outputs_collections, sc, outputs) + + @add_arg_scope def max_pool2d(inputs, kernel_size, @@ -2160,13 +2252,14 @@ def max_pool2d(inputs, raise ValueError('data_format has to be either NCHW or NHWC.') with ops.name_scope(scope, 'MaxPool2D', [inputs]) as sc: inputs = ops.convert_to_tensor(inputs) - df = ('channels_first' if data_format and data_format.startswith('NC') - else 'channels_last') - layer = pooling_layers.MaxPooling2D(pool_size=kernel_size, - strides=stride, - padding=padding, - data_format=df, - _scope=sc) + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') + layer = pooling_layers.MaxPooling2D( + pool_size=kernel_size, + strides=stride, + padding=padding, + data_format=df, + _scope=sc) outputs = layer.apply(inputs) return utils.collect_named_outputs(outputs_collections, sc, outputs) @@ -2209,13 +2302,14 @@ def max_pool3d(inputs, raise ValueError('data_format has to be either NCDHW or NDHWC.') with ops.name_scope(scope, 'MaxPool3D', [inputs]) as sc: inputs = ops.convert_to_tensor(inputs) - df = ('channels_first' if data_format and data_format.startswith('NC') - else 'channels_last') - layer = pooling_layers.MaxPooling3D(pool_size=kernel_size, - strides=stride, - padding=padding, - data_format=df, - _scope=sc) + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') + layer = pooling_layers.MaxPooling3D( + pool_size=kernel_size, + strides=stride, + padding=padding, + data_format=df, + _scope=sc) outputs = layer.apply(inputs) return utils.collect_named_outputs(outputs_collections, sc, outputs) @@ -2268,8 +2362,8 @@ def pool(inputs, """ # pylint: enable=line-too-long - with ops.name_scope(scope, '%s_pool' % - (pooling_type.lower()), [inputs]) as sc: + with ops.name_scope(scope, '%s_pool' % (pooling_type.lower()), + [inputs]) as sc: inputs = ops.convert_to_tensor(inputs) input_rank = inputs.get_shape().ndims if input_rank is None: @@ -2314,18 +2408,16 @@ def one_hot_encoding(labels, labels = ops.convert_to_tensor(labels) if labels.dtype == dtypes.int32: labels = standard_ops.to_int64(labels) - outputs = standard_ops.one_hot(labels, - num_classes, - on_value=on_value, - off_value=off_value) + outputs = standard_ops.one_hot( + labels, num_classes, on_value=on_value, off_value=off_value) return utils.collect_named_outputs(outputs_collections, sc, outputs) def _apply_activation(y, activation_fn, output_collections): if activation_fn is not None: y = activation_fn(y) - ops.add_to_collections(list(output_collections or []) + - [ops.GraphKeys.ACTIVATIONS], y) + ops.add_to_collections( + list(output_collections or []) + [ops.GraphKeys.ACTIVATIONS], y) return y @@ -2370,7 +2462,7 @@ def repeat(inputs, repetitions, layer, *args, **kwargs): scope = 'repeat' outputs = inputs for i in range(repetitions): - kwargs['scope'] = scope + '_' + str(i+1) + kwargs['scope'] = scope + '_' + str(i + 1) outputs = layer(outputs, *args, **kwargs) return outputs @@ -2385,8 +2477,8 @@ def _scale_gradient_grad(op, grad): return [grad * op.inputs[1], None] -@function.Defun(python_grad_func=_scale_gradient_grad, - shape_func=_scale_gradient_shape) +@function.Defun( + python_grad_func=_scale_gradient_grad, shape_func=_scale_gradient_shape) def scale_gradient(inputs, gradient_multiplier): """Identity operation, but with the gradient multiplied by a tensor. @@ -2491,18 +2583,21 @@ def separable_convolution2d( """ if data_format not in (DATA_FORMAT_NCHW, DATA_FORMAT_NHWC): raise ValueError('data_format has to be either NCHW or NHWC.') - layer_variable_getter = _build_variable_getter( - {'bias': 'biases', - 'depthwise_kernel': 'depthwise_weights', - 'pointwise_kernel': 'pointwise_weights'}) + layer_variable_getter = _build_variable_getter({ + 'bias': 'biases', + 'depthwise_kernel': 'depthwise_weights', + 'pointwise_kernel': 'pointwise_weights' + }) with variable_scope.variable_scope( - scope, 'SeparableConv2d', [inputs], reuse=reuse, + scope, + 'SeparableConv2d', [inputs], + reuse=reuse, custom_getter=layer_variable_getter) as sc: inputs = ops.convert_to_tensor(inputs) - df = ('channels_first' if data_format and data_format.startswith('NC') - else 'channels_last') + df = ('channels_first' + if data_format and data_format.startswith('NC') else 'channels_last') if num_outputs is not None: # Apply separable conv using the SeparableConvolution2D layer. layer = convolutional_layers.SeparableConvolution2D( @@ -2535,8 +2630,8 @@ def separable_convolution2d( _add_variable_to_collections(layer.pointwise_kernel, variables_collections, 'weights') if layer.bias is not None: - _add_variable_to_collections(layer.bias, - variables_collections, 'biases') + _add_variable_to_collections(layer.bias, variables_collections, + 'biases') if normalizer_fn is not None: normalizer_params = normalizer_params or {} @@ -2551,8 +2646,7 @@ def separable_convolution2d( weights_collections = utils.get_variable_collections( variables_collections, 'weights') - depthwise_shape = [kernel_h, kernel_w, - num_filters_in, depth_multiplier] + depthwise_shape = [kernel_h, kernel_w, num_filters_in, depth_multiplier] depthwise_weights = variables.model_variable( 'depthwise_weights', shape=depthwise_shape, @@ -2566,9 +2660,13 @@ def separable_convolution2d( 1, stride_h, stride_w, 1 ] - outputs = nn.depthwise_conv2d(inputs, depthwise_weights, strides, padding, - rate=utils.two_element_tuple(rate), - data_format=data_format) + outputs = nn.depthwise_conv2d( + inputs, + depthwise_weights, + strides, + padding, + rate=utils.two_element_tuple(rate), + data_format=data_format) num_outputs = depth_multiplier * num_filters_in if normalizer_fn is not None: @@ -2578,13 +2676,16 @@ def separable_convolution2d( if biases_initializer is not None: biases_collections = utils.get_variable_collections( variables_collections, 'biases') - biases = variables.model_variable('biases', - shape=[num_outputs,], - dtype=dtype, - initializer=biases_initializer, - regularizer=biases_regularizer, - trainable=trainable, - collections=biases_collections) + biases = variables.model_variable( + 'biases', + shape=[ + num_outputs, + ], + dtype=dtype, + initializer=biases_initializer, + regularizer=biases_regularizer, + trainable=trainable, + collections=biases_collections) outputs = nn.bias_add(outputs, biases, data_format=data_format) if activation_fn is not None: @@ -2592,6 +2693,36 @@ def separable_convolution2d( return utils.collect_named_outputs(outputs_collections, sc.name, outputs) +@add_arg_scope +def sequence_to_images(inputs, height, output_data_format='channels_last', + outputs_collections=None, scope=None): + """Convert a batch of sequences into a batch of images. + Args: + inputs: (num_steps, num_batches, depth) sequence tensor + height: the height of the images + output_data_format: Format of output tensor. + Currently supports `'channels_first'` and `'channels_last'`. + outputs_collections: The collections to which the outputs are added. + scope: Optional scope for name_scope. + Returns: + A tensor representing the output of the operation. + """ + with ops.name_scope(scope, 'SequenceToImages', [inputs]) as sc: + inputs = ops.convert_to_tensor(inputs) + width, num_batches, depth = inputs.get_shape().as_list() + if num_batches is None: + num_batches = -1 + else: + num_batches = num_batches // height + 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: + outputs = array_ops.transpose(reshaped, [1, 2, 0, 3]) + return utils.collect_named_outputs(outputs_collections, sc, outputs) + + @add_arg_scope def softmax(logits, scope=None): """Performs softmax on Nth dimension of N-dimensional logit tensor. @@ -2669,35 +2800,45 @@ def spatial_softmax(features, with ops.name_scope('spatial_softmax_op', 'spatial_softmax_op', [features]): # Create tensors for x and y coordinate values, scaled to range [-1, 1]. - pos_x, pos_y = array_ops.meshgrid(math_ops.lin_space(-1., 1., num=height), - math_ops.lin_space(-1., 1., num=width), - indexing='ij') + pos_x, pos_y = array_ops.meshgrid( + math_ops.lin_space(-1., 1., num=height), + math_ops.lin_space(-1., 1., num=width), + indexing='ij') pos_x = array_ops.reshape(pos_x, [height * width]) pos_y = array_ops.reshape(pos_y, [height * width]) + if temperature is None: - temperature_collections = utils.get_variable_collections( + temp_initializer = init_ops.ones_initializer() + else: + temp_initializer = init_ops.constant_initializer(temperature) + + if not trainable: + temp_collections = None + else: + temp_collections = utils.get_variable_collections( variables_collections, 'temperature') - temperature = variables.model_variable( - 'temperature', - shape=(), - dtype=dtypes.float32, - initializer=init_ops.ones_initializer(), - collections=temperature_collections, - trainable=trainable) + + temperature = variables.model_variable( + 'temperature', + shape=(), + dtype=dtypes.float32, + initializer=temp_initializer, + collections=temp_collections, + trainable=trainable) if data_format == 'NCHW': features = array_ops.reshape(features, [-1, height * width]) else: features = array_ops.reshape( array_ops.transpose(features, [0, 3, 1, 2]), [-1, height * width]) - softmax_attention = nn.softmax(features/temperature) + softmax_attention = nn.softmax(features / temperature) expected_x = math_ops.reduce_sum( - pos_x * softmax_attention, [1], keep_dims=True) + pos_x * softmax_attention, [1], keepdims=True) expected_y = math_ops.reduce_sum( - pos_y * softmax_attention, [1], keep_dims=True) + pos_y * softmax_attention, [1], keepdims=True) expected_xy = array_ops.concat([expected_x, expected_y], 1) - feature_keypoints = array_ops.reshape( - expected_xy, [-1, num_channels.value * 2]) + feature_keypoints = array_ops.reshape(expected_xy, + [-1, num_channels.value * 2]) feature_keypoints.set_shape([None, num_channels.value * 2]) return feature_keypoints @@ -2749,7 +2890,7 @@ def stack(inputs, layer, stack_args, **kwargs): scope = 'stack' outputs = inputs for i in range(len(stack_args)): - kwargs['scope'] = scope + '_' + str(i+1) + kwargs['scope'] = scope + '_' + str(i + 1) layer_args = stack_args[i] if not isinstance(layer_args, (list, tuple)): layer_args = [layer_args] @@ -2780,11 +2921,10 @@ def unit_norm(inputs, dim, epsilon=1e-7, scope=None): raise ValueError('The input rank must be known.') input_rank = len(inputs.get_shape().as_list()) if dim < 0 or dim >= input_rank: - raise ValueError( - 'dim must be positive but smaller than the input rank.') + raise ValueError('dim must be positive but smaller than the input rank.') - lengths = math_ops.sqrt(epsilon + math_ops.reduce_sum( - math_ops.square(inputs), dim, True)) + lengths = math_ops.sqrt( + epsilon + math_ops.reduce_sum(math_ops.square(inputs), dim, True)) multiples = [] if dim > 0: multiples.append(array_ops.ones([dim], dtypes.int32)) @@ -2828,7 +2968,7 @@ def poincare_normalize(x, axis=1, epsilon=1e-5, name=None): """ with ops.name_scope(name, 'poincare_normalize', [x]) as name: x = ops.convert_to_tensor(x, name='x') - square_sum = math_ops.reduce_sum(math_ops.square(x), axis, keep_dims=True) + square_sum = math_ops.reduce_sum(math_ops.square(x), axis, keepdims=True) x_inv_norm = math_ops.rsqrt(square_sum) x_inv_norm = math_ops.minimum((1. - epsilon) * x_inv_norm, 1.) return math_ops.multiply(x, x_inv_norm, name=name) @@ -2925,29 +3065,31 @@ def legacy_fully_connected(x, raise ValueError('last dimension of x must be known but is None') dtype = x.dtype.base_dtype - weight_collections = set(list(weight_collections or []) + - [ops.GraphKeys.GLOBAL_VARIABLES]) - w = variable_scope.get_variable('weights', - shape=[num_input_units, num_output_units], - dtype=dtype, - initializer=weight_init, - collections=weight_collections, - regularizer=weight_regularizer, - trainable=trainable) - x_2_dim = x if len(dims) <= 2 else array_ops.reshape(x, - [-1, num_input_units]) + weight_collections = set( + list(weight_collections or []) + [ops.GraphKeys.GLOBAL_VARIABLES]) + w = variable_scope.get_variable( + 'weights', + shape=[num_input_units, num_output_units], + dtype=dtype, + initializer=weight_init, + collections=weight_collections, + regularizer=weight_regularizer, + trainable=trainable) + x_2_dim = x if len(dims) <= 2 else array_ops.reshape( + x, [-1, num_input_units]) y = standard_ops.matmul(x_2_dim, w) if bias_init is not None: - bias_collections = set(list(bias_collections or []) + - [ops.GraphKeys.GLOBAL_VARIABLES]) - b = variable_scope.get_variable('bias', - shape=[num_output_units], - dtype=dtype, - initializer=bias_init, - collections=bias_collections, - regularizer=bias_regularizer, - trainable=trainable) + bias_collections = set( + list(bias_collections or []) + [ops.GraphKeys.GLOBAL_VARIABLES]) + b = variable_scope.get_variable( + 'bias', + shape=[num_output_units], + dtype=dtype, + initializer=bias_init, + collections=bias_collections, + regularizer=bias_regularizer, + trainable=trainable) y = nn.bias_add(y, b) diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py index 1150328b7aaab92e9dccbee09a10de51e717ce26..0f062adbab3ca9acfb89543b69c7c957bbdf5dd8 100644 --- a/tensorflow/contrib/layers/python/layers/layers_test.py +++ b/tensorflow/contrib/layers/python/layers/layers_test.py @@ -44,6 +44,7 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops from tensorflow.python.ops import partitioned_variables from tensorflow.python.ops import random_ops +from tensorflow.python.ops import sparse_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import template from tensorflow.python.ops import variable_scope @@ -126,8 +127,8 @@ class AvgPool3DTest(test.TestCase): def testInvalidDataFormat(self): depth, height, width = 3, 6, 9 images = np.random.uniform(size=(5, depth, height, width, 3)) - with self.assertRaisesRegexp(ValueError, - 'data_format has to be either NCDHW or NDHWC.'): + with self.assertRaisesRegexp( + ValueError, 'data_format has to be either NCDHW or NDHWC.'): _layers.avg_pool3d(images, [3, 3, 3], data_format='CDHWN') def testCreateAvgPool(self): @@ -147,7 +148,8 @@ class AvgPool3DTest(test.TestCase): def testCollectOutputs(self): depth, height, width = 3, 6, 9 images = random_ops.random_uniform((5, depth, height, width, 3), seed=1) - output = _layers.avg_pool3d(images, [3, 3, 3], outputs_collections='outputs') + output = _layers.avg_pool3d( + images, [3, 3, 3], outputs_collections='outputs') output_collected = ops.get_collection('outputs')[0] self.assertEqual(output_collected.aliases, ['AvgPool3D']) self.assertEqual(output_collected, output) @@ -182,7 +184,8 @@ class AvgPool3DTest(test.TestCase): depth, height, width = 3, 6, 9 images = random_ops.random_uniform((5, depth, height, width, 3), seed=1) output = _layers.avg_pool3d(images, [3, 3, 3], stride=1, padding='SAME') - self.assertListEqual(output.get_shape().as_list(), [5, depth, height, width, 3]) + self.assertListEqual(output.get_shape().as_list(), + [5, depth, height, width, 3]) def testGlobalAvgPool(self): depth, height, width = 3, 6, 9 @@ -514,7 +517,9 @@ class ConvolutionTest(test.TestCase): with arg_scope( [layers_lib.convolution2d], normalizer_fn=_layers.batch_norm, - normalizer_params={'decay': 0.9}): + normalizer_params={ + 'decay': 0.9 + }): net = layers_lib.convolution2d(images, 32, [3, 3]) net = layers_lib.convolution2d(net, 32, [3, 3]) self.assertEqual(len(variables.get_variables()), 8) @@ -528,7 +533,9 @@ class ConvolutionTest(test.TestCase): with arg_scope( [layers_lib.convolution2d], normalizer_fn=_layers.batch_norm, - normalizer_params={'decay': 0.9}): + normalizer_params={ + 'decay': 0.9 + }): net = layers_lib.convolution2d(images, 32, [3, 3], scope='Conv') net = layers_lib.convolution2d( net, 32, [3, 3], scope='Conv', reuse=True) @@ -701,7 +708,7 @@ class Convolution2dTransposeTests(test.TestCase): _layers.convolution2d_transpose(images, 32, 3, data_format='CHWN') def testOutputSizeWithStrideOneSamePaddingNCHW(self): - # `NCHW` data fomat is only supported for `GPU` device. + # `NCHW` data format is only supported for `GPU` device. if test.is_gpu_available(cuda_only=True): with self.test_session(use_gpu=True) as sess: num_filters = 32 @@ -1030,7 +1037,8 @@ class Convolution2dTransposeTests(test.TestCase): for _ in range(10): num_filters = 1 input_size = [ - 1, np.random.randint(1, max_image_size), + 1, + np.random.randint(1, max_image_size), np.random.randint(1, max_image_size), 1 ] filter_size = [ @@ -1184,8 +1192,10 @@ class ConvolutionInPlaneTest(test.TestCase): with self.test_session() as sess: sess.run(init_op) - result = sess.run(horz_gradients, - feed_dict={image: np.ones((1, 10, 10, 1))}) + result = sess.run( + horz_gradients, feed_dict={ + image: np.ones((1, 10, 10, 1)) + }) expected = np.zeros((1, 10, 9, 1)) self.assertAllEqual(result, expected) @@ -1292,6 +1302,19 @@ class ConvolutionInPlaneTest(test.TestCase): self.assertAllClose(result, expected, rtol=1e-5, atol=1e-5) +class DenseToSparseTest(test.TestCase): + + def testDenseFromConstantToSparse(self): + expected_constant = np.reshape(np.arange(24, dtype=np.int64), (3, 4, 2)) + tensor = constant_op.constant(expected_constant) + sparse = _layers.dense_to_sparse(tensor) + dense = sparse_ops.sparse_to_dense(sparse.indices, sparse.dense_shape, + sparse.values) + with self.test_session() as sess: + constant = sess.run(dense) + self.assertAllEqual(expected_constant, constant) + + class DropoutTest(test.TestCase): def testCreateDropout(self): @@ -1406,8 +1429,7 @@ class FlattenTest(test.TestCase): with ops.Graph().as_default() as g, self.test_session(g): inputs = array_ops.placeholder(dtype=dtypes.float32) inputs.set_shape(tensor_shape.TensorShape((5,))) - with self.assertRaisesRegexp(ValueError, - 'incompatible with the layer'): + with self.assertRaisesRegexp(ValueError, 'incompatible with the layer'): _layers.flatten(inputs) def testUnknownLastDim(self): @@ -1717,7 +1739,9 @@ class FCTest(test.TestCase): with arg_scope( [_layers.fully_connected], normalizer_fn=_layers.batch_norm, - normalizer_params={'decay': 0.9}): + normalizer_params={ + 'decay': 0.9 + }): net = _layers.fully_connected(images, 27) net = _layers.fully_connected(net, 27) self.assertEqual(len(variables.get_variables()), 8) @@ -1733,7 +1757,9 @@ class FCTest(test.TestCase): with arg_scope( [_layers.fully_connected], normalizer_fn=_layers.batch_norm, - normalizer_params={'decay': 0.9}): + normalizer_params={ + 'decay': 0.9 + }): net = _layers.fully_connected(images, 27, scope='fc1') net = _layers.fully_connected(net, 27, scope='fc1', reuse=True) self.assertEqual(len(variables.get_variables()), 4) @@ -1750,8 +1776,8 @@ class BatchNormTest(test.TestCase): def testBatchNormCenterFalse(self): a = array_ops.placeholder(dtype=dtypes.float32, shape=(10, 10, 10, 10)) # Test that center=False builds a valid graph. - _layers.batch_norm(a, center=False, data_format='NCHW', - zero_debias_moving_mean=True) + _layers.batch_norm( + a, center=False, data_format='NCHW', zero_debias_moving_mean=True) def testUnknownShape(self): with ops.Graph().as_default() as g, self.test_session(g): @@ -1788,8 +1814,8 @@ class BatchNormTest(test.TestCase): images = np.random.uniform(size=(5, height, width, 3)).astype( dtype.as_numpy_dtype) output = _layers.batch_norm(images, fused=fused) - expected_name = ('BatchNorm/FusedBatchNorm' if fused else - 'BatchNorm/batchnorm') + expected_name = ('BatchNorm/FusedBatchNorm' + if fused else 'BatchNorm/batchnorm') self.assertTrue(output.op.name.startswith(expected_name)) self.assertListEqual(output.get_shape().as_list(), [5, height, width, 3]) self.assertEqual( @@ -2008,8 +2034,8 @@ class BatchNormTest(test.TestCase): expected_var = np.var(image_values, axis=axis) if fused: # Add Bessel's correction - expected_var, _ = self._addBesselsCorrection(batch_size * height * - width, expected_var) + expected_var, _ = self._addBesselsCorrection( + batch_size * height * width, expected_var) images = constant_op.constant( image_values, shape=image_shape, dtype=dtypes.float32) output = _layers.batch_norm( @@ -2170,7 +2196,7 @@ class BatchNormTest(test.TestCase): # After initialization moving_mean == 0 and moving_variance == 1. self.assertAllClose(mean, [0] * 3) self.assertAllClose(variance, [1] * 3) - # Simulate assigment from saver restore. + # Simulate assignment from saver restore. init_assigns = [ state_ops.assign(moving_mean, expected_mean), state_ops.assign(moving_variance, expected_var) @@ -2528,8 +2554,8 @@ class BatchNormTest(test.TestCase): expected_var = np.var(image_values, axis=axis) if fused: # Add Bessel's correction - expected_var, _ = self._addBesselsCorrection(batch_size * height * - width, expected_var) + expected_var, _ = self._addBesselsCorrection( + batch_size * height * width, expected_var) images = constant_op.constant( image_values, shape=image_shape, dtype=dtypes.float32) output = _layers.batch_norm( @@ -2559,8 +2585,9 @@ class BatchNormTest(test.TestCase): np_output, new_images_gradients = sess.run([output, images_gradients]) # The outputs should be close to 0.0 mean and 1.0 variance self.assertAllClose( - np.mean( - np_output, axis=axis), [0] * channels, rtol=0.001, atol=0.001) + np.mean(np_output, axis=axis), [0] * channels, + rtol=0.001, + atol=0.001) self.assertAllClose( np.var(np_output, axis=axis), [1] * channels, rtol=0.01, atol=0.01) # The gradients should change slowly while updating moving_mean. @@ -2588,14 +2615,14 @@ class BatchNormTest(test.TestCase): channels = 3 with self.test_session() as sess: images = (np.ones((5, height, width, channels)) * 9.0).astype('f') - beta = init_ops.constant_initializer((np.ones(channels) * 5.0).astype( - 'f')) - gamma = init_ops.constant_initializer((np.ones(channels) * 2.0).astype( - 'f')) - mean = init_ops.constant_initializer((np.ones(channels) * 5.0).astype( - 'f')) - variance = init_ops.constant_initializer((np.ones(channels) * 4.0).astype( - 'f')) + beta = init_ops.constant_initializer( + (np.ones(channels) * 5.0).astype('f')) + gamma = init_ops.constant_initializer( + (np.ones(channels) * 2.0).astype('f')) + mean = init_ops.constant_initializer( + (np.ones(channels) * 5.0).astype('f')) + variance = init_ops.constant_initializer( + (np.ones(channels) * 4.0).astype('f')) output = _layers.batch_norm( images, is_training=False, @@ -2616,21 +2643,18 @@ class BatchNormTest(test.TestCase): with self.test_session(use_gpu=True) as sess: images = np.arange(np.product(shape), dtype=np.float32).reshape(shape) beta = init_ops.constant_initializer( - np.arange( - 2, channels + 2, dtype=np.float32)) + np.arange(2, channels + 2, dtype=np.float32)) gamma = init_ops.constant_initializer( - np.arange( - 10, channels + 10, dtype=np.float32) * 2.0) + np.arange(10, channels + 10, dtype=np.float32) * 2.0) mean = init_ops.constant_initializer( - np.arange( - 3, channels + 3, dtype=np.float32) * 5.0) + np.arange(3, channels + 3, dtype=np.float32) * 5.0) variance = init_ops.constant_initializer( - np.arange( - 1, channels + 1, dtype=np.float32) * 4.0) + np.arange(1, channels + 1, dtype=np.float32) * 4.0) if data_format == 'NCHW': # Reshape inputs from NHWC to NCHW format. images = array_ops.transpose( - images, [0, len(shape) - 1] + list(range(1, len(shape) - 1))) + images, [0, len(shape) - 1] + list(range(1, + len(shape) - 1))) output = _layers.batch_norm( images, is_training=is_training, @@ -2733,16 +2757,16 @@ class BatchNormTest(test.TestCase): # Tests that the adjustment is appropriately passed to and used by the core # BN layer. all_adjustments = [] + def _create_adjustment(shape): adjustments = [array_ops.ones(shape[-1:]), array_ops.zeros(shape[-1:])] all_adjustments.extend(adjustments) return adjustments + depth = 8 images = array_ops.zeros([10, 5, 5, depth]) output = _layers.batch_norm( - images, - is_training=True, - adjustment=_create_adjustment) + images, is_training=True, adjustment=_create_adjustment) self.assertListEqual(output.shape.as_list(), images.shape.as_list()) self.assertEqual(len(all_adjustments), 2) self.assertListEqual(all_adjustments[0].shape.as_list(), [depth]) @@ -2807,7 +2831,10 @@ class LayerNormTest(test.TestCase): # output_train and output_eval should be the same. self.assertAllClose(sess.run([output_train]), sess.run([output_eval])) - def doOutputTest(self, input_shape, tol=1e-5, begin_norm_axis=1, + def doOutputTest(self, + input_shape, + tol=1e-5, + begin_norm_axis=1, dtype=dtypes.float64): expected_mean = np.zeros(input_shape[:begin_norm_axis]) expected_var = np.ones(input_shape[:begin_norm_axis]) @@ -2838,13 +2865,10 @@ class LayerNormTest(test.TestCase): # Layer-norm implemented in numpy eps = 1e-12 expected_out = ( - (gamma * ( - input_values - - np.mean(input_values, axis=moments_axis, keepdims=True)) - / np.sqrt( - eps - + np.var(input_values, axis=moments_axis, keepdims=True))) - + beta) + (gamma * (input_values - np.mean( + input_values, axis=moments_axis, keepdims=True)) / + np.sqrt(eps + np.var( + input_values, axis=moments_axis, keepdims=True))) + beta) self.assertAllClose(expected_mean, mean, atol=tol, rtol=tol) self.assertAllClose(expected_var, var, atol=tol) # The full computation gets a bigger tolerance @@ -2862,10 +2886,10 @@ class LayerNormTest(test.TestCase): def testOutput4DInputNormOnInnermostAxis(self): # Equivalent tests - self.doOutputTest((100, 10, 10, 3), begin_norm_axis=3, tol=1e-4, - dtype=dtypes.float64) - self.doOutputTest((100, 10, 10, 3), begin_norm_axis=-1, tol=1e-4, - dtype=dtypes.float64) + self.doOutputTest( + (100, 10, 10, 3), begin_norm_axis=3, tol=1e-4, dtype=dtypes.float64) + self.doOutputTest( + (100, 10, 10, 3), begin_norm_axis=-1, tol=1e-4, dtype=dtypes.float64) def testOutputSmallInput(self): self.doOutputTest((10, 10, 10, 30)) @@ -2902,7 +2926,7 @@ class GDNTest(test.TestCase): x = np.random.uniform(size=(1, 2, 3, 4)[:ndim]) y = self._runGDN(x, x.shape, False, 'channels_last') self.assertEqual(x.shape, y.shape) - self.assertAllClose(y, x / np.sqrt(1 + .1 * (x ** 2)), rtol=0, atol=1e-6) + self.assertAllClose(y, x / np.sqrt(1 + .1 * (x**2)), rtol=0, atol=1e-6) def testChannelsFirst(self): # `bias_add` doesn't support NCHW on CPU. @@ -2911,8 +2935,7 @@ class GDNTest(test.TestCase): x = np.random.uniform(size=(4, 3, 2, 1)[:ndim]) y = self._runGDN(x, x.shape, False, 'channels_first') self.assertEqual(x.shape, y.shape) - self.assertAllClose( - y, x / np.sqrt(1 + .1 * (x ** 2)), rtol=0, atol=1e-6) + self.assertAllClose(y, x / np.sqrt(1 + .1 * (x**2)), rtol=0, atol=1e-6) def testWrongDims(self): for ndim in [1, 2, 6]: @@ -2924,7 +2947,29 @@ class GDNTest(test.TestCase): x = np.random.uniform(size=(1, 2, 3, 4)) y = self._runGDN(x, x.shape, True, 'channels_last') self.assertEqual(x.shape, y.shape) - self.assertAllClose(y, x * np.sqrt(1 + .1 * (x ** 2)), rtol=0, atol=1e-6) + self.assertAllClose(y, x * np.sqrt(1 + .1 * (x**2)), rtol=0, atol=1e-6) + + +class ImagesToSequenceTest(test.TestCase): + + def testInvalidDataFormat(self): + height, width = 7, 11 + images = np.random.uniform(size=(5, height, width, 2)) + with self.assertRaisesRegexp(ValueError, + 'data_format has to be either NCHW or NHWC.'): + _layers.images_to_sequence(images, data_format='CHWN') + + def testImagesToSequenceDims(self): + height, width = 7, 11 + images = np.random.uniform(size=(2, height, width, 5)).astype(np.float32) + output = _layers.images_to_sequence(images) + self.assertListEqual(output.get_shape().as_list(), [11, 14, 5]) + + def testImagesToSequenceNCHW(self): + height, width = 7, 11 + images = np.random.uniform(size=(2, 5, height, width)).astype(np.float32) + output = _layers.images_to_sequence(images, data_format='NCHW') + self.assertListEqual(output.get_shape().as_list(), [11, 14, 5]) class MaxPool2DTest(test.TestCase): @@ -3001,20 +3046,22 @@ class MaxPool3DTest(test.TestCase): def testInvalidDataFormat(self): depth, height, width = 3, 6, 9 images = np.random.uniform(size=(5, depth, height, width, 3)) - with self.assertRaisesRegexp(ValueError, - 'data_format has to be either NCDHW or NDHWC.'): + with self.assertRaisesRegexp( + ValueError, 'data_format has to be either NCDHW or NDHWC.'): _layers.max_pool3d(images, [3, 3, 3], data_format='CDHWN') def testCreateMaxPool(self): depth, height, width = 3, 6, 9 - images = np.random.uniform(size=(5, depth, height, width, 3)).astype(np.float32) + images = np.random.uniform(size=(5, depth, height, width, 3)).astype( + np.float32) output = _layers.max_pool3d(images, [3, 3, 3]) self.assertEqual(output.op.name, 'MaxPool3D/MaxPool3D') self.assertListEqual(output.get_shape().as_list(), [5, 1, 2, 4, 3]) def testCreateMaxPoolNCDHW(self): depth, height, width = 3, 6, 9 - images = np.random.uniform(size=(5, 3, depth, height, width)).astype(np.float32) + images = np.random.uniform(size=(5, 3, depth, height, width)).astype( + np.float32) output = _layers.max_pool3d(images, [3, 3, 3], data_format='NCDHW') self.assertEquals(output.op.name, 'MaxPool3D/transpose_1') self.assertListEqual(output.get_shape().as_list(), [5, 3, 1, 2, 4]) @@ -3022,7 +3069,8 @@ class MaxPool3DTest(test.TestCase): def testCollectOutputs(self): depth, height, width = 3, 6, 9 images = random_ops.random_uniform((5, depth, height, width, 3), seed=1) - output = _layers.max_pool3d(images, [3, 3, 3], outputs_collections='outputs') + output = _layers.max_pool3d( + images, [3, 3, 3], outputs_collections='outputs') output_collected = ops.get_collection('outputs')[0] self.assertEqual(output_collected.aliases, ['MaxPool3D']) self.assertEqual(output_collected, output) @@ -3057,7 +3105,8 @@ class MaxPool3DTest(test.TestCase): depth, height, width = 3, 6, 9 images = random_ops.random_uniform((5, depth, height, width, 3), seed=1) output = _layers.max_pool3d(images, [3, 3, 3], stride=1, padding='SAME') - self.assertListEqual(output.get_shape().as_list(), [5, depth, height, width, 3]) + self.assertListEqual(output.get_shape().as_list(), + [5, depth, height, width, 3]) def testGlobalMaxPool(self): depth, height, width = 3, 6, 9 @@ -3237,7 +3286,11 @@ class SeparableConv2dTest(test.TestCase): images = random_ops.random_uniform((5, height, width, 3), seed=1) regularizer = regularizers.l2_regularizer(0.01) layers_lib.separable_conv2d( - images, 32, [3, 3], 2, weights_regularizer=regularizer) + images, + 32, [3, 3], + 2, + weights_regularizer=regularizer, + weights_initializer=init_ops.ones_initializer()) self.assertEqual( len(ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)), 2) weight_decay = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)[0] @@ -3245,12 +3298,31 @@ class SeparableConv2dTest(test.TestCase): weight_decay.op.name, 'SeparableConv2d/depthwise_kernel/Regularizer/l2_regularizer') sess.run(variables_lib.global_variables_initializer()) - self.assertLessEqual(sess.run(weight_decay), 0.05) + depth_weight_one = sess.run(weight_decay) weight_decay = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)[1] self.assertEqual( weight_decay.op.name, 'SeparableConv2d/pointwise_kernel/Regularizer/l2_regularizer') - self.assertLessEqual(sess.run(weight_decay), 0.05) + pointwise_weight_one = sess.run(weight_decay) + + regularizer = regularizers.l2_regularizer(1.0) + layers_lib.separable_conv2d( + images, + 32, [3, 3], + 2, + weights_regularizer=regularizer, + weights_initializer=init_ops.ones_initializer()) + self.assertEqual( + len(ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)), 4) + weight_decay = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)[2] + sess.run(variables_lib.global_variables_initializer()) + depth_weight_two = sess.run(weight_decay) + weight_decay = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)[3] + pointwise_weight_two = sess.run(weight_decay) + + self.assertAllClose( + [100.0 * depth_weight_one, 100.0 * pointwise_weight_one], + [depth_weight_two, pointwise_weight_two]) def testReuseConvWithWeightDecay(self): height, width = 3, 3 @@ -3368,6 +3440,33 @@ class ScaleGradientTests(test.TestCase): np.testing.assert_array_equal([3 * 2], g_x.eval()) +class SequenceToImagesTest(test.TestCase): + + def testImagesToSequenceDims(self): + num_batches = 14 + 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) + self.assertListEqual(output.get_shape().as_list(), [2, 7, 11, 5]) + + def testImagesToSequenceNCHW(self): + num_batches = 14 + 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') + self.assertListEqual(output.get_shape().as_list(), [2, 5, 7, 11]) + + class SoftmaxTests(test.TestCase): def setUp(self): @@ -3446,8 +3545,7 @@ class SpatialSoftmaxTests(test.TestCase): sess.run(variables_lib.global_variables_initializer()) feed_dict = {features: np_features} keypoints = sess.run(spatial_softmax, feed_dict) - self.assertAllEqual(keypoints.shape, - (batch_shape[0], batch_shape[3] * 2)) + self.assertAllEqual(keypoints.shape, (batch_shape[0], batch_shape[3] * 2)) def testSpatialSoftmaxShapeNCHW(self): batch_shape = (2, 2, 35, 35) @@ -3458,8 +3556,7 @@ class SpatialSoftmaxTests(test.TestCase): sess.run(variables_lib.global_variables_initializer()) feed_dict = {features: np_features} keypoints = sess.run(spatial_softmax, feed_dict) - self.assertAllEqual(keypoints.shape, - (batch_shape[0], batch_shape[1] * 2)) + self.assertAllEqual(keypoints.shape, (batch_shape[0], batch_shape[1] * 2)) def testTwoMaxActivationsSameChannel(self): batch_size, height, width, nchannels = (2, 35, 35, 1) @@ -3478,8 +3575,8 @@ class SpatialSoftmaxTests(test.TestCase): x_loc = [avg_x] y_loc = [avg_y] - np_keypoints = self._SpatialSoftmax( - x_loc, y_loc, height, width, batch_size, nchannels) + np_keypoints = self._SpatialSoftmax(x_loc, y_loc, height, width, batch_size, + nchannels) # Make sure expected location keypoints matches actual location keypoints. with self.test_session() as sess: @@ -3497,13 +3594,13 @@ class SpatialSoftmaxTests(test.TestCase): spatial_softmax = _layers.spatial_softmax(features) np_features = np.zeros(batch_shape, dtype=np.float32) - edges = [(0, 0), (0, width-1), (height-1, 0), (height-1, width-1)] + edges = [(0, 0), (0, width - 1), (height - 1, 0), (height - 1, width - 1)] x_loc, y_loc = zip(*edges) for c in range(nchannels): np_features[:, x_loc[c], y_loc[c], c] = 100. - np_keypoints = self._SpatialSoftmax( - x_loc, y_loc, height, width, batch_size, nchannels) + np_keypoints = self._SpatialSoftmax(x_loc, y_loc, height, width, batch_size, + nchannels) # Make sure expected location keypoints matches actual location keypoints. with self.test_session() as sess: @@ -3532,10 +3629,10 @@ class SpatialSoftmaxTests(test.TestCase): np_features1[:, x_loc[c], y_loc[c], c] = 100. np_features2[:, x_loc[c], y_loc[c], c] = 100. - np_keypoints1 = self._SpatialSoftmax( - x_loc, y_loc, height1, width1, batch_size, nchannels) - np_keypoints2 = self._SpatialSoftmax( - x_loc, y_loc, height2, width2, batch_size, nchannels) + np_keypoints1 = self._SpatialSoftmax(x_loc, y_loc, height1, width1, + batch_size, nchannels) + np_keypoints2 = self._SpatialSoftmax(x_loc, y_loc, height2, width2, + batch_size, nchannels) # Make sure expected location keypoints matches actual location keypoints. with self.test_session() as sess: @@ -3561,8 +3658,8 @@ class SpatialSoftmaxTests(test.TestCase): for c in range(nchannels): np_features[:, x_loc[c], y_loc[c], c] = 100. - np_keypoints = self._SpatialSoftmax( - x_loc, y_loc, height, width, batch_size, nchannels) + np_keypoints = self._SpatialSoftmax(x_loc, y_loc, height, width, batch_size, + nchannels) # Make sure expected location keypoints matches actual location keypoints. with self.test_session() as sess: @@ -3584,8 +3681,8 @@ class SpatialSoftmaxTests(test.TestCase): for c in range(nchannels): np_features[:, c, x_loc[c], y_loc[c]] = 100. - np_keypoints = self._SpatialSoftmax( - x_loc, y_loc, height, width, batch_size, nchannels) + np_keypoints = self._SpatialSoftmax(x_loc, y_loc, height, width, batch_size, + nchannels) # Make sure expected location keypoints matches actual location keypoints. with self.test_session() as sess: @@ -3680,8 +3777,7 @@ class UnitNormTests(test.TestCase): image = random_ops.random_uniform((height, width, 3)) output = _layers.unit_norm(image, dim=dim, epsilon=1e-6) norms = math_ops.sqrt( - math_ops.reduce_sum( - math_ops.square(output), reduction_indices=dim)) + math_ops.reduce_sum(math_ops.square(output), reduction_indices=dim)) shape = [height, width, 3] del shape[dim] @@ -3717,8 +3813,7 @@ class UnitNormTests(test.TestCase): image = array_ops.placeholder(dtypes.float32, (None, None, 3)) output = _layers.unit_norm(image, dim=dim, epsilon=1e-6) norms = math_ops.sqrt( - math_ops.reduce_sum( - math_ops.square(output), reduction_indices=dim)) + math_ops.reduce_sum(math_ops.square(output), reduction_indices=dim)) with self.test_session(): actual = norms.eval({image: placeholder_value}) @@ -3782,8 +3877,8 @@ class PoincareNormalizeTest(test.TestCase): with self.test_session(): x_tf = constant_op.constant(x_np, name='x') y_tf = _layers.poincare_normalize(x_tf, dim) - err = gradient_checker.compute_gradient_error(x_tf, x_shape, - y_tf, x_shape) + err = gradient_checker.compute_gradient_error(x_tf, x_shape, y_tf, + x_shape) print('PoinCareNormalize gradient err = %g ' % err) self.assertLess(err, 1e-4) @@ -3795,14 +3890,9 @@ class LegacyFullyConnectedTest(test.TestCase): test.TestCase.setUp(self) random_seed.set_random_seed(1234) self.input = constant_op.constant([[1., 2., 3.], [-4., 15., -6.]]) - self.input_3_dim_arr = [[[1., 1.1, 1.2], - [2., 2.1, 2.2], - [3., 3.1, 3.2], - [4., 4.1, 4.2]], - [[5., 5.1, 5.2], - [6., 6.1, 6.2], - [7., 7.1, 7.2], - [8., 8.1, 8.2]]] + self.input_3_dim_arr = [[[1., 1.1, 1.2], [2., 2.1, 2.2], [3., 3.1, 3.2], + [4., 4.1, 4.2]], [[5., 5.1, 5.2], [6., 6.1, 6.2], + [7., 7.1, 7.2], [8., 8.1, 8.2]]] self.input_3_dim = constant_op.constant(self.input_3_dim_arr) assert not ops.get_collection(ops.GraphKeys.SUMMARIES) @@ -3897,15 +3987,10 @@ class LegacyFullyConnectedTest(test.TestCase): self._custom_initializers(self.input, 2, [[13.0, 13.0], [11.0, 11.0]]) def test_custom_initializers_multi_dim(self): - self._custom_initializers(self.input_3_dim, 2, - [[[7.6, 7.6], - [13.6, 13.6], - [19.6, 19.6], - [25.6, 25.6]], - [[31.6, 31.6], - [37.6, 37.6], - [43.6, 43.6], - [49.6, 49.6]]]) + self._custom_initializers( + self.input_3_dim, 2, + [[[7.6, 7.6], [13.6, 13.6], [19.6, 19.6], [25.6, 25.6]], + [[31.6, 31.6], [37.6, 37.6], [43.6, 43.6], [49.6, 49.6]]]) def test_custom_collections(self): layers_lib.legacy_relu( @@ -4015,12 +4100,16 @@ class LegacyFullyConnectedTest(test.TestCase): with self.test_session() as sess: variables_lib.global_variables_initializer().run() # we can feed in input with first dimension 2 - shape_value = sess.run(array_ops.shape(y), - feed_dict={x: self.input_3_dim_arr}) + shape_value = sess.run( + array_ops.shape(y), feed_dict={ + x: self.input_3_dim_arr + }) self.assertAllClose(shape_value, [2, 4, 1]) # we can feed in input with first dimension 1 - shape_value = sess.run(array_ops.shape(y), - feed_dict={x: [self.input_3_dim_arr[0]]}) + shape_value = sess.run( + array_ops.shape(y), feed_dict={ + x: [self.input_3_dim_arr[0]] + }) self.assertAllClose(shape_value, [1, 4, 1]) # we cannot feed in input with inconsistent dimensions with self.assertRaises(ValueError): diff --git a/tensorflow/contrib/learn/BUILD b/tensorflow/contrib/learn/BUILD index 5df2c77249b81434125d838f896f0ace2a5ee130..abf6e393bb0fbbce4e43f6d209e9b30517df36c3 100644 --- a/tensorflow/contrib/learn/BUILD +++ b/tensorflow/contrib/learn/BUILD @@ -172,6 +172,17 @@ tf_py_test( ], ) +py_test( + name = "export_strategy_test", + size = "small", + srcs = ["python/learn/export_strategy_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":learn", + "//tensorflow/python:client_testlib", + ], +) + py_test( name = "graph_actions_test", size = "small", @@ -377,6 +388,7 @@ py_test( "//tensorflow/python:client_testlib", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:lookup_ops", + "//tensorflow/python:math_ops", "//tensorflow/python:session", "//tensorflow/python:sparse_tensor", "//tensorflow/python:variables", @@ -483,7 +495,7 @@ py_test( name = "linear_test", size = "medium", srcs = ["python/learn/estimators/linear_test.py"], - shard_count = 4, + shard_count = 20, srcs_version = "PY2AND3", tags = ["no_pip"], deps = [ @@ -739,20 +751,6 @@ tf_py_test( grpc_enabled = True, ) -py_test( - name = "numpy_io_test", - size = "small", - srcs = ["python/learn/learn_io/numpy_io_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":learn", - "//tensorflow/python:client_testlib", - "//tensorflow/python:errors", - "//tensorflow/python:training", - "//third_party/py/numpy", - ], -) - py_test( name = "pandas_io_test", size = "small", diff --git a/tensorflow/contrib/learn/python/learn/datasets/__init__.py b/tensorflow/contrib/learn/python/learn/datasets/__init__.py index a3521b4109ab40d8478f20afc317cf5154da2b43..7240b0de149051afa045a8113f9e9b212840c311 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/__init__.py +++ b/tensorflow/contrib/learn/python/learn/datasets/__init__.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Dataset utilities and synthetic/reference datasets.""" from __future__ import absolute_import @@ -46,11 +45,12 @@ DATASETS = { # List of all synthetic datasets SYNTHETIC = { - # All of these will return ['data', 'target'] -> base.Dataset - 'circles': synthetic.circles, - 'spirals': synthetic.spirals + # All of these will return ['data', 'target'] -> base.Dataset + 'circles': synthetic.circles, + 'spirals': synthetic.spirals } + def load_dataset(name, size='small', test_with_fake_data=False): """Loads dataset by name. @@ -83,23 +83,28 @@ def make_dataset(name, n_samples=100, noise=None, seed=42, *args, **kwargs): seed: int or None, seed for noise Returns: - Shuffled features and labels for given synthetic dataset of type `base.Dataset` + Shuffled features and labels for given synthetic dataset of type + `base.Dataset` Raises: ValueError: Raised if `name` not found Note: - - This is a generic synthetic data generator - individual generators might have more parameters! + - This is a generic synthetic data generator - individual generators might + have more parameters! See documentation for individual parameters - - Note that the `noise` parameter uses `numpy.random.normal` and depends on `numpy`'s seed + - Note that the `noise` parameter uses `numpy.random.normal` and depends on + `numpy`'s seed TODO: - Support multiclass datasets - - Need shuffling routine. Currently synthetic datasets are reshuffled to avoid train/test correlation, + - Need shuffling routine. Currently synthetic datasets are reshuffled to + avoid train/test correlation, but that hurts reprodusability """ # seed = kwargs.pop('seed', None) if name not in SYNTHETIC: raise ValueError('Synthetic dataset not found or not implemeted: %s' % name) else: - return SYNTHETIC[name](n_samples=n_samples, noise=noise, seed=seed, *args, **kwargs) + return SYNTHETIC[name]( + n_samples=n_samples, noise=noise, seed=seed, *args, **kwargs) diff --git a/tensorflow/contrib/learn/python/learn/datasets/base.py b/tensorflow/contrib/learn/python/learn/datasets/base.py index 71978d439449e29c7cb907b18bab5d6659a972b6..ca720ae5ed26e74da12bd6c5a37231b41442f76f 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/base.py +++ b/tensorflow/contrib/learn/python/learn/datasets/base.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Base utilities for loading datasets.""" from __future__ import absolute_import @@ -24,13 +23,11 @@ import csv import os from os import path import random -import tempfile import time import numpy as np from six.moves import urllib -from tensorflow.contrib.framework import deprecated from tensorflow.python.platform import gfile Dataset = collections.namedtuple('Dataset', ['data', 'target']) @@ -100,9 +97,7 @@ def load_iris(data_path=None): module_path = path.dirname(__file__) data_path = path.join(module_path, 'data', 'iris.csv') return load_csv_with_header( - data_path, - target_dtype=np.int, - features_dtype=np.float) + data_path, target_dtype=np.int, features_dtype=np.float) def load_boston(data_path=None): @@ -118,16 +113,10 @@ def load_boston(data_path=None): module_path = path.dirname(__file__) data_path = path.join(module_path, 'data', 'boston_house_prices.csv') return load_csv_with_header( - data_path, - target_dtype=np.float, - features_dtype=np.float) + data_path, target_dtype=np.float, features_dtype=np.float) -def retry(initial_delay, - max_delay, - factor=2.0, - jitter=0.25, - is_retriable=None): +def retry(initial_delay, max_delay, factor=2.0, jitter=0.25, is_retriable=None): """Simple decorator for wrapping retriable functions. Args: @@ -152,7 +141,7 @@ def retry(initial_delay, def delays(): delay = initial_delay while delay <= max_delay: - yield delay * random.uniform(1 - jitter, 1 + jitter) + yield delay * random.uniform(1 - jitter, 1 + jitter) delay *= factor def wrap(fn): @@ -172,7 +161,9 @@ def retry(initial_delay, else: raise return fn(*args, **kwargs) + return wrapped_fn + return wrap diff --git a/tensorflow/contrib/learn/python/learn/datasets/mnist.py b/tensorflow/contrib/learn/python/learn/datasets/mnist.py index 1f3295747e141760445b021bf4f59cc47b88b8b2..37f9175015a239f763c7721cf36ab8063c0a3e32 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/mnist.py +++ b/tensorflow/contrib/learn/python/learn/datasets/mnist.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Functions for downloading and reading MNIST data.""" from __future__ import absolute_import @@ -123,8 +122,8 @@ class DataSet(object): numpy.random.seed(seed1 if seed is None else seed2) dtype = dtypes.as_dtype(dtype).base_dtype if dtype not in (dtypes.uint8, dtypes.float32): - raise TypeError('Invalid image dtype %r, expected uint8 or float32' % - dtype) + raise TypeError( + 'Invalid image dtype %r, expected uint8 or float32' % dtype) if fake_data: self._num_examples = 10000 self.one_hot = one_hot @@ -202,7 +201,9 @@ class DataSet(object): end = self._index_in_epoch images_new_part = self._images[start:end] labels_new_part = self._labels[start:end] - return numpy.concatenate((images_rest_part, images_new_part), axis=0) , numpy.concatenate((labels_rest_part, labels_new_part), axis=0) + return numpy.concatenate( + (images_rest_part, images_new_part), axis=0), numpy.concatenate( + (labels_rest_part, labels_new_part), axis=0) else: self._index_in_epoch += batch_size end = self._index_in_epoch @@ -257,16 +258,14 @@ def read_data_sets(train_dir, test_labels = extract_labels(f, one_hot=one_hot) if not 0 <= validation_size <= len(train_images): - raise ValueError( - 'Validation size should be between 0 and {}. Received: {}.' - .format(len(train_images), validation_size)) + raise ValueError('Validation size should be between 0 and {}. Received: {}.' + .format(len(train_images), validation_size)) validation_images = train_images[:validation_size] validation_labels = train_labels[:validation_size] train_images = train_images[validation_size:] train_labels = train_labels[validation_size:] - options = dict(dtype=dtype, reshape=reshape, seed=seed) train = DataSet(train_images, train_labels, **options) diff --git a/tensorflow/contrib/learn/python/learn/datasets/synthetic.py b/tensorflow/contrib/learn/python/learn/datasets/synthetic.py index 907dc0f3dfced7e55c5f46711fbe93f6400e1de7..9a843168c27d9cae3f55efe4fe4c688d86c745f3 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/synthetic.py +++ b/tensorflow/contrib/learn/python/learn/datasets/synthetic.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Synthetic dataset generators.""" from __future__ import absolute_import @@ -23,18 +22,27 @@ import numpy as np from tensorflow.contrib.learn.python.learn.datasets.base import Dataset -def circles(n_samples=100, noise=None, seed=None, factor=0.8, n_classes=2, *args, **kwargs): + +def circles(n_samples=100, + noise=None, + seed=None, + factor=0.8, + n_classes=2, + *args, + **kwargs): """Create circles separated by some value Args: n_samples: int, number of datapoints to generate noise: float or None, standard deviation of the Gaussian noise added seed: int or None, seed for the noise - factor: float, size factor of the inner circles with respect to the outer ones + factor: float, size factor of the inner circles with respect to the outer + ones n_classes: int, number of classes to generate Returns: - Shuffled features and labels for 'circles' synthetic dataset of type `base.Dataset` + Shuffled features and labels for 'circles' synthetic dataset of type + `base.Dataset` Note: The multi-class support might not work as expected if `noise` is enabled @@ -54,7 +62,7 @@ def circles(n_samples=100, noise=None, seed=None, factor=0.8, n_classes=2, *args if seed is not None: np.random.seed(seed) # Algo: 1) Generate initial circle, 2) For ever class generate a smaller radius circle - linspace = np.linspace(0, 2*np.pi, n_samples // n_classes) + linspace = np.linspace(0, 2 * np.pi, n_samples // n_classes) circ_x = np.empty(0, dtype=np.int32) circ_y = np.empty(0, dtype=np.int32) base_cos = np.cos(linspace) @@ -66,12 +74,12 @@ def circles(n_samples=100, noise=None, seed=None, factor=0.8, n_classes=2, *args circ_y = np.append(circ_y, base_sin) base_cos *= factor base_sin *= factor - y = np.append(y, label*np.ones(n_samples // n_classes, dtype=np.int32)) + y = np.append(y, label * np.ones(n_samples // n_classes, dtype=np.int32)) # Add more points if n_samples is not divisible by n_classes (unbalanced!) extras = n_samples % n_classes - circ_x = np.append(circ_x, np.cos(np.random.rand(extras)*2*np.pi)) - circ_y = np.append(circ_y, np.sin(np.random.rand(extras)*2*np.pi)) + circ_x = np.append(circ_x, np.cos(np.random.rand(extras) * 2 * np.pi)) + circ_y = np.append(circ_y, np.sin(np.random.rand(extras) * 2 * np.pi)) y = np.append(y, np.zeros(extras, dtype=np.int32)) # Reshape the features/labels @@ -85,10 +93,13 @@ def circles(n_samples=100, noise=None, seed=None, factor=0.8, n_classes=2, *args return Dataset(data=X[indices], target=y[indices]) -def spirals(n_samples=100, noise=None, seed=None, - mode = 'archimedes', - n_loops = 2, - *args, **kwargs): +def spirals(n_samples=100, + noise=None, + seed=None, + mode='archimedes', + n_loops=2, + *args, + **kwargs): """Create spirals Currently only binary classification is supported for spiral generation @@ -104,7 +115,8 @@ def spirals(n_samples=100, noise=None, seed=None, 'fermat': a spiral with branch distances decreasing (sqrt) Returns: - Shuffled features and labels for 'spirals' synthetic dataset of type `base.Dataset` + Shuffled features and labels for 'spirals' synthetic dataset of type + `base.Dataset` Raises: ValueError: If the generation `mode` is not valid @@ -112,34 +124,35 @@ def spirals(n_samples=100, noise=None, seed=None, TODO: - Generation of unbalanced data """ - n_classes = 2 # I am not sure how to make it multiclass + n_classes = 2 # I am not sure how to make it multiclass _modes = { - 'archimedes': _archimedes_spiral, - 'bernoulli': _bernoulli_spiral, - 'fermat': _fermat_spiral + 'archimedes': _archimedes_spiral, + 'bernoulli': _bernoulli_spiral, + 'fermat': _fermat_spiral } if mode is None or mode not in _modes: - raise ValueError("Cannot generate spiral with mode %s"%mode) + raise ValueError('Cannot generate spiral with mode %s' % mode) if seed is not None: np.random.seed(seed) - linspace = np.linspace(0, 2*n_loops*np.pi, n_samples // n_classes) + linspace = np.linspace(0, 2 * n_loops * np.pi, n_samples // n_classes) spir_x = np.empty(0, dtype=np.int32) spir_y = np.empty(0, dtype=np.int32) y = np.empty(0, dtype=np.int32) for label in range(n_classes): - base_cos, base_sin = _modes[mode](linspace, label*np.pi, *args, **kwargs) + base_cos, base_sin = _modes[mode](linspace, label * np.pi, *args, **kwargs) spir_x = np.append(spir_x, base_cos) spir_y = np.append(spir_y, base_sin) - y = np.append(y, label*np.ones(n_samples // n_classes, dtype=np.int32)) + y = np.append(y, label * np.ones(n_samples // n_classes, dtype=np.int32)) # Add more points if n_samples is not divisible by n_classes (unbalanced!) extras = n_samples % n_classes if extras > 0: - x_exrta, y_extra = _modes[mode](np.random.rand(extras)*2*np.pi, *args, **kwargs) + x_extra, y_extra = _modes[mode](np.random.rand(extras) * 2 * np.pi, *args, + **kwargs) spir_x = np.append(spir_x, x_extra) spir_y = np.append(spir_y, y_extra) y = np.append(y, np.zeros(extras, dtype=np.int32)) @@ -162,7 +175,8 @@ def _archimedes_spiral(theta, theta_offset=0., *args, **kwargs): theta: array-like, angles from polar coordinates to be converted theta_offset: float, angle offset in radians (2*pi = 0) """ - x, y = theta*np.cos(theta + theta_offset), theta*np.sin(theta + theta_offset) + x, y = theta * np.cos(theta + theta_offset), theta * np.sin( + theta + theta_offset) x_norm = np.max(np.abs(x)) y_norm = np.max(np.abs(y)) x, y = x / x_norm, y / y_norm @@ -181,7 +195,8 @@ def _bernoulli_spiral(theta, theta_offset=0., *args, **kwargs): """ exp_scale = kwargs.pop('exp_scale', 0.1) - x, y = np.exp(exp_scale*theta)*np.cos(theta + theta_offset), np.exp(exp_scale*theta)*np.sin(theta + theta_offset) + x, y = np.exp(exp_scale * theta) * np.cos(theta + theta_offset), np.exp( + exp_scale * theta) * np.sin(theta + theta_offset) x_norm = np.max(np.abs(x)) y_norm = np.max(np.abs(y)) x, y = x / x_norm, y / y_norm @@ -195,7 +210,8 @@ def _fermat_spiral(theta, theta_offset=0., *args, **kwargs): theta: array-like, angles from polar coordinates to be converted theta_offset: float, angle offset in radians (2*pi = 0) """ - x, y = np.sqrt(theta)*np.cos(theta + theta_offset), np.sqrt(theta)*np.sin(theta + theta_offset) + x, y = np.sqrt(theta) * np.cos(theta + theta_offset), np.sqrt(theta) * np.sin( + theta + theta_offset) x_norm = np.max(np.abs(x)) y_norm = np.max(np.abs(y)) x, y = x / x_norm, y / y_norm diff --git a/tensorflow/contrib/learn/python/learn/datasets/synthetic_test.py b/tensorflow/contrib/learn/python/learn/datasets/synthetic_test.py index 5340afab46eba957d6d612bb583983b627537547..5809995c8c7d8e72eb47ee88a72547bae7fd3594 100644 --- a/tensorflow/contrib/learn/python/learn/datasets/synthetic_test.py +++ b/tensorflow/contrib/learn/python/learn/datasets/synthetic_test.py @@ -24,12 +24,14 @@ from tensorflow.python.platform import test from tensorflow.contrib.learn.python.learn import datasets from tensorflow.contrib.learn.python.learn.datasets import synthetic + class SyntheticTest(test.TestCase): """Test synthetic dataset generation""" def test_make_dataset(self): """Test if the synthetic routine wrapper complains about the name""" - self.assertRaises(ValueError, datasets.make_dataset, name='_non_existing_name') + self.assertRaises( + ValueError, datasets.make_dataset, name='_non_existing_name') def test_all_datasets_callable(self): """Test if all methods inside the `SYNTHETIC` are callable""" @@ -52,9 +54,10 @@ class SyntheticTest(test.TestCase): """ n_samples = 100 n_classes = 2 - circ = synthetic.circles(n_samples = n_samples, noise = None, n_classes = n_classes) + circ = synthetic.circles( + n_samples=n_samples, noise=None, n_classes=n_classes) self.assertIsInstance(circ, datasets.base.Dataset) - self.assertTupleEqual(circ.data.shape, (n_samples,2)) + self.assertTupleEqual(circ.data.shape, (n_samples, 2)) self.assertTupleEqual(circ.target.shape, (n_samples,)) self.assertSetEqual(set(circ.target), set(range(n_classes))) @@ -67,17 +70,24 @@ class SyntheticTest(test.TestCase): """ seed = 42 noise = 0.1 - circ0 = synthetic.circles(n_samples = 100, noise = noise, n_classes = 2, seed = seed) - circ1 = synthetic.circles(n_samples = 100, noise = noise, n_classes = 2, seed = seed) + circ0 = synthetic.circles( + n_samples=100, noise=noise, n_classes=2, seed=seed) + circ1 = synthetic.circles( + n_samples=100, noise=noise, n_classes=2, seed=seed) np.testing.assert_array_equal(circ0.data, circ1.data) np.testing.assert_array_equal(circ0.target, circ1.target) - circ1 = synthetic.circles(n_samples = 100, noise = noise, n_classes = 2, seed = seed+1) - self.assertRaises(AssertionError, np.testing.assert_array_equal, circ0.data, circ1.data) - self.assertRaises(AssertionError, np.testing.assert_array_equal, circ0.target, circ1.target) + circ1 = synthetic.circles( + n_samples=100, noise=noise, n_classes=2, seed=seed + 1) + self.assertRaises(AssertionError, np.testing.assert_array_equal, circ0.data, + circ1.data) + self.assertRaises(AssertionError, np.testing.assert_array_equal, + circ0.target, circ1.target) - circ1 = synthetic.circles(n_samples = 100, noise = noise/2., n_classes = 2, seed = seed) - self.assertRaises(AssertionError, np.testing.assert_array_equal, circ0.data, circ1.data) + circ1 = synthetic.circles( + n_samples=100, noise=noise / 2., n_classes=2, seed=seed) + self.assertRaises(AssertionError, np.testing.assert_array_equal, circ0.data, + circ1.data) def test_spirals(self): """Test if the circles are generated correctly @@ -89,13 +99,14 @@ class SyntheticTest(test.TestCase): - returned `target` shape is (n_samples,) - set of unique classes range is [0, n_classes) """ - self.assertRaises(ValueError, synthetic.spirals, mode='_unknown_mode_spiral_') + self.assertRaises( + ValueError, synthetic.spirals, mode='_unknown_mode_spiral_') n_samples = 100 modes = ('archimedes', 'bernoulli', 'fermat') for mode in modes: - spir = synthetic.spirals(n_samples = n_samples, noise = None, mode = mode) + spir = synthetic.spirals(n_samples=n_samples, noise=None, mode=mode) self.assertIsInstance(spir, datasets.base.Dataset) - self.assertTupleEqual(spir.data.shape, (n_samples,2)) + self.assertTupleEqual(spir.data.shape, (n_samples, 2)) self.assertTupleEqual(spir.target.shape, (n_samples,)) self.assertSetEqual(set(spir.target), set(range(2))) @@ -110,18 +121,24 @@ class SyntheticTest(test.TestCase): noise = 0.1 modes = ('archimedes', 'bernoulli', 'fermat') for mode in modes: - spir0 = synthetic.spirals(n_samples = 1000, noise = noise, seed = seed) - spir1 = synthetic.spirals(n_samples = 1000, noise = noise, seed = seed) + spir0 = synthetic.spirals(n_samples=1000, noise=noise, seed=seed) + spir1 = synthetic.spirals(n_samples=1000, noise=noise, seed=seed) np.testing.assert_array_equal(spir0.data, spir1.data) np.testing.assert_array_equal(spir0.target, spir1.target) - spir1 = synthetic.spirals(n_samples = 1000, noise = noise, seed = seed+1) - self.assertRaises(AssertionError, np.testing.assert_array_equal, spir0.data, spir1.data) - self.assertRaises(AssertionError, np.testing.assert_array_equal, spir0.target, spir1.target) + spir1 = synthetic.spirals(n_samples=1000, noise=noise, seed=seed + 1) + self.assertRaises(AssertionError, np.testing.assert_array_equal, + spir0.data, spir1.data) + self.assertRaises(AssertionError, np.testing.assert_array_equal, + spir0.target, spir1.target) + + spir1 = synthetic.spirals(n_samples=1000, noise=noise / 2., seed=seed) + self.assertRaises(AssertionError, np.testing.assert_array_equal, + spir0.data, spir1.data) - spir1 = synthetic.spirals(n_samples = 1000, noise = noise/2., seed = seed) - self.assertRaises(AssertionError, np.testing.assert_array_equal, spir0.data, spir1.data) + def test_spirals_synthetic(self): + synthetic.spirals(3) -if __name__ == "__main__": +if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/learn/python/learn/estimators/debug_test.py b/tensorflow/contrib/learn/python/learn/estimators/debug_test.py index 6b125534a42c5cdde69773d99cefd6e7b2d60c9c..b968aeed1b7a11d522b531783f04f0104b37904f 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/debug_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/debug_test.py @@ -44,7 +44,6 @@ from tensorflow.python.ops import math_ops from tensorflow.python.platform import test from tensorflow.python.training import input as input_lib - NUM_EXAMPLES = 100 N_CLASSES = 5 # Cardinality of multiclass labels. LABEL_DIMENSION = 3 # Dimensionality of regression labels. @@ -52,8 +51,10 @@ LABEL_DIMENSION = 3 # Dimensionality of regression labels. def _train_test_split(features_and_labels): features, labels = features_and_labels - train_set = (features[:int(len(features) / 2)], labels[:int(len(features) / 2)]) - test_set = (features[int(len(features) / 2):], labels[int(len(features) / 2):]) + train_set = (features[:int(len(features) / 2)], + labels[:int(len(features) / 2)]) + test_set = (features[int(len(features) / 2):], + labels[int(len(features) / 2):]) return train_set, test_set @@ -86,17 +87,17 @@ class DebugClassifierTest(test.TestCase): (train_features, train_labels), (test_features, test_labels) = _train_test_split( [self.features, self.labels]) - majority_class, _ = max(collections.Counter(train_labels).items(), - key=operator.itemgetter(1)) + majority_class, _ = max( + collections.Counter(train_labels).items(), key=operator.itemgetter(1)) expected_prediction = np.vstack( [[majority_class] for _ in range(test_labels.shape[0])]) classifier = debug.DebugClassifier(n_classes=N_CLASSES) - classifier.fit(input_fn=_input_fn_builder(train_features, train_labels), - steps=50) + classifier.fit( + input_fn=_input_fn_builder(train_features, train_labels), steps=50) - pred = classifier.predict_classes(input_fn=_input_fn_builder(test_features, - None)) + pred = classifier.predict_classes( + input_fn=_input_fn_builder(test_features, None)) self.assertAllEqual(expected_prediction, np.vstack(pred)) def testPredictBinary(self): @@ -105,34 +106,34 @@ class DebugClassifierTest(test.TestCase): test_labels) = _train_test_split( [self.features, self.binary_labels]) - majority_class, _ = max(collections.Counter(train_labels).items(), - key=operator.itemgetter(1)) + majority_class, _ = max( + collections.Counter(train_labels).items(), key=operator.itemgetter(1)) expected_prediction = np.vstack( [[majority_class] for _ in range(test_labels.shape[0])]) classifier = debug.DebugClassifier(n_classes=2) - classifier.fit(input_fn=_input_fn_builder(train_features, train_labels), - steps=50) + classifier.fit( + input_fn=_input_fn_builder(train_features, train_labels), steps=50) - pred = classifier.predict_classes(input_fn=_input_fn_builder(test_features, - None)) + pred = classifier.predict_classes( + input_fn=_input_fn_builder(test_features, None)) self.assertAllEqual(expected_prediction, np.vstack(pred)) - (train_features, train_labels), ( - test_features, test_labels) = _train_test_split( - [self.features, self.binary_float_labels]) + (train_features, + train_labels), (test_features, test_labels) = _train_test_split( + [self.features, self.binary_float_labels]) - majority_class, _ = max(collections.Counter(train_labels).items(), - key=operator.itemgetter(1)) + majority_class, _ = max( + collections.Counter(train_labels).items(), key=operator.itemgetter(1)) expected_prediction = np.vstack( [[majority_class] for _ in range(test_labels.shape[0])]) classifier = debug.DebugClassifier(n_classes=2) - classifier.fit(input_fn=_input_fn_builder(train_features, train_labels), - steps=50) + classifier.fit( + input_fn=_input_fn_builder(train_features, train_labels), steps=50) - pred = classifier.predict_classes(input_fn=_input_fn_builder(test_features, - None)) + pred = classifier.predict_classes( + input_fn=_input_fn_builder(test_features, None)) self.assertAllEqual(expected_prediction, np.vstack(pred)) def testPredictProba(self): @@ -150,8 +151,8 @@ class DebugClassifierTest(test.TestCase): [class_distribution for _ in range(test_labels.shape[0])]) classifier = debug.DebugClassifier(n_classes=N_CLASSES) - classifier.fit(input_fn=_input_fn_builder(train_features, train_labels), - steps=50) + classifier.fit( + input_fn=_input_fn_builder(train_features, train_labels), steps=50) pred = classifier.predict_proba( input_fn=_input_fn_builder(test_features, None)) @@ -173,17 +174,17 @@ class DebugClassifierTest(test.TestCase): [class_distribution for _ in range(test_labels.shape[0])]) classifier = debug.DebugClassifier(n_classes=2) - classifier.fit(input_fn=_input_fn_builder(train_features, train_labels), - steps=50) + classifier.fit( + input_fn=_input_fn_builder(train_features, train_labels), steps=50) pred = classifier.predict_proba( input_fn=_input_fn_builder(test_features, None)) self.assertAllClose(expected_prediction, np.vstack(pred), atol=0.1) - (train_features, train_labels), ( - test_features, test_labels) = _train_test_split( - [self.features, self.binary_float_labels]) + (train_features, + train_labels), (test_features, test_labels) = _train_test_split( + [self.features, self.binary_float_labels]) class_distribution = np.zeros((1, 2)) for label in train_labels: @@ -194,8 +195,8 @@ class DebugClassifierTest(test.TestCase): [class_distribution for _ in range(test_labels.shape[0])]) classifier = debug.DebugClassifier(n_classes=2) - classifier.fit(input_fn=_input_fn_builder(train_features, train_labels), - steps=50) + classifier.fit( + input_fn=_input_fn_builder(train_features, train_labels), steps=50) pred = classifier.predict_proba( input_fn=_input_fn_builder(test_features, None)) @@ -232,13 +233,12 @@ class DebugClassifierTest(test.TestCase): def _input_fn(): iris = test_data.prepare_iris_data_for_logistic_regression() return { - 'feature': constant_op.constant( - iris.data, dtype=dtypes.float32) + 'feature': constant_op.constant(iris.data, dtype=dtypes.float32) }, constant_op.constant( iris.target, shape=[100], dtype=dtypes.int32) - classifier = debug.DebugClassifier(config=run_config.RunConfig( - tf_random_seed=1)) + classifier = debug.DebugClassifier( + config=run_config.RunConfig(tf_random_seed=1)) classifier.fit(input_fn=_input_fn, steps=5) scores = classifier.evaluate(input_fn=_input_fn, steps=1) self.assertIn('loss', scores) @@ -342,8 +342,7 @@ class DebugClassifierTest(test.TestCase): def _input_fn(): iris = base.load_iris() return { - 'feature': constant_op.constant( - iris.data, dtype=dtypes.float32) + 'feature': constant_op.constant(iris.data, dtype=dtypes.float32) }, constant_op.constant( iris.target, shape=[150], dtype=dtypes.int32) @@ -387,7 +386,9 @@ class DebugClassifierTest(test.TestCase): # Create 4 rows, one of them (y = x), three of them (y=Not(x)) # The logistic prediction should be (y = 0.25). labels = constant_op.constant([[1], [0], [0], [0]]) - features = {'x': array_ops.ones(shape=[4, 1], dtype=dtypes.float32),} + features = { + 'x': array_ops.ones(shape=[4, 1], dtype=dtypes.float32), + } return features, labels classifier = debug.DebugClassifier(n_classes=2) @@ -404,8 +405,7 @@ class DebugClassifierTest(test.TestCase): # The logistic prediction should be (y = 0.25). labels = constant_op.constant([[1.], [0.], [0.], [0.]]) features = { - 'x': array_ops.ones( - shape=[4, 1], dtype=dtypes.float32), + 'x': array_ops.ones(shape=[4, 1], dtype=dtypes.float32), 'w': constant_op.constant([[1.], [1.], [1.], [1.]]) } return features, labels @@ -414,8 +414,7 @@ class DebugClassifierTest(test.TestCase): # 4 rows, with different weights. labels = constant_op.constant([[1.], [0.], [0.], [0.]]) features = { - 'x': array_ops.ones( - shape=[4, 1], dtype=dtypes.float32), + 'x': array_ops.ones(shape=[4, 1], dtype=dtypes.float32), 'w': constant_op.constant([[7.], [1.], [1.], [1.]]) } return features, labels @@ -438,8 +437,7 @@ class DebugClassifierTest(test.TestCase): # than (y=Not(x)) due to the relative higher weight of the first row. labels = constant_op.constant([[1], [0], [0], [0]]) features = { - 'x': array_ops.ones( - shape=[4, 1], dtype=dtypes.float32), + 'x': array_ops.ones(shape=[4, 1], dtype=dtypes.float32), 'w': constant_op.constant([[100.], [3.], [2.], [2.]]) } return features, labels @@ -448,8 +446,7 @@ class DebugClassifierTest(test.TestCase): # Create 4 rows (y = x) labels = constant_op.constant([[1], [1], [1], [1]]) features = { - 'x': array_ops.ones( - shape=[4, 1], dtype=dtypes.float32), + 'x': array_ops.ones(shape=[4, 1], dtype=dtypes.float32), 'w': constant_op.constant([[1.], [1.], [1.], [1.]]) } return features, labels @@ -469,8 +466,7 @@ class DebugClassifierTest(test.TestCase): features = { 'x': input_lib.limit_epochs( - array_ops.ones( - shape=[4, 1], dtype=dtypes.float32), + array_ops.ones(shape=[4, 1], dtype=dtypes.float32), num_epochs=num_epochs), } return features, labels @@ -578,12 +574,11 @@ class DebugClassifierTest(test.TestCase): language = feature_column.sparse_column_with_hash_bucket('language', 100) feature_columns = [ feature_column.real_valued_column('age'), - feature_column.embedding_column( - language, dimension=1) + feature_column.embedding_column(language, dimension=1) ] - classifier = debug.DebugClassifier(config=run_config.RunConfig( - tf_random_seed=1)) + classifier = debug.DebugClassifier( + config=run_config.RunConfig(tf_random_seed=1)) classifier.fit(input_fn=input_fn, steps=5) def default_input_fn(unused_estimator, examples): @@ -614,8 +609,8 @@ class DebugRegressorTest(test.TestCase): classifier.fit( input_fn=_input_fn_builder(train_features, train_labels), steps=50) - pred = classifier.predict_scores(input_fn=_input_fn_builder(test_features, - None)) + pred = classifier.predict_scores( + input_fn=_input_fn_builder(test_features, None)) self.assertAllClose(expected_prediction, np.vstack(pred), atol=0.1) def testExperimentIntegration(self): @@ -698,7 +693,9 @@ class DebugRegressorTest(test.TestCase): # Create 4 rows, one of them (y = x), three of them (y=Not(x)) # The algorithm should learn (y = 0.25). labels = constant_op.constant([[1.], [0.], [0.], [0.]]) - features = {'x': array_ops.ones(shape=[4, 1], dtype=dtypes.float32),} + features = { + 'x': array_ops.ones(shape=[4, 1], dtype=dtypes.float32), + } return features, labels regressor = debug.DebugRegressor( @@ -853,5 +850,6 @@ class DebugRegressorTest(test.TestCase): predictions2 = list(regressor2.predict_scores(input_fn=predict_input_fn)) self.assertAllClose(predictions, predictions2) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined_test.py b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined_test.py index 57e70e169ca9d6fb2adc4e50bf387cc7cf330aed..4e65c180d8bee9ab8fe9b1fbf32edc229c31af09 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined_test.py @@ -1046,11 +1046,14 @@ class DNNLinearCombinedClassifierTest(test.TestCase): if global_step == 100: # Expected is 100, but because of the global step increment bug, is 50. - self.assertEqual(50, step_counter.steps) + # Occasionally, step increments one more time due to a race condition, + # reaching 51 steps. + self.assertIn(step_counter.steps, [50, 51]) else: - # Occasionally, training stops when global_step == 101, due to a race - # condition. - self.assertEqual(51, step_counter.steps) + # Occasionally, training stops when global_step == 102, due to a race + # condition. In addition, occasionally step increments one more time due + # to a race condition reaching 52 steps. + self.assertIn(step_counter.steps, [51, 52]) def testGlobalStepDNNLinearCombinedBugFixed(self): """Tests global step update for dnn-linear combined model.""" diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn_test.py b/tensorflow/contrib/learn/python/learn/estimators/dnn_test.py index 12f9bba531a296a00d17956b8ce32e5d7dead380..2bd57597c2e9444b51b1dacfbe4180b443c95a3d 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dnn_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dnn_test.py @@ -1224,7 +1224,7 @@ class DNNRegressorTest(test.TestCase): self, predictions, expected_shape): predictions_nparray = np.array(predictions) self.assertAllEqual(expected_shape, predictions_nparray.shape) - self.assertTrue(np.issubdtype(predictions_nparray.dtype, np.float)) + self.assertTrue(np.issubdtype(predictions_nparray.dtype, np.floating)) def testPredict_AsIterableFalse(self): """Tests predict method with as_iterable=False.""" diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimator.py b/tensorflow/contrib/learn/python/learn/estimators/estimator.py index 05ed8b3409e68ae54e5ef89b3a1592a6f285565b..4b63e08ab3372849309ee5d28d754de82e9632f4 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimator.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimator.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Base Estimator class.""" from __future__ import absolute_import @@ -76,7 +75,6 @@ from tensorflow.python.util import compat from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect - AS_ITERABLE_DATE = '2016-09-15' AS_ITERABLE_INSTRUCTIONS = ( 'The default behavior of predict() is changing. The default value for\n' @@ -213,7 +211,7 @@ def _get_replica_device_setter(config): 'Variable', 'VariableV2', 'AutoReloadVariable', 'MutableHashTable', 'MutableHashTableV2', 'MutableHashTableOfTensors', 'MutableHashTableOfTensorsV2', 'MutableDenseHashTable', - 'MutableDenseHashTableV2' + 'MutableDenseHashTableV2', 'VarHandleOp' ] if config.task_type: @@ -223,8 +221,11 @@ def _get_replica_device_setter(config): if config.num_ps_replicas > 0: return device_setter.replica_device_setter( - ps_tasks=config.num_ps_replicas, worker_device=worker_device, - merge_devices=True, ps_ops=ps_ops, cluster=config.cluster_spec) + ps_tasks=config.num_ps_replicas, + worker_device=worker_device, + merge_devices=True, + ps_ops=ps_ops, + cluster=config.cluster_spec) else: return None @@ -284,10 +285,10 @@ def _make_metrics_ops(metrics, features, labels, predictions): raise ValueError('Invalid metric for {}. It returned a tuple with ' 'len {}, expected 2.'.format(name, len(name))) if not isinstance(predictions, dict): - raise ValueError( - 'Metrics passed provide (name, prediction), ' - 'but predictions are not dict. ' - 'Metrics: %s, Predictions: %s.' % (metrics, predictions)) + raise ValueError('Metrics passed provide (name, prediction), ' + 'but predictions are not dict. ' + 'Metrics: %s, Predictions: %s.' % (metrics, + predictions)) # Here are two options: labels are single Tensor or a dict. if isinstance(labels, dict) and name[1] in labels: # If labels are dict and the prediction name is in it, apply metric. @@ -298,10 +299,10 @@ def _make_metrics_ops(metrics, features, labels, predictions): else: # Single head metrics. if isinstance(predictions, dict): - raise ValueError( - 'Metrics passed provide only name, no prediction, ' - 'but predictions are dict. ' - 'Metrics: %s, Labels: %s.' % (metrics, labels_tensor_or_dict)) + raise ValueError('Metrics passed provide only name, no prediction, ' + 'but predictions are dict. ' + 'Metrics: %s, Labels: %s.' % (metrics, + labels_tensor_or_dict)) result[name] = metric(predictions, labels_tensor_or_dict) return result @@ -360,10 +361,22 @@ def _write_dict_to_summary(output_dir, dictionary, current_global_step): logging.warn('Skipping summary for %s, cannot parse string to Summary.', key) continue + elif isinstance(dictionary[key], np.ndarray): + value = summary_proto.value.add() + value.tag = key + value.node_name = key + tensor_proto = tensor_util.make_tensor_proto(dictionary[key]) + value.tensor.CopyFrom(tensor_proto) + logging.info( + 'Summary for np.ndarray is not visible in Tensorboard by default. ' + 'Consider using a Tensorboard plugin for visualization (see ' + 'https://github.com/tensorflow/tensorboard-plugin-example/blob/master/README.md' + ' for more information).') else: logging.warn( 'Skipping summary for %s, must be a float, np.float32, np.int64, ' - 'np.int32 or int or a serialized string of Summary.', key) + 'np.int32 or int or np.ndarray or a serialized string of Summary.', + key) summary_writer.add_summary(summary_proto, current_global_step) summary_writer.flush() @@ -372,8 +385,8 @@ GraphRewriteSpec = collections.namedtuple('GraphRewriteSpec', ['tags', 'transforms']) -class BaseEstimator( - sklearn.BaseEstimator, evaluable.Evaluable, trainable.Trainable): +class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, + trainable.Trainable): """Abstract BaseEstimator class to train and evaluate TensorFlow models. Users should not instantiate or subclass this class. Instead, use an @@ -415,7 +428,7 @@ class BaseEstimator( # necessary. # pylint: disable=g-doc-exception raise ValueError( - "model_dir are set both in constructor and RunConfig, but with " + 'model_dir are set both in constructor and RunConfig, but with ' "different values. In constructor: '{}', in RunConfig: " "'{}' ".format(model_dir, self._config.model_dir)) # pylint: enable=g-doc-exception @@ -444,12 +457,16 @@ class BaseEstimator( # TODO(wicke): make RunConfig immutable, and then return it without a copy. return copy.deepcopy(self._config) - @deprecated_args( - SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, ('x', None), - ('y', None), ('batch_size', None) - ) - def fit(self, x=None, y=None, input_fn=None, steps=None, batch_size=None, - monitors=None, max_steps=None): + @deprecated_args(SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, + ('x', None), ('y', None), ('batch_size', None)) + def fit(self, + x=None, + y=None, + input_fn=None, + steps=None, + batch_size=None, + monitors=None, + max_steps=None): # pylint: disable=g-doc-args,g-doc-return-or-yield """See `Trainable`. @@ -481,13 +498,15 @@ class BaseEstimator( logging.info('Loss for final step: %s.', loss) return self - @deprecated_args( - SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, ('x', None), - ('y', None), ('batch_size', None) - ) - def partial_fit( - self, x=None, y=None, input_fn=None, steps=1, batch_size=None, - monitors=None): + @deprecated_args(SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, + ('x', None), ('y', None), ('batch_size', None)) + def partial_fit(self, + x=None, + y=None, + input_fn=None, + steps=1, + batch_size=None, + monitors=None): """Incremental fit on a batch of samples. This method is expected to be called several times consecutively @@ -523,13 +542,16 @@ class BaseEstimator( """ logging.warning('The current implementation of partial_fit is not optimized' ' for use in a loop. Consider using fit() instead.') - return self.fit(x=x, y=y, input_fn=input_fn, steps=steps, - batch_size=batch_size, monitors=monitors) + return self.fit( + x=x, + y=y, + input_fn=input_fn, + steps=steps, + batch_size=batch_size, + monitors=monitors) - @deprecated_args( - SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, ('x', None), - ('y', None), ('batch_size', None) - ) + @deprecated_args(SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, + ('x', None), ('y', None), ('batch_size', None)) def evaluate(self, x=None, y=None, @@ -571,13 +593,15 @@ class BaseEstimator( eval_results.update({'global_step': global_step}) return eval_results - @deprecated_args( - SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, ('x', None), - ('batch_size', None), ('as_iterable', True) - ) - def predict( - self, x=None, input_fn=None, batch_size=None, outputs=None, - as_iterable=True): + @deprecated_args(SCIKIT_DECOUPLE_DATE, SCIKIT_DECOUPLE_INSTRUCTIONS, + ('x', None), ('batch_size', None), ('as_iterable', True)) + def predict(self, + x=None, + input_fn=None, + batch_size=None, + outputs=None, + as_iterable=True, + iterate_batches=False): """Returns predictions for given features. Args: @@ -593,6 +617,9 @@ class BaseEstimator( for each example until inputs are exhausted. Note: The inputs must terminate if you want the iterable to terminate (e.g. be sure to pass num_epochs=1 if you are using something like read_batch_features). + iterate_batches: If True, yield the whole batch at once instead of + decomposing the batch into individual samples. Only relevant when + as_iterable is True. Returns: A numpy array of predicted classes or regression values if the @@ -612,7 +639,8 @@ class BaseEstimator( input_fn=input_fn, feed_fn=feed_fn, outputs=outputs, - as_iterable=as_iterable) + as_iterable=as_iterable, + iterate_batches=iterate_batches) def get_variable_value(self, name): """Returns value of the variable given by name. @@ -638,16 +666,17 @@ class BaseEstimator( return self._model_dir @deprecated('2017-03-25', 'Please use Estimator.export_savedmodel() instead.') - def export(self, - export_dir, - input_fn=export._default_input_fn, # pylint: disable=protected-access - input_feature_key=None, - use_deprecated_input_fn=True, - signature_fn=None, - prediction_key=None, - default_batch_size=1, - exports_to_keep=None, - checkpoint_path=None): + def export( + self, + export_dir, + input_fn=export._default_input_fn, # pylint: disable=protected-access + input_feature_key=None, + use_deprecated_input_fn=True, + signature_fn=None, + prediction_key=None, + default_batch_size=1, + exports_to_keep=None, + checkpoint_path=None): """Exports inference graph into given dir. Args: @@ -785,8 +814,8 @@ class BaseEstimator( logging.debug('Setting feature info to %s.', str(self._features_info)) if labels is not None: if self._labels_info is not None: - logging.debug('Given labels: %s, required signatures: %s.', - str(labels), str(self._labels_info)) + logging.debug('Given labels: %s, required signatures: %s.', str(labels), + str(self._labels_info)) if not tensor_signature.tensors_compatible(labels, self._labels_info): raise ValueError('Labels are incompatible with given information. ' 'Given labels: %s, required signatures: %s.' % @@ -837,13 +866,13 @@ class BaseEstimator( if not checkpoint_path: latest_path = saver.latest_checkpoint(self._model_dir) if not latest_path: - raise NotFittedError("Couldn't find trained model at %s." - % self._model_dir) + raise NotFittedError( + "Couldn't find trained model at %s." % self._model_dir) checkpoint_path = latest_path # Setup output directory. - eval_dir = os.path.join(self._model_dir, 'eval' if not name else - 'eval_' + name) + eval_dir = os.path.join(self._model_dir, 'eval' + if not name else 'eval_' + name) with ops.Graph().as_default() as g: random_seed.set_random_seed(self._config.tf_random_seed) @@ -866,8 +895,7 @@ class BaseEstimator( 'Use steps=None if intended.') if steps: hooks.append( - evaluation.StopAfterNEvalsHook( - steps, log_progress=log_progress)) + evaluation.StopAfterNEvalsHook(steps, log_progress=log_progress)) global_step_key = 'global_step' while global_step_key in eval_dict: @@ -903,8 +931,8 @@ class BaseEstimator( # Check that model has been trained. checkpoint_path = saver.latest_checkpoint(self._model_dir) if not checkpoint_path: - raise NotFittedError("Couldn't find trained model at %s." - % self._model_dir) + raise NotFittedError( + "Couldn't find trained model at %s." % self._model_dir) with ops.Graph().as_default() as g: random_seed.set_random_seed(self._config.tf_random_seed) @@ -966,7 +994,8 @@ class BaseEstimator( existing_keys = predictions.keys() predictions = { key: value - for key, value in six.iteritems(predictions) if key in outputs + for key, value in six.iteritems(predictions) + if key in outputs } if not predictions: raise ValueError('Expected to run at least one output from %s, ' @@ -1032,8 +1061,7 @@ class BaseEstimator( chief_only_hooks=chief_hooks + model_fn_ops.training_chief_hooks, save_checkpoint_secs=0, # Saving is handled by a hook. save_summaries_steps=self._config.save_summary_steps, - config=self._session_config - ) as mon_sess: + config=self._session_config) as mon_sess: loss = None while not mon_sess.should_stop(): _, loss = mon_sess.run([model_fn_ops.train_op, model_fn_ops.loss]) @@ -1124,8 +1152,7 @@ class Estimator(BaseEstimator): if params is not None and 'params' not in model_fn_args: raise ValueError('Estimator\'s model_fn (%s) does not have a params ' 'argument, but params (%s) were passed to the ' - 'Estimator\'s constructor.' % - (model_fn, params)) + 'Estimator\'s constructor.' % (model_fn, params)) if params is None and 'params' in model_fn_args: logging.warning('Estimator\'s model_fn (%s) includes params ' 'argument, but params are not passed to Estimator.', @@ -1179,8 +1206,9 @@ class Estimator(BaseEstimator): # Custom metrics should overwrite defaults. if metrics: - model_fn_ops.eval_metric_ops.update(_make_metrics_ops( - metrics, features, labels, model_fn_ops.predictions)) + model_fn_ops.eval_metric_ops.update( + _make_metrics_ops(metrics, features, labels, + model_fn_ops.predictions)) return model_fn_ops @@ -1225,8 +1253,8 @@ class Estimator(BaseEstimator): Raises: ValueError: if `metrics` don't match `labels`. """ - model_fn_ops = self._call_model_fn( - features, labels, model_fn_lib.ModeKeys.EVAL, metrics) + model_fn_ops = self._call_model_fn(features, labels, + model_fn_lib.ModeKeys.EVAL, metrics) if metric_key.MetricKey.LOSS not in model_fn_ops.eval_metric_ops: model_fn_ops.eval_metric_ops[metric_key.MetricKey.LOSS] = ( @@ -1250,13 +1278,17 @@ class Estimator(BaseEstimator): self._labels_info) return self._call_model_fn(features, labels, model_fn_lib.ModeKeys.INFER) - def export_savedmodel( - self, export_dir_base, serving_input_fn, - default_output_alternative_key=None, - assets_extra=None, - as_text=False, - checkpoint_path=None, - graph_rewrite_specs=(GraphRewriteSpec((tag_constants.SERVING,), ()),)): + def export_savedmodel(self, + export_dir_base, + serving_input_fn, + default_output_alternative_key=None, + assets_extra=None, + as_text=False, + checkpoint_path=None, + graph_rewrite_specs=(GraphRewriteSpec( + (tag_constants.SERVING,), ()),), + strip_default_attrs=False): + # pylint: disable=line-too-long """Exports inference graph as a SavedModel into given dir. Args: @@ -1280,6 +1312,10 @@ class Estimator(BaseEstimator): produce a separate MetaGraphDef within the exported SavedModel, tagged and rewritten as specified. Defaults to a single entry using the default serving tag ("serve") and no rewriting. + strip_default_attrs: Boolean. If `True`, default-valued attributes will be + removed from the NodeDefs. For a detailed guide, see + [Stripping Default-Valued + Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes). Returns: The string path to the exported directory. @@ -1287,6 +1323,7 @@ class Estimator(BaseEstimator): Raises: ValueError: if an unrecognized export_type is requested. """ + # pylint: enable=line-too-long if serving_input_fn is None: raise ValueError('serving_input_fn must be defined.') @@ -1294,8 +1331,8 @@ class Estimator(BaseEstimator): # Locate the latest checkpoint checkpoint_path = saver.latest_checkpoint(self._model_dir) if not checkpoint_path: - raise NotFittedError("Couldn't find trained model at %s." - % self._model_dir) + raise NotFittedError( + "Couldn't find trained model at %s." % self._model_dir) export_dir = saved_model_export_utils.get_timestamped_export_dir( export_dir_base) @@ -1329,10 +1366,10 @@ class Estimator(BaseEstimator): saved_model_export_utils.get_output_alternatives( model_fn_ops, default_output_alternative_key)) - init_op = control_flow_ops.group( - variables.local_variables_initializer(), - resources.initialize_resources(resources.shared_resources()), - lookup_ops.tables_initializer()) + init_op = control_flow_ops.group(variables.local_variables_initializer(), + resources.initialize_resources( + resources.shared_resources()), + lookup_ops.tables_initializer()) # Build the SignatureDefs from all pairs of input and output alternatives signature_def_map = saved_model_export_utils.build_all_signature_defs( @@ -1362,11 +1399,12 @@ class Estimator(BaseEstimator): # TODO(soergel): switch to main_op or otherwise update when dust settles builder.add_meta_graph_and_variables( - session, untransformed_tags, + session, + untransformed_tags, signature_def_map=signature_def_map, - assets_collection=ops.get_collection( - ops.GraphKeys.ASSET_FILEPATHS), - legacy_init_op=init_op) + assets_collection=ops.get_collection(ops.GraphKeys.ASSET_FILEPATHS), + legacy_init_op=init_op, + strip_default_attrs=strip_default_attrs) # pylint: disable=protected-access base_meta_graph_def = builder._saved_model.meta_graphs[0] @@ -1375,12 +1413,16 @@ class Estimator(BaseEstimator): if graph_rewrite_specs[1:]: # Prepare the input_names and output_names needed for the # meta_graph_transform call below. - input_names = [tensor.name - for input_dict in input_alternatives.values() - for tensor in input_dict.values()] - output_names = [tensor.name - for output_alternative in output_alternatives.values() - for tensor in output_alternative[1].values()] + input_names = [ + tensor.name + for input_dict in input_alternatives.values() + for tensor in input_dict.values() + ] + output_names = [ + tensor.name + for output_alternative in output_alternatives.values() + for tensor in output_alternative[1].values() + ] # Write the additional MetaGraphDefs for graph_rewrite_spec in graph_rewrite_specs[1:]: @@ -1399,11 +1441,11 @@ class Estimator(BaseEstimator): # Add the extra assets if assets_extra: - assets_extra_path = os.path.join(compat.as_bytes(temp_export_dir), - compat.as_bytes('assets.extra')) + assets_extra_path = os.path.join( + compat.as_bytes(temp_export_dir), compat.as_bytes('assets.extra')) for dest_relative, source in assets_extra.items(): - dest_absolute = os.path.join(compat.as_bytes(assets_extra_path), - compat.as_bytes(dest_relative)) + dest_absolute = os.path.join( + compat.as_bytes(assets_extra_path), compat.as_bytes(dest_relative)) dest_path = os.path.dirname(dest_absolute) gfile.MakeDirs(dest_path) gfile.Copy(source, dest_absolute) @@ -1423,25 +1465,36 @@ class SKCompat(sklearn.BaseEstimator): def fit(self, x, y, batch_size=128, steps=None, max_steps=None, monitors=None): - input_fn, feed_fn = _get_input_fn(x, y, input_fn=None, feed_fn=None, - batch_size=batch_size, shuffle=True, - epochs=None) + input_fn, feed_fn = _get_input_fn( + x, + y, + input_fn=None, + feed_fn=None, + batch_size=batch_size, + shuffle=True, + epochs=None) all_monitors = [] if feed_fn: all_monitors = [basic_session_run_hooks.FeedFnHook(feed_fn)] if monitors: all_monitors.extend(monitors) - self._estimator.fit(input_fn=input_fn, - steps=steps, - max_steps=max_steps, - monitors=all_monitors) + self._estimator.fit( + input_fn=input_fn, + steps=steps, + max_steps=max_steps, + monitors=all_monitors) return self def score(self, x, y, batch_size=128, steps=None, metrics=None, name=None): - input_fn, feed_fn = _get_input_fn(x, y, input_fn=None, - feed_fn=None, batch_size=batch_size, - shuffle=False, epochs=1) + input_fn, feed_fn = _get_input_fn( + x, + y, + input_fn=None, + feed_fn=None, + batch_size=batch_size, + shuffle=False, + epochs=1) if metrics is not None and not isinstance(metrics, dict): raise ValueError('Metrics argument should be None or dict. ' 'Got %s.' % metrics) @@ -1457,8 +1510,13 @@ class SKCompat(sklearn.BaseEstimator): def predict(self, x, batch_size=128, outputs=None): input_fn, feed_fn = _get_input_fn( - x, None, input_fn=None, feed_fn=None, batch_size=batch_size, - shuffle=False, epochs=1) + x, + None, + input_fn=None, + feed_fn=None, + batch_size=batch_size, + shuffle=False, + epochs=1) results = list( self._estimator._infer_model( input_fn=input_fn, @@ -1469,7 +1527,6 @@ class SKCompat(sklearn.BaseEstimator): if not isinstance(results[0], dict): return np.concatenate([output for output in results], axis=0) return { - key: np.concatenate( - [output[key] for output in results], axis=0) + key: np.concatenate([output[key] for output in results], axis=0) for key in results[0] } diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimator_input_test.py b/tensorflow/contrib/learn/python/learn/estimators/estimator_input_test.py index 9d7c1a099aa4be64ca0296fa5b870597dabec7b4..d4a46b41d0c93ef58d5db8c433cbf348fec10f5e 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimator_input_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimator_input_test.py @@ -41,7 +41,6 @@ from tensorflow.python.platform import test from tensorflow.python.training import input as input_lib from tensorflow.python.training import queue_runner_impl - _BOSTON_INPUT_DIM = 13 _IRIS_INPUT_DIM = 4 @@ -93,8 +92,8 @@ def boston_eval_fn(): constant_op.constant(boston.data), [n_examples, _BOSTON_INPUT_DIM]) labels = array_ops.reshape( constant_op.constant(boston.target), [n_examples, 1]) - return array_ops.concat([features, features], 0), array_ops.concat( - [labels, labels], 0) + return array_ops.concat([features, features], + 0), array_ops.concat([labels, labels], 0) def extract(data, key): @@ -129,7 +128,10 @@ def linear_model_fn(features, labels, mode): (_, features), = features.items() prediction, loss = (models.linear_regression_zero_init(features, labels)) train_op = optimizers.optimize_loss( - loss, training_util.get_global_step(), optimizer='Adagrad', learning_rate=0.1) + loss, + training_util.get_global_step(), + optimizer='Adagrad', + learning_rate=0.1) return prediction, loss, train_op @@ -139,7 +141,10 @@ def linear_model_fn_with_model_fn_ops(features, labels, mode): model_fn.ModeKeys.INFER) prediction, loss = (models.linear_regression_zero_init(features, labels)) train_op = optimizers.optimize_loss( - loss, training_util.get_global_step(), optimizer='Adagrad', learning_rate=0.1) + loss, + training_util.get_global_step(), + optimizer='Adagrad', + learning_rate=0.1) return model_fn.ModelFnOps( mode=mode, predictions=prediction, loss=loss, train_op=train_op) @@ -150,7 +155,10 @@ def logistic_model_no_mode_fn(features, labels): labels = array_ops.one_hot(labels, 3, 1, 0) prediction, loss = (models.logistic_regression_zero_init(features, labels)) train_op = optimizers.optimize_loss( - loss, training_util.get_global_step(), optimizer='Adagrad', learning_rate=0.1) + loss, + training_util.get_global_step(), + optimizer='Adagrad', + learning_rate=0.1) return { 'class': math_ops.argmax(prediction, 1), 'prob': prediction @@ -173,7 +181,9 @@ class EstimatorInputTest(test.TestCase): scores = est.evaluate( x=boston_input, y=float64_target, - metrics={'MSE': metric_ops.streaming_mean_squared_error}) + metrics={ + 'MSE': metric_ops.streaming_mean_squared_error + }) del est # Create another estimator object with the same output dir. est2 = estimator.Estimator(model_fn=linear_model_fn, model_dir=output_dir) @@ -182,7 +192,9 @@ class EstimatorInputTest(test.TestCase): scores2 = est2.evaluate( x=boston_input, y=float64_target, - metrics={'MSE': metric_ops.streaming_mean_squared_error}) + metrics={ + 'MSE': metric_ops.streaming_mean_squared_error + }) self.assertAllClose(scores2['MSE'], scores['MSE']) predictions = np.array(list(est2.predict(x=boston_input))) other_score = _sklearn.mean_squared_error(predictions, @@ -197,7 +209,9 @@ class EstimatorInputTest(test.TestCase): scores = est.score( x=boston.data, y=float64_labels, - metrics={'MSE': metric_ops.streaming_mean_squared_error}) + metrics={ + 'MSE': metric_ops.streaming_mean_squared_error + }) predictions = np.array(list(est.predict(x=boston.data))) other_score = _sklearn.mean_squared_error(predictions, boston.target) self.assertAllClose(scores['MSE'], other_score) @@ -213,7 +227,9 @@ class EstimatorInputTest(test.TestCase): scores = est.evaluate( x=boston_input, y=float64_target, - metrics={'MSE': metric_ops.streaming_mean_squared_error}) + metrics={ + 'MSE': metric_ops.streaming_mean_squared_error + }) predictions = np.array(list(est.predict(x=boston_input))) other_score = _sklearn.mean_squared_error(predictions, boston.target) self.assertAllClose(other_score, scores['MSE']) @@ -228,14 +244,15 @@ class EstimatorInputTest(test.TestCase): scores = est.score( x=iris.data, y=iris.target, - metrics={('accuracy', 'class'): metric_ops.streaming_accuracy}) + metrics={ + ('accuracy', 'class'): metric_ops.streaming_accuracy + }) predictions = est.predict(x=iris.data) predictions_class = est.predict(x=iris.data, outputs=['class'])['class'] self.assertEqual(predictions['prob'].shape[0], iris.target.shape[0]) self.assertAllClose(predictions['class'], predictions_class) - self.assertAllClose( - predictions['class'], np.argmax( - predictions['prob'], axis=1)) + self.assertAllClose(predictions['class'], + np.argmax(predictions['prob'], axis=1)) other_score = _sklearn.accuracy_score(iris.target, predictions['class']) self.assertAllClose(scores['accuracy'], other_score) self.assertTrue('global_step' in scores) @@ -250,17 +267,18 @@ class EstimatorInputTest(test.TestCase): scores = est.evaluate( x=iris_data, y=iris_target, - metrics={('accuracy', 'class'): metric_ops.streaming_accuracy}) + metrics={ + ('accuracy', 'class'): metric_ops.streaming_accuracy + }) predictions = list(est.predict(x=iris_data)) predictions_class = list(est.predict(x=iris_data, outputs=['class'])) self.assertEqual(len(predictions), iris.target.shape[0]) classes_batch = np.array([p['class'] for p in predictions]) self.assertAllClose(classes_batch, np.array([p['class'] for p in predictions_class])) - self.assertAllClose( - classes_batch, - np.argmax( - np.array([p['prob'] for p in predictions]), axis=1)) + self.assertAllClose(classes_batch, + np.argmax( + np.array([p['prob'] for p in predictions]), axis=1)) other_score = _sklearn.accuracy_score(iris.target, classes_batch) self.assertAllClose(other_score, scores['accuracy']) self.assertTrue('global_step' in scores) diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py b/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py index 2a13a84627df35a68a4f04b25ab26ceecad0db0d..d81a534b79bc90fe91ffd3cb97a7865a7cb4c2a9 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py @@ -111,8 +111,8 @@ def boston_eval_fn(): constant_op.constant(boston.data), [n_examples, _BOSTON_INPUT_DIM]) labels = array_ops.reshape( constant_op.constant(boston.target), [n_examples, 1]) - return array_ops.concat([features, features], 0), array_ops.concat( - [labels, labels], 0) + return array_ops.concat([features, features], + 0), array_ops.concat([labels, labels], 0) def extract(data, key): @@ -147,7 +147,10 @@ def linear_model_fn(features, labels, mode): (_, features), = features.items() prediction, loss = (models.linear_regression_zero_init(features, labels)) train_op = optimizers.optimize_loss( - loss, training_util.get_global_step(), optimizer='Adagrad', learning_rate=0.1) + loss, + training_util.get_global_step(), + optimizer='Adagrad', + learning_rate=0.1) return prediction, loss, train_op @@ -157,7 +160,10 @@ def linear_model_fn_with_model_fn_ops(features, labels, mode): model_fn.ModeKeys.INFER) prediction, loss = (models.linear_regression_zero_init(features, labels)) train_op = optimizers.optimize_loss( - loss, training_util.get_global_step(), optimizer='Adagrad', learning_rate=0.1) + loss, + training_util.get_global_step(), + optimizer='Adagrad', + learning_rate=0.1) return model_fn.ModelFnOps( mode=mode, predictions=prediction, loss=loss, train_op=train_op) @@ -168,7 +174,10 @@ def logistic_model_no_mode_fn(features, labels): labels = array_ops.one_hot(labels, 3, 1, 0) prediction, loss = (models.logistic_regression_zero_init(features, labels)) train_op = optimizers.optimize_loss( - loss, training_util.get_global_step(), optimizer='Adagrad', learning_rate=0.1) + loss, + training_util.get_global_step(), + optimizer='Adagrad', + learning_rate=0.1) return { 'class': math_ops.argmax(prediction, 1), 'prob': prediction @@ -184,14 +193,12 @@ def _build_estimator_for_export_tests(tmpdir): def _input_fn(): iris = base.load_iris() return { - 'feature': constant_op.constant( - iris.data, dtype=dtypes.float32) + 'feature': constant_op.constant(iris.data, dtype=dtypes.float32) }, constant_op.constant( iris.target, shape=[150], dtype=dtypes.int32) feature_columns = [ - feature_column_lib.real_valued_column( - 'feature', dimension=4) + feature_column_lib.real_valued_column('feature', dimension=4) ] est = linear.LinearRegressor(feature_columns) @@ -291,8 +298,8 @@ class CheckCallsMonitor(monitors_lib.BaseMonitor): self.begin_calls == self.expect_calls) -def _model_fn_ops( - expected_features, expected_labels, actual_features, actual_labels, mode): +def _model_fn_ops(expected_features, expected_labels, actual_features, + actual_labels, mode): assert_ops = tuple([ check_ops.assert_equal( expected_features[k], actual_features[k], name='assert_%s' % k) @@ -310,11 +317,11 @@ def _model_fn_ops( def _make_input_fn(features, labels): + def _input_fn(): - return { - k: constant_op.constant(v) - for k, v in six.iteritems(features) - }, constant_op.constant(labels) + return {k: constant_op.constant(v) + for k, v in six.iteritems(features)}, constant_op.constant(labels) + return _input_fn @@ -369,11 +376,13 @@ class EstimatorModelFnTest(test.TestCase): self.assertEqual(expected_params, params) self.assertTrue(config.i_am_test) return _model_fn_ops(features, labels, arg0, arg1, mode) + partial_model_fn = functools.partial( _model_fn, foo=expected_foo, bar=expected_bar) est = estimator.Estimator( - model_fn=partial_model_fn, params=expected_params, + model_fn=partial_model_fn, + params=expected_params, config=expected_config) self.assertEqual(0, model_fn_call_count[0]) est.fit(input_fn=_make_input_fn(features, labels), steps=1) @@ -382,7 +391,12 @@ class EstimatorModelFnTest(test.TestCase): def testModelFnWithModelDir(self): expected_param = {'some_param': 'some_value'} expected_model_dir = tempfile.mkdtemp() - def _argument_checker(features, labels, mode, params, config=None, + + def _argument_checker(features, + labels, + mode, + params, + config=None, model_dir=None): _, _, _ = features, labels, config self.assertEqual(model_fn.ModeKeys.TRAIN, mode) @@ -390,9 +404,11 @@ class EstimatorModelFnTest(test.TestCase): self.assertEqual(model_dir, expected_model_dir) return (constant_op.constant(0.), constant_op.constant(0.), training_util.get_global_step().assign_add(1)) - est = estimator.Estimator(model_fn=_argument_checker, - params=expected_param, - model_dir=expected_model_dir) + + est = estimator.Estimator( + model_fn=_argument_checker, + params=expected_param, + model_dir=expected_model_dir) est.fit(input_fn=boston_input_fn, steps=1) def testInvalidModelFn_no_train_op(self): @@ -447,8 +463,7 @@ class EstimatorModelFnTest(test.TestCase): est.predict(input_fn=boston_input_fn) with self.assertRaisesRegexp(ValueError, 'Missing prediction'): est.predict( - input_fn=functools.partial( - boston_input_fn, num_epochs=1), + input_fn=functools.partial(boston_input_fn, num_epochs=1), as_iterable=True) def testModelFnScaffoldInTraining(self): @@ -498,15 +513,17 @@ class EstimatorModelFnTest(test.TestCase): self.assertTrue(self.mock_saver.restore.called) est.predict(input_fn=input_fn) self.assertTrue(self.mock_saver.restore.called) + def serving_input_fn(): - serialized_tf_example = array_ops.placeholder(dtype=dtypes.string, - shape=[None], - name='input_example_tensor') + serialized_tf_example = array_ops.placeholder( + dtype=dtypes.string, shape=[None], name='input_example_tensor') features, labels = input_fn() - return input_fn_utils.InputFnOps( - features, labels, {'examples': serialized_tf_example}) + return input_fn_utils.InputFnOps(features, labels, { + 'examples': serialized_tf_example + }) - est.export_savedmodel(os.path.join(est.model_dir, 'export'), serving_input_fn) + est.export_savedmodel( + os.path.join(est.model_dir, 'export'), serving_input_fn) self.assertTrue(self.mock_saver.restore.called) @@ -550,33 +567,28 @@ class EstimatorTest(test.TestCase): def testRunConfigModelDir(self): config = run_config.RunConfig(model_dir='test_dir') - est = estimator.Estimator(model_fn=linear_model_fn, - config=config) + est = estimator.Estimator(model_fn=linear_model_fn, config=config) self.assertEqual('test_dir', est.config.model_dir) self.assertEqual('test_dir', est.model_dir) def testModelDirAndRunConfigModelDir(self): config = run_config.RunConfig(model_dir='test_dir') - est = estimator.Estimator(model_fn=linear_model_fn, - config=config, - model_dir='test_dir') + est = estimator.Estimator( + model_fn=linear_model_fn, config=config, model_dir='test_dir') self.assertEqual('test_dir', est.config.model_dir) with self.assertRaisesRegexp( - ValueError, - 'model_dir are set both in constructor and RunConfig, ' + ValueError, 'model_dir are set both in constructor and RunConfig, ' 'but with different'): - estimator.Estimator(model_fn=linear_model_fn, - config=config, - model_dir='different_dir') + estimator.Estimator( + model_fn=linear_model_fn, config=config, model_dir='different_dir') def testModelDirIsCopiedToRunConfig(self): config = run_config.RunConfig() self.assertIsNone(config.model_dir) - est = estimator.Estimator(model_fn=linear_model_fn, - model_dir='test_dir', - config=config) + est = estimator.Estimator( + model_fn=linear_model_fn, model_dir='test_dir', config=config) self.assertEqual('test_dir', est.config.model_dir) self.assertEqual('test_dir', est.model_dir) @@ -656,25 +668,27 @@ class EstimatorTest(test.TestCase): boston = base.load_boston() output_dir = tempfile.mkdtemp() est = estimator.SKCompat( - estimator.Estimator( - model_fn=linear_model_fn, model_dir=output_dir)) + estimator.Estimator(model_fn=linear_model_fn, model_dir=output_dir)) float64_labels = boston.target.astype(np.float64) est.fit(x=boston.data, y=float64_labels, steps=50) scores = est.score( x=boston.data, y=float64_labels, - metrics={'MSE': metric_ops.streaming_mean_squared_error}) + metrics={ + 'MSE': metric_ops.streaming_mean_squared_error + }) del est # Create another estimator object with the same output dir. est2 = estimator.SKCompat( - estimator.Estimator( - model_fn=linear_model_fn, model_dir=output_dir)) + estimator.Estimator(model_fn=linear_model_fn, model_dir=output_dir)) # Check we can evaluate and predict. scores2 = est2.score( x=boston.data, y=float64_labels, - metrics={'MSE': metric_ops.streaming_mean_squared_error}) + metrics={ + 'MSE': metric_ops.streaming_mean_squared_error + }) self.assertAllClose(scores['MSE'], scores2['MSE']) predictions = np.array(list(est2.predict(x=boston.data))) other_score = _sklearn.mean_squared_error(predictions, float64_labels) @@ -685,14 +699,15 @@ class EstimatorTest(test.TestCase): scores3 = est2.score( x=boston.data, y=float64_labels, - metrics={'MSE': metric_ops.streaming_mean_squared_error}) + metrics={ + 'MSE': metric_ops.streaming_mean_squared_error + }) self.assertLess(scores3['MSE'], scores['MSE']) def test_checkpoint_contains_relative_paths(self): tmpdir = tempfile.mkdtemp() est = estimator.Estimator( - model_dir=tmpdir, - model_fn=linear_model_fn_with_model_fn_ops) + model_dir=tmpdir, model_fn=linear_model_fn_with_model_fn_ops) est.fit(input_fn=boston_input_fn, steps=5) checkpoint_file_content = file_io.read_file_to_string( @@ -700,22 +715,20 @@ class EstimatorTest(test.TestCase): ckpt = checkpoint_state_pb2.CheckpointState() text_format.Merge(checkpoint_file_content, ckpt) self.assertEqual(ckpt.model_checkpoint_path, 'model.ckpt-5') - self.assertAllEqual( - ['model.ckpt-1', 'model.ckpt-5'], ckpt.all_model_checkpoint_paths) + self.assertAllEqual(['model.ckpt-1', 'model.ckpt-5'], + ckpt.all_model_checkpoint_paths) def test_train_save_copy_reload(self): tmpdir = tempfile.mkdtemp() model_dir1 = os.path.join(tmpdir, 'model_dir1') est1 = estimator.Estimator( - model_dir=model_dir1, - model_fn=linear_model_fn_with_model_fn_ops) + model_dir=model_dir1, model_fn=linear_model_fn_with_model_fn_ops) est1.fit(input_fn=boston_input_fn, steps=5) model_dir2 = os.path.join(tmpdir, 'model_dir2') os.renames(model_dir1, model_dir2) est2 = estimator.Estimator( - model_dir=model_dir2, - model_fn=linear_model_fn_with_model_fn_ops) + model_dir=model_dir2, model_fn=linear_model_fn_with_model_fn_ops) self.assertEqual(5, est2.get_variable_value('global_step')) est2.fit(input_fn=boston_input_fn, steps=5) self.assertEqual(10, est2.get_variable_value('global_step')) @@ -724,7 +737,9 @@ class EstimatorTest(test.TestCase): boston = base.load_boston() est = estimator.SKCompat( estimator.Estimator( - model_fn=linear_model_params_fn, params={'learning_rate': 0.01})) + model_fn=linear_model_params_fn, params={ + 'learning_rate': 0.01 + })) est.fit(x=boston.data, y=boston.target, steps=100) def testHooksNotChanged(self): @@ -824,11 +839,13 @@ class EstimatorTest(test.TestCase): def testMonitorsForFit(self): est = estimator.Estimator(model_fn=linear_model_fn) - est.fit(input_fn=boston_input_fn, - steps=21, - monitors=[CheckCallsMonitor(expect_calls=21)]) + est.fit( + input_fn=boston_input_fn, + steps=21, + monitors=[CheckCallsMonitor(expect_calls=21)]) def testHooksForEvaluate(self): + class CheckCallHook(session_run_hook.SessionRunHook): def __init__(self): @@ -874,7 +891,9 @@ class EstimatorTest(test.TestCase): est.evaluate( input_fn=boston_input_fn, steps=200, - metrics={'MSE': _streaming_mean_squared_error_histogram}) + metrics={ + 'MSE': _streaming_mean_squared_error_histogram + }) events = util_test.latest_events(est.model_dir + '/eval') output_values = {} for e in events: @@ -884,6 +903,37 @@ class EstimatorTest(test.TestCase): self.assertTrue('MSE' in output_values) self.assertTrue(output_values['MSE'].HasField('histo')) + def testSummaryWritingWithTensor(self): + + def _streaming_precition_mean_tensor(predictions, + weights=None, + metrics_collections=None, + updates_collections=None, + name=None): + return metric_ops.streaming_mean_tensor( + predictions, + weights=weights, + metrics_collections=metrics_collections, + updates_collections=updates_collections, + name=name) + + est = estimator.Estimator(model_fn=linear_model_fn) + est.fit(input_fn=boston_input_fn, steps=200) + est.evaluate( + input_fn=boston_input_fn, + steps=200, + metrics={ + 'PMT': _streaming_precition_mean_tensor + }) + events = util_test.latest_events(est.model_dir + '/eval') + output_values = {} + for e in events: + if e.HasField('summary'): + for v in e.summary.value: + output_values[v.tag] = v + self.assertTrue('PMT' in output_values) + self.assertTrue(output_values['PMT'].HasField('tensor')) + def testLossInGraphCollection(self): class _LossCheckerHook(session_run_hook.SessionRunHook): @@ -927,8 +977,8 @@ class EstimatorTest(test.TestCase): self.assertTrue( gfile.Exists( os.path.join( - compat.as_bytes(export_dir), compat.as_bytes( - 'saved_model.pb')))) + compat.as_bytes(export_dir), + compat.as_bytes('saved_model.pb')))) self.assertTrue( gfile.Exists( os.path.join( @@ -988,11 +1038,11 @@ class EstimatorTest(test.TestCase): self.assertTrue('input_example_tensor' in graph_ops) self.assertTrue('ParseExample/ParseExample' in graph_ops) self.assertTrue('linear/linear/feature/matmul' in graph_ops) - self.assertItemsEqual( - ['bogus_lookup', 'feature'], - [compat.as_str_any(x) for x in graph.get_collection( - constants.COLLECTION_DEF_KEY_FOR_INPUT_FEATURE_KEYS)]) - + self.assertItemsEqual(['bogus_lookup', 'feature'], [ + compat.as_str_any(x) + for x in graph.get_collection( + constants.COLLECTION_DEF_KEY_FOR_INPUT_FEATURE_KEYS) + ]) # cleanup gfile.DeleteRecursively(tmpdir) @@ -1010,8 +1060,8 @@ class EstimatorTest(test.TestCase): self.assertTrue( gfile.Exists( os.path.join( - compat.as_bytes(export_dir), compat.as_bytes( - 'saved_model.pb')))) + compat.as_bytes(export_dir), + compat.as_bytes('saved_model.pb')))) self.assertTrue( gfile.Exists( os.path.join( @@ -1054,19 +1104,22 @@ class EstimatorTest(test.TestCase): export_dir_base = os.path.join( compat.as_bytes(tmpdir), compat.as_bytes('export')) export_dir = est.export_savedmodel( - export_dir_base, serving_input_fn, assets_extra=assets_extra, + export_dir_base, + serving_input_fn, + assets_extra=assets_extra, graph_rewrite_specs=[ estimator.GraphRewriteSpec(['tag_1'], []), estimator.GraphRewriteSpec(['tag_2', 'tag_3'], - ['strip_unused_nodes'])]) + ['strip_unused_nodes']) + ]) self.assertTrue(gfile.Exists(export_dir_base)) self.assertTrue(gfile.Exists(export_dir)) self.assertTrue( gfile.Exists( os.path.join( - compat.as_bytes(export_dir), compat.as_bytes( - 'saved_model.pb')))) + compat.as_bytes(export_dir), + compat.as_bytes('saved_model.pb')))) self.assertTrue( gfile.Exists( os.path.join( @@ -1179,18 +1232,15 @@ class InferRealValuedColumnsTest(test.TestCase): self.assertEqual(1, len(feature_columns)) feature_column = feature_columns[0] self.assertEqual('', feature_column.name) - self.assertEqual( - { - '': - parsing_ops.FixedLenFeature( - shape=expected_shape, dtype=expected_dtype) - }, - feature_column.config) + self.assertEqual({ + '': + parsing_ops.FixedLenFeature( + shape=expected_shape, dtype=expected_dtype) + }, feature_column.config) def testInt32Input(self): feature_columns = estimator.infer_real_valued_columns_from_input( - np.ones( - shape=[7, 8], dtype=np.int32)) + np.ones(shape=[7, 8], dtype=np.int32)) self._assert_single_feature_column([8], dtypes.int32, feature_columns) def testInt32InputFn(self): @@ -1200,8 +1250,7 @@ class InferRealValuedColumnsTest(test.TestCase): def testInt64Input(self): feature_columns = estimator.infer_real_valued_columns_from_input( - np.ones( - shape=[7, 8], dtype=np.int64)) + np.ones(shape=[7, 8], dtype=np.int64)) self._assert_single_feature_column([8], dtypes.int64, feature_columns) def testInt64InputFn(self): @@ -1211,8 +1260,7 @@ class InferRealValuedColumnsTest(test.TestCase): def testFloat32Input(self): feature_columns = estimator.infer_real_valued_columns_from_input( - np.ones( - shape=[7, 8], dtype=np.float32)) + np.ones(shape=[7, 8], dtype=np.float32)) self._assert_single_feature_column([8], dtypes.float32, feature_columns) def testFloat32InputFn(self): @@ -1222,8 +1270,7 @@ class InferRealValuedColumnsTest(test.TestCase): def testFloat64Input(self): feature_columns = estimator.infer_real_valued_columns_from_input( - np.ones( - shape=[7, 8], dtype=np.float64)) + np.ones(shape=[7, 8], dtype=np.float64)) self._assert_single_feature_column([8], dtypes.float64, feature_columns) def testFloat64InputFn(self): @@ -1242,8 +1289,8 @@ class InferRealValuedColumnsTest(test.TestCase): ValueError, 'on integer or non floating types are not supported'): # pylint: disable=g-long-lambda estimator.infer_real_valued_columns_from_input_fn( - lambda: (constant_op.constant(False, shape=[7, 8], dtype=dtypes.bool), - None)) + lambda: (constant_op.constant(False, shape=[7, 8], dtype=dtypes.bool), None) + ) def testStringInput(self): with self.assertRaisesRegexp( @@ -1280,8 +1327,9 @@ class ReplicaDeviceSetterTest(test.TestCase): def testVariablesAreOnPs(self): tf_config = {'cluster': {run_config.TaskType.PS: ['fake_ps_0']}} - with test.mock.patch.dict('os.environ', - {'TF_CONFIG': json.dumps(tf_config)}): + with test.mock.patch.dict('os.environ', { + 'TF_CONFIG': json.dumps(tf_config) + }): config = run_config.RunConfig() with ops.device(estimator._get_replica_device_setter(config)): @@ -1308,14 +1356,14 @@ class ReplicaDeviceSetterTest(test.TestCase): def testMutableHashTableIsOnPs(self): tf_config = {'cluster': {run_config.TaskType.PS: ['fake_ps_0']}} - with test.mock.patch.dict('os.environ', - {'TF_CONFIG': json.dumps(tf_config)}): + with test.mock.patch.dict('os.environ', { + 'TF_CONFIG': json.dumps(tf_config) + }): config = run_config.RunConfig() with ops.device(estimator._get_replica_device_setter(config)): default_val = constant_op.constant([-1, -1], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) + table = lookup.MutableHashTable(dtypes.string, dtypes.int64, default_val) input_string = constant_op.constant(['brain', 'salad', 'tank']) output = table.lookup(input_string) self.assertDeviceEqual('/job:ps/task:0', table._table_ref.device) @@ -1325,8 +1373,7 @@ class ReplicaDeviceSetterTest(test.TestCase): with ops.device( estimator._get_replica_device_setter(run_config.RunConfig())): default_val = constant_op.constant([-1, -1], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) + table = lookup.MutableHashTable(dtypes.string, dtypes.int64, default_val) input_string = constant_op.constant(['brain', 'salad', 'tank']) output = table.lookup(input_string) self.assertDeviceEqual('', table._table_ref.device) @@ -1342,8 +1389,9 @@ class ReplicaDeviceSetterTest(test.TestCase): 'index': 3 } } - with test.mock.patch.dict('os.environ', - {'TF_CONFIG': json.dumps(tf_config)}): + with test.mock.patch.dict('os.environ', { + 'TF_CONFIG': json.dumps(tf_config) + }): config = run_config.RunConfig() with ops.device(estimator._get_replica_device_setter(config)): diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimators_test.py b/tensorflow/contrib/learn/python/learn/estimators/estimators_test.py index 8131e0fde6fea5501cacc4714f53ed8d867ca70f..2113fae3940f14c8ca07e5f76986408ae8a33831 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimators_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimators_test.py @@ -72,9 +72,11 @@ class FeatureEngineeringFunctionTest(test.TestCase): # predictions = transformed_x (9) self.assertEqual(9., prediction) metrics = estimator.evaluate( - input_fn=input_fn, steps=1, - metrics={"label": - metric_spec.MetricSpec(lambda predictions, labels: labels)}) + input_fn=input_fn, + steps=1, + metrics={ + "label": metric_spec.MetricSpec(lambda predictions, labels: labels) + }) # labels = transformed_y (99) self.assertEqual(99., metrics["label"]) @@ -82,10 +84,10 @@ class FeatureEngineeringFunctionTest(test.TestCase): def input_fn(): return { - "x": constant_op.constant(["9."]) - }, { - "y": constant_op.constant(["99."]) - } + "x": constant_op.constant(["9."]) + }, { + "y": constant_op.constant(["99."]) + } def feature_engineering_fn(features, labels): # Github #12205: raise a TypeError if called twice. @@ -104,15 +106,17 @@ class FeatureEngineeringFunctionTest(test.TestCase): return predictions, loss, update_global_step estimator = estimator_lib.Estimator( - model_fn=model_fn, feature_engineering_fn=feature_engineering_fn) + model_fn=model_fn, feature_engineering_fn=feature_engineering_fn) estimator.fit(input_fn=input_fn, steps=1) prediction = next(estimator.predict(input_fn=input_fn, as_iterable=True)) # predictions = transformed_x (9) self.assertEqual(9., prediction) metrics = estimator.evaluate( - input_fn=input_fn, steps=1, - metrics={"label": - metric_spec.MetricSpec(lambda predictions, labels: labels)}) + input_fn=input_fn, + steps=1, + metrics={ + "label": metric_spec.MetricSpec(lambda predictions, labels: labels) + }) # labels = transformed_y (99) self.assertEqual(99., metrics["label"]) @@ -150,12 +154,10 @@ class FeatureEngineeringFunctionTest(test.TestCase): # predictions = x prediction_with_fe_fn = next( - estimator_with_fe_fn.predict( - input_fn=input_fn, as_iterable=True)) + estimator_with_fe_fn.predict(input_fn=input_fn, as_iterable=True)) self.assertEqual(9., prediction_with_fe_fn) prediction_without_fe_fn = next( - estimator_without_fe_fn.predict( - input_fn=input_fn, as_iterable=True)) + estimator_without_fe_fn.predict(input_fn=input_fn, as_iterable=True)) self.assertEqual(1., prediction_without_fe_fn) diff --git a/tensorflow/contrib/learn/python/learn/estimators/head.py b/tensorflow/contrib/learn/python/learn/estimators/head.py index bc0e6fc0091c9b5419ab526855b404eb4a927e97..9b124b2c19f16bbc9b2afeadb82a32006e1a0ae9 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/head.py +++ b/tensorflow/contrib/learn/python/learn/estimators/head.py @@ -181,7 +181,8 @@ def regression_head(label_name=None, weight_column_name=None, label_dimension=1, enable_centered_bias=False, - head_name=None): + head_name=None, + link_fn=None): """Creates a `Head` for linear regression. Args: @@ -199,6 +200,8 @@ def regression_head(label_name=None, head_name: name of the head. If provided, predictions, summary and metrics keys will be suffixed by `"/" + head_name` and the default variable scope will be `head_name`. + link_fn: link function to convert logits to predictions. If provided, + this link function will be used instead of identity. Returns: An instance of `Head` for linear regression. @@ -210,7 +213,7 @@ def regression_head(label_name=None, enable_centered_bias=enable_centered_bias, head_name=head_name, loss_fn=_mean_squared_loss, - link_fn=array_ops.identity) + link_fn=(link_fn if link_fn is not None else array_ops.identity)) def poisson_regression_head(label_name=None, diff --git a/tensorflow/contrib/learn/python/learn/estimators/head_test.py b/tensorflow/contrib/learn/python/learn/estimators/head_test.py index 3881bf533d642bef68fa9ab4ba908bbb8f7f8091..6d5da81b4c2087fb9c5307902e452a6220a17cd0 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/head_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/head_test.py @@ -33,6 +33,7 @@ from tensorflow.python.client import session from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import lookup_ops +from tensorflow.python.ops import math_ops from tensorflow.python.ops import variables from tensorflow.python.ops.losses import losses as losses_lib from tensorflow.python.platform import test @@ -153,6 +154,25 @@ class RegressionHeadTest(test.TestCase): _assert_no_variables(self) _assert_metrics(self, 5. / 3, {"loss": 5. / 3}, model_fn_ops) + def testRegressionWithLogitFn(self): + head = head_lib.regression_head(link_fn=math_ops.square) + def _assert_preditions(test_case, expected_predictions, model_fn_ops): + variables.initialize_local_variables().run() + test_case.assertAllClose(expected_predictions, + model_fn_ops.predictions["scores"].eval()) + with ops.Graph().as_default(), session.Session(): + model_fn_ops = head.create_model_fn_ops( + {}, + labels=((0.,), (1.,), (1.,)), + mode=model_fn.ModeKeys.TRAIN, + train_op_fn=head_lib.no_op_train_fn, + logits=((1.,), (1.,), (3.,))) + self._assert_output_alternatives(model_fn_ops) + _assert_summary_tags(self, ["loss"]) + _assert_no_variables(self) + _assert_metrics(self, 5. / 3, {"loss": 5. / 3}, model_fn_ops) + _assert_preditions(self, ([1.0, 1.0, 9.0]), model_fn_ops) + def testRegressionWithInvalidLogits(self): head = head_lib.regression_head() with ops.Graph().as_default(), session.Session(): @@ -342,7 +362,7 @@ class MultiLabelHeadTest(test.TestCase): "auc_precision_recall": 0.166667, "auc_precision_recall/class0": 0, "auc_precision_recall/class1": 0., - "auc_precision_recall/class2": 1., + "auc_precision_recall/class2": 0.49999, "labels/actual_label_mean/class0": self._labels[0][0], "labels/actual_label_mean/class1": self._labels[0][1], "labels/actual_label_mean/class2": self._labels[0][2], @@ -728,7 +748,7 @@ class BinaryClassificationHeadTest(test.TestCase): "accuracy/baseline_label_mean": label_mean, "accuracy/threshold_0.500000_mean": 1. / 2, "auc": 1. / 2, - "auc_precision_recall": 0.749999, + "auc_precision_recall": 0.25, "labels/actual_label_mean": label_mean, "labels/prediction_mean": .731059, # softmax "loss": expected_loss, diff --git a/tensorflow/contrib/learn/python/learn/estimators/logistic_regressor_test.py b/tensorflow/contrib/learn/python/learn/estimators/logistic_regressor_test.py index 656d68b76888d9319c0b9be481f9b0478ac4314c..ac2d10011e222eb9c534d7fbae3c0cb5f4820945 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/logistic_regressor_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/logistic_regressor_test.py @@ -57,7 +57,10 @@ def _logistic_regression_model_fn(features, labels, mode): predictions = math_ops.sigmoid(logits) loss = losses.sigmoid_cross_entropy(labels, logits) train_op = optimizers.optimize_loss( - loss, training_util.get_global_step(), optimizer='Adagrad', learning_rate=0.1) + loss, + training_util.get_global_step(), + optimizer='Adagrad', + learning_rate=0.1) return predictions, loss, train_op diff --git a/tensorflow/contrib/learn/python/learn/evaluable.py b/tensorflow/contrib/learn/python/learn/evaluable.py index 66e15265171679dcd710fdf05bed3105de6bab99..8f6cd39864b437f163dd7c1140dc88755ce98529 100644 --- a/tensorflow/contrib/learn/python/learn/evaluable.py +++ b/tensorflow/contrib/learn/python/learn/evaluable.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """`Evaluable` interface.""" from __future__ import absolute_import @@ -59,9 +58,12 @@ class Evaluable(object): for which this evaluation was performed. Args: - x: Matrix of shape [n_samples, n_features...] or dictionary of many matrices - containing the input samples for fitting the model. Can be iterator that returns - arrays of features or dictionary of array of features. If set, `input_fn` must + x: Matrix of shape [n_samples, n_features...] or dictionary of many + matrices + containing the input samples for fitting the model. Can be iterator that + returns + arrays of features or dictionary of array of features. If set, + `input_fn` must be `None`. y: Vector or matrix [n_samples] or [n_samples, n_outputs] containing the label values (class labels in classification, real numbers in diff --git a/tensorflow/contrib/learn/python/learn/experiment.py b/tensorflow/contrib/learn/python/learn/experiment.py index 9576ff21c243022276bb0641882dfaf0decf05c0..bec976afd2719138117976381669ca3292360480 100644 --- a/tensorflow/contrib/learn/python/learn/experiment.py +++ b/tensorflow/contrib/learn/python/learn/experiment.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Experiment class collecting information needed for a single training run.""" from __future__ import absolute_import @@ -43,7 +42,6 @@ from tensorflow.python.training import saver from tensorflow.python.training import server_lib from tensorflow.python.util import compat - __all__ = ["Experiment"] @@ -278,8 +276,7 @@ class Experiment(object): self._train_steps_per_iteration = train_steps_per_iteration if (self._train_steps_per_iteration is not None and not isinstance(self._train_steps_per_iteration, int)): - raise ValueError( - "`train_steps_per_iteration` must be an integer.") + raise ValueError("`train_steps_per_iteration` must be an integer.") @property def estimator(self): @@ -359,9 +356,10 @@ class Experiment(object): config.cluster_spec and config.master): self._start_server() elif config.cluster_spec and config.master: - raise ValueError('For distributed runtime, Experiment class only works with' - 'tf.contrib.learn.RunConfig for now, but provided {}' - .format(type(config))) + raise ValueError( + "For distributed runtime, Experiment class only works with" + "tf.contrib.learn.RunConfig for now, but provided {}".format( + type(config))) extra_hooks = [] if delay_secs is None: @@ -414,11 +412,12 @@ class Experiment(object): logging.info("Waiting %d secs before starting eval.", delay_secs) time.sleep(delay_secs) - return self._call_evaluate(input_fn=self._eval_input_fn, - steps=self._eval_steps, - metrics=self._eval_metrics, - name=(name or "one_pass"), - hooks=self._eval_hooks) + return self._call_evaluate( + input_fn=self._eval_input_fn, + steps=self._eval_steps, + metrics=self._eval_metrics, + name=(name or "one_pass"), + hooks=self._eval_hooks) @deprecated( "2016-10-23", @@ -499,15 +498,12 @@ class Experiment(object): previous_path = None eval_result = None last_warning_time = 0 - while (not predicate_fn or - predicate_fn( - eval_result, - checkpoint_path=previous_path if eval_result else None)): + while (not predicate_fn or predicate_fn( + eval_result, checkpoint_path=previous_path if eval_result else None)): # Exit if we have already reached number of steps to train. if self._has_training_stopped(eval_result): logging.info("Exiting continuous eval, global_step=%s >= " - "train_step=%s", - eval_result[ops.GraphKeys.GLOBAL_STEP], + "train_step=%s", eval_result[ops.GraphKeys.GLOBAL_STEP], self._train_steps) return @@ -528,12 +524,13 @@ class Experiment(object): logging.warning(error_msg) last_warning_time = time.time() else: - eval_result = self._call_evaluate(input_fn=input_fn, - steps=self._eval_steps, - metrics=self._eval_metrics, - name=name, - checkpoint_path=latest_path, - hooks=self._eval_hooks) + eval_result = self._call_evaluate( + input_fn=input_fn, + steps=self._eval_steps, + metrics=self._eval_metrics, + name=name, + checkpoint_path=latest_path, + hooks=self._eval_hooks) # Ensure eval result is not None for next round of evaluation. if not eval_result: eval_result = {} @@ -558,8 +555,8 @@ class Experiment(object): return False global_step = eval_result.get(ops.GraphKeys.GLOBAL_STEP) - return global_step and self._train_steps and ( - global_step >= self._train_steps) + return global_step and self._train_steps and (global_step >= + self._train_steps) def continuous_eval(self, delay_secs=None, @@ -678,8 +675,7 @@ class Experiment(object): return eval_result, export_results @experimental - def continuous_train_and_eval(self, - continuous_eval_predicate_fn=None): + def continuous_train_and_eval(self, continuous_eval_predicate_fn=None): """Interleaves training and evaluation. The frequency of evaluation is controlled by the `train_steps_per_iteration` @@ -752,10 +748,9 @@ class Experiment(object): elif self._train_steps is not None: train_steps_per_iteration = int(self._train_steps / 10) - while (not predicate_fn or - predicate_fn( - eval_result, - checkpoint_path=latest_checkpoint if eval_result else None)): + while (not predicate_fn or predicate_fn( + eval_result, checkpoint_path=latest_checkpoint + if eval_result else None)): if self._has_training_stopped(eval_result): # Exits once max steps of training is satisfied. @@ -785,8 +780,7 @@ class Experiment(object): def _maybe_export(self, eval_result, checkpoint_path=None): """Export the Estimator using export_fn, if defined.""" export_dir_base = os.path.join( - compat.as_bytes(self._estimator.model_dir), - compat.as_bytes("export")) + compat.as_bytes(self._estimator.model_dir), compat.as_bytes("export")) export_results = [] for strategy in self._export_strategies: @@ -824,10 +818,11 @@ class Experiment(object): hooks=self._train_monitors, saving_listeners=self._saving_listeners) - eval_result = self._call_evaluate(input_fn=self._eval_input_fn, - steps=1, - metrics=self._eval_metrics, - name="one_pass") + eval_result = self._call_evaluate( + input_fn=self._eval_input_fn, + steps=1, + metrics=self._eval_metrics, + name="one_pass") _ = self._maybe_export(eval_result) return eval_result @@ -849,9 +844,14 @@ class Experiment(object): server.start() return server - def _call_train(self, _sentinel=None, # pylint: disable=invalid-name, - input_fn=None, steps=None, hooks=None, max_steps=None, - saving_listeners=None): + def _call_train( + self, + _sentinel=None, # pylint: disable=invalid-name, + input_fn=None, + steps=None, + hooks=None, + max_steps=None, + saving_listeners=None): if _sentinel is not None: raise ValueError("_call_train should be called with keyword args only") @@ -867,14 +867,18 @@ class Experiment(object): hooks=hooks, saving_listeners=saving_listeners) else: - return self._estimator.fit(input_fn=input_fn, - steps=steps, - max_steps=max_steps, - monitors=hooks) - - def _call_evaluate(self, _sentinel=None, # pylint: disable=invalid-name, - input_fn=None, steps=None, metrics=None, name=None, - checkpoint_path=None, hooks=None): + return self._estimator.fit( + input_fn=input_fn, steps=steps, max_steps=max_steps, monitors=hooks) + + def _call_evaluate( + self, + _sentinel=None, # pylint: disable=invalid-name, + input_fn=None, + steps=None, + metrics=None, + name=None, + checkpoint_path=None, + hooks=None): if _sentinel is not None: raise ValueError("_call_evaluate should be called with keyword args only") @@ -882,18 +886,20 @@ class Experiment(object): if metrics is not None: raise ValueError( "`eval_metrics` must be `None` with `tf.estimator.Estimator`") - return self._estimator.evaluate(input_fn=input_fn, - steps=steps, - name=name, - checkpoint_path=checkpoint_path, - hooks=hooks) + return self._estimator.evaluate( + input_fn=input_fn, + steps=steps, + name=name, + checkpoint_path=checkpoint_path, + hooks=hooks) else: - return self._estimator.evaluate(input_fn=input_fn, - steps=steps, - metrics=metrics, - name=name, - checkpoint_path=checkpoint_path, - hooks=hooks) + return self._estimator.evaluate( + input_fn=input_fn, + steps=steps, + metrics=metrics, + name=name, + checkpoint_path=checkpoint_path, + hooks=hooks) @contextlib.contextmanager diff --git a/tensorflow/contrib/learn/python/learn/export_strategy.py b/tensorflow/contrib/learn/python/learn/export_strategy.py index f276aab0e6beb011a21c20fa194dd5212db796d1..55a8b824312b89e0ac66513242191f4201ac212a 100644 --- a/tensorflow/contrib/learn/python/learn/export_strategy.py +++ b/tensorflow/contrib/learn/python/learn/export_strategy.py @@ -26,13 +26,14 @@ __all__ = ['ExportStrategy'] class ExportStrategy( - collections.namedtuple('ExportStrategy', ['name', 'export_fn'])): + collections.namedtuple('ExportStrategy', + ['name', 'export_fn', 'strip_default_attrs'])): """A class representing a type of model export. Typically constructed by a utility function specific to the exporter, such as `saved_model_export_utils.make_export_strategy()`. - The fields are: + Attributes: name: The directory name under the export base directory where exports of this type will be written. export_fn: A function that writes an export, given an estimator, a @@ -45,11 +46,20 @@ class ExportStrategy( The signature of this function must be one of: - * `(estimator, export_path) -> export_path` - * `(estimator, export_path, checkpoint_path) -> export_path` - * `(estimator, export_path, checkpoint_path, eval_result) -> export_path` + * `(estimator, export_path) -> export_path` + * `(estimator, export_path, checkpoint_path) -> export_path` + * `(estimator, export_path, checkpoint_path, eval_result) -> export_path` + * `(estimator, export_path, checkpoint_path, eval_result, + strip_default_attrs) -> export_path` + strip_default_attrs: (Optional) Boolean. If set as True, default attrs in + the `GraphDef` will be stripped on write. This is recommended for better + forward compatibility of the resulting `SavedModel`. """ + def __new__(cls, name, export_fn, strip_default_attrs=None): + return super(ExportStrategy, cls).__new__( + cls, name, export_fn, strip_default_attrs) + def export(self, estimator, export_path, @@ -83,5 +93,6 @@ class ExportStrategy( raise ValueError('An export_fn accepting eval_result must also accept ' 'checkpoint_path.') kwargs['eval_result'] = eval_result - + if 'strip_default_attrs' in export_fn_args: + kwargs['strip_default_attrs'] = self.strip_default_attrs return self.export_fn(estimator, export_path, **kwargs) diff --git a/tensorflow/contrib/learn/python/learn/export_strategy_test.py b/tensorflow/contrib/learn/python/learn/export_strategy_test.py new file mode 100644 index 0000000000000000000000000000000000000000..43c3551cccc3b8e6b66bd2b36839a3dfc5fe8eea --- /dev/null +++ b/tensorflow/contrib/learn/python/learn/export_strategy_test.py @@ -0,0 +1,89 @@ +# 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 ExportStrategy.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.learn.python.learn import export_strategy +from tensorflow.python.platform import test + + +class ExportStrategyTest(test.TestCase): + + def test_no_optional_args_export(self): + model_path = '/path/to/model' + def _export_fn(estimator, export_path): + self.assertTupleEqual((estimator, export_path), (None, None)) + return model_path + + strategy = export_strategy.ExportStrategy('foo', _export_fn) + self.assertTupleEqual(strategy, ('foo', _export_fn, None)) + self.assertIs(strategy.export(None, None), model_path) + + def test_checkpoint_export(self): + ckpt_model_path = '/path/to/checkpoint_model' + def _ckpt_export_fn(estimator, export_path, checkpoint_path): + self.assertTupleEqual((estimator, export_path), (None, None)) + self.assertEqual(checkpoint_path, 'checkpoint') + return ckpt_model_path + + strategy = export_strategy.ExportStrategy('foo', _ckpt_export_fn) + self.assertTupleEqual(strategy, ('foo', _ckpt_export_fn, None)) + self.assertIs(strategy.export(None, None, 'checkpoint'), ckpt_model_path) + + def test_checkpoint_eval_export(self): + ckpt_eval_model_path = '/path/to/checkpoint_eval_model' + def _ckpt_eval_export_fn(estimator, export_path, checkpoint_path, + eval_result): + self.assertTupleEqual((estimator, export_path), (None, None)) + self.assertEqual(checkpoint_path, 'checkpoint') + self.assertEqual(eval_result, 'eval') + return ckpt_eval_model_path + + strategy = export_strategy.ExportStrategy('foo', _ckpt_eval_export_fn) + self.assertTupleEqual(strategy, ('foo', _ckpt_eval_export_fn, None)) + self.assertIs(strategy.export(None, None, 'checkpoint', 'eval'), + ckpt_eval_model_path) + + def test_eval_only_export(self): + def _eval_export_fn(estimator, export_path, eval_result): + del estimator, export_path, eval_result + + strategy = export_strategy.ExportStrategy('foo', _eval_export_fn) + self.assertTupleEqual(strategy, ('foo', _eval_export_fn, None)) + with self.assertRaisesRegexp(ValueError, 'An export_fn accepting ' + 'eval_result must also accept ' + 'checkpoint_path'): + strategy.export(None, None, eval_result='eval') + + def test_strip_default_attr_export(self): + strip_default_attrs_model_path = '/path/to/strip_default_attrs_model' + def _strip_default_attrs_export_fn(estimator, export_path, + strip_default_attrs): + self.assertTupleEqual((estimator, export_path), (None, None)) + self.assertTrue(strip_default_attrs) + return strip_default_attrs_model_path + + strategy = export_strategy.ExportStrategy('foo', + _strip_default_attrs_export_fn, + True) + self.assertTupleEqual(strategy, + ('foo', _strip_default_attrs_export_fn, True)) + self.assertIs(strategy.export(None, None), strip_default_attrs_model_path) + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/learn/python/learn/learn_io/data_feeder.py b/tensorflow/contrib/learn/python/learn/learn_io/data_feeder.py index f36a778b529a83f158241ddb060959c4b33e2e95..96be8b1bc402479d5611965f27abb197363cb939 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/data_feeder.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/data_feeder.py @@ -35,6 +35,7 @@ from tensorflow.python.platform import tf_logging as logging # pylint: disable=g-multiple-import,g-bad-import-order from .pandas_io import HAS_PANDAS, extract_pandas_data, extract_pandas_matrix, extract_pandas_labels from .dask_io import HAS_DASK, extract_dask_data, extract_dask_labels + # pylint: enable=g-multiple-import,g-bad-import-order @@ -74,11 +75,11 @@ def _get_in_out_shape(x_shape, y_shape, n_classes, batch_size=None): if not y_is_dict: output_shape = out_el_shape(y_shape, n_classes) else: - output_shape = dict([ - (k, out_el_shape(v, n_classes[k] - if n_classes is not None and k in n_classes else None)) - for k, v in list(y_shape.items()) - ]) + output_shape = dict([(k, + out_el_shape(v, n_classes[k] + if n_classes is not None and + k in n_classes else None)) + for k, v in list(y_shape.items())]) return input_shape, output_shape, batch_size @@ -314,23 +315,23 @@ class DataFeeder(object): input_dtype: DType of input (or dictionary of shapes). output_dtype: DType of output (or dictionary of shapes. """ - x_is_dict, y_is_dict = isinstance(x, dict), y is not None and isinstance( - y, dict) + x_is_dict, y_is_dict = isinstance( + x, dict), y is not None and isinstance(y, dict) if isinstance(y, list): y = np.array(y) self._x = dict([(k, check_array(v, v.dtype)) for k, v in list(x.items()) ]) if x_is_dict else check_array(x, x.dtype) - self._y = None if y is None else ( - dict([(k, check_array(v, v.dtype)) for k, v in list(y.items())]) - if y_is_dict else check_array(y, y.dtype)) + self._y = None if y is None else (dict( + [(k, check_array(v, v.dtype)) for k, v in list(y.items())]) + if y_is_dict else check_array(y, y.dtype)) # self.n_classes is not None means we're converting raw target indices # to one-hot. if n_classes is not None: if not y_is_dict: - y_dtype = (np.int64 - if n_classes is not None and n_classes > 1 else np.float32) + y_dtype = ( + np.int64 if n_classes is not None and n_classes > 1 else np.float32) self._y = (None if y is None else check_array(y, dtype=y_dtype)) self.n_classes = n_classes @@ -352,8 +353,8 @@ class DataFeeder(object): # self._output_dtype == np.float32 when y is None self._output_dtype = ( dict([(k, _check_dtype(v.dtype)) for k, v in list(self._y.items())]) - if y_is_dict else ( - _check_dtype(self._y.dtype) if y is not None else np.float32)) + if y_is_dict else (_check_dtype(self._y.dtype) + if y is not None else np.float32)) # self.n_classes is None means we're passing in raw target indices if n_classes is not None and y_is_dict: @@ -478,8 +479,8 @@ class DataFeeder(object): # Assign input features from random indices. def extract(data, indices): - return (np.array(_access(data, indices)).reshape((indices.shape[0], 1)) if - len(data.shape) == 1 else _access(data, indices)) + return (np.array(_access(data, indices)).reshape((indices.shape[0], 1)) + if len(data.shape) == 1 else _access(data, indices)) # assign labels from random indices def assign_label(data, shape, dtype, n_classes, indices): @@ -511,16 +512,18 @@ class DataFeeder(object): feed_dict[self._epoch_placeholder.name] = [self.epoch] # Take next batch of indices. - x_len = list(self._x.values())[0].shape[ - 0] if x_is_dict else self._x.shape[0] + x_len = list( + self._x.values())[0].shape[0] if x_is_dict else self._x.shape[0] end = min(x_len, self.offset + self._batch_size) batch_indices = self.indices[self.offset:end] # adding input placeholder feed_dict.update( dict([(self._input_placeholder[k].name, extract(v, batch_indices)) - for k, v in list(self._x.items())]) if x_is_dict else - {self._input_placeholder.name: extract(self._x, batch_indices)}) + for k, v in list(self._x.items())]) if x_is_dict else { + self._input_placeholder.name: + extract(self._x, batch_indices) + }) # move offset and reset it if necessary self.offset += self._batch_size @@ -545,7 +548,8 @@ class DataFeeder(object): assign_label(v, shape, dtype, n_classes, batch_indices) }) else: - shape, dtype, n_classes = self.output_shape, self._output_dtype, self.n_classes + shape, dtype, n_classes = (self.output_shape, self._output_dtype, + self.n_classes) feed_dict.update({ self._output_placeholder.name: assign_label(self._y, shape, dtype, n_classes, batch_indices) @@ -621,8 +625,9 @@ class StreamingDataFeeder(DataFeeder): elif y is None: y_first_el_shape = None else: - y_first_el_shape = ([1] + list(y_first_el[0].shape if isinstance( - y_first_el, list) else y_first_el.shape)) + y_first_el_shape = ( + [1] + list(y_first_el[0].shape + if isinstance(y_first_el, list) else y_first_el.shape)) self.input_shape, self.output_shape, self._batch_size = _get_in_out_shape( x_first_el_shape, y_first_el_shape, n_classes, batch_size) @@ -683,8 +688,8 @@ class StreamingDataFeeder(DataFeeder): if shape is None: return None elif isinstance(shape, dict): - return dict([(k, np.zeros(shape[k], dtype[k])) - for k in list(shape.keys())]) + return dict( + [(k, np.zeros(shape[k], dtype[k])) for k in list(shape.keys())]) else: return np.zeros(shape, dtype=dtype) diff --git a/tensorflow/contrib/learn/python/learn/learn_io/numpy_io_test.py b/tensorflow/contrib/learn/python/learn/learn_io/numpy_io_test.py deleted file mode 100644 index 6fe8de8705b8854e5861879d2a505fe03fddc7e5..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/learn/python/learn/learn_io/numpy_io_test.py +++ /dev/null @@ -1,280 +0,0 @@ -# Copyright 2016 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 numpy_io.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -from tensorflow.contrib.learn.python.learn.learn_io import numpy_io -from tensorflow.python.framework import errors -from tensorflow.python.platform import test -from tensorflow.python.training import coordinator -from tensorflow.python.training import queue_runner_impl - - -class NumpyIoTest(test.TestCase): - - def testNumpyInputFn(self): - a = np.arange(4) * 1.0 - b = np.arange(32, 36) - x = {'a': a, 'b': b} - y = np.arange(-32, -28) - - with self.test_session() as session: - input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=2, shuffle=False, num_epochs=1) - features, target = input_fn() - - coord = coordinator.Coordinator() - threads = queue_runner_impl.start_queue_runners(session, coord=coord) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [0, 1]) - self.assertAllEqual(res[0]['b'], [32, 33]) - self.assertAllEqual(res[1], [-32, -31]) - - session.run([features, target]) - with self.assertRaises(errors.OutOfRangeError): - session.run([features, target]) - - coord.request_stop() - coord.join(threads) - - def testNumpyInputFnWithVeryLargeBatchSizeAndMultipleEpochs(self): - a = np.arange(2) * 1.0 - b = np.arange(32, 34) - x = {'a': a, 'b': b} - y = np.arange(-32, -30) - - with self.test_session() as session: - input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=128, shuffle=False, num_epochs=2) - features, target = input_fn() - - coord = coordinator.Coordinator() - threads = queue_runner_impl.start_queue_runners(session, coord=coord) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [0, 1, 0, 1]) - self.assertAllEqual(res[0]['b'], [32, 33, 32, 33]) - self.assertAllEqual(res[1], [-32, -31, -32, -31]) - - with self.assertRaises(errors.OutOfRangeError): - session.run([features, target]) - - coord.request_stop() - coord.join(threads) - - def testNumpyInputFnWithZeroEpochs(self): - a = np.arange(4) * 1.0 - b = np.arange(32, 36) - x = {'a': a, 'b': b} - y = np.arange(-32, -28) - - with self.test_session() as session: - input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=2, shuffle=False, num_epochs=0) - features, target = input_fn() - - coord = coordinator.Coordinator() - threads = queue_runner_impl.start_queue_runners(session, coord=coord) - - with self.assertRaises(errors.OutOfRangeError): - session.run([features, target]) - - coord.request_stop() - coord.join(threads) - - def testNumpyInputFnWithBatchSizeNotDividedByDataSize(self): - batch_size = 2 - a = np.arange(5) * 1.0 - b = np.arange(32, 37) - x = {'a': a, 'b': b} - y = np.arange(-32, -27) - - with self.test_session() as session: - input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=batch_size, shuffle=False, num_epochs=1) - features, target = input_fn() - - coord = coordinator.Coordinator() - threads = queue_runner_impl.start_queue_runners(session, coord=coord) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [0, 1]) - self.assertAllEqual(res[0]['b'], [32, 33]) - self.assertAllEqual(res[1], [-32, -31]) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [2, 3]) - self.assertAllEqual(res[0]['b'], [34, 35]) - self.assertAllEqual(res[1], [-30, -29]) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [4]) - self.assertAllEqual(res[0]['b'], [36]) - self.assertAllEqual(res[1], [-28]) - - with self.assertRaises(errors.OutOfRangeError): - session.run([features, target]) - - coord.request_stop() - coord.join(threads) - - def testNumpyInputFnWithBatchSizeNotDividedByDataSizeAndMultipleEpochs(self): - batch_size = 2 - a = np.arange(3) * 1.0 - b = np.arange(32, 35) - x = {'a': a, 'b': b} - y = np.arange(-32, -29) - - with self.test_session() as session: - input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=batch_size, shuffle=False, num_epochs=3) - features, target = input_fn() - - coord = coordinator.Coordinator() - threads = queue_runner_impl.start_queue_runners(session, coord=coord) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [0, 1]) - self.assertAllEqual(res[0]['b'], [32, 33]) - self.assertAllEqual(res[1], [-32, -31]) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [2, 0]) - self.assertAllEqual(res[0]['b'], [34, 32]) - self.assertAllEqual(res[1], [-30, -32]) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [1, 2]) - self.assertAllEqual(res[0]['b'], [33, 34]) - self.assertAllEqual(res[1], [-31, -30]) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [0, 1]) - self.assertAllEqual(res[0]['b'], [32, 33]) - self.assertAllEqual(res[1], [-32, -31]) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [2]) - self.assertAllEqual(res[0]['b'], [34]) - self.assertAllEqual(res[1], [-30]) - - with self.assertRaises(errors.OutOfRangeError): - session.run([features, target]) - - coord.request_stop() - coord.join(threads) - - def testNumpyInputFnWithBatchSizeLargerThanDataSize(self): - batch_size = 10 - a = np.arange(4) * 1.0 - b = np.arange(32, 36) - x = {'a': a, 'b': b} - y = np.arange(-32, -28) - - with self.test_session() as session: - input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=batch_size, shuffle=False, num_epochs=1) - features, target = input_fn() - - coord = coordinator.Coordinator() - threads = queue_runner_impl.start_queue_runners(session, coord=coord) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [0, 1, 2, 3]) - self.assertAllEqual(res[0]['b'], [32, 33, 34, 35]) - self.assertAllEqual(res[1], [-32, -31, -30, -29]) - - with self.assertRaises(errors.OutOfRangeError): - session.run([features, target]) - - coord.request_stop() - coord.join(threads) - - def testNumpyInputFnWithDifferentDimensionsOfFeatures(self): - a = np.array([[1, 2], [3, 4]]) - b = np.array([5, 6]) - x = {'a': a, 'b': b} - y = np.arange(-32, -30) - - with self.test_session() as session: - input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=2, shuffle=False, num_epochs=1) - features, target = input_fn() - - coord = coordinator.Coordinator() - threads = queue_runner_impl.start_queue_runners(session, coord=coord) - - res = session.run([features, target]) - self.assertAllEqual(res[0]['a'], [[1, 2], [3, 4]]) - self.assertAllEqual(res[0]['b'], [5, 6]) - self.assertAllEqual(res[1], [-32, -31]) - - coord.request_stop() - coord.join(threads) - - def testNumpyInputFnWithXAsNonDict(self): - x = np.arange(32, 36) - y = np.arange(4) - with self.test_session(): - with self.assertRaisesRegexp(TypeError, 'x must be dict'): - failing_input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=2, shuffle=False, num_epochs=1) - failing_input_fn() - - def testNumpyInputFnWithTargetKeyAlreadyInX(self): - array = np.arange(32, 36) - x = {'__target_key__': array} - y = np.arange(4) - - with self.test_session(): - input_fn = numpy_io.numpy_input_fn( - x, y, batch_size=2, shuffle=False, num_epochs=1) - input_fn() - self.assertAllEqual(x['__target_key__'], array) - self.assertItemsEqual(x.keys(), ['__target_key__']) - - def testNumpyInputFnWithMismatchLengthOfInputs(self): - a = np.arange(4) * 1.0 - b = np.arange(32, 36) - x = {'a': a, 'b': b} - x_mismatch_length = {'a': np.arange(1), 'b': b} - y_longer_length = np.arange(10) - - with self.test_session(): - with self.assertRaisesRegexp( - ValueError, 'Length of tensors in x and y is mismatched.'): - failing_input_fn = numpy_io.numpy_input_fn( - x, y_longer_length, batch_size=2, shuffle=False, num_epochs=1) - failing_input_fn() - - with self.assertRaisesRegexp( - ValueError, 'Length of tensors in x and y is mismatched.'): - failing_input_fn = numpy_io.numpy_input_fn( - x=x_mismatch_length, - y=None, - batch_size=2, - shuffle=False, - num_epochs=1) - failing_input_fn() - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/learn/python/learn/monitors.py b/tensorflow/contrib/learn/python/learn/monitors.py index 3e0b1ad21a9a4a08fa94c8e9796f2b0dd5f8d622..51381a7427c919592b8e818c4b46dba974992610 100644 --- a/tensorflow/contrib/learn/python/learn/monitors.py +++ b/tensorflow/contrib/learn/python/learn/monitors.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Monitors instrument the training process. @@get_default_monitors @@ -151,8 +150,8 @@ class BaseMonitor(object): ValueError: if we've not begun an epoch, or `epoch` number does not match. """ if self._current_epoch != epoch: - raise ValueError( - "epoch_end expected %s but got %s.", self._current_epoch, epoch) + raise ValueError("epoch_end expected %s but got %s.", self._current_epoch, + epoch) self._current_epoch = None def step_begin(self, step): @@ -171,8 +170,8 @@ class BaseMonitor(object): ValueError: if we've already begun a step, or `step` < 0, or `step` > `max_steps`. """ - if (step < 0) or ( - (self._max_steps is not None) and (step > self._max_steps)): + if (step < 0) or ((self._max_steps is not None) and + (step > self._max_steps)): raise ValueError("Invalid step %s." % step) self._current_step = step return [] @@ -203,8 +202,8 @@ class BaseMonitor(object): ValueError: if we've not begun a step, or `step` number does not match. """ if self._current_step != step: - raise ValueError( - "step_end expected %s but got %s.", self._current_step, step) + raise ValueError("step_end expected %s but got %s.", self._current_step, + step) self._current_step = None return False @@ -253,6 +252,7 @@ class EveryN(BaseMonitor): treatment. """ + # TODO(ipolosukhin): Add also every n seconds. def __init__(self, every_n_steps=100, first_n_steps=1): @@ -475,8 +475,8 @@ class LoggingTrainable(EveryN): super(LoggingTrainable, self).every_n_step_begin(step) # Get a list of trainable variables at the beginning of every N steps. # We cannot get this in __init__ because train_op has not been generated. - trainables = ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES, - scope=self._scope) + trainables = ops.get_collection( + ops.GraphKeys.TRAINABLE_VARIABLES, scope=self._scope) self._names = {} for var in trainables: self._names[var.name] = var.value().name @@ -561,12 +561,19 @@ class ValidationMonitor(EveryN): provided. """ - def __init__(self, x=None, y=None, input_fn=None, batch_size=None, + def __init__(self, + x=None, + y=None, + input_fn=None, + batch_size=None, eval_steps=None, - every_n_steps=100, metrics=None, hooks=None, + every_n_steps=100, + metrics=None, + hooks=None, early_stopping_rounds=None, early_stopping_metric="loss", - early_stopping_metric_minimize=True, name=None): + early_stopping_metric_minimize=True, + name=None): """Initializes a ValidationMonitor. Args: @@ -597,8 +604,8 @@ class ValidationMonitor(EveryN): Raises: ValueError: If both x and input_fn are provided. """ - super(ValidationMonitor, self).__init__(every_n_steps=every_n_steps, - first_n_steps=-1) + super(ValidationMonitor, self).__init__( + every_n_steps=every_n_steps, first_n_steps=-1) # TODO(mdan): Checks like this are already done by evaluate. if x is None and input_fn is None: raise ValueError("Either x or input_fn should be provided.") @@ -654,20 +661,27 @@ class ValidationMonitor(EveryN): def _evaluate_estimator(self): if isinstance(self._estimator, core_estimator.Estimator): - if any((x is not None for x in - [self.x, self.y, self.batch_size, self.metrics])): + if any((x is not None + for x in [self.x, self.y, self.batch_size, self.metrics])): raise ValueError( "tf.estimator.Estimator does not support following " "arguments: x, y, batch_size, metrics. Should set as `None` " "in ValidationMonitor") return self._estimator.evaluate( - input_fn=self.input_fn, steps=self.eval_steps, hooks=self.hooks, + input_fn=self.input_fn, + steps=self.eval_steps, + hooks=self.hooks, name=self.name) else: return self._estimator.evaluate( - x=self.x, y=self.y, input_fn=self.input_fn, - batch_size=self.batch_size, steps=self.eval_steps, - metrics=self.metrics, hooks=self.hooks, name=self.name) + x=self.x, + y=self.y, + input_fn=self.input_fn, + batch_size=self.batch_size, + steps=self.eval_steps, + metrics=self.metrics, + hooks=self.hooks, + name=self.name) def every_n_step_end(self, step, outputs): super(ValidationMonitor, self).every_n_step_end(step, outputs) @@ -700,8 +714,9 @@ class ValidationMonitor(EveryN): # Early stopping logic. if self.early_stopping_rounds is not None: if self.early_stopping_metric not in validation_outputs: - raise ValueError("Metric %s missing from outputs %s." % ( - self.early_stopping_metric, set(validation_outputs.keys()))) + raise ValueError("Metric %s missing from outputs %s." % + (self.early_stopping_metric, + set(validation_outputs.keys()))) current_value = validation_outputs[self.early_stopping_metric] if (self._best_value is None or (self.early_stopping_metric_minimize and (current_value < self._best_value)) or @@ -712,9 +727,9 @@ class ValidationMonitor(EveryN): self._best_value_step = step stop_now = (step - self._best_value_step >= self.early_stopping_rounds) if stop_now: - logging.info("Stopping. Best step: {} with {} = {}." - .format(self._best_value_step, - self.early_stopping_metric, self._best_value)) + logging.info("Stopping. Best step: {} with {} = {}.".format( + self._best_value_step, self.early_stopping_metric, + self._best_value)) self._early_stopped = True return True return False @@ -763,8 +778,11 @@ class CaptureVariable(EveryN): self._var_values[step] = _extract_output(outputs, self._var_name) -def get_default_monitors(loss_op=None, summary_op=None, save_summary_steps=100, - output_dir=None, summary_writer=None): +def get_default_monitors(loss_op=None, + summary_op=None, + save_summary_steps=100, + output_dir=None, + summary_writer=None): """Returns a default set of typically-used monitors. Args: @@ -782,9 +800,12 @@ def get_default_monitors(loss_op=None, summary_op=None, save_summary_steps=100, if loss_op is not None: monitors.append(PrintTensor(tensor_names={"loss": loss_op.name})) if summary_op is not None: - monitors.append(SummarySaver(summary_op, save_steps=save_summary_steps, - output_dir=output_dir, - summary_writer=summary_writer)) + monitors.append( + SummarySaver( + summary_op, + save_steps=save_summary_steps, + output_dir=output_dir, + summary_writer=summary_writer)) return monitors @@ -794,8 +815,10 @@ class GraphDump(BaseMonitor): Note, this is very expensive, prefer `PrintTensor` in production. """ - IGNORE_OPS = ["Const", "Assign", "Identity", "Placeholder", - "RandomUniform", "Cast", "RestoreSlice"] + IGNORE_OPS = [ + "Const", "Assign", "Identity", "Placeholder", "RandomUniform", "Cast", + "RestoreSlice" + ] def __init__(self, ignore_ops=None): """Initializes GraphDump monitor. @@ -856,7 +879,7 @@ class GraphDump(BaseMonitor): this_output = self.data[step] if step in self.data else {} other_output = other_dump.data[step] if step in other_dump.data else {} for key in this_output: - if not isinstance(key, str) and not isinstance(key, unicode): + if not isinstance(key, six.string_types): continue if key not in other_output: raise ValueError("%s missing at step %s.", (key, step)) @@ -881,8 +904,8 @@ class ExportMonitor(EveryN): """Monitor that exports Estimator every N steps.""" @deprecation.deprecated("2017-03-25", - "ExportMonitor is deprecated. Please pass an " - "ExportStrategy to Experiment instead.") + "ExportMonitor is deprecated. Please pass an " + "ExportStrategy to Experiment instead.") def __init__(self, every_n_steps, export_dir, @@ -1088,8 +1111,7 @@ class CheckpointSaver(BaseMonitor): class StepCounter(EveryN): """Steps per second monitor.""" - def __init__(self, every_n_steps=100, output_dir=None, - summary_writer=None): + def __init__(self, every_n_steps=100, output_dir=None, summary_writer=None): super(StepCounter, self).__init__(every_n_steps=every_n_steps) self._summary_tag = "global_step/sec" self._last_reported_step = None @@ -1101,7 +1123,8 @@ class StepCounter(EveryN): def set_estimator(self, estimator): super(StepCounter, self).set_estimator(estimator) if self._summary_writer is None: - self._summary_writer = core_summary.FileWriterCache.get(estimator.model_dir) + self._summary_writer = core_summary.FileWriterCache.get( + estimator.model_dir) def every_n_step_end(self, current_step, outputs): current_time = time.time() @@ -1109,8 +1132,9 @@ class StepCounter(EveryN): added_steps = current_step - self._last_reported_step elapsed_time = current_time - self._last_reported_time steps_per_sec = added_steps / elapsed_time - summary = Summary(value=[Summary.Value(tag=self._summary_tag, - simple_value=steps_per_sec)]) + summary = Summary(value=[ + Summary.Value(tag=self._summary_tag, simple_value=steps_per_sec) + ]) self._summary_writer.add_summary(summary, current_step) self._last_reported_step = current_step self._last_reported_time = current_time diff --git a/tensorflow/contrib/learn/python/learn/ops/ops_test.py b/tensorflow/contrib/learn/python/learn/ops/ops_test.py index d0b9eb8abcbee187b6c53b7b419882f0a1e7da51..80d4923db37feb2a1304218f501ab51f9e0d9a14 100644 --- a/tensorflow/contrib/learn/python/learn/ops/ops_test.py +++ b/tensorflow/contrib/learn/python/learn/ops/ops_test.py @@ -20,7 +20,6 @@ from __future__ import print_function import numpy as np -from tensorflow.contrib.layers import conv2d from tensorflow.contrib.learn.python.learn import ops from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes diff --git a/tensorflow/contrib/learn/python/learn/trainable.py b/tensorflow/contrib/learn/python/learn/trainable.py index 972fec026f25d39dca75e8c5bafffb57fcd323fa..429b6040be21d8cbe1f2bba58090366552fdfbe7 100644 --- a/tensorflow/contrib/learn/python/learn/trainable.py +++ b/tensorflow/contrib/learn/python/learn/trainable.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """`Trainable` interface.""" from __future__ import absolute_import @@ -28,18 +27,31 @@ class Trainable(object): __metaclass__ = abc.ABCMeta @abc.abstractmethod - def fit(self, x=None, y=None, input_fn=None, steps=None, batch_size=None, - monitors=None, max_steps=None): + def fit(self, + x=None, + y=None, + input_fn=None, + steps=None, + batch_size=None, + monitors=None, + max_steps=None): """Trains a model given training data `x` predictions and `y` labels. Args: - x: Matrix of shape [n_samples, n_features...] or the dictionary of Matrices. - Can be iterator that returns arrays of features or dictionary of arrays of features. - The training input samples for fitting the model. If set, `input_fn` must be `None`. - y: Vector or matrix [n_samples] or [n_samples, n_outputs] or the dictionary of same. - Can be iterator that returns array of labels or dictionary of array of labels. - The training label values (class labels in classification, real numbers in regression). - If set, `input_fn` must be `None`. Note: For classification, label values must + x: Matrix of shape [n_samples, n_features...] or the dictionary of + Matrices. + Can be iterator that returns arrays of features or dictionary of arrays + of features. + The training input samples for fitting the model. If set, `input_fn` + must be `None`. + y: Vector or matrix [n_samples] or [n_samples, n_outputs] or the + dictionary of same. + Can be iterator that returns array of labels or dictionary of array of + labels. + The training label values (class labels in classification, real numbers + in regression). + If set, `input_fn` must be `None`. Note: For classification, label + values must be integers representing the class index (i.e. values from 0 to n_classes-1). input_fn: Input function returning a tuple of: diff --git a/tensorflow/contrib/learn/python/learn/utils/export_test.py b/tensorflow/contrib/learn/python/learn/utils/export_test.py index 95070ada3b9d3ccb00009bd9b885e8163d7fbed4..9bfb1fc952c07bd6c09d1f1074e8dc5539dc0529 100644 --- a/tensorflow/contrib/learn/python/learn/utils/export_test.py +++ b/tensorflow/contrib/learn/python/learn/utils/export_test.py @@ -50,6 +50,7 @@ def _training_input_fn(): class ExportTest(test.TestCase): + def _get_default_signature(self, export_meta_filename): """ Gets the default signature from the export.meta file. """ with session.Session(): @@ -69,18 +70,18 @@ class ExportTest(test.TestCase): # Only the written checkpoints are exported. self.assertTrue( saver.checkpoint_exists(os.path.join(export_dir, '00000001', 'export')), - 'Exported checkpoint expected but not found: %s' % - os.path.join(export_dir, '00000001', 'export')) + 'Exported checkpoint expected but not found: %s' % os.path.join( + export_dir, '00000001', 'export')) self.assertTrue( saver.checkpoint_exists(os.path.join(export_dir, '00000010', 'export')), - 'Exported checkpoint expected but not found: %s' % - os.path.join(export_dir, '00000010', 'export')) + 'Exported checkpoint expected but not found: %s' % os.path.join( + export_dir, '00000010', 'export')) self.assertEquals( six.b(os.path.join(export_dir, '00000010')), export_monitor.last_export_dir) # Validate the signature signature = self._get_default_signature( - os.path.join(export_dir, '00000010', 'export.meta')) + os.path.join(export_dir, '00000010', 'export.meta')) self.assertTrue(signature.HasField(expected_signature)) def testExportMonitor_EstimatorProvidesSignature(self): @@ -116,8 +117,7 @@ class ExportTest(test.TestCase): def _serving_input_fn(): return { _X_KEY: - random_ops.random_uniform( - shape=(1,), minval=0.0, maxval=1000.0) + random_ops.random_uniform(shape=(1,), minval=0.0, maxval=1000.0) }, None input_feature_key = 'my_example_key' @@ -160,8 +160,7 @@ class ExportTest(test.TestCase): input_feature_key: None, _X_KEY: - random_ops.random_uniform( - shape=(1,), minval=0.0, maxval=1000.0) + random_ops.random_uniform(shape=(1,), minval=0.0, maxval=1000.0) }, None monitor = learn.monitors.ExportMonitor( @@ -182,8 +181,7 @@ class ExportTest(test.TestCase): def _serving_input_fn(): return { input_feature_key: - array_ops.placeholder( - dtype=dtypes.string, shape=(1,)) + array_ops.placeholder(dtype=dtypes.string, shape=(1,)) }, None monitor = learn.monitors.ExportMonitor( @@ -204,11 +202,9 @@ class ExportTest(test.TestCase): def _serving_input_fn(): return { input_feature_key: - array_ops.placeholder( - dtype=dtypes.string, shape=(1,)), + array_ops.placeholder(dtype=dtypes.string, shape=(1,)), _X_KEY: - random_ops.random_uniform( - shape=(1,), minval=0.0, maxval=1000.0) + random_ops.random_uniform(shape=(1,), minval=0.0, maxval=1000.0) }, None export_dir = os.path.join(tempfile.mkdtemp(), 'export') @@ -227,8 +223,8 @@ class ExportTest(test.TestCase): def _regression_signature(examples, unused_features, predictions): signatures = {} - signatures['regression'] = (exporter.regression_signature(examples, - predictions)) + signatures['regression'] = ( + exporter.regression_signature(examples, predictions)) return signatures['regression'], signatures random.seed(42) @@ -248,10 +244,10 @@ class ExportTest(test.TestCase): with self.assertRaises(errors.NotFoundError): saver.checkpoint_exists(os.path.join(export_dir, '00000000', 'export')) self.assertTrue( - saver.checkpoint_exists(os.path.join(export_dir, '00000010', 'export'))) + saver.checkpoint_exists(os.path.join(export_dir, '00000010', 'export'))) # Validate the signature signature = self._get_default_signature( - os.path.join(export_dir, '00000010', 'export.meta')) + os.path.join(export_dir, '00000010', 'export.meta')) self.assertTrue(signature.HasField('regression_signature')) diff --git a/tensorflow/contrib/learn/python/learn/utils/gc_test.py b/tensorflow/contrib/learn/python/learn/utils/gc_test.py index 76cfd88e1d68856907131f7e2bae65d4c9fcc4b1..e7d091e18a8f186f89f5217442c24fb106c5cdab 100644 --- a/tensorflow/contrib/learn/python/learn/utils/gc_test.py +++ b/tensorflow/contrib/learn/python/learn/utils/gc_test.py @@ -34,12 +34,13 @@ def _create_parser(base_dir): # create a simple parser that pulls the export_version from the directory. def parser(path): # Modify the path object for RegEx match for Windows Paths - if os.name == 'nt': - match = re.match("^" + compat.as_str_any(base_dir).replace('\\','/') + "/(\\d+)$", - compat.as_str_any(path.path).replace('\\','/')) + if os.name == "nt": + match = re.match( + "^" + compat.as_str_any(base_dir).replace("\\", "/") + "/(\\d+)$", + compat.as_str_any(path.path).replace("\\", "/")) else: match = re.match("^" + compat.as_str_any(base_dir) + "/(\\d+)$", - compat.as_str_any(path.path)) + compat.as_str_any(path.path)) if not match: return None return path._replace(export_version=int(match.group(1))) @@ -63,7 +64,9 @@ class GcTest(test_util.TensorFlowTestCase): def testModExportVersion(self): paths = [ - gc.Path("/foo", 4), gc.Path("/foo", 5), gc.Path("/foo", 6), + gc.Path("/foo", 4), + gc.Path("/foo", 5), + gc.Path("/foo", 6), gc.Path("/foo", 9) ] mod = gc.mod_export_version(2) @@ -73,14 +76,21 @@ class GcTest(test_util.TensorFlowTestCase): def testOneOfEveryNExportVersions(self): paths = [ - gc.Path("/foo", 0), gc.Path("/foo", 1), gc.Path("/foo", 3), - gc.Path("/foo", 5), gc.Path("/foo", 6), gc.Path("/foo", 7), - gc.Path("/foo", 8), gc.Path("/foo", 33) + gc.Path("/foo", 0), + gc.Path("/foo", 1), + gc.Path("/foo", 3), + gc.Path("/foo", 5), + gc.Path("/foo", 6), + gc.Path("/foo", 7), + gc.Path("/foo", 8), + gc.Path("/foo", 33) ] one_of = gc.one_of_every_n_export_versions(3) self.assertEqual( one_of(paths), [ - gc.Path("/foo", 3), gc.Path("/foo", 6), gc.Path("/foo", 8), + gc.Path("/foo", 3), + gc.Path("/foo", 6), + gc.Path("/foo", 8), gc.Path("/foo", 33) ]) @@ -98,13 +108,19 @@ class GcTest(test_util.TensorFlowTestCase): f = gc.union(gc.largest_export_versions(3), gc.mod_export_version(3)) self.assertEqual( f(paths), [ - gc.Path("/foo", 0), gc.Path("/foo", 3), gc.Path("/foo", 6), - gc.Path("/foo", 7), gc.Path("/foo", 8), gc.Path("/foo", 9) + gc.Path("/foo", 0), + gc.Path("/foo", 3), + gc.Path("/foo", 6), + gc.Path("/foo", 7), + gc.Path("/foo", 8), + gc.Path("/foo", 9) ]) def testNegation(self): paths = [ - gc.Path("/foo", 4), gc.Path("/foo", 5), gc.Path("/foo", 6), + gc.Path("/foo", 4), + gc.Path("/foo", 5), + gc.Path("/foo", 6), gc.Path("/foo", 9) ] mod = gc.negation(gc.mod_export_version(2)) @@ -121,8 +137,7 @@ class GcTest(test_util.TensorFlowTestCase): gfile.MakeDirs(os.path.join(base_dir, "ignore")) self.assertEqual( - gc.get_paths(base_dir, _create_parser(base_dir)), - [ + gc.get_paths(base_dir, _create_parser(base_dir)), [ gc.Path(os.path.join(base_dir, "0"), 0), gc.Path(os.path.join(base_dir, "1"), 1), gc.Path(os.path.join(base_dir, "2"), 2) @@ -131,10 +146,10 @@ class GcTest(test_util.TensorFlowTestCase): def testMixedStrTypes(self): temp_dir = compat.as_bytes(test.get_temp_dir()) - for sub_dir in ['str', b'bytes', u'unicode']: + for sub_dir in ["str", b"bytes", u"unicode"]: base_dir = os.path.join( - (temp_dir if isinstance(sub_dir, bytes) else temp_dir.decode()), - sub_dir) + (temp_dir + if isinstance(sub_dir, bytes) else temp_dir.decode()), sub_dir) self.assertFalse(gfile.Exists(base_dir)) gfile.MakeDirs(os.path.join(compat.as_str_any(base_dir), "42")) gc.get_paths(base_dir, _create_parser(base_dir)) diff --git a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py index 4b404a8e20e33a17a0d5f857e4220f90c7bc799f..1593380007b2799fb1d17e92408ab19a7b47fe1e 100644 --- a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py +++ b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py @@ -50,6 +50,7 @@ from tensorflow.python.platform import gfile from tensorflow.python.platform import tf_logging as logging from tensorflow.python.saved_model import signature_constants from tensorflow.python.saved_model import signature_def_utils +from tensorflow.python.summary import summary_iterator from tensorflow.python.training import saver from tensorflow.python.util import compat @@ -390,7 +391,8 @@ def make_export_strategy(serving_input_fn, default_output_alternative_key=None, assets_extra=None, as_text=False, - exports_to_keep=5): + exports_to_keep=5, + strip_default_attrs=None): """Create an ExportStrategy for use with Experiment. Args: @@ -411,12 +413,16 @@ def make_export_strategy(serving_input_fn, exports_to_keep: Number of exports to keep. Older exports will be garbage-collected. Defaults to 5. Set to None to disable garbage collection. + strip_default_attrs: Boolean. If True, default attrs in the + `GraphDef` will be stripped on write. This is recommended for better + forward compatibility of the resulting `SavedModel`. Returns: An ExportStrategy that can be passed to the Experiment constructor. """ - def export_fn(estimator, export_dir_base, checkpoint_path=None): + def export_fn(estimator, export_dir_base, checkpoint_path=None, + strip_default_attrs=False): """Exports the given Estimator as a SavedModel. Args: @@ -425,6 +431,8 @@ def make_export_strategy(serving_input_fn, graph and checkpoints. checkpoint_path: The checkpoint path to export. If None (the default), the most recent checkpoint found within the model directory is chosen. + strip_default_attrs: Boolean. If `True`, default-valued attributes will + be removed from the NodeDefs. Returns: The string path to the exported directory. @@ -443,7 +451,8 @@ def make_export_strategy(serving_input_fn, serving_input_fn, assets_extra=assets_extra, as_text=as_text, - checkpoint_path=checkpoint_path) + checkpoint_path=checkpoint_path, + strip_default_attrs=strip_default_attrs) else: export_result = estimator.export_savedmodel( export_dir_base, @@ -451,12 +460,13 @@ def make_export_strategy(serving_input_fn, default_output_alternative_key=default_output_alternative_key, assets_extra=assets_extra, as_text=as_text, - checkpoint_path=checkpoint_path) + checkpoint_path=checkpoint_path, + strip_default_attrs=strip_default_attrs) garbage_collect_exports(export_dir_base, exports_to_keep) return export_result - return export_strategy.ExportStrategy('Servo', export_fn) + return export_strategy.ExportStrategy('Servo', export_fn, strip_default_attrs) def make_parsing_export_strategy(feature_columns, @@ -464,7 +474,8 @@ def make_parsing_export_strategy(feature_columns, assets_extra=None, as_text=False, exports_to_keep=5, - target_core=False): + target_core=False, + strip_default_attrs=None): """Create an ExportStrategy for use with Experiment, using `FeatureColumn`s. Creates a SavedModel export that expects to be fed with a single string @@ -492,6 +503,9 @@ def make_parsing_export_strategy(feature_columns, target_core: If True, prepare an ExportStrategy for use with tensorflow.python.estimator.*. If False (default), prepare an ExportStrategy for use with tensorflow.contrib.learn.python.learn.*. + strip_default_attrs: Boolean. If True, default attrs in the + `GraphDef` will be stripped on write. This is recommended for better + forward compatibility of the resulting `SavedModel`. Returns: An ExportStrategy that can be passed to the Experiment constructor. @@ -508,7 +522,8 @@ def make_parsing_export_strategy(feature_columns, default_output_alternative_key=default_output_alternative_key, assets_extra=assets_extra, as_text=as_text, - exports_to_keep=exports_to_keep) + exports_to_keep=exports_to_keep, + strip_default_attrs=strip_default_attrs) def _default_compare_fn(curr_best_eval_result, cand_eval_result): @@ -542,15 +557,16 @@ def _default_compare_fn(curr_best_eval_result, cand_eval_result): class BestModelSelector(object): """A helper that keeps track of export selection candidates.""" - def __init__(self, compare_fn=None): + def __init__(self, event_file_pattern=None, compare_fn=None): """Constructor of this class. Args: + event_file_pattern: absolute event file name pattern. compare_fn: a function that returns true if the candidate is better than the current best model. """ - self._best_eval_result = None self._compare_fn = compare_fn or _default_compare_fn + self._best_eval_result = self._get_best_eval_result(event_file_pattern) def update(self, checkpoint_path, eval_result): """Records a given checkpoint and exports if this is the best model. @@ -580,11 +596,40 @@ class BestModelSelector(object): else: return '', None + def _get_best_eval_result(self, event_files): + """Get the best eval result from event files. + + Args: + event_files: Absolute pattern of event files. -def make_best_model_export_strategy(serving_input_fn, - exports_to_keep=1, - compare_fn=None, - default_output_alternative_key=None): + Returns: + The best eval result. + """ + if not event_files: + return None + + best_eval_result = None + for event_file in gfile.Glob(os.path.join(event_files)): + for event in summary_iterator.summary_iterator(event_file): + if event.HasField('summary'): + event_eval_result = {} + for value in event.summary.value: + if value.HasField('simple_value'): + event_eval_result[value.tag] = value.simple_value + if best_eval_result is None or self._compare_fn( + best_eval_result, event_eval_result): + best_eval_result = event_eval_result + return best_eval_result + + +def make_best_model_export_strategy( + serving_input_fn, + exports_to_keep=1, + model_dir=None, + event_file_pattern=None, + compare_fn=None, + default_output_alternative_key=None, + strip_default_attrs=None): """Creates an custom ExportStrategy for use with tf.contrib.learn.Experiment. Args: @@ -592,10 +637,24 @@ def make_best_model_export_strategy(serving_input_fn, `InputFnOps`. exports_to_keep: an integer indicating how many historical best models need to be preserved. + model_dir: Directory where model parameters, graph etc. are saved. This will + be used to load eval metrics from the directory when the export strategy + is created. So the best metrics would not be lost even if the export + strategy got preempted, which guarantees that only the best model would + be exported regardless of preemption. If None, however, the export + strategy would not be preemption-safe. To be preemption-safe, both + model_dir and event_file_pattern would be needed. + event_file_pattern: event file name pattern relative to model_dir, e.g. + "eval_continuous/*.tfevents.*". If None, however, the export strategy + would not be preemption-safe. To be preemption-safe, both + model_dir and event_file_pattern would be needed. compare_fn: a function that select the 'best' candidate from a dictionary of evaluation result keyed by corresponding checkpoint path. default_output_alternative_key: the key for default serving signature for multi-headed inference graphs. + strip_default_attrs: Boolean. If True, default attrs in the + `GraphDef` will be stripped on write. This is recommended for better + forward compatibility of the resulting `SavedModel`. Returns: An ExportStrategy that can be passed to the Experiment constructor. @@ -603,9 +662,13 @@ def make_best_model_export_strategy(serving_input_fn, best_model_export_strategy = make_export_strategy( serving_input_fn, exports_to_keep=exports_to_keep, - default_output_alternative_key=default_output_alternative_key) + default_output_alternative_key=default_output_alternative_key, + strip_default_attrs=strip_default_attrs) - best_model_selector = BestModelSelector(compare_fn) + full_event_file_pattern = os.path.join( + model_dir, + event_file_pattern) if model_dir and event_file_pattern else None + best_model_selector = BestModelSelector(full_event_file_pattern, compare_fn) def export_fn(estimator, export_dir_base, checkpoint_path, eval_result=None): """Exports the given Estimator as a SavedModel. diff --git a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py index 628eb254c3b1129648c453dc47f0c0919891de6f..14bf1136e8e9ab1488c4850d458382028ec5583d 100644 --- a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py +++ b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils_test.py @@ -24,13 +24,14 @@ import time from tensorflow.contrib.layers.python.layers import feature_column as fc from tensorflow.contrib.learn.python.learn import export_strategy as export_strategy_lib from tensorflow.contrib.learn.python.learn.estimators import constants -from tensorflow.contrib.learn.python.learn.estimators import estimator as core_estimator +from tensorflow.contrib.learn.python.learn.estimators import estimator from tensorflow.contrib.learn.python.learn.estimators import model_fn from tensorflow.contrib.learn.python.learn.utils import input_fn_utils from tensorflow.contrib.learn.python.learn.utils import saved_model_export_utils from tensorflow.core.framework import tensor_shape_pb2 from tensorflow.core.framework import types_pb2 from tensorflow.core.protobuf import meta_graph_pb2 +from tensorflow.python.estimator import estimator as core_estimator from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops @@ -41,7 +42,7 @@ from tensorflow.python.saved_model import signature_def_utils from tensorflow.python.util import compat -class TestEstimator(core_estimator.Estimator): +class TestEstimator(estimator.Estimator): def __init__(self, *args, **kwargs): super(TestEstimator, self).__init__(*args, **kwargs) @@ -55,7 +56,8 @@ class TestEstimator(core_estimator.Estimator): default_output_alternative_key=None, assets_extra=None, as_text=False, - checkpoint_path=None): + checkpoint_path=None, + strip_default_attrs=False): if not os.path.exists(export_dir): os.makedirs(export_dir) @@ -93,9 +95,9 @@ class SavedModelExportUtilsTest(test.TestCase): name="input-tensor-1:0", dtype=dtype_string, tensor_shape=shape)) expected_signature_def.outputs[ signature_constants.REGRESS_OUTPUTS].CopyFrom( - meta_graph_pb2.TensorInfo(name="output-tensor-1:0", - dtype=dtype_float, - tensor_shape=shape)) + meta_graph_pb2.TensorInfo( + name="output-tensor-1:0", dtype=dtype_float, + tensor_shape=shape)) expected_signature_def.method_name = signature_constants.REGRESS_METHOD_NAME self.assertEqual(actual_signature_def, expected_signature_def) @@ -506,7 +508,9 @@ class SavedModelExportUtilsTest(test.TestCase): input_example = constant_op.constant(["input string"]) input_ops = input_fn_utils.InputFnOps({ "features": input_features - }, None, {"default input": input_example}) + }, None, { + "default input": input_example + }) input_alternatives, _ = ( saved_model_export_utils.get_input_alternatives(input_ops)) output_1 = constant_op.constant([1.0]) @@ -527,8 +531,9 @@ class SavedModelExportUtilsTest(test.TestCase): model_fn.ModeKeys.INFER, predictions={"some_output": constant_op.constant(["4"])}, output_alternatives=provided_output_alternatives) - output_alternatives, _ = (saved_model_export_utils.get_output_alternatives( - model_fn_ops, "head-1")) + output_alternatives, _ = ( + saved_model_export_utils.get_output_alternatives( + model_fn_ops, "head-1")) signature_defs = saved_model_export_utils.build_all_signature_defs( input_alternatives, output_alternatives, "head-1") @@ -546,7 +551,9 @@ class SavedModelExportUtilsTest(test.TestCase): "default_input_alternative:head-3": signature_def_utils.predict_signature_def({ "default input": input_example - }, {"some_output_3": output_3}), + }, { + "some_output_3": output_3 + }), # "features_input_alternative:head-1": # signature_def_utils.regression_signature_def(input_features, # output_1), @@ -589,8 +596,9 @@ class SavedModelExportUtilsTest(test.TestCase): model_fn.ModeKeys.INFER, predictions={"some_output": constant_op.constant(["4"])}, output_alternatives=provided_output_alternatives) - output_alternatives, _ = (saved_model_export_utils.get_output_alternatives( - model_fn_ops, "head-1")) + output_alternatives, _ = ( + saved_model_export_utils.get_output_alternatives( + model_fn_ops, "head-1")) with self.assertRaisesRegexp( ValueError, "A default input_alternative must be provided"): @@ -706,25 +714,72 @@ class SavedModelExportUtilsTest(test.TestCase): self.assertNotEqual("", export_strategy.export(test_estimator, export_dir_base, - "fake_ckpt_0", {"loss": 100})) + "fake_ckpt_0", { + "loss": 100 + })) self.assertNotEqual("", test_estimator.last_exported_dir) self.assertNotEqual("", test_estimator.last_exported_checkpoint) self.assertEqual("", export_strategy.export(test_estimator, export_dir_base, - "fake_ckpt_1", {"loss": 101})) + "fake_ckpt_1", { + "loss": 101 + })) self.assertEqual(test_estimator.last_exported_dir, os.path.join(export_dir_base, "fake_ckpt_0")) self.assertNotEqual("", export_strategy.export(test_estimator, export_dir_base, - "fake_ckpt_2", {"loss": 10})) + "fake_ckpt_2", { + "loss": 10 + })) + self.assertEqual(test_estimator.last_exported_dir, + os.path.join(export_dir_base, "fake_ckpt_2")) + + self.assertEqual("", + export_strategy.export(test_estimator, export_dir_base, + "fake_ckpt_3", { + "loss": 20 + })) + self.assertEqual(test_estimator.last_exported_dir, + os.path.join(export_dir_base, "fake_ckpt_2")) + + def test_make_best_model_export_strategy_with_preemption(self): + model_dir = self.get_temp_dir() + eval_dir_base = os.path.join(model_dir, "eval_continuous") + core_estimator._write_dict_to_summary(eval_dir_base, {"loss": 50}, 1) + core_estimator._write_dict_to_summary(eval_dir_base, {"loss": 60}, 2) + + test_estimator = TestEstimator() + export_strategy = saved_model_export_utils.make_best_model_export_strategy( + serving_input_fn=None, + exports_to_keep=3, + model_dir=model_dir, + event_file_pattern="eval_continuous/*.tfevents.*", + compare_fn=None) + + export_dir_base = os.path.join(self.get_temp_dir(), "export") + self.assertEqual("", + export_strategy.export(test_estimator, export_dir_base, + "fake_ckpt_0", { + "loss": 100 + })) + self.assertEqual("", test_estimator.last_exported_dir) + self.assertEqual("", test_estimator.last_exported_checkpoint) + + self.assertNotEqual("", + export_strategy.export(test_estimator, export_dir_base, + "fake_ckpt_2", { + "loss": 10 + })) self.assertEqual(test_estimator.last_exported_dir, os.path.join(export_dir_base, "fake_ckpt_2")) self.assertEqual("", export_strategy.export(test_estimator, export_dir_base, - "fake_ckpt_3", {"loss": 20})) + "fake_ckpt_3", { + "loss": 20 + })) self.assertEqual(test_estimator.last_exported_dir, os.path.join(export_dir_base, "fake_ckpt_2")) diff --git a/tensorflow/contrib/libsvm/kernels/decode_libsvm_op.cc b/tensorflow/contrib/libsvm/kernels/decode_libsvm_op.cc index 27ce55d5687b77052e14e89d550758aeb7dfe79a..720c74e3de5907fa006227d1278c45fd2175fe5f 100644 --- a/tensorflow/contrib/libsvm/kernels/decode_libsvm_op.cc +++ b/tensorflow/contrib/libsvm/kernels/decode_libsvm_op.cc @@ -22,10 +22,6 @@ limitations under the License. #include "tensorflow/core/lib/strings/str_util.h" namespace tensorflow { -namespace { -template -bool ConvertHelper(const string& s, T* value); -} template class DecodeLibsvmOp : public OpKernel { @@ -50,34 +46,47 @@ class DecodeLibsvmOp : public OpKernel { std::vector out_values; std::vector> out_indices; for (int i = 0; i < input_flat.size(); ++i) { - std::vector entries = - str_util::Split(input_flat(i), " ", str_util::SkipEmpty()); - OP_REQUIRES(ctx, !entries.empty(), - errors::InvalidArgument("No entries found for input[", i, + StringPiece line(input_flat(i)); + str_util::RemoveWhitespaceContext(&line); + + StringPiece piece; + OP_REQUIRES(ctx, str_util::ConsumeNonWhitespace(&line, &piece), + errors::InvalidArgument("No label found for input[", i, "]: \"", input_flat(i), "\"")); + Tlabel label_value; - OP_REQUIRES( - ctx, ConvertHelper(entries[0], &label_value), - errors::InvalidArgument("Label format incorrect: ", entries[0])); + OP_REQUIRES(ctx, + strings::SafeStringToNumeric(piece, &label_value), + errors::InvalidArgument("Label format incorrect: ", piece)); + label(i) = label_value; - for (int j = 1; j < entries.size(); j++) { - std::vector pair = str_util::Split(entries[j], ":"); - OP_REQUIRES( - ctx, (pair.size() == 2), - errors::InvalidArgument("Invalid feature \"", entries[j], "\"")); + + str_util::RemoveLeadingWhitespace(&line); + while (str_util::ConsumeNonWhitespace(&line, &piece)) { + size_t p = piece.find(':'); + OP_REQUIRES(ctx, (p != StringPiece::npos), + errors::InvalidArgument("Invalid feature \"", piece, "\"")); + int64 feature_index; OP_REQUIRES( - ctx, strings::safe_strto64(pair[0].c_str(), &feature_index), - errors::InvalidArgument("Feature format incorrect: ", entries[j])); + ctx, strings::safe_strto64(piece.substr(0, p), &feature_index), + errors::InvalidArgument("Feature format incorrect: ", piece)); OP_REQUIRES(ctx, (feature_index >= 0), errors::InvalidArgument( "Feature index should be >= 0, got ", feature_index)); + T feature_value; OP_REQUIRES( - ctx, ConvertHelper(pair[1], &feature_value), - errors::InvalidArgument("Feature format incorrect: ", entries[j])); + + ctx, + strings::SafeStringToNumeric(piece.substr(p + 1), + &feature_value), + errors::InvalidArgument("Feature format incorrect: ", piece)); + out_values.emplace_back(feature_value); out_indices.emplace_back(std::pair(i, feature_index)); + + str_util::RemoveLeadingWhitespace(&line); } } @@ -128,25 +137,6 @@ class DecodeLibsvmOp : public OpKernel { int64 num_features_; }; -namespace { -template <> -bool ConvertHelper(const string& s, float* value) { - return strings::safe_strtof(s.c_str(), value); -} -template <> -bool ConvertHelper(const string& s, double* value) { - return strings::safe_strtod(s.c_str(), value); -} -template <> -bool ConvertHelper(const string& s, int32* value) { - return strings::safe_strto32(s.c_str(), value); -} -template <> -bool ConvertHelper(const string& s, int64* value) { - return strings::safe_strto64(s.c_str(), value); -} -} // namespace - #define REGISTER_KERNEL(type) \ REGISTER_KERNEL_BUILDER(Name("DecodeLibsvm") \ .Device(DEVICE_CPU) \ diff --git a/tensorflow/contrib/linear_optimizer/BUILD b/tensorflow/contrib/linear_optimizer/BUILD index fe2f183ac970cef4ebf6ca1a927b5a48eefb7d7b..cea3627ed565f0de86d8d9bb6b45c4b19c5b5558 100644 --- a/tensorflow/contrib/linear_optimizer/BUILD +++ b/tensorflow/contrib/linear_optimizer/BUILD @@ -126,6 +126,7 @@ py_library( py_test( name = "sdca_estimator_test", srcs = ["python/sdca_estimator_test.py"], + shard_count = 4, srcs_version = "PY2AND3", deps = [ ":sdca_estimator_py", diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py b/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py index 7526f3ae0dbdb3d6827e9d7f690090b8438e4f6e..3f5fdc18bb8f47cceee8f81dd5ded02059344b8b 100644 --- a/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py +++ b/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py @@ -211,9 +211,8 @@ class SdcaModel(object): sums.append( math_ops.reduce_sum( math_ops.abs(math_ops.cast(weights, dtypes.float64)))) - sum = math_ops.add_n(sums) # SDCA L1 regularization cost is: l1 * sum(|weights|) - return self._options['symmetric_l1_regularization'] * sum + return self._options['symmetric_l1_regularization'] * math_ops.add_n(sums) def _l2_loss(self, l2): """Computes the (un-normalized) l2 loss of the model.""" @@ -225,9 +224,8 @@ class SdcaModel(object): sums.append( math_ops.reduce_sum( math_ops.square(math_ops.cast(weights, dtypes.float64)))) - sum = math_ops.add_n(sums) # SDCA L2 regularization cost is: l2 * sum(weights^2) / 2 - return l2 * sum / 2.0 + return l2 * math_ops.add_n(sums) / 2.0 def _convert_n_to_tensor(self, input_list, as_ref=False): """Converts input list to a set of tensors.""" diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py index 7e214905b13db6a7e2f54f15873f5a9aedb4f44f..ec726bbed41a86eb314e3591ecaedaa6bf0e5e9b 100644 --- a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py +++ b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py @@ -102,7 +102,7 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): keys.get_shape()) def lookup(self, keys, name=None): - if keys.dtype != self._key_dtype: + if keys.dtype.base_dtype != self._key_dtype: raise TypeError('Signature mismatch. Keys must be dtype %s, got %s.' % (self._key_dtype, keys.dtype)) self._check_keys(keys) diff --git a/tensorflow/contrib/lite/BUILD b/tensorflow/contrib/lite/BUILD index 3f1b0be1a73a3ff1da3452f4ee1a9125f9e26178..44c4a7e2ca8d019ca602c7f2b492cd1e70b17561 100644 --- a/tensorflow/contrib/lite/BUILD +++ b/tensorflow/contrib/lite/BUILD @@ -6,8 +6,11 @@ licenses(["notice"]) # Apache 2.0 load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts", "gen_selected_ops") +exports_files(["LICENSE"]) + exports_files(glob([ "testdata/*.bin", + "testdata/*.pb", "models/testdata/*", ])) @@ -25,16 +28,35 @@ config_setting( }, ) -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_test", -) - cc_library( name = "schema_fbs_version", hdrs = ["version.h"], ) +cc_library( + name = "arena_planner", + srcs = ["arena_planner.cc"], + hdrs = ["arena_planner.h"], + deps = [ + ":context", + ":graph_info", + ":memory_planner", + ":simple_memory_arena", + ], +) + +cc_test( + name = "arena_planner_test", + size = "small", + srcs = ["arena_planner_test.cc"], + deps = [ + ":arena_planner", + "//tensorflow/contrib/lite/testing:util", + "//tensorflow/core:lib", + "@com_google_googletest//:gtest", + ], +) + # Main library. No ops are included here. # TODO(aselle): Resolve problems preventing C99 usage. cc_library( @@ -43,6 +65,25 @@ cc_library( hdrs = ["context.h"], ) +cc_library( + name = "graph_info", + hdrs = ["graph_info.h"], + deps = [":context"], +) + +cc_library( + name = "memory_planner", + hdrs = ["memory_planner.h"], + deps = [":context"], +) + +cc_library( + name = "simple_memory_arena", + srcs = ["simple_memory_arena.cc"], + hdrs = ["simple_memory_arena.h"], + deps = [":context"], +) + cc_library( name = "builtin_op_data", hdrs = [ @@ -66,27 +107,31 @@ cc_library( srcs = [ "allocation.cc", "error_reporter.cc", + "graph_info.cc", "interpreter.cc", "model.cc", "nnapi_delegate.cc", "optional_debug_tools.cc", - "simple_memory_arena.cc", ], hdrs = [ "allocation.h", "context.h", "error_reporter.h", + "graph_info.h", "interpreter.h", "model.h", "nnapi_delegate.h", "optional_debug_tools.h", - "simple_memory_arena.h", ], copts = tflite_copts(), deps = [ + ":arena_planner", ":builtin_op_data", ":context", + ":graph_info", + ":memory_planner", ":schema_fbs_version", + ":simple_memory_arena", "//tensorflow/contrib/lite/kernels:gemm_support", "//tensorflow/contrib/lite/nnapi:nnapi_lib", "//tensorflow/contrib/lite/schema:schema_fbs", @@ -124,6 +169,22 @@ cc_test( deps = [ ":framework", ":string_util", + "//tensorflow/contrib/lite/kernels/internal:tensor_utils", + "//tensorflow/contrib/lite/schema:schema_fbs", + "//tensorflow/contrib/lite/testing:util", + "@com_google_googletest//:gtest", + ], +) + +# Test graph utils +cc_test( + name = "graph_info_test", + size = "small", + srcs = ["graph_info_test.cc"], + deps = [ + ":framework", + ":string_util", + "//tensorflow/contrib/lite/testing:util", "@com_google_googletest//:gtest", ], ) @@ -134,7 +195,7 @@ cc_test( size = "small", srcs = ["simple_memory_arena_test.cc"], deps = [ - ":framework", + ":simple_memory_arena", "//tensorflow/contrib/lite/testing:util", "@com_google_googletest//:gtest", ], @@ -175,18 +236,18 @@ cc_test( # Model tests -cc_library( - name = "models_test_utils", - testonly = 1, - hdrs = ["models/test_utils.h"], - deps = select({ - "//tensorflow:android": [], - "//conditions:default": [ - "@com_google_absl//absl/strings", - "//tensorflow/core:test", - ], - }), -) +#cc_library( +# name = "models_test_utils", +# testonly = 1, +# hdrs = ["models/test_utils.h"], +# deps = select({ +# "//tensorflow:android": [], +# "//conditions:default": [ +# "@com_google_absl//absl/strings", +# "//tensorflow/core:test", +# ], +# }), +#) filegroup( name = "all_files", diff --git a/tensorflow/contrib/lite/Makefile b/tensorflow/contrib/lite/Makefile index 78402727abdd2742ffff54bf59ca076d8b97b042..7f316292724ea0baaf034d4e914773ad97a957d4 100644 --- a/tensorflow/contrib/lite/Makefile +++ b/tensorflow/contrib/lite/Makefile @@ -56,7 +56,7 @@ LIBS := \ -lz # If we're on Linux, also link in the dl library. -ifeq ($(OS),LINUX) +ifeq ($(HOST_OS),LINUX) LIBS += -ldl -lpthread endif diff --git a/tensorflow/contrib/lite/README.md b/tensorflow/contrib/lite/README.md index 852284cbc7f33b5d9c0f7774bca89c1dff3fa3ec..00e93d2c4f3ab27057b855fba6fccf2ec8d7a1c1 100644 --- a/tensorflow/contrib/lite/README.md +++ b/tensorflow/contrib/lite/README.md @@ -6,7 +6,7 @@ TensorFlow Lite uses many techniques for achieving low latency like optimizing t ![image](g3doc/TFLite-Architecture.jpg) # Getting Started with an Android Demo App -This section contains an example application using TensorFlow Lite for Android devices. The demo is a sample camera app that classifies images continuously using a quantized Mobilenet model. A device running Android 5.0 ( API 21) or higher is required to run the demo. +This section contains an example application using TensorFlow Lite for Android devices. The demo is a sample camera app that classifies images continuously using either a quantized Mobilenet model or a floating point Inception-v3 model. A device running Android 5.0 ( API 21) or higher is required to run the demo. There are 3 ways to get the demo app to your device - Download the prebuilt binary or @@ -29,9 +29,16 @@ The simplest way to compile the demo app, and try out changes to the project cod - Make sure the Android SDK version is greater than 26 and NDK version is greater than 14 (in the Android Studio Settings). - Import the `tensorflow/contrib/lite/java/demo` directory as a new Android Studio project. - Click through installing all the Gradle extensions it requests. - - Download the quantized Mobilenet TensorFlow Lite model from [here](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip) - - unzip and copy mobilenet_quant_v1_224.tflite to the assets directory: - `tensorflow/contrib/lite/java/demo/app/src/main/assets/` + - Either + - Download the quantized Mobilenet TensorFlow Lite model from [here](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip) + - unzip and copy mobilenet_quant_v1_224.tflite to the assets directory: + `tensorflow/contrib/lite/java/demo/app/src/main/assets/` + - Or download the floating point Inception-v3 model from [here](https://storage.googleapis.com/download.tensorflow.org/models/tflite/inception_v3_slim_2016_android_2017_11_10.zip) + - unzip and copy inceptionv3_non_slim_2015.tflite to the assets directory + - change the chosen classifier in [Camera2BasicFragment.java](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java) from + `classifier = new ImageClassifierQuantizedMobileNet(getActivity());` + to + `classifier = new ImageClassifierFloatInception(getActivity());` - Build and run the demo app ## Building TensorFlow Lite and the demo app from source @@ -84,7 +91,7 @@ Currently, we only support building the Android demo app within a Python 2 environment (due to a Bazel bug). ### More about the demo -The demo is resizing each camera image frame to (224 width * 224 height) to match the quantized Mobilenet model being used. The resized image is converted into a ByteBuffer row by row of size 1 * 224 * 224 * 3 bytes, where 1 is the number of images in a batch 224 * 224 is the width and height of the image 3 bytes represents three colors of a pixel. This demo uses the TensorFlow Lite Java inference API for models which take a single input and provide a single output. This outputs a two-dimensional array, with the first dimension being the category index and the second dimension being the confidence of classification. The Mobilenet model has 1001 unique categories and the app sorts the probabilities of all the categories and displays the top three. The Mobilenet quantized model is bundled within the assets directory of the app. +The demo is resizing each camera image frame to (224 width * 224 height) to match the quantized Mobilenet model being used (229 * 229 for Inception-v3). The resized image is converted into a ByteBuffer row by row of size 1 * 224 * 224 * 3 bytes, where 1 is the number of images in a batch. 224 * 224 (299 * 299) is the width and height of the image. 3 bytes represents three colors of a pixel. This demo uses the TensorFlow Lite Java inference API for models which take a single input and provide a single output. This outputs a two-dimensional array, with the first dimension being the category index and the second dimension being the confidence of classification. Both models have 1001 unique categories and the app sorts the probabilities of all the categories and displays the top three. The model file must be downloaded and bundled within the assets directory of the app. # iOS Demo App @@ -92,7 +99,7 @@ Similar to the Android demo app, there's an iOS camera app that uses exactly the This demo app requires a camera so it doesn't work with simulators. It need to be executed on a real iOS device. Follow the instructions to build and run the demo app: -1. Follow the Building section [here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/g3doc/ios.md#building) to build the universal iOS library for TensorFlow Lite. +1. Run `third_party/tensorflow/contrib/lite/examples/ios/download_models.sh` to download the model files used by the demo app. 1. Install [CocoaPods](https://cocoapods.org/) if it wasn't installed yet: `sudo gem install cocoapods`. 1. Run `pod install` in `tensorflow/contrib/lite/examples/ios/camera` to generate the workspace file. 1. Open the project by running `open tflite_camera_example.xcworkspace`, and build the app in XCode. @@ -142,7 +149,7 @@ Since we employ several formats, the following definitions may be useful: - SavedModel - A collection of GraphDef and CheckPoint together with a signature that labels input and output arguments to a model. A GraphDef and Checkpoint can be extracted from a saved model. - - TensorFlow lite model (.lite) - a serialized flatbuffer, containing TensorFlow lite operators and Tensors for the TensorFlow lite interpreter. This is most analogous to TensorFlow frozen GraphDefs. + - TensorFlow lite model (.tflite) - a serialized flatbuffer, containing TensorFlow lite operators and Tensors for the TensorFlow lite interpreter. This is most analogous to TensorFlow frozen GraphDefs. ### Freeze Graph To use this .pb GraphDef file within TensorFlow Lite, the application developer will need checkpoints containing trained weight parameters. The .pb contains only the structure of the graph. The process of merging the checkpoint values with the graph structure is known as "freezing" the graph. @@ -164,18 +171,18 @@ bazel-bin/tensorflow/python/tools/freeze_graph\ The user has to first build the freeze_graph script using bazel and then run the script. The input_binary flag has to be enabled to ensure that the protobuf is read and written in binary format. The user has to input the .pb and the .ckpt files to freeze the graph The output_node_names may not be obvious outside of the code that built the model. The easiest way to find them is to visualize the graph, either with graphviz, or [in tensorboard](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets-2/#3). -This frozen Graphdef is now ready to be converted to flatbuffer format (.lite) for use on Android or iOS. On Android users have the flexibility to use either the float or quantized versions of the frozen graphdef, if available, using the Tensorflow Optimizing Converter tool. +This frozen Graphdef is now ready to be converted to flatbuffer format (.tflite) for use on Android or iOS. On Android users have the flexibility to use either the float or quantized versions of the frozen graphdef, if available, using the Tensorflow Optimizing Converter tool. -Here is a sample command line to convert the frozen Graphdef to '.lite' format for The Tensorflow Optimizing Converter supports both float and quantized models, however, different configuration parameters are needed depending on whether a FLOAT or QUANTIZED mode is being used. +Here is a sample command line to convert the frozen Graphdef to '.tflite' format for The Tensorflow Optimizing Converter supports both float and quantized models, however, different configuration parameters are needed depending on whether a FLOAT or QUANTIZED mode is being used. (Here is a link to the pb [file](https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_1.0_224_frozen.tgz)). ``` bazel build tensorflow/contrib/lite/toco:toco -bazel-bin/tensorflow/contrib/lite/toco/toco -- \ +bazel-bin/tensorflow/contrib/lite/toco/toco \ --input_file=$(pwd)/mobilenet_v1_1.0_224/frozen_graph.pb \ --input_format=TENSORFLOW_GRAPHDEF --output_format=TFLITE \ - --output_file=/tmp/mobilenet_v1_1.0_224.lite --inference_type=FLOAT \ + --output_file=/tmp/mobilenet_v1_1.0_224.tflite --inference_type=FLOAT \ --input_type=FLOAT --input_arrays=input \ --output_arrays=MobilenetV1/Predictions/Reshape_1 --input_shapes=1,224,224,3 ``` @@ -188,7 +195,7 @@ bazel-bin/tensorflow/contrib/lite/toco/toco -- \ Note, it is also possible to use the Tensorflow Optimizing Converter through protos either from Python or from the command line see the documentation [here](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/toco/python/toco_from_protos.py). A developer can then integrate the conversion step into their model design workflow to ensure that a model will be easily convertible to a mobile inference graph. For example, -``` +```python import tensorflow as tf img = tf.placeholder(name="img", dtype=tf.float32, shape=(1, 64, 64, 3)) @@ -203,9 +210,15 @@ For detailed instructions on how to use the Tensorflow Optimizing Converter, ple You may refer to the [Ops compatibility guide](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md) for troubleshooting help. If that doesn't help, please file an [issue](https://github.com/tensorflow/tensorflow/issues). +If you would like to see a visual description of your TensorFlow Lite model after conversion, you can use tensorflow/contrib/lite/tools/visualize.py by running +```sh +bazel run tensorflow/contrib/lite/tools:visualize -- model.tflite model_viz.html +``` +and then visualize the resulting HTML file in a browser. + ## Step 3. Use the TensorFlow Lite model for inference in a mobile app -After completion of Step 2 the developer should have a .lite model. +After completion of Step 2 the developer should have a .tflite model. ### For Android Because Android apps need to be written in Java, and core TensorFlow is in C++, a JNI library is provided to interface between the two. Its interface is aimed only at inference, so it provides the ability to load a graph, set up inputs, and run the model to calculate particular outputs. The full documentation for the set of methods can be seen [here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/g3doc/). The demo app is also open sourced on [github](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/java/demo/app). diff --git a/tensorflow/contrib/lite/allocation.h b/tensorflow/contrib/lite/allocation.h index ee8a7ccd0b232f9e48095567fd4aefe94f595bc3..68aee2e64473320c461ec8b3f194904e7b8da43c 100644 --- a/tensorflow/contrib/lite/allocation.h +++ b/tensorflow/contrib/lite/allocation.h @@ -14,8 +14,8 @@ limitations under the License. ==============================================================================*/ // Main abstraction controlling the tflite interpreter. // See context.h for the API for defining operations (TfLiteRegistration). -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_ALLOCATION_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_ALLOCATION_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_ALLOCATION_H_ +#define TENSORFLOW_CONTRIB_LITE_ALLOCATION_H_ #include #include @@ -91,4 +91,4 @@ class MemoryAllocation : public Allocation { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_ALLOCATION_H_ +#endif // TENSORFLOW_CONTRIB_LITE_ALLOCATION_H_ diff --git a/tensorflow/contrib/lite/arena_planner.cc b/tensorflow/contrib/lite/arena_planner.cc new file mode 100644 index 0000000000000000000000000000000000000000..87b17c338e7afc33d32dd9688cc0825ac319dd19 --- /dev/null +++ b/tensorflow/contrib/lite/arena_planner.cc @@ -0,0 +1,251 @@ +/* Copyright 2017 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/lite/arena_planner.h" + +namespace tflite { + +namespace { + +// Memory allocation tuning +constexpr const int kDefaultArenaAlignment = 64; +constexpr const int kDefaultTensorAlignment = 4; + +} // namespace + +struct AllocationInfo { + // The node index requesting this allocation. + int node; + // The tensor index to be allocated or deallocated. + int tensor; + // Whether to allocate or deallocate + enum { ALLOC, DEALLOC } type; +}; + +ArenaPlanner::ArenaPlanner(TfLiteContext* context, + std::unique_ptr graph_info) + : context_(context), + graph_info_(std::move(graph_info)), + arena_(kDefaultArenaAlignment), + persistent_arena_(kDefaultArenaAlignment) {} + +ArenaPlanner::~ArenaPlanner() {} + +int64_t ArenaPlanner::BasePointer(TfLiteAllocationType type) { + if (type == kTfLiteArenaRwPersistent) { + return persistent_arena_.BasePointer(); + } + if (type == kTfLiteArenaRw) { + return arena_.BasePointer(); + } + return 0; +} + +TfLiteStatus ArenaPlanner::ResetAllocations() { + TF_LITE_ENSURE_STATUS(arena_.Clear()); + TF_LITE_ENSURE_STATUS(persistent_arena_.Clear()); + allocs_.clear(); + allocs_.resize(graph_info_->num_tensors()); + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::PlanAllocations() { + // Invalidate any existing data. + TF_LITE_ENSURE_STATUS(ResetAllocations()); + + // Keeps track of references to each tensor. + std::vector refcounts(graph_info_->num_tensors(), 0); + + // There will be an entry in alloc_queue_ for the allocation of each tensor + // and another for their deallocation. + alloc_queue_.reserve(2 * graph_info_->num_tensors()); + + // We must make sure the output tensors are never overwritten. We do that by + // artificially adding one to their ref-counts so they are never selected + // for deallocation. + for (int tensor_index : graph_info_->outputs()) { + refcounts[tensor_index]++; + } + + // Count references to node input tensors. + for (int i = 0; i < graph_info_->num_nodes(); ++i) { + const TfLiteNode& node = graph_info_->node(i); + TfLiteIntArray* node_inputs = node.inputs; + for (int j = 0; j < node_inputs->size; ++j) { + int tensor_index = node_inputs->data[j]; + if (tensor_index != kOptionalTensor) { + refcounts[tensor_index]++; + } + } + } + + // Queue all graph inputs for allocation. + for (int tensor_index : graph_info_->inputs()) { + if (tensor_index != kOptionalTensor) { + alloc_queue_.push_back({0, tensor_index, AllocationInfo::ALLOC}); + } + } + + // Go through the graph in execution order. + for (int i = 0; i < graph_info_->num_nodes(); ++i) { + const TfLiteNode& node = graph_info_->node(i); + + // First queue output tensors for allocation. + TfLiteIntArray* node_outputs = node.outputs; + for (int j = 0; j < node_outputs->size; ++j) { + int tensor_index = node_outputs->data[j]; + alloc_queue_.push_back({i, tensor_index, AllocationInfo::ALLOC}); + } + + // Then update the ref-counts of the node's inputs, and if necessary queue + // them for deallocation. + TfLiteIntArray* node_inputs = node.inputs; + for (int j = 0; j < node_inputs->size; ++j) { + int tensor_index = node_inputs->data[j]; + if (tensor_index != kOptionalTensor) { + refcounts[tensor_index]--; + if (refcounts[tensor_index] == 0) { + alloc_queue_.push_back({i, tensor_index, AllocationInfo::DEALLOC}); + } + } + } + } + + // Note that graph outputs will never be scheduled for deallocation. We + // could do that here for completeness, but it won't have any effect. + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::ExecuteAllocations(int first_node, int last_node) { + TF_LITE_ENSURE_STATUS(CalculateAllocations(first_node, last_node)); + TF_LITE_ENSURE_STATUS(Commit()); + + for (int i = 0; i < graph_info_->num_tensors(); ++i) { + // TODO(ahentz): we could do this only for the tensors that were modified + // in CalculateAllocations(), instead of redoing it for tensors that + // already had proper pointers. However we must be very careful, because + // SimpleMemoryArena::Commit() could move the base pointer. + TF_LITE_ENSURE_STATUS(ResolveTensorAllocation(i)); + } + + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::Commit() { + TF_LITE_ENSURE_STATUS(arena_.Commit(context_)); + TF_LITE_ENSURE_STATUS(persistent_arena_.Commit(context_)); + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::CalculateAllocations(int first_node, int last_node) { + int active_node = first_node; + // When dynamic tensors are present this method is called multiple times. + // The items in the alloc_queue_ referring to nodes before first_node were + // processed previously and should be skipped. Entries after last_node are + // not yet ready to be handled. + for (const auto& alloc_info : alloc_queue_) { + if (alloc_info.node < first_node) continue; + if (alloc_info.node > last_node) break; + if (alloc_info.node == active_node) { + // This is the first allocation/deallocation for a given node. It is + // time to deallocate the previous temporaries and allocate new ones. + if (active_node != first_node) { + TF_LITE_ENSURE_STATUS( + CalculateDeallocationOfInternalTensors(active_node - 1)); + } + TF_LITE_ENSURE_STATUS(CalculateAllocationOfInternalTensors(active_node)); + ++active_node; + } + // Handle the current item. + if (alloc_info.type == AllocationInfo::ALLOC) { + TF_LITE_ENSURE_STATUS(CalculateTensorAllocation(alloc_info.tensor)); + } else { + TF_LITE_ENSURE_STATUS(CalculateTensorDeallocation(alloc_info.tensor)); + } + } + + // Don't forget to deallocate temporaries of last node. + TF_LITE_ENSURE_STATUS( + CalculateDeallocationOfInternalTensors(active_node - 1)); + + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::ResolveTensorAllocation(int tensor_index) { + TfLiteTensor& tensor = *graph_info_->tensor(tensor_index); + if (tensor.allocation_type == kTfLiteArenaRw) { + // Skip resolution if the size of the tensor is zero, leaving it as a + // nullptr. + if (allocs_[tensor_index].size != 0) { + TF_LITE_ENSURE_STATUS(arena_.ResolveAlloc(context_, allocs_[tensor_index], + &tensor.data.raw)); + } + } + if (tensor.allocation_type == kTfLiteArenaRwPersistent) { + TF_LITE_ENSURE_STATUS(persistent_arena_.ResolveAlloc( + context_, allocs_[tensor_index], &tensor.data.raw)); + } + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::CalculateTensorAllocation(int tensor_index) { + TfLiteTensor& tensor = *graph_info_->tensor(tensor_index); + if (tensor.allocation_type == kTfLiteArenaRw) { + TF_LITE_ENSURE_STATUS(arena_.Allocate(context_, kDefaultTensorAlignment, + tensor.bytes, + &allocs_[tensor_index])); + } + if (tensor.allocation_type == kTfLiteArenaRwPersistent) { + TF_LITE_ENSURE_STATUS( + persistent_arena_.Allocate(context_, kDefaultTensorAlignment, + tensor.bytes, &allocs_[tensor_index])); + } + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::CalculateTensorDeallocation(int tensor_index) { + TfLiteTensor& tensor = *graph_info_->tensor(tensor_index); + if (tensor.allocation_type == kTfLiteArenaRw) { + TF_LITE_ENSURE_STATUS(arena_.Deallocate(context_, allocs_[tensor_index])); + } + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::CalculateAllocationOfInternalTensors( + int node_index) { + if (node_index < graph_info_->num_nodes()) { + const TfLiteNode& node = graph_info_->node(node_index); + TfLiteIntArray* node_temporaries = node.temporaries; + for (int i = 0; i < node_temporaries->size; ++i) { + int tensor_index = node_temporaries->data[i]; + TF_LITE_ENSURE_STATUS(CalculateTensorAllocation(tensor_index)); + } + } + return kTfLiteOk; +} + +TfLiteStatus ArenaPlanner::CalculateDeallocationOfInternalTensors( + int node_index) { + if (node_index < graph_info_->num_nodes()) { + const TfLiteNode& node = graph_info_->node(node_index); + TfLiteIntArray* node_temporaries = node.temporaries; + for (int i = 0; i < node_temporaries->size; ++i) { + int tensor_index = node_temporaries->data[i]; + TF_LITE_ENSURE_STATUS(CalculateTensorDeallocation(tensor_index)); + } + } + return kTfLiteOk; +} + +} // namespace tflite diff --git a/tensorflow/contrib/lite/arena_planner.h b/tensorflow/contrib/lite/arena_planner.h new file mode 100644 index 0000000000000000000000000000000000000000..58bc164619c2c053b9492e9a0e5de2da30e199af --- /dev/null +++ b/tensorflow/contrib/lite/arena_planner.h @@ -0,0 +1,107 @@ +/* Copyright 2017 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_LITE_ARENA_PLANNER_H_ +#define TENSORFLOW_CONTRIB_LITE_ARENA_PLANNER_H_ + +#include +#include + +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/graph_info.h" +#include "tensorflow/contrib/lite/memory_planner.h" +#include "tensorflow/contrib/lite/simple_memory_arena.h" + +namespace tflite { + +class AllocationInfo; + +// A memory planner that makes all the allocations using arenas. +// +// Before a model is executed by the interpreter, this class determines when +// each tensor needs to be allocated and deallocated, and preallocates all the +// necessary memory (the PlanAllocations phase). It then assigns portions of +// this memory buffer to each tensor (the ExecuteAllocations phase). Tensors may +// share some of the bufer if a tensor B is to be allocated after another tensor +// A has been deallocated. +// +// If dynamic tensors are used the planning steps can be repeated during model +// execution. Since dynamic tensors don't have sizes until after the +// corresponding operation is executed, this class supports incremental +// planning. +class ArenaPlanner : public MemoryPlanner { + public: + // Ownership of 'context' is not taken and it must remain util the + // ArenaPlanner is destroyed. + ArenaPlanner(TfLiteContext* context, std::unique_ptr graph_info); + ~ArenaPlanner() override; + ArenaPlanner(const ArenaPlanner&) = delete; + ArenaPlanner& operator=(const ArenaPlanner&) = delete; + + TfLiteStatus ResetAllocations() override; + TfLiteStatus PlanAllocations() override; + TfLiteStatus ExecuteAllocations(int first_node, int last_node) override; + + // Returns the base arena location for a given allocation type. + int64_t BasePointer(TfLiteAllocationType type); + + private: + // Make sure all the arenas have reserved enough memory to store all their + // tensors. + TfLiteStatus Commit(); + + // Traverse the allocation queue and reserve space in the appropriate arena + // for all tensors affected by ops in the interval [first_node, last_node]. + TfLiteStatus CalculateAllocations(int first_node, int last_node); + + // Assign absolute memory location to a tensor, based on its relative + // position inside the corresponding arena buffer. + TfLiteStatus ResolveTensorAllocation(int tensor_index); + + // Register an allocation for the given tensor. + TfLiteStatus CalculateTensorAllocation(int tensor_index); + + // Register a deallocation for the given tensor. + TfLiteStatus CalculateTensorDeallocation(int tensor_index); + + // Register an allocation for all internal (temporary) tensors of + // 'node_index'. + TfLiteStatus CalculateAllocationOfInternalTensors(int node_index); + + // Register a deallocation for all internal (temporary) tensors of + // 'node_index'. + TfLiteStatus CalculateDeallocationOfInternalTensors(int node_index); + + TfLiteContext* context_; + std::unique_ptr graph_info_; + + // Stores allocation data for all tensors. + std::vector allocs_; + + // A chronological list of instructions to allocated and deallocate tensors, + // reflecting the way they are used in the graph. + std::vector alloc_queue_; + + // Raw memory buffer that is allocated for all temporary and graph outputs. + // that are declared kTfLiteArenaRw. + SimpleMemoryArena arena_; + + // Raw memory buffer that is allocated for persistent tensors that are + // declared as kTfLiteArenaRwPersistent. + SimpleMemoryArena persistent_arena_; +}; + +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_ARENA_PLANNER_H_ diff --git a/tensorflow/contrib/lite/arena_planner_test.cc b/tensorflow/contrib/lite/arena_planner_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..a8a8755e2c9e81474f2ff9cd2b85c0eb3d5c3441 --- /dev/null +++ b/tensorflow/contrib/lite/arena_planner_test.cc @@ -0,0 +1,468 @@ +/* Copyright 2017 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/lite/arena_planner.h" + +#include + +#include +#include +#include "tensorflow/contrib/lite/testing/util.h" +#include "tensorflow/core/platform/logging.h" + +namespace tflite { +namespace { + +// A simple op to be used in tests, as syntactic sugar. +class TestOp { + public: + TestOp(std::initializer_list inputs, std::initializer_list outputs, + std::initializer_list temporaries) + : inputs_(inputs), outputs_(outputs), temporaries_(temporaries) {} + + const std::vector& inputs() const { return inputs_; } + const std::vector& outputs() const { return outputs_; } + const std::vector& temporaries() const { return temporaries_; } + + private: + std::vector inputs_; + std::vector outputs_; + std::vector temporaries_; +}; + +// A test graph where inputs are processed by the given nodes to produce +// outputs. +class TestGraph { + public: + TestGraph(std::initializer_list inputs, + std::initializer_list nodes, + std::initializer_list outputs) + : inputs_(inputs), outputs_(outputs) { + int max_tensor_index = 0; + + for (int t : inputs) { + max_tensor_index = std::max(max_tensor_index, t); + } + for (int t : outputs) { + max_tensor_index = std::max(max_tensor_index, t); + } + for (const auto& node : nodes) { + auto int_array = [](const std::vector& x) { + TfLiteIntArray* lite = TfLiteIntArrayCreate(x.size()); + for (size_t i = 0; i < x.size(); i++) lite->data[i] = x[i]; + return lite; + }; + + nodes_.push_back(TfLiteNode()); + nodes_.back().inputs = int_array(node.inputs()); + for (int t : node.inputs()) { + max_tensor_index = std::max(max_tensor_index, t); + } + nodes_.back().outputs = int_array(node.outputs()); + for (int t : node.outputs()) { + max_tensor_index = std::max(max_tensor_index, t); + } + nodes_.back().temporaries = int_array(node.temporaries()); + for (int t : node.temporaries()) { + max_tensor_index = std::max(max_tensor_index, t); + } + } + + for (int i = 0; i <= max_tensor_index; ++i) { + tensors_.push_back(TfLiteTensor()); + // Set some default values for allocation_type and bytes, which are the + // only fields used by the arena planner. + tensors_.back().allocation_type = kTfLiteArenaRw; + tensors_.back().bytes = (i + 1) * 3; + } + } + + ~TestGraph() { + for (auto node : nodes_) { + TfLiteIntArrayFree(node.inputs); + TfLiteIntArrayFree(node.outputs); + TfLiteIntArrayFree(node.temporaries); + } + } + + const std::vector& nodes() { return nodes_; } + std::vector* tensors() { return &tensors_; } + const std::vector& inputs() { return inputs_; } + const std::vector& outputs() { return outputs_; } + + private: + std::vector nodes_; + std::vector tensors_; + std::vector inputs_; + std::vector outputs_; +}; + +// The GraphInfo for a TestGraph. +class TestGraphInfo : public GraphInfo { + public: + explicit TestGraphInfo(TestGraph* graph) : graph_(graph) {} + + size_t num_tensors() const override { return graph_->tensors()->size(); } + TfLiteTensor* tensor(size_t index) override { + return &graph_->tensors()->at(index); + } + size_t num_nodes() const override { return graph_->nodes().size(); } + const TfLiteNode& node(size_t index) const override { + return graph_->nodes()[index]; + } + const std::vector& inputs() const override { return graph_->inputs(); } + const std::vector& outputs() const override { return graph_->outputs(); } + + private: + TestGraph* graph_; +}; + +void ReportError(TfLiteContext* context, const char* format, ...) { + const size_t kBufferSize = 1024; + char temp_buffer[kBufferSize]; + + va_list args; + va_start(args, format); + vsnprintf(temp_buffer, kBufferSize, format, args); + va_end(args); + + LOG(INFO) << temp_buffer; +} + +class ArenaPlannerTest : public ::testing::Test { + protected: + void SetGraph(TestGraph* graph) { + graph_ = graph; + context_.ReportError = ReportError; + planner_.reset(new ArenaPlanner( + &context_, std::unique_ptr(new TestGraphInfo(graph)))); + CHECK(planner_->ResetAllocations() == kTfLiteOk); + CHECK(planner_->PlanAllocations() == kTfLiteOk); + } + + void Execute(int start, int end) { + CHECK(planner_->ExecuteAllocations(start, end) == kTfLiteOk); + } + + // Returns the actual offset of a given tensor, relative to the start of its + // arena. + int64_t GetOffset(int tensor_index) { + const TfLiteTensor& tensor = (*graph_->tensors())[tensor_index]; + return reinterpret_cast(tensor.data.raw) - + planner_->BasePointer(tensor.allocation_type); + } + + // Returns the first aligned offset after a given tensor. + int64_t GetOffsetAfter(int tensor_index) { + const TfLiteTensor& tensor = (*graph_->tensors())[tensor_index]; + int64_t offset = GetOffset(tensor_index) + tensor.bytes; + // We must make sure the offset is aligned to kDefaultArenaAlignment. + if (offset % 4 != 0) { + offset += 4 - offset % 4; + } + return offset; + }; + + TfLiteContext context_; + TestGraph* graph_; + std::unique_ptr planner_; +}; + +TEST_F(ArenaPlannerTest, EmptyGraph) { + TestGraph graph({}, {}, {}); + SetGraph(&graph); + Execute(0, 10); +} + +TEST_F(ArenaPlannerTest, GraphWithNoOps) { + TestGraph graph({0, 10}, {}, {5, 11}); + SetGraph(&graph); + Execute(0, 10); + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(10), GetOffsetAfter(0)); + // The outputs are never allocated because they are not connected to any + // inputs. + EXPECT_TRUE((*graph.tensors())[5].data.raw == nullptr); + EXPECT_TRUE((*graph.tensors())[11].data.raw == nullptr); +} + +TEST_F(ArenaPlannerTest, GraphWithOneOp) { + TestGraph graph({1}, {{{1}, {2}, {}}}, {2}); + SetGraph(&graph); + Execute(0, 10); + EXPECT_EQ(GetOffset(1), 0); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); +} + +TEST_F(ArenaPlannerTest, ZeroSizedTensors) { + TestGraph graph({1}, {{{1}, {2}, {}}}, {2}); + (*graph.tensors())[1].bytes = 0; + SetGraph(&graph); + // TODO(ahentz): this is currently broken because the arena finds two + // allocations with the same offset and returns an error. + ASSERT_FALSE(planner_->ExecuteAllocations(0, 10) == kTfLiteOk); + // EXPECT_EQ(GetOffset(1), 0); + // EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); +} + +TEST_F(ArenaPlannerTest, SimpleGraph) { + TestGraph graph({0, 1}, + { + /* in, out, tmp */ + {{0, 1}, {2}, {}}, // First op + {{2, 0}, {4, 5}, {}}, // Second op + {{4, 5}, {3}, {}} // Third op + }, + {3}); + SetGraph(&graph); + Execute(0, 10); + + // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +4 +5 -2 -0 +3 -4 -5 + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); + EXPECT_EQ(GetOffset(3), 0); +} + +TEST_F(ArenaPlannerTest, SimpleGraphWithTemporary) { + TestGraph graph({0, 1}, + { + /* in, out, tmp */ + {{0, 1}, {2}, {}}, // First op + {{2, 0}, {4}, {5}}, // Second op, with temporary + {{4}, {3}, {}} // Third op + }, + {3}); + SetGraph(&graph); + Execute(0, 10); + + // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +5 +4 -2 -0 -5 +3 -4 + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(5)); + EXPECT_EQ(GetOffset(3), 0); +} + +TEST_F(ArenaPlannerTest, SimpleGraphWithOptionals) { + TestGraph graph({0, -1, 1}, + { + /* in, out, tmp */ + {{0, 1}, {2}, {}}, // First op + {{2, 0}, {4, 5}, {}}, // Second op + {{4, -1, 5}, {3}, {}} // Third op, with optional + }, + {3}); + SetGraph(&graph); + Execute(0, 10); + + // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +4 +5 -2 -0 +3 -4 -5 + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); + EXPECT_EQ(GetOffset(3), 0); +} + +TEST_F(ArenaPlannerTest, SimpleGraphWithLargeTensor) { + TestGraph graph({0, -1, 1}, + { + /* in, out, tmp */ + {{0, 1}, {2}, {}}, // First op + {{2, 0}, {4}, {5}}, // Second op, with temporary + {{4, -1}, {3}, {}} // Third op, with optional + }, + {3}); + + // Make #1 very large so its vacancy can be filled with #5 and #4. + (*graph.tensors())[1].bytes = 40; + + SetGraph(&graph); + Execute(0, 10); + + // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +5 +4 -2 -0 -5 +3 -4 + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(5)); + EXPECT_EQ(GetOffset(3), 0); +} + +TEST_F(ArenaPlannerTest, SimpleGraphWithPersistentTensor) { + TestGraph graph({0, -1, 1}, + { + /* in, out, tmp */ + {{0, 1}, {2}, {}}, // First op + {{2, 0}, {4}, {5}}, // Second op, with temporary + {{4, -1}, {3}, {}} // Third op, with optional + }, + {3}); + + // Make #1 persistent so it goes into its own arena. + (*graph.tensors())[1].allocation_type = kTfLiteArenaRwPersistent; + + SetGraph(&graph); + Execute(0, 10); + + // Make sure #0 and #1 were given different memory locations (because they + // will both have offset=0, in different arenas.) + EXPECT_NE((*graph.tensors())[0].data.raw, (*graph.tensors())[1].data.raw); + + // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +5 +4 -2 -0 -5 +3 -4 + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), 0); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(5)); + EXPECT_EQ(GetOffset(3), 0); +} + +TEST_F(ArenaPlannerTest, SimpleGraphWithDynamicTensor) { + TestGraph graph({0, -1, 1}, + { + /* in, out, tmp */ + {{0, 1}, {2}, {}}, // First op + {{2, 0}, {4}, {5}}, // Second op, with temporary + {{4, -1}, {3}, {}} // Third op, with optional + }, + {3}); + + // Make #1 dynaic so it does not get allocated. + (*graph.tensors())[1].allocation_type = kTfLiteDynamic; + + SetGraph(&graph); + Execute(0, 10); + + EXPECT_EQ((*graph.tensors())[1].data.raw, nullptr); + + // Alloc(+) and dealloc(-) order: +0 +1 +2 -1 +5 +4 -2 -0 -5 +3 -4 + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(5)); + EXPECT_EQ(GetOffset(3), 0); +} + +TEST_F(ArenaPlannerTest, LargerGraphAndStepwiseAllocation) { + TestGraph graph({0, 1}, + { + /* in, out, tmp */ + {{0, 1}, {2, 3}, {}}, + {{2, 0}, {4, 5}, {6}}, + {{1, -1}, {7}, {}}, + {{7, 3}, {8}, {9}}, + {{4, 5, 8}, {10}, {}}, + }, + {10}); + SetGraph(&graph); + + auto is_unallocated = [&](int tensor_index) { + return (*graph.tensors())[tensor_index].data.raw == nullptr; + }; + + // The allocation plan is made at the beginning and is independent of + // the execution steps. Here's the allocation order: + // Op0: +0 +1 +2 +3 + // Op1: +6 +4 +5 -6 -0 -2 + // Op2: +7 -1 + // Op3: +9 +8 -9 -3 -7 + // Op4: +10 -4 -5 -8 + + Execute(0, 0); + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); + EXPECT_TRUE(is_unallocated(6)); + EXPECT_TRUE(is_unallocated(4)); + EXPECT_TRUE(is_unallocated(5)); + EXPECT_TRUE(is_unallocated(7)); + EXPECT_TRUE(is_unallocated(9)); + EXPECT_TRUE(is_unallocated(8)); + EXPECT_TRUE(is_unallocated(10)); + + Execute(1, 1); + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(6), GetOffsetAfter(3)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(6)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); + EXPECT_TRUE(is_unallocated(7)); + EXPECT_TRUE(is_unallocated(9)); + EXPECT_TRUE(is_unallocated(8)); + EXPECT_TRUE(is_unallocated(10)); + + Execute(2, 2); + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(6), GetOffsetAfter(3)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(6)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); + // Here's an interesting allocation. Even though #6 requires only 21 bytes, + // its deallocation freed up 24 bytes due to the alignment requirements in + // the arena. That means we can fit #7 in the same space! + EXPECT_EQ(GetOffset(7), GetOffsetAfter(3)); + EXPECT_TRUE(is_unallocated(9)); + EXPECT_TRUE(is_unallocated(8)); + EXPECT_TRUE(is_unallocated(10)); + + Execute(3, 3); + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(6), GetOffsetAfter(3)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(6)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); + EXPECT_EQ(GetOffset(7), GetOffsetAfter(3)); + // The deallocation of #0, #1 and #2 freed up 24 bytes but that's not enough + // for #9, so it goes at the end. + EXPECT_EQ(GetOffset(9), GetOffsetAfter(5)); + EXPECT_EQ(GetOffset(8), GetOffsetAfter(9)); + EXPECT_TRUE(is_unallocated(10)); + + Execute(4, 4); + EXPECT_EQ(GetOffset(0), 0); + EXPECT_EQ(GetOffset(1), GetOffsetAfter(0)); + EXPECT_EQ(GetOffset(2), GetOffsetAfter(1)); + EXPECT_EQ(GetOffset(3), GetOffsetAfter(2)); + EXPECT_EQ(GetOffset(6), GetOffsetAfter(3)); + EXPECT_EQ(GetOffset(4), GetOffsetAfter(6)); + EXPECT_EQ(GetOffset(5), GetOffsetAfter(4)); + EXPECT_EQ(GetOffset(7), GetOffsetAfter(3)); + EXPECT_EQ(GetOffset(9), GetOffsetAfter(5)); + EXPECT_EQ(GetOffset(8), GetOffsetAfter(9)); + // There's just enough space at the beginning for #10 due to the + // deallocation of #0, #1, #2 and #3 (total 36 bytes, #10 needs + // only 33.) + EXPECT_EQ(GetOffset(10), 0); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/build_def.bzl b/tensorflow/contrib/lite/build_def.bzl index d1fcdce70a34393defce0f2d0f6d5bb53f21c45e..19829e4991651111e13fc1805f97daef8bc016a7 100644 --- a/tensorflow/contrib/lite/build_def.bzl +++ b/tensorflow/contrib/lite/build_def.bzl @@ -5,25 +5,25 @@ def tflite_copts(): copts = [ "-DFARMHASH_NO_CXX_STRING", ] + select({ - "//tensorflow:android_arm64": [ + str(Label("//tensorflow:android_arm64")): [ "-std=c++11", "-O3", ], - "//tensorflow:android_arm": [ + str(Label("//tensorflow:android_arm")): [ "-mfpu=neon", "-mfloat-abi=softfp", "-std=c++11", "-O3", ], - "//tensorflow:android_x86": [ + str(Label("//tensorflow:android_x86")): [ "-DGEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK", ], - "//tensorflow:ios_x86_64": [ + str(Label("//tensorflow:ios_x86_64")): [ "-msse4.1", ], "//conditions:default": [], }) + select({ - "//tensorflow:with_default_optimizations": [], + str(Label("//tensorflow:with_default_optimizations")): [], "//conditions:default": ["-DGEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK"], }) @@ -94,7 +94,6 @@ def tflite_jni_linkopts(): "//conditions:default": [], }) - def tflite_jni_binary(name, copts=tflite_copts(), linkopts=tflite_jni_linkopts(), diff --git a/tensorflow/contrib/lite/build_ios_universal_lib.sh b/tensorflow/contrib/lite/build_ios_universal_lib.sh index cbc96e6edd4358f6666731caa4c208c77d9c6c54..4a9023ff33de15dd384531d51e39de4ffeecdb8b 100755 --- a/tensorflow/contrib/lite/build_ios_universal_lib.sh +++ b/tensorflow/contrib/lite/build_ios_universal_lib.sh @@ -15,6 +15,10 @@ # ============================================================================== set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR/../../.." + make -f tensorflow/contrib/lite/Makefile TARGET=IOS IOS_ARCH=x86_64 -j 8 make -f tensorflow/contrib/lite/Makefile TARGET=IOS IOS_ARCH=i386 -j 8 make -f tensorflow/contrib/lite/Makefile TARGET=IOS IOS_ARCH=armv7 -j 8 diff --git a/tensorflow/contrib/lite/builtin_op_data.h b/tensorflow/contrib/lite/builtin_op_data.h index 5c6f3016b1c7d06ba35229faeff9cec32e168ef2..5fc8954743e5b3b458e5c2004f4378cbad6056c0 100644 --- a/tensorflow/contrib/lite/builtin_op_data.h +++ b/tensorflow/contrib/lite/builtin_op_data.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ +#define TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ #include @@ -83,7 +83,14 @@ typedef struct { TfLiteFusedActivation activation; } TfLiteRNNParams; -typedef struct { TfLiteFusedActivation activation; } TfLiteFullyConnectedParams; +typedef struct { + bool time_major; + TfLiteFusedActivation activation; +} TfLiteSequenceRNNParams; + +typedef struct { + TfLiteFusedActivation activation; +} TfLiteFullyConnectedParams; typedef enum { kTfLiteLshProjectionUnknown = 0, @@ -91,9 +98,13 @@ typedef enum { kTfLiteLshProjectionDense = 2, } TfLiteLSHProjectionType; -typedef struct { TfLiteLSHProjectionType type; } TfLiteLSHProjectionParams; +typedef struct { + TfLiteLSHProjectionType type; +} TfLiteLSHProjectionParams; -typedef struct { float beta; } TfLiteSoftmaxParams; +typedef struct { + float beta; +} TfLiteSoftmaxParams; typedef struct { int axis; @@ -105,20 +116,23 @@ typedef struct { } TfLiteAddParams; typedef struct { - // Number of spatial dimensions. - // For now only NHWC is supported, and the value should always be 2. - int num_spatial_dimensions; - // TODO(ahentz): We can't have dynamic data in this struct, at least not yet. - // For now we will fix the maximum possible number of dimensions. - int block_shape[2]; - int before_crops[2]; - int after_crops[2]; +} TfLiteSpaceToBatchNDParams; + +typedef struct { } TfLiteBatchToSpaceNDParams; typedef struct { TfLiteFusedActivation activation; } TfLiteMulParams; +typedef struct { + TfLiteFusedActivation activation; +} TfLiteSubParams; + +typedef struct { + TfLiteFusedActivation activation; +} TfLiteDivParams; + typedef struct { TfLiteFusedActivation activation; } TfLiteL2NormParams; @@ -137,16 +151,10 @@ typedef struct { } TfLiteLSTMParams; typedef struct { - int new_height; - int new_width; + bool align_corners; } TfLiteResizeBilinearParams; typedef struct { - // TODO(ahentz): We can't have dynamic data in this struct, at least not yet. - // For now we will fix the maximum possible number of dimensions. - int before_padding[8]; - int after_padding[8]; - int num_dimensions; } TfLitePadParams; typedef struct { @@ -180,8 +188,34 @@ typedef struct { int axis; } TfLiteGatherParams; +typedef struct { +} TfLiteTransposeParams; + +typedef struct { + bool keep_dims; +} TfLiteMeanParams; + +typedef struct { + int num_splits; +} TfLiteSplitParams; + +typedef struct { + // TODO(ahentz): We can't have dynamic data in this struct, at least not yet. + // For now we will fix the maximum possible number of dimensions. + int squeeze_dims[8]; + int num_squeeze_dims; +} TfLiteSqueezeParams; + +typedef struct { + int begin_mask; + int end_mask; + int ellipsis_mask; + int new_axis_mask; + int shrink_axis_mask; +} TfLiteStridedSliceParams; + #ifdef __cplusplus } // extern "C" #endif // __cplusplus -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ +#endif // TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ diff --git a/tensorflow/contrib/lite/builtin_ops.h b/tensorflow/contrib/lite/builtin_ops.h new file mode 100644 index 0000000000000000000000000000000000000000..4ebd1586de791eecf0304637bde76232d9f0a11d --- /dev/null +++ b/tensorflow/contrib/lite/builtin_ops.h @@ -0,0 +1,80 @@ +/* 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_LITE_BUILTIN_OPS_H_ +#define TENSORFLOW_CONTRIB_LITE_BUILTIN_OPS_H_ + +// DO NOT EDIT MANUALLY: This file is automatically generated by +// `schema_builtin_ops_header_generator.py`. + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +typedef enum { + kTfLiteBuiltinAdd = 0, + kTfLiteBuiltinAveragePool2d = 1, + kTfLiteBuiltinConcatenation = 2, + kTfLiteBuiltinConv2d = 3, + kTfLiteBuiltinDepthwiseConv2d = 4, + kTfLiteBuiltinEmbeddingLookup = 7, + kTfLiteBuiltinFullyConnected = 9, + kTfLiteBuiltinHashtableLookup = 10, + kTfLiteBuiltinL2Normalization = 11, + kTfLiteBuiltinL2Pool2d = 12, + kTfLiteBuiltinLocalResponseNormalization = 13, + kTfLiteBuiltinLogistic = 14, + kTfLiteBuiltinLshProjection = 15, + kTfLiteBuiltinLstm = 16, + kTfLiteBuiltinMaxPool2d = 17, + kTfLiteBuiltinMul = 18, + kTfLiteBuiltinRelu = 19, + kTfLiteBuiltinReluN1To1 = 20, + kTfLiteBuiltinRelu6 = 21, + kTfLiteBuiltinReshape = 22, + kTfLiteBuiltinResizeBilinear = 23, + kTfLiteBuiltinRnn = 24, + kTfLiteBuiltinSoftmax = 25, + kTfLiteBuiltinSpaceToDepth = 26, + kTfLiteBuiltinSvdf = 27, + kTfLiteBuiltinTanh = 28, + kTfLiteBuiltinConcatEmbeddings = 29, + kTfLiteBuiltinSkipGram = 30, + kTfLiteBuiltinCall = 31, + kTfLiteBuiltinCustom = 32, + kTfLiteBuiltinEmbeddingLookupSparse = 33, + kTfLiteBuiltinPad = 34, + kTfLiteBuiltinUnidirectionalSequenceRnn = 35, + kTfLiteBuiltinGather = 36, + kTfLiteBuiltinBatchToSpaceNd = 37, + kTfLiteBuiltinSpaceToBatchNd = 38, + kTfLiteBuiltinTranspose = 39, + kTfLiteBuiltinMean = 40, + kTfLiteBuiltinSub = 41, + kTfLiteBuiltinDiv = 42, + kTfLiteBuiltinSqueeze = 43, + kTfLiteBuiltinUnidirectionalSequenceLstm = 44, + kTfLiteBuiltinStridedSlice = 45, + kTfLiteBuiltinBidirectionalSequenceRnn = 46, + kTfLiteBuiltinExp = 47, + kTfLiteBuiltinTopkV2 = 48, + kTfLiteBuiltinSplit = 49, +} TfLiteBuiltinOperator; + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus +#endif // TENSORFLOW_CONTRIB_LITE_BUILTIN_OPS_H_ +} diff --git a/tensorflow/contrib/lite/context.h b/tensorflow/contrib/lite/context.h index 41257a53b145cbe7e252c9d4de6ea7ef654431b5..b0c4d3431f9a67bc87d51ada91ed73f1661023a2 100644 --- a/tensorflow/contrib/lite/context.h +++ b/tensorflow/contrib/lite/context.h @@ -26,8 +26,8 @@ limitations under the License. // TfLiteRegistration - the implementation of a conceptual operation. // // Some abstractions in this file are created and managed by Interpreter. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_CONTEXT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_CONTEXT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_CONTEXT_H_ +#define TENSORFLOW_CONTRIB_LITE_CONTEXT_H_ #include #include @@ -38,6 +38,9 @@ extern "C" { typedef enum { kTfLiteOk = 0, kTfLiteError = 1 } TfLiteStatus; +// Forward declare so GetNode can use this is in Context. +typedef struct _TfLiteRegistration TfLiteRegistration; + #define kOptionalTensor (-1) // Fixed size list of integers. Used for dimensions and inputs/outputs tensor @@ -141,6 +144,7 @@ typedef struct { // A union of points that points to memory for a given tensor. typedef union { int* i32; + int64_t* i64; float* f; char* raw; const char* raw_const; @@ -204,9 +208,56 @@ void TfLiteTensorReset(TfLiteType type, const char* name, TfLiteIntArray* dims, // Resize the allocated data of a (dynamic) tensor. void TfLiteTensorRealloc(size_t num_bytes, TfLiteTensor* tensor); +// A structure representing an instance of a node. +// This structure only exhibits the inputs, outputs and user defined data, not +// other features like the type. +typedef struct { + // Inputs to this node expressed as indices into the simulator's tensors. + TfLiteIntArray* inputs; + + // Outputs to this node expressed as indices into the simulator's tensors. + TfLiteIntArray* outputs; + + // Temporary tensors uses during the computations. This usually contains no + // tensors, but ops are allowed to change that if they need scratch space of + // any sort. + TfLiteIntArray* temporaries; + + // Opaque data provided by the node implementer through `Registration.init`. + void* user_data; + + // Opaque data provided to the node if the node is a builtin. This is usually + // a structure defined in builtin_op_data.h + void* builtin_data; + + // Custom initial data. This is the opaque data provided in the flatbuffer. + // WARNING: This is an experimental interface that is subject to change. + const void* custom_initial_data; + int custom_initial_data_size; +} TfLiteNode; + typedef struct TfLiteContext { // Number of tensors in the context. int tensors_size; + + // The execution plan contains a list of the node indices in execution + // order. execution_plan->size is the current number of nodes. And, + // execution_plan->data[0] is the first node that needs to be run. + // TfLiteDelegates can traverse the current execution plan by iterating + // through each member of this array and using GetNodeAndRegistration() to + // access details about a node. i.e. + // TfLiteIntArray* execution_plan; + // TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &execution_plan)); + // for (int exec_index = 0; exec_index < execution_plan->size; exec_index++) { + // int node_index = execution_plan->data[exec_index]; + // TfLiteNode* node; + // TfLiteRegistration* reg; + // context->GetNodeAndRegistration(context, node_index, &node, ®); + // } + // WARNING: This is an experimental interface that is subject to change. + TfLiteStatus (*GetExecutionPlan)(struct TfLiteContext* context, + TfLiteIntArray** execution_plan); + // An tensor of tensors in the interpreter context (of length `tensors_size`) TfLiteTensor* tensors; @@ -226,34 +277,23 @@ typedef struct TfLiteContext { TfLiteStatus (*AddTensors)(struct TfLiteContext*, int tensors_to_add, int* first_new_tensor_index); + // Get a Tensor node by node_index. + // WARNING: This is an experimental interface that is subject to change. + TfLiteStatus (*GetNodeAndRegistration)(struct TfLiteContext*, int node_index, + TfLiteNode** node, + TfLiteRegistration** registration); + + // Replace ops with delegate. + TfLiteStatus (*ReplaceSubgraphsWithDelegateKernels)( + struct TfLiteContext*, TfLiteRegistration registration, + const TfLiteIntArray* nodes_to_replace); + // TODO(ahentz): we should create a more general mechanism for this sort of // library-global objects. void* gemm_context; } TfLiteContext; -// A structure representing an instance of a node. -// This structure only exhibits the inputs, outputs and user defined data, not -// other features like the type. -typedef struct { - // Inputs to this node expressed as indices into the simulator's tensors. - TfLiteIntArray* inputs; - - // Outputs to this node expressed as indices into the simulator's tensors. - TfLiteIntArray* outputs; - - // Temporary tensors uses during the computations. This usually contains no - // tensors, but ops are allowed to change that if they need scratch space of - // any sort. - TfLiteIntArray* temporaries; - - // Opaque data provided by the node implementer through `Registration.init`. - void* user_data; - - // Opaque data provided to the node if the node is a builtin. - void* builtin_data; -} TfLiteNode; - -typedef struct { +typedef struct _TfLiteRegistration { // Initializes the op from serialized data. // If a built-in op: // `buffer` is the op's params data (TfLiteLSTMParams*). @@ -290,9 +330,27 @@ typedef struct { // NN API. Note, it is the responsibility of the registration binder to // set this properly. int32_t builtin_code; + + // Custom op name. If the op is a builtin, this will be null. + // WARNING: This is an experimental interface that is subject to change. + const char* custom_name; } TfLiteRegistration; +// WARNING: This is an experimental interface that is subject to change. +typedef struct { + // Data that delegate needs to identify itself. This data is owned by the + // delegate. The delegate is owned in the user code, so the delegate is + // responsible for doing this when it is destroyed. + void* data_; + // Invoked by ModifyGraphWithDelegate. This prepare is called, giving the + // delegate a view of the current graph through TfLiteContext*. It typically + // will look at the nodes and call ReplaceSubgraphsWithDelegateKernels() + // to ask the TensorFlow lite runtime to create macro-nodes to represent + // delegated subgraphs of the original graph. + TfLiteStatus (*Prepare)(TfLiteContext* context, void* data); +} TfLiteDelegate; + #ifdef __cplusplus } // extern "C" #endif // __cplusplus -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_CONTEXT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_CONTEXT_H_ diff --git a/tensorflow/contrib/lite/download_dependencies.sh b/tensorflow/contrib/lite/download_dependencies.sh index 7fce1ba3461066e6dada95246781440258d844c1..a93ed201d647ddf2359a57254a959871c13fb94f 100755 --- a/tensorflow/contrib/lite/download_dependencies.sh +++ b/tensorflow/contrib/lite/download_dependencies.sh @@ -16,6 +16,9 @@ set -e +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR/../../.." + DOWNLOADS_DIR=tensorflow/contrib/lite/downloads BZL_FILE_PATH=tensorflow/workspace.bzl @@ -26,15 +29,13 @@ if [ ! -f $BZL_FILE_PATH ]; then exit 1; fi -EIGEN_URL="$(grep -o 'http.*bitbucket.org/eigen/eigen/get/.*tar\.gz' "${BZL_FILE_PATH}" | grep -v bazel-mirror | head -n1)" +EIGEN_URL="$(grep -o 'http.*bitbucket.org/eigen/eigen/get/.*tar\.gz' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" GEMMLOWP_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)" GOOGLETEST_URL="https://github.com/google/googletest/archive/release-1.8.0.tar.gz" ABSL_URL="$(grep -o 'https://github.com/abseil/abseil-cpp/.*tar.gz' "${BZL_FILE_PATH}" | head -n1)" NEON_2_SSE_URL="https://github.com/intel/ARM_NEON_2_x86_SSE/archive/master.zip" FARMHASH_URL="https://mirror.bazel.build/github.com/google/farmhash/archive/816a4ae622e964763ca0862d9dbd19324a1eaf45.tar.gz" FLATBUFFERS_URL="https://github.com/google/flatbuffers/archive/master.zip" -MODELS_URL="https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_224_ios_lite_float_2017_11_08.zip" -QUANTIZED_MODELS_URL="https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip" # TODO(petewarden): Some new code in Eigen triggers a clang bug with iOS arm64, # so work around it by patching the source. @@ -90,8 +91,6 @@ download_and_extract "${ABSL_URL}" "${DOWNLOADS_DIR}/absl" download_and_extract "${NEON_2_SSE_URL}" "${DOWNLOADS_DIR}/neon_2_sse" download_and_extract "${FARMHASH_URL}" "${DOWNLOADS_DIR}/farmhash" download_and_extract "${FLATBUFFERS_URL}" "${DOWNLOADS_DIR}/flatbuffers" -download_and_extract "${MODELS_URL}" "${DOWNLOADS_DIR}/models" -download_and_extract "${QUANTIZED_MODELS_URL}" "${DOWNLOADS_DIR}/quantized_models" replace_by_sed 's#static uint32x4_t p4ui_CONJ_XOR = vld1q_u32( conj_XOR_DATA );#static uint32x4_t p4ui_CONJ_XOR; // = vld1q_u32( conj_XOR_DATA ); - Removed by script#' \ "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h" @@ -100,7 +99,4 @@ replace_by_sed 's#static uint32x2_t p2ui_CONJ_XOR = vld1_u32( conj_XOR_DATA );#s replace_by_sed 's#static uint64x2_t p2ul_CONJ_XOR = vld1q_u64( p2ul_conj_XOR_DATA );#static uint64x2_t p2ul_CONJ_XOR;// = vld1q_u64( p2ul_conj_XOR_DATA ); - Removed by script#' \ "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h" -cp ${DOWNLOADS_DIR}/models/models/* tensorflow/contrib/lite/examples/ios/simple/data/ -cp ${DOWNLOADS_DIR}/quantized_models/* tensorflow/contrib/lite/examples/ios/camera/data/ - echo "download_dependencies.sh completed successfully." >&2 diff --git a/tensorflow/contrib/lite/error_reporter.h b/tensorflow/contrib/lite/error_reporter.h index d5715e4f90aead79a617fe4576bfe5100d5e121a..da193d2586e9123341b9a41be049ee2a4382017a 100644 --- a/tensorflow/contrib/lite/error_reporter.h +++ b/tensorflow/contrib/lite/error_reporter.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ +#define TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ #include #include "tensorflow/contrib/lite/context.h" @@ -51,4 +51,4 @@ ErrorReporter* DefaultErrorReporter(); } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ +#endif // TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ diff --git a/tensorflow/contrib/lite/examples/ios/camera/CameraExampleViewController.mm b/tensorflow/contrib/lite/examples/ios/camera/CameraExampleViewController.mm index 10f31bb6f17242c9f7f70f0648ec643f99c5ac86..d74e275f0439b1ce56b29e0eadff5f211f6a4faa 100644 --- a/tensorflow/contrib/lite/examples/ios/camera/CameraExampleViewController.mm +++ b/tensorflow/contrib/lite/examples/ios/camera/CameraExampleViewController.mm @@ -225,14 +225,8 @@ static void GetTopN(const uint8_t* prediction, const int prediction_size, const assert(pixelBuffer != NULL); OSType sourcePixelFormat = CVPixelBufferGetPixelFormatType(pixelBuffer); - int doReverseChannels; - if (kCVPixelFormatType_32ARGB == sourcePixelFormat) { - doReverseChannels = 1; - } else if (kCVPixelFormatType_32BGRA == sourcePixelFormat) { - doReverseChannels = 0; - } else { - assert(false); // Unknown source format - } + assert(sourcePixelFormat == kCVPixelFormatType_32ARGB || + sourcePixelFormat == kCVPixelFormatType_32BGRA); const int sourceRowBytes = (int)CVPixelBufferGetBytesPerRow(pixelBuffer); const int image_width = (int)CVPixelBufferGetWidth(pixelBuffer); diff --git a/tensorflow/contrib/lite/examples/ios/camera/Podfile b/tensorflow/contrib/lite/examples/ios/camera/Podfile index 4ae6fb6b94e4489f63506b05a2f348b7daafd3b7..c7d3b1c966eaa0de71f5c37a6a77b3881e30ddd7 100644 --- a/tensorflow/contrib/lite/examples/ios/camera/Podfile +++ b/tensorflow/contrib/lite/examples/ios/camera/Podfile @@ -2,4 +2,4 @@ platform :ios, '8.0' inhibit_all_warnings! target 'tflite_camera_example' - pod 'TensorFlow-experimental' + pod 'TensorFlowLite' diff --git a/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj b/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj index c98183276bd60d2a0ad023ba26aad12572a02786..b0236e9c608ec35437bcfe79c51149a76f9f416e 100644 --- a/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj +++ b/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj @@ -16,7 +16,6 @@ 1CDB2D4E1ED3AA35007929E9 /* Info.plist in Resources */ = {isa = PBXBuildFile; fileRef = 1CDB2D4D1ED3AA35007929E9 /* Info.plist */; }; 54DC6C3C5F734F3A58069F0C /* libPods-tflite_camera_example.a in Frameworks */ = {isa = PBXBuildFile; fileRef = 3BA8BF92C84895BFE59D8236 /* libPods-tflite_camera_example.a */; }; AC1F82661FBA3CBD0052BA77 /* labels.txt in Resources */ = {isa = PBXBuildFile; fileRef = AC1F82641FBA3CBD0052BA77 /* labels.txt */; }; - AC1F82691FBA3F930052BA77 /* libtensorflow-lite.a in Frameworks */ = {isa = PBXBuildFile; fileRef = AC1F82681FBA3F930052BA77 /* libtensorflow-lite.a */; }; ACA1A4CA1FBB6C28009B8D86 /* mobilenet_quant_v1_224.tflite in Resources */ = {isa = PBXBuildFile; fileRef = ACA1A4C91FBB6C28009B8D86 /* mobilenet_quant_v1_224.tflite */; }; /* End PBXBuildFile section */ @@ -38,7 +37,6 @@ 3BC5BE4BBD09374D3E98F082 /* Pods-tflite_camera_example.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-tflite_camera_example.debug.xcconfig"; path = "Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example.debug.xcconfig"; sourceTree = ""; }; 55ED318E8D29C8AFEF03DF1E /* Pods-tflite_camera_example.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-tflite_camera_example.release.xcconfig"; path = "Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example.release.xcconfig"; sourceTree = ""; }; AC1F82641FBA3CBD0052BA77 /* labels.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = labels.txt; sourceTree = ""; }; - AC1F82681FBA3F930052BA77 /* libtensorflow-lite.a */ = {isa = PBXFileReference; lastKnownFileType = archive.ar; name = "libtensorflow-lite.a"; path = "../../../gen/lib/libtensorflow-lite.a"; sourceTree = ""; }; ACA1A4C91FBB6C28009B8D86 /* mobilenet_quant_v1_224.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = mobilenet_quant_v1_224.tflite; sourceTree = ""; }; /* End PBXFileReference section */ @@ -47,7 +45,6 @@ isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( - AC1F82691FBA3F930052BA77 /* libtensorflow-lite.a in Frameworks */, 1CB47D491ED3AD1700DF7666 /* AVFoundation.framework in Frameworks */, 1CA5EB931ED3ABFB00247A34 /* CoreMedia.framework in Frameworks */, 54DC6C3C5F734F3A58069F0C /* libPods-tflite_camera_example.a in Frameworks */, @@ -60,7 +57,6 @@ 24D7686C331131624F4454A0 /* Frameworks */ = { isa = PBXGroup; children = ( - AC1F82681FBA3F930052BA77 /* libtensorflow-lite.a */, 1CB47D481ED3AD1700DF7666 /* AVFoundation.framework */, 1CA5EB921ED3ABFB00247A34 /* CoreMedia.framework */, 1C0D734A1ECCC460008C1DAB /* CoreGraphics.framework */, @@ -336,7 +332,6 @@ ../../../downloads/, ); IPHONEOS_DEPLOYMENT_TARGET = 8.0; - LIBRARY_SEARCH_PATHS = ../../../gen/lib/; MTL_ENABLE_DEBUG_INFO = YES; ONLY_ACTIVE_ARCH = YES; SDKROOT = iphoneos; @@ -384,7 +379,6 @@ ../../../downloads/, ); IPHONEOS_DEPLOYMENT_TARGET = 8.0; - LIBRARY_SEARCH_PATHS = ../../../gen/lib/; MTL_ENABLE_DEBUG_INFO = NO; SDKROOT = iphoneos; TARGETED_DEVICE_FAMILY = "1,2"; diff --git a/tensorflow/contrib/lite/examples/ios/download_models.sh b/tensorflow/contrib/lite/examples/ios/download_models.sh new file mode 100755 index 0000000000000000000000000000000000000000..ccd163758c5830dc9367e023dcb3a604e07ca5db --- /dev/null +++ b/tensorflow/contrib/lite/examples/ios/download_models.sh @@ -0,0 +1,57 @@ +#!/bin/bash +# Copyright 2017 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. +# ============================================================================== + +set -ex + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MODELS_URL="https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_224_ios_lite_float_2017_11_08.zip" +QUANTIZED_MODELS_URL="https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip" +DOWNLOADS_DIR=$(mktemp -d) + +cd $SCRIPT_DIR + +download_and_extract() { + local usage="Usage: download_and_extract URL DIR" + local url="${1:?${usage}}" + local dir="${2:?${usage}}" + echo "downloading ${url}" >&2 + mkdir -p "${dir}" + tempdir=$(mktemp -d) + tempdir2=$(mktemp -d) + + curl -L ${url} > ${tempdir}/zipped.zip + unzip ${tempdir}/zipped.zip -d ${tempdir2} + + # If the zip file contains nested directories, extract the files from the + # inner directory. + if ls ${tempdir2}/*/* 1> /dev/null 2>&1; then + # unzip has no strip components, so unzip to a temp dir, and move the + # files we want from the tempdir to destination. + cp -R ${tempdir2}/*/* ${dir}/ + else + cp -R ${tempdir2}/* ${dir}/ + fi + rm -rf ${tempdir2} ${tempdir} +} + +download_and_extract "${MODELS_URL}" "${DOWNLOADS_DIR}/models" +download_and_extract "${QUANTIZED_MODELS_URL}" "${DOWNLOADS_DIR}/quantized_models" + +file ${DOWNLOADS_DIR}/models + +cp ${DOWNLOADS_DIR}/models/models/* simple/data/ +cp ${DOWNLOADS_DIR}/quantized_models/* camera/data/ + diff --git a/tensorflow/contrib/lite/examples/ios/simple/Podfile b/tensorflow/contrib/lite/examples/ios/simple/Podfile index 1740ad64573a84fae6de0fcf284eb06afec67e25..e4aca2be82d437a0225d2c15d3e486b0344aa978 100644 --- a/tensorflow/contrib/lite/examples/ios/simple/Podfile +++ b/tensorflow/contrib/lite/examples/ios/simple/Podfile @@ -1,5 +1,5 @@ platform :ios, '8.0' inhibit_all_warnings! -target 'tf_simple_example' - pod 'TensorFlow-experimental' +target 'tflite_simple_example' + pod 'TensorFlowLite' diff --git a/tensorflow/contrib/lite/examples/ios/simple/RunModel-Info.plist b/tensorflow/contrib/lite/examples/ios/simple/RunModel-Info.plist index 1a3eaa8a2c18d1cd24dfd475d396b00ec4d86c9d..a19a43a7541e3d751116e868dbcbdd607d15ab4a 100644 --- a/tensorflow/contrib/lite/examples/ios/simple/RunModel-Info.plist +++ b/tensorflow/contrib/lite/examples/ios/simple/RunModel-Info.plist @@ -7,7 +7,7 @@ CFBundleDisplayName tflite-simple-example CFBundleExecutable - tf_simple_example + tflite_simple_example CFBundleIdentifier $(PRODUCT_BUNDLE_IDENTIFIER) CFBundleInfoDictionaryVersion diff --git a/tensorflow/contrib/lite/examples/ios/simple/RunModelViewController.mm b/tensorflow/contrib/lite/examples/ios/simple/RunModelViewController.mm index 0dafb1f61e19f46bb3b17f07c55e09f5813ed560..0ab7aa25d0b4e6d2c02e61ec1d82b85258b3dfbc 100644 --- a/tensorflow/contrib/lite/examples/ios/simple/RunModelViewController.mm +++ b/tensorflow/contrib/lite/examples/ios/simple/RunModelViewController.mm @@ -29,13 +29,6 @@ #include "ios_image_load.h" -#define LOG(x) std::cerr -#define CHECK(x) \ - if (!(x)) { \ - LOG(ERROR) << #x << "failed"; \ - exit(1); \ - } - NSString* RunInferenceOnImage(); @interface RunModelViewController () @@ -89,28 +82,29 @@ static void GetTopN(const float* prediction, const int prediction_size, const in NSString* FilePathForResourceName(NSString* name, NSString* extension) { NSString* file_path = [[NSBundle mainBundle] pathForResource:name ofType:extension]; if (file_path == NULL) { - LOG(FATAL) << "Couldn't find '" << [name UTF8String] << "." << [extension UTF8String] - << "' in bundle."; + NSLog(@"Couldn't find '%@.%@' in bundle.", name, extension); + exit(-1); } return file_path; } NSString* RunInferenceOnImage() { - std::string graph; + NSString* graph = @"mobilenet_v1_1.0_224"; const int num_threads = 1; std::string input_layer_type = "float"; std::vector sizes = {1, 224, 224, 3}; - NSString* graph_path = FilePathForResourceName(@"mobilenet_v1_1.0_224", @"tflite"); + const NSString* graph_path = FilePathForResourceName(graph, @"tflite"); std::unique_ptr model( tflite::FlatBufferModel::BuildFromFile([graph_path UTF8String])); if (!model) { - LOG(FATAL) << "Failed to mmap model " << graph; + NSLog(@"Failed to mmap model %@.", graph); + exit(-1); } - LOG(INFO) << "Loaded model " << graph; + NSLog(@"Loaded model %@.", graph); model->error_reporter(); - LOG(INFO) << "resolved reporter"; + NSLog(@"Resolved reporter."); #ifdef TFLITE_CUSTOM_OPS_HEADER tflite::MutableOpResolver resolver; @@ -122,7 +116,8 @@ NSString* RunInferenceOnImage() { std::unique_ptr interpreter; tflite::InterpreterBuilder(*model, resolver)(&interpreter); if (!interpreter) { - LOG(FATAL) << "Failed to construct interpreter"; + NSLog(@"Failed to construct interpreter."); + exit(-1); } if (num_threads != -1) { @@ -136,7 +131,8 @@ NSString* RunInferenceOnImage() { } if (interpreter->AllocateTensors() != kTfLiteOk) { - LOG(FATAL) << "Failed to allocate tensors!"; + NSLog(@"Failed to allocate tensors."); + exit(-1); } // Read the label list @@ -181,7 +177,8 @@ NSString* RunInferenceOnImage() { } if (interpreter->Invoke() != kTfLiteOk) { - LOG(FATAL) << "Failed to invoke!"; + NSLog(@"Failed to invoke!"); + exit(-1); } float* output = interpreter->typed_output_tensor(0); @@ -211,11 +208,9 @@ NSString* RunInferenceOnImage() { ss << "\n"; } - LOG(INFO) << "Predictions: " << ss.str(); - std::string predictions = ss.str(); NSString* result = @""; result = [NSString stringWithFormat:@"%@ - %s", result, predictions.c_str()]; - + NSLog(@"Predictions: %@", result); return result; } diff --git a/tensorflow/contrib/lite/examples/ios/simple/simple.xcodeproj/project.pbxproj b/tensorflow/contrib/lite/examples/ios/simple/simple.xcodeproj/project.pbxproj index 9277c230b8cce1b5673a50d32d7640d52e2e8f9d..f5b8382d5ae4ac80a7edb52c34ebaf12ad65f4db 100644 --- a/tensorflow/contrib/lite/examples/ios/simple/simple.xcodeproj/project.pbxproj +++ b/tensorflow/contrib/lite/examples/ios/simple/simple.xcodeproj/project.pbxproj @@ -9,7 +9,7 @@ /* Begin PBXBuildFile section */ 1C0D734B1ECCC460008C1DAB /* CoreGraphics.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 1C0D734A1ECCC460008C1DAB /* CoreGraphics.framework */; }; 1CA45FFF1ECCC356002FA6A4 /* UIKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 1CA45FFE1ECCC356002FA6A4 /* UIKit.framework */; }; - 594C14AE1FB8F9B500EE8BFE /* libtensorflow-lite.a in Frameworks */ = {isa = PBXBuildFile; fileRef = 594C14AD1FB8F9B500EE8BFE /* libtensorflow-lite.a */; }; + 1E6F42DBB39A4A3871D4F848 /* libPods-tflite_simple_example.a in Frameworks */ = {isa = PBXBuildFile; fileRef = 73DBC33C5DD9A526EE6D1EF2 /* libPods-tflite_simple_example.a */; }; 594C14B11FB9037100EE8BFE /* labels.txt in Resources */ = {isa = PBXBuildFile; fileRef = 594C14AF1FB9037100EE8BFE /* labels.txt */; }; 594C14B21FB9037100EE8BFE /* mobilenet_v1_1.0_224.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 594C14B01FB9037100EE8BFE /* mobilenet_v1_1.0_224.tflite */; }; 59A3D0011CF4E68100C4259F /* AppDelegate.mm in Sources */ = {isa = PBXBuildFile; fileRef = 59A3CFF21CF4E68100C4259F /* AppDelegate.mm */; }; @@ -24,8 +24,7 @@ 1C0D73481ECCC41B008C1DAB /* CoreImage.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreImage.framework; path = System/Library/Frameworks/CoreImage.framework; sourceTree = SDKROOT; }; 1C0D734A1ECCC460008C1DAB /* CoreGraphics.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreGraphics.framework; path = System/Library/Frameworks/CoreGraphics.framework; sourceTree = SDKROOT; }; 1CA45FFE1ECCC356002FA6A4 /* UIKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = UIKit.framework; path = System/Library/Frameworks/UIKit.framework; sourceTree = SDKROOT; }; - 5911579B1CF4011C00C31E3A /* tf_simple_example.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = tf_simple_example.app; sourceTree = BUILT_PRODUCTS_DIR; }; - 594C14AD1FB8F9B500EE8BFE /* libtensorflow-lite.a */ = {isa = PBXFileReference; lastKnownFileType = archive.ar; name = "libtensorflow-lite.a"; path = "../../../gen/lib/libtensorflow-lite.a"; sourceTree = ""; }; + 5911579B1CF4011C00C31E3A /* tflite_simple_example.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = tflite_simple_example.app; sourceTree = BUILT_PRODUCTS_DIR; }; 594C14AF1FB9037100EE8BFE /* labels.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = labels.txt; sourceTree = ""; }; 594C14B01FB9037100EE8BFE /* mobilenet_v1_1.0_224.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = mobilenet_v1_1.0_224.tflite; sourceTree = ""; }; 59A3CFF11CF4E68100C4259F /* AppDelegate.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = AppDelegate.h; sourceTree = ""; }; @@ -38,7 +37,9 @@ 59A3CFFE1CF4E68100C4259F /* RunModelViewController.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = RunModelViewController.h; sourceTree = ""; }; 59A3CFFF1CF4E68100C4259F /* RunModelViewController.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = RunModelViewController.mm; sourceTree = ""; }; 59A3D0001CF4E68100C4259F /* RunModelViewController.xib */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = file.xib; path = RunModelViewController.xib; sourceTree = ""; }; - 73DBC33C5DD9A526EE6D1EF2 /* libPods-tf_simple_example.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPods-tf_simple_example.a"; sourceTree = BUILT_PRODUCTS_DIR; }; + 5D6203B9FAEEB9824194DBE8 /* Pods-tflite_simple_example.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-tflite_simple_example.release.xcconfig"; path = "Pods/Target Support Files/Pods-tflite_simple_example/Pods-tflite_simple_example.release.xcconfig"; sourceTree = ""; }; + 73DBC33C5DD9A526EE6D1EF2 /* libPods-tflite_simple_example.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPods-tflite_simple_example.a"; sourceTree = BUILT_PRODUCTS_DIR; }; + 987DD5BCAB2DD8B682674E20 /* Pods-tflite_simple_example.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-tflite_simple_example.debug.xcconfig"; path = "Pods/Target Support Files/Pods-tflite_simple_example/Pods-tflite_simple_example.debug.xcconfig"; sourceTree = ""; }; /* End PBXFileReference section */ /* Begin PBXFrameworksBuildPhase section */ @@ -46,9 +47,9 @@ isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( - 594C14AE1FB8F9B500EE8BFE /* libtensorflow-lite.a in Frameworks */, 1C0D734B1ECCC460008C1DAB /* CoreGraphics.framework in Frameworks */, 1CA45FFF1ECCC356002FA6A4 /* UIKit.framework in Frameworks */, + 1E6F42DBB39A4A3871D4F848 /* libPods-tflite_simple_example.a in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; }; @@ -58,11 +59,10 @@ 24D7686C331131624F4454A0 /* Frameworks */ = { isa = PBXGroup; children = ( - 594C14AD1FB8F9B500EE8BFE /* libtensorflow-lite.a */, 1C0D734A1ECCC460008C1DAB /* CoreGraphics.framework */, 1C0D73481ECCC41B008C1DAB /* CoreImage.framework */, 1CA45FFE1ECCC356002FA6A4 /* UIKit.framework */, - 73DBC33C5DD9A526EE6D1EF2 /* libPods-tf_simple_example.a */, + 73DBC33C5DD9A526EE6D1EF2 /* libPods-tflite_simple_example.a */, ); name = Frameworks; sourceTree = ""; @@ -82,13 +82,14 @@ 59A3D0001CF4E68100C4259F /* RunModelViewController.xib */, 5911579C1CF4011C00C31E3A /* Products */, 24D7686C331131624F4454A0 /* Frameworks */, + 5CE7E4179B26BF77944D8637 /* Pods */, ); sourceTree = ""; }; 5911579C1CF4011C00C31E3A /* Products */ = { isa = PBXGroup; children = ( - 5911579B1CF4011C00C31E3A /* tf_simple_example.app */, + 5911579B1CF4011C00C31E3A /* tflite_simple_example.app */, ); name = Products; sourceTree = ""; @@ -103,24 +104,36 @@ path = data; sourceTree = ""; }; + 5CE7E4179B26BF77944D8637 /* Pods */ = { + isa = PBXGroup; + children = ( + 987DD5BCAB2DD8B682674E20 /* Pods-tflite_simple_example.debug.xcconfig */, + 5D6203B9FAEEB9824194DBE8 /* Pods-tflite_simple_example.release.xcconfig */, + ); + name = Pods; + sourceTree = ""; + }; /* End PBXGroup section */ /* Begin PBXNativeTarget section */ - 5911579A1CF4011C00C31E3A /* tf_simple_example */ = { + 5911579A1CF4011C00C31E3A /* tflite_simple_example */ = { isa = PBXNativeTarget; - buildConfigurationList = 591157B21CF4011D00C31E3A /* Build configuration list for PBXNativeTarget "tf_simple_example" */; + buildConfigurationList = 591157B21CF4011D00C31E3A /* Build configuration list for PBXNativeTarget "tflite_simple_example" */; buildPhases = ( + A507411BCC70190B9ABD2721 /* [CP] Check Pods Manifest.lock */, 591157971CF4011C00C31E3A /* Sources */, 591157981CF4011C00C31E3A /* Frameworks */, 591157991CF4011C00C31E3A /* Resources */, + 25E1671BDC7334C678FB5DFB /* [CP] Embed Pods Frameworks */, + 10976C49D86B7F8A59157601 /* [CP] Copy Pods Resources */, ); buildRules = ( ); dependencies = ( ); - name = tf_simple_example; + name = tflite_simple_example; productName = tf_ios_makefile_example; - productReference = 5911579B1CF4011C00C31E3A /* tf_simple_example.app */; + productReference = 5911579B1CF4011C00C31E3A /* tflite_simple_example.app */; productType = "com.apple.product-type.application"; }; /* End PBXNativeTarget section */ @@ -152,7 +165,7 @@ projectDirPath = ""; projectRoot = ""; targets = ( - 5911579A1CF4011C00C31E3A /* tf_simple_example */, + 5911579A1CF4011C00C31E3A /* tflite_simple_example */, ); }; /* End PBXProject section */ @@ -171,6 +184,57 @@ }; /* End PBXResourcesBuildPhase section */ +/* Begin PBXShellScriptBuildPhase section */ + 10976C49D86B7F8A59157601 /* [CP] Copy Pods Resources */ = { + isa = PBXShellScriptBuildPhase; + buildActionMask = 2147483647; + files = ( + ); + inputPaths = ( + ); + name = "[CP] Copy Pods Resources"; + outputPaths = ( + ); + runOnlyForDeploymentPostprocessing = 0; + shellPath = /bin/sh; + shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-tflite_simple_example/Pods-tflite_simple_example-resources.sh\"\n"; + showEnvVarsInLog = 0; + }; + 25E1671BDC7334C678FB5DFB /* [CP] Embed Pods Frameworks */ = { + isa = PBXShellScriptBuildPhase; + buildActionMask = 2147483647; + files = ( + ); + inputPaths = ( + ); + name = "[CP] Embed Pods Frameworks"; + outputPaths = ( + ); + runOnlyForDeploymentPostprocessing = 0; + shellPath = /bin/sh; + shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-tflite_simple_example/Pods-tflite_simple_example-frameworks.sh\"\n"; + showEnvVarsInLog = 0; + }; + A507411BCC70190B9ABD2721 /* [CP] Check Pods Manifest.lock */ = { + isa = PBXShellScriptBuildPhase; + buildActionMask = 2147483647; + files = ( + ); + inputPaths = ( + "${PODS_PODFILE_DIR_PATH}/Podfile.lock", + "${PODS_ROOT}/Manifest.lock", + ); + name = "[CP] Check Pods Manifest.lock"; + outputPaths = ( + "$(DERIVED_FILE_DIR)/Pods-tflite_simple_example-checkManifestLockResult.txt", + ); + runOnlyForDeploymentPostprocessing = 0; + shellPath = /bin/sh; + shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n # print error to STDERR\n echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n"; + showEnvVarsInLog = 0; + }; +/* End PBXShellScriptBuildPhase section */ + /* Begin PBXSourcesBuildPhase section */ 591157971CF4011C00C31E3A /* Sources */ = { isa = PBXSourcesBuildPhase; @@ -274,6 +338,7 @@ }; 591157B31CF4011D00C31E3A /* Debug */ = { isa = XCBuildConfiguration; + baseConfigurationReference = 987DD5BCAB2DD8B682674E20 /* Pods-tflite_simple_example.debug.xcconfig */; buildSettings = { CLANG_DEBUG_INFORMATION_LEVEL = default; CODE_SIGN_IDENTITY = "iPhone Developer"; @@ -283,15 +348,10 @@ GCC_ENABLE_CPP_RTTI = YES; HEADER_SEARCH_PATHS = ( "$(inherited)", - ../../../../../../, - ../../../downloads/flatbuffers/include/, - ../../../downloads/eigen/, - ../../../downloads/, ); INFOPLIST_FILE = "$(SRCROOT)/RunModel-Info.plist"; IPHONEOS_DEPLOYMENT_TARGET = 9.2; LD_RUNPATH_SEARCH_PATHS = "$(inherited) @executable_path/Frameworks"; - LIBRARY_SEARCH_PATHS = ../../../gen/lib/; OTHER_CPLUSPLUSFLAGS = "$(OTHER_CFLAGS)"; OTHER_LDFLAGS = "$(inherited)"; PRODUCT_BUNDLE_IDENTIFIER = "com.google.tflite-simple-example"; @@ -304,6 +364,7 @@ }; 591157B41CF4011D00C31E3A /* Release */ = { isa = XCBuildConfiguration; + baseConfigurationReference = 5D6203B9FAEEB9824194DBE8 /* Pods-tflite_simple_example.release.xcconfig */; buildSettings = { CLANG_DEBUG_INFORMATION_LEVEL = default; CODE_SIGN_IDENTITY = "iPhone Developer"; @@ -313,15 +374,10 @@ GCC_ENABLE_CPP_RTTI = YES; HEADER_SEARCH_PATHS = ( "$(inherited)", - ../../../../../../, - ../../../downloads/flatbuffers/include/, - ../../../downloads/eigen/, - ../../../downloads/, ); INFOPLIST_FILE = "$(SRCROOT)/RunModel-Info.plist"; IPHONEOS_DEPLOYMENT_TARGET = 9.2; LD_RUNPATH_SEARCH_PATHS = "$(inherited) @executable_path/Frameworks"; - LIBRARY_SEARCH_PATHS = ../../../gen/lib/; ONLY_ACTIVE_ARCH = YES; OTHER_CPLUSPLUSFLAGS = "$(OTHER_CFLAGS)"; OTHER_LDFLAGS = "$(inherited)"; @@ -344,7 +400,7 @@ defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; - 591157B21CF4011D00C31E3A /* Build configuration list for PBXNativeTarget "tf_simple_example" */ = { + 591157B21CF4011D00C31E3A /* Build configuration list for PBXNativeTarget "tflite_simple_example" */ = { isa = XCConfigurationList; buildConfigurations = ( 591157B31CF4011D00C31E3A /* Debug */, diff --git a/tensorflow/contrib/lite/examples/label_image/BUILD b/tensorflow/contrib/lite/examples/label_image/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..959347b5491514ddc13af57ea6f7385a0d39e418 --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/BUILD @@ -0,0 +1,83 @@ +# Description: +# TensorFlow Lite Example Label Image. + +package(default_visibility = ["//visibility:public"]) + +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "tf_cc_binary") +load("//tensorflow/contrib/lite:build_def.bzl", "tflite_linkopts") + +exports_files(glob([ + "testdata/*.bmp", +])) + +tf_cc_binary( + name = "label_image", + srcs = [ + "get_top_n.h", + "get_top_n_impl.h", + "label_image.cc", + ], + linkopts = tflite_linkopts() + select({ + "//tensorflow:android": [ + "-pie", # Android 5.0 and later supports only PIE + "-lm", # some builtin ops, e.g., tanh, need -lm + ], + "//conditions:default": [], + }), + deps = [ + ":bitmap_helpers", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite:string_util", + "//tensorflow/contrib/lite/kernels:builtin_ops", + ], +) + +cc_library( + name = "bitmap_helpers", + srcs = ["bitmap_helpers.cc"], + hdrs = [ + "bitmap_helpers.h", + "bitmap_helpers_impl.h", + "label_image.h", + ], + deps = [ + "//tensorflow/contrib/lite:builtin_op_data", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite:schema_fbs_version", + "//tensorflow/contrib/lite:string", + "//tensorflow/contrib/lite:string_util", + "//tensorflow/contrib/lite/kernels:builtin_ops", + "//tensorflow/contrib/lite/schema:schema_fbs", + ], +) + +# TODO(ahentz): Test disabled as it has a memory leek from read_bmp +# cc_test( +# name = "label_image_test", +# srcs = [ +# "get_top_n.h", +# "get_top_n_impl.h", +# "label_image_test.cc", +# ], +# data = [ +# "testdata/grace_hopper.bmp", +# ], +# deps = [ +# ":bitmap_helpers", +# "//testing/base/public:gunit", +# ], +# ) + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.cc b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.cc new file mode 100644 index 0000000000000000000000000000000000000000..0b38cd38c83927c65d251b9356301b6bef7521f2 --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.cc @@ -0,0 +1,120 @@ +/* Copyright 2017 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 +#include +#include +#include +#include + +#include // NOLINT(build/include_order) + +#include "tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h" + +#define LOG(x) std::cerr + +namespace tflite { +namespace label_image { + +uint8_t* decode_bmp(const uint8_t* input, int row_size, uint8_t* const output, + int width, int height, int channels, bool top_down) { + for (int i = 0; i < height; i++) { + int src_pos; + int dst_pos; + + for (int j = 0; j < width; j++) { + if (!top_down) { + src_pos = ((height - 1 - i) * row_size) + j * channels; + } else { + src_pos = i * row_size + j * channels; + } + + dst_pos = (i * width + j) * channels; + + switch (channels) { + case 1: + output[dst_pos] = input[src_pos]; + break; + case 3: + // BGR -> RGB + output[dst_pos] = input[src_pos + 2]; + output[dst_pos + 1] = input[src_pos + 1]; + output[dst_pos + 2] = input[src_pos]; + break; + case 4: + // BGRA -> RGBA + output[dst_pos] = input[src_pos + 2]; + output[dst_pos + 1] = input[src_pos + 1]; + output[dst_pos + 2] = input[src_pos]; + output[dst_pos + 3] = input[src_pos + 3]; + break; + default: + LOG(FATAL) << "Unexpected number of channels: " << channels; + break; + } + } + } + + return output; +} + +uint8_t* read_bmp(const std::string& input_bmp_name, int* width, int* height, + int* channels, Settings* s) { + int begin, end; + + std::ifstream file(input_bmp_name, std::ios::in | std::ios::binary); + if (!file) { + LOG(FATAL) << "input file " << input_bmp_name << " not found\n"; + exit(-1); + } + + begin = file.tellg(); + file.seekg(0, std::ios::end); + end = file.tellg(); + size_t len = end - begin; + + if (s->verbose) LOG(INFO) << "len: " << len << "\n"; + + const uint8_t* img_bytes = new uint8_t[len]; + file.seekg(0, std::ios::beg); + file.read((char*)img_bytes, len); + const int32_t header_size = + *(reinterpret_cast(img_bytes + 10)); + *width = *(reinterpret_cast(img_bytes + 18)); + *height = *(reinterpret_cast(img_bytes + 22)); + const int32_t bpp = *(reinterpret_cast(img_bytes + 28)); + *channels = bpp / 8; + + if (s->verbose) + LOG(INFO) << "width, height, channels: " << *width << ", " << *height + << ", " << *channels << "\n"; + + // there may be padding bytes when the width is not a multiple of 4 bytes + // 8 * channels == bits per pixel + const int row_size = (8 * *channels * *width + 31) / 32 * 4; + + // if height is negative, data layout is top down + // otherwise, it's bottom up + bool top_down = (*height < 0); + + // Decode image, allocating tensor once the image size is known + uint8_t* output = new uint8_t[abs(*height) * *width * *channels]; + const uint8_t* bmp_pixels = &img_bytes[header_size]; + return decode_bmp(bmp_pixels, row_size, output, *width, abs(*height), + *channels, top_down); +} + +} // namespace label_image +} // namespace tflite diff --git a/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h new file mode 100644 index 0000000000000000000000000000000000000000..97343dde6b31694e5b2de20b35a7083fb8fe4a0e --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h @@ -0,0 +1,42 @@ +/* Copyright 2017 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_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_H_ +#define TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_H_ + +#include "tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h" +#include "tensorflow/contrib/lite/examples/label_image/label_image.h" + +namespace tflite { +namespace label_image { + +uint8_t* read_bmp(const std::string& input_bmp_name, int* width, int* height, + int* channels, Settings* s); + +template +void resize(T* out, uint8_t* in, int image_height, int image_width, + int image_channels, int wanted_height, int wanted_width, + int wanted_channels, Settings* s); + +// explicit instantiation +template void resize(uint8_t*, unsigned char*, int, int, int, int, int, + int, Settings*); +template void resize(float*, unsigned char*, int, int, int, int, int, + int, Settings*); + +} // namespace label_image +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_H diff --git a/tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h new file mode 100644 index 0000000000000000000000000000000000000000..2a64c1de725b601e9b6e9325d9faacb37df0e626 --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h @@ -0,0 +1,103 @@ +/* Copyright 2017 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_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_IMPL_H_ +#define TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_IMPL_H_ + +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/string_util.h" +#include "tensorflow/contrib/lite/version.h" + +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/string_util.h" +#include "tensorflow/contrib/lite/version.h" + +#include "tensorflow/contrib/lite/examples/label_image/label_image.h" + +namespace tflite { +namespace label_image { + +template +void resize(T* out, uint8_t* in, int image_height, int image_width, + int image_channels, int wanted_height, int wanted_width, + int wanted_channels, Settings* s) { + int number_of_pixels = image_height * image_width * image_channels; + std::unique_ptr interpreter(new Interpreter); + + int base_index = 0; + + // two inputs: input and new_sizes + interpreter->AddTensors(2, &base_index); + // one output + interpreter->AddTensors(1, &base_index); + // set input and output tensors + interpreter->SetInputs({0, 1}); + interpreter->SetOutputs({2}); + + // set parameters of tensors + TfLiteQuantizationParams quant; + interpreter->SetTensorParametersReadWrite( + 0, kTfLiteFloat32, "input", + {1, image_height, image_width, image_channels}, quant); + interpreter->SetTensorParametersReadWrite(1, kTfLiteInt32, "new_size", {2}, + quant); + interpreter->SetTensorParametersReadWrite( + 2, kTfLiteFloat32, "output", + {1, wanted_height, wanted_width, wanted_channels}, quant); + + ops::builtin::BuiltinOpResolver resolver; + TfLiteRegistration* resize_op = + resolver.FindOp(BuiltinOperator_RESIZE_BILINEAR); + auto* params = reinterpret_cast( + malloc(sizeof(TfLiteResizeBilinearParams))); + params->align_corners = false; + interpreter->AddNodeWithParameters({0, 1}, {2}, nullptr, 0, params, resize_op, + nullptr); + + interpreter->AllocateTensors(); + + // fill input image + // in[] are integers, cannot do memcpy() directly + auto input = interpreter->typed_tensor(0); + for (int i = 0; i < number_of_pixels; i++) { + input[i] = in[i]; + } + + // fill new_sizes + interpreter->typed_tensor(1)[0] = wanted_height; + interpreter->typed_tensor(1)[1] = wanted_width; + + interpreter->Invoke(); + + auto output = interpreter->typed_tensor(2); + auto output_number_of_pixels = + wanted_height * wanted_height * wanted_channels; + + for (int i = 0; i < output_number_of_pixels; i++) { + if (s->input_floating) + out[i] = (output[i] - s->input_mean) / s->input_std; + else + out[i] = (uint8_t)output[i]; + } +} + +} // namespace label_image +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_IMPL_H_ diff --git a/tensorflow/contrib/lite/examples/label_image/get_top_n.h b/tensorflow/contrib/lite/examples/label_image/get_top_n.h new file mode 100644 index 0000000000000000000000000000000000000000..70a7586fe6a008f0da20a7bac928ca676e5914ab --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/get_top_n.h @@ -0,0 +1,38 @@ +/* Copyright 2017 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_LITE_EXAMPLES_LABEL_IMAGE_GET_TOP_N_H +#define TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_GET_TOP_N_H + +#include "tensorflow/contrib/lite/examples/label_image/get_top_n_impl.h" + +namespace tflite { +namespace label_image { + +template +void get_top_n(T* prediction, int prediction_size, size_t num_results, + float threshold, std::vector>* top_results, + bool input_floating); + +// explicit instantiation so that we can use them otherwhere +template void get_top_n(uint8_t*, int, size_t, float, + std::vector>*, bool); +template void get_top_n(float*, int, size_t, float, + std::vector>*, bool); + +} // namespace label_image +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_GET_TOP_N_H diff --git a/tensorflow/contrib/lite/examples/label_image/get_top_n_impl.h b/tensorflow/contrib/lite/examples/label_image/get_top_n_impl.h new file mode 100644 index 0000000000000000000000000000000000000000..e416fbd39b125ea65d1155b19ab0967a9062e71a --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/get_top_n_impl.h @@ -0,0 +1,70 @@ +/* Copyright 2017 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_LITE_EXAMPLES_LABEL_IMAGE_GET_TOP_N_IMPL_H +#define TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_GET_TOP_N_IMPL_H + +#include +#include + +namespace tflite { +namespace label_image { + +extern bool input_floating; + +// Returns the top N confidence values over threshold in the provided vector, +// sorted by confidence in descending order. +template +void get_top_n(T* prediction, int prediction_size, size_t num_results, + float threshold, std::vector>* top_results, + bool input_floating) { + // Will contain top N results in ascending order. + std::priority_queue, std::vector>, + std::greater>> + top_result_pq; + + const long count = prediction_size; // NOLINT(runtime/int) + for (int i = 0; i < count; ++i) { + float value; + if (input_floating) + value = prediction[i]; + else + value = prediction[i] / 255.0; + // Only add it if it beats the threshold and has a chance at being in + // the top N. + if (value < threshold) { + continue; + } + + top_result_pq.push(std::pair(value, i)); + + // If at capacity, kick the smallest value out. + if (top_result_pq.size() > num_results) { + top_result_pq.pop(); + } + } + + // Copy to output vector and reverse into descending order. + while (!top_result_pq.empty()) { + top_results->push_back(top_result_pq.top()); + top_result_pq.pop(); + } + std::reverse(top_results->begin(), top_results->end()); +} + +} // namespace label_image +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_GET_TOP_N_IMPL_H diff --git a/tensorflow/contrib/lite/examples/label_image/label_image.cc b/tensorflow/contrib/lite/examples/label_image/label_image.cc new file mode 100644 index 0000000000000000000000000000000000000000..a91467d345fdce1268635a69a96939921dc170e8 --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/label_image.cc @@ -0,0 +1,308 @@ +/* Copyright 2017 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include // NOLINT(build/include_order) +#include // NOLINT(build/include_order) +#include // NOLINT(build/include_order) +#include // NOLINT(build/include_order) +#include // NOLINT(build/include_order) +#include // NOLINT(build/include_order) + +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/model.h" +#include "tensorflow/contrib/lite/optional_debug_tools.h" +#include "tensorflow/contrib/lite/string_util.h" + +#include "tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h" +#include "tensorflow/contrib/lite/examples/label_image/get_top_n.h" + +#define LOG(x) std::cerr + +namespace tflite { +namespace label_image { + +double get_us(struct timeval t) { return (t.tv_sec * 1000000 + t.tv_usec); } + +// Takes a file name, and loads a list of labels from it, one per line, and +// returns a vector of the strings. It pads with empty strings so the length +// of the result is a multiple of 16, because our model expects that. +TfLiteStatus ReadLabelsFile(const string& file_name, + std::vector* result, + size_t* found_label_count) { + std::ifstream file(file_name); + if (!file) { + LOG(FATAL) << "Labels file " << file_name << " not found\n"; + return kTfLiteError; + } + result->clear(); + string line; + while (std::getline(file, line)) { + result->push_back(line); + } + *found_label_count = result->size(); + const int padding = 16; + while (result->size() % padding) { + result->emplace_back(); + } + return kTfLiteOk; +} + +void RunInference(Settings* s) { + if (!s->model_name.c_str()) { + LOG(ERROR) << "no model file name\n"; + exit(-1); + } + + std::unique_ptr model; + std::unique_ptr interpreter; + model = tflite::FlatBufferModel::BuildFromFile(s->model_name.c_str()); + if (!model) { + LOG(FATAL) << "\nFailed to mmap model " << s->model_name << "\n"; + exit(-1); + } + LOG(INFO) << "Loaded model " << s->model_name << "\n"; + model->error_reporter(); + LOG(INFO) << "resolved reporter\n"; + + tflite::ops::builtin::BuiltinOpResolver resolver; + + tflite::InterpreterBuilder(*model, resolver)(&interpreter); + if (!interpreter) { + LOG(FATAL) << "Failed to construct interpreter\n"; + exit(-1); + } + + interpreter->UseNNAPI(s->accel); + + if (s->verbose) { + LOG(INFO) << "tensors size: " << interpreter->tensors_size() << "\n"; + LOG(INFO) << "nodes size: " << interpreter->nodes_size() << "\n"; + LOG(INFO) << "inputs: " << interpreter->inputs().size() << "\n"; + LOG(INFO) << "input(0) name: " << interpreter->GetInputName(0) << "\n"; + + int t_size = interpreter->tensors_size(); + for (int i = 0; i < t_size; i++) { + if (interpreter->tensor(i)->name) + LOG(INFO) << i << ": " << interpreter->tensor(i)->name << ", " + << interpreter->tensor(i)->bytes << ", " + << interpreter->tensor(i)->type << ", " + << interpreter->tensor(i)->params.scale << ", " + << interpreter->tensor(i)->params.zero_point << "\n"; + } + } + + if (s->number_of_threads != -1) { + interpreter->SetNumThreads(s->number_of_threads); + } + + int image_width = 224; + int image_height = 224; + int image_channels = 3; + uint8_t* in = read_bmp(s->input_bmp_name, &image_width, &image_height, + &image_channels, s); + + int input = interpreter->inputs()[0]; + if (s->verbose) LOG(INFO) << "input: " << input << "\n"; + + const std::vector inputs = interpreter->inputs(); + const std::vector outputs = interpreter->outputs(); + + if (s->verbose) { + LOG(INFO) << "number of inputs: " << inputs.size() << "\n"; + LOG(INFO) << "number of outputs: " << outputs.size() << "\n"; + } + + if (interpreter->AllocateTensors() != kTfLiteOk) { + LOG(FATAL) << "Failed to allocate tensors!"; + } + + if (s->verbose) PrintInterpreterState(interpreter.get()); + + // get input dimension from the input tensor metadata + // assuming one input only + TfLiteIntArray* dims = interpreter->tensor(input)->dims; + int wanted_height = dims->data[1]; + int wanted_width = dims->data[2]; + int wanted_channels = dims->data[3]; + + switch (interpreter->tensor(input)->type) { + case kTfLiteFloat32: + s->input_floating = true; + resize(interpreter->typed_tensor(input), in, image_height, + image_width, image_channels, wanted_height, wanted_width, + wanted_channels, s); + break; + case kTfLiteUInt8: + resize(interpreter->typed_tensor(input), in, + image_height, image_width, image_channels, wanted_height, + wanted_width, wanted_channels, s); + break; + default: + LOG(FATAL) << "cannot handle input type " + << interpreter->tensor(input)->type << " yet"; + exit(-1); + } + + struct timeval start_time, stop_time; + gettimeofday(&start_time, NULL); + for (int i = 0; i < s->loop_count; i++) { + if (interpreter->Invoke() != kTfLiteOk) { + LOG(FATAL) << "Failed to invoke tflite!\n"; + } + } + gettimeofday(&stop_time, NULL); + LOG(INFO) << "invoked \n"; + LOG(INFO) << "average time: " + << (get_us(stop_time) - get_us(start_time)) / (s->loop_count * 1000) + << " ms \n"; + + const int output_size = 1000; + const size_t num_results = 5; + const float threshold = 0.001f; + + std::vector> top_results; + + int output = interpreter->outputs()[0]; + switch (interpreter->tensor(output)->type) { + case kTfLiteFloat32: + get_top_n(interpreter->typed_output_tensor(0), output_size, + num_results, threshold, &top_results, true); + break; + case kTfLiteUInt8: + get_top_n(interpreter->typed_output_tensor(0), + output_size, num_results, threshold, &top_results, + false); + break; + default: + LOG(FATAL) << "cannot handle output type " + << interpreter->tensor(input)->type << " yet"; + exit(-1); + } + + std::vector labels; + size_t label_count; + + if (ReadLabelsFile(s->labels_file_name, &labels, &label_count) != kTfLiteOk) + exit(-1); + + for (const auto& result : top_results) { + const float confidence = result.first; + const int index = result.second; + LOG(INFO) << confidence << ": " << index << " " << labels[index] << "\n"; + } +} + +void display_usage() { + LOG(INFO) << "label_image\n" + << "--accelerated, -a: [0|1], use Android NNAPI or note\n" + << "--count, -c: loop interpreter->Invoke() for certain times\n" + << "--input_mean, -b: input mean\n" + << "--input_std, -s: input standard deviation\n" + << "--image, -i: image_name.bmp\n" + << "--labels, -l: labels for the model\n" + << "--tflite_model, -m: model_name.tflite\n" + << "--threads, -t: number of threads\n" + << "--verbose, -v: [0|1] print more information\n" + << "\n"; +} + +int Main(int argc, char** argv) { + Settings s; + + int c; + while (1) { + static struct option long_options[] = { + {"accelerated", required_argument, 0, 'a'}, + {"count", required_argument, 0, 'c'}, + {"verbose", required_argument, 0, 'v'}, + {"image", required_argument, 0, 'i'}, + {"labels", required_argument, 0, 'l'}, + {"tflite_model", required_argument, 0, 'm'}, + {"threads", required_argument, 0, 't'}, + {"input_mean", required_argument, 0, 'b'}, + {"input_std", required_argument, 0, 's'}, + {0, 0, 0, 0}}; + + /* getopt_long stores the option index here. */ + int option_index = 0; + + c = getopt_long(argc, argv, "a:b:c:f:i:l:m:s:t:v:", long_options, + &option_index); + + /* Detect the end of the options. */ + if (c == -1) break; + + switch (c) { + case 'a': + s.accel = strtol( // NOLINT(runtime/deprecated_fn) + optarg, (char**)NULL, 10); + break; + case 'b': + s.input_mean = strtod(optarg, NULL); + break; + case 'c': + s.loop_count = strtol( // NOLINT(runtime/deprecated_fn) + optarg, (char**)NULL, 10); + break; + case 'i': + s.input_bmp_name = optarg; + break; + case 'l': + s.labels_file_name = optarg; + break; + case 'm': + s.model_name = optarg; + break; + case 's': + s.input_std = strtod(optarg, NULL); + break; + case 't': + s.number_of_threads = strtol( // NOLINT(runtime/deprecated_fn) + optarg, (char**)NULL, 10); + break; + case 'v': + s.verbose = strtol( // NOLINT(runtime/deprecated_fn) + optarg, (char**)NULL, 10); + break; + case 'h': + case '?': + /* getopt_long already printed an error message. */ + display_usage(); + exit(-1); + default: + exit(-1); + } + } + RunInference(&s); + return 0; +} + +} // namespace label_image +} // namespace tflite + +int main(int argc, char** argv) { + return tflite::label_image::Main(argc, argv); +} diff --git a/tensorflow/contrib/lite/examples/label_image/label_image.h b/tensorflow/contrib/lite/examples/label_image/label_image.h new file mode 100644 index 0000000000000000000000000000000000000000..4de32e33fb4ef2ab5d0e111886cdc737398147e9 --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/label_image.h @@ -0,0 +1,41 @@ +/* Copyright 2017 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_LITE_EXAMPLES_LABEL_IMAGE_LABEL_IMAGE_H +#define TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_LABEL_IMAGE_H + +#include "tensorflow/contrib/lite/string.h" + +namespace tflite { +namespace label_image { + +struct Settings { + bool verbose = false; + bool accel = false; + bool input_floating = false; + int loop_count = 1; + float input_mean = 127.5f; + float input_std = 127.5f; + string model_name = "./mobilenet_quant_v1_224.tflite"; + string input_bmp_name = "./grace_hopper.bmp"; + string labels_file_name = "./labels.txt"; + string input_layer_type = "uint8_t"; + int number_of_threads = 4; +}; + +} // namespace label_image +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_LABEL_IMAGE_H diff --git a/tensorflow/contrib/lite/examples/label_image/label_image.md b/tensorflow/contrib/lite/examples/label_image/label_image.md new file mode 100644 index 0000000000000000000000000000000000000000..9ce32cf101897f2d41cd14a485aeb432344928a0 --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/label_image.md @@ -0,0 +1,78 @@ +label_image for TensorFlow Lite inspired by TensorFlow's label_image. + +To build label_image for Android, run $TENSORFLOW_ROOT/configure +and set Android NDK or configure NDK setting in +$TENSORFLOW_ROOT/WORKSPACE first. + +To build it for android ARMv8: +``` +> bazel build --config monolithic --cxxopt=-std=c++11 \ + --crosstool_top=//external:android/crosstool \ + --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ + --cpu=arm64-v8a \ + //tensorflow/contrib/lite/examples/label_image:label_image +``` +or +``` +> bazel build --config android_arm64 --config monolithic --cxxopt=-std=c++11 \ + //tensorflow/contrib/lite/examples/label_image:label_image +``` + +To build it for android arm-v7a: +``` +> bazel build --config monolithic --cxxopt=-std=c++11 \ + --crosstool_top=//external:android/crosstool \ + --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ + --cpu=armeabi-v7a \ + //tensorflow/contrib/lite/examples/label_image:label_image +``` +or +``` +> bazel build --config android_arm --config monolithic --cxxopt=-std=c++11 \ + //tensorflow/contrib/lite/examples/label_image:label_image +``` + +Build it for desktop machines (tested on Ubuntu and OS X) +``` +> bazel build --config opt --cxxopt=-std=c++11 //tensorflow/contrib/lite/examples/label_image:label_image +``` +To run it. Prepare `./mobilenet_quant_v1_224.tflite`, `./grace_hopper.bmp`, and `./labels.txt`. + +Run it: +``` +> ./label_image +Loaded model ./mobilenet_quant_v1_224.tflite +resolved reporter +invoked +average time: 100.986 ms +0.439216: 653 military uniform +0.372549: 458 bow tie +0.0705882: 466 bulletproof vest +0.0235294: 514 cornet +0.0196078: 835 suit +``` +Run `interpreter->Invoker()` 100 times: +``` +> ./label_image -c 100 +Loaded model ./mobilenet_quant_v1_224.tflite +resolved reporter +invoked +average time: 33.4694 ms +... +``` + +Run a floating point (`mobilenet_v1_1.0_224.tflite`) model, +``` +> ./label_image -f 1 -m mobilenet_v1_1.0_224.tflite +Loaded model mobilenet_v1_1.0_224.tflite +resolved reporter +invoked +average time: 263.493 ms +0.88615: 653 military uniform +0.0422316: 440 bearskin +0.0109948: 466 bulletproof vest +0.0105327: 401 academic gown +0.00947104: 723 ping-pong bal +``` + +See the source code for other command line options. diff --git a/tensorflow/contrib/lite/examples/label_image/label_image_test.cc b/tensorflow/contrib/lite/examples/label_image/label_image_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..ce35483f76e8f40ced79e1ee30774c62d0eba94e --- /dev/null +++ b/tensorflow/contrib/lite/examples/label_image/label_image_test.cc @@ -0,0 +1,61 @@ +/* Copyright 2016 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 +#include + +#include "tensorflow/contrib/lite/examples/label_image/bitmap_helpers.h" +#include "tensorflow/contrib/lite/examples/label_image/get_top_n.h" +#include "tensorflow/contrib/lite/examples/label_image/label_image.h" + +using ::testing::ElementsAreArray; + +namespace tflite { +namespace label_image { + +TEST(LabelImageTest, GraceHopper) { + std::string lena_file = + "tensorflow/contrib/lite/examples/label_image/testdata/grace_hopper.bmp"; + int height, width, channels; + Settings s; + uint8_t *data; + + data = read_bmp(lena_file, &width, &height, &channels, &s); + ASSERT_EQ(height, 606); + ASSERT_EQ(width, 517); + ASSERT_EQ(channels, 3); + + uint8_t *out = new uint8_t[606 * 517 * 3]; + downsize(out, data, 606, 517, 3, 214, 214, 3, &s); + ASSERT_EQ(out[0], 0x15); + ASSERT_EQ(out[214 * 214 * 3 - 1], 0x12); +} + +TEST(LabelImageTest, GetTopN) { + uint8_t in[] = {1, 1, 2, 2, 4, 4, 16, 32, 128, 64}; + + std::vector> top_results; + get_top_n(in, 10, 5, 0.025, &top_results, false); + ASSERT_EQ(top_results.size(), 4); + ASSERT_EQ(top_results[0].second, 8); +} + +} // namespace label_image +} // namespace tflite + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/examples/label_image/testdata/grace_hopper.bmp b/tensorflow/contrib/lite/examples/label_image/testdata/grace_hopper.bmp new file mode 100644 index 0000000000000000000000000000000000000000..0d94cd3e930a138b7c20308f5ba375576484d48b Binary files /dev/null and b/tensorflow/contrib/lite/examples/label_image/testdata/grace_hopper.bmp differ diff --git a/tensorflow/contrib/lite/g3doc/custom_operators.md b/tensorflow/contrib/lite/g3doc/custom_operators.md index 204a489a93519309bb09238f1b2c8bbd4f1f19e4..d7cc854ebac08e79d346df0aca6e1fa56b490156 100644 --- a/tensorflow/contrib/lite/g3doc/custom_operators.md +++ b/tensorflow/contrib/lite/g3doc/custom_operators.md @@ -73,7 +73,7 @@ TfLiteStatus SinEval(TfLiteContext* context, TfLiteNode* node) { } TfLiteRegistration* Register_SIN() { - static TfLiteRegistration r = {nullptr, nullptr, SinResize, SinEval}; + static TfLiteRegistration r = {nullptr, nullptr, SinPrepare, SinEval}; return &r; } ``` diff --git a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md index 9ade04eb8c696d7e0e39a8104e02b6e5feec95eb..b1bbb7c67013acfb575cc1e9f9390ba191cbd08e 100644 --- a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md +++ b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md @@ -1,4 +1,4 @@ -# TensorFlow Compatibility Guide +# TensorFlow Lite & TensorFlow Compatibility Guide TensorFlow Lite supports a number of TensorFlow operations used in common inference models. As they are processed by the TensorFlow Lite Optimizing @@ -329,18 +329,18 @@ Inputs { 0: a tensor } Outputs { - 0: a tensor equivalent to max(0, min(input, 1) + 0: a tensor equivalent to max(0, input) } ``` -**RELU1** +**RELU_N1_TO_1** ``` Inputs { 0: a tensor } Outputs { - 0: a tensor equivalent to max(-1, min(input, 6) + 0: a tensor equivalent to max(-1, min(input, 1) } ``` diff --git a/tensorflow/contrib/lite/graph_info.cc b/tensorflow/contrib/lite/graph_info.cc new file mode 100644 index 0000000000000000000000000000000000000000..e60ed2c2463cb621015ba725ca030e8d8c02f3c7 --- /dev/null +++ b/tensorflow/contrib/lite/graph_info.cc @@ -0,0 +1,224 @@ +/* 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/lite/graph_info.h" +#include + +namespace tflite { + +namespace { + +// Provide a range iterable wrapper for TfLiteIntArray* (C lists that TfLite +// C api uses. Can't use the google array_view, since we can't depend on even +// absl for embedded device reasons. +// TODO(aselle): Move this into central utilities. +class TfLiteIntArrayView { + public: + // Construct a view of a TfLiteIntArray*. Note, `int_array` should be non-null + // and this view does not take ownership of it. + explicit TfLiteIntArrayView(const TfLiteIntArray* int_array) + : int_array_(int_array) {} + + typedef const int* const_iterator; + const_iterator begin() const { return int_array_->data; } + const_iterator end() const { return &int_array_->data[int_array_->size]; } + + TfLiteIntArrayView(const TfLiteIntArrayView&) = default; + TfLiteIntArrayView& operator=(const TfLiteIntArrayView& rhs) = default; + + private: + const TfLiteIntArray* int_array_; +}; + +// Helper class that actually performs partitioning by subgraph. +// Outputs to a provided `subgraphs` structure. +// +// Example usage: +// PartitionGraphIntoIndependentSubgraphsImpl partitioner( +// info, nodes_to_part, subgraphs); +// partitioner.Partition(); +class PartitionGraphIntoIndependentSubgraphsImpl { + public: + PartitionGraphIntoIndependentSubgraphsImpl( + const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, + std::vector* subgraphs) + : info_(info), + subgraphs_(subgraphs), + node_type_(info->num_nodes(), Subgraph::kTfNonPartition) { + // Populate the node_type_ map. + for (auto node_index : TfLiteIntArrayView(nodes_to_partition)) { + node_type_[node_index] = Subgraph::kTfPartition; + } + } + + // Actually partition the graph. + void Partition() { + // Initialize here to make Partition() re-entrant. + subgraphs_->clear(); + tensor_epochs_.clear(); + tensor_epochs_.resize(info_->num_tensors(), kEpochAlwaysReady); + node_epochs_.clear(); + node_epochs_.resize(info_->num_nodes(), kEpochNotReady); + // Set computed tensors to be kEpochNotReady (initializer set everything to + // AlwaysReady). + for (int node_index = 0; node_index < info_->num_nodes(); node_index++) { + const TfLiteNode& node = info_->node(node_index); + for (int output_tensor_index : TfLiteIntArrayView(node.outputs)) { + tensor_epochs_[output_tensor_index] = kEpochNotReady; + } + } + + // Do a graph traversal where each iteration in the loop is an epoch + // that corresponds to a subgraph that only contains nodes that are of + // the same node_type_. + while (true) { + BuildSubgraph(); + if (subgraphs_->back().nodes.empty()) { + subgraphs_->pop_back(); + break; + } + } + + // Mark model outputs as subgraph outputs. All the rest have already been + // identified. + for (int output_index : info_->outputs()) { + int output_epoch = tensor_epochs_[output_index]; + Subgraph& output_subgraph = (*subgraphs_)[output_epoch]; + output_subgraph.output_tensors.push_back(output_index); + } + // Make sure every subgraph's inputs and outputs are unique. Since the + // list of inputs and outputs is generated in a way that produces + // duplicates. + for (Subgraph& subgraph : *subgraphs_) { + // Sort and uniquefy using standard library algorithms. + auto uniquefy = [](std::vector* items) { + std::sort(items->begin(), items->end()); + auto last = std::unique(items->begin(), items->end()); + items->erase(last, items->end()); + }; + uniquefy(&subgraph.input_tensors); + uniquefy(&subgraph.output_tensors); + } + } + + private: + // Special integer values needed for tensor_epochs_ and node_epochs_. + enum { + // The node or tensor is not ready to be assigned an epoch. e.g. a node's + // inputs have not all been assigned epochs. + kEpochNotReady = -1, + // Used for tensor_epochs_. This means that the tensor is always ready. + // e.g. an input to the whole model or a constant that has no dependencies. + kEpochAlwaysReady = -2 + }; + + // Updates the node `node_index` and returns true if it is assigned to an + // epoch. False is returned if the node is already set to an epoch, its inputs + // are not all assigned to epochs, or if it cannot be assigned to the current + // epoch since the epoch's node_type doesn't match. + bool UpdateNode(int node_index) { + const TfLiteNode& node = info_->node(node_index); + Subgraph& current_subgraph = subgraphs_->back(); + int current_epoch = subgraphs_->size() - 1; + // Check if node is already done. + if (node_epochs_[node_index] != kEpochNotReady) { + return false; + } + // See if all dependencies of this node are already assigned to a + // subgraph. + for (int input_tensor_index : TfLiteIntArrayView(node.inputs)) { + if (tensor_epochs_[input_tensor_index] == kEpochNotReady) { + return false; + } + } + // When we are starting a new epoch, the first ready node defines + // the type of that epoch. + if (current_subgraph.type == Subgraph::kTfUnexplored) { + current_subgraph.type = node_type_[node_index]; + } + // The node gets assigned to this epoch if it is the same type as + // the epoch's assigned type. Note, if this is the current ready + // node encountered during this epoch, this condition will be + // automatically true. + if (current_subgraph.type == node_type_[node_index]) { + node_epochs_[node_index] = current_epoch; + current_subgraph.nodes.push_back(node_index); + // All outputs of this node now are assigned to this epoch as + // well. + for (int output_tensor_index : TfLiteIntArrayView(node.outputs)) { + tensor_epochs_[output_tensor_index] = current_epoch; + } + // Look at our inputs one more time to update that tensor's + // epochs' outputs + for (int input_tensor_index : TfLiteIntArrayView(node.inputs)) { + int input_epoch = tensor_epochs_[input_tensor_index]; + int node_epoch = current_epoch; + if (input_epoch != node_epoch) { + current_subgraph.input_tensors.push_back(input_tensor_index); + // Set inputs to be outputs of the subgraph where they reside. + // the if condition makes sure inputs to the whole computation + // are not included (i.e. those initialized to -2 above). + if (input_epoch >= 0) { + Subgraph& input_subgraph = (*subgraphs_)[input_epoch]; + input_subgraph.output_tensors.push_back(input_tensor_index); + } + } + } + return true; + } else { + return false; + } + } + + // Completely populates the current subgraph by doing graph traversal + void BuildSubgraph() { + subgraphs_->emplace_back(Subgraph()); + // loop until no more nodes can be updated. + while (true) { + bool did_something = false; + for (int node_index = 0; node_index < info_->num_nodes(); node_index++) { + if (UpdateNode(node_index)) { + did_something = true; + } + } + if (!did_something) return; + } + } + + // Temporary data needed for partitioning. + const GraphInfo* info_; + // List of subgraphs to populate + std::vector* subgraphs_; + std::vector node_type_; + // Maps from tensor index to the epoch in which it is assigned. Also special + // negative values of kEpochNotAssigned if not assigned, kEpochNotReady if it + // is an input or constant. + std::vector tensor_epochs_; + // Maps from tensor index to the epoch in which it is assigned. Also special + // negative values of kEpochNotAssigned if not assigned. + std::vector node_epochs_; +}; + +} // namespace + +TfLiteStatus PartitionGraphIntoIndependentSubgraphs( + const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, + std::vector* subgraphs) { + PartitionGraphIntoIndependentSubgraphsImpl(info, nodes_to_partition, + subgraphs) + .Partition(); + return kTfLiteOk; +} + +} // namespace tflite diff --git a/tensorflow/contrib/lite/graph_info.h b/tensorflow/contrib/lite/graph_info.h new file mode 100644 index 0000000000000000000000000000000000000000..313af5fb7574b42bcdd53b4baad06e4ccfb34053 --- /dev/null +++ b/tensorflow/contrib/lite/graph_info.h @@ -0,0 +1,79 @@ +/* Copyright 2017 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_LITE_GRAPH_INFO_H_ +#define TENSORFLOW_CONTRIB_LITE_GRAPH_INFO_H_ + +#include + +#include "tensorflow/contrib/lite/context.h" + +namespace tflite { + +// Basic information about an inference graph, where execution nodes +// are connected via tensors. +class GraphInfo { + public: + virtual ~GraphInfo() {} + + // Total number of tensors in the graph. + virtual size_t num_tensors() const = 0; + + // Returns a tensor given its index which is expected to be between 0 and + // num_tensors(). + virtual TfLiteTensor* tensor(size_t index) = 0; + + // Total number of nodes in the graph. + virtual size_t num_nodes() const = 0; + + // Returns a node given its index which is expected to be between 0 and + // num_nodes(). + virtual const TfLiteNode& node(size_t index) const = 0; + + // Returns the indices of the input tensors. + virtual const std::vector& inputs() const = 0; + + // Returns the indices of the output tensors. + virtual const std::vector& outputs() const = 0; +}; + +// Represents a subgraph of a TensorFlow Lite graph. +struct Subgraph { + enum Type { + kTfUnexplored = 0, // temporarily used during creation + kTfPartition, + kTfNonPartition + }; + Type type = kTfUnexplored; + // Nodes within the subgraph + std::vector nodes; + // Tensors that stride output from another subgraph that this depends on, + // or global inputs to the TensorFlow Lite full graph. + std::vector input_tensors; + // Outputs that are consumed by other subgraphs or are global output tensors. + // All output tensors of the nodes in the subgraph that do not appear in this + // list are intermediate results that can be potentially elided. + std::vector output_tensors; +}; + +// Partitions a list of node indices `nodes_to_partition` into subgraphs. +// Each subgraph is in dependency order (i.e. all members of the subgraph). +// `subgraphs` is assumed to be empty. +TfLiteStatus PartitionGraphIntoIndependentSubgraphs( + const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, + std::vector* subgraphs); + +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_GRAPH_INFO_H_ diff --git a/tensorflow/contrib/lite/graph_info_test.cc b/tensorflow/contrib/lite/graph_info_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..ea38b43993fef71c6820c7a978351d92d5420287 --- /dev/null +++ b/tensorflow/contrib/lite/graph_info_test.cc @@ -0,0 +1,270 @@ +/* Copyright 2017 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 +#include + +#include "tensorflow/contrib/lite/graph_info.h" +#include "tensorflow/contrib/lite/testing/util.h" + +namespace tflite { +namespace { + +// Makes a TfLiteIntArray* from std::vector, must free with TfLiteIntFree(). +TfLiteIntArray* ConvertVector(const std::vector& x) { + TfLiteIntArray* lite = TfLiteIntArrayCreate(x.size()); + for (size_t i = 0; i < x.size(); i++) lite->data[i] = x[i]; + return lite; +} + +// A very simple test graph that supports setting in/out tensors on nodes. +class SimpleTestGraph : public GraphInfo { + public: + ~SimpleTestGraph() override { + for (auto& node : nodes_) { + TfLiteIntArrayFree(node.inputs); + TfLiteIntArrayFree(node.outputs); + } + } + + size_t num_tensors() const override { return tensors_.size(); } + size_t num_nodes() const override { return nodes_.size(); } + const TfLiteNode& node(size_t index) const override { return nodes_[index]; } + TfLiteTensor* tensor(size_t index) override { return &tensors_[index]; } + const std::vector& inputs() const override { return inputs_; } + const std::vector& outputs() const override { return outputs_; } + + void AddNode(const std::vector& inputs, + const std::vector& outputs) { + nodes_.push_back(TfLiteNode()); + TfLiteNode& node = nodes_.back(); + node.inputs = ConvertVector(inputs); + node.outputs = ConvertVector(outputs); + } + + void AddTensors(int count) { tensors_.resize(count + tensors_.size()); } + + void SetInputsAndOutputs(const std::vector& inputs, + const std::vector& outputs) { + inputs_ = inputs; + outputs_ = outputs; + } + + private: + std::vector nodes_; + std::vector tensors_; + std::vector inputs_; + std::vector outputs_; +}; + +// Partition a graph to generate a list of subgraphs. This wraps the API call +// we are testing and handles memory management and conversion to +// TfLiteIntArray. Populates `subgraphs` with resulting generated subgraphs. +void PartitionGraph(const SimpleTestGraph& graph, + const std::vector& nodes_to_partition, + std::vector* subgraphs) { + TfLiteIntArray* nodes_to_partition_int_array = + ConvertVector(nodes_to_partition); + PartitionGraphIntoIndependentSubgraphs(&graph, nodes_to_partition_int_array, + subgraphs); + TfLiteIntArrayFree(nodes_to_partition_int_array); +} + +// Check a generated list of subgraphs against the expected list of subgraphs. +void CheckPartitionSubgraphs(const std::vector& generated_subgraphs, + const std::vector& expected_subgraphs) { + ASSERT_EQ(generated_subgraphs.size(), expected_subgraphs.size()); + for (int subgraph_index = 0; subgraph_index < generated_subgraphs.size(); + subgraph_index++) { + EXPECT_EQ(generated_subgraphs[subgraph_index].nodes, + expected_subgraphs[subgraph_index].nodes); + EXPECT_EQ(generated_subgraphs[subgraph_index].input_tensors, + expected_subgraphs[subgraph_index].input_tensors); + EXPECT_EQ(generated_subgraphs[subgraph_index].output_tensors, + expected_subgraphs[subgraph_index].output_tensors); + } +} + +// Test an empty trivial graph with no partitions. +TEST(PartitionTest, Nodes0_PartitionNodes0) { + SimpleTestGraph graph; + std::vector nodes_to_partition = {}; + std::vector generated_subgraphs; + PartitionGraph(graph, nodes_to_partition, &generated_subgraphs); + CheckPartitionSubgraphs(generated_subgraphs, {}); +} + +// Test a 1 node graph with no partitions. +// Input: tensor(0) -> node(0) -> tensor(1), nodes_to_partition=[] +// Output: [kTfNoPartition, tensor(0) -> node(0) -> tensor(1)] +TEST(PartitionTest, Nodes1PartitionNodes0) { + SimpleTestGraph graph; + graph.AddTensors(2); + graph.AddNode({0}, {1}); + graph.SetInputsAndOutputs({0}, {1}); + std::vector nodes_to_partition = {}; + std::vector generated_subgraphs; + PartitionGraph(graph, nodes_to_partition, &generated_subgraphs); + + Subgraph expected_subgraph; + expected_subgraph.type = Subgraph::kTfNonPartition; + expected_subgraph.nodes = {0}; + expected_subgraph.input_tensors = {0}; + expected_subgraph.output_tensors = {1}; + CheckPartitionSubgraphs(generated_subgraphs, {expected_subgraph}); +} + +// Test a 1 node graph with no inputs that is fully partitioned. +// Input: node(0) -> tensor(1), nodes_to_partition=[node0] +// Output: [kTfPartition, node(0) -> tensor(1)] +TEST(PartitionTest, Nodes1PartitionNodes0Inputs0) { + SimpleTestGraph graph; + graph.AddTensors(1); + graph.AddNode({}, {0}); + graph.SetInputsAndOutputs({}, {0}); + std::vector generated_subgraphs; + std::vector nodes_to_partition = {0}; + PartitionGraph(graph, nodes_to_partition, &generated_subgraphs); + + Subgraph expected_subgraph; + expected_subgraph.type = Subgraph::kTfPartition; + expected_subgraph.nodes = {0}; + expected_subgraph.input_tensors = {}; + expected_subgraph.output_tensors = {0}; + CheckPartitionSubgraphs(generated_subgraphs, {expected_subgraph}); +} + +// Test a 1 node graph that is partitioned completely. +// Input: tensor(0) -> node(0) -> tensor(1), nodes_to_partition=[node0] +// Output: [kTfPartition, tensor(0) -> node(0) -> tensor(1)] +TEST(PartitionTest, Nodes1PartitionNodes1) { + SimpleTestGraph graph; + graph.AddTensors(2); + graph.AddNode({0}, {1}); + graph.SetInputsAndOutputs({0}, {1}); + std::vector nodes_to_partition = {0}; + std::vector generated_subgraphs; + PartitionGraph(graph, nodes_to_partition, &generated_subgraphs); + + Subgraph expected_subgraph; + expected_subgraph.type = Subgraph::kTfPartition; + expected_subgraph.nodes = {0}; + expected_subgraph.input_tensors = {0}; + expected_subgraph.output_tensors = {1}; + CheckPartitionSubgraphs(generated_subgraphs, {expected_subgraph}); +} + +// Test a 2 node graph where 1 node is partitioned and the other is not. +// Input: tensor(0) -> node(0) -> tensor(1) -> node(1) -> tensor(2), +// nodes_to_partition = [1] +// Output: [kTfNonPartition, tensor(0) -> node(0) -> tensor(1), +// kTfPartition, tensor(1) -> node(1), tensor(2)] +TEST(PartitionTest, Nodes2PartitionNodes1) { + SimpleTestGraph graph; + graph.AddTensors(3); + graph.AddNode({0}, {1}); + graph.AddNode({1}, {2}); + graph.SetInputsAndOutputs({0}, {2}); + std::vector nodes_to_partition = {1}; + std::vector generated_subgraphs; + PartitionGraph(graph, nodes_to_partition, &generated_subgraphs); + + Subgraph expected_subgraph0; + expected_subgraph0.type = Subgraph::kTfPartition; + expected_subgraph0.nodes = {0}; + expected_subgraph0.input_tensors = {0}; + expected_subgraph0.output_tensors = {1}; + Subgraph expected_subgraph1; + expected_subgraph1.type = Subgraph::kTfPartition; + expected_subgraph1.nodes = {1}; + expected_subgraph1.input_tensors = {1}; + expected_subgraph1.output_tensors = {2}; + CheckPartitionSubgraphs(generated_subgraphs, + {expected_subgraph0, expected_subgraph1}); +} + +// Test a 2 node graph where both nodes are fully partitioned. +// Input: tensor(0) -> node(0) -> tensor(1) -> node(1) -> tensor(2), +// nodes_to_partition = [0, 1] +// Output: [kTfPartition, tensor(0) -> node(0) -> node(1) -> tensor(1)] +TEST(PartitionTest, Nodes2PartitionNodes2) { + SimpleTestGraph graph; + graph.AddTensors(3); + graph.AddNode({0}, {1}); + graph.AddNode({1}, {2}); + graph.SetInputsAndOutputs({0}, {2}); + std::vector nodes_to_partition = {0, 1}; + std::vector generated_subgraphs; + PartitionGraph(graph, nodes_to_partition, &generated_subgraphs); + + Subgraph expected_subgraph0; + expected_subgraph0.type = Subgraph::kTfPartition; + expected_subgraph0.nodes = {0, 1}; + expected_subgraph0.input_tensors = {0}; + expected_subgraph0.output_tensors = {2}; + CheckPartitionSubgraphs(generated_subgraphs, {expected_subgraph0}); +} + +// Test a three node model where we want to partition nodes 0 and nodes +// 2, but nodes 0 and nodes 2 cannot be in the same subgraph since node 2 +// depends on node 1 which depends on node 0. Thus, we need to produce three +// subgraphs. +// +// Input: tensor(0) -> node(0) -> tensor(1) +// tensor(1) -> node(1) -> tensor(2) +// [tensor(2), tensor(1)] -> node(2) -> tensor(3) +// nodes_to_partition = [0, 2] +// Output: [[kTfPartition, tensor(0) -> node(0) -> tensor(1), +// [kTfNonPartition, tensor(1) -> node(1) -> tensor(2)], +// [kTfPartition, [tensor(2), tensor(1)] -> node(2) -> node(3)] +TEST(PartitionTest, Nodes3PartitionNodes2) { + SimpleTestGraph graph; + graph.AddTensors(4); + graph.AddNode({0}, {1}); + graph.AddNode({1}, {2}); + graph.AddNode({1, 2}, {3}); + graph.SetInputsAndOutputs({0}, {3}); + std::vector nodes_to_partition = {0, 2}; + std::vector generated_subgraphs; + PartitionGraph(graph, nodes_to_partition, &generated_subgraphs); + + Subgraph expected_subgraph0; + expected_subgraph0.type = Subgraph::kTfPartition; + expected_subgraph0.nodes = {0}; + expected_subgraph0.input_tensors = {0}; + expected_subgraph0.output_tensors = {1}; + Subgraph expected_subgraph1; + expected_subgraph1.type = Subgraph::kTfNonPartition; + expected_subgraph1.nodes = {1}; + expected_subgraph1.input_tensors = {1}; + expected_subgraph1.output_tensors = {2}; + Subgraph expected_subgraph2; + expected_subgraph2.type = Subgraph::kTfPartition; + expected_subgraph2.nodes = {2}; + expected_subgraph2.input_tensors = {1, 2}; + expected_subgraph2.output_tensors = {3}; + CheckPartitionSubgraphs( + generated_subgraphs, + {expected_subgraph0, expected_subgraph1, expected_subgraph2}); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/interpreter.cc b/tensorflow/contrib/lite/interpreter.cc index 954e236ac8f0c8c59a9d20d62e66b3aa1164ecc1..028449211b8108d004df4d1cd8a58b4a08df6604 100644 --- a/tensorflow/contrib/lite/interpreter.cc +++ b/tensorflow/contrib/lite/interpreter.cc @@ -18,16 +18,16 @@ limitations under the License. #include #include #include +#include "tensorflow/contrib/lite/arena_planner.h" #include "tensorflow/contrib/lite/context.h" #include "tensorflow/contrib/lite/error_reporter.h" +#include "tensorflow/contrib/lite/graph_info.h" #include "tensorflow/contrib/lite/kernels/gemm_support.h" +#include "tensorflow/contrib/lite/memory_planner.h" #include "tensorflow/contrib/lite/nnapi_delegate.h" namespace { -// Memory allocation tuning -constexpr const int kDefaultArenaAlignment = 64; -constexpr const int kDefaultTensorAlignment = 4; // std::vector preallocation tuning. constexpr const int kSlotsToReserve = 128; @@ -35,10 +35,40 @@ constexpr const int kSlotsToReserve = 128; namespace tflite { +// A trivial implementation of GraphInfo around the Interpreter. +// NOTE: this interpreter info represents the subset of the +// graph that is executed according to execution plan. Thus, +// the indices are execution plan indices rather than raw node +// indices. +class InterpreterInfo : public GraphInfo { + public: + explicit InterpreterInfo(Interpreter* interpreter) + : interpreter_(interpreter) {} + + size_t num_tensors() const override { return interpreter_->tensors_size(); } + TfLiteTensor* tensor(size_t index) override { + return interpreter_->tensor(index); + } + size_t num_nodes() const override { + return interpreter_->execution_plan().size(); + } + const TfLiteNode& node(size_t index) const override { + int node_index = interpreter_->execution_plan()[index]; + return interpreter_->node_and_registration(node_index)->first; + } + const std::vector& inputs() const override { + return interpreter_->inputs(); + } + const std::vector& outputs() const override { + return interpreter_->outputs(); + } + + public: + Interpreter* interpreter_; +}; + Interpreter::Interpreter(ErrorReporter* error_reporter) - : arena_(kDefaultArenaAlignment), - persistent_arena_(kDefaultArenaAlignment), - error_reporter_(error_reporter ? error_reporter + : error_reporter_(error_reporter ? error_reporter : DefaultErrorReporter()) { context_.impl_ = static_cast(this); context_.ResizeTensor = ResizeTensor; @@ -47,10 +77,16 @@ Interpreter::Interpreter(ErrorReporter* error_reporter) context_.tensors = nullptr; context_.tensors_size = 0; context_.gemm_context = nullptr; + + // Invalid to call these these except from TfLiteDelegate + context_.GetNodeAndRegistration = nullptr; + context_.ReplaceSubgraphsWithDelegateKernels = nullptr; + context_.GetExecutionPlan = nullptr; + // Reserve some space for the tensors to avoid excessive resizing. tensors_.reserve(kSlotsToReserve); nodes_and_registration_.reserve(kSlotsToReserve); - next_allocate_node_id_ = 0; + next_execution_plan_index_to_prepare_ = 0; UseNNAPI(false); } @@ -70,6 +106,78 @@ Interpreter::~Interpreter() { } } +TfLiteStatus Interpreter::ReplaceSubgraphsWithDelegateKernels( + TfLiteContext* context, TfLiteRegistration registration, + const TfLiteIntArray* nodes_to_replace) { + return static_cast(context->impl_) + ->ReplaceSubgraphsWithDelegateKernels(registration, nodes_to_replace); +} + +TfLiteStatus Interpreter::ReplaceSubgraphsWithDelegateKernels( + TfLiteRegistration registration, const TfLiteIntArray* nodes_to_replace) { + // Analyze the graph to find all independent subgraphs that are either + // fully not-this-delegate or this-delegate computation. + InterpreterInfo info(this); + std::vector subgraphs; + PartitionGraphIntoIndependentSubgraphs(&info, nodes_to_replace, &subgraphs); + + execution_plan_.clear(); + for (auto& subgraph : subgraphs) { + // Turn subgraph.nodes into a TfLiteIntArray compatible data structure. + // TODO(aselle): Avoid this copy by constructing subgraph.nodes that way + // in the first place + subgraph.nodes.insert(subgraph.nodes.begin(), + static_cast(subgraph.nodes.size())); + // Subgraphs calimed by the delegate should have a "macro" op created, the + // other subgraphs (kTfNonPartition) just have their nodes added back to + // the execution plan. + switch (subgraph.type) { + case Subgraph::kTfNonPartition: + for (auto it = subgraph.nodes.begin() + 1; it != subgraph.nodes.end(); + ++it) { + execution_plan_.push_back(*it); + } + break; + case Subgraph::kTfPartition: { + void* builtin_data = nullptr; + int node_index; + // Create a node that represents computation of this subgraph. + AddNodeWithParameters( + subgraph.input_tensors, subgraph.output_tensors, + reinterpret_cast(subgraph.nodes.data()), + subgraph.nodes.size() * sizeof(subgraph.nodes[0]), builtin_data, + ®istration, &node_index); + } break; + case Subgraph::kTfUnexplored: + return kTfLiteError; + break; + } + } + return kTfLiteOk; +} + +// Gets an TfLiteIntArray* representing the execution plan. The interpreter owns +// this memory and it is only guaranteed to exist during the invocation of the +// delegate prepare. +TfLiteStatus Interpreter::GetExecutionPlan(TfLiteIntArray** execution_plan) { + // TODO(aselle): Do not make a copy here + plan_cache_.reset(TfLiteIntArrayCreate(execution_plan_.size())); + *execution_plan = plan_cache_.get(); + static_assert(sizeof(plan_cache_->data[0]) == sizeof(execution_plan_[0]), + "TfLiteIntArray and execution_plan do not contain same type."); + memcpy(plan_cache_->data, execution_plan_.data(), + sizeof(plan_cache_->data[0]) * execution_plan_.size()); + return kTfLiteOk; +} + +// WARNING: This is an experimental interface that is subject to change. +// Entry point for C node plugin API to get the execution plan +TfLiteStatus Interpreter::GetExecutionPlan(struct TfLiteContext* context, + TfLiteIntArray** execution_plan) { + return static_cast(context->impl_) + ->GetExecutionPlan(execution_plan); +} + TfLiteStatus Interpreter::SetInputs(std::vector inputs) { TF_LITE_ENSURE_OK(&context_, CheckTensorIndices("inputs", inputs.data(), inputs.size())); @@ -128,181 +236,6 @@ TfLiteStatus Interpreter::BytesRequired(TfLiteType type, const int* dims, return kTfLiteOk; } -TfLiteStatus Interpreter::AllocateTensorsWhoseSizesAreKnown() { - if (!consistent_) { - ReportError(&context_, "AllocateTensors() called on inconsistent model."); - return kTfLiteError; - } - if (next_allocate_node_id_ == nodes_and_registration_.size() && invokable_) { - return kTfLiteOk; - } - allocs_and_refcounts_.resize(context_.tensors_size); - - int new_next_allocate_node_id = next_allocate_node_id_; - invokable_ = false; - - // Allocate graph input nodes. - if (next_allocate_node_id_ == 0) { - for (int i = 0; i < inputs_.size(); ++i) { - int tensor_index = inputs_[i]; - if (tensor_index == kOptionalTensor) { - continue; - } - TfLiteTensor& tensor = context_.tensors[tensor_index]; - if (tensor.allocation_type == kTfLiteArenaRw) { - TF_LITE_ENSURE_OK( - &context_, - arena_.Allocate(&context_, kDefaultTensorAlignment, tensor.bytes, - &allocs_and_refcounts_[tensor_index].alloc)); - } - } - // Add 1 to output tensors, so they will not get overwritten. - for (int i = 0; i < outputs_.size(); ++i) { - allocs_and_refcounts_[outputs_[i]].count++; - } - } - - // Count references to node input tensors, and resize node-referenced tensors - // until we encounter a node that has a dynamic output tensor. - for (int k = next_allocate_node_id_; k < nodes_and_registration_.size(); - k++) { - new_next_allocate_node_id++; - TfLiteNode& node = nodes_and_registration_[k].first; - const TfLiteRegistration& registration = nodes_and_registration_[k].second; - if (OpPrepare(registration, &node) == kTfLiteError) { - return kTfLiteError; - } - - TfLiteIntArray* node_inputs = node.inputs; - for (int i = 0; i < node_inputs->size; ++i) { - int tensor_index = node_inputs->data[i]; - if (tensor_index != kOptionalTensor) { - allocs_and_refcounts_[node_inputs->data[i]].count++; - } - } - - // Discontinue if the node has dynamic outputs. - bool has_unallocated_dynamic_tensor = false; - TfLiteIntArray* node_outputs = node.outputs; - for (int i = 0; i < node_outputs->size; ++i) { - TfLiteTensor& tensor = context_.tensors[node_outputs->data[i]]; - if (tensor.allocation_type == kTfLiteDynamic) { - has_unallocated_dynamic_tensor = true; - break; - } - } - if (has_unallocated_dynamic_tensor) { - break; - } - } - - // Allocate graph persistent outputs, e.g. RNN cell states, etc. - for (int k = next_allocate_node_id_; k < new_next_allocate_node_id; k++) { - TfLiteNode& node = nodes_and_registration_[k].first; - - // Go through output tensors and allocate the persistent ones first. - TfLiteIntArray* node_outputs = node.outputs; - for (int i = 0; i < node_outputs->size; ++i) { - int tensor_index = node_outputs->data[i]; - TfLiteTensor& tensor = context_.tensors[tensor_index]; - if (tensor.allocation_type == kTfLiteArenaRwPersistent) { - TF_LITE_ENSURE_OK(&context_, - persistent_arena_.Allocate( - &context_, kDefaultTensorAlignment, tensor.bytes, - &allocs_and_refcounts_[tensor_index].alloc)); - } - } - } - - // Go through the graph in execution order. - for (int k = next_allocate_node_id_; k < new_next_allocate_node_id; k++) { - TfLiteNode& node = nodes_and_registration_[k].first; - - // First allocate output tensors. - TfLiteIntArray* node_outputs = node.outputs; - for (int i = 0; i < node_outputs->size; ++i) { - int tensor_index = node_outputs->data[i]; - TfLiteTensor& tensor = context_.tensors[tensor_index]; - if (tensor.allocation_type == kTfLiteArenaRw) { - TF_LITE_ENSURE_OK( - &context_, - arena_.Allocate(&context_, kDefaultTensorAlignment, tensor.bytes, - &allocs_and_refcounts_[tensor_index].alloc)); - } - } - // Then the temporaries, in two passes. First allocate them all, them - // deallocate them. - TfLiteIntArray* node_temporaries = node.temporaries; - for (int i = 0; i < node_temporaries->size; ++i) { - int tensor_index = node_temporaries->data[i]; - TfLiteTensor& tensor = context_.tensors[tensor_index]; - if (tensor.allocation_type == kTfLiteArenaRw) { - TF_LITE_ENSURE_OK( - &context_, - arena_.Allocate(&context_, kDefaultTensorAlignment, tensor.bytes, - &allocs_and_refcounts_[tensor_index].alloc)); - } - } - for (int i = 0; i < node_temporaries->size; ++i) { - int tensor_index = node_temporaries->data[i]; - TfLiteTensor& tensor = context_.tensors[tensor_index]; - allocs_and_refcounts_[tensor_index].count--; - if (tensor.allocation_type == kTfLiteArenaRw && - allocs_and_refcounts_[tensor_index].count == 0) { - TF_LITE_ENSURE_OK( - &context_, - arena_.Deallocate(&context_, - allocs_and_refcounts_[tensor_index].alloc)); - } - } - - // Then process the node's inputs. - TfLiteIntArray* node_inputs = node.inputs; - for (int i = 0; i < node_inputs->size; ++i) { - int tensor_index = node_inputs->data[i]; - if (tensor_index == kOptionalTensor) { - continue; - } - TfLiteTensor& tensor = context_.tensors[tensor_index]; - - // Decrease reference count and deallocate if not needed anymore. - allocs_and_refcounts_[tensor_index].count--; - if (tensor.allocation_type == kTfLiteArenaRw && - allocs_and_refcounts_[tensor_index].count == 0) { - TF_LITE_ENSURE_OK( - &context_, - arena_.Deallocate(&context_, - allocs_and_refcounts_[tensor_index].alloc)); - } - } - } - - // Resize the buffer and commit the arena. - TF_LITE_ENSURE_OK(&context_, arena_.Commit(&context_)); - TF_LITE_ENSURE_OK(&context_, persistent_arena_.Commit(&context_)); - - // Rewire the tensors to use the underlying arena buffer. - for (int i = 0; i < context_.tensors_size; ++i) { - TfLiteTensor& tensor = context_.tensors[i]; - if (tensor.allocation_type == kTfLiteArenaRw) { - TF_LITE_ENSURE_OK( - &context_, - arena_.ResolveAlloc(&context_, allocs_and_refcounts_[i].alloc, - &tensor.data.raw)); - } - if (tensor.allocation_type == kTfLiteArenaRwPersistent) { - TF_LITE_ENSURE_OK( - &context_, - persistent_arena_.ResolveAlloc( - &context_, allocs_and_refcounts_[i].alloc, &tensor.data.raw)); - } - } - - invokable_ = true; - next_allocate_node_id_ = new_next_allocate_node_id; - return kTfLiteOk; -} - namespace { TfLiteIntArray* convertVectorToTfLiteIntArray(const std::vector& x) { TfLiteIntArray* lite = TfLiteIntArrayCreate(x.size()); @@ -312,11 +245,19 @@ TfLiteIntArray* convertVectorToTfLiteIntArray(const std::vector& x) { } // namespace TfLiteStatus Interpreter::AllocateTensors() { - next_allocate_node_id_ = 0; - TF_LITE_ENSURE_OK(&context_, arena_.Clear()); - TF_LITE_ENSURE_OK(&context_, persistent_arena_.Clear()); - allocs_and_refcounts_.clear(); - return AllocateTensorsWhoseSizesAreKnown(); + next_execution_plan_index_to_prepare_ = 0; + if (memory_planner_) { + TF_LITE_ENSURE_STATUS(memory_planner_->ResetAllocations()); + } + + if (!consistent_) { + ReportError(&context_, "AllocateTensors() called on inconsistent model."); + return kTfLiteError; + } + + TF_LITE_ENSURE_STATUS(PrepareOpsAndTensors()); + invokable_ = true; + return kTfLiteOk; } TfLiteStatus Interpreter::AddNodeWithParameters( @@ -334,8 +275,10 @@ TfLiteStatus Interpreter::AddNodeWithParameters( &context_, CheckTensorIndices("node outputs", outputs.data(), outputs.size())); - if (node_index) *node_index = nodes_and_registration_.size(); + int new_node_index = nodes_and_registration_.size(); + if (node_index) *node_index = new_node_index; nodes_and_registration_.resize(nodes_and_registration_.size() + 1); + auto& node_and_reg = nodes_and_registration_.back(); TfLiteNode& node = node_and_reg.first; if (node.inputs) TfLiteIntArrayFree(node.inputs); @@ -357,6 +300,7 @@ TfLiteStatus Interpreter::AddNodeWithParameters( } node.builtin_data = builtin_data_deleter.release(); node_and_reg.second = *registration; + execution_plan_.push_back(new_node_index); return kTfLiteOk; } @@ -372,6 +316,60 @@ TfLiteStatus Interpreter::ResizeInputTensor(int tensor_index, return ResizeTensorImpl(&context_.tensors[tensor_index], dims_lite); } +// Returns true if at least one tensor in the given list is kTfLiteDynamic. +bool HasDynamicTensor(const TfLiteContext& context, + const TfLiteIntArray* tensors) { + for (int i = 0; i < tensors->size; ++i) { + const TfLiteTensor& tensor = context.tensors[tensors->data[i]]; + if (tensor.allocation_type == kTfLiteDynamic) { + return true; + } + } + return false; +} + +TfLiteStatus Interpreter::PrepareOpsStartingAt( + int first_execution_plan_index, int* last_execution_plan_index_prepared) { + for (int execution_plan_index = first_execution_plan_index; + execution_plan_index < execution_plan_.size(); execution_plan_index++) { + int node_index = execution_plan_[execution_plan_index]; + TfLiteNode& node = nodes_and_registration_[node_index].first; + const TfLiteRegistration& registration = + nodes_and_registration_[node_index].second; + if (OpPrepare(registration, &node) == kTfLiteError) { + return kTfLiteError; + } + + *last_execution_plan_index_prepared = execution_plan_index; + + // Discontinue if the node has dynamic outputs. Note that we don't + // stop for dynamic temporary tensors since they won't affect the + // sizes of other tensors in the graph. + if (HasDynamicTensor(context_, node.outputs)) { + break; + } + } + return kTfLiteOk; +} + +TfLiteStatus Interpreter::PrepareOpsAndTensors() { + if (!memory_planner_) { + memory_planner_.reset(new ArenaPlanner( + &context_, std::unique_ptr(new InterpreterInfo(this)))); + memory_planner_->PlanAllocations(); + } + + int last_exec_plan_index_prepared = 0; + + TF_LITE_ENSURE_STATUS(PrepareOpsStartingAt( + next_execution_plan_index_to_prepare_, &last_exec_plan_index_prepared)); + TF_LITE_ENSURE_STATUS(memory_planner_->ExecuteAllocations( + next_execution_plan_index_to_prepare_, last_exec_plan_index_prepared)); + + next_execution_plan_index_to_prepare_ = last_exec_plan_index_prepared + 1; + return kTfLiteOk; +} + TfLiteStatus Interpreter::Invoke() { if (!consistent_) { ReportError(&context_, "Invoke called on model that is not consistent."); @@ -384,10 +382,7 @@ TfLiteStatus Interpreter::Invoke() { TfLiteStatus status = kTfLiteOk; if (nnapi_delegate_) { - if (AllocateTensorsWhoseSizesAreKnown() == kTfLiteError) { - return kTfLiteError; - } - if (next_allocate_node_id_ == nodes_and_registration_.size()) { + if (next_execution_plan_index_to_prepare_ == execution_plan_.size()) { TF_LITE_ENSURE_OK(&context_, nnapi_delegate_->Invoke(this)); return kTfLiteOk; } else { @@ -400,17 +395,24 @@ TfLiteStatus Interpreter::Invoke() { } } - for (int i = 0; i < nodes_and_registration_.size(); i++) { - // Ensure we have allocated up to this node. The point of this is to - // allocate as much as possible before running any evaluation, but - // dynamic shapes can prevent this from being possible. - if (i >= next_allocate_node_id_) { - if (AllocateTensorsWhoseSizesAreKnown() == kTfLiteError) { - return kTfLiteError; - } + // Invocations are always done in node order. + // Note that calling Invoke repeatedly will cause the original memory plan to + // be reused, unless either ResizeInputTensor() or AllocateTensors() has been + // called. + // TODO(b/71913981): we should force recalculation in the presence of dynamic + // tensors, because they may have new value which in turn may affect shapes + // and allocations. + for (int execution_plan_index = 0; + execution_plan_index < execution_plan_.size(); execution_plan_index++) { + if (execution_plan_index == next_execution_plan_index_to_prepare_) { + TF_LITE_ENSURE_STATUS(PrepareOpsAndTensors()); + TF_LITE_ENSURE(&context_, next_execution_plan_index_to_prepare_ >= + execution_plan_index); } - TfLiteNode& node = nodes_and_registration_[i].first; - const TfLiteRegistration& registration = nodes_and_registration_[i].second; + int node_index = execution_plan_[execution_plan_index]; + TfLiteNode& node = nodes_and_registration_[node_index].first; + const TfLiteRegistration& registration = + nodes_and_registration_[node_index].second; if (OpInvoke(registration, &node) == kTfLiteError) { status = kTfLiteError; } @@ -465,6 +467,22 @@ TfLiteStatus Interpreter::AddTensors(TfLiteContext* context, int tensors_to_add, ->AddTensors(tensors_to_add, first_new_tensor_index); } +TfLiteStatus Interpreter::GetNodeAndRegistration( + int node_index, TfLiteNode** node, TfLiteRegistration** registration) { + TF_LITE_ENSURE(&context_, node_index < nodes_size() && node_index >= 0); + TF_LITE_ENSURE(&context_, node != nullptr && registration != nullptr); + *node = &nodes_and_registration_[node_index].first; + *registration = &nodes_and_registration_[node_index].second; + return kTfLiteOk; +} + +TfLiteStatus Interpreter::GetNodeAndRegistration( + struct TfLiteContext* context, int node_index, TfLiteNode** node, + TfLiteRegistration** registration) { + return static_cast(context->impl_) + ->GetNodeAndRegistration(node_index, node, registration); +} + TfLiteStatus Interpreter::SetTensorParametersReadOnly( int tensor_index, TfLiteType type, const char* name, const std::vector& dims, TfLiteQuantizationParams quantization, @@ -514,6 +532,14 @@ TfLiteStatus Interpreter::SetTensorParametersReadWrite( return kTfLiteOk; } +TfLiteStatus Interpreter::SetExecutionPlan(const std::vector& new_plan) { + for (int node_index : new_plan) { + TF_LITE_ENSURE(&context_, node_index >= 0 && node_index < nodes_size()); + } + execution_plan_ = new_plan; + return kTfLiteOk; +} + TfLiteStatus Interpreter::ResizeTensorImpl(TfLiteTensor* tensor, TfLiteIntArray* new_size) { // Note that in theory we could resize kTfLiteArenaRwPersistent tensors too. @@ -527,6 +553,9 @@ TfLiteStatus Interpreter::ResizeTensorImpl(TfLiteTensor* tensor, TfLiteIntArrayFree(new_size); return kTfLiteError; } + + // Realloc space for kTfLiteDynamic tensors. + TfLiteTensorRealloc(bytesRequired, tensor); tensor->bytes = bytesRequired; } if (tensor->dims) TfLiteIntArrayFree(tensor->dims); @@ -564,4 +593,20 @@ void Interpreter::SetNumThreads(int num_threads) { tflite::gemm_support::SetMaxNumThreads(&context_, num_threads); } +TfLiteStatus Interpreter::ModifyGraphWithDelegate(TfLiteDelegate* delegate) { + // TODO(aselle): Consider if it is worth storing pointers to delegates. + // Setup additional context interface + context_.GetNodeAndRegistration = GetNodeAndRegistration; + context_.ReplaceSubgraphsWithDelegateKernels = + ReplaceSubgraphsWithDelegateKernels; + context_.GetExecutionPlan = GetExecutionPlan; + + TfLiteStatus status = delegate->Prepare(&context_, delegate->data_); + // Remove additional context info. + context_.GetNodeAndRegistration = nullptr; + context_.ReplaceSubgraphsWithDelegateKernels = nullptr; + context_.GetExecutionPlan = nullptr; + return status; +} + } // namespace tflite diff --git a/tensorflow/contrib/lite/interpreter.h b/tensorflow/contrib/lite/interpreter.h index 65c61e44bee48535f884a3afaddc691972f5e04b..bab56a9d72f8992a9d8af23f92133c7c918fd46d 100644 --- a/tensorflow/contrib/lite/interpreter.h +++ b/tensorflow/contrib/lite/interpreter.h @@ -14,8 +14,8 @@ limitations under the License. ==============================================================================*/ // Main abstraction controlling the tflite interpreter. // See context.h for the API for defining operations (TfLiteRegistration). -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ +#define TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ #include #include @@ -23,7 +23,7 @@ limitations under the License. #include "tensorflow/contrib/lite/allocation.h" #include "tensorflow/contrib/lite/context.h" #include "tensorflow/contrib/lite/error_reporter.h" -#include "tensorflow/contrib/lite/simple_memory_arena.h" +#include "tensorflow/contrib/lite/memory_planner.h" namespace tflite { @@ -49,13 +49,6 @@ constexpr TfLiteType typeToTfLiteType() { return kTfLiteUInt8; } -struct ArenaAllocRefCount { - ArenaAllocRefCount() : alloc(), count(0) {} - - ArenaAlloc alloc; - int count; -}; - // Forward declare since NNAPIDelegate uses Interpreter. class NNAPIDelegate; @@ -87,6 +80,12 @@ class NNAPIDelegate; // foo.Invoke(); // +struct TfLiteIntArrayDeleter { + void operator()(TfLiteIntArray* a) { + if (a) TfLiteIntArrayFree(a); + } +}; + class Interpreter { public: // Instantiate an interpreter. All errors associated with reading and @@ -115,7 +114,7 @@ class Interpreter { // Adds a node with the given parameters and returns the index of the new // node in `node_index` (optionally). Interpreter will take ownership of - // `builtin_data` and destroy it with `delete`. Ownership of 'init_data' + // `builtin_data` and destroy it with `free`. Ownership of 'init_data' // remains with the caller. TfLiteStatus AddNodeWithParameters(const std::vector& inputs, const std::vector& outputs, @@ -173,12 +172,19 @@ class Interpreter { // Return the number of ops in the model. int nodes_size() const { return nodes_and_registration_.size(); } + // WARNING: Experimental interface, subject to change + const std::vector& execution_plan() const { return execution_plan_; } + + // WARNING: Experimental interface, subject to change + // Overrides execution plan. This bounds checks indices sent in. + TfLiteStatus SetExecutionPlan(const std::vector& new_plan); + // Get a tensor data structure. // TODO(aselle): Create a safe ArrayHandle interface to avoid exposing this // read/write access to structure TfLiteTensor* tensor(int tensor_index) { if (tensor_index >= context_.tensors_size || tensor_index < 0) - return nullptr; + return nullptr; return &context_.tensors[tensor_index]; } @@ -247,6 +253,11 @@ class Interpreter { // Set the number of threads available to the interpreter. void SetNumThreads(int num_threads); + // Allow a delegate to look at the graph and modify the graph to handle + // parts of the graph themselves. After this is called, the graph may + // contain new nodes that replace 1 more nodes. + TfLiteStatus ModifyGraphWithDelegate(TfLiteDelegate* delegate); + private: // Give 'op_reg' a chance to initialize itself using the contents of // 'buffer'. @@ -276,9 +287,18 @@ class Interpreter { return op_reg.invoke(&context_, node); } - // Allocate tensors whose sizes are known in order of nodes. Discontinue when - // we encounter a node that has a dynamic output tensor. - TfLiteStatus AllocateTensorsWhoseSizesAreKnown(); + // Call OpPrepare() for as many ops as possible, allocating memory for their + // tensors. If an op containing dynamic tensors is found, preparation will be + // postponed until this function is called again. This allows the interpreter + // to wait until Invoke() to resolve the sizes of dynamic tensors. + TfLiteStatus PrepareOpsAndTensors(); + + // Call OpPrepare() for all ops starting at 'first_node'. Stop when a + // dynamic tensors is found or all ops have been prepared. Fill + // 'last_node_prepared' with the id of the op containing dynamic tensors, or + // the last in the graph. + TfLiteStatus PrepareOpsStartingAt(int first_execution_plan_index, + int* last_execution_plan_index_prepared); // Tensors needed by the interpreter. Use `AddTensors` to add more blank // tensor entries. Note, `tensors_.data()` needs to be synchronized to the @@ -298,7 +318,8 @@ class Interpreter { TfLiteStatus BytesRequired(TfLiteType type, const int* dims, int dims_size, size_t* bytes); - // Request an tensor be resized implementation. + // Request an tensor be resized implementation. If the given tensor is of + // type kTfLiteDynamic it will also be allocated new memory. TfLiteStatus ResizeTensorImpl(TfLiteTensor* tensor, TfLiteIntArray* new_size); // Report a detailed error string (will be printed to stderr). @@ -315,6 +336,40 @@ class Interpreter { static TfLiteStatus AddTensors(TfLiteContext* context, int tensors_to_add, int* first_new_tensor_index); + // WARNING: This is an experimental API and subject to change. + // Entry point for C API ReplaceSubgraphsWithDelegateKernels + static TfLiteStatus ReplaceSubgraphsWithDelegateKernels( + TfLiteContext* context, TfLiteRegistration registration, + const TfLiteIntArray* nodes_to_replace); + + // Update the execution graph to replace some of the nodes with stub + // nodes. Specifically any node index that has `nodes[index]==1` will be + // slated for replacement with a delegate kernel specified by registration. + // WARNING: This is an experimental interface that is subject to change. + TfLiteStatus ReplaceSubgraphsWithDelegateKernels( + TfLiteRegistration registration, const TfLiteIntArray* nodes_to_replace); + + // WARNING: This is an experimental interface that is subject to change. + // Gets the internal pointer to a TensorFlow lite node by node_index. + TfLiteStatus GetNodeAndRegistration(int node_index, TfLiteNode** node, + TfLiteRegistration** registration); + + // WARNING: This is an experimental interface that is subject to change. + // Entry point for C node plugin API to get a node by index. + static TfLiteStatus GetNodeAndRegistration(struct TfLiteContext*, + int node_index, TfLiteNode** node, + TfLiteRegistration** registration); + + // WARNING: This is an experimental interface that is subject to change. + // Gets an TfLiteIntArray* representing the execution plan. The caller owns + // this memory and must free it with TfLiteIntArrayFree(). + TfLiteStatus GetExecutionPlan(TfLiteIntArray** execution_plan); + + // WARNING: This is an experimental interface that is subject to change. + // Entry point for C node plugin API to get the execution plan + static TfLiteStatus GetExecutionPlan(struct TfLiteContext* context, + TfLiteIntArray** execution_plan); + // A pure C data structure used to communicate with the pure C plugin // interface. To avoid copying tensor metadata, this is also the definitive // structure to store tensors. @@ -325,17 +380,6 @@ class Interpreter { std::vector> nodes_and_registration_; - // Raw memory buffer that is allocated for all temporary and graph outputs. - // that are declared kTfLiteArenaRw. - SimpleMemoryArena arena_; - - // Raw memory buffer that is allocated for persistent tensors that are - // declared as kTfLiteArenaRwPersistent. - SimpleMemoryArena persistent_arena_; - - // Stores allocation and reference counts of all tensors. - std::vector allocs_and_refcounts_; - // Whether the model is consistent. That is to say if the inputs and outputs // of every node and the global inputs and outputs are valid indexes into // the tensor array. @@ -356,7 +400,7 @@ class Interpreter { // The error reporter delegate that tflite will forward queries errors to. ErrorReporter* error_reporter_; - // Next node to allocate output tensors. + // Index of the next node to prepare. // During Invoke(), Interpreter will allocate input tensors first, which are // known to be fixed size. Then it will allocate outputs from nodes as many // as possible. When there is a node that produces dynamic sized tensor. @@ -364,11 +408,24 @@ class Interpreter { // node id, and execute the node to generate the output tensor before continue // to allocate successors. This process repeats until all nodes are executed. // NOTE: this relies on the order of nodes that is in topological order. - int next_allocate_node_id_; + int next_execution_plan_index_to_prepare_; + + // WARNING: This is an experimental interface that is subject to change. + // This is a list of node indices (to index into nodes_and_registration). + // This represents a valid topological sort (dependency ordered) execution + // plan. In particular, it is valid for this ordering to contain only a + // subset of the node indices. + std::vector execution_plan_; + + // In the future, we'd like a TfLiteIntArray compatible representation. + // TODO(aselle): replace execution_plan_ with this. + std::unique_ptr plan_cache_; // Whether to delegate to NN API std::unique_ptr nnapi_delegate_; + + std::unique_ptr memory_planner_; }; } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ +#endif // TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ diff --git a/tensorflow/contrib/lite/interpreter_test.cc b/tensorflow/contrib/lite/interpreter_test.cc index edff2109430c6e1ec6c481619ed7772237a3301d..28c96e5dde6ffa62bb073db9716a00f91c6e0bdf 100644 --- a/tensorflow/contrib/lite/interpreter_test.cc +++ b/tensorflow/contrib/lite/interpreter_test.cc @@ -16,8 +16,10 @@ limitations under the License. #include "tensorflow/contrib/lite/interpreter.h" #include #include "tensorflow/contrib/lite/error_reporter.h" +#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" +#include "tensorflow/contrib/lite/schema/schema_generated.h" #include "tensorflow/contrib/lite/string_util.h" - +#include "tensorflow/contrib/lite/testing/util.h" namespace tflite { namespace { @@ -282,6 +284,51 @@ TEST(BasicInterpreter, NoOpInterpreter) { ASSERT_EQ(interpreter.Invoke(), kTfLiteOk); } +TEST(BasicInterpreter, ResizingTensors) { + Interpreter interpreter; + ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk); + ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk); + ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk); + + ASSERT_EQ(interpreter.SetTensorParametersReadWrite( + 0, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()), + kTfLiteOk); + + int t = interpreter.inputs()[0]; + TfLiteTensor* tensor = interpreter.tensor(t); + + ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 3}), kTfLiteOk); + EXPECT_EQ(tensor->bytes, 6 * sizeof(float)); + ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk); + + tensor->data.f[5] = 0.123f; + + // Changing from kTfLiteArenaRw to kTfLiteDynamic is quite complicate: we need + // to unset data.raw, otherwise Realloc will try to free that memory. + tensor->data.raw = nullptr; + tensor->allocation_type = kTfLiteDynamic; + + ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 4}), kTfLiteOk); + EXPECT_EQ(tensor->bytes, 8 * sizeof(float)); + ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk); + + // TODO(ahentz): We shouldn't have to force reallocation, but + // ResizeInputTensor doesn't realloc dynamic tensors. Also note that + // TfLiteTensorRealloc(tensor->bytes, tensor) is a no-op. + TfLiteTensorRealloc(9 * sizeof(float), tensor); + tensor->data.f[7] = 0.123f; + + ASSERT_EQ(interpreter.ResizeInputTensor(t, {2, 2, 4}), kTfLiteOk); + EXPECT_EQ(tensor->bytes, 16 * sizeof(float)); + ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk); + + // TODO(ahentz): We shouldn't have to force reallocation, but + // ResizeInputTensor doesn't realloc dynamic tensors. Also note that + // TfLiteTensorRealloc(tensor->bytes, tensor) is a no-op. + TfLiteTensorRealloc(17 * sizeof(float), tensor); + tensor->data.f[15] = 0.123f; +} + TEST(BasicInterpreter, OneOpInterpreter) { Interpreter interpreter; ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk); @@ -514,13 +561,283 @@ TEST(BasicInterpreter, TestCustomErrorReporter) { ASSERT_EQ(reporter.calls, 1); } +// Test fixture that allows playing with execution plans. It creates a two +// node graph that can be executed in either [0,1] order or [1,0] order. +// The CopyOp records when it is invoked in the class member run_order_ +// so we can test whether the execution plan was honored. +class TestExecutionPlan : public ::testing::Test { + // Encapsulates the node ids and provides them to a C primitive data type + // Allocatable with placement new, but never destructed, so make sure this + // doesn't own any heap allocated data. This is then is used as op local + // data to allow access to the test fixture data. + class CallReporting { + public: + CallReporting(int node_id, std::vector* run_order) + : node_id_(node_id), run_order_(run_order) {} + + void Record() { run_order_->push_back(node_id_); } + + private: + // The node id for this particular node + int node_id_; + // A pointer to the global run-order + std::vector* run_order_; + }; + + // Build a kernel registration for an op that copies its one input + // to an output + TfLiteRegistration CopyOpRegistration() { + TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr}; + + reg.prepare = [](TfLiteContext* context, TfLiteNode* node) { + // Set output size to input size + TfLiteTensor* tensor0 = &context->tensors[node->inputs->data[0]]; + TfLiteTensor* tensor1 = &context->tensors[node->outputs->data[0]]; + TfLiteIntArray* newSize = TfLiteIntArrayCopy(tensor0->dims); + return context->ResizeTensor(context, tensor1, newSize); + }; + + reg.invoke = [](TfLiteContext* context, TfLiteNode* node) { + CallReporting* call_reporting = + reinterpret_cast(node->builtin_data); + // Copy input data to output data. + TfLiteTensor* a0 = &context->tensors[node->inputs->data[0]]; + TfLiteTensor* a1 = &context->tensors[node->outputs->data[0]]; + int num = a0->dims->data[0]; + for (int i = 0; i < num; i++) { + a1->data.f[i] = a0->data.f[i]; + } + call_reporting->Record(); + return kTfLiteOk; + }; + return reg; + } + + // Adds a copy node going from tensor `input` to output tensor `output`. + // Note, input is used as the node_id. Inject run_order as op accessible + // data. Note: this is a little strange of a way to do this, but it is + // using op functionality to avoid static global variables. + void MakeCopyNode(int input, int output) { + // Ownership of call_reporting is taken by interpreter (malloc is used due + // to nodes being a C99 interface so free() is used). + TfLiteRegistration copy_op = CopyOpRegistration(); + CallReporting* call_reporting_1 = + reinterpret_cast(malloc(sizeof(CallReporting))); + new (call_reporting_1) CallReporting(input, &run_order_); + ASSERT_EQ(interpreter_.AddNodeWithParameters( + {0}, {2}, nullptr, 0, + reinterpret_cast(call_reporting_1), ©_op), + kTfLiteOk); + ASSERT_EQ(interpreter_.ResizeInputTensor(input, {3}), kTfLiteOk); + } + + void SetUp() final { + // Add two inputs and two outputs that don't depend on each other + ASSERT_EQ(interpreter_.AddTensors(4), kTfLiteOk); + interpreter_.SetInputs({0, 1}); + interpreter_.SetOutputs({2, 3}); + TfLiteQuantizationParams quantized; + for (int tensor_index = 0; tensor_index < 4; tensor_index++) { + ASSERT_EQ(interpreter_.SetTensorParametersReadWrite( + tensor_index, kTfLiteFloat32, "", {3}, quantized), + kTfLiteOk); + } + + // Define two copy functions that also use the user_data to report that + // they were called. + // i.e. tensor[2] = copy(tensor[0]); tensor[3] = copy(tensor[1]); + // thus we can reorder the two nodes arbitrary and still satisfy dependency + // order. + MakeCopyNode(0, 2); + MakeCopyNode(1, 3); + + ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk); + } + + protected: + Interpreter interpreter_; + + // list of node_ids that were run + std::vector run_order_; +}; + +TEST_F(TestExecutionPlan, DefaultExecutionPlan) { + // Check default order + ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk); + ASSERT_EQ(run_order_, std::vector({0, 1})); +} + +TEST_F(TestExecutionPlan, ReversedExecutionPlan) { + // Check reversed order + interpreter_.SetExecutionPlan({1, 0}); + ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk); + ASSERT_EQ(run_order_, std::vector({1, 0})); +} + +TEST_F(TestExecutionPlan, SubsetExecutionPlan) { + // Check running only node index 1 + interpreter_.SetExecutionPlan({1}); + ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk); + ASSERT_EQ(run_order_, std::vector({1})); +} + +TEST_F(TestExecutionPlan, NullExecutionPlan) { + // Check nothing executed. + interpreter_.SetExecutionPlan({}); + ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk); + ASSERT_EQ(run_order_, std::vector()); +} + +// Build a kernel registration for an op that copies its one input +// to an output +TfLiteRegistration AddOpRegistration() { + TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr}; + + reg.custom_name = "my_add"; + reg.builtin_code = tflite::BuiltinOperator_CUSTOM; + + reg.prepare = [](TfLiteContext* context, TfLiteNode* node) { + // Set output size to input size + TfLiteTensor* tensor0 = &context->tensors[node->inputs->data[0]]; + TfLiteTensor* tensor1 = &context->tensors[node->inputs->data[1]]; + TfLiteTensor* tensor2 = &context->tensors[node->outputs->data[0]]; + TfLiteIntArray* newSize = TfLiteIntArrayCopy(tensor0->dims); + TfLiteIntArray* newSizeOther = TfLiteIntArrayCopy(tensor1->dims); + TF_LITE_ENSURE_EQ(context, newSize->size, newSizeOther->size); + TF_LITE_ENSURE_STATUS(context->ResizeTensor(context, tensor2, newSize)); + return kTfLiteOk; + }; + + reg.invoke = [](TfLiteContext* context, TfLiteNode* node) { + // Copy input data to output data. + TfLiteTensor* a0 = &context->tensors[node->inputs->data[0]]; + TfLiteTensor* a1 = &context->tensors[node->inputs->data[1]]; + TfLiteTensor* out = &context->tensors[node->outputs->data[0]]; + int num = a0->dims->data[0]; + for (int i = 0; i < num; i++) { + out->data.f[i] = a0->data.f[i] + a1->data.f[i]; + } + return kTfLiteOk; + }; + return reg; +} + +class TestDelegate : public ::testing::Test { + public: + TestDelegate() { + interpreter_.AddTensors(5); + interpreter_.SetInputs({0, 1}); + interpreter_.SetOutputs({3, 4}); + TfLiteQuantizationParams quant; + interpreter_.SetTensorParametersReadWrite(0, kTfLiteFloat32, "", {3}, + quant); + interpreter_.SetTensorParametersReadWrite(1, kTfLiteFloat32, "", {3}, + quant); + interpreter_.SetTensorParametersReadWrite(2, kTfLiteFloat32, "", {3}, + quant); + interpreter_.SetTensorParametersReadWrite(3, kTfLiteFloat32, "", {3}, + quant); + TfLiteRegistration reg = AddOpRegistration(); + interpreter_.AddNodeWithParameters({0, 0}, {2}, nullptr, 0, nullptr, ®); + interpreter_.AddNodeWithParameters({1, 1}, {3}, nullptr, 0, nullptr, ®); + interpreter_.AddNodeWithParameters({2, 1}, {4}, nullptr, 0, nullptr, ®); + } + + protected: + class SimpleDelegate { + public: + // Create a simple implementation of a TfLiteDelegate. We use the C++ class + // SimpleDelegate and it can produce a handle TfLiteDelegate that is + // value-copyable and compatible with TfLite. + explicit SimpleDelegate(const std::vector& nodes) : nodes_(nodes) { + delegate_.Prepare = [](TfLiteContext* context, + void* data) -> TfLiteStatus { + auto* simple = reinterpret_cast(data); + TfLiteIntArray* nodes_to_separate = + TfLiteIntArrayCreate(simple->nodes_.size()); + // Mark nodes that we want in TfLiteIntArray* structure. + int index = 0; + for (auto node_index : simple->nodes_) { + nodes_to_separate->data[index++] = node_index; + // make sure node is add + TfLiteNode* node; + TfLiteRegistration* reg; + context->GetNodeAndRegistration(context, node_index, &node, ®); + TFLITE_CHECK_EQ(reg->builtin_code, tflite::BuiltinOperator_CUSTOM); + TFLITE_CHECK_EQ(strcmp(reg->custom_name, "my_add"), 0); + } + // Check that all nodes are available + TfLiteIntArray* execution_plan; + TF_LITE_ENSURE_STATUS( + context->GetExecutionPlan(context, &execution_plan)); + for (int exec_index = 0; exec_index < execution_plan->size; + exec_index++) { + int node_index = execution_plan->data[exec_index]; + // Check that we are an identity map to start. + TFLITE_CHECK_EQ(exec_index, node_index); + TfLiteNode* node; + TfLiteRegistration* reg; + context->GetNodeAndRegistration(context, node_index, &node, ®); + TFLITE_CHECK_EQ(reg->builtin_code, tflite::BuiltinOperator_CUSTOM); + TFLITE_CHECK_EQ(strcmp(reg->custom_name, "my_add"), 0); + } + + context->ReplaceSubgraphsWithDelegateKernels( + context, FakeFusedRegistration(), nodes_to_separate); + TfLiteIntArrayFree(nodes_to_separate); + return kTfLiteOk; + }; + // Store type-punned data SimpleDelegate structure. + delegate_.data_ = reinterpret_cast(this); + } + + static TfLiteRegistration FakeFusedRegistration() { + TfLiteRegistration reg = {nullptr}; + reg.custom_name = "fake_fused_op"; + return reg; + } + + TfLiteDelegate* get_tf_lite_delegate() { return &delegate_; } + + private: + std::vector nodes_; + TfLiteDelegate delegate_; + }; + Interpreter interpreter_; +}; + +TEST_F(TestDelegate, BasicDelegate) { + interpreter_.Invoke(); + SimpleDelegate simple({0, 1, 2}); + interpreter_.ModifyGraphWithDelegate(simple.get_tf_lite_delegate()); + + ASSERT_EQ(interpreter_.execution_plan().size(), 1); + int node = interpreter_.execution_plan()[0]; + const auto* node_and_reg = interpreter_.node_and_registration(node); + ASSERT_EQ(node_and_reg->second.custom_name, + SimpleDelegate::FakeFusedRegistration().custom_name); +} + +TEST_F(TestDelegate, ComplexDeligate) { + interpreter_.Invoke(); + SimpleDelegate simple({1, 2}); + interpreter_.ModifyGraphWithDelegate(simple.get_tf_lite_delegate()); + + ASSERT_EQ(interpreter_.execution_plan().size(), 2); + // 0th should be a non-delegated original op + ASSERT_EQ(interpreter_.execution_plan()[0], 0); + // 1st should be a new macro op (3) which didn't exist) + ASSERT_EQ(interpreter_.execution_plan()[1], 3); + const auto* node_and_reg = interpreter_.node_and_registration(3); + ASSERT_EQ(node_and_reg->second.custom_name, + SimpleDelegate::FakeFusedRegistration().custom_name); +} + } // namespace } // namespace tflite int main(int argc, char** argv) { -#ifdef OS_LINUX - FLAGS_logtostderr = true; -#endif + ::tflite::LogToStderr(); ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); } diff --git a/tensorflow/contrib/lite/ios_makefile.inc b/tensorflow/contrib/lite/ios_makefile.inc index bcff7ed9889e95c13294b6cf0d0f4788991a04df..fc6594c3a04ba6aabba99bb631f85737baf389f1 100644 --- a/tensorflow/contrib/lite/ios_makefile.inc +++ b/tensorflow/contrib/lite/ios_makefile.inc @@ -22,6 +22,7 @@ ifeq ($(TARGET), IOS) IOS_ARCH := x86_64 CXXFLAGS += -miphoneos-version-min=$(MIN_SDK_VERSION) \ -DGEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK \ + -DTFLITE_USE_APPLE_ACCELERATE_FOR_CONV \ -fembed-bitcode \ -Wno-c++11-narrowing \ -mno-thumb \ @@ -30,6 +31,9 @@ ifeq ($(TARGET), IOS) ${IPHONEOS_SYSROOT} \ -arch $(IOS_ARCH) \ -O3 + ifeq ($(IOS_ARCH), x86_64) + CXXFLAGS += -msse4.1 + endif CCFLAGS += -miphoneos-version-min=$(MIN_SDK_VERSION) \ -fembed-bitcode \ -mno-thumb \ @@ -39,6 +43,7 @@ ifeq ($(TARGET), IOS) -O3 LDFLAGS := -fembed-bitcode \ -miphoneos-version-min=${MIN_SDK_VERSION} \ + -framework Accelerate \ -arch $(IOS_ARCH) OBJDIR := $(OBJDIR)ios_$(IOS_ARCH)/ LIBDIR := $(LIBDIR)ios_$(IOS_ARCH)/ diff --git a/tensorflow/contrib/lite/java/AndroidManifest.xml b/tensorflow/contrib/lite/java/AndroidManifest.xml new file mode 100644 index 0000000000000000000000000000000000000000..f705feacbec38ab5152ce52b701320d8f1cd8d3d --- /dev/null +++ b/tensorflow/contrib/lite/java/AndroidManifest.xml @@ -0,0 +1,7 @@ + + + + + + diff --git a/tensorflow/contrib/lite/java/BUILD b/tensorflow/contrib/lite/java/BUILD index 1de28eb52ddb458df0be0a8f9ef453f7caf68654..35aacb70002d1d454f675484e4398bcdffc4acf1 100644 --- a/tensorflow/contrib/lite/java/BUILD +++ b/tensorflow/contrib/lite/java/BUILD @@ -7,6 +7,16 @@ licenses(["notice"]) # Apache 2.0 load("//tensorflow/java:build_defs.bzl", "JAVACOPTS") load("//tensorflow/contrib/lite:build_def.bzl", "tflite_jni_binary") +load("//tensorflow/contrib/lite/java:aar_with_jni.bzl", "aar_with_jni") + +# Building tensorflow-lite.aar including 4 variants of .so +# To build an aar for release, run below command: +# bazel build --cxxopt='--std=c++11' -c opt --fat_apk_cpu=x86,x86_64,arm64-v8a,armeabi-v7a \ +# tensorflow/contrib/lite/java:tensorflow-lite +aar_with_jni( + name = "tensorflow-lite", + android_library = ":tensorflowlite", +) android_library( name = "tensorflowlite", @@ -15,6 +25,7 @@ android_library( "src/main/java/org/tensorflow/lite/*.java", ], ), + manifest = "AndroidManifest.xml", visibility = ["//visibility:public"], deps = [ ":tflite_runtime", @@ -100,6 +111,26 @@ java_test( ], ) +# TODO: generate large models at runtime, instead of storing them. +java_test( + name = "InterpreterTest", + size = "small", + srcs = ["src/test/java/org/tensorflow/lite/InterpreterTest.java"], + data = [ + "src/testdata/add.bin", + "src/testdata/mobilenet.tflite.bin", + ], + javacopts = JAVACOPTS, + test_class = "org.tensorflow.lite.InterpreterTest", + visibility = ["//visibility:private"], + deps = [ + ":libtensorflowlite_jni.so", + ":tensorflowlitelib", + "@com_google_truth", + "@junit", + ], +) + java_test( name = "TensorTest", size = "small", diff --git a/tensorflow/contrib/lite/java/aar_with_jni.bzl b/tensorflow/contrib/lite/java/aar_with_jni.bzl new file mode 100644 index 0000000000000000000000000000000000000000..4450bc9085555b3416f51bac07ea94a1240e919c --- /dev/null +++ b/tensorflow/contrib/lite/java/aar_with_jni.bzl @@ -0,0 +1,47 @@ +"""Generate zipped aar file including different variants of .so in jni folder.""" + +def aar_with_jni(name, android_library): + # Generate dummy AndroidManifest.xml for dummy apk usage + # (dummy apk is generated by _dummy_app_for_so target below) + native.genrule( + name = name + "_binary_manifest_generator", + outs = [name + "_generated_AndroidManifest.xml"], + cmd = """ +cat > $(OUTS) < + + +EOF +""", + ) + + # Generate dummy apk including .so files and later we extract out + # .so files and throw away the apk. + native.android_binary( + name = name + "_dummy_app_for_so", + manifest = name + "_generated_AndroidManifest.xml", + custom_package = "dummy.package.for.so", + deps = [android_library], + # In some platforms we don't have an Android SDK/NDK and this target + # can't be built. We need to prevent the build system from trying to + # use the target in that case. + tags = ["manual"], + ) + + native.genrule( + name = name, + srcs = [android_library + ".aar", name + "_dummy_app_for_so_unsigned.apk"], + outs = [name + ".aar"], + tags = ["manual"], + cmd = """ +cp $(location {}.aar) $(location :{}.aar) +chmod +w $(location :{}.aar) +origdir=$$PWD +cd $$(mktemp -d) +unzip $$origdir/$(location :{}_dummy_app_for_so_unsigned.apk) "lib/*" +cp -r lib jni +zip -r $$origdir/$(location :{}.aar) jni/*/*.so +""".format(android_library, name, name, name, name), + ) diff --git a/tensorflow/contrib/lite/java/build_aar_for_release.sh b/tensorflow/contrib/lite/java/build_aar_for_release.sh new file mode 100755 index 0000000000000000000000000000000000000000..fbcb1e7db9a3f9b885505e989b7ff7224f2d2b15 --- /dev/null +++ b/tensorflow/contrib/lite/java/build_aar_for_release.sh @@ -0,0 +1,66 @@ +#!/bin/bash +# Copyright 2017 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. +# ============================================================================== +set -e +set -x + +TMPDIR=`mktemp -d` +trap "rm -rf $TMPDIR" EXIT + +VERSION=1.0 + +BUILDER=bazel +BASEDIR=tensorflow/contrib/lite +CROSSTOOL="//external:android/crosstool" +HOST_CROSSTOOL="@bazel_tools//tools/cpp:toolchain" + +BUILD_OPTS="--cxxopt=--std=c++11 -c opt" +CROSSTOOL_OPTS="--crosstool_top=$CROSSTOOL --host_crosstool_top=$HOST_CROSSTOOL" + +test -d $BASEDIR || (echo "Aborting: not at top-level build directory"; exit 1) + +function build_basic_aar() { + local OUTDIR=$1 + $BUILDER build $BUILD_OPTS $BASEDIR/java:tensorflowlite.aar + unzip -d $OUTDIR $BUILDER-bin/$BASEDIR/java/tensorflowlite.aar + # targetSdkVersion is here to prevent the app from requesting spurious + # permissions, such as permission to make phone calls. It worked for v1.0, + # but minSdkVersion might be the preferred way to handle this. + sed -i -e 's///' $OUTDIR/AndroidManifest.xml +} + +function build_arch() { + local ARCH=$1 + local CONFIG=$2 + local OUTDIR=$3 + mkdir -p $OUTDIR/jni/$ARCH/ + $BUILDER build $BUILD_OPTS $CROSSTOOL_OPTS --cpu=$CONFIG \ + $BASEDIR/java:libtensorflowlite_jni.so + cp $BUILDER-bin/$BASEDIR/java/libtensorflowlite_jni.so $OUTDIR/jni/$ARCH/ +} + +rm -rf $TMPDIR +mkdir -p $TMPDIR/jni + +build_basic_aar $TMPDIR +build_arch arm64-v8a arm64-v8a $TMPDIR +build_arch armeabi-v7a armeabi-v7a $TMPDIR +build_arch x86 x86 $TMPDIR +build_arch x86_64 x86_64 $TMPDIR + +AAR_FILE=`realpath tflite-${VERSION}.aar` +(cd $TMPDIR && zip $AAR_FILE -r *) +echo "New AAR file is $AAR_FILE" + diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/assets/labels_imagenet_slim.txt b/tensorflow/contrib/lite/java/demo/app/src/main/assets/labels_imagenet_slim.txt new file mode 100644 index 0000000000000000000000000000000000000000..572eccf90087c1c19874e40b950c1610f59cc9c2 --- /dev/null +++ b/tensorflow/contrib/lite/java/demo/app/src/main/assets/labels_imagenet_slim.txt @@ -0,0 +1,1001 @@ +dummy +tench +goldfish +great white shark +tiger shark +hammerhead +electric ray +stingray +cock +hen +ostrich +brambling +goldfinch +house finch +junco +indigo bunting +robin +bulbul +jay +magpie +chickadee +water ouzel +kite +bald eagle +vulture +great grey owl +European fire salamander +common newt +eft +spotted salamander +axolotl +bullfrog +tree frog +tailed frog +loggerhead +leatherback turtle +mud turtle +terrapin +box turtle +banded gecko +common iguana +American chameleon +whiptail +agama +frilled lizard +alligator lizard +Gila monster +green lizard +African chameleon +Komodo dragon +African crocodile +American alligator +triceratops +thunder snake +ringneck snake +hognose snake +green snake +king snake +garter snake +water snake +vine snake +night snake +boa constrictor +rock python +Indian cobra +green mamba +sea snake +horned viper +diamondback +sidewinder +trilobite +harvestman +scorpion +black and gold garden spider +barn spider +garden spider +black widow +tarantula +wolf spider +tick +centipede +black grouse +ptarmigan +ruffed grouse +prairie chicken +peacock +quail +partridge +African grey +macaw +sulphur-crested cockatoo +lorikeet +coucal +bee eater +hornbill +hummingbird +jacamar +toucan +drake +red-breasted merganser +goose +black swan +tusker +echidna +platypus +wallaby +koala +wombat +jellyfish +sea anemone +brain coral +flatworm +nematode +conch +snail +slug +sea slug +chiton +chambered nautilus +Dungeness crab +rock crab +fiddler crab +king crab +American lobster +spiny lobster +crayfish +hermit crab +isopod +white stork +black stork +spoonbill +flamingo +little blue heron +American egret +bittern +crane +limpkin +European gallinule +American coot +bustard +ruddy turnstone +red-backed sandpiper +redshank +dowitcher +oystercatcher +pelican +king penguin +albatross +grey whale +killer whale +dugong +sea lion +Chihuahua +Japanese spaniel +Maltese dog +Pekinese +Shih-Tzu +Blenheim spaniel +papillon +toy terrier +Rhodesian ridgeback +Afghan hound +basset +beagle +bloodhound +bluetick +black-and-tan coonhound +Walker hound +English foxhound +redbone +borzoi +Irish wolfhound +Italian greyhound +whippet +Ibizan hound +Norwegian elkhound +otterhound +Saluki +Scottish deerhound +Weimaraner +Staffordshire bullterrier +American Staffordshire terrier +Bedlington terrier +Border terrier +Kerry blue terrier +Irish terrier +Norfolk terrier +Norwich terrier +Yorkshire terrier +wire-haired fox terrier +Lakeland terrier +Sealyham terrier +Airedale +cairn +Australian terrier +Dandie Dinmont +Boston bull +miniature schnauzer +giant schnauzer +standard schnauzer +Scotch terrier +Tibetan terrier +silky terrier +soft-coated wheaten terrier +West Highland white terrier +Lhasa +flat-coated retriever +curly-coated retriever +golden retriever +Labrador retriever +Chesapeake Bay retriever +German short-haired pointer +vizsla +English setter +Irish setter +Gordon setter +Brittany spaniel +clumber +English springer +Welsh springer spaniel +cocker spaniel +Sussex spaniel +Irish water spaniel +kuvasz +schipperke +groenendael +malinois +briard +kelpie +komondor +Old English sheepdog +Shetland sheepdog +collie +Border collie +Bouvier des Flandres +Rottweiler +German shepherd +Doberman +miniature pinscher +Greater Swiss Mountain dog +Bernese mountain dog +Appenzeller +EntleBucher +boxer +bull mastiff +Tibetan mastiff +French bulldog +Great Dane +Saint Bernard +Eskimo dog +malamute +Siberian husky +dalmatian +affenpinscher +basenji +pug +Leonberg +Newfoundland +Great Pyrenees +Samoyed +Pomeranian +chow +keeshond +Brabancon griffon +Pembroke +Cardigan +toy poodle +miniature poodle +standard poodle +Mexican hairless +timber wolf +white wolf +red wolf +coyote +dingo +dhole +African hunting dog +hyena +red fox +kit fox +Arctic fox +grey fox +tabby +tiger cat +Persian cat +Siamese cat +Egyptian cat +cougar +lynx +leopard +snow leopard +jaguar +lion +tiger +cheetah +brown bear +American black bear +ice bear +sloth bear +mongoose +meerkat +tiger beetle +ladybug +ground beetle +long-horned beetle +leaf beetle +dung beetle +rhinoceros beetle +weevil +fly +bee +ant +grasshopper +cricket +walking stick +cockroach +mantis +cicada +leafhopper +lacewing +dragonfly +damselfly +admiral +ringlet +monarch +cabbage butterfly +sulphur butterfly +lycaenid +starfish +sea urchin +sea cucumber +wood rabbit +hare +Angora +hamster +porcupine +fox squirrel +marmot +beaver +guinea pig +sorrel +zebra +hog +wild boar +warthog +hippopotamus +ox +water buffalo +bison +ram +bighorn +ibex +hartebeest +impala +gazelle +Arabian camel +llama +weasel +mink +polecat +black-footed ferret +otter +skunk +badger +armadillo +three-toed sloth +orangutan +gorilla +chimpanzee +gibbon +siamang +guenon +patas +baboon +macaque +langur +colobus +proboscis monkey +marmoset +capuchin +howler monkey +titi +spider monkey +squirrel monkey +Madagascar cat +indri +Indian elephant +African elephant +lesser panda +giant panda +barracouta +eel +coho +rock beauty +anemone fish +sturgeon +gar +lionfish +puffer +abacus +abaya +academic gown +accordion +acoustic guitar +aircraft carrier +airliner +airship +altar +ambulance +amphibian +analog clock +apiary +apron +ashcan +assault rifle +backpack +bakery +balance beam +balloon +ballpoint +Band Aid +banjo +bannister +barbell +barber chair +barbershop +barn +barometer +barrel +barrow +baseball +basketball +bassinet +bassoon +bathing cap +bath towel +bathtub +beach wagon +beacon +beaker +bearskin +beer bottle +beer glass +bell cote +bib +bicycle-built-for-two +bikini +binder +binoculars +birdhouse +boathouse +bobsled +bolo tie +bonnet +bookcase +bookshop +bottlecap +bow +bow tie +brass +brassiere +breakwater +breastplate +broom +bucket +buckle +bulletproof vest +bullet train +butcher shop +cab +caldron +candle +cannon +canoe +can opener +cardigan +car mirror +carousel +carpenter's kit +carton +car wheel +cash machine +cassette +cassette player +castle +catamaran +CD player +cello +cellular telephone +chain +chainlink fence +chain mail +chain saw +chest +chiffonier +chime +china cabinet +Christmas stocking +church +cinema +cleaver +cliff dwelling +cloak +clog +cocktail shaker +coffee mug +coffeepot +coil +combination lock +computer keyboard +confectionery +container ship +convertible +corkscrew +cornet +cowboy boot +cowboy hat +cradle +crane +crash helmet +crate +crib +Crock Pot +croquet ball +crutch +cuirass +dam +desk +desktop computer +dial telephone +diaper +digital clock +digital watch +dining table +dishrag +dishwasher +disk brake +dock +dogsled +dome +doormat +drilling platform +drum +drumstick +dumbbell +Dutch oven +electric fan +electric guitar +electric locomotive +entertainment center +envelope +espresso maker +face powder +feather boa +file +fireboat +fire engine +fire screen +flagpole +flute +folding chair +football helmet +forklift +fountain +fountain pen +four-poster +freight car +French horn +frying pan +fur coat +garbage truck +gasmask +gas pump +goblet +go-kart +golf ball +golfcart +gondola +gong +gown +grand piano +greenhouse +grille +grocery store +guillotine +hair slide +hair spray +half track +hammer +hamper +hand blower +hand-held computer +handkerchief +hard disc +harmonica +harp +harvester +hatchet +holster +home theater +honeycomb +hook +hoopskirt +horizontal bar +horse cart +hourglass +iPod +iron +jack-o'-lantern +jean +jeep +jersey +jigsaw puzzle +jinrikisha +joystick +kimono +knee pad +knot +lab coat +ladle +lampshade +laptop +lawn mower +lens cap +letter opener +library +lifeboat +lighter +limousine +liner +lipstick +Loafer +lotion +loudspeaker +loupe +lumbermill +magnetic compass +mailbag +mailbox +maillot +maillot +manhole cover +maraca +marimba +mask +matchstick +maypole +maze +measuring cup +medicine chest +megalith +microphone +microwave +military uniform +milk can +minibus +miniskirt +minivan +missile +mitten +mixing bowl +mobile home +Model T +modem +monastery +monitor +moped +mortar +mortarboard +mosque +mosquito net +motor scooter +mountain bike +mountain tent +mouse +mousetrap +moving van +muzzle +nail +neck brace +necklace +nipple +notebook +obelisk +oboe +ocarina +odometer +oil filter +organ +oscilloscope +overskirt +oxcart +oxygen mask +packet +paddle +paddlewheel +padlock +paintbrush +pajama +palace +panpipe +paper towel +parachute +parallel bars +park bench +parking meter +passenger car +patio +pay-phone +pedestal +pencil box +pencil sharpener +perfume +Petri dish +photocopier +pick +pickelhaube +picket fence +pickup +pier +piggy bank +pill bottle +pillow +ping-pong ball +pinwheel +pirate +pitcher +plane +planetarium +plastic bag +plate rack +plow +plunger +Polaroid camera +pole +police van +poncho +pool table +pop bottle +pot +potter's wheel +power drill +prayer rug +printer +prison +projectile +projector +puck +punching bag +purse +quill +quilt +racer +racket +radiator +radio +radio telescope +rain barrel +recreational vehicle +reel +reflex camera +refrigerator +remote control +restaurant +revolver +rifle +rocking chair +rotisserie +rubber eraser +rugby ball +rule +running shoe +safe +safety pin +saltshaker +sandal +sarong +sax +scabbard +scale +school bus +schooner +scoreboard +screen +screw +screwdriver +seat belt +sewing machine +shield +shoe shop +shoji +shopping basket +shopping cart +shovel +shower cap +shower curtain +ski +ski mask +sleeping bag +slide rule +sliding door +slot +snorkel +snowmobile +snowplow +soap dispenser +soccer ball +sock +solar dish +sombrero +soup bowl +space bar +space heater +space shuttle +spatula +speedboat +spider web +spindle +sports car +spotlight +stage +steam locomotive +steel arch bridge +steel drum +stethoscope +stole +stone wall +stopwatch +stove +strainer +streetcar +stretcher +studio couch +stupa +submarine +suit +sundial +sunglass +sunglasses +sunscreen +suspension bridge +swab +sweatshirt +swimming trunks +swing +switch +syringe +table lamp +tank +tape player +teapot +teddy +television +tennis ball +thatch +theater curtain +thimble +thresher +throne +tile roof +toaster +tobacco shop +toilet seat +torch +totem pole +tow truck +toyshop +tractor +trailer truck +tray +trench coat +tricycle +trimaran +tripod +triumphal arch +trolleybus +trombone +tub +turnstile +typewriter keyboard +umbrella +unicycle +upright +vacuum +vase +vault +velvet +vending machine +vestment +viaduct +violin +volleyball +waffle iron +wall clock +wallet +wardrobe +warplane +washbasin +washer +water bottle +water jug +water tower +whiskey jug +whistle +wig +window screen +window shade +Windsor tie +wine bottle +wing +wok +wooden spoon +wool +worm fence +wreck +yawl +yurt +web site +comic book +crossword puzzle +street sign +traffic light +book jacket +menu +plate +guacamole +consomme +hot pot +trifle +ice cream +ice lolly +French loaf +bagel +pretzel +cheeseburger +hotdog +mashed potato +head cabbage +broccoli +cauliflower +zucchini +spaghetti squash +acorn squash +butternut squash +cucumber +artichoke +bell pepper +cardoon +mushroom +Granny Smith +strawberry +orange +lemon +fig +pineapple +banana +jackfruit +custard apple +pomegranate +hay +carbonara +chocolate sauce +dough +meat loaf +pizza +potpie +burrito +red wine +espresso +cup +eggnog +alp +bubble +cliff +coral reef +geyser +lakeside +promontory +sandbar +seashore +valley +volcano +ballplayer +groom +scuba diver +rapeseed +daisy +yellow lady's slipper +corn +acorn +hip +buckeye +coral fungus +agaric +gyromitra +stinkhorn +earthstar +hen-of-the-woods +bolete +ear +toilet tissue diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/assets/labels.txt b/tensorflow/contrib/lite/java/demo/app/src/main/assets/labels_mobilenet_quant_v1_224.txt similarity index 100% rename from tensorflow/contrib/lite/java/demo/app/src/main/assets/labels.txt rename to tensorflow/contrib/lite/java/demo/app/src/main/assets/labels_mobilenet_quant_v1_224.txt diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java index 74737a8b883d23684220dd32bbd7a9e8ab4b2123..d2048b41b1e76fe42c919c9b889df5be8a94957f 100644 --- a/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java +++ b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java @@ -296,7 +296,8 @@ public class Camera2BasicFragment extends Fragment public void onActivityCreated(Bundle savedInstanceState) { super.onActivityCreated(savedInstanceState); try { - classifier = new ImageClassifier(getActivity()); + // create either a new ImageClassifierQuantizedMobileNet or an ImageClassifierFloatInception + classifier = new ImageClassifierQuantizedMobileNet(getActivity()); } catch (IOException e) { Log.e(TAG, "Failed to initialize an image classifier."); } @@ -659,7 +660,7 @@ public class Camera2BasicFragment extends Fragment return; } Bitmap bitmap = - textureView.getBitmap(ImageClassifier.DIM_IMG_SIZE_X, ImageClassifier.DIM_IMG_SIZE_Y); + textureView.getBitmap(classifier.getImageSizeX(), classifier.getImageSizeY()); String textToShow = classifier.classifyFrame(bitmap); bitmap.recycle(); showToast(textToShow); diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifier.java b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifier.java index e7bad4637041d003c1e507d81c0c30404c587653..c319bff9f11546ac4d49c6b34c5ecdbc41547d58 100644 --- a/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifier.java +++ b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifier.java @@ -20,6 +20,9 @@ import android.content.res.AssetFileDescriptor; 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; @@ -34,20 +37,15 @@ import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.PriorityQueue; -import org.tensorflow.lite.Interpreter; -/** Classifies images with Tensorflow Lite. */ -public class ImageClassifier { +/** + * Classifies images with Tensorflow Lite. + */ +public abstract class ImageClassifier { /** Tag for the {@link Log}. */ private static final String TAG = "TfLiteCameraDemo"; - /** Name of the model file stored in Assets. */ - private static final String MODEL_PATH = "mobilenet_quant_v1_224.tflite"; - - /** Name of the label file stored in Assets. */ - private static final String LABEL_PATH = "labels.txt"; - /** Number of results to show in the UI. */ private static final int RESULTS_TO_SHOW = 3; @@ -56,23 +54,23 @@ public class ImageClassifier { private static final int DIM_PIXEL_SIZE = 3; - static final int DIM_IMG_SIZE_X = 224; - static final int DIM_IMG_SIZE_Y = 224; - /* Preallocated buffers for storing image data in. */ - private int[] intValues = new int[DIM_IMG_SIZE_X * DIM_IMG_SIZE_Y]; + private int[] intValues = new int[getImageSizeX() * getImageSizeY()]; /** An instance of the driver class to run model inference with Tensorflow Lite. */ - private Interpreter tflite; + protected Interpreter tflite; /** Labels corresponding to the output of the vision model. */ private List labelList; /** A ByteBuffer to hold image data, to be feed into Tensorflow Lite as inputs. */ - private ByteBuffer imgData = null; + protected ByteBuffer imgData = null; + + /** multi-stage low pass filter * */ + private float[][] filterLabelProbArray = null; - /** An array to hold inference results, to be feed into Tensorflow Lite as outputs. */ - private byte[][] labelProbArray = null; + private static final int FILTER_STAGES = 3; + private static final float FILTER_FACTOR = 0.4f; private PriorityQueue> sortedLabels = new PriorityQueue<>( @@ -90,9 +88,10 @@ public class ImageClassifier { labelList = loadLabelList(activity); imgData = ByteBuffer.allocateDirect( - DIM_BATCH_SIZE * DIM_IMG_SIZE_X * DIM_IMG_SIZE_Y * DIM_PIXEL_SIZE); + DIM_BATCH_SIZE * getImageSizeX() * getImageSizeY() * DIM_PIXEL_SIZE * + getNumBytesPerChannel()); imgData.order(ByteOrder.nativeOrder()); - labelProbArray = new byte[1][labelList.size()]; + filterLabelProbArray = new float[FILTER_STAGES][getNumLabels()]; Log.d(TAG, "Created a Tensorflow Lite Image Classifier."); } @@ -105,14 +104,41 @@ public class ImageClassifier { convertBitmapToByteBuffer(bitmap); // Here's where the magic happens!!! long startTime = SystemClock.uptimeMillis(); - tflite.run(imgData, labelProbArray); + runInference(); long endTime = SystemClock.uptimeMillis(); Log.d(TAG, "Timecost to run model inference: " + Long.toString(endTime - startTime)); + + // Smooth the results across frames. + applyFilter(); + + // Print the results. String textToShow = printTopKLabels(); textToShow = Long.toString(endTime - startTime) + "ms" + textToShow; return textToShow; } + void applyFilter() { + int numLabels = getNumLabels(); + + // Low pass filter `labelProbArray` into the first stage of the filter. + for (int j = 0; j < numLabels; ++j) { + filterLabelProbArray[0][j] += + FILTER_FACTOR * (getProbability(j) - filterLabelProbArray[0][j]); + } + // Low pass filter each stage into the next. + for (int i = 1; i < FILTER_STAGES; ++i) { + for (int j = 0; j < numLabels; ++j) { + filterLabelProbArray[i][j] += + FILTER_FACTOR * (filterLabelProbArray[i - 1][j] - filterLabelProbArray[i][j]); + } + } + + // Copy the last stage filter output back to `labelProbArray`. + for (int j = 0; j < numLabels; ++j) { + setProbability(j, filterLabelProbArray[FILTER_STAGES - 1][j]); + } + } + /** Closes tflite to release resources. */ public void close() { tflite.close(); @@ -123,7 +149,7 @@ public class ImageClassifier { private List loadLabelList(Activity activity) throws IOException { List labelList = new ArrayList(); BufferedReader reader = - new BufferedReader(new InputStreamReader(activity.getAssets().open(LABEL_PATH))); + new BufferedReader(new InputStreamReader(activity.getAssets().open(getLabelPath()))); String line; while ((line = reader.readLine()) != null) { labelList.add(line); @@ -134,7 +160,7 @@ public class ImageClassifier { /** Memory-map the model file in Assets. */ private MappedByteBuffer loadModelFile(Activity activity) throws IOException { - AssetFileDescriptor fileDescriptor = activity.getAssets().openFd(MODEL_PATH); + AssetFileDescriptor fileDescriptor = activity.getAssets().openFd(getModelPath()); FileInputStream inputStream = new FileInputStream(fileDescriptor.getFileDescriptor()); FileChannel fileChannel = inputStream.getChannel(); long startOffset = fileDescriptor.getStartOffset(); @@ -152,12 +178,10 @@ public class ImageClassifier { // Convert the image to floating point. int pixel = 0; long startTime = SystemClock.uptimeMillis(); - for (int i = 0; i < DIM_IMG_SIZE_X; ++i) { - for (int j = 0; j < DIM_IMG_SIZE_Y; ++j) { + for (int i = 0; i < getImageSizeX(); ++i) { + for (int j = 0; j < getImageSizeY(); ++j) { final int val = intValues[pixel++]; - imgData.put((byte) ((val >> 16) & 0xFF)); - imgData.put((byte) ((val >> 8) & 0xFF)); - imgData.put((byte) (val & 0xFF)); + addPixelValue(val); } } long endTime = SystemClock.uptimeMillis(); @@ -166,9 +190,9 @@ public class ImageClassifier { /** Prints top-K labels, to be shown in UI as the results. */ private String printTopKLabels() { - for (int i = 0; i < labelList.size(); ++i) { + for (int i = 0; i < getNumLabels(); ++i) { sortedLabels.add( - new AbstractMap.SimpleEntry<>(labelList.get(i), (labelProbArray[0][i] & 0xff) / 255.0f)); + new AbstractMap.SimpleEntry<>(labelList.get(i), getNormalizedProbability(i))); if (sortedLabels.size() > RESULTS_TO_SHOW) { sortedLabels.poll(); } @@ -177,8 +201,84 @@ public class ImageClassifier { final int size = sortedLabels.size(); for (int i = 0; i < size; ++i) { Map.Entry label = sortedLabels.poll(); - textToShow = "\n" + label.getKey() + ":" + Float.toString(label.getValue()) + textToShow; + textToShow = String.format("\n%s: %4.2f", label.getKey(), label.getValue()) + textToShow; } return textToShow; } + + /** + * Get the name of the model file stored in Assets. + * @return + */ + protected abstract String getModelPath(); + + /** + * Get the name of the label file stored in Assets. + * @return + */ + protected abstract String getLabelPath(); + + /** + * Get the image size along the x axis. + * @return + */ + protected abstract int getImageSizeX(); + + /** + * Get the image size along the y axis. + * @return + */ + protected abstract int getImageSizeY(); + + /** + * Get the number of bytes that is used to store a single color channel value. + * @return + */ + protected abstract int getNumBytesPerChannel(); + + /** + * Add pixelValue to byteBuffer. + * @param pixelValue + */ + protected abstract void addPixelValue(int pixelValue); + + /** + * Read the probability value for the specified label + * This is either the original value as it was read from the net's output or the updated value + * after the filter was applied. + * @param labelIndex + * @return + */ + protected abstract float getProbability(int labelIndex); + + /** + * Set the probability value for the specified label. + * @param labelIndex + * @param value + */ + protected abstract void setProbability(int labelIndex, Number value); + + /** + * Get the normalized probability value for the specified label. + * This is the final value as it will be shown to the user. + * @return + */ + protected abstract float getNormalizedProbability(int labelIndex); + + /** + * Run inference using the prepared input in {@link #imgData}. + * Afterwards, the result will be provided by getProbability(). + * + * This additional method is necessary, because we don't have a common base for different + * primitive data types. + */ + protected abstract void runInference(); + + /** + * Get the total number of labels. + * @return + */ + protected int getNumLabels() { + return labelList.size(); + } } diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierFloatInception.java b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierFloatInception.java new file mode 100644 index 0000000000000000000000000000000000000000..be17b85e0cd93778fd123663595c43b730fb44f7 --- /dev/null +++ b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierFloatInception.java @@ -0,0 +1,105 @@ +/* Copyright 2017 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. +==============================================================================*/ + +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. + */ +public class ImageClassifierFloatInception extends ImageClassifier { + + /** + * 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. + */ + private float[][] labelProbArray = null; + + /** + * Initializes an {@code ImageClassifier}. + * + * @param activity + */ + ImageClassifierFloatInception(Activity activity) throws IOException { + super(activity); + labelProbArray = new float[1][getNumLabels()]; + } + + @Override + protected String getModelPath() { + // you can download this file from + // https://storage.googleapis.com/download.tensorflow.org/models/tflite/inception_v3_slim_2016_android_2017_11_10.zip + return "inceptionv3_slim_2016.tflite"; + } + + @Override + protected String getLabelPath() { + return "labels_imagenet_slim.txt"; + } + + @Override + protected int getImageSizeX() { + return 299; + } + + @Override + protected int getImageSizeY() { + return 299; + } + + @Override + protected int getNumBytesPerChannel() { + // a 32bit float value requires 4 bytes + return 4; + } + + @Override + protected void addPixelValue(int pixelValue) { + imgData.putFloat((((pixelValue >> 16) & 0xFF) - IMAGE_MEAN) / IMAGE_STD); + imgData.putFloat((((pixelValue >> 8) & 0xFF) - IMAGE_MEAN) / IMAGE_STD); + imgData.putFloat(((pixelValue & 0xFF) - IMAGE_MEAN) / IMAGE_STD); + } + + @Override + protected float getProbability(int labelIndex) { + return labelProbArray[0][labelIndex]; + } + + @Override + protected void setProbability(int labelIndex, Number value) { + labelProbArray[0][labelIndex] = value.floatValue(); + } + + @Override + protected float getNormalizedProbability(int labelIndex) { + // TODO the following value isn't in [0,1] yet, but may be greater. Why? + return getProbability(labelIndex); + } + + @Override + protected void runInference() { + tflite.run(imgData, labelProbArray); + } +} diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierQuantizedMobileNet.java b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierQuantizedMobileNet.java new file mode 100644 index 0000000000000000000000000000000000000000..156c895146940adfe71f111be6e354e02b75ea48 --- /dev/null +++ b/tensorflow/contrib/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierQuantizedMobileNet.java @@ -0,0 +1,97 @@ +/* Copyright 2017 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. +==============================================================================*/ + +package com.example.android.tflitecamerademo; + +import android.app.Activity; + +import java.io.IOException; + +/** + * 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. + */ + private byte[][] labelProbArray = null; + + /** + * Initializes an {@code ImageClassifier}. + * + * @param activity + */ + ImageClassifierQuantizedMobileNet(Activity activity) throws IOException { + super(activity); + labelProbArray = new byte[1][getNumLabels()]; + } + + @Override + protected String getModelPath() { + // you can download this file from + // https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip + return "mobilenet_quant_v1_224.tflite"; + } + + @Override + protected String getLabelPath() { + return "labels_mobilenet_quant_v1_224.txt"; + } + + @Override + protected int getImageSizeX() { + return 224; + } + + @Override + protected int getImageSizeY() { + return 224; + } + + @Override + protected int getNumBytesPerChannel() { + // the quantized model uses a single byte only + return 1; + } + + @Override + protected void addPixelValue(int pixelValue) { + imgData.put((byte) ((pixelValue >> 16) & 0xFF)); + imgData.put((byte) ((pixelValue >> 8) & 0xFF)); + imgData.put((byte) (pixelValue & 0xFF)); + } + + @Override + protected float getProbability(int labelIndex) { + return labelProbArray[0][labelIndex]; + } + + @Override + protected void setProbability(int labelIndex, Number value) { + labelProbArray[0][labelIndex] = value.byteValue(); + } + + @Override + protected float getNormalizedProbability(int labelIndex) { + return (labelProbArray[0][labelIndex] & 0xff) / 255.0f; + } + + @Override + protected void runInference() { + tflite.run(imgData, labelProbArray); + } +} diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java index 1939a078ad8031b99620773c9b91335c4e8f7b22..5ee594dec492ad2fee22e603a6de311b3fed4cac 100644 --- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java +++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java @@ -34,7 +34,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { NativeInterpreterWrapper(String modelPath) { errorHandle = createErrorReporter(ERROR_BUFFER_SIZE); modelHandle = createModel(modelPath, errorHandle); - interpreterHandle = createInterpreter(modelHandle); + interpreterHandle = createInterpreter(modelHandle, errorHandle); } /** @@ -46,7 +46,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { modelByteBuffer = mappedByteBuffer; errorHandle = createErrorReporter(ERROR_BUFFER_SIZE); modelHandle = createModelWithBuffer(modelByteBuffer, errorHandle); - interpreterHandle = createInterpreter(modelHandle); + interpreterHandle = createInterpreter(modelHandle, errorHandle); } /** Releases resources associated with this {@code NativeInterpreterWrapper}. */ @@ -103,11 +103,22 @@ final class NativeInterpreterWrapper implements AutoCloseable { return outputs; } + private static native long[] run( + long interpreterHandle, + long errorHandle, + Object[] sizes, + int[] dtypes, + int[] numsOfBytes, + Object[] values); + /** Resizes dimensions of a specific input. */ void resizeInput(int idx, int[] dims) { resizeInput(interpreterHandle, errorHandle, idx, dims); } + private static native void resizeInput( + long interpreterHandle, long errorHandle, int inputIdx, int[] dims); + void setUseNNAPI(boolean useNNAPI) { useNNAPI(interpreterHandle, useNNAPI); } @@ -245,9 +256,6 @@ final class NativeInterpreterWrapper implements AutoCloseable { private static native String[] getOutputNames(long interpreterHandle); - private static native void resizeInput( - long interpreterHandle, long errorHandle, int inputIdx, int[] dims); - private static native void useNNAPI(long interpreterHandle, boolean state); private static native long createErrorReporter(int size); @@ -256,15 +264,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { private static native long createModelWithBuffer(MappedByteBuffer modelBuffer, long errorHandle); - private static native long createInterpreter(long modelHandle); - - private static native long[] run( - long interpreterHandle, - long errorHandle, - Object[] sizes, - int[] dtypes, - int[] numsOfBytes, - Object[] values); + private static native long createInterpreter(long modelHandle, long errorHandle); private static native void delete(long errorHandle, long modelHandle, long interpreterHandle); diff --git a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc index bc6462eb5466e14769f94c5103984f5201b4b8dc..c346f9f92e360c0722ebac440d790da6441ceecf 100644 --- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc +++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc @@ -200,6 +200,12 @@ TfLiteStatus setInputs(JNIEnv* env, tflite::Interpreter* interpreter, return kTfLiteOk; } +// TODO(yichengfan): evaluate the benefit to use tflite verifier. +bool VerifyModel(const void* buf, size_t len) { + flatbuffers::Verifier verifier(static_cast(buf), len); + return tflite::VerifyModelBuffer(verifier); +} + } // namespace JNIEXPORT jobjectArray JNICALL @@ -271,6 +277,17 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createModel( convertLongToErrorReporter(env, error_handle); if (error_reporter == nullptr) return 0; const char* path = env->GetStringUTFChars(model_file, nullptr); + + { + tflite::FileCopyAllocation allocation(path, nullptr); + if (!VerifyModel(allocation.base(), allocation.bytes())) { + throwException(env, kIllegalArgumentException, + "Contents of %s is not a valid flatbuffer model", path); + env->ReleaseStringUTFChars(model_file, path); + return 0; + } + } + auto model = tflite::FlatBufferModel::BuildFromFile(path, error_reporter); if (!model) { throwException(env, kIllegalArgumentException, @@ -293,6 +310,12 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createModelWithBuffer( const char* buf = static_cast(env->GetDirectBufferAddress(model_buffer)); jlong capacity = env->GetDirectBufferCapacity(model_buffer); + if (!VerifyModel(buf, capacity)) { + throwException(env, kIllegalArgumentException, + "MappedByteBuffer is not a valid flatbuffer model"); + return 0; + } + auto model = tflite::FlatBufferModel::BuildFromBuffer( buf, static_cast(capacity), error_reporter); if (!model) { @@ -307,12 +330,21 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createModelWithBuffer( JNIEXPORT jlong JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_createInterpreter( - JNIEnv* env, jclass clazz, jlong model_handle) { + JNIEnv* env, jclass clazz, jlong model_handle, jlong error_handle) { tflite::FlatBufferModel* model = convertLongToModel(env, model_handle); if (model == nullptr) return 0; + BufferErrorReporter* error_reporter = + convertLongToErrorReporter(env, error_handle); + if (error_reporter == nullptr) return 0; auto resolver = ::tflite::CreateOpResolver(); std::unique_ptr interpreter; - tflite::InterpreterBuilder(*model, *(resolver.get()))(&interpreter); + TfLiteStatus status = + tflite::InterpreterBuilder(*model, *(resolver.get()))(&interpreter); + if (status != kTfLiteOk) { + throwException(env, kIllegalArgumentException, + "Cannot create interpreter: %s", + error_reporter->CachedErrorMessage()); + } return reinterpret_cast(interpreter.release()); } diff --git a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h index 430886b7cc04a356d1826843acc1bbebf4189bf7..c52a7e4e439936344be26d5761fb5747db64794a 100644 --- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h +++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h @@ -95,11 +95,11 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createModelWithBuffer( /* * Class: org_tensorflow_lite_NativeInterpreterWrapper * Method: - * Signature: (J)J + * Signature: (JJ)J */ JNIEXPORT jlong JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_createInterpreter( - JNIEnv* env, jclass clazz, jlong model_handle); + JNIEnv* env, jclass clazz, jlong model_handle, jlong error_handle); /* * Class: org_tensorflow_lite_NativeInterpreterWrapper diff --git a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/NativeInterpreterWrapperTest.java b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/NativeInterpreterWrapperTest.java index 9a6894f49c0b7278511717d2671648c6d1763e00..90323555d88419d837a76bca7de6d9998e388fca 100644 --- a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/NativeInterpreterWrapperTest.java +++ b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/NativeInterpreterWrapperTest.java @@ -25,6 +25,7 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; /** Unit tests for {@link org.tensorflow.lite.NativeInterpreterWrapper}. */ +// TODO(b/71818425): Generates model files dynamically. @RunWith(JUnit4.class) public final class NativeInterpreterWrapperTest { @@ -43,6 +44,9 @@ public final class NativeInterpreterWrapperTest { private static final String INVALID_MODEL_PATH = "tensorflow/contrib/lite/java/src/testdata/invalid_model.bin"; + private static final String MODEL_WITH_CUSTOM_OP_PATH = + "tensorflow/contrib/lite/java/src/testdata/with_custom_op.lite"; + @Test public void testConstructor() { NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(FLOAT_MODEL_PATH); @@ -55,10 +59,20 @@ public final class NativeInterpreterWrapperTest { try { NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(INVALID_MODEL_PATH); fail(); + } catch (IllegalArgumentException e) { + assertThat(e).hasMessageThat().contains("is not a valid flatbuffer model"); + } + } + + @Test + public void testConstructorWithUnresolableCustomOp() { + try { + NativeInterpreterWrapper wrapper = new NativeInterpreterWrapper(MODEL_WITH_CUSTOM_OP_PATH); + fail(); } catch (IllegalArgumentException e) { assertThat(e) .hasMessageThat() - .contains("Model provided has model identifier ' is ', should be 'TFL3'"); + .contains("Cannot create interpreter: Didn't find custom op for name 'Assign'"); } } diff --git a/tensorflow/contrib/lite/java/src/testdata/with_custom_op.lite b/tensorflow/contrib/lite/java/src/testdata/with_custom_op.lite new file mode 100644 index 0000000000000000000000000000000000000000..e775d56d88854ecdf70233262ff5884d224f4373 Binary files /dev/null and b/tensorflow/contrib/lite/java/src/testdata/with_custom_op.lite differ diff --git a/tensorflow/contrib/lite/kernels/BUILD b/tensorflow/contrib/lite/kernels/BUILD index cc02cddb3d6cce3787fd15ee1734a490389fb9b3..b59dc5ffb339caade28626d1954d41bc821fae41 100644 --- a/tensorflow/contrib/lite/kernels/BUILD +++ b/tensorflow/contrib/lite/kernels/BUILD @@ -50,7 +50,7 @@ cc_library( deps = [ ":op_macros", "//tensorflow/contrib/lite:context", - "@gemmlowp//:gemmlowp", + "@gemmlowp", ], ) @@ -71,6 +71,32 @@ cc_library( ], ) +cc_library( + name = "kernel_util", + srcs = [ + "kernel_util.cc", + ], + hdrs = [ + "kernel_util.h", + ], + deps = [ + "//tensorflow/contrib/lite:builtin_op_data", + "//tensorflow/contrib/lite:context", + "//tensorflow/contrib/lite/kernels/internal:round", + ], +) + +tf_cc_test( + name = "kernel_util_test", + size = "small", + srcs = ["kernel_util_test.cc"], + deps = [ + ":kernel_util", + "//tensorflow/contrib/lite/testing:util", + "@com_google_googletest//:gtest", + ], +) + cc_library( name = "builtin_ops", srcs = [ @@ -78,19 +104,22 @@ cc_library( "add.cc", "basic_rnn.cc", "batch_to_space_nd.cc", + "bidirectional_sequence_rnn.cc", "concatenation.cc", "conv.cc", "depthwise_conv.cc", + "div.cc", "embedding_lookup.cc", "embedding_lookup_sparse.cc", + "exp.cc", "fully_connected.cc", "gather.cc", "hashtable_lookup.cc", - "kernel_util.cc", "l2norm.cc", "local_response_norm.cc", "lsh_projection.cc", "lstm.cc", + "mean.cc", "mul.cc", "pad.cc", "pooling.cc", @@ -98,12 +127,19 @@ cc_library( "reshape.cc", "resize_bilinear.cc", "skip_gram.cc", + "space_to_batch_nd.cc", "space_to_depth.cc", + "split.cc", + "squeeze.cc", + "strided_slice.cc", + "sub.cc", "svdf.cc", + "topk_v2.cc", + "transpose.cc", + "unidirectional_sequence_lstm.cc", "unidirectional_sequence_rnn.cc", ], hdrs = [ - "kernel_util.h", "padding.h", "register.h", ], @@ -117,11 +153,13 @@ cc_library( }), deps = [ ":activation_functor", + ":kernel_util", ":op_macros", "//tensorflow/contrib/lite:builtin_op_data", "//tensorflow/contrib/lite:framework", "//tensorflow/contrib/lite:string_util", "//tensorflow/contrib/lite/kernels:gemm_support", + "//tensorflow/contrib/lite/kernels/internal:kernel_utils", "//tensorflow/contrib/lite/kernels/internal:optimized", "//tensorflow/contrib/lite/kernels/internal:optimized_base", "//tensorflow/contrib/lite/kernels/internal:quantization_util", @@ -157,6 +195,32 @@ tf_cc_test( ], ) +tf_cc_test( + name = "transpose_test", + size = "small", + srcs = ["transpose_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "//tensorflow/contrib/lite/kernels/internal:reference", + "//tensorflow/contrib/lite/kernels/internal:reference_base", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "space_to_batch_nd_test", + size = "small", + srcs = ["space_to_batch_nd_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + tf_cc_test( name = "batch_to_space_nd_test", size = "small", @@ -189,6 +253,7 @@ tf_cc_test( ":builtin_ops", "//tensorflow/contrib/lite:framework", "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_absl//absl/memory", "@com_google_googletest//:gtest", ], ) @@ -217,6 +282,30 @@ tf_cc_test( ], ) +tf_cc_test( + name = "unidirectional_sequence_lstm_test", + size = "small", + srcs = ["unidirectional_sequence_lstm_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "bidirectional_sequence_rnn_test", + size = "small", + srcs = ["bidirectional_sequence_rnn_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + tf_cc_test( name = "unidirectional_sequence_rnn_test", size = "small", @@ -241,6 +330,30 @@ tf_cc_test( ], ) +tf_cc_test( + name = "exp_test", + size = "small", + srcs = ["exp_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "mean_test", + size = "small", + srcs = ["mean_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + tf_cc_test( name = "mul_test", size = "small", @@ -290,6 +403,19 @@ tf_cc_test( ], ) +tf_cc_test( + name = "topk_v2_test", + size = "small", + srcs = ["topk_v2_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:builtin_op_data", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + tf_cc_test( name = "resize_bilinear_test", size = "small", @@ -449,6 +575,42 @@ tf_cc_test( ], ) +tf_cc_test( + name = "split_test", + size = "small", + srcs = ["split_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "squeeze_test", + size = "small", + srcs = ["squeeze_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "strided_slice_test", + size = "small", + srcs = ["strided_slice_test.cc"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/contrib/lite/kernels/activation_functor.h b/tensorflow/contrib/lite/kernels/activation_functor.h index cfb3369e991a474315424423fe655ba214edabbc..41ec3cca33ae1c6bb3f7c43dd1923f104c2ab6a2 100644 --- a/tensorflow/contrib/lite/kernels/activation_functor.h +++ b/tensorflow/contrib/lite/kernels/activation_functor.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_ACTIVATION_FUNCTOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_ACTIVATION_FUNCTOR_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_ACTIVATION_FUNCTOR_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_ACTIVATION_FUNCTOR_H_ #include #include @@ -55,4 +55,4 @@ class ActivationFunctor { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_ACTIVATION_FUNCTOR_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_ACTIVATION_FUNCTOR_H_ diff --git a/tensorflow/contrib/lite/kernels/activations.cc b/tensorflow/contrib/lite/kernels/activations.cc index 7ab60a33e5e2ff61bae5f4c6db85ab9c47a391bc..3c5c77815d0f2592ab549152b4d77f45b967a660 100644 --- a/tensorflow/contrib/lite/kernels/activations.cc +++ b/tensorflow/contrib/lite/kernels/activations.cc @@ -15,8 +15,8 @@ limitations under the License. #include #include #include -#include #include +#include #include #include @@ -134,8 +134,7 @@ TfLiteStatus ReluEval(TfLiteContext* context, TfLiteNode* node) { float* out = output->data.f; for (; in < in_end; in++, out++) *out = std::max(0.f, *in); return kTfLiteOk; - } - break; + } break; default: context->ReportError(context, "Only float32 supported currently."); return kTfLiteError; @@ -173,8 +172,7 @@ TfLiteStatus Relu6Eval(TfLiteContext* context, TfLiteNode* node) { float* out = output->data.f; for (; in < in_end; in++, out++) *out = std::min(std::max(0.f, *in), 6.f); return kTfLiteOk; - } - break; + } break; default: context->ReportError(context, "Only float32 supported currently."); return kTfLiteError; @@ -192,8 +190,7 @@ TfLiteStatus TanhEval(TfLiteContext* context, TfLiteNode* node) { float* out = output->data.f; for (; in < in_end; in++, out++) *out = std::tanh(*in); return kTfLiteOk; - } - break; + } break; default: context->ReportError(context, "Only float32 supported currently."); return kTfLiteError; @@ -349,7 +346,7 @@ TfLiteRegistration* Register_RELU() { return &r; } -TfLiteRegistration* Register_RELU1() { +TfLiteRegistration* Register_RELU_N1_TO_1() { static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, activations::GenericPrepare, activations::Relu1Eval}; diff --git a/tensorflow/contrib/lite/kernels/activations_test.cc b/tensorflow/contrib/lite/kernels/activations_test.cc index 33ca56e745c043efd12b851af14f273fb273d577..68d49944e51b043b6b82aa1589d22f6ebed37574 100644 --- a/tensorflow/contrib/lite/kernels/activations_test.cc +++ b/tensorflow/contrib/lite/kernels/activations_test.cc @@ -102,7 +102,7 @@ TEST(FloatActivationsOpTest, Relu) { } TEST(FloatActivationsOpTest, Relu1) { - FloatActivationsOpModel m(BuiltinOperator_RELU1, + FloatActivationsOpModel m(BuiltinOperator_RELU_N1_TO_1, /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}); m.SetInput({ 0.0, -0.6, 0.2, -0.4, // diff --git a/tensorflow/contrib/lite/kernels/add.cc b/tensorflow/contrib/lite/kernels/add.cc index 0e10a249abac3ba19cf107e055aa71d1eee00122..63ea89df56bafa995950afec3a58267681af304f 100644 --- a/tensorflow/contrib/lite/kernels/add.cc +++ b/tensorflow/contrib/lite/kernels/add.cc @@ -37,7 +37,23 @@ constexpr int kInputTensor1 = 0; constexpr int kInputTensor2 = 1; constexpr int kOutputTensor = 0; +struct OpData { + bool requires_broadcast; +}; + +void* Init(TfLiteContext* context, const char* buffer, size_t length) { + auto* data = new OpData; + data->requires_broadcast = false; + return data; +} + +void Free(TfLiteContext* context, void* buffer) { + delete reinterpret_cast(buffer); +} + TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + OpData* data = reinterpret_cast(node->user_data); + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); @@ -45,43 +61,56 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - TF_LITE_ENSURE_EQ(context, NumDimensions(input1), NumDimensions(input2)); - for (int i = 0; i < NumDimensions(input1); ++i) { - TF_LITE_ENSURE_EQ(context, SizeOfDimension(input1, i), - SizeOfDimension(input2, i)); - } + TF_LITE_ENSURE_EQ(context, input1->type, input2->type); + output->type = input2->type; - TF_LITE_ENSURE_EQ(context, input1->type, output->type); - TF_LITE_ENSURE_EQ(context, input2->type, output->type); + data->requires_broadcast = !HaveSameShapes(input1, input2); + + TfLiteIntArray* output_size = nullptr; + if (data->requires_broadcast) { + TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast( + context, input1, input2, &output_size)); + } else { + output_size = TfLiteIntArrayCopy(input1->dims); + } - TfLiteIntArray* output_size = TfLiteIntArrayCopy(input1->dims); return context->ResizeTensor(context, output, output_size); } template void EvalAddFloat(TfLiteContext* context, TfLiteNode* node, - TfLiteAddParams* params, TfLiteTensor* input1, - TfLiteTensor* input2, TfLiteTensor* output) { + TfLiteAddParams* params, const OpData* data, + TfLiteTensor* input1, TfLiteTensor* input2, + TfLiteTensor* output) { float output_activation_min, output_activation_max; CalculateActivationRangeFloat(params->activation, &output_activation_min, &output_activation_max); -#define TF_LITE_ADD(type) \ - type::Add(GetTensorData(input1), GetTensorDims(input1), \ - GetTensorData(input2), GetTensorDims(input2), \ - output_activation_min, output_activation_max, \ - GetTensorData(output), GetTensorDims(output)) - if (kernel_type == kReference) { - TF_LITE_ADD(reference_ops); +#define TF_LITE_ADD(type, opname) \ + type::opname(GetTensorData(input1), GetTensorDims(input1), \ + GetTensorData(input2), GetTensorDims(input2), \ + output_activation_min, output_activation_max, \ + GetTensorData(output), GetTensorDims(output)) + if (kernel_type == kReference) { + if (data->requires_broadcast) { + TF_LITE_ADD(reference_ops, BroadcastAdd); } else { - TF_LITE_ADD(optimized_ops); + TF_LITE_ADD(reference_ops, Add); + } + } else { + if (data->requires_broadcast) { + TF_LITE_ADD(optimized_ops, BroadcastAdd); + } else { + TF_LITE_ADD(optimized_ops, Add); + } } #undef TF_LITE_ADD } template void EvalAddQuantized(TfLiteContext* context, TfLiteNode* node, - TfLiteAddParams* params, TfLiteTensor* input1, - TfLiteTensor* input2, TfLiteTensor* output) { + TfLiteAddParams* params, const OpData* data, + TfLiteTensor* input1, TfLiteTensor* input2, + TfLiteTensor* output) { auto input1_offset = -input1->params.zero_point; auto input2_offset = -input2->params.zero_point; auto output_offset = output->params.zero_point; @@ -112,19 +141,20 @@ void EvalAddQuantized(TfLiteContext* context, TfLiteNode* node, CalculateActivationRangeUint8(params->activation, output, &output_activation_min, &output_activation_max); -#define TF_LITE_ADD(type) \ - type::BroadcastAdd( \ - left_shift, GetTensorData(input1), GetTensorDims(input1), \ - input1_offset, input1_multiplier, input1_shift, \ - GetTensorData(input2), GetTensorDims(input2), input2_offset, \ - input2_multiplier, input2_shift, output_offset, output_multiplier, \ - output_shift, output_activation_min, output_activation_max, \ - GetTensorData(output), GetTensorDims(output)); - +#define TF_LITE_ADD(type, opname) \ + type::opname(left_shift, GetTensorData(input1), \ + GetTensorDims(input1), input1_offset, input1_multiplier, \ + input1_shift, GetTensorData(input2), \ + GetTensorDims(input2), input2_offset, input2_multiplier, \ + input2_shift, output_offset, output_multiplier, output_shift, \ + output_activation_min, output_activation_max, \ + GetTensorData(output), GetTensorDims(output)); + // The quantized version of Add doesn't support activations, so we + // always use BroadcastAdd. if (kernel_type == kReference) { - TF_LITE_ADD(reference_ops); + TF_LITE_ADD(reference_ops, BroadcastAdd); } else { - TF_LITE_ADD(optimized_ops); + TF_LITE_ADD(optimized_ops, BroadcastAdd); } #undef TF_LITE_ADD } @@ -132,15 +162,17 @@ void EvalAddQuantized(TfLiteContext* context, TfLiteNode* node, template TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->builtin_data); + OpData* data = reinterpret_cast(node->user_data); TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); if (output->type == kTfLiteFloat32) { - EvalAddFloat(context, node, params, input1, input2, output); + EvalAddFloat(context, node, params, data, input1, input2, + output); } else if (output->type == kTfLiteUInt8) { - EvalAddQuantized(context, node, params, input1, input2, + EvalAddQuantized(context, node, params, data, input1, input2, output); } else { context->ReportError(context, @@ -154,19 +186,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { } // namespace add TfLiteRegistration* Register_ADD_REF() { - static TfLiteRegistration r = {nullptr, nullptr, add::Prepare, + static TfLiteRegistration r = {add::Init, add::Free, add::Prepare, add::Eval}; return &r; } TfLiteRegistration* Register_ADD_GENERIC_OPT() { - static TfLiteRegistration r = {nullptr, nullptr, add::Prepare, + static TfLiteRegistration r = {add::Init, add::Free, add::Prepare, add::Eval}; return &r; } TfLiteRegistration* Register_ADD_NEON_OPT() { - static TfLiteRegistration r = {nullptr, nullptr, add::Prepare, + static TfLiteRegistration r = {add::Init, add::Free, add::Prepare, add::Eval}; return &r; } diff --git a/tensorflow/contrib/lite/kernels/add_test.cc b/tensorflow/contrib/lite/kernels/add_test.cc index ddf45bb576755d57d50c9e6e01bf50f15612c56d..956d05bed5162f6ce59705d59aad77ff056dda77 100644 --- a/tensorflow/contrib/lite/kernels/add_test.cc +++ b/tensorflow/contrib/lite/kernels/add_test.cc @@ -25,10 +25,11 @@ using ::testing::ElementsAreArray; class BaseAddOpModel : public SingleOpModel { public: - BaseAddOpModel(const TensorData& input, const TensorData& output, + BaseAddOpModel(const TensorData& input1, const TensorData& input2, + const TensorData& output, ActivationFunctionType activation_type) { - input1_ = AddInput(input); - input2_ = AddInput(input); + input1_ = AddInput(input1); + input2_ = AddInput(input2); output_ = AddOutput(output); SetBuiltinOp(BuiltinOperator_ADD, BuiltinOptions_AddOptions, CreateAddOptions(builder_, activation_type).Union()); @@ -70,6 +71,7 @@ float GetTolerance(int min, int max) { TEST(FloatAddOpModel, NoActivation) { FloatAddOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, + {TensorType_FLOAT32, {1, 2, 2, 1}}, {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8}); m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.5}); @@ -77,9 +79,10 @@ TEST(FloatAddOpModel, NoActivation) { EXPECT_THAT(m.GetOutput(), ElementsAreArray({-1.9, 0.4, 1.0, 1.3})); } -TEST(FloatAddOpModel, ActivationRELU1) { - FloatAddOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, - {TensorType_FLOAT32, {}}, ActivationFunctionType_RELU1); +TEST(FloatAddOpModel, ActivationRELU_N1_TO_1) { + FloatAddOpModel m( + {TensorType_FLOAT32, {1, 2, 2, 1}}, {TensorType_FLOAT32, {1, 2, 2, 1}}, + {TensorType_FLOAT32, {}}, ActivationFunctionType_RELU_N1_TO_1); m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8}); m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.5}); m.Invoke(); @@ -91,6 +94,7 @@ TEST(FloatAddOpModel, VariousInputShapes) { {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; for (int i = 0; i < test_shapes.size(); ++i) { FloatAddOpModel m({TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, test_shapes[i]}, {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.5, 1.1, 0.1}); @@ -101,6 +105,23 @@ TEST(FloatAddOpModel, VariousInputShapes) { } } +TEST(FloatAddOpModel, WithBroadcast) { + std::vector> test_shapes = { + {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; + for (int i = 0; i < test_shapes.size(); ++i) { + FloatAddOpModel m({TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, {}}, // always a scalar + {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); + m.PopulateTensor(m.input2(), {0.1}); + m.Invoke(); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-1.9, 0.3, 0.8, 0.9, 1.2, 2.1}))) + << "With shape number " << i; + } +} + TEST(QuantizedAddOpModel, QuantizedTestsNoActivation) { float kQuantizedTolerance = GetTolerance(-1.0, 1.0); std::vector> inputs1 = { @@ -111,6 +132,7 @@ TEST(QuantizedAddOpModel, QuantizedTestsNoActivation) { {0.7, 0.6, 0.6, 0.5}, {-0.2, 0.6, 0.9, -0.1}, {-0.2, 0.6, -0.1, 0.8}}; for (int i = 0; i < inputs1.size(); ++i) { QuantizedAddOpModel m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, + {TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, {TensorType_UINT8, {}, -1.0, 1.0}, ActivationFunctionType_NONE); m.QuantizeAndPopulate(m.input1(), inputs1[i]); @@ -122,7 +144,7 @@ TEST(QuantizedAddOpModel, QuantizedTestsNoActivation) { } } -TEST(QuantizedAddOpModel, QuantizedTestsActivationRELU1) { +TEST(QuantizedAddOpModel, QuantizedTestsActivationRELU_N1_TO_1) { float kQuantizedTolerance = GetTolerance(-1.0, 1.0); std::vector> inputs1 = {{-0.8, 0.2, 0.9, 0.7}, {-0.8, 0.2, 0.7, 0.3}}; @@ -132,8 +154,9 @@ TEST(QuantizedAddOpModel, QuantizedTestsActivationRELU1) { {-0.2, 0.6, -0.1, 0.8}}; for (int i = 0; i < inputs1.size(); ++i) { QuantizedAddOpModel m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, + {TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, {TensorType_UINT8, {}, -1.0, 1.0}, - ActivationFunctionType_RELU1); + ActivationFunctionType_RELU_N1_TO_1); m.QuantizeAndPopulate(m.input1(), inputs1[i]); m.QuantizeAndPopulate(m.input2(), inputs2[i]); m.Invoke(); @@ -149,6 +172,7 @@ TEST(QuantizedAddOpModel, QuantizedVariousInputShapes) { {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; for (int i = 0; i < test_shapes.size(); ++i) { QuantizedAddOpModel m({TensorType_UINT8, test_shapes[i], -3.0, 3.0}, + {TensorType_UINT8, test_shapes[i], -3.0, 3.0}, {TensorType_UINT8, {}, -3.0, 3.0}, ActivationFunctionType_NONE); m.QuantizeAndPopulate(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); @@ -161,6 +185,25 @@ TEST(QuantizedAddOpModel, QuantizedVariousInputShapes) { } } +TEST(QuantizedAddOpModel, QuantizedWithBroadcast) { + float kQuantizedTolerance = GetTolerance(-3.0, 3.0); + std::vector> test_shapes = { + {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; + for (int i = 0; i < test_shapes.size(); ++i) { + QuantizedAddOpModel m({TensorType_UINT8, test_shapes[i], -3.0, 3.0}, + {TensorType_UINT8, {}, -3.0, 3.0}, + {TensorType_UINT8, {}, -3.0, 3.0}, + ActivationFunctionType_NONE); + m.QuantizeAndPopulate(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); + m.QuantizeAndPopulate(m.input2(), {0.1}); + m.Invoke(); + EXPECT_THAT(m.GetDequantizedOutput(), + ElementsAreArray(ArrayFloatNear({-1.9, 0.3, 0.8, 0.9, 1.2, 2.1}, + kQuantizedTolerance))) + << "With shape number " << i; + } +} + } // namespace } // namespace tflite int main(int argc, char** argv) { diff --git a/tensorflow/contrib/lite/kernels/basic_rnn.cc b/tensorflow/contrib/lite/kernels/basic_rnn.cc index 3cee43c68b2a0af5a3fd84b33a980b74bb8f0cb4..2c5074eca3176c7f33a6f051b492dc41333257ed 100644 --- a/tensorflow/contrib/lite/kernels/basic_rnn.cc +++ b/tensorflow/contrib/lite/kernels/basic_rnn.cc @@ -15,14 +15,15 @@ limitations under the License. #include #include #include -#include #include +#include #include #include #include "tensorflow/contrib/lite/builtin_op_data.h" #include "tensorflow/contrib/lite/context.h" #include "tensorflow/contrib/lite/kernels/activation_functor.h" +#include "tensorflow/contrib/lite/kernels/internal/kernel_utils.h" #include "tensorflow/contrib/lite/kernels/op_macros.h" namespace tflite { @@ -76,8 +77,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TfLiteIntArray* output_size_array = TfLiteIntArrayCreate(2); output_size_array->data[0] = batch_size; output_size_array->data[1] = num_units; - TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, output, - output_size_array)); + TF_LITE_ENSURE_OK(context, + context->ResizeTensor(context, output, output_size_array)); return kTfLiteOk; } @@ -101,50 +102,20 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { const int batch_size = input->dims->data[0]; const int num_units = input_weights->dims->data[0]; const int input_size = input->dims->data[1]; - const int input_weights_stride = input_weights->dims->data[1]; - const int recurrent_weights_stride = recurrent_weights->dims->data[1]; - - // For each batch - for (int b = 0; b < batch_size; b++) { - // Initialize the pointer to input, output and bias. - const float* input_ptr_batch = input->data.f + b * input_size; - float* output_ptr_batch = output->data.f + b * num_units; - float* hidden_state_ptr_batch = hidden_state->data.f + b * num_units; - - // Initialize input_weights and recurrent_weights. - const float* input_weights_ptr = input_weights->data.f; - const float* recurrent_weights_ptr = recurrent_weights->data.f; - - // Output = bias - for (int o = 0; o < num_units; o++) { - output_ptr_batch[o] = bias_ptr[o]; - } - - // Output += input * input_weights - for (int o = 0; o < num_units; o++) { - for (int i = 0; i < input_size; i++) { - output_ptr_batch[o] += input_ptr_batch[i] * input_weights_ptr[i]; - } - input_weights_ptr += input_weights_stride; - } - - // Output += recurrent_weights * hidden_state - for (int o = 0; o < num_units; o++) { - for (int h = 0; h < num_units; h++) { - output_ptr_batch[o] += - hidden_state_ptr_batch[h] * recurrent_weights_ptr[h]; - } - recurrent_weights_ptr += recurrent_weights_stride; - } - - // Output = activation(Output) and update hidden_state - for (int o = 0; o < num_units; o++) { - output_ptr_batch[o] = - (ActivationFunctor(params->activation))(output_ptr_batch[o]); - hidden_state_ptr_batch[o] = output_ptr_batch[o]; - } - } + // Initialize the pointer to hidden state. + float* hidden_state_ptr_batch = hidden_state->data.f; + // Initialize the pointer to input and output. + const float* input_ptr_batch = input->data.f; + float* output_ptr_batch = output->data.f; + // Initialize input_weights and recurrent_weights. + const float* input_weights_ptr = input_weights->data.f; + const float* recurrent_weights_ptr = recurrent_weights->data.f; + + kernel_utils::RnnBatchStep(input_ptr_batch, input_weights_ptr, + recurrent_weights_ptr, bias_ptr, input_size, + num_units, batch_size, params->activation, + hidden_state_ptr_batch, output_ptr_batch); return kTfLiteOk; } diff --git a/tensorflow/contrib/lite/kernels/basic_rnn_test.cc b/tensorflow/contrib/lite/kernels/basic_rnn_test.cc index 5ecccb985e91238f1183c8f94a2b5f468758ce55..fa7ef525db47c93f98951604cd04da66196422d7 100644 --- a/tensorflow/contrib/lite/kernels/basic_rnn_test.cc +++ b/tensorflow/contrib/lite/kernels/basic_rnn_test.cc @@ -14,8 +14,8 @@ limitations under the License. ==============================================================================*/ // Unit test for TFLite RNN op. -#include #include +#include #include #include @@ -120,8 +120,7 @@ static float rnn_golden_output[] = { 0.415153, 0.210318, 0, 0, 0, 0, 0, 2.02616, 0, 0.728256, 0.84183, 0.0907453, - 0.628881, 3.58099, 1.49974, 0 -}; + 0.628881, 3.58099, 1.49974, 0}; class RNNOpModel : public SingleOpModel { public: diff --git a/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc b/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc index 0eed680fdcc2afc4bc72be55a5e7722310fa4538..bc438f99c6a72fdbc2794dee03524db6a7523834 100644 --- a/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc +++ b/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc @@ -35,12 +35,14 @@ enum KernelType { struct BatchToSpaceNDContext { BatchToSpaceNDContext(TfLiteContext* context, TfLiteNode* node) { - params = reinterpret_cast(node->builtin_data); input = GetInput(context, node, 0); + block_shape = GetInput(context, node, 1); + crops = GetInput(context, node, 2); output = GetOutput(context, node, 0); } - TfLiteBatchToSpaceNDParams* params; TfLiteTensor* input; + TfLiteTensor* block_shape; + TfLiteTensor* crops; TfLiteTensor* output; }; @@ -48,23 +50,28 @@ struct BatchToSpaceNDContext { // The 4D array need to have exactly 2 spatial dimensions. // TODO(ycling): Support arbitrary dimension in BatchToSpaceND. const int kInputDimensionNum = 4; -const int kOutputDimensionNum = 4; +const int kBlockSizeDimensionNum = 1; const int kSpatialDimensionNum = 2; -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - // The 2nd tensor (block_shape) and the 3rd tensor (crops) are ignored now. - TF_LITE_ENSURE(context, NumInputs(node) >= 1 && NumInputs(node) <= 3); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); +TfLiteStatus ResizeOutputTensor(TfLiteContext* context, + BatchToSpaceNDContext* op_context) { + TfLiteIntArray* input_size = op_context->input->dims; + const int* block_shape = GetTensorData(op_context->block_shape); + const int* crops = GetTensorData(op_context->crops); - BatchToSpaceNDContext op_context(context, node); - TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.input), - kInputDimensionNum); - TF_LITE_ENSURE_EQ(context, op_context.params->num_spatial_dimensions, + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->block_shape), + kBlockSizeDimensionNum); + TF_LITE_ENSURE_EQ(context, op_context->block_shape->dims->data[0], + kSpatialDimensionNum); + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->crops), kSpatialDimensionNum); - TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type); - const TfLiteIntArray* input_size = op_context.input->dims; - const int* block_shape = op_context.params->block_shape; + // TODO(ycling): Add crops as part of calculation. Remove check for a crops + // containing all zeroes. + TF_LITE_ENSURE_EQ(context, crops[0], 0); + TF_LITE_ENSURE_EQ(context, crops[1], 0); + TF_LITE_ENSURE_EQ(context, crops[2], 0); + TF_LITE_ENSURE_EQ(context, crops[3], 0); // Number of batch must be multiple of (block_shape[0] * block_shape[1]). TF_LITE_ENSURE_EQ(context, @@ -76,27 +83,47 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { const int output_width = input_size->data[2] * block_shape[1]; const int output_channel_size = input_size->data[3]; - TfLiteIntArray* output_size = TfLiteIntArrayCreate(kOutputDimensionNum); + TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_size); output_size->data[0] = output_batch_size; output_size->data[1] = output_height; output_size->data[2] = output_width; output_size->data[3] = output_channel_size; - return context->ResizeTensor(context, op_context.output, output_size); + return context->ResizeTensor(context, op_context->output, output_size); +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 3); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + BatchToSpaceNDContext op_context(context, node); + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.input), + kInputDimensionNum); + TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type); + + if (!IsConstantTensor(op_context.block_shape) || + !IsConstantTensor(op_context.crops)) { + SetTensorToDynamic(op_context.output); + return kTfLiteOk; + } + return ResizeOutputTensor(context, &op_context); } template TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { BatchToSpaceNDContext op_context(context, node); - int block_shape_dims_array[1] = {kSpatialDimensionNum}; - Dims<4> block_shape_dims = GetTensorDims(block_shape_dims_array, 1); + // Resize the output tensor if the output tensor is dynamic. + if (IsDynamicTensor(op_context.output)) { + TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); + } -#define TF_LITE_BATCH_TO_SPACE_ND(type, scalar) \ - type::BatchToSpaceND(GetTensorData(op_context.input), \ - GetTensorDims(op_context.input), \ - op_context.params->block_shape, block_shape_dims, \ - GetTensorData(op_context.output), \ +#define TF_LITE_BATCH_TO_SPACE_ND(type, scalar) \ + type::BatchToSpaceND(GetTensorData(op_context.input), \ + GetTensorDims(op_context.input), \ + GetTensorData(op_context.block_shape), \ + GetTensorDims(op_context.block_shape), \ + GetTensorData(op_context.output), \ GetTensorDims(op_context.output)) switch (op_context.input->type) { // Already know in/out types are same. case kTfLiteFloat32: diff --git a/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc b/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc index 3ec4efbebcef9d55d0042d93007018c9f6ee3b58..8485cde1b40066f2070855bca91ea78a9f80e83c 100644 --- a/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc +++ b/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc @@ -26,36 +26,76 @@ using ::testing::ElementsAreArray; class BatchToSpaceNDOpModel : public SingleOpModel { public: - BatchToSpaceNDOpModel(std::initializer_list input_shape, - std::initializer_list block_shape, - std::initializer_list before_crops, - std::initializer_list after_crops) { - input_ = AddInput(TensorType_FLOAT32); - output_ = AddOutput(TensorType_FLOAT32); - SetBuiltinOp(BuiltinOperator_BATCH_TO_SPACE_ND, - BuiltinOptions_BatchToSpaceNDOptions, - CreateBatchToSpaceNDOptions( - builder_, builder_.CreateVector(block_shape), - builder_.CreateVector(before_crops), - builder_.CreateVector(after_crops)) - .Union()); - BuildInterpreter({input_shape}); - } - void SetInput(std::initializer_list data) { PopulateTensor(input_, data); } + void SetBlockShape(std::initializer_list data) { + PopulateTensor(block_shape_, data); + } + + void SetCrops(std::initializer_list data) { + PopulateTensor(crops_, data); + } + std::vector GetOutput() { return ExtractVector(output_); } std::vector GetOutputShape() { return GetTensorShape(output_); } - private: + protected: int input_; + int block_shape_; + int crops_; int output_; }; -TEST(BatchToSpaceNDOpTest, SimpleTest) { - BatchToSpaceNDOpModel m({4, 2, 2, 1}, {2, 2}, {0, 0}, {0, 0}); +// Tests case where block_shape and crops are const tensors. +// +// Example usage is as follows: +// BatchToSpaceNDOpConstModel m(input_shape, block_shape, crops); +// m.SetInput(input_data); +// m.Invoke(); +class BatchToSpaceNDOpConstModel : public BatchToSpaceNDOpModel { + public: + BatchToSpaceNDOpConstModel(std::initializer_list input_shape, + std::initializer_list block_shape, + std::initializer_list crops) { + input_ = AddInput(TensorType_FLOAT32); + block_shape_ = AddConstInput(TensorType_INT32, block_shape, {2}); + crops_ = AddConstInput(TensorType_INT32, crops, {2, 2}); + output_ = AddOutput(TensorType_FLOAT32); + + SetBuiltinOp(BuiltinOperator_BATCH_TO_SPACE_ND, + BuiltinOptions_BatchToSpaceNDOptions, + CreateBatchToSpaceNDOptions(builder_).Union()); + BuildInterpreter({input_shape}); + } +}; + +// Tests case where block_shape and crops are non-const tensors. +// +// Example usage is as follows: +// BatchToSpaceNDOpDynamicModel m(input_shape); +// m.SetInput(input_data); +// m.SetBlockShape(block_shape); +// m.SetPaddings(crops); +// m.Invoke(); +class BatchToSpaceNDOpDynamicModel : public BatchToSpaceNDOpModel { + public: + BatchToSpaceNDOpDynamicModel(std::initializer_list input_shape) { + input_ = AddInput(TensorType_FLOAT32); + block_shape_ = AddInput(TensorType_INT32); + crops_ = AddInput(TensorType_INT32); + output_ = AddOutput(TensorType_FLOAT32); + + SetBuiltinOp(BuiltinOperator_BATCH_TO_SPACE_ND, + BuiltinOptions_BatchToSpaceNDOptions, + CreateBatchToSpaceNDOptions(builder_).Union()); + BuildInterpreter({input_shape, {2}, {2, 2}}); + } +}; + +TEST(BatchToSpaceNDOpTest, SimpleConstTest) { + BatchToSpaceNDOpConstModel m({4, 2, 2, 1}, {2, 2}, {0, 0, 0, 0}); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); @@ -63,11 +103,35 @@ TEST(BatchToSpaceNDOpTest, SimpleTest) { 4, 8, 11, 15, 12, 16})); } +TEST(BatchToSpaceNDOpTest, SimpleDynamicTest) { + BatchToSpaceNDOpDynamicModel m({4, 2, 2, 1}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + m.SetBlockShape({2, 2}); + m.SetCrops({0, 0, 0, 0}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 5, 2, 6, 9, 13, 10, 14, 3, 7, + 4, 8, 11, 15, 12, 16})); +} + TEST(BatchToSpaceNDOpTest, InvalidShapeTest) { - EXPECT_DEATH(BatchToSpaceNDOpModel({3, 2, 2, 1}, {2, 2}, {0, 0}, {0, 0}), + EXPECT_DEATH(BatchToSpaceNDOpConstModel({3, 2, 2, 1}, {2, 2}, {0, 0, 0, 0}), "Cannot allocate tensors"); } +TEST(BatchToSpaceNDOpTest, InvalidCropsConstTest) { + EXPECT_DEATH(BatchToSpaceNDOpConstModel({3, 2, 2, 1}, {2, 2}, {0, 0, 0, 1}), + "1 != 0"); +} + +TEST(BatchToSpaceNDOpTest, InvalidCropsDynamicTest) { + BatchToSpaceNDOpDynamicModel m({4, 2, 2, 1}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + m.SetBlockShape({2, 2}); + m.SetCrops({0, 0, 1, 0}); + EXPECT_DEATH(m.Invoke(), "1 != 0"); +} + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn.cc new file mode 100644 index 0000000000000000000000000000000000000000..aa24c1f34cd1e8c02a6a75b62fbe5f3c629498ca --- /dev/null +++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn.cc @@ -0,0 +1,205 @@ +/* Copyright 2017 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 +#include +#include +#include +#include +#include +#include + +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/activation_functor.h" +#include "tensorflow/contrib/lite/kernels/internal/kernel_utils.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace bidirectional_sequence_rnn { + +constexpr int kInputTensor = 0; +// Forward and backward cell tensors. +constexpr int kFwWeightsTensor = 1; +constexpr int kFwRecurrentWeightsTensor = 2; +constexpr int kFwBiasTensor = 3; +constexpr int kBwWeightsTensor = 4; +constexpr int kBwRecurrentWeightsTensor = 5; +constexpr int kBwBiasTensor = 6; +// State and output tensors. +constexpr int kFwHiddenStateTensor = 0; +constexpr int kFwOutputTensor = 1; +constexpr int kBwHiddenStateTensor = 2; +constexpr int kBwOutputTensor = 3; + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + // Check we have all the inputs and outputs we need. + TF_LITE_ENSURE_EQ(context, node->inputs->size, 7); + TF_LITE_ENSURE_EQ(context, node->outputs->size, 4); + + TfLiteTensor* input = &context->tensors[node->inputs->data[kInputTensor]]; + TfLiteTensor* fw_input_weights = + &context->tensors[node->inputs->data[kFwWeightsTensor]]; + TfLiteTensor* fw_recurrent_weights = + &context->tensors[node->inputs->data[kFwRecurrentWeightsTensor]]; + TfLiteTensor* fw_bias = &context->tensors[node->inputs->data[kFwBiasTensor]]; + TfLiteTensor* bw_input_weights = + &context->tensors[node->inputs->data[kBwWeightsTensor]]; + TfLiteTensor* bw_recurrent_weights = + &context->tensors[node->inputs->data[kBwRecurrentWeightsTensor]]; + TfLiteTensor* bw_bias = &context->tensors[node->inputs->data[kBwBiasTensor]]; + + // Check all the parameters of tensor match within themselves and match the + // input configuration. + const int batch_size = input->dims->data[0]; + const int max_time = input->dims->data[1]; + const int fw_num_units = fw_input_weights->dims->data[0]; + const int bw_num_units = bw_input_weights->dims->data[0]; + TF_LITE_ASSERT_EQ(input->dims->data[2], fw_input_weights->dims->data[1]); + TF_LITE_ASSERT_EQ(input->dims->data[2], bw_input_weights->dims->data[1]); + TF_LITE_ASSERT_EQ(fw_input_weights->dims->data[0], fw_bias->dims->data[0]); + TF_LITE_ASSERT_EQ(bw_input_weights->dims->data[0], bw_bias->dims->data[0]); + TF_LITE_ASSERT_EQ(fw_recurrent_weights->dims->data[0], + fw_bias->dims->data[0]); + TF_LITE_ASSERT_EQ(bw_recurrent_weights->dims->data[1], + bw_bias->dims->data[0]); + + TfLiteTensor* fw_output = + &context->tensors[node->outputs->data[kFwOutputTensor]]; + TfLiteTensor* bw_output = + &context->tensors[node->outputs->data[kBwOutputTensor]]; + + // Resize hidden states. + TfLiteIntArray* fw_hidden_state_size_array = TfLiteIntArrayCreate(2); + fw_hidden_state_size_array->data[0] = batch_size; + fw_hidden_state_size_array->data[1] = fw_num_units; + TfLiteTensor* fw_hidden_state = + &context->tensors[node->outputs->data[kFwHiddenStateTensor]]; + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, fw_hidden_state, + fw_hidden_state_size_array)); + + TfLiteIntArray* bw_hidden_state_size_array = TfLiteIntArrayCreate(2); + bw_hidden_state_size_array->data[0] = batch_size; + bw_hidden_state_size_array->data[1] = fw_num_units; + TfLiteTensor* bw_hidden_state = + &context->tensors[node->outputs->data[kBwHiddenStateTensor]]; + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, bw_hidden_state, + bw_hidden_state_size_array)); + + // Mark hidden states as a persistent tensor. + fw_hidden_state->allocation_type = kTfLiteArenaRwPersistent; + bw_hidden_state->allocation_type = kTfLiteArenaRwPersistent; + + // Resize outputs. + TfLiteIntArray* fw_output_size_array = TfLiteIntArrayCreate(3); + fw_output_size_array->data[0] = batch_size; + fw_output_size_array->data[1] = max_time; + fw_output_size_array->data[2] = fw_num_units; + TF_LITE_ENSURE_OK( + context, context->ResizeTensor(context, fw_output, fw_output_size_array)); + TfLiteIntArray* bw_output_size_array = TfLiteIntArrayCreate(3); + bw_output_size_array->data[0] = batch_size; + bw_output_size_array->data[1] = max_time; + bw_output_size_array->data[2] = bw_num_units; + TF_LITE_ENSURE_OK( + context, context->ResizeTensor(context, bw_output, bw_output_size_array)); + + return kTfLiteOk; +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + auto* params = reinterpret_cast(node->builtin_data); + + TfLiteTensor* input = &context->tensors[node->inputs->data[kInputTensor]]; + TfLiteTensor* fw_input_weights = + &context->tensors[node->inputs->data[kFwWeightsTensor]]; + TfLiteTensor* fw_recurrent_weights = + &context->tensors[node->inputs->data[kFwRecurrentWeightsTensor]]; + TfLiteTensor* fw_bias = &context->tensors[node->inputs->data[kFwBiasTensor]]; + TfLiteTensor* fw_hidden_state = + &context->tensors[node->outputs->data[kFwHiddenStateTensor]]; + TfLiteTensor* fw_output = + &context->tensors[node->outputs->data[kFwOutputTensor]]; + + TfLiteTensor* bw_input_weights = + &context->tensors[node->inputs->data[kBwWeightsTensor]]; + TfLiteTensor* bw_recurrent_weights = + &context->tensors[node->inputs->data[kBwRecurrentWeightsTensor]]; + TfLiteTensor* bw_bias = &context->tensors[node->inputs->data[kBwBiasTensor]]; + TfLiteTensor* bw_hidden_state = + &context->tensors[node->outputs->data[kBwHiddenStateTensor]]; + TfLiteTensor* bw_output = + &context->tensors[node->outputs->data[kBwOutputTensor]]; + + const int batch_size = input->dims->data[0]; + const int max_time = input->dims->data[1]; + const int input_size = input->dims->data[2]; + + const int fw_num_units = fw_input_weights->dims->data[0]; + const float* fw_bias_ptr = fw_bias->data.f; + const float* fw_input_weights_ptr = fw_input_weights->data.f; + const float* fw_recurrent_weights_ptr = fw_recurrent_weights->data.f; + + const int bw_num_units = bw_input_weights->dims->data[0]; + const float* bw_bias_ptr = bw_bias->data.f; + const float* bw_input_weights_ptr = bw_input_weights->data.f; + const float* bw_recurrent_weights_ptr = bw_recurrent_weights->data.f; + + for (int b = 0; b < batch_size; b++) { + // Forward cell. + float* fw_hidden_state_ptr_batch = + fw_hidden_state->data.f + b * fw_num_units; + for (int s = 0; s < max_time; s++) { + const float* input_ptr_batch = + input->data.f + b * input_size * max_time + s * input_size; + float* output_ptr_batch = + fw_output->data.f + b * fw_num_units * max_time + s * fw_num_units; + + kernel_utils::RnnBatchStep( + input_ptr_batch, fw_input_weights_ptr, fw_recurrent_weights_ptr, + fw_bias_ptr, input_size, fw_num_units, /*batch_size=*/1, + params->activation, fw_hidden_state_ptr_batch, output_ptr_batch); + } + // Backward cell. + float* bw_hidden_state_ptr_batch = + bw_hidden_state->data.f + b * bw_num_units; + for (int s = max_time - 1; s >= 0; s--) { + const float* input_ptr_batch = + input->data.f + b * input_size * max_time + s * input_size; + float* output_ptr_batch = + bw_output->data.f + b * bw_num_units * max_time + s * bw_num_units; + + kernel_utils::RnnBatchStep( + input_ptr_batch, bw_input_weights_ptr, bw_recurrent_weights_ptr, + bw_bias_ptr, input_size, bw_num_units, /*batch_size=*/1, + params->activation, bw_hidden_state_ptr_batch, output_ptr_batch); + } + } + return kTfLiteOk; +} + +} // namespace bidirectional_sequence_rnn + +TfLiteRegistration* Register_BIDIRECTIONAL_SEQUENCE_RNN() { + static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, + bidirectional_sequence_rnn::Prepare, + bidirectional_sequence_rnn::Eval}; + return &r; +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn_test.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..12f4ff97cfd90e3a6894a24d15fcbc356f96cde2 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_rnn_test.cc @@ -0,0 +1,931 @@ +/* Copyright 2017 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. +==============================================================================*/ +// Unit test for TFLite Bidirectional RNN op. + +#include +#include + +#include +#include +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +static float rnn_input[] = { + 0.23689353, 0.285385, 0.037029743, -0.19858193, -0.27569133, + 0.43773448, 0.60379338, 0.35562468, -0.69424844, -0.93421471, + -0.87287879, 0.37144363, -0.62476718, 0.23791671, 0.40060222, + 0.1356622, -0.99774903, -0.98858172, -0.38952237, -0.47685933, + 0.31073618, 0.71511042, -0.63767755, -0.31729108, 0.33468103, + 0.75801885, 0.30660987, -0.37354088, 0.77002847, -0.62747043, + -0.68572164, 0.0069220066, 0.65791464, 0.35130811, 0.80834007, + -0.61777675, -0.21095741, 0.41213346, 0.73784804, 0.094794154, + 0.47791874, 0.86496925, -0.53376222, 0.85315156, 0.10288584, + 0.86684, -0.011186242, 0.10513687, 0.87825835, 0.59929144, + 0.62827742, 0.18899453, 0.31440187, 0.99059987, 0.87170351, + -0.35091716, 0.74861872, 0.17831337, 0.2755419, 0.51864719, + 0.55084288, 0.58982027, -0.47443086, 0.20875752, -0.058871567, + -0.66609079, 0.59098077, 0.73017097, 0.74604273, 0.32882881, + -0.17503482, 0.22396147, 0.19379807, 0.29120302, 0.077113032, + -0.70331609, 0.15804303, -0.93407321, 0.40182066, 0.036301374, + 0.66521823, 0.0300982, -0.7747041, -0.02038002, 0.020698071, + -0.90300065, 0.62870288, -0.23068321, 0.27531278, -0.095755219, + -0.712036, -0.17384434, -0.50593495, -0.18646687, -0.96508682, + 0.43519354, 0.14744234, 0.62589407, 0.1653645, -0.10651493, + -0.045277178, 0.99032974, -0.88255352, -0.85147917, 0.28153265, + 0.19455957, -0.55479527, -0.56042433, 0.26048636, 0.84702539, + 0.47587705, -0.074295521, -0.12287641, 0.70117295, 0.90532446, + 0.89782166, 0.79817224, 0.53402734, -0.33286154, 0.073485017, + -0.56172788, -0.044897556, 0.89964068, -0.067662835, 0.76863563, + 0.93455386, -0.6324693, -0.083922029}; + +static float rnn_golden_fw_output[] = { + 0.496726, 0, 0.965996, 0, 0.0584254, 0, + 0, 0.12315, 0, 0, 0.612266, 0.456601, + 0, 0.52286, 1.16099, 0.0291232, + + 0, 0, 0.524901, 0, 0, 0, + 0, 1.02116, 0, 1.35762, 0, 0.356909, + 0.436415, 0.0355727, 0, 0, + + 0, 0, 0, 0.262335, 0, 0, + 0, 1.33992, 0, 2.9739, 0, 0, + 1.31914, 2.66147, 0, 0, + + 0.942568, 0, 0, 0, 0.025507, 0, + 0, 0, 0.321429, 0.569141, 1.25274, 1.57719, + 0.8158, 1.21805, 0.586239, 0.25427, + + 1.04436, 0, 0.630725, 0, 0.133801, 0.210693, + 0.363026, 0, 0.533426, 0, 1.25926, 0.722707, + 0, 1.22031, 1.30117, 0.495867, + + 0.222187, 0, 0.72725, 0, 0.767003, 0, + 0, 0.147835, 0, 0, 0, 0.608758, + 0.469394, 0.00720298, 0.927537, 0, + + 0.856974, 0.424257, 0, 0, 0.937329, 0, + 0, 0, 0.476425, 0, 0.566017, 0.418462, + 0.141911, 0.996214, 1.13063, 0, + + 0.967899, 0, 0, 0, 0.0831304, 0, + 0, 1.00378, 0, 0, 0, 1.44818, + 1.01768, 0.943891, 0.502745, 0, + + 0.940135, 0, 0, 0, 0, 0, + 0, 2.13243, 0, 0.71208, 0.123918, 1.53907, + 1.30225, 1.59644, 0.70222, 0, + + 0.804329, 0, 0.430576, 0, 0.505872, 0.509603, + 0.343448, 0, 0.107756, 0.614544, 1.44549, 1.52311, + 0.0454298, 0.300267, 0.562784, 0.395095, + + 0.228154, 0, 0.675323, 0, 1.70536, 0.766217, + 0, 0, 0, 0.735363, 0.0759267, 1.91017, + 0.941888, 0, 0, 0, + + 0, 0, 1.5909, 0, 0, 0, + 0, 0.5755, 0, 0.184687, 0, 1.56296, + 0.625285, 0, 0, 0, + + 0, 0, 0.0857888, 0, 0, 0, + 0, 0.488383, 0.252786, 0, 0, 0, + 1.02817, 1.85665, 0, 0, + + 0.00981836, 0, 1.06371, 0, 0, 0, + 0, 0, 0, 0.290445, 0.316406, 0, + 0.304161, 1.25079, 0.0707152, 0, + + 0.986264, 0.309201, 0, 0, 0, 0, + 0, 1.64896, 0.346248, 0, 0.918175, 0.78884, + 0.524981, 1.92076, 2.07013, 0.333244, + + 0.415153, 0.210318, 0, 0, 0, 0, + 0, 2.02616, 0, 0.728256, 0.84183, 0.0907453, + 0.628881, 3.58099, 1.49974, 0}; + +static float rnn_golden_bw_output[] = { + 0.496726, 0, 1.00883, 0, 0.0584256, 0, 0, + 0.236412, 0, 0, 0.612267, 0.487726, 0, 0.54883, + 1.16099, 0.0291233, 0, 0, 0.428302, 0, 0, + 0, 0, 1.13262, 0, 1.64415, 0, 0.311249, + 0.570804, 0.259696, 0, 0, 0, 0, 0, + 0.262334, 0, 0, 0, 1.23781, 0, 2.86532, + 0, 0, 1.34389, 2.76409, 0, 0, 1.03969, + 0, 0.00410865, 0, 0.0470295, 0, 0, 0, + 0.371556, 0.27175, 1.36614, 1.63956, 0.683887, 1.06176, 0.719552, + 0.301314, 0.971195, 0, 0.697143, 0, 0.215219, 0.210693, + 0.363027, 0, 0.501283, 0, 1.13399, 0.623774, 0, + 1.09851, 1.33313, 0.470441, 0.210965, 0, 0.664178, 0, + 0.839686, 0, 0, 0.147834, 0, 0, 0, + 0.58786, 0.490128, 0, 0.905806, 0, 0.932134, 0.424257, + 0, 0, 0.860629, 0, 0, 0, 0.476425, + 0, 0.566017, 0.513721, 0.207341, 1.09508, 1.08385, 0, + 0.973787, 0, 0, 0, 0, 0, 0, + 1.20698, 0, 0, 0, 1.56135, 1.12369, 0.99588, + 0.459803, 0, 0.915854, 0, 0, 0, 0, + 0, 0, 2.03206, 0, 0.773264, 0.267228, 1.55012, + 1.202, 1.51611, 0.701202, 0, 0.725088, 0, 0.509069, + 0, 0.671349, 0.581129, 0.343447, 0, 0.107755, 0.611838, + 1.4331, 1.55871, 0.015242, 0.140624, 0.492562, 0.395095, 0.147722, + 0, 0.784925, 0, 1.65477, 0.715257, 0, 0, + 0, 0.685024, 0, 1.89505, 1.00037, 0, 0, + 0, 0, 0, 1.52659, 0, 0, 0, + 0, 0.618583, 0, 0.11115, 0, 1.37194, 0.630225, + 0, 0, 0, 0, 0, 0.0322124, 0, + 0, 0, 0, 0.430834, 0.252786, 0, 0, + 0, 0.991297, 1.98451, 0, 0, 0.111511, 0, + 1.05513, 0, 0, 0, 0, 0, 0, + 0.290445, 0.412559, 0.0429958, 0.256564, 1.27858, 0.289948, 0, + 1.01693, 0.327141, 0, 0, 0, 0, 0, + 1.83508, 0.346248, 0, 0.961535, 0.790026, 0.552203, 2.13457, + 2.19233, 0.333244, 0.316526, 0.179398, 0, 0, 0, + 0, 0, 1.86126, 0, 0.728256, 0.750013, 0.011861, + 0.576383, 3.38891, 1.29273, 0}; + +constexpr std::initializer_list weights = { + 0.461459, 0.153381, 0.529743, -0.00371218, 0.676267, -0.211346, + 0.317493, 0.969689, -0.343251, 0.186423, 0.398151, 0.152399, + 0.448504, 0.317662, 0.523556, -0.323514, 0.480877, 0.333113, + -0.757714, -0.674487, -0.643585, 0.217766, -0.0251462, 0.79512, + -0.595574, -0.422444, 0.371572, -0.452178, -0.556069, -0.482188, + -0.685456, -0.727851, 0.841829, 0.551535, -0.232336, 0.729158, + -0.00294906, -0.69754, 0.766073, -0.178424, 0.369513, -0.423241, + 0.548547, -0.0152023, -0.757482, -0.85491, 0.251331, -0.989183, + 0.306261, -0.340716, 0.886103, -0.0726757, -0.723523, -0.784303, + 0.0354295, 0.566564, -0.485469, -0.620498, 0.832546, 0.697884, + -0.279115, 0.294415, -0.584313, 0.548772, 0.0648819, 0.968726, + 0.723834, -0.0080452, -0.350386, -0.272803, 0.115121, -0.412644, + -0.824713, -0.992843, -0.592904, -0.417893, 0.863791, -0.423461, + -0.147601, -0.770664, -0.479006, 0.654782, 0.587314, -0.639158, + 0.816969, -0.337228, 0.659878, 0.73107, 0.754768, -0.337042, + 0.0960841, 0.368357, 0.244191, -0.817703, -0.211223, 0.442012, + 0.37225, -0.623598, -0.405423, 0.455101, 0.673656, -0.145345, + -0.511346, -0.901675, -0.81252, -0.127006, 0.809865, -0.721884, + 0.636255, 0.868989, -0.347973, -0.10179, -0.777449, 0.917274, + 0.819286, 0.206218, -0.00785118, 0.167141, 0.45872, 0.972934, + -0.276798, 0.837861, 0.747958, -0.0151566, -0.330057, -0.469077, + 0.277308, 0.415818}; + +static float endtoend_input[] = { + 0.996808, 0.060710, 0.981855, 0.570017, 0.525164, 0.796859, 0.696547, + 0.505925, 0.991844, 0.461208, 0.949371, 0.027624, 0.539236, 0.841854, + 0.915222, 0.538569, 0.069375, 0.237905, 0.903700, 0.441703, 0.536196, + 0.402724, 0.761635, 0.025063, 0.082592, 0.688245, 0.239310, 0.256931, + 0.658900, 0.105695, 0.301983, 0.655708, 0.166405, 0.283837, 0.225725, + 0.691569, 0.080696, 0.922272, 0.197494, 0.072540, 0.383481, 0.146865, + 0.100163, 0.922717, 0.988720, 0.015386, 0.461286, 0.058095, 0.253290, + 0.364986, 0.499797, 0.789487, 0.767709, 0.261433, 0.814549, 0.850302, + 0.949678, 0.053859, 0.107233, 0.608577, 0.159554, 0.409215, 0.264285, + 0.325960, 0.693053, 0.490011, 0.017529, 0.773749, 0.412283, 0.215023, + 0.846288, 0.795764, 0.361889, 0.946452, 0.718481, 0.350608, 0.961837, + 0.179767, 0.408703, 0.215128, 0.544753, 0.908500, 0.004614, 0.312462, + 0.169933, 0.819163, 0.162764, 0.119611, 0.873022, 0.269997, 0.728188, + 0.032576, 0.679212, 0.992474, 0.358536, 0.372265, 0.482484, 0.376065, + 0.146014, 0.894767, 0.591088, 0.992302, 0.690531, 0.952977, 0.938754, + 0.409012, 0.303585, 0.900591, 0.588780, 0.712287, 0.115719, 0.133533, + 0.620788, 0.120334, 0.445995, 0.790720, 0.939497, 0.608759, 0.910331, + 0.812519, 0.878756, 0.638519, 0.845096, 0.557968, 0.630993, 0.203632, + 0.930233, 0.113477, 0.579697, 0.076247, 0.008244, 0.170785, 0.068549, + 0.698776, 0.123761, 0.007303, 0.107788, 0.427346, 0.907894, 0.696568, + 0.139633, 0.023613, 0.830100, 0.760421, 0.143947, 0.276096, 0.551141, + 0.083444, 0.884855, 0.461472, 0.895963, 0.763611, 0.099992, 0.741059, + 0.321579, 0.730984, 0.944691, 0.251812, 0.844461, 0.524388, 0.328059, + 0.852706, 0.695172, 0.396607, 0.551482, 0.818934, 0.403910, 0.659270, + 0.246280, 0.311804, 0.355838, 0.385913, 0.335418, 0.185938, 0.146334, + 0.479364, 0.462034, 0.697475, 0.562808, 0.346888, 0.158948, 0.458771, + 0.110499, 0.258939, 0.199830, 0.432078, 0.989924, 0.144521, 0.683890, + 0.834385, 0.668908, 0.011949, 0.687091, 0.364081, 0.408556, 0.238572, + 0.183015, 0.812466, 0.897842, 0.429294, 0.124271, 0.253680, 0.815207, + 0.459688, 0.439618, 0.961541, 0.939053, 0.901651, 0.659016, 0.501861, + 0.248539, 0.817964, 0.960632, 0.359038, 0.076903, 0.160462, 0.791117, + 0.066826, 0.304983, 0.475007, 0.901211, 0.973891, 0.486955, 0.588302, + 0.337972, 0.895512, 0.826874, 0.520987, 0.707978, 0.724716, 0.950281, + 0.832249, 0.978396, 0.765488, 0.291937, 0.418014, 0.727029, 0.230990, + 0.319665, 0.386045, 0.732850, 0.568204, 0.204009, 0.693482, 0.927242, + 0.280912, 0.853944, 0.718359, 0.347738, 0.158927, 0.193366, 0.248950, + 0.132818, 0.680321, 0.837252, 0.470790, 0.575833, 0.664126, 0.991777, + 0.283811, 0.388843, 0.942058, 0.116060, 0.367239, 0.707546, 0.407997, + 0.785253, 0.434575, 0.638986, 0.104917, 0.820620, 0.371837, 0.673121, + 0.024629, 0.065319, 0.600363, 0.305541, 0.919263, 0.318722, 0.653279, + 0.078190, 0.512088, 0.902229, 0.211009, 0.192409, 0.739480, 0.681799, + 0.768242, 0.403607, 0.673576, 0.052052, 0.792450, 0.615634, 0.168112, + 0.159689, 0.323180, 0.576109, 0.944941, 0.757755, 0.215095, 0.049858, + 0.578375, 0.586932, 0.722979, 0.603003, 0.652251, 0.323343, 0.908544, + 0.571514, 0.642065, 0.561823, 0.649704, 0.154153, 0.464051, 0.860713, + 0.346562, 0.203532, 0.542512, 0.114804, 0.607139, 0.216088, 0.166856, + 0.399588, 0.831722, 0.334968, 0.559277, 0.154902, 0.911077, 0.504218, + 0.912656, 0.126172, 0.554076, 0.491031, 0.713104, 0.277055, 0.094034, + 0.365355, 0.600398, 0.002578, 0.936869, 0.242463, 0.564401, 0.586574, + 0.396616, 0.028452, 0.447287, 0.743178, 0.231984, 0.989799, 0.857982, + 0.839122, 0.205887, 0.024838, 0.238711, 0.037608, 0.359806, 0.797987, + 0.192510, 0.270883, 0.302205, 0.105166, 0.397055, 0.856281, 0.596197, + 0.110160, 0.133336, 0.690231, 0.475515, 0.733734, 0.692809, 0.412384, + 0.976196, 0.257209, 0.998958, 0.372812, 0.285661, 0.446245, 0.115990, + 0.517645, 0.436044, 0.973972, 0.356767, 0.641930, 0.998810, 0.595478, + 0.679539, 0.358617, 0.393465, 0.872049, 0.629500, 0.695670, 0.977215, + 0.026555, 0.551951, 0.573412, 0.136715, 0.685287, 0.263643, 0.612229, + 0.419020, 0.956451, 0.024613, 0.395216, 0.213661, 0.023572, 0.768029, + 0.499322, 0.469816, 0.884019, 0.016967, 0.905860, 0.857991, 0.373734, + 0.547791, 0.856802, 0.969211, 0.227330, 0.215418, 0.362676, 0.099378, + 0.844918, 0.058346, 0.076594, 0.871473, 0.610297, 0.650006, 0.008188, + 0.295583, 0.913648, 0.620417, 0.714603, 0.870100, 0.645031, 0.109820, + 0.083760, 0.668602, 0.877849, 0.583082, 0.138419, 0.761868, 0.600049, + 0.044279, 0.619859, 0.973783, 0.592069, 0.476661, 0.942994, 0.819399, + 0.692079, 0.305670, 0.918778, 0.536997, 0.364016, 0.995371, 0.408470, + 0.974313, 0.645377, 0.416658, 0.269896, 0.559025, 0.037075, 0.984499, + 0.429125, 0.682105, 0.094319, 0.512885, 0.350707, 0.972168, 0.095967, + 0.489126, 0.734035, 0.696016, 0.533405, 0.353894, 0.669799, 0.125474, + 0.830555, 0.612793, 0.944873, 0.522634, 0.918463, 0.863651, 0.059631, + 0.282479, 0.859022, 0.468101, 0.256791, 0.504398, 0.884758, 0.526687, + 0.063423, 0.921833, 0.511186, 0.492548, 0.603939, 0.605505, 0.005433, + 0.954646, 0.577673, 0.101400, 0.443772, 0.311708, 0.797417, 0.977176, + 0.665602, 0.467216, 0.102650, 0.496157, 0.080009, 0.047524, 0.018791, + 0.998471, 0.911174, 0.078422, 0.280950, 0.770196, 0.546523, 0.537741, + 0.274594, 0.431281, 0.064428, 0.338017, 0.353115, 0.575615, 0.830565, + 0.957053, 0.181120, 0.835998, 0.911699, 0.758793, 0.937398, 0.355471, + 0.070501, 0.734815, 0.332647, 0.736103, 0.202031, 0.435297, 0.232261, + 0.282039, 0.482821, 0.251052, 0.280511, 0.393995, 0.329474, 0.561460, + 0.164191, 0.875997, 0.099202, 0.438785, 0.307278, 0.163630, 0.776802, + 0.660393, 0.739244, 0.607367, 0.617446, 0.920364, 0.443365, 0.529145, + 0.679157, 0.380763, 0.884616, 0.749658, 0.115578, 0.217263, 0.485761, + 0.317609, 0.652560, 0.718021, 0.599648, 0.135381, 0.969073, 0.880159, + 0.529376, 0.298547, 0.441619, 0.693567, 0.174544, 0.540821, 0.132351, + 0.481822, 0.704450, 0.909153, 0.142215, 0.443695, 0.516520, 0.759661, + 0.364059, 0.959885, 0.288806, 0.043216, 0.340648, 0.173422, 0.792874, + 0.456226, 0.390685, 0.278634, 0.773834, 0.043245, 0.996656, 0.373483, + 0.178625, 0.965729, 0.253641, 0.708001, 0.264276, 0.695260, 0.401568, + 0.438820, 0.236081, 0.533919, 0.920642, 0.940531, 0.443072, 0.062857, + 0.384226, 0.959592, 0.822518, 0.748285, 0.919477, 0.111325, 0.791501, + 0.260124, 0.284747, 0.584375, 0.716350, 0.675431, 0.863009, 0.490184, + 0.718676, 0.859665, 0.863666, 0.897301, 0.825393, 0.117308, 0.605302, + 0.089669, 0.812568, 0.006870, 0.528489, 0.048649, 0.540788, 0.449131, + 0.989180, 0.983860, 0.511988, 0.373407, 0.943452, 0.334506, 0.121692, + 0.862929, 0.445831, 0.913193, 0.123053, 0.730578, 0.497568, 0.839402, + 0.406009, 0.360577, 0.329586, 0.124685, 0.220241, 0.193253, 0.021986, + 0.045634, 0.310560, 0.627288, 0.135303, 0.123128, 0.634158, 0.663792, + 0.171777, 0.174946, 0.112923, 0.160958, 0.158806, 0.624911, 0.534364, + 0.102259, 0.959418, 0.656056, 0.965187, 0.405249, 0.569249, 0.088240, + 0.135827, 0.066817, 0.927642, 0.541836, 0.427393, 0.257229, 0.666520, + 0.647634, 0.450481, 0.688506, 0.693269, 0.761042, 0.315794, 0.828572, + 0.884170, 0.949952, 0.492364, 0.055947, 0.124898, 0.605288, 0.216905, + 0.283705, 0.230199, 0.751269, 0.385963, 0.189616, 0.407326, 0.351151, + 0.594865, 0.976575, 0.439391, 0.730692, 0.043392, 0.367033, 0.272527, + 0.470785, 0.624261, 0.939048, 0.118419, 0.074743, 0.627554, 0.811688, + 0.835784, 0.943348, 0.640260, 0.719954, 0.893300, 0.132625, 0.775901, + 0.018199, 0.737913, 0.992806, 0.301903, 0.968111, 0.744076, 0.687867, + 0.157728, 0.151401, 0.039017, 0.752593, 0.127976, 0.478408, 0.483284, + 0.171368, 0.845441, 0.755811, 0.642153, 0.469702, 0.694859, 0.760572, + 0.544445, 0.322413, 0.572260, 0.380229, 0.265761, 0.212521, 0.100183, + 0.159062, 0.345146, 0.876084, 0.177261, 0.083058, 0.868891, 0.479164, + 0.051169, 0.612966, 0.167030, 0.208897, 0.764367, 0.206048, 0.961490, + 0.892343, 0.684456, 0.444774, 0.063711, 0.529896, 0.200585, 0.705863, + 0.999598, 0.895444, 0.466435, 0.544043, 0.217857, 0.038696, 0.924272, + 0.483618, 0.251217, 0.024455, 0.642680, 0.596362, 0.900539, 0.819941, + 0.679420, 0.769430, 0.299105, 0.730590, 0.382396, 0.466135, 0.939487, + 0.146763, 0.672183, 0.900977, 0.039106, 0.356638, 0.345750, 0.102817, + 0.886535, 0.546336, 0.808681, 0.886133, 0.441780, 0.275116, 0.430176, + 0.659637, 0.313812, 0.354448, 0.143255, 0.565028, 0.378903, 0.785935, + 0.161391, 0.279443, 0.605876, 0.840811, 0.048873, 0.904980, 0.571401, + 0.431269, 0.371115, 0.510887, 0.578032, 0.043298, 0.411864, 0.617138, + 0.399936, 0.757614, 0.719955, 0.286471, 0.303950, 0.528636, 0.172604, + 0.745730, 0.803752, 0.602780, 0.405367, 0.117564, 0.957228, 0.548622, + 0.682592, 0.336131, 0.334557, 0.843983, 0.615574, 0.940433, 0.684794, + 0.664447, 0.845413, 0.256194, 0.095715, 0.216529, 0.767082, 0.673747, + 0.259827, 0.178946, 0.290885, 0.659763, 0.936560, 0.010840, 0.946234, + 0.240510, 0.539476, 0.118838, 0.986240, 0.343228, 0.721618, 0.391606, + 0.460792, 0.678846, 0.940228, 0.143384, 0.014977, 0.274785, 0.987367, + 0.630551, 0.215218, 0.672161, 0.294998, 0.060631, 0.928355, 0.390713, + 0.277160, 0.695436, 0.064460, 0.536987, 0.874382, 0.355345, 0.196751, + 0.810942, 0.366185, 0.142985, 0.051452, 0.905661, 0.261823, 0.037691, + 0.248889, 0.983441, 0.429297, 0.709681, 0.662286, 0.369525, 0.853066, + 0.677263, 0.644310, 0.840433, 0.307814, 0.859528, 0.512593, 0.602812, + 0.920160, 0.440948, 0.993525, 0.197320, 0.136384, 0.057984, 0.734307, + 0.010766, 0.413329, 0.931058, 0.821707, 0.779514, 0.074043, 0.873159, + 0.685175, 0.335865, 0.910850, 0.934065, 0.319306, 0.340147, 0.643746, + 0.981592, 0.709673, 0.496812, 0.658856, 0.353983, 0.337245, 0.966670, + 0.213511, 0.849838, 0.569482, 0.133671, 0.290786, 0.563007, 0.330991, + 0.427170, 0.620991, 0.065299, 0.437936, 0.034320, 0.996356, 0.259643, + 0.813834, 0.070399, 0.132802, 0.499009, 0.406265, 0.043652, 0.433074, + 0.725570, 0.383800, 0.076820, 0.707163, 0.093473, 0.573632, 0.366018, + 0.447456, 0.910877, 0.332688, 0.660967, 0.760714, 0.902170, 0.794638, + 0.051500, 0.465177, 0.125630, 0.478670, 0.086168, 0.190928, 0.916605, + 0.120488, 0.187285, 0.176248, 0.934322, 0.257684, 0.309050, 0.433331, + 0.663949, 0.352703, 0.866405, 0.389519, 0.736502, 0.943226, 0.096682, + 0.829975, 0.516858, 0.462700, 0.277430, 0.427734, 0.795388, 0.938398, + 0.188449, 0.697558, 0.733036, 0.239948, 0.162735, 0.858666, 0.718618, + 0.248903, 0.049594, 0.635223, 0.369391, 0.236879, 0.811472, 0.303713, + 0.494563, 0.120522, 0.737044, 0.158511, 0.473225, 0.603450, 0.548030, + 0.209727, 0.546675, 0.644712, 0.039702, 0.063533, 0.107412, 0.317132, + 0.491267, 0.902800, 0.255530, 0.679716, 0.600359, 0.988566, 0.919664, + 0.763094, 0.847232, 0.638283, 0.011997, 0.896825, 0.273506, 0.381388, + 0.133704, 0.084978, 0.685101, 0.628267, 0.205500, 0.422145, 0.786778, + 0.678725, 0.025595, 0.334808, 0.888452, 0.572271, 0.979520, 0.928154, + 0.635804, 0.086932, 0.245286, 0.127071, 0.989732, 0.500816, 0.806787, + 0.590091, 0.489382, 0.726451, 0.353185, 0.336614, 0.364734, 0.365182, + 0.233439, 0.638240, 0.746570, 0.367143, 0.723218, 0.431671, 0.995410, + 0.928718, 0.853816, 0.782188, 0.607442, 0.879411, 0.116995, 0.495894, + 0.451682, 0.096515, 0.424048, 0.087485, 0.183447, 0.669334, 0.214556, + 0.173179, 0.170151, 0.021343, 0.763269, 0.659533, 0.747794, 0.116454, + 0.996147, 0.112528, 0.481635, 0.229586, 0.750768, 0.228205, 0.596730, + 0.473985, 0.659876, 0.592139, 0.402703, 0.513692, 0.374327, 0.010145, + 0.393103, 0.491322, 0.506039, 0.844785, 0.587837, 0.930088, 0.932270, + 0.771284, 0.599422, 0.146826, 0.944463, 0.769573, 0.168169, 0.707732, + 0.429106, 0.915964, 0.824186, 0.425253, 0.028492, 0.305821, 0.654839, + 0.779259, 0.534026, 0.251569, 0.253245, 0.193901, 0.843708, 0.655947, + 0.707593, 0.218035, 0.666093, 0.100696, 0.709357, 0.172132, 0.945481, + 0.297195, 0.102220, 0.877751, 0.068479, 0.701642, 0.024577, 0.012941, + 0.471215, 0.192747, 0.720673, 0.900321, 0.108710, 0.544859, 0.325574, + 0.137202, 0.850679, 0.980413, 0.916462, 0.384705, 0.231982, 0.169706, + 0.578607, 0.075690, 0.825654, 0.286200, 0.293725, 0.491746, 0.386896, + 0.003083, 0.663878, 0.332377, 0.300278, 0.766098, 0.210128, 0.368756, + 0.467740, 0.234705, 0.381697, 0.938955, 0.427451, 0.102370, 0.839275, + 0.536162, 0.647229, 0.164849, 0.673364, 0.497908, 0.145262, 0.589825, + 0.882613, 0.377244, 0.759532, 0.461220, 0.452934, 0.585185, 0.747420, + 0.746660, 0.076932, 0.134316, 0.749743, 0.740810, 0.466692, 0.050020, + 0.506908, 0.676820, 0.418776, 0.974648, 0.911525, 0.800474, 0.913602, + 0.338976, 0.902844, 0.752878, 0.875138, 0.550072, 0.917727, 0.548502, + 0.047981, 0.062989, 0.138327, 0.930594, 0.440233, 0.897859, 0.391814, + 0.893168, 0.483044, 0.139234, 0.639828, 0.559975, 0.273549, 0.389570, + 0.300785, 0.740242, 0.439590, 0.807693, 0.417062, 0.858367, 0.782341, + 0.328586, 0.658840, 0.695943, 0.667562, 0.561684, 0.448821, 0.542700, + 0.111756, 0.366548, 0.091202, 0.159737, 0.429537, 0.229529, 0.090331, + 0.869770, 0.127388, 0.482145, 0.762938, 0.610432, 0.621379, 0.402765, + 0.170407, 0.894928, 0.792336, 0.471192, 0.635170, 0.231926, 0.278886, + 0.052232, 0.090293, 0.061226, 0.380818, 0.749133, 0.757170, 0.048380, + 0.310817, 0.205990, 0.591080, 0.422573, 0.572538, 0.682282, 0.582310, + 0.002075, 0.911812, 0.672641, 0.871845, 0.039199, 0.154786, 0.634783, + 0.649631, 0.776165, 0.037548, 0.820038, 0.671093, 0.829884, 0.291231, + 0.306263, 0.061810, 0.570116, 0.358495, 0.152103, 0.631343, 0.739313, + 0.901236, 0.388512, 0.787693, 0.212053, 0.594503, 0.378773, 0.634626, + 0.167040, 0.061056, 0.216937, 0.169115, 0.972867, 0.889578, 0.040960, + 0.012067, 0.044364, 0.675743, 0.661698, 0.820529, 0.713291, 0.481736, + 0.491623, 0.543175, 0.772966, 0.797886, 0.604985, 0.343083, 0.156380, + 0.757088, 0.974425, 0.895693, 0.658324, 0.362938, 0.683386, 0.870376, + 0.957440, 0.062159, 0.505002, 0.124481, 0.123215, 0.721939, 0.293596, + 0.096082, 0.611517, 0.334556, 0.108149, 0.655881, 0.010299, 0.769846, + 0.476411, 0.723590, 0.251582, 0.968033, 0.266765, 0.024548, 0.765919, + 0.871750, 0.367631, 0.922299, 0.628838, 0.342056, 0.817992, 0.287162, + 0.704994, 0.501378, 0.157538, 0.662434, 0.563537, 0.662541, 0.786915, + 0.686752, 0.384480, 0.080511, 0.782834, 0.995997, 0.415067, 0.890983, + 0.651878, 0.425365, 0.660829, 0.128289, 0.148956, 0.912411, 0.096322, + 0.415721, 0.936959, 0.862241, 0.287471, 0.304590, 0.784540, 0.916309, + 0.646646, 0.602533, 0.203471, 0.351640, 0.103911, 0.361009, 0.014074, + 0.667448, 0.023550, 0.800989, 0.354200, 0.408030, 0.881500, 0.137034, + 0.404026, 0.296566, 0.028017, 0.055904, 0.721932, 0.688846, 0.184193, + 0.870887, 0.601257, 0.280515, 0.286608, 0.538216, 0.142755, 0.574079, + 0.842806, 0.927296, 0.490388, 0.489452, 0.529828, 0.693859, 0.841092, + 0.633739, 0.054869, 0.855167, 0.301187, 0.078419, 0.656156, 0.655388, + 0.486448, 0.537656, 0.792422, 0.890475, 0.834222, 0.820439, 0.946379, + 0.556153, 0.509285, 0.130571, 0.427041, 0.110542, 0.411086, 0.713648, + 0.648758, 0.553842, 0.287727, 0.491563, 0.481137, 0.778116, 0.981015, + 0.010966, 0.471975, 0.822107, 0.644705, 0.526844, 0.677274, 0.945892, + 0.605263, 0.333430, 0.601280, 0.091711, 0.871086, 0.393702, 0.982186, + 0.705307, 0.214141, 0.928564, 0.261461, 0.723426, 0.059136, 0.688501, + 0.833968, 0.470222, 0.402150, 0.482725, 0.024063, 0.689877, 0.974289, + 0.505201, 0.467993, 0.955304, 0.516166, 0.939968, 0.777411, 0.160871, + 0.466812, 0.454685, 0.106763, 0.072075, 0.788115, 0.708043, 0.163786, + 0.659201, 0.101744, 0.145971, 0.364508, 0.315885, 0.074536, 0.625969, + 0.039311, 0.133672, 0.314471, 0.873279, 0.603893, 0.716620, 0.356004, + 0.627957, 0.406498, 0.330292, 0.133157, 0.874490, 0.285596, 0.649324, + 0.814458, 0.063007, 0.810195, 0.281270, 0.517693, 0.916958, 0.353345, + 0.305808, 0.625000, 0.517131, 0.965009, 0.726745, 0.663102, 0.329518, + 0.042630, 0.737638, 0.955487, 0.081940, 0.871310, 0.269957, 0.955219, + 0.475203, 0.986578, 0.311223, 0.103160, 0.393075, 0.641515, 0.236317, + 0.267566, 0.927112, 0.885641, 0.082024, 0.990119, 0.695835, 0.363295, + 0.507812, 0.612793, 0.716640, 0.813620, 0.237793, 0.233770, 0.778629, + 0.964538, 0.896872, 0.108147, 0.007167, 0.634510, 0.063633, 0.089108, + 0.505820, 0.333591, 0.044327, 0.981023, 0.320168, 0.355550, 0.084182, + 0.713244, 0.997065, 0.320499, 0.980810, 0.924177, 0.206140, 0.062834, + 0.914296, 0.901975, 0.426129, 0.422107, 0.514768, 0.142768, 0.235727, + 0.752561, 0.376539, 0.014356, 0.717099, 0.273411, 0.122502, 0.724266, + 0.907921, 0.186136, 0.813374, 0.413741, 0.519726, 0.857701, 0.394764, + 0.839895, 0.213251, 0.478946, 0.553139, 0.210317, 0.799446, 0.533948, + 0.134493, 0.005586, 0.596782, 0.048789, 0.907561, 0.022911, 0.470896, + 0.422329, 0.165679, 0.706623, 0.174890, 0.542218, 0.720979, 0.891989, + 0.815629, 0.843481, 0.616255, 0.723551, 0.029617, 0.429630, 0.137292, + 0.549343, 0.287331, 0.532056, 0.389238, 0.500583, 0.011002, 0.942377, + 0.710899, 0.810448, 0.476326, 0.845392, 0.816033, 0.073108, 0.894181, + 0.723594, 0.096019, 0.365077, 0.145923, 0.261699, 0.071700, 0.320813, + 0.803917, 0.792679, 0.212802, 0.619546, 0.636160, 0.829057, 0.343096, + 0.665777, 0.258687, 0.480388, 0.215121, 0.546018, 0.012444, 0.604359, + 0.046601, 0.023446, 0.546736, 0.757500, 0.833893, 0.023062, 0.602892, + 0.649927, 0.096170, 0.497074, 0.373521, 0.192189, 0.862151, 0.519444, + 0.453887, 0.933851, 0.840257, 0.257804, 0.726531, 0.053058, 0.877350, + 0.362691, 0.882115, 0.220446, 0.028468, 0.140802, 0.700834, 0.243589, + 0.686821, 0.713278, 0.847948, 0.733421, 0.736723, 0.394684, 0.490921, + 0.570617, 0.417746, 0.093813, 0.220543, 0.513916, 0.590887, 0.594064, + 0.706105, 0.453038, 0.113508, 0.159992, 0.386889, 0.953765, 0.417796, + 0.113420, 0.006823, 0.295146, 0.476111, 0.888938, 0.515592, 0.504579, + 0.029741, 0.216426, 0.748168, 0.716561, 0.929703, 0.596117, 0.449982, + 0.666427, 0.990801, 0.940903, 0.237043, 0.408547, 0.034717, 0.457587, + 0.922463, 0.625603, 0.051651, 0.628568, 0.078641, 0.165159, 0.788560, + 0.465530, 0.118923, 0.206356, 0.578950, 0.125746, 0.501502, 0.055060, + 0.014685, 0.017094, 0.559640, 0.044425, 0.233519, 0.307808, 0.760986, + 0.163223, 0.903925, 0.210969, 0.829650, 0.894726, 0.151872, 0.066693, + 0.303273, 0.186589, 0.524279, 0.225736, 0.812192, 0.575930, 0.854304, + 0.890833, 0.741089, 0.642864, 0.356363, 0.860012, 0.849220, 0.935313, + 0.985758, 0.350722, 0.990373, 0.000443, 0.367815, 0.550013, 0.044868, + 0.601335, 0.857820, 0.805855, 0.764557, 0.761745, 0.016823, 0.594207, + 0.656471, 0.168696, 0.660900, 0.959744, 0.355284, 0.185179, 0.185480, + 0.167477, 0.761110, 0.039784, 0.058310, 0.502199, 0.682648, 0.414673, + 0.362211, 0.531868, 0.349985, 0.347969, 0.882589, 0.340358, 0.348412, + 0.250404, 0.890371, 0.393280, 0.851739, 0.748191, 0.199135, 0.616297, + 0.509936, 0.215958, 0.210504, 0.166407, 0.384654, 0.871404, 0.126151, + 0.739938, 0.056583, 0.311631, 0.907415, 0.817693, 0.351415, 0.965724, + 0.319891, 0.034062, 0.380397, 0.682102, 0.565930, 0.730382, 0.030072, + 0.448519, 0.070741, 0.378484, 0.698924, 0.961112, 0.771764, 0.550663, + 0.709303, 0.970899, 0.166959, 0.219239, 0.186857, 0.377463, 0.385647, + 0.571511, 0.248867, 0.511798, 0.311449, 0.305450, 0.823429, 0.218864, + 0.123142, 0.174844, 0.184588, 0.443034, 0.208906, 0.564986, 0.125136, + 0.774836, 0.295368, 0.155207, 0.223355, 0.366109, 0.533691, 0.922279, + 0.327221, 0.305455, 0.472942, 0.036524, 0.276354, 0.639901, 0.255763, + 0.463211, 0.017364, 0.641410, 0.034722, 0.266231, 0.153207, 0.346171, + 0.571680, 0.976636, 0.565036, 0.694822, 0.151480, 0.749624, 0.137856, + 0.360386, 0.314610, 0.262992, 0.135222, 0.609978, 0.418200, 0.358578, + 0.976087, 0.951891, 0.280856, 0.303307, 0.257346, 0.753798, 0.339831, + 0.533700, 0.393699, 0.595594, 0.996911, 0.411063, 0.237003, 0.031634, + 0.677294, 0.390211, 0.377805, 0.248974, 0.366847, 0.942841, 0.943796, + 0.518327, 0.692465, 0.081653, 0.878713, 0.007074, 0.344645, 0.013936, + 0.617052, 0.762845, 0.372513, 0.593138, 0.714736, 0.653370, 0.896446, + 0.972082, 0.407168, 0.236276, 0.505782, 0.800867, 0.831870, 0.502693, + 0.211930, 0.068873, 0.534327, 0.889224, 0.459084, 0.912132, 0.138197, + 0.825931, 0.854972, 0.081994, 0.344259, 0.547437, 0.163646, 0.222972, + 0.554511, 0.508291, 0.236908, 0.171563, 0.271135, 0.609421, 0.764701, + 0.985871, 0.262790, 0.661147, 0.957953, 0.669958, 0.897423, 0.463734, + 0.470825, 0.729293, 0.966427, 0.682755, 0.798166, 0.500754, 0.571978, + 0.257251, 0.412886, 0.710176, 0.083182, 0.267858, 0.792169, 0.427441, + 0.815295, 0.955815, 0.650413, 0.369805, 0.464106, 0.887320, 0.541368, + 0.735242, 0.496741, 0.306069, 0.721113, 0.759531, 0.967216, 0.679065, + 0.429489, 0.864639, 0.142799, 0.900314, 0.593932, 0.109227, 0.583069, + 0.392098, 0.609981, 0.155047, 0.649349, 0.022867, 0.865222, 0.732531, + 0.290725, 0.657392, 0.159972, 0.106019, 0.613207, 0.810384, 0.475824, + 0.077313, 0.697704, 0.017192, 0.812555}; + +static float golden_endtoend_output[] = { + -1.881211, -0.028385, -3.585066, 1.939770, -3.461155, 1.280415, -4.408978, + 0.608663, -2.704937, 1.859742, -5.777429, 2.691839, -1.049012, 1.640870, + -4.856245, 1.604236, 0.992707, 0.422858, -4.307465, 1.887332, -0.884831, + -0.154277, -2.634801, 0.586827, -1.849960, 1.399608, -4.531559, 1.943591, + 0.271676, -2.893054, -2.066826, 0.235467, -1.248263, -1.164534, -2.640174, + -0.112878, -4.386484, 1.253024, -4.135623, 1.068984, -0.043579, -0.832957, + -3.257258, -0.514396, -1.651174, 0.638630, -4.364372, 1.548441, -0.289455, + 0.539845, -4.097627, 0.635001, -0.465071, -0.927701, -2.481498, 0.356616, + -2.355012, 0.728806, -3.340283, 1.609038, -4.786268, -0.532272, -1.886150, + 0.254797, 0.746620, -1.657134, -3.264265, 0.525551, -1.756837, 0.845446, + -5.572190, 1.715797, -2.856942, 3.394245, -5.803662, 2.281806, -3.014739, + 2.616136, -4.728482, 1.659984, -2.106307, 2.711709, -6.173832, 1.352869, + -0.038035, 0.107619, -4.279774, 2.341930, -0.980413, -0.119538, -4.049717, + 1.172128, -3.477744, 2.602274, -6.231380, 2.537300, -0.862214, 0.568722, + -3.858362, 0.197867, -1.725885, 3.687312, -7.067363, 2.403544, -0.944963, + 0.235639, -3.250094, 0.659117, -1.459576, 0.426128, -3.637207, 1.030386, + -4.224351, 3.516220, -6.053367, 0.993473, -2.182416, -0.762625, -1.884405, + -0.113736, -2.572602, 0.329290, -1.913233, 0.517418, -0.019757, 0.203176, + -3.715881, 0.482136, -1.912823, 1.357907, -5.473043, 1.714658, -3.177160, + 0.089285, -3.127669, 1.268076, 0.772498, -1.622712, -3.850314, 0.436124, + -1.495983, 3.439982, -7.623405, 1.726721, -0.423979, 0.180201, -2.902406, + 0.986457, -1.845638, 0.460903, -5.359343, -1.133931, -1.074456, 0.717304, + -3.519856, 1.012126, -0.562301, 1.881967, -6.716627, 2.525036, 0.945480, + 0.337081, -5.210562, 2.572035, -0.943370, 0.442026, -2.666313, 0.411296, + 0.002787, -0.000735, -2.498933, 0.771719, -3.568153, 3.833721, -6.617026, + 2.813922, -0.573970, 1.025208, -3.909923, 1.722648, -1.406849, 0.719783, + -5.207438, 1.819442, -0.530895, -0.010887, -2.939614, 0.971225, -1.660297, + 1.345243, -4.454571, 2.244876, -2.021213, 1.756090, -4.880947, 0.364597, + -2.380270, 2.763117, -5.613013, 2.137534, 0.289101, -2.279400, -3.365582, + 0.170028, -1.142254, -0.709604, -3.656223, 1.804870, -0.854690, 0.592102, + -5.010415, 2.462687, -1.474710, 0.566002, -3.621819, -0.391946, -0.423524, + -0.631428, -3.513310, 0.962825, -1.480262, 0.319791, -3.610137, 1.842339, + -0.250073, 1.182022, -6.249267, 1.604172, 1.153759, -0.734054, -4.620415, + -0.030858, 0.050911, 1.524406, -4.724010, 1.451846, -3.277104, 2.414182, + -4.605285, 1.846092, -1.503047, -0.618200, -2.746546, -0.459332, -0.980326, + -1.199977, -2.043865, -0.165793, -2.214698, 3.108281, -7.127830, -0.123065, + 1.244948, -3.039923, -4.660061, -0.225957, -0.307210, -1.513205, -2.456005, + 0.840048, -0.741445, 2.328635, -6.015267, 2.723240, -1.381171, -0.728878, + -5.114925, -0.362034, -0.574923, 0.518080, -3.892457, 1.798948, 0.435119, + -0.371696, -2.807571, 1.302864, -2.063052, 1.036388, -4.232038, 1.397059, + -1.615668, -1.511019, -3.095508, 1.290955, -3.428723, 2.000287, -4.196487, + 1.566983, 0.196957, 0.224343, -4.926359, -0.691975, -0.214941, 1.546821, + -5.384868, 2.290820, -1.878865, 0.493692, -4.129823, 2.112036, 0.516558, + -2.553077, -2.717338, 0.017146, -2.016057, 1.628995, -4.240602, 1.189533, + -5.460220, 1.254738, -4.214903, 0.755659, -2.893235, 2.937762, -6.169453, + 2.035456, -5.613212, -0.122254, -1.973646, -0.060619, -2.119598, 1.413512, + -4.938738, 1.890244, 0.544169, -2.062413, -3.329637, -0.062515, -1.855805, + -0.791297, -2.570353, 0.607615, 0.305812, 0.338930, -4.150270, 2.274937, + 0.042653, 0.133825, -3.538155, 1.523639, -3.173690, -1.496599, -2.414655, + 0.464687, -1.448998, -0.368907, -3.520129, 0.203382, -2.443626, 1.266233, + -3.393848, 0.605911, -0.015353, 1.402006, -4.441003, 1.419281, 0.603587, + 0.434146, -4.966566, 2.171872, -0.688264, -0.009981, -4.461103, 1.538354, + -5.029816, -0.264424, -1.713510, -0.315258, -1.891606, 0.252074, -2.419428, + 0.043970, -1.291143, 2.048704, -4.590105, 0.524734, -1.889576, 0.134836, + -3.462745, 1.390663, -0.112773, 0.402735, -4.203784, 1.381043, -1.201634, + -1.968277, -1.425637, -0.181725, -1.250742, -2.102041, -3.925464, -1.256797, + -3.701354, -1.754610, -1.917231, -1.455910, -1.838006, 2.041781, -5.666212, + 2.752957, -2.659553, 2.553637, -4.872212, 1.443437, -2.081846, 3.311263, + -5.912457, 1.871049, 0.196148, -0.307044, -4.024967, 2.149149, 0.361809, + 0.620415, -5.939984, 0.180672, -1.209180, -0.269122, -3.240285, 1.460315, + -1.040803, 1.125700, -6.060366, 0.887767, -3.214111, 1.314368, -3.026808, + 1.023640, -3.815175, 1.795642, -4.355603, 1.064454, -0.046472, 0.618463, + -5.941646, 2.861891, -2.852155, -0.990457, -2.624445, 1.794494, -1.176747, + -0.358159, -3.206776, 1.138721, -2.819523, -1.825522, -1.450902, -0.187312, + -0.808727, 0.636872, -4.120567, 1.192623, 0.810731, -1.768519, -3.699450, + 1.527116, -2.772720, 3.012835, -5.912736, 1.599365, -4.696381, 2.234591, + -4.139552, 1.061768, -1.880089, 3.596274, -7.006379, 2.382152, -3.158115, + 3.844430, -7.044156, 2.307596, -2.473970, 1.312644, -5.467269, 0.197154, + -1.530040, 1.762275, -5.550757, 0.630276, -3.048947, 1.043777, -3.096658, + 1.345893, -1.329494, 2.065748, -4.711032, 2.227600, -0.413321, -0.032428, + -4.599650, 1.668734, -4.351490, -0.200022, -2.359903, 0.021997, 0.116028, + 1.159718, -5.093972, -0.142951, -2.409895, 0.906133, -2.728812, 0.809932, + -2.597363, 0.494130, -2.357861, 0.369825, -2.165235, 1.148522, -3.130562, + 0.759034, 0.646335, -1.463660, -3.508299, 1.059679, -1.485465, 1.007319, + -4.340716, 1.789864, -1.590654, 1.612324, -4.452007, 2.389805, -5.200148, + -1.068398, -1.306923, -0.472408, -0.392165, -0.524996, -2.933478, 1.518430, + -1.287781, 0.113422, -3.020525, 1.338359, -0.105982, 0.936014, -4.132197, + 1.836807, -0.616589, -1.029716, -3.271347, 0.284889, -2.653359, 2.135829, + -4.643613, 1.627981, 0.287733, -2.017263, -2.776574, 1.184792, 1.004161, + -1.483019, -4.339290, -0.787322, 0.582420, 1.137839, -5.673941, -0.001862, + -1.219142, 0.532561, -4.457245, 1.826807, -3.343291, 3.034610, -6.179855, + 2.235917, -4.369989, 4.018128, -6.632714, 0.926585, -0.485469, 0.536073, + -4.179557, 1.489637, -0.521762, 1.636089, -6.137912, 1.500867, -4.086009, + 1.961372, -3.688977, 1.358220, -1.544034, 1.763837, -4.357567, 1.852201, + -2.018725, 1.046264, -6.211127, 1.609419, -0.118441, 1.602284, -6.242423, + 1.518578, -0.604078, 1.106613, -5.393445, 2.595629, 0.142712, -1.903953, + -2.821177, 0.032758, -0.009152, 0.184628, -4.227636, 2.046843, -2.240138, + 1.256176, -5.108516, -0.308447, -2.998571, 4.657396, -7.582112, 2.510951, + -3.535784, 1.704560, -5.068484, 1.318466, -3.058265, 3.073172, -6.998089, + 3.178849, -2.420286, 2.277806, -4.999528, 1.423890, -1.672914, 0.447460, + -4.088940, 1.351087, -1.051546, -0.417955, -4.042147, 1.604102, -1.700931, + 2.796663, -6.497579, 2.857974, -0.240828, 0.858001, -5.778933, 2.778508, + -0.406211, 1.300766, -5.073671, 2.089362, -0.201673, 1.588396, -6.000150, + 2.185055, -2.332125, 0.768216, -2.609184, 0.327277, -3.358943, -1.020736, + -2.389984, 0.315512, -0.561905, 1.948740, -6.408485, 2.231985, -0.603652, + 0.661829, -5.070386, -1.063058, -0.624796, 1.375772, -4.379606, 1.929358, + -1.047263, 0.739100, -5.217857, 2.127625, -5.025338, 0.650344, -2.068460, + 0.076936, -0.457505, -1.050984, -1.917765, 1.150908, 0.782625, 0.855595, + -5.321719, 0.787209, -0.460232, 1.106736, -5.552326, 2.801043, -0.360217, + -0.434432, -4.273378, 0.967556, -0.972652, 0.874811, -5.429918, -0.331039, + 0.115477, 0.111883, -5.418786, 1.240546, -1.842794, 0.505880, -3.676064, + -0.682369, 1.858984, -0.742566, -5.784060, 0.673239, -1.280398, 0.280842, + -4.848077, 2.214860, -0.785100, -0.588488, -2.438206, 0.786651, -1.568752, + 1.935400, -6.320256, 2.125338, -1.476457, -1.651941, -2.695734, 0.007338, + -3.280860, 2.310385, -5.319578, 1.890123, -0.775723, 0.630606, -4.321582, + 1.085521, -1.847371, 1.188521, -4.596577, 2.056443, -2.340172, -0.108501, + -3.156392, 0.933279, -0.495331, 0.122405, -5.171133, 1.763245, -0.796913, + 2.310487, -7.247197, 2.401678, -1.908860, 0.043798, -2.393796, 0.573806, + -0.608531, 0.154710, -4.669001, 0.750680, 0.468380, 0.392591, -4.755001, + 2.615217, -1.957774, 1.153513, -4.530099, 1.124362, -3.569415, 1.697154, + -3.536335, 0.910758, -2.976264, 1.833129, -4.287203, -0.547050, -2.409768, + 0.061585, -1.324116, 0.268497, -2.962222, -1.524245, -2.063413, 0.442058, + -4.292337, 3.538863, -6.699603, 1.718664, -2.290363, 1.994596, -6.245037, + -0.433084, -0.367059, 1.020297, -4.940721, 2.902264, -0.577056, -0.709887, + -5.001413, -0.268316, -1.112048, -1.083307, -1.753492, 0.209973, 0.139540, + 0.917602, -5.232745, 2.538467, -2.139234, -0.187388, -1.837249, -0.478582, + -0.731653, -0.481550, -2.531261, 1.044770, 0.707750, 0.279971, -3.221119, + 1.552074, -2.373144, 0.859518, -3.665156, 1.620278, -1.440871, -0.525581, + -2.758271, 1.491873, -2.302013, 1.119935, -5.257080, 2.627170, -3.174739, + 1.363282, -4.831639, 1.101076, -4.337008, 2.689639, -5.165915, 1.069201, + -1.882078, -0.120370, -2.287967, 1.147619, -1.403616, 1.077150, -5.084296, + 1.658236, -0.919642, 0.487423, -3.001075, 0.741268, 0.107300, 0.943556, + -3.544311, 1.000239, -1.627171, 2.871253, -5.179172, 1.429893, -0.826040, + 0.188670, -4.499894, 1.013447, -2.101299, 0.317516, -3.452141, -0.833776, + -1.362144, 1.272437, -4.449355, 1.613591, -2.039873, 2.613175, -6.229640, + 1.659790, -1.595520, -0.237462, -2.744997, 0.337841, 0.148981, -1.703771, + -2.388023, 1.276469, 1.058508, -0.401642, -4.680769, 0.861881, -1.336381, + 1.153080, -2.834378, 0.721075, 0.900115, 1.360511, -5.573611, 0.949182, + -2.970844, 2.017563, -5.186108, -0.201038, -1.192824, 0.610142, -4.450919, + -0.897114, -1.812093, 0.422310, -5.245487, 0.256549, 0.320275, -2.324150, + -2.967040, -0.260536, -0.721467, 0.454148, -5.058031, 0.526370, -0.895656, + 0.732240, -3.327363, 1.353953, -1.277912, -0.483171, -1.926713, 0.065044, + -2.167506, -0.196606, -1.923437, 0.604962, -2.088319, 1.406834, -5.227296, + 2.247351, -4.421744, 1.729791, -5.007922, 1.264769, -0.897019, 0.922902, + -3.887108, 2.087432, -1.310226, -0.101938, -3.359082, -0.079662, -0.514988, + -0.963179, -4.038209, 2.223278, -0.590083, -2.310458, -1.748338, 0.363406, + -0.540731, -0.885913, -4.179595, 2.216781, -3.044339, -0.447100, -2.446098, + 0.931101, -1.676190, 2.096175, -4.980755, 2.262151, -1.095047, 1.897516, + -5.996138, 2.191038, 0.297128, -0.780974, -2.884299, 1.195408, -0.521065, + -1.955837, -3.091064, -0.404183, -1.961519, 4.076096, -7.521851, 2.242064, + -1.988043, 0.303300, -2.422585, 0.322230, -3.377634, 3.499955, -7.084434, + 2.375587, -0.718851, 2.150076, -5.412241, 2.374280, -2.006088, 2.229828, + -5.848188, 2.543077, -2.171042, 2.096026, -5.300007, 0.141405, -1.187745, + 0.105340, -4.003816, 1.034281, -3.980804, 1.856709, -5.103042, 0.623737, + -2.080307, 0.896140, -3.104050, 0.983158, -0.424898, -1.154270, -3.805728, + 1.978917, -1.314387, 1.235096, -3.148906, 1.113173, 0.111713, 2.055213, + -7.565283, 2.100342}; +constexpr std::initializer_list biases = { + 0.065691948, -0.69055247, 0.1107955, -0.97084129, -0.23957068, -0.23566568, + -0.389184, 0.47481549, -0.4791103, 0.29931796, 0.10463274, 0.83918178, + 0.37197268, 0.61957061, 0.3956964, -0.37609905}; + +constexpr std::initializer_list recurrent_weights = { + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1}; + +class BidirectionalRNNOpModel : public SingleOpModel { + public: + BidirectionalRNNOpModel(int batches, int sequence_len, int fw_units, + int bw_units, int input_size) + : batches_(batches), + sequence_len_(sequence_len), + fw_units_(fw_units), + bw_units_(bw_units), + input_size_(input_size) { + input_ = AddInput(TensorType_FLOAT32); + fw_weights_ = AddInput(TensorType_FLOAT32); + fw_recurrent_weights_ = AddInput(TensorType_FLOAT32); + fw_bias_ = AddInput(TensorType_FLOAT32); + fw_hidden_state_ = AddOutput(TensorType_FLOAT32); + fw_output_ = AddOutput(TensorType_FLOAT32); + bw_weights_ = AddInput(TensorType_FLOAT32); + bw_recurrent_weights_ = AddInput(TensorType_FLOAT32); + bw_bias_ = AddInput(TensorType_FLOAT32); + bw_hidden_state_ = AddOutput(TensorType_FLOAT32); + bw_output_ = AddOutput(TensorType_FLOAT32); + SetBuiltinOp(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN, + BuiltinOptions_SequenceRNNOptions, + CreateSequenceRNNOptions(builder_, /*time_major=*/false, + ActivationFunctionType_RELU) + .Union()); + BuildInterpreter({ + {batches_, sequence_len_, input_size_}, // input + {fw_units_, input_size_}, // fw_weights + {fw_units_, fw_units_}, // fw_recurrent_weights + {fw_units_}, // fw_bias + {bw_units_, input_size_}, // bw_weights + {bw_units_, bw_units_}, // bw_recurrent_weights + {bw_units_} // bw_bias + }); + } + + void SetFwBias(std::initializer_list f) { + PopulateTensor(fw_bias_, f); + } + + void SetBwBias(std::initializer_list f) { + PopulateTensor(bw_bias_, f); + } + + void SetFwWeights(std::initializer_list f) { + PopulateTensor(fw_weights_, f); + } + + void SetBwWeights(std::initializer_list f) { + PopulateTensor(bw_weights_, f); + } + + void SetFwRecurrentWeights(std::initializer_list f) { + PopulateTensor(fw_recurrent_weights_, f); + } + + void SetBwRecurrentWeights(std::initializer_list f) { + PopulateTensor(bw_recurrent_weights_, f); + } + + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + + void SetInput(int offset, float* begin, float* end) { + PopulateTensor(input_, offset, begin, end); + } + + void ResetHiddenStates() { + const int fw_zero_buffer_size = fw_units_ * batches_; + std::unique_ptr fw_zero_buffer(new float[fw_zero_buffer_size]); + memset(fw_zero_buffer.get(), 0, fw_zero_buffer_size * sizeof(float)); + PopulateTensor(fw_hidden_state_, 0, fw_zero_buffer.get(), + fw_zero_buffer.get() + fw_zero_buffer_size); + const int bw_zero_buffer_size = bw_units_ * batches_; + std::unique_ptr bw_zero_buffer(new float[bw_zero_buffer_size]); + memset(bw_zero_buffer.get(), 0, bw_zero_buffer_size * sizeof(float)); + PopulateTensor(bw_hidden_state_, 0, bw_zero_buffer.get(), + bw_zero_buffer.get() + bw_zero_buffer_size); + } + + std::vector GetFwOutput() { return ExtractVector(fw_output_); } + std::vector GetBwOutput() { return ExtractVector(bw_output_); } + + int input_size() { return input_size_; } + int num_fw_units() { return fw_units_; } + int num_bw_units() { return bw_units_; } + int num_batches() { return batches_; } + int sequence_len() { return sequence_len_; } + + private: + int input_; + int fw_weights_; + int fw_recurrent_weights_; + int fw_bias_; + int fw_hidden_state_; + int fw_output_; + int bw_weights_; + int bw_recurrent_weights_; + int bw_bias_; + int bw_hidden_state_; + int bw_output_; + + int batches_; + int sequence_len_; + int fw_units_; + int bw_units_; + int input_size_; +}; + +// TODO(mirkov): add another test which directly compares to TF once TOCO +// supports the conversion from dynamic_rnn with BasicRNNCell. +TEST(BidirectionalRNNOpTest, BlackBoxTest) { + BidirectionalRNNOpModel rnn(/*batches=*/2, /*sequence_len=*/16, + /*fw_units=*/16, /*bw_units=*/16, + /*input_size=*/8); + rnn.SetFwWeights(weights); + rnn.SetBwWeights(weights); + rnn.SetFwBias(biases); + rnn.SetBwBias(biases); + rnn.SetFwRecurrentWeights(recurrent_weights); + rnn.SetBwRecurrentWeights(recurrent_weights); + + rnn.ResetHiddenStates(); + const int input_sequence_size = rnn.input_size() * rnn.sequence_len(); + float* batch_start = rnn_input; + float* batch_end = batch_start + input_sequence_size; + rnn.SetInput(0, batch_start, batch_end); + rnn.SetInput(input_sequence_size, batch_start, batch_end); + + rnn.Invoke(); + + float* golden_fw_start = rnn_golden_fw_output; + float* golden_fw_end = + golden_fw_start + rnn.num_fw_units() * rnn.sequence_len(); + std::vector fw_expected; + fw_expected.insert(fw_expected.end(), golden_fw_start, golden_fw_end); + fw_expected.insert(fw_expected.end(), golden_fw_start, golden_fw_end); + EXPECT_THAT(rnn.GetFwOutput(), ElementsAreArray(ArrayFloatNear(fw_expected))); + + float* golden_bw_start = rnn_golden_bw_output; + float* golden_bw_end = + golden_bw_start + rnn.num_bw_units() * rnn.sequence_len(); + std::vector bw_expected; + bw_expected.insert(bw_expected.end(), golden_bw_start, golden_bw_end); + bw_expected.insert(bw_expected.end(), golden_bw_start, golden_bw_end); + EXPECT_THAT(rnn.GetBwOutput(), ElementsAreArray(ArrayFloatNear(bw_expected))); +} + +// Check that if the input sequence is reversed the outputs are the same just +// forward and backward are swapped (and reversed). +TEST(BidirectionalRNNOpTest, BlackBoxTestReverseInputs) { + BidirectionalRNNOpModel rnn(/*batches=*/2, /*sequence_len=*/16, + /*fw_units=*/16, /*bw_units=*/16, + /*input_size=*/8); + rnn.SetFwWeights(weights); + rnn.SetBwWeights(weights); + rnn.SetFwBias(biases); + rnn.SetBwBias(biases); + rnn.SetFwRecurrentWeights(recurrent_weights); + rnn.SetBwRecurrentWeights(recurrent_weights); + + rnn.ResetHiddenStates(); + + // Reverse inputs in each batch: in_1, in_2,..., in_k is inserted in the + // following order: [in_k,..., in_2, in_1, in_k,...,in_2, in_1]. + for (int i = 0; i < rnn.sequence_len(); i++) { + float* batch_start = rnn_input + i * rnn.input_size(); + float* batch_end = batch_start + rnn.input_size(); + const int reverse_idx = rnn.sequence_len() - i - 1; + rnn.SetInput(reverse_idx * rnn.input_size(), batch_start, batch_end); + rnn.SetInput((rnn.sequence_len() + reverse_idx) * rnn.input_size(), + batch_start, batch_end); + } + + rnn.Invoke(); + + // The forward and backward outputs are swapped. + std::vector fw_expected; // consider using std::deque instead. + for (int i = 0; i < rnn.sequence_len(); i++) { + float* golden_fw_start = rnn_golden_bw_output + i * rnn.num_fw_units(); + float* golden_fw_end = golden_fw_start + rnn.num_fw_units(); + fw_expected.insert(fw_expected.begin(), golden_fw_start, golden_fw_end); + } + fw_expected.insert(fw_expected.end(), fw_expected.begin(), fw_expected.end()); + EXPECT_THAT(rnn.GetFwOutput(), ElementsAreArray(ArrayFloatNear(fw_expected))); + + std::vector bw_expected; + for (int i = 0; i < rnn.sequence_len(); i++) { + float* golden_bw_start = rnn_golden_fw_output + i * rnn.num_bw_units(); + float* golden_bw_end = golden_bw_start + rnn.num_bw_units(); + bw_expected.insert(bw_expected.begin(), golden_bw_start, golden_bw_end); + } + bw_expected.insert(bw_expected.end(), bw_expected.begin(), bw_expected.end()); + EXPECT_THAT(rnn.GetBwOutput(), ElementsAreArray(ArrayFloatNear(bw_expected))); +} + +// Tests an end-to-end neural network with a Bidirectional RNN followed by a +// DNN that aggregates the outputs from the two sequences. +TEST(BidirectionalRNNOpTest, EndToEndTest) { + BidirectionalRNNOpModel rnn(/*batches=*/1, /*sequence_len=*/4, + /*fw_units=*/16, /*bw_units=*/16, + /*input_size=*/8); + const int output_size = 4; + float dnn_weights[] = { + -0.5782342, -0.052212059, 0.73036242, -0.81216097, -0.80088139, + -0.23420811, -0.39647382, 0.31423986, 0.61819065, -0.73659575, + -0.89698344, -0.8931554, -0.0845688, 0.5617367, 0.38415289, + -0.11487955, -0.7617774, 0.17927337, 0.15726972, 0.059798479, + 0.19009054, -0.27616632, -0.39142907, 0.77744663, -0.046830714, + -0.6603595, 0.21945822, 0.051494241, 0.23785079, 0.19239247, + -0.53268754, 0.65961659, -0.85981959, -0.80232513, 0.84745562, + -0.66070104, -0.036533296, -0.54901814, 0.65353882, -0.41834265, + -0.28561389, 0.75655544, -0.31149811, 0.62981737, 0.31829214, + -0.92734522, -0.48506218, 0.55651462, 0.25192821, 0.67220747, + -0.3836869, -0.55798125, -0.60395885, 0.22488403, -0.78053463, + 0.3492105, 0.56452453, 0.4389236, -0.59929526, -0.19762468, + -0.36868393, -0.13198286, -0.53800809, -0.22850353}; + + std::initializer_list dnn_biases = { + 0.29177809, -0.98799044, 0.065919638, 0.68781924}; + + rnn.SetFwWeights(weights); + rnn.SetBwWeights(weights); + rnn.SetFwBias(biases); + rnn.SetBwBias(biases); + rnn.SetFwRecurrentWeights(recurrent_weights); + rnn.SetBwRecurrentWeights(recurrent_weights); + + rnn.ResetHiddenStates(); + + const int input_sequence_size = rnn.input_size() * rnn.sequence_len(); + const int output_sequence_size = output_size * rnn.sequence_len(); + const int num_examples = 64; + for (int k = 0; k < num_examples; k++) { + float* batch_start = endtoend_input + k * input_sequence_size; + float* batch_end = batch_start + input_sequence_size; + rnn.SetInput(0, batch_start, batch_end); + + rnn.Invoke(); + + std::vector fw_output = rnn.GetFwOutput(); + std::vector bw_output = rnn.GetBwOutput(); + EXPECT_EQ(fw_output.size(), bw_output.size()); + + std::transform(fw_output.begin(), fw_output.end(), bw_output.begin(), + fw_output.begin(), std::plus()); + + std::vector sequence_result; + for (int s = 0; s < rnn.sequence_len(); s++) { + const float* rnn_output = fw_output.data() + s * rnn.num_fw_units(); + std::vector results(dnn_biases); + for (int i = 0; i < output_size; i++) { + for (int j = 0; j < rnn.num_fw_units(); j++) { + results[i] += *(rnn_output + j) * dnn_weights[output_size * j + i]; + } + } + sequence_result.insert(sequence_result.end(), results.begin(), + results.end()); + } + + float* golden_start = golden_endtoend_output + k * output_sequence_size; + float* golden_end = golden_start + output_sequence_size; + + std::vector expected; + expected.insert(expected.end(), golden_start, golden_end); + EXPECT_THAT(sequence_result, ElementsAreArray(ArrayFloatNear(expected))); + } +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + // On Linux, add: tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/concatenation.cc b/tensorflow/contrib/lite/kernels/concatenation.cc index 9e7a1233dac0f3cd02dc386f9d194597f38ca3b8..a619ada86af64c299f8e518a7493db20f1011a50 100644 --- a/tensorflow/contrib/lite/kernels/concatenation.cc +++ b/tensorflow/contrib/lite/kernels/concatenation.cc @@ -49,6 +49,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // dimensions except 'axis' must be equal. TfLiteTensor* t0 = &context->tensors[node->inputs->data[0]]; TfLiteType input_type = t0->type; + if (axis < 0) axis += t0->dims->size; TF_LITE_ENSURE(context, axis >= 0); TF_LITE_ENSURE(context, axis < t0->dims->size); @@ -95,53 +96,22 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { return context->ResizeTensor(context, output, output_size); } -template -class VectorOfInputs { - public: - VectorOfInputs(const TfLiteContext& context, const TfLiteIntArray& inputs) { - int num_inputs = inputs.size; - - all_data_.reserve(num_inputs); - all_dims_.reserve(num_inputs); - all_dims_ptr_.reserve(num_inputs); - - for (int i = 0; i < num_inputs; ++i) { - TfLiteTensor* input = &context.tensors[inputs.data[i]]; - all_data_.push_back(GetTensorData(input)); - all_dims_.push_back(GetTensorDims(input)); - } - - // Taking the pointer from inside a std::vector is only OK if the vector is - // never modified, so we populate all_dims in the previous loop and then we - // are free to grab iterators here. - for (int i = 0; i < num_inputs; ++i) { - all_dims_ptr_.push_back(&all_dims_[i]); - } - } - const T* const* data() const { return all_data_.data(); } - const Dims<4>* const* dims() const { return all_dims_ptr_.data(); } - - private: - std::vector all_data_; - std::vector> all_dims_; - std::vector*> all_dims_ptr_; -}; - template TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->builtin_data); - + int axis = params->axis; TfLiteTensor* output = &context->tensors[node->outputs->data[0]]; + if (axis < 0) axis += output->dims->size; // TODO(ahentz): Creating 'all_inputs' below is not very efficient. We should // allocate and populate these during Prepare(). // TODO(ycling): Activation function parameter is ignored. For now we dont have // a model with a Concatenation with fused activation function. #define TF_LITE_CONCATENATION(type, scalar) \ - VectorOfInputs all_inputs(*context, *node->inputs); \ + VectorOfTensors all_inputs(*context, *node->inputs); \ type::Concatenation( \ - RemapDim(NumDimensions(output), params->axis), all_inputs.data(), \ + RemapDim(NumDimensions(output), axis), all_inputs.data(), \ all_inputs.dims(), node->inputs->size, GetTensorData(output), \ GetTensorDims(output)) diff --git a/tensorflow/contrib/lite/kernels/concatenation_test.cc b/tensorflow/contrib/lite/kernels/concatenation_test.cc index 499856a93cbbfbf9aa1a326912e52ce32bbbdf83..ba1ffc5f8423b9626c9c8e2a1086ea0dcca43f50 100644 --- a/tensorflow/contrib/lite/kernels/concatenation_test.cc +++ b/tensorflow/contrib/lite/kernels/concatenation_test.cc @@ -94,7 +94,7 @@ TEST(ConcatenationOpTest, TwoDimensionalOneInput) { EXPECT_THAT(m0.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6})); } -TEST(ConcatenationOpTest, TwoInputsTwoAxis) { +TEST(ConcatenationOpTest, TwoInputsTwoAxesNegativeAxes) { // We will concatenate two tensors along different dimensions. auto tensor0 = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; auto tensor1 = {7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f}; @@ -107,6 +107,14 @@ TEST(ConcatenationOpTest, TwoInputsTwoAxis) { EXPECT_THAT(m0.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + ConcatenationOpModel m0_negative({TensorType_FLOAT32, {2, 3}}, /*axis=*/-2, + /*num_inputs=*/2); + m0_negative.SetInput(0, tensor0); + m0_negative.SetInput(1, tensor1); + m0_negative.Invoke(); + EXPECT_THAT(m0_negative.GetOutput(), + ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + ConcatenationOpModel m1({TensorType_FLOAT32, {2, 3}}, /*axis=*/1, /*num_inputs=*/2); m1.SetInput(0, tensor0); @@ -114,6 +122,14 @@ TEST(ConcatenationOpTest, TwoInputsTwoAxis) { m1.Invoke(); EXPECT_THAT(m1.GetOutput(), ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12})); + + ConcatenationOpModel m1_negative({TensorType_FLOAT32, {2, 3}}, /*axis=*/-1, + /*num_inputs=*/2); + m1_negative.SetInput(0, tensor0); + m1_negative.SetInput(1, tensor1); + m1_negative.Invoke(); + EXPECT_THAT(m1_negative.GetOutput(), + ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12})); } TEST(ConcatenationOpTest, FourInputs) { diff --git a/tensorflow/contrib/lite/kernels/conv.cc b/tensorflow/contrib/lite/kernels/conv.cc index c75c04baeac2ce53c6261d677dca8d72fafa0da5..66d2c04bba4a164bbcdcf4b1a097d9aac0b3aeeb 100644 --- a/tensorflow/contrib/lite/kernels/conv.cc +++ b/tensorflow/contrib/lite/kernels/conv.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/contrib/lite/builtin_op_data.h" #include "tensorflow/contrib/lite/context.h" #include "tensorflow/contrib/lite/kernels/gemm_support.h" +#include "tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h" #include "tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h" #include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" #include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" @@ -38,11 +39,16 @@ namespace ops { namespace builtin { namespace conv { -// This file has three implementation of Conv. +// This file has 4 implementation of Conv. enum KernelType { kReference, kGenericOptimized, // Neon-free - kNeonOptimized, + kMultithreadOptimized, + // The kernel uses use CBLAS interface for matrix multiplication. + // It's fast when an optimized CBLAS implementation is available (e.g. Apple + // Accelerate Framework), and it's slow when falling back to naive + // implementation. + kCblasOptimized, }; struct OpData { @@ -265,10 +271,13 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { free(hwcn_weights->data.raw); hwcn_weights->data.raw = nullptr; } + + // Note that hwcn_weights_status is a kTfLiteDynamic tensor, and + // ResizeTensor will actually allocate space for it. The would be more + // efficient if we placed hwcn_weights_status in the persistent arena. auto hwcn_weights_status = context->ResizeTensor(context, hwcn_weights, hwcn_weights_size); if (hwcn_weights_status != kTfLiteOk) return hwcn_weights_status; - hwcn_weights->data.raw = static_cast(malloc(hwcn_weights->bytes)); // TODO(petewarden): If Resize() is called when the size hasn't actually // changed, this will do extra redundant work. @@ -290,26 +299,34 @@ void EvalQuantized(TfLiteContext* context, TfLiteNode* node, auto filter_offset = -filter->params.zero_point; auto output_offset = output->params.zero_point; - if (kernel_type == kReference) { - reference_ops::Conv( - GetTensorData(input), GetTensorDims(input), input_offset, - GetTensorData(filter), GetTensorDims(filter), filter_offset, - GetTensorData(bias), GetTensorDims(bias), params->stride_width, - params->stride_height, data->padding.width, data->padding.height, - output_offset, data->output_multiplier, data->output_shift, - data->output_activation_min, data->output_activation_max, - GetTensorData(output), GetTensorDims(output), - GetTensorData(im2col), GetTensorDims(im2col), gemm_context); - } else { - optimized_ops::Conv( - GetTensorData(input), GetTensorDims(input), input_offset, - GetTensorData(filter), GetTensorDims(filter), filter_offset, - GetTensorData(bias), GetTensorDims(bias), params->stride_width, - params->stride_height, data->padding.width, data->padding.height, - output_offset, data->output_multiplier, data->output_shift, - data->output_activation_min, data->output_activation_max, - GetTensorData(output), GetTensorDims(output), - GetTensorData(im2col), GetTensorDims(im2col), gemm_context); + switch (kernel_type) { + case kReference: + reference_ops::Conv( + GetTensorData(input), GetTensorDims(input), input_offset, + GetTensorData(filter), GetTensorDims(filter), filter_offset, + GetTensorData(bias), GetTensorDims(bias), + params->stride_width, params->stride_height, data->padding.width, + data->padding.height, output_offset, data->output_multiplier, + data->output_shift, data->output_activation_min, + data->output_activation_max, GetTensorData(output), + GetTensorDims(output), GetTensorData(im2col), + GetTensorDims(im2col), gemm_context); + break; + case kGenericOptimized: + case kMultithreadOptimized: + case kCblasOptimized: + // There is only one optimized implementation for Quantized Conv. + optimized_ops::Conv( + GetTensorData(input), GetTensorDims(input), input_offset, + GetTensorData(filter), GetTensorDims(filter), filter_offset, + GetTensorData(bias), GetTensorDims(bias), + params->stride_width, params->stride_height, data->padding.width, + data->padding.height, output_offset, data->output_multiplier, + data->output_shift, data->output_activation_min, + data->output_activation_max, GetTensorData(output), + GetTensorDims(output), GetTensorData(im2col), + GetTensorDims(im2col), gemm_context); + break; } } @@ -322,30 +339,57 @@ void EvalFloat(TfLiteContext* context, TfLiteNode* node, CalculateActivationRangeFloat(params->activation, &output_activation_min, &output_activation_max); - const float* filter_data; - if (data->need_hwcn_weights) { - filter_data = GetTensorData(hwcn_weights); - } else { - filter_data = GetTensorData(filter); - } - - if (kernel_type == kReference) { - reference_ops::Conv( - GetTensorData(input), GetTensorDims(input), filter_data, - GetTensorDims(filter), GetTensorData(bias), GetTensorDims(bias), - params->stride_width, params->stride_height, data->padding.width, - data->padding.height, output_activation_min, output_activation_max, - GetTensorData(output), GetTensorDims(output), - GetTensorData(im2col), GetTensorDims(im2col)); - } else { - multithreaded_ops::Conv( - GetTensorData(input), GetTensorDims(input), filter_data, - GetTensorDims(filter), GetTensorData(bias), GetTensorDims(bias), - params->stride_width, params->stride_height, data->padding.width, - data->padding.height, params->padding, output_activation_min, - output_activation_max, GetTensorData(output), - GetTensorDims(output), GetTensorData(im2col), - GetTensorDims(im2col)); + switch (kernel_type) { + case kReference: { + reference_ops::Conv(GetTensorData(input), GetTensorDims(input), + GetTensorData(filter), GetTensorDims(filter), + GetTensorData(bias), GetTensorDims(bias), + params->stride_width, params->stride_height, + data->padding.width, data->padding.height, + output_activation_min, output_activation_max, + GetTensorData(output), GetTensorDims(output), + GetTensorData(im2col), GetTensorDims(im2col)); + break; + } + case kGenericOptimized: { + optimized_ops::Conv(GetTensorData(input), GetTensorDims(input), + GetTensorData(filter), GetTensorDims(filter), + GetTensorData(bias), GetTensorDims(bias), + params->stride_width, params->stride_height, + data->padding.width, data->padding.height, + output_activation_min, output_activation_max, + GetTensorData(output), GetTensorDims(output), + GetTensorData(im2col), GetTensorDims(im2col)); + break; + } + case kMultithreadOptimized: { + const float* filter_data; + if (data->need_hwcn_weights) { + filter_data = GetTensorData(hwcn_weights); + } else { + filter_data = GetTensorData(filter); + } + multithreaded_ops::Conv( + GetTensorData(input), GetTensorDims(input), filter_data, + GetTensorDims(filter), GetTensorData(bias), + GetTensorDims(bias), params->stride_width, params->stride_height, + data->padding.width, data->padding.height, params->padding, + output_activation_min, output_activation_max, + GetTensorData(output), GetTensorDims(output), + GetTensorData(im2col), GetTensorDims(im2col)); + break; + } + case kCblasOptimized: { + cblas_ops::Conv(GetTensorData(input), GetTensorDims(input), + GetTensorData(filter), GetTensorDims(filter), + GetTensorData(bias), GetTensorDims(bias), + params->stride_width, params->stride_height, + data->padding.width, data->padding.height, + output_activation_min, output_activation_max, + GetTensorData(output), GetTensorDims(output), + GetTensorData(im2col), GetTensorDims(im2col)); + break; + } } } @@ -406,17 +450,23 @@ TfLiteRegistration* Register_CONVOLUTION_GENERIC_OPT() { return &r; } -TfLiteRegistration* Register_CONVOLUTION_NEON_OPT() { +TfLiteRegistration* Register_CONVOLUTION_MULTITHREADED_OPT() { + static TfLiteRegistration r = {conv::Init, conv::Free, conv::Prepare, + conv::Eval}; + return &r; +} + +TfLiteRegistration* Register_CONVOLUTION_CBLAS_OPT() { static TfLiteRegistration r = {conv::Init, conv::Free, conv::Prepare, - conv::Eval}; + conv::Eval}; return &r; } TfLiteRegistration* Register_CONV_2D() { -#ifdef USE_NEON - return Register_CONVOLUTION_NEON_OPT(); +#ifdef TFLITE_USE_APPLE_ACCELERATE_FOR_CONV + return Register_CONVOLUTION_CBLAS_OPT(); #else - return Register_CONVOLUTION_GENERIC_OPT(); + return Register_CONVOLUTION_MULTITHREADED_OPT(); #endif } diff --git a/tensorflow/contrib/lite/kernels/conv_test.cc b/tensorflow/contrib/lite/kernels/conv_test.cc index 1d0a81c3135625c07a3566f5f9a8e5401f0d4db7..d2393c3c97bb9516e2b8a6c8ae037dc0dfdfe64b 100644 --- a/tensorflow/contrib/lite/kernels/conv_test.cc +++ b/tensorflow/contrib/lite/kernels/conv_test.cc @@ -15,12 +15,25 @@ limitations under the License. #include #include +#include "absl/memory/memory.h" #include "tensorflow/contrib/lite/interpreter.h" #include "tensorflow/contrib/lite/kernels/register.h" #include "tensorflow/contrib/lite/kernels/test_util.h" #include "tensorflow/contrib/lite/model.h" namespace tflite { + +namespace ops { +namespace builtin { + +TfLiteRegistration* Register_CONVOLUTION_REF(); +TfLiteRegistration* Register_CONVOLUTION_GENERIC_OPT(); +TfLiteRegistration* Register_CONVOLUTION_MULTITHREADED_OPT(); +TfLiteRegistration* Register_CONVOLUTION_CBLAS_OPT(); + +} // namespace builtin +} // namespace ops + namespace { using ::testing::ElementsAreArray; @@ -30,9 +43,9 @@ class BaseConvolutionOpModel : public SingleOpModel { // TODO(ahentz): Also test different activation types, bias, padding types, // stride values. BaseConvolutionOpModel( - const TensorData& input, const TensorData& filter, - const TensorData& output, int stride_width = 2, int stride_height = 2, - enum Padding padding = Padding_VALID, + TfLiteRegistration* registration, const TensorData& input, + const TensorData& filter, const TensorData& output, int stride_width = 2, + int stride_height = 2, enum Padding padding = Padding_VALID, enum ActivationFunctionType activation = ActivationFunctionType_NONE) { input_ = AddInput(input); filter_ = AddInput(filter); @@ -62,6 +75,8 @@ class BaseConvolutionOpModel : public SingleOpModel { stride_height, activation) .Union()); + resolver_ = absl::make_unique(BuiltinOperator_CONV_2D, + registration); BuildInterpreter({GetShape(input_), GetShape(filter_), GetShape(bias_)}); } @@ -83,12 +98,26 @@ class ConvolutionOpModel : public BaseConvolutionOpModel { void SetInput(std::initializer_list data) { PopulateTensor(input_, data); } - std::vector GetOutput() { return ExtractVector(output_); } }; -TEST(ConvolutionOpTest, SimpleTestFloat32) { - ConvolutionOpModel m({TensorType_FLOAT32, {2, 2, 4, 1}}, +const auto kKernelMap = new std::map({ + {"Reference", ops::builtin::Register_CONVOLUTION_REF()}, + {"GenericOptimized", ops::builtin::Register_CONVOLUTION_GENERIC_OPT()}, + {"MultithreadedOptimized", + ops::builtin::Register_CONVOLUTION_MULTITHREADED_OPT()}, + {"CblasOptimized", ops::builtin::Register_CONVOLUTION_CBLAS_OPT()}, +}); + +class ConvolutionOpTest : public SingleOpTest { + protected: + const std::map& GetKernelMap() override { + return *kKernelMap; + } +}; + +TEST_P(ConvolutionOpTest, SimpleTestFloat32) { + ConvolutionOpModel m(GetRegistration(), {TensorType_FLOAT32, {2, 2, 4, 1}}, {TensorType_FLOAT32, {3, 2, 2, 1}}, {TensorType_FLOAT32, {}}); @@ -117,8 +146,8 @@ TEST(ConvolutionOpTest, SimpleTestFloat32) { })); } -TEST(ConvolutionOpTest, SimpleTestFloat32WithAnisotropicStrides) { - ConvolutionOpModel m({TensorType_FLOAT32, {1, 3, 6, 1}}, +TEST_P(ConvolutionOpTest, SimpleTestFloat32WithAnisotropicStrides) { + ConvolutionOpModel m(GetRegistration(), {TensorType_FLOAT32, {1, 3, 6, 1}}, {TensorType_FLOAT32, {1, 2, 2, 1}}, {TensorType_FLOAT32, {}}, /*stride_width=*/3, /*stride_height=*/1); @@ -139,7 +168,7 @@ TEST(ConvolutionOpTest, SimpleTestFloat32WithAnisotropicStrides) { })); } -TEST(ConvolutionOpTest, HandCalculatedFloat32) { +TEST_P(ConvolutionOpTest, HandCalculatedFloat32) { const int depth = 1; const int image_width = 4; const int image_height = 3; @@ -150,6 +179,7 @@ TEST(ConvolutionOpTest, HandCalculatedFloat32) { const int stride_height = 1; const Padding padding = Padding_SAME; ConvolutionOpModel m( + GetRegistration(), {TensorType_FLOAT32, {image_batch_count, image_height, image_width, depth}}, {TensorType_FLOAT32, {depth, filter_size, filter_size, filter_count}}, @@ -192,7 +222,7 @@ TEST(ConvolutionOpTest, HandCalculatedFloat32) { 178, 187, 234, 261, 121})); } -TEST(ConvolutionOpTest, HandCalculatedWithBiasFloat32) { +TEST_P(ConvolutionOpTest, HandCalculatedWithBiasFloat32) { const int depth = 1; const int image_width = 4; const int image_height = 3; @@ -203,6 +233,7 @@ TEST(ConvolutionOpTest, HandCalculatedWithBiasFloat32) { const int stride_height = 1; const Padding padding = Padding_SAME; ConvolutionOpModel m( + GetRegistration(), {TensorType_FLOAT32, {image_batch_count, image_height, image_width, depth}}, {TensorType_FLOAT32, {depth, filter_size, filter_size, filter_count}}, @@ -245,7 +276,7 @@ TEST(ConvolutionOpTest, HandCalculatedWithBiasFloat32) { 367, 188, 197, 244, 271, 131})); } -TEST(ConvolutionOpTest, HandCalculatedWithReluFloat32) { +TEST_P(ConvolutionOpTest, HandCalculatedWithReluFloat32) { const int depth = 1; const int image_width = 4; const int image_height = 3; @@ -256,6 +287,7 @@ TEST(ConvolutionOpTest, HandCalculatedWithReluFloat32) { const int stride_height = 1; const Padding padding = Padding_SAME; ConvolutionOpModel m( + GetRegistration(), {TensorType_FLOAT32, {image_batch_count, image_height, image_width, depth}}, {TensorType_FLOAT32, {depth, filter_size, filter_size, filter_count}}, @@ -300,7 +332,7 @@ TEST(ConvolutionOpTest, HandCalculatedWithReluFloat32) { ElementsAreArray({0, 0, 0, 0, 35, 112, 157, 0, 0, 34, 61, 0})); } -TEST(ConvolutionOpTest, HandCalculatedValidFloat32) { +TEST_P(ConvolutionOpTest, HandCalculatedValidFloat32) { const int depth = 1; const int image_width = 4; const int image_height = 3; @@ -311,6 +343,7 @@ TEST(ConvolutionOpTest, HandCalculatedValidFloat32) { const int stride_height = 1; const Padding padding = Padding_VALID; ConvolutionOpModel m( + GetRegistration(), {TensorType_FLOAT32, {image_batch_count, image_height, image_width, depth}}, {TensorType_FLOAT32, {depth, filter_size, filter_size, filter_count}}, @@ -366,8 +399,9 @@ class QuantizedConvolutionOpModel : public BaseConvolutionOpModel { // In this tests we set the input and output scales so that the results // match exactly the 'non-quantized' version. -TEST(ConvolutionOpTest, SimpleTestQuantized) { - QuantizedConvolutionOpModel m({TensorType_UINT8, {2, 2, 4, 1}, -63.5, 64}, +TEST_P(ConvolutionOpTest, SimpleTestQuantized) { + QuantizedConvolutionOpModel m(GetRegistration(), + {TensorType_UINT8, {2, 2, 4, 1}, -63.5, 64}, {TensorType_UINT8, {3, 2, 2, 1}, -63.5, 64}, {TensorType_UINT8, {}, -127, 128}); m.SetInput({ @@ -405,8 +439,9 @@ TEST(ConvolutionOpTest, SimpleTestQuantized) { })); } -TEST(ConvolutionOpTest, SimpleTestQuantizedWithAnisotropicStrides) { - QuantizedConvolutionOpModel m({TensorType_UINT8, {1, 3, 6, 1}, -63.5, 64}, +TEST_P(ConvolutionOpTest, SimpleTestQuantizedWithAnisotropicStrides) { + QuantizedConvolutionOpModel m(GetRegistration(), + {TensorType_UINT8, {1, 3, 6, 1}, -63.5, 64}, {TensorType_UINT8, {1, 2, 2, 1}, -63.5, 64}, {TensorType_UINT8, {}, -127, 128}, /*stride_width=*/3, /*stride_height=*/1); @@ -430,6 +465,11 @@ TEST(ConvolutionOpTest, SimpleTestQuantizedWithAnisotropicStrides) { 167, 93, // })); } + +INSTANTIATE_TEST_CASE_P( + ConvolutionOpTest, ConvolutionOpTest, + ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap))); + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/div.cc b/tensorflow/contrib/lite/kernels/div.cc new file mode 100644 index 0000000000000000000000000000000000000000..44bd0dc85d50c98ec6b6888e05064a8f2e2731c0 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/div.cc @@ -0,0 +1,129 @@ +/* Copyright 2017 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/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace div { + +// This file has three implementation of Div. +enum KernelType { + kReference, + kGenericOptimized, // Neon-free + kNeonOptimized, +}; + +constexpr int kInputTensor1 = 0; +constexpr int kInputTensor2 = 1; +constexpr int kOutputTensor = 0; + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); + TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + TF_LITE_ENSURE_EQ(context, NumDimensions(input1), NumDimensions(input2)); + for (int i = 0; i < NumDimensions(input1); ++i) { + TF_LITE_ENSURE_EQ(context, SizeOfDimension(input1, i), + SizeOfDimension(input2, i)); + } + + TF_LITE_ENSURE_EQ(context, input1->type, output->type); + TF_LITE_ENSURE_EQ(context, input2->type, output->type); + + TfLiteIntArray* output_size = TfLiteIntArrayCopy(input1->dims); + return context->ResizeTensor(context, output, output_size); +} + +template +void EvalDivFloat(TfLiteContext* context, TfLiteNode* node, + TfLiteDivParams* params, TfLiteTensor* input1, + TfLiteTensor* input2, TfLiteTensor* output) { + float output_activation_min, output_activation_max; + CalculateActivationRangeFloat(params->activation, &output_activation_min, + &output_activation_max); +#define TF_LITE_DIV(type) \ + type::Div(GetTensorData(input1), GetTensorDims(input1), \ + GetTensorData(input2), GetTensorDims(input2), \ + output_activation_min, output_activation_max, \ + GetTensorData(output), GetTensorDims(output)) + if (kernel_type == kReference) { + TF_LITE_DIV(reference_ops); + } else { + TF_LITE_DIV(optimized_ops); + } +#undef TF_LITE_DIV +} + +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + auto* params = reinterpret_cast(node->builtin_data); + + TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); + TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + if (output->type == kTfLiteFloat32) { + EvalDivFloat(context, node, params, input1, input2, output); + } else { + context->ReportError(context, "Inputs and outputs not all float types."); + return kTfLiteError; + } + + return kTfLiteOk; +} + +} // namespace div + +TfLiteRegistration* Register_DIV_REF() { + static TfLiteRegistration r = {nullptr, nullptr, div::Prepare, + div::Eval}; + return &r; +} + +TfLiteRegistration* Register_DIV_GENERIC_OPT() { + static TfLiteRegistration r = {nullptr, nullptr, div::Prepare, + div::Eval}; + return &r; +} + +TfLiteRegistration* Register_DIV_NEON_OPT() { + static TfLiteRegistration r = {nullptr, nullptr, div::Prepare, + div::Eval}; + return &r; +} + +TfLiteRegistration* Register_DIV() { +#ifdef USE_NEON + return Register_DIV_NEON_OPT(); +#else + return Register_DIV_GENERIC_OPT(); +#endif +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/embedding_lookup_sparse_test.cc b/tensorflow/contrib/lite/kernels/embedding_lookup_sparse_test.cc index dcdc5fffad9ceac1a9d23a4e91637a9ff92a8dda..ef2b5422253ea880a9ded4d3c0efc5cec07178a9 100644 --- a/tensorflow/contrib/lite/kernels/embedding_lookup_sparse_test.cc +++ b/tensorflow/contrib/lite/kernels/embedding_lookup_sparse_test.cc @@ -123,18 +123,16 @@ TEST(EmbeddingLookupOpTest, SimpleTestSqrtn) { [](int i, int j, int k) { return i + j / 10.0f + k / 100.0f; }); m.Invoke(); - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray(ArrayFloatNear({ - 1.00, 1.01, 1.10, 1.11, 1.20, 1.21, // Row 1 - 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, // - - 6.00f / std::sqrt(20.0f), 6.06f / std::sqrt(20.0f), - 6.60f / std::sqrt(20.0f), 6.66f / std::sqrt(20.0f), - 7.20f / std::sqrt(20.0f), - 7.26f / - std::sqrt( - 20.0f), // 2 * Row 3 + 4 * Row 0, // 2 * Row 3 + 4 * Row 0 - }))); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear({ + 1.00, 1.01, 1.10, 1.11, 1.20, 1.21, // Row 1 + 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, // - + 6.00f / std::sqrt(20.0f), 6.06f / std::sqrt(20.0f), + 6.60f / std::sqrt(20.0f), 6.66f / std::sqrt(20.0f), + 7.20f / std::sqrt(20.0f), + 7.26f / std::sqrt(20.0f), // 2 * Row 3 + 4 * Row 0, // 2 * + // Row 3 + 4 * Row 0 + }))); } TEST(EmbeddingLookupOpTest, Indices3DTest) { diff --git a/tensorflow/contrib/lite/kernels/exp.cc b/tensorflow/contrib/lite/kernels/exp.cc new file mode 100644 index 0000000000000000000000000000000000000000..a9e79b742dc2c80ce4ed9a3aa786814265dcb660 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/exp.cc @@ -0,0 +1,92 @@ +/* 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 +#include +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace exp { + +// This file has reference implementation of Exp. +enum KernelType { + kReference, +}; + +struct ExpContext { + ExpContext(TfLiteContext* context, TfLiteNode* node) { + input = GetInput(context, node, 0); + output = GetOutput(context, node, 0); + } + TfLiteTensor* input; + TfLiteTensor* output; +}; + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + ExpContext op_context(context, node); + TfLiteIntArray* output_dims = TfLiteIntArrayCopy(op_context.input->dims); + op_context.output->type = op_context.input->type; + return context->ResizeTensor(context, op_context.output, output_dims); +} + +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + ExpContext op_context(context, node); + +#define TF_LITE_EXP(kernel_type, data_type) \ + kernel_type::Exp(GetTensorData(op_context.input), \ + NumElements(op_context.input), \ + GetTensorData(op_context.output)) + + // TODO(kanlig): supports half, bfloat16, float64, complex64, and complex128. + if (kernel_type == kReference) { + switch (op_context.input->type) { + case kTfLiteFloat32: + TF_LITE_EXP(reference_ops, float); + break; + default: + context->ReportError(context, + "Type %d is currently not supported by Exp.", + op_context.input->type); + return kTfLiteError; + } + } +#undef TF_LITE_EXP + return kTfLiteOk; +} + +} // namespace exp + +TfLiteRegistration* Register_EXP_REF() { + static TfLiteRegistration r = {nullptr, nullptr, exp::Prepare, + exp::Eval}; + return &r; +} + +// TODO(kanlig): add optimized implementation of Exp. +TfLiteRegistration* Register_EXP() { return Register_EXP_REF(); } + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/exp_test.cc b/tensorflow/contrib/lite/kernels/exp_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..eed67369a1f30e57cd29a3975a899db41938def0 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/exp_test.cc @@ -0,0 +1,70 @@ +/* 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 +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +class ExpOpModel : public SingleOpModel { + public: + ExpOpModel(const TensorData& input, const TensorType& output) { + input_ = AddInput(input); + output_ = AddOutput(output); + SetBuiltinOp(BuiltinOperator_EXP, BuiltinOptions_ExpOptions, + CreateExpOptions(builder_).Union()); + BuildInterpreter({GetShape(input_)}); + } + + template + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + + template + std::vector GetOutput() { + return ExtractVector(output_); + } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + protected: + int input_; + int output_; +}; + +TEST(ExpOpTest, FloatTest) { + std::initializer_list data = {1.0, 0.0, -1.0, 1.0, 1.0, -1.0}; + ExpOpModel m({TensorType_FLOAT32, {3, 1, 2}}, TensorType_FLOAT32); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 1, 2})); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear( + {2.71828, 1, 0.367879, 2.71828, 2.71828, 0.367879}))); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/gather.cc b/tensorflow/contrib/lite/kernels/gather.cc index f8df797daf7338e33b16508c21fc61cd9836db1e..0e4187d1eac64636a2e2b25e9a1cc45c3a4da557 100644 --- a/tensorflow/contrib/lite/kernels/gather.cc +++ b/tensorflow/contrib/lite/kernels/gather.cc @@ -42,9 +42,10 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_EQ(context, positions->type, kTfLiteInt32); // Check that input and output types match. TF_LITE_ENSURE_EQ(context, input->type, output->type); - // TODO(mgubin): only 1D positions are currently supported. - TF_LITE_ENSURE_EQ(context, NumDimensions(positions), 1); + // TODO(mgubin): only 0D or 1D positions are currently supported. + TF_LITE_ENSURE(context, NumDimensions(positions) <= 1); // TODO(mgubin): Only default axis == 0 is supported. + TF_LITE_ENSURE_EQ(context, params->axis, 0); // Check conditions for different types. switch (input->type) { case kTfLiteFloat32: @@ -64,7 +65,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { } const int num_dimensions = NumDimensions(input) + NumDimensions(positions) - 1; - TF_LITE_ENSURE(context, params->axis < num_dimensions); + TF_LITE_ENSURE(context, params->axis <= num_dimensions); TfLiteIntArray* output_shape = TfLiteIntArrayCreate(num_dimensions); int output_index = 0; for (int i = 0; i < params->axis; ++i) { diff --git a/tensorflow/contrib/lite/kernels/gather_test.cc b/tensorflow/contrib/lite/kernels/gather_test.cc index 6343d3b4ef20ae3e030396ec1b6adbcf83a3e45f..cdadbeda1884ba0186846826dd16be6ff69878d9 100644 --- a/tensorflow/contrib/lite/kernels/gather_test.cc +++ b/tensorflow/contrib/lite/kernels/gather_test.cc @@ -48,8 +48,8 @@ class GatherOpModel : public SingleOpModel { PopulateStringTensor(input_, data); } - void SetPositions(std::initializer_list data) { - PopulateTensor(positions_, data); + void SetPositions(std::initializer_list data) { + PopulateTensor(positions_, data); } std::vector GetOutputFloat() { return ExtractVector(output_); } @@ -76,6 +76,26 @@ TEST(GatherOpTest, Shuffle) { ElementsAreArray(ArrayFloatNear({0.7, 0.8, -2, 0.2}))); } +TEST(GatherOpTest, Test0DIndex) { + GatherOpModel m({2, 2}, TensorType_FLOAT32, {}); + m.SetInputFloat({-2.0, 0.2, 0.7, 0.8}); + m.SetPositions({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputFloat(), ElementsAreArray(ArrayFloatNear({0.7, 0.8}))); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); +} + +TEST(GatherOpTest, Test0DIndexWith0DResult) { + // 0D tensor is special case in current TFLite. Test it once to make sure + // existing workarounds are fine with it. + GatherOpModel m({3}, TensorType_FLOAT32, {}); + m.SetInputFloat({1.0, 2.0, 3.0}); + m.SetPositions({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputFloat(), ElementsAreArray(ArrayFloatNear({2.0}))); + EXPECT_TRUE(m.GetOutputShape().empty()); +} + TEST(FloatGatherOpTest, Duplicate) { GatherOpModel m({1, 2, 2}, TensorType_FLOAT32, {2}); m.SetInputFloat({-2.0, 0.2, 0.7, 0.8}); diff --git a/tensorflow/contrib/lite/kernels/gemm_support.h b/tensorflow/contrib/lite/kernels/gemm_support.h index b531959ffb143c774ee715743480b03ebfbdc114..466781cbcecc7fb851d9078c450cc6c12364d2bb 100644 --- a/tensorflow/contrib/lite/kernels/gemm_support.h +++ b/tensorflow/contrib/lite/kernels/gemm_support.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_GEMM_SUPPORT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_GEMM_SUPPORT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_GEMM_SUPPORT_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_GEMM_SUPPORT_H_ #include "public/gemmlowp.h" #include "tensorflow/contrib/lite/context.h" @@ -51,4 +51,4 @@ void SetMaxNumThreads(TfLiteContext* context, int num_threads); } // namespace gemm_support } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_GEMM_SUPPORT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_GEMM_SUPPORT_H_ diff --git a/tensorflow/contrib/lite/kernels/hashtable_lookup_test.cc b/tensorflow/contrib/lite/kernels/hashtable_lookup_test.cc index cb6038f9009a3865661e7b4f075c3033166d0f91..ba0ed5ce06392613238b757308dddc2b22e7eb30 100644 --- a/tensorflow/contrib/lite/kernels/hashtable_lookup_test.cc +++ b/tensorflow/contrib/lite/kernels/hashtable_lookup_test.cc @@ -116,7 +116,10 @@ TEST(HashtableLookupOpTest, Test2DInput) { 1.0, 1.1, // 1-st item }))); EXPECT_THAT(m.GetHit(), ElementsAreArray({ - 1, 0, 1, 1, + 1, + 0, + 1, + 1, })); } diff --git a/tensorflow/contrib/lite/kernels/internal/BUILD b/tensorflow/contrib/lite/kernels/internal/BUILD index a3ecb2ebf6a889729954d1e447997c510e8ff6d4..f47fb04cbaa688b75e763ff9d3cb7df44ac3f166 100644 --- a/tensorflow/contrib/lite/kernels/internal/BUILD +++ b/tensorflow/contrib/lite/kernels/internal/BUILD @@ -124,6 +124,13 @@ config_setting( }, ) +config_setting( + name = "darwin_x86_64", + values = { + "cpu": "darwin_x86_64", + }, +) + config_setting( name = "freebsd", values = { @@ -145,7 +152,7 @@ cc_library( ":types", ":round", "//third_party/eigen3", - "@gemmlowp//:gemmlowp", + "@gemmlowp", "//tensorflow/contrib/lite:builtin_op_data", ] + select({ ":haswell": tflite_deps_intel, @@ -154,6 +161,7 @@ cc_library( ":x86": tflite_deps_intel, ":x86_64": tflite_deps_intel, ":darwin": tflite_deps_intel, + ":darwin_x86_64": tflite_deps_intel, ":freebsd": tflite_deps_intel, "//conditions:default": [], }), @@ -162,6 +170,8 @@ cc_library( cc_library( name = "optimized", hdrs = [ + "optimized/cblas_conv.h", + "optimized/cblas_reference.h", "optimized/eigen_spatial_convolutions.h", "optimized/eigen_tensor_reduced_instantiations_oss.h", "optimized/multithreaded_conv.h", @@ -223,7 +233,7 @@ cc_library( ":round", ":types", "//third_party/eigen3", - "@gemmlowp//:gemmlowp", + "@gemmlowp", "//tensorflow/contrib/lite:builtin_op_data", ] + select({ ":haswell": tflite_deps_intel, @@ -232,6 +242,7 @@ cc_library( ":x86": tflite_deps_intel, ":x86_64": tflite_deps_intel, ":darwin": tflite_deps_intel, + ":darwin_x86_64": tflite_deps_intel, ":freebsd": tflite_deps_intel, "//conditions:default": [], }), @@ -267,6 +278,8 @@ cc_library( "optimized/neon_tensor_utils.cc", ], hdrs = [ + "common.h", + "optimized/cpu_check.h", "optimized/neon_tensor_utils.h", "optimized/tensor_utils_impl.h", ], @@ -274,8 +287,21 @@ cc_library( deps = [ ":cpu_check", ":portable_tensor_utils", + ":types", "//tensorflow/contrib/lite:builtin_op_data", "//tensorflow/contrib/lite/kernels:activation_functor", + "@arm_neon_2_x86_sse", + "@gemmlowp", + ], +) + +cc_library( + name = "kernel_utils", + srcs = ["kernel_utils.cc"], + hdrs = ["kernel_utils.h"], + deps = [ + ":tensor_utils", + "//tensorflow/contrib/lite:builtin_op_data", ], ) @@ -285,14 +311,21 @@ cc_library( "tensor_utils.cc", ], hdrs = [ + "common.h", + "compatibility.h", + "optimized/cpu_check.h", + "optimized/neon_tensor_utils.h", "optimized/tensor_utils_impl.h", "reference/portable_tensor_utils.h", "tensor_utils.h", + "types.h", ], copts = NEON_FLAGS_IF_APPLICABLE, deps = [ "//tensorflow/contrib/lite/kernels:activation_functor", "//tensorflow/contrib/lite:builtin_op_data", + "@arm_neon_2_x86_sse", + "@gemmlowp", ] + select({ ":arm": [ ":neon_tensor_utils", @@ -312,6 +345,21 @@ cc_library( ":ios_arm64": [ ":neon_tensor_utils", ], + ":ios_x86_64": [ + ":neon_tensor_utils", + ], + ":x86_64": [ + ":neon_tensor_utils", + ], + ":x86": [ + ":neon_tensor_utils", + ], + ":k8": [ + ":neon_tensor_utils", + ], + ":darwin": [ + ":neon_tensor_utils", + ], "//conditions:default": [ ":portable_tensor_utils", ], diff --git a/tensorflow/contrib/lite/kernels/internal/common.h b/tensorflow/contrib/lite/kernels/internal/common.h index 28f19a250629aec4d03aa71df57d31d8a5014e9f..18601df22c1894dea6ce51f46ba815cd12dab095 100644 --- a/tensorflow/contrib/lite/kernels/internal/common.h +++ b/tensorflow/contrib/lite/kernels/internal/common.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ #ifndef ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK #ifdef GEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK @@ -102,6 +102,17 @@ inline int32 MultiplyByQuantizedMultiplierGreaterThanOne( quantized_multiplier); } +inline int32 MultiplyByQuantizedMultiplier(int32 x, int32 quantized_multiplier, + int shift) { + using gemmlowp::RoundingDivideByPOT; + using gemmlowp::SaturatingRoundingDoublingHighMul; + int left_shift = shift > 0 ? shift : 0; + int right_shift = shift > 0 ? 0 : -shift; + return RoundingDivideByPOT(SaturatingRoundingDoublingHighMul( + x * (1 << left_shift), quantized_multiplier), + right_shift); +} + } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/compatibility.h b/tensorflow/contrib/lite/kernels/internal/compatibility.h index 796a03566a4bf971294dd2375f590dfd20d600f7..51426bb1c584b82af7b1a2ffaf5a675a1dd9a6fd 100644 --- a/tensorflow/contrib/lite/kernels/internal/compatibility.h +++ b/tensorflow/contrib/lite/kernels/internal/compatibility.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ #include #include @@ -27,6 +27,10 @@ limitations under the License. #define TFLITE_DCHECK_EQ(x, y) ((x) == (y)) ? (void)0 : assert(false) #endif +#ifndef TFLITE_DCHECK_NE +#define TFLITE_DCHECK_NE(x, y) ((x) != (y)) ? (void)0 : assert(false) +#endif + #ifndef TFLITE_DCHECK_GE #define TFLITE_DCHECK_GE(x, y) ((x) >= (y)) ? (void)0 : assert(false) #endif @@ -52,6 +56,10 @@ limitations under the License. #define TFLITE_CHECK_EQ(x, y) ((x) == (y)) ? (void)0 : abort() #endif +#ifndef TFLITE_CHECK_NE +#define TFLITE_CHECK_NE(x, y) ((x) != (y)) ? (void)0 : abort() +#endif + #ifndef TFLITE_CHECK_GE #define TFLITE_CHECK_GE(x, y) ((x) >= (y)) ? (void)0 : abort() #endif @@ -75,4 +83,4 @@ using uint16 = std::uint16_t; using int32 = std::int32_t; using uint32 = std::uint32_t; -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/kernel_utils.cc b/tensorflow/contrib/lite/kernels/internal/kernel_utils.cc new file mode 100644 index 0000000000000000000000000000000000000000..510395126ce3785b1d44fec1e0eb994c29ff0db7 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/kernel_utils.cc @@ -0,0 +1,44 @@ +/* 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/lite/kernels/internal/tensor_utils.h" + +namespace tflite { +namespace kernel_utils { + +void RnnBatchStep(const float* input_ptr_batch, const float* input_weights_ptr, + const float* recurrent_weights_ptr, const float* bias_ptr, + int input_size, int num_units, int batch_size, + TfLiteFusedActivation activation, + float* hidden_state_ptr_batch, float* output_ptr_batch) { + // Output = bias + tensor_utils::VectorBatchVectorAssign(bias_ptr, num_units, batch_size, + output_ptr_batch); + // Output += input * input_weights + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + input_weights_ptr, num_units, input_size, input_ptr_batch, batch_size, + output_ptr_batch, /*result_stride=*/1); + // Output += recurrent_weights * hidden_state + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + recurrent_weights_ptr, num_units, num_units, hidden_state_ptr_batch, + batch_size, output_ptr_batch, /*result_stride=*/1); + // Output = activation(Output) and update hidden_state + tensor_utils::ApplyActivationToVector( + output_ptr_batch, num_units * batch_size, activation, output_ptr_batch); + tensor_utils::VectorBatchVectorAssign(output_ptr_batch, num_units, batch_size, + hidden_state_ptr_batch); +} + +} // namespace kernel_utils +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/kernel_utils.h b/tensorflow/contrib/lite/kernels/internal/kernel_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..9872d4500b862388ed4b96c97e3755f548e35d35 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/kernel_utils.h @@ -0,0 +1,40 @@ +/* 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_LITE_KERNELS_INTERNAL_KERNEL_UTILS_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_KERNEL_UTILS_H_ + +#include "tensorflow/contrib/lite/builtin_op_data.h" + +namespace tflite { +namespace kernel_utils { + +// Performs an RNN batch inference step for inputs specified by input_ptr_batch. +// The RNN cell is specified by the pointers to its input and recurrent weights, +// and biases, along with the input size, number of units, activation. +// +// The pointers to the hidden state and the output are updated as a result. +// +// The pointers with the suffix "_batch" point to data aligned in batch_major +// order, and each step processes batch_size many inputs from input_ptr_batch, +// and updates batch_size many outputs and hidden states. +void RnnBatchStep(const float* input_ptr_batch, const float* input_weights_ptr, + const float* recurrent_weights_ptr, const float* bias_ptr, + int input_size, int num_units, int batch_size, + TfLiteFusedActivation activation, + float* hidden_state_ptr_batch, float* output_ptr_batch); + +} // namespace kernel_utils +} // namespace tflite +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_KERNEL_UTILS_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h b/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h new file mode 100644 index 0000000000000000000000000000000000000000..4a90e7e640ef29b675c236d8bbb479aa16560761 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h @@ -0,0 +1,92 @@ +/* Copyright 2017 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_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_CONV_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_CONV_H_ + +// The Conv implementation based on CBLAS interface. This is only used on iOS +// for now, utilizing Apple's Accelerate framework. + +#if TFLITE_USE_APPLE_ACCELERATE_FOR_CONV +#include +#else +#include "tensorflow/contrib/lite/kernels/internal/optimized/cblas_reference.h" +#endif + +#include "tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h" +#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" + +namespace tflite { +namespace cblas_ops { + +inline void Conv(const float* input_data, const Dims<4>& input_dims, + const float* filter_data, const Dims<4>& filter_dims, + const float* bias_data, const Dims<4>& bias_dims, + int stride_width, int stride_height, int pad_width, + int pad_height, float output_activation_min, + float output_activation_max, float* output_data, + const Dims<4>& output_dims, float* im2col_data, + const Dims<4>& im2col_dims) { + gemmlowp::ScopedProfilingLabel label("Conv/cblas"); + + const float* gemm_input_data = nullptr; + const Dims<4>* gemm_input_dims = nullptr; + const int filter_width = ArraySize(filter_dims, 1); + const int filter_height = ArraySize(filter_dims, 2); + const bool need_im2col = stride_width != 1 || stride_height != 1 || + filter_width != 1 || filter_height != 1; + if (need_im2col) { + TFLITE_DCHECK(im2col_data); + optimized_ops::Im2col(input_data, input_dims, stride_width, stride_height, + pad_width, pad_height, filter_height, filter_width, 0, + im2col_data, im2col_dims); + gemm_input_data = im2col_data; + gemm_input_dims = &im2col_dims; + } else { + TFLITE_DCHECK(!im2col_data); + gemm_input_data = input_data; + gemm_input_dims = &input_dims; + } + + // The following code computes matrix multiplication c = a * transponse(b) + // with CBLAS, where: + // * `a` is a matrix with dimensions (m, k). + // * `b` is a matrix with dimensions (n, k), so transpose(b) is (k, n). + // * `c` is a matrix with dimensions (m, n). + // The naming of variables are aligned with CBLAS specification here. + const float* a = gemm_input_data; + const float* b = filter_data; + float* c = output_data; + int m = gemm_input_dims->sizes[1] * gemm_input_dims->sizes[2] * + gemm_input_dims->sizes[3]; + int n = output_dims.sizes[0]; + int k = gemm_input_dims->sizes[0]; + // The stride of matrix a, b and c respectively. + int stride_a = k; + int stride_b = k; + int stride_c = n; + + cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, n, k, 1.0f, a, + stride_a, b, stride_b, 0.0f, c, stride_c); + + optimized_ops::AddBiasAndEvalActivationFunction( + bias_data, bias_dims, output_data, output_dims, output_activation_min, + output_activation_max); +} + +} // namespace cblas_ops +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_CONV_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/cblas_reference.h b/tensorflow/contrib/lite/kernels/internal/optimized/cblas_reference.h new file mode 100644 index 0000000000000000000000000000000000000000..6acc513805c9398c304f3e24175d3bd6c96938f6 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/optimized/cblas_reference.h @@ -0,0 +1,69 @@ +/* Copyright 2017 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_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_REFERENCE_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_REFERENCE_H_ + +#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" + +// The reference implementation for a small subset of CBLAS interface. +// This is only used for testing CBLAS implementation, and should never be used +// in production code. + +namespace tflite { +namespace cblas_ops { + +// The following code follows the original CBLAS specification, and it might +// conflict with the TensorFlow naming convention. +// TODO(ycling): Find another way to test CBLAS with bazel, without writing +// a reference implementation by ourselves. +enum CBLAS_ORDER { CblasRowMajor = 0, CblasColMajor = 1 }; + +enum CBLAS_TRANSPOSE { CblasNoTrans = 0, CblasTrans = 1, CblasConjTrans = 2 }; + +// A reference implementation for matrix multiplication. +// The following code computes, c = a * transponse(b) matrix multiplication +// with CBLAS, where: +// * `a` is a matrix with dimensions (m, k). +// * `b` is a matrix with dimensions (n, k), so transpose(b) is (k, n). +// * `c` is a matrix with dimensions (m, n). +// The naming of variables is aligned with CBLAS specification here. +void cblas_sgemm(const enum CBLAS_ORDER order, + const enum CBLAS_TRANSPOSE trans_a, + const enum CBLAS_TRANSPOSE trans_b, const int m, const int n, + const int k, const float alpha, const float *a, + const int stride_a, const float *b, const int stride_b, + const float beta, float *c, const int stride_c) { + TFLITE_DCHECK(order == CblasRowMajor); + TFLITE_DCHECK(trans_a == CblasNoTrans); + TFLITE_DCHECK(trans_b == CblasTrans); + TFLITE_DCHECK(beta == 0.0f); + for (int row = 0; row < m; ++row) { + for (int col = 0; col < n; ++col) { + // If `beta` non-zero, multiple it with the original values in output. + // Otherwise, ignore the original value in output completely. + float value = 0.0f; + for (int idx = 0; idx < k; ++idx) { + value += alpha * a[stride_a * row + idx] * b[stride_b * col + idx]; + } + c[stride_c * row + col] = value; + } + } +} + +} // namespace cblas_ops +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_REFERENCE_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/cpu_check.h b/tensorflow/contrib/lite/kernels/internal/optimized/cpu_check.h index dea46cc12065ed34cf681916a46a55bd7a86f463..3a53d3ab07faf63250fc18fc846e0b8f5a39d9c4 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/cpu_check.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/cpu_check.h @@ -34,17 +34,13 @@ inline bool TestCPUFeatureNeon() { #endif // __aarch64__ } -#elif __ARM_NEON +#elif defined USE_NEON || defined __ARM_NEON -inline bool TestCPUFeatureNeon() { - return true; -} +inline bool TestCPUFeatureNeon() { return true; } #else -inline bool TestCPUFeatureNeon() { - return false; -} +inline bool TestCPUFeatureNeon() { return false; } #endif diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h index da34c8aef94b1c69e661bd33fcb518e73034c4bd..7f6eea2d5d1cfd6f4e2a569760ecbe0d96f754c8 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ #include "public/gemmlowp.h" #include "tensorflow/contrib/lite/kernels/internal/common.h" @@ -573,6 +573,46 @@ struct FloatDepthwiseConvKernel { } }; +template <> +struct FloatDepthwiseConvKernel { + static void Run(int num_output_pixels, int input_depth, int depth_multiplier, + const float* input_ptr, int input_ptr_increment, + const float* filter_ptr, float* acc_buffer_ptr) { + // Load the filters + float32x4_t filter_0 = vld1q_f32(filter_ptr + 4 * 0); + float32x4_t filter_1 = vld1q_f32(filter_ptr + 4 * 1); + float32x4_t filter_2 = vld1q_f32(filter_ptr + 4 * 2); + float32x4_t filter_3 = vld1q_f32(filter_ptr + 4 * 3); + float32x4_t filter_4 = vld1q_f32(filter_ptr + 4 * 4); + + // Handle one output pixel at a time. + for (int outp = 0; outp < num_output_pixels; outp++) { + // Load the inputs + const float input_val = *input_ptr; + input_ptr += input_ptr_increment; + // Load the accumulators from acc_buffer + float32x4_t acc_0 = vld1q_f32(acc_buffer_ptr + 4 * 0); + float32x4_t acc_1 = vld1q_f32(acc_buffer_ptr + 4 * 1); + float32x4_t acc_2 = vld1q_f32(acc_buffer_ptr + 4 * 2); + float32x4_t acc_3 = vld1q_f32(acc_buffer_ptr + 4 * 3); + float32x4_t acc_4 = vld1q_f32(acc_buffer_ptr + 4 * 4); + // Multiply-accumulate + acc_0 = vmlaq_n_f32(acc_0, filter_0, input_val); + acc_1 = vmlaq_n_f32(acc_1, filter_1, input_val); + acc_2 = vmlaq_n_f32(acc_2, filter_2, input_val); + acc_3 = vmlaq_n_f32(acc_3, filter_3, input_val); + acc_4 = vmlaq_n_f32(acc_4, filter_4, input_val); + // Store the accumulators back to acc_buffer + vst1q_f32(acc_buffer_ptr + 4 * 0, acc_0); + vst1q_f32(acc_buffer_ptr + 4 * 1, acc_1); + vst1q_f32(acc_buffer_ptr + 4 * 2, acc_2); + vst1q_f32(acc_buffer_ptr + 4 * 3, acc_3); + vst1q_f32(acc_buffer_ptr + 4 * 4, acc_4); + acc_buffer_ptr += 20; + } + } +}; + template <> struct FloatDepthwiseConvKernel { static void Run(int num_output_pixels, int input_depth, int depth_multiplier, @@ -926,6 +966,7 @@ inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims, TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 1) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 8) + TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 20) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 32) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 2, 1) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 3, 2) @@ -992,11 +1033,11 @@ inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims, for (int k = 0; k < 4; k++) { acc[k] = vld1q_f32(acc_buffer + i + 4 * k); } - for (int k = 0; k < 4; k++) { - acc[k] = vmaxq_f32( - vdupq_n_f32(output_activation_min), - vminq_f32(vdupq_n_f32(output_activation_max), acc[k])); - } + for (int k = 0; k < 4; k++) { + acc[k] = vmaxq_f32( + vdupq_n_f32(output_activation_min), + vminq_f32(vdupq_n_f32(output_activation_max), acc[k])); + } for (int k = 0; k < 4; k++) { vst1q_f32(output_ptr + 4 * k, acc[k]); } @@ -1057,4 +1098,4 @@ void DepthwiseConv(const float* input_data, const Dims<4>& input_dims, } // namespace optimized_ops } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h index 051ed2a2c44a04f0473dfd26637e53865a5a51ac..dbc4f0d6fdca8279072d6ea225334722d6a89eb2 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ #include "fixedpoint/fixedpoint.h" #include "public/gemmlowp.h" @@ -1205,6 +1205,55 @@ struct QuantizedDepthwiseConvKernel { } }; +template <> +struct QuantizedDepthwiseConvKernel { + static void Run(int num_output_pixels, int input_depth, int depth_multiplier, + const uint8* input_ptr, int16 input_offset, + int input_ptr_increment, const uint8* filter_ptr, + int16 filter_offset, int32* acc_buffer_ptr) { + // Load the filters, add filter_offset. + // NEON wants to load 8 bytes at a time, but 20 is not divisible by 8. + // We load the first 16 bytes into filter_u8_{0,1} as usual. + // Then we load the 8 last bytes into filter_u8_x (x for 'extra'). + // This is redundant: the first 4 bytes of filter_u8_x are the same + // as the last 4 bytes of filter_u8_x. + uint8x8_t filter_u8_0 = vld1_u8(filter_ptr + 8 * 0); + uint8x8_t filter_u8_1 = vld1_u8(filter_ptr + 8 * 1); + uint8x8_t filter_u8_x = vld1_u8(filter_ptr + 8 * 1 + 4); + int16x8_t filter_0 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_0)); + int16x8_t filter_1 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_1)); + int16x8_t filter_x = vreinterpretq_s16_u16(vmovl_u8(filter_u8_x)); + filter_0 = vaddq_s16(filter_0, vdupq_n_s16(filter_offset)); + filter_1 = vaddq_s16(filter_1, vdupq_n_s16(filter_offset)); + filter_x = vaddq_s16(filter_x, vdupq_n_s16(filter_offset)); + // Handle one output pixel at a time. + for (int outp = 0; outp < num_output_pixels; outp++) { + uint8 input_u8 = *input_ptr; + input_ptr += input_ptr_increment; + int16 input = static_cast(input_u8 + input_offset); + // Load the accumulators from acc_buffer + int32x4_t acc_0 = vld1q_s32(acc_buffer_ptr + 4 * 0); + int32x4_t acc_1 = vld1q_s32(acc_buffer_ptr + 4 * 1); + int32x4_t acc_2 = vld1q_s32(acc_buffer_ptr + 4 * 2); + int32x4_t acc_3 = vld1q_s32(acc_buffer_ptr + 4 * 3); + int32x4_t acc_4 = vld1q_s32(acc_buffer_ptr + 4 * 4); + // Multiply-accumulate + acc_0 = vmlal_n_s16(acc_0, vget_low_s16(filter_0), input); + acc_1 = vmlal_n_s16(acc_1, vget_high_s16(filter_0), input); + acc_2 = vmlal_n_s16(acc_2, vget_low_s16(filter_1), input); + acc_3 = vmlal_n_s16(acc_3, vget_high_s16(filter_1), input); + acc_4 = vmlal_n_s16(acc_4, vget_high_s16(filter_x), input); + // Store the accumulators back to acc_buffer + vst1q_s32(acc_buffer_ptr + 4 * 0, acc_0); + vst1q_s32(acc_buffer_ptr + 4 * 1, acc_1); + vst1q_s32(acc_buffer_ptr + 4 * 2, acc_2); + vst1q_s32(acc_buffer_ptr + 4 * 3, acc_3); + vst1q_s32(acc_buffer_ptr + 4 * 4, acc_4); + acc_buffer_ptr += 20; + } + } +}; + template <> struct QuantizedDepthwiseConvKernel { static void Run(int num_output_pixels, int input_depth, int depth_multiplier, @@ -1504,7 +1553,7 @@ inline void QuantizedDepthwiseConvAccumRowGeneric( << "*\n" << "* If you would like to carry on with the slow code, compile\n" << "* with this preprocessor token defined:\n" - << "* TFLITE_ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK.\n" + << "* ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK.\n" << "*\n" << "* The right thing to do, if you care about performance, is to add\n" << "* a new DepthwiseConv kernel to tfmini to cover your case.\n" @@ -1691,6 +1740,7 @@ inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims, TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 2) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 16, 1) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 16) + TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 20) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 32) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 8) TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 1) @@ -1913,4 +1963,4 @@ void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims, } // namespace optimized_ops } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/eigen_spatial_convolutions.h b/tensorflow/contrib/lite/kernels/internal/optimized/eigen_spatial_convolutions.h index 8004c24a9914e216974539930853d0aadf61e324..ce3cde76999c77e1f9bf1eaccdba7e84ed508dda 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/eigen_spatial_convolutions.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/eigen_spatial_convolutions.h @@ -16,8 +16,8 @@ limitations under the License. // Copied from tensorflow/core/kernels/eigen_spatial_convolutions.h. // TODO(petewarden) - move this to a common location in Eigen itself. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_ #define EIGEN_USE_CUSTOM_THREAD_POOL #define EIGEN_USE_THREADS @@ -39,7 +39,6 @@ limitations under the License. #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #endif - namespace Eigen { /** SpatialConvolution @@ -215,17 +214,16 @@ EIGEN_DEVICE_FUNC } // TODO(yangke): choose() is defined in TensorContraction.h -- consider // moving it to somewhere more "common". - return - input - .extract_image_patches(kernelRows, kernelCols, row_stride, col_stride, - row_in_stride, col_in_stride, padding_type) - .reshape(pre_contract_dims) - .contract(kernel.reshape(kernel_dims), contract_dims) - .reshape(post_contract_dims); + return input + .extract_image_patches(kernelRows, kernelCols, row_stride, col_stride, + row_in_stride, col_in_stride, padding_type) + .reshape(pre_contract_dims) + .contract(kernel.reshape(kernel_dims), contract_dims) + .reshape(post_contract_dims); } } // end namespace Eigen // clang-format on -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_google.h b/tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_google.h index 7f78f69360b1ebbfb08600c8bc427f1ba9d5244d..d85e06a5d5af8d23235a08592d49754e4f493d34 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_google.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_google.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_GOOGLE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_GOOGLE_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_GOOGLE_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_GOOGLE_H_ #define EIGEN_USE_CUSTOM_THREAD_POOL #define EIGEN_USE_THREADS @@ -140,4 +140,4 @@ limitations under the License. #include "third_party/eigen3/unsupported/Eigen/CXX11/src/Tensor/TensorIO.h" #include "Eigen/src/Core/util/ReenableStupidWarnings.h" -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_H +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_H diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_oss.h b/tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_oss.h index 1d5c316194df0b87ee7eecbdd04bd5ce9e2e40b5..d34708b8fd0c0732c13ddbd8d70c87a278c40ff8 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_oss.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/eigen_tensor_reduced_instantiations_oss.h @@ -19,8 +19,8 @@ limitations under the License. // clang-format off -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_OSS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_OSS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_OSS_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_OSS_H_ #include "Eigen/Core" @@ -164,4 +164,4 @@ typedef unsigned __int64 uint64_t; #include "Eigen/src/Core/util/ReenableStupidWarnings.h" -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_OSS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_TENSOR_REDUCED_INSTANTIATIONS_OSS_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h b/tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h index b3615f4658a1a70284cc9d386a868a87aa09819b..0bfb4e9b1f8ee4167cfb629645a38538be1d73d4 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MULTITHREAD_CONV -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MULTITHREAD_CONV +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MULTITHREAD_CONV +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MULTITHREAD_CONV #include #include @@ -192,4 +192,4 @@ inline void Conv(const float* input_data, const Dims<4>& input_dims, } // namespace multithreaded_ops } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MULTITHREAD_CONV +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_MULTITHREAD_CONV diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc index bf0bdfb1fb875c4b54c55e25d4a17541507ecd4c..780401e052733cccae0cc34f495df090c1530624 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc +++ b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc @@ -15,12 +15,13 @@ limitations under the License. #include #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" #ifdef USE_NEON -#include #define kFloatWeightsPerNeonLane 4 namespace tflite { diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h index 3a4af87304eaf33489b38bd9b15ad9789e091d24..b7e317dc60e2c68e9e993ff45c9090a01bd13b94 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_ // TODO(ghodrat): Remove this header file and the dependency to internal data // structure. @@ -110,4 +110,4 @@ void ReductionSumVector(const float* input_vector, float* output_vector, } // namespace tensor_utils } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h index 2df919e579efaaa283f191df91cd433374b31567..dec58fea4facf0fbd9af6ac0cc916bfe5778e1a1 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_OPS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_OPS_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_OPS_H_ #include #include @@ -1538,9 +1538,10 @@ void Add(const int32* input1_data, const Dims<4>& input1_dims, // reference_ops.h. Once an optimized version is implemented and NdArrayDesc // is no longer referenced in this file, move NdArrayDesc from types.h to // reference_ops.h. -template +template void BroadcastAdd(const T* input1_data, const Dims<4>& input1_dims, const T* input2_data, const Dims<4>& input2_dims, + T output_activation_min, T output_activation_max, T* output_data, const Dims<4>& output_dims) { gemmlowp::ScopedProfilingLabel label("BroadcastAdd"); @@ -1563,15 +1564,30 @@ void BroadcastAdd(const T* input1_data, const Dims<4>& input1_dims, for (int y = 0; y < ArraySize(output_dims, 2); ++y) { for (int x = 0; x < ArraySize(output_dims, 1); ++x) { for (int c = 0; c < ArraySize(output_dims, 0); ++c) { - output_data[Offset(output_dims, c, x, y, b)] = ActivationFunction( - input1_data[SubscriptToIndex(desc1, c, x, y, b)] + - input2_data[SubscriptToIndex(desc2, c, x, y, b)]); + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[SubscriptToIndex(desc1, c, x, y, b)] + + input2_data[SubscriptToIndex(desc2, c, x, y, b)], + output_activation_min, output_activation_max); } } } } } +// legacy, for compatibility with old checked-in code +template +void BroadcastAdd(const T* input1_data, const Dims<4>& input1_dims, + const T* input2_data, const Dims<4>& input2_dims, + T* output_data, const Dims<4>& output_dims) { + T output_activation_min, output_activation_max; + GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); + + BroadcastAdd(input1_data, input1_dims, input2_data, input2_dims, + output_activation_min, output_activation_max, output_data, + output_dims); +} + inline void BroadcastAdd(int left_shift, const uint8* input1_data, const Dims<4>& input1_dims, int32 input1_offset, int32 input1_multiplier, int input1_shift, @@ -1772,9 +1788,10 @@ void Mul(const int32* input1_data, const Dims<4>& input1_dims, // reference_ops.h. Once an optimized version is implemented and NdArrayDesc // is no longer referenced in this file, move NdArrayDesc from types.h to // reference_ops.h. -template +template void BroadcastMul(const T* input1_data, const Dims<4>& input1_dims, const T* input2_data, const Dims<4>& input2_dims, + T output_activation_min, T output_activation_max, T* output_data, const Dims<4>& output_dims) { gemmlowp::ScopedProfilingLabel label("BroadcastMul"); @@ -1797,15 +1814,30 @@ void BroadcastMul(const T* input1_data, const Dims<4>& input1_dims, for (int y = 0; y < ArraySize(output_dims, 2); ++y) { for (int x = 0; x < ArraySize(output_dims, 1); ++x) { for (int c = 0; c < ArraySize(output_dims, 0); ++c) { - output_data[Offset(output_dims, c, x, y, b)] = ActivationFunction( - input1_data[SubscriptToIndex(desc1, c, x, y, b)] * - input2_data[SubscriptToIndex(desc2, c, x, y, b)]); + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[SubscriptToIndex(desc1, c, x, y, b)] * + input2_data[SubscriptToIndex(desc2, c, x, y, b)], + output_activation_min, output_activation_max); } } } } } +// legacy, for compatibility with old checked-in code +template +void BroadcastMul(const T* input1_data, const Dims<4>& input1_dims, + const T* input2_data, const Dims<4>& input2_dims, + T* output_data, const Dims<4>& output_dims) { + T output_activation_min, output_activation_max; + GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); + + BroadcastMul(input1_data, input1_dims, input2_data, input2_dims, + output_activation_min, output_activation_max, output_data, + output_dims); +} + inline void BroadcastMul(const uint8* input1_data, const Dims<4>& input1_dims, int32 input1_offset, const uint8* input2_data, const Dims<4>& input2_dims, int32 input2_offset, @@ -1868,6 +1900,61 @@ inline void BroadcastMul(const uint8* input1_data, const Dims<4>& input1_dims, output_data, output_dims); } +// TODO(aselle): This is not actually optimized yet. +inline void Div(const float* input1_data, const Dims<4>& input1_dims, + const float* input2_data, const Dims<4>& input2_dims, + float output_activation_min, float output_activation_max, + float* output_data, const Dims<4>& output_dims) { + const int batches = + MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); + const int height = + MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); + const int width = + MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); + const int depth = + MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); + for (int b = 0; b < batches; ++b) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + for (int c = 0; c < depth; ++c) { + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[Offset(input1_dims, c, x, y, b)] / + input2_data[Offset(input2_dims, c, x, y, b)], + output_activation_min, output_activation_max); + } + } + } + } +} + +// TODO(aselle): This is not actually optimized yet. +inline void Sub(const float* input1_data, const Dims<4>& input1_dims, + const float* input2_data, const Dims<4>& input2_dims, + float output_activation_min, float output_activation_max, + float* output_data, const Dims<4>& output_dims) { + const int batches = + MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); + const int height = + MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); + const int width = + MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); + const int depth = + MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); + for (int b = 0; b < batches; ++b) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + for (int c = 0; c < depth; ++c) { + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[Offset(input1_dims, c, x, y, b)] - + input2_data[Offset(input2_dims, c, x, y, b)], + output_activation_min, output_activation_max); + } + } + } + } +} template void Concatenation(int concat_dim, const Scalar* const* input_data, const Dims<4>* const* input_dims, int inputs_count, @@ -1994,6 +2081,166 @@ inline void LstmCell(const float* input_data, const Dims<4>& input_dims, output_state_map.tanh(); } +// Quantized LSTM cell. Currently just a copy of the reference impl in +// reference_ops.h. See the big function comment there, not replicating it +// here. +template +void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims, + const uint8* prev_activ_data_uint8, + const Dims<4>& prev_activ_dims, const uint8* weights_data_uint8, + const Dims<4>& weights_dims, const int32* bias_data_int32, + const Dims<4>& bias_dims, const int16* prev_state_data_int16, + const Dims<4>& prev_state_dims, int16* output_state_data_int16, + const Dims<4>& output_state_dims, uint8* output_activ_data_uint8, + const Dims<4>& output_activ_dims, uint8* concat_temp_data_uint8, + const Dims<4>& concat_temp_dims, int16* activ_temp_data_int16, + const Dims<4>& activ_temp_dims, int32 weights_zero_point, + int32 accum_multiplier, int accum_shift) { + gemmlowp::ScopedProfilingLabel label( + "LstmCell/quantized (8bit external, 16bit internal)"); + // Gather dimensions information, and perform consistency checks. + const int batches = + MatchingArraySize(input_dims, 3, prev_activ_dims, 3, prev_state_dims, 3, + output_state_dims, 3, output_activ_dims, 3); + const int height = + MatchingArraySize(input_dims, 2, prev_activ_dims, 2, prev_state_dims, 2, + output_state_dims, 2, output_activ_dims, 2); + const int width = + MatchingArraySize(input_dims, 1, prev_activ_dims, 1, prev_state_dims, 1, + output_state_dims, 1, output_activ_dims, 1); + TFLITE_CHECK_EQ(ArraySize(weights_dims, 2), 1); + TFLITE_CHECK_EQ(ArraySize(weights_dims, 3), 1); + const int input_depth = ArraySize(input_dims, 0); + const int prev_activ_depth = ArraySize(prev_activ_dims, 0); + const int total_input_depth = prev_activ_depth + input_depth; + TFLITE_CHECK_EQ(ArraySize(weights_dims, 0), total_input_depth); + TFLITE_CHECK_EQ(MatchingArraySize(bias_dims, 1, bias_dims, 2, bias_dims, 3), + 1); + const int intern_activ_depth = + MatchingArraySize(weights_dims, 1, bias_dims, 0); + TFLITE_CHECK_EQ(intern_activ_depth % 4, 0); + const int output_depth = + MatchingArraySize(prev_state_dims, 0, prev_activ_dims, 0, + output_state_dims, 0, output_activ_dims, 0); + TFLITE_CHECK_EQ(output_depth, intern_activ_depth / 4); + const int fc_batches = ArraySize(activ_temp_dims, 1) * + ArraySize(activ_temp_dims, 2) * + ArraySize(activ_temp_dims, 3); + const int fc_output_depth = + MatchingArraySize(weights_dims, 1, activ_temp_dims, 0); + const int fc_accum_depth = ArraySize(weights_dims, 0); + TFLITE_CHECK_EQ(fc_output_depth, 4 * output_depth); + + // Depth-concatenate prev_activ and input data together. + uint8 const* concat_input_arrays_data[2] = {input_data_uint8, + prev_activ_data_uint8}; + Dims<4> const* concat_input_arrays_dims[2] = {&input_dims, &prev_activ_dims}; + Concatenation( + 0, concat_input_arrays_data, concat_input_arrays_dims, 2, + concat_temp_data_uint8, concat_temp_dims); + + // Implementation of the fully connected node inside the LSTM cell. + // The operands are 8-bit integers, the accumulators are internally 32bit + // integers, and the output is 16-bit fixed-point with 3 integer bits so + // the output range is [-2^3, 2^3] == [-8, 8]. The rationale for that + // is explained in the function comment above. + for (int b = 0; b < fc_batches; ++b) { + for (int out_c = 0; out_c < fc_output_depth; ++out_c) { + // Internal accumulation. + // Initialize accumulator with the bias-value. + int32 accum = bias_data_int32[out_c]; + // Accumulation loop. + for (int d = 0; d < fc_accum_depth; ++d) { + int16 input_val = concat_temp_data_uint8[b * fc_accum_depth + d] - 128; + int16 weights_val = + weights_data_uint8[out_c * fc_accum_depth + d] - weights_zero_point; + accum += input_val * weights_val; + } + // Down-scale the final int32 accumulator to the scale used by our + // (16-bit, using 3 integer bits) fixed-point format. The quantized + // multiplier and shift here have been pre-computed offline + // (e.g. by toco). + // Note that the implicit assumption here, that this multiplier is smaller + // than one, is equivalent to the assumption that the fully-connected + // weights min-max is enclosed within [-4, 4] (it may be narrower). + // If that eventually fails, offline tools (e.g. toco) will fail early + // and that will be easy to support as needed. For now, assuming that + // this multiplier is less than one allows us to use a simpler, more + // accurate implementation. + accum = + MultiplyByQuantizedMultiplier(accum, accum_multiplier, accum_shift); + // Saturate, cast to int16, and store to the temporary activations array. + accum = std::max(-32768, std::min(32767, accum)); + activ_temp_data_int16[out_c + fc_output_depth * b] = accum; + } + } + + // Rest of the LSTM cell: tanh and logistic math functions, and some adds + // and muls, all done in 16-bit fixed-point. + const int outer_size = batches * width * height; + for (int b = 0; b < outer_size; ++b) { + for (int c = 0; c < output_depth; ++c) { + // Define the fixed-point data types that we will use here. All use + // int16 as the underlying integer type i.e. all are 16-bit fixed-point. + // They only differ by the number of integral vs. fractional bits, + // determining the range of values that they can represent. + // + // F0 uses 0 integer bits, range [-1, 1]. + // This is the return type of math functions such as tanh, logistic, + // whose range is in [-1, 1]. + using F0 = gemmlowp::FixedPoint; + // F3 uses 3 integer bits, range [-8, 8]. + // This is the range of the previous fully-connected node's output, + // which is our input here. + using F3 = gemmlowp::FixedPoint; + // FS uses StateIntegerBits integer bits, range [-2^StateIntegerBits, + // 2^StateIntegerBits]. It's used to represent the internal state, whose + // number of integer bits is currently dictated by the model. See comment + // on the StateIntegerBits template parameter above. + using FS = gemmlowp::FixedPoint; + // Implementation of input gate, using fixed-point logistic function. + F3 input_gate_input = F3::FromRaw( + activ_temp_data_int16[b * fc_output_depth + 0 * output_depth + c]); + F0 input_gate_output = gemmlowp::logistic(input_gate_input); + // Implementation of input modulation gate, using fixed-point tanh + // function. + F3 input_modulation_gate_input = F3::FromRaw( + activ_temp_data_int16[b * fc_output_depth + 1 * output_depth + c]); + F0 input_modulation_gate_output = + gemmlowp::tanh(input_modulation_gate_input); + // Implementation of forget gate, using fixed-point logistic function. + F3 forget_gate_input = F3::FromRaw( + activ_temp_data_int16[b * fc_output_depth + 2 * output_depth + c]); + F0 forget_gate_output = gemmlowp::logistic(forget_gate_input); + // Implementation of output gate, using fixed-point logistic function. + F3 output_gate_input = F3::FromRaw( + activ_temp_data_int16[b * fc_output_depth + 3 * output_depth + c]); + F0 output_gate_output = gemmlowp::logistic(output_gate_input); + // Implementation of internal multiplication nodes, still in fixed-point. + F0 input_times_input_modulation = + input_gate_output * input_modulation_gate_output; + FS prev_state = FS::FromRaw(prev_state_data_int16[b * output_depth + c]); + FS prev_state_times_forget_state = forget_gate_output * prev_state; + // Implementation of internal addition node, saturating. + FS new_state = gemmlowp::SaturatingAdd( + gemmlowp::Rescale(input_times_input_modulation), + prev_state_times_forget_state); + // Implementation of last internal tanh node, still in fixed-point. + F0 output_activ_int16 = output_gate_output * gemmlowp::tanh(new_state); + // Store the new internal state back to memory, as 16-bit integers. + output_state_data_int16[b * output_depth + c] = new_state.raw(); + // Down-scale the output activations to 8-bit integers, saturating, + // and store back to memory. + int16 rescaled_output_activ = + gemmlowp::RoundingDivideByPOT(output_activ_int16.raw(), 8); + int16 clamped_output_activ = + std::max(-128, std::min(127, rescaled_output_activ)); + output_activ_data_uint8[b * output_depth + c] = + 128 + clamped_output_activ; + } + } +} + template void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims, int outputs_count, Scalar* const* output_data, @@ -2619,74 +2866,194 @@ inline void Softmax(const uint8* input_data, const Dims<4>& input_dims, using FixedPointAccum = gemmlowp::FixedPoint; using FixedPoint0 = gemmlowp::FixedPoint; - gemmlowp::ScopedProfilingLabel label("Softmax"); + gemmlowp::ScopedProfilingLabel label("Softmax/8bit"); const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); const int height = MatchingArraySize(input_dims, 2, output_dims, 2); const int width = MatchingArraySize(input_dims, 1, output_dims, 1); const int depth = MatchingArraySize(input_dims, 0, output_dims, 0); - for (int b = 0; b < batches; ++b) { - for (int x = 0; x < width; ++x) { - for (int y = 0; y < height; ++y) { - uint8 max_in_row = 0; - for (int c = 0; c < depth; ++c) { - max_in_row = - std::max(max_in_row, input_data[Offset(input_dims, c, x, y, b)]); - } + const int outer_size = batches * height * width; - FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); - for (int c = 0; c < depth; ++c) { - int32 input_diff = - static_cast(input_data[Offset(input_dims, c, x, y, b)]) - - max_in_row; - if (input_diff >= diff_min) { - const int32 input_diff_rescaled = - MultiplyByQuantizedMultiplierGreaterThanOne( - input_diff, input_beta_multiplier, input_beta_left_shift); - const FixedPointScaledDiff scaled_diff_f8 = - FixedPointScaledDiff::FromRaw(input_diff_rescaled); - sum_of_exps = - sum_of_exps + gemmlowp::Rescale( - exp_on_negative_values(scaled_diff_f8)); - } + for (int b = 0; b < outer_size; ++b) { + const uint8* input_data_ptr = input_data + b * depth; + uint8* output_data_ptr = output_data + b * depth; + + // Determine the largest entry in the current row + uint8 max_in_row = 0; + { + int c = 0; +#ifdef USE_NEON + uint8x16_t max16_0 = vdupq_n_u8(0); + uint8x16_t max16_1 = vdupq_n_u8(0); + for (; c <= depth - 32; c += 32) { + max16_0 = vmaxq_u8(max16_0, vld1q_u8(input_data_ptr + c + 0)); + max16_1 = vmaxq_u8(max16_1, vld1q_u8(input_data_ptr + c + 16)); + } + uint8x16_t max16 = vmaxq_u8(max16_0, max16_1); + if (c <= depth - 16) { + max16 = vmaxq_u8(max16, vld1q_u8(input_data_ptr + c)); + c += 16; + } + uint8x8_t max8 = vmax_u8(vget_low_u8(max16), vget_high_u8(max16)); + if (c <= depth - 8) { + max8 = vmax_u8(max8, vld1_u8(input_data_ptr + c)); + c += 8; + } + uint8x8_t max4 = vmax_u8(max8, vext_u8(max8, max8, 4)); + uint8x8_t max2 = vmax_u8(max4, vext_u8(max4, max4, 2)); + uint8x8_t max1 = vpmax_u8(max2, max2); + max_in_row = vget_lane_u8(max1, 0); +#endif + for (; c < depth; ++c) { + max_in_row = std::max(max_in_row, input_data_ptr[c]); + } + } + +#ifdef USE_NEON + using FixedPointAccumInt32x4 = + gemmlowp::FixedPoint; + using FixedPointScaledDiffInt32x4 = + gemmlowp::FixedPoint; + using FixedPoint0Int32x4 = gemmlowp::FixedPoint; + FixedPoint0Int32x4 input_beta_multiplier_f0 = + FixedPoint0Int32x4::FromScalarRaw(input_beta_multiplier); + int16x8_t max_in_row_s16 = vdupq_n_s16(max_in_row); +#endif + + // Compute the sum of exponentials of the differences of entries in the + // current row from the largest entry in the current row. + FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); + { + int c = 0; +#ifdef USE_NEON + int32x4_t diff_min_s32 = vdupq_n_s32(diff_min); + FixedPointAccumInt32x4 sum_of_exps_0 = FixedPointAccumInt32x4::Zero(); + FixedPointAccumInt32x4 sum_of_exps_1 = FixedPointAccumInt32x4::Zero(); + FixedPointAccumInt32x4 zeros = FixedPointAccumInt32x4::Zero(); + for (; c <= depth - 8; c += 8) { + uint16x8_t input_u16 = vmovl_u8(vld1_u8(input_data_ptr + c)); + int16x8_t input_diff_s16 = + vsubq_s16(vreinterpretq_s16_u16(input_u16), max_in_row_s16); + int32x4_t input_diff_s32_0 = vmovl_s16(vget_low_s16(input_diff_s16)); + int32x4_t input_diff_s32_1 = vmovl_s16(vget_high_s16(input_diff_s16)); + int32x4_t mask_0 = + gemmlowp::MaskIfGreaterThanOrEqual(input_diff_s32_0, diff_min_s32); + int32x4_t mask_1 = + gemmlowp::MaskIfGreaterThanOrEqual(input_diff_s32_1, diff_min_s32); + FixedPointScaledDiffInt32x4 scaled_diff_0 = + input_beta_multiplier_f0 * + FixedPointScaledDiffInt32x4::FromRaw( + gemmlowp::ShiftLeft(input_diff_s32_0, input_beta_left_shift)); + FixedPointScaledDiffInt32x4 scaled_diff_1 = + input_beta_multiplier_f0 * + FixedPointScaledDiffInt32x4::FromRaw( + gemmlowp::ShiftLeft(input_diff_s32_1, input_beta_left_shift)); + FixedPointAccumInt32x4 exps_0 = + gemmlowp::Rescale( + exp_on_negative_values(scaled_diff_0)); + FixedPointAccumInt32x4 exps_1 = + gemmlowp::Rescale( + exp_on_negative_values(scaled_diff_1)); + FixedPointAccumInt32x4 masked_exps_0 = + SelectUsingMask(mask_0, exps_0, zeros); + FixedPointAccumInt32x4 masked_exps_1 = + SelectUsingMask(mask_1, exps_1, zeros); + sum_of_exps_0 = sum_of_exps_0 + masked_exps_0; + sum_of_exps_1 = sum_of_exps_1 + masked_exps_1; + } + int32x4_t sum_of_exps_reduced_4 = (sum_of_exps_0 + sum_of_exps_1).raw(); + int32x2_t sum_of_exps_reduced_2 = + vadd_s32(vget_low_s32(sum_of_exps_reduced_4), + vget_high_s32(sum_of_exps_reduced_4)); + int32x2_t sum_of_exps_reduced_1 = + vpadd_s32(sum_of_exps_reduced_2, sum_of_exps_reduced_2); + sum_of_exps = + FixedPointAccum::FromRaw(vget_lane_s32(sum_of_exps_reduced_1, 0)); +#endif + for (; c < depth; ++c) { + int32 input_diff = static_cast(input_data_ptr[c]) - max_in_row; + if (input_diff >= diff_min) { + const int32 input_diff_rescaled = + MultiplyByQuantizedMultiplierGreaterThanOne( + input_diff, input_beta_multiplier, input_beta_left_shift); + const FixedPointScaledDiff scaled_diff_f8 = + FixedPointScaledDiff::FromRaw(input_diff_rescaled); + sum_of_exps = + sum_of_exps + gemmlowp::Rescale( + exp_on_negative_values(scaled_diff_f8)); } + } + } + + // Compute the fixed-point multiplier and shift that we need to apply to + // perform a division by the above-computed sum-of-exponentials. + int32 fixed_sum_of_exps = sum_of_exps.raw(); + int headroom_plus_one = + __builtin_clz(static_cast(fixed_sum_of_exps)); + // This is the number of bits to the left of the binary point above 1.0. + // Consider fixed_sum_of_exps=1.25. In that case shifted_scale=0.8 and + // no later adjustment will be needed. + int num_bits_over_unit = kAccumulationIntegerBits - headroom_plus_one; + int32 shifted_sum_minus_one = static_cast( + (static_cast(fixed_sum_of_exps) << headroom_plus_one) - + (static_cast(1) << 31)); + FixedPoint0 shifted_scale = gemmlowp::one_over_one_plus_x_for_x_in_0_1( + FixedPoint0::FromRaw(shifted_sum_minus_one)); + + // Compute the quotients of exponentials of differences of entries in the + // current row from the largest entry, over the previously-computed sum of + // exponentials. + { + int c = 0; +#ifdef USE_NEON + int16x8_t diff_min_s16 = vdupq_n_s16(diff_min); + for (; c <= depth - 8; c += 8) { + uint16x8_t input_u16 = vmovl_u8(vld1_u8(input_data_ptr + c)); + int16x8_t input_diff_s16 = + vsubq_s16(vreinterpretq_s16_u16(input_u16), max_in_row_s16); + int32x4_t input_diff_s32_0 = vmovl_s16(vget_low_s16(input_diff_s16)); + int32x4_t input_diff_s32_1 = vmovl_s16(vget_high_s16(input_diff_s16)); + uint8x8_t mask = vmovn_u16(vcgeq_s16(input_diff_s16, diff_min_s16)); + FixedPointScaledDiffInt32x4 scaled_diff_0 = + input_beta_multiplier_f0 * + FixedPointScaledDiffInt32x4::FromRaw( + gemmlowp::ShiftLeft(input_diff_s32_0, input_beta_left_shift)); + FixedPointScaledDiffInt32x4 scaled_diff_1 = + input_beta_multiplier_f0 * + FixedPointScaledDiffInt32x4::FromRaw( + gemmlowp::ShiftLeft(input_diff_s32_1, input_beta_left_shift)); + FixedPoint0Int32x4 exp_0 = exp_on_negative_values(scaled_diff_0); + FixedPoint0Int32x4 exp_1 = exp_on_negative_values(scaled_diff_1); + int32x4_t output_s32_0 = gemmlowp::RoundingDivideByPOT( + vqrdmulhq_n_s32(exp_0.raw(), shifted_scale.raw()), + num_bits_over_unit + 31 - 8); + int32x4_t output_s32_1 = gemmlowp::RoundingDivideByPOT( + vqrdmulhq_n_s32(exp_1.raw(), shifted_scale.raw()), + num_bits_over_unit + 31 - 8); + int16x8_t output_s16 = + vcombine_s16(vqmovn_s32(output_s32_0), vqmovn_s32(output_s32_1)); + uint8x8_t output_u8 = vqmovun_s16(output_s16); + uint8x8_t masked_output = vbsl_u8(mask, output_u8, vdup_n_u8(0)); + vst1_u8(output_data_ptr + c, masked_output); + } +#endif + for (; c < depth; ++c) { + int32 input_diff = static_cast(input_data_ptr[c]) - max_in_row; + if (input_diff >= diff_min) { + const int32 input_diff_rescaled = + MultiplyByQuantizedMultiplierGreaterThanOne( + input_diff, input_beta_multiplier, input_beta_left_shift); + const FixedPointScaledDiff scaled_diff_f8 = + FixedPointScaledDiff::FromRaw(input_diff_rescaled); - int32 fixed_sum_of_exps = sum_of_exps.raw(); - // TODO(starka): Use a NEON intrinsic like vclzq_u32 instead. - int headroom_plus_one = - __builtin_clz(static_cast(fixed_sum_of_exps)); - // This is the number of bits to the left of the binary point above 1.0. - // Consider fixed_sum_of_exps=1.25. In that case shifted_scale=0.8 and - // no later adjustment will be needed. - int num_bits_over_unit = kAccumulationIntegerBits - headroom_plus_one; - int32 shifted_sum_minus_one = static_cast( - (static_cast(fixed_sum_of_exps) << headroom_plus_one) - - (static_cast(1) << 31)); - - FixedPoint0 shifted_scale = gemmlowp::one_over_one_plus_x_for_x_in_0_1( - FixedPoint0::FromRaw(shifted_sum_minus_one)); + FixedPoint0 exp_in_0 = exp_on_negative_values(scaled_diff_f8); + int32 unsat_output = gemmlowp::RoundingDivideByPOT( + (shifted_scale * exp_in_0).raw(), num_bits_over_unit + 31 - 8); - for (int c = 0; c < depth; ++c) { - int32 input_diff = - static_cast(input_data[Offset(input_dims, c, x, y, b)]) - - max_in_row; - if (input_diff >= diff_min) { - const int32 input_diff_rescaled = - MultiplyByQuantizedMultiplierGreaterThanOne( - input_diff, input_beta_multiplier, input_beta_left_shift); - const FixedPointScaledDiff scaled_diff_f8 = - FixedPointScaledDiff::FromRaw(input_diff_rescaled); - - FixedPoint0 exp_in_0 = exp_on_negative_values(scaled_diff_f8); - int32 unsat_output = gemmlowp::RoundingDivideByPOT( - (shifted_scale * exp_in_0).raw(), num_bits_over_unit + 31 - 8); - - output_data[Offset(output_dims, c, x, y, b)] = - std::max(std::min(unsat_output, 255), 0); - - } else { - output_data[Offset(output_dims, c, x, y, b)] = 0; - } + output_data_ptr[c] = std::max(std::min(unsat_output, 255), 0); + + } else { + output_data_ptr[c] = 0; } } } @@ -2851,6 +3218,156 @@ inline void Tanh(const float* input_data, const Dims<4>& input_dims, output_map.array() = input_map.array().tanh(); } +inline void Tanh(const uint8* input_data, const Dims<4>& input_dims, + int32 input_zero_point, int32 input_range_radius, + int32 input_multiplier, int input_left_shift, + uint8* output_data, const Dims<4>& output_dims) { + // Note that this is almost the exact same code as in Logistic(). + gemmlowp::ScopedProfilingLabel label("Tanh"); + /* batches */ MatchingArraySize(input_dims, 3, output_dims, 3); + /* height */ MatchingArraySize(input_dims, 2, output_dims, 2); + /* width */ MatchingArraySize(input_dims, 1, output_dims, 1); + /* depth */ MatchingArraySize(input_dims, 0, output_dims, 0); + const int size = RequiredBufferSizeForDims(input_dims); + + int c = 0; + int32_t output_zero_point = 128; +#ifdef USE_NEON + // Handle 16 values at a time + for (; c <= size - 16; c += 16) { + // Read input uint8 values, cast to int16 and subtract input_zero_point + uint8x16_t input_val_u8 = vld1q_u8(input_data + c); + int16x8_t input_val_centered_0 = + vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(input_val_u8))), + vdupq_n_s16(input_zero_point)); + int16x8_t input_val_centered_1 = + vsubq_s16(vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(input_val_u8))), + vdupq_n_s16(input_zero_point)); + + // Prepare the bit masks that we will use at the end to implement the logic + // that was expressed in the scalar code with branching: + // if (input_val_centered < -input_range_radius) { + // output_val = 0; + // } else if (input_val_centered > input_range_radius) { + // output_val = 255; + // } else { + // ... + uint16x8_t mask_rightclamp_0 = + vcgtq_s16(input_val_centered_0, vdupq_n_s16(input_range_radius)); + uint16x8_t mask_rightclamp_1 = + vcgtq_s16(input_val_centered_1, vdupq_n_s16(input_range_radius)); + uint16x8_t mask_leftclamp_0 = + vcgeq_s16(input_val_centered_0, vdupq_n_s16(-input_range_radius)); + uint16x8_t mask_leftclamp_1 = + vcgeq_s16(input_val_centered_1, vdupq_n_s16(-input_range_radius)); + uint8x16_t mask_rightclamp = vcombine_u8(vshrn_n_u16(mask_rightclamp_0, 8), + vshrn_n_u16(mask_rightclamp_1, 8)); + uint8x16_t mask_leftclamp = vcombine_u8(vshrn_n_u16(mask_leftclamp_0, 8), + vshrn_n_u16(mask_leftclamp_1, 8)); + + // This performs what is expressed in the scalar code as + // const int32 input_val_rescaled = + // MultiplyByQuantizedMultiplierGreaterThanOne( + // input_val_centered, input_multiplier, input_left_shift); + int32x4_t input_val_rescaled_0 = + vshlq_s32(vmovl_s16(vget_low_s16(input_val_centered_0)), + vdupq_n_s32(input_left_shift)); + int32x4_t input_val_rescaled_1 = + vshlq_s32(vmovl_s16(vget_high_s16(input_val_centered_0)), + vdupq_n_s32(input_left_shift)); + int32x4_t input_val_rescaled_2 = + vshlq_s32(vmovl_s16(vget_low_s16(input_val_centered_1)), + vdupq_n_s32(input_left_shift)); + int32x4_t input_val_rescaled_3 = + vshlq_s32(vmovl_s16(vget_high_s16(input_val_centered_1)), + vdupq_n_s32(input_left_shift)); + input_val_rescaled_0 = + vqrdmulhq_n_s32(input_val_rescaled_0, input_multiplier); + input_val_rescaled_1 = + vqrdmulhq_n_s32(input_val_rescaled_1, input_multiplier); + input_val_rescaled_2 = + vqrdmulhq_n_s32(input_val_rescaled_2, input_multiplier); + input_val_rescaled_3 = + vqrdmulhq_n_s32(input_val_rescaled_3, input_multiplier); + + // Invoke gemmlowp::tanh on FixedPoint wrapping int32x4_t + using FixedPoint4 = gemmlowp::FixedPoint; + using FixedPoint0 = gemmlowp::FixedPoint; + const FixedPoint4 input_val_f4_0 = + FixedPoint4::FromRaw(input_val_rescaled_0); + const FixedPoint4 input_val_f4_1 = + FixedPoint4::FromRaw(input_val_rescaled_1); + const FixedPoint4 input_val_f4_2 = + FixedPoint4::FromRaw(input_val_rescaled_2); + const FixedPoint4 input_val_f4_3 = + FixedPoint4::FromRaw(input_val_rescaled_3); + const FixedPoint0 output_val_f0_0 = gemmlowp::tanh(input_val_f4_0); + const FixedPoint0 output_val_f0_1 = gemmlowp::tanh(input_val_f4_1); + const FixedPoint0 output_val_f0_2 = gemmlowp::tanh(input_val_f4_2); + const FixedPoint0 output_val_f0_3 = gemmlowp::tanh(input_val_f4_3); + + // Divide by 2^24 as in the scalar code + using gemmlowp::RoundingDivideByPOT; + int32x4_t output_val_s32_0 = RoundingDivideByPOT(output_val_f0_0.raw(), 24); + int32x4_t output_val_s32_1 = RoundingDivideByPOT(output_val_f0_1.raw(), 24); + int32x4_t output_val_s32_2 = RoundingDivideByPOT(output_val_f0_2.raw(), 24); + int32x4_t output_val_s32_3 = RoundingDivideByPOT(output_val_f0_3.raw(), 24); + + // Add the output zero point + int32x4_t output_zero_point_s32 = vdupq_n_s32(output_zero_point); + output_val_s32_0 = vaddq_s32(output_val_s32_0, output_zero_point_s32); + output_val_s32_1 = vaddq_s32(output_val_s32_1, output_zero_point_s32); + output_val_s32_2 = vaddq_s32(output_val_s32_2, output_zero_point_s32); + output_val_s32_3 = vaddq_s32(output_val_s32_3, output_zero_point_s32); + + // Cast output values to uint8, saturating + int16x8_t output_val_s16_0 = vcombine_s16(vqmovn_s32(output_val_s32_0), + vqmovn_s32(output_val_s32_1)); + int16x8_t output_val_s16_1 = vcombine_s16(vqmovn_s32(output_val_s32_2), + vqmovn_s32(output_val_s32_3)); + uint8x16_t output_val_u8 = vcombine_u8(vqmovun_s16(output_val_s16_0), + vqmovun_s16(output_val_s16_1)); + + // Perform the bit-masking with the bit masks computed at the beginning, + // see the comment there. + output_val_u8 = vorrq_u8(output_val_u8, mask_rightclamp); + output_val_u8 = vandq_u8(output_val_u8, mask_leftclamp); + + // Store back to memory + vst1q_u8(output_data + c, output_val_u8); + } +#endif + // Leftover loop: handle one value at a time with scalar code. + for (; c < size; ++c) { + const uint8 input_val_u8 = input_data[c]; + const int32 input_val_centered = + static_cast(input_val_u8) - input_zero_point; + uint8 output_val; + if (input_val_centered < -input_range_radius) { + output_val = 0; + } else if (input_val_centered > input_range_radius) { + output_val = 255; + } else { + const int32 input_val_rescaled = + MultiplyByQuantizedMultiplierGreaterThanOne( + input_val_centered, input_multiplier, input_left_shift); + using FixedPoint4 = gemmlowp::FixedPoint; + using FixedPoint0 = gemmlowp::FixedPoint; + const FixedPoint4 input_val_f4 = FixedPoint4::FromRaw(input_val_rescaled); + const FixedPoint0 output_val_f0 = gemmlowp::tanh(input_val_f4); + using gemmlowp::RoundingDivideByPOT; + int32 output_val_s32 = RoundingDivideByPOT(output_val_f0.raw(), 24); + output_val_s32 += output_zero_point; + if (output_val_s32 == 256) { + output_val_s32 = 255; + } + TFLITE_DCHECK_GE(output_val_s32, 0); + TFLITE_DCHECK_LE(output_val_s32, 255); + output_val = static_cast(output_val_s32); + } + output_data[c] = output_val; + } +} inline void Dequantize(const uint8* input_data, const Dims<4>& input_dims, int32 zero_point, double scale, float* output_data, const Dims<4>& output_dims) { @@ -3323,7 +3840,7 @@ inline void ResizeBilinearGeneric(const float* input_data, inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims, const int32* output_size_data, const Dims<4>& output_size_dims, float* output_data, - const Dims<4>& output_dims) { + const Dims<4>& output_dims, bool align_corners) { gemmlowp::ScopedProfilingLabel label("ResizeBilinear"); int32 batches = MatchingArraySize(input_dims, 3, output_dims, 3); int32 input_height = ArraySize(input_dims, 2); @@ -3338,13 +3855,20 @@ inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims, int32 output_width = output_size_data[Offset(output_size_dims, 1, 0, 0, 0)]; // Specialize for 2x2 upsample. - if (output_height == 2 * input_height && output_width == 2 * input_width) { + if (!align_corners && output_height == 2 * input_height && + output_width == 2 * input_width) { ResizeBilinear2x2(input_data, input_dims, output_data, output_dims, batches, input_height, input_width, depth, output_height, output_width); } else { float height_scale = static_cast(input_height) / output_height; float width_scale = static_cast(input_width) / output_width; + if (align_corners && output_height > 1) { + height_scale = static_cast(input_height - 1) / (output_height - 1); + } + if (align_corners && output_width > 1) { + width_scale = static_cast(input_width - 1) / (output_width - 1); + } ResizeBilinearGeneric(input_data, input_dims, output_data, output_dims, batches, input_height, input_width, depth, @@ -3353,6 +3877,15 @@ inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims, } } +// legacy, for compatibility with old checked-in code +inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims, + const int32* output_size_data, + const Dims<4>& output_size_dims, float* output_data, + const Dims<4>& output_dims) { + ResizeBilinear(input_data, input_dims, output_size_data, output_size_dims, + output_data, output_dims, /*align_corners=*/false); +} + template inline void SpaceToBatchND(const T* input_data, const Dims<4>& input_dims, const int32* block_shape_data, @@ -3381,10 +3914,11 @@ inline void SpaceToBatchND(const T* input_data, const Dims<4>& input_dims, for (int out_h = 0; out_h < output_height; ++out_h) { for (int out_w = 0; out_w < output_width; ++out_w) { T* out = output_data + Offset(output_dims, 0, out_w, out_h, out_b); - if (out_h * block_shape_height < padding_top || - out_h * block_shape_height >= padding_top + input_height || - out_w * block_shape_width < padding_left || - out_w * block_shape_width >= padding_left + input_width) { + if (out_h * block_shape_height + shift_h < padding_top || + out_h * block_shape_height + shift_h >= + padding_top + input_height || + out_w * block_shape_width + shift_w < padding_left || + out_w * block_shape_width + shift_w >= padding_left + input_width) { memset(out, 0, depth * sizeof(T)); } else { const T* in = @@ -3749,4 +4283,4 @@ void ArgMax(const T3* axis, const T1* input_data, const Dims<4>& input_dims, #pragma GCC diagnostic pop #endif -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_OPS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_OPS_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h b/tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h index f8be99e82fb8721ced7a3e5da686b20ce241ea2d..4e324a5e107cf5a90c0042331899edab831c8e51 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h @@ -15,7 +15,7 @@ limitations under the License. #ifndef TF_LITE_KERNELS_INTERNAL_OPTIMIZED_TENSOR_UTILS_IMPL_H_ #define TF_LITE_KERNELS_INTERNAL_OPTIMIZED_TENSOR_UTILS_IMPL_H_ -// TDOD(ghodrat): Remove this header file and the dependency to internal data +// TODO(ghodrat): Remove this header file and the dependency to internal data // structure. #include "tensorflow/contrib/lite/builtin_op_data.h" diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util.cc b/tensorflow/contrib/lite/kernels/internal/quantization_util.cc index 98f2e365c5249a6c28673fc185ebec34cc2105b2..18be6777a5caeb45a4ffabd8b7f1793de7b053f8 100644 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util.cc +++ b/tensorflow/contrib/lite/kernels/internal/quantization_util.cc @@ -22,27 +22,20 @@ limitations under the License. namespace tflite { -void QuantizeMultiplierSmallerThanOne(double double_multiplier, - int32_t* quantized_multiplier, - int* right_shift) { - TFLITE_CHECK(double_multiplier >= 0.); - TFLITE_CHECK(double_multiplier < 1.); +void QuantizeMultiplier(double double_multiplier, int32_t* quantized_multiplier, + int* shift) { if (double_multiplier == 0.) { *quantized_multiplier = 0; - *right_shift = 0; + *shift = 0; return; } - TFLITE_CHECK(double_multiplier > 0.); - const double q = std::frexp(double_multiplier, right_shift); - *right_shift *= -1; - + const double q = std::frexp(double_multiplier, shift); auto q_fixed = static_cast(TfLiteRound(q * (1ll << 31))); TFLITE_CHECK(q_fixed <= (1ll << 31)); if (q_fixed == (1ll << 31)) { q_fixed /= 2; - --*right_shift; + ++*shift; } - TFLITE_CHECK_GE(*right_shift, 0); TFLITE_CHECK_LE(q_fixed, std::numeric_limits::max()); *quantized_multiplier = static_cast(q_fixed); } @@ -50,17 +43,20 @@ void QuantizeMultiplierSmallerThanOne(double double_multiplier, void QuantizeMultiplierGreaterThanOne(double double_multiplier, int32_t* quantized_multiplier, int* left_shift) { - TFLITE_CHECK(double_multiplier > 1.); - const double q = std::frexp(double_multiplier, left_shift); - auto q_fixed = static_cast(TfLiteRound(q * (1ll << 31))); - TFLITE_CHECK(q_fixed <= (1ll << 31)); - if (q_fixed == (1ll << 31)) { - q_fixed /= 2; - ++*left_shift; - } + TFLITE_CHECK_GT(double_multiplier, 1.); + QuantizeMultiplier(double_multiplier, quantized_multiplier, left_shift); TFLITE_CHECK_GE(*left_shift, 0); - TFLITE_CHECK_LE(q_fixed, std::numeric_limits::max()); - *quantized_multiplier = static_cast(q_fixed); +} + +void QuantizeMultiplierSmallerThanOne(double double_multiplier, + int32_t* quantized_multiplier, + int* right_shift) { + TFLITE_CHECK_LT(double_multiplier, 1.); + TFLITE_CHECK_GT(double_multiplier, 0.); + int shift; + QuantizeMultiplier(double_multiplier, quantized_multiplier, &shift); + TFLITE_CHECK_LE(shift, 0); + *right_shift = -shift; } void PreprocessSoftmaxScaling(double beta, double input_scale, diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util.h b/tensorflow/contrib/lite/kernels/internal/quantization_util.h index efb7191c8deb2a23ea5473ab131d2b6537202765..ba06bc0975b6847b24592daa60efe99983d03707 100644 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util.h +++ b/tensorflow/contrib/lite/kernels/internal/quantization_util.h @@ -20,7 +20,8 @@ limitations under the License. namespace tflite { // Decompose a double multiplier into a Q0.31 int32 representation of its -// significand, and shift representation of its exponent. +// significand, and shift representation of NEGATIVE its exponent --- +// this is intended as a RIGHT-shift. // // Restricted to the case where the multiplier < 1 (and non-negative). void QuantizeMultiplierSmallerThanOne(double double_multiplier, @@ -35,6 +36,16 @@ void QuantizeMultiplierGreaterThanOne(double double_multiplier, int32_t* quantized_multiplier, int* left_shift); +// Decompose a double multiplier into a Q0.31 int32 representation of its +// significand, and shift representation of its exponent. +// +// Handles an arbitrary positive multiplier. The 'shift' output-value is +// basically the 'floating-point exponent' of the multiplier: +// Negative for a right-shift (when the multiplier is <1), positive for a +// left-shift (when the multiplier is >1) +void QuantizeMultiplier(double double_multiplier, int32_t* quantized_multiplier, + int* shift); + // This first creates a multiplier in a double equivalent of // Q(input_integer_bits).(31-input_integer_bits) representation, with extra // precision in the double's fractional bits. It then splits the result into diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc b/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc index d6f306e2cbae3c780b3d773638ba46cd2abf02f5..19b1b408ec74b0939065b0ad10b91ecfc2cd4765 100644 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc +++ b/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc @@ -31,7 +31,7 @@ TEST(QuantizationUtilTest, QuantizeMultiplierSmallerThanOne) { }; EXPECT_DEATH(quantize(-0.1), ""); - EXPECT_THAT(quantize(0.0), Pair(0, 0)); + EXPECT_DEATH(quantize(0.0), ""); EXPECT_THAT(quantize(0.25), Pair(1073741824, 1)); // Around 0.5 we can see the change in exponent and how we try hard to diff --git a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h index 8e0f234545e43dd8b2412e065aaecad8325a1182..9aabee5000c29ed97fcf7e874d661e72fd768f84 100644 --- a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h +++ b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ #include "tensorflow/contrib/lite/kernels/internal/common.h" #include "tensorflow/contrib/lite/kernels/internal/compatibility.h" @@ -112,4 +112,4 @@ void DepthwiseConv(const float* input_data, const Dims<4>& input_dims, } // end namespace reference_ops } // end namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h index 8a80558b32f2858778460956cd9f57617674e21e..e9b6baeaee87d22aef238410bc9f447509a81c47 100644 --- a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h +++ b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ #include @@ -135,4 +135,4 @@ void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims, } // end namespace reference_ops } // end namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h b/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h index 7f90d731b8454a020ab273e6b5591ed90aab14c7..c05c21b472b05f2cbe133adf94d91ab0c6d9ef40 100644 --- a/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h +++ b/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h @@ -12,10 +12,10 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_PORTABLE_TENSOR_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_PORTABLE_TENSOR_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_PORTABLE_TENSOR_UTILS_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_PORTABLE_TENSOR_UTILS_H_ -// TDOD(ghodrat): Remove this header file and the dependency to internal data +// TODO(ghodrat): Remove this header file and the dependency to internal data // structure. #include "tensorflow/contrib/lite/builtin_op_data.h" @@ -186,4 +186,4 @@ void ReductionSumVector(const float* input_vector, float* output_vector, } // namespace tensor_utils } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_PORTABLE_TENSOR_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_PORTABLE_TENSOR_UTILS_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h index 14c430258740b65dce65816f7c5c41fccf6dd5cf..5f4d5be3235433787e3eed6d37a7f403348d9eee 100644 --- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h +++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_REFERENCE_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_REFERENCE_OPS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_REFERENCE_OPS_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_REFERENCE_OPS_H_ #include #include @@ -889,10 +889,11 @@ inline void Add(int left_shift, const uint8* input1_data, // dimensionality if the runtime code does a single loop over one dimension // that handles broadcasting as the base case. The code generator would then // generate max(D1, D2) nested for loops. -template -void BroadcastAdd(const float* input1_data, const Dims<4>& input1_dims, - const float* input2_data, const Dims<4>& input2_dims, - float* output_data, const Dims<4>& output_dims) { +template +void BroadcastAdd(const T* input1_data, const Dims<4>& input1_dims, + const T* input2_data, const Dims<4>& input2_dims, + T output_activation_min, T output_activation_max, + T* output_data, const Dims<4>& output_dims) { gemmlowp::ScopedProfilingLabel label("BroadcastAdd"); NdArrayDesc<4> desc1; @@ -914,15 +915,30 @@ void BroadcastAdd(const float* input1_data, const Dims<4>& input1_dims, for (int y = 0; y < ArraySize(output_dims, 2); ++y) { for (int x = 0; x < ArraySize(output_dims, 1); ++x) { for (int c = 0; c < ArraySize(output_dims, 0); ++c) { - output_data[Offset(output_dims, c, x, y, b)] = ActivationFunction( - input1_data[SubscriptToIndex(desc1, c, x, y, b)] + - input2_data[SubscriptToIndex(desc2, c, x, y, b)]); + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[SubscriptToIndex(desc1, c, x, y, b)] + + input2_data[SubscriptToIndex(desc2, c, x, y, b)], + output_activation_min, output_activation_max); } } } } } +// legacy, for compatibility with old checked-in code +template +void BroadcastAdd(const T* input1_data, const Dims<4>& input1_dims, + const T* input2_data, const Dims<4>& input2_dims, + T* output_data, const Dims<4>& output_dims) { + T output_activation_min, output_activation_max; + GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); + + BroadcastAdd(input1_data, input1_dims, input2_data, input2_dims, + output_activation_min, output_activation_max, output_data, + output_dims); +} + inline void BroadcastAdd(int left_shift, const uint8* input1_data, const Dims<4>& input1_dims, int32 input1_offset, int32 input1_multiplier, int input1_shift, @@ -1053,10 +1069,11 @@ void Mul(const float* input1_data, const Dims<4>& input1_dims, // dimensionality if the runtime code does a single loop over one dimension // that handles broadcasting as the base case. The code generator would then // generate max(D1, D2) nested for loops. -template -void BroadcastMul(const float* input1_data, const Dims<4>& input1_dims, - const float* input2_data, const Dims<4>& input2_dims, - float* output_data, const Dims<4>& output_dims) { +template +void BroadcastMul(const T* input1_data, const Dims<4>& input1_dims, + const T* input2_data, const Dims<4>& input2_dims, + T output_activation_min, T output_activation_max, + T* output_data, const Dims<4>& output_dims) { gemmlowp::ScopedProfilingLabel label("BroadcastMul"); NdArrayDesc<4> desc1; @@ -1078,15 +1095,30 @@ void BroadcastMul(const float* input1_data, const Dims<4>& input1_dims, for (int y = 0; y < ArraySize(output_dims, 2); ++y) { for (int x = 0; x < ArraySize(output_dims, 1); ++x) { for (int c = 0; c < ArraySize(output_dims, 0); ++c) { - output_data[Offset(output_dims, c, x, y, b)] = ActivationFunction( - input1_data[SubscriptToIndex(desc1, c, x, y, b)] * - input2_data[SubscriptToIndex(desc2, c, x, y, b)]); + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[SubscriptToIndex(desc1, c, x, y, b)] * + input2_data[SubscriptToIndex(desc2, c, x, y, b)], + output_activation_min, output_activation_max); } } } } } +// legacy, for compatibility with old checked-in code +template +void BroadcastMul(const T* input1_data, const Dims<4>& input1_dims, + const T* input2_data, const Dims<4>& input2_dims, + T* output_data, const Dims<4>& output_dims) { + T output_activation_min, output_activation_max; + GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); + + BroadcastMul(input1_data, input1_dims, input2_data, input2_dims, + output_activation_min, output_activation_max, output_data, + output_dims); +} + inline void BroadcastMul(const uint8* input1_data, const Dims<4>& input1_dims, int32 input1_offset, const uint8* input2_data, const Dims<4>& input2_dims, int32 input2_offset, @@ -1149,6 +1181,60 @@ inline void BroadcastMul(const uint8* input1_data, const Dims<4>& input1_dims, output_data, output_dims); } +inline void Div(const float* input1_data, const Dims<4>& input1_dims, + const float* input2_data, const Dims<4>& input2_dims, + float output_activation_min, float output_activation_max, + float* output_data, const Dims<4>& output_dims) { + const int batches = + MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); + const int height = + MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); + const int width = + MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); + const int depth = + MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); + for (int b = 0; b < batches; ++b) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + for (int c = 0; c < depth; ++c) { + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[Offset(input1_dims, c, x, y, b)] / + input2_data[Offset(input2_dims, c, x, y, b)], + output_activation_min, output_activation_max); + } + } + } + } +} + +inline void Sub(const float* input1_data, const Dims<4>& input1_dims, + const float* input2_data, const Dims<4>& input2_dims, + float output_activation_min, float output_activation_max, + float* output_data, const Dims<4>& output_dims) { + const int batches = + MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); + const int height = + MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); + const int width = + MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); + const int depth = + MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); + for (int b = 0; b < batches; ++b) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + for (int c = 0; c < depth; ++c) { + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[Offset(input1_dims, c, x, y, b)] - + input2_data[Offset(input2_dims, c, x, y, b)], + output_activation_min, output_activation_max); + } + } + } + } +} + template void Concatenation(int concat_dim, const Scalar* const* input_data, const Dims<4>* const* input_dims, int inputs_count, @@ -1272,6 +1358,265 @@ inline void LstmCell(const float* input_data, const Dims<4>& input_dims, } } +// Quantized LSTM cell implementation. +// The quantization of the input, output arrays is as follows: +// - The input activations are quantized as uint8 on the interval +// [-1, 127/128]. +// The rationale for that is that that is the natural interval for output +// activations (see next point) and these need to be concatenated together. +// We could accommodate different ranges by re-scaling, but we empirically +// found that setting the input activations range to be [-1, 127/128] in the +// first place, removing the need for re-scaling, greatly improves accuracy. +// - The output activations are quantized as uint8 on the interval +// [-1, 127/128]. +// The rationale for that is that the definition of a LSTM cell makes them +// intrinsically constrained in [-1, 1]; tweaking that to [-1, 127/128] +// makes for simpler, more accurate fixed-point arithmetic. +// - The output-at-previous-timestep state array is obviously quantized as +// the output activations. +// - The internal LSTM memory (not the output-at-previous-timestep, the other +// internal state array) is int16-quantized and may use any power-of-two, +// symmetric range i.e. [-2^N, 2^N * 32767/32768] for any N, which we call +// StateIntegerBits below, see the below discussion of that template +// parameter ("The StateIntegerBits template parameter"). +// - The output of the internal fully-connected node is int16-quantized +// on the interval [-8, 8 * 32767/32768], the rationale for which is +// explained just below ("Why [-8, 8] for fully-connected output?"). +// +// +// === The StateIntegerBits template parameter === +// +// The StateIntegerBits template parameter controls the fixed-point format used +// to represent the internal memory of the LSTM cell (not the +// output-at-previous-timestep, the other internal state array). It's currently +// a template parameter so that the model can control that. The most typical +// value for StateIntegerBits is 4. Other plausible values are anywhere between +// 3 and 5. We might eventually standardize on a single supported value, e.g. 4, +// and drop that template parameter. The reason why it can't be a runtime +// parameter is that this controls the fixed-point format used, i.e. we need to +// generate actually different code based on it. In particular, we generate code +// for a fixed-point tanh() implementation for that format, which internally +// uses a fixed-point exp() implementation, which internally uses a +// barrel-shifter with a number of steps that depends on StateIntegerBits. +// Another consequence of that is that a higher value of StateIntegerBits +// results in a more expensive implementation (more barrel shifter steps +// needed). +// +// +// === Why [-8, 8] for fully-connected output? === +// +// This array is only fed to Logistic and Tanh functions, for which +// the quantized implementation will want to use fixed-point arithmetic, +// requiring a power-of-two representation interval. Thus, we should right +// away quantize this array to a power-of-two interval; otherwise, +// implementation will need to rescale that, losing any benefit that a tighter +// representation interval might otherwise yield, while introducting some +// numerical error and computational overhead. +// +// Now, Logistic and Tanh +// are nearly constant (nearly equal to their horizontal asymptotes) +// outside of a small bounded interval around 0: +// +// Logistic(4) = 1 - 1.8e-2 Tanh(4) = 1 - 6.7e-4 +// Logistic(8) = 1 - 3.4e-4 Tanh(8) = 1 - 2.3e-7 +// Logistic(16) = 1 - 1.1e-7 Tanh(16) = 1 - 2.5e-14 +// +// From this, we see that clamping to [-4, 4] would be too inaccurate +// (the error of 1.8e-2 on Logistic would be felt even in 8bit precision) +// while clamping to [-16, 16] would make no difference even in float32. +// However, for a fixed-point implementation in 16-bit integers, using 5 +// integer bits to represent the [-16, 16] range would leave only 11 +// fractional bits, giving an increment of 2^-11 = 4.9e-4 between consecutive +// representable values. Notice that that is higher than the +// worst-case clamping error with clamping to [-8, 8]: 3.4e-4 for Logistic. +// Using [-8, 8] thus seems like the better compromise overall, enjoying +// an increment of 2.4e-4 between representable values and a worst-case +// clamping error of 3.4e-4, both better than the increment of 4.9e-4 with +// [-16, 16]. +// +// Moreover, all other things being equal, it is nice to choose the narrower +// representation range, as that makes the implementation of fixed-point +// math functions a little cheaper (each integer bit requires an additional +// barrel-shifter atep in the implementation of exp(-x)). That is further +// reason to prefer [-8, 8] over [-16, 16]. The choice of [-16, 16] would make +// sense for 32-bit float or 32-bit fixed-point quantization, but we are +// aiming for 16-bit fixed-point quantization of these internal nodes here. +// +template +void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims, + const uint8* prev_activ_data_uint8, + const Dims<4>& prev_activ_dims, const uint8* weights_data_uint8, + const Dims<4>& weights_dims, const int32* bias_data_int32, + const Dims<4>& bias_dims, const int16* prev_state_data_int16, + const Dims<4>& prev_state_dims, int16* output_state_data_int16, + const Dims<4>& output_state_dims, uint8* output_activ_data_uint8, + const Dims<4>& output_activ_dims, uint8* concat_temp_data_uint8, + const Dims<4>& concat_temp_dims, int16* activ_temp_data_int16, + const Dims<4>& activ_temp_dims, int32 weights_zero_point, + int32 accum_multiplier, int accum_shift) { + // Gather dimensions information, and perform consistency checks. + const int batches = + MatchingArraySize(input_dims, 3, prev_activ_dims, 3, prev_state_dims, 3, + output_state_dims, 3, output_activ_dims, 3); + const int height = + MatchingArraySize(input_dims, 2, prev_activ_dims, 2, prev_state_dims, 2, + output_state_dims, 2, output_activ_dims, 2); + const int width = + MatchingArraySize(input_dims, 1, prev_activ_dims, 1, prev_state_dims, 1, + output_state_dims, 1, output_activ_dims, 1); + TFLITE_CHECK_EQ(ArraySize(weights_dims, 2), 1); + TFLITE_CHECK_EQ(ArraySize(weights_dims, 3), 1); + const int input_depth = ArraySize(input_dims, 0); + const int prev_activ_depth = ArraySize(prev_activ_dims, 0); + const int total_input_depth = prev_activ_depth + input_depth; + TFLITE_CHECK_EQ(ArraySize(weights_dims, 0), total_input_depth); + TFLITE_CHECK_EQ(MatchingArraySize(bias_dims, 1, bias_dims, 2, bias_dims, 3), + 1); + const int intern_activ_depth = + MatchingArraySize(weights_dims, 1, bias_dims, 0); + TFLITE_CHECK_EQ(intern_activ_depth % 4, 0); + const int output_depth = + MatchingArraySize(prev_state_dims, 0, prev_activ_dims, 0, + output_state_dims, 0, output_activ_dims, 0); + TFLITE_CHECK_EQ(output_depth, intern_activ_depth / 4); + const int fc_batches = ArraySize(activ_temp_dims, 1) * + ArraySize(activ_temp_dims, 2) * + ArraySize(activ_temp_dims, 3); + const int fc_output_depth = + MatchingArraySize(weights_dims, 1, activ_temp_dims, 0); + const int fc_accum_depth = ArraySize(weights_dims, 0); + TFLITE_CHECK_EQ(fc_output_depth, 4 * output_depth); + + // Depth-concatenate prev_activ and input data together. + uint8 const* concat_input_arrays_data[2] = {input_data_uint8, + prev_activ_data_uint8}; + Dims<4> const* concat_input_arrays_dims[2] = {&input_dims, &prev_activ_dims}; + Concatenation( + 0, concat_input_arrays_data, concat_input_arrays_dims, 2, + concat_temp_data_uint8, concat_temp_dims); + + // Implementation of the fully connected node inside the LSTM cell. + // The operands are 8-bit integers, the accumulators are internally 32bit + // integers, and the output is 16-bit fixed-point with 3 integer bits so + // the output range is [-2^3, 2^3] == [-8, 8]. The rationale for that + // is explained in the function comment above. + for (int b = 0; b < fc_batches; ++b) { + for (int out_c = 0; out_c < fc_output_depth; ++out_c) { + // Internal accumulation. + // Initialize accumulator with the bias-value. + int32 accum = bias_data_int32[out_c]; + // Accumulation loop. + for (int d = 0; d < fc_accum_depth; ++d) { + int16 input_val = concat_temp_data_uint8[b * fc_accum_depth + d] - 128; + int16 weights_val = + weights_data_uint8[out_c * fc_accum_depth + d] - weights_zero_point; + accum += input_val * weights_val; + } + // Down-scale the final int32 accumulator to the scale used by our + // (16-bit, using 3 integer bits) fixed-point format. The quantized + // multiplier and shift here have been pre-computed offline + // (e.g. by toco). + accum = + MultiplyByQuantizedMultiplier(accum, accum_multiplier, accum_shift); + // Saturate, cast to int16, and store to the temporary activations array. + accum = std::max(-32768, std::min(32767, accum)); + activ_temp_data_int16[out_c + fc_output_depth * b] = accum; + } + } + + // Rest of the LSTM cell: tanh and logistic math functions, and some adds + // and muls, all done in 16-bit fixed-point. + const int outer_size = batches * width * height; + for (int b = 0; b < outer_size; ++b) { + for (int c = 0; c < output_depth; ++c) { + // Define the fixed-point data types that we will use here. All use + // int16 as the underlying integer type i.e. all are 16-bit fixed-point. + // They only differ by the number of integral vs. fractional bits, + // determining the range of values that they can represent. + // + // F0 uses 0 integer bits, range [-1, 1]. + // This is the return type of math functions such as tanh, logistic, + // whose range is in [-1, 1]. + using F0 = gemmlowp::FixedPoint; + // F3 uses 3 integer bits, range [-8, 8]. + // This is the range of the previous fully-connected node's output, + // which is our input here. + using F3 = gemmlowp::FixedPoint; + // FS uses StateIntegerBits integer bits, range [-2^StateIntegerBits, + // 2^StateIntegerBits]. It's used to represent the internal state, whose + // number of integer bits is currently dictated by the model. See comment + // on the StateIntegerBits template parameter above. + using FS = gemmlowp::FixedPoint; + // Implementation of input gate, using fixed-point logistic function. + F3 input_gate_input = F3::FromRaw( + activ_temp_data_int16[b * fc_output_depth + 0 * output_depth + c]); + F0 input_gate_output = gemmlowp::logistic(input_gate_input); + // Implementation of input modulation gate, using fixed-point tanh + // function. + F3 input_modulation_gate_input = F3::FromRaw( + activ_temp_data_int16[b * fc_output_depth + 1 * output_depth + c]); + F0 input_modulation_gate_output = + gemmlowp::tanh(input_modulation_gate_input); + // Implementation of forget gate, using fixed-point logistic function. + F3 forget_gate_input = F3::FromRaw( + activ_temp_data_int16[b * fc_output_depth + 2 * output_depth + c]); + F0 forget_gate_output = gemmlowp::logistic(forget_gate_input); + // Implementation of output gate, using fixed-point logistic function. + F3 output_gate_input = F3::FromRaw( + activ_temp_data_int16[b * fc_output_depth + 3 * output_depth + c]); + F0 output_gate_output = gemmlowp::logistic(output_gate_input); + // Implementation of internal multiplication nodes, still in fixed-point. + F0 input_times_input_modulation = + input_gate_output * input_modulation_gate_output; + FS prev_state = FS::FromRaw(prev_state_data_int16[b * output_depth + c]); + FS prev_state_times_forget_state = forget_gate_output * prev_state; + // Implementation of internal addition node, saturating. + FS new_state = gemmlowp::SaturatingAdd( + gemmlowp::Rescale(input_times_input_modulation), + prev_state_times_forget_state); + // Implementation of last internal tanh node, still in fixed-point. + F0 output_activ_int16 = output_gate_output * gemmlowp::tanh(new_state); + // Store the new internal state back to memory, as 16-bit integers. + output_state_data_int16[b * output_depth + c] = new_state.raw(); + // Down-scale the output activations to 8-bit integers, saturating, + // and store back to memory. + int16 rescaled_output_activ = + gemmlowp::RoundingDivideByPOT(output_activ_int16.raw(), 8); + int16 clamped_output_activ = + std::max(-128, std::min(127, rescaled_output_activ)); + output_activ_data_uint8[b * output_depth + c] = + 128 + clamped_output_activ; + } + } +} + +template +void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims, + int axis, int outputs_count, Scalar* const* output_data, + const Dims<4>* const* output_dims) { + const int batches = ArraySize(*output_dims[0], 3); + const int height = ArraySize(*output_dims[0], 2); + const int width = ArraySize(*output_dims[0], 1); + const int depth = ArraySize(*output_dims[0], 0); + + const int slice_size = ArraySize(*output_dims[0], axis); + + for (int i = 0; i < outputs_count; ++i) { + int offset = i * slice_size * input_dims.strides[axis]; + for (int b = 0; b < batches; ++b) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + for (int c = 0; c < depth; ++c) { + auto out = Offset(*output_dims[i], c, x, y, b); + auto in = Offset(input_dims, c, x, y, b); + output_data[i][out] = input_data[offset + in]; + } + } + } + } + } +} + template void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims, int outputs_count, Scalar* const* output_data, @@ -1282,28 +1627,12 @@ void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims, /* height = */ MatchingArraySize(*output_dims[i], 2, input_dims, 2); /* width = */ MatchingArraySize(*output_dims[i], 1, input_dims, 1); } - const int batches = MatchingArraySize(*output_dims[0], 3, input_dims, 3); - const int height = MatchingArraySize(*output_dims[0], 2, input_dims, 2); - const int width = MatchingArraySize(*output_dims[0], 1, input_dims, 1); // for now we dont have a model with a TensorFlowSplit // with fused activation function. TFLITE_DCHECK(Ac == FusedActivationFunctionType::kNone); - for (int b = 0; b < batches; ++b) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - int in_c = 0; - for (int i = 0; i < outputs_count; ++i) { - const int depth = ArraySize(*output_dims[i], 0); - for (int c = 0; c < depth; ++c) { - output_data[i][Offset(*output_dims[i], c, x, y, b)] = - input_data[Offset(input_dims, in_c, x, y, b)]; - in_c++; - } - } - TFLITE_DCHECK(in_c == ArraySize(input_dims, 0)); - } - } - } + + TensorFlowSplit(input_data, input_dims, /*axis=*/0, outputs_count, + output_data, output_dims); } // TODO(benoitjacob) make this a proper reference impl without Eigen! @@ -1957,6 +2286,54 @@ inline void Tanh(const float* input_data, const Dims<4>& input_dims, } } +inline void Tanh(const uint8* input_data, const Dims<4>& input_dims, + int32 input_zero_point, int32 input_range_radius, + int32 input_multiplier, int input_left_shift, + uint8* output_data, const Dims<4>& output_dims) { + const int32 output_zero_point = 128; + const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); + const int height = MatchingArraySize(input_dims, 2, output_dims, 2); + const int width = MatchingArraySize(input_dims, 1, output_dims, 1); + const int depth = MatchingArraySize(input_dims, 0, output_dims, 0); + for (int b = 0; b < batches; ++b) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + for (int c = 0; c < depth; ++c) { + const uint8 input_val_u8 = input_data[Offset(input_dims, c, x, y, b)]; + const int32 input_val_centered = + static_cast(input_val_u8) - input_zero_point; + uint8 output_val; + if (input_val_centered <= -input_range_radius) { + output_val = 0; + } else if (input_val_centered >= input_range_radius) { + output_val = 255; + } else { + const int32 input_val_rescaled = + MultiplyByQuantizedMultiplierGreaterThanOne( + input_val_centered, input_multiplier, input_left_shift); + using FixedPoint4 = gemmlowp::FixedPoint; + using FixedPoint0 = gemmlowp::FixedPoint; + const FixedPoint4 input_val_f4 = + FixedPoint4::FromRaw(input_val_rescaled); + const FixedPoint0 output_val_f0 = gemmlowp::tanh(input_val_f4); + + using gemmlowp::RoundingDivideByPOT; + int32 output_val_s32 = RoundingDivideByPOT(output_val_f0.raw(), 24); + output_val_s32 += output_zero_point; + if (output_val_s32 == 256) { + output_val_s32 = 255; + } + TFLITE_DCHECK_GE(output_val_s32, 0); + TFLITE_DCHECK_LE(output_val_s32, 255); + output_val = static_cast(output_val_s32); + } + output_data[Offset(output_dims, c, x, y, b)] = output_val; + } + } + } + } +} + inline void Dequantize(const uint8* input_data, const Dims<4>& input_dims, int32 zero_point, double scale, float* output_data, const Dims<4>& output_dims) { @@ -2116,7 +2493,7 @@ inline void Gather(const T* input_data, const Dims<4>& input_dims, inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims, const int32* output_size_data, const Dims<4>& output_size_dims, float* output_data, - const Dims<4>& output_dims) { + const Dims<4>& output_dims, bool align_corners) { int32 batches = MatchingArraySize(input_dims, 3, output_dims, 3); int32 input_height = ArraySize(input_dims, 2); int32 input_width = ArraySize(input_dims, 1); @@ -2130,6 +2507,12 @@ inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims, int32 output_width = output_size_data[Offset(output_size_dims, 1, 0, 0, 0)]; float height_scale = static_cast(input_height) / output_height; float width_scale = static_cast(input_width) / output_width; + if (align_corners && output_height > 1) { + height_scale = static_cast(input_height - 1) / (output_height - 1); + } + if (align_corners && output_width > 1) { + width_scale = static_cast(input_width - 1) / (output_width - 1); + } for (int b = 0; b < batches; ++b) { for (int y = 0; y < output_height; ++y) { @@ -2157,6 +2540,15 @@ inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims, } } +// legacy, for compatibility with old checked-in code +inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims, + const int32* output_size_data, + const Dims<4>& output_size_dims, float* output_data, + const Dims<4>& output_dims) { + ResizeBilinear(input_data, input_dims, output_size_data, output_size_dims, + output_data, output_dims, /*align_corners=*/false); +} + template inline void SpaceToBatchND(const T* input_data, const Dims<4>& input_dims, const int32* block_shape_data, @@ -2183,10 +2575,11 @@ inline void SpaceToBatchND(const T* input_data, const Dims<4>& input_dims, for (int out_h = 0; out_h < output_height; ++out_h) { for (int out_w = 0; out_w < output_width; ++out_w) { T* out = output_data + Offset(output_dims, 0, out_w, out_h, out_b); - if (out_h * block_shape_height < padding_top || - out_h * block_shape_height >= padding_top + input_height || - out_w * block_shape_width < padding_left || - out_w * block_shape_width >= padding_left + input_width) { + if (out_h * block_shape_height + shift_h < padding_top || + out_h * block_shape_height + shift_h >= + padding_top + input_height || + out_w * block_shape_width + shift_w < padding_left || + out_w * block_shape_width + shift_w >= padding_left + input_width) { memset(out, 0, depth * sizeof(T)); } else { const T* in = @@ -2275,27 +2668,60 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims, } } +inline bool LoopCondition(int index, int stop, int stride) { + return stride > 0 ? index < stop : index > stop; +} + +inline int StartIndex(int start, int stride, int dim, bool masked) { + return masked ? (stride > 0 ? 0 : dim - 1) : start; +} + +inline int StopIndex(int start, int stop, int stride, int dim, bool masked, + bool shrink_axis_masked) { + return shrink_axis_masked ? stride > 0 ? start + 1 : start - 1 + : masked ? (stride > 0 ? dim : -1) : stop; +} + template inline void StridedSlice(const T* input_data, const Dims<4>& input_dims, - int begin_mask, int end_mask, + int begin_mask, int end_mask, int shrink_axis_mask, const std::vector& starts, const std::vector& stops, const std::vector& strides, T* output_data, const Dims<4>& output_dims) { - const int start_b = (begin_mask & 8) ? 0 : starts[3]; - const int stop_b = (end_mask & 8) ? input_dims.sizes[3] : stops[3]; - const int start_h = (begin_mask & 4) ? 0 : starts[2]; - const int stop_h = (end_mask & 4) ? input_dims.sizes[2] : stops[2]; - const int start_w = (begin_mask & 2) ? 0 : starts[1]; - const int stop_w = (end_mask & 2) ? input_dims.sizes[1] : stops[1]; - const int start_d = (begin_mask & 1) ? 0 : starts[0]; - const int stop_d = (end_mask & 1) ? input_dims.sizes[0] : stops[0]; + TFLITE_DCHECK_EQ(starts.size(), 4); + TFLITE_DCHECK_EQ(stops.size(), 4); + TFLITE_DCHECK_EQ(strides.size(), 4); + const int start_b = + StartIndex(starts[3], strides[3], input_dims.sizes[3], begin_mask & 8); + const int stop_b = + StopIndex(start_b, stops[3], strides[3], input_dims.sizes[3], + end_mask & 8, shrink_axis_mask & 8); + const int start_h = + StartIndex(starts[2], strides[2], input_dims.sizes[2], begin_mask & 4); + const int stop_h = + StopIndex(start_h, stops[2], strides[2], input_dims.sizes[2], + end_mask & 4, shrink_axis_mask & 4); + const int start_w = + StartIndex(starts[1], strides[1], input_dims.sizes[1], begin_mask & 2); + const int stop_w = + StopIndex(start_w, stops[1], strides[1], input_dims.sizes[1], + end_mask & 2, shrink_axis_mask & 2); + const int start_d = + StartIndex(starts[0], strides[0], input_dims.sizes[0], begin_mask & 1); + const int stop_d = + StopIndex(start_d, stops[0], strides[0], input_dims.sizes[0], + end_mask & 1, shrink_axis_mask & 1); T* out_ptr = output_data; - for (int in_b = start_b; in_b < stop_b; in_b += strides[3]) { - for (int in_h = start_h; in_h < stop_h; in_h += strides[2]) { - for (int in_w = start_w; in_w < stop_w; in_w += strides[1]) { - for (int in_d = start_d; in_d < stop_d; in_d += strides[0]) { + for (int in_b = start_b; LoopCondition(in_b, stop_b, strides[3]); + in_b += strides[3]) { + for (int in_h = start_h; LoopCondition(in_h, stop_h, strides[2]); + in_h += strides[2]) { + for (int in_w = start_w; LoopCondition(in_w, stop_w, strides[1]); + in_w += strides[1]) { + for (int in_d = start_d; LoopCondition(in_d, stop_d, strides[0]); + in_d += strides[0]) { *out_ptr++ = input_data[Offset(input_dims, in_d, in_w, in_h, in_b)]; } } @@ -2303,6 +2729,18 @@ inline void StridedSlice(const T* input_data, const Dims<4>& input_dims, } } +template +inline void StridedSlice(const T* input_data, const Dims<4>& input_dims, + int begin_mask, int end_mask, + const std::vector& starts, + const std::vector& stops, + const std::vector& strides, T* output_data, + const Dims<4>& output_dims) { + StridedSlice(input_data, input_dims, begin_mask, end_mask, + /*shrink_axis_mask=*/0, starts, stops, strides, output_data, + output_dims); +} + template inline void Slice(const T* input_data, const Dims<4>& input_dims, const std::vector& begin, const std::vector& size, @@ -2335,6 +2773,72 @@ inline void Slice(const T* input_data, const Dims<4>& input_dims, } } +template +inline void Exp(const T* input_data, const size_t num_elements, + T* output_data) { + for (size_t idx = 0; idx < num_elements; ++idx) { + output_data[idx] = exp(input_data[idx]); + } +} + +template +inline void Mean(T* input_data, const int* input_dims, const int input_num_dims, + T* output_data, const int* output_dims, + const int output_num_dims, const int* axis, + const int num_axis_dimensions, bool keep_dims, int* temp_index, + int* resolved_axis) { + // resets output data. + size_t num_outputs = 1; + for (int idx = 0; idx < output_num_dims; ++idx) { + num_outputs *= static_cast(output_dims[idx]); + } + for (size_t idx = 0; idx < num_outputs; ++idx) { + output_data[idx] = 0; + } + // resets temp index. + for (int idx = 0; idx < input_num_dims; ++idx) { + temp_index[idx] = 0; + } + // resolves axis. + int num_resolved_axis = 0; + for (int idx = 0; idx < num_axis_dimensions; ++idx) { + int current = axis[idx]; + TFLITE_DCHECK(current < input_num_dims && current + input_num_dims >= 0); + if (current < 0) { + current += input_num_dims; + } + bool is_dup = false; + for (int j = 0; j < num_resolved_axis; ++j) { + if (resolved_axis[j] == current) { + is_dup = true; + break; + } + } + if (!is_dup) { + resolved_axis[num_resolved_axis++] = current; + } + } + // iterates through input_data. + for (bool has_next = true; has_next; + has_next = NextIndex(input_num_dims, input_dims, temp_index)) { + size_t input_offset = + ReducedOutputOffset(input_num_dims, input_dims, temp_index, 0, nullptr); + size_t output_offset = + ReducedOutputOffset(input_num_dims, input_dims, temp_index, + num_resolved_axis, resolved_axis); + output_data[output_offset] += input_data[input_offset]; + } + // takes average by num of elements added to get mean. + size_t num_elements_in_axis = 1; + for (int idx = 0; idx < num_resolved_axis; ++idx) { + num_elements_in_axis *= static_cast(input_dims[resolved_axis[idx]]); + } + for (size_t idx = 0; idx < num_outputs; ++idx) { + output_data[idx] = static_cast(static_cast(output_data[idx]) / + num_elements_in_axis); + } +} + template inline void Mean(const T* input_data, const Dims<4>& input_dims, const std::vector& reduction_indices, T* output_data, @@ -2483,7 +2987,36 @@ void ArgMax(const T3* axis, const T1* input_data, const Dims<4>& input_dims, } } +template +void Transpose(const T* input, const Dims<4>& input_dims, T* output, + const Dims<4>& output_dims, int* permuted_axes) { + int out_sizes[4]; + // Compute the inverse permutation array so we can do an output centered + // transpose. Also, check to make sure output_dims is matching input_dims. + for (int k = 0; k < 4; k++) { + out_sizes[k] = + MatchingArraySize(input_dims, permuted_axes[k], output_dims, k); + } + + // Naive transpose loop (iterate on output index and compute input index). + int o[4]; // loop index (on output). + int i[4]; + for (o[3] = 0; o[3] < out_sizes[3]; o[3]++) { + i[permuted_axes[3]] = o[3]; + for (o[2] = 0; o[2] < out_sizes[2]; o[2]++) { + i[permuted_axes[2]] = o[2]; + for (o[1] = 0; o[1] < out_sizes[1]; o[1]++) { + i[permuted_axes[1]] = o[1]; + for (o[0] = 0; o[0] < out_sizes[0]; o[0]++) { + i[permuted_axes[0]] = o[0]; + output[Offset(output_dims, o)] = input[Offset(input_dims, i)]; + } + } + } + } +} + } // namespace reference_ops } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_REFERENCE_OPS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_REFERENCE_OPS_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/round.h b/tensorflow/contrib/lite/kernels/internal/round.h index 38525b0e208b852343849096ac68cbfc9ef3e389..f299d0bd8733dc603c4950091c8ac3d7890548a7 100644 --- a/tensorflow/contrib/lite/kernels/internal/round.h +++ b/tensorflow/contrib/lite/kernels/internal/round.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_ROUND_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_ROUND_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_ROUND_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_ROUND_H_ #include @@ -36,4 +36,4 @@ inline T TfLiteRound(const T x) { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_ROUND_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_ROUND_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/tensor.h b/tensorflow/contrib/lite/kernels/internal/tensor.h index ee4111e0416560d94d513c528971bdf3bf819662..62e38e0d4c3e023d0ed2242fc9438b096b86dc59 100644 --- a/tensorflow/contrib/lite/kernels/internal/tensor.h +++ b/tensorflow/contrib/lite/kernels/internal/tensor.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_H_ #include #include "tensorflow/contrib/lite/context.h" @@ -41,8 +41,7 @@ inline int32_t* GetTensorData(TfLiteTensor* tensor) { template <> inline int64_t* GetTensorData(TfLiteTensor* tensor) { - return tensor != nullptr ? reinterpret_cast(tensor->data.raw) - : nullptr; + return tensor != nullptr ? tensor->data.i64 : nullptr; } inline int RemapDim(int max_dimensions, int d) { @@ -82,6 +81,51 @@ inline Dims<4> GetTensorDims(const TfLiteTensor* tensor) { return GetTensorDims(dims->data, dims->size); } +// A list of tensors in a format that can be used by kernels like split and +// concatenation. +template +class VectorOfTensors { + public: + // Build with the tensors in 'tensor_list'. + VectorOfTensors(const TfLiteContext& context, + const TfLiteIntArray& tensor_list) { + int num_tensors = tensor_list.size; + + all_data_.reserve(num_tensors); + all_dims_.reserve(num_tensors); + all_dims_ptr_.reserve(num_tensors); + + for (int i = 0; i < num_tensors; ++i) { + TfLiteTensor* t = &context.tensors[tensor_list.data[i]]; + all_data_.push_back(GetTensorData(t)); + all_dims_.push_back(GetTensorDims(t)); + } + + // Taking the pointer from inside a std::vector is only OK if the vector is + // never modified, so we populate all_dims in the previous loop and then we + // are free to grab iterators here. + for (int i = 0; i < num_tensors; ++i) { + all_dims_ptr_.push_back(&all_dims_[i]); + } + } + // Return a pointer to the data pointers of all tensors in the list. For + // example: + // float* const* f = v.data(); + // f[0][1] is the second element of the first tensor. + T* const* data() const { return all_data_.data(); } + + // Return a pointer the dim pointers of all tensors in the list. For + // example: + // const Dims<4>* const* d = v.dims(); + // dims[1] are the dimensions of the second tensor in the list. + const Dims<4>* const* dims() const { return all_dims_ptr_.data(); } + + private: + std::vector all_data_; + std::vector> all_dims_; + std::vector*> all_dims_ptr_; +}; + } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/tensor_utils.cc b/tensorflow/contrib/lite/kernels/internal/tensor_utils.cc index 904a97803a6a9ba369c1e64c711b12d19ffc10c4..f4181b18a8f46fd9bef4b81a210a6b8134a4e9d0 100644 --- a/tensorflow/contrib/lite/kernels/internal/tensor_utils.cc +++ b/tensorflow/contrib/lite/kernels/internal/tensor_utils.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" +#include "tensorflow/contrib/lite/kernels/internal/common.h" #ifndef USE_NEON #if defined(__ARM_NEON__) || defined(__ARM_NEON) diff --git a/tensorflow/contrib/lite/kernels/internal/tensor_utils.h b/tensorflow/contrib/lite/kernels/internal/tensor_utils.h index e7e2994397650004c7ba442fa1803290e6b12302..40d144979b2f965725db86ff311e90f39438802f 100644 --- a/tensorflow/contrib/lite/kernels/internal/tensor_utils.h +++ b/tensorflow/contrib/lite/kernels/internal/tensor_utils.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_ #include "tensorflow/contrib/lite/builtin_op_data.h" @@ -113,4 +113,4 @@ void ReductionSumVector(const float* input_vector, float* output_vector, } // namespace tensor_utils } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/types.h b/tensorflow/contrib/lite/kernels/internal/types.h index 07f1cb40045fff3ae47ed4efa6ec43b0cb88a0a7..afe131b06ec41201395e80aa5415fd7db990f8d4 100644 --- a/tensorflow/contrib/lite/kernels/internal/types.h +++ b/tensorflow/contrib/lite/kernels/internal/types.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TYPES_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TYPES_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TYPES_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TYPES_H_ #include "tensorflow/contrib/lite/kernels/internal/compatibility.h" @@ -27,6 +27,58 @@ struct Dims { int strides[N]; }; +// Gets next index to iterate through a multidimensional array. +inline bool NextIndex(const int num_dims, const int* dims, int* current) { + TFLITE_DCHECK_GT(num_dims, 0); + TFLITE_DCHECK(dims != nullptr); + TFLITE_DCHECK(current != nullptr); + int carry = 1; + for (int idx = num_dims - 1; idx >= 0; --idx) { + int current_val = current[idx] + carry; + TFLITE_DCHECK_GE(dims[idx], current_val); + if (dims[idx] == current_val) { + current[idx] = 0; + } else { + current[idx] = current_val; + carry = 0; + break; + } + } + return (carry == 0); +} + +// Gets offset of index if reducing on axis. When reducing, the flattened offset +// will not change, if the input index changes on the given axis. For example, +// if you have a 3D tensor and you are reducing to 2D by eliminating axis 0, +// then index (0, 1, 2) and index (1, 1, 2) will map to the same flattened +// offset. +// TODO(kanlig): uses Dims to represent dimensions. +inline size_t ReducedOutputOffset(const int num_dims, const int* dims, + const int* index, const int num_axis, + const int* axis) { + TFLITE_DCHECK_GT(num_dims, 0); + TFLITE_DCHECK(dims != nullptr); + TFLITE_DCHECK(index != nullptr); + size_t offset = 0; + for (int idx = 0; idx < num_dims; ++idx) { + // if we need to skip this axis + bool is_axis = false; + if (axis != nullptr) { + for (int axis_idx = 0; axis_idx < num_axis; ++axis_idx) { + if (idx == axis[axis_idx]) { + is_axis = true; + break; + } + } + } + if (!is_axis) { + offset = offset * static_cast(dims[idx]) + + static_cast(index[idx]); + } + } + return offset; +} + inline int Offset(const Dims<4>& dims, int i0, int i1, int i2, int i3) { TFLITE_DCHECK(i0 >= 0 && i0 < dims.sizes[0]); TFLITE_DCHECK(i1 >= 0 && i1 < dims.sizes[1]); @@ -36,6 +88,10 @@ inline int Offset(const Dims<4>& dims, int i0, int i1, int i2, int i3) { i3 * dims.strides[3]; } +inline int Offset(const Dims<4>& dims, int* index) { + return Offset(dims, index[0], index[1], index[2], index[3]); +} + // Get array size, DCHECKing that the dim index is in range. template int ArraySize(const Dims& array, int index) { @@ -78,4 +134,4 @@ bool IsPackedWithoutStrides(const Dims& dims) { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TYPES_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TYPES_H_ diff --git a/tensorflow/contrib/lite/kernels/kernel_util.cc b/tensorflow/contrib/lite/kernels/kernel_util.cc index b0546c00cf977af5f722a802866448b0cb293b8d..955e8c5764c6adad37a0009f4ddf8accb437b174 100644 --- a/tensorflow/contrib/lite/kernels/kernel_util.cc +++ b/tensorflow/contrib/lite/kernels/kernel_util.cc @@ -13,8 +13,11 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/lite/kernels/kernel_util.h" + #include #include +#include + #include "tensorflow/contrib/lite/kernels/internal/round.h" namespace tflite { @@ -84,4 +87,27 @@ void CalculateActivationRangeFloat(TfLiteFusedActivation activation, } } +bool HaveSameShapes(TfLiteTensor* input1, TfLiteTensor* input2) { + return TfLiteIntArrayEqual(input1->dims, input2->dims); +} + +TfLiteStatus CalculateShapeForBroadcast(TfLiteContext* context, + TfLiteTensor* input1, + TfLiteTensor* input2, + TfLiteIntArray** output_shape) { + int64_t dims1 = NumDimensions(input1); + int64_t dims2 = NumDimensions(input2); + int64_t out_dims = std::max(dims1, dims2); + std::unique_ptr shape( + TfLiteIntArrayCreate(out_dims), TfLiteIntArrayFree); + for (int i = 0; i < out_dims; ++i) { + int64_t d1 = i >= dims1 ? 1 : SizeOfDimension(input1, dims1 - i - 1); + int64_t d2 = i >= dims2 ? 1 : SizeOfDimension(input2, dims2 - i - 1); + TF_LITE_ENSURE(context, d1 == d2 || d1 == 1 || d2 == 1); + shape->data[out_dims - i - 1] = std::max(d1, d2); + } + *output_shape = shape.release(); + return kTfLiteOk; +} + } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/kernel_util.h b/tensorflow/contrib/lite/kernels/kernel_util.h index 25556ae4567aca45b3bfe4ba02b1cb58331d239d..28f53b9fbbc5620f2fab5c73e40bed8af4af5f1e 100644 --- a/tensorflow/contrib/lite/kernels/kernel_util.h +++ b/tensorflow/contrib/lite/kernels/kernel_util.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_KERNEL_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_KERNEL_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_KERNEL_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_KERNEL_UTIL_H_ #include "tensorflow/contrib/lite/builtin_op_data.h" #include "tensorflow/contrib/lite/context.h" @@ -35,6 +35,14 @@ inline TfLiteTensor* GetOutput(TfLiteContext* context, TfLiteNode* node, inline int NumInputs(const TfLiteNode* node) { return node->inputs->size; } inline int NumOutputs(const TfLiteNode* node) { return node->outputs->size; } +inline int64_t NumElements(const TfLiteTensor* t) { + int64_t count = 1; + for (int i = 0; i < NumDimensions(t); ++i) { + count *= SizeOfDimension(t, i); + } + return count; +} + inline TfLiteTensor* GetOptionalInputTensor(TfLiteContext* context, const TfLiteNode* node, int index) { const bool use_tensor = node->inputs->data[index] != kOptionalTensor; @@ -44,6 +52,25 @@ inline TfLiteTensor* GetOptionalInputTensor(TfLiteContext* context, return nullptr; } +// Determines whether tensor is constant. +inline bool IsConstantTensor(TfLiteTensor* tensor) { + return tensor->allocation_type == kTfLiteMmapRo; +} + +// Determines whether tensor is dynamic. Note that a tensor can be non-const and +// not dynamic. This function specificially checks for a dynamic tensor. +inline bool IsDynamicTensor(TfLiteTensor* tensor) { + return tensor->allocation_type == kTfLiteDynamic; +} + +// Sets tensor to dynamic. +inline void SetTensorToDynamic(TfLiteTensor* tensor) { + if (tensor->allocation_type != kTfLiteDynamic) { + tensor->allocation_type = kTfLiteDynamic; + tensor->data.raw = nullptr; + } +} + // Calculates the multiplication factor for a quantized convolution (or // quantized depthwise convolution) involving the given tensors. Returns an // error if the scales of the tensors are not compatible. @@ -60,6 +87,15 @@ void CalculateActivationRangeFloat(TfLiteFusedActivation activation, float* activation_min, float* activation_max); +// Return true if the given tensors have the same shape. +bool HaveSameShapes(TfLiteTensor* input1, TfLiteTensor* input2); + +// Calculate the output_shape that is necessary for element-wise operations +// with broadcasting involving the two input tensors. +TfLiteStatus CalculateShapeForBroadcast(TfLiteContext* context, + TfLiteTensor* input1, + TfLiteTensor* input2, + TfLiteIntArray** output_shape); } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_KERNEL_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_KERNEL_UTIL_H_ diff --git a/tensorflow/contrib/lite/kernels/kernel_util_test.cc b/tensorflow/contrib/lite/kernels/kernel_util_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..c65b68970f6853e17af3a70aad7a2bc982a1ee60 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/kernel_util_test.cc @@ -0,0 +1,152 @@ +/* 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/lite/kernels/kernel_util.h" + +#include +#include +#include "tensorflow/contrib/lite/testing/util.h" + +namespace tflite { +namespace { + +void ReportError(TfLiteContext* context, const char* format, ...) {} + +class KernelUtilTest : public ::testing::Test { + public: + KernelUtilTest() { + context_.ReportError = ReportError; + + tensor1_.dims = nullptr; + tensor2_.dims = nullptr; + tensor1_.allocation_type = kTfLiteMmapRo; + tensor2_.allocation_type = kTfLiteMmapRo; + } + ~KernelUtilTest() { + TfLiteTensorFree(&tensor1_); + TfLiteTensorFree(&tensor2_); + } + + void SetShape(TfLiteTensor* tensor, std::initializer_list dims) { + TfLiteTensorFree(tensor); + tensor->dims = TfLiteIntArrayCreate(dims.size()); + int i = 0; + for (int d : dims) { + tensor->dims->data[i] = d; + ++i; + } + } + + std::vector GetShape(TfLiteIntArray* dims) { + std::vector result; + for (int i = 0; i < dims->size; ++i) { + result.push_back(dims->data[i]); + } + return result; + } + + protected: + TfLiteContext context_; + TfLiteTensor tensor1_; + TfLiteTensor tensor2_; +}; + +TEST_F(KernelUtilTest, SameShapeEmpty) { + EXPECT_TRUE(HaveSameShapes(&tensor1_, &tensor2_)); + + SetShape(&tensor1_, {1, 2, 3}); + EXPECT_FALSE(HaveSameShapes(&tensor1_, &tensor2_)); + + SetShape(&tensor2_, {1, 2}); + EXPECT_FALSE(HaveSameShapes(&tensor1_, &tensor2_)); + + SetShape(&tensor2_, {1, 2, 3, 4}); + EXPECT_FALSE(HaveSameShapes(&tensor1_, &tensor2_)); + + SetShape(&tensor2_, {1, 2, 3}); + EXPECT_TRUE(HaveSameShapes(&tensor1_, &tensor2_)); + + SetShape(&tensor2_, {}); + EXPECT_FALSE(HaveSameShapes(&tensor1_, &tensor2_)); + + SetShape(&tensor1_, {}); + EXPECT_TRUE(HaveSameShapes(&tensor1_, &tensor2_)); +} + +TEST_F(KernelUtilTest, BroadcastShapeIncompatibleDim) { + TfLiteIntArray* output = nullptr; + SetShape(&tensor1_, {1, 2}); + SetShape(&tensor2_, {1, 3}); + EXPECT_NE(kTfLiteOk, CalculateShapeForBroadcast(&context_, &tensor1_, + &tensor2_, &output)); + EXPECT_EQ(output, nullptr); +} + +TEST_F(KernelUtilTest, BroadcastShapeOnes) { + TfLiteIntArray* output = nullptr; + SetShape(&tensor1_, {1, 1}); + SetShape(&tensor2_, {1, 3}); + EXPECT_EQ(kTfLiteOk, CalculateShapeForBroadcast(&context_, &tensor1_, + &tensor2_, &output)); + TfLiteIntArrayFree(output); + + SetShape(&tensor1_, {1, 2}); + SetShape(&tensor2_, {1, 1}); + EXPECT_EQ(kTfLiteOk, CalculateShapeForBroadcast(&context_, &tensor1_, + &tensor2_, &output)); + TfLiteIntArrayFree(output); +} + +TEST_F(KernelUtilTest, BroadcastShapeScalars) { + TfLiteIntArray* output = nullptr; + SetShape(&tensor1_, {1, 2}); + SetShape(&tensor2_, {}); + EXPECT_EQ(kTfLiteOk, CalculateShapeForBroadcast(&context_, &tensor1_, + &tensor2_, &output)); + EXPECT_THAT(GetShape(output), ::testing::ElementsAre(1, 2)); + TfLiteIntArrayFree(output); + + SetShape(&tensor1_, {}); + SetShape(&tensor2_, {2}); + EXPECT_EQ(kTfLiteOk, CalculateShapeForBroadcast(&context_, &tensor1_, + &tensor2_, &output)); + EXPECT_THAT(GetShape(output), ::testing::ElementsAre(2)); + TfLiteIntArrayFree(output); +} + +TEST_F(KernelUtilTest, BroadcastShapeDifferentSizes) { + TfLiteIntArray* output = nullptr; + SetShape(&tensor1_, {1, 2}); + SetShape(&tensor2_, {3, 1, 1}); + EXPECT_EQ(kTfLiteOk, CalculateShapeForBroadcast(&context_, &tensor1_, + &tensor2_, &output)); + EXPECT_THAT(GetShape(output), ::testing::ElementsAre(3, 1, 2)); + TfLiteIntArrayFree(output); + + SetShape(&tensor1_, {1, 2, 3, 4}); + SetShape(&tensor2_, {1, 3, 1}); + EXPECT_EQ(kTfLiteOk, CalculateShapeForBroadcast(&context_, &tensor1_, + &tensor2_, &output)); + EXPECT_THAT(GetShape(output), ::testing::ElementsAre(1, 2, 3, 4)); + TfLiteIntArrayFree(output); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/l2norm.cc b/tensorflow/contrib/lite/kernels/l2norm.cc index f43aa372b6398a38e57dd38f3d7c7db2bd3aefc1..ee8bfe56d95e9f383ef49b40b8f58b63d61da3e1 100644 --- a/tensorflow/contrib/lite/kernels/l2norm.cc +++ b/tensorflow/contrib/lite/kernels/l2norm.cc @@ -43,8 +43,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TfLiteTensor* input = GetInput(context, node, kInputTensor); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - // TODO(ahentz): Our current implementations rely on the inputs being 4D. - TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4); + TF_LITE_ENSURE(context, NumDimensions(input) <= 4); // TODO(ahentz): Our current implementations only support float32. TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32); @@ -54,12 +53,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // activations. TF_LITE_ENSURE_EQ(context, params->activation, kTfLiteActNone); - TfLiteIntArray* output_size = TfLiteIntArrayCreate(4); - output_size->data[0] = input->dims->data[0]; - output_size->data[1] = input->dims->data[1]; - output_size->data[2] = input->dims->data[2]; - output_size->data[3] = input->dims->data[3]; - + TfLiteIntArray* output_size = TfLiteIntArrayCopy(input->dims); return context->ResizeTensor(context, output, output_size); } diff --git a/tensorflow/contrib/lite/kernels/mean.cc b/tensorflow/contrib/lite/kernels/mean.cc new file mode 100644 index 0000000000000000000000000000000000000000..aff19581ea56f94c08638b7b388ae181f566cf4f --- /dev/null +++ b/tensorflow/contrib/lite/kernels/mean.cc @@ -0,0 +1,233 @@ +/* Copyright 2017 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 +#include +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace mean { + +// This file has reference implementation of Mean. +enum KernelType { + kReference, +}; + +struct MeanContext { + MeanContext(TfLiteContext* context, TfLiteNode* node) { + params = reinterpret_cast(node->builtin_data); + input = GetInput(context, node, 0); + axis = GetInput(context, node, 1); + output = GetOutput(context, node, 0); + } + TfLiteMeanParams* params; + TfLiteTensor* input; + TfLiteTensor* axis; + TfLiteTensor* output; +}; + +void* Init(TfLiteContext* context, const char* buffer, size_t length) { + // Creates two temp tensors to store index and axis for internal + // implementation only. + auto* scratch_tensor_index = new int; + context->AddTensors(context, 2, scratch_tensor_index); + return scratch_tensor_index; +} + +void Free(TfLiteContext* context, void* buffer) { + delete reinterpret_cast(buffer); +} + +// Resizes the temp tensor that stores resolved axis. +TfLiteStatus ResizeTempAxis(TfLiteContext* context, MeanContext* op_context, + TfLiteTensor* resolved_axis) { + TfLiteIntArray* axis_size = TfLiteIntArrayCreate(1); + axis_size->data[0] = static_cast(NumElements(op_context->axis)); + return context->ResizeTensor(context, resolved_axis, axis_size); +} + +// Resizes output array based on the input size and resolved axis. +TfLiteStatus ResizeOutputTensor(TfLiteContext* context, + MeanContext* op_context) { + size_t num_axis = NumElements(op_context->axis); + const TfLiteIntArray* input_dims = op_context->input->dims; + int input_num_dims = NumDimensions(op_context->input); + const int* axis = GetTensorData(op_context->axis); + if (op_context->params->keep_dims) { + TfLiteIntArray* output_dims = TfLiteIntArrayCreate(input_num_dims); + for (int idx = 0; idx < input_num_dims; ++idx) { + bool is_axis = false; + for (int axis_idx = 0; axis_idx < num_axis; ++axis_idx) { + if (axis[axis_idx] == idx || axis[axis_idx] + input_num_dims == idx) { + is_axis = true; + break; + } + } + if (is_axis) { + output_dims->data[idx] = 1; + } else { + output_dims->data[idx] = input_dims->data[idx]; + } + } + return context->ResizeTensor(context, op_context->output, output_dims); + } else { + // Calculates size of reducing axis. + int num_reduce_axis = num_axis; + for (int i = 0; i < num_axis; ++i) { + int current = axis[i]; + if (current < 0) { + current += input_num_dims; + } + TF_LITE_ENSURE(context, current >= 0 && current < input_num_dims); + for (int j = 0; j < i; ++j) { + int previous = axis[j]; + if (previous < 0) { + previous += input_num_dims; + } + if (current == previous) { + --num_reduce_axis; + break; + } + } + } + // Determines output dimensions. + TfLiteIntArray* output_dims = + TfLiteIntArrayCreate(input_num_dims - num_reduce_axis); + int num_skip_axis = 0; + for (int idx = 0; idx < input_num_dims; ++idx) { + bool is_axis = false; + for (int axis_idx = 0; axis_idx < num_axis; ++axis_idx) { + if (axis[axis_idx] == idx || axis[axis_idx] + input_num_dims == idx) { + ++num_skip_axis; + is_axis = true; + break; + } + } + if (!is_axis) { + output_dims->data[idx - num_skip_axis] = input_dims->data[idx]; + } + } + return context->ResizeTensor(context, op_context->output, output_dims); + } +} + +// Initializes temp tensors to store index and resolved axis. +TfLiteStatus InitializeTemporaries(TfLiteContext* context, TfLiteNode* node, + MeanContext* op_context) { + // Creates a temp index to iterate through input data. + int* scratch_tensor_index = reinterpret_cast(node->user_data); + TfLiteIntArrayFree(node->temporaries); + node->temporaries = TfLiteIntArrayCreate(2); + node->temporaries->data[0] = *scratch_tensor_index; + TfLiteTensor* scratch_tensor = &context->tensors[node->temporaries->data[0]]; + scratch_tensor->type = kTfLiteInt32; + scratch_tensor->allocation_type = kTfLiteArenaRw; + TfLiteIntArray* index_size = TfLiteIntArrayCreate(1); + index_size->data[0] = NumDimensions(op_context->input); + TF_LITE_ENSURE_OK(context, + context->ResizeTensor(context, scratch_tensor, index_size)); + + // Creates a temp tensor to store resolved axis given input data. + node->temporaries->data[1] = *scratch_tensor_index + 1; + TfLiteTensor* resolved_axis = &context->tensors[node->temporaries->data[1]]; + resolved_axis->type = kTfLiteInt32; + return kTfLiteOk; +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + MeanContext op_context(context, node); + TF_LITE_ENSURE_OK(context, InitializeTemporaries(context, node, &op_context)); + + TfLiteTensor* resolved_axis = &context->tensors[node->temporaries->data[1]]; + // Leaves work to Eval if axis is not constant; else resizes output. + if (!IsConstantTensor(op_context.axis)) { + SetTensorToDynamic(op_context.output); + SetTensorToDynamic(resolved_axis); + return kTfLiteOk; + } + resolved_axis->allocation_type = kTfLiteArenaRw; + TF_LITE_ENSURE_OK(context, + ResizeTempAxis(context, &op_context, resolved_axis)); + return ResizeOutputTensor(context, &op_context); +} + +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + MeanContext op_context(context, node); + int num_axis = static_cast(NumElements(op_context.axis)); + TfLiteTensor* temp_index = &context->tensors[node->temporaries->data[0]]; + TfLiteTensor* resolved_axis = &context->tensors[node->temporaries->data[1]]; + // Resize the output tensor if the output tensor is dynamic. + if (IsDynamicTensor(op_context.output)) { + TF_LITE_ENSURE_OK(context, + ResizeTempAxis(context, &op_context, resolved_axis)); + TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); + } + +#define TF_LITE_MEAN(kernel_type, data_type) \ + kernel_type::Mean<>( \ + GetTensorData(op_context.input), \ + op_context.input->dims->data, op_context.input->dims->size, \ + GetTensorData(op_context.output), \ + op_context.output->dims->data, op_context.output->dims->size, \ + GetTensorData(op_context.axis), num_axis, \ + op_context.params->keep_dims, GetTensorData(temp_index), \ + GetTensorData(resolved_axis)) + + if (kernel_type == kReference) { + switch (op_context.input->type) { + case kTfLiteFloat32: + TF_LITE_MEAN(reference_ops, float); + break; + case kTfLiteInt32: + TF_LITE_MEAN(reference_ops, int); + break; + case kTfLiteUInt8: + TF_LITE_MEAN(reference_ops, uint8_t); + break; + case kTfLiteInt64: + TF_LITE_MEAN(reference_ops, int64_t); + break; + default: + return kTfLiteError; + } + } +#undef TF_LITE_MEAN + return kTfLiteOk; +} + +} // namespace mean + +TfLiteRegistration* Register_MEAN_REF() { + static TfLiteRegistration r = {mean::Init, mean::Free, mean::Prepare, + mean::Eval}; + return &r; +} + +// TODO(kanlig): add optimized implementation of Mean. +TfLiteRegistration* Register_MEAN() { return Register_MEAN_REF(); } + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/mean_test.cc b/tensorflow/contrib/lite/kernels/mean_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..c4c53c2ded351849e7c458fc754c36395a25ebd0 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/mean_test.cc @@ -0,0 +1,140 @@ +/* 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 +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +class BaseMeanOpModel : public SingleOpModel { + public: + void SetAxis(std::initializer_list data) { PopulateTensor(axis_, data); } + + template + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + + template + std::vector GetOutput() { + return ExtractVector(output_); + } + + std::vector GetOutputShape() { return GetTensorShape(output_); } + + protected: + int input_; + int axis_; + int output_; +}; + +// Model for the tests case where axis is a const tensor. +class MeanOpConstModel : public BaseMeanOpModel { + public: + MeanOpConstModel(const TensorData& input, const TensorData& output, + std::initializer_list axis_shape, + std::initializer_list axis, bool keep_dims) { + input_ = AddInput(input); + axis_ = AddConstInput(TensorType_INT32, axis, axis_shape); + output_ = AddOutput(output); + SetBuiltinOp(BuiltinOperator_MEAN, BuiltinOptions_MeanOptions, + CreateMeanOptions(builder_, keep_dims).Union()); + BuildInterpreter({GetShape(input_)}); + } +}; + +// Model for the tests case where axis is a dynamic tensor. +class MeanOpDynamicModel : public BaseMeanOpModel { + public: + MeanOpDynamicModel(const TensorData& input, const TensorData& output, + const TensorData& axis, bool keep_dims) { + input_ = AddInput(input); + axis_ = AddInput(axis); + output_ = AddOutput(output); + SetBuiltinOp(BuiltinOperator_MEAN, BuiltinOptions_MeanOptions, + CreateMeanOptions(builder_, keep_dims).Union()); + BuildInterpreter({GetShape(input_)}); + } +}; + +TEST(ConstMeanOpTest, NotKeepDims) { + std::initializer_list data = { + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, + 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0}; + MeanOpConstModel m({TensorType_FLOAT32, {4, 3, 2}}, {TensorType_FLOAT32, {2}}, + {4}, {1, 0, -3, -3}, false); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({12, 13}))); +} + +TEST(ConstMeanOpTest, KeepDims) { + std::initializer_list data = { + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, + 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0}; + MeanOpConstModel m({TensorType_FLOAT32, {4, 3, 2}}, {TensorType_FLOAT32, {3}}, + {2}, {0, 2}, true); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 3, 1})); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear({10.5, 12.5, 14.5}))); +} + +TEST(DynamicMeanOpTest, NotKeepDims) { + std::initializer_list data = { + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, + 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0}; + MeanOpDynamicModel m({TensorType_FLOAT32, {4, 3, 2}}, + {TensorType_FLOAT32, {2}}, {TensorType_INT32, {4}}, + false); + std::initializer_list axis = {1, 0, -3, -3}; + m.SetAxis(axis); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({12, 13}))); +} + +TEST(DynamicMeanOpTest, KeepDims) { + std::initializer_list data = { + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, + 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0}; + MeanOpDynamicModel m({TensorType_FLOAT32, {4, 3, 2}}, + {TensorType_FLOAT32, {3}}, {TensorType_INT32, {2}}, + true); + std::initializer_list axis = {0, 2}; + m.SetAxis(axis); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 3, 1})); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear({10.5, 12.5, 14.5}))); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/mul.cc b/tensorflow/contrib/lite/kernels/mul.cc index 81c73f2523186c2d4072d56bdc8980fcdbb588a3..54575019de4c678ce25561cf2ac8dc80c9973363 100644 --- a/tensorflow/contrib/lite/kernels/mul.cc +++ b/tensorflow/contrib/lite/kernels/mul.cc @@ -37,7 +37,23 @@ constexpr int kInputTensor1 = 0; constexpr int kInputTensor2 = 1; constexpr int kOutputTensor = 0; +struct OpData { + bool requires_broadcast; +}; + +void* Init(TfLiteContext* context, const char* buffer, size_t length) { + auto* data = new OpData; + data->requires_broadcast = false; + return data; +} + +void Free(TfLiteContext* context, void* buffer) { + delete reinterpret_cast(buffer); +} + TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + OpData* data = reinterpret_cast(node->user_data); + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); @@ -45,43 +61,56 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - TF_LITE_ENSURE_EQ(context, NumDimensions(input1), NumDimensions(input2)); - for (int i = 0; i < NumDimensions(input1); ++i) { - TF_LITE_ENSURE_EQ(context, SizeOfDimension(input1, i), - SizeOfDimension(input2, i)); - } + TF_LITE_ENSURE_EQ(context, input1->type, input2->type); + output->type = input2->type; + + data->requires_broadcast = !HaveSameShapes(input1, input2); - TF_LITE_ENSURE_EQ(context, input1->type, output->type); - TF_LITE_ENSURE_EQ(context, input2->type, output->type); + TfLiteIntArray* output_size = nullptr; + if (data->requires_broadcast) { + TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast( + context, input1, input2, &output_size)); + } else { + output_size = TfLiteIntArrayCopy(input1->dims); + } - TfLiteIntArray* output_size = TfLiteIntArrayCopy(input1->dims); return context->ResizeTensor(context, output, output_size); } template void EvalFloat(TfLiteContext* context, TfLiteNode* node, - TfLiteMulParams* params, TfLiteTensor* input1, - TfLiteTensor* input2, TfLiteTensor* output) { + TfLiteMulParams* params, const OpData* data, + TfLiteTensor* input1, TfLiteTensor* input2, + TfLiteTensor* output) { float output_activation_min, output_activation_max; CalculateActivationRangeFloat(params->activation, &output_activation_min, &output_activation_max); -#define TF_LITE_MUL(type) \ - type::Mul(GetTensorData(input1), GetTensorDims(input1), \ - GetTensorData(input2), GetTensorDims(input2), \ - output_activation_min, output_activation_max, \ - GetTensorData(output), GetTensorDims(output)) +#define TF_LITE_MUL(type, opname) \ + type::opname(GetTensorData(input1), GetTensorDims(input1), \ + GetTensorData(input2), GetTensorDims(input2), \ + output_activation_min, output_activation_max, \ + GetTensorData(output), GetTensorDims(output)) if (kernel_type == kReference) { - TF_LITE_MUL(reference_ops); + if (data->requires_broadcast) { + TF_LITE_MUL(reference_ops, BroadcastMul); + } else { + TF_LITE_MUL(reference_ops, Mul); + } } else { - TF_LITE_MUL(optimized_ops); + if (data->requires_broadcast) { + TF_LITE_MUL(optimized_ops, BroadcastMul); + } else { + TF_LITE_MUL(optimized_ops, Mul); + } } #undef TF_LITE_MUL } template void EvalQuantized(TfLiteContext* context, TfLiteNode* node, - TfLiteMulParams* params, TfLiteTensor* input1, - TfLiteTensor* input2, TfLiteTensor* output) { + TfLiteMulParams* params, const OpData* data, + TfLiteTensor* input1, TfLiteTensor* input2, + TfLiteTensor* output) { auto input1_offset = -input1->params.zero_point; auto input2_offset = -input2->params.zero_point; auto output_offset = output->params.zero_point; @@ -98,17 +127,19 @@ void EvalQuantized(TfLiteContext* context, TfLiteNode* node, CalculateActivationRangeUint8(params->activation, output, &output_activation_min, &output_activation_max); -#define TF_LITE_MUL(type) \ - type::BroadcastMul(GetTensorData(input1), GetTensorDims(input1), \ - input1_offset, GetTensorData(input2), \ - GetTensorDims(input2), input2_offset, output_offset, \ - output_multiplier, output_shift, output_activation_min, \ - output_activation_max, GetTensorData(output), \ - GetTensorDims(output)); +#define TF_LITE_MUL(type, opname) \ + type::opname(GetTensorData(input1), GetTensorDims(input1), \ + input1_offset, GetTensorData(input2), \ + GetTensorDims(input2), input2_offset, output_offset, \ + output_multiplier, output_shift, output_activation_min, \ + output_activation_max, GetTensorData(output), \ + GetTensorDims(output)); + // The quantized version of Mul doesn't support activations, so we + // always use BroadcastMul. if (kernel_type == kReference) { - TF_LITE_MUL(reference_ops); + TF_LITE_MUL(reference_ops, BroadcastMul); } else { - TF_LITE_MUL(optimized_ops); + TF_LITE_MUL(optimized_ops, BroadcastMul); } #undef TF_LITE_MUL } @@ -116,15 +147,17 @@ void EvalQuantized(TfLiteContext* context, TfLiteNode* node, template TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->builtin_data); + OpData* data = reinterpret_cast(node->user_data); TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); if (output->type == kTfLiteFloat32) { - EvalFloat(context, node, params, input1, input2, output); + EvalFloat(context, node, params, data, input1, input2, output); } else if (output->type == kTfLiteUInt8) { - EvalQuantized(context, node, params, input1, input2, output); + EvalQuantized(context, node, params, data, input1, input2, + output); } else { context->ReportError(context, "Mul only supports FLOAT32 and quantized UINT8 now."); @@ -137,19 +170,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { } // namespace mul TfLiteRegistration* Register_MUL_REF() { - static TfLiteRegistration r = {nullptr, nullptr, mul::Prepare, + static TfLiteRegistration r = {mul::Init, mul::Free, mul::Prepare, mul::Eval}; return &r; } TfLiteRegistration* Register_MUL_GENERIC_OPT() { - static TfLiteRegistration r = {nullptr, nullptr, mul::Prepare, + static TfLiteRegistration r = {mul::Init, mul::Free, mul::Prepare, mul::Eval}; return &r; } TfLiteRegistration* Register_MUL_NEON_OPT() { - static TfLiteRegistration r = {nullptr, nullptr, mul::Prepare, + static TfLiteRegistration r = {mul::Init, mul::Free, mul::Prepare, mul::Eval}; return &r; } diff --git a/tensorflow/contrib/lite/kernels/mul_test.cc b/tensorflow/contrib/lite/kernels/mul_test.cc index 4255cfe18a043c55f3ce7292afdedb6e988a28a2..f1a30f82634631ba8320421d5b36ffe446f443fa 100644 --- a/tensorflow/contrib/lite/kernels/mul_test.cc +++ b/tensorflow/contrib/lite/kernels/mul_test.cc @@ -25,10 +25,11 @@ using ::testing::ElementsAreArray; class BaseMulOpModel : public SingleOpModel { public: - BaseMulOpModel(TensorData input, TensorData output, + BaseMulOpModel(const TensorData& input1, const TensorData& input2, + const TensorData& output, ActivationFunctionType activation_type) { - input1_ = AddInput(input); - input2_ = AddInput(input); + input1_ = AddInput(input1); + input2_ = AddInput(input2); output_ = AddOutput(output); SetBuiltinOp(BuiltinOperator_MUL, BuiltinOptions_MulOptions, CreateMulOptions(builder_, activation_type).Union()); @@ -70,6 +71,7 @@ class QuantizedMulOpModel : public BaseMulOpModel { TEST(FloatMulOpTest, NoActivation) { FloatMulOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, + {TensorType_FLOAT32, {1, 2, 2, 1}}, {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8}); m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.5}); @@ -78,9 +80,10 @@ TEST(FloatMulOpTest, NoActivation) { ElementsAreArray(ArrayFloatNear({-0.2, 0.04, 0.21, 0.4}))); } -TEST(FloatMulOpTest, ActivationRELU1) { - FloatMulOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, - {TensorType_FLOAT32, {}}, ActivationFunctionType_RELU1); +TEST(FloatMulOpTest, ActivationRELU_N1_TO_1) { + FloatMulOpModel m( + {TensorType_FLOAT32, {1, 2, 2, 1}}, {TensorType_FLOAT32, {1, 2, 2, 1}}, + {TensorType_FLOAT32, {}}, ActivationFunctionType_RELU_N1_TO_1); m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8}); m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 5}); m.Invoke(); @@ -93,6 +96,7 @@ TEST(FloatMulOpTest, VariousInputShapes) { {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; for (int i = 0; i < test_shapes.size(); ++i) { FloatMulOpModel m({TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, test_shapes[i]}, {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.5, 1.1, 0.1}); @@ -104,8 +108,26 @@ TEST(FloatMulOpTest, VariousInputShapes) { } } +TEST(FloatMulOpTest, WithBroadcast) { + std::vector> test_shapes = { + {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; + for (int i = 0; i < test_shapes.size(); ++i) { + FloatMulOpModel m({TensorType_FLOAT32, test_shapes[i]}, + {TensorType_FLOAT32, {}}, // always a scalar + {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); + m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); + m.PopulateTensor(m.input2(), {0.1}); + m.Invoke(); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray(ArrayFloatNear({-0.2, 0.02, 0.07, 0.08, 0.11, 0.2}))) + << "With shape number " << i; + } +} + TEST(QuantizedMulOpTest, NoActivation) { QuantizedMulOpModel m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, + {TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, {TensorType_UINT8, {}, -1.0, 1.0}, ActivationFunctionType_NONE); m.QuantizeAndPopulate(m.input1(), {-0.8, 0.2, 0.9, 0.7}); @@ -116,6 +138,32 @@ TEST(QuantizedMulOpTest, NoActivation) { kQuantizedTolerance))); } +// for quantized Mul, the error shouldn't exceed 2*step +float GetTolerance(int min, int max) { + float kQuantizedStep = (max - min) / 255.0; + float kQuantizedTolerance = 2.0 * kQuantizedStep; + return kQuantizedTolerance; +} + +TEST(QuantizedMulOpTest, WithBroadcast) { + float kQuantizedTolerance = GetTolerance(-3.0, 3.0); + std::vector> test_shapes = { + {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; + for (int i = 0; i < test_shapes.size(); ++i) { + QuantizedMulOpModel m({TensorType_UINT8, test_shapes[i], -3.0, 3.0}, + {TensorType_UINT8, {}, -3.0, 3.0}, // always a scalar + {TensorType_UINT8, {}, -3.0, 3.0}, + ActivationFunctionType_NONE); + m.QuantizeAndPopulate(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); + m.QuantizeAndPopulate(m.input2(), {0.1}); + m.Invoke(); + EXPECT_THAT(m.GetDequantizedOutput(), + ElementsAreArray(ArrayFloatNear( + {-0.2, 0.02, 0.07, 0.08, 0.11, 0.2}, kQuantizedTolerance))) + << "With shape number " << i; + } +} + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/op_macros.h b/tensorflow/contrib/lite/kernels/op_macros.h index 63670efcb1e6349317aa5c75756707fb7a7fa2aa..7568eaa88edfa3260964e16f03299aecb97da6be 100644 --- a/tensorflow/contrib/lite/kernels/op_macros.h +++ b/tensorflow/contrib/lite/kernels/op_macros.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_OP_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_OP_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_OP_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_OP_UTIL_H_ #include @@ -31,4 +31,4 @@ limitations under the License. if ((x) != (y)) TF_LITE_FATAL(#x " didn't equal " #y); \ } while (0) -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_OP_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_OP_UTIL_H_ diff --git a/tensorflow/contrib/lite/kernels/optional_tensor_test.cc b/tensorflow/contrib/lite/kernels/optional_tensor_test.cc index 17166715ca30ff3d8ba3d384110e403f8910e39d..cee3ec6197c698a11004d42dccdfe2bcca088015 100644 --- a/tensorflow/contrib/lite/kernels/optional_tensor_test.cc +++ b/tensorflow/contrib/lite/kernels/optional_tensor_test.cc @@ -243,7 +243,6 @@ class LSTMOpModel : public SingleOpModel { int n_output_; }; - TEST(LSTMOpTest, BlackBoxTestWithCifgWithPeepholeNoProjectionNoClipping) { const int n_batch = 1; const int n_input = 2; @@ -282,7 +281,6 @@ TEST(LSTMOpTest, BlackBoxTestWithCifgWithPeepholeNoProjectionNoClipping) { {0}, // projection_bias tensor }); - lstm.SetInputToCellWeights({-0.49770179, -0.27711356, -0.09624726, 0.05100781, 0.04717243, 0.48944736, -0.38535351, -0.17212132}); diff --git a/tensorflow/contrib/lite/kernels/pad.cc b/tensorflow/contrib/lite/kernels/pad.cc index 5e90282a43b1b6caf7918b3874fd4273f59e31b7..c29da3862e84d6756bf5ef34b2ca06307b0a065d 100644 --- a/tensorflow/contrib/lite/kernels/pad.cc +++ b/tensorflow/contrib/lite/kernels/pad.cc @@ -33,85 +33,132 @@ enum KernelType { kGenericOptimized, }; -// TODO(nupurgarg): Padding represented as a tensor is ignored. Only use the -// `left_padding` and `right_padding` specified in `params`. struct PadContext { PadContext(TfLiteContext* context, TfLiteNode* node) { - params = reinterpret_cast(node->builtin_data); input = GetInput(context, node, 0); + paddings = GetInput(context, node, 1); output = GetOutput(context, node, 0); + dims = NumDimensions(input); } - TfLitePadParams* params; TfLiteTensor* input; + TfLiteTensor* paddings; TfLiteTensor* output; + int dims; }; +// Resizes output array based on the input size and padding size. This function +// is callable from both Prepare() and Eval() as long as the caller ensures the +// paddings data is present. +TfLiteStatus ResizeOutputTensor(TfLiteContext* context, + PadContext* op_context) { + // Ensures the paddings array is dims x 2. + TF_LITE_ENSURE_EQ(context, SizeOfDimension(op_context->paddings, 0), + op_context->dims); + TF_LITE_ENSURE_EQ(context, SizeOfDimension(op_context->paddings, 1), 2); + + // Determines the size of the output tensor. + TfLiteIntArray* input_size = op_context->input->dims; + TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_size); + const int32* paddings_data = GetTensorData(op_context->paddings); + + for (int idx = 0; idx < op_context->dims; ++idx) { + int before_padding = *paddings_data++; + int after_padding = *paddings_data++; + + TF_LITE_ENSURE_MSG(context, (before_padding >= 0 && after_padding >= 0), + "Pad value has to be greater than equal to 0."); + + output_size->data[idx] = + (input_size->data[idx] + before_padding + after_padding); + } + + return context->ResizeTensor(context, op_context->output, output_size); +} + TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - TF_LITE_ENSURE(context, NumInputs(node) == 1 || NumInputs(node) == 2); + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - // Determines size of output tensor. PadContext op_context(context, node); - int dims = NumDimensions(op_context.input); - TF_LITE_ENSURE_EQ(context, dims, op_context.params->num_dimensions); + TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type); // TODO(nupurgarg): Our current implementations rely on the inputs being 4D. - TF_LITE_ENSURE_EQ(context, dims, 4); - - const TfLiteIntArray* input_size = op_context.input->dims; - TfLiteIntArray* output_size = TfLiteIntArrayCreate(dims); - for (int idx = 0; idx < dims; ++idx) { - TF_LITE_ENSURE_MSG(context, - (op_context.params->before_padding[idx] >= 0 && - op_context.params->after_padding[idx] >= 0), - "Pad value has to be greater than equal to 0."); - output_size->data[idx] = - (input_size->data[idx] + op_context.params->before_padding[idx] + - op_context.params->after_padding[idx]); - } + TF_LITE_ENSURE_EQ(context, op_context.dims, 4); - return context->ResizeTensor(context, op_context.output, output_size); + // Exit early if paddings is a non-const tensor. Set output tensor to + // dynamic so output size can be determined in Eval. + if (!IsConstantTensor(op_context.paddings)) { + SetTensorToDynamic(op_context.output); + return kTfLiteOk; + } + return ResizeOutputTensor(context, &op_context); } template TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { PadContext op_context(context, node); - // TODO(nupurgarg): Support different data types. - if (op_context.output->type == kTfLiteFloat32) { - std::vector before_padding( - op_context.params->before_padding, - op_context.params->before_padding + op_context.params->num_dimensions); - std::vector after_padding( - op_context.params->after_padding, - op_context.params->after_padding + op_context.params->num_dimensions); - - // TODO(nupurgarg): Change TOCO's implementation to use padding arrays - // in forward order (depth, width, height, batch). - // Converts from int[] = {depth, width, height, batch} to int[] = {batch, - // height, width, depth} to match TOCO's implementation of pad in - // referenced_ops.h and optimized_ops.h. - std::reverse(before_padding.begin(), before_padding.end()); - std::reverse(after_padding.begin(), after_padding.end()); - -#define TF_LITE_PAD(type) \ - type::Pad(GetTensorData(op_context.input), \ + // Resize the output tensor if the output tensor is dynamic. + if (IsDynamicTensor(op_context.output)) { + TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); + } + + // TODO(nupurgarg): Change kernel implementation to take in int* instead of + // vector to remove malloc from Eval(). + // Create before and after padding arrays that are accepted by the kernel. + std::vector before_padding; + std::vector after_padding; + const int32* paddings_data = GetTensorData(op_context.paddings); + + // TODO(nupurgarg): Change kernel implementation to use padding arrays in + // forward order (depth, width, height, batch). + // Build paddings in order of int[] = {batch, height, width, depth} to match + // kernel implementation of Pad in referenced_ops.h and optimized_ops.h. + for (int idx = op_context.dims - 1; idx >= 0; --idx) { + before_padding.push_back(paddings_data[idx * 2]); + after_padding.push_back(paddings_data[idx * 2 + 1]); + } + +#define TF_LITE_PAD(type, scalar) \ + type::Pad(GetTensorData(op_context.input), \ GetTensorDims(op_context.input), before_padding, after_padding, \ - GetTensorData(op_context.output), \ + GetTensorData(op_context.output), \ GetTensorDims(op_context.output)) - if (kernel_type == kReference) { - TF_LITE_PAD(reference_ops); - } - if (kernel_type == kGenericOptimized) { - TF_LITE_PAD(optimized_ops); - } -#undef TF_LITE_PAD - } else { - context->ReportError(context, "Inputs and outputs not all float types."); - return kTfLiteError; + switch (op_context.input->type) { + case kTfLiteFloat32: + if (kernel_type == kReference) { + TF_LITE_PAD(reference_ops, float); + } else if (kernel_type == kGenericOptimized) { + TF_LITE_PAD(optimized_ops, float); + } + break; + case kTfLiteUInt8: + if (kernel_type == kReference) { + TF_LITE_PAD(reference_ops, uint8_t); + } else if (kernel_type == kGenericOptimized) { + TF_LITE_PAD(optimized_ops, uint8_t); + } + break; + case kTfLiteInt32: + if (kernel_type == kReference) { + TF_LITE_PAD(reference_ops, int32_t); + } else if (kernel_type == kGenericOptimized) { + TF_LITE_PAD(optimized_ops, int32_t); + } + break; + case kTfLiteInt64: + if (kernel_type == kReference) { + TF_LITE_PAD(reference_ops, int64_t); + } else if (kernel_type == kGenericOptimized) { + TF_LITE_PAD(optimized_ops, int64_t); + } + break; + default: + context->ReportError(context, "Type is currently not supported by Pad."); + return kTfLiteError; } - +#undef TF_LITE_PAD return kTfLiteOk; } @@ -129,10 +176,7 @@ TfLiteRegistration* Register_PAD_GENERIC_OPT() { return &r; } -TfLiteRegistration* Register_PAD() { - return Register_PAD_GENERIC_OPT(); - // return Register_PAD_REF(); -} +TfLiteRegistration* Register_PAD() { return Register_PAD_GENERIC_OPT(); } } // namespace builtin } // namespace ops diff --git a/tensorflow/contrib/lite/kernels/pad_test.cc b/tensorflow/contrib/lite/kernels/pad_test.cc index f3ea9417df0e61dcff7a877726ab91c9b22691ba..28834ad0719291b2e868bca2d86a6685e6eb9962 100644 --- a/tensorflow/contrib/lite/kernels/pad_test.cc +++ b/tensorflow/contrib/lite/kernels/pad_test.cc @@ -25,52 +25,87 @@ using ::testing::ElementsAreArray; class PadOpModel : public SingleOpModel { public: - PadOpModel(std::initializer_list input_shape, - std::initializer_list before_padding, - std::initializer_list after_padding) { - input_ = AddInput(TensorType_FLOAT32); - output_ = AddOutput(TensorType_FLOAT32); - SetBuiltinOp( - BuiltinOperator_PAD, BuiltinOptions_PadOptions, - CreatePadOptions(builder_, builder_.CreateVector(before_padding), - builder_.CreateVector(after_padding)) - .Union()); - BuildInterpreter({input_shape}); - } - void SetInput(std::initializer_list data) { PopulateTensor(input_, data); } + void SetPaddings(std::initializer_list paddings) { + PopulateTensor(paddings_, paddings); + } + std::vector GetOutput() { return ExtractVector(output_); } std::vector GetOutputShape() { return GetTensorShape(output_); } - private: + protected: int input_; int output_; + int paddings_; +}; + +// Tests case where paddings is a const tensor. +// +// Example usage is as follows: +// PadOpDynamicModel m(input_shape, paddings_shape, paddings_data); +// m.SetInput(input_data); +// m.Invoke(); +class PadOpConstModel : public PadOpModel { + public: + PadOpConstModel(std::initializer_list input_shape, + std::initializer_list paddings_shape, + std::initializer_list paddings) { + input_ = AddInput(TensorType_FLOAT32); + paddings_ = AddConstInput(TensorType_INT32, paddings, paddings_shape); + output_ = AddOutput(TensorType_FLOAT32); + + SetBuiltinOp(BuiltinOperator_PAD, BuiltinOptions_PadOptions, + CreatePadOptions(builder_).Union()); + BuildInterpreter({input_shape}); + } +}; + +// Test case where paddings is a non-const tensor. +// +// Example usage is as follows: +// PadOpDynamicModel m(input_shape, paddings_shape); +// m.SetInput(input_data); +// m.SetPaddings(paddings_data); +// m.Invoke(); +class PadOpDynamicModel : public PadOpModel { + public: + PadOpDynamicModel(std::initializer_list input_shape, + std::initializer_list paddings_shape) { + input_ = AddInput(TensorType_FLOAT32); + paddings_ = AddInput(TensorType_INT32); + output_ = AddOutput(TensorType_FLOAT32); + + SetBuiltinOp(BuiltinOperator_PAD, BuiltinOptions_PadOptions, + CreatePadOptions(builder_).Union()); + BuildInterpreter({input_shape, paddings_shape}); + } }; TEST(PadOpTest, TooManyDimensions) { EXPECT_DEATH( - PadOpModel({1, 2, 3, 4, 5, 6, 7, 8, 9}, {1, 2, 3, 4, 5, 6, 7, 8, 9}, - {1, 2, 3, 4, 5, 6, 7, 8, 9}), + PadOpConstModel({1, 2, 3, 4, 5, 6, 7, 8, 9}, {9, 2}, + {1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9}), "dims != 4"); } -// TODO(nupurgarg): Test case where before padding and after padding arrays -// don't contain the same number of dimensions. TEST(PadOpTest, UnequalDimensions) { - EXPECT_DEATH(PadOpModel({1, 1, 2, 1}, {1, 2, 3}, {1, 2, 3}), - "dims != op_context.params->num_dimensions"); + EXPECT_DEATH(PadOpConstModel({1, 1, 2, 1}, {3, 2}, {1, 1, 2, 2, 3, 3}), + "3 != 4"); } TEST(PadOpTest, InvalidPadValue) { - EXPECT_DEATH(PadOpModel({1, 1, 2, 1}, {0, 1, 2, 0}, {0, -1, -1, 0}), - "Pad value has to be greater than equal to 0."); + EXPECT_DEATH( + PadOpConstModel({1, 1, 2, 1}, {4, 2}, {0, 0, 1, -1, 2, -1, 0, 0}), + "Pad value has to be greater than equal to 0."); } -TEST(PadOpTest, SimpleTest) { - PadOpModel m({1, 2, 2, 1}, {0, 1, 1, 0}, {0, 1, 1, 0}); +TEST(PadOpTest, SimpleConstTest) { + // Padding is represented as four 2-D lists representing above padding and + // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}). + PadOpConstModel m({1, 2, 2, 1}, {4, 2}, {0, 0, 1, 1, 1, 1, 0, 0}); m.SetInput({1, 2, 3, 4}); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray({0, 0, 0, 0, 0, 1, 2, 0, 0, 3, 4, @@ -78,10 +113,30 @@ TEST(PadOpTest, SimpleTest) { EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); } -TEST(PadOpTest, AdvancedTest) { - // The padding is input in the order of batch, height, width, depth. - PadOpModel m({1, 2, 3, 1}, {0, 0, 1, 0}, {0, 2, 3, 0}); +TEST(PadOpTest, SimpleDynamicTest) { + PadOpDynamicModel m({1, 2, 2, 1}, {4, 2}); + m.SetInput({1, 2, 3, 4}); + m.SetPaddings({0, 0, 1, 1, 1, 1, 0, 0}); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({0, 0, 0, 0, 0, 1, 2, 0, 0, 3, 4, + 0, 0, 0, 0, 0})); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); +} + +TEST(PadOpTest, AdvancedConstTest) { + PadOpConstModel m({1, 2, 3, 1}, {4, 2}, {0, 0, 0, 2, 1, 3, 0, 0}); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.Invoke(); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray({0, 1, 2, 3, 0, 0, 0, 0, 4, 5, 6, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1})); +} + +TEST(PadOpTest, AdvancedDynamicTest) { + PadOpDynamicModel m({1, 2, 3, 1}, {4, 2}); m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetPaddings({0, 0, 0, 2, 1, 3, 0, 0}); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray({0, 1, 2, 3, 0, 0, 0, 0, 4, 5, 6, 0, 0, 0, diff --git a/tensorflow/contrib/lite/kernels/padding.h b/tensorflow/contrib/lite/kernels/padding.h index 3a60274524c468ef29e522de5569e0d8354974c2..40b8476b3779c66e31a04856bce8aebd378f1e5f 100644 --- a/tensorflow/contrib/lite/kernels/padding.h +++ b/tensorflow/contrib/lite/kernels/padding.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_PADDING_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_PADDING_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_PADDING_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_PADDING_H_ namespace tflite { @@ -25,4 +25,4 @@ inline int ComputePadding(int stride, int in_size, int filter_size, } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_PADDING_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_PADDING_H_ diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc index d4e7503f48debbdc092ad7950ee4c0e52854c432..edc4e26edbd44784f8604e7da32156a8e695d2e2 100644 --- a/tensorflow/contrib/lite/kernels/register.cc +++ b/tensorflow/contrib/lite/kernels/register.cc @@ -20,7 +20,7 @@ namespace ops { namespace builtin { TfLiteRegistration* Register_RELU(); -TfLiteRegistration* Register_RELU1(); +TfLiteRegistration* Register_RELU_N1_TO_1(); TfLiteRegistration* Register_RELU6(); TfLiteRegistration* Register_TANH(); TfLiteRegistration* Register_LOGISTIC(); @@ -31,6 +31,7 @@ TfLiteRegistration* Register_CONV_2D(); TfLiteRegistration* Register_DEPTHWISE_CONV_2D(); TfLiteRegistration* Register_SVDF(); TfLiteRegistration* Register_RNN(); +TfLiteRegistration* Register_BIDIRECTIONAL_SEQUENCE_RNN(); TfLiteRegistration* Register_UNIDIRECTIONAL_SEQUENCE_RNN(); TfLiteRegistration* Register_EMBEDDING_LOOKUP(); TfLiteRegistration* Register_EMBEDDING_LOOKUP_SPARSE(); @@ -40,21 +41,32 @@ TfLiteRegistration* Register_HASHTABLE_LOOKUP(); TfLiteRegistration* Register_SOFTMAX(); TfLiteRegistration* Register_CONCATENATION(); TfLiteRegistration* Register_ADD(); +TfLiteRegistration* Register_SPACE_TO_BATCH_ND(); +TfLiteRegistration* Register_DIV(); +TfLiteRegistration* Register_SUB(); TfLiteRegistration* Register_BATCH_TO_SPACE_ND(); TfLiteRegistration* Register_MUL(); TfLiteRegistration* Register_L2_NORMALIZATION(); TfLiteRegistration* Register_LOCAL_RESPONSE_NORMALIZATION(); TfLiteRegistration* Register_LSTM(); +TfLiteRegistration* Register_UNIDIRECTIONAL_SEQUENCE_LSTM(); TfLiteRegistration* Register_PAD(); TfLiteRegistration* Register_RESHAPE(); TfLiteRegistration* Register_RESIZE_BILINEAR(); TfLiteRegistration* Register_SKIP_GRAM(); TfLiteRegistration* Register_SPACE_TO_DEPTH(); TfLiteRegistration* Register_GATHER(); +TfLiteRegistration* Register_TRANSPOSE(); +TfLiteRegistration* Register_MEAN(); +TfLiteRegistration* Register_SPLIT(); +TfLiteRegistration* Register_SQUEEZE(); +TfLiteRegistration* Register_STRIDED_SLICE(); +TfLiteRegistration* Register_EXP(); +TfLiteRegistration* Register_TOPK_V2(); BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_RELU, Register_RELU()); - AddBuiltin(BuiltinOperator_RELU1, Register_RELU1()); + AddBuiltin(BuiltinOperator_RELU_N1_TO_1, Register_RELU_N1_TO_1()); AddBuiltin(BuiltinOperator_RELU6, Register_RELU6()); AddBuiltin(BuiltinOperator_TANH, Register_TANH()); AddBuiltin(BuiltinOperator_LOGISTIC, Register_LOGISTIC()); @@ -65,6 +77,8 @@ BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_DEPTHWISE_CONV_2D, Register_DEPTHWISE_CONV_2D()); AddBuiltin(BuiltinOperator_SVDF, Register_SVDF()); AddBuiltin(BuiltinOperator_RNN, Register_RNN()); + AddBuiltin(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN, + Register_BIDIRECTIONAL_SEQUENCE_RNN()); AddBuiltin(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, Register_UNIDIRECTIONAL_SEQUENCE_RNN()); AddBuiltin(BuiltinOperator_EMBEDDING_LOOKUP, Register_EMBEDDING_LOOKUP()); @@ -76,18 +90,30 @@ BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_SOFTMAX, Register_SOFTMAX()); AddBuiltin(BuiltinOperator_CONCATENATION, Register_CONCATENATION()); AddBuiltin(BuiltinOperator_ADD, Register_ADD()); + AddBuiltin(BuiltinOperator_SPACE_TO_BATCH_ND, Register_SPACE_TO_BATCH_ND()); AddBuiltin(BuiltinOperator_BATCH_TO_SPACE_ND, Register_BATCH_TO_SPACE_ND()); AddBuiltin(BuiltinOperator_MUL, Register_MUL()); AddBuiltin(BuiltinOperator_L2_NORMALIZATION, Register_L2_NORMALIZATION()); AddBuiltin(BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION, Register_LOCAL_RESPONSE_NORMALIZATION()); AddBuiltin(BuiltinOperator_LSTM, Register_LSTM()); + AddBuiltin(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, + Register_UNIDIRECTIONAL_SEQUENCE_LSTM()); AddBuiltin(BuiltinOperator_PAD, Register_PAD()); AddBuiltin(BuiltinOperator_RESHAPE, Register_RESHAPE()); AddBuiltin(BuiltinOperator_RESIZE_BILINEAR, Register_RESIZE_BILINEAR()); AddBuiltin(BuiltinOperator_SKIP_GRAM, Register_SKIP_GRAM()); AddBuiltin(BuiltinOperator_SPACE_TO_DEPTH, Register_SPACE_TO_DEPTH()); AddBuiltin(BuiltinOperator_GATHER, Register_GATHER()); + AddBuiltin(BuiltinOperator_TRANSPOSE, Register_TRANSPOSE()); + AddBuiltin(BuiltinOperator_MEAN, Register_MEAN()); + AddBuiltin(BuiltinOperator_DIV, Register_DIV()); + AddBuiltin(BuiltinOperator_SUB, Register_SUB()); + AddBuiltin(BuiltinOperator_SPLIT, Register_SPLIT()); + AddBuiltin(BuiltinOperator_SQUEEZE, Register_SQUEEZE()); + AddBuiltin(BuiltinOperator_STRIDED_SLICE, Register_STRIDED_SLICE()); + AddBuiltin(BuiltinOperator_EXP, Register_EXP()); + AddBuiltin(BuiltinOperator_TOPK_V2, Register_TOPK_V2()); } TfLiteRegistration* BuiltinOpResolver::FindOp( diff --git a/tensorflow/contrib/lite/kernels/register.h b/tensorflow/contrib/lite/kernels/register.h index 28f5e0fcc80a14cf9fb6fb19b795d0c0d55e0df9..b9cff0ae21086b44e0c920095d5f6c9668346f38 100644 --- a/tensorflow/contrib/lite/kernels/register.h +++ b/tensorflow/contrib/lite/kernels/register.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_REGISTER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_REGISTER_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_REGISTER_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_REGISTER_H_ #include #include "tensorflow/contrib/lite/context.h" @@ -47,4 +47,4 @@ class BuiltinOpResolver : public OpResolver { } // namespace ops } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_BUILTIN_KERNELS_H +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_BUILTIN_KERNELS_H diff --git a/tensorflow/contrib/lite/kernels/resize_bilinear.cc b/tensorflow/contrib/lite/kernels/resize_bilinear.cc index 1613c9a89faa3579b913408cc09cdad7f942cb99..9e3e19c09a4012ebdadbc2a7c2ba06c4bfefd206 100644 --- a/tensorflow/contrib/lite/kernels/resize_bilinear.cc +++ b/tensorflow/contrib/lite/kernels/resize_bilinear.cc @@ -33,32 +33,44 @@ enum KernelType { }; constexpr int kInputTensor = 0; +constexpr int kSizeTensor = 1; constexpr int kOutputTensor = 0; -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - auto* params = - reinterpret_cast(node->builtin_data); +TfLiteStatus ResizeOutputTensor(TfLiteContext* context, TfLiteTensor* input, + TfLiteTensor* size, TfLiteTensor* output) { + TfLiteIntArray* output_size = TfLiteIntArrayCreate(4); + output_size->data[0] = input->dims->data[0]; + const int32* size_data = GetTensorData(size); + output_size->data[1] = size_data[0]; + output_size->data[2] = size_data[1]; + output_size->data[3] = input->dims->data[3]; + return context->ResizeTensor(context, output, output_size); +} - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); TfLiteTensor* input = GetInput(context, node, kInputTensor); + TfLiteTensor* size = GetInput(context, node, kSizeTensor); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); // TODO(ahentz): Our current implementations rely on the inputs being 4D. TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4); + TF_LITE_ENSURE_EQ(context, NumDimensions(size), 1); // TODO(ahentz): Our current implementations only support float32. - TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, input->type, output->type); - - TfLiteIntArray* output_size = TfLiteIntArrayCreate(4); - output_size->data[0] = input->dims->data[0]; - output_size->data[1] = params->new_height; - output_size->data[2] = params->new_width; - output_size->data[3] = input->dims->data[3]; - - return context->ResizeTensor(context, output, output_size); + TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32); + TF_LITE_ENSURE_EQ(context, size->type, kTfLiteInt32); + // ResizeBilinear creates a float tensor even when the input is made of + // integers. + output->type = kTfLiteFloat32; + + if (!IsConstantTensor(size)) { + SetTensorToDynamic(output); + return kTfLiteOk; + } + return ResizeOutputTensor(context, input, size, output); } template @@ -68,15 +80,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { TfLiteTensor* input = GetInput(context, node, kInputTensor); TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + TfLiteTensor* size = GetInput(context, node, kSizeTensor); - // We have to fake a tensor here, to satisfy ResizeBilinear(). - int32 output_size_data[2] = {params->new_height, params->new_width}; + if (IsDynamicTensor(output)) { + TF_LITE_ENSURE_OK(context, + ResizeOutputTensor(context, input, size, output)); + } if (output->type == kTfLiteFloat32) { -#define TF_LITE_RESIZE_BILINEAR(type) \ - type::ResizeBilinear(GetTensorData(input), GetTensorDims(input), \ - output_size_data, GetTensorDims({1, 1, 1, 2}), \ - GetTensorData(output), GetTensorDims(output)) +#define TF_LITE_RESIZE_BILINEAR(type) \ + type::ResizeBilinear(GetTensorData(input), GetTensorDims(input), \ + GetTensorData(size), GetTensorDims(size), \ + GetTensorData(output), GetTensorDims(output), \ + params->align_corners) if (kernel_type == kReference) { TF_LITE_RESIZE_BILINEAR(reference_ops); diff --git a/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc b/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc index 314a71e210d9b5ea75bb137ef228273ef48f28b5..4e03f3820a5c14ee1692c553db61e385716b1723 100644 --- a/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc +++ b/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc @@ -25,63 +25,101 @@ using ::testing::ElementsAreArray; class ResizeBilinearOpModel : public SingleOpModel { public: - ResizeBilinearOpModel(std::initializer_list input_shape, int new_height, - int new_width) { - input_ = AddInput(TensorType_FLOAT32); - output_ = AddOutput(TensorType_FLOAT32); - SetBuiltinOp( - BuiltinOperator_RESIZE_BILINEAR, BuiltinOptions_ResizeBilinearOptions, - CreateResizeBilinearOptions(builder_, new_height, new_width).Union()); - BuildInterpreter({input_shape}); + ResizeBilinearOpModel(const TensorData& input, + std::initializer_list size_data = {}) { + bool const_size = size_data.size() != 0; + input_ = AddInput(input); + if (const_size) { + size_ = AddConstInput(TensorType_INT32, size_data, {2}); + } else { + size_ = AddInput({TensorType_INT32, {2}}); + } + output_ = AddOutput(TensorType_FLOAT32); // Always float. + SetBuiltinOp(BuiltinOperator_RESIZE_BILINEAR, + BuiltinOptions_ResizeBilinearOptions, + CreateResizeBilinearOptions(builder_).Union()); + if (const_size) { + BuildInterpreter({GetShape(input_)}); + } else { + BuildInterpreter({GetShape(input_), GetShape(size_)}); + } } void SetInput(std::initializer_list data) { PopulateTensor(input_, data); } + void SetSize(std::initializer_list data) { PopulateTensor(size_, data); } std::vector GetOutput() { return ExtractVector(output_); } private: int input_; + int size_; int output_; }; TEST(ResizeBilinearOpTest, HorizontalResize) { - ResizeBilinearOpModel m({1, 1, 2, 1}, 1, 3); + ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 1, 2, 1}}); m.SetInput({3, 6}); + m.SetSize({1, 3}); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({3, 5, 6}))); + + ResizeBilinearOpModel const_m({TensorType_FLOAT32, {1, 1, 2, 1}}, {1, 3}); + const_m.SetInput({3, 6}); + const_m.Invoke(); + EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({3, 5, 6}))); } TEST(ResizeBilinearOpTest, VerticalResize) { - ResizeBilinearOpModel m({1, 2, 1, 1}, 3, 1); + ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 2, 1, 1}}); m.SetInput({3, 9}); + m.SetSize({3, 1}); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({3, 7, 9}))); + + ResizeBilinearOpModel const_m({TensorType_FLOAT32, {1, 2, 1, 1}}, {3, 1}); + const_m.SetInput({3, 9}); + const_m.Invoke(); + EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({3, 7, 9}))); } TEST(ResizeBilinearOpTest, TwoDimensionalResize) { - ResizeBilinearOpModel m({1, 2, 2, 1}, 3, 3); + ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}); m.SetInput({ 3, 6, // 9, 12 // }); + m.SetSize({3, 3}); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ 3, 5, 6, // 7, 9, 10, // 9, 11, 12, // }))); + + ResizeBilinearOpModel const_m({TensorType_FLOAT32, {1, 2, 2, 1}}, {3, 3}); + const_m.SetInput({ + 3, 6, // + 9, 12 // + }); + const_m.Invoke(); + EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ + 3, 5, 6, // + 7, 9, 10, // + 9, 11, 12, // + }))); } TEST(ResizeBilinearOpTest, TwoDimensionalResizeWithTwoBatches) { - ResizeBilinearOpModel m({2, 2, 2, 1}, 3, 3); + ResizeBilinearOpModel m({TensorType_FLOAT32, {2, 2, 2, 1}}); m.SetInput({ 3, 6, // 9, 12, // 4, 10, // 10, 16 // }); + m.SetSize({3, 3}); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ 3, 5, 6, // @@ -91,20 +129,50 @@ TEST(ResizeBilinearOpTest, TwoDimensionalResizeWithTwoBatches) { 8, 12, 14, // 10, 14, 16, // }))); + + ResizeBilinearOpModel const_m({TensorType_FLOAT32, {2, 2, 2, 1}}, {3, 3}); + const_m.SetInput({ + 3, 6, // + 9, 12, // + 4, 10, // + 10, 16 // + }); + const_m.Invoke(); + EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ + 3, 5, 6, // + 7, 9, 10, // + 9, 11, 12, // + 4, 8, 10, // + 8, 12, 14, // + 10, 14, 16, // + }))); } TEST(ResizeBilinearOpTest, ThreeDimensionalResize) { - ResizeBilinearOpModel m({1, 2, 2, 2}, 3, 3); + ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 2, 2, 2}}); m.SetInput({ 3, 4, 6, 10, // 9, 10, 12, 16, // }); + m.SetSize({3, 3}); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ 3, 4, 5, 8, 6, 10, // 7, 8, 9, 12, 10, 14, // 9, 10, 11, 14, 12, 16, // }))); + + ResizeBilinearOpModel const_m({TensorType_FLOAT32, {1, 2, 2, 2}}, {3, 3}); + const_m.SetInput({ + 3, 4, 6, 10, // + 9, 10, 12, 16, // + }); + const_m.Invoke(); + EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ + 3, 4, 5, 8, 6, 10, // + 7, 8, 9, 12, 10, 14, // + 9, 10, 11, 14, 12, 16, // + }))); } } // namespace diff --git a/tensorflow/contrib/lite/kernels/space_to_batch_nd.cc b/tensorflow/contrib/lite/kernels/space_to_batch_nd.cc new file mode 100644 index 0000000000000000000000000000000000000000..d8c9e352f00627eee45ae836b720f2af77140538 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/space_to_batch_nd.cc @@ -0,0 +1,186 @@ +/* Copyright 2017 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 +#include +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace space_to_batch_nd { + +// This file has two implementations of SpaceToBatchND. +enum KernelType { + kReference, + kGenericOptimized, +}; + +struct SpaceToBatchNDContext { + SpaceToBatchNDContext(TfLiteContext* context, TfLiteNode* node) { + input = GetInput(context, node, 0); + block_shape = GetInput(context, node, 1); + paddings = GetInput(context, node, 2); + output = GetOutput(context, node, 0); + } + TfLiteTensor* input; + TfLiteTensor* block_shape; + TfLiteTensor* paddings; + TfLiteTensor* output; +}; + +// Currently, only 4D NHWC input/output op_context are supported. +// The 4D array need to have exactly 2 spatial dimensions. +// TODO(nupurgarg): Support arbitrary dimension in SpaceToBatchND. +const int kInputDimensionNum = 4; +const int kBlockSizeDimensionNum = 1; +const int kSpatialDimensionNum = 2; + +TfLiteStatus ResizeOutputTensor(TfLiteContext* context, + SpaceToBatchNDContext* op_context) { + TfLiteIntArray* input_size = op_context->input->dims; + const int32* block_shape = GetTensorData(op_context->block_shape); + const int32* paddings_data = GetTensorData(op_context->paddings); + + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->block_shape), + kBlockSizeDimensionNum); + TF_LITE_ENSURE_EQ(context, op_context->block_shape->dims->data[0], + kSpatialDimensionNum); + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->paddings), + kSpatialDimensionNum); + + TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_size); + + // Ensures the input height and width (with padding) is a multiple of block + // shape height and width. + for (int dim = 0; dim < kSpatialDimensionNum; ++dim) { + int final_dim_size = (input_size->data[dim + 1] + paddings_data[dim * 2] + + paddings_data[dim * 2 + 1]); + TF_LITE_ENSURE_EQ(context, final_dim_size % block_shape[dim], 0); + output_size->data[dim + 1] = final_dim_size / block_shape[dim]; + } + + const int output_batch_size = + input_size->data[0] * block_shape[0] * block_shape[1]; + const int output_channel_size = input_size->data[3]; + + output_size->data[0] = output_batch_size; + output_size->data[3] = output_channel_size; + + return context->ResizeTensor(context, op_context->output, output_size); +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 3); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + SpaceToBatchNDContext op_context(context, node); + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.input), + kInputDimensionNum); + TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type); + + if (!IsConstantTensor(op_context.block_shape) || + !IsConstantTensor(op_context.paddings)) { + SetTensorToDynamic(op_context.output); + return kTfLiteOk; + } + return ResizeOutputTensor(context, &op_context); +} + +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + SpaceToBatchNDContext op_context(context, node); + + // Resize the output tensor if the output tensor is dynamic. + if (IsDynamicTensor(op_context.output)) { + TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); + } + +#define TF_LITE_SPACE_TO_BATCH_ND(type, scalar) \ + type::SpaceToBatchND(GetTensorData(op_context.input), \ + GetTensorDims(op_context.input), \ + GetTensorData(op_context.block_shape), \ + GetTensorDims(op_context.block_shape), \ + GetTensorData(op_context.paddings), \ + GetTensorDims(op_context.paddings), \ + GetTensorData(op_context.output), \ + GetTensorDims(op_context.output)) + switch (op_context.input->type) { // Already know in/out types are same. + case kTfLiteFloat32: + if (kernel_type == kReference) { + TF_LITE_SPACE_TO_BATCH_ND(reference_ops, float); + } else { + TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, float); + } + break; + case kTfLiteUInt8: + if (kernel_type == kReference) { + TF_LITE_SPACE_TO_BATCH_ND(reference_ops, uint8_t); + } else { + TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, uint8_t); + } + break; + case kTfLiteInt32: + if (kernel_type == kReference) { + TF_LITE_SPACE_TO_BATCH_ND(reference_ops, int32_t); + } else { + TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, int32_t); + } + break; + case kTfLiteInt64: + if (kernel_type == kReference) { + TF_LITE_SPACE_TO_BATCH_ND(reference_ops, int64_t); + } else { + TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, int64_t); + } + break; + default: + context->ReportError(context, + "Type is currently not supported by SpaceToBatch."); + return kTfLiteError; + } +#undef TF_LITE_SPACE_TO_BATCH_ND + return kTfLiteOk; +} + +} // namespace space_to_batch_nd + +TfLiteRegistration* Register_SPACE_TO_BATCH_ND_REF() { + static TfLiteRegistration r = { + nullptr, nullptr, space_to_batch_nd::Prepare, + space_to_batch_nd::Eval}; + return &r; +} + +TfLiteRegistration* Register_SPACE_TO_BATCH_ND_GENERIC_OPT() { + static TfLiteRegistration r = { + nullptr, nullptr, space_to_batch_nd::Prepare, + space_to_batch_nd::Eval}; + return &r; +} + +TfLiteRegistration* Register_SPACE_TO_BATCH_ND() { + // return Register_SPACE_TO_BATCH_ND_REF(); + return Register_SPACE_TO_BATCH_ND_GENERIC_OPT(); +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/space_to_batch_nd_test.cc b/tensorflow/contrib/lite/kernels/space_to_batch_nd_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..92a4a037d5873e608ee7bdbdfc5eaa5e9b62bc8c --- /dev/null +++ b/tensorflow/contrib/lite/kernels/space_to_batch_nd_test.cc @@ -0,0 +1,199 @@ +/* Copyright 2017 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 +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +class SpaceToBatchNDOpModel : public SingleOpModel { + public: + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + + void SetBlockShape(std::initializer_list data) { + PopulateTensor(block_shape_, data); + } + + void SetPaddings(std::initializer_list data) { + PopulateTensor(paddings_, data); + } + + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + protected: + int input_; + int block_shape_; + int paddings_; + int output_; +}; + +// Tests case where block_shape and paddings are const tensors. +// +// Example usage is as follows: +// SpaceToBatchNDOpConstModel m(input_shape, block_shape, paddings); +// m.SetInput(input_data); +// m.Invoke(); +class SpaceToBatchNDOpConstModel : public SpaceToBatchNDOpModel { + public: + SpaceToBatchNDOpConstModel(std::initializer_list input_shape, + std::initializer_list block_shape, + std::initializer_list paddings) { + input_ = AddInput(TensorType_FLOAT32); + block_shape_ = AddConstInput(TensorType_INT32, block_shape, {2}); + paddings_ = AddConstInput(TensorType_INT32, paddings, {2, 2}); + output_ = AddOutput(TensorType_FLOAT32); + + SetBuiltinOp(BuiltinOperator_SPACE_TO_BATCH_ND, + BuiltinOptions_SpaceToBatchNDOptions, + CreateSpaceToBatchNDOptions(builder_).Union()); + BuildInterpreter({input_shape}); + } +}; + +// Tests case where block_shape and paddings are non-const tensors. +// +// Example usage is as follows: +// SpaceToBatchNDOpDynamicModel m(input_shape); +// m.SetInput(input_data); +// m.SetBlockShape(block_shape); +// m.SetPaddings(paddings); +// m.Invoke(); +class SpaceToBatchNDOpDynamicModel : public SpaceToBatchNDOpModel { + public: + SpaceToBatchNDOpDynamicModel(std::initializer_list input_shape) { + input_ = AddInput(TensorType_FLOAT32); + block_shape_ = AddInput(TensorType_INT32); + paddings_ = AddInput(TensorType_INT32); + output_ = AddOutput(TensorType_FLOAT32); + + SetBuiltinOp(BuiltinOperator_SPACE_TO_BATCH_ND, + BuiltinOptions_SpaceToBatchNDOptions, + CreateSpaceToBatchNDOptions(builder_).Union()); + BuildInterpreter({input_shape, {2}, {2, 2}}); + } +}; + +TEST(SpaceToBatchNDOpTest, InvalidShapeTest) { + EXPECT_DEATH(SpaceToBatchNDOpConstModel({1, 3, 3, 1}, {2, 2}, {0, 0, 0, 0}), + "Cannot allocate tensors"); +} + +TEST(SpaceToBatchNDOpTest, SimpleConstTest) { + SpaceToBatchNDOpConstModel m({1, 4, 4, 1}, {2, 2}, {0, 0, 0, 0}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 2, 2, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 9, 11, 2, 4, 10, 12, 5, 7, + 13, 15, 6, 8, 14, 16})); +} + +TEST(SpaceToBatchNDOpTest, SimpleDynamicTest) { + SpaceToBatchNDOpDynamicModel m({1, 4, 4, 1}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + m.SetBlockShape({2, 2}); + m.SetPaddings({0, 0, 0, 0}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 2, 2, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 9, 11, 2, 4, 10, 12, 5, 7, + 13, 15, 6, 8, 14, 16})); +} + +TEST(SpaceToBatchNDOpTest, MultipleInputBatchesConstTest) { + SpaceToBatchNDOpConstModel m({2, 2, 4, 1}, {2, 2}, {0, 0, 0, 0}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({8, 1, 2, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 9, 11, 2, 4, 10, 12, 5, 7, + 13, 15, 6, 8, 14, 16})); +} + +TEST(SpaceToBatchNDOpTest, MultipleInputBatchesDynamicTest) { + SpaceToBatchNDOpDynamicModel m({2, 2, 4, 1}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + m.SetBlockShape({2, 2}); + m.SetPaddings({0, 0, 0, 0}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({8, 1, 2, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 9, 11, 2, 4, 10, 12, 5, 7, + 13, 15, 6, 8, 14, 16})); +} + +TEST(SpaceToBatchNDOpTest, SimplePaddingConstTest) { + SpaceToBatchNDOpConstModel m({1, 5, 2, 1}, {3, 2}, {1, 0, 2, 0}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({6, 2, 2, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({ + 0, 0, 0, 5, 0, 0, 0, 6, 0, 1, 0, 7, + 0, 2, 0, 8, 0, 3, 0, 9, 0, 4, 0, 10, + })); +} + +TEST(SpaceToBatchNDOpTest, SimplePaddingDynamicTest) { + SpaceToBatchNDOpDynamicModel m({1, 5, 2, 1}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10}); + m.SetBlockShape({3, 2}); + m.SetPaddings({1, 0, 2, 0}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({6, 2, 2, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({ + 0, 0, 0, 5, 0, 0, 0, 6, 0, 1, 0, 7, + 0, 2, 0, 8, 0, 3, 0, 9, 0, 4, 0, 10, + })); +} + +TEST(SpaceToBatchNDOpTest, ComplexPaddingConstTest) { + SpaceToBatchNDOpConstModel m({1, 4, 2, 1}, {3, 2}, {1, 1, 2, 4}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({6, 2, 4, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({ + 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, + 0, 1, 0, 0, 0, 7, 0, 0, 0, 2, 0, 0, 0, 8, 0, 0, + 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, + })); +} + +TEST(SpaceToBatchNDOpTest, ComplexPaddingDynamicTest) { + SpaceToBatchNDOpDynamicModel m({1, 4, 2, 1}); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8}); + m.SetBlockShape({3, 2}); + m.SetPaddings({1, 1, 2, 4}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({6, 2, 4, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({ + 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, + 0, 1, 0, 0, 0, 7, 0, 0, 0, 2, 0, 0, 0, 8, 0, 0, + 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, + })); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/split.cc b/tensorflow/contrib/lite/kernels/split.cc new file mode 100644 index 0000000000000000000000000000000000000000..b524c79f8779b0119781679c0af9fe354e38ad4f --- /dev/null +++ b/tensorflow/contrib/lite/kernels/split.cc @@ -0,0 +1,159 @@ +/* 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 +#include +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace split { + +struct OpContext { + OpContext(TfLiteContext* context, TfLiteNode* node) { + params = reinterpret_cast(node->builtin_data); + axis = GetInput(context, node, 0); + input = GetInput(context, node, 1); + } + TfLiteSplitParams* params; + TfLiteTensor* axis; + TfLiteTensor* input; +}; + +TfLiteStatus UseDynamicOutputTensors(TfLiteContext* context, TfLiteNode* node) { + for (int i = 0; i < NumOutputs(node); ++i) { + SetTensorToDynamic(GetOutput(context, node, i)); + } + return kTfLiteOk; +} + +TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node, + TfLiteTensor* axis, TfLiteTensor* input, + int num_splits) { + int axis_value = GetTensorData(axis)[0]; + if (axis_value < 0) { + axis_value += NumDimensions(input); + } + + const int input_size = SizeOfDimension(input, axis_value); + TF_LITE_ENSURE_MSG(context, input_size % num_splits == 0, + "Not an even split"); + const int slice_size = input_size / num_splits; + + for (int i = 0; i < NumOutputs(node); ++i) { + TfLiteIntArray* output_dims = TfLiteIntArrayCopy(input->dims); + output_dims->data[axis_value] = slice_size; + TfLiteTensor* output = GetOutput(context, node, i); + TF_LITE_ENSURE_STATUS(context->ResizeTensor(context, output, output_dims)); + } + + return kTfLiteOk; +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + + OpContext op_context(context, node); + + TF_LITE_ENSURE_EQ(context, NumOutputs(node), op_context.params->num_splits); + + auto input_type = op_context.input->type; + TF_LITE_ENSURE(context, + input_type == kTfLiteFloat32 || input_type == kTfLiteUInt8); + for (int i = 0; i < NumOutputs(node); ++i) { + GetOutput(context, node, i)->type = input_type; + } + + // If we know the contents of the 'axis' tensor, resize all outputs. + // Otherwise, wait until Eval(). + if (IsConstantTensor(op_context.axis)) { + return ResizeOutputTensors(context, node, op_context.axis, op_context.input, + op_context.params->num_splits); + } else { + return UseDynamicOutputTensors(context, node); + } +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + OpContext op_context(context, node); + + // When the 'axis' tensor is non-const we can't resize output tensors in + // Prepare(), and we have to do it now. + if (!IsConstantTensor(op_context.axis)) { + TF_LITE_ENSURE_OK( + context, + ResizeOutputTensors(context, node, op_context.axis, op_context.input, + op_context.params->num_splits)); + } + + int axis_value = GetTensorData(op_context.axis)[0]; + if (axis_value < 0) { + axis_value += NumDimensions(op_context.input); + } + axis_value = RemapDim(NumDimensions(op_context.input), axis_value); + + // TODO(ahentz): Our usage of VectorOfTensors could be optimized by + // calculating it in Prepare, unless we defer shape calculation. + // TODO(ahentz): We can improve the optimized_ops version to handle other + // cases too. +#define TF_LITE_SPLIT(scalar) \ + VectorOfTensors all_outputs(*context, *node->outputs); \ + if (axis_value == NumDimensions(op_context.input)) { \ + optimized_ops::TensorFlowSplit( \ + GetTensorData(op_context.input), \ + GetTensorDims(op_context.input), NumOutputs(node), all_outputs.data(), \ + all_outputs.dims()); \ + } else { \ + reference_ops::TensorFlowSplit( \ + GetTensorData(op_context.input), \ + GetTensorDims(op_context.input), axis_value, NumOutputs(node), \ + all_outputs.data(), all_outputs.dims()); \ + } + switch (op_context.input->type) { + case kTfLiteFloat32: { + TF_LITE_SPLIT(float); + break; + } + case kTfLiteUInt8: { + TF_LITE_SPLIT(uint8_t); + break; + } + default: + context->ReportError(context, + "Only float32 and uint8 are currently supported."); + return kTfLiteError; + } +#undef TF_LITE_SPLIT + + return kTfLiteOk; +} + +} // namespace split + +TfLiteRegistration* Register_SPLIT() { + static TfLiteRegistration r = {nullptr, nullptr, split::Prepare, split::Eval}; + return &r; +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/split_test.cc b/tensorflow/contrib/lite/kernels/split_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..61a0759c6475795c06a9b55d3586d2b818f298b2 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/split_test.cc @@ -0,0 +1,147 @@ +/* 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 +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +constexpr int kAxisIsATensor = -1000; + +class SplitOpModel : public SingleOpModel { + public: + SplitOpModel(const TensorData& input, int num_splits, + int axis = kAxisIsATensor) { + if (axis == kAxisIsATensor) { + axis_ = AddInput({TensorType_INT32, {1}}); + } else { + axis_ = AddConstInput(TensorType_INT32, {axis}, {1}); + } + input_ = AddInput(input); + for (int i = 0; i < num_splits; ++i) { + outputs_.push_back(AddOutput(input.type)); + } + SetBuiltinOp(BuiltinOperator_SPLIT, BuiltinOptions_SplitOptions, + CreateSplitOptions(builder_, num_splits).Union()); + if (axis == kAxisIsATensor) { + BuildInterpreter({GetShape(axis_), GetShape(input_)}); + } else { + BuildInterpreter({{}, GetShape(input_)}); + } + } + + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + void SetAxis(int axis) { PopulateTensor(axis_, {axis}); } + + std::vector GetOutput(int i) { + return ExtractVector(outputs_[i]); + } + std::vector GetOutputShape(int i) { return GetTensorShape(outputs_[i]); } + + private: + int input_; + int axis_; + std::vector outputs_; +}; + +using TensorValues = std::initializer_list; + +void Check(int axis, int num_splits, std::initializer_list input_shape, + std::initializer_list output_shape, + const TensorValues& input_data, + const std::vector& output_data) { + auto debug = [&](int i) { + std::stringstream ss; + ss << "for output tensor " << i << " axis=" << axis + << " and num_splits=" << num_splits; + return ss.str(); + }; + SplitOpModel m({TensorType_FLOAT32, input_shape}, num_splits); + m.SetInput(input_data); + m.SetAxis(axis); + m.Invoke(); + for (int i = 0; i < num_splits; ++i) { + EXPECT_THAT(m.GetOutput(i), ElementsAreArray(output_data[i])) << debug(i); + EXPECT_THAT(m.GetOutputShape(i), ElementsAreArray(output_shape)) + << debug(i); + } + + SplitOpModel const_m({TensorType_FLOAT32, input_shape}, num_splits, axis); + const_m.SetInput(input_data); + const_m.Invoke(); + for (int i = 0; i < num_splits; ++i) { + EXPECT_THAT(const_m.GetOutput(i), ElementsAreArray(output_data[i])) + << debug(i); + EXPECT_THAT(const_m.GetOutputShape(i), ElementsAreArray(output_shape)) + << debug(i); + } +} + +TEST(SplitOpTest, FourDimensional) { + Check(/*axis=*/0, /*num_splits=*/2, {2, 2, 2, 2}, {1, 2, 2, 2}, + {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + { + {1, 2, 3, 4, 5, 6, 7, 8}, + {9, 10, 11, 12, 13, 14, 15, 16}, + }); + Check(/*axis=*/1, /*num_splits=*/2, {2, 2, 2, 2}, {2, 1, 2, 2}, + {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + { + {1, 2, 3, 4, 9, 10, 11, 12}, + {5, 6, 7, 8, 13, 14, 15, 16}, + }); + Check(/*axis=*/2, /*num_splits=*/2, {2, 2, 2, 2}, {2, 2, 1, 2}, + {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + { + {1, 2, 5, 6, 9, 10, 13, 14}, + {3, 4, 7, 8, 11, 12, 15, 16}, + }); + Check(/*axis=*/3, /*num_splits=*/2, {2, 2, 2, 2}, {2, 2, 2, 1}, + {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + { + {1, 3, 5, 7, 9, 11, 13, 15}, + {2, 4, 6, 8, 10, 12, 14, 16}, + }); +} + +TEST(SplitOpTest, OneDimensional) { + Check(/*axis=*/0, /*num_splits=*/8, {8}, {1}, {1, 2, 3, 4, 5, 6, 7, 8}, + {{1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}}); +} + +TEST(SplitOpTest, NegativeAxis) { + Check(/*axis=*/-4, /*num_splits=*/2, {2, 2, 2, 2}, {1, 2, 2, 2}, + {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + { + {1, 2, 3, 4, 5, 6, 7, 8}, + {9, 10, 11, 12, 13, 14, 15, 16}, + }); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/squeeze.cc b/tensorflow/contrib/lite/kernels/squeeze.cc new file mode 100644 index 0000000000000000000000000000000000000000..29447ab021c7b68ff51070d35262402e08dc7ab9 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/squeeze.cc @@ -0,0 +1,99 @@ +/* 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 +#include +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace squeeze { + +struct SqueezeContext { + SqueezeContext(TfLiteContext* context, TfLiteNode* node) { + params = reinterpret_cast(node->builtin_data); + input = GetInput(context, node, 0); + output = GetOutput(context, node, 0); + } + TfLiteSqueezeParams* params; + TfLiteTensor* input; + TfLiteTensor* output; +}; + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + SqueezeContext op_context(context, node); + int input_num_dims = NumDimensions(op_context.input); + int num_squeeze_dims = op_context.params->num_squeeze_dims; + + // Determines number of dimensions of output tensor after squeeze. + const TfLiteIntArray* input_dims = op_context.input->dims; + const int* squeeze_dims = op_context.params->squeeze_dims; + TF_LITE_ENSURE(context, input_num_dims <= 8); + bool should_squeeze[8] = {false}; + int num_squeezed_dims = 0; + if (num_squeeze_dims == 0) { + for (int idx = 0; idx < input_num_dims; ++idx) { + if (input_dims->data[idx] == 1) { + should_squeeze[idx] = true; + ++num_squeezed_dims; + } + } + } else { + for (int idx = 0; idx < num_squeeze_dims; ++idx) { + int current = squeeze_dims[idx] < 0 ? squeeze_dims[idx] + input_num_dims + : squeeze_dims[idx]; + TF_LITE_ENSURE(context, current >= 0 && current < input_num_dims && + input_dims->data[current] == 1); + if (!should_squeeze[current]) ++num_squeezed_dims; + should_squeeze[current] = true; + } + } + // Sets output dimensions. + TfLiteIntArray* output_dims = + TfLiteIntArrayCreate(input_num_dims - num_squeezed_dims); + for (int in_idx = 0, out_idx = 0; in_idx < input_num_dims; ++in_idx) { + if (!should_squeeze[in_idx]) { + output_dims->data[out_idx++] = input_dims->data[in_idx]; + } + } + return context->ResizeTensor(context, op_context.output, output_dims); +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + SqueezeContext op_context(context, node); + TF_LITE_ENSURE_EQ(context, op_context.input->bytes, op_context.output->bytes); + memcpy(op_context.output->data.raw, op_context.input->data.raw, + op_context.input->bytes); + return kTfLiteOk; +} + +} // namespace squeeze + +TfLiteRegistration* Register_SQUEEZE() { + static TfLiteRegistration r = {nullptr, nullptr, squeeze::Prepare, + squeeze::Eval}; + return &r; +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/squeeze_test.cc b/tensorflow/contrib/lite/kernels/squeeze_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..a8aab88357cacbb72784a4bc6e860aeb47783eb3 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/squeeze_test.cc @@ -0,0 +1,124 @@ +/* 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 +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; +using ::testing::IsEmpty; + +class BaseSqueezeOpModel : public SingleOpModel { + public: + BaseSqueezeOpModel(const TensorData& input, const TensorData& output, + std::initializer_list axis) { + input_ = AddInput(input); + output_ = AddOutput(output); + SetBuiltinOp( + BuiltinOperator_SQUEEZE, BuiltinOptions_SqueezeOptions, + CreateSqueezeOptions(builder_, builder_.CreateVector(axis)) + .Union()); + BuildInterpreter({GetShape(input_)}); + } + + int input() { return input_; } + + protected: + int input_; + int output_; +}; + +class FloatSqueezeOpModel : public BaseSqueezeOpModel { + public: + using BaseSqueezeOpModel::BaseSqueezeOpModel; + + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } +}; + +TEST(FloatSqueezeOpTest, SqueezeAll) { + std::initializer_list data = { + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, + 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0}; + FloatSqueezeOpModel m({TensorType_FLOAT32, {1, 24, 1}}, + {TensorType_FLOAT32, {24}}, {}); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({24})); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray({1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, + 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, + 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0})); +} + +TEST(FloatSqueezeOpTest, SqueezeSelectedAxis) { + std::initializer_list data = { + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, + 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0}; + FloatSqueezeOpModel m({TensorType_FLOAT32, {1, 24, 1}}, + {TensorType_FLOAT32, {24}}, {2}); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 24})); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray({1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, + 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, + 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0})); +} + +TEST(FloatSqueezeOpTest, SqueezeNegativeAxis) { + std::initializer_list data = { + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, + 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0}; + FloatSqueezeOpModel m({TensorType_FLOAT32, {1, 24, 1}}, + {TensorType_FLOAT32, {24}}, {-1, 0}); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({24})); + EXPECT_THAT( + m.GetOutput(), + ElementsAreArray({1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, + 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, + 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0})); +} + +TEST(FloatSqueezeOpTest, SqueezeAllDims) { + std::initializer_list data = {3.85}; + FloatSqueezeOpModel m({TensorType_FLOAT32, {1, 1, 1, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {1}}, {}); + m.SetInput(data); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), IsEmpty()); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({3.85})); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/strided_slice.cc b/tensorflow/contrib/lite/kernels/strided_slice.cc new file mode 100644 index 0000000000000000000000000000000000000000..fb1e11e0ca00abb36d7f29d562711a7bbcbeca1c --- /dev/null +++ b/tensorflow/contrib/lite/kernels/strided_slice.cc @@ -0,0 +1,259 @@ +/* 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 +#include +#include +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace strided_slice { + +enum KernelType { + kReference, + // TODO(soroosh): add kGenericOptimized +}; + +constexpr int kInputTensor = 0; +constexpr int kBeginTensor = 1; +constexpr int kEndTensor = 2; +constexpr int kStridesTensor = 3; +constexpr int kOutputTensor = 0; + +struct StridedSliceContext { + StridedSliceContext(TfLiteContext* context, TfLiteNode* node) { + params = reinterpret_cast(node->builtin_data); + input = GetInput(context, node, kInputTensor); + begin = GetInput(context, node, kBeginTensor); + end = GetInput(context, node, kEndTensor); + strides = GetInput(context, node, kStridesTensor); + output = GetOutput(context, node, kOutputTensor); + dims = NumDimensions(input); + } + TfLiteStridedSliceParams* params; + TfLiteTensor* input; + TfLiteTensor* begin; + TfLiteTensor* end; + TfLiteTensor* strides; + TfLiteTensor* output; + int dims; +}; + +// Reverse order of bits in the mask to match the expected order in kernel +inline int ReverseMaskBits(int mask, int num_dimensions) { + int out = 0; + for (int dim = 0; dim < num_dimensions; dim++) { + out <<= 1; + out += (mask & 1); + mask >>= 1; + } + return out; +} + +// This Op only supports 1-4D cases and since we use the reference 4D +// implementation, the 1-3D tensors are mapped to 4D. +const int kMaxDim = 4; + +inline int32_t PositiveRemainder(int32_t dividend, int32_t divisor) { + return (divisor + (dividend % divisor)) % divisor; +} + +inline int32_t ClampedIndex(int32_t index, int dim, bool pos_stride) { + return pos_stride + ? (index >= dim ? dim + : PositiveRemainder( + std::min(std::max(index, -dim), dim), dim)) + : (index < -dim + ? -1 + : PositiveRemainder( + std::min(std::max(index, -dim), dim - 1), dim)); +} + +inline int32_t GetBeginValueAtIndex(StridedSliceContext* op_context, int idx) { + const int dim = op_context->input->dims->data[idx]; + const bool pos_stride = GetTensorData(op_context->strides)[idx] > 0; + return op_context->params->begin_mask & (1 << idx) + ? pos_stride ? 0 : dim - 1 + : ClampedIndex(GetTensorData(op_context->begin)[idx], dim, + pos_stride); +} + +inline int32_t GetEndValueAtIndex(StridedSliceContext* op_context, int idx) { + const int dim = op_context->input->dims->data[idx]; + const bool pos_stride = GetTensorData(op_context->strides)[idx] > 0; + return op_context->params->end_mask & (1 << idx) + ? pos_stride ? dim : -1 + : ClampedIndex(GetTensorData(op_context->end)[idx], dim, + pos_stride); +} + +// Processes the indexing tensors (begin, end and strides) to resize the +// output tensor. This function is callable from both Prepare() and Eval() as +// long as the caller ensures the indexing tensors are present. +TfLiteStatus ResizeOutputTensor(TfLiteContext* context, + StridedSliceContext* op_context) { + std::vector output_shape_vector; + + for (int idx = op_context->dims - 1; idx >= 0; --idx) { + int32_t stride = GetTensorData(op_context->strides)[idx]; + TF_LITE_ENSURE_MSG(context, stride != 0, "stride value has to be non-zero"); + + int32_t begin = GetBeginValueAtIndex(op_context, idx); + int32_t end = GetEndValueAtIndex(op_context, idx); + + // This is valid for both positive and negative strides + int32_t dim_shape = ceil((end - begin) / static_cast(stride)); + dim_shape = dim_shape < 0 ? 0 : dim_shape; + if (!(op_context->params->shrink_axis_mask & (1 << idx))) { + output_shape_vector.push_back(dim_shape); + } + } + + TfLiteIntArray* output_shape = + TfLiteIntArrayCreate(output_shape_vector.size()); + + std::reverse_copy(output_shape_vector.begin(), output_shape_vector.end(), + output_shape->data); + + TF_LITE_ENSURE_STATUS( + context->ResizeTensor(context, op_context->output, output_shape)); + + return kTfLiteOk; +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 4); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + StridedSliceContext op_context(context, node); + + // Ensure validity of input tensor and its dimension + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.begin), 1); + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.end), 1); + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.strides), 1); + TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type); + // Only INT32 begin/end/strides are supported + // TODO(soroosh) add support for INT64 + TF_LITE_ENSURE_EQ(context, op_context.begin->type, kTfLiteInt32); + TF_LITE_ENSURE_EQ(context, op_context.end->type, kTfLiteInt32); + TF_LITE_ENSURE_EQ(context, op_context.strides->type, kTfLiteInt32); + TF_LITE_ENSURE_MSG(context, op_context.dims <= 4, + "StridedSlice op only supports 1D-4D input arrays."); + + // TODO(soroosh): add the following missing functionalities + TF_LITE_ENSURE_MSG(context, op_context.params->ellipsis_mask == 0, + "ellipsis_mask is not implemented yet."); + TF_LITE_ENSURE_MSG(context, op_context.params->new_axis_mask == 0, + "new_axis_mask is not implemented yet."); + + // Postpone allocation of output if any of the indexing tensors is not + // constant + if (!(IsConstantTensor(op_context.begin) && + IsConstantTensor(op_context.end) && + IsConstantTensor(op_context.strides))) { + SetTensorToDynamic(op_context.output); + return kTfLiteOk; + } + return ResizeOutputTensor(context, &op_context); +} + +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + StridedSliceContext op_context(context, node); + + if (IsDynamicTensor(op_context.output)) { + TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); + } + + std::vector starts; + std::vector stops; + std::vector strides; + + for (int idx = op_context.dims - 1; idx >= 0; --idx) { + starts.emplace_back(GetBeginValueAtIndex(&op_context, idx)); + stops.emplace_back(GetEndValueAtIndex(&op_context, idx)); + strides.emplace_back(GetTensorData(op_context.strides)[idx]); + } + + for (int i = op_context.dims; i < kMaxDim; i++) { + starts.emplace_back(0); + stops.emplace_back(1); + strides.emplace_back(1); + } + + op_context.params->begin_mask = + ReverseMaskBits(op_context.params->begin_mask, op_context.dims); + op_context.params->end_mask = + ReverseMaskBits(op_context.params->end_mask, op_context.dims); + op_context.params->shrink_axis_mask = + ReverseMaskBits(op_context.params->shrink_axis_mask, op_context.dims); + +#define TF_LITE_STRIDED_SLICE(kernel_type, data_type) \ + kernel_type::StridedSlice( \ + GetTensorData(op_context.input), \ + GetTensorDims(op_context.input), op_context.params->begin_mask, \ + op_context.params->end_mask, op_context.params->shrink_axis_mask, \ + starts, stops, strides, GetTensorData(op_context.output), \ + GetTensorDims(op_context.output)) + + switch (op_context.input->type) { + case kTfLiteFloat32: + if (kernel_type == kReference) { + TF_LITE_STRIDED_SLICE(reference_ops, float); + } + break; + case kTfLiteInt32: + if (kernel_type == kReference) { + TF_LITE_STRIDED_SLICE(reference_ops, int32_t); + } + break; + case kTfLiteInt64: + if (kernel_type == kReference) { + TF_LITE_STRIDED_SLICE(reference_ops, int64_t); + } + break; + default: + context->ReportError(context, + "Type is currently not supported " + "by StridedSlice."); + return kTfLiteError; + } +#undef TF_LITE_STRIDED_SLICE + return kTfLiteOk; +} + +} // namespace strided_slice + +TfLiteRegistration* Register_STRIDED_SLICE_REF() { + static TfLiteRegistration r = { + nullptr, nullptr, strided_slice::Prepare, + strided_slice::Eval}; + return &r; +} + +// TODO(soroosh): add optimized +TfLiteRegistration* Register_STRIDED_SLICE() { + return Register_STRIDED_SLICE_REF(); +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/strided_slice_test.cc b/tensorflow/contrib/lite/kernels/strided_slice_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..5cac04b38364958c5b0794c21742e8b592372ae9 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/strided_slice_test.cc @@ -0,0 +1,532 @@ +/* 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 +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::int32; +using ::testing::ElementsAreArray; + +class StridedSliceOpModel : public SingleOpModel { + public: + StridedSliceOpModel(std::initializer_list input_shape, + std::initializer_list begin_shape, + std::initializer_list end_shape, + std::initializer_list strides_shape, int begin_mask, + int end_mask, int ellipsis_mask, int new_axis_mask, + int shrink_axis_mask) { + input_ = AddInput(TensorType_FLOAT32); + begin_ = AddInput(TensorType_INT32); + end_ = AddInput(TensorType_INT32); + strides_ = AddInput(TensorType_INT32); + output_ = AddOutput(TensorType_FLOAT32); + SetBuiltinOp( + BuiltinOperator_STRIDED_SLICE, BuiltinOptions_StridedSliceOptions, + CreateStridedSliceOptions(builder_, begin_mask, end_mask, ellipsis_mask, + new_axis_mask, shrink_axis_mask) + .Union()); + BuildInterpreter({input_shape, begin_shape, end_shape, strides_shape}); + } + + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + void SetBegin(std::initializer_list data) { + PopulateTensor(begin_, data); + } + void SetEnd(std::initializer_list data) { + PopulateTensor(end_, data); + } + void SetStrides(std::initializer_list data) { + PopulateTensor(strides_, data); + } + + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + private: + int input_; + int begin_; + int end_; + int strides_; + int output_; +}; + +TEST(StridedSliceOpTest, UnsupportedInputSize) { + EXPECT_DEATH( + StridedSliceOpModel({2, 2, 2, 2, 2}, {5}, {5}, {5}, 0, 0, 0, 0, 0), + "StridedSlice op only supports 1D-4D input arrays."); +} + +TEST(StridedSliceOpTest, UnssupportedArgs) { + EXPECT_DEATH(StridedSliceOpModel({3, 2}, {2}, {2}, {2}, 0, 0, 1, 0, 0), + "ellipsis_mask is not implemented yet."); + EXPECT_DEATH(StridedSliceOpModel({3, 2}, {2}, {2}, {2}, 0, 0, 0, 1, 0), + "new_axis_mask is not implemented yet."); +} + +TEST(StridedSliceOpTest, In1D) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({1}); + m.SetEnd({3}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({2, 3})); +} + +TEST(StridedSliceOpTest, In1D_EmptyOutput) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({10}); + m.SetEnd({3}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({0})); +} + +TEST(StridedSliceOpTest, In1D_NegativeBegin) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({-3}); + m.SetEnd({3}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({2, 3})); +} + +TEST(StridedSliceOpTest, In1D_OutOfRangeBegin) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({-5}); + m.SetEnd({3}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3})); +} + +TEST(StridedSliceOpTest, In1D_NegativeEnd) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({1}); + m.SetEnd({-2}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({2})); +} + +TEST(StridedSliceOpTest, In1D_OutOfRangeEnd) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({-3}); + m.SetEnd({5}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({2, 3, 4})); +} + +TEST(StridedSliceOpTest, In1D_BeginMask) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 1, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({1}); + m.SetEnd({3}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3})); +} + +TEST(StridedSliceOpTest, In1D_NegativeBeginNegativeStride) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({-2}); + m.SetEnd({-3}); + m.SetStrides({-1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({3})); +} + +TEST(StridedSliceOpTest, In1D_OutOfRangeBeginNegativeStride) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({5}); + m.SetEnd({2}); + m.SetStrides({-1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({4})); +} + +TEST(StridedSliceOpTest, In1D_NegativeEndNegativeStride) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({2}); + m.SetEnd({-4}); + m.SetStrides({-1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 2})); +} + +TEST(StridedSliceOpTest, In1D_OutOfRangeEndNegativeStride) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({-3}); + m.SetEnd({-5}); + m.SetStrides({-1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({2, 1})); +} + +TEST(StridedSliceOpTest, In1D_EndMask) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 1, 0, 0, 0); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({1}); + m.SetEnd({3}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({2, 3, 4})); +} + +TEST(StridedSliceOpTest, In1D_NegStride) { + StridedSliceOpModel m({3}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3}); + m.SetBegin({-1}); + m.SetEnd({-4}); + m.SetStrides({-1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 2, 1})); +} + +TEST(StridedSliceOpTest, In1D_EvenLenStride2) { + StridedSliceOpModel m({2}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2}); + m.SetBegin({0}); + m.SetEnd({2}); + m.SetStrides({2}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1})); +} + +TEST(StridedSliceOpTest, In1D_OddLenStride2) { + StridedSliceOpModel m({3}, {1}, {1}, {1}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3}); + m.SetBegin({0}); + m.SetEnd({3}); + m.SetStrides({2}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3})); +} + +TEST(StridedSliceOpTest, In2D_Identity) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({0, 0}); + m.SetEnd({2, 3}); + m.SetStrides({1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6})); +} + +TEST(StridedSliceOpTest, In2D) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({1, 0}); + m.SetEnd({2, 2}); + m.SetStrides({1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({4, 5})); +} + +TEST(StridedSliceOpTest, In2D_Stride2) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({0, 0}); + m.SetEnd({2, 3}); + m.SetStrides({2, 2}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3})); +} + +TEST(StridedSliceOpTest, In2D_NegStride) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({1, -1}); + m.SetEnd({2, -4}); + m.SetStrides({2, -1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({6, 5, 4})); +} + +TEST(StridedSliceOpTest, In2D_BeginMask) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 1, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({1, 0}); + m.SetEnd({2, 2}); + m.SetStrides({1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 4, 5})); +} + +TEST(StridedSliceOpTest, In2D_EndMask) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 2, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({1, 0}); + m.SetEnd({2, 2}); + m.SetStrides({1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({4, 5, 6})); +} + +TEST(StridedSliceOpTest, In2D_NegStrideBeginMask) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 2, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({1, -2}); + m.SetEnd({2, -4}); + m.SetStrides({1, -1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({6, 5, 4})); +} + +TEST(StridedSliceOpTest, In2D_NegStrideEndMask) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 2, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({1, -2}); + m.SetEnd({2, -3}); + m.SetStrides({1, -1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 4})); +} + +TEST(StridedSliceOpTest, In3D_Identity) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({1, 1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3, 2})); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); +} + +TEST(StridedSliceOpTest, In3D_NegStride) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({-1, -1, -1}); + m.SetEnd({-3, -4, -3}); + m.SetStrides({-1, -1, -1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3, 2})); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray({12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1})); +} + +TEST(StridedSliceOpTest, In3D_Strided2) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 0); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({2, 2, 2}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 5})); +} + +TEST(StridedSliceOpTest, In1D_ShrinkAxisMask1) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 1); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({1}); + m.SetEnd({3}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_TRUE(m.GetOutputShape().empty()); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({2})); +} + +TEST(StridedSliceOpTest, In1D_EmptyOutputShrinkAxisMask1) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 1); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({2}); + m.SetEnd({1}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_TRUE(m.GetOutputShape().empty()); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({3})); +} + +TEST(StridedSliceOpTest, In1D_BeginMaskShrinkAxisMask1) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 1, 0, 0, 0, 1); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({1}); + m.SetEnd({3}); + m.SetStrides({1}); + m.Invoke(); + EXPECT_TRUE(m.GetOutputShape().empty()); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1})); +} + +TEST(StridedSliceOpTest, In1D_NegativeBeginNegativeStrideShrinkAxisMask1) { + StridedSliceOpModel m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 1); + m.SetInput({1, 2, 3, 4}); + m.SetBegin({-2}); + m.SetEnd({-3}); + m.SetStrides({-1}); + m.Invoke(); + EXPECT_TRUE(m.GetOutputShape().empty()); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({3})); +} + +TEST(StridedSliceOpTest, In2D_ShrinkAxisMask1) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 1); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({0, 0}); + m.SetEnd({2, 3}); + m.SetStrides({1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3})); +} + +TEST(StridedSliceOpTest, In2D_ShrinkAxisMask2) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 2); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({0, 0}); + m.SetEnd({2, 3}); + m.SetStrides({1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 4})); +} + +TEST(StridedSliceOpTest, In2D_ShrinkAxisMask3) { + StridedSliceOpModel m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 3); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetBegin({0, 0}); + m.SetEnd({2, 3}); + m.SetStrides({1, 1}); + m.Invoke(); + EXPECT_TRUE(m.GetOutputShape().empty()); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1})); +} + +TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis1) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 1); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({1, 1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6})); +} + +TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis2) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 2); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({1, 1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 7, 8})); +} + +TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis3) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 3); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({1, 1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2})); +} + +TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis4) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 4); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({1, 1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 5, 7, 9, 11})); +} + +TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis5) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 5); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({1, 1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 5})); +} + +TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis6) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 6); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({1, 1, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 7})); +} + +TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis7) { + StridedSliceOpModel m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 7); + m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + m.SetBegin({0, 0, 0}); + m.SetEnd({2, 3, 2}); + m.SetStrides({1, 1, 1}); + m.Invoke(); + EXPECT_TRUE(m.GetOutputShape().empty()); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1})); +} +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/sub.cc b/tensorflow/contrib/lite/kernels/sub.cc new file mode 100644 index 0000000000000000000000000000000000000000..ddaf498d5bac0109429224e7cf66cb3debcabc22 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/sub.cc @@ -0,0 +1,129 @@ +/* Copyright 2017 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/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace sub { + +// This file has three implementation of Div. +enum KernelType { + kReference, + kGenericOptimized, // Neon-free + kNeonOptimized, +}; + +constexpr int kInputTensor1 = 0; +constexpr int kInputTensor2 = 1; +constexpr int kOutputTensor = 0; + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); + TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + TF_LITE_ENSURE_EQ(context, NumDimensions(input1), NumDimensions(input2)); + for (int i = 0; i < NumDimensions(input1); ++i) { + TF_LITE_ENSURE_EQ(context, SizeOfDimension(input1, i), + SizeOfDimension(input2, i)); + } + + TF_LITE_ENSURE_EQ(context, input1->type, output->type); + TF_LITE_ENSURE_EQ(context, input2->type, output->type); + + TfLiteIntArray* output_size = TfLiteIntArrayCopy(input1->dims); + return context->ResizeTensor(context, output, output_size); +} + +template +void EvalSubFloat(TfLiteContext* context, TfLiteNode* node, + TfLiteSubParams* params, TfLiteTensor* input1, + TfLiteTensor* input2, TfLiteTensor* output) { + float output_activation_min, output_activation_max; + CalculateActivationRangeFloat(params->activation, &output_activation_min, + &output_activation_max); +#define TF_LITE_Sub(type) \ + type::Sub(GetTensorData(input1), GetTensorDims(input1), \ + GetTensorData(input2), GetTensorDims(input2), \ + output_activation_min, output_activation_max, \ + GetTensorData(output), GetTensorDims(output)) + if (kernel_type == kReference) { + TF_LITE_Sub(reference_ops); + } else { + TF_LITE_Sub(optimized_ops); + } +#undef TF_LITE_Sub +} + +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + auto* params = reinterpret_cast(node->builtin_data); + + TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); + TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + if (output->type == kTfLiteFloat32) { + EvalSubFloat(context, node, params, input1, input2, output); + } else { + context->ReportError(context, "Inputs and outputs not all float types."); + return kTfLiteError; + } + + return kTfLiteOk; +} + +} // namespace sub + +TfLiteRegistration* Register_SUB_REF() { + static TfLiteRegistration r = {nullptr, nullptr, sub::Prepare, + sub::Eval}; + return &r; +} + +TfLiteRegistration* Register_SUB_GENERIC_OPT() { + static TfLiteRegistration r = {nullptr, nullptr, sub::Prepare, + sub::Eval}; + return &r; +} + +TfLiteRegistration* Register_SUB_NEON_OPT() { + static TfLiteRegistration r = {nullptr, nullptr, sub::Prepare, + sub::Eval}; + return &r; +} + +TfLiteRegistration* Register_SUB() { +#ifdef USE_NEON + return Register_SUB_NEON_OPT(); +#else + return Register_SUB_GENERIC_OPT(); +#endif +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/svdf.cc b/tensorflow/contrib/lite/kernels/svdf.cc index 72f705fe4242b01c1516c99d3500484e8729fd9a..c69755447d5093e25d408eb6dea80750937465e7 100644 --- a/tensorflow/contrib/lite/kernels/svdf.cc +++ b/tensorflow/contrib/lite/kernels/svdf.cc @@ -15,8 +15,8 @@ limitations under the License. #include #include #include -#include #include +#include #include #include diff --git a/tensorflow/contrib/lite/kernels/svdf_test.cc b/tensorflow/contrib/lite/kernels/svdf_test.cc index 4de2ceaf053df31a4bc857fb250db416c071e80f..0f166dc69b95f3459388135b3a6c4d9b73a31cb4 100644 --- a/tensorflow/contrib/lite/kernels/svdf_test.cc +++ b/tensorflow/contrib/lite/kernels/svdf_test.cc @@ -14,8 +14,8 @@ limitations under the License. ==============================================================================*/ // Unit test for TFLite SVDF op. -#include #include +#include #include #include diff --git a/tensorflow/contrib/lite/kernels/test_util.cc b/tensorflow/contrib/lite/kernels/test_util.cc index b69f2b3e4bc66c94fdfc7ed4c244151be63a1711..373310bd87370a670a847cf5328633956028a850 100644 --- a/tensorflow/contrib/lite/kernels/test_util.cc +++ b/tensorflow/contrib/lite/kernels/test_util.cc @@ -49,7 +49,7 @@ std::vector> ArrayFloatNear(const std::vector& values, return matchers; } -int SingleOpModel::AddTensor(TensorData t) { +int SingleOpModel::AddTensor(TensorData t, std::initializer_list data) { int id = tensors_.size(); // This is slightly different depending on whether we are adding a @@ -78,8 +78,23 @@ int SingleOpModel::AddTensor(TensorData t) { builder_.CreateVector({t.zero_point})); } - tensors_.push_back(CreateTensor(builder_, builder_.CreateVector({}), - t.type, /*buffer=*/0, + int buffer_id = 0; + if (data.size()) { + // Initialize buffers list with empty buffer to allow for non-const tensors. + if (buffers_.empty()) { + buffers_.push_back(CreateBuffer(builder_, builder_.CreateVector({}))); + } + + // Add data as a Buffer to buffers list. + buffer_id = buffers_.size(); + auto data_buffer = + builder_.CreateVector(reinterpret_cast(data.begin()), + sizeof(int) * data.size()); + buffers_.push_back(CreateBuffer(builder_, data_buffer)); + } + + tensors_.push_back(CreateTensor(builder_, builder_.CreateVector(t.shape), + t.type, /*buffer=*/buffer_id, /*name=*/0, q_params)); tensor_data_[id] = t; @@ -88,7 +103,15 @@ int SingleOpModel::AddTensor(TensorData t) { } int SingleOpModel::AddInput(const TensorData& t) { - int id = AddTensor(t); + int id = AddTensor(t, {}); + inputs_.push_back(id); + return id; +} + +int SingleOpModel::AddConstInput(TensorType type, + std::initializer_list data, + std::initializer_list shape) { + int id = AddTensor(TensorData{type, shape}, data); inputs_.push_back(id); return id; } @@ -100,7 +123,7 @@ int SingleOpModel::AddNullInput() { } int SingleOpModel::AddOutput(const TensorData& t) { - int id = AddTensor(t); + int id = AddTensor(t, {}); outputs_.push_back(id); return id; } @@ -142,19 +165,21 @@ void SingleOpModel::BuildInterpreter( subgraphs.push_back(subgraph); auto subgraphs_flatbuffer = builder_.CreateVector(subgraphs); - std::vector> buffers_vec; - auto buffers = builder_.CreateVector(buffers_vec); + auto buffers = builder_.CreateVector(buffers_); auto description = builder_.CreateString("programmatic model"); builder_.Finish(CreateModel(builder_, TFLITE_SCHEMA_VERSION, opcodes, subgraphs_flatbuffer, description, buffers)); auto* model = GetModel(builder_.GetBufferPointer()); - ops::builtin::BuiltinOpResolver builtins; - for (const auto& reg : custom_registrations_) { - builtins.AddCustom(reg.first.data(), reg.second()); + if (!resolver_) { + auto resolver = new ops::builtin::BuiltinOpResolver(); + for (const auto& reg : custom_registrations_) { + resolver->AddCustom(reg.first.data(), reg.second()); + } + resolver_ = std::unique_ptr(resolver); } - InterpreterBuilder(model, builtins)(&interpreter_); + InterpreterBuilder(model, *resolver_)(&interpreter_); CHECK(interpreter_ != nullptr); @@ -162,6 +187,7 @@ void SingleOpModel::BuildInterpreter( for (const auto& shape : input_shapes) { int input_idx = interpreter_->inputs()[i++]; if (input_idx == kOptionalTensor) continue; + if (shape.empty()) continue; CHECK(interpreter_->ResizeInputTensor(input_idx, shape) == kTfLiteOk); } CHECK(interpreter_->AllocateTensors() == kTfLiteOk) diff --git a/tensorflow/contrib/lite/kernels/test_util.h b/tensorflow/contrib/lite/kernels/test_util.h index 531c1366a87e20e140e779b767e29b1fd1111f97..7d476ba1eaffbb24fb77390c0e71c32d60b6411e 100644 --- a/tensorflow/contrib/lite/kernels/test_util.h +++ b/tensorflow/contrib/lite/kernels/test_util.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ #include @@ -85,6 +85,23 @@ struct TensorData { int32_t zero_point; }; +class SingleOpResolver : public OpResolver { + public: + SingleOpResolver(const BuiltinOperator op, TfLiteRegistration* registration) + : op_(op), registration_(registration) {} + TfLiteRegistration* FindOp(BuiltinOperator op) const override { + if (op == op_) { + return registration_; + } + return nullptr; + } + TfLiteRegistration* FindOp(const char* op) const override { return nullptr; } + + private: + const BuiltinOperator op_; + TfLiteRegistration* registration_; +}; + class SingleOpModel { public: SingleOpModel() {} @@ -98,6 +115,10 @@ class SingleOpModel { int AddInput(TensorType type) { return AddInput(TensorData{type}); } int AddInput(const TensorData& t); + // Add a Tensor containing const data and return the tensor id. + int AddConstInput(TensorType type, std::initializer_list data, + std::initializer_list shape); + // Add a null input tensor (optional input) and return kOptionalTensor. int AddNullInput(); @@ -174,14 +195,19 @@ class SingleOpModel { return result; } + void SetResolver(std::unique_ptr resolver) { + resolver_ = std::move(resolver); + } + protected: int32_t GetTensorSize(int index) const; flatbuffers::FlatBufferBuilder builder_; std::unique_ptr interpreter_; + std::unique_ptr resolver_; private: - int AddTensor(TensorData t); + int AddTensor(TensorData t, std::initializer_list data); std::map tensor_data_; std::vector inputs_; @@ -189,12 +215,43 @@ class SingleOpModel { std::vector> tensors_; std::vector> opcodes_; std::vector> operators_; + std::vector> buffers_; std::map> custom_registrations_; }; +// Base class for single op unit tests. +// The tests are parameterized to test multiple kernels for a single op. +// The parameters are strings like "optimized" and "reference" to have better +// readability in test reports. +// +// To use this class: +// * Define a constant map from strings to TfLiteRegistration. +// * Implement a test class that inherits SingleOpTest. +// * Instantiate the test cases with SingleOpTest::GetKernelTags helper +// function. +// * Call GetRegistration to get the TfLiteRegistration to be used before +// building the interpreter. +class SingleOpTest : public ::testing::TestWithParam { + public: + static std::vector GetKernelTags( + const std::map& kernel_map) { + std::vector tags; + for (auto it : kernel_map) { + tags.push_back(it.first); + } + return tags; + } + + protected: + virtual const std::map& GetKernelMap() = 0; + TfLiteRegistration* GetRegistration() { + return GetKernelMap().at(GetParam()); + } +}; + // Strings have a special implementation that is in test_util.cc template <> std::vector SingleOpModel::ExtractVector(int index); } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ diff --git a/tensorflow/contrib/lite/kernels/topk_v2.cc b/tensorflow/contrib/lite/kernels/topk_v2.cc new file mode 100644 index 0000000000000000000000000000000000000000..807e84609f8b23d25324d99d26086331d78a0684 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/topk_v2.cc @@ -0,0 +1,232 @@ +/* Copyright 2017 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 + +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" +namespace tflite { +namespace ops { +namespace builtin { +namespace topk_v2 { +constexpr int kInputTensor = 0; +constexpr int kInputTopK = 1; +constexpr int kOutputIndexes = 0; +constexpr int kOutputValues = 1; + +namespace { +TfLiteStatus ResizeOutput(TfLiteContext* context, TfLiteNode* node) { + TfLiteTensor* top_k = GetInput(context, node, kInputTopK); + // INT32 number of top results is supported. + TF_LITE_ENSURE_EQ(context, top_k->type, kTfLiteInt32); + // Check that the tensor contains only one value. + TF_LITE_ENSURE_EQ(context, NumDimensions(top_k), 1); + TF_LITE_ENSURE_EQ(context, NumElements(top_k), 1); + const int32 k = top_k->data.i32[0]; + + TfLiteTensor* input = GetInput(context, node, kInputTensor); + const int num_dimensions = NumDimensions(input); + // Check that input has one or more dimensions. + TF_LITE_ENSURE_MSG(context, input->dims->size >= 1, + "TopK k input must have 1 or more dimensions."); + // Check that k is less or equal the internal dimension. + TF_LITE_ENSURE_MSG(context, k <= input->dims->data[num_dimensions - 1], + "TopK k is higher than the internal dimension."); + + TfLiteIntArray* output_indexes_shape = TfLiteIntArrayCreate(num_dimensions); + TfLiteIntArray* output_values_shape = TfLiteIntArrayCreate(num_dimensions); + for (int i = 0; i < num_dimensions - 1; ++i) { + output_indexes_shape->data[i] = input->dims->data[i]; + output_values_shape->data[i] = input->dims->data[i]; + } + output_indexes_shape->data[num_dimensions - 1] = k; + output_values_shape->data[num_dimensions - 1] = k; + TfLiteTensor* output_indexes = GetOutput(context, node, kOutputIndexes); + TfLiteTensor* output_values = GetOutput(context, node, kOutputValues); + auto resize_tensor = [context](TfLiteTensor* tensor, TfLiteIntArray* new_size, + TfLiteIntArray* delete_on_error) { + TfLiteStatus status = context->ResizeTensor(context, tensor, new_size); + if (status != kTfLiteOk) { + TfLiteIntArrayFree(new_size); + if (delete_on_error != nullptr) { + TfLiteIntArrayFree(delete_on_error); + } + } + return status; + }; + TF_LITE_ENSURE_OK(context, resize_tensor(output_indexes, output_indexes_shape, + output_values_shape)); + TF_LITE_ENSURE_OK(context, + resize_tensor(output_values, output_values_shape, nullptr)); + return kTfLiteOk; +} + +// The class that collects top indexes of k values. Based on template +// tensorflow::gtl::TopN<> but, for optimization, +// it re-uses the same container. +template +class TopContainer { + public: + TopContainer() = delete; + TopContainer(int32 k, int32 row_size) : k_(k) { + container_.reserve(std::min(k, row_size) + 1); + } + + void start_collecting(const T* values) { + values_ = values; + container_.clear(); + } + void push(int32 a) { + auto comparator = [this](int32 a, int32 b) { return compare_fun(a, b); }; + if (container_.size() <= k_) { + container_.push_back(a); + if (container_.size() == k_ + 1) { + std::make_heap(container_.begin(), container_.end(), comparator); + std::pop_heap(container_.begin(), container_.end(), comparator); + } + } else if (comparator(a, container_.front())) { + container_.back() = a; + std::push_heap(container_.begin(), container_.end(), comparator); + std::pop_heap(container_.begin(), container_.end(), comparator); + } + } + + const std::vector& sorted_result() { + auto comparator = [this](int32 a, int32 b) { return compare_fun(a, b); }; + if (container_.size() <= k_) { + std::sort(container_.begin(), container_.end(), comparator); + } else { + std::sort_heap(container_.begin(), container_.end() - 1, comparator); + container_.resize(k_); + } + return container_; + } + + private: + int32 k_; + std::vector container_; + const T* values_ = nullptr; + + bool compare_fun(int32 a, int32 b) const { + if (values_[b] < values_[a]) { + return true; + } else if (values_[b] > values_[a]) { + return false; + } else { + return a < b; + } + } +}; + +// Mostly modeled on tensorflow/core/kernels/topk_op.cc for CPU. +template +void TopK(int32 row_size, int32 num_rows, const T* data, int32 k, + int32* output_indexes, T* output_values) { + TopContainer topc(k, row_size); + for (int row = 0; row < num_rows; ++row) { + const T* values_row = data + row * row_size; + topc.start_collecting(values_row); + for (int32 c = 0; c < row_size; ++c) { + topc.push(c); + } + + // Prepare output buffers. + int32* indexes_row = output_indexes + row * k; + T* output_row = output_values + row * k; + // We always assume that the output is sorted. + const auto& top_k = topc.sorted_result(); + std::copy(top_k.begin(), top_k.end(), indexes_row); + std::transform(top_k.begin(), top_k.end(), output_row, + [values_row](const int32 loc) { return values_row[loc]; }); + } +} + +} // namespace + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + // Check that the inputs and outputs have the right sizes and types. + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2); + + TfLiteTensor* input = GetInput(context, node, kInputTensor); + TfLiteTensor* output_values = GetOutput(context, node, kOutputValues); + TF_LITE_ENSURE_EQ(context, input->type, output_values->type); + + TfLiteTensor* top_k = GetInput(context, node, kInputTopK); + TF_LITE_ENSURE_EQ(context, top_k->type, kTfLiteInt32); + + // Set output dynamic if the input is not const. + if (IsConstantTensor(top_k)) { + TF_LITE_ENSURE_OK(context, ResizeOutput(context, node)); + } else { + TfLiteTensor* output_indexes = GetOutput(context, node, kOutputIndexes); + TfLiteTensor* output_values = GetOutput(context, node, kOutputValues); + SetTensorToDynamic(output_indexes); + SetTensorToDynamic(output_values); + } + return kTfLiteOk; +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + TfLiteTensor* output_values = GetOutput(context, node, kOutputValues); + TfLiteTensor* output_indexes = GetOutput(context, node, kOutputIndexes); + if (IsDynamicTensor(output_values)) { + TF_LITE_ENSURE_OK(context, ResizeOutput(context, node)); + } + TfLiteTensor* top_k = GetInput(context, node, kInputTopK); + const int32 k = top_k->data.i32[0]; + // The tensor can have more than 2 dimensions or even be a vector, the code + // anyway calls the internal dimension as row; + TfLiteTensor* input = GetInput(context, node, kInputTensor); + const int32 row_size = input->dims->data[input->dims->size - 1]; + int32 num_rows = 1; + for (int i = 0; i < input->dims->size - 1; ++i) { + num_rows *= input->dims->data[i]; + } + switch (output_values->type) { + case kTfLiteFloat32: + TopK(row_size, num_rows, input->data.f, k, output_indexes->data.i32, + output_values->data.f); + break; + case kTfLiteUInt8: + TopK(row_size, num_rows, input->data.uint8, k, output_indexes->data.i32, + output_values->data.uint8); + break; + case kTfLiteInt32: + TopK(row_size, num_rows, input->data.i32, k, output_indexes->data.i32, + output_values->data.i32); + break; + case kTfLiteInt64: + TopK(row_size, num_rows, input->data.i64, k, output_indexes->data.i32, + output_values->data.i64); + break; + default: + context->ReportError(context, "Type is currently not supported by TopK."); + return kTfLiteError; + } + + return kTfLiteOk; +} +} // namespace topk_v2 +TfLiteRegistration* Register_TOPK_V2() { + static TfLiteRegistration r = {nullptr, nullptr, topk_v2::Prepare, + topk_v2::Eval}; + return &r; +} +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/topk_v2_test.cc b/tensorflow/contrib/lite/kernels/topk_v2_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..29f2a057cd45e1cded3ff1aa0f0fdcad666ce2fa --- /dev/null +++ b/tensorflow/contrib/lite/kernels/topk_v2_test.cc @@ -0,0 +1,155 @@ + +/* Copyright 2017 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 +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +class TopKV2OpModel : public SingleOpModel { + public: + TopKV2OpModel(std::initializer_list input_shape, TensorType input_type, + int top_k) { + input_ = AddInput(input_type); + top_k_ = AddInput(TensorType_INT32); + output_indexes_ = AddOutput(TensorType_INT32); + output_values_ = AddOutput(input_type); + SetBuiltinOp(BuiltinOperator_TOPK_V2, BuiltinOptions_TopKV2Options, 0); + BuildInterpreter({input_shape, {1}}); + PopulateTensor(top_k_, {top_k}); + } + + void SetInputFloat(std::initializer_list data) { + PopulateTensor(input_, data); + } + + void SetInputUInt8(std::initializer_list data) { + PopulateTensor(input_, data); + } + + void SetInputInt32(std::initializer_list data) { + PopulateTensor(input_, data); + } + + void SetInputInt64(std::initializer_list data) { + PopulateTensor(input_, data); + } + + std::vector GetIndexes() { + return ExtractVector(output_indexes_); + } + + std::vector GetValuesFloat() { + return ExtractVector(output_values_); + } + + std::vector GetValuesUInt8() { + return ExtractVector(output_values_); + } + + std::vector GetValuesInt32() { + return ExtractVector(output_values_); + } + + std::vector GetValuesInt64() { + return ExtractVector(output_values_); + } + + protected: + int input_; + int top_k_; + int output_indexes_; + int output_values_; +}; + +// The test where the tensor dimension is equal to top. +TEST(TopKV2OpTest, EqualFloat) { + TopKV2OpModel m({2, 2}, TensorType_FLOAT32, 2); + m.SetInputFloat({-2.0, 0.2, 0.8, 0.1}); + m.Invoke(); + EXPECT_THAT(m.GetIndexes(), ElementsAreArray({1, 0, 0, 1})); + EXPECT_THAT(m.GetValuesFloat(), + ElementsAreArray(ArrayFloatNear({0.2, -2.0, 0.8, 0.1}))); +} + +// Test when internal dimension is k+1. +TEST(TopKV2OpTest, BorderFloat) { + TopKV2OpModel m({2, 3}, TensorType_FLOAT32, 2); + m.SetInputFloat({-2.0, -3.0, 0.2, 0.8, 0.1, -0.1}); + m.Invoke(); + EXPECT_THAT(m.GetIndexes(), ElementsAreArray({2, 0, 0, 1})); + EXPECT_THAT(m.GetValuesFloat(), + ElementsAreArray(ArrayFloatNear({0.2, -2.0, 0.8, 0.1}))); +} +// Test when internal dimension is higher than k. +TEST(TopKV2OpTest, LargeFloat) { + TopKV2OpModel m({2, 4}, TensorType_FLOAT32, 2); + m.SetInputFloat({-2.0, -3.0, -4.0, 0.2, 0.8, 0.1, -0.1, -0.8}); + m.Invoke(); + EXPECT_THAT(m.GetIndexes(), ElementsAreArray({3, 0, 0, 1})); + EXPECT_THAT(m.GetValuesFloat(), + ElementsAreArray(ArrayFloatNear({0.2, -2.0, 0.8, 0.1}))); +} + +// Test 1D case. +TEST(TopKV2OpTest, VectorFloat) { + TopKV2OpModel m({8}, TensorType_FLOAT32, 2); + m.SetInputFloat({-2.0, -3.0, -4.0, 0.2, 0.8, 0.1, -0.1, -0.8}); + m.Invoke(); + EXPECT_THAT(m.GetIndexes(), ElementsAreArray({4, 3})); + EXPECT_THAT(m.GetValuesFloat(), ElementsAreArray(ArrayFloatNear({0.8, 0.2}))); +} + +// Check that uint8 works. +TEST(TopKV2OpTest, TypeUint8) { + TopKV2OpModel m({2, 3}, TensorType_UINT8, 2); + m.SetInputUInt8({1, 2, 3, 251, 250, 249}); + m.Invoke(); + EXPECT_THAT(m.GetIndexes(), ElementsAreArray({2, 1, 0, 1})); + EXPECT_THAT(m.GetValuesUInt8(), ElementsAreArray({3, 2, 251, 250})); +} + +// Check that int32 works. +TEST(TopKV2OpTest, TypeInt32) { + TopKV2OpModel m({2, 3}, TensorType_INT32, 2); + m.SetInputInt32({1, 2, 3, 10251, 10250, 10249}); + m.Invoke(); + EXPECT_THAT(m.GetIndexes(), ElementsAreArray({2, 1, 0, 1})); + EXPECT_THAT(m.GetValuesInt32(), ElementsAreArray({3, 2, 10251, 10250})); +} + +// Check that int64 works. +TEST(TopKV2OpTest, TypeInt64) { + TopKV2OpModel m({2, 3}, TensorType_INT64, 2); + m.SetInputInt64({1, 2, 3, -1, -2, -3}); + m.Invoke(); + EXPECT_THAT(m.GetIndexes(), ElementsAreArray({2, 1, 0, 1})); + EXPECT_THAT(m.GetValuesInt64(), ElementsAreArray({3, 2, -1, -2})); +} +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/transpose.cc b/tensorflow/contrib/lite/kernels/transpose.cc new file mode 100644 index 0000000000000000000000000000000000000000..d3c10a9bb7b07404ccd8cfe2636473a622b91787 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/transpose.cc @@ -0,0 +1,159 @@ +/* Copyright 2017 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 +#include +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace transpose { + +// This file has two implementations of Transpose. +enum KernelType { + kReference, +}; + +struct TransposeContext { + TransposeContext(TfLiteContext* context, TfLiteNode* node) { + input = GetInput(context, node, 0); + perm = GetInput(context, node, 1); + output = GetOutput(context, node, 0); + } + TfLiteTensor* input; + TfLiteTensor* perm; + TfLiteTensor* output; +}; + +TfLiteStatus ResizeOutputTensor(TfLiteContext* context, + TransposeContext* op_context) { + int dims = NumDimensions(op_context->input); + const int* perm_data = GetTensorData(op_context->perm); + + // Ensure validity of the permutations tensor as a 1D tensor. + TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->perm), 1); + TF_LITE_ENSURE_EQ(context, op_context->perm->dims->data[0], dims); + for (int idx = 0; idx < dims; ++idx) { + TF_LITE_ENSURE_MSG(context, (perm_data[idx] >= 0 && perm_data[idx] < dims), + "Transpose op permutations array is out of bounds."); + } + + // Determine size of output tensor. + TfLiteIntArray* input_size = op_context->input->dims; + TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_size); + for (int idx = 0; idx < dims; ++idx) { + output_size->data[idx] = input_size->data[perm_data[idx]]; + } + + return context->ResizeTensor(context, op_context->output, output_size); +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + TransposeContext op_context(context, node); + + // Ensure validity of input tensor. + TF_LITE_ENSURE_MSG(context, NumDimensions(op_context.input) <= 4, + "Transpose op only supports 1D-4D input arrays."); + TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type); + + if (!IsConstantTensor(op_context.perm)) { + SetTensorToDynamic(op_context.output); + return kTfLiteOk; + } + return ResizeOutputTensor(context, &op_context); +} + +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + TransposeContext op_context(context, node); + + // Resize the output tensor if the output tensor is dynamic. + if (IsDynamicTensor(op_context.output)) { + TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context)); + } + + // Reverse the permuted axes and convert to 4D due to the way Dims are + // constructed in GetTensorDims. + const int* perm_data = GetTensorData(op_context.perm); + const int size = op_context.perm->dims->data[0]; + const int kOutputDimensionNum = 4; + int reversed_perm[kOutputDimensionNum]; + + for (int output_k = 0, input_k = size - 1; output_k < size; + ++output_k, --input_k) { + reversed_perm[output_k] = size - perm_data[input_k] - 1; + } + for (int k = size; k < kOutputDimensionNum; ++k) { + reversed_perm[k] = k; + } + +#define TF_LITE_TRANSPOSE(type, scalar) \ + type::Transpose(GetTensorData(op_context.input), \ + GetTensorDims(op_context.input), \ + GetTensorData(op_context.output), \ + GetTensorDims(op_context.output), reversed_perm) + + switch (op_context.input->type) { + case kTfLiteFloat32: + if (kernel_type == kReference) { + TF_LITE_TRANSPOSE(reference_ops, float); + } + break; + case kTfLiteUInt8: + if (kernel_type == kReference) { + TF_LITE_TRANSPOSE(reference_ops, uint8_t); + } + break; + case kTfLiteInt32: + if (kernel_type == kReference) { + TF_LITE_TRANSPOSE(reference_ops, int32_t); + } + break; + case kTfLiteInt64: + if (kernel_type == kReference) { + TF_LITE_TRANSPOSE(reference_ops, int64_t); + } + break; + default: + context->ReportError(context, + "Type is currently not supported by Transpose."); + return kTfLiteError; + } +#undef TF_LITE_TRANSPOSE + + return kTfLiteOk; +} + +} // namespace transpose + +TfLiteRegistration* Register_TRANSPOSE_REF() { + static TfLiteRegistration r = {nullptr, nullptr, transpose::Prepare, + transpose::Eval}; + return &r; +} + +TfLiteRegistration* Register_TRANSPOSE() { return Register_TRANSPOSE_REF(); } + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/transpose_test.cc b/tensorflow/contrib/lite/kernels/transpose_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..337bc144b967392523bf784603cca4c1b968cdf2 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/transpose_test.cc @@ -0,0 +1,347 @@ +/* Copyright 2017 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 +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +void RunTestPermutation(const std::vector& shape, + const std::vector& perms, + std::vector* input_transposed) { + // Count elements and allocate output. + int count = 1; + for (auto factor : shape) count *= factor; + input_transposed->resize(count); + + // Create the dummy data + std::vector input(count); + for (int i = 0; i < input.size(); i++) { + input[i] = i; + } + + // Create reversed and padded perms. + int reversed_perms[4]; + for (int output_k = 0, input_k = shape.size() - 1; output_k < shape.size(); + output_k++, input_k--) { + reversed_perms[output_k] = shape.size() - perms[input_k] - 1; + } + // Unused dimensions should not be permuted so pad with identity transform + // subset. + for (int k = shape.size(); k < 4; k++) { + reversed_perms[k] = k; + } + + // Make input and output dims (i.e. reversed shape and dest_shape). + Dims<4> input_dims = GetTensorDims(shape); + Dims<4> output_dims; + for (int i = 0; i < 4; i++) { + output_dims.sizes[i] = input_dims.sizes[reversed_perms[i]]; + } + output_dims.strides[0] = 1; + for (int k = 1; k < 4; k++) { + output_dims.strides[k] = + output_dims.strides[k - 1] * output_dims.sizes[k - 1]; + } + + reference_ops::Transpose(input.data(), input_dims, + input_transposed->data(), output_dims, + reversed_perms); +} + +TEST(TransposeTest, TestRefOps1D) { + // Basic 1D identity. + std::vector out; + RunTestPermutation({3}, {0}, &out); + ASSERT_EQ(out, std::vector({0, 1, 2})); +} + +TEST(TransposeTest, TestRefOps2D) { + std::vector out; + // Basic 2D. + RunTestPermutation({3, 2}, {1, 0}, &out); + ASSERT_EQ(out, std::vector({0, 2, 4, 1, 3, 5})); + // Identity. + RunTestPermutation({3, 2}, {0, 1}, &out); + ASSERT_EQ(out, std::vector({0, 1, 2, 3, 4, 5})); +} + +TEST(TransposeTest, TestRefOps3D) { + std::vector out; + // Test 3 dimensional + { + std::vector ref({0, 4, 8, 12, 16, 20, 1, 5, 9, 13, 17, 21, + 2, 6, 10, 14, 18, 22, 3, 7, 11, 15, 19, 23}); + RunTestPermutation({2, 3, 4}, {2, 0, 1}, &out); + ASSERT_EQ(out, ref); + } + // Test 3 dimensional identity transform + { + RunTestPermutation({2, 3, 4}, {0, 1, 2}, &out); + std::vector ref(out.size()); + for (int k = 0; k < ref.size(); k++) ref[k] = k; + ASSERT_EQ(out, ref); + } +} + +TEST(TransposeTest, TestRefOps4D) { + std::vector out; + // Basic 4d. + RunTestPermutation({2, 3, 4, 5}, {2, 0, 1, 3}, &out); + ASSERT_EQ( + out, + std::vector( + {0, 1, 2, 3, 4, 20, 21, 22, 23, 24, 40, 41, 42, 43, 44, + 60, 61, 62, 63, 64, 80, 81, 82, 83, 84, 100, 101, 102, 103, 104, + 5, 6, 7, 8, 9, 25, 26, 27, 28, 29, 45, 46, 47, 48, 49, + 65, 66, 67, 68, 69, 85, 86, 87, 88, 89, 105, 106, 107, 108, 109, + 10, 11, 12, 13, 14, 30, 31, 32, 33, 34, 50, 51, 52, 53, 54, + 70, 71, 72, 73, 74, 90, 91, 92, 93, 94, 110, 111, 112, 113, 114, + 15, 16, 17, 18, 19, 35, 36, 37, 38, 39, 55, 56, 57, 58, 59, + 75, 76, 77, 78, 79, 95, 96, 97, 98, 99, 115, 116, 117, 118, 119})); + RunTestPermutation({2, 3, 4, 5}, {0, 1, 2, 3}, &out); + // Basic identity. + std::vector ref(out.size()); + for (int k = 0; k < ref.size(); k++) ref[k] = k; + ASSERT_EQ(out, ref); +} + +class TransposeOpModel : public SingleOpModel { + public: + void SetInput(std::initializer_list data) { + PopulateTensor(input_, data); + } + + void SetPerm(std::initializer_list data) { + PopulateTensor(perm_, data); + } + + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + protected: + int input_; + int perm_; + int output_; +}; + +// Tests case where perm is a const tensor. +// +// Example usage is as follows: +// SpaceToBatchNDOpConstModel m(input_shape, perm_shape, perm_data); +// m.SetInput(input_data); +// m.Invoke(); +class TransposeOpConstModel : public TransposeOpModel { + public: + TransposeOpConstModel(std::initializer_list input_shape, + std::initializer_list perm_shape, + std::initializer_list perm) { + input_ = AddInput(TensorType_FLOAT32); + perm_ = AddConstInput(TensorType_INT32, perm, perm_shape); + output_ = AddOutput(TensorType_FLOAT32); + SetBuiltinOp(BuiltinOperator_TRANSPOSE, BuiltinOptions_TransposeOptions, + CreateTransposeOptions(builder_).Union()); + BuildInterpreter({input_shape}); + } +}; + +// Tests case where perm is a non-const tensor. +// +// Example usage is as follows: +// TransposeOpDynamicModel m(input_shape, perm_shape); +// m.SetInput(input_data); +// m.SetPerm(perm_data); +// m.Invoke(); +class TransposeOpDynamicModel : public TransposeOpModel { + public: + TransposeOpDynamicModel(std::initializer_list input_shape, + std::initializer_list perm_shape) { + input_ = AddInput(TensorType_FLOAT32); + perm_ = AddInput(TensorType_INT32); + output_ = AddOutput(TensorType_FLOAT32); + SetBuiltinOp(BuiltinOperator_TRANSPOSE, BuiltinOptions_TransposeOptions, + CreateTransposeOptions(builder_).Union()); + BuildInterpreter({input_shape, perm_shape}); + } +}; + +TEST(TransposeTest, TestUnequalPermSize) { + EXPECT_DEATH(TransposeOpConstModel({1, 3, 3, 1}, {2}, {2, 2}), "2 != 4"); +} + +TEST(TransposeTest, TestPermOutOfBounds) { + EXPECT_DEATH(TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, -1, -2, -3}), + "Transpose op permutations array is out of bounds."); + EXPECT_DEATH(TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 4}), + "Transpose op permutations array is out of bounds."); +} + +TEST(TransposeTest, Test1DInputConstTensor) { + TransposeOpConstModel m({3}, {1}, {0}); + m.SetInput({1, 2, 3}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3})); +} + +TEST(TransposeTest, Test1DInputDynamicTensor) { + TransposeOpDynamicModel m({3}, {1}); + m.SetInput({1, 2, 3}); + m.SetPerm({0}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3})); +} + +TEST(TransposeTest, Test2DInputConstTensor) { + TransposeOpConstModel m({3, 2}, {2}, {1, 0}); + m.SetInput({0, 1, 2, 3, 4, 5}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({0, 2, 4, 1, 3, 5})); +} + +TEST(TransposeTest, Test2DInputDynamicTensor) { + TransposeOpDynamicModel m({3, 2}, {2}); + m.SetInput({0, 1, 2, 3, 4, 5}); + m.SetPerm({1, 0}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({0, 2, 4, 1, 3, 5})); +} + +TEST(TransposeTest, Test3DInputConstTensor) { + TransposeOpConstModel m({2, 3, 4}, {3}, {2, 0, 1}); + m.SetInput({0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 2, 3})); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray({0, 4, 8, 12, 16, 20, 1, 5, 9, 13, 17, 21, + 2, 6, 10, 14, 18, 22, 3, 7, 11, 15, 19, 23})); +} + +TEST(TransposeTest, Test3DInputDynamicTensor) { + TransposeOpDynamicModel m({2, 3, 4}, {3}); + m.SetInput({0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}); + m.SetPerm({2, 0, 1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 2, 3})); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray({0, 4, 8, 12, 16, 20, 1, 5, 9, 13, 17, 21, + 2, 6, 10, 14, 18, 22, 3, 7, 11, 15, 19, 23})); +} + +TEST(TransposeTest, Test5DInputTensor) { + EXPECT_DEATH(TransposeOpConstModel({1, 2, 3, 4, 5}, {5}, {0, 1, 2, 3, 4}), + "Transpose op only supports 1D-4D input arrays."); +} + +TEST(TransposeTest, SimpleTestNoReorderConstTensor) { + TransposeOpConstModel m({1, 2, 3, 1}, {4}, {0, 1, 2, 3}); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 3, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6})); +} + +TEST(TransposeTest, SimpleTestNoReorderDynamicTensor) { + TransposeOpDynamicModel m({1, 2, 3, 1}, {4}); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.SetPerm({0, 1, 2, 3}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 3, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6})); +} + +TEST(TransposeTest, SimpleTestWithReorderConstTensor) { + TransposeOpConstModel m({1, 2, 3, 1}, {4}, {2, 1, 3, 0}); + m.SetInput({1, 2, 3, 4, 5, 6}); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 2, 1, 1})); + EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 4, 2, 5, 3, 6})); +} + +TEST(TransposeTest, ComplexTestWithReorderConstTensor) { + TransposeOpConstModel m({2, 3, 4, 5}, {4}, {2, 0, 1, 3}); + m.SetInput({0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, + 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, + 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, + 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, + 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, + 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, + 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, + 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119}); + m.Invoke(); + + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 2, 3, 5})); + auto result = ElementsAreArray( + {0, 1, 2, 3, 4, 20, 21, 22, 23, 24, 40, 41, 42, 43, 44, + 60, 61, 62, 63, 64, 80, 81, 82, 83, 84, 100, 101, 102, 103, 104, + 5, 6, 7, 8, 9, 25, 26, 27, 28, 29, 45, 46, 47, 48, 49, + 65, 66, 67, 68, 69, 85, 86, 87, 88, 89, 105, 106, 107, 108, 109, + 10, 11, 12, 13, 14, 30, 31, 32, 33, 34, 50, 51, 52, 53, 54, + 70, 71, 72, 73, 74, 90, 91, 92, 93, 94, 110, 111, 112, 113, 114, + 15, 16, 17, 18, 19, 35, 36, 37, 38, 39, 55, 56, 57, 58, 59, + 75, 76, 77, 78, 79, 95, 96, 97, 98, 99, 115, 116, 117, 118, 119}); + EXPECT_THAT(m.GetOutput(), result); +} + +TEST(TransposeTest, ComplexTestWithReorderDynamicTensor) { + TransposeOpDynamicModel m({2, 3, 4, 5}, {4}); + m.SetInput({0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, + 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, + 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, + 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, + 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, + 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, + 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, + 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119}); + m.SetPerm({2, 0, 1, 3}); + m.Invoke(); + + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 2, 3, 5})); + auto result = ElementsAreArray( + {0, 1, 2, 3, 4, 20, 21, 22, 23, 24, 40, 41, 42, 43, 44, + 60, 61, 62, 63, 64, 80, 81, 82, 83, 84, 100, 101, 102, 103, 104, + 5, 6, 7, 8, 9, 25, 26, 27, 28, 29, 45, 46, 47, 48, 49, + 65, 66, 67, 68, 69, 85, 86, 87, 88, 89, 105, 106, 107, 108, 109, + 10, 11, 12, 13, 14, 30, 31, 32, 33, 34, 50, 51, 52, 53, 54, + 70, 71, 72, 73, 74, 90, 91, 92, 93, 94, 110, 111, 112, 113, 114, + 15, 16, 17, 18, 19, 35, 36, 37, 38, 39, 55, 56, 57, 58, 59, + 75, 76, 77, 78, 79, 95, 96, 97, 98, 99, 115, 116, 117, 118, 119}); + EXPECT_THAT(m.GetOutput(), result); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc new file mode 100644 index 0000000000000000000000000000000000000000..9cdb58714edb5fee771fc45f3c53a570f8fb28d1 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc @@ -0,0 +1,527 @@ +/* Copyright 2017 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 +#include +#include +#include +#include +#include +#include + +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/activation_functor.h" +#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" +#include "tensorflow/contrib/lite/kernels/kernel_util.h" +#include "tensorflow/contrib/lite/kernels/op_macros.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace unidirectional_sequence_lstm { + +// Input Tensors of size {max_time, n_batch, n_input} +constexpr int kInputTensor = 0; + +// Input weight tensors of size: {n_cell, n_input} +constexpr int kInputToInputWeightsTensor = 1; // Optional +constexpr int kInputToForgetWeightsTensor = 2; +constexpr int kInputToCellWeightsTensor = 3; +constexpr int kInputToOutputWeightsTensor = 4; + +// Recurrent weight tensors of size {n_cell, n_output} +constexpr int kRecurrentToInputWeightsTensor = 5; // Optional +constexpr int kRecurrentToForgetWeightsTensor = 6; +constexpr int kRecurrentToCellWeightsTensor = 7; +constexpr int kRecurrentToOutputWeightsTensor = 8; + +// Peephole weights tensors of size {n_cell}, representing a diagonal matrix. +constexpr int kCellToInputWeightsTensor = 9; // Optional +constexpr int kCellToForgetWeightsTensor = 10; // Optional +constexpr int kCellToOutputWeightsTensor = 11; // Optional + +// Gates bias tensors of size {n_cell} +constexpr int kInputGateBiasTensor = 12; // Optional +constexpr int kForgetGateBiasTensor = 13; +constexpr int kCellGateBiasTensor = 14; +constexpr int kOutputGateBiasTensor = 15; + +// Projection weight tensor of size {n_output, n_cell} +constexpr int kProjectionWeightsTensor = 16; // Optional +// Projection bias tensor of size {n_output} +constexpr int kProjectionBiasTensor = 17; // Optional + +// Output tensors. +constexpr int kScratchBufferTensor = 0; +constexpr int kOutputStateTensor = 1; +constexpr int kCellStateTensor = 2; +constexpr int kOutputTensor = 3; + +// Check that input tensor dimensions matches with each other. +TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context, + TfLiteNode* node, int n_input, + int n_output, int n_cell) { + auto* params = reinterpret_cast(node->builtin_data); + + // Making sure clipping parameters have valid values. + // == 0 means no clipping + // > 0 means clipping + TF_LITE_ENSURE(context, params->cell_clip >= 0); + TF_LITE_ENSURE(context, params->proj_clip >= 0); + + TfLiteTensor* input_to_input_weights = + GetOptionalInputTensor(context, node, kInputToInputWeightsTensor); + if (input_to_input_weights) { + TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->data[0], n_cell); + TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->data[1], n_input); + } + + TfLiteTensor* input_to_forget_weights = + GetInput(context, node, kInputToForgetWeightsTensor); + TF_LITE_ENSURE_EQ(context, input_to_forget_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, input_to_forget_weights->dims->data[0], n_cell); + TF_LITE_ENSURE_EQ(context, input_to_forget_weights->dims->data[1], n_input); + + TfLiteTensor* input_to_cell_weights = + GetInput(context, node, kInputToCellWeightsTensor); + TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->data[0], n_cell); + TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->data[1], n_input); + + TfLiteTensor* recurrent_to_input_weights = + GetOptionalInputTensor(context, node, kRecurrentToInputWeightsTensor); + if (recurrent_to_input_weights) { + TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->data[0], + n_cell); + TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->data[1], + n_output); + } + + TfLiteTensor* recurrent_to_forget_weights = + GetInput(context, node, kRecurrentToForgetWeightsTensor); + TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->dims->data[0], + n_cell); + TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->dims->data[1], + n_output); + + TfLiteTensor* recurrent_to_cell_weights = + GetInput(context, node, kRecurrentToCellWeightsTensor); + TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[0], n_cell); + TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[1], + n_output); + + // We make sure the input-gate's parameters are either both present (regular + // LSTM) or not at all (CIFG-LSTM). + const bool cifg_weights_all_or_none = + ((input_to_input_weights != nullptr) && + (recurrent_to_input_weights != nullptr)) || + ((input_to_input_weights == nullptr) && + (recurrent_to_input_weights == nullptr)); + TF_LITE_ENSURE(context, cifg_weights_all_or_none == true); + + TfLiteTensor* cell_to_input_weights = + GetOptionalInputTensor(context, node, kCellToInputWeightsTensor); + if (cell_to_input_weights) { + TF_LITE_ENSURE_EQ(context, cell_to_input_weights->dims->size, 1); + TF_LITE_ENSURE_EQ(context, cell_to_input_weights->dims->data[0], n_cell); + } + + TfLiteTensor* cell_to_forget_weights = + GetOptionalInputTensor(context, node, kCellToForgetWeightsTensor); + if (cell_to_forget_weights) { + TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->dims->size, 1); + TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->dims->data[0], n_cell); + } + + TfLiteTensor* cell_to_output_weights = + GetOptionalInputTensor(context, node, kCellToOutputWeightsTensor); + if (cell_to_output_weights) { + TF_LITE_ENSURE_EQ(context, cell_to_output_weights->dims->size, 1); + TF_LITE_ENSURE_EQ(context, cell_to_output_weights->dims->data[0], n_cell); + } + + // Making sure the peephole weights are there all or none. + const bool use_cifg = (input_to_input_weights == nullptr); + const bool peephole_weights_all_or_none = + ((cell_to_input_weights != nullptr || use_cifg) && + (cell_to_forget_weights != nullptr) && + (cell_to_output_weights != nullptr)) || + ((cell_to_input_weights == nullptr) && + (cell_to_forget_weights == nullptr) && + (cell_to_output_weights == nullptr)); + TF_LITE_ENSURE(context, peephole_weights_all_or_none == true); + + // Make sure the input gate bias is present only when not a CIFG-LSTM. + TfLiteTensor* input_gate_bias = + GetOptionalInputTensor(context, node, kInputGateBiasTensor); + if (use_cifg) { + TF_LITE_ENSURE_EQ(context, input_gate_bias, nullptr); + } else { + TF_LITE_ENSURE_EQ(context, input_gate_bias->dims->size, 1); + TF_LITE_ENSURE_EQ(context, input_gate_bias->dims->data[0], n_cell); + } + + TfLiteTensor* forget_gate_bias = + GetInput(context, node, kForgetGateBiasTensor); + TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->size, 1); + TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->data[0], n_cell); + + TfLiteTensor* cell_bias = GetInput(context, node, kCellGateBiasTensor); + TF_LITE_ENSURE_EQ(context, cell_bias->dims->size, 1); + TF_LITE_ENSURE_EQ(context, cell_bias->dims->data[0], n_cell); + + TfLiteTensor* output_gate_bias = + GetInput(context, node, kOutputGateBiasTensor); + TF_LITE_ENSURE_EQ(context, output_gate_bias->dims->size, 1); + TF_LITE_ENSURE_EQ(context, output_gate_bias->dims->data[0], n_cell); + + TfLiteTensor* projection_weights = + GetOptionalInputTensor(context, node, kProjectionWeightsTensor); + if (projection_weights) { + TF_LITE_ENSURE_EQ(context, projection_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, projection_weights->dims->data[0], n_output); + TF_LITE_ENSURE_EQ(context, projection_weights->dims->data[1], n_cell); + } + + TfLiteTensor* projection_bias = + GetOptionalInputTensor(context, node, kProjectionBiasTensor); + if (projection_bias) { + TF_LITE_ENSURE_EQ(context, projection_bias->dims->size, 1); + TF_LITE_ENSURE_EQ(context, projection_bias->dims->data[0], n_output); + } + + // Making sure the projection tensors are consistent: + // 1) If projection weight is not present, then projection bias should not be + // present. + // 2) If projection weight is present, then projection bias is optional. + // TODO(ghodrat): make sure this is correct. + const bool projecton_tensors_consistent = + ((projection_weights != nullptr) || (projection_bias == nullptr)); + TF_LITE_ENSURE(context, projecton_tensors_consistent == true); + + return kTfLiteOk; +} + +// Resize the output, state and scratch tensors based on the sizes of the input +// tensors. Also check that the size of the input tensors match each other. +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + // Check we have all the inputs and outputs we need. + TF_LITE_ENSURE_EQ(context, node->inputs->size, 18); + TF_LITE_ENSURE_EQ(context, node->outputs->size, 4); + + // Inferring batch size, number of outputs and sequence length and + // number of cells from the input tensors. + TfLiteTensor* input = GetInput(context, node, kInputTensor); + TF_LITE_ENSURE(context, input->dims->size > 1); + const int max_time = input->dims->data[0]; + const int n_batch = input->dims->data[1]; + const int n_input = input->dims->data[2]; + + TfLiteTensor* input_to_output_weights = + GetInput(context, node, kInputToOutputWeightsTensor); + const int n_cell = input_to_output_weights->dims->data[0]; + TF_LITE_ENSURE_EQ(context, input_to_output_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, input_to_output_weights->dims->data[1], n_input); + + TfLiteTensor* recurrent_to_output_weights = + GetInput(context, node, kRecurrentToOutputWeightsTensor); + TF_LITE_ENSURE_EQ(context, recurrent_to_output_weights->dims->size, 2); + TF_LITE_ENSURE_EQ(context, recurrent_to_output_weights->dims->data[0], + n_cell); + const int n_output = recurrent_to_output_weights->dims->data[1]; + + // Check that input tensor dimensions matches with each other. + CheckInputTensorDimensions(context, node, n_input, n_output, n_cell); + + // Get the pointer to output, state and scratch buffer tensors. + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + TfLiteTensor* output_state = GetOutput(context, node, kOutputStateTensor); + TfLiteTensor* cell_state = GetOutput(context, node, kCellStateTensor); + // TODO(ghodrat): Modify this as soon as we have a finalized method for + // scratch buffers. + TfLiteTensor* scratch_buffer = GetOutput(context, node, kScratchBufferTensor); + + // Resize the output and output_state tensors. + TfLiteIntArray* output_size = TfLiteIntArrayCreate(3); + output_size->data[0] = max_time; + output_size->data[1] = n_batch; + output_size->data[2] = n_output; + TF_LITE_ENSURE_OK(context, + context->ResizeTensor(context, output, output_size)); + + TfLiteIntArray* output_state_size = TfLiteIntArrayCreate(2); + output_state_size->data[0] = n_batch; + output_state_size->data[1] = n_output; + TF_LITE_ENSURE_OK( + context, context->ResizeTensor(context, output_state, output_state_size)); + + // Resize the scratch buffer tensor. + TfLiteIntArray* cell_size = TfLiteIntArrayCreate(2); + cell_size->data[0] = n_batch; + cell_size->data[1] = n_cell; + TF_LITE_ENSURE_OK(context, + context->ResizeTensor(context, cell_state, cell_size)); + + // Mark state tensors as persistent tensors. + output_state->allocation_type = kTfLiteArenaRwPersistent; + cell_state->allocation_type = kTfLiteArenaRwPersistent; + + TfLiteTensor* input_to_input_weights = + GetOptionalInputTensor(context, node, kInputToInputWeightsTensor); + const bool use_cifg = (input_to_input_weights == nullptr); + if (use_cifg) { + TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(2); + scratch_buffer_size->data[0] = n_batch; + // Reserving space for Cell, Forget, Output gates + scratch_buffer_size->data[1] = n_cell * 3; + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer, + scratch_buffer_size)); + } else { + TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(2); + scratch_buffer_size->data[0] = n_batch; + // Reserving space for Input, Cell, Forget, Output gates + scratch_buffer_size->data[1] = n_cell * 4; + TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer, + scratch_buffer_size)); + } + return kTfLiteOk; +} + +// The LSTM Op engine. +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + auto* params = reinterpret_cast(node->builtin_data); + TfLiteTensor* input = GetInput(context, node, kInputTensor); + + TfLiteTensor* input_to_input_weights = + GetOptionalInputTensor(context, node, kInputToInputWeightsTensor); + TfLiteTensor* input_to_forget_weights = + GetInput(context, node, kInputToForgetWeightsTensor); + TfLiteTensor* input_to_cell_weights = + GetInput(context, node, kInputToCellWeightsTensor); + TfLiteTensor* input_to_output_weights = + GetInput(context, node, kInputToOutputWeightsTensor); + + TfLiteTensor* recurrent_to_input_weights = + GetOptionalInputTensor(context, node, kRecurrentToInputWeightsTensor); + TfLiteTensor* recurrent_to_forget_weights = + GetInput(context, node, kRecurrentToForgetWeightsTensor); + TfLiteTensor* recurrent_to_cell_weights = + GetInput(context, node, kRecurrentToCellWeightsTensor); + TfLiteTensor* recurrent_to_output_weights = + GetInput(context, node, kRecurrentToOutputWeightsTensor); + + TfLiteTensor* cell_to_input_weights = + GetOptionalInputTensor(context, node, kCellToInputWeightsTensor); + TfLiteTensor* cell_to_forget_weights = + GetOptionalInputTensor(context, node, kCellToForgetWeightsTensor); + TfLiteTensor* cell_to_output_weights = + GetOptionalInputTensor(context, node, kCellToOutputWeightsTensor); + + TfLiteTensor* input_gate_bias = + GetOptionalInputTensor(context, node, kInputGateBiasTensor); + TfLiteTensor* forget_gate_bias = + GetInput(context, node, kForgetGateBiasTensor); + TfLiteTensor* cell_bias = GetInput(context, node, kCellGateBiasTensor); + TfLiteTensor* output_gate_bias = + GetInput(context, node, kOutputGateBiasTensor); + + TfLiteTensor* projection_weights = + GetOptionalInputTensor(context, node, kProjectionWeightsTensor); + TfLiteTensor* projection_bias = + GetOptionalInputTensor(context, node, kProjectionBiasTensor); + + TfLiteTensor* output_state = GetOutput(context, node, kOutputStateTensor); + TfLiteTensor* cell_state = GetOutput(context, node, kCellStateTensor); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + const int max_time = input->dims->data[0]; + const int n_batch = input->dims->data[1]; + const int n_input = input->dims->data[2]; + // n_cell and n_output will be the same size when there is no projection. + const int n_cell = input_to_output_weights->dims->data[0]; + const int n_output = recurrent_to_output_weights->dims->data[1]; + + // Since we have already checked that weights are all there or none, we can + // check the existense of only one to the get the condition. + const bool use_cifg = (input_to_input_weights == nullptr); + const bool use_peephole = (cell_to_output_weights != nullptr); + + // Index the scratch buffers pointers to the global scratch buffer. + TfLiteTensor* scratch_buffer = GetOutput(context, node, kScratchBufferTensor); + float* input_gate_scratch = nullptr; + float* cell_scratch = nullptr; + float* forget_gate_scratch = nullptr; + float* output_gate_scratch = nullptr; + if (use_cifg) { + cell_scratch = scratch_buffer->data.f; + forget_gate_scratch = scratch_buffer->data.f + n_cell * n_batch; + output_gate_scratch = scratch_buffer->data.f + 2 * n_cell * n_batch; + } else { + input_gate_scratch = scratch_buffer->data.f; + cell_scratch = scratch_buffer->data.f + n_cell * n_batch; + forget_gate_scratch = scratch_buffer->data.f + 2 * n_cell * n_batch; + output_gate_scratch = scratch_buffer->data.f + 3 * n_cell * n_batch; + } + + for (int t = 0; t < max_time; t++) { + const float* input_ptr_time = input->data.f + t * n_batch * n_input; + // Initialize scratch buffers with bias. + if (!use_cifg) { + tensor_utils::VectorBatchVectorAssign(input_gate_bias->data.f, n_cell, + n_batch, input_gate_scratch); + } + tensor_utils::VectorBatchVectorAssign(forget_gate_bias->data.f, n_cell, + n_batch, forget_gate_scratch); + tensor_utils::VectorBatchVectorAssign(cell_bias->data.f, n_cell, n_batch, + cell_scratch); + tensor_utils::VectorBatchVectorAssign(output_gate_bias->data.f, n_cell, + n_batch, output_gate_scratch); + + // For each batch and cell: compute input_weight * input. + if (!use_cifg) { + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + input_to_input_weights->data.f, n_cell, n_input, input_ptr_time, + n_batch, input_gate_scratch, /*result_stride=*/1); + } + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + input_to_forget_weights->data.f, n_cell, n_input, input_ptr_time, + n_batch, forget_gate_scratch, /*result_stride=*/1); + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + input_to_cell_weights->data.f, n_cell, n_input, input_ptr_time, n_batch, + cell_scratch, /*result_stride=*/1); + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + input_to_output_weights->data.f, n_cell, n_input, input_ptr_time, + n_batch, output_gate_scratch, /*result_stride=*/1); + + // For each batch and cell: compute recurrent_weight * output_state. + if (!use_cifg) { + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + recurrent_to_input_weights->data.f, n_cell, n_output, + output_state->data.f, n_batch, input_gate_scratch, + /*result_stride=*/1); + } + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + recurrent_to_forget_weights->data.f, n_cell, n_output, + output_state->data.f, n_batch, forget_gate_scratch, + /*result_stride=*/1); + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + recurrent_to_cell_weights->data.f, n_cell, n_output, + output_state->data.f, n_batch, cell_scratch, /*result_stride=*/1); + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + recurrent_to_output_weights->data.f, n_cell, n_output, + output_state->data.f, n_batch, output_gate_scratch, + /*result_stride=*/1); + + // For each batch and cell: update input gate. + if (!use_cifg) { + if (use_peephole) { + tensor_utils::VectorBatchVectorCwiseProductAccumulate( + cell_to_input_weights->data.f, n_cell, cell_state->data.f, n_batch, + input_gate_scratch); + } + tensor_utils::ApplySigmoidToVector(input_gate_scratch, n_cell * n_batch, + input_gate_scratch); + } + + // For each batch and cell: update forget gate. + if (use_peephole) { + tensor_utils::VectorBatchVectorCwiseProductAccumulate( + cell_to_forget_weights->data.f, n_cell, cell_state->data.f, n_batch, + forget_gate_scratch); + } + tensor_utils::ApplySigmoidToVector(forget_gate_scratch, n_cell * n_batch, + forget_gate_scratch); + + // For each batch and cell: update the cell. + tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, + cell_state->data.f, n_batch * n_cell, + cell_state->data.f); + tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell, + params->activation, cell_scratch); + if (use_cifg) { + tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell, + forget_gate_scratch); + tensor_utils::VectorVectorCwiseProductAccumulate( + cell_scratch, forget_gate_scratch, n_batch * n_cell, + cell_state->data.f); + } else { + tensor_utils::VectorVectorCwiseProductAccumulate( + cell_scratch, input_gate_scratch, n_batch * n_cell, + cell_state->data.f); + } + if (params->cell_clip > 0.0) { + tensor_utils::ClipVector(cell_state->data.f, n_batch * n_cell, + params->cell_clip, cell_state->data.f); + } + + // For each batch and cell: update the output gate. + if (use_peephole) { + tensor_utils::VectorBatchVectorCwiseProductAccumulate( + cell_to_output_weights->data.f, n_cell, cell_state->data.f, n_batch, + output_gate_scratch); + } + tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell, + output_gate_scratch); + tensor_utils::ApplyActivationToVector(cell_state->data.f, n_batch * n_cell, + params->activation, cell_scratch); + tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch, + n_batch * n_cell, + output_gate_scratch); + + // For each batch: update the projection and output_state. + const bool use_projection_weight = (projection_weights != nullptr); + const bool use_projection_bias = (projection_bias != nullptr); + float* output_ptr_time = output->data.f + t * n_batch * n_output; + if (use_projection_weight) { + if (use_projection_bias) { + tensor_utils::VectorBatchVectorAssign(projection_bias->data.f, n_output, + n_batch, output_ptr_time); + } else { + tensor_utils::ZeroVector(output_ptr_time, n_batch * n_output); + } + tensor_utils::MatrixBatchVectorMultiplyAccumulate( + projection_weights->data.f, n_output, n_cell, output_gate_scratch, + n_batch, output_ptr_time, /*result_stride=*/1); + if (params->proj_clip > 0.0) { + tensor_utils::ClipVector(output_ptr_time, n_batch * n_output, + params->proj_clip, output_ptr_time); + } + } else { + tensor_utils::CopyVector(output_gate_scratch, n_batch * n_output, + output_ptr_time); + } + tensor_utils::CopyVector(output_ptr_time, n_batch * n_output, + output_state->data.f); + } + return kTfLiteOk; +} + +} // namespace unidirectional_sequence_lstm + +TfLiteRegistration* Register_UNIDIRECTIONAL_SEQUENCE_LSTM() { + static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, + unidirectional_sequence_lstm::Prepare, + unidirectional_sequence_lstm::Eval}; + return &r; +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm_test.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..93b635ae576e99854796d9fa997e5bf355b20534 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm_test.cc @@ -0,0 +1,1089 @@ +/* Copyright 2017 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. +==============================================================================*/ +// Unit test for TFLite Sequential LSTM op. + +#include +#include +#include + +#include +#include +#include "tensorflow/contrib/lite/interpreter.h" +#include "tensorflow/contrib/lite/kernels/register.h" +#include "tensorflow/contrib/lite/kernels/test_util.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { +namespace { + +using ::testing::ElementsAreArray; + +class UnidirectionalLSTMOpModel : public SingleOpModel { + public: + UnidirectionalLSTMOpModel(int n_batch, int n_input, int n_cell, int n_output, + int sequence_length, bool use_cifg, + bool use_peephole, bool use_projection_weights, + bool use_projection_bias, float cell_clip, + float proj_clip, + const std::vector>& input_shapes) + : n_batch_(n_batch), + n_input_(n_input), + n_cell_(n_cell), + n_output_(n_output), + sequence_length_(sequence_length) { + input_ = AddInput(TensorType_FLOAT32); + + if (use_cifg) { + input_to_input_weights_ = AddNullInput(); + } else { + input_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + + input_to_forget_weights_ = AddInput(TensorType_FLOAT32); + input_to_cell_weights_ = AddInput(TensorType_FLOAT32); + input_to_output_weights_ = AddInput(TensorType_FLOAT32); + + if (use_cifg) { + recurrent_to_input_weights_ = AddNullInput(); + } else { + recurrent_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + + recurrent_to_forget_weights_ = AddInput(TensorType_FLOAT32); + recurrent_to_cell_weights_ = AddInput(TensorType_FLOAT32); + recurrent_to_output_weights_ = AddInput(TensorType_FLOAT32); + + if (use_peephole) { + if (use_cifg) { + cell_to_input_weights_ = AddNullInput(); + } else { + cell_to_input_weights_ = AddInput(TensorType_FLOAT32); + } + cell_to_forget_weights_ = AddInput(TensorType_FLOAT32); + cell_to_output_weights_ = AddInput(TensorType_FLOAT32); + } else { + cell_to_input_weights_ = AddNullInput(); + cell_to_forget_weights_ = AddNullInput(); + cell_to_output_weights_ = AddNullInput(); + } + + if (use_cifg) { + input_gate_bias_ = AddNullInput(); + } else { + input_gate_bias_ = AddInput(TensorType_FLOAT32); + } + forget_gate_bias_ = AddInput(TensorType_FLOAT32); + cell_bias_ = AddInput(TensorType_FLOAT32); + output_gate_bias_ = AddInput(TensorType_FLOAT32); + + if (use_projection_weights) { + projection_weights_ = AddInput(TensorType_FLOAT32); + if (use_projection_bias) { + projection_bias_ = AddInput(TensorType_FLOAT32); + } else { + projection_bias_ = AddNullInput(); + } + } else { + projection_weights_ = AddNullInput(); + projection_bias_ = AddNullInput(); + } + + scratch_buffer_ = AddOutput(TensorType_FLOAT32); + // TODO(ghodrat): Modify these states when we have a permanent solution for + // persistent buffer. + output_state_ = AddOutput(TensorType_FLOAT32); + cell_state_ = AddOutput(TensorType_FLOAT32); + output_ = AddOutput(TensorType_FLOAT32); + + SetBuiltinOp(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, + BuiltinOptions_LSTMOptions, + CreateLSTMOptions(builder_, ActivationFunctionType_TANH, + cell_clip, proj_clip) + .Union()); + BuildInterpreter(input_shapes); + } + + void SetInputToInputWeights(std::initializer_list f) { + PopulateTensor(input_to_input_weights_, f); + } + + void SetInputToForgetWeights(std::initializer_list f) { + PopulateTensor(input_to_forget_weights_, f); + } + + void SetInputToCellWeights(std::initializer_list f) { + PopulateTensor(input_to_cell_weights_, f); + } + + void SetInputToOutputWeights(std::initializer_list f) { + PopulateTensor(input_to_output_weights_, f); + } + + void SetRecurrentToInputWeights(std::initializer_list f) { + PopulateTensor(recurrent_to_input_weights_, f); + } + + void SetRecurrentToForgetWeights(std::initializer_list f) { + PopulateTensor(recurrent_to_forget_weights_, f); + } + + void SetRecurrentToCellWeights(std::initializer_list f) { + PopulateTensor(recurrent_to_cell_weights_, f); + } + + void SetRecurrentToOutputWeights(std::initializer_list f) { + PopulateTensor(recurrent_to_output_weights_, f); + } + + void SetCellToInputWeights(std::initializer_list f) { + PopulateTensor(cell_to_input_weights_, f); + } + + void SetCellToForgetWeights(std::initializer_list f) { + PopulateTensor(cell_to_forget_weights_, f); + } + + void SetCellToOutputWeights(std::initializer_list f) { + PopulateTensor(cell_to_output_weights_, f); + } + + void SetInputGateBias(std::initializer_list f) { + PopulateTensor(input_gate_bias_, f); + } + + void SetForgetGateBias(std::initializer_list f) { + PopulateTensor(forget_gate_bias_, f); + } + + void SetCellBias(std::initializer_list f) { + PopulateTensor(cell_bias_, f); + } + + void SetOutputGateBias(std::initializer_list f) { + PopulateTensor(output_gate_bias_, f); + } + + void SetProjectionWeights(std::initializer_list f) { + PopulateTensor(projection_weights_, f); + } + + void SetProjectionBias(std::initializer_list f) { + PopulateTensor(projection_bias_, f); + } + + void ResetOutputState() { + const int zero_buffer_size = n_cell_ * n_batch_; + std::unique_ptr zero_buffer(new float[zero_buffer_size]); + memset(zero_buffer.get(), 0, zero_buffer_size * sizeof(float)); + PopulateTensor(output_state_, 0, zero_buffer.get(), + zero_buffer.get() + zero_buffer_size); + } + + void ResetCellState() { + const int zero_buffer_size = n_cell_ * n_batch_; + std::unique_ptr zero_buffer(new float[zero_buffer_size]); + memset(zero_buffer.get(), 0, zero_buffer_size * sizeof(float)); + PopulateTensor(cell_state_, 0, zero_buffer.get(), + zero_buffer.get() + zero_buffer_size); + } + + void SetInput(int offset, float* begin, float* end) { + PopulateTensor(input_, offset, begin, end); + } + + std::vector GetOutput() { return ExtractVector(output_); } + + int num_inputs() { return n_input_; } + int num_outputs() { return n_output_; } + int num_cells() { return n_cell_; } + int num_batches() { return n_batch_; } + int sequence_length() { return sequence_length_; } + + private: + int input_; + int input_to_input_weights_; + int input_to_forget_weights_; + int input_to_cell_weights_; + int input_to_output_weights_; + + int recurrent_to_input_weights_; + int recurrent_to_forget_weights_; + int recurrent_to_cell_weights_; + int recurrent_to_output_weights_; + + int cell_to_input_weights_; + int cell_to_forget_weights_; + int cell_to_output_weights_; + + int input_gate_bias_; + int forget_gate_bias_; + int cell_bias_; + int output_gate_bias_; + + int projection_weights_; + int projection_bias_; + + int output_; + int output_state_; + int cell_state_; + int scratch_buffer_; + + int n_batch_; + int n_input_; + int n_cell_; + int n_output_; + int sequence_length_; +}; + +TEST(LSTMOpTest, BlackBoxTestNoCifgNoPeepholeNoProjectionNoClipping) { + const int n_batch = 1; + const int n_input = 2; + // n_cell and n_output have the same size when there is no projection. + const int n_cell = 4; + const int n_output = 4; + const int sequence_length = 3; + + UnidirectionalLSTMOpModel lstm( + n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/false, + /*use_peephole=*/false, /*use_projection_weights=*/false, + /*use_projection_bias=*/false, /*cell_clip=*/0.0, /*proj_clip=*/0.0, + { + {sequence_length, n_batch, n_input}, // input tensor + + {n_cell, n_input}, // input_to_input_weight tensor + {n_cell, n_input}, // input_to_forget_weight tensor + {n_cell, n_input}, // input_to_cell_weight tensor + {n_cell, n_input}, // input_to_output_weight tensor + + {n_cell, n_output}, // recurrent_to_input_weight tensor + {n_cell, n_output}, // recurrent_to_forget_weight tensor + {n_cell, n_output}, // recurrent_to_cell_weight tensor + {n_cell, n_output}, // recurrent_to_output_weight tensor + + {0}, // cell_to_input_weight tensor + {0}, // cell_to_forget_weight tensor + {0}, // cell_to_output_weight tensor + + {n_cell}, // input_gate_bias tensor + {n_cell}, // forget_gate_bias tensor + {n_cell}, // cell_bias tensor + {n_cell}, // output_gate_bias tensor + + {0, 0}, // projection_weight tensor + {0}, // projection_bias tensor + }); + + lstm.SetInputToInputWeights({-0.45018822, -0.02338299, -0.0870589, + -0.34550029, 0.04266912, -0.15680569, + -0.34856534, 0.43890524}); + + lstm.SetInputToCellWeights({-0.50013041, 0.1370284, 0.11810488, 0.2013163, + -0.20583314, 0.44344562, 0.22077113, + -0.29909778}); + + lstm.SetInputToForgetWeights({0.09701663, 0.20334584, -0.50592935, + -0.31343272, -0.40032279, 0.44781327, + 0.01387155, -0.35593212}); + + lstm.SetInputToOutputWeights({-0.25065863, -0.28290087, 0.04613829, + 0.40525138, 0.44272184, 0.03897077, -0.1556896, + 0.19487578}); + + lstm.SetInputGateBias({0., 0., 0., 0.}); + + lstm.SetCellBias({0., 0., 0., 0.}); + + lstm.SetForgetGateBias({1., 1., 1., 1.}); + + lstm.SetOutputGateBias({0., 0., 0., 0.}); + + lstm.SetRecurrentToInputWeights( + {-0.0063535, -0.2042388, 0.31454784, -0.35746509, 0.28902304, 0.08183324, + -0.16555229, 0.02286911, -0.13566875, 0.03034258, 0.48091322, + -0.12528998, 0.24077177, -0.51332325, -0.33502164, 0.10629296}); + + lstm.SetRecurrentToCellWeights( + {-0.3407414, 0.24443203, -0.2078532, 0.26320225, 0.05695659, -0.00123841, + -0.4744786, -0.35869038, -0.06418842, -0.13502428, -0.501764, 0.22830659, + -0.46367589, 0.26016325, -0.03894562, -0.16368064}); + + lstm.SetRecurrentToForgetWeights( + {-0.48684245, -0.06655136, 0.42224967, 0.2112639, 0.27654213, 0.20864892, + -0.07646349, 0.45877004, 0.00141793, -0.14609534, 0.36447752, 0.09196436, + 0.28053468, 0.01560611, -0.20127171, -0.01140004}); + + lstm.SetRecurrentToOutputWeights( + {0.43385774, -0.17194885, 0.2718237, 0.09215671, 0.24107647, -0.39835793, + 0.18212086, 0.01301402, 0.48572797, -0.50656658, 0.20047462, -0.20607421, + -0.51818722, -0.15390486, 0.0468148, 0.39922136}); + + // Input should have n_input * sequence_length many values. + static float lstm_input[] = {2., 3., 3., 4., 1., 1.}; + static float lstm_golden_output[] = {-0.02973187, 0.1229473, 0.20885126, + -0.15358765, -0.03716109, 0.12507336, + 0.41193449, -0.20860538, -0.15053082, + 0.09120187, 0.24278517, -0.12222792}; + + // Resetting cell_state and output_state + lstm.ResetCellState(); + lstm.ResetOutputState(); + + float* batch0_start = lstm_input; + float* batch0_end = batch0_start + lstm.num_inputs() * lstm.sequence_length(); + + lstm.SetInput(0, batch0_start, batch0_end); + + lstm.Invoke(); + + float* golden_start = lstm_golden_output; + float* golden_end = + golden_start + lstm.num_outputs() * lstm.sequence_length(); + std::vector expected; + expected.insert(expected.end(), golden_start, golden_end); + EXPECT_THAT(lstm.GetOutput(), ElementsAreArray(ArrayFloatNear(expected))); +} + +TEST(LSTMOpTest, BlackBoxTestWithCifgWithPeepholeNoProjectionNoClipping) { + const int n_batch = 1; + const int n_input = 2; + // n_cell and n_output have the same size when there is no projection. + const int n_cell = 4; + const int n_output = 4; + const int sequence_length = 3; + + UnidirectionalLSTMOpModel lstm( + n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/true, + /*use_peephole=*/true, /*use_projection_weights=*/false, + /*use_projection_bias=*/false, /*cell_clip=*/0.0, /*proj_clip=*/0.0, + { + {sequence_length, n_batch, n_input}, // input tensor + + {0, 0}, // input_to_input_weight tensor + {n_cell, n_input}, // input_to_forget_weight tensor + {n_cell, n_input}, // input_to_cell_weight tensor + {n_cell, n_input}, // input_to_output_weight tensor + + {0, 0}, // recurrent_to_input_weight tensor + {n_cell, n_output}, // recurrent_to_forget_weight tensor + {n_cell, n_output}, // recurrent_to_cell_weight tensor + {n_cell, n_output}, // recurrent_to_output_weight tensor + + {0}, // cell_to_input_weight tensor + {n_cell}, // cell_to_forget_weight tensor + {n_cell}, // cell_to_output_weight tensor + + {0}, // input_gate_bias tensor + {n_cell}, // forget_gate_bias tensor + {n_cell}, // cell_bias tensor + {n_cell}, // output_gate_bias tensor + + {0, 0}, // projection_weight tensor + {0}, // projection_bias tensor + }); + + lstm.SetInputToCellWeights({-0.49770179, -0.27711356, -0.09624726, 0.05100781, + 0.04717243, 0.48944736, -0.38535351, + -0.17212132}); + + lstm.SetInputToForgetWeights({-0.55291498, -0.42866567, 0.13056988, + -0.3633365, -0.22755712, 0.28253698, 0.24407166, + 0.33826375}); + + lstm.SetInputToOutputWeights({0.10725588, -0.02335852, -0.55932593, + -0.09426838, -0.44257352, 0.54939759, + 0.01533556, 0.42751634}); + + lstm.SetCellBias({0., 0., 0., 0.}); + + lstm.SetForgetGateBias({1., 1., 1., 1.}); + + lstm.SetOutputGateBias({0., 0., 0., 0.}); + + lstm.SetRecurrentToCellWeights( + {0.54066205, -0.32668582, -0.43562764, -0.56094903, 0.42957711, + 0.01841056, -0.32764608, -0.33027974, -0.10826075, 0.20675004, + 0.19069612, -0.03026325, -0.54532051, 0.33003211, 0.44901288, + 0.21193194}); + + lstm.SetRecurrentToForgetWeights( + {-0.13832897, -0.0515101, -0.2359007, -0.16661474, -0.14340827, + 0.36986142, 0.23414481, 0.55899, 0.10798943, -0.41174671, 0.17751795, + -0.34484994, -0.35874045, -0.11352962, 0.27268326, 0.54058349}); + + lstm.SetRecurrentToOutputWeights( + {0.41613156, 0.42610586, -0.16495961, -0.5663873, 0.30579174, -0.05115908, + -0.33941799, 0.23364776, 0.11178309, 0.09481031, -0.26424935, 0.46261835, + 0.50248802, 0.26114327, -0.43736315, 0.33149987}); + + lstm.SetCellToForgetWeights( + {0.47485286, -0.51955009, -0.24458408, 0.31544167}); + lstm.SetCellToOutputWeights( + {-0.17135078, 0.82760304, 0.85573703, -0.77109635}); + + static float lstm_input[] = {2., 3., 3., 4., 1., 1.}; + static float lstm_golden_output[] = {-0.36444446, -0.00352185, 0.12886585, + -0.05163646, -0.42312205, -0.01218222, + 0.24201041, -0.08124574, -0.358325, + -0.04621704, 0.21641694, -0.06471302}; + + // Resetting cell_state and output_state + lstm.ResetCellState(); + lstm.ResetOutputState(); + + float* batch0_start = lstm_input; + float* batch0_end = batch0_start + lstm.num_inputs() * lstm.sequence_length(); + + lstm.SetInput(0, batch0_start, batch0_end); + + lstm.Invoke(); + + float* golden_start = lstm_golden_output; + float* golden_end = + golden_start + lstm.num_outputs() * lstm.sequence_length(); + std::vector expected; + expected.insert(expected.end(), golden_start, golden_end); + EXPECT_THAT(lstm.GetOutput(), ElementsAreArray(ArrayFloatNear(expected))); +} + +TEST(LSTMOpTest, BlackBoxTestWithPeepholeWithProjectionNoClipping) { + const int n_batch = 2; + const int n_input = 5; + const int n_cell = 20; + const int n_output = 16; + const int sequence_length = 4; + + UnidirectionalLSTMOpModel lstm( + n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/false, + /*use_peephole=*/true, /*use_projection_weights=*/true, + /*use_projection_bias=*/false, + /*cell_clip=*/0.0, /*proj_clip=*/0.0, + { + {sequence_length, n_batch, n_input}, // input tensor + + {n_cell, n_input}, // input_to_input_weight tensor + {n_cell, n_input}, // input_to_forget_weight tensor + {n_cell, n_input}, // input_to_cell_weight tensor + {n_cell, n_input}, // input_to_output_weight tensor + + {n_cell, n_output}, // recurrent_to_input_weight tensor + {n_cell, n_output}, // recurrent_to_forget_weight tensor + {n_cell, n_output}, // recurrent_to_cell_weight tensor + {n_cell, n_output}, // recurrent_to_output_weight tensor + + {n_cell}, // cell_to_input_weight tensor + {n_cell}, // cell_to_forget_weight tensor + {n_cell}, // cell_to_output_weight tensor + + {n_cell}, // input_gate_bias tensor + {n_cell}, // forget_gate_bias tensor + {n_cell}, // cell_bias tensor + {n_cell}, // output_gate_bias tensor + + {n_output, n_cell}, // projection_weight tensor + {0}, // projection_bias tensor + }); + + lstm.SetInputToInputWeights( + {0.021393683, 0.06124551, 0.046905167, -0.014657677, -0.03149463, + 0.09171803, 0.14647801, 0.10797193, -0.0057968358, 0.0019193048, + -0.2726754, 0.10154029, -0.018539885, 0.080349885, -0.10262385, + -0.022599787, -0.09121155, -0.008675967, -0.045206103, -0.0821282, + -0.008045952, 0.015478081, 0.055217247, 0.038719587, 0.044153627, + -0.06453243, 0.05031825, -0.046935108, -0.008164439, 0.014574226, + -0.1671009, -0.15519552, -0.16819797, -0.13971269, -0.11953059, + 0.25005487, -0.22790983, 0.009855087, -0.028140958, -0.11200698, + 0.11295408, -0.0035217577, 0.054485075, 0.05184695, 0.064711206, + 0.10989193, 0.11674786, 0.03490607, 0.07727357, 0.11390585, + -0.1863375, -0.1034451, -0.13945189, -0.049401227, -0.18767063, + 0.042483903, 0.14233552, 0.13832581, 0.18350165, 0.14545603, + -0.028545704, 0.024939531, 0.050929718, 0.0076203286, -0.0029723682, + -0.042484224, -0.11827596, -0.09171104, -0.10808628, -0.16327988, + -0.2273378, -0.0993647, -0.017155107, 0.0023917493, 0.049272764, + 0.0038534778, 0.054764505, 0.089753784, 0.06947234, 0.08014476, + -0.04544234, -0.0497073, -0.07135631, -0.048929106, -0.004042012, + -0.009284026, 0.018042054, 0.0036860977, -0.07427302, -0.11434604, + -0.018995456, 0.031487543, 0.012834908, 0.019977754, 0.044256654, + -0.39292613, -0.18519334, -0.11651281, -0.06809892, 0.011373677}); + + lstm.SetInputToForgetWeights( + {-0.0018401089, -0.004852237, 0.03698424, 0.014181704, 0.028273236, + -0.016726194, -0.05249759, -0.10204261, 0.00861066, -0.040979505, + -0.009899187, 0.01923892, -0.028177269, -0.08535103, -0.14585495, + 0.10662567, -0.01909731, -0.017883534, -0.0047269356, -0.045103323, + 0.0030784295, 0.076784775, 0.07463696, 0.094531395, 0.0814421, + -0.12257899, -0.033945758, -0.031303465, 0.045630626, 0.06843887, + -0.13492945, -0.012480007, -0.0811829, -0.07224499, -0.09628791, + 0.045100946, 0.0012300825, 0.013964662, 0.099372394, 0.02543059, + 0.06958324, 0.034257296, 0.0482646, 0.06267997, 0.052625068, + 0.12784666, 0.07077897, 0.025725935, 0.04165009, 0.07241905, + 0.018668644, -0.037377294, -0.06277783, -0.08833636, -0.040120605, + -0.011405586, -0.007808335, -0.010301386, -0.005102167, 0.027717464, + 0.05483423, 0.11449111, 0.11289652, 0.10939839, 0.13396506, + -0.08402166, -0.01901462, -0.044678304, -0.07720565, 0.014350063, + -0.11757958, -0.0652038, -0.08185733, -0.076754324, -0.092614375, + 0.10405491, 0.052960336, 0.035755895, 0.035839386, -0.012540553, + 0.036881298, 0.02913376, 0.03420159, 0.05448447, -0.054523353, + 0.02582715, 0.02327355, -0.011857179, -0.0011980024, -0.034641717, + -0.026125094, -0.17582615, -0.15923657, -0.27486774, -0.0006143371, + 0.0001771948, -8.470171e-05, 0.02651807, 0.045790765, 0.06956496}); + + lstm.SetInputToCellWeights( + {-0.04580283, -0.09549462, -0.032418985, -0.06454633, + -0.043528453, 0.043018587, -0.049152344, -0.12418144, + -0.078985475, -0.07596889, 0.019484362, -0.11434962, + -0.0074034138, -0.06314844, -0.092981495, 0.0062155537, + -0.025034338, -0.0028890965, 0.048929527, 0.06235075, + 0.10665918, -0.032036792, -0.08505916, -0.10843358, + -0.13002433, -0.036816437, -0.02130134, -0.016518239, + 0.0047691227, -0.0025825808, 0.066017866, 0.029991534, + -0.10652836, -0.1037554, -0.13056071, -0.03266643, + -0.033702414, -0.006473424, -0.04611692, 0.014419339, + -0.025174323, 0.0396852, 0.081777506, 0.06157468, + 0.10210095, -0.009658194, 0.046511717, 0.03603906, + 0.0069369148, 0.015960095, -0.06507666, 0.09551598, + 0.053568836, 0.06408714, 0.12835667, -0.008714329, + -0.20211966, -0.12093674, 0.029450472, 0.2849013, + -0.029227901, 0.1164364, -0.08560263, 0.09941786, + -0.036999565, -0.028842626, -0.0033637602, -0.017012902, + -0.09720865, -0.11193351, -0.029155117, -0.017936034, + -0.009768936, -0.04223324, -0.036159635, 0.06505112, + -0.021742892, -0.023377212, -0.07221364, -0.06430552, + 0.05453865, 0.091149814, 0.06387331, 0.007518393, + 0.055960953, 0.069779344, 0.046411168, 0.10509911, + 0.07463894, 0.0075130584, 0.012850982, 0.04555431, + 0.056955688, 0.06555285, 0.050801456, -0.009862683, + 0.00826772, -0.026555609, -0.0073611983, -0.0014897042}); + + lstm.SetInputToOutputWeights( + {-0.0998932, -0.07201956, -0.052803773, -0.15629593, -0.15001918, + -0.07650751, 0.02359855, -0.075155355, -0.08037709, -0.15093534, + 0.029517552, -0.04751393, 0.010350531, -0.02664851, -0.016839722, + -0.023121163, 0.0077019283, 0.012851257, -0.05040649, -0.0129761, + -0.021737747, -0.038305793, -0.06870586, -0.01481247, -0.001285394, + 0.10124236, 0.083122835, 0.053313006, -0.062235646, -0.075637154, + -0.027833903, 0.029774971, 0.1130802, 0.09218906, 0.09506135, + -0.086665764, -0.037162706, -0.038880914, -0.035832845, -0.014481564, + -0.09825003, -0.12048569, -0.097665586, -0.05287633, -0.0964047, + -0.11366429, 0.035777505, 0.13568819, 0.052451383, 0.050649304, + 0.05798951, -0.021852335, -0.099848844, 0.014740475, -0.078897946, + 0.04974699, 0.014160473, 0.06973932, 0.04964942, 0.033364646, + 0.08190124, 0.025535367, 0.050893165, 0.048514254, 0.06945813, + -0.078907564, -0.06707616, -0.11844508, -0.09986688, -0.07509403, + 0.06263226, 0.14925587, 0.20188436, 0.12098451, 0.14639415, + 0.0015017595, -0.014267382, -0.03417257, 0.012711468, 0.0028300495, + -0.024758482, -0.05098548, -0.0821182, 0.014225672, 0.021544158, + 0.08949725, 0.07505268, -0.0020780868, 0.04908258, 0.06476295, + -0.022907063, 0.027562456, 0.040185735, 0.019567577, -0.015598739, + -0.049097303, -0.017121866, -0.083368234, -0.02332002, -0.0840956}); + + lstm.SetInputGateBias( + {0.02234832, 0.14757581, 0.18176508, 0.10380666, 0.053110216, + -0.06928846, -0.13942584, -0.11816189, 0.19483899, 0.03652339, + -0.10250295, 0.036714908, -0.18426876, 0.036065217, 0.21810818, + 0.02383196, -0.043370757, 0.08690144, -0.04444982, 0.00030581196}); + + lstm.SetForgetGateBias({0.035185695, -0.042891346, -0.03032477, 0.23027696, + 0.11098921, 0.15378423, 0.09263801, 0.09790885, + 0.09508917, 0.061199076, 0.07665568, -0.015443159, + -0.03499149, 0.046190713, 0.08895977, 0.10899629, + 0.40694186, 0.06030037, 0.012413437, -0.06108739}); + + lstm.SetCellBias({-0.024379363, 0.0055531194, 0.23377132, 0.033463873, + -0.1483596, -0.10639995, -0.091433935, 0.058573797, + -0.06809782, -0.07889636, -0.043246906, -0.09829136, + -0.4279842, 0.034901652, 0.18797937, 0.0075234566, + 0.016178843, 0.1749513, 0.13975595, 0.92058027}); + + lstm.SetOutputGateBias( + {0.046159424, -0.0012809046, 0.03563469, 0.12648113, 0.027195795, + 0.35373217, -0.018957434, 0.008907322, -0.0762701, 0.12018895, + 0.04216877, 0.0022856654, 0.040952638, 0.3147856, 0.08225149, + -0.057416286, -0.14995944, -0.008040261, 0.13208859, 0.029760877}); + + lstm.SetRecurrentToInputWeights( + {-0.001374326, -0.078856036, 0.10672688, 0.029162422, + -0.11585556, 0.02557986, -0.13446963, -0.035785314, + -0.01244275, 0.025961924, -0.02337298, -0.044228926, + -0.055839065, -0.046598054, -0.010546039, -0.06900766, + 0.027239809, 0.022582639, -0.013296484, -0.05459212, + 0.08981, -0.045407712, 0.08682226, -0.06867011, + -0.14390695, -0.02916037, 0.000996957, 0.091420636, + 0.14283475, -0.07390571, -0.06402044, 0.062524505, + -0.093129106, 0.04860203, -0.08364217, -0.08119002, + 0.009352075, 0.22920375, 0.0016303885, 0.11583097, + -0.13732095, 0.012405723, -0.07551853, 0.06343048, + 0.12162708, -0.031923793, -0.014335606, 0.01790974, + -0.10650317, -0.0724401, 0.08554849, -0.05727212, + 0.06556731, -0.042729504, -0.043227166, 0.011683251, + -0.013082158, -0.029302018, -0.010899579, -0.062036745, + -0.022509435, -0.00964907, -0.01567329, 0.04260106, + -0.07787477, -0.11576462, 0.017356863, 0.048673786, + -0.017577527, -0.05527947, -0.082487635, -0.040137455, + -0.10820036, -0.04666372, 0.022746278, -0.07851417, + 0.01068115, 0.032956902, 0.022433773, 0.0026891115, + 0.08944216, -0.0685835, 0.010513544, 0.07228705, + 0.02032331, -0.059686817, -0.0005566496, -0.086984694, + 0.040414046, -0.1380399, 0.094208956, -0.05722982, + 0.012092817, -0.04989123, -0.086576, -0.003399834, + -0.04696032, -0.045747425, 0.10091314, 0.048676282, + -0.029037097, 0.031399418, -0.0040285117, 0.047237843, + 0.09504992, 0.041799378, -0.049185462, -0.031518843, + -0.10516937, 0.026374253, 0.10058866, -0.0033195973, + -0.041975245, 0.0073591834, 0.0033782164, -0.004325073, + -0.10167381, 0.042500053, -0.01447153, 0.06464186, + -0.017142897, 0.03312627, 0.009205989, 0.024138335, + -0.011337001, 0.035530265, -0.010912711, 0.0706555, + -0.005894094, 0.051841937, -0.1401738, -0.02351249, + 0.0365468, 0.07590991, 0.08838724, 0.021681072, + -0.10086113, 0.019608743, -0.06195883, 0.077335775, + 0.023646897, -0.095322326, 0.02233014, 0.09756986, + -0.048691444, -0.009579111, 0.07595467, 0.11480546, + -0.09801813, 0.019894179, 0.08502348, 0.004032281, + 0.037211012, 0.068537936, -0.048005626, -0.091520436, + -0.028379958, -0.01556313, 0.06554592, -0.045599163, + -0.01672207, -0.020169014, -0.011877351, -0.20212261, + 0.010889619, 0.0047078193, 0.038385306, 0.08540671, + -0.017140968, -0.0035865551, 0.016678626, 0.005633034, + 0.015963363, 0.00871737, 0.060130805, 0.028611384, + 0.10109069, -0.015060172, -0.07894427, 0.06401885, + 0.011584063, -0.024466386, 0.0047652307, -0.09041358, + 0.030737216, -0.0046374933, 0.14215417, -0.11823516, + 0.019899689, 0.006106124, -0.027092824, 0.0786356, + 0.05052217, -0.058925, -0.011402121, -0.024987547, + -0.0013661642, -0.06832946, -0.015667673, -0.1083353, + -0.00096863037, -0.06988685, -0.053350925, -0.027275559, + -0.033664223, -0.07978348, -0.025200296, -0.017207067, + -0.058403496, -0.055697463, 0.005798788, 0.12965427, + -0.062582195, 0.0013350133, -0.10482091, 0.0379771, + 0.072521195, -0.0029455067, -0.13797039, -0.03628521, + 0.013806405, -0.017858358, -0.01008298, -0.07700066, + -0.017081132, 0.019358726, 0.0027079724, 0.004635139, + 0.062634714, -0.02338735, -0.039547626, -0.02050681, + 0.03385117, -0.083611414, 0.002862572, -0.09421313, + 0.058618143, -0.08598433, 0.00972939, 0.023867095, + -0.053934585, -0.023203006, 0.07452513, -0.048767887, + -0.07314807, -0.056307215, -0.10433547, -0.06440842, + 0.04328182, 0.04389765, -0.020006588, -0.09076438, + -0.11652589, -0.021705797, 0.03345259, -0.010329105, + -0.025767034, 0.013057034, -0.07316461, -0.10145612, + 0.06358255, 0.18531723, 0.07759293, 0.12006465, + 0.1305557, 0.058638252, -0.03393652, 0.09622831, + -0.16253184, -2.4580743e-06, 0.079869635, -0.070196845, + -0.005644518, 0.06857898, -0.12598175, -0.035084512, + 0.03156317, -0.12794146, -0.031963028, 0.04692781, + 0.030070418, 0.0071660685, -0.095516115, -0.004643372, + 0.040170413, -0.062104587, -0.0037324072, 0.0554317, + 0.08184801, -0.019164372, 0.06791302, 0.034257166, + -0.10307039, 0.021943003, 0.046745934, 0.0790918, + -0.0265588, -0.007824208, 0.042546265, -0.00977924, + -0.0002440307, -0.017384544, -0.017990116, 0.12252321, + -0.014512694, -0.08251313, 0.08861942, 0.13589665, + 0.026351685, 0.012641483, 0.07466548, 0.044301085, + -0.045414884, -0.051112458, 0.03444247, -0.08502782, + -0.04106223, -0.028126027, 0.028473156, 0.10467447}); + + lstm.SetRecurrentToForgetWeights( + {-0.057784554, -0.026057621, -0.068447545, -0.022581743, + 0.14811787, 0.10826372, 0.09471067, 0.03987225, + -0.0039523416, 0.00030638507, 0.053185795, 0.10572994, + 0.08414449, -0.022036452, -0.00066928595, -0.09203576, + 0.032950465, -0.10985798, -0.023809856, 0.0021431844, + -0.02196096, -0.00326074, 0.00058621005, -0.074678116, + -0.06193199, 0.055729095, 0.03736828, 0.020123724, + 0.061878487, -0.04729229, 0.034919553, -0.07585433, + -0.04421272, -0.044019096, 0.085488975, 0.04058006, + -0.06890133, -0.030951202, -0.024628663, -0.07672815, + 0.034293607, 0.08556707, -0.05293577, -0.033561368, + -0.04899627, 0.0241671, 0.015736353, -0.095442444, + -0.029564252, 0.016493602, -0.035026584, 0.022337519, + -0.026871363, 0.004780428, 0.0077918363, -0.03601621, + 0.016435321, -0.03263031, -0.09543275, -0.047392778, + 0.013454138, 0.028934088, 0.01685226, -0.086110644, + -0.046250615, -0.01847454, 0.047608484, 0.07339695, + 0.034546845, -0.04881143, 0.009128804, -0.08802852, + 0.03761666, 0.008096139, -0.014454086, 0.014361001, + -0.023502491, -0.0011840804, -0.07607001, 0.001856849, + -0.06509276, -0.006021153, -0.08570962, -0.1451793, + 0.060212336, 0.055259194, 0.06974018, 0.049454916, + -0.027794661, -0.08077226, -0.016179763, 0.1169753, + 0.17213494, -0.0056326236, -0.053934924, -0.0124349, + -0.11520337, 0.05409887, 0.088759385, 0.0019655675, + 0.0042065294, 0.03881498, 0.019844765, 0.041858196, + -0.05695512, 0.047233116, 0.038937137, -0.06542224, + 0.014429736, -0.09719407, 0.13908425, -0.05379757, + 0.012321099, 0.082840554, -0.029899208, 0.044217527, + 0.059855383, 0.07711018, -0.045319796, 0.0948846, + -0.011724666, -0.0033288454, -0.033542685, -0.04764985, + -0.13873616, 0.040668588, 0.034832682, -0.015319203, + -0.018715994, 0.046002675, 0.0599172, -0.043107376, + 0.0294216, -0.002314414, -0.022424703, 0.0030315618, + 0.0014641669, 0.0029166266, -0.11878115, 0.013738511, + 0.12375372, -0.0006038222, 0.029104086, 0.087442465, + 0.052958444, 0.07558703, 0.04817258, 0.044462286, + -0.015213451, -0.08783778, -0.0561384, -0.003008196, + 0.047060397, -0.002058388, 0.03429439, -0.018839769, + 0.024734668, 0.024614193, -0.042046934, 0.09597743, + -0.0043254104, 0.04320769, 0.0064070094, -0.0019131786, + -0.02558259, -0.022822596, -0.023273505, -0.02464396, + -0.10991725, -0.006240552, 0.0074488563, 0.024044557, + 0.04383914, -0.046476185, 0.028658995, 0.060410924, + 0.050786525, 0.009452605, -0.0073054377, -0.024810238, + 0.0052906186, 0.0066939713, -0.0020913032, 0.014515517, + 0.015898481, 0.021362653, -0.030262267, 0.016587038, + -0.011442813, 0.041154444, -0.007631438, -0.03423484, + -0.010977775, 0.036152758, 0.0066366293, 0.11915515, + 0.02318443, -0.041350313, 0.021485701, -0.10906167, + -0.028218046, -0.00954771, 0.020531068, -0.11995105, + -0.03672871, 0.024019798, 0.014255957, -0.05221243, + -0.00661567, -0.04630967, 0.033188973, 0.10107534, + -0.014027541, 0.030796422, -0.10270911, -0.035999842, + 0.15443139, 0.07684145, 0.036571592, -0.035900835, + -0.0034699554, 0.06209149, 0.015920248, -0.031122351, + -0.03858649, 0.01849943, 0.13872518, 0.01503974, + 0.069941424, -0.06948533, -0.0088794185, 0.061282158, + -0.047401894, 0.03100163, -0.041533746, -0.10430945, + 0.044574402, -0.01425562, -0.024290353, 0.034563623, + 0.05866852, 0.023947537, -0.09445152, 0.035450947, + 0.02247216, -0.0042998926, 0.061146557, -0.10250651, + 0.020881841, -0.06747029, 0.10062043, -0.0023941975, + 0.03532124, -0.016341697, 0.09685456, -0.016764693, + 0.051808182, 0.05875331, -0.04536488, 0.001626336, + -0.028892258, -0.01048663, -0.009793449, -0.017093895, + 0.010987891, 0.02357273, -0.00010856845, 0.0099760275, + -0.001845119, -0.03551521, 0.0018358806, 0.05763657, + -0.01769146, 0.040995963, 0.02235177, -0.060430344, + 0.11475477, -0.023854522, 0.10071741, 0.0686208, + -0.014250481, 0.034261297, 0.047418304, 0.08562733, + -0.030519066, 0.0060542435, 0.014653856, -0.038836084, + 0.04096551, 0.032249358, -0.08355519, -0.026823482, + 0.056386515, -0.010401743, -0.028396193, 0.08507674, + 0.014410365, 0.020995233, 0.17040324, 0.11511526, + 0.02459721, 0.0066619175, 0.025853224, -0.023133837, + -0.081302024, 0.017264642, -0.009585969, 0.09491168, + -0.051313367, 0.054532815, -0.014298593, 0.10657464, + 0.007076659, 0.10964551, 0.0409152, 0.008275321, + -0.07283536, 0.07937492, 0.04192024, -0.1075027}); + + lstm.SetRecurrentToCellWeights( + {-0.037322544, 0.018592842, 0.0056175636, -0.06253426, + 0.055647098, -0.05713207, -0.05626563, 0.005559383, + 0.03375411, -0.025757805, -0.088049285, 0.06017052, + -0.06570978, 0.007384076, 0.035123326, -0.07920549, + 0.053676967, 0.044480428, -0.07663568, 0.0071805613, + 0.08089997, 0.05143358, 0.038261272, 0.03339287, + -0.027673481, 0.044746667, 0.028349208, 0.020090483, + -0.019443132, -0.030755889, -0.0040000007, 0.04465846, + -0.021585021, 0.0031670958, 0.0053199246, -0.056117613, + -0.10893326, 0.076739706, -0.08509834, -0.027997585, + 0.037871376, 0.01449768, -0.09002357, -0.06111149, + -0.046195522, 0.0422062, -0.005683705, -0.1253618, + -0.012925729, -0.04890792, 0.06985068, 0.037654128, + 0.03398274, -0.004781977, 0.007032333, -0.031787455, + 0.010868644, -0.031489216, 0.09525667, 0.013939797, + 0.0058680447, 0.0167067, 0.02668468, -0.04797466, + -0.048885044, -0.12722108, 0.035304096, 0.06554885, + 0.00972396, -0.039238118, -0.05159735, -0.11329045, + 0.1613692, -0.03750952, 0.06529313, -0.071974665, + -0.11769596, 0.015524369, -0.0013754242, -0.12446318, + 0.02786344, -0.014179351, 0.005264273, 0.14376344, + 0.015983658, 0.03406988, -0.06939408, 0.040699873, + 0.02111075, 0.09669095, 0.041345075, -0.08316494, + -0.07684199, -0.045768797, 0.032298047, -0.041805092, + 0.0119405, 0.0061010392, 0.12652606, 0.0064572375, + -0.024950314, 0.11574242, 0.04508852, -0.04335324, + 0.06760663, -0.027437469, 0.07216407, 0.06977076, + -0.05438599, 0.034033038, -0.028602652, 0.05346137, + 0.043184172, -0.037189785, 0.10420091, 0.00882477, + -0.054019816, -0.074273005, -0.030617684, -0.0028467078, + 0.024302477, -0.0038869337, 0.005332455, 0.0013399826, + 0.04361412, -0.007001822, 0.09631092, -0.06702025, + -0.042049985, -0.035070654, -0.04103342, -0.10273396, + 0.0544271, 0.037184782, -0.13150354, -0.0058036847, + -0.008264958, 0.042035464, 0.05891794, 0.029673764, + 0.0063542654, 0.044788733, 0.054816857, 0.062257513, + -0.00093483756, 0.048938446, -0.004952862, -0.007730018, + -0.04043371, -0.017094059, 0.07229206, -0.023670016, + -0.052195564, -0.025616996, -0.01520939, 0.045104615, + -0.007376126, 0.003533447, 0.006570588, 0.056037236, + 0.12436656, 0.051817212, 0.028532185, -0.08686856, + 0.11868599, 0.07663395, -0.07323171, 0.03463402, + -0.050708205, -0.04458982, -0.11590894, 0.021273347, + 0.1251325, -0.15313013, -0.12224372, 0.17228661, + 0.023029093, 0.086124025, 0.006445803, -0.03496501, + 0.028332196, 0.04449512, -0.042436164, -0.026587414, + -0.006041347, -0.09292539, -0.05678812, 0.03897832, + 0.09465633, 0.008115513, -0.02171956, 0.08304309, + 0.071401566, 0.019622514, 0.032163795, -0.004167056, + 0.02295182, 0.030739572, 0.056506045, 0.004612461, + 0.06524936, 0.059999723, 0.046395954, -0.0045512207, + -0.1335546, -0.030136576, 0.11584653, -0.014678886, + 0.0020118146, -0.09688814, -0.0790206, 0.039770417, + -0.0329582, 0.07922767, 0.029322514, 0.026405897, + 0.04207835, -0.07073373, 0.063781224, 0.0859677, + -0.10925287, -0.07011058, 0.048005477, 0.03438226, + -0.09606514, -0.006669445, -0.043381985, 0.04240257, + -0.06955775, -0.06769346, 0.043903265, -0.026784198, + -0.017840602, 0.024307009, -0.040079936, -0.019946516, + 0.045318738, -0.12233574, 0.026170589, 0.0074471775, + 0.15978073, 0.10185836, 0.10298046, -0.015476589, + -0.039390966, -0.072174534, 0.0739445, -0.1211869, + -0.0347889, -0.07943156, 0.014809798, -0.12412325, + -0.0030663363, 0.039695457, 0.0647603, -0.08291318, + -0.018529687, -0.004423833, 0.0037507233, 0.084633216, + -0.01514876, -0.056505352, -0.012800942, -0.06994386, + 0.012962922, -0.031234352, 0.07029052, 0.016418684, + 0.03618972, 0.055686004, -0.08663945, -0.017404709, + -0.054761406, 0.029065743, 0.052404847, 0.020238016, + 0.0048197987, -0.0214882, 0.07078733, 0.013016777, + 0.06262858, 0.009184685, 0.020785125, -0.043904778, + -0.0270329, -0.03299152, -0.060088247, -0.015162964, + -0.001828936, 0.12642565, -0.056757294, 0.013586685, + 0.09232601, -0.035886683, 0.06000002, 0.05229691, + -0.052580316, -0.082029596, -0.010794592, 0.012947712, + -0.036429964, -0.085508935, -0.13127148, -0.017744139, + 0.031502828, 0.036232427, -0.031581745, 0.023051167, + -0.05325106, -0.03421577, 0.028793324, -0.034633752, + -0.009881397, -0.043551125, -0.018609839, 0.0019097115, + -0.008799762, 0.056595087, 0.0022273948, 0.055752404}); + + lstm.SetRecurrentToOutputWeights({ + 0.025825322, -0.05813119, 0.09495884, -0.045984812, -0.01255415, + -0.0026479573, -0.08196161, -0.054914974, -0.0046604523, -0.029587349, + -0.044576716, -0.07480124, -0.082868785, 0.023254942, 0.027502948, + -0.0039728214, -0.08683098, -0.08116779, -0.014675607, -0.037924774, + -0.023314456, -0.007401714, -0.09255757, 0.029460307, -0.08829125, + -0.005139627, -0.08989442, -0.0555066, 0.13596267, -0.025062224, + -0.048351806, -0.03850004, 0.07266485, -0.022414139, 0.05940088, + 0.075114764, 0.09597592, -0.010211725, -0.0049794707, -0.011523867, + -0.025980417, 0.072999895, 0.11091378, -0.081685916, 0.014416728, + 0.043229222, 0.034178585, -0.07530371, 0.035837382, -0.085607, + -0.007721233, -0.03287832, -0.043848954, -0.06404588, -0.06632928, + -0.073643476, 0.008214239, -0.045984086, 0.039764922, 0.03474462, + 0.060612556, -0.080590084, 0.049127717, 0.04151091, -0.030063879, + 0.008801774, -0.023021035, -0.019558564, 0.05158114, -0.010947698, + -0.011825728, 0.0075720972, 0.0699727, -0.0039981045, 0.069350146, + 0.08799282, 0.016156472, 0.035502106, 0.11695009, 0.006217345, + 0.13392477, -0.037875112, 0.025745004, 0.08940699, -0.00924166, + 0.0046702605, -0.036598757, -0.08811812, 0.10522024, -0.032441203, + 0.008176899, -0.04454919, 0.07058152, 0.0067963637, 0.039206743, + 0.03259838, 0.03725492, -0.09515802, 0.013326398, -0.052055415, + -0.025676316, 0.03198509, -0.015951829, -0.058556724, 0.036879618, + 0.043357447, 0.028362012, -0.05908629, 0.0059240665, -0.04995891, + -0.019187413, 0.0276265, -0.01628143, 0.0025863599, 0.08800015, + 0.035250366, -0.022165963, -0.07328642, -0.009415526, -0.07455109, + 0.11690406, 0.0363299, 0.07411125, 0.042103454, -0.009660886, + 0.019076364, 0.018299393, -0.046004917, 0.08891175, 0.0431396, + -0.026327137, -0.051502608, 0.08979574, -0.051670972, 0.04940282, + -0.07491107, -0.021240504, 0.022596184, -0.034280192, 0.060163025, + -0.058211457, -0.051837247, -0.01349775, -0.04639988, -0.035936575, + -0.011681591, 0.064818054, 0.0073146066, -0.021745546, -0.043124277, + -0.06471268, -0.07053354, -0.029321948, -0.05330136, 0.016933719, + -0.053782392, 0.13747959, -0.1361751, -0.11569455, 0.0033329215, + 0.05693899, -0.053219706, 0.063698, 0.07977434, -0.07924483, + 0.06936997, 0.0034815092, -0.007305279, -0.037325785, -0.07251102, + -0.033633437, -0.08677009, 0.091591336, -0.14165086, 0.021752775, + 0.019683983, 0.0011612234, -0.058154266, 0.049996935, 0.0288841, + -0.0024567875, -0.14345716, 0.010955264, -0.10234828, 0.1183656, + -0.0010731248, -0.023590032, -0.072285876, -0.0724771, -0.026382286, + -0.0014920527, 0.042667855, 0.0018776858, 0.02986552, 0.009814309, + 0.0733756, 0.12289186, 0.018043943, -0.0458958, 0.049412545, + 0.033632483, 0.05495232, 0.036686596, -0.013781798, -0.010036754, + 0.02576849, -0.08307328, 0.010112348, 0.042521734, -0.05869831, + -0.071689695, 0.03876447, -0.13275425, -0.0352966, -0.023077697, + 0.10285965, 0.084736146, 0.15568255, -0.00040734606, 0.027835453, + -0.10292561, -0.032401145, 0.10053256, -0.026142767, -0.08271222, + -0.0030240538, -0.016368777, 0.1070414, 0.042672627, 0.013456989, + -0.0437609, -0.022309763, 0.11576483, 0.04108048, 0.061026827, + -0.0190714, -0.0869359, 0.037901703, 0.0610107, 0.07202949, + 0.01675338, 0.086139716, -0.08795751, -0.014898893, -0.023771819, + -0.01965048, 0.007955471, -0.043740474, 0.03346837, -0.10549954, + 0.090567775, 0.042013682, -0.03176985, 0.12569028, -0.02421228, + -0.029526481, 0.023851605, 0.031539805, 0.05292009, -0.02344001, + -0.07811758, -0.08834428, 0.10094801, 0.16594367, -0.06861939, + -0.021256343, -0.041093912, -0.06669611, 0.035498552, 0.021757556, + -0.09302526, -0.015403468, -0.06614931, -0.051798206, -0.013874718, + 0.03630673, 0.010412845, -0.08077351, 0.046185967, 0.0035662893, + 0.03541868, -0.094149634, -0.034814864, 0.003128424, -0.020674974, + -0.03944324, -0.008110165, -0.11113267, 0.08484226, 0.043586485, + 0.040582247, 0.0968012, -0.065249965, -0.028036479, 0.0050708856, + 0.0017462453, 0.0326779, 0.041296225, 0.09164146, -0.047743853, + -0.015952192, -0.034451712, 0.084197424, -0.05347844, -0.11768019, + 0.085926116, -0.08251791, -0.045081906, 0.0948852, 0.068401024, + 0.024856757, 0.06978981, -0.057309967, -0.012775832, -0.0032452994, + 0.01977615, -0.041040014, -0.024264973, 0.063464895, 0.05431621, + }); + + lstm.SetCellToInputWeights( + {0.040369894, 0.030746894, 0.24704495, 0.018586371, -0.037586458, + -0.15312155, -0.11812848, -0.11465643, 0.20259799, 0.11418174, + -0.10116027, -0.011334949, 0.12411352, -0.076769054, -0.052169047, + 0.21198851, -0.38871562, -0.09061183, -0.09683246, -0.21929175}); + + lstm.SetCellToForgetWeights( + {-0.01998659, -0.15568835, -0.24248174, -0.012770197, 0.041331276, + -0.072311886, -0.052123554, -0.0066330447, -0.043891653, 0.036225766, + -0.047248036, 0.021479502, 0.033189066, 0.11952997, -0.020432774, + 0.64658105, -0.06650122, -0.03467612, 0.095340036, 0.23647355}); + + lstm.SetCellToOutputWeights( + {0.08286371, -0.08261836, -0.51210177, 0.002913762, 0.17764764, + -0.5495371, -0.08460716, -0.24552552, 0.030037103, 0.04123544, + -0.11940523, 0.007358328, 0.1890978, 0.4833202, -0.34441817, + 0.36312827, -0.26375428, 0.1457655, -0.19724406, 0.15548733}); + + lstm.SetProjectionWeights( + {-0.009802181, 0.09401916, 0.0717386, -0.13895074, 0.09641832, + 0.060420845, 0.08539281, 0.054285463, 0.061395317, 0.034448683, + -0.042991187, 0.019801661, -0.16840284, -0.015726732, -0.23041931, + -0.024478018, -0.10959692, -0.013875541, 0.18600968, -0.061274476, + 0.0138165, -0.08160894, -0.07661644, 0.032372914, 0.16169067, + 0.22465782, -0.03993472, -0.004017731, 0.08633481, -0.28869787, + 0.08682067, 0.17240396, 0.014975425, 0.056431185, 0.031037588, + 0.16702051, 0.0077946745, 0.15140012, 0.29405436, 0.120285, + -0.188994, -0.027265169, 0.043389652, -0.022061434, 0.014777949, + -0.20203483, 0.094781205, 0.19100232, 0.13987629, -0.036132768, + -0.06426278, -0.05108664, 0.13221376, 0.009441198, -0.16715929, + 0.15859416, -0.040437475, 0.050779544, -0.022187516, 0.012166504, + 0.027685808, -0.07675938, -0.0055694645, -0.09444123, 0.0046453946, + 0.050794356, 0.10770313, -0.20790008, -0.07149004, -0.11425117, + 0.008225835, -0.035802525, 0.14374903, 0.15262283, 0.048710253, + 0.1847461, -0.007487823, 0.11000021, -0.09542012, 0.22619456, + -0.029149994, 0.08527916, 0.009043713, 0.0042746216, 0.016261552, + 0.022461696, 0.12689082, -0.043589946, -0.12035478, -0.08361797, + -0.050666027, -0.1248618, -0.1275799, -0.071875185, 0.07377272, + 0.09944291, -0.18897448, -0.1593054, -0.06526116, -0.040107165, + -0.004618631, -0.067624845, -0.007576253, 0.10727444, 0.041546922, + -0.20424393, 0.06907816, 0.050412357, 0.00724631, 0.039827548, + 0.12449835, 0.10747581, 0.13708383, 0.09134148, -0.12617786, + -0.06428341, 0.09956831, 0.1208086, -0.14676677, -0.0727722, + 0.1126304, 0.010139365, 0.015571211, -0.038128063, 0.022913318, + -0.042050496, 0.16842307, -0.060597885, 0.10531834, -0.06411776, + -0.07451711, -0.03410368, -0.13393489, 0.06534304, 0.003620307, + 0.04490757, 0.05970546, 0.05197996, 0.02839995, 0.10434969, + -0.013699693, -0.028353551, -0.07260381, 0.047201227, -0.024575593, + -0.036445823, 0.07155557, 0.009672501, -0.02328883, 0.009533515, + -0.03606021, -0.07421458, -0.028082801, -0.2678904, -0.13221288, + 0.18419984, -0.13012612, -0.014588381, -0.035059117, -0.04824723, + 0.07830115, -0.056184657, 0.03277091, 0.025466874, 0.14494097, + -0.12522776, -0.098633975, -0.10766018, -0.08317623, 0.08594209, + 0.07749552, 0.039474737, 0.1776665, -0.07409566, -0.0477268, + 0.29323658, 0.10801441, 0.1154011, 0.013952499, 0.10739139, + 0.10708251, -0.051456142, 0.0074137426, -0.10430189, 0.10034707, + 0.045594677, 0.0635285, -0.0715442, -0.089667566, -0.10811871, + 0.00026344223, 0.08298446, -0.009525053, 0.006585689, -0.24567553, + -0.09450807, 0.09648481, 0.026996298, -0.06419476, -0.04752702, + -0.11063944, -0.23441927, -0.17608605, -0.052156363, 0.067035615, + 0.19271925, -0.0032889997, -0.043264326, 0.09663576, -0.057112187, + -0.10100678, 0.0628376, 0.04447668, 0.017961001, -0.10094388, + -0.10190601, 0.18335468, 0.10494553, -0.052095775, -0.0026118709, + 0.10539724, -0.04383912, -0.042349473, 0.08438151, -0.1947263, + 0.02251204, 0.11216432, -0.10307853, 0.17351969, -0.039091777, + 0.08066188, -0.00561982, 0.12633002, 0.11335965, -0.0088127935, + -0.019777594, 0.06864014, -0.059751723, 0.016233567, -0.06894641, + -0.28651384, -0.004228674, 0.019708522, -0.16305895, -0.07468996, + -0.0855457, 0.099339016, -0.07580735, -0.13775392, 0.08434318, + 0.08330512, -0.12131499, 0.031935584, 0.09180414, -0.08876437, + -0.08049874, 0.008753825, 0.03498998, 0.030215185, 0.03907079, + 0.089751154, 0.029194152, -0.03337423, -0.019092513, 0.04331237, + 0.04299654, -0.036394123, -0.12915532, 0.09793732, 0.07512415, + -0.11319543, -0.032502122, 0.15661901, 0.07671967, -0.005491124, + -0.19379048, -0.218606, 0.21448623, 0.017840758, 0.1416943, + -0.07051762, 0.19488361, 0.02664691, -0.18104725, -0.09334311, + 0.15026465, -0.15493552, -0.057762887, -0.11604192, -0.262013, + -0.01391798, 0.012185008, 0.11156489, -0.07483202, 0.06693364, + -0.26151478, 0.046425626, 0.036540434, -0.16435726, 0.17338543, + -0.21401681, -0.11385144, -0.08283257, -0.069031075, 0.030635102, + 0.010969227, 0.11109743, 0.010919218, 0.027526086, 0.13519906, + 0.01891392, -0.046839405, -0.040167913, 0.017953383, -0.09700955, + 0.0061885654, -0.07000971, 0.026893595, -0.038844477, 0.14543656}); + + static float lstm_input[][20] = { + {// Batch0: 4 (input_sequence_size) * 5 (n_input) + 0.787926, 0.151646, 0.071352, 0.118426, 0.458058, 0.596268, 0.998386, + 0.568695, 0.864524, 0.571277, 0.073204, 0.296072, 0.743333, 0.069199, + 0.045348, 0.867394, 0.291279, 0.013714, 0.482521, 0.626339}, + + {// Batch1: 4 (input_sequence_size) * 5 (n_input) + 0.295743, 0.544053, 0.690064, 0.858138, 0.497181, 0.642421, 0.524260, + 0.134799, 0.003639, 0.162482, 0.640394, 0.930399, 0.050782, 0.432485, + 0.988078, 0.082922, 0.563329, 0.865614, 0.333232, 0.259916}}; + + static float lstm_golden_output[][64] = { + {// Batch0: 4 (input_sequence_size) * 16 (n_output) + -0.00396806, 0.029352, -0.00279226, 0.0159977, -0.00835576, + -0.0211779, 0.0283512, -0.0114597, 0.00907307, -0.0244004, + -0.0152191, -0.0259063, 0.00914318, 0.00415118, 0.017147, + 0.0134203, -0.0166936, 0.0381209, 0.000889694, 0.0143363, + -0.0328911, -0.0234288, 0.0333051, -0.012229, 0.0110322, + -0.0457725, -0.000832209, -0.0202817, 0.0327257, 0.0121308, + 0.0155969, 0.0312091, -0.0213783, 0.0350169, 0.000324794, + 0.0276012, -0.0263374, -0.0371449, 0.0446149, -0.0205474, + 0.0103729, -0.0576349, -0.0150052, -0.0292043, 0.0376827, + 0.0136115, 0.0243435, 0.0354492, -0.0189322, 0.0464512, + -0.00251373, 0.0225745, -0.0308346, -0.0317124, 0.0460407, + -0.0189395, 0.0149363, -0.0530162, -0.0150767, -0.0340193, + 0.0286833, 0.00824207, 0.0264887, 0.0305169}, + {// Batch1: 4 (input_sequence_size) * 16 (n_output) + -0.013869, 0.0287268, -0.00334693, 0.00733398, -0.0287926, + -0.0186926, 0.0193662, -0.0115437, 0.00422612, -0.0345232, + 0.00223253, -0.00957321, 0.0210624, 0.013331, 0.0150954, + 0.02168, -0.0141913, 0.0322082, 0.00227024, 0.0260507, + -0.0188721, -0.0296489, 0.0399134, -0.0160509, 0.0116039, + -0.0447318, -0.0150515, -0.0277406, 0.0316596, 0.0118233, + 0.0214762, 0.0293641, -0.0204549, 0.0450315, -0.00117378, + 0.0167673, -0.0375007, -0.0238314, 0.038784, -0.0174034, + 0.0131743, -0.0506589, -0.0048447, -0.0240239, 0.0325789, + 0.00790065, 0.0220157, 0.0333314, -0.0264787, 0.0387855, + -0.000764675, 0.0217599, -0.037537, -0.0335206, 0.0431679, + -0.0211424, 0.010203, -0.062785, -0.00832363, -0.025181, + 0.0412031, 0.0118723, 0.0239643, 0.0394009}}; + + // Resetting cell_state and output_state + lstm.ResetCellState(); + lstm.ResetOutputState(); + + for (int i = 0; i < lstm.sequence_length(); i++) { + float* batch0_start = lstm_input[0] + i * lstm.num_inputs(); + float* batch0_end = batch0_start + lstm.num_inputs(); + + lstm.SetInput(2 * i * lstm.num_inputs(), batch0_start, batch0_end); + + float* batch1_start = lstm_input[1] + i * lstm.num_inputs(); + float* batch1_end = batch1_start + lstm.num_inputs(); + lstm.SetInput((2 * i + 1) * lstm.num_inputs(), batch1_start, batch1_end); + } + + lstm.Invoke(); + + std::vector expected; + for (int i = 0; i < lstm.sequence_length(); i++) { + float* golden_start_batch0 = lstm_golden_output[0] + i * lstm.num_outputs(); + float* golden_end_batch0 = golden_start_batch0 + lstm.num_outputs(); + float* golden_start_batch1 = lstm_golden_output[1] + i * lstm.num_outputs(); + float* golden_end_batch1 = golden_start_batch1 + lstm.num_outputs(); + expected.insert(expected.end(), golden_start_batch0, golden_end_batch0); + expected.insert(expected.end(), golden_start_batch1, golden_end_batch1); + } + EXPECT_THAT(lstm.GetOutput(), ElementsAreArray(ArrayFloatNear(expected))); +} + +} // namespace +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn.cc index 85e09049eea5f66a2bb854990bf80e9ed5dcc88a..ac00c37b67dcbe77023a2495a698967ca555b1d5 100644 --- a/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn.cc +++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn.cc @@ -15,14 +15,15 @@ limitations under the License. #include #include #include -#include #include +#include #include #include #include "tensorflow/contrib/lite/builtin_op_data.h" #include "tensorflow/contrib/lite/context.h" #include "tensorflow/contrib/lite/kernels/activation_functor.h" +#include "tensorflow/contrib/lite/kernels/internal/kernel_utils.h" #include "tensorflow/contrib/lite/kernels/op_macros.h" namespace tflite { @@ -34,7 +35,7 @@ constexpr int kInputTensor = 0; constexpr int kWeightsTensor = 1; constexpr int kRecurrentWeightsTensor = 2; constexpr int kBiasTensor = 3; -constexpr int KHiddenStateTensor = 0; +constexpr int kHiddenStateTensor = 0; constexpr int kOutputTensor = 1; TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { @@ -51,8 +52,12 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // Check all the parameters of tensor match within themselves and match the // input configuration. - const int batch_size = input->dims->data[0]; - const int max_time = input->dims->data[1]; + auto* params = reinterpret_cast(node->builtin_data); + const bool time_major = params->time_major; + const int batch_size = + (time_major) ? input->dims->data[1] : input->dims->data[0]; + const int max_time = + (time_major) ? input->dims->data[0] : input->dims->data[1]; const int num_units = input_weights->dims->data[0]; TF_LITE_ASSERT_EQ(input->dims->data[2], input_weights->dims->data[1]); TF_LITE_ASSERT_EQ(input_weights->dims->data[0], bias->dims->data[0]); @@ -60,7 +65,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ASSERT_EQ(recurrent_weights->dims->data[1], bias->dims->data[0]); TfLiteTensor* hidden_state = - &context->tensors[node->outputs->data[KHiddenStateTensor]]; + &context->tensors[node->outputs->data[kHiddenStateTensor]]; TfLiteTensor* output = &context->tensors[node->outputs->data[kOutputTensor]]; // Resize state. @@ -75,17 +80,17 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // Resize output. TfLiteIntArray* output_size_array = TfLiteIntArrayCreate(3); - output_size_array->data[0] = batch_size; - output_size_array->data[1] = max_time; + output_size_array->data[0] = (time_major) ? max_time : batch_size; + output_size_array->data[1] = (time_major) ? batch_size : max_time; output_size_array->data[2] = num_units; - TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, output, - output_size_array)); + TF_LITE_ENSURE_OK(context, + context->ResizeTensor(context, output, output_size_array)); return kTfLiteOk; } TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - auto* params = reinterpret_cast(node->builtin_data); + auto* params = reinterpret_cast(node->builtin_data); TfLiteTensor* input = &context->tensors[node->inputs->data[kInputTensor]]; TfLiteTensor* input_weights = @@ -94,61 +99,55 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { &context->tensors[node->inputs->data[kRecurrentWeightsTensor]]; TfLiteTensor* bias = &context->tensors[node->inputs->data[kBiasTensor]]; TfLiteTensor* hidden_state = - &context->tensors[node->outputs->data[KHiddenStateTensor]]; + &context->tensors[node->outputs->data[kHiddenStateTensor]]; TfLiteTensor* output = &context->tensors[node->outputs->data[kOutputTensor]]; // Initialize the pointer bias. const float* bias_ptr = bias->data.f; - const int batch_size = input->dims->data[0]; - const int max_time = input->dims->data[1]; + const bool time_major = params->time_major; + const int batch_size = + (time_major) ? input->dims->data[1] : input->dims->data[0]; + const int max_time = + (time_major) ? input->dims->data[0] : input->dims->data[1]; const int num_units = input_weights->dims->data[0]; const int input_size = input->dims->data[2]; - const int input_weights_stride = input_weights->dims->data[1]; - const int recurrent_weights_stride = recurrent_weights->dims->data[1]; - // For each batch - for (int b = 0; b < batch_size; b++) { + // Initialize input_weights and recurrent_weights. + const float* input_weights_ptr = input_weights->data.f; + const float* recurrent_weights_ptr = recurrent_weights->data.f; + + if (time_major) { // Initialize the pointer to hidden state. - float* hidden_state_ptr_batch = hidden_state->data.f + b * num_units; + float* hidden_state_ptr_batch = hidden_state->data.f; + // Unroll the sequence and use batch batch operations for efficiency. for (int s = 0; s < max_time; s++) { // Initialize the pointer to input and output. const float* input_ptr_batch = - input->data.f + b * input_size * max_time + s * input_size; - float* output_ptr_batch = - output->data.f + b * num_units * max_time + s * num_units; - - // Initialize input_weights and recurrent_weights. - const float* input_weights_ptr = input_weights->data.f; - const float* recurrent_weights_ptr = recurrent_weights->data.f; - - // Output = bias - for (int o = 0; o < num_units; o++) { - output_ptr_batch[o] = bias_ptr[o]; - } + input->data.f + s * input_size * batch_size; + float* output_ptr_batch = output->data.f + s * num_units * batch_size; - // Output += input * input_weights - for (int o = 0; o < num_units; o++) { - for (int i = 0; i < input_size; i++) { - output_ptr_batch[o] += input_ptr_batch[i] * input_weights_ptr[i]; - } - input_weights_ptr += input_weights_stride; - } - - // Output += recurrent_weights * hidden_state - for (int o = 0; o < num_units; o++) { - for (int h = 0; h < num_units; h++) { - output_ptr_batch[o] += - hidden_state_ptr_batch[h] * recurrent_weights_ptr[h]; - } - recurrent_weights_ptr += recurrent_weights_stride; - } - - // Output = activation(Output) and update hidden_state - for (int o = 0; o < num_units; o++) { - output_ptr_batch[o] = - (ActivationFunctor(params->activation))(output_ptr_batch[o]); - hidden_state_ptr_batch[o] = output_ptr_batch[o]; + kernel_utils::RnnBatchStep(input_ptr_batch, input_weights_ptr, + recurrent_weights_ptr, bias_ptr, input_size, + num_units, batch_size, params->activation, + hidden_state_ptr_batch, output_ptr_batch); + } + } else { + // For each batch + for (int b = 0; b < batch_size; b++) { + // Initialize the pointer to hidden state. + float* hidden_state_ptr_batch = hidden_state->data.f + b * num_units; + for (int s = 0; s < max_time; s++) { + // Initialize the pointer to input and output. + const float* input_ptr_batch = + input->data.f + b * input_size * max_time + s * input_size; + float* output_ptr_batch = + output->data.f + b * num_units * max_time + s * num_units; + + kernel_utils::RnnBatchStep( + input_ptr_batch, input_weights_ptr, recurrent_weights_ptr, bias_ptr, + input_size, num_units, /*batch_size=*/1, params->activation, + hidden_state_ptr_batch, output_ptr_batch); } } } diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn_test.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn_test.cc index a1c1eda16034f83ca5c79fc18f4fa495a3e73f90..7e32969763b59620dc3534708f965750680002d2 100644 --- a/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn_test.cc +++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_rnn_test.cc @@ -12,10 +12,10 @@ 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. ==============================================================================*/ -// Unit test for TFLite RNN op. +// Unit test for TFLite Sequential RNN op. -#include #include +#include #include #include @@ -120,12 +120,12 @@ static float rnn_golden_output[] = { 0.415153, 0.210318, 0, 0, 0, 0, 0, 2.02616, 0, 0.728256, 0.84183, 0.0907453, - 0.628881, 3.58099, 1.49974, 0 -}; + 0.628881, 3.58099, 1.49974, 0}; class UnidirectionalRNNOpModel : public SingleOpModel { public: - UnidirectionalRNNOpModel(int batches, int sequence_len, int units, int size) + UnidirectionalRNNOpModel(int batches, int sequence_len, int units, int size, + bool time_major) : batches_(batches), sequence_len_(sequence_len), units_(units), @@ -136,13 +136,22 @@ class UnidirectionalRNNOpModel : public SingleOpModel { bias_ = AddInput(TensorType_FLOAT32); hidden_state_ = AddOutput(TensorType_FLOAT32); output_ = AddOutput(TensorType_FLOAT32); - SetBuiltinOp( - BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, BuiltinOptions_RNNOptions, - CreateRNNOptions(builder_, ActivationFunctionType_RELU).Union()); - BuildInterpreter({{batches_, sequence_len_, input_size_}, - {units_, input_size_}, - {units_, units_}, - {units_}}); + SetBuiltinOp(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, + BuiltinOptions_SequenceRNNOptions, + CreateSequenceRNNOptions(builder_, time_major, + ActivationFunctionType_RELU) + .Union()); + if (time_major) { + BuildInterpreter({{sequence_len_, batches_, input_size_}, + {units_, input_size_}, + {units_, units_}, + {units_}}); + } else { + BuildInterpreter({{batches_, sequence_len_, input_size_}, + {units_, input_size_}, + {units_, units_}, + {units_}}); + } } void SetBias(std::initializer_list f) { PopulateTensor(bias_, f); } @@ -195,7 +204,8 @@ class UnidirectionalRNNOpModel : public SingleOpModel { // TODO(mirkov): add another test which directly compares to TF once TOCO // supports the conversion from dynamic_rnn with BasicRNNCell. TEST(FullyConnectedOpTest, BlackBoxTest) { - UnidirectionalRNNOpModel rnn(2, 16, 16, 8); + UnidirectionalRNNOpModel rnn(/*batches=*/2, /*sequence_len=*/16, + /*units=*/16, /*size=*/8, /*time_major=*/false); rnn.SetWeights( {0.461459, 0.153381, 0.529743, -0.00371218, 0.676267, -0.211346, 0.317493, 0.969689, -0.343251, 0.186423, 0.398151, 0.152399, @@ -260,6 +270,77 @@ TEST(FullyConnectedOpTest, BlackBoxTest) { EXPECT_THAT(rnn.GetOutput(), ElementsAreArray(ArrayFloatNear(expected))); } +TEST(FullyConnectedOpTest, TimeMajorBlackBoxTest) { + UnidirectionalRNNOpModel rnn(/*batches=*/2, /*sequence_len=*/16, + /*units=*/16, /*size=*/8, /*time_major=*/true); + rnn.SetWeights( + {0.461459, 0.153381, 0.529743, -0.00371218, 0.676267, -0.211346, + 0.317493, 0.969689, -0.343251, 0.186423, 0.398151, 0.152399, + 0.448504, 0.317662, 0.523556, -0.323514, 0.480877, 0.333113, + -0.757714, -0.674487, -0.643585, 0.217766, -0.0251462, 0.79512, + -0.595574, -0.422444, 0.371572, -0.452178, -0.556069, -0.482188, + -0.685456, -0.727851, 0.841829, 0.551535, -0.232336, 0.729158, + -0.00294906, -0.69754, 0.766073, -0.178424, 0.369513, -0.423241, + 0.548547, -0.0152023, -0.757482, -0.85491, 0.251331, -0.989183, + 0.306261, -0.340716, 0.886103, -0.0726757, -0.723523, -0.784303, + 0.0354295, 0.566564, -0.485469, -0.620498, 0.832546, 0.697884, + -0.279115, 0.294415, -0.584313, 0.548772, 0.0648819, 0.968726, + 0.723834, -0.0080452, -0.350386, -0.272803, 0.115121, -0.412644, + -0.824713, -0.992843, -0.592904, -0.417893, 0.863791, -0.423461, + -0.147601, -0.770664, -0.479006, 0.654782, 0.587314, -0.639158, + 0.816969, -0.337228, 0.659878, 0.73107, 0.754768, -0.337042, + 0.0960841, 0.368357, 0.244191, -0.817703, -0.211223, 0.442012, + 0.37225, -0.623598, -0.405423, 0.455101, 0.673656, -0.145345, + -0.511346, -0.901675, -0.81252, -0.127006, 0.809865, -0.721884, + 0.636255, 0.868989, -0.347973, -0.10179, -0.777449, 0.917274, + 0.819286, 0.206218, -0.00785118, 0.167141, 0.45872, 0.972934, + -0.276798, 0.837861, 0.747958, -0.0151566, -0.330057, -0.469077, + 0.277308, 0.415818}); + + rnn.SetBias({0.065691948, -0.69055247, 0.1107955, -0.97084129, -0.23957068, + -0.23566568, -0.389184, 0.47481549, -0.4791103, 0.29931796, + 0.10463274, 0.83918178, 0.37197268, 0.61957061, 0.3956964, + -0.37609905}); + + rnn.SetRecurrentWeights({0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0.1}); + + rnn.ResetHiddenState(); + for (int i = 0; i < rnn.sequence_len(); i++) { + float* batch_start = rnn_input + i * rnn.input_size(); + float* batch_end = batch_start + rnn.input_size(); + // The two batches are identical. + rnn.SetInput(2 * i * rnn.input_size(), batch_start, batch_end); + rnn.SetInput((2 * i + 1) * rnn.input_size(), batch_start, batch_end); + } + + rnn.Invoke(); + + std::vector expected; + for (int i = 0; i < rnn.sequence_len(); i++) { + float* golden_batch_start = rnn_golden_output + i * rnn.num_units(); + float* golden_batch_end = golden_batch_start + rnn.num_units(); + expected.insert(expected.end(), golden_batch_start, golden_batch_end); + expected.insert(expected.end(), golden_batch_start, golden_batch_end); + } + + EXPECT_THAT(rnn.GetOutput(), ElementsAreArray(ArrayFloatNear(expected))); +} + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/lib_package/BUILD b/tensorflow/contrib/lite/lib_package/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..3c1b8d3d45f2bb382bbe6b789ec6ac7ec89ebc66 --- /dev/null +++ b/tensorflow/contrib/lite/lib_package/BUILD @@ -0,0 +1,16 @@ +package(default_visibility = ["//visibility:private"]) + +# Create the LICENSE file for libraries that are used by TensorFlow Lite +# C library. +genrule( + name = "clicenses_generate", + srcs = [ + "//third_party/eigen3:LICENSE", + "@arm_neon_2_x86_sse//:LICENSE", + "@farmhash_archive//:COPYING", + "@gemmlowp//:LICENSE", + ], + outs = ["LICENSE"], + cmd = "$(location :concat_licenses.sh) $(SRCS) >$@", + tools = [":concat_licenses.sh"], +) diff --git a/tensorflow/contrib/lite/lib_package/concat_licenses.sh b/tensorflow/contrib/lite/lib_package/concat_licenses.sh new file mode 100755 index 0000000000000000000000000000000000000000..2070f64e9fa4384234361556da0ed6f5089319b3 --- /dev/null +++ b/tensorflow/contrib/lite/lib_package/concat_licenses.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# Copyright 2016 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. +# ============================================================================== +# +# Script aimed to combining multiple license files into a single one. + +for f in $@ +do + echo "--------------------------------------------------------------------------------" + echo "BEGIN LICENSE FOR $f" + echo "--------------------------------------------------------------------------------" + cat $f + echo "--------------------------------------------------------------------------------" + echo "END LICENSE FOR $f" + echo "--------------------------------------------------------------------------------" +done diff --git a/tensorflow/contrib/lite/lib_package/create_ios_frameworks.sh b/tensorflow/contrib/lite/lib_package/create_ios_frameworks.sh new file mode 100755 index 0000000000000000000000000000000000000000..b58ae266017caf8781c28331f49a8f5bc1550767 --- /dev/null +++ b/tensorflow/contrib/lite/lib_package/create_ios_frameworks.sh @@ -0,0 +1,81 @@ +#!/bin/bash -x +# Copyright 2017 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. +# ============================================================================== + +set -e + +echo "Starting" +TFLITE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/.." + +TMP_DIR=$(mktemp -d) +echo "Package dir: " $TMP_DIR +FW_DIR=$TMP_DIR/tensorflow_lite_ios_frameworks +FW_DIR_TFLITE=$FW_DIR/tensorflow_lite.framework +FW_DIR_TFLITE_HDRS=$FW_DIR_TFLITE/Headers + +echo "Creating target Headers directories" +mkdir -p $FW_DIR_TFLITE_HDRS + +echo "Headers, populating: TensorFlow Lite" +cd $TFLITE_DIR/../../.. + +find tensorflow/contrib/lite -name '*.h' \ + -not -path 'tensorflow/contrib/lite/downloads/*' \ + -not -path 'tensorflow/contrib/lite/examples/*' \ + -not -path 'tensorflow/contrib/lite/gen/*' \ + -not -path 'tensorflow/contrib/lite/toco/*' \ + -not -path 'tensorflow/contrib/lite/nnapi/*' \ + -not -path 'tensorflow/contrib/lite/java/*' \ + | tar -cf $FW_DIR_TFLITE_HDRS/tmp.tar -T - +cd $FW_DIR_TFLITE_HDRS +tar xf tmp.tar +rm -f tmp.tar + +echo "Headers, populating: Flatbuffer" +cd $TFLITE_DIR/downloads/flatbuffers/include/ +find . -name '*.h' | tar -cf $FW_DIR_TFLITE_HDRS/tmp.tar -T - +cd $FW_DIR_TFLITE_HDRS +tar xf tmp.tar +rm -f tmp.tar + +cd $TFLITE_DIR/../../.. +echo "Generate master LICENSE file and copy to target" +bazel build //tensorflow/tools/lib_package:clicenses_generate +cp $TFLITE_DIR/../../../bazel-genfiles/tensorflow/tools/lib_package/include/tensorflow/c/LICENSE \ + $FW_DIR_TFLITE + +echo "Copying static libraries" +cp $TFLITE_DIR/gen/lib/libtensorflow-lite.a \ + $FW_DIR_TFLITE/tensorflow_lite + +# This is required, otherwise they interfere with the documentation of the +# pod at cocoapods.org. +echo "Remove all README files" +cd $FW_DIR_TFLITE_HDRS +find . -type f -name README\* -exec rm -f {} \; +find . -type f -name readme\* -exec rm -f {} \; + +TARGET_GEN_LOCATION="$TFLITE_DIR/gen/ios_frameworks" +echo "Moving results to target: " $TARGET_GEN_LOCATION +cd $FW_DIR +zip -q -r tensorflow_lite.framework.zip tensorflow_lite.framework -x .DS_Store +rm -rf $TARGET_GEN_LOCATION +mkdir -p $TARGET_GEN_LOCATION +cp -r tensorflow_lite.framework.zip $TARGET_GEN_LOCATION + +echo "Cleaning up" +rm -rf $TMP_DIR + +echo "Finished" diff --git a/tensorflow/contrib/lite/memory_planner.h b/tensorflow/contrib/lite/memory_planner.h new file mode 100644 index 0000000000000000000000000000000000000000..5cd6c208500f3ea84ab8146f7f136e8b7851ff03 --- /dev/null +++ b/tensorflow/contrib/lite/memory_planner.h @@ -0,0 +1,45 @@ +/* Copyright 2017 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_LITE_MEMORY_PLANNER_H_ +#define TENSORFLOW_CONTRIB_LITE_MEMORY_PLANNER_H_ + +#include "tensorflow/contrib/lite/context.h" + +namespace tflite { + +// A MemoryPlanner is responsible for planning and executing a number of +// memory-related operations that are necessary in TF Lite. +class MemoryPlanner { + public: + virtual ~MemoryPlanner() {} + + // Plans the necessary memory allocations. This is the MemoryPlanner's + // pre-processing step and is called when the graph structure is known but + // actual size of the tensors is not. + virtual TfLiteStatus PlanAllocations() = 0; + + // Allocates the necessary memory to execute all nodes in the interval + // [first_node, last_node]. + virtual TfLiteStatus ExecuteAllocations(int first_node, int last_node) = 0; + + // Invalidates allocations made earliers. This is called when tensors sizes + // have change. All planned allocations remain, but can't be used until + // ExecuteAllocations() is called. + virtual TfLiteStatus ResetAllocations() = 0; +}; + +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_MEMORY_PLANNER_H_ diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc index 94e22b265964b300c862a9ee52511d479c20c64d..d6522fc077d03bb49fe54b7a04fa6341ccf4cf3a 100644 --- a/tensorflow/contrib/lite/model.cc +++ b/tensorflow/contrib/lite/model.cc @@ -30,17 +30,6 @@ limitations under the License. namespace tflite { -namespace { -inline const tflite::Model* VerifyAndGetModel(const void* buf, size_t len) { - ::flatbuffers::Verifier verifier(static_cast(buf), len); - if (VerifyModelBuffer(verifier)) { - return ::tflite::GetModel(buf); - } else { - return nullptr; - } -} -} // namespace - const char* kEmptyTensorName = ""; std::unique_ptr FlatBufferModel::BuildFromFile( @@ -80,10 +69,9 @@ FlatBufferModel::FlatBufferModel(const char* filename, bool mmap_file, } else { allocation_ = new FileCopyAllocation(filename, error_reporter); } - if (!allocation_->valid()) return; - if (!CheckModelIdentifier()) return; + if (!allocation_->valid() || !CheckModelIdentifier()) return; - model_ = VerifyAndGetModel(allocation_->base(), allocation_->bytes()); + model_ = ::tflite::GetModel(allocation_->base()); } bool FlatBufferModel::CheckModelIdentifier() const { @@ -104,7 +92,7 @@ FlatBufferModel::FlatBufferModel(const char* ptr, size_t num_bytes, allocation_ = new MemoryAllocation(ptr, num_bytes, error_reporter); if (!allocation_->valid()) return; - model_ = VerifyAndGetModel(allocation_->base(), allocation_->bytes()); + model_ = ::tflite::GetModel(allocation_->base()); } FlatBufferModel::FlatBufferModel(const Model* model, @@ -148,7 +136,7 @@ TfLiteStatus InterpreterBuilder::BuildLocalIndexToRegistrationMapping() { } } else if (!opcode->custom_code()) { error_reporter_->Report( - "Operator with builtin_code==0 has no custom_code.\n"); + "Operator with CUSTOM builtin_code has no custom_code.\n"); status = kTfLiteError; } else { const char* name = opcode->custom_code()->c_str(); @@ -231,7 +219,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, return kTfLiteActNone; case ActivationFunctionType_RELU: return kTfLiteActRelu; - case ActivationFunctionType_RELU1: + case ActivationFunctionType_RELU_N1_TO_1: return kTfLiteActRelu1; case ActivationFunctionType_RELU6: return kTfLiteActRelu6; @@ -287,9 +275,11 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, case BuiltinOperator_TANH: case BuiltinOperator_LOGISTIC: case BuiltinOperator_RELU: - case BuiltinOperator_RELU1: + case BuiltinOperator_RELU_N1_TO_1: case BuiltinOperator_RELU6: case BuiltinOperator_CONCAT_EMBEDDINGS: + case BuiltinOperator_EXP: + case BuiltinOperator_TOPK_V2: break; case BuiltinOperator_LSH_PROJECTION: { TfLiteLSHProjectionParams* params = @@ -340,7 +330,18 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, builtin_data = reinterpret_cast(params); break; } - case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN: + case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN: + case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN: { + TfLiteSequenceRNNParams* params = MallocPOD(); + if (auto* sequence_rnn_params = + op->builtin_options_as_SequenceRNNOptions()) { + params->activation = + parse_activation(sequence_rnn_params->fused_activation_function()); + params->time_major = sequence_rnn_params->time_major(); + } + builtin_data = reinterpret_cast(params); + break; + } case BuiltinOperator_RNN: { TfLiteRNNParams* params = MallocPOD(); if (auto* rnn_params = op->builtin_options_as_RNNOptions()) { @@ -415,6 +416,24 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, builtin_data = reinterpret_cast(params); break; } + case BuiltinOperator_DIV: { + auto* params = MallocPOD(); + if (auto* schema_params = op->builtin_options_as_DivOptions()) { + params->activation = + parse_activation(schema_params->fused_activation_function()); + } + builtin_data = reinterpret_cast(params); + break; + } + case BuiltinOperator_SUB: { + auto* params = MallocPOD(); + if (auto* schema_params = op->builtin_options_as_SubOptions()) { + params->activation = + parse_activation(schema_params->fused_activation_function()); + } + builtin_data = reinterpret_cast(params); + break; + } case BuiltinOperator_L2_NORMALIZATION: { auto* params = MallocPOD(); if (auto* schema_params = op->builtin_options_as_L2NormOptions()) { @@ -436,6 +455,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, builtin_data = reinterpret_cast(params); break; } + case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM: case BuiltinOperator_LSTM: { TfLiteLSTMParams* params = MallocPOD(); if (auto* lstm_params = op->builtin_options_as_LSTMOptions()) { @@ -451,32 +471,12 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, auto* params = MallocPOD(); if (auto* schema_params = op->builtin_options_as_ResizeBilinearOptions()) { - params->new_height = schema_params->new_height(); - params->new_width = schema_params->new_width(); + params->align_corners = schema_params->align_corners(); } builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_PAD: { - auto* params = MallocPOD(); - if (auto* schema_params = op->builtin_options_as_PadOptions()) { - auto* before_padding = schema_params->before_padding(); - FlatBufferIntVectorToArray(sizeof(params->before_padding), - before_padding, params->before_padding, - error_reporter); - - auto* after_padding = schema_params->after_padding(); - FlatBufferIntVectorToArray(sizeof(params->after_padding), after_padding, - params->after_padding, error_reporter); - - if (before_padding->Length() != after_padding->Length()) { - error_reporter->Report( - "Before padding and after padding arrays need to contain the " - "same number of dimensions.\n"); - } - params->num_dimensions = after_padding->Length(); - } - builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_RESHAPE: { @@ -518,20 +518,50 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, builtin_data = reinterpret_cast(params); break; } + case BuiltinOperator_SPACE_TO_BATCH_ND: { + break; + } case BuiltinOperator_BATCH_TO_SPACE_ND: { - auto* params = MallocPOD(); - if (auto* schema_params = - op->builtin_options_as_BatchToSpaceNDOptions()) { - const auto& block_shape = schema_params->block_shape(); - FlatBufferIntVectorToArray(sizeof(params->block_shape), block_shape, - params->block_shape, error_reporter); - const auto& before_crops = schema_params->before_crops(); - FlatBufferIntVectorToArray(sizeof(params->before_crops), before_crops, - params->before_crops, error_reporter); - const auto& after_crops = schema_params->after_crops(); - FlatBufferIntVectorToArray(sizeof(params->after_crops), after_crops, - params->after_crops, error_reporter); - params->num_spatial_dimensions = block_shape->Length(); + break; + } + case BuiltinOperator_TRANSPOSE: { + break; + } + case BuiltinOperator_MEAN: { + auto* params = MallocPOD(); + if (auto* schema_params = op->builtin_options_as_MeanOptions()) { + params->keep_dims = schema_params->keep_dims(); + } + builtin_data = reinterpret_cast(params); + break; + } + case BuiltinOperator_SPLIT: { + auto* params = MallocPOD(); + if (auto* schema_params = op->builtin_options_as_SplitOptions()) { + params->num_splits = schema_params->num_splits(); + } + builtin_data = reinterpret_cast(params); + break; + } + case BuiltinOperator_SQUEEZE: { + auto* params = MallocPOD(); + if (auto* schema_params = op->builtin_options_as_SqueezeOptions()) { + const auto& squeeze_dims = schema_params->squeeze_dims(); + FlatBufferIntVectorToArray(sizeof(params->squeeze_dims), squeeze_dims, + params->squeeze_dims, error_reporter); + params->num_squeeze_dims = squeeze_dims->Length(); + } + builtin_data = reinterpret_cast(params); + break; + } + case BuiltinOperator_STRIDED_SLICE: { + auto* params = MallocPOD(); + if (auto* schema_params = op->builtin_options_as_StridedSliceOptions()) { + params->begin_mask = schema_params->begin_mask(); + params->end_mask = schema_params->end_mask(); + params->ellipsis_mask = schema_params->ellipsis_mask(); + params->new_axis_mask = schema_params->new_axis_mask(); + params->shrink_axis_mask = schema_params->shrink_axis_mask(); } builtin_data = reinterpret_cast(params); break; diff --git a/tensorflow/contrib/lite/model.h b/tensorflow/contrib/lite/model.h index e0c96f7f0480cd3146f95a22957477809cf0096d..a467df5bb4eee3f6ce814512cb8b74bf09a6a4e7 100644 --- a/tensorflow/contrib/lite/model.h +++ b/tensorflow/contrib/lite/model.h @@ -31,8 +31,8 @@ limitations under the License. // OpResolver must be defined to provide your kernel implementations to the // interpreter. This is environment specific and may consist of just the builtin // ops, or some custom operators you defined to extend tflite. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODEL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODEL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_MODEL_H_ +#define TENSORFLOW_CONTRIB_LITE_MODEL_H_ #include #include "tensorflow/contrib/lite/error_reporter.h" @@ -173,4 +173,4 @@ class InterpreterBuilder { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODEL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_MODEL_H_ diff --git a/tensorflow/contrib/lite/model_test.cc b/tensorflow/contrib/lite/model_test.cc index 5330c8f594593655b2a8776cf6b399c0d16cdc19..66f22fd66a9ae0d35553a1f780ef73a5c5994c99 100644 --- a/tensorflow/contrib/lite/model_test.cc +++ b/tensorflow/contrib/lite/model_test.cc @@ -20,7 +20,6 @@ limitations under the License. #include #include #include -#include #include "tensorflow/contrib/lite/model.h" @@ -247,14 +246,6 @@ TEST(BasicFlatBufferModel, TestNullErrorReporter) { ASSERT_NE(interpreter->Invoke(), kTfLiteOk); } -// Test what happens if we cannot bind any of the ops. -TEST(BasicFlatBufferModel, TestBuildModelFromCorruptedData) { - std::string corrupted_data = "123"; - auto model = FlatBufferModel::BuildFromBuffer(corrupted_data.c_str(), - corrupted_data.length()); - ASSERT_FALSE(model); -} - // Test that loading model directly from a Model flatbuffer works. TEST(BasicFlatBufferModel, TestBuildFromModel) { TestErrorReporter reporter; diff --git a/tensorflow/contrib/lite/models/BUILD b/tensorflow/contrib/lite/models/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..6a1255b586ef04b80159156a78f0c4569a4661c5 --- /dev/null +++ b/tensorflow/contrib/lite/models/BUILD @@ -0,0 +1,26 @@ +# Model tests +package( + default_visibility = ["//visibility:public"], +) + +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + +load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") + +exports_files(glob([ + "testdata/*", +])) + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/contrib/lite/models/smartreply/g3doc/README.md b/tensorflow/contrib/lite/models/smartreply/g3doc/README.md index cab5dcca43a31ec3cf824f00d6794ea9e66d9bf8..a6d75648b3f3da98afd85daad6c2234e73a802e8 100644 --- a/tensorflow/contrib/lite/models/smartreply/g3doc/README.md +++ b/tensorflow/contrib/lite/models/smartreply/g3doc/README.md @@ -137,8 +137,8 @@ Following are the ops supported for using On-Device Smart Reply model: * **HASHTABLE_LOOKUP** - This is a custom op that uses label id from predict op and looks up the - response text from the given label id. + This is an op inside TensorFlow Lite that uses label id from predict op and + looks up the response text from the given label id. ## Further Information diff --git a/tensorflow/contrib/lite/models/smartreply/predictor.h b/tensorflow/contrib/lite/models/smartreply/predictor.h index d17323a3f9a0ea80ad5e215b0a4700e625d0c590..90260c8d620b0e756f72089d3f4d8d9f92d44fbe 100644 --- a/tensorflow/contrib/lite/models/smartreply/predictor.h +++ b/tensorflow/contrib/lite/models/smartreply/predictor.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODELS_SMARTREPLY_PREDICTOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODELS_SMARTREPLY_PREDICTOR_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_MODELS_SMARTREPLY_PREDICTOR_H_ +#define TENSORFLOW_CONTRIB_LITE_MODELS_SMARTREPLY_PREDICTOR_H_ #include #include @@ -77,4 +77,4 @@ struct SmartReplyConfig { } // namespace custom } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODELS_SMARTREPLY_PREDICTOR_H_ +#endif // TENSORFLOW_CONTRIB_LITE_MODELS_SMARTREPLY_PREDICTOR_H_ diff --git a/tensorflow/contrib/lite/models/smartreply/predictor_test.cc b/tensorflow/contrib/lite/models/smartreply/predictor_test.cc index 97d3c650e21c3cb4bef1db09df93f4bf24f38ba5..e6c8d966f1aff5a867f9469f8fcdec526df84763 100644 --- a/tensorflow/contrib/lite/models/smartreply/predictor_test.cc +++ b/tensorflow/contrib/lite/models/smartreply/predictor_test.cc @@ -22,8 +22,9 @@ limitations under the License. #include #include "absl/strings/str_cat.h" #include "absl/strings/str_split.h" -#include "tensorflow/contrib/lite/models/test_utils.h" +//#include "tensorflow/contrib/lite/models/test_utils.h" #include "tensorflow/contrib/lite/string_util.h" +#include "tensorflow/core/platform/test.h" namespace tflite { namespace custom { @@ -33,6 +34,11 @@ namespace { const char kModelName[] = "smartreply_ondevice_model.bin"; const char kSamples[] = "smartreply_samples.tsv"; +string TestDataPath() { + return string(StrCat(tensorflow::testing::TensorFlowSrcRoot(), "/", + "contrib/lite/models/testdata/")); +} + MATCHER_P(IncludeAnyResponesIn, expected_response, "contains the response") { bool has_expected_response = false; for (const auto &item : *arg) { diff --git a/tensorflow/contrib/lite/models/speech_asr_am_model_test.cc b/tensorflow/contrib/lite/models/speech_asr_am_model_test.cc deleted file mode 100644 index bf95b313f31c2f76046727353a9a7b0658dbf067..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/speech_asr_am_model_test.cc +++ /dev/null @@ -1,127 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ -// Unit test for speech ASR AM model using TFLite Ops. - -#include - -#include -#include - -#include "base/logging.h" -#include "file/base/path.h" -#include "testing/base/public/googletest.h" -#include -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/models/test_utils.h" - -namespace tflite { -namespace models { - -constexpr int kModelInputTensor = 0; -constexpr int kLstmLayer1OutputStateTensor = 19; -constexpr int kLstmLayer1CellStateTensor = 20; -constexpr int kLstmLayer2OutputStateTensor = 40; -constexpr int kLstmLayer2CellStateTensor = 41; -constexpr int kLstmLayer3OutputStateTensor = 61; -constexpr int kLstmLayer3CellStateTensor = 62; -constexpr int kLstmLayer4OutputStateTensor = 82; -constexpr int kLstmLayer4CellStateTensor = 83; -constexpr int kLstmLayer5OutputStateTensor = 103; -constexpr int kLstmLayer5CellStateTensor = 104; -constexpr int kModelOutputTensor = 109; - -TEST(SpeechAsrAm, RandomIOTest) { - // Read the model. - string tflite_file_path = - file::JoinPath(TestDataPath(), "speech_asr_am_model.tflite"); - auto model = FlatBufferModel::BuildFromFile(tflite_file_path.c_str()); - CHECK(model) << "Failed to mmap model " << tflite_file_path; - - // Initialize the interpreter. - ops::builtin::BuiltinOpResolver builtins; - std::unique_ptr interpreter; - InterpreterBuilder(*model, builtins)(&interpreter); - CHECK(interpreter != nullptr); - interpreter->AllocateTensors(); - - // Load the input frames. - Frames input_frames; - const string input_file_path = - file::JoinPath(TestDataPath(), "speech_asr_am_model_in.csv"); - ReadFrames(input_file_path, &input_frames); - - // Load the golden output results. - Frames output_frames; - const string output_file_path = - file::JoinPath(TestDataPath(), "speech_asr_am_model_out.csv"); - ReadFrames(output_file_path, &output_frames); - - const int speech_batch_size = - interpreter->tensor(kModelInputTensor)->dims->data[0]; - const int speech_input_size = - interpreter->tensor(kModelInputTensor)->dims->data[1]; - const int speech_output_size = - interpreter->tensor(kModelOutputTensor)->dims->data[1]; - - float* input_ptr = interpreter->tensor(kModelInputTensor)->data.f; - float* output_ptr = interpreter->tensor(kModelOutputTensor)->data.f; - - // Clear the LSTM state for layers. - memset(interpreter->tensor(kLstmLayer1OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer1OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer1CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer1CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer2OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer2OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer2CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer2CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer3OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer3OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer3CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer3CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer4OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer4OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer4CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer4CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer5OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer5OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer5CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer5CellStateTensor)->bytes); - - - for (int i = 0; i < input_frames.size(); i++) { - // Feed the input to model. - int frame_ptr = 0; - for (int k = 0; k < speech_input_size * speech_batch_size; k++) { - input_ptr[k] = input_frames[i][frame_ptr++]; - } - // Run the model. - interpreter->Invoke(); - // Validate the output. - for (int k = 0; k < speech_output_size; k++) { - ASSERT_NEAR(output_ptr[k], output_frames[i][k], 5.2e-4); - } - } -} - -} // namespace models -} // namespace tflite diff --git a/tensorflow/contrib/lite/models/speech_asr_lm_model_test.cc b/tensorflow/contrib/lite/models/speech_asr_lm_model_test.cc deleted file mode 100644 index 53f2b66da492f8fe56fa9e234f0951cf61c35037..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/speech_asr_lm_model_test.cc +++ /dev/null @@ -1,122 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ -// Unit test for speech ASR LM model using TFLite Ops. - -#include - -#include -#include - -#include "base/logging.h" -#include "file/base/path.h" -#include "testing/base/public/googletest.h" -#include -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/models/test_utils.h" - -namespace tflite { -namespace models { - -constexpr int kModelInput1Tensor = 0; -constexpr int kModelInput2Tensor = 66; -constexpr int kLstmLayer1OutputStateTensor = 21; -constexpr int kLstmLayer1CellStateTensor = 22; -constexpr int kLstmLayer2OutputStateTensor = 42; -constexpr int kLstmLayer2CellStateTensor = 43; -constexpr int kLstmLayer3OutputStateTensor = 63; -constexpr int kLstmLayer3CellStateTensor = 64; -constexpr int kModelOutputTensor = 75; - -static void ClearLstmStates(Interpreter* interpreter) { - memset(interpreter->tensor(kLstmLayer1OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer1OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer1CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer1CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer2OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer2OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer2CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer2CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer3OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer3OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer3CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer3CellStateTensor)->bytes); -} - -TEST(SpeechAsrLm, EndToEndTest) { - // Read the model. - string tflite_file_path = - file::JoinPath(TestDataPath(), "speech_asr_lm_model.tflite"); - auto model = FlatBufferModel::BuildFromFile(tflite_file_path.c_str()); - CHECK(model) << "Failed to mmap model " << tflite_file_path; - - // Initialize the interpreter. - ops::builtin::BuiltinOpResolver builtins; - std::unique_ptr interpreter; - InterpreterBuilder(*model, builtins)(&interpreter); - CHECK(interpreter != nullptr); - interpreter->AllocateTensors(); - - // Load the input frames. - Frames input_frames; - const string input_file_path = - file::JoinPath(TestDataPath(), "speech_asr_lm_model_in.csv"); - ReadFrames(input_file_path, &input_frames); - - // Load the golden output results. - Frames output_frames; - const string output_file_path = - file::JoinPath(TestDataPath(), "speech_asr_lm_model_out.csv"); - ReadFrames(output_file_path, &output_frames); - - CHECK_EQ(interpreter->tensor(kModelInput1Tensor)->dims->size, 1); - const int input1_size = - interpreter->tensor(kModelInput1Tensor)->dims->data[0]; - CHECK_EQ(input1_size, 1); - CHECK_EQ(interpreter->tensor(kModelInput2Tensor)->dims->size, 1); - const int output_size = - interpreter->tensor(kModelOutputTensor)->dims->data[0]; - CHECK_EQ(output_size, 1); - - int* input_lookup_ptr = interpreter->tensor(kModelInput1Tensor)->data.i32; - int* output_lookup_ptr = interpreter->tensor(kModelInput2Tensor)->data.i32; - float* output_ptr = interpreter->tensor(kModelOutputTensor)->data.f; - - - for (int i = 0; i < input_frames.size(); i++) { - float output_score = 0.0f; - // Reset LSTM states for each sequence. - ClearLstmStates(interpreter.get()); - // For subsequent inputs feed them sequentially, one-by-one. - for (int k = 1; k < input_frames[i].size(); k++) { - // Feed the inputs to model. - input_lookup_ptr[0] = static_cast(input_frames[i][k - 1]); - output_lookup_ptr[0] = static_cast(input_frames[i][k]); - // Run the model. - interpreter->Invoke(); - // Sum up the outputs. - output_score += output_ptr[0]; - } - // Validate the output. - ASSERT_NEAR(output_score, output_frames[i][0], 1.4e-5); - } -} - -} // namespace models -} // namespace tflite diff --git a/tensorflow/contrib/lite/models/speech_hotword_model_test.cc b/tensorflow/contrib/lite/models/speech_hotword_model_test.cc deleted file mode 100644 index f69cae8d2cb08678f9eec8c9b9d653cfce55bd2e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/speech_hotword_model_test.cc +++ /dev/null @@ -1,114 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ -// Unit test for speech Hotword model using TFLite Ops. - -#include - -#include -#include - -#include "base/logging.h" -#include "testing/base/public/googletest.h" -#include -#include "absl/strings/str_cat.h" -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/models/test_utils.h" - -namespace tflite { -namespace models { - -void RunTest(int model_input_tensor, int svdf_layer_state_tensor, - int model_output_tensor, const string& model_name, - const string& golden_in_name, const string& golden_out_name) { - // Read the model. - string tflite_file_path = StrCat(TestDataPath(), "/", model_name); - auto model = FlatBufferModel::BuildFromFile(tflite_file_path.c_str()); - CHECK(model) << "Failed to read model from file " << tflite_file_path; - - // Initialize the interpreter. - ops::builtin::BuiltinOpResolver builtins; - std::unique_ptr interpreter; - InterpreterBuilder(*model, builtins)(&interpreter); - CHECK(interpreter != nullptr); - interpreter->AllocateTensors(); - - // Reset the SVDF layer state. - memset(interpreter->tensor(svdf_layer_state_tensor)->data.raw, 0, - interpreter->tensor(svdf_layer_state_tensor)->bytes); - - // Load the input frames. - Frames input_frames; - const string input_file_path = StrCat(TestDataPath(), "/", golden_in_name); - ReadFrames(input_file_path, &input_frames); - - // Load the golden output results. - Frames output_frames; - const string output_file_path = StrCat(TestDataPath(), "/", golden_out_name); - ReadFrames(output_file_path, &output_frames); - - const int speech_batch_size = - interpreter->tensor(model_input_tensor)->dims->data[0]; - const int speech_input_size = - interpreter->tensor(model_input_tensor)->dims->data[1]; - const int speech_output_size = - interpreter->tensor(model_output_tensor)->dims->data[1]; - const int input_sequence_size = - input_frames[0].size() / (speech_input_size * speech_batch_size); - float* input_ptr = interpreter->tensor(model_input_tensor)->data.f; - float* output_ptr = interpreter->tensor(model_output_tensor)->data.f; - - // The first layer (SVDF) input size is 40 (speech_input_size). Each speech - // input frames for this model is 1600 floats, which can be fed to input in a - // sequence of size 40 (input_sequence_size). - for (int i = 0; i < TestInputSize(input_frames); i++) { - int frame_ptr = 0; - for (int s = 0; s < input_sequence_size; s++) { - for (int k = 0; k < speech_input_size * speech_batch_size; k++) { - input_ptr[k] = input_frames[i][frame_ptr++]; - } - interpreter->Invoke(); - } - // After the whole frame (1280 floats) is fed, we can check the output frame - // matches with the golden output frame. - for (int k = 0; k < speech_output_size; k++) { - ASSERT_NEAR(output_ptr[k], output_frames[i][k], 1e-5); - } - } -} - -TEST(SpeechHotword, OkGoogleTestRank1) { - constexpr int kModelInputTensor = 0; - constexpr int kSvdfLayerStateTensor = 4; - constexpr int kModelOutputTensor = 18; - - RunTest(kModelInputTensor, kSvdfLayerStateTensor, kModelOutputTensor, - "speech_hotword_model_rank1.tflite", "speech_hotword_model_in.csv", - "speech_hotword_model_out_rank1.csv"); -} - -TEST(SpeechHotword, OkGoogleTestRank2) { - constexpr int kModelInputTensor = 17; - constexpr int kSvdfLayerStateTensor = 1; - constexpr int kModelOutputTensor = 18; - RunTest(kModelInputTensor, kSvdfLayerStateTensor, kModelOutputTensor, - "speech_hotword_model_rank2.tflite", "speech_hotword_model_in.csv", - "speech_hotword_model_out_rank2.csv"); -} - -} // namespace models -} // namespace tflite diff --git a/tensorflow/contrib/lite/models/speech_speakerid_model_test.cc b/tensorflow/contrib/lite/models/speech_speakerid_model_test.cc deleted file mode 100644 index 9da0fb1fc62360dcf584c4a08f99b0cef9964a0d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/speech_speakerid_model_test.cc +++ /dev/null @@ -1,114 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ -// Unit test for speech SpeakerId model using TFLite Ops. - -#include - -#include -#include - -#include "base/logging.h" -#include "testing/base/public/googletest.h" -#include -#include "absl/strings/str_cat.h" -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/models/test_utils.h" -#include "tensorflow/contrib/lite/tools/mutable_op_resolver.h" - -void RegisterSelectedOps(::tflite::MutableOpResolver* resolver); - -namespace tflite { -namespace models { - -constexpr int kModelInputTensor = 0; -constexpr int kLstmLayer1OutputStateTensor = 19; -constexpr int kLstmLayer1CellStateTensor = 20; -constexpr int kLstmLayer2OutputStateTensor = 40; -constexpr int kLstmLayer2CellStateTensor = 41; -constexpr int kLstmLayer3OutputStateTensor = 61; -constexpr int kLstmLayer3CellStateTensor = 62; -constexpr int kModelOutputTensor = 66; - -TEST(SpeechSpeakerId, OkGoogleTest) { - // Read the model. - string tflite_file_path = - StrCat(TestDataPath(), "/", "speech_speakerid_model.tflite"); - auto model = FlatBufferModel::BuildFromFile(tflite_file_path.c_str()); - CHECK(model) << "Failed to read model from file " << tflite_file_path; - - // Initialize the interpreter. - ::tflite::MutableOpResolver resolver; - RegisterSelectedOps(&resolver); - std::unique_ptr interpreter; - InterpreterBuilder(*model, resolver)(&interpreter); - CHECK(interpreter != nullptr); - interpreter->AllocateTensors(); - - // Load the input frames. - Frames input_frames; - const string input_file_path = - StrCat(TestDataPath(), "/", "speech_speakerid_model_in.csv"); - ReadFrames(input_file_path, &input_frames); - - // Load the golden output results. - Frames output_frames; - const string output_file_path = - StrCat(TestDataPath(), "/", "speech_speakerid_model_out.csv"); - ReadFrames(output_file_path, &output_frames); - - const int speech_batch_size = - interpreter->tensor(kModelInputTensor)->dims->data[0]; - const int speech_input_size = - interpreter->tensor(kModelInputTensor)->dims->data[1]; - const int speech_output_size = - interpreter->tensor(kModelOutputTensor)->dims->data[1]; - - float* input_ptr = interpreter->tensor(kModelInputTensor)->data.f; - float* output_ptr = interpreter->tensor(kModelOutputTensor)->data.f; - - // Clear the LSTM state for layers. - memset(interpreter->tensor(kLstmLayer1OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer1OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer1CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer1CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer2OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer2OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer2CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer2CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer3OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer3OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer3CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer3CellStateTensor)->bytes); - for (int i = 0; i < input_frames.size(); i++) { - // Feed the input to model. - int frame_ptr = 0; - for (int k = 0; k < speech_input_size * speech_batch_size; k++) { - input_ptr[k] = input_frames[i][frame_ptr++]; - } - // Run the model. - interpreter->Invoke(); - // Validate the output. - for (int k = 0; k < speech_output_size; k++) { - ASSERT_NEAR(output_ptr[k], output_frames[i][k], 1e-5); - } - } -} - -} // namespace models -} // namespace tflite diff --git a/tensorflow/contrib/lite/models/speech_test.cc b/tensorflow/contrib/lite/models/speech_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..daa8c3100b64e9290256aa14a6ab641f19174a0a --- /dev/null +++ b/tensorflow/contrib/lite/models/speech_test.cc @@ -0,0 +1,189 @@ +/* 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. +==============================================================================*/ +// Unit test for speech models (Hotword, SpeakerId) using TFLite Ops. + +#include +#include + +#include + +#include "testing/base/public/googletest.h" +#include +#include "tensorflow/contrib/lite/testing/parse_testdata.h" +#include "tensorflow/contrib/lite/testing/split.h" +#include "tensorflow/contrib/lite/testing/tflite_driver.h" + +namespace tflite { +namespace { + +const char kDataPath[] = "third_party/tensorflow/contrib/lite/models/testdata/"; + +bool Init(const string& in_file_name, testing::TfLiteDriver* driver, + std::ifstream* in_file) { + driver->SetModelBaseDir(kDataPath); + in_file->open(string(kDataPath) + in_file_name, std::ifstream::in); + return in_file->is_open(); +} + +// Converts a set of test files provided by the speech team into a single +// test_spec. Input CSV files are supposed to contain a number of sequences per +// line. Each sequence maps to a single invocation of the interpreter and the +// output tensor after all sequences have run is compared to the corresponding +// line in the output CSV file. +bool ConvertCsvData(const string& model_name, const string& in_name, + const string& out_name, const string& input_tensor, + const string& output_tensor, + const string& persistent_tensors, int sequence_size, + std::ostream* out) { + auto data_path = [](const string& s) { return string(kDataPath) + s; }; + + *out << "load_model: \"" << data_path(model_name) << "\"" << std::endl; + + *out << "init_state: \"" << persistent_tensors << "\"" << std::endl; + + string in_file_name = data_path(in_name); + std::ifstream in_file(in_file_name); + if (!in_file.is_open()) { + std::cerr << "Failed to open " << in_file_name << std::endl; + return false; + } + string out_file_name = data_path(out_name); + std::ifstream out_file(out_file_name); + if (!out_file.is_open()) { + std::cerr << "Failed to open " << out_file_name << std::endl; + return false; + } + + int invocation_count = 0; + string in_values; + while (std::getline(in_file, in_values, '\n')) { + std::vector input = testing::Split(in_values, ","); + int num_sequences = input.size() / sequence_size; + + for (int j = 0; j < num_sequences; ++j) { + *out << "invoke {" << std::endl; + *out << " id: " << invocation_count << std::endl; + *out << " input: \""; + for (int k = 0; k < sequence_size; ++k) { + *out << input[k + j * sequence_size] << ","; + } + *out << "\"" << std::endl; + + if (j == num_sequences - 1) { + string out_values; + if (!std::getline(out_file, out_values, '\n')) { + std::cerr << "Not enough lines in " << out_file_name << std::endl; + return false; + } + *out << " output: \"" << out_values << "\"" << std::endl; + } + + *out << "}" << std::endl; + ++invocation_count; + } + } + return true; +} + +TEST(SpeechTest, HotwordOkGoogleRank1Test) { + std::stringstream os; + ASSERT_TRUE(ConvertCsvData( + "speech_hotword_model_rank1.tflite", "speech_hotword_model_in.csv", + "speech_hotword_model_out_rank1.csv", /*input_tensor=*/"0", + /*output_tensor=*/"18", /*persistent_tensors=*/"4", + /*sequence_size=*/40, &os)); + testing::TfLiteDriver test_driver(/*use_nnapi=*/false); + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + << test_driver.GetErrorMessage(); +} + +TEST(SpeechTest, HotwordOkGoogleRank2Test) { + std::stringstream os; + ASSERT_TRUE(ConvertCsvData( + "speech_hotword_model_rank2.tflite", "speech_hotword_model_in.csv", + "speech_hotword_model_out_rank2.csv", /*input_tensor=*/"17", + /*output_tensor=*/"18", /*persistent_tensors=*/"1", + /*sequence_size=*/40, &os)); + testing::TfLiteDriver test_driver(/*use_nnapi=*/false); + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + << test_driver.GetErrorMessage(); +} + +TEST(SpeechTest, SpeakerIdOkGoogleTest) { + std::stringstream os; + ASSERT_TRUE(ConvertCsvData( + "speech_speakerid_model.tflite", "speech_speakerid_model_in.csv", + "speech_speakerid_model_out.csv", /*input_tensor=*/"0", + /*output_tensor=*/"66", + /*persistent_tensors=*/"19,20,40,41,61,62", + /*sequence_size=*/80, &os)); + testing::TfLiteDriver test_driver(/*use_nnapi=*/false); + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + << test_driver.GetErrorMessage(); +} + +TEST(SpeechTest, AsrAmTest) { + std::stringstream os; + ASSERT_TRUE( + ConvertCsvData("speech_asr_am_model.tflite", "speech_asr_am_model_in.csv", + "speech_asr_am_model_out.csv", /*input_tensor=*/"0", + /*output_tensor=*/"109", + /*persistent_tensors=*/"19,20,40,41,61,62,82,83,103,104", + /*sequence_size=*/320, &os)); + testing::TfLiteDriver test_driver(/*use_nnapi=*/false); + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + << test_driver.GetErrorMessage(); +} + +// The original version of speech_asr_lm_model_test.cc ran a few sequences +// through the interpreter and stored the sum of all the output, which was them +// compared for correctness. In this test we are comparing all the intermediate +// results. +TEST(SpeechTest, AsrLmTest) { + std::ifstream in_file; + testing::TfLiteDriver test_driver(/*use_nnapi=*/false); + ASSERT_TRUE(Init("speech_asr_lm_model.test_spec", &test_driver, &in_file)); + ASSERT_TRUE(testing::ParseAndRunTests(&in_file, &test_driver)) + << test_driver.GetErrorMessage(); +} + +TEST(SpeechTest, EndpointerTest) { + std::stringstream os; + ASSERT_TRUE(ConvertCsvData( + "speech_endpointer_model.tflite", "speech_endpointer_model_in.csv", + "speech_endpointer_model_out.csv", /*input_tensor=*/"0", + /*output_tensor=*/"58", + /*persistent_tensors=*/"28,29,49,50", + /*sequence_size=*/320, &os)); + testing::TfLiteDriver test_driver(/*use_nnapi=*/false); + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + << test_driver.GetErrorMessage(); +} + +TEST(SpeechTest, TtsTest) { + std::stringstream os; + ASSERT_TRUE(ConvertCsvData("speech_tts_model.tflite", + "speech_tts_model_in.csv", + "speech_tts_model_out.csv", /*input_tensor=*/"0", + /*output_tensor=*/"74", + /*persistent_tensors=*/"25,26,46,47,67,68,73", + /*sequence_size=*/334, &os)); + testing::TfLiteDriver test_driver(/*use_nnapi=*/false); + ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver)) + << test_driver.GetErrorMessage(); +} + +} // namespace +} // namespace tflite diff --git a/tensorflow/contrib/lite/models/speech_tts_model_test.cc b/tensorflow/contrib/lite/models/speech_tts_model_test.cc deleted file mode 100644 index 88291776892f3186ca5bfc726e814f8d23d73b11..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/speech_tts_model_test.cc +++ /dev/null @@ -1,116 +0,0 @@ -/* Copyright 2017 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. -==============================================================================*/ -// Unit test for speech TTS model using TFLite Ops. - -#include - -#include -#include - -#include "base/logging.h" -#include "testing/base/public/googletest.h" -#include -#include "absl/strings/str_cat.h" -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/models/test_utils.h" - -namespace tflite { -namespace models { - -constexpr int kModelInputTensor = 0; -constexpr int kLstmLayer1OutputStateTensor = 25; -constexpr int kLstmLayer1CellStateTensor = 26; -constexpr int kLstmLayer2OutputStateTensor = 46; -constexpr int kLstmLayer2CellStateTensor = 47; -constexpr int kLstmLayer3OutputStateTensor = 67; -constexpr int kLstmLayer3CellStateTensor = 68; -constexpr int kRnnLayerHiddenStateTensor = 73; -constexpr int kModelOutputTensor = 74; - -TEST(SpeechTTS, RandomIOTest) { - // Read the model. - string tflite_file_path = - StrCat(TestDataPath(), "/", "speech_tts_model.tflite"); - auto model = FlatBufferModel::BuildFromFile(tflite_file_path.c_str()); - CHECK(model) << "Failed to mmap model " << tflite_file_path; - - // Initialize the interpreter. - ops::builtin::BuiltinOpResolver builtins; - std::unique_ptr interpreter; - InterpreterBuilder(*model, builtins)(&interpreter); - CHECK(interpreter != nullptr); - interpreter->AllocateTensors(); - - // Load the input frames. - Frames input_frames; - const string input_file_path = - StrCat(TestDataPath(), "/", "speech_tts_model_in.csv"); - ReadFrames(input_file_path, &input_frames); - - // Load the golden output results. - Frames output_frames; - const string output_file_path = - StrCat(TestDataPath(), "/", "speech_tts_model_out.csv"); - ReadFrames(output_file_path, &output_frames); - - const int speech_batch_size = - interpreter->tensor(kModelInputTensor)->dims->data[0]; - const int speech_input_size = - interpreter->tensor(kModelInputTensor)->dims->data[1]; - const int speech_output_size = - interpreter->tensor(kModelOutputTensor)->dims->data[1]; - - float* input_ptr = interpreter->tensor(kModelInputTensor)->data.f; - float* output_ptr = interpreter->tensor(kModelOutputTensor)->data.f; - - // Clear the LSTM state for layers. - memset(interpreter->tensor(kLstmLayer1OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer1OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer1CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer1CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer2OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer2OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer2CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer2CellStateTensor)->bytes); - - memset(interpreter->tensor(kLstmLayer3OutputStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer3OutputStateTensor)->bytes); - memset(interpreter->tensor(kLstmLayer3CellStateTensor)->data.raw, 0, - interpreter->tensor(kLstmLayer3CellStateTensor)->bytes); - - memset(interpreter->tensor(kRnnLayerHiddenStateTensor)->data.raw, 0, - interpreter->tensor(kRnnLayerHiddenStateTensor)->bytes); - - for (int i = 0; i < input_frames.size(); i++) { - // Feed the input to model. - int frame_ptr = 0; - for (int k = 0; k < speech_input_size * speech_batch_size; k++) { - input_ptr[k] = input_frames[i][frame_ptr++]; - } - // Run the model. - interpreter->Invoke(); - // Validate the output. - for (int k = 0; k < speech_output_size; k++) { - ASSERT_NEAR(output_ptr[k], output_frames[i][k], 1e-5); - } - } -} - -} // namespace models -} // namespace tflite diff --git a/tensorflow/contrib/lite/models/test_utils.h b/tensorflow/contrib/lite/models/test_utils.h deleted file mode 100644 index 1e14c26a3544ed44f9395ff3b59a70551a1a6394..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/test_utils.h +++ /dev/null @@ -1,84 +0,0 @@ -/* Copyright 2017 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODELS_TEST_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODELS_TEST_UTILS_H_ - -#include -#include - -#include -#include -#include -#include - -namespace tflite { -namespace models { -using Frames = std::vector>; -} // namespace models -} // namespace tflite - -#ifndef __ANDROID__ -#include "absl/strings/str_cat.h" -#include "tensorflow/core/platform/test.h" - -inline string TestDataPath() { - return string(StrCat(tensorflow::testing::TensorFlowSrcRoot(), "/", - "contrib/lite/models/testdata/")); -} -inline int TestInputSize(const tflite::models::Frames& input_frames) { - return input_frames.size(); -} -#else -inline string TestDataPath() { - return string("third_party/tensorflow/contrib/lite/models/testdata/"); -} - -inline int TestInputSize(const tflite::models::Frames& input_frames) { - // Android TAP is very slow, we only test the first 20 frames. - return 20; -} -#endif - -namespace tflite { -namespace models { - -// Read float data from a comma-separated file: -// Each line will be read into a float vector. -// The return result will be a vector of float vectors. -void ReadFrames(const string& csv_file_path, Frames* frames) { - std::ifstream csv_file(csv_file_path); - string line; - while (std::getline(csv_file, line, '\n')) { - std::vector fields; - // Used by strtok_r internaly for successive calls on the same string. - char* save_ptr = nullptr; - - // Tokenize the line. - char* next_token = - strtok_r(const_cast(line.c_str()), ",", &save_ptr); - while (next_token != nullptr) { - float f = strtod(next_token, nullptr); - fields.push_back(f); - next_token = strtok_r(nullptr, ",", &save_ptr); - } - frames->push_back(fields); - } - csv_file.close(); -} - -} // namespace models -} // namespace tflite - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_MODELS_TEST_UTILS_H_ diff --git a/tensorflow/contrib/lite/models/testdata/g3doc/README.md b/tensorflow/contrib/lite/models/testdata/g3doc/README.md index 46b24248f002b8a1a30a2ac614c3874dfd2207db..1c47e00aae2a0e76ba04004a2fc3cc02ec4536f7 100644 --- a/tensorflow/contrib/lite/models/testdata/g3doc/README.md +++ b/tensorflow/contrib/lite/models/testdata/g3doc/README.md @@ -53,7 +53,7 @@ with the corresponding parameters as shown in the figure. ### Automatic Speech Recognizer (ASR) Acoustic Model (AM) The acoustic model for automatic speech recognition is the neural network model -for matching phonemes to the input autio features. It generates posterior +for matching phonemes to the input audio features. It generates posterior probabilities of phonemes from speech frontend features (log-mel filterbanks). It has an input size of 320 (float), an output size of 42 (float), five LSTM layers and one fully connected layers with a Softmax activation function, with @@ -68,13 +68,27 @@ for predicting the probability of a word given previous words in a sentence. It generates posterior probabilities of the next word based from a sequence of words. The words are encoded as indices in a fixed size dictionary. The model has two inputs both of size one (integer): the current word index and -next word index, an output size of one (float): the log probability. It consits +next word index, an output size of one (float): the log probability. It consists of three embedding layer, three LSTM layers, followed by a multiplication, a fully connected layers and an addition. The corresponding parameters as shown in the figure. ![asr_lm_model](asr_lm.svg "ASR LM model") +### Endpointer Model + +The endpointer model is the neural network model for predicting end of speech +in an utterance. More precisely, it generates posterior probabilities of various +events that allow detection of speech start and end events. +It has an input size of 40 (float) which are speech frontend features +(log-mel filterbanks), and an output size of four corresponding to: +speech, intermediate non-speech, initial non-speech, and final non-speech. +The model consists of a convolutional layer, followed by a fully-connected +layer, two LSTM layers, and two additional fully-connected layers. +The corresponding parameters as shown in the figure. +![endpointer_model](endpointer.svg "Endpointer model") + + ## Speech models test input/output generation As mentioned above the input to models are generated from a pre-processing @@ -113,7 +127,13 @@ test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/li test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_tts_model_test.cc) [ASR AM model -test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_terse_am_model_test.cc) +test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_asr_am_model_test.cc) + +[ASR LM model +test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_asr_lm_model_test.cc) + +[Endpointer model +test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_endpointer_model_test.cc) ## Android Support The models have been tested on Android phones, using the following tests: diff --git a/tensorflow/contrib/lite/models/testdata/g3doc/endpointer.svg b/tensorflow/contrib/lite/models/testdata/g3doc/endpointer.svg new file mode 100644 index 0000000000000000000000000000000000000000..6033bdc529e18355131965a26c49b6f17d671f27 --- /dev/null +++ b/tensorflow/contrib/lite/models/testdata/g3doc/endpointer.svg @@ -0,0 +1,4 @@ + + + + diff --git a/tensorflow/contrib/lite/models/testdata/speech_asr_lm_model.test_spec b/tensorflow/contrib/lite/models/testdata/speech_asr_lm_model.test_spec new file mode 100644 index 0000000000000000000000000000000000000000..5812de4b30382f6b031c907bf8bd12a34ac9e0b3 --- /dev/null +++ b/tensorflow/contrib/lite/models/testdata/speech_asr_lm_model.test_spec @@ -0,0 +1,202 @@ +load_model: "speech_asr_lm_model.tflite" +init_state: "21,22,42,43,63,64" +invoke { + id: 3 + input: "63982" + input: "8409" + output: "-2.75389" +} +invoke { + id: 4 + input: "8409" + input: "1488" + output: "0.601841" +} +invoke { + id: 5 + input: "1488" + input: "63981" + output: "-0.314846" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 6 + input: "63982" + input: "8409" + output: "-2.75389" +} +invoke { + id: 7 + input: "8409" + input: "3082" + output: "-3.63721" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 8 + input: "63982" + input: "8409" + output: "-2.75389" +} +invoke { + id: 9 + input: "8409" + input: "18965" + output: "-6.93985" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 13 + input: "63982" + input: "12516" + output: "-6.20867" +} +invoke { + id: 14 + input: "12516" + input: "914" + output: "-0.407277" +} +invoke { + id: 15 + input: "914" + input: "63981" + output: "-3.82091" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 19 + input: "63982" + input: "12516" + output: "-6.20867" +} +invoke { + id: 20 + input: "12516" + input: "914" + output: "-0.407277" +} +invoke { + id: 21 + input: "914" + input: "48619" + output: "-4.02131" +} +invoke { + id: 22 + input: "48619" + input: "63981" + output: "-0.677399" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 26 + input: "63982" + input: "12516" + output: "-6.20867" +} +invoke { + id: 27 + input: "12516" + input: "914" + output: "-0.407277" +} +invoke { + id: 28 + input: "914" + input: "4700" + output: "-4.056" +} +invoke { + id: 29 + input: "4700" + input: "63981" + output: "0.415889" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 30 + input: "63982" + input: "12516" + output: "-6.20867" +} +invoke { + id: 31 + input: "12516" + input: "914" + output: "-0.407277" +invoke { + id: 32 + input: "914" + input: "51923" + output: "-14.1147" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 34 + input: "63982" + input: "5520" + output: "-4.56971" +} +invoke { + id: 35 + input: "5520" + input: "16318" + output: "-1.54815" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 36 + input: "63982" + input: "5520" + output: "-4.56971" +} +invoke { + id: 37 + input: "5520" + input: "28303" + output: "-14.0947" +} +init_state: "21,22,42,43,63,64" +invoke { + id: 38 + input: "63982" + input: "12451" + output: "-6.24243" +} +invoke { + id: 39 + input: "12451" + input: "752" + output: "0.0700736" +} +invoke { + id: 40 + input: "752" + input: "11" + output: "-1.72744" +} +invoke { + id: 41 + input: "11" + input: "19454" + output: "-3.19211" +} +invoke { + id: 42 + input: "19454" + input: "16989" + output: "-4.01684" +} +invoke { + id: 43 + input: "16989" + input: "40168" + output: "-8.91317" +} +invoke { + id: 44 + input: "40168" + input: "63981" + output: "-0.675377" +} diff --git a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h index 80668890786becd161f9fd07317970b199ddb044..76032771af2c8e099aed498b2071816646f3b606 100644 --- a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h +++ b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h @@ -370,7 +370,7 @@ enum { * Looks up items from a given tensor. * * Each item in the output is a raw copy of the corresponding item in - * the input “values”. If the the given “lookup” indices are out of bounds, + * the input “values”. If the given “lookup” indices are out of bounds, * the op will fail and an error will be reported. * * Inputs: @@ -1571,7 +1571,7 @@ inline int ANeuralNetworksModel_addOperation(ANeuralNetworksModel* model, } /** - * Specfifies which operands will be the model's inputs and outputs. + * Specifies which operands will be the model's inputs and outputs. * * An operand cannot be used for both input and output. Doing so will * return an error. @@ -1774,7 +1774,7 @@ inline int ANeuralNetworksExecution_setInput( * model. If the type is the same as specified when the model * was built, NULL can be passed. * @param memory The memory containing the data. - * @param offset This specifies the location of the data whithin the memory. + * @param offset This specifies the location of the data within the memory. * The offset is in bytes from the start of memory. * @param length The size in bytes of the data value. * @@ -1841,7 +1841,7 @@ inline int ANeuralNetworksExecution_setOutput( * model. If the type is the same as specified when the model * was built, NULL can be passed. * @param memory The memory where the data is to be stored. - * @param offset This specifies the location of the data whithin the memory. + * @param offset This specifies the location of the data within the memory. * The offset is in bytes from the start of memory. * @param length The length in bytes of the data value. * diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc index 5cb0afcea0286b847ca8365548db419ad619d3e6..02e8499f61c6a3d5fceb978aa0e63a4ee90cf19a 100644 --- a/tensorflow/contrib/lite/nnapi_delegate.cc +++ b/tensorflow/contrib/lite/nnapi_delegate.cc @@ -161,6 +161,14 @@ void AddOpsAndParams(tflite::Interpreter* interpreter, augmented_inputs.push_back(next_id++); }; + auto duplicate_state_tensor_float32 = + [interpreter, &nn_model, &augmented_inputs, &next_id](int tensor_id) { + const TfLiteTensor* tensor = interpreter->tensor(tensor_id); + CHECK_NN(ANeuralNetworksModel_setOperandValue( + nn_model, tensor_id, tensor->data.raw, tensor->bytes)); + augmented_inputs.push_back(tensor_id); + }; + auto add_add_params = [&add_scalar_int32]() { add_scalar_int32(0); }; auto add_pooling_params = [&add_scalar_int32](void* data) { @@ -213,6 +221,14 @@ void AddOpsAndParams(tflite::Interpreter* interpreter, add_scalar_int32(builtin->block_size); }; + auto add_lstm_params = [&add_scalar_int32, + &add_scalar_float32](void* data) { + auto builtin = reinterpret_cast(data); + add_scalar_int32(builtin->activation); + add_scalar_float32(builtin->cell_clip); + add_scalar_float32(builtin->proj_clip); + }; + #if 0 auto add_reshape_params = [&](void* data) { auto builtin = reinterpret_cast(data); @@ -289,15 +305,25 @@ void AddOpsAndParams(tflite::Interpreter* interpreter, add_space_to_depth_params(node.builtin_data); nn_op_type = ANEURALNETWORKS_SPACE_TO_DEPTH; break; + case tflite::BuiltinOperator_LSTM: { + duplicate_state_tensor_float32( + node.outputs->data[/*kOutputStateTensor*/ 1]); + duplicate_state_tensor_float32( + node.outputs->data[/*kCellStateTensor*/ 2]); + add_lstm_params(node.builtin_data); + nn_op_type = ANEURALNETWORKS_LSTM; + break; + } case tflite::BuiltinOperator_CONCAT_EMBEDDINGS: case tflite::BuiltinOperator_LSH_PROJECTION: case tflite::BuiltinOperator_SVDF: case tflite::BuiltinOperator_HASHTABLE_LOOKUP: case tflite::BuiltinOperator_RNN: + case tflite::BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN: case tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN: case tflite::BuiltinOperator_EMBEDDING_LOOKUP: case tflite::BuiltinOperator_EMBEDDING_LOOKUP_SPARSE: - case tflite::BuiltinOperator_LSTM: + case tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM: case tflite::BuiltinOperator_L2_NORMALIZATION: case tflite::BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION: case tflite::BuiltinOperator_MUL: @@ -305,9 +331,19 @@ void AddOpsAndParams(tflite::Interpreter* interpreter, case tflite::BuiltinOperator_RESIZE_BILINEAR: case tflite::BuiltinOperator_CALL: case tflite::BuiltinOperator_SKIP_GRAM: - case tflite::BuiltinOperator_RELU1: + case tflite::BuiltinOperator_RELU_N1_TO_1: case tflite::BuiltinOperator_GATHER: + case tflite::BuiltinOperator_SPACE_TO_BATCH_ND: case tflite::BuiltinOperator_BATCH_TO_SPACE_ND: + case tflite::BuiltinOperator_TOPK_V2: + case tflite::BuiltinOperator_TRANSPOSE: + case tflite::BuiltinOperator_MEAN: + case tflite::BuiltinOperator_DIV: + case tflite::BuiltinOperator_SUB: + case tflite::BuiltinOperator_SPLIT: + case tflite::BuiltinOperator_SQUEEZE: + case tflite::BuiltinOperator_STRIDED_SLICE: + case tflite::BuiltinOperator_EXP: FATAL("Op code %d is currently not delegated to NNAPI", builtin); nn_op_type = -1; // set to invalid break; diff --git a/tensorflow/contrib/lite/nnapi_delegate.h b/tensorflow/contrib/lite/nnapi_delegate.h index f29aa9e18e605ef0b5d246b2a672639c64391646..e98000929a1168c786f6c18f498f9d1d72311ada 100644 --- a/tensorflow/contrib/lite/nnapi_delegate.h +++ b/tensorflow/contrib/lite/nnapi_delegate.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_NNAPI_DELEGATE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_NNAPI_DELEGATE_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_NNAPI_DELEGATE_H_ +#define TENSORFLOW_CONTRIB_LITE_NNAPI_DELEGATE_H_ #include "tensorflow/contrib/lite/allocation.h" #include "tensorflow/contrib/lite/context.h" @@ -63,4 +63,4 @@ class NNAPIDelegate { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_NNAPI_DELEGATE_H_ +#endif // TENSORFLOW_CONTRIB_LITE_NNAPI_DELEGATE_H_ diff --git a/tensorflow/contrib/lite/optional_debug_tools.h b/tensorflow/contrib/lite/optional_debug_tools.h index 54d48760951c946d0493a86961348df25e53bd1f..1b6998cda382782b974bea3d18ffb6217e8f780c 100644 --- a/tensorflow/contrib/lite/optional_debug_tools.h +++ b/tensorflow/contrib/lite/optional_debug_tools.h @@ -14,8 +14,8 @@ limitations under the License. ==============================================================================*/ // Optional debugging functionality. For small sized binaries, these are not // needed. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_DEBUG_TOOLS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_DEBUG_TOOLS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_DEBUG_TOOLS_H_ +#define TENSORFLOW_CONTRIB_LITE_DEBUG_TOOLS_H_ #include "tensorflow/contrib/lite/interpreter.h" @@ -29,4 +29,4 @@ TfLiteStatus ValidateInterpreterState(const Interpreter* interpreter); } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_DEBUG_TOOLS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_DEBUG_TOOLS_H_ diff --git a/tensorflow/contrib/lite/python/BUILD b/tensorflow/contrib/lite/python/BUILD index 3d6a3ec0fd4c673f601254b19452bbf8b9454e27..82feae0f0041997949212613c654a5695f468d56 100644 --- a/tensorflow/contrib/lite/python/BUILD +++ b/tensorflow/contrib/lite/python/BUILD @@ -13,6 +13,7 @@ py_library( srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ + ":op_hint", "//tensorflow/contrib/lite/toco:model_flags_proto_py", "//tensorflow/contrib/lite/toco:toco_flags_proto_py", "//tensorflow/contrib/lite/toco/python:tensorflow_wrap_toco", @@ -20,6 +21,18 @@ py_library( ], ) +py_library( + name = "op_hint", + srcs = ["op_hint.py"], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/contrib/framework:framework_py", + "//tensorflow/core:protos_all_py", + "//tensorflow/python:platform", + ], +) + py_test( name = "lite_test", srcs = ["lite_test.py"], @@ -27,6 +40,7 @@ py_test( tags = ["no_oss"], deps = [ ":lite", + ":op_hint", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", diff --git a/tensorflow/contrib/lite/python/lite.py b/tensorflow/contrib/lite/python/lite.py index 95309478a6f9791e3510736c45f9c5cfab88703b..5d2f21653762a405a57288a7ba38323e5e42b3e1 100644 --- a/tensorflow/contrib/lite/python/lite.py +++ b/tensorflow/contrib/lite/python/lite.py @@ -18,23 +18,36 @@ EXPERIMENTAL: APIs here are unstable and likely to change without notice. @@toco_convert @@toco_convert_protos +@@OpHint +@@convert_op_hints_to_stubs """ from __future__ import absolute_import from __future__ import division from __future__ import print_function - import os import subprocess import tempfile +# pylint: disable=unused-import +from tensorflow.contrib.lite.python.op_hint import convert_op_hints_to_stubs +from tensorflow.contrib.lite.python.op_hint import OpHint +# pylint: enable=unused-import from tensorflow.contrib.lite.toco import model_flags_pb2 as _model_flags_pb2 from tensorflow.contrib.lite.toco import toco_flags_pb2 as _toco_flags_pb2 from tensorflow.contrib.lite.toco import types_pb2 as _types_pb2 -from tensorflow.contrib.lite.toco.python.tensorflow_wrap_toco import TocoConvert as _toco_convert_protos from tensorflow.python.framework import dtypes as _dtypes from tensorflow.python.platform import resource_loader as _resource_loader from tensorflow.python.util.all_util import remove_undocumented +from tensorflow.python.util.lazy_loader import LazyLoader + +# Lazy load since some of the performance benchmark skylark rules +# break dependencies. +_toco_python = LazyLoader( + "tensorflow_wrap_toco", globals(), + "tensorflow.contrib.lite.toco.python." + "tensorflow_wrap_toco") +del LazyLoader # Enum types from the protobuf promoted to the API FLOAT = _types_pb2.FLOAT @@ -86,7 +99,8 @@ def toco_convert_protos(model_flags_str, toco_flags_str, input_data_str): # TODO(aselle): When toco does not use fatal errors for failure, we can # switch this on. if not _toco_from_proto_bin: - return _toco_convert_protos(model_flags_str, toco_flags_str, input_data_str) + return _toco_python.TocoConvert( + model_flags_str, toco_flags_str, input_data_str) with tempfile.NamedTemporaryFile() as fp_toco, \ tempfile.NamedTemporaryFile() as fp_model, \ @@ -184,7 +198,7 @@ def toco_convert(input_data, if inference_type == QUANTIZED_UINT8: if tflite_input_type == FLOAT: tflite_input_type = QUANTIZED_UINT8 - input_array.mean, input_array.std = quantized_input_stats[idx] + input_array.mean_value, input_array.std_value = quantized_input_stats[idx] input_array.name = _tensor_name(input_tensor) input_array.shape.dims.extend(map(int, input_tensor.get_shape())) diff --git a/tensorflow/contrib/lite/python/lite_test.py b/tensorflow/contrib/lite/python/lite_test.py index da360aeb344ab9c4eb183d84e9b5f60ba715c6e8..b8b4510188bee867b32ffde714b27f41a1df778a 100644 --- a/tensorflow/contrib/lite/python/lite_test.py +++ b/tensorflow/contrib/lite/python/lite_test.py @@ -18,10 +18,14 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib.lite.python import lite +from tensorflow.contrib.lite.python.op_hint import _tensor_name_base as _tensor_name_base from tensorflow.python.client import session from tensorflow.python.framework import dtypes from tensorflow.python.framework import test_util +from tensorflow.python.framework.graph_util_impl import _bfs_for_reachable_nodes +from tensorflow.python.framework.graph_util_impl import _extract_graph_summary from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops from tensorflow.python.platform import test @@ -35,11 +39,133 @@ class LiteTest(test_util.TensorFlowTestCase): # Try running on valid graph result = lite.toco_convert(sess.graph_def, [in_tensor], [out_tensor]) self.assertTrue(result) - # TODO(aselle): remove tests that fail. + # TODO(aselle): remove tests that fail (we must get TOCO to not fatal + # all the time). # Try running on identity graph (known fail) # with self.assertRaisesRegexp(RuntimeError, "!model->operators.empty()"): # result = lite.toco_convert(sess.graph_def, [in_tensor], [in_tensor]) + def testQuantization(self): + in_tensor = array_ops.placeholder(shape=[1, 16, 16, 3], + dtype=dtypes.float32) + out_tensor = array_ops.fake_quant_with_min_max_args(in_tensor + in_tensor, + min=0., max=1.) + sess = session.Session() + result = lite.toco_convert(sess.graph_def, [in_tensor], [out_tensor], + inference_type=lite.QUANTIZED_UINT8, + quantized_input_stats=[(0., 1.)]) + self.assertTrue(result) + + +class LiteTestOpHint(test_util.TensorFlowTestCase): + """Test the hint to stub functionality.""" + + def _getGraphOpTypes(self, graphdef, output_nodes): + """Returns used op types in `graphdef` reachable from `output_nodes`. + + This is used to check that after the stub transformation the expected + nodes are there. Typically use this with self.assertCountEqual(...). + + NOTE: this is not a exact test that the graph is the correct output, but + it balances compact expressibility of test with sanity checking. + + Args: + graphdef: TensorFlow proto graphdef. + output_nodes: A list of output node names that we need to reach. + + Returns: + A set of node types reachable from `output_nodes`. + """ + name_to_input_name, name_to_node, _ = ( + _extract_graph_summary(graphdef)) + # Find all nodes that are needed by the outputs + used_node_names = _bfs_for_reachable_nodes(output_nodes, name_to_input_name) + return set([name_to_node[node_name].op for node_name in used_node_names]) + + def _countIdentities(self, nodes): + """Count the number of "Identity" op types in the list of proto nodes. + + Args: + nodes: NodeDefs of the graph. + + Returns: + The number of nodes with op type "Identity" found. + """ + return len([x for x in nodes if x.op == "Identity"]) + + def testSwishLiteHint(self): + """Makes a custom op swish and makes sure it gets converted as a unit.""" + image = array_ops.constant([1., 2., 3., 4.]) + swish_scale = array_ops.constant(1.0) + + def _swish(input_tensor, scale): + custom = lite.OpHint("cool_activation") + input_tensor, scale = custom.add_inputs(input_tensor, scale) + output = math_ops.sigmoid(input_tensor) * input_tensor * scale + output, = custom.add_outputs(output) + return output + output = array_ops.identity(_swish(image, swish_scale), name="ModelOutput") + + with self.test_session() as sess: + # check if identities have been put into the graph (2 input, 1 output, + # and 1 final output). + self.assertEqual(self._countIdentities(sess.graph_def.node), 4) + + stubbed_graphdef = lite.convert_op_hints_to_stubs(sess) + + self.assertCountEqual( + self._getGraphOpTypes( + stubbed_graphdef, output_nodes=[_tensor_name_base(output)]), + ["cool_activation", "Const", "Identity"]) + + def testScaleAndBiasAndIdentity(self): + """This tests a scaled add which has 3 inputs and 2 outputs.""" + a = array_ops.constant(1.) + x = array_ops.constant([2., 3.]) + b = array_ops.constant([4., 5.]) + + def _scaled_and_bias_and_identity(a, x, b): + custom = lite.OpHint("scale_and_bias_and_identity") + a, x, b = custom.add_inputs(a, x, b) + return custom.add_outputs(a * x + b, x) + output = array_ops.identity(_scaled_and_bias_and_identity(a, x, b), + name="ModelOutput") + + with self.test_session() as sess: + # make sure one identity for each input (3) and output (2) => 3 + 2 = 5 + # +1 for the final output + self.assertEqual(self._countIdentities(sess.graph_def.node), 6) + + stubbed_graphdef = lite.convert_op_hints_to_stubs(sess) + + self.assertCountEqual( + self._getGraphOpTypes( + stubbed_graphdef, output_nodes=[_tensor_name_base(output)]), + ["scale_and_bias_and_identity", "Const", "Identity", "Pack"]) + + def testTwoFunctions(self): + """Tests if two functions are converted correctly.""" + a = array_ops.constant([1.]) + b = array_ops.constant([1.]) + def _double_values(x): + custom = lite.OpHint("add_test") + x = custom.add_inputs(x) + output = math_ops.multiply(x, x) + output, = custom.add_outputs(output) + return output + output = array_ops.identity( + math_ops.add(_double_values(a), _double_values(b)), name="ModelOutput") + + with self.test_session() as sess: + # make sure one identity for each input (2) and output (2) => 2 + 2 + # +1 for the final output + self.assertEqual(self._countIdentities(sess.graph_def.node), 5) + stubbed_graphdef = lite.convert_op_hints_to_stubs(sess) + self.assertCountEqual( + self._getGraphOpTypes( + stubbed_graphdef, output_nodes=[_tensor_name_base(output)]), + ["add_test", "Const", "Identity", "Add"]) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/lite/python/op_hint.py b/tensorflow/contrib/lite/python/op_hint.py new file mode 100644 index 0000000000000000000000000000000000000000..9a3971228a683211e84b4c55d3a3e8d574b5ed94 --- /dev/null +++ b/tensorflow/contrib/lite/python/op_hint.py @@ -0,0 +1,306 @@ +# Copyright 2017 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. +# ============================================================================== +"""Define tflite op hints (intrinsic operations). + +This essentially allows defining a TensorFlow API for tflite operations in +Python with hints on how they are represented in TensorFlow Lite. This basically +is a form of tflite intrinsic. It wraps a subpart of a TensorFlow execution +graph and is useful for LSTMs and other complicated TensorFlow constructions +that are difficult to pattern match in TOCO, but are represented by a single +accelerated tflite op. + +Example: + def tflite_cool_activation(input): + # A cool activation function. + custom = tf.contrib.lite.OpHint("cool_activation") + input = custom.add_inputs(input) + output = tf.sigmoid(input) * input + custom.add_outputs(output) + return output + + image = tf.placeholder(tf.float32, (1, 16, 16, 1)) + output = tf.identity(tflite_cool_activation(image)) + + session = tf.Session() + + graphdef_to_convert = tf.contrib.lite.convert_op_hints_to_stubs(session) + tflite_graph = tf.contrib.lite.toco_convert(graphdef_to_convert, + [image], [output]) + [image], [output]) + with open("/tmp/graph.fb", "wb") as fp: + fp.write(tflite_graph) + +How does it work?: + +OpHint is a helper that you use when defining a vanilla python function. +It allows you to wrap arguments with tf.identities with some custom attributes. +These attributes allow you to find the original block of ops that was created. +For example, if you use cool_activation above you essentially get: + +a_input = tf.identity() +result = tf.multiply(tf.sigmoid(a_input), a_input) +output = tf.identity() + +a_input, output are identities that have parameters representing +what argument they are, what the name of the function they should turn into +in tf lite as well as a guid that uniquely identifies a particular invocation. + +Once you have built your whole tensorflow graph, you can run it and train it +as usual, but after you have done that, you need to convert the graph into +a form that replaces these subgraphs wrapped in identities to stub ops. These +ops don't actually exist in the normal TensorFlow runtime, but will be +understood by toco later. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import collections as _collections +import itertools as _itertools +import uuid as _uuid + +from tensorflow.contrib import framework as _framework +from tensorflow.core.framework import attr_value_pb2 as _attr_value_pb2 +from tensorflow.python.framework import ops as _ops +from tensorflow.python.ops import array_ops as _array_ops +from tensorflow.python.util.all_util import remove_undocumented + + +class OpHint(object): + """A class that helps build tflite function invocations. + + It allows you to take a bunch of TensorFlow ops and annotate the construction + such that toco knows how to convert it to tflite. This embeds a pseudo + function in a TensorFlow graph. This allows embedding high-level API usage + information in a lower level TensorFlow implementation so that an alternative + implementation can be substituted later. + + Essentially, any "input" into this pseudo op is fed into an identity, and + attributes are added to that input before being used by the constituent ops + that make up the pseudo op. A similar process is done to any output that + is to be exported from the current op. + + TODO(aselle): When TensorFlow functions functionality works for arbitrary + constructs, this mechanism can be retired and changed to use python defun's. + """ + + # Attr constants that are used for representation in the GraphDef + FUNCTION_NAME_ATTR = "_tflite_function_name" + FUNCTION_UUID_ATTR = "_tflite_function_uuid" + FUNCTION_INPUT_INDEX_ATTR = "_tflite_function_input_index" + FUNCTION_OUTPUT_INDEX_ATTR = "_tflite_function_output_index" + + def __init__(self, function_name, **kwargs): + """Create a OpHint. + + Args: + function_name: Name of the function (the custom op name in tflite) + **kwargs: Keyword arguments of any constant attributes for the function. + """ + self._function_name = function_name + self._unique_function_id = _uuid.uuid1().hex # TODO(aselle): Unique enough? + self._curr_input_index = 0 + self._curr_output_index = 0 + self._attrs_to_store_later = kwargs + self._stored_attrs = False + + def _setattr(self, dest_op, name, value): + tensor_value = _ops.convert_to_tensor(value) + dest_op.op.node_def.attr[name].tensor.CopyFrom( + tensor_value.op.node_def.attr["value"].tensor) + + def add_inputs(self, *args): + """Add a sequence of inputs to the function invocation. + + Args: + *args: List of inputs to be converted (should be Tf.Tensor). + Returns: + Wrapped inputs (identity standins that have additional metadata). These + are also are also tf.Tensor's. + """ + + def augmented_identity(arg): + identity_op = _array_ops.identity(arg) + # pylint: disable=protected-access + identity_op.op._set_attr( + OpHint.FUNCTION_NAME_ATTR, + _attr_value_pb2.AttrValue(s=self._function_name)) + identity_op.op._set_attr( + OpHint.FUNCTION_UUID_ATTR, + _attr_value_pb2.AttrValue(s=self._unique_function_id)) + identity_op.op._set_attr( + OpHint.FUNCTION_INPUT_INDEX_ATTR, + _attr_value_pb2.AttrValue(i=self._curr_input_index)) + # pylint: enable=protected-access + self._curr_input_index += 1 + return identity_op + + return [augmented_identity(arg) for arg in args] + + def add_outputs(self, *args): + """Add a sequence of outputs to the function invocation. + + Args: + *args: List of outputs to be converted (should be tf.Tensor). + Returns: + Wrapped outputs (identity standins that have additional metadata). These + are also tf.Tensor's. + """ + + def augmented_identity(arg): + identity_op = _array_ops.identity(arg) + # pylint: disable=protected-access + identity_op.op._set_attr( + OpHint.FUNCTION_NAME_ATTR, + _attr_value_pb2.AttrValue(s=self._function_name)) + identity_op.op._set_attr( + OpHint.FUNCTION_UUID_ATTR, + _attr_value_pb2.AttrValue(s=self._unique_function_id)) + identity_op.op._set_attr( + OpHint.FUNCTION_OUTPUT_INDEX_ATTR, + _attr_value_pb2.AttrValue(i=self._curr_output_index)) + # pylint: enable=protected-access + self._curr_output_index += 1 + return identity_op + + wrapped_outputs = [augmented_identity(arg) for arg in args] + + if not self._stored_attrs: + for key, value in self._attrs_to_store_later.iteritems(): + self._setattr(wrapped_outputs[0], "_tflite_attr_" + key, value) + self._stored_attrs = True + + return wrapped_outputs + + +class _LiteFuncCall(object): + """Represent a TensorFlow Lite custom function. + + This is uses to accumulate found hints in the graphdef into a single + conceptual unit. + + Properties: + self.inputs: inputs to the op (hash from index # to argument) + self.outputs: outputs to the op (hash from index # to argument) + self.function_name: the tflite custom op name to use + self.uuid: a unique call id for this particular call (i.e. + multiple function calls would have the same function_name but different + uuids. + self.params: A param name to key value for op constant data. I.e. for + axis on a reduction, strides on a convolution, etc. + """ + + def __init__(self): + self.inputs = {} + self.outputs = {} + self.function_name = None + self.uuid = None + self.params = {} + + def __str__(self): + return "tflite function %s call %s\n\tinputs: %r\n\toutputs: %r" % ( + self.function_name, self.uuid, self.inputs, self.outputs) + + +def _find_all_hints_in_graph_def(session): + """Look at the current default graph and return a list of LiteFuncCall objs. + + Args: + session: A TensorFlow session that contains the graph to convert. + Returns: + a list of `LifeFuncCall` objects in the form + + """ + func_calls = _collections.defaultdict(_LiteFuncCall) + seen_ops = set() + + for op in session.graph.get_operations(): + for operand in _itertools.chain(op.inputs, op.outputs): + if operand in seen_ops: + continue + seen_ops.add(operand) + attr = operand.op.node_def.attr + uuid = attr[OpHint.FUNCTION_UUID_ATTR].s + if OpHint.FUNCTION_UUID_ATTR not in attr: + continue + call_def = func_calls[uuid] + call_def.uuid = uuid + if OpHint.FUNCTION_UUID_ATTR in attr: + call_def.function_name = attr[OpHint.FUNCTION_NAME_ATTR].s + if OpHint.FUNCTION_INPUT_INDEX_ATTR in attr: + call_def.inputs[attr[OpHint.FUNCTION_INPUT_INDEX_ATTR].i] = operand + if OpHint.FUNCTION_OUTPUT_INDEX_ATTR in attr: + call_def.outputs[attr[OpHint.FUNCTION_OUTPUT_INDEX_ATTR].i] = operand + + for a in attr: + if a.startswith("_tflite_attr_"): + # TODO(aselle): Remember the attribute tensors so we can put them + # in collapse. + call_def.params[a.replace("_tflite_attr_,", "")] = attr[a].tensor + + return func_calls + + +def _tensor_name_base(full_tensor_name): + """Removes the device assignment code from a tensor. + + e.g. _tensor_name_base("foo:3") => "foo" + + Args: + full_tensor_name: A tensor name that is annotated with a device placement + (this is what tensor flow introspection gives). + Returns: + A name without any device assignment. + """ + return full_tensor_name.name.split(":")[0] + + +def convert_op_hints_to_stubs(session): + """Converts a graphdef with LiteOp hints into stub operations. + + This is used to prepare for toco conversion of complex intrinsic usages. + + Args: + session: A TensorFlow session that contains the graph to convert. + Returns: + A new graphdef with all ops contained in OpHints being replaced by + a single op call with the right parameters. + """ + hints = _find_all_hints_in_graph_def(session) + current_graph_def = session.graph_def + for call in hints.values(): + input_names = [None] * len(call.inputs) + output_names = [None] * len(call.outputs) + output_dtypes = [None] * len(call.outputs) + output_quantized = False + for input_index, tensor in call.inputs.items(): + input_names[input_index] = _tensor_name_base(tensor) + for output_index, tensor in call.outputs.items(): + output_names[output_index] = _tensor_name_base(tensor) + output_dtypes[output_index] = tensor.dtype.as_datatype_enum + # TODO(aselle): Support quantized flag properly + current_graph_def = _framework.fuse_op( + current_graph_def, input_names, output_names, output_dtypes, + output_quantized, call.uuid, call.function_name) + for node in current_graph_def.node: + if node.name == call.uuid: + for param, tensor in call.params.items(): + node.attr[param].tensor.CopyFrom(tensor) + return current_graph_def + + +_allowed_symbols = ["OpHint", "convert_op_hints_to_stubs"] +remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/BUILD b/tensorflow/contrib/lite/schema/builtin_ops_header/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..0148149a6adc141d67e82808f7e8c72ddb7e309a --- /dev/null +++ b/tensorflow/contrib/lite/schema/builtin_ops_header/BUILD @@ -0,0 +1,43 @@ +package(default_visibility = [ + "//visibility:public", +]) + +licenses(["notice"]) # Apache 2.0 + +cc_library( + name = "generator", + srcs = ["generator.cc"], + hdrs = ["generator.h"], + deps = [ + "//tensorflow/contrib/lite/schema:schema_fbs", + ], +) + +cc_binary( + name = "generate", + srcs = ["generate.cc"], + deps = [ + ":generator", + ], +) + +cc_test( + name = "generator_test", + srcs = ["generator_test.cc"], + deps = [ + ":generator", + "@com_google_googletest//:gtest", + ], +) + +cc_test( + name = "consistency_test", + srcs = ["consistency_test.cc"], + data = [ + "//tensorflow/contrib/lite:builtin_ops.h", + ], + deps = [ + ":generator", + "@com_google_googletest//:gtest", + ], +) diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/README.md b/tensorflow/contrib/lite/schema/builtin_ops_header/README.md new file mode 100644 index 0000000000000000000000000000000000000000..f20d4f664e62fdd52e55339e45b9603307a2b671 --- /dev/null +++ b/tensorflow/contrib/lite/schema/builtin_ops_header/README.md @@ -0,0 +1,12 @@ +# Builtin Ops Header Generator. + +This directory contains a code generator to generate a pure C header for +builtin op definition. + +Whenever you add a new builtin op, please execute: + +```sh +bazel run \ + //tensorflow/contrib/lite/schema/builtin_ops_header:generate > \ + tensorflow/contrib/lite/builtin_ops.h +``` diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/consistency_test.cc b/tensorflow/contrib/lite/schema/builtin_ops_header/consistency_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..d55c125c117db3c1b8d67ab0b674abe2e7c39d94 --- /dev/null +++ b/tensorflow/contrib/lite/schema/builtin_ops_header/consistency_test.cc @@ -0,0 +1,47 @@ +/* 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 +#include +#include "tensorflow/contrib/lite/schema/builtin_ops_header/generator.h" + +namespace { + +const char* kHeaderFileName = + "tensorflow/contrib/lite/builtin_ops.h"; + +// The test ensures that `builtin_ops.h` is consistent with the FlatBuffer +// schema definition. When the schema is modified, it's required to run the +// generator to re-generate the header. +// Please see README.md for more details. +TEST(BuiltinOpsHeaderTest, TestConsistency) { + std::ifstream input_stream(kHeaderFileName, std::ios::binary); + ASSERT_TRUE(input_stream); + std::string file_content((std::istreambuf_iterator(input_stream)), + std::istreambuf_iterator()); + + std::ostringstream output_stream; + tflite::builtin_ops_header::GenerateHeader(output_stream); + std::string generated_content = output_stream.str(); + + EXPECT_EQ(file_content, generated_content); +} + +} // anonymous namespace + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/generate.cc b/tensorflow/contrib/lite/schema/builtin_ops_header/generate.cc new file mode 100644 index 0000000000000000000000000000000000000000..72a28987b8d4863b0f03f7861177940177edd884 --- /dev/null +++ b/tensorflow/contrib/lite/schema/builtin_ops_header/generate.cc @@ -0,0 +1,25 @@ +/* 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 +#include "tensorflow/contrib/lite/schema/builtin_ops_header/generator.h" + +// This executable is used to generate builtin_ops.h in TensorFlow Lite. +// Please see README.md for more details. +int main() { + if (!tflite::builtin_ops_header::GenerateHeader(std::cout)) { + std::cerr << "Failed to generate the header file.\n"; + } + return 0; +} diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/generator.cc b/tensorflow/contrib/lite/schema/builtin_ops_header/generator.cc new file mode 100644 index 0000000000000000000000000000000000000000..b983d59d85955b241a22012b4e9adbeea346f80d --- /dev/null +++ b/tensorflow/contrib/lite/schema/builtin_ops_header/generator.cc @@ -0,0 +1,132 @@ +/* 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/lite/schema/builtin_ops_header/generator.h" +#include "tensorflow/contrib/lite/schema/schema_generated.h" + +namespace tflite { +namespace builtin_ops_header { + +namespace { +const char* kFileHeader = + R"(/* 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_LITE_BUILTIN_OPS_H_ +#define TENSORFLOW_CONTRIB_LITE_BUILTIN_OPS_H_ + +// DO NOT EDIT MANUALLY: This file is automatically generated by +// `schema_builtin_ops_header_generator.py`. + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +typedef enum { +)"; + +const char* kFileFooter = + R"(} TfLiteBuiltinOperator; + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus +#endif // TENSORFLOW_CONTRIB_LITE_BUILTIN_OPS_H_ +} +)"; +} // anonymous namespace + +bool IsValidInputEnumName(const std::string& name) { + const char* begin = name.c_str(); + const char* ch = begin; + while (*ch != '\0') { + // If it's not the first character, expect an underscore. + if (ch != begin) { + if (*ch != '_') { + return false; + } + ++ch; + } + + // Expecting a word with upper case letters or digits, like "CONV", + // "CONV2D", "2D"...etc. + bool empty = true; + while (isupper(*ch) || isdigit(*ch)) { + // It's not empty if at least one character is consumed. + empty = false; + ++ch; + } + if (empty) { + return false; + } + } + return true; +} + +std::string ConstantizeVariableName(const std::string& name) { + std::string result = "kTfLiteBuiltin"; + bool uppercase = true; + for (char input_char : name) { + if (input_char == '_') { + uppercase = true; + } else if (uppercase) { + result += toupper(input_char); + uppercase = false; + } else { + result += tolower(input_char); + } + } + + return result; +} + +bool GenerateHeader(std::ostream& os) { + auto enum_names = tflite::EnumNamesBuiltinOperator(); + + // Check if all the input enum names are valid. + for (auto enum_value : EnumValuesBuiltinOperator()) { + auto enum_name = enum_names[enum_value]; + if (!IsValidInputEnumName(enum_name)) { + std::cerr << "Invalid input enum name: " << enum_name << std::endl; + return false; + } + } + + os << kFileHeader; + for (auto enum_value : EnumValuesBuiltinOperator()) { + auto enum_name = enum_names[enum_value]; + os << " "; + os << ConstantizeVariableName(enum_name); + os << " = "; + os << enum_value; + os << ",\n"; + } + os << kFileFooter; + return true; +} + +} // namespace builtin_ops_header +} // namespace tflite diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/generator.h b/tensorflow/contrib/lite/schema/builtin_ops_header/generator.h new file mode 100644 index 0000000000000000000000000000000000000000..3241ff83d599ed8a476fc1d5a88c26143ebfbaf2 --- /dev/null +++ b/tensorflow/contrib/lite/schema/builtin_ops_header/generator.h @@ -0,0 +1,38 @@ +/* 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. +==============================================================================*/ +// An utility library to generate pure C header for builtin ops definition. +#ifndef TENSORFLOW_CONTRIB_LITE_SCHEMA_BUILTIN_OPS_HEADER_GENERATOR_H_ +#define TENSORFLOW_CONTRIB_LITE_SCHEMA_BUILTIN_OPS_HEADER_GENERATOR_H_ + +#include + +namespace tflite { +namespace builtin_ops_header { + +// Check if the input enum name (from the Flatbuffer definition) is valid. +bool IsValidInputEnumName(const std::string& name); + +// Convert the enum name from Flatbuffer convention to C enum name convention. +// E.g. `L2_POOL_2D` becomes `kTfLiteBuiltinL2Pool2d`. +std::string ConstantizeVariableName(const std::string& name); + +// The function generates a pure C header for builtin ops definition, and write +// it to the output stream. +bool GenerateHeader(std::ostream& os); + +} // namespace builtin_ops_header +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_SCHEMA_BUILTIN_OPS_HEADER_GENERATOR_H_ diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/generator_test.cc b/tensorflow/contrib/lite/schema/builtin_ops_header/generator_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..a7dc8e1b0486eda6e09f38a209dca95c0317a1fb --- /dev/null +++ b/tensorflow/contrib/lite/schema/builtin_ops_header/generator_test.cc @@ -0,0 +1,63 @@ + +/* 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/lite/schema/builtin_ops_header/generator.h" +#include +#include + +namespace { + +using tflite::builtin_ops_header::ConstantizeVariableName; +using tflite::builtin_ops_header::IsValidInputEnumName; + +TEST(TestIsValidInputEnumName, TestWithValidInputNames) { + EXPECT_TRUE(IsValidInputEnumName("ADD")); + EXPECT_TRUE(IsValidInputEnumName("CONV_2D")); + EXPECT_TRUE(IsValidInputEnumName("L2_POOL_2D")); +} + +TEST(TestIsValidInputEnumName, TestWithLeadingUnderscore) { + EXPECT_FALSE(IsValidInputEnumName("_ADD")); + EXPECT_FALSE(IsValidInputEnumName("_CONV_2D")); +} + +TEST(TestIsValidInputEnumName, TestWithLowerCase) { + EXPECT_FALSE(IsValidInputEnumName("_AdD")); + EXPECT_FALSE(IsValidInputEnumName("_COnV_2D")); +} + +TEST(TestIsValidInputEnumName, TestWithOtherCharacters) { + EXPECT_FALSE(IsValidInputEnumName("_AdD!2D")); + EXPECT_FALSE(IsValidInputEnumName("_COnV?2D")); +} + +TEST(TestIsValidInputEnumName, TestWithDoubleUnderscores) { + EXPECT_FALSE(IsValidInputEnumName("ADD__2D")); + EXPECT_FALSE(IsValidInputEnumName("CONV__2D")); +} + +TEST(TestConstantizeVariableName, TestWithValidInputNames) { + EXPECT_EQ(ConstantizeVariableName("ADD"), "kTfLiteBuiltinAdd"); + EXPECT_EQ(ConstantizeVariableName("CONV_2D"), "kTfLiteBuiltinConv2d"); + EXPECT_EQ(ConstantizeVariableName("L2_POOL_2D"), "kTfLiteBuiltinL2Pool2d"); +} + +} // anonymous namespace + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/schema/schema.fbs b/tensorflow/contrib/lite/schema/schema.fbs index cc31e03dfc9c0449cb404ee64d6a9964738d7bfb..75970b41267613058199c22a0fcb0c80a1c8f04f 100644 --- a/tensorflow/contrib/lite/schema/schema.fbs +++ b/tensorflow/contrib/lite/schema/schema.fbs @@ -47,18 +47,18 @@ table QuantizationParameters { table Tensor { // The tensor shape. The meaning of each entry is operator-specific but - // builtin ops use: [batch size, number of channels, height, width] (That's - // Tensorflow's NCHW). + // builtin ops use: [batch size, height, width, number of channels] (That's + // Tensorflow's NHWC). shape:[int]; type:TensorType; // An index that refers to the buffers table at the root of the model. Or, // if there is no data buffer associated (i.e. intermediate results), then - // this is 0 (which refers to an always existant empty buffer). + // this is 0 (which refers to an always existent empty buffer). // // The data_buffer itself is an opaque container, with the assumption that the // target device is little-endian. In addition, all builtin operators assume // the memory is ordered such that if `shape` is [4, 3, 2], then index - // [i, j, k] maps to data_buffer[i*3*2 + j*3 + k]. + // [i, j, k] maps to data_buffer[i*3*2 + j*2 + k]. buffer:uint; name:string; // For debugging and importing back into tensorflow. quantization:QuantizationParameters; // Optional. @@ -89,7 +89,10 @@ enum BuiltinOperator : byte { MAX_POOL_2D = 17, MUL = 18, RELU = 19, - RELU1 = 20, + // NOTE(aselle): RELU_N1_TO_1 used to be called RELU1, but it was renamed + // since different model developers use RELU1 in different ways. Never + // create another op called RELU1. + RELU_N1_TO_1 = 20, RELU6 = 21, RESHAPE = 22, RESIZE_BILINEAR = 23, @@ -108,6 +111,18 @@ enum BuiltinOperator : byte { UNIDIRECTIONAL_SEQUENCE_RNN = 35, GATHER = 36, BATCH_TO_SPACE_ND = 37, + SPACE_TO_BATCH_ND = 38, + TRANSPOSE = 39, + MEAN = 40, + SUB = 41, + DIV = 42, + SQUEEZE = 43, + UNIDIRECTIONAL_SEQUENCE_LSTM = 44, + STRIDED_SLICE = 45, + BIDIRECTIONAL_SEQUENCE_RNN = 46, + EXP = 47, + TOPK_V2 = 48, + SPLIT = 49, } // Options for the builtin operators. @@ -136,6 +151,17 @@ union BuiltinOptions { PadOptions, GatherOptions, BatchToSpaceNDOptions, + SpaceToBatchNDOptions, + TransposeOptions, + MeanOptions, + SubOptions, + DivOptions, + SqueezeOptions, + SequenceRNNOptions, + StridedSliceOptions, + ExpOptions, + TopKV2Options, + SplitOptions, } enum Padding : byte { SAME, VALID } @@ -143,7 +169,7 @@ enum Padding : byte { SAME, VALID } enum ActivationFunctionType : byte { NONE = 0, RELU = 1, - RELU1 = 2, + RELU_N1_TO_1 = 2, RELU6 = 3, TANH = 4, SIGN_BIT = 5, @@ -199,6 +225,18 @@ table RNNOptions { fused_activation_function:ActivationFunctionType; } +// An implementation of TensorFlow dynamic_rnn with RNNCell. +table SequenceRNNOptions { + time_major:bool; + fused_activation_function:ActivationFunctionType; +} + +// An implementation of TensorFlow bidrectional_dynamic_rnn with RNNCell. +table BidirectionalSequenceRNNOptions { + time_major:bool; + fused_activation_function:ActivationFunctionType; +} + // An implementation of TensorFlow fully_connected (a.k.a Dense) layer. table FullyConnectedOptions { fused_activation_function:ActivationFunctionType; @@ -241,8 +279,9 @@ table LSTMOptions { } table ResizeBilinearOptions { - new_height:int; - new_width:int; + new_height: int (deprecated); + new_width: int (deprecated); + align_corners: bool; } // A call operation options @@ -252,18 +291,16 @@ table CallOptions { } table PadOptions { - before_padding:[int]; - after_padding:[int]; } table ReshapeOptions { new_shape:[int]; } +table SpaceToBatchNDOptions { +} + table BatchToSpaceNDOptions { - block_shape:[int]; - before_crops:[int]; - after_crops:[int]; } table SkipGramOptions { @@ -276,6 +313,17 @@ table SpaceToDepthOptions { block_size: int; } +table SubOptions { + fused_activation_function:ActivationFunctionType; +} + +table DivOptions { + fused_activation_function:ActivationFunctionType; +} + +table TopKV2Options { +} + enum CombinerType : byte { SUM = 0, MEAN = 1, @@ -290,6 +338,32 @@ table GatherOptions { axis: int; } +table TransposeOptions { +} + +table ExpOptions { +} + +table MeanOptions { + keep_dims: bool; +} + +table SqueezeOptions { + squeeze_dims:[int]; +} + +table SplitOptions { + num_splits: int; +} + +table StridedSliceOptions { + begin_mask: int; + end_mask: int; + ellipsis_mask: int; + new_axis_mask: int; + shrink_axis_mask: int; +} + // An OperatorCode can be an enum value (BuiltinOperator) if the operator is a // builtin, or a string if the operator is custom. table OperatorCode { diff --git a/tensorflow/contrib/lite/schema/schema_generated.h b/tensorflow/contrib/lite/schema/schema_generated.h index aa169198fe96b538addf091a8c7569e9127d0a01..06989c7b61dc9904bff380e7f1cdc11097cb340d 100755 --- a/tensorflow/contrib/lite/schema/schema_generated.h +++ b/tensorflow/contrib/lite/schema/schema_generated.h @@ -1,5 +1,4 @@ - -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* 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. @@ -13,9 +12,9 @@ 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. ==============================================================================*/ - // automatically generated by the FlatBuffers compiler, do not modify + #ifndef FLATBUFFERS_GENERATED_SCHEMA_TFLITE_H_ #define FLATBUFFERS_GENERATED_SCHEMA_TFLITE_H_ @@ -50,6 +49,12 @@ struct SVDFOptionsT; struct RNNOptions; struct RNNOptionsT; +struct SequenceRNNOptions; +struct SequenceRNNOptionsT; + +struct BidirectionalSequenceRNNOptions; +struct BidirectionalSequenceRNNOptionsT; + struct FullyConnectedOptions; struct FullyConnectedOptionsT; @@ -86,6 +91,9 @@ struct PadOptionsT; struct ReshapeOptions; struct ReshapeOptionsT; +struct SpaceToBatchNDOptions; +struct SpaceToBatchNDOptionsT; + struct BatchToSpaceNDOptions; struct BatchToSpaceNDOptionsT; @@ -95,12 +103,39 @@ struct SkipGramOptionsT; struct SpaceToDepthOptions; struct SpaceToDepthOptionsT; +struct SubOptions; +struct SubOptionsT; + +struct DivOptions; +struct DivOptionsT; + +struct TopKV2Options; +struct TopKV2OptionsT; + struct EmbeddingLookupSparseOptions; struct EmbeddingLookupSparseOptionsT; struct GatherOptions; struct GatherOptionsT; +struct TransposeOptions; +struct TransposeOptionsT; + +struct ExpOptions; +struct ExpOptionsT; + +struct MeanOptions; +struct MeanOptionsT; + +struct SqueezeOptions; +struct SqueezeOptionsT; + +struct SplitOptions; +struct SplitOptionsT; + +struct StridedSliceOptions; +struct StridedSliceOptionsT; + struct OperatorCode; struct OperatorCodeT; @@ -128,15 +163,27 @@ enum TensorType { }; inline TensorType (&EnumValuesTensorType())[6] { - static TensorType values[] = {TensorType_FLOAT32, TensorType_FLOAT16, - TensorType_INT32, TensorType_UINT8, - TensorType_INT64, TensorType_STRING}; + static TensorType values[] = { + TensorType_FLOAT32, + TensorType_FLOAT16, + TensorType_INT32, + TensorType_UINT8, + TensorType_INT64, + TensorType_STRING + }; return values; } inline const char **EnumNamesTensorType() { - static const char *names[] = {"FLOAT32", "FLOAT16", "INT32", "UINT8", - "INT64", "STRING", nullptr}; + static const char *names[] = { + "FLOAT32", + "FLOAT16", + "INT32", + "UINT8", + "INT64", + "STRING", + nullptr + }; return names; } @@ -163,7 +210,7 @@ enum BuiltinOperator { BuiltinOperator_MAX_POOL_2D = 17, BuiltinOperator_MUL = 18, BuiltinOperator_RELU = 19, - BuiltinOperator_RELU1 = 20, + BuiltinOperator_RELU_N1_TO_1 = 20, BuiltinOperator_RELU6 = 21, BuiltinOperator_RESHAPE = 22, BuiltinOperator_RESIZE_BILINEAR = 23, @@ -181,90 +228,129 @@ enum BuiltinOperator { BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN = 35, BuiltinOperator_GATHER = 36, BuiltinOperator_BATCH_TO_SPACE_ND = 37, + BuiltinOperator_SPACE_TO_BATCH_ND = 38, + BuiltinOperator_TRANSPOSE = 39, + BuiltinOperator_MEAN = 40, + BuiltinOperator_SUB = 41, + BuiltinOperator_DIV = 42, + BuiltinOperator_SQUEEZE = 43, + BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM = 44, + BuiltinOperator_STRIDED_SLICE = 45, + BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN = 46, + BuiltinOperator_EXP = 47, + BuiltinOperator_TOPK_V2 = 48, + BuiltinOperator_SPLIT = 49, BuiltinOperator_MIN = BuiltinOperator_ADD, - BuiltinOperator_MAX = BuiltinOperator_BATCH_TO_SPACE_ND + BuiltinOperator_MAX = BuiltinOperator_SPLIT }; -inline BuiltinOperator (&EnumValuesBuiltinOperator())[35] { +inline BuiltinOperator (&EnumValuesBuiltinOperator())[47] { static BuiltinOperator values[] = { - BuiltinOperator_ADD, - BuiltinOperator_AVERAGE_POOL_2D, - BuiltinOperator_CONCATENATION, - BuiltinOperator_CONV_2D, - BuiltinOperator_DEPTHWISE_CONV_2D, - BuiltinOperator_EMBEDDING_LOOKUP, - BuiltinOperator_FULLY_CONNECTED, - BuiltinOperator_HASHTABLE_LOOKUP, - BuiltinOperator_L2_NORMALIZATION, - BuiltinOperator_L2_POOL_2D, - BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION, - BuiltinOperator_LOGISTIC, - BuiltinOperator_LSH_PROJECTION, - BuiltinOperator_LSTM, - BuiltinOperator_MAX_POOL_2D, - BuiltinOperator_MUL, - BuiltinOperator_RELU, - BuiltinOperator_RELU1, - BuiltinOperator_RELU6, - BuiltinOperator_RESHAPE, - BuiltinOperator_RESIZE_BILINEAR, - BuiltinOperator_RNN, - BuiltinOperator_SOFTMAX, - BuiltinOperator_SPACE_TO_DEPTH, - BuiltinOperator_SVDF, - BuiltinOperator_TANH, - BuiltinOperator_CONCAT_EMBEDDINGS, - BuiltinOperator_SKIP_GRAM, - BuiltinOperator_CALL, - BuiltinOperator_CUSTOM, - BuiltinOperator_EMBEDDING_LOOKUP_SPARSE, - BuiltinOperator_PAD, - BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, - BuiltinOperator_GATHER, - BuiltinOperator_BATCH_TO_SPACE_ND}; + BuiltinOperator_ADD, + BuiltinOperator_AVERAGE_POOL_2D, + BuiltinOperator_CONCATENATION, + BuiltinOperator_CONV_2D, + BuiltinOperator_DEPTHWISE_CONV_2D, + BuiltinOperator_EMBEDDING_LOOKUP, + BuiltinOperator_FULLY_CONNECTED, + BuiltinOperator_HASHTABLE_LOOKUP, + BuiltinOperator_L2_NORMALIZATION, + BuiltinOperator_L2_POOL_2D, + BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION, + BuiltinOperator_LOGISTIC, + BuiltinOperator_LSH_PROJECTION, + BuiltinOperator_LSTM, + BuiltinOperator_MAX_POOL_2D, + BuiltinOperator_MUL, + BuiltinOperator_RELU, + BuiltinOperator_RELU_N1_TO_1, + BuiltinOperator_RELU6, + BuiltinOperator_RESHAPE, + BuiltinOperator_RESIZE_BILINEAR, + BuiltinOperator_RNN, + BuiltinOperator_SOFTMAX, + BuiltinOperator_SPACE_TO_DEPTH, + BuiltinOperator_SVDF, + BuiltinOperator_TANH, + BuiltinOperator_CONCAT_EMBEDDINGS, + BuiltinOperator_SKIP_GRAM, + BuiltinOperator_CALL, + BuiltinOperator_CUSTOM, + BuiltinOperator_EMBEDDING_LOOKUP_SPARSE, + BuiltinOperator_PAD, + BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, + BuiltinOperator_GATHER, + BuiltinOperator_BATCH_TO_SPACE_ND, + BuiltinOperator_SPACE_TO_BATCH_ND, + BuiltinOperator_TRANSPOSE, + BuiltinOperator_MEAN, + BuiltinOperator_SUB, + BuiltinOperator_DIV, + BuiltinOperator_SQUEEZE, + BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, + BuiltinOperator_STRIDED_SLICE, + BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN, + BuiltinOperator_EXP, + BuiltinOperator_TOPK_V2, + BuiltinOperator_SPLIT + }; return values; } inline const char **EnumNamesBuiltinOperator() { - static const char *names[] = {"ADD", - "AVERAGE_POOL_2D", - "CONCATENATION", - "CONV_2D", - "DEPTHWISE_CONV_2D", - "", - "", - "EMBEDDING_LOOKUP", - "", - "FULLY_CONNECTED", - "HASHTABLE_LOOKUP", - "L2_NORMALIZATION", - "L2_POOL_2D", - "LOCAL_RESPONSE_NORMALIZATION", - "LOGISTIC", - "LSH_PROJECTION", - "LSTM", - "MAX_POOL_2D", - "MUL", - "RELU", - "RELU1", - "RELU6", - "RESHAPE", - "RESIZE_BILINEAR", - "RNN", - "SOFTMAX", - "SPACE_TO_DEPTH", - "SVDF", - "TANH", - "CONCAT_EMBEDDINGS", - "SKIP_GRAM", - "CALL", - "CUSTOM", - "EMBEDDING_LOOKUP_SPARSE", - "PAD", - "UNIDIRECTIONAL_SEQUENCE_RNN", - "GATHER", - "BATCH_TO_SPACE_ND", - nullptr}; + static const char *names[] = { + "ADD", + "AVERAGE_POOL_2D", + "CONCATENATION", + "CONV_2D", + "DEPTHWISE_CONV_2D", + "", + "", + "EMBEDDING_LOOKUP", + "", + "FULLY_CONNECTED", + "HASHTABLE_LOOKUP", + "L2_NORMALIZATION", + "L2_POOL_2D", + "LOCAL_RESPONSE_NORMALIZATION", + "LOGISTIC", + "LSH_PROJECTION", + "LSTM", + "MAX_POOL_2D", + "MUL", + "RELU", + "RELU_N1_TO_1", + "RELU6", + "RESHAPE", + "RESIZE_BILINEAR", + "RNN", + "SOFTMAX", + "SPACE_TO_DEPTH", + "SVDF", + "TANH", + "CONCAT_EMBEDDINGS", + "SKIP_GRAM", + "CALL", + "CUSTOM", + "EMBEDDING_LOOKUP_SPARSE", + "PAD", + "UNIDIRECTIONAL_SEQUENCE_RNN", + "GATHER", + "BATCH_TO_SPACE_ND", + "SPACE_TO_BATCH_ND", + "TRANSPOSE", + "MEAN", + "SUB", + "DIV", + "SQUEEZE", + "UNIDIRECTIONAL_SEQUENCE_LSTM", + "STRIDED_SLICE", + "BIDIRECTIONAL_SEQUENCE_RNN", + "EXP", + "TOPK_V2", + "SPLIT", + nullptr + }; return names; } @@ -299,67 +385,103 @@ enum BuiltinOptions { BuiltinOptions_PadOptions = 22, BuiltinOptions_GatherOptions = 23, BuiltinOptions_BatchToSpaceNDOptions = 24, + BuiltinOptions_SpaceToBatchNDOptions = 25, + BuiltinOptions_TransposeOptions = 26, + BuiltinOptions_MeanOptions = 27, + BuiltinOptions_SubOptions = 28, + BuiltinOptions_DivOptions = 29, + BuiltinOptions_SqueezeOptions = 30, + BuiltinOptions_SequenceRNNOptions = 31, + BuiltinOptions_StridedSliceOptions = 32, + BuiltinOptions_ExpOptions = 33, + BuiltinOptions_TopKV2Options = 34, + BuiltinOptions_SplitOptions = 35, BuiltinOptions_MIN = BuiltinOptions_NONE, - BuiltinOptions_MAX = BuiltinOptions_BatchToSpaceNDOptions + BuiltinOptions_MAX = BuiltinOptions_SplitOptions }; -inline BuiltinOptions (&EnumValuesBuiltinOptions())[25] { +inline BuiltinOptions (&EnumValuesBuiltinOptions())[36] { static BuiltinOptions values[] = { - BuiltinOptions_NONE, - BuiltinOptions_Conv2DOptions, - BuiltinOptions_DepthwiseConv2DOptions, - BuiltinOptions_ConcatEmbeddingsOptions, - BuiltinOptions_LSHProjectionOptions, - BuiltinOptions_Pool2DOptions, - BuiltinOptions_SVDFOptions, - BuiltinOptions_RNNOptions, - BuiltinOptions_FullyConnectedOptions, - BuiltinOptions_SoftmaxOptions, - BuiltinOptions_ConcatenationOptions, - BuiltinOptions_AddOptions, - BuiltinOptions_L2NormOptions, - BuiltinOptions_LocalResponseNormalizationOptions, - BuiltinOptions_LSTMOptions, - BuiltinOptions_ResizeBilinearOptions, - BuiltinOptions_CallOptions, - BuiltinOptions_ReshapeOptions, - BuiltinOptions_SkipGramOptions, - BuiltinOptions_SpaceToDepthOptions, - BuiltinOptions_EmbeddingLookupSparseOptions, - BuiltinOptions_MulOptions, - BuiltinOptions_PadOptions, - BuiltinOptions_GatherOptions, - BuiltinOptions_BatchToSpaceNDOptions}; + BuiltinOptions_NONE, + BuiltinOptions_Conv2DOptions, + BuiltinOptions_DepthwiseConv2DOptions, + BuiltinOptions_ConcatEmbeddingsOptions, + BuiltinOptions_LSHProjectionOptions, + BuiltinOptions_Pool2DOptions, + BuiltinOptions_SVDFOptions, + BuiltinOptions_RNNOptions, + BuiltinOptions_FullyConnectedOptions, + BuiltinOptions_SoftmaxOptions, + BuiltinOptions_ConcatenationOptions, + BuiltinOptions_AddOptions, + BuiltinOptions_L2NormOptions, + BuiltinOptions_LocalResponseNormalizationOptions, + BuiltinOptions_LSTMOptions, + BuiltinOptions_ResizeBilinearOptions, + BuiltinOptions_CallOptions, + BuiltinOptions_ReshapeOptions, + BuiltinOptions_SkipGramOptions, + BuiltinOptions_SpaceToDepthOptions, + BuiltinOptions_EmbeddingLookupSparseOptions, + BuiltinOptions_MulOptions, + BuiltinOptions_PadOptions, + BuiltinOptions_GatherOptions, + BuiltinOptions_BatchToSpaceNDOptions, + BuiltinOptions_SpaceToBatchNDOptions, + BuiltinOptions_TransposeOptions, + BuiltinOptions_MeanOptions, + BuiltinOptions_SubOptions, + BuiltinOptions_DivOptions, + BuiltinOptions_SqueezeOptions, + BuiltinOptions_SequenceRNNOptions, + BuiltinOptions_StridedSliceOptions, + BuiltinOptions_ExpOptions, + BuiltinOptions_TopKV2Options, + BuiltinOptions_SplitOptions + }; return values; } inline const char **EnumNamesBuiltinOptions() { - static const char *names[] = {"NONE", - "Conv2DOptions", - "DepthwiseConv2DOptions", - "ConcatEmbeddingsOptions", - "LSHProjectionOptions", - "Pool2DOptions", - "SVDFOptions", - "RNNOptions", - "FullyConnectedOptions", - "SoftmaxOptions", - "ConcatenationOptions", - "AddOptions", - "L2NormOptions", - "LocalResponseNormalizationOptions", - "LSTMOptions", - "ResizeBilinearOptions", - "CallOptions", - "ReshapeOptions", - "SkipGramOptions", - "SpaceToDepthOptions", - "EmbeddingLookupSparseOptions", - "MulOptions", - "PadOptions", - "GatherOptions", - "BatchToSpaceNDOptions", - nullptr}; + static const char *names[] = { + "NONE", + "Conv2DOptions", + "DepthwiseConv2DOptions", + "ConcatEmbeddingsOptions", + "LSHProjectionOptions", + "Pool2DOptions", + "SVDFOptions", + "RNNOptions", + "FullyConnectedOptions", + "SoftmaxOptions", + "ConcatenationOptions", + "AddOptions", + "L2NormOptions", + "LocalResponseNormalizationOptions", + "LSTMOptions", + "ResizeBilinearOptions", + "CallOptions", + "ReshapeOptions", + "SkipGramOptions", + "SpaceToDepthOptions", + "EmbeddingLookupSparseOptions", + "MulOptions", + "PadOptions", + "GatherOptions", + "BatchToSpaceNDOptions", + "SpaceToBatchNDOptions", + "TransposeOptions", + "MeanOptions", + "SubOptions", + "DivOptions", + "SqueezeOptions", + "SequenceRNNOptions", + "StridedSliceOptions", + "ExpOptions", + "TopKV2Options", + "SplitOptions", + nullptr + }; return names; } @@ -368,166 +490,170 @@ inline const char *EnumNameBuiltinOptions(BuiltinOptions e) { return EnumNamesBuiltinOptions()[index]; } -template -struct BuiltinOptionsTraits { +template struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_NONE; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_Conv2DOptions; }; -template <> -struct BuiltinOptionsTraits { - static const BuiltinOptions enum_value = - BuiltinOptions_DepthwiseConv2DOptions; +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_DepthwiseConv2DOptions; }; -template <> -struct BuiltinOptionsTraits { - static const BuiltinOptions enum_value = - BuiltinOptions_ConcatEmbeddingsOptions; +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_ConcatEmbeddingsOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_LSHProjectionOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_Pool2DOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_SVDFOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_RNNOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_FullyConnectedOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_SoftmaxOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_ConcatenationOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_AddOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_L2NormOptions; }; -template <> -struct BuiltinOptionsTraits { - static const BuiltinOptions enum_value = - BuiltinOptions_LocalResponseNormalizationOptions; +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_LocalResponseNormalizationOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_LSTMOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_ResizeBilinearOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_CallOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_ReshapeOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_SkipGramOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_SpaceToDepthOptions; }; -template <> -struct BuiltinOptionsTraits { - static const BuiltinOptions enum_value = - BuiltinOptions_EmbeddingLookupSparseOptions; +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_EmbeddingLookupSparseOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_MulOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_PadOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_GatherOptions; }; -template <> -struct BuiltinOptionsTraits { +template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_BatchToSpaceNDOptions; }; +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_SpaceToBatchNDOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_TransposeOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_MeanOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_SubOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_DivOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_SqueezeOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_SequenceRNNOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_StridedSliceOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_ExpOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_TopKV2Options; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_SplitOptions; +}; + struct BuiltinOptionsUnion { BuiltinOptions type; void *value; BuiltinOptionsUnion() : type(BuiltinOptions_NONE), value(nullptr) {} - BuiltinOptionsUnion(BuiltinOptionsUnion &&u) FLATBUFFERS_NOEXCEPT - : type(BuiltinOptions_NONE), - value(nullptr) { - std::swap(type, u.type); - std::swap(value, u.value); - } + BuiltinOptionsUnion(BuiltinOptionsUnion&& u) FLATBUFFERS_NOEXCEPT : + type(BuiltinOptions_NONE), value(nullptr) + { std::swap(type, u.type); std::swap(value, u.value); } BuiltinOptionsUnion(const BuiltinOptionsUnion &) FLATBUFFERS_NOEXCEPT; - BuiltinOptionsUnion &operator=(const BuiltinOptionsUnion &u) - FLATBUFFERS_NOEXCEPT { - BuiltinOptionsUnion t(u); - std::swap(type, t.type); - std::swap(value, t.value); - return *this; - } - BuiltinOptionsUnion &operator=(BuiltinOptionsUnion &&u) FLATBUFFERS_NOEXCEPT { - std::swap(type, u.type); - std::swap(value, u.value); - return *this; - } + BuiltinOptionsUnion &operator=(const BuiltinOptionsUnion &u) FLATBUFFERS_NOEXCEPT + { BuiltinOptionsUnion t(u); std::swap(type, t.type); std::swap(value, t.value); return *this; } + BuiltinOptionsUnion &operator=(BuiltinOptionsUnion &&u) FLATBUFFERS_NOEXCEPT + { std::swap(type, u.type); std::swap(value, u.value); return *this; } ~BuiltinOptionsUnion() { Reset(); } void Reset(); #ifndef FLATBUFFERS_CPP98_STL template - void Set(T &&val) { + void Set(T&& val) { Reset(); type = BuiltinOptionsTraits::enum_value; if (type != BuiltinOptions_NONE) { @@ -536,262 +662,293 @@ struct BuiltinOptionsUnion { } #endif // FLATBUFFERS_CPP98_STL - static void *UnPack(const void *obj, BuiltinOptions type, - const flatbuffers::resolver_function_t *resolver); - flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const flatbuffers::rehasher_function_t *_rehasher = nullptr) const; + static void *UnPack(const void *obj, BuiltinOptions type, const flatbuffers::resolver_function_t *resolver); + flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const flatbuffers::rehasher_function_t *_rehasher = nullptr) const; Conv2DOptionsT *AsConv2DOptions() { - return type == BuiltinOptions_Conv2DOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_Conv2DOptions ? + reinterpret_cast(value) : nullptr; } const Conv2DOptionsT *AsConv2DOptions() const { - return type == BuiltinOptions_Conv2DOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_Conv2DOptions ? + reinterpret_cast(value) : nullptr; } DepthwiseConv2DOptionsT *AsDepthwiseConv2DOptions() { - return type == BuiltinOptions_DepthwiseConv2DOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_DepthwiseConv2DOptions ? + reinterpret_cast(value) : nullptr; } const DepthwiseConv2DOptionsT *AsDepthwiseConv2DOptions() const { - return type == BuiltinOptions_DepthwiseConv2DOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_DepthwiseConv2DOptions ? + reinterpret_cast(value) : nullptr; } ConcatEmbeddingsOptionsT *AsConcatEmbeddingsOptions() { - return type == BuiltinOptions_ConcatEmbeddingsOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_ConcatEmbeddingsOptions ? + reinterpret_cast(value) : nullptr; } const ConcatEmbeddingsOptionsT *AsConcatEmbeddingsOptions() const { - return type == BuiltinOptions_ConcatEmbeddingsOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_ConcatEmbeddingsOptions ? + reinterpret_cast(value) : nullptr; } LSHProjectionOptionsT *AsLSHProjectionOptions() { - return type == BuiltinOptions_LSHProjectionOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_LSHProjectionOptions ? + reinterpret_cast(value) : nullptr; } const LSHProjectionOptionsT *AsLSHProjectionOptions() const { - return type == BuiltinOptions_LSHProjectionOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_LSHProjectionOptions ? + reinterpret_cast(value) : nullptr; } Pool2DOptionsT *AsPool2DOptions() { - return type == BuiltinOptions_Pool2DOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_Pool2DOptions ? + reinterpret_cast(value) : nullptr; } const Pool2DOptionsT *AsPool2DOptions() const { - return type == BuiltinOptions_Pool2DOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_Pool2DOptions ? + reinterpret_cast(value) : nullptr; } SVDFOptionsT *AsSVDFOptions() { - return type == BuiltinOptions_SVDFOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SVDFOptions ? + reinterpret_cast(value) : nullptr; } const SVDFOptionsT *AsSVDFOptions() const { - return type == BuiltinOptions_SVDFOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SVDFOptions ? + reinterpret_cast(value) : nullptr; } RNNOptionsT *AsRNNOptions() { - return type == BuiltinOptions_RNNOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_RNNOptions ? + reinterpret_cast(value) : nullptr; } const RNNOptionsT *AsRNNOptions() const { - return type == BuiltinOptions_RNNOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_RNNOptions ? + reinterpret_cast(value) : nullptr; } FullyConnectedOptionsT *AsFullyConnectedOptions() { - return type == BuiltinOptions_FullyConnectedOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_FullyConnectedOptions ? + reinterpret_cast(value) : nullptr; } const FullyConnectedOptionsT *AsFullyConnectedOptions() const { - return type == BuiltinOptions_FullyConnectedOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_FullyConnectedOptions ? + reinterpret_cast(value) : nullptr; } SoftmaxOptionsT *AsSoftmaxOptions() { - return type == BuiltinOptions_SoftmaxOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SoftmaxOptions ? + reinterpret_cast(value) : nullptr; } const SoftmaxOptionsT *AsSoftmaxOptions() const { - return type == BuiltinOptions_SoftmaxOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SoftmaxOptions ? + reinterpret_cast(value) : nullptr; } ConcatenationOptionsT *AsConcatenationOptions() { - return type == BuiltinOptions_ConcatenationOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_ConcatenationOptions ? + reinterpret_cast(value) : nullptr; } const ConcatenationOptionsT *AsConcatenationOptions() const { - return type == BuiltinOptions_ConcatenationOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_ConcatenationOptions ? + reinterpret_cast(value) : nullptr; } AddOptionsT *AsAddOptions() { - return type == BuiltinOptions_AddOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_AddOptions ? + reinterpret_cast(value) : nullptr; } const AddOptionsT *AsAddOptions() const { - return type == BuiltinOptions_AddOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_AddOptions ? + reinterpret_cast(value) : nullptr; } L2NormOptionsT *AsL2NormOptions() { - return type == BuiltinOptions_L2NormOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_L2NormOptions ? + reinterpret_cast(value) : nullptr; } const L2NormOptionsT *AsL2NormOptions() const { - return type == BuiltinOptions_L2NormOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_L2NormOptions ? + reinterpret_cast(value) : nullptr; } LocalResponseNormalizationOptionsT *AsLocalResponseNormalizationOptions() { - return type == BuiltinOptions_LocalResponseNormalizationOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_LocalResponseNormalizationOptions ? + reinterpret_cast(value) : nullptr; } - const LocalResponseNormalizationOptionsT * - AsLocalResponseNormalizationOptions() const { - return type == BuiltinOptions_LocalResponseNormalizationOptions - ? reinterpret_cast( - value) - : nullptr; + const LocalResponseNormalizationOptionsT *AsLocalResponseNormalizationOptions() const { + return type == BuiltinOptions_LocalResponseNormalizationOptions ? + reinterpret_cast(value) : nullptr; } LSTMOptionsT *AsLSTMOptions() { - return type == BuiltinOptions_LSTMOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_LSTMOptions ? + reinterpret_cast(value) : nullptr; } const LSTMOptionsT *AsLSTMOptions() const { - return type == BuiltinOptions_LSTMOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_LSTMOptions ? + reinterpret_cast(value) : nullptr; } ResizeBilinearOptionsT *AsResizeBilinearOptions() { - return type == BuiltinOptions_ResizeBilinearOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_ResizeBilinearOptions ? + reinterpret_cast(value) : nullptr; } const ResizeBilinearOptionsT *AsResizeBilinearOptions() const { - return type == BuiltinOptions_ResizeBilinearOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_ResizeBilinearOptions ? + reinterpret_cast(value) : nullptr; } CallOptionsT *AsCallOptions() { - return type == BuiltinOptions_CallOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_CallOptions ? + reinterpret_cast(value) : nullptr; } const CallOptionsT *AsCallOptions() const { - return type == BuiltinOptions_CallOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_CallOptions ? + reinterpret_cast(value) : nullptr; } ReshapeOptionsT *AsReshapeOptions() { - return type == BuiltinOptions_ReshapeOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_ReshapeOptions ? + reinterpret_cast(value) : nullptr; } const ReshapeOptionsT *AsReshapeOptions() const { - return type == BuiltinOptions_ReshapeOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_ReshapeOptions ? + reinterpret_cast(value) : nullptr; } SkipGramOptionsT *AsSkipGramOptions() { - return type == BuiltinOptions_SkipGramOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SkipGramOptions ? + reinterpret_cast(value) : nullptr; } const SkipGramOptionsT *AsSkipGramOptions() const { - return type == BuiltinOptions_SkipGramOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SkipGramOptions ? + reinterpret_cast(value) : nullptr; } SpaceToDepthOptionsT *AsSpaceToDepthOptions() { - return type == BuiltinOptions_SpaceToDepthOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SpaceToDepthOptions ? + reinterpret_cast(value) : nullptr; } const SpaceToDepthOptionsT *AsSpaceToDepthOptions() const { - return type == BuiltinOptions_SpaceToDepthOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_SpaceToDepthOptions ? + reinterpret_cast(value) : nullptr; } EmbeddingLookupSparseOptionsT *AsEmbeddingLookupSparseOptions() { - return type == BuiltinOptions_EmbeddingLookupSparseOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_EmbeddingLookupSparseOptions ? + reinterpret_cast(value) : nullptr; } const EmbeddingLookupSparseOptionsT *AsEmbeddingLookupSparseOptions() const { - return type == BuiltinOptions_EmbeddingLookupSparseOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_EmbeddingLookupSparseOptions ? + reinterpret_cast(value) : nullptr; } MulOptionsT *AsMulOptions() { - return type == BuiltinOptions_MulOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_MulOptions ? + reinterpret_cast(value) : nullptr; } const MulOptionsT *AsMulOptions() const { - return type == BuiltinOptions_MulOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_MulOptions ? + reinterpret_cast(value) : nullptr; } PadOptionsT *AsPadOptions() { - return type == BuiltinOptions_PadOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_PadOptions ? + reinterpret_cast(value) : nullptr; } const PadOptionsT *AsPadOptions() const { - return type == BuiltinOptions_PadOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_PadOptions ? + reinterpret_cast(value) : nullptr; } GatherOptionsT *AsGatherOptions() { - return type == BuiltinOptions_GatherOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_GatherOptions ? + reinterpret_cast(value) : nullptr; } const GatherOptionsT *AsGatherOptions() const { - return type == BuiltinOptions_GatherOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_GatherOptions ? + reinterpret_cast(value) : nullptr; } BatchToSpaceNDOptionsT *AsBatchToSpaceNDOptions() { - return type == BuiltinOptions_BatchToSpaceNDOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_BatchToSpaceNDOptions ? + reinterpret_cast(value) : nullptr; } const BatchToSpaceNDOptionsT *AsBatchToSpaceNDOptions() const { - return type == BuiltinOptions_BatchToSpaceNDOptions - ? reinterpret_cast(value) - : nullptr; + return type == BuiltinOptions_BatchToSpaceNDOptions ? + reinterpret_cast(value) : nullptr; + } + SpaceToBatchNDOptionsT *AsSpaceToBatchNDOptions() { + return type == BuiltinOptions_SpaceToBatchNDOptions ? + reinterpret_cast(value) : nullptr; + } + const SpaceToBatchNDOptionsT *AsSpaceToBatchNDOptions() const { + return type == BuiltinOptions_SpaceToBatchNDOptions ? + reinterpret_cast(value) : nullptr; + } + TransposeOptionsT *AsTransposeOptions() { + return type == BuiltinOptions_TransposeOptions ? + reinterpret_cast(value) : nullptr; + } + const TransposeOptionsT *AsTransposeOptions() const { + return type == BuiltinOptions_TransposeOptions ? + reinterpret_cast(value) : nullptr; + } + MeanOptionsT *AsMeanOptions() { + return type == BuiltinOptions_MeanOptions ? + reinterpret_cast(value) : nullptr; + } + const MeanOptionsT *AsMeanOptions() const { + return type == BuiltinOptions_MeanOptions ? + reinterpret_cast(value) : nullptr; + } + SubOptionsT *AsSubOptions() { + return type == BuiltinOptions_SubOptions ? + reinterpret_cast(value) : nullptr; + } + const SubOptionsT *AsSubOptions() const { + return type == BuiltinOptions_SubOptions ? + reinterpret_cast(value) : nullptr; + } + DivOptionsT *AsDivOptions() { + return type == BuiltinOptions_DivOptions ? + reinterpret_cast(value) : nullptr; + } + const DivOptionsT *AsDivOptions() const { + return type == BuiltinOptions_DivOptions ? + reinterpret_cast(value) : nullptr; + } + SqueezeOptionsT *AsSqueezeOptions() { + return type == BuiltinOptions_SqueezeOptions ? + reinterpret_cast(value) : nullptr; + } + const SqueezeOptionsT *AsSqueezeOptions() const { + return type == BuiltinOptions_SqueezeOptions ? + reinterpret_cast(value) : nullptr; + } + SequenceRNNOptionsT *AsSequenceRNNOptions() { + return type == BuiltinOptions_SequenceRNNOptions ? + reinterpret_cast(value) : nullptr; + } + const SequenceRNNOptionsT *AsSequenceRNNOptions() const { + return type == BuiltinOptions_SequenceRNNOptions ? + reinterpret_cast(value) : nullptr; + } + StridedSliceOptionsT *AsStridedSliceOptions() { + return type == BuiltinOptions_StridedSliceOptions ? + reinterpret_cast(value) : nullptr; + } + const StridedSliceOptionsT *AsStridedSliceOptions() const { + return type == BuiltinOptions_StridedSliceOptions ? + reinterpret_cast(value) : nullptr; + } + ExpOptionsT *AsExpOptions() { + return type == BuiltinOptions_ExpOptions ? + reinterpret_cast(value) : nullptr; + } + const ExpOptionsT *AsExpOptions() const { + return type == BuiltinOptions_ExpOptions ? + reinterpret_cast(value) : nullptr; + } + TopKV2OptionsT *AsTopKV2Options() { + return type == BuiltinOptions_TopKV2Options ? + reinterpret_cast(value) : nullptr; + } + const TopKV2OptionsT *AsTopKV2Options() const { + return type == BuiltinOptions_TopKV2Options ? + reinterpret_cast(value) : nullptr; + } + SplitOptionsT *AsSplitOptions() { + return type == BuiltinOptions_SplitOptions ? + reinterpret_cast(value) : nullptr; + } + const SplitOptionsT *AsSplitOptions() const { + return type == BuiltinOptions_SplitOptions ? + reinterpret_cast(value) : nullptr; } }; -bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *obj, - BuiltinOptions type); -bool VerifyBuiltinOptionsVector( - flatbuffers::Verifier &verifier, - const flatbuffers::Vector> *values, - const flatbuffers::Vector *types); +bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *obj, BuiltinOptions type); +bool VerifyBuiltinOptionsVector(flatbuffers::Verifier &verifier, const flatbuffers::Vector> *values, const flatbuffers::Vector *types); enum Padding { Padding_SAME = 0, @@ -801,12 +958,19 @@ enum Padding { }; inline Padding (&EnumValuesPadding())[2] { - static Padding values[] = {Padding_SAME, Padding_VALID}; + static Padding values[] = { + Padding_SAME, + Padding_VALID + }; return values; } inline const char **EnumNamesPadding() { - static const char *names[] = {"SAME", "VALID", nullptr}; + static const char *names[] = { + "SAME", + "VALID", + nullptr + }; return names; } @@ -818,7 +982,7 @@ inline const char *EnumNamePadding(Padding e) { enum ActivationFunctionType { ActivationFunctionType_NONE = 0, ActivationFunctionType_RELU = 1, - ActivationFunctionType_RELU1 = 2, + ActivationFunctionType_RELU_N1_TO_1 = 2, ActivationFunctionType_RELU6 = 3, ActivationFunctionType_TANH = 4, ActivationFunctionType_SIGN_BIT = 5, @@ -828,15 +992,26 @@ enum ActivationFunctionType { inline ActivationFunctionType (&EnumValuesActivationFunctionType())[6] { static ActivationFunctionType values[] = { - ActivationFunctionType_NONE, ActivationFunctionType_RELU, - ActivationFunctionType_RELU1, ActivationFunctionType_RELU6, - ActivationFunctionType_TANH, ActivationFunctionType_SIGN_BIT}; + ActivationFunctionType_NONE, + ActivationFunctionType_RELU, + ActivationFunctionType_RELU_N1_TO_1, + ActivationFunctionType_RELU6, + ActivationFunctionType_TANH, + ActivationFunctionType_SIGN_BIT + }; return values; } inline const char **EnumNamesActivationFunctionType() { - static const char *names[] = {"NONE", "RELU", "RELU1", "RELU6", - "TANH", "SIGN_BIT", nullptr}; + static const char *names[] = { + "NONE", + "RELU", + "RELU_N1_TO_1", + "RELU6", + "TANH", + "SIGN_BIT", + nullptr + }; return names; } @@ -854,14 +1029,21 @@ enum LSHProjectionType { }; inline LSHProjectionType (&EnumValuesLSHProjectionType())[3] { - static LSHProjectionType values[] = {LSHProjectionType_UNKNOWN, - LSHProjectionType_SPARSE, - LSHProjectionType_DENSE}; + static LSHProjectionType values[] = { + LSHProjectionType_UNKNOWN, + LSHProjectionType_SPARSE, + LSHProjectionType_DENSE + }; return values; } inline const char **EnumNamesLSHProjectionType() { - static const char *names[] = {"UNKNOWN", "SPARSE", "DENSE", nullptr}; + static const char *names[] = { + "UNKNOWN", + "SPARSE", + "DENSE", + nullptr + }; return names; } @@ -879,13 +1061,21 @@ enum CombinerType { }; inline CombinerType (&EnumValuesCombinerType())[3] { - static CombinerType values[] = {CombinerType_SUM, CombinerType_MEAN, - CombinerType_SQRTN}; + static CombinerType values[] = { + CombinerType_SUM, + CombinerType_MEAN, + CombinerType_SQRTN + }; return values; } inline const char **EnumNamesCombinerType() { - static const char *names[] = {"SUM", "MEAN", "SQRTN", nullptr}; + static const char *names[] = { + "SUM", + "MEAN", + "SQRTN", + nullptr + }; return names; } @@ -901,12 +1091,17 @@ enum CustomOptionsFormat { }; inline CustomOptionsFormat (&EnumValuesCustomOptionsFormat())[1] { - static CustomOptionsFormat values[] = {CustomOptionsFormat_FLEXBUFFERS}; + static CustomOptionsFormat values[] = { + CustomOptionsFormat_FLEXBUFFERS + }; return values; } inline const char **EnumNamesCustomOptionsFormat() { - static const char *names[] = {"FLEXBUFFERS", nullptr}; + static const char *names[] = { + "FLEXBUFFERS", + nullptr + }; return names; } @@ -921,13 +1116,18 @@ struct QuantizationParametersT : public flatbuffers::NativeTable { std::vector max; std::vector scale; std::vector zero_point; - QuantizationParametersT() {} + QuantizationParametersT() { + } }; -struct QuantizationParameters FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct QuantizationParameters FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef QuantizationParametersT NativeTableType; - enum { VT_MIN = 4, VT_MAX = 6, VT_SCALE = 8, VT_ZERO_POINT = 10 }; + enum { + VT_MIN = 4, + VT_MAX = 6, + VT_SCALE = 8, + VT_ZERO_POINT = 10 + }; const flatbuffers::Vector *min() const { return GetPointer *>(VT_MIN); } @@ -941,20 +1141,20 @@ struct QuantizationParameters FLATBUFFERS_FINAL_CLASS return GetPointer *>(VT_ZERO_POINT); } bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && VerifyOffset(verifier, VT_MIN) && - verifier.Verify(min()) && VerifyOffset(verifier, VT_MAX) && - verifier.Verify(max()) && VerifyOffset(verifier, VT_SCALE) && - verifier.Verify(scale()) && VerifyOffset(verifier, VT_ZERO_POINT) && - verifier.Verify(zero_point()) && verifier.EndTable(); - } - QuantizationParametersT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - QuantizationParametersT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const QuantizationParametersT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + return VerifyTableStart(verifier) && + VerifyOffset(verifier, VT_MIN) && + verifier.Verify(min()) && + VerifyOffset(verifier, VT_MAX) && + verifier.Verify(max()) && + VerifyOffset(verifier, VT_SCALE) && + verifier.Verify(scale()) && + VerifyOffset(verifier, VT_ZERO_POINT) && + verifier.Verify(zero_point()) && + verifier.EndTable(); + } + QuantizationParametersT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(QuantizationParametersT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const QuantizationParametersT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct QuantizationParametersBuilder { @@ -969,16 +1169,14 @@ struct QuantizationParametersBuilder { void add_scale(flatbuffers::Offset> scale) { fbb_.AddOffset(QuantizationParameters::VT_SCALE, scale); } - void add_zero_point( - flatbuffers::Offset> zero_point) { + void add_zero_point(flatbuffers::Offset> zero_point) { fbb_.AddOffset(QuantizationParameters::VT_ZERO_POINT, zero_point); } explicit QuantizationParametersBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } - QuantizationParametersBuilder &operator=( - const QuantizationParametersBuilder &); + QuantizationParametersBuilder &operator=(const QuantizationParametersBuilder &); flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); auto o = flatbuffers::Offset(end); @@ -1000,23 +1198,21 @@ inline flatbuffers::Offset CreateQuantizationParameters( return builder_.Finish(); } -inline flatbuffers::Offset -CreateQuantizationParametersDirect( +inline flatbuffers::Offset CreateQuantizationParametersDirect( flatbuffers::FlatBufferBuilder &_fbb, const std::vector *min = nullptr, const std::vector *max = nullptr, const std::vector *scale = nullptr, const std::vector *zero_point = nullptr) { return tflite::CreateQuantizationParameters( - _fbb, min ? _fbb.CreateVector(*min) : 0, + _fbb, + min ? _fbb.CreateVector(*min) : 0, max ? _fbb.CreateVector(*max) : 0, scale ? _fbb.CreateVector(*scale) : 0, zero_point ? _fbb.CreateVector(*zero_point) : 0); } -flatbuffers::Offset CreateQuantizationParameters( - flatbuffers::FlatBufferBuilder &_fbb, const QuantizationParametersT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateQuantizationParameters(flatbuffers::FlatBufferBuilder &_fbb, const QuantizationParametersT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct TensorT : public flatbuffers::NativeTable { typedef Tensor TableType; @@ -1025,7 +1221,10 @@ struct TensorT : public flatbuffers::NativeTable { uint32_t buffer; std::string name; std::unique_ptr quantization; - TensorT() : type(TensorType_FLOAT32), buffer(0) {} + TensorT() + : type(TensorType_FLOAT32), + buffer(0) { + } }; struct Tensor FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { @@ -1043,7 +1242,9 @@ struct Tensor FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { TensorType type() const { return static_cast(GetField(VT_TYPE, 0)); } - uint32_t buffer() const { return GetField(VT_BUFFER, 0); } + uint32_t buffer() const { + return GetField(VT_BUFFER, 0); + } const flatbuffers::String *name() const { return GetPointer(VT_NAME); } @@ -1051,20 +1252,20 @@ struct Tensor FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { return GetPointer(VT_QUANTIZATION); } bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && VerifyOffset(verifier, VT_SHAPE) && - verifier.Verify(shape()) && VerifyField(verifier, VT_TYPE) && + return VerifyTableStart(verifier) && + VerifyOffset(verifier, VT_SHAPE) && + verifier.Verify(shape()) && + VerifyField(verifier, VT_TYPE) && VerifyField(verifier, VT_BUFFER) && - VerifyOffset(verifier, VT_NAME) && verifier.Verify(name()) && + VerifyOffset(verifier, VT_NAME) && + verifier.Verify(name()) && VerifyOffset(verifier, VT_QUANTIZATION) && - verifier.VerifyTable(quantization()) && verifier.EndTable(); + verifier.VerifyTable(quantization()) && + verifier.EndTable(); } - TensorT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo(TensorT *_o, const flatbuffers::resolver_function_t *_resolver = - nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const TensorT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + TensorT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(TensorT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const TensorT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct TensorBuilder { @@ -1082,11 +1283,11 @@ struct TensorBuilder { void add_name(flatbuffers::Offset name) { fbb_.AddOffset(Tensor::VT_NAME, name); } - void add_quantization( - flatbuffers::Offset quantization) { + void add_quantization(flatbuffers::Offset quantization) { fbb_.AddOffset(Tensor::VT_QUANTIZATION, quantization); } - explicit TensorBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { + explicit TensorBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } TensorBuilder &operator=(const TensorBuilder &); @@ -1100,7 +1301,8 @@ struct TensorBuilder { inline flatbuffers::Offset CreateTensor( flatbuffers::FlatBufferBuilder &_fbb, flatbuffers::Offset> shape = 0, - TensorType type = TensorType_FLOAT32, uint32_t buffer = 0, + TensorType type = TensorType_FLOAT32, + uint32_t buffer = 0, flatbuffers::Offset name = 0, flatbuffers::Offset quantization = 0) { TensorBuilder builder_(_fbb); @@ -1115,17 +1317,20 @@ inline flatbuffers::Offset CreateTensor( inline flatbuffers::Offset CreateTensorDirect( flatbuffers::FlatBufferBuilder &_fbb, const std::vector *shape = nullptr, - TensorType type = TensorType_FLOAT32, uint32_t buffer = 0, + TensorType type = TensorType_FLOAT32, + uint32_t buffer = 0, const char *name = nullptr, flatbuffers::Offset quantization = 0) { return tflite::CreateTensor( - _fbb, shape ? _fbb.CreateVector(*shape) : 0, type, buffer, - name ? _fbb.CreateString(name) : 0, quantization); + _fbb, + shape ? _fbb.CreateVector(*shape) : 0, + type, + buffer, + name ? _fbb.CreateString(name) : 0, + quantization); } -flatbuffers::Offset CreateTensor( - flatbuffers::FlatBufferBuilder &_fbb, const TensorT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateTensor(flatbuffers::FlatBufferBuilder &_fbb, const TensorT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct Conv2DOptionsT : public flatbuffers::NativeTable { typedef Conv2DOptions TableType; @@ -1137,7 +1342,8 @@ struct Conv2DOptionsT : public flatbuffers::NativeTable { : padding(Padding_SAME), stride_w(0), stride_h(0), - fused_activation_function(ActivationFunctionType_NONE) {} + fused_activation_function(ActivationFunctionType_NONE) { + } }; struct Conv2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { @@ -1151,11 +1357,14 @@ struct Conv2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { Padding padding() const { return static_cast(GetField(VT_PADDING, 0)); } - int32_t stride_w() const { return GetField(VT_STRIDE_W, 0); } - int32_t stride_h() const { return GetField(VT_STRIDE_H, 0); } + int32_t stride_w() const { + return GetField(VT_STRIDE_W, 0); + } + int32_t stride_h() const { + return GetField(VT_STRIDE_H, 0); + } ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -1165,22 +1374,16 @@ struct Conv2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - Conv2DOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - Conv2DOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + Conv2DOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(Conv2DOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct Conv2DOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; void add_padding(Padding padding) { - fbb_.AddElement(Conv2DOptions::VT_PADDING, - static_cast(padding), 0); + fbb_.AddElement(Conv2DOptions::VT_PADDING, static_cast(padding), 0); } void add_stride_w(int32_t stride_w) { fbb_.AddElement(Conv2DOptions::VT_STRIDE_W, stride_w, 0); @@ -1188,13 +1391,11 @@ struct Conv2DOptionsBuilder { void add_stride_h(int32_t stride_h) { fbb_.AddElement(Conv2DOptions::VT_STRIDE_H, stride_h, 0); } - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(Conv2DOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(Conv2DOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit Conv2DOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } Conv2DOptionsBuilder &operator=(const Conv2DOptionsBuilder &); @@ -1206,10 +1407,11 @@ struct Conv2DOptionsBuilder { }; inline flatbuffers::Offset CreateConv2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, Padding padding = Padding_SAME, - int32_t stride_w = 0, int32_t stride_h = 0, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + flatbuffers::FlatBufferBuilder &_fbb, + Padding padding = Padding_SAME, + int32_t stride_w = 0, + int32_t stride_h = 0, + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { Conv2DOptionsBuilder builder_(_fbb); builder_.add_stride_h(stride_h); builder_.add_stride_w(stride_w); @@ -1218,9 +1420,7 @@ inline flatbuffers::Offset CreateConv2DOptions( return builder_.Finish(); } -flatbuffers::Offset CreateConv2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateConv2DOptions(flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct Pool2DOptionsT : public flatbuffers::NativeTable { typedef Pool2DOptions TableType; @@ -1236,7 +1436,8 @@ struct Pool2DOptionsT : public flatbuffers::NativeTable { stride_h(0), filter_width(0), filter_height(0), - fused_activation_function(ActivationFunctionType_NONE) {} + fused_activation_function(ActivationFunctionType_NONE) { + } }; struct Pool2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { @@ -1252,15 +1453,20 @@ struct Pool2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { Padding padding() const { return static_cast(GetField(VT_PADDING, 0)); } - int32_t stride_w() const { return GetField(VT_STRIDE_W, 0); } - int32_t stride_h() const { return GetField(VT_STRIDE_H, 0); } - int32_t filter_width() const { return GetField(VT_FILTER_WIDTH, 0); } + int32_t stride_w() const { + return GetField(VT_STRIDE_W, 0); + } + int32_t stride_h() const { + return GetField(VT_STRIDE_H, 0); + } + int32_t filter_width() const { + return GetField(VT_FILTER_WIDTH, 0); + } int32_t filter_height() const { return GetField(VT_FILTER_HEIGHT, 0); } ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -1272,22 +1478,16 @@ struct Pool2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - Pool2DOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - Pool2DOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const Pool2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + Pool2DOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(Pool2DOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const Pool2DOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct Pool2DOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; void add_padding(Padding padding) { - fbb_.AddElement(Pool2DOptions::VT_PADDING, - static_cast(padding), 0); + fbb_.AddElement(Pool2DOptions::VT_PADDING, static_cast(padding), 0); } void add_stride_w(int32_t stride_w) { fbb_.AddElement(Pool2DOptions::VT_STRIDE_W, stride_w, 0); @@ -1301,13 +1501,11 @@ struct Pool2DOptionsBuilder { void add_filter_height(int32_t filter_height) { fbb_.AddElement(Pool2DOptions::VT_FILTER_HEIGHT, filter_height, 0); } - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(Pool2DOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(Pool2DOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit Pool2DOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } Pool2DOptionsBuilder &operator=(const Pool2DOptionsBuilder &); @@ -1319,11 +1517,13 @@ struct Pool2DOptionsBuilder { }; inline flatbuffers::Offset CreatePool2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, Padding padding = Padding_SAME, - int32_t stride_w = 0, int32_t stride_h = 0, int32_t filter_width = 0, + flatbuffers::FlatBufferBuilder &_fbb, + Padding padding = Padding_SAME, + int32_t stride_w = 0, + int32_t stride_h = 0, + int32_t filter_width = 0, int32_t filter_height = 0, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { Pool2DOptionsBuilder builder_(_fbb); builder_.add_filter_height(filter_height); builder_.add_filter_width(filter_width); @@ -1334,9 +1534,7 @@ inline flatbuffers::Offset CreatePool2DOptions( return builder_.Finish(); } -flatbuffers::Offset CreatePool2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, const Pool2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreatePool2DOptions(flatbuffers::FlatBufferBuilder &_fbb, const Pool2DOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct DepthwiseConv2DOptionsT : public flatbuffers::NativeTable { typedef DepthwiseConv2DOptions TableType; @@ -1350,11 +1548,11 @@ struct DepthwiseConv2DOptionsT : public flatbuffers::NativeTable { stride_w(0), stride_h(0), depth_multiplier(0), - fused_activation_function(ActivationFunctionType_NONE) {} + fused_activation_function(ActivationFunctionType_NONE) { + } }; -struct DepthwiseConv2DOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct DepthwiseConv2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef DepthwiseConv2DOptionsT NativeTableType; enum { VT_PADDING = 4, @@ -1366,14 +1564,17 @@ struct DepthwiseConv2DOptions FLATBUFFERS_FINAL_CLASS Padding padding() const { return static_cast(GetField(VT_PADDING, 0)); } - int32_t stride_w() const { return GetField(VT_STRIDE_W, 0); } - int32_t stride_h() const { return GetField(VT_STRIDE_H, 0); } + int32_t stride_w() const { + return GetField(VT_STRIDE_W, 0); + } + int32_t stride_h() const { + return GetField(VT_STRIDE_H, 0); + } int32_t depth_multiplier() const { return GetField(VT_DEPTH_MULTIPLIER, 0); } ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -1384,22 +1585,16 @@ struct DepthwiseConv2DOptions FLATBUFFERS_FINAL_CLASS VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - DepthwiseConv2DOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - DepthwiseConv2DOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const DepthwiseConv2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + DepthwiseConv2DOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(DepthwiseConv2DOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const DepthwiseConv2DOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct DepthwiseConv2DOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; void add_padding(Padding padding) { - fbb_.AddElement(DepthwiseConv2DOptions::VT_PADDING, - static_cast(padding), 0); + fbb_.AddElement(DepthwiseConv2DOptions::VT_PADDING, static_cast(padding), 0); } void add_stride_w(int32_t stride_w) { fbb_.AddElement(DepthwiseConv2DOptions::VT_STRIDE_W, stride_w, 0); @@ -1408,21 +1603,16 @@ struct DepthwiseConv2DOptionsBuilder { fbb_.AddElement(DepthwiseConv2DOptions::VT_STRIDE_H, stride_h, 0); } void add_depth_multiplier(int32_t depth_multiplier) { - fbb_.AddElement(DepthwiseConv2DOptions::VT_DEPTH_MULTIPLIER, - depth_multiplier, 0); + fbb_.AddElement(DepthwiseConv2DOptions::VT_DEPTH_MULTIPLIER, depth_multiplier, 0); } - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement( - DepthwiseConv2DOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(DepthwiseConv2DOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit DepthwiseConv2DOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } - DepthwiseConv2DOptionsBuilder &operator=( - const DepthwiseConv2DOptionsBuilder &); + DepthwiseConv2DOptionsBuilder &operator=(const DepthwiseConv2DOptionsBuilder &); flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); auto o = flatbuffers::Offset(end); @@ -1431,10 +1621,12 @@ struct DepthwiseConv2DOptionsBuilder { }; inline flatbuffers::Offset CreateDepthwiseConv2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, Padding padding = Padding_SAME, - int32_t stride_w = 0, int32_t stride_h = 0, int32_t depth_multiplier = 0, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + flatbuffers::FlatBufferBuilder &_fbb, + Padding padding = Padding_SAME, + int32_t stride_w = 0, + int32_t stride_h = 0, + int32_t depth_multiplier = 0, + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { DepthwiseConv2DOptionsBuilder builder_(_fbb); builder_.add_depth_multiplier(depth_multiplier); builder_.add_stride_h(stride_h); @@ -1444,34 +1636,33 @@ inline flatbuffers::Offset CreateDepthwiseConv2DOptions( return builder_.Finish(); } -flatbuffers::Offset CreateDepthwiseConv2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, const DepthwiseConv2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateDepthwiseConv2DOptions(flatbuffers::FlatBufferBuilder &_fbb, const DepthwiseConv2DOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct ConcatEmbeddingsOptionsT : public flatbuffers::NativeTable { typedef ConcatEmbeddingsOptions TableType; int32_t num_channels; std::vector num_columns_per_channel; std::vector embedding_dim_per_channel; - ConcatEmbeddingsOptionsT() : num_channels(0) {} + ConcatEmbeddingsOptionsT() + : num_channels(0) { + } }; -struct ConcatEmbeddingsOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct ConcatEmbeddingsOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef ConcatEmbeddingsOptionsT NativeTableType; enum { VT_NUM_CHANNELS = 4, VT_NUM_COLUMNS_PER_CHANNEL = 6, VT_EMBEDDING_DIM_PER_CHANNEL = 8 }; - int32_t num_channels() const { return GetField(VT_NUM_CHANNELS, 0); } + int32_t num_channels() const { + return GetField(VT_NUM_CHANNELS, 0); + } const flatbuffers::Vector *num_columns_per_channel() const { - return GetPointer *>( - VT_NUM_COLUMNS_PER_CHANNEL); + return GetPointer *>(VT_NUM_COLUMNS_PER_CHANNEL); } const flatbuffers::Vector *embedding_dim_per_channel() const { - return GetPointer *>( - VT_EMBEDDING_DIM_PER_CHANNEL); + return GetPointer *>(VT_EMBEDDING_DIM_PER_CHANNEL); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -1479,43 +1670,31 @@ struct ConcatEmbeddingsOptions FLATBUFFERS_FINAL_CLASS VerifyOffset(verifier, VT_NUM_COLUMNS_PER_CHANNEL) && verifier.Verify(num_columns_per_channel()) && VerifyOffset(verifier, VT_EMBEDDING_DIM_PER_CHANNEL) && - verifier.Verify(embedding_dim_per_channel()) && verifier.EndTable(); + verifier.Verify(embedding_dim_per_channel()) && + verifier.EndTable(); } - ConcatEmbeddingsOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - ConcatEmbeddingsOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatEmbeddingsOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + ConcatEmbeddingsOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(ConcatEmbeddingsOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const ConcatEmbeddingsOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct ConcatEmbeddingsOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; void add_num_channels(int32_t num_channels) { - fbb_.AddElement(ConcatEmbeddingsOptions::VT_NUM_CHANNELS, - num_channels, 0); + fbb_.AddElement(ConcatEmbeddingsOptions::VT_NUM_CHANNELS, num_channels, 0); } - void add_num_columns_per_channel( - flatbuffers::Offset> - num_columns_per_channel) { - fbb_.AddOffset(ConcatEmbeddingsOptions::VT_NUM_COLUMNS_PER_CHANNEL, - num_columns_per_channel); + void add_num_columns_per_channel(flatbuffers::Offset> num_columns_per_channel) { + fbb_.AddOffset(ConcatEmbeddingsOptions::VT_NUM_COLUMNS_PER_CHANNEL, num_columns_per_channel); } - void add_embedding_dim_per_channel( - flatbuffers::Offset> - embedding_dim_per_channel) { - fbb_.AddOffset(ConcatEmbeddingsOptions::VT_EMBEDDING_DIM_PER_CHANNEL, - embedding_dim_per_channel); + void add_embedding_dim_per_channel(flatbuffers::Offset> embedding_dim_per_channel) { + fbb_.AddOffset(ConcatEmbeddingsOptions::VT_EMBEDDING_DIM_PER_CHANNEL, embedding_dim_per_channel); } explicit ConcatEmbeddingsOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } - ConcatEmbeddingsOptionsBuilder &operator=( - const ConcatEmbeddingsOptionsBuilder &); + ConcatEmbeddingsOptionsBuilder &operator=(const ConcatEmbeddingsOptionsBuilder &); flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); auto o = flatbuffers::Offset(end); @@ -1523,13 +1702,11 @@ struct ConcatEmbeddingsOptionsBuilder { } }; -inline flatbuffers::Offset -CreateConcatEmbeddingsOptions(flatbuffers::FlatBufferBuilder &_fbb, - int32_t num_channels = 0, - flatbuffers::Offset> - num_columns_per_channel = 0, - flatbuffers::Offset> - embedding_dim_per_channel = 0) { +inline flatbuffers::Offset CreateConcatEmbeddingsOptions( + flatbuffers::FlatBufferBuilder &_fbb, + int32_t num_channels = 0, + flatbuffers::Offset> num_columns_per_channel = 0, + flatbuffers::Offset> embedding_dim_per_channel = 0) { ConcatEmbeddingsOptionsBuilder builder_(_fbb); builder_.add_embedding_dim_per_channel(embedding_dim_per_channel); builder_.add_num_columns_per_channel(num_columns_per_channel); @@ -1537,61 +1714,54 @@ CreateConcatEmbeddingsOptions(flatbuffers::FlatBufferBuilder &_fbb, return builder_.Finish(); } -inline flatbuffers::Offset -CreateConcatEmbeddingsOptionsDirect( - flatbuffers::FlatBufferBuilder &_fbb, int32_t num_channels = 0, +inline flatbuffers::Offset CreateConcatEmbeddingsOptionsDirect( + flatbuffers::FlatBufferBuilder &_fbb, + int32_t num_channels = 0, const std::vector *num_columns_per_channel = nullptr, const std::vector *embedding_dim_per_channel = nullptr) { return tflite::CreateConcatEmbeddingsOptions( - _fbb, num_channels, - num_columns_per_channel - ? _fbb.CreateVector(*num_columns_per_channel) - : 0, - embedding_dim_per_channel - ? _fbb.CreateVector(*embedding_dim_per_channel) - : 0); + _fbb, + num_channels, + num_columns_per_channel ? _fbb.CreateVector(*num_columns_per_channel) : 0, + embedding_dim_per_channel ? _fbb.CreateVector(*embedding_dim_per_channel) : 0); } -flatbuffers::Offset CreateConcatEmbeddingsOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatEmbeddingsOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateConcatEmbeddingsOptions(flatbuffers::FlatBufferBuilder &_fbb, const ConcatEmbeddingsOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct LSHProjectionOptionsT : public flatbuffers::NativeTable { typedef LSHProjectionOptions TableType; LSHProjectionType type; - LSHProjectionOptionsT() : type(LSHProjectionType_UNKNOWN) {} + LSHProjectionOptionsT() + : type(LSHProjectionType_UNKNOWN) { + } }; -struct LSHProjectionOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct LSHProjectionOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef LSHProjectionOptionsT NativeTableType; - enum { VT_TYPE = 4 }; + enum { + VT_TYPE = 4 + }; LSHProjectionType type() const { return static_cast(GetField(VT_TYPE, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_TYPE) && verifier.EndTable(); + VerifyField(verifier, VT_TYPE) && + verifier.EndTable(); } - LSHProjectionOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - LSHProjectionOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + LSHProjectionOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(LSHProjectionOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct LSHProjectionOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; void add_type(LSHProjectionType type) { - fbb_.AddElement(LSHProjectionOptions::VT_TYPE, - static_cast(type), 0); + fbb_.AddElement(LSHProjectionOptions::VT_TYPE, static_cast(type), 0); } explicit LSHProjectionOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } LSHProjectionOptionsBuilder &operator=(const LSHProjectionOptionsBuilder &); @@ -1610,25 +1780,29 @@ inline flatbuffers::Offset CreateLSHProjectionOptions( return builder_.Finish(); } -flatbuffers::Offset CreateLSHProjectionOptions( - flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateLSHProjectionOptions(flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct SVDFOptionsT : public flatbuffers::NativeTable { typedef SVDFOptions TableType; int32_t rank; ActivationFunctionType fused_activation_function; SVDFOptionsT() - : rank(0), fused_activation_function(ActivationFunctionType_NONE) {} + : rank(0), + fused_activation_function(ActivationFunctionType_NONE) { + } }; struct SVDFOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef SVDFOptionsT NativeTableType; - enum { VT_RANK = 4, VT_FUSED_ACTIVATION_FUNCTION = 6 }; - int32_t rank() const { return GetField(VT_RANK, 0); } + enum { + VT_RANK = 4, + VT_FUSED_ACTIVATION_FUNCTION = 6 + }; + int32_t rank() const { + return GetField(VT_RANK, 0); + } ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -1636,14 +1810,9 @@ struct SVDFOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - SVDFOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - SVDFOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SVDFOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + SVDFOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SVDFOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SVDFOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct SVDFOptionsBuilder { @@ -1652,13 +1821,11 @@ struct SVDFOptionsBuilder { void add_rank(int32_t rank) { fbb_.AddElement(SVDFOptions::VT_RANK, rank, 0); } - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(SVDFOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(SVDFOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit SVDFOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SVDFOptionsBuilder &operator=(const SVDFOptionsBuilder &); @@ -1670,57 +1837,51 @@ struct SVDFOptionsBuilder { }; inline flatbuffers::Offset CreateSVDFOptions( - flatbuffers::FlatBufferBuilder &_fbb, int32_t rank = 0, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + flatbuffers::FlatBufferBuilder &_fbb, + int32_t rank = 0, + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { SVDFOptionsBuilder builder_(_fbb); builder_.add_rank(rank); builder_.add_fused_activation_function(fused_activation_function); return builder_.Finish(); } -flatbuffers::Offset CreateSVDFOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SVDFOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateSVDFOptions(flatbuffers::FlatBufferBuilder &_fbb, const SVDFOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct RNNOptionsT : public flatbuffers::NativeTable { typedef RNNOptions TableType; ActivationFunctionType fused_activation_function; - RNNOptionsT() : fused_activation_function(ActivationFunctionType_NONE) {} + RNNOptionsT() + : fused_activation_function(ActivationFunctionType_NONE) { + } }; struct RNNOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef RNNOptionsT NativeTableType; - enum { VT_FUSED_ACTIVATION_FUNCTION = 4 }; + enum { + VT_FUSED_ACTIVATION_FUNCTION = 4 + }; ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - RNNOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - RNNOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const RNNOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + RNNOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(RNNOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const RNNOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct RNNOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(RNNOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(RNNOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit RNNOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } RNNOptionsBuilder &operator=(const RNNOptionsBuilder &); @@ -1733,57 +1894,180 @@ struct RNNOptionsBuilder { inline flatbuffers::Offset CreateRNNOptions( flatbuffers::FlatBufferBuilder &_fbb, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { RNNOptionsBuilder builder_(_fbb); builder_.add_fused_activation_function(fused_activation_function); return builder_.Finish(); } -flatbuffers::Offset CreateRNNOptions( - flatbuffers::FlatBufferBuilder &_fbb, const RNNOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateRNNOptions(flatbuffers::FlatBufferBuilder &_fbb, const RNNOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct SequenceRNNOptionsT : public flatbuffers::NativeTable { + typedef SequenceRNNOptions TableType; + bool time_major; + ActivationFunctionType fused_activation_function; + SequenceRNNOptionsT() + : time_major(false), + fused_activation_function(ActivationFunctionType_NONE) { + } +}; + +struct SequenceRNNOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef SequenceRNNOptionsT NativeTableType; + enum { + VT_TIME_MAJOR = 4, + VT_FUSED_ACTIVATION_FUNCTION = 6 + }; + bool time_major() const { + return GetField(VT_TIME_MAJOR, 0) != 0; + } + ActivationFunctionType fused_activation_function() const { + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_TIME_MAJOR) && + VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && + verifier.EndTable(); + } + SequenceRNNOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SequenceRNNOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SequenceRNNOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct SequenceRNNOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_time_major(bool time_major) { + fbb_.AddElement(SequenceRNNOptions::VT_TIME_MAJOR, static_cast(time_major), 0); + } + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(SequenceRNNOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); + } + explicit SequenceRNNOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + SequenceRNNOptionsBuilder &operator=(const SequenceRNNOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateSequenceRNNOptions( + flatbuffers::FlatBufferBuilder &_fbb, + bool time_major = false, + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { + SequenceRNNOptionsBuilder builder_(_fbb); + builder_.add_fused_activation_function(fused_activation_function); + builder_.add_time_major(time_major); + return builder_.Finish(); +} + +flatbuffers::Offset CreateSequenceRNNOptions(flatbuffers::FlatBufferBuilder &_fbb, const SequenceRNNOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct BidirectionalSequenceRNNOptionsT : public flatbuffers::NativeTable { + typedef BidirectionalSequenceRNNOptions TableType; + bool time_major; + ActivationFunctionType fused_activation_function; + BidirectionalSequenceRNNOptionsT() + : time_major(false), + fused_activation_function(ActivationFunctionType_NONE) { + } +}; + +struct BidirectionalSequenceRNNOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef BidirectionalSequenceRNNOptionsT NativeTableType; + enum { + VT_TIME_MAJOR = 4, + VT_FUSED_ACTIVATION_FUNCTION = 6 + }; + bool time_major() const { + return GetField(VT_TIME_MAJOR, 0) != 0; + } + ActivationFunctionType fused_activation_function() const { + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_TIME_MAJOR) && + VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && + verifier.EndTable(); + } + BidirectionalSequenceRNNOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(BidirectionalSequenceRNNOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const BidirectionalSequenceRNNOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct BidirectionalSequenceRNNOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_time_major(bool time_major) { + fbb_.AddElement(BidirectionalSequenceRNNOptions::VT_TIME_MAJOR, static_cast(time_major), 0); + } + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(BidirectionalSequenceRNNOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); + } + explicit BidirectionalSequenceRNNOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + BidirectionalSequenceRNNOptionsBuilder &operator=(const BidirectionalSequenceRNNOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateBidirectionalSequenceRNNOptions( + flatbuffers::FlatBufferBuilder &_fbb, + bool time_major = false, + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { + BidirectionalSequenceRNNOptionsBuilder builder_(_fbb); + builder_.add_fused_activation_function(fused_activation_function); + builder_.add_time_major(time_major); + return builder_.Finish(); +} + +flatbuffers::Offset CreateBidirectionalSequenceRNNOptions(flatbuffers::FlatBufferBuilder &_fbb, const BidirectionalSequenceRNNOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct FullyConnectedOptionsT : public flatbuffers::NativeTable { typedef FullyConnectedOptions TableType; ActivationFunctionType fused_activation_function; FullyConnectedOptionsT() - : fused_activation_function(ActivationFunctionType_NONE) {} + : fused_activation_function(ActivationFunctionType_NONE) { + } }; -struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef FullyConnectedOptionsT NativeTableType; - enum { VT_FUSED_ACTIVATION_FUNCTION = 4 }; + enum { + VT_FUSED_ACTIVATION_FUNCTION = 4 + }; ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - FullyConnectedOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - FullyConnectedOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + FullyConnectedOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(FullyConnectedOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct FullyConnectedOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(FullyConnectedOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(FullyConnectedOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit FullyConnectedOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } FullyConnectedOptionsBuilder &operator=(const FullyConnectedOptionsBuilder &); @@ -1796,39 +2080,38 @@ struct FullyConnectedOptionsBuilder { inline flatbuffers::Offset CreateFullyConnectedOptions( flatbuffers::FlatBufferBuilder &_fbb, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { FullyConnectedOptionsBuilder builder_(_fbb); builder_.add_fused_activation_function(fused_activation_function); return builder_.Finish(); } -flatbuffers::Offset CreateFullyConnectedOptions( - flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateFullyConnectedOptions(flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct SoftmaxOptionsT : public flatbuffers::NativeTable { typedef SoftmaxOptions TableType; float beta; - SoftmaxOptionsT() : beta(0.0f) {} + SoftmaxOptionsT() + : beta(0.0f) { + } }; struct SoftmaxOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef SoftmaxOptionsT NativeTableType; - enum { VT_BETA = 4 }; - float beta() const { return GetField(VT_BETA, 0.0f); } + enum { + VT_BETA = 4 + }; + float beta() const { + return GetField(VT_BETA, 0.0f); + } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_BETA) && verifier.EndTable(); + VerifyField(verifier, VT_BETA) && + verifier.EndTable(); } - SoftmaxOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - SoftmaxOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + SoftmaxOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SoftmaxOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct SoftmaxOptionsBuilder { @@ -1838,7 +2121,7 @@ struct SoftmaxOptionsBuilder { fbb_.AddElement(SoftmaxOptions::VT_BETA, beta, 0.0f); } explicit SoftmaxOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SoftmaxOptionsBuilder &operator=(const SoftmaxOptionsBuilder &); @@ -1850,32 +2133,36 @@ struct SoftmaxOptionsBuilder { }; inline flatbuffers::Offset CreateSoftmaxOptions( - flatbuffers::FlatBufferBuilder &_fbb, float beta = 0.0f) { + flatbuffers::FlatBufferBuilder &_fbb, + float beta = 0.0f) { SoftmaxOptionsBuilder builder_(_fbb); builder_.add_beta(beta); return builder_.Finish(); } -flatbuffers::Offset CreateSoftmaxOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateSoftmaxOptions(flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct ConcatenationOptionsT : public flatbuffers::NativeTable { typedef ConcatenationOptions TableType; int32_t axis; ActivationFunctionType fused_activation_function; ConcatenationOptionsT() - : axis(0), fused_activation_function(ActivationFunctionType_NONE) {} + : axis(0), + fused_activation_function(ActivationFunctionType_NONE) { + } }; -struct ConcatenationOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct ConcatenationOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef ConcatenationOptionsT NativeTableType; - enum { VT_AXIS = 4, VT_FUSED_ACTIVATION_FUNCTION = 6 }; - int32_t axis() const { return GetField(VT_AXIS, 0); } + enum { + VT_AXIS = 4, + VT_FUSED_ACTIVATION_FUNCTION = 6 + }; + int32_t axis() const { + return GetField(VT_AXIS, 0); + } ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -1883,14 +2170,9 @@ struct ConcatenationOptions FLATBUFFERS_FINAL_CLASS VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - ConcatenationOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - ConcatenationOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + ConcatenationOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(ConcatenationOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct ConcatenationOptionsBuilder { @@ -1899,13 +2181,11 @@ struct ConcatenationOptionsBuilder { void add_axis(int32_t axis) { fbb_.AddElement(ConcatenationOptions::VT_AXIS, axis, 0); } - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(ConcatenationOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(ConcatenationOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit ConcatenationOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } ConcatenationOptionsBuilder &operator=(const ConcatenationOptionsBuilder &); @@ -1917,57 +2197,51 @@ struct ConcatenationOptionsBuilder { }; inline flatbuffers::Offset CreateConcatenationOptions( - flatbuffers::FlatBufferBuilder &_fbb, int32_t axis = 0, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + flatbuffers::FlatBufferBuilder &_fbb, + int32_t axis = 0, + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { ConcatenationOptionsBuilder builder_(_fbb); builder_.add_axis(axis); builder_.add_fused_activation_function(fused_activation_function); return builder_.Finish(); } -flatbuffers::Offset CreateConcatenationOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateConcatenationOptions(flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct AddOptionsT : public flatbuffers::NativeTable { typedef AddOptions TableType; ActivationFunctionType fused_activation_function; - AddOptionsT() : fused_activation_function(ActivationFunctionType_NONE) {} + AddOptionsT() + : fused_activation_function(ActivationFunctionType_NONE) { + } }; struct AddOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef AddOptionsT NativeTableType; - enum { VT_FUSED_ACTIVATION_FUNCTION = 4 }; + enum { + VT_FUSED_ACTIVATION_FUNCTION = 4 + }; ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - AddOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - AddOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + AddOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(AddOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct AddOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(AddOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(AddOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit AddOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } AddOptionsBuilder &operator=(const AddOptionsBuilder &); @@ -1980,55 +2254,48 @@ struct AddOptionsBuilder { inline flatbuffers::Offset CreateAddOptions( flatbuffers::FlatBufferBuilder &_fbb, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { AddOptionsBuilder builder_(_fbb); builder_.add_fused_activation_function(fused_activation_function); return builder_.Finish(); } -flatbuffers::Offset CreateAddOptions( - flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateAddOptions(flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct MulOptionsT : public flatbuffers::NativeTable { typedef MulOptions TableType; ActivationFunctionType fused_activation_function; - MulOptionsT() : fused_activation_function(ActivationFunctionType_NONE) {} + MulOptionsT() + : fused_activation_function(ActivationFunctionType_NONE) { + } }; struct MulOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef MulOptionsT NativeTableType; - enum { VT_FUSED_ACTIVATION_FUNCTION = 4 }; + enum { + VT_FUSED_ACTIVATION_FUNCTION = 4 + }; ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - MulOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - MulOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + MulOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(MulOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct MulOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(MulOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(MulOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit MulOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } MulOptionsBuilder &operator=(const MulOptionsBuilder &); @@ -2041,55 +2308,48 @@ struct MulOptionsBuilder { inline flatbuffers::Offset CreateMulOptions( flatbuffers::FlatBufferBuilder &_fbb, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { MulOptionsBuilder builder_(_fbb); builder_.add_fused_activation_function(fused_activation_function); return builder_.Finish(); } -flatbuffers::Offset CreateMulOptions( - flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateMulOptions(flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct L2NormOptionsT : public flatbuffers::NativeTable { typedef L2NormOptions TableType; ActivationFunctionType fused_activation_function; - L2NormOptionsT() : fused_activation_function(ActivationFunctionType_NONE) {} + L2NormOptionsT() + : fused_activation_function(ActivationFunctionType_NONE) { + } }; struct L2NormOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef L2NormOptionsT NativeTableType; - enum { VT_FUSED_ACTIVATION_FUNCTION = 4 }; + enum { + VT_FUSED_ACTIVATION_FUNCTION = 4 + }; ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && verifier.EndTable(); } - L2NormOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - L2NormOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + L2NormOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(L2NormOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct L2NormOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(L2NormOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(L2NormOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } explicit L2NormOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } L2NormOptionsBuilder &operator=(const L2NormOptionsBuilder &); @@ -2102,16 +2362,13 @@ struct L2NormOptionsBuilder { inline flatbuffers::Offset CreateL2NormOptions( flatbuffers::FlatBufferBuilder &_fbb, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE) { + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { L2NormOptionsBuilder builder_(_fbb); builder_.add_fused_activation_function(fused_activation_function); return builder_.Finish(); } -flatbuffers::Offset CreateL2NormOptions( - flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateL2NormOptions(flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct LocalResponseNormalizationOptionsT : public flatbuffers::NativeTable { typedef LocalResponseNormalizationOptions TableType; @@ -2120,61 +2377,66 @@ struct LocalResponseNormalizationOptionsT : public flatbuffers::NativeTable { float alpha; float beta; LocalResponseNormalizationOptionsT() - : radius(0), bias(0.0f), alpha(0.0f), beta(0.0f) {} + : radius(0), + bias(0.0f), + alpha(0.0f), + beta(0.0f) { + } }; -struct LocalResponseNormalizationOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct LocalResponseNormalizationOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef LocalResponseNormalizationOptionsT NativeTableType; - enum { VT_RADIUS = 4, VT_BIAS = 6, VT_ALPHA = 8, VT_BETA = 10 }; - int32_t radius() const { return GetField(VT_RADIUS, 0); } - float bias() const { return GetField(VT_BIAS, 0.0f); } - float alpha() const { return GetField(VT_ALPHA, 0.0f); } - float beta() const { return GetField(VT_BETA, 0.0f); } + enum { + VT_RADIUS = 4, + VT_BIAS = 6, + VT_ALPHA = 8, + VT_BETA = 10 + }; + int32_t radius() const { + return GetField(VT_RADIUS, 0); + } + float bias() const { + return GetField(VT_BIAS, 0.0f); + } + float alpha() const { + return GetField(VT_ALPHA, 0.0f); + } + float beta() const { + return GetField(VT_BETA, 0.0f); + } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_RADIUS) && VerifyField(verifier, VT_BIAS) && VerifyField(verifier, VT_ALPHA) && - VerifyField(verifier, VT_BETA) && verifier.EndTable(); + VerifyField(verifier, VT_BETA) && + verifier.EndTable(); } - LocalResponseNormalizationOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - LocalResponseNormalizationOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const LocalResponseNormalizationOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + LocalResponseNormalizationOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(LocalResponseNormalizationOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const LocalResponseNormalizationOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct LocalResponseNormalizationOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; void add_radius(int32_t radius) { - fbb_.AddElement(LocalResponseNormalizationOptions::VT_RADIUS, - radius, 0); + fbb_.AddElement(LocalResponseNormalizationOptions::VT_RADIUS, radius, 0); } void add_bias(float bias) { - fbb_.AddElement(LocalResponseNormalizationOptions::VT_BIAS, bias, - 0.0f); + fbb_.AddElement(LocalResponseNormalizationOptions::VT_BIAS, bias, 0.0f); } void add_alpha(float alpha) { - fbb_.AddElement(LocalResponseNormalizationOptions::VT_ALPHA, alpha, - 0.0f); + fbb_.AddElement(LocalResponseNormalizationOptions::VT_ALPHA, alpha, 0.0f); } void add_beta(float beta) { - fbb_.AddElement(LocalResponseNormalizationOptions::VT_BETA, beta, - 0.0f); + fbb_.AddElement(LocalResponseNormalizationOptions::VT_BETA, beta, 0.0f); } - explicit LocalResponseNormalizationOptionsBuilder( - flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + explicit LocalResponseNormalizationOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } - LocalResponseNormalizationOptionsBuilder &operator=( - const LocalResponseNormalizationOptionsBuilder &); + LocalResponseNormalizationOptionsBuilder &operator=(const LocalResponseNormalizationOptionsBuilder &); flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); auto o = flatbuffers::Offset(end); @@ -2182,10 +2444,12 @@ struct LocalResponseNormalizationOptionsBuilder { } }; -inline flatbuffers::Offset -CreateLocalResponseNormalizationOptions(flatbuffers::FlatBufferBuilder &_fbb, - int32_t radius = 0, float bias = 0.0f, - float alpha = 0.0f, float beta = 0.0f) { +inline flatbuffers::Offset CreateLocalResponseNormalizationOptions( + flatbuffers::FlatBufferBuilder &_fbb, + int32_t radius = 0, + float bias = 0.0f, + float alpha = 0.0f, + float beta = 0.0f) { LocalResponseNormalizationOptionsBuilder builder_(_fbb); builder_.add_beta(beta); builder_.add_alpha(alpha); @@ -2194,11 +2458,7 @@ CreateLocalResponseNormalizationOptions(flatbuffers::FlatBufferBuilder &_fbb, return builder_.Finish(); } -flatbuffers::Offset -CreateLocalResponseNormalizationOptions( - flatbuffers::FlatBufferBuilder &_fbb, - const LocalResponseNormalizationOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateLocalResponseNormalizationOptions(flatbuffers::FlatBufferBuilder &_fbb, const LocalResponseNormalizationOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct LSTMOptionsT : public flatbuffers::NativeTable { typedef LSTMOptions TableType; @@ -2208,41 +2468,43 @@ struct LSTMOptionsT : public flatbuffers::NativeTable { LSTMOptionsT() : fused_activation_function(ActivationFunctionType_NONE), cell_clip(0.0f), - proj_clip(0.0f) {} + proj_clip(0.0f) { + } }; struct LSTMOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef LSTMOptionsT NativeTableType; - enum { VT_FUSED_ACTIVATION_FUNCTION = 4, VT_CELL_CLIP = 6, VT_PROJ_CLIP = 8 }; + enum { + VT_FUSED_ACTIVATION_FUNCTION = 4, + VT_CELL_CLIP = 6, + VT_PROJ_CLIP = 8 + }; ActivationFunctionType fused_activation_function() const { - return static_cast( - GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + } + float cell_clip() const { + return GetField(VT_CELL_CLIP, 0.0f); + } + float proj_clip() const { + return GetField(VT_PROJ_CLIP, 0.0f); } - float cell_clip() const { return GetField(VT_CELL_CLIP, 0.0f); } - float proj_clip() const { return GetField(VT_PROJ_CLIP, 0.0f); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && VerifyField(verifier, VT_CELL_CLIP) && - VerifyField(verifier, VT_PROJ_CLIP) && verifier.EndTable(); + VerifyField(verifier, VT_PROJ_CLIP) && + verifier.EndTable(); } - LSTMOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - LSTMOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + LSTMOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(LSTMOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct LSTMOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_fused_activation_function( - ActivationFunctionType fused_activation_function) { - fbb_.AddElement(LSTMOptions::VT_FUSED_ACTIVATION_FUNCTION, - static_cast(fused_activation_function), 0); + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(LSTMOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } void add_cell_clip(float cell_clip) { fbb_.AddElement(LSTMOptions::VT_CELL_CLIP, cell_clip, 0.0f); @@ -2251,7 +2513,7 @@ struct LSTMOptionsBuilder { fbb_.AddElement(LSTMOptions::VT_PROJ_CLIP, proj_clip, 0.0f); } explicit LSTMOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } LSTMOptionsBuilder &operator=(const LSTMOptionsBuilder &); @@ -2264,9 +2526,9 @@ struct LSTMOptionsBuilder { inline flatbuffers::Offset CreateLSTMOptions( flatbuffers::FlatBufferBuilder &_fbb, - ActivationFunctionType fused_activation_function = - ActivationFunctionType_NONE, - float cell_clip = 0.0f, float proj_clip = 0.0f) { + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE, + float cell_clip = 0.0f, + float proj_clip = 0.0f) { LSTMOptionsBuilder builder_(_fbb); builder_.add_proj_clip(proj_clip); builder_.add_cell_clip(cell_clip); @@ -2274,50 +2536,42 @@ inline flatbuffers::Offset CreateLSTMOptions( return builder_.Finish(); } -flatbuffers::Offset CreateLSTMOptions( - flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateLSTMOptions(flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct ResizeBilinearOptionsT : public flatbuffers::NativeTable { typedef ResizeBilinearOptions TableType; - int32_t new_height; - int32_t new_width; - ResizeBilinearOptionsT() : new_height(0), new_width(0) {} + bool align_corners; + ResizeBilinearOptionsT() + : align_corners(false) { + } }; -struct ResizeBilinearOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct ResizeBilinearOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef ResizeBilinearOptionsT NativeTableType; - enum { VT_NEW_HEIGHT = 4, VT_NEW_WIDTH = 6 }; - int32_t new_height() const { return GetField(VT_NEW_HEIGHT, 0); } - int32_t new_width() const { return GetField(VT_NEW_WIDTH, 0); } + enum { + VT_ALIGN_CORNERS = 8 + }; + bool align_corners() const { + return GetField(VT_ALIGN_CORNERS, 0) != 0; + } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_NEW_HEIGHT) && - VerifyField(verifier, VT_NEW_WIDTH) && verifier.EndTable(); + VerifyField(verifier, VT_ALIGN_CORNERS) && + verifier.EndTable(); } - ResizeBilinearOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - ResizeBilinearOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + ResizeBilinearOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(ResizeBilinearOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct ResizeBilinearOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_new_height(int32_t new_height) { - fbb_.AddElement(ResizeBilinearOptions::VT_NEW_HEIGHT, new_height, - 0); - } - void add_new_width(int32_t new_width) { - fbb_.AddElement(ResizeBilinearOptions::VT_NEW_WIDTH, new_width, 0); + void add_align_corners(bool align_corners) { + fbb_.AddElement(ResizeBilinearOptions::VT_ALIGN_CORNERS, static_cast(align_corners), 0); } explicit ResizeBilinearOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } ResizeBilinearOptionsBuilder &operator=(const ResizeBilinearOptionsBuilder &); @@ -2329,40 +2583,39 @@ struct ResizeBilinearOptionsBuilder { }; inline flatbuffers::Offset CreateResizeBilinearOptions( - flatbuffers::FlatBufferBuilder &_fbb, int32_t new_height = 0, - int32_t new_width = 0) { + flatbuffers::FlatBufferBuilder &_fbb, + bool align_corners = false) { ResizeBilinearOptionsBuilder builder_(_fbb); - builder_.add_new_width(new_width); - builder_.add_new_height(new_height); + builder_.add_align_corners(align_corners); return builder_.Finish(); } -flatbuffers::Offset CreateResizeBilinearOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateResizeBilinearOptions(flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct CallOptionsT : public flatbuffers::NativeTable { typedef CallOptions TableType; uint32_t subgraph; - CallOptionsT() : subgraph(0) {} + CallOptionsT() + : subgraph(0) { + } }; struct CallOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef CallOptionsT NativeTableType; - enum { VT_SUBGRAPH = 4 }; - uint32_t subgraph() const { return GetField(VT_SUBGRAPH, 0); } + enum { + VT_SUBGRAPH = 4 + }; + uint32_t subgraph() const { + return GetField(VT_SUBGRAPH, 0); + } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_SUBGRAPH) && verifier.EndTable(); + VerifyField(verifier, VT_SUBGRAPH) && + verifier.EndTable(); } - CallOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - CallOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + CallOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(CallOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct CallOptionsBuilder { @@ -2372,7 +2625,7 @@ struct CallOptionsBuilder { fbb_.AddElement(CallOptions::VT_SUBGRAPH, subgraph, 0); } explicit CallOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } CallOptionsBuilder &operator=(const CallOptionsBuilder &); @@ -2384,62 +2637,37 @@ struct CallOptionsBuilder { }; inline flatbuffers::Offset CreateCallOptions( - flatbuffers::FlatBufferBuilder &_fbb, uint32_t subgraph = 0) { + flatbuffers::FlatBufferBuilder &_fbb, + uint32_t subgraph = 0) { CallOptionsBuilder builder_(_fbb); builder_.add_subgraph(subgraph); return builder_.Finish(); } -flatbuffers::Offset CreateCallOptions( - flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateCallOptions(flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct PadOptionsT : public flatbuffers::NativeTable { typedef PadOptions TableType; - std::vector before_padding; - std::vector after_padding; - PadOptionsT() {} + PadOptionsT() { + } }; struct PadOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef PadOptionsT NativeTableType; - enum { VT_BEFORE_PADDING = 4, VT_AFTER_PADDING = 6 }; - const flatbuffers::Vector *before_padding() const { - return GetPointer *>(VT_BEFORE_PADDING); - } - const flatbuffers::Vector *after_padding() const { - return GetPointer *>(VT_AFTER_PADDING); - } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyOffset(verifier, VT_BEFORE_PADDING) && - verifier.Verify(before_padding()) && - VerifyOffset(verifier, VT_AFTER_PADDING) && - verifier.Verify(after_padding()) && verifier.EndTable(); + verifier.EndTable(); } - PadOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - PadOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + PadOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(PadOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct PadOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_before_padding( - flatbuffers::Offset> before_padding) { - fbb_.AddOffset(PadOptions::VT_BEFORE_PADDING, before_padding); - } - void add_after_padding( - flatbuffers::Offset> after_padding) { - fbb_.AddOffset(PadOptions::VT_AFTER_PADDING, after_padding); - } explicit PadOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } PadOptionsBuilder &operator=(const PadOptionsBuilder &); @@ -2451,63 +2679,47 @@ struct PadOptionsBuilder { }; inline flatbuffers::Offset CreatePadOptions( - flatbuffers::FlatBufferBuilder &_fbb, - flatbuffers::Offset> before_padding = 0, - flatbuffers::Offset> after_padding = 0) { + flatbuffers::FlatBufferBuilder &_fbb) { PadOptionsBuilder builder_(_fbb); - builder_.add_after_padding(after_padding); - builder_.add_before_padding(before_padding); return builder_.Finish(); } -inline flatbuffers::Offset CreatePadOptionsDirect( - flatbuffers::FlatBufferBuilder &_fbb, - const std::vector *before_padding = nullptr, - const std::vector *after_padding = nullptr) { - return tflite::CreatePadOptions( - _fbb, before_padding ? _fbb.CreateVector(*before_padding) : 0, - after_padding ? _fbb.CreateVector(*after_padding) : 0); -} - -flatbuffers::Offset CreatePadOptions( - flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreatePadOptions(flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct ReshapeOptionsT : public flatbuffers::NativeTable { typedef ReshapeOptions TableType; std::vector new_shape; - ReshapeOptionsT() {} + ReshapeOptionsT() { + } }; struct ReshapeOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef ReshapeOptionsT NativeTableType; - enum { VT_NEW_SHAPE = 4 }; + enum { + VT_NEW_SHAPE = 4 + }; const flatbuffers::Vector *new_shape() const { return GetPointer *>(VT_NEW_SHAPE); } bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && VerifyOffset(verifier, VT_NEW_SHAPE) && - verifier.Verify(new_shape()) && verifier.EndTable(); + return VerifyTableStart(verifier) && + VerifyOffset(verifier, VT_NEW_SHAPE) && + verifier.Verify(new_shape()) && + verifier.EndTable(); } - ReshapeOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - ReshapeOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + ReshapeOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(ReshapeOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct ReshapeOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_new_shape( - flatbuffers::Offset> new_shape) { + void add_new_shape(flatbuffers::Offset> new_shape) { fbb_.AddOffset(ReshapeOptions::VT_NEW_SHAPE, new_shape); } explicit ReshapeOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } ReshapeOptionsBuilder &operator=(const ReshapeOptionsBuilder &); @@ -2530,70 +2742,74 @@ inline flatbuffers::Offset CreateReshapeOptionsDirect( flatbuffers::FlatBufferBuilder &_fbb, const std::vector *new_shape = nullptr) { return tflite::CreateReshapeOptions( - _fbb, new_shape ? _fbb.CreateVector(*new_shape) : 0); + _fbb, + new_shape ? _fbb.CreateVector(*new_shape) : 0); } -flatbuffers::Offset CreateReshapeOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateReshapeOptions(flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); -struct BatchToSpaceNDOptionsT : public flatbuffers::NativeTable { - typedef BatchToSpaceNDOptions TableType; - std::vector block_shape; - std::vector before_crops; - std::vector after_crops; - BatchToSpaceNDOptionsT() {} +struct SpaceToBatchNDOptionsT : public flatbuffers::NativeTable { + typedef SpaceToBatchNDOptions TableType; + SpaceToBatchNDOptionsT() { + } }; -struct BatchToSpaceNDOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { - typedef BatchToSpaceNDOptionsT NativeTableType; - enum { VT_BLOCK_SHAPE = 4, VT_BEFORE_CROPS = 6, VT_AFTER_CROPS = 8 }; - const flatbuffers::Vector *block_shape() const { - return GetPointer *>(VT_BLOCK_SHAPE); +struct SpaceToBatchNDOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef SpaceToBatchNDOptionsT NativeTableType; + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + verifier.EndTable(); + } + SpaceToBatchNDOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SpaceToBatchNDOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SpaceToBatchNDOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct SpaceToBatchNDOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + explicit SpaceToBatchNDOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); } - const flatbuffers::Vector *before_crops() const { - return GetPointer *>(VT_BEFORE_CROPS); + SpaceToBatchNDOptionsBuilder &operator=(const SpaceToBatchNDOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; } - const flatbuffers::Vector *after_crops() const { - return GetPointer *>(VT_AFTER_CROPS); +}; + +inline flatbuffers::Offset CreateSpaceToBatchNDOptions( + flatbuffers::FlatBufferBuilder &_fbb) { + SpaceToBatchNDOptionsBuilder builder_(_fbb); + return builder_.Finish(); +} + +flatbuffers::Offset CreateSpaceToBatchNDOptions(flatbuffers::FlatBufferBuilder &_fbb, const SpaceToBatchNDOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct BatchToSpaceNDOptionsT : public flatbuffers::NativeTable { + typedef BatchToSpaceNDOptions TableType; + BatchToSpaceNDOptionsT() { } +}; + +struct BatchToSpaceNDOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef BatchToSpaceNDOptionsT NativeTableType; bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyOffset(verifier, VT_BLOCK_SHAPE) && - verifier.Verify(block_shape()) && - VerifyOffset(verifier, VT_BEFORE_CROPS) && - verifier.Verify(before_crops()) && - VerifyOffset(verifier, VT_AFTER_CROPS) && - verifier.Verify(after_crops()) && verifier.EndTable(); - } - BatchToSpaceNDOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - BatchToSpaceNDOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + verifier.EndTable(); + } + BatchToSpaceNDOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(BatchToSpaceNDOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct BatchToSpaceNDOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_block_shape( - flatbuffers::Offset> block_shape) { - fbb_.AddOffset(BatchToSpaceNDOptions::VT_BLOCK_SHAPE, block_shape); - } - void add_before_crops( - flatbuffers::Offset> before_crops) { - fbb_.AddOffset(BatchToSpaceNDOptions::VT_BEFORE_CROPS, before_crops); - } - void add_after_crops( - flatbuffers::Offset> after_crops) { - fbb_.AddOffset(BatchToSpaceNDOptions::VT_AFTER_CROPS, after_crops); - } explicit BatchToSpaceNDOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } BatchToSpaceNDOptionsBuilder &operator=(const BatchToSpaceNDOptionsBuilder &); @@ -2605,32 +2821,12 @@ struct BatchToSpaceNDOptionsBuilder { }; inline flatbuffers::Offset CreateBatchToSpaceNDOptions( - flatbuffers::FlatBufferBuilder &_fbb, - flatbuffers::Offset> block_shape = 0, - flatbuffers::Offset> before_crops = 0, - flatbuffers::Offset> after_crops = 0) { + flatbuffers::FlatBufferBuilder &_fbb) { BatchToSpaceNDOptionsBuilder builder_(_fbb); - builder_.add_after_crops(after_crops); - builder_.add_before_crops(before_crops); - builder_.add_block_shape(block_shape); return builder_.Finish(); } -inline flatbuffers::Offset -CreateBatchToSpaceNDOptionsDirect( - flatbuffers::FlatBufferBuilder &_fbb, - const std::vector *block_shape = nullptr, - const std::vector *before_crops = nullptr, - const std::vector *after_crops = nullptr) { - return tflite::CreateBatchToSpaceNDOptions( - _fbb, block_shape ? _fbb.CreateVector(*block_shape) : 0, - before_crops ? _fbb.CreateVector(*before_crops) : 0, - after_crops ? _fbb.CreateVector(*after_crops) : 0); -} - -flatbuffers::Offset CreateBatchToSpaceNDOptions( - flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateBatchToSpaceNDOptions(flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct SkipGramOptionsT : public flatbuffers::NativeTable { typedef SkipGramOptions TableType; @@ -2638,13 +2834,22 @@ struct SkipGramOptionsT : public flatbuffers::NativeTable { int32_t max_skip_size; bool include_all_ngrams; SkipGramOptionsT() - : ngram_size(0), max_skip_size(0), include_all_ngrams(false) {} + : ngram_size(0), + max_skip_size(0), + include_all_ngrams(false) { + } }; struct SkipGramOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef SkipGramOptionsT NativeTableType; - enum { VT_NGRAM_SIZE = 4, VT_MAX_SKIP_SIZE = 6, VT_INCLUDE_ALL_NGRAMS = 8 }; - int32_t ngram_size() const { return GetField(VT_NGRAM_SIZE, 0); } + enum { + VT_NGRAM_SIZE = 4, + VT_MAX_SKIP_SIZE = 6, + VT_INCLUDE_ALL_NGRAMS = 8 + }; + int32_t ngram_size() const { + return GetField(VT_NGRAM_SIZE, 0); + } int32_t max_skip_size() const { return GetField(VT_MAX_SKIP_SIZE, 0); } @@ -2658,14 +2863,9 @@ struct SkipGramOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { VerifyField(verifier, VT_INCLUDE_ALL_NGRAMS) && verifier.EndTable(); } - SkipGramOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - SkipGramOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + SkipGramOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SkipGramOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct SkipGramOptionsBuilder { @@ -2675,15 +2875,13 @@ struct SkipGramOptionsBuilder { fbb_.AddElement(SkipGramOptions::VT_NGRAM_SIZE, ngram_size, 0); } void add_max_skip_size(int32_t max_skip_size) { - fbb_.AddElement(SkipGramOptions::VT_MAX_SKIP_SIZE, max_skip_size, - 0); + fbb_.AddElement(SkipGramOptions::VT_MAX_SKIP_SIZE, max_skip_size, 0); } void add_include_all_ngrams(bool include_all_ngrams) { - fbb_.AddElement(SkipGramOptions::VT_INCLUDE_ALL_NGRAMS, - static_cast(include_all_ngrams), 0); + fbb_.AddElement(SkipGramOptions::VT_INCLUDE_ALL_NGRAMS, static_cast(include_all_ngrams), 0); } explicit SkipGramOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SkipGramOptionsBuilder &operator=(const SkipGramOptionsBuilder &); @@ -2695,8 +2893,10 @@ struct SkipGramOptionsBuilder { }; inline flatbuffers::Offset CreateSkipGramOptions( - flatbuffers::FlatBufferBuilder &_fbb, int32_t ngram_size = 0, - int32_t max_skip_size = 0, bool include_all_ngrams = false) { + flatbuffers::FlatBufferBuilder &_fbb, + int32_t ngram_size = 0, + int32_t max_skip_size = 0, + bool include_all_ngrams = false) { SkipGramOptionsBuilder builder_(_fbb); builder_.add_max_skip_size(max_skip_size); builder_.add_ngram_size(ngram_size); @@ -2704,33 +2904,32 @@ inline flatbuffers::Offset CreateSkipGramOptions( return builder_.Finish(); } -flatbuffers::Offset CreateSkipGramOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateSkipGramOptions(flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct SpaceToDepthOptionsT : public flatbuffers::NativeTable { typedef SpaceToDepthOptions TableType; int32_t block_size; - SpaceToDepthOptionsT() : block_size(0) {} + SpaceToDepthOptionsT() + : block_size(0) { + } }; -struct SpaceToDepthOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct SpaceToDepthOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef SpaceToDepthOptionsT NativeTableType; - enum { VT_BLOCK_SIZE = 4 }; - int32_t block_size() const { return GetField(VT_BLOCK_SIZE, 0); } + enum { + VT_BLOCK_SIZE = 4 + }; + int32_t block_size() const { + return GetField(VT_BLOCK_SIZE, 0); + } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_BLOCK_SIZE) && verifier.EndTable(); + VerifyField(verifier, VT_BLOCK_SIZE) && + verifier.EndTable(); } - SpaceToDepthOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - SpaceToDepthOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + SpaceToDepthOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SpaceToDepthOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct SpaceToDepthOptionsBuilder { @@ -2740,7 +2939,7 @@ struct SpaceToDepthOptionsBuilder { fbb_.AddElement(SpaceToDepthOptions::VT_BLOCK_SIZE, block_size, 0); } explicit SpaceToDepthOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SpaceToDepthOptionsBuilder &operator=(const SpaceToDepthOptionsBuilder &); @@ -2752,58 +2951,200 @@ struct SpaceToDepthOptionsBuilder { }; inline flatbuffers::Offset CreateSpaceToDepthOptions( - flatbuffers::FlatBufferBuilder &_fbb, int32_t block_size = 0) { + flatbuffers::FlatBufferBuilder &_fbb, + int32_t block_size = 0) { SpaceToDepthOptionsBuilder builder_(_fbb); builder_.add_block_size(block_size); return builder_.Finish(); } -flatbuffers::Offset CreateSpaceToDepthOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateSpaceToDepthOptions(flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct SubOptionsT : public flatbuffers::NativeTable { + typedef SubOptions TableType; + ActivationFunctionType fused_activation_function; + SubOptionsT() + : fused_activation_function(ActivationFunctionType_NONE) { + } +}; + +struct SubOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef SubOptionsT NativeTableType; + enum { + VT_FUSED_ACTIVATION_FUNCTION = 4 + }; + ActivationFunctionType fused_activation_function() const { + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && + verifier.EndTable(); + } + SubOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SubOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SubOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct SubOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(SubOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); + } + explicit SubOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + SubOptionsBuilder &operator=(const SubOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateSubOptions( + flatbuffers::FlatBufferBuilder &_fbb, + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { + SubOptionsBuilder builder_(_fbb); + builder_.add_fused_activation_function(fused_activation_function); + return builder_.Finish(); +} + +flatbuffers::Offset CreateSubOptions(flatbuffers::FlatBufferBuilder &_fbb, const SubOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct DivOptionsT : public flatbuffers::NativeTable { + typedef DivOptions TableType; + ActivationFunctionType fused_activation_function; + DivOptionsT() + : fused_activation_function(ActivationFunctionType_NONE) { + } +}; + +struct DivOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef DivOptionsT NativeTableType; + enum { + VT_FUSED_ACTIVATION_FUNCTION = 4 + }; + ActivationFunctionType fused_activation_function() const { + return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && + verifier.EndTable(); + } + DivOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(DivOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const DivOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct DivOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_fused_activation_function(ActivationFunctionType fused_activation_function) { + fbb_.AddElement(DivOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); + } + explicit DivOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + DivOptionsBuilder &operator=(const DivOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateDivOptions( + flatbuffers::FlatBufferBuilder &_fbb, + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { + DivOptionsBuilder builder_(_fbb); + builder_.add_fused_activation_function(fused_activation_function); + return builder_.Finish(); +} + +flatbuffers::Offset CreateDivOptions(flatbuffers::FlatBufferBuilder &_fbb, const DivOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct TopKV2OptionsT : public flatbuffers::NativeTable { + typedef TopKV2Options TableType; + TopKV2OptionsT() { + } +}; + +struct TopKV2Options FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef TopKV2OptionsT NativeTableType; + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + verifier.EndTable(); + } + TopKV2OptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(TopKV2OptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const TopKV2OptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct TopKV2OptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + explicit TopKV2OptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + TopKV2OptionsBuilder &operator=(const TopKV2OptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateTopKV2Options( + flatbuffers::FlatBufferBuilder &_fbb) { + TopKV2OptionsBuilder builder_(_fbb); + return builder_.Finish(); +} + +flatbuffers::Offset CreateTopKV2Options(flatbuffers::FlatBufferBuilder &_fbb, const TopKV2OptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct EmbeddingLookupSparseOptionsT : public flatbuffers::NativeTable { typedef EmbeddingLookupSparseOptions TableType; CombinerType combiner; - EmbeddingLookupSparseOptionsT() : combiner(CombinerType_SUM) {} + EmbeddingLookupSparseOptionsT() + : combiner(CombinerType_SUM) { + } }; -struct EmbeddingLookupSparseOptions FLATBUFFERS_FINAL_CLASS - : private flatbuffers::Table { +struct EmbeddingLookupSparseOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef EmbeddingLookupSparseOptionsT NativeTableType; - enum { VT_COMBINER = 4 }; + enum { + VT_COMBINER = 4 + }; CombinerType combiner() const { return static_cast(GetField(VT_COMBINER, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_COMBINER) && verifier.EndTable(); + VerifyField(verifier, VT_COMBINER) && + verifier.EndTable(); } - EmbeddingLookupSparseOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - EmbeddingLookupSparseOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const EmbeddingLookupSparseOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + EmbeddingLookupSparseOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(EmbeddingLookupSparseOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const EmbeddingLookupSparseOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct EmbeddingLookupSparseOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; void add_combiner(CombinerType combiner) { - fbb_.AddElement(EmbeddingLookupSparseOptions::VT_COMBINER, - static_cast(combiner), 0); + fbb_.AddElement(EmbeddingLookupSparseOptions::VT_COMBINER, static_cast(combiner), 0); } - explicit EmbeddingLookupSparseOptionsBuilder( - flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + explicit EmbeddingLookupSparseOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } - EmbeddingLookupSparseOptionsBuilder &operator=( - const EmbeddingLookupSparseOptionsBuilder &); + EmbeddingLookupSparseOptionsBuilder &operator=(const EmbeddingLookupSparseOptionsBuilder &); flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); auto o = flatbuffers::Offset(end); @@ -2811,42 +3152,40 @@ struct EmbeddingLookupSparseOptionsBuilder { } }; -inline flatbuffers::Offset -CreateEmbeddingLookupSparseOptions(flatbuffers::FlatBufferBuilder &_fbb, - CombinerType combiner = CombinerType_SUM) { +inline flatbuffers::Offset CreateEmbeddingLookupSparseOptions( + flatbuffers::FlatBufferBuilder &_fbb, + CombinerType combiner = CombinerType_SUM) { EmbeddingLookupSparseOptionsBuilder builder_(_fbb); builder_.add_combiner(combiner); return builder_.Finish(); } -flatbuffers::Offset -CreateEmbeddingLookupSparseOptions( - flatbuffers::FlatBufferBuilder &_fbb, - const EmbeddingLookupSparseOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateEmbeddingLookupSparseOptions(flatbuffers::FlatBufferBuilder &_fbb, const EmbeddingLookupSparseOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct GatherOptionsT : public flatbuffers::NativeTable { typedef GatherOptions TableType; int32_t axis; - GatherOptionsT() : axis(0) {} + GatherOptionsT() + : axis(0) { + } }; struct GatherOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef GatherOptionsT NativeTableType; - enum { VT_AXIS = 4 }; - int32_t axis() const { return GetField(VT_AXIS, 0); } + enum { + VT_AXIS = 4 + }; + int32_t axis() const { + return GetField(VT_AXIS, 0); + } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_AXIS) && verifier.EndTable(); + VerifyField(verifier, VT_AXIS) && + verifier.EndTable(); } - GatherOptionsT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - GatherOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + GatherOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(GatherOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct GatherOptionsBuilder { @@ -2856,7 +3195,7 @@ struct GatherOptionsBuilder { fbb_.AddElement(GatherOptions::VT_AXIS, axis, 0); } explicit GatherOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + : fbb_(_fbb) { start_ = fbb_.StartTable(); } GatherOptionsBuilder &operator=(const GatherOptionsBuilder &); @@ -2868,432 +3207,755 @@ struct GatherOptionsBuilder { }; inline flatbuffers::Offset CreateGatherOptions( - flatbuffers::FlatBufferBuilder &_fbb, int32_t axis = 0) { + flatbuffers::FlatBufferBuilder &_fbb, + int32_t axis = 0) { GatherOptionsBuilder builder_(_fbb); builder_.add_axis(axis); return builder_.Finish(); } -flatbuffers::Offset CreateGatherOptions( - flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateGatherOptions(flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); -struct OperatorCodeT : public flatbuffers::NativeTable { - typedef OperatorCode TableType; - BuiltinOperator builtin_code; - std::string custom_code; - OperatorCodeT() : builtin_code(BuiltinOperator_ADD) {} +struct TransposeOptionsT : public flatbuffers::NativeTable { + typedef TransposeOptions TableType; + TransposeOptionsT() { + } }; -struct OperatorCode FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { - typedef OperatorCodeT NativeTableType; - enum { VT_BUILTIN_CODE = 4, VT_CUSTOM_CODE = 6 }; - BuiltinOperator builtin_code() const { - return static_cast(GetField(VT_BUILTIN_CODE, 0)); - } - const flatbuffers::String *custom_code() const { - return GetPointer(VT_CUSTOM_CODE); - } +struct TransposeOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef TransposeOptionsT NativeTableType; bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && - VerifyField(verifier, VT_BUILTIN_CODE) && - VerifyOffset(verifier, VT_CUSTOM_CODE) && - verifier.Verify(custom_code()) && verifier.EndTable(); + verifier.EndTable(); } - OperatorCodeT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - OperatorCodeT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + TransposeOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(TransposeOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const TransposeOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; -struct OperatorCodeBuilder { +struct TransposeOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_builtin_code(BuiltinOperator builtin_code) { - fbb_.AddElement(OperatorCode::VT_BUILTIN_CODE, - static_cast(builtin_code), 0); - } - void add_custom_code(flatbuffers::Offset custom_code) { - fbb_.AddOffset(OperatorCode::VT_CUSTOM_CODE, custom_code); - } - explicit OperatorCodeBuilder(flatbuffers::FlatBufferBuilder &_fbb) - : fbb_(_fbb) { + explicit TransposeOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } - OperatorCodeBuilder &operator=(const OperatorCodeBuilder &); - flatbuffers::Offset Finish() { + TransposeOptionsBuilder &operator=(const TransposeOptionsBuilder &); + flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); - auto o = flatbuffers::Offset(end); + auto o = flatbuffers::Offset(end); return o; } }; -inline flatbuffers::Offset CreateOperatorCode( - flatbuffers::FlatBufferBuilder &_fbb, - BuiltinOperator builtin_code = BuiltinOperator_ADD, - flatbuffers::Offset custom_code = 0) { - OperatorCodeBuilder builder_(_fbb); - builder_.add_custom_code(custom_code); - builder_.add_builtin_code(builtin_code); +inline flatbuffers::Offset CreateTransposeOptions( + flatbuffers::FlatBufferBuilder &_fbb) { + TransposeOptionsBuilder builder_(_fbb); return builder_.Finish(); } -inline flatbuffers::Offset CreateOperatorCodeDirect( - flatbuffers::FlatBufferBuilder &_fbb, - BuiltinOperator builtin_code = BuiltinOperator_ADD, - const char *custom_code = nullptr) { - return tflite::CreateOperatorCode( - _fbb, builtin_code, custom_code ? _fbb.CreateString(custom_code) : 0); +flatbuffers::Offset CreateTransposeOptions(flatbuffers::FlatBufferBuilder &_fbb, const TransposeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct ExpOptionsT : public flatbuffers::NativeTable { + typedef ExpOptions TableType; + ExpOptionsT() { + } +}; + +struct ExpOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef ExpOptionsT NativeTableType; + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + verifier.EndTable(); + } + ExpOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(ExpOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const ExpOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct ExpOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + explicit ExpOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + ExpOptionsBuilder &operator=(const ExpOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateExpOptions( + flatbuffers::FlatBufferBuilder &_fbb) { + ExpOptionsBuilder builder_(_fbb); + return builder_.Finish(); } -flatbuffers::Offset CreateOperatorCode( - flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateExpOptions(flatbuffers::FlatBufferBuilder &_fbb, const ExpOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); -struct OperatorT : public flatbuffers::NativeTable { - typedef Operator TableType; - uint32_t opcode_index; - std::vector inputs; - std::vector outputs; - BuiltinOptionsUnion builtin_options; - std::vector custom_options; - CustomOptionsFormat custom_options_format; - OperatorT() - : opcode_index(0), - custom_options_format(CustomOptionsFormat_FLEXBUFFERS) {} +struct MeanOptionsT : public flatbuffers::NativeTable { + typedef MeanOptions TableType; + bool keep_dims; + MeanOptionsT() + : keep_dims(false) { + } }; -struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { - typedef OperatorT NativeTableType; +struct MeanOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef MeanOptionsT NativeTableType; enum { - VT_OPCODE_INDEX = 4, - VT_INPUTS = 6, - VT_OUTPUTS = 8, - VT_BUILTIN_OPTIONS_TYPE = 10, - VT_BUILTIN_OPTIONS = 12, - VT_CUSTOM_OPTIONS = 14, - VT_CUSTOM_OPTIONS_FORMAT = 16 + VT_KEEP_DIMS = 4 }; - uint32_t opcode_index() const { - return GetField(VT_OPCODE_INDEX, 0); - } - const flatbuffers::Vector *inputs() const { - return GetPointer *>(VT_INPUTS); + bool keep_dims() const { + return GetField(VT_KEEP_DIMS, 0) != 0; } - const flatbuffers::Vector *outputs() const { - return GetPointer *>(VT_OUTPUTS); + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_KEEP_DIMS) && + verifier.EndTable(); } - BuiltinOptions builtin_options_type() const { - return static_cast( - GetField(VT_BUILTIN_OPTIONS_TYPE, 0)); + MeanOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(MeanOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const MeanOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct MeanOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_keep_dims(bool keep_dims) { + fbb_.AddElement(MeanOptions::VT_KEEP_DIMS, static_cast(keep_dims), 0); } - const void *builtin_options() const { - return GetPointer(VT_BUILTIN_OPTIONS); + explicit MeanOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); } - template - const T *builtin_options_as() const; - const Conv2DOptions *builtin_options_as_Conv2DOptions() const { - return builtin_options_type() == BuiltinOptions_Conv2DOptions - ? static_cast(builtin_options()) - : nullptr; - } - const DepthwiseConv2DOptions *builtin_options_as_DepthwiseConv2DOptions() - const { - return builtin_options_type() == BuiltinOptions_DepthwiseConv2DOptions - ? static_cast(builtin_options()) - : nullptr; - } - const ConcatEmbeddingsOptions *builtin_options_as_ConcatEmbeddingsOptions() - const { - return builtin_options_type() == BuiltinOptions_ConcatEmbeddingsOptions - ? static_cast(builtin_options()) - : nullptr; + MeanOptionsBuilder &operator=(const MeanOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; } - const LSHProjectionOptions *builtin_options_as_LSHProjectionOptions() const { - return builtin_options_type() == BuiltinOptions_LSHProjectionOptions - ? static_cast(builtin_options()) - : nullptr; +}; + +inline flatbuffers::Offset CreateMeanOptions( + flatbuffers::FlatBufferBuilder &_fbb, + bool keep_dims = false) { + MeanOptionsBuilder builder_(_fbb); + builder_.add_keep_dims(keep_dims); + return builder_.Finish(); +} + +flatbuffers::Offset CreateMeanOptions(flatbuffers::FlatBufferBuilder &_fbb, const MeanOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct SqueezeOptionsT : public flatbuffers::NativeTable { + typedef SqueezeOptions TableType; + std::vector squeeze_dims; + SqueezeOptionsT() { } - const Pool2DOptions *builtin_options_as_Pool2DOptions() const { - return builtin_options_type() == BuiltinOptions_Pool2DOptions - ? static_cast(builtin_options()) - : nullptr; +}; + +struct SqueezeOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef SqueezeOptionsT NativeTableType; + enum { + VT_SQUEEZE_DIMS = 4 + }; + const flatbuffers::Vector *squeeze_dims() const { + return GetPointer *>(VT_SQUEEZE_DIMS); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyOffset(verifier, VT_SQUEEZE_DIMS) && + verifier.Verify(squeeze_dims()) && + verifier.EndTable(); + } + SqueezeOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SqueezeOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SqueezeOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct SqueezeOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_squeeze_dims(flatbuffers::Offset> squeeze_dims) { + fbb_.AddOffset(SqueezeOptions::VT_SQUEEZE_DIMS, squeeze_dims); + } + explicit SqueezeOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + SqueezeOptionsBuilder &operator=(const SqueezeOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateSqueezeOptions( + flatbuffers::FlatBufferBuilder &_fbb, + flatbuffers::Offset> squeeze_dims = 0) { + SqueezeOptionsBuilder builder_(_fbb); + builder_.add_squeeze_dims(squeeze_dims); + return builder_.Finish(); +} + +inline flatbuffers::Offset CreateSqueezeOptionsDirect( + flatbuffers::FlatBufferBuilder &_fbb, + const std::vector *squeeze_dims = nullptr) { + return tflite::CreateSqueezeOptions( + _fbb, + squeeze_dims ? _fbb.CreateVector(*squeeze_dims) : 0); +} + +flatbuffers::Offset CreateSqueezeOptions(flatbuffers::FlatBufferBuilder &_fbb, const SqueezeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct SplitOptionsT : public flatbuffers::NativeTable { + typedef SplitOptions TableType; + int32_t num_splits; + SplitOptionsT() + : num_splits(0) { + } +}; + +struct SplitOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef SplitOptionsT NativeTableType; + enum { + VT_NUM_SPLITS = 4 + }; + int32_t num_splits() const { + return GetField(VT_NUM_SPLITS, 0); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_NUM_SPLITS) && + verifier.EndTable(); + } + SplitOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SplitOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SplitOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct SplitOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_num_splits(int32_t num_splits) { + fbb_.AddElement(SplitOptions::VT_NUM_SPLITS, num_splits, 0); + } + explicit SplitOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + SplitOptionsBuilder &operator=(const SplitOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateSplitOptions( + flatbuffers::FlatBufferBuilder &_fbb, + int32_t num_splits = 0) { + SplitOptionsBuilder builder_(_fbb); + builder_.add_num_splits(num_splits); + return builder_.Finish(); +} + +flatbuffers::Offset CreateSplitOptions(flatbuffers::FlatBufferBuilder &_fbb, const SplitOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct StridedSliceOptionsT : public flatbuffers::NativeTable { + typedef StridedSliceOptions TableType; + int32_t begin_mask; + int32_t end_mask; + int32_t ellipsis_mask; + int32_t new_axis_mask; + int32_t shrink_axis_mask; + StridedSliceOptionsT() + : begin_mask(0), + end_mask(0), + ellipsis_mask(0), + new_axis_mask(0), + shrink_axis_mask(0) { + } +}; + +struct StridedSliceOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef StridedSliceOptionsT NativeTableType; + enum { + VT_BEGIN_MASK = 4, + VT_END_MASK = 6, + VT_ELLIPSIS_MASK = 8, + VT_NEW_AXIS_MASK = 10, + VT_SHRINK_AXIS_MASK = 12 + }; + int32_t begin_mask() const { + return GetField(VT_BEGIN_MASK, 0); + } + int32_t end_mask() const { + return GetField(VT_END_MASK, 0); + } + int32_t ellipsis_mask() const { + return GetField(VT_ELLIPSIS_MASK, 0); + } + int32_t new_axis_mask() const { + return GetField(VT_NEW_AXIS_MASK, 0); + } + int32_t shrink_axis_mask() const { + return GetField(VT_SHRINK_AXIS_MASK, 0); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_BEGIN_MASK) && + VerifyField(verifier, VT_END_MASK) && + VerifyField(verifier, VT_ELLIPSIS_MASK) && + VerifyField(verifier, VT_NEW_AXIS_MASK) && + VerifyField(verifier, VT_SHRINK_AXIS_MASK) && + verifier.EndTable(); + } + StridedSliceOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(StridedSliceOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const StridedSliceOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct StridedSliceOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_begin_mask(int32_t begin_mask) { + fbb_.AddElement(StridedSliceOptions::VT_BEGIN_MASK, begin_mask, 0); + } + void add_end_mask(int32_t end_mask) { + fbb_.AddElement(StridedSliceOptions::VT_END_MASK, end_mask, 0); + } + void add_ellipsis_mask(int32_t ellipsis_mask) { + fbb_.AddElement(StridedSliceOptions::VT_ELLIPSIS_MASK, ellipsis_mask, 0); + } + void add_new_axis_mask(int32_t new_axis_mask) { + fbb_.AddElement(StridedSliceOptions::VT_NEW_AXIS_MASK, new_axis_mask, 0); + } + void add_shrink_axis_mask(int32_t shrink_axis_mask) { + fbb_.AddElement(StridedSliceOptions::VT_SHRINK_AXIS_MASK, shrink_axis_mask, 0); + } + explicit StridedSliceOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + StridedSliceOptionsBuilder &operator=(const StridedSliceOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateStridedSliceOptions( + flatbuffers::FlatBufferBuilder &_fbb, + int32_t begin_mask = 0, + int32_t end_mask = 0, + int32_t ellipsis_mask = 0, + int32_t new_axis_mask = 0, + int32_t shrink_axis_mask = 0) { + StridedSliceOptionsBuilder builder_(_fbb); + builder_.add_shrink_axis_mask(shrink_axis_mask); + builder_.add_new_axis_mask(new_axis_mask); + builder_.add_ellipsis_mask(ellipsis_mask); + builder_.add_end_mask(end_mask); + builder_.add_begin_mask(begin_mask); + return builder_.Finish(); +} + +flatbuffers::Offset CreateStridedSliceOptions(flatbuffers::FlatBufferBuilder &_fbb, const StridedSliceOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct OperatorCodeT : public flatbuffers::NativeTable { + typedef OperatorCode TableType; + BuiltinOperator builtin_code; + std::string custom_code; + OperatorCodeT() + : builtin_code(BuiltinOperator_ADD) { + } +}; + +struct OperatorCode FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef OperatorCodeT NativeTableType; + enum { + VT_BUILTIN_CODE = 4, + VT_CUSTOM_CODE = 6 + }; + BuiltinOperator builtin_code() const { + return static_cast(GetField(VT_BUILTIN_CODE, 0)); + } + const flatbuffers::String *custom_code() const { + return GetPointer(VT_CUSTOM_CODE); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_BUILTIN_CODE) && + VerifyOffset(verifier, VT_CUSTOM_CODE) && + verifier.Verify(custom_code()) && + verifier.EndTable(); + } + OperatorCodeT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(OperatorCodeT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct OperatorCodeBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_builtin_code(BuiltinOperator builtin_code) { + fbb_.AddElement(OperatorCode::VT_BUILTIN_CODE, static_cast(builtin_code), 0); + } + void add_custom_code(flatbuffers::Offset custom_code) { + fbb_.AddOffset(OperatorCode::VT_CUSTOM_CODE, custom_code); + } + explicit OperatorCodeBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + OperatorCodeBuilder &operator=(const OperatorCodeBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateOperatorCode( + flatbuffers::FlatBufferBuilder &_fbb, + BuiltinOperator builtin_code = BuiltinOperator_ADD, + flatbuffers::Offset custom_code = 0) { + OperatorCodeBuilder builder_(_fbb); + builder_.add_custom_code(custom_code); + builder_.add_builtin_code(builtin_code); + return builder_.Finish(); +} + +inline flatbuffers::Offset CreateOperatorCodeDirect( + flatbuffers::FlatBufferBuilder &_fbb, + BuiltinOperator builtin_code = BuiltinOperator_ADD, + const char *custom_code = nullptr) { + return tflite::CreateOperatorCode( + _fbb, + builtin_code, + custom_code ? _fbb.CreateString(custom_code) : 0); +} + +flatbuffers::Offset CreateOperatorCode(flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct OperatorT : public flatbuffers::NativeTable { + typedef Operator TableType; + uint32_t opcode_index; + std::vector inputs; + std::vector outputs; + BuiltinOptionsUnion builtin_options; + std::vector custom_options; + CustomOptionsFormat custom_options_format; + OperatorT() + : opcode_index(0), + custom_options_format(CustomOptionsFormat_FLEXBUFFERS) { + } +}; + +struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef OperatorT NativeTableType; + enum { + VT_OPCODE_INDEX = 4, + VT_INPUTS = 6, + VT_OUTPUTS = 8, + VT_BUILTIN_OPTIONS_TYPE = 10, + VT_BUILTIN_OPTIONS = 12, + VT_CUSTOM_OPTIONS = 14, + VT_CUSTOM_OPTIONS_FORMAT = 16 + }; + uint32_t opcode_index() const { + return GetField(VT_OPCODE_INDEX, 0); + } + const flatbuffers::Vector *inputs() const { + return GetPointer *>(VT_INPUTS); + } + const flatbuffers::Vector *outputs() const { + return GetPointer *>(VT_OUTPUTS); + } + BuiltinOptions builtin_options_type() const { + return static_cast(GetField(VT_BUILTIN_OPTIONS_TYPE, 0)); + } + const void *builtin_options() const { + return GetPointer(VT_BUILTIN_OPTIONS); + } + template const T *builtin_options_as() const; + const Conv2DOptions *builtin_options_as_Conv2DOptions() const { + return builtin_options_type() == BuiltinOptions_Conv2DOptions ? static_cast(builtin_options()) : nullptr; + } + const DepthwiseConv2DOptions *builtin_options_as_DepthwiseConv2DOptions() const { + return builtin_options_type() == BuiltinOptions_DepthwiseConv2DOptions ? static_cast(builtin_options()) : nullptr; + } + const ConcatEmbeddingsOptions *builtin_options_as_ConcatEmbeddingsOptions() const { + return builtin_options_type() == BuiltinOptions_ConcatEmbeddingsOptions ? static_cast(builtin_options()) : nullptr; + } + const LSHProjectionOptions *builtin_options_as_LSHProjectionOptions() const { + return builtin_options_type() == BuiltinOptions_LSHProjectionOptions ? static_cast(builtin_options()) : nullptr; + } + const Pool2DOptions *builtin_options_as_Pool2DOptions() const { + return builtin_options_type() == BuiltinOptions_Pool2DOptions ? static_cast(builtin_options()) : nullptr; } const SVDFOptions *builtin_options_as_SVDFOptions() const { - return builtin_options_type() == BuiltinOptions_SVDFOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_SVDFOptions ? static_cast(builtin_options()) : nullptr; } const RNNOptions *builtin_options_as_RNNOptions() const { - return builtin_options_type() == BuiltinOptions_RNNOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_RNNOptions ? static_cast(builtin_options()) : nullptr; } - const FullyConnectedOptions *builtin_options_as_FullyConnectedOptions() - const { - return builtin_options_type() == BuiltinOptions_FullyConnectedOptions - ? static_cast(builtin_options()) - : nullptr; + const FullyConnectedOptions *builtin_options_as_FullyConnectedOptions() const { + return builtin_options_type() == BuiltinOptions_FullyConnectedOptions ? static_cast(builtin_options()) : nullptr; } const SoftmaxOptions *builtin_options_as_SoftmaxOptions() const { - return builtin_options_type() == BuiltinOptions_SoftmaxOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_SoftmaxOptions ? static_cast(builtin_options()) : nullptr; } const ConcatenationOptions *builtin_options_as_ConcatenationOptions() const { - return builtin_options_type() == BuiltinOptions_ConcatenationOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_ConcatenationOptions ? static_cast(builtin_options()) : nullptr; } const AddOptions *builtin_options_as_AddOptions() const { - return builtin_options_type() == BuiltinOptions_AddOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_AddOptions ? static_cast(builtin_options()) : nullptr; } const L2NormOptions *builtin_options_as_L2NormOptions() const { - return builtin_options_type() == BuiltinOptions_L2NormOptions - ? static_cast(builtin_options()) - : nullptr; - } - const LocalResponseNormalizationOptions * - builtin_options_as_LocalResponseNormalizationOptions() const { - return builtin_options_type() == - BuiltinOptions_LocalResponseNormalizationOptions - ? static_cast( - builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_L2NormOptions ? static_cast(builtin_options()) : nullptr; + } + const LocalResponseNormalizationOptions *builtin_options_as_LocalResponseNormalizationOptions() const { + return builtin_options_type() == BuiltinOptions_LocalResponseNormalizationOptions ? static_cast(builtin_options()) : nullptr; } const LSTMOptions *builtin_options_as_LSTMOptions() const { - return builtin_options_type() == BuiltinOptions_LSTMOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_LSTMOptions ? static_cast(builtin_options()) : nullptr; } - const ResizeBilinearOptions *builtin_options_as_ResizeBilinearOptions() - const { - return builtin_options_type() == BuiltinOptions_ResizeBilinearOptions - ? static_cast(builtin_options()) - : nullptr; + const ResizeBilinearOptions *builtin_options_as_ResizeBilinearOptions() const { + return builtin_options_type() == BuiltinOptions_ResizeBilinearOptions ? static_cast(builtin_options()) : nullptr; } const CallOptions *builtin_options_as_CallOptions() const { - return builtin_options_type() == BuiltinOptions_CallOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_CallOptions ? static_cast(builtin_options()) : nullptr; } const ReshapeOptions *builtin_options_as_ReshapeOptions() const { - return builtin_options_type() == BuiltinOptions_ReshapeOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_ReshapeOptions ? static_cast(builtin_options()) : nullptr; } const SkipGramOptions *builtin_options_as_SkipGramOptions() const { - return builtin_options_type() == BuiltinOptions_SkipGramOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_SkipGramOptions ? static_cast(builtin_options()) : nullptr; } const SpaceToDepthOptions *builtin_options_as_SpaceToDepthOptions() const { - return builtin_options_type() == BuiltinOptions_SpaceToDepthOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_SpaceToDepthOptions ? static_cast(builtin_options()) : nullptr; } - const EmbeddingLookupSparseOptions * - builtin_options_as_EmbeddingLookupSparseOptions() const { - return builtin_options_type() == BuiltinOptions_EmbeddingLookupSparseOptions - ? static_cast( - builtin_options()) - : nullptr; + const EmbeddingLookupSparseOptions *builtin_options_as_EmbeddingLookupSparseOptions() const { + return builtin_options_type() == BuiltinOptions_EmbeddingLookupSparseOptions ? static_cast(builtin_options()) : nullptr; } const MulOptions *builtin_options_as_MulOptions() const { - return builtin_options_type() == BuiltinOptions_MulOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_MulOptions ? static_cast(builtin_options()) : nullptr; } const PadOptions *builtin_options_as_PadOptions() const { - return builtin_options_type() == BuiltinOptions_PadOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_PadOptions ? static_cast(builtin_options()) : nullptr; } const GatherOptions *builtin_options_as_GatherOptions() const { - return builtin_options_type() == BuiltinOptions_GatherOptions - ? static_cast(builtin_options()) - : nullptr; + return builtin_options_type() == BuiltinOptions_GatherOptions ? static_cast(builtin_options()) : nullptr; + } + const BatchToSpaceNDOptions *builtin_options_as_BatchToSpaceNDOptions() const { + return builtin_options_type() == BuiltinOptions_BatchToSpaceNDOptions ? static_cast(builtin_options()) : nullptr; + } + const SpaceToBatchNDOptions *builtin_options_as_SpaceToBatchNDOptions() const { + return builtin_options_type() == BuiltinOptions_SpaceToBatchNDOptions ? static_cast(builtin_options()) : nullptr; + } + const TransposeOptions *builtin_options_as_TransposeOptions() const { + return builtin_options_type() == BuiltinOptions_TransposeOptions ? static_cast(builtin_options()) : nullptr; + } + const MeanOptions *builtin_options_as_MeanOptions() const { + return builtin_options_type() == BuiltinOptions_MeanOptions ? static_cast(builtin_options()) : nullptr; } - const BatchToSpaceNDOptions *builtin_options_as_BatchToSpaceNDOptions() - const { - return builtin_options_type() == BuiltinOptions_BatchToSpaceNDOptions - ? static_cast(builtin_options()) - : nullptr; + const SubOptions *builtin_options_as_SubOptions() const { + return builtin_options_type() == BuiltinOptions_SubOptions ? static_cast(builtin_options()) : nullptr; + } + const DivOptions *builtin_options_as_DivOptions() const { + return builtin_options_type() == BuiltinOptions_DivOptions ? static_cast(builtin_options()) : nullptr; + } + const SqueezeOptions *builtin_options_as_SqueezeOptions() const { + return builtin_options_type() == BuiltinOptions_SqueezeOptions ? static_cast(builtin_options()) : nullptr; + } + const SequenceRNNOptions *builtin_options_as_SequenceRNNOptions() const { + return builtin_options_type() == BuiltinOptions_SequenceRNNOptions ? static_cast(builtin_options()) : nullptr; + } + const StridedSliceOptions *builtin_options_as_StridedSliceOptions() const { + return builtin_options_type() == BuiltinOptions_StridedSliceOptions ? static_cast(builtin_options()) : nullptr; + } + const ExpOptions *builtin_options_as_ExpOptions() const { + return builtin_options_type() == BuiltinOptions_ExpOptions ? static_cast(builtin_options()) : nullptr; + } + const TopKV2Options *builtin_options_as_TopKV2Options() const { + return builtin_options_type() == BuiltinOptions_TopKV2Options ? static_cast(builtin_options()) : nullptr; + } + const SplitOptions *builtin_options_as_SplitOptions() const { + return builtin_options_type() == BuiltinOptions_SplitOptions ? static_cast(builtin_options()) : nullptr; } const flatbuffers::Vector *custom_options() const { return GetPointer *>(VT_CUSTOM_OPTIONS); } CustomOptionsFormat custom_options_format() const { - return static_cast( - GetField(VT_CUSTOM_OPTIONS_FORMAT, 0)); + return static_cast(GetField(VT_CUSTOM_OPTIONS_FORMAT, 0)); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_OPCODE_INDEX) && - VerifyOffset(verifier, VT_INPUTS) && verifier.Verify(inputs()) && - VerifyOffset(verifier, VT_OUTPUTS) && verifier.Verify(outputs()) && + VerifyOffset(verifier, VT_INPUTS) && + verifier.Verify(inputs()) && + VerifyOffset(verifier, VT_OUTPUTS) && + verifier.Verify(outputs()) && VerifyField(verifier, VT_BUILTIN_OPTIONS_TYPE) && VerifyOffset(verifier, VT_BUILTIN_OPTIONS) && - VerifyBuiltinOptions(verifier, builtin_options(), - builtin_options_type()) && + VerifyBuiltinOptions(verifier, builtin_options(), builtin_options_type()) && VerifyOffset(verifier, VT_CUSTOM_OPTIONS) && verifier.Verify(custom_options()) && VerifyField(verifier, VT_CUSTOM_OPTIONS_FORMAT) && verifier.EndTable(); } - OperatorT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - OperatorT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const OperatorT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + OperatorT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(OperatorT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const OperatorT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; -template <> -inline const Conv2DOptions *Operator::builtin_options_as() - const { +template<> inline const Conv2DOptions *Operator::builtin_options_as() const { return builtin_options_as_Conv2DOptions(); } -template <> -inline const DepthwiseConv2DOptions * -Operator::builtin_options_as() const { +template<> inline const DepthwiseConv2DOptions *Operator::builtin_options_as() const { return builtin_options_as_DepthwiseConv2DOptions(); } -template <> -inline const ConcatEmbeddingsOptions * -Operator::builtin_options_as() const { +template<> inline const ConcatEmbeddingsOptions *Operator::builtin_options_as() const { return builtin_options_as_ConcatEmbeddingsOptions(); } -template <> -inline const LSHProjectionOptions * -Operator::builtin_options_as() const { +template<> inline const LSHProjectionOptions *Operator::builtin_options_as() const { return builtin_options_as_LSHProjectionOptions(); } -template <> -inline const Pool2DOptions *Operator::builtin_options_as() - const { +template<> inline const Pool2DOptions *Operator::builtin_options_as() const { return builtin_options_as_Pool2DOptions(); } -template <> -inline const SVDFOptions *Operator::builtin_options_as() const { +template<> inline const SVDFOptions *Operator::builtin_options_as() const { return builtin_options_as_SVDFOptions(); } -template <> -inline const RNNOptions *Operator::builtin_options_as() const { +template<> inline const RNNOptions *Operator::builtin_options_as() const { return builtin_options_as_RNNOptions(); } -template <> -inline const FullyConnectedOptions * -Operator::builtin_options_as() const { +template<> inline const FullyConnectedOptions *Operator::builtin_options_as() const { return builtin_options_as_FullyConnectedOptions(); } -template <> -inline const SoftmaxOptions *Operator::builtin_options_as() - const { +template<> inline const SoftmaxOptions *Operator::builtin_options_as() const { return builtin_options_as_SoftmaxOptions(); } -template <> -inline const ConcatenationOptions * -Operator::builtin_options_as() const { +template<> inline const ConcatenationOptions *Operator::builtin_options_as() const { return builtin_options_as_ConcatenationOptions(); } -template <> -inline const AddOptions *Operator::builtin_options_as() const { +template<> inline const AddOptions *Operator::builtin_options_as() const { return builtin_options_as_AddOptions(); } -template <> -inline const L2NormOptions *Operator::builtin_options_as() - const { +template<> inline const L2NormOptions *Operator::builtin_options_as() const { return builtin_options_as_L2NormOptions(); } -template <> -inline const LocalResponseNormalizationOptions * -Operator::builtin_options_as() const { +template<> inline const LocalResponseNormalizationOptions *Operator::builtin_options_as() const { return builtin_options_as_LocalResponseNormalizationOptions(); } -template <> -inline const LSTMOptions *Operator::builtin_options_as() const { +template<> inline const LSTMOptions *Operator::builtin_options_as() const { return builtin_options_as_LSTMOptions(); } -template <> -inline const ResizeBilinearOptions * -Operator::builtin_options_as() const { +template<> inline const ResizeBilinearOptions *Operator::builtin_options_as() const { return builtin_options_as_ResizeBilinearOptions(); } -template <> -inline const CallOptions *Operator::builtin_options_as() const { +template<> inline const CallOptions *Operator::builtin_options_as() const { return builtin_options_as_CallOptions(); } -template <> -inline const ReshapeOptions *Operator::builtin_options_as() - const { +template<> inline const ReshapeOptions *Operator::builtin_options_as() const { return builtin_options_as_ReshapeOptions(); } -template <> -inline const SkipGramOptions *Operator::builtin_options_as() - const { +template<> inline const SkipGramOptions *Operator::builtin_options_as() const { return builtin_options_as_SkipGramOptions(); } -template <> -inline const SpaceToDepthOptions * -Operator::builtin_options_as() const { +template<> inline const SpaceToDepthOptions *Operator::builtin_options_as() const { return builtin_options_as_SpaceToDepthOptions(); } -template <> -inline const EmbeddingLookupSparseOptions * -Operator::builtin_options_as() const { +template<> inline const EmbeddingLookupSparseOptions *Operator::builtin_options_as() const { return builtin_options_as_EmbeddingLookupSparseOptions(); } -template <> -inline const MulOptions *Operator::builtin_options_as() const { +template<> inline const MulOptions *Operator::builtin_options_as() const { return builtin_options_as_MulOptions(); } -template <> -inline const PadOptions *Operator::builtin_options_as() const { +template<> inline const PadOptions *Operator::builtin_options_as() const { return builtin_options_as_PadOptions(); } -template <> -inline const GatherOptions *Operator::builtin_options_as() - const { +template<> inline const GatherOptions *Operator::builtin_options_as() const { return builtin_options_as_GatherOptions(); } -template <> -inline const BatchToSpaceNDOptions * -Operator::builtin_options_as() const { +template<> inline const BatchToSpaceNDOptions *Operator::builtin_options_as() const { return builtin_options_as_BatchToSpaceNDOptions(); } +template<> inline const SpaceToBatchNDOptions *Operator::builtin_options_as() const { + return builtin_options_as_SpaceToBatchNDOptions(); +} + +template<> inline const TransposeOptions *Operator::builtin_options_as() const { + return builtin_options_as_TransposeOptions(); +} + +template<> inline const MeanOptions *Operator::builtin_options_as() const { + return builtin_options_as_MeanOptions(); +} + +template<> inline const SubOptions *Operator::builtin_options_as() const { + return builtin_options_as_SubOptions(); +} + +template<> inline const DivOptions *Operator::builtin_options_as() const { + return builtin_options_as_DivOptions(); +} + +template<> inline const SqueezeOptions *Operator::builtin_options_as() const { + return builtin_options_as_SqueezeOptions(); +} + +template<> inline const SequenceRNNOptions *Operator::builtin_options_as() const { + return builtin_options_as_SequenceRNNOptions(); +} + +template<> inline const StridedSliceOptions *Operator::builtin_options_as() const { + return builtin_options_as_StridedSliceOptions(); +} + +template<> inline const ExpOptions *Operator::builtin_options_as() const { + return builtin_options_as_ExpOptions(); +} + +template<> inline const TopKV2Options *Operator::builtin_options_as() const { + return builtin_options_as_TopKV2Options(); +} + +template<> inline const SplitOptions *Operator::builtin_options_as() const { + return builtin_options_as_SplitOptions(); +} + struct OperatorBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; @@ -3307,21 +3969,19 @@ struct OperatorBuilder { fbb_.AddOffset(Operator::VT_OUTPUTS, outputs); } void add_builtin_options_type(BuiltinOptions builtin_options_type) { - fbb_.AddElement(Operator::VT_BUILTIN_OPTIONS_TYPE, - static_cast(builtin_options_type), 0); + fbb_.AddElement(Operator::VT_BUILTIN_OPTIONS_TYPE, static_cast(builtin_options_type), 0); } void add_builtin_options(flatbuffers::Offset builtin_options) { fbb_.AddOffset(Operator::VT_BUILTIN_OPTIONS, builtin_options); } - void add_custom_options( - flatbuffers::Offset> custom_options) { + void add_custom_options(flatbuffers::Offset> custom_options) { fbb_.AddOffset(Operator::VT_CUSTOM_OPTIONS, custom_options); } void add_custom_options_format(CustomOptionsFormat custom_options_format) { - fbb_.AddElement(Operator::VT_CUSTOM_OPTIONS_FORMAT, - static_cast(custom_options_format), 0); + fbb_.AddElement(Operator::VT_CUSTOM_OPTIONS_FORMAT, static_cast(custom_options_format), 0); } - explicit OperatorBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { + explicit OperatorBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } OperatorBuilder &operator=(const OperatorBuilder &); @@ -3333,14 +3993,14 @@ struct OperatorBuilder { }; inline flatbuffers::Offset CreateOperator( - flatbuffers::FlatBufferBuilder &_fbb, uint32_t opcode_index = 0, + flatbuffers::FlatBufferBuilder &_fbb, + uint32_t opcode_index = 0, flatbuffers::Offset> inputs = 0, flatbuffers::Offset> outputs = 0, BuiltinOptions builtin_options_type = BuiltinOptions_NONE, flatbuffers::Offset builtin_options = 0, flatbuffers::Offset> custom_options = 0, - CustomOptionsFormat custom_options_format = - CustomOptionsFormat_FLEXBUFFERS) { + CustomOptionsFormat custom_options_format = CustomOptionsFormat_FLEXBUFFERS) { OperatorBuilder builder_(_fbb); builder_.add_custom_options(custom_options); builder_.add_builtin_options(builtin_options); @@ -3353,25 +4013,26 @@ inline flatbuffers::Offset CreateOperator( } inline flatbuffers::Offset CreateOperatorDirect( - flatbuffers::FlatBufferBuilder &_fbb, uint32_t opcode_index = 0, + flatbuffers::FlatBufferBuilder &_fbb, + uint32_t opcode_index = 0, const std::vector *inputs = nullptr, const std::vector *outputs = nullptr, BuiltinOptions builtin_options_type = BuiltinOptions_NONE, flatbuffers::Offset builtin_options = 0, const std::vector *custom_options = nullptr, - CustomOptionsFormat custom_options_format = - CustomOptionsFormat_FLEXBUFFERS) { + CustomOptionsFormat custom_options_format = CustomOptionsFormat_FLEXBUFFERS) { return tflite::CreateOperator( - _fbb, opcode_index, inputs ? _fbb.CreateVector(*inputs) : 0, - outputs ? _fbb.CreateVector(*outputs) : 0, builtin_options_type, + _fbb, + opcode_index, + inputs ? _fbb.CreateVector(*inputs) : 0, + outputs ? _fbb.CreateVector(*outputs) : 0, + builtin_options_type, builtin_options, custom_options ? _fbb.CreateVector(*custom_options) : 0, custom_options_format); } -flatbuffers::Offset CreateOperator( - flatbuffers::FlatBufferBuilder &_fbb, const OperatorT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateOperator(flatbuffers::FlatBufferBuilder &_fbb, const OperatorT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct SubGraphT : public flatbuffers::NativeTable { typedef SubGraph TableType; @@ -3380,7 +4041,8 @@ struct SubGraphT : public flatbuffers::NativeTable { std::vector outputs; std::vector> operators; std::string name; - SubGraphT() {} + SubGraphT() { + } }; struct SubGraph FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { @@ -3393,8 +4055,7 @@ struct SubGraph FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { VT_NAME = 12 }; const flatbuffers::Vector> *tensors() const { - return GetPointer> *>( - VT_TENSORS); + return GetPointer> *>(VT_TENSORS); } const flatbuffers::Vector *inputs() const { return GetPointer *>(VT_INPUTS); @@ -3403,41 +4064,36 @@ struct SubGraph FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { return GetPointer *>(VT_OUTPUTS); } const flatbuffers::Vector> *operators() const { - return GetPointer< - const flatbuffers::Vector> *>( - VT_OPERATORS); + return GetPointer> *>(VT_OPERATORS); } const flatbuffers::String *name() const { return GetPointer(VT_NAME); } bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && VerifyOffset(verifier, VT_TENSORS) && + return VerifyTableStart(verifier) && + VerifyOffset(verifier, VT_TENSORS) && verifier.Verify(tensors()) && verifier.VerifyVectorOfTables(tensors()) && - VerifyOffset(verifier, VT_INPUTS) && verifier.Verify(inputs()) && - VerifyOffset(verifier, VT_OUTPUTS) && verifier.Verify(outputs()) && + VerifyOffset(verifier, VT_INPUTS) && + verifier.Verify(inputs()) && + VerifyOffset(verifier, VT_OUTPUTS) && + verifier.Verify(outputs()) && VerifyOffset(verifier, VT_OPERATORS) && verifier.Verify(operators()) && verifier.VerifyVectorOfTables(operators()) && - VerifyOffset(verifier, VT_NAME) && verifier.Verify(name()) && + VerifyOffset(verifier, VT_NAME) && + verifier.Verify(name()) && verifier.EndTable(); } - SubGraphT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo( - SubGraphT *_o, - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SubGraphT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + SubGraphT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(SubGraphT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const SubGraphT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct SubGraphBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - void add_tensors( - flatbuffers::Offset>> - tensors) { + void add_tensors(flatbuffers::Offset>> tensors) { fbb_.AddOffset(SubGraph::VT_TENSORS, tensors); } void add_inputs(flatbuffers::Offset> inputs) { @@ -3446,15 +4102,14 @@ struct SubGraphBuilder { void add_outputs(flatbuffers::Offset> outputs) { fbb_.AddOffset(SubGraph::VT_OUTPUTS, outputs); } - void add_operators( - flatbuffers::Offset>> - operators) { + void add_operators(flatbuffers::Offset>> operators) { fbb_.AddOffset(SubGraph::VT_OPERATORS, operators); } void add_name(flatbuffers::Offset name) { fbb_.AddOffset(SubGraph::VT_NAME, name); } - explicit SubGraphBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { + explicit SubGraphBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } SubGraphBuilder &operator=(const SubGraphBuilder &); @@ -3467,12 +4122,10 @@ struct SubGraphBuilder { inline flatbuffers::Offset CreateSubGraph( flatbuffers::FlatBufferBuilder &_fbb, - flatbuffers::Offset>> - tensors = 0, + flatbuffers::Offset>> tensors = 0, flatbuffers::Offset> inputs = 0, flatbuffers::Offset> outputs = 0, - flatbuffers::Offset>> - operators = 0, + flatbuffers::Offset>> operators = 0, flatbuffers::Offset name = 0) { SubGraphBuilder builder_(_fbb); builder_.add_name(name); @@ -3495,38 +4148,36 @@ inline flatbuffers::Offset CreateSubGraphDirect( tensors ? _fbb.CreateVector>(*tensors) : 0, inputs ? _fbb.CreateVector(*inputs) : 0, outputs ? _fbb.CreateVector(*outputs) : 0, - operators ? _fbb.CreateVector>(*operators) - : 0, + operators ? _fbb.CreateVector>(*operators) : 0, name ? _fbb.CreateString(name) : 0); } -flatbuffers::Offset CreateSubGraph( - flatbuffers::FlatBufferBuilder &_fbb, const SubGraphT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateSubGraph(flatbuffers::FlatBufferBuilder &_fbb, const SubGraphT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct BufferT : public flatbuffers::NativeTable { typedef Buffer TableType; std::vector data; - BufferT() {} + BufferT() { + } }; struct Buffer FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { typedef BufferT NativeTableType; - enum { VT_DATA = 4 }; + enum { + VT_DATA = 4 + }; const flatbuffers::Vector *data() const { return GetPointer *>(VT_DATA); } bool Verify(flatbuffers::Verifier &verifier) const { - return VerifyTableStart(verifier) && VerifyOffset(verifier, VT_DATA) && - verifier.Verify(data()) && verifier.EndTable(); + return VerifyTableStart(verifier) && + VerifyOffset(verifier, VT_DATA) && + verifier.Verify(data()) && + verifier.EndTable(); } - BufferT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo(BufferT *_o, const flatbuffers::resolver_function_t *_resolver = - nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const BufferT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + BufferT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(BufferT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const BufferT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; struct BufferBuilder { @@ -3535,7 +4186,8 @@ struct BufferBuilder { void add_data(flatbuffers::Offset> data) { fbb_.AddOffset(Buffer::VT_DATA, data); } - explicit BufferBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { + explicit BufferBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { start_ = fbb_.StartTable(); } BufferBuilder &operator=(const BufferBuilder &); @@ -3557,13 +4209,12 @@ inline flatbuffers::Offset CreateBuffer( inline flatbuffers::Offset CreateBufferDirect( flatbuffers::FlatBufferBuilder &_fbb, const std::vector *data = nullptr) { - return tflite::CreateBuffer(_fbb, - data ? _fbb.CreateVector(*data) : 0); + return tflite::CreateBuffer( + _fbb, + data ? _fbb.CreateVector(*data) : 0); } -flatbuffers::Offset CreateBuffer( - flatbuffers::FlatBufferBuilder &_fbb, const BufferT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateBuffer(flatbuffers::FlatBufferBuilder &_fbb, const BufferT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct ModelT : public flatbuffers::NativeTable { typedef Model TableType; @@ -3572,7 +4223,9 @@ struct ModelT : public flatbuffers::NativeTable { std::vector> subgraphs; std::string description; std::vector> buffers; - ModelT() : version(0) {} + ModelT() + : version(0) { + } }; struct Model FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { @@ -3584,24 +4237,20 @@ struct Model FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { VT_DESCRIPTION = 10, VT_BUFFERS = 12 }; - uint32_t version() const { return GetField(VT_VERSION, 0); } - const flatbuffers::Vector> *operator_codes() - const { - return GetPointer< - const flatbuffers::Vector> *>( - VT_OPERATOR_CODES); + uint32_t version() const { + return GetField(VT_VERSION, 0); + } + const flatbuffers::Vector> *operator_codes() const { + return GetPointer> *>(VT_OPERATOR_CODES); } const flatbuffers::Vector> *subgraphs() const { - return GetPointer< - const flatbuffers::Vector> *>( - VT_SUBGRAPHS); + return GetPointer> *>(VT_SUBGRAPHS); } const flatbuffers::String *description() const { return GetPointer(VT_DESCRIPTION); } const flatbuffers::Vector> *buffers() const { - return GetPointer> *>( - VT_BUFFERS); + return GetPointer> *>(VT_BUFFERS); } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && @@ -3614,1760 +4263,1334 @@ struct Model FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { verifier.VerifyVectorOfTables(subgraphs()) && VerifyOffset(verifier, VT_DESCRIPTION) && verifier.Verify(description()) && - VerifyOffset(verifier, VT_BUFFERS) && verifier.Verify(buffers()) && - verifier.VerifyVectorOfTables(buffers()) && verifier.EndTable(); + VerifyOffset(verifier, VT_BUFFERS) && + verifier.Verify(buffers()) && + verifier.VerifyVectorOfTables(buffers()) && + verifier.EndTable(); } - ModelT *UnPack( - const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo(ModelT *_o, const flatbuffers::resolver_function_t *_resolver = - nullptr) const; - static flatbuffers::Offset Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ModelT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); + ModelT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(ModelT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const ModelT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; -struct ModelBuilder { - flatbuffers::FlatBufferBuilder &fbb_; - flatbuffers::uoffset_t start_; - void add_version(uint32_t version) { - fbb_.AddElement(Model::VT_VERSION, version, 0); - } - void add_operator_codes( - flatbuffers::Offset< - flatbuffers::Vector>> - operator_codes) { - fbb_.AddOffset(Model::VT_OPERATOR_CODES, operator_codes); - } - void add_subgraphs( - flatbuffers::Offset>> - subgraphs) { - fbb_.AddOffset(Model::VT_SUBGRAPHS, subgraphs); - } - void add_description(flatbuffers::Offset description) { - fbb_.AddOffset(Model::VT_DESCRIPTION, description); - } - void add_buffers( - flatbuffers::Offset>> - buffers) { - fbb_.AddOffset(Model::VT_BUFFERS, buffers); - } - explicit ModelBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { - start_ = fbb_.StartTable(); - } - ModelBuilder &operator=(const ModelBuilder &); - flatbuffers::Offset Finish() { - const auto end = fbb_.EndTable(start_); - auto o = flatbuffers::Offset(end); - return o; - } -}; +struct ModelBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_version(uint32_t version) { + fbb_.AddElement(Model::VT_VERSION, version, 0); + } + void add_operator_codes(flatbuffers::Offset>> operator_codes) { + fbb_.AddOffset(Model::VT_OPERATOR_CODES, operator_codes); + } + void add_subgraphs(flatbuffers::Offset>> subgraphs) { + fbb_.AddOffset(Model::VT_SUBGRAPHS, subgraphs); + } + void add_description(flatbuffers::Offset description) { + fbb_.AddOffset(Model::VT_DESCRIPTION, description); + } + void add_buffers(flatbuffers::Offset>> buffers) { + fbb_.AddOffset(Model::VT_BUFFERS, buffers); + } + explicit ModelBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + ModelBuilder &operator=(const ModelBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateModel( + flatbuffers::FlatBufferBuilder &_fbb, + uint32_t version = 0, + flatbuffers::Offset>> operator_codes = 0, + flatbuffers::Offset>> subgraphs = 0, + flatbuffers::Offset description = 0, + flatbuffers::Offset>> buffers = 0) { + ModelBuilder builder_(_fbb); + builder_.add_buffers(buffers); + builder_.add_description(description); + builder_.add_subgraphs(subgraphs); + builder_.add_operator_codes(operator_codes); + builder_.add_version(version); + return builder_.Finish(); +} + +inline flatbuffers::Offset CreateModelDirect( + flatbuffers::FlatBufferBuilder &_fbb, + uint32_t version = 0, + const std::vector> *operator_codes = nullptr, + const std::vector> *subgraphs = nullptr, + const char *description = nullptr, + const std::vector> *buffers = nullptr) { + return tflite::CreateModel( + _fbb, + version, + operator_codes ? _fbb.CreateVector>(*operator_codes) : 0, + subgraphs ? _fbb.CreateVector>(*subgraphs) : 0, + description ? _fbb.CreateString(description) : 0, + buffers ? _fbb.CreateVector>(*buffers) : 0); +} + +flatbuffers::Offset CreateModel(flatbuffers::FlatBufferBuilder &_fbb, const ModelT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +inline QuantizationParametersT *QuantizationParameters::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new QuantizationParametersT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void QuantizationParameters::UnPackTo(QuantizationParametersT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = min(); if (_e) { _o->min.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->min[_i] = _e->Get(_i); } } }; + { auto _e = max(); if (_e) { _o->max.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->max[_i] = _e->Get(_i); } } }; + { auto _e = scale(); if (_e) { _o->scale.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->scale[_i] = _e->Get(_i); } } }; + { auto _e = zero_point(); if (_e) { _o->zero_point.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->zero_point[_i] = _e->Get(_i); } } }; +} + +inline flatbuffers::Offset QuantizationParameters::Pack(flatbuffers::FlatBufferBuilder &_fbb, const QuantizationParametersT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateQuantizationParameters(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateQuantizationParameters(flatbuffers::FlatBufferBuilder &_fbb, const QuantizationParametersT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const QuantizationParametersT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _min = _o->min.size() ? _fbb.CreateVector(_o->min) : 0; + auto _max = _o->max.size() ? _fbb.CreateVector(_o->max) : 0; + auto _scale = _o->scale.size() ? _fbb.CreateVector(_o->scale) : 0; + auto _zero_point = _o->zero_point.size() ? _fbb.CreateVector(_o->zero_point) : 0; + return tflite::CreateQuantizationParameters( + _fbb, + _min, + _max, + _scale, + _zero_point); +} + +inline TensorT *Tensor::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new TensorT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void Tensor::UnPackTo(TensorT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = shape(); if (_e) { _o->shape.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->shape[_i] = _e->Get(_i); } } }; + { auto _e = type(); _o->type = _e; }; + { auto _e = buffer(); _o->buffer = _e; }; + { auto _e = name(); if (_e) _o->name = _e->str(); }; + { auto _e = quantization(); if (_e) _o->quantization = std::unique_ptr(_e->UnPack(_resolver)); }; +} + +inline flatbuffers::Offset Tensor::Pack(flatbuffers::FlatBufferBuilder &_fbb, const TensorT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateTensor(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateTensor(flatbuffers::FlatBufferBuilder &_fbb, const TensorT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const TensorT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _shape = _o->shape.size() ? _fbb.CreateVector(_o->shape) : 0; + auto _type = _o->type; + auto _buffer = _o->buffer; + auto _name = _o->name.empty() ? 0 : _fbb.CreateString(_o->name); + auto _quantization = _o->quantization ? CreateQuantizationParameters(_fbb, _o->quantization.get(), _rehasher) : 0; + return tflite::CreateTensor( + _fbb, + _shape, + _type, + _buffer, + _name, + _quantization); +} + +inline Conv2DOptionsT *Conv2DOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new Conv2DOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void Conv2DOptions::UnPackTo(Conv2DOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = padding(); _o->padding = _e; }; + { auto _e = stride_w(); _o->stride_w = _e; }; + { auto _e = stride_h(); _o->stride_h = _e; }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; +} + +inline flatbuffers::Offset Conv2DOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateConv2DOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateConv2DOptions(flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const Conv2DOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _padding = _o->padding; + auto _stride_w = _o->stride_w; + auto _stride_h = _o->stride_h; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateConv2DOptions( + _fbb, + _padding, + _stride_w, + _stride_h, + _fused_activation_function); +} + +inline Pool2DOptionsT *Pool2DOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new Pool2DOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void Pool2DOptions::UnPackTo(Pool2DOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = padding(); _o->padding = _e; }; + { auto _e = stride_w(); _o->stride_w = _e; }; + { auto _e = stride_h(); _o->stride_h = _e; }; + { auto _e = filter_width(); _o->filter_width = _e; }; + { auto _e = filter_height(); _o->filter_height = _e; }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; +} + +inline flatbuffers::Offset Pool2DOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const Pool2DOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreatePool2DOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreatePool2DOptions(flatbuffers::FlatBufferBuilder &_fbb, const Pool2DOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const Pool2DOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _padding = _o->padding; + auto _stride_w = _o->stride_w; + auto _stride_h = _o->stride_h; + auto _filter_width = _o->filter_width; + auto _filter_height = _o->filter_height; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreatePool2DOptions( + _fbb, + _padding, + _stride_w, + _stride_h, + _filter_width, + _filter_height, + _fused_activation_function); +} + +inline DepthwiseConv2DOptionsT *DepthwiseConv2DOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new DepthwiseConv2DOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void DepthwiseConv2DOptions::UnPackTo(DepthwiseConv2DOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = padding(); _o->padding = _e; }; + { auto _e = stride_w(); _o->stride_w = _e; }; + { auto _e = stride_h(); _o->stride_h = _e; }; + { auto _e = depth_multiplier(); _o->depth_multiplier = _e; }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; +} + +inline flatbuffers::Offset DepthwiseConv2DOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const DepthwiseConv2DOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateDepthwiseConv2DOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateDepthwiseConv2DOptions(flatbuffers::FlatBufferBuilder &_fbb, const DepthwiseConv2DOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const DepthwiseConv2DOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _padding = _o->padding; + auto _stride_w = _o->stride_w; + auto _stride_h = _o->stride_h; + auto _depth_multiplier = _o->depth_multiplier; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateDepthwiseConv2DOptions( + _fbb, + _padding, + _stride_w, + _stride_h, + _depth_multiplier, + _fused_activation_function); +} + +inline ConcatEmbeddingsOptionsT *ConcatEmbeddingsOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new ConcatEmbeddingsOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void ConcatEmbeddingsOptions::UnPackTo(ConcatEmbeddingsOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = num_channels(); _o->num_channels = _e; }; + { auto _e = num_columns_per_channel(); if (_e) { _o->num_columns_per_channel.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->num_columns_per_channel[_i] = _e->Get(_i); } } }; + { auto _e = embedding_dim_per_channel(); if (_e) { _o->embedding_dim_per_channel.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->embedding_dim_per_channel[_i] = _e->Get(_i); } } }; +} + +inline flatbuffers::Offset ConcatEmbeddingsOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const ConcatEmbeddingsOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateConcatEmbeddingsOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateConcatEmbeddingsOptions(flatbuffers::FlatBufferBuilder &_fbb, const ConcatEmbeddingsOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const ConcatEmbeddingsOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _num_channels = _o->num_channels; + auto _num_columns_per_channel = _o->num_columns_per_channel.size() ? _fbb.CreateVector(_o->num_columns_per_channel) : 0; + auto _embedding_dim_per_channel = _o->embedding_dim_per_channel.size() ? _fbb.CreateVector(_o->embedding_dim_per_channel) : 0; + return tflite::CreateConcatEmbeddingsOptions( + _fbb, + _num_channels, + _num_columns_per_channel, + _embedding_dim_per_channel); +} + +inline LSHProjectionOptionsT *LSHProjectionOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new LSHProjectionOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void LSHProjectionOptions::UnPackTo(LSHProjectionOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = type(); _o->type = _e; }; +} + +inline flatbuffers::Offset LSHProjectionOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateLSHProjectionOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateLSHProjectionOptions(flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const LSHProjectionOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _type = _o->type; + return tflite::CreateLSHProjectionOptions( + _fbb, + _type); +} + +inline SVDFOptionsT *SVDFOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SVDFOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void SVDFOptions::UnPackTo(SVDFOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = rank(); _o->rank = _e; }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; +} + +inline flatbuffers::Offset SVDFOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SVDFOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSVDFOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateSVDFOptions(flatbuffers::FlatBufferBuilder &_fbb, const SVDFOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SVDFOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _rank = _o->rank; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateSVDFOptions( + _fbb, + _rank, + _fused_activation_function); +} + +inline RNNOptionsT *RNNOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new RNNOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void RNNOptions::UnPackTo(RNNOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; +} + +inline flatbuffers::Offset RNNOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const RNNOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateRNNOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateRNNOptions(flatbuffers::FlatBufferBuilder &_fbb, const RNNOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const RNNOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateRNNOptions( + _fbb, + _fused_activation_function); +} + +inline SequenceRNNOptionsT *SequenceRNNOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SequenceRNNOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void SequenceRNNOptions::UnPackTo(SequenceRNNOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = time_major(); _o->time_major = _e; }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; +} + +inline flatbuffers::Offset SequenceRNNOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SequenceRNNOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSequenceRNNOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateSequenceRNNOptions(flatbuffers::FlatBufferBuilder &_fbb, const SequenceRNNOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SequenceRNNOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _time_major = _o->time_major; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateSequenceRNNOptions( + _fbb, + _time_major, + _fused_activation_function); +} + +inline BidirectionalSequenceRNNOptionsT *BidirectionalSequenceRNNOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new BidirectionalSequenceRNNOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} -inline flatbuffers::Offset CreateModel( - flatbuffers::FlatBufferBuilder &_fbb, uint32_t version = 0, - flatbuffers::Offset>> - operator_codes = 0, - flatbuffers::Offset>> - subgraphs = 0, - flatbuffers::Offset description = 0, - flatbuffers::Offset>> - buffers = 0) { - ModelBuilder builder_(_fbb); - builder_.add_buffers(buffers); - builder_.add_description(description); - builder_.add_subgraphs(subgraphs); - builder_.add_operator_codes(operator_codes); - builder_.add_version(version); - return builder_.Finish(); +inline void BidirectionalSequenceRNNOptions::UnPackTo(BidirectionalSequenceRNNOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = time_major(); _o->time_major = _e; }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset CreateModelDirect( - flatbuffers::FlatBufferBuilder &_fbb, uint32_t version = 0, - const std::vector> *operator_codes = - nullptr, - const std::vector> *subgraphs = nullptr, - const char *description = nullptr, - const std::vector> *buffers = nullptr) { - return tflite::CreateModel( - _fbb, version, - operator_codes ? _fbb.CreateVector>( - *operator_codes) - : 0, - subgraphs ? _fbb.CreateVector>(*subgraphs) - : 0, - description ? _fbb.CreateString(description) : 0, - buffers ? _fbb.CreateVector>(*buffers) : 0); +inline flatbuffers::Offset BidirectionalSequenceRNNOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const BidirectionalSequenceRNNOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateBidirectionalSequenceRNNOptions(_fbb, _o, _rehasher); } -flatbuffers::Offset CreateModel( - flatbuffers::FlatBufferBuilder &_fbb, const ModelT *_o, - const flatbuffers::rehasher_function_t *_rehasher = nullptr); +inline flatbuffers::Offset CreateBidirectionalSequenceRNNOptions(flatbuffers::FlatBufferBuilder &_fbb, const BidirectionalSequenceRNNOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const BidirectionalSequenceRNNOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _time_major = _o->time_major; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateBidirectionalSequenceRNNOptions( + _fbb, + _time_major, + _fused_activation_function); +} -inline QuantizationParametersT *QuantizationParameters::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new QuantizationParametersT(); +inline FullyConnectedOptionsT *FullyConnectedOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new FullyConnectedOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void QuantizationParameters::UnPackTo( - QuantizationParametersT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void FullyConnectedOptions::UnPackTo(FullyConnectedOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = min(); - if (_e) { - _o->min.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->min[_i] = _e->Get(_i); - } - } - }; - { - auto _e = max(); - if (_e) { - _o->max.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->max[_i] = _e->Get(_i); - } - } - }; - { - auto _e = scale(); - if (_e) { - _o->scale.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->scale[_i] = _e->Get(_i); - } - } - }; - { - auto _e = zero_point(); - if (_e) { - _o->zero_point.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->zero_point[_i] = _e->Get(_i); - } - } - }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset QuantizationParameters::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const QuantizationParametersT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateQuantizationParameters(_fbb, _o, _rehasher); +inline flatbuffers::Offset FullyConnectedOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateFullyConnectedOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateQuantizationParameters( - flatbuffers::FlatBufferBuilder &_fbb, const QuantizationParametersT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateFullyConnectedOptions(flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const QuantizationParametersT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _min = _o->min.size() ? _fbb.CreateVector(_o->min) : 0; - auto _max = _o->max.size() ? _fbb.CreateVector(_o->max) : 0; - auto _scale = _o->scale.size() ? _fbb.CreateVector(_o->scale) : 0; - auto _zero_point = - _o->zero_point.size() ? _fbb.CreateVector(_o->zero_point) : 0; - return tflite::CreateQuantizationParameters(_fbb, _min, _max, _scale, - _zero_point); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const FullyConnectedOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateFullyConnectedOptions( + _fbb, + _fused_activation_function); } -inline TensorT *Tensor::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new TensorT(); +inline SoftmaxOptionsT *SoftmaxOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SoftmaxOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void Tensor::UnPackTo( - TensorT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void SoftmaxOptions::UnPackTo(SoftmaxOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = shape(); - if (_e) { - _o->shape.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->shape[_i] = _e->Get(_i); - } - } - }; - { - auto _e = type(); - _o->type = _e; - }; - { - auto _e = buffer(); - _o->buffer = _e; - }; - { - auto _e = name(); - if (_e) _o->name = _e->str(); - }; - { - auto _e = quantization(); - if (_e) - _o->quantization = - std::unique_ptr(_e->UnPack(_resolver)); - }; + { auto _e = beta(); _o->beta = _e; }; } -inline flatbuffers::Offset Tensor::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const TensorT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateTensor(_fbb, _o, _rehasher); +inline flatbuffers::Offset SoftmaxOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSoftmaxOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateTensor( - flatbuffers::FlatBufferBuilder &_fbb, const TensorT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateSoftmaxOptions(flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const TensorT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _shape = _o->shape.size() ? _fbb.CreateVector(_o->shape) : 0; - auto _type = _o->type; - auto _buffer = _o->buffer; - auto _name = _o->name.empty() ? 0 : _fbb.CreateString(_o->name); - auto _quantization = _o->quantization - ? CreateQuantizationParameters( - _fbb, _o->quantization.get(), _rehasher) - : 0; - return tflite::CreateTensor(_fbb, _shape, _type, _buffer, _name, - _quantization); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SoftmaxOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _beta = _o->beta; + return tflite::CreateSoftmaxOptions( + _fbb, + _beta); } -inline Conv2DOptionsT *Conv2DOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new Conv2DOptionsT(); +inline ConcatenationOptionsT *ConcatenationOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new ConcatenationOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void Conv2DOptions::UnPackTo( - Conv2DOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void ConcatenationOptions::UnPackTo(ConcatenationOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = padding(); - _o->padding = _e; - }; - { - auto _e = stride_w(); - _o->stride_w = _e; - }; - { - auto _e = stride_h(); - _o->stride_h = _e; - }; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = axis(); _o->axis = _e; }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset Conv2DOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateConv2DOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset ConcatenationOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateConcatenationOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateConv2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateConcatenationOptions(flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const Conv2DOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _padding = _o->padding; - auto _stride_w = _o->stride_w; - auto _stride_h = _o->stride_h; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const ConcatenationOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _axis = _o->axis; auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateConv2DOptions(_fbb, _padding, _stride_w, _stride_h, - _fused_activation_function); + return tflite::CreateConcatenationOptions( + _fbb, + _axis, + _fused_activation_function); } -inline Pool2DOptionsT *Pool2DOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new Pool2DOptionsT(); +inline AddOptionsT *AddOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new AddOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void Pool2DOptions::UnPackTo( - Pool2DOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void AddOptions::UnPackTo(AddOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = padding(); - _o->padding = _e; - }; - { - auto _e = stride_w(); - _o->stride_w = _e; - }; - { - auto _e = stride_h(); - _o->stride_h = _e; - }; - { - auto _e = filter_width(); - _o->filter_width = _e; - }; - { - auto _e = filter_height(); - _o->filter_height = _e; - }; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset Pool2DOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const Pool2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreatePool2DOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset AddOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateAddOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreatePool2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, const Pool2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateAddOptions(flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const Pool2DOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _padding = _o->padding; - auto _stride_w = _o->stride_w; - auto _stride_h = _o->stride_h; - auto _filter_width = _o->filter_width; - auto _filter_height = _o->filter_height; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const AddOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreatePool2DOptions(_fbb, _padding, _stride_w, _stride_h, - _filter_width, _filter_height, - _fused_activation_function); + return tflite::CreateAddOptions( + _fbb, + _fused_activation_function); } -inline DepthwiseConv2DOptionsT *DepthwiseConv2DOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new DepthwiseConv2DOptionsT(); +inline MulOptionsT *MulOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new MulOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void DepthwiseConv2DOptions::UnPackTo( - DepthwiseConv2DOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void MulOptions::UnPackTo(MulOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = padding(); - _o->padding = _e; - }; - { - auto _e = stride_w(); - _o->stride_w = _e; - }; - { - auto _e = stride_h(); - _o->stride_h = _e; - }; - { - auto _e = depth_multiplier(); - _o->depth_multiplier = _e; - }; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset DepthwiseConv2DOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const DepthwiseConv2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateDepthwiseConv2DOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset MulOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateMulOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateDepthwiseConv2DOptions( - flatbuffers::FlatBufferBuilder &_fbb, const DepthwiseConv2DOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateMulOptions(flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const DepthwiseConv2DOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _padding = _o->padding; - auto _stride_w = _o->stride_w; - auto _stride_h = _o->stride_h; - auto _depth_multiplier = _o->depth_multiplier; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const MulOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateDepthwiseConv2DOptions(_fbb, _padding, _stride_w, - _stride_h, _depth_multiplier, - _fused_activation_function); + return tflite::CreateMulOptions( + _fbb, + _fused_activation_function); } -inline ConcatEmbeddingsOptionsT *ConcatEmbeddingsOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new ConcatEmbeddingsOptionsT(); +inline L2NormOptionsT *L2NormOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new L2NormOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void ConcatEmbeddingsOptions::UnPackTo( - ConcatEmbeddingsOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void L2NormOptions::UnPackTo(L2NormOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = num_channels(); - _o->num_channels = _e; - }; - { - auto _e = num_columns_per_channel(); - if (_e) { - _o->num_columns_per_channel.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->num_columns_per_channel[_i] = _e->Get(_i); - } - } - }; - { - auto _e = embedding_dim_per_channel(); - if (_e) { - _o->embedding_dim_per_channel.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->embedding_dim_per_channel[_i] = _e->Get(_i); - } - } - }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset -ConcatEmbeddingsOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatEmbeddingsOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateConcatEmbeddingsOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset L2NormOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateL2NormOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset -CreateConcatEmbeddingsOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatEmbeddingsOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateL2NormOptions(flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const ConcatEmbeddingsOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _num_channels = _o->num_channels; - auto _num_columns_per_channel = - _o->num_columns_per_channel.size() - ? _fbb.CreateVector(_o->num_columns_per_channel) - : 0; - auto _embedding_dim_per_channel = - _o->embedding_dim_per_channel.size() - ? _fbb.CreateVector(_o->embedding_dim_per_channel) - : 0; - return tflite::CreateConcatEmbeddingsOptions(_fbb, _num_channels, - _num_columns_per_channel, - _embedding_dim_per_channel); -} - -inline LSHProjectionOptionsT *LSHProjectionOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new LSHProjectionOptionsT(); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const L2NormOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateL2NormOptions( + _fbb, + _fused_activation_function); +} + +inline LocalResponseNormalizationOptionsT *LocalResponseNormalizationOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new LocalResponseNormalizationOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void LSHProjectionOptions::UnPackTo( - LSHProjectionOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void LocalResponseNormalizationOptions::UnPackTo(LocalResponseNormalizationOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = type(); - _o->type = _e; - }; + { auto _e = radius(); _o->radius = _e; }; + { auto _e = bias(); _o->bias = _e; }; + { auto _e = alpha(); _o->alpha = _e; }; + { auto _e = beta(); _o->beta = _e; }; } -inline flatbuffers::Offset LSHProjectionOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateLSHProjectionOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset LocalResponseNormalizationOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const LocalResponseNormalizationOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateLocalResponseNormalizationOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateLSHProjectionOptions( - flatbuffers::FlatBufferBuilder &_fbb, const LSHProjectionOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateLocalResponseNormalizationOptions(flatbuffers::FlatBufferBuilder &_fbb, const LocalResponseNormalizationOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const LSHProjectionOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _type = _o->type; - return tflite::CreateLSHProjectionOptions(_fbb, _type); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const LocalResponseNormalizationOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _radius = _o->radius; + auto _bias = _o->bias; + auto _alpha = _o->alpha; + auto _beta = _o->beta; + return tflite::CreateLocalResponseNormalizationOptions( + _fbb, + _radius, + _bias, + _alpha, + _beta); } -inline SVDFOptionsT *SVDFOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new SVDFOptionsT(); +inline LSTMOptionsT *LSTMOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new LSTMOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SVDFOptions::UnPackTo( - SVDFOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void LSTMOptions::UnPackTo(LSTMOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = rank(); - _o->rank = _e; - }; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; + { auto _e = cell_clip(); _o->cell_clip = _e; }; + { auto _e = proj_clip(); _o->proj_clip = _e; }; } -inline flatbuffers::Offset SVDFOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SVDFOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateSVDFOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset LSTMOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateLSTMOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateSVDFOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SVDFOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateLSTMOptions(flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const SVDFOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _rank = _o->rank; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const LSTMOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateSVDFOptions(_fbb, _rank, _fused_activation_function); + auto _cell_clip = _o->cell_clip; + auto _proj_clip = _o->proj_clip; + return tflite::CreateLSTMOptions( + _fbb, + _fused_activation_function, + _cell_clip, + _proj_clip); } -inline RNNOptionsT *RNNOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new RNNOptionsT(); +inline ResizeBilinearOptionsT *ResizeBilinearOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new ResizeBilinearOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void RNNOptions::UnPackTo( - RNNOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void ResizeBilinearOptions::UnPackTo(ResizeBilinearOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = align_corners(); _o->align_corners = _e; }; } -inline flatbuffers::Offset RNNOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const RNNOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateRNNOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset ResizeBilinearOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateResizeBilinearOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateRNNOptions( - flatbuffers::FlatBufferBuilder &_fbb, const RNNOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateResizeBilinearOptions(flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const RNNOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateRNNOptions(_fbb, _fused_activation_function); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const ResizeBilinearOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _align_corners = _o->align_corners; + return tflite::CreateResizeBilinearOptions( + _fbb, + _align_corners); } -inline FullyConnectedOptionsT *FullyConnectedOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new FullyConnectedOptionsT(); +inline CallOptionsT *CallOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new CallOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void FullyConnectedOptions::UnPackTo( - FullyConnectedOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void CallOptions::UnPackTo(CallOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = subgraph(); _o->subgraph = _e; }; } -inline flatbuffers::Offset FullyConnectedOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateFullyConnectedOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset CallOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateCallOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateFullyConnectedOptions( - flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateCallOptions(flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const FullyConnectedOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateFullyConnectedOptions(_fbb, _fused_activation_function); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const CallOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _subgraph = _o->subgraph; + return tflite::CreateCallOptions( + _fbb, + _subgraph); } -inline SoftmaxOptionsT *SoftmaxOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new SoftmaxOptionsT(); +inline PadOptionsT *PadOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new PadOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SoftmaxOptions::UnPackTo( - SoftmaxOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void PadOptions::UnPackTo(PadOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = beta(); - _o->beta = _e; - }; } -inline flatbuffers::Offset SoftmaxOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateSoftmaxOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset PadOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreatePadOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateSoftmaxOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SoftmaxOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreatePadOptions(flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const SoftmaxOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _beta = _o->beta; - return tflite::CreateSoftmaxOptions(_fbb, _beta); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const PadOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreatePadOptions( + _fbb); } -inline ConcatenationOptionsT *ConcatenationOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new ConcatenationOptionsT(); +inline ReshapeOptionsT *ReshapeOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new ReshapeOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void ConcatenationOptions::UnPackTo( - ConcatenationOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void ReshapeOptions::UnPackTo(ReshapeOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = axis(); - _o->axis = _e; - }; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = new_shape(); if (_e) { _o->new_shape.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->new_shape[_i] = _e->Get(_i); } } }; } -inline flatbuffers::Offset ConcatenationOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateConcatenationOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset ReshapeOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateReshapeOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateConcatenationOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ConcatenationOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateReshapeOptions(flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const ConcatenationOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _axis = _o->axis; - auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateConcatenationOptions(_fbb, _axis, - _fused_activation_function); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const ReshapeOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _new_shape = _o->new_shape.size() ? _fbb.CreateVector(_o->new_shape) : 0; + return tflite::CreateReshapeOptions( + _fbb, + _new_shape); } -inline AddOptionsT *AddOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new AddOptionsT(); +inline SpaceToBatchNDOptionsT *SpaceToBatchNDOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SpaceToBatchNDOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void AddOptions::UnPackTo( - AddOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void SpaceToBatchNDOptions::UnPackTo(SpaceToBatchNDOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; } -inline flatbuffers::Offset AddOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateAddOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset SpaceToBatchNDOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SpaceToBatchNDOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSpaceToBatchNDOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateAddOptions( - flatbuffers::FlatBufferBuilder &_fbb, const AddOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateSpaceToBatchNDOptions(flatbuffers::FlatBufferBuilder &_fbb, const SpaceToBatchNDOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const AddOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateAddOptions(_fbb, _fused_activation_function); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SpaceToBatchNDOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateSpaceToBatchNDOptions( + _fbb); } -inline MulOptionsT *MulOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new MulOptionsT(); +inline BatchToSpaceNDOptionsT *BatchToSpaceNDOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new BatchToSpaceNDOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void MulOptions::UnPackTo( - MulOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void BatchToSpaceNDOptions::UnPackTo(BatchToSpaceNDOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; } -inline flatbuffers::Offset MulOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateMulOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset BatchToSpaceNDOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateBatchToSpaceNDOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateMulOptions( - flatbuffers::FlatBufferBuilder &_fbb, const MulOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateBatchToSpaceNDOptions(flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const MulOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateMulOptions(_fbb, _fused_activation_function); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const BatchToSpaceNDOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateBatchToSpaceNDOptions( + _fbb); } -inline L2NormOptionsT *L2NormOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new L2NormOptionsT(); +inline SkipGramOptionsT *SkipGramOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SkipGramOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void L2NormOptions::UnPackTo( - L2NormOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void SkipGramOptions::UnPackTo(SkipGramOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; + { auto _e = ngram_size(); _o->ngram_size = _e; }; + { auto _e = max_skip_size(); _o->max_skip_size = _e; }; + { auto _e = include_all_ngrams(); _o->include_all_ngrams = _e; }; } -inline flatbuffers::Offset L2NormOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateL2NormOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset SkipGramOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSkipGramOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateL2NormOptions( - flatbuffers::FlatBufferBuilder &_fbb, const L2NormOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateSkipGramOptions(flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const L2NormOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _fused_activation_function = _o->fused_activation_function; - return tflite::CreateL2NormOptions(_fbb, _fused_activation_function); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SkipGramOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _ngram_size = _o->ngram_size; + auto _max_skip_size = _o->max_skip_size; + auto _include_all_ngrams = _o->include_all_ngrams; + return tflite::CreateSkipGramOptions( + _fbb, + _ngram_size, + _max_skip_size, + _include_all_ngrams); } -inline LocalResponseNormalizationOptionsT * -LocalResponseNormalizationOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new LocalResponseNormalizationOptionsT(); +inline SpaceToDepthOptionsT *SpaceToDepthOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SpaceToDepthOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void LocalResponseNormalizationOptions::UnPackTo( - LocalResponseNormalizationOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void SpaceToDepthOptions::UnPackTo(SpaceToDepthOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = radius(); - _o->radius = _e; - }; - { - auto _e = bias(); - _o->bias = _e; - }; - { - auto _e = alpha(); - _o->alpha = _e; - }; - { - auto _e = beta(); - _o->beta = _e; - }; + { auto _e = block_size(); _o->block_size = _e; }; } -inline flatbuffers::Offset -LocalResponseNormalizationOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const LocalResponseNormalizationOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateLocalResponseNormalizationOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset SpaceToDepthOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSpaceToDepthOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateSpaceToDepthOptions(flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SpaceToDepthOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _block_size = _o->block_size; + return tflite::CreateSpaceToDepthOptions( + _fbb, + _block_size); +} + +inline SubOptionsT *SubOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SubOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void SubOptions::UnPackTo(SubOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; +} + +inline flatbuffers::Offset SubOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SubOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSubOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset -CreateLocalResponseNormalizationOptions( - flatbuffers::FlatBufferBuilder &_fbb, - const LocalResponseNormalizationOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateSubOptions(flatbuffers::FlatBufferBuilder &_fbb, const SubOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const LocalResponseNormalizationOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _radius = _o->radius; - auto _bias = _o->bias; - auto _alpha = _o->alpha; - auto _beta = _o->beta; - return tflite::CreateLocalResponseNormalizationOptions(_fbb, _radius, _bias, - _alpha, _beta); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SubOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _fused_activation_function = _o->fused_activation_function; + return tflite::CreateSubOptions( + _fbb, + _fused_activation_function); } -inline LSTMOptionsT *LSTMOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new LSTMOptionsT(); +inline DivOptionsT *DivOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new DivOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void LSTMOptions::UnPackTo( - LSTMOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void DivOptions::UnPackTo(DivOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = fused_activation_function(); - _o->fused_activation_function = _e; - }; - { - auto _e = cell_clip(); - _o->cell_clip = _e; - }; - { - auto _e = proj_clip(); - _o->proj_clip = _e; - }; + { auto _e = fused_activation_function(); _o->fused_activation_function = _e; }; } -inline flatbuffers::Offset LSTMOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateLSTMOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset DivOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const DivOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateDivOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateLSTMOptions( - flatbuffers::FlatBufferBuilder &_fbb, const LSTMOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateDivOptions(flatbuffers::FlatBufferBuilder &_fbb, const DivOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const LSTMOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const DivOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; auto _fused_activation_function = _o->fused_activation_function; - auto _cell_clip = _o->cell_clip; - auto _proj_clip = _o->proj_clip; - return tflite::CreateLSTMOptions(_fbb, _fused_activation_function, _cell_clip, - _proj_clip); + return tflite::CreateDivOptions( + _fbb, + _fused_activation_function); } -inline ResizeBilinearOptionsT *ResizeBilinearOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new ResizeBilinearOptionsT(); +inline TopKV2OptionsT *TopKV2Options::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new TopKV2OptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void ResizeBilinearOptions::UnPackTo( - ResizeBilinearOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void TopKV2Options::UnPackTo(TopKV2OptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = new_height(); - _o->new_height = _e; - }; - { - auto _e = new_width(); - _o->new_width = _e; - }; } -inline flatbuffers::Offset ResizeBilinearOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateResizeBilinearOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset TopKV2Options::Pack(flatbuffers::FlatBufferBuilder &_fbb, const TopKV2OptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateTopKV2Options(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateResizeBilinearOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ResizeBilinearOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateTopKV2Options(flatbuffers::FlatBufferBuilder &_fbb, const TopKV2OptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const ResizeBilinearOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _new_height = _o->new_height; - auto _new_width = _o->new_width; - return tflite::CreateResizeBilinearOptions(_fbb, _new_height, _new_width); -} - -inline CallOptionsT *CallOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new CallOptionsT(); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const TopKV2OptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateTopKV2Options( + _fbb); +} + +inline EmbeddingLookupSparseOptionsT *EmbeddingLookupSparseOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new EmbeddingLookupSparseOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void CallOptions::UnPackTo( - CallOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void EmbeddingLookupSparseOptions::UnPackTo(EmbeddingLookupSparseOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = subgraph(); - _o->subgraph = _e; - }; + { auto _e = combiner(); _o->combiner = _e; }; } -inline flatbuffers::Offset CallOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateCallOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset EmbeddingLookupSparseOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const EmbeddingLookupSparseOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateEmbeddingLookupSparseOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateCallOptions( - flatbuffers::FlatBufferBuilder &_fbb, const CallOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateEmbeddingLookupSparseOptions(flatbuffers::FlatBufferBuilder &_fbb, const EmbeddingLookupSparseOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const CallOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _subgraph = _o->subgraph; - return tflite::CreateCallOptions(_fbb, _subgraph); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const EmbeddingLookupSparseOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _combiner = _o->combiner; + return tflite::CreateEmbeddingLookupSparseOptions( + _fbb, + _combiner); } -inline PadOptionsT *PadOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new PadOptionsT(); +inline GatherOptionsT *GatherOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new GatherOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void PadOptions::UnPackTo( - PadOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void GatherOptions::UnPackTo(GatherOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = before_padding(); - if (_e) { - _o->before_padding.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->before_padding[_i] = _e->Get(_i); - } - } - }; - { - auto _e = after_padding(); - if (_e) { - _o->after_padding.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->after_padding[_i] = _e->Get(_i); - } - } - }; + { auto _e = axis(); _o->axis = _e; }; } -inline flatbuffers::Offset PadOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreatePadOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset GatherOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateGatherOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreatePadOptions( - flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateGatherOptions(flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const PadOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _before_padding = - _o->before_padding.size() ? _fbb.CreateVector(_o->before_padding) : 0; - auto _after_padding = - _o->after_padding.size() ? _fbb.CreateVector(_o->after_padding) : 0; - return tflite::CreatePadOptions(_fbb, _before_padding, _after_padding); -} - -inline ReshapeOptionsT *ReshapeOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new ReshapeOptionsT(); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const GatherOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _axis = _o->axis; + return tflite::CreateGatherOptions( + _fbb, + _axis); +} + +inline TransposeOptionsT *TransposeOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new TransposeOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void ReshapeOptions::UnPackTo( - ReshapeOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void TransposeOptions::UnPackTo(TransposeOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = new_shape(); - if (_e) { - _o->new_shape.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->new_shape[_i] = _e->Get(_i); - } - } - }; } -inline flatbuffers::Offset ReshapeOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateReshapeOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset TransposeOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const TransposeOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateTransposeOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateReshapeOptions( - flatbuffers::FlatBufferBuilder &_fbb, const ReshapeOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateTransposeOptions(flatbuffers::FlatBufferBuilder &_fbb, const TransposeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const ReshapeOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _new_shape = _o->new_shape.size() ? _fbb.CreateVector(_o->new_shape) : 0; - return tflite::CreateReshapeOptions(_fbb, _new_shape); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const TransposeOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateTransposeOptions( + _fbb); } -inline BatchToSpaceNDOptionsT *BatchToSpaceNDOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new BatchToSpaceNDOptionsT(); +inline ExpOptionsT *ExpOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new ExpOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void BatchToSpaceNDOptions::UnPackTo( - BatchToSpaceNDOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void ExpOptions::UnPackTo(ExpOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = block_shape(); - if (_e) { - _o->block_shape.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->block_shape[_i] = _e->Get(_i); - } - } - }; - { - auto _e = before_crops(); - if (_e) { - _o->before_crops.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->before_crops[_i] = _e->Get(_i); - } - } - }; - { - auto _e = after_crops(); - if (_e) { - _o->after_crops.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->after_crops[_i] = _e->Get(_i); - } - } - }; } -inline flatbuffers::Offset BatchToSpaceNDOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateBatchToSpaceNDOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset ExpOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const ExpOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateExpOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateBatchToSpaceNDOptions( - flatbuffers::FlatBufferBuilder &_fbb, const BatchToSpaceNDOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateExpOptions(flatbuffers::FlatBufferBuilder &_fbb, const ExpOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const BatchToSpaceNDOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _block_shape = - _o->block_shape.size() ? _fbb.CreateVector(_o->block_shape) : 0; - auto _before_crops = - _o->before_crops.size() ? _fbb.CreateVector(_o->before_crops) : 0; - auto _after_crops = - _o->after_crops.size() ? _fbb.CreateVector(_o->after_crops) : 0; - return tflite::CreateBatchToSpaceNDOptions(_fbb, _block_shape, _before_crops, - _after_crops); -} - -inline SkipGramOptionsT *SkipGramOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new SkipGramOptionsT(); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const ExpOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateExpOptions( + _fbb); +} + +inline MeanOptionsT *MeanOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new MeanOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SkipGramOptions::UnPackTo( - SkipGramOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void MeanOptions::UnPackTo(MeanOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = ngram_size(); - _o->ngram_size = _e; - }; - { - auto _e = max_skip_size(); - _o->max_skip_size = _e; - }; - { - auto _e = include_all_ngrams(); - _o->include_all_ngrams = _e; - }; + { auto _e = keep_dims(); _o->keep_dims = _e; }; } -inline flatbuffers::Offset SkipGramOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateSkipGramOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset MeanOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const MeanOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateMeanOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateSkipGramOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SkipGramOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateMeanOptions(flatbuffers::FlatBufferBuilder &_fbb, const MeanOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const SkipGramOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _ngram_size = _o->ngram_size; - auto _max_skip_size = _o->max_skip_size; - auto _include_all_ngrams = _o->include_all_ngrams; - return tflite::CreateSkipGramOptions(_fbb, _ngram_size, _max_skip_size, - _include_all_ngrams); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const MeanOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _keep_dims = _o->keep_dims; + return tflite::CreateMeanOptions( + _fbb, + _keep_dims); } -inline SpaceToDepthOptionsT *SpaceToDepthOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new SpaceToDepthOptionsT(); +inline SqueezeOptionsT *SqueezeOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SqueezeOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void SpaceToDepthOptions::UnPackTo( - SpaceToDepthOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void SqueezeOptions::UnPackTo(SqueezeOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = block_size(); - _o->block_size = _e; - }; + { auto _e = squeeze_dims(); if (_e) { _o->squeeze_dims.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->squeeze_dims[_i] = _e->Get(_i); } } }; } -inline flatbuffers::Offset SpaceToDepthOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateSpaceToDepthOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset SqueezeOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SqueezeOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSqueezeOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateSpaceToDepthOptions( - flatbuffers::FlatBufferBuilder &_fbb, const SpaceToDepthOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateSqueezeOptions(flatbuffers::FlatBufferBuilder &_fbb, const SqueezeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const SpaceToDepthOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _block_size = _o->block_size; - return tflite::CreateSpaceToDepthOptions(_fbb, _block_size); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SqueezeOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _squeeze_dims = _o->squeeze_dims.size() ? _fbb.CreateVector(_o->squeeze_dims) : 0; + return tflite::CreateSqueezeOptions( + _fbb, + _squeeze_dims); } -inline EmbeddingLookupSparseOptionsT *EmbeddingLookupSparseOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new EmbeddingLookupSparseOptionsT(); +inline SplitOptionsT *SplitOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new SplitOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void EmbeddingLookupSparseOptions::UnPackTo( - EmbeddingLookupSparseOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void SplitOptions::UnPackTo(SplitOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = combiner(); - _o->combiner = _e; - }; + { auto _e = num_splits(); _o->num_splits = _e; }; } -inline flatbuffers::Offset -EmbeddingLookupSparseOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const EmbeddingLookupSparseOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateEmbeddingLookupSparseOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset SplitOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SplitOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateSplitOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset -CreateEmbeddingLookupSparseOptions( - flatbuffers::FlatBufferBuilder &_fbb, - const EmbeddingLookupSparseOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateSplitOptions(flatbuffers::FlatBufferBuilder &_fbb, const SplitOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const EmbeddingLookupSparseOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _combiner = _o->combiner; - return tflite::CreateEmbeddingLookupSparseOptions(_fbb, _combiner); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SplitOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _num_splits = _o->num_splits; + return tflite::CreateSplitOptions( + _fbb, + _num_splits); } -inline GatherOptionsT *GatherOptions::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new GatherOptionsT(); +inline StridedSliceOptionsT *StridedSliceOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new StridedSliceOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void GatherOptions::UnPackTo( - GatherOptionsT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void StridedSliceOptions::UnPackTo(StridedSliceOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = axis(); - _o->axis = _e; - }; + { auto _e = begin_mask(); _o->begin_mask = _e; }; + { auto _e = end_mask(); _o->end_mask = _e; }; + { auto _e = ellipsis_mask(); _o->ellipsis_mask = _e; }; + { auto _e = new_axis_mask(); _o->new_axis_mask = _e; }; + { auto _e = shrink_axis_mask(); _o->shrink_axis_mask = _e; }; } -inline flatbuffers::Offset GatherOptions::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { - return CreateGatherOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset StridedSliceOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const StridedSliceOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateStridedSliceOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateGatherOptions( - flatbuffers::FlatBufferBuilder &_fbb, const GatherOptionsT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateStridedSliceOptions(flatbuffers::FlatBufferBuilder &_fbb, const StridedSliceOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const GatherOptionsT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _axis = _o->axis; - return tflite::CreateGatherOptions(_fbb, _axis); + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const StridedSliceOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _begin_mask = _o->begin_mask; + auto _end_mask = _o->end_mask; + auto _ellipsis_mask = _o->ellipsis_mask; + auto _new_axis_mask = _o->new_axis_mask; + auto _shrink_axis_mask = _o->shrink_axis_mask; + return tflite::CreateStridedSliceOptions( + _fbb, + _begin_mask, + _end_mask, + _ellipsis_mask, + _new_axis_mask, + _shrink_axis_mask); } -inline OperatorCodeT *OperatorCode::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline OperatorCodeT *OperatorCode::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new OperatorCodeT(); UnPackTo(_o, _resolver); return _o; } -inline void OperatorCode::UnPackTo( - OperatorCodeT *_o, - const flatbuffers::resolver_function_t *_resolver) const { +inline void OperatorCode::UnPackTo(OperatorCodeT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = builtin_code(); - _o->builtin_code = _e; - }; - { - auto _e = custom_code(); - if (_e) _o->custom_code = _e->str(); - }; + { auto _e = builtin_code(); _o->builtin_code = _e; }; + { auto _e = custom_code(); if (_e) _o->custom_code = _e->str(); }; } -inline flatbuffers::Offset OperatorCode::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset OperatorCode::Pack(flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT* _o, const flatbuffers::rehasher_function_t *_rehasher) { return CreateOperatorCode(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateOperatorCode( - flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateOperatorCode(flatbuffers::FlatBufferBuilder &_fbb, const OperatorCodeT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const OperatorCodeT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const OperatorCodeT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; auto _builtin_code = _o->builtin_code; - auto _custom_code = - _o->custom_code.empty() ? 0 : _fbb.CreateString(_o->custom_code); - return tflite::CreateOperatorCode(_fbb, _builtin_code, _custom_code); + auto _custom_code = _o->custom_code.empty() ? 0 : _fbb.CreateString(_o->custom_code); + return tflite::CreateOperatorCode( + _fbb, + _builtin_code, + _custom_code); } -inline OperatorT *Operator::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline OperatorT *Operator::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new OperatorT(); UnPackTo(_o, _resolver); return _o; } -inline void Operator::UnPackTo( - OperatorT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void Operator::UnPackTo(OperatorT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = opcode_index(); - _o->opcode_index = _e; - }; - { - auto _e = inputs(); - if (_e) { - _o->inputs.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->inputs[_i] = _e->Get(_i); - } - } - }; - { - auto _e = outputs(); - if (_e) { - _o->outputs.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->outputs[_i] = _e->Get(_i); - } - } - }; - { - auto _e = builtin_options_type(); - _o->builtin_options.type = _e; - }; - { - auto _e = builtin_options(); - if (_e) - _o->builtin_options.value = - BuiltinOptionsUnion::UnPack(_e, builtin_options_type(), _resolver); - }; - { - auto _e = custom_options(); - if (_e) { - _o->custom_options.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->custom_options[_i] = _e->Get(_i); - } - } - }; - { - auto _e = custom_options_format(); - _o->custom_options_format = _e; - }; + { auto _e = opcode_index(); _o->opcode_index = _e; }; + { auto _e = inputs(); if (_e) { _o->inputs.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->inputs[_i] = _e->Get(_i); } } }; + { auto _e = outputs(); if (_e) { _o->outputs.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->outputs[_i] = _e->Get(_i); } } }; + { auto _e = builtin_options_type(); _o->builtin_options.type = _e; }; + { auto _e = builtin_options(); if (_e) _o->builtin_options.value = BuiltinOptionsUnion::UnPack(_e, builtin_options_type(), _resolver); }; + { auto _e = custom_options(); if (_e) { _o->custom_options.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->custom_options[_i] = _e->Get(_i); } } }; + { auto _e = custom_options_format(); _o->custom_options_format = _e; }; } -inline flatbuffers::Offset Operator::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const OperatorT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset Operator::Pack(flatbuffers::FlatBufferBuilder &_fbb, const OperatorT* _o, const flatbuffers::rehasher_function_t *_rehasher) { return CreateOperator(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateOperator( - flatbuffers::FlatBufferBuilder &_fbb, const OperatorT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateOperator(flatbuffers::FlatBufferBuilder &_fbb, const OperatorT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const OperatorT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const OperatorT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; auto _opcode_index = _o->opcode_index; auto _inputs = _o->inputs.size() ? _fbb.CreateVector(_o->inputs) : 0; auto _outputs = _o->outputs.size() ? _fbb.CreateVector(_o->outputs) : 0; auto _builtin_options_type = _o->builtin_options.type; auto _builtin_options = _o->builtin_options.Pack(_fbb); - auto _custom_options = - _o->custom_options.size() ? _fbb.CreateVector(_o->custom_options) : 0; + auto _custom_options = _o->custom_options.size() ? _fbb.CreateVector(_o->custom_options) : 0; auto _custom_options_format = _o->custom_options_format; - return tflite::CreateOperator(_fbb, _opcode_index, _inputs, _outputs, - _builtin_options_type, _builtin_options, - _custom_options, _custom_options_format); + return tflite::CreateOperator( + _fbb, + _opcode_index, + _inputs, + _outputs, + _builtin_options_type, + _builtin_options, + _custom_options, + _custom_options_format); } -inline SubGraphT *SubGraph::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline SubGraphT *SubGraph::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new SubGraphT(); UnPackTo(_o, _resolver); return _o; } -inline void SubGraph::UnPackTo( - SubGraphT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void SubGraph::UnPackTo(SubGraphT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = tensors(); - if (_e) { - _o->tensors.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->tensors[_i] = - std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); - } - } - }; - { - auto _e = inputs(); - if (_e) { - _o->inputs.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->inputs[_i] = _e->Get(_i); - } - } - }; - { - auto _e = outputs(); - if (_e) { - _o->outputs.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->outputs[_i] = _e->Get(_i); - } - } - }; - { - auto _e = operators(); - if (_e) { - _o->operators.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->operators[_i] = - std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); - } - } - }; - { - auto _e = name(); - if (_e) _o->name = _e->str(); - }; + { auto _e = tensors(); if (_e) { _o->tensors.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->tensors[_i] = std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); } } }; + { auto _e = inputs(); if (_e) { _o->inputs.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->inputs[_i] = _e->Get(_i); } } }; + { auto _e = outputs(); if (_e) { _o->outputs.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->outputs[_i] = _e->Get(_i); } } }; + { auto _e = operators(); if (_e) { _o->operators.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->operators[_i] = std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); } } }; + { auto _e = name(); if (_e) _o->name = _e->str(); }; } -inline flatbuffers::Offset SubGraph::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const SubGraphT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset SubGraph::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SubGraphT* _o, const flatbuffers::rehasher_function_t *_rehasher) { return CreateSubGraph(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateSubGraph( - flatbuffers::FlatBufferBuilder &_fbb, const SubGraphT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateSubGraph(flatbuffers::FlatBufferBuilder &_fbb, const SubGraphT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const SubGraphT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; - auto _tensors = - _o->tensors.size() - ? _fbb.CreateVector>( - _o->tensors.size(), - [](size_t i, _VectorArgs *__va) { - return CreateTensor(*__va->__fbb, __va->__o->tensors[i].get(), - __va->__rehasher); - }, - &_va) - : 0; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SubGraphT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _tensors = _o->tensors.size() ? _fbb.CreateVector> (_o->tensors.size(), [](size_t i, _VectorArgs *__va) { return CreateTensor(*__va->__fbb, __va->__o->tensors[i].get(), __va->__rehasher); }, &_va ) : 0; auto _inputs = _o->inputs.size() ? _fbb.CreateVector(_o->inputs) : 0; auto _outputs = _o->outputs.size() ? _fbb.CreateVector(_o->outputs) : 0; - auto _operators = _o->operators.size() - ? _fbb.CreateVector>( - _o->operators.size(), - [](size_t i, _VectorArgs *__va) { - return CreateOperator( - *__va->__fbb, __va->__o->operators[i].get(), - __va->__rehasher); - }, - &_va) - : 0; + auto _operators = _o->operators.size() ? _fbb.CreateVector> (_o->operators.size(), [](size_t i, _VectorArgs *__va) { return CreateOperator(*__va->__fbb, __va->__o->operators[i].get(), __va->__rehasher); }, &_va ) : 0; auto _name = _o->name.empty() ? 0 : _fbb.CreateString(_o->name); - return tflite::CreateSubGraph(_fbb, _tensors, _inputs, _outputs, _operators, - _name); + return tflite::CreateSubGraph( + _fbb, + _tensors, + _inputs, + _outputs, + _operators, + _name); } -inline BufferT *Buffer::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline BufferT *Buffer::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new BufferT(); UnPackTo(_o, _resolver); return _o; } -inline void Buffer::UnPackTo( - BufferT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void Buffer::UnPackTo(BufferT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = data(); - if (_e) { - _o->data.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->data[_i] = _e->Get(_i); - } - } - }; + { auto _e = data(); if (_e) { _o->data.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->data[_i] = _e->Get(_i); } } }; } -inline flatbuffers::Offset Buffer::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const BufferT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset Buffer::Pack(flatbuffers::FlatBufferBuilder &_fbb, const BufferT* _o, const flatbuffers::rehasher_function_t *_rehasher) { return CreateBuffer(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateBuffer( - flatbuffers::FlatBufferBuilder &_fbb, const BufferT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateBuffer(flatbuffers::FlatBufferBuilder &_fbb, const BufferT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const BufferT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const BufferT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; auto _data = _o->data.size() ? _fbb.CreateVector(_o->data) : 0; - return tflite::CreateBuffer(_fbb, _data); + return tflite::CreateBuffer( + _fbb, + _data); } -inline ModelT *Model::UnPack( - const flatbuffers::resolver_function_t *_resolver) const { +inline ModelT *Model::UnPack(const flatbuffers::resolver_function_t *_resolver) const { auto _o = new ModelT(); UnPackTo(_o, _resolver); return _o; } -inline void Model::UnPackTo( - ModelT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void Model::UnPackTo(ModelT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; - { - auto _e = version(); - _o->version = _e; - }; - { - auto _e = operator_codes(); - if (_e) { - _o->operator_codes.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->operator_codes[_i] = - std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); - } - } - }; - { - auto _e = subgraphs(); - if (_e) { - _o->subgraphs.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->subgraphs[_i] = - std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); - } - } - }; - { - auto _e = description(); - if (_e) _o->description = _e->str(); - }; - { - auto _e = buffers(); - if (_e) { - _o->buffers.resize(_e->size()); - for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { - _o->buffers[_i] = - std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); - } - } - }; + { auto _e = version(); _o->version = _e; }; + { auto _e = operator_codes(); if (_e) { _o->operator_codes.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->operator_codes[_i] = std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); } } }; + { auto _e = subgraphs(); if (_e) { _o->subgraphs.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->subgraphs[_i] = std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); } } }; + { auto _e = description(); if (_e) _o->description = _e->str(); }; + { auto _e = buffers(); if (_e) { _o->buffers.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->buffers[_i] = std::unique_ptr(_e->Get(_i)->UnPack(_resolver)); } } }; } -inline flatbuffers::Offset Model::Pack( - flatbuffers::FlatBufferBuilder &_fbb, const ModelT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset Model::Pack(flatbuffers::FlatBufferBuilder &_fbb, const ModelT* _o, const flatbuffers::rehasher_function_t *_rehasher) { return CreateModel(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateModel( - flatbuffers::FlatBufferBuilder &_fbb, const ModelT *_o, - const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateModel(flatbuffers::FlatBufferBuilder &_fbb, const ModelT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { - flatbuffers::FlatBufferBuilder *__fbb; - const ModelT *__o; - const flatbuffers::rehasher_function_t *__rehasher; - } _va = {&_fbb, _o, _rehasher}; - (void)_va; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const ModelT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; auto _version = _o->version; - auto _operator_codes = - _o->operator_codes.size() - ? _fbb.CreateVector>( - _o->operator_codes.size(), - [](size_t i, _VectorArgs *__va) { - return CreateOperatorCode(*__va->__fbb, - __va->__o->operator_codes[i].get(), - __va->__rehasher); - }, - &_va) - : 0; - auto _subgraphs = _o->subgraphs.size() - ? _fbb.CreateVector>( - _o->subgraphs.size(), - [](size_t i, _VectorArgs *__va) { - return CreateSubGraph( - *__va->__fbb, __va->__o->subgraphs[i].get(), - __va->__rehasher); - }, - &_va) - : 0; - auto _description = - _o->description.empty() ? 0 : _fbb.CreateString(_o->description); - auto _buffers = - _o->buffers.size() - ? _fbb.CreateVector>( - _o->buffers.size(), - [](size_t i, _VectorArgs *__va) { - return CreateBuffer(*__va->__fbb, __va->__o->buffers[i].get(), - __va->__rehasher); - }, - &_va) - : 0; - return tflite::CreateModel(_fbb, _version, _operator_codes, _subgraphs, - _description, _buffers); -} - -inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, - const void *obj, BuiltinOptions type) { + auto _operator_codes = _o->operator_codes.size() ? _fbb.CreateVector> (_o->operator_codes.size(), [](size_t i, _VectorArgs *__va) { return CreateOperatorCode(*__va->__fbb, __va->__o->operator_codes[i].get(), __va->__rehasher); }, &_va ) : 0; + auto _subgraphs = _o->subgraphs.size() ? _fbb.CreateVector> (_o->subgraphs.size(), [](size_t i, _VectorArgs *__va) { return CreateSubGraph(*__va->__fbb, __va->__o->subgraphs[i].get(), __va->__rehasher); }, &_va ) : 0; + auto _description = _o->description.empty() ? 0 : _fbb.CreateString(_o->description); + auto _buffers = _o->buffers.size() ? _fbb.CreateVector> (_o->buffers.size(), [](size_t i, _VectorArgs *__va) { return CreateBuffer(*__va->__fbb, __va->__o->buffers[i].get(), __va->__rehasher); }, &_va ) : 0; + return tflite::CreateModel( + _fbb, + _version, + _operator_codes, + _subgraphs, + _description, + _buffers); +} + +inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *obj, BuiltinOptions type) { switch (type) { case BuiltinOptions_NONE: { return true; @@ -5421,8 +5644,7 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, return verifier.VerifyTable(ptr); } case BuiltinOptions_LocalResponseNormalizationOptions: { - auto ptr = - reinterpret_cast(obj); + auto ptr = reinterpret_cast(obj); return verifier.VerifyTable(ptr); } case BuiltinOptions_LSTMOptions: { @@ -5469,28 +5691,67 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, auto ptr = reinterpret_cast(obj); return verifier.VerifyTable(ptr); } - default: - return false; + case BuiltinOptions_SpaceToBatchNDOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_TransposeOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_MeanOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_SubOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_DivOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_SqueezeOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_SequenceRNNOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_StridedSliceOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_ExpOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_TopKV2Options: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_SplitOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + default: return false; } } -inline bool VerifyBuiltinOptionsVector( - flatbuffers::Verifier &verifier, - const flatbuffers::Vector> *values, - const flatbuffers::Vector *types) { +inline bool VerifyBuiltinOptionsVector(flatbuffers::Verifier &verifier, const flatbuffers::Vector> *values, const flatbuffers::Vector *types) { + if (!values || !types) return !values && !types; if (values->size() != types->size()) return false; for (flatbuffers::uoffset_t i = 0; i < values->size(); ++i) { - if (!VerifyBuiltinOptions(verifier, values->Get(i), - types->GetEnum(i))) { + if (!VerifyBuiltinOptions( + verifier, values->Get(i), types->GetEnum(i))) { return false; } } return true; } -inline void *BuiltinOptionsUnion::UnPack( - const void *obj, BuiltinOptions type, - const flatbuffers::resolver_function_t *resolver) { +inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, const flatbuffers::resolver_function_t *resolver) { switch (type) { case BuiltinOptions_Conv2DOptions: { auto ptr = reinterpret_cast(obj); @@ -5541,8 +5802,7 @@ inline void *BuiltinOptionsUnion::UnPack( return ptr->UnPack(resolver); } case BuiltinOptions_LocalResponseNormalizationOptions: { - auto ptr = - reinterpret_cast(obj); + auto ptr = reinterpret_cast(obj); return ptr->UnPack(resolver); } case BuiltinOptions_LSTMOptions: { @@ -5589,14 +5849,55 @@ inline void *BuiltinOptionsUnion::UnPack( auto ptr = reinterpret_cast(obj); return ptr->UnPack(resolver); } - default: - return nullptr; + case BuiltinOptions_SpaceToBatchNDOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_TransposeOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_MeanOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_SubOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_DivOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_SqueezeOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_SequenceRNNOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_StridedSliceOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_ExpOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_TopKV2Options: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_SplitOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + default: return nullptr; } } -inline flatbuffers::Offset BuiltinOptionsUnion::Pack( - flatbuffers::FlatBufferBuilder &_fbb, - const flatbuffers::rehasher_function_t *_rehasher) const { +inline flatbuffers::Offset BuiltinOptionsUnion::Pack(flatbuffers::FlatBufferBuilder &_fbb, const flatbuffers::rehasher_function_t *_rehasher) const { switch (type) { case BuiltinOptions_Conv2DOptions: { auto ptr = reinterpret_cast(value); @@ -5647,10 +5948,8 @@ inline flatbuffers::Offset BuiltinOptionsUnion::Pack( return CreateL2NormOptions(_fbb, ptr, _rehasher).Union(); } case BuiltinOptions_LocalResponseNormalizationOptions: { - auto ptr = - reinterpret_cast(value); - return CreateLocalResponseNormalizationOptions(_fbb, ptr, _rehasher) - .Union(); + auto ptr = reinterpret_cast(value); + return CreateLocalResponseNormalizationOptions(_fbb, ptr, _rehasher).Union(); } case BuiltinOptions_LSTMOptions: { auto ptr = reinterpret_cast(value); @@ -5696,32 +5995,70 @@ inline flatbuffers::Offset BuiltinOptionsUnion::Pack( auto ptr = reinterpret_cast(value); return CreateBatchToSpaceNDOptions(_fbb, ptr, _rehasher).Union(); } - default: - return 0; + case BuiltinOptions_SpaceToBatchNDOptions: { + auto ptr = reinterpret_cast(value); + return CreateSpaceToBatchNDOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_TransposeOptions: { + auto ptr = reinterpret_cast(value); + return CreateTransposeOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_MeanOptions: { + auto ptr = reinterpret_cast(value); + return CreateMeanOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_SubOptions: { + auto ptr = reinterpret_cast(value); + return CreateSubOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_DivOptions: { + auto ptr = reinterpret_cast(value); + return CreateDivOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_SqueezeOptions: { + auto ptr = reinterpret_cast(value); + return CreateSqueezeOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_SequenceRNNOptions: { + auto ptr = reinterpret_cast(value); + return CreateSequenceRNNOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_StridedSliceOptions: { + auto ptr = reinterpret_cast(value); + return CreateStridedSliceOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_ExpOptions: { + auto ptr = reinterpret_cast(value); + return CreateExpOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_TopKV2Options: { + auto ptr = reinterpret_cast(value); + return CreateTopKV2Options(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_SplitOptions: { + auto ptr = reinterpret_cast(value); + return CreateSplitOptions(_fbb, ptr, _rehasher).Union(); + } + default: return 0; } } -inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) - FLATBUFFERS_NOEXCEPT : type(u.type), - value(nullptr) { +inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FLATBUFFERS_NOEXCEPT : type(u.type), value(nullptr) { switch (type) { case BuiltinOptions_Conv2DOptions: { value = new Conv2DOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_DepthwiseConv2DOptions: { - value = new DepthwiseConv2DOptionsT( - *reinterpret_cast(u.value)); + value = new DepthwiseConv2DOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_ConcatEmbeddingsOptions: { - value = new ConcatEmbeddingsOptionsT( - *reinterpret_cast(u.value)); + value = new ConcatEmbeddingsOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_LSHProjectionOptions: { - value = new LSHProjectionOptionsT( - *reinterpret_cast(u.value)); + value = new LSHProjectionOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_Pool2DOptions: { @@ -5737,18 +6074,15 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) break; } case BuiltinOptions_FullyConnectedOptions: { - value = new FullyConnectedOptionsT( - *reinterpret_cast(u.value)); + value = new FullyConnectedOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_SoftmaxOptions: { - value = - new SoftmaxOptionsT(*reinterpret_cast(u.value)); + value = new SoftmaxOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_ConcatenationOptions: { - value = new ConcatenationOptionsT( - *reinterpret_cast(u.value)); + value = new ConcatenationOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_AddOptions: { @@ -5760,8 +6094,7 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) break; } case BuiltinOptions_LocalResponseNormalizationOptions: { - value = new LocalResponseNormalizationOptionsT( - *reinterpret_cast(u.value)); + value = new LocalResponseNormalizationOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_LSTMOptions: { @@ -5769,8 +6102,7 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) break; } case BuiltinOptions_ResizeBilinearOptions: { - value = new ResizeBilinearOptionsT( - *reinterpret_cast(u.value)); + value = new ResizeBilinearOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_CallOptions: { @@ -5778,23 +6110,19 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) break; } case BuiltinOptions_ReshapeOptions: { - value = - new ReshapeOptionsT(*reinterpret_cast(u.value)); + value = new ReshapeOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_SkipGramOptions: { - value = - new SkipGramOptionsT(*reinterpret_cast(u.value)); + value = new SkipGramOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_SpaceToDepthOptions: { - value = new SpaceToDepthOptionsT( - *reinterpret_cast(u.value)); + value = new SpaceToDepthOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_EmbeddingLookupSparseOptions: { - value = new EmbeddingLookupSparseOptionsT( - *reinterpret_cast(u.value)); + value = new EmbeddingLookupSparseOptionsT(*reinterpret_cast(u.value)); break; } case BuiltinOptions_MulOptions: { @@ -5810,8 +6138,51 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) break; } case BuiltinOptions_BatchToSpaceNDOptions: { - value = new BatchToSpaceNDOptionsT( - *reinterpret_cast(u.value)); + value = new BatchToSpaceNDOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_SpaceToBatchNDOptions: { + value = new SpaceToBatchNDOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_TransposeOptions: { + value = new TransposeOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_MeanOptions: { + value = new MeanOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_SubOptions: { + value = new SubOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_DivOptions: { + value = new DivOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_SqueezeOptions: { + value = new SqueezeOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_SequenceRNNOptions: { + value = new SequenceRNNOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_StridedSliceOptions: { + value = new StridedSliceOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_ExpOptions: { + value = new ExpOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_TopKV2Options: { + value = new TopKV2OptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_SplitOptions: { + value = new SplitOptionsT(*reinterpret_cast(u.value)); break; } default: @@ -5941,8 +6312,62 @@ inline void BuiltinOptionsUnion::Reset() { delete ptr; break; } - default: + case BuiltinOptions_SpaceToBatchNDOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_TransposeOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_MeanOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_SubOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_DivOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_SqueezeOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_SequenceRNNOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_StridedSliceOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_ExpOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_TopKV2Options: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_SplitOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; break; + } + default: break; } value = nullptr; type = BuiltinOptions_NONE; @@ -5952,25 +6377,33 @@ inline const tflite::Model *GetModel(const void *buf) { return flatbuffers::GetRoot(buf); } -inline const char *ModelIdentifier() { return "TFL3"; } +inline const char *ModelIdentifier() { + return "TFL3"; +} inline bool ModelBufferHasIdentifier(const void *buf) { - return flatbuffers::BufferHasIdentifier(buf, ModelIdentifier()); + return flatbuffers::BufferHasIdentifier( + buf, ModelIdentifier()); } -inline bool VerifyModelBuffer(flatbuffers::Verifier &verifier) { +inline bool VerifyModelBuffer( + flatbuffers::Verifier &verifier) { return verifier.VerifyBuffer(ModelIdentifier()); } -inline const char *ModelExtension() { return "tflite"; } +inline const char *ModelExtension() { + return "tflite"; +} -inline void FinishModelBuffer(flatbuffers::FlatBufferBuilder &fbb, - flatbuffers::Offset root) { +inline void FinishModelBuffer( + flatbuffers::FlatBufferBuilder &fbb, + flatbuffers::Offset root) { fbb.Finish(root, ModelIdentifier()); } inline std::unique_ptr UnPackModel( - const void *buf, const flatbuffers::resolver_function_t *res = nullptr) { + const void *buf, + const flatbuffers::resolver_function_t *res = nullptr) { return std::unique_ptr(GetModel(buf)->UnPack(res)); } diff --git a/tensorflow/contrib/lite/simple_memory_arena.h b/tensorflow/contrib/lite/simple_memory_arena.h index 0d0b7f9ff79bf9fd8a60dbc057d63f44eeaa6396..0535522374c63459d029c252ebe94628cf3122d5 100644 --- a/tensorflow/contrib/lite/simple_memory_arena.h +++ b/tensorflow/contrib/lite/simple_memory_arena.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_ +#define TENSORFLOW_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_ #include #include @@ -36,9 +36,9 @@ struct ArenaAlloc { } }; -// This small class is responsible for allocating, dealocating and reusing +// This small class is responsible for allocating, deallocating and reusing // dynamic memory from a common underlying buffer. The arena can be used in -// scenarios when the pattern of memory allocations and dealocations is +// scenarios when the pattern of memory allocations and deallocations is // repetitive, e.g. running NN inference in multiple iterations. class SimpleMemoryArena { public: @@ -68,6 +68,10 @@ class SimpleMemoryArena { TfLiteStatus Clear(); + int64_t BasePointer() const { + return reinterpret_cast(underlying_buffer_aligned_ptr_); + } + private: bool commited_; size_t arena_alignment_; @@ -81,4 +85,4 @@ class SimpleMemoryArena { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_ +#endif // TENSORFLOW_CONTRIB_LITE_SIMPLE_MEMORY_ARENA_H_ diff --git a/tensorflow/contrib/lite/string_util.h b/tensorflow/contrib/lite/string_util.h index 12872d11232e2a32527d660be8acce3e09f00125..c35a2fff3c23b17515323b65d08df6f6da288834 100644 --- a/tensorflow/contrib/lite/string_util.h +++ b/tensorflow/contrib/lite/string_util.h @@ -37,8 +37,8 @@ limitations under the License. // # described above. // buf.WriteToTensor(tensor) -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_STRING_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_STRING_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_STRING_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_STRING_UTIL_H_ #include @@ -88,4 +88,4 @@ int GetStringCount(const TfLiteTensor* tensor); StringRef GetString(const TfLiteTensor* tensor, int string_index); } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_STRING_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_STRING_UTIL_H_ diff --git a/tensorflow/contrib/lite/testdata/multi_add.pb b/tensorflow/contrib/lite/testdata/multi_add.pb new file mode 100644 index 0000000000000000000000000000000000000000..e95a20841fb2b320bd77994d9dda157d79311dd6 --- /dev/null +++ b/tensorflow/contrib/lite/testdata/multi_add.pb @@ -0,0 +1,26 @@ + +I +a Placeholder" /device:CPU:0* +shape:* +dtype0 +I +b Placeholder" /device:CPU:0* +dtype0* +shape: +I +c Placeholder" /device:CPU:0* +dtype0* +shape: +I +d Placeholder" /device:CPU:0* +dtype0* +shape: +& +iAddbc" /device:CPU:0* +T0 +& +xAddai" /device:CPU:0* +T0 +& +yAdddi" /device:CPU:0* +T0" \ No newline at end of file diff --git a/tensorflow/contrib/lite/testing/BUILD b/tensorflow/contrib/lite/testing/BUILD index 96800304e56a2517b7c1222f47d0207725edc15f..06570ae9aa3d10c3cb73ab362e30244ec0b78a35 100644 --- a/tensorflow/contrib/lite/testing/BUILD +++ b/tensorflow/contrib/lite/testing/BUILD @@ -24,6 +24,8 @@ gen_zipped_test_files( "control_dep.zip", "conv.zip", "depthwiseconv.zip", + "div.zip", + "exp.zip", "fully_connected.zip", "fused_batch_norm.zip", "gather.zip", @@ -32,6 +34,7 @@ gen_zipped_test_files( "l2norm.zip", "local_response_norm.zip", "max_pool.zip", + "mean.zip", "mul.zip", "pad.zip", "relu.zip", @@ -41,7 +44,14 @@ gen_zipped_test_files( "resize_bilinear.zip", "sigmoid.zip", "softmax.zip", + "space_to_batch_nd.zip", "space_to_depth.zip", + "split.zip", + "squeeze.zip", + "strided_slice.zip", + "sub.zip", + "topk.zip", + "transpose.zip", ], ) @@ -114,6 +124,21 @@ cc_test( ], ) +cc_library( + name = "join", + hdrs = ["join.h"], +) + +cc_test( + name = "join_test", + size = "small", + srcs = ["join_test.cc"], + deps = [ + ":join", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "tflite_driver", srcs = ["tflite_driver.cc"], @@ -183,31 +208,71 @@ cc_binary( srcs = ["nnapi_example.cc"], deps = [ ":parse_testdata_lib", - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", + ":tflite_driver", "//tensorflow/contrib/lite/nnapi:nnapi_lib", ], ) +cc_library( + name = "tf_driver", + srcs = ["tf_driver.cc"], + hdrs = ["tf_driver.h"], + deps = [ + ":join", + ":split", + ":test_runner", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:tensorflow", + ], +) + +cc_test( + name = "tf_driver_test", + size = "small", + srcs = ["tf_driver_test.cc"], + data = ["//tensorflow/contrib/lite:testdata/multi_add.pb"], + deps = [ + ":tf_driver", + "@com_google_googletest//:gtest_main", + ], +) + tf_cc_test( name = "generated_examples_zip_test", - size = "medium", + size = "large", srcs = ["generated_examples_zip_test.cc"], + args = [ + "--zip_files_dir=tensorflow/contrib/lite/testing/optest", + # TODO(angerson) We may be able to add an external unzip binary instead + # of relying on an existing one for OSS builds. + "--unzip_binary_path=/usr/bin/unzip", + ], data = [":optest"], - shard_count = 10, + shard_count = 20, tags = ["no_oss"], deps = [ ":parse_testdata_lib", + ":tflite_driver", + ":util", + "@com_google_googletest//:gtest", + "@com_googlesource_code_re2//:re2", "//tensorflow/contrib/lite:builtin_op_data", "//tensorflow/contrib/lite:framework", "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "@com_google_googletest//:gtest", - "@com_googlesource_code_re2//:re2", - ], + ] + select({ + "//conditions:default": [ + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], + "//tensorflow:android": [ + "//tensorflow/core:android_tensorflow_lib", + "//tensorflow/core:android_tensorflow_test_lib", + ], + }), ) filegroup( diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py index 02f59438cdb08180ded1098181613067ec4c7c67..b6c09306d6adb8e54d5108dac850f0249ffcb838 100644 --- a/tensorflow/contrib/lite/testing/generate_examples.py +++ b/tensorflow/contrib/lite/testing/generate_examples.py @@ -94,12 +94,17 @@ KNOWN_BUGS = { r"softmax.*input_shape=\[1,3,4,3\]": "67749831", # SpaceToDepth only supports float32. r"space_to_depth.*(float16|int32|uint8|int64)": "68018134", - # Gather doesn't support int64 indices. - r"gather.*indices_dtype=int64": "XXXX", - # BatchToSpaceND doesn't support cropping. + # BatchToSpaceND doesn't support cropping. This catches test cases with + # const tensors as crops. r"batch_to_space_nd.*crops=\[\[1,1\],\[1,1\]\]": "70594634", # BatchToSpaceND only supports 4D tensors. r"batch_to_space_nd.*input_shape=\[8,2,2,2,1,1\]": "70594733", + # Div will use floordiv. + r"div.*int32": "72051395", + # TOCO require matching dimensions in strided_slice. + r"strided_slice.*begin=\[0\].*end=\[1\].*": "73170889", + # No support for SplitV + r"split.*num_or_size_splits=\[2,2\]": "73377559", } @@ -240,7 +245,7 @@ def create_tensor_data(dtype, shape, min_value=-100, max_value=100): if dtype in (tf.float32, tf.float16): value = (max_value-min_value)*np.random.random_sample(shape)+min_value elif dtype in (tf.int32, tf.uint8, tf.int64): - value = np.random.random_integers(min_value, max_value, shape) + value = np.random.randint(min_value, max_value+1, shape) return value.astype(dtype) @@ -326,6 +331,12 @@ def toco_convert(graph_def_str, input_tensors, output_tensors, return (None if exit_code != 0 else output_file.read()), log +def normalize_output_name(output_name): + """Remove :0 suffix from tensor names.""" + return output_name.split(":")[0] if output_name.endswith( + ":0") else output_name + + def make_zip_of_tests(zip_path, test_parameters, make_graph, @@ -414,8 +425,8 @@ def make_zip_of_tests(zip_path, sess.graph_def.SerializeToString(), [(input_tensor.name.split(":")[0], input_tensor.get_shape(), input_tensor.dtype) for input_tensor in inputs], - [out.name.split(":")[0] - for out in outputs], drop_control_dependency) + [normalize_output_name(out.name) for out in outputs], + drop_control_dependency) report["toco"] = (report_lib.SUCCESS if tflite_model_binary is not None else report_lib.FAILED) report["toco_log"] = toco_log @@ -618,7 +629,7 @@ def make_constant_tests(zip_path): def build_graph(parameters): # Since Toco & Tflite can't have a single constant op in the entire graph, - # this test adds a zero tesnor with a constant op tensor. + # this test adds a zero tensor with a constant op tensor. input1 = tf.placeholder(dtype=parameters["dtype"], name="input1", shape=parameters["input_shape"]) out = tf.ones(parameters["input_shape"], dtype=parameters["dtype"]) + input1 @@ -632,7 +643,7 @@ def make_constant_tests(zip_path): make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) -def make_add_tests(zip_path): +def make_binary_op_tests(zip_path, binary_operator): """Make a set of tests to do add with and without broadcast.""" # These parameters are split because we don't support broadcasting. @@ -640,25 +651,36 @@ def make_add_tests(zip_path): "dtype": [tf.float32, tf.int32], "input_shape_1": [[1, 3, 4, 3]], "input_shape_2": [[1, 3, 4, 3]], + "activation": [True] }, { "dtype": [tf.float32], "input_shape_1": [[5]], "input_shape_2": [[5]], + "activation": [False, True] }, { "dtype": [tf.float32], "input_shape_1": [[1, 3, 4, 3]], "input_shape_2": [[3]], + "activation": [True] }] def build_graph(parameters): - input1 = tf.placeholder(dtype=parameters["dtype"], name="input1", - shape=parameters["input_shape_1"]) - input2 = tf.placeholder(dtype=parameters["dtype"], name="input2", - shape=parameters["input_shape_2"]) - out = tf.add(input1, input2) + """Builds the graph given the current parameters.""" + input1 = tf.placeholder( + dtype=parameters["dtype"], + name="input1", + shape=parameters["input_shape_1"]) + input2 = tf.placeholder( + dtype=parameters["dtype"], + name="input2", + shape=parameters["input_shape_2"]) + out = binary_operator(input1, input2) + if parameters["activation"]: + out = tf.nn.relu(out) return [input1, input2], [out] def build_inputs(parameters, sess, inputs, outputs): + """Builds operand inputs for op.""" input1 = create_tensor_data(parameters["dtype"], parameters["input_shape_1"]) input2 = create_tensor_data(parameters["dtype"], @@ -672,44 +694,100 @@ def make_add_tests(zip_path): make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) -def make_mul_tests(zip_path): - """Make a set of tests to do mul with and without broadcast.""" +def make_mean_tests(zip_path): + """Make a set of tests to do mean.""" - # These parameters are split because we don't support broadcasting. test_parameters = [{ - "dtype": [tf.float32, tf.int32], - "input_shape_1": [[1, 3, 4, 3]], - "input_shape_2": [[1, 3, 4, 3]], + "input_dtype": [tf.float32, tf.int32, tf.int64], + "input_shape": [[3, 2, 4]], + "axis": [ + None, 0, 1, 2, [0, 1], [0, 2], [1, 2], [0, 1, 2], [1, 0], [2, 0], + [2, 1], [2, 1, 0], [2, 0, 1], -1, -2, -3, [1, -1], [0, -1], [-1, 0], + [-1, -2, -3], [0, 0, 0], [2, 2, 0], [1, 0, -3, -3] + ], + "const_axis": [True, False], + "keep_dims": [True, False], }, { - "dtype": [tf.float32], - "input_shape_1": [[5]], - "input_shape_2": [[5]], - }, { - "dtype": [tf.float32], - "input_shape_1": [[1, 3, 4, 3]], - "input_shape_2": [[3]], + "input_dtype": [tf.float32, tf.int32, tf.int64], + "input_shape": [[1, 224, 224, 3]], + "axis": [ + None, 0, 1, 2, 3, [1, 2], [0, 3], [1, 2, 3], [0, 1, 2, 3], + [3, 2, 1, 0], [3, 1, 0, 2], [2, 0], [3, 0], [3, 1], [1, 0], -1, -2, + -3, -4, [0, -2], [2, 3, -1, 0], [3, 1, 2, -3], [3, -4], [2, 2, 2], + [2, 2, 3], [-3, -3, -4], [-3, 2, 1] + ], + "const_axis": [True, False], + "keep_dims": [True, False], }] def build_graph(parameters): - input1 = tf.placeholder(dtype=parameters["dtype"], name="input1", - shape=parameters["input_shape_1"]) - input2 = tf.placeholder(dtype=parameters["dtype"], name="input2", - shape=parameters["input_shape_2"]) - out = tf.multiply(input1, input2) - return [input1, input2], [out] + """Build the mean op testing graph.""" + input_tensor = tf.placeholder( + dtype=parameters["input_dtype"], + name="input", + shape=parameters["input_shape"]) + + # Get axis as either a placeholder or constants. + if parameters["const_axis"]: + axis = parameters["axis"] + input_tensors = [input_tensor] + else: + if isinstance(parameters["axis"], list): + shape = [len(parameters["axis"])] + else: + shape = [0] # shape for None or integers. + axis = tf.placeholder(dtype=tf.int32, name="axis", shape=shape) + input_tensors = [input_tensor, axis] + + out = tf.reduce_mean( + input_tensor, axis=axis, keep_dims=parameters["keep_dims"]) + return input_tensors, [out] def build_inputs(parameters, sess, inputs, outputs): - input1 = create_tensor_data(parameters["dtype"], - parameters["input_shape_1"]) - input2 = create_tensor_data(parameters["dtype"], - parameters["input_shape_2"]) - return [input1, input2], sess.run( - outputs, feed_dict={inputs[0]: input1, - inputs[1]: input2}) + values = [ + create_tensor_data(parameters["input_dtype"], parameters["input_shape"]) + ] + if not parameters["const_axis"]: + if parameters["axis"]: + values.append(np.array(parameters["axis"])) + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) + + make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + + +def make_exp_tests(zip_path): + """Make a set of tests to do exp.""" + + test_parameters = [{ + "input_dtype": [tf.float32], + "input_shape": [[3], [1, 100], [4, 2, 3], [5, 224, 224, 3]], + }] + + def build_graph(parameters): + """Build the exp op testing graph.""" + input_tensor = tf.placeholder( + dtype=parameters["input_dtype"], + name="input", + shape=parameters["input_shape"]) + + out = tf.exp(input_tensor) + return [input_tensor], [out] + + def build_inputs(parameters, sess, inputs, outputs): + values = [ + create_tensor_data(parameters["input_dtype"], parameters["input_shape"], + min_value=-100, max_value=9) + ] + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) +def make_binary_op_tests_func(binary_operator): + """Return a function that does a test on a binary operator.""" + return lambda zip_path: make_binary_op_tests(zip_path, binary_operator) + + def make_gather_tests(zip_path): """Make a set of tests to do gather.""" @@ -830,34 +908,55 @@ def make_fused_batch_norm_tests(zip_path): def make_conv_tests(zip_path): """Make a set of tests to do convolution.""" - test_parameters = [{ - "input_shape": [[1, 3, 4, 3]], - "filter_shape": [[1, 1, 3, 2]], - "strides": [[1, 1, 1, 1], [1, 2, 3, 1]], - "padding": ["SAME", "VALID"], - "data_format": ["NHWC"], # TODO(aselle): NCHW would be good - }, { - "input_shape": [[2, 14, 14, 2]], - "filter_shape": [[6, 6, 2, 2]], - "strides": [[1, 1, 1, 1], [1, 2, 3, 1]], - "padding": ["SAME", "VALID"], - "data_format": ["NHWC"], # TODO(aselle): NCHW would be good - }] + test_parameters = [ + { + "input_shape": [[1, 3, 4, 3]], + "filter_shape": [[1, 1, 3, 2]], + "strides": [[1, 1, 1, 1], [1, 2, 3, 1]], + "padding": ["SAME", "VALID"], + "data_format": ["NHWC"], # TODO(aselle): NCHW would be good + "constant_filter": [True, False], + }, + { + "input_shape": [[2, 14, 14, 2]], + "filter_shape": [[6, 6, 2, 2]], + "strides": [[1, 1, 1, 1], [1, 2, 3, 1]], + "padding": ["SAME", "VALID"], + "data_format": ["NHWC"], # TODO(aselle): NCHW would be good + "constant_filter": [True, False], + } + ] def build_graph(parameters): + """Build a conv graph given `parameters`.""" input_tensor = tf.placeholder( dtype=tf.float32, name="input", shape=parameters["input_shape"]) - filter_values = create_tensor_data(np.float32, parameters["filter_shape"]) - out = tf.nn.conv2d(input_tensor, filter_values, - strides=parameters["strides"], - padding=parameters["padding"], - data_format=parameters["data_format"]) - return [input_tensor], [out] + + # Get filter input either as a placeholder or constants. Also get a list of + # the input tensors that are represented as placeholders. + if parameters["constant_filter"]: + filter_input = create_tensor_data(np.float32, parameters["filter_shape"]) + input_tensors = [input_tensor] + else: + filter_input = tf.placeholder( + dtype=tf.float32, name="filter", shape=parameters["filter_shape"]) + input_tensors = [input_tensor, filter_input] + + out = tf.nn.conv2d( + input_tensor, + filter_input, + strides=parameters["strides"], + padding=parameters["padding"], + data_format=parameters["data_format"]) + return input_tensors, [out] def build_inputs(parameters, sess, inputs, outputs): - input_values = create_tensor_data(np.float32, parameters["input_shape"]) - return [input_values], sess.run( - outputs, feed_dict=dict(zip(inputs, [input_values]))) + # Build list of input values either containing 1 tensor (input) or 2 tensors + # (input, filter) based on whether filter is constant or variable input. + values = [create_tensor_data(np.float32, parameters["input_shape"])] + if not parameters["constant_filter"]: + values.append(create_tensor_data(np.float32, parameters["filter_shape"])) + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) @@ -866,62 +965,112 @@ def make_depthwiseconv_tests(zip_path): """Make a set of tests to do convolution.""" # Tensorflow only supports equal strides - test_parameters = [{ - "input_shape": [[1, 3, 4, 3], [1, 10, 10, 3]], - "filter_size": [[1, 1], [1, 2], [3, 3]], - "strides": [[1, 1, 1, 1], [1, 3, 3, 1]], - "channel_multiplier": [1, 2], - "rate": [[1, 1]], - "padding": ["SAME", "VALID"], - "data_format": ["NHWC"], - }, { - "input_shape": [[1, 3, 4, 3]], - "filter_size": [[1, 1]], - "strides": [[1, 1, 2, 1]], # TF needs [1, x, x, 1] - "channel_multiplier": [2], - "rate": [[2, 2]], # Only [1, 1] is supported - "padding": ["SAME"], - "data_format": ["NHWC"], - }] + test_parameters = [ + { + "input_shape": [[1, 3, 4, 3], [1, 10, 10, 3]], + "filter_size": [[1, 1], [1, 2], [3, 3]], + "strides": [[1, 1, 1, 1], [1, 3, 3, 1]], + "channel_multiplier": [1, 2], + "rate": [[1, 1]], + "padding": ["SAME", "VALID"], + "data_format": ["NHWC"], + "constant_filter": [True, False], + }, + { + "input_shape": [[1, 3, 4, 3]], + "filter_size": [[1, 1]], + "strides": [[1, 1, 2, 1]], # TF needs [1, x, x, 1] + "channel_multiplier": [2], + "rate": [[2, 2]], # Only [1, 1] is supported + "padding": ["SAME"], + "data_format": ["NHWC"], + "constant_filter": [True, False], + } + ] - def build_graph(parameters): - """Build a depthwise conv graph given `parameters`.""" + def get_tensor_shapes(parameters): input_shape = parameters["input_shape"] filter_size = parameters["filter_size"] + filter_shape = filter_size + [ + input_shape[3], parameters["channel_multiplier"] + ] + return [input_shape, filter_shape] + + def build_graph(parameters): + """Build a depthwise conv graph given `parameters`.""" + input_shape, filter_shape = get_tensor_shapes(parameters) input_tensor = tf.placeholder( dtype=tf.float32, name="input", shape=input_shape) - filter_shape = filter_size + [ - input_shape[3], parameters["channel_multiplier"]] - filter_values = create_tensor_data(np.float32, filter_shape) + + # Get filter input either as a placeholder or constants. Also get a list of + # the input tensors that are represented as placeholders. + if parameters["constant_filter"]: + filter_input = create_tensor_data(np.float32, filter_shape) + input_tensors = [input_tensor] + else: + filter_input = tf.placeholder( + dtype=tf.float32, name="filter", shape=filter_shape) + input_tensors = [input_tensor, filter_input] + out = tf.nn.depthwise_conv2d( - input_tensor, filter_values, + input_tensor, + filter_input, strides=parameters["strides"], rate=parameters["rate"], padding=parameters["padding"], data_format=parameters["data_format"]) - return [input_tensor], [out] + return input_tensors, [out] def build_inputs(parameters, sess, inputs, outputs): - input_values = create_tensor_data(np.float32, parameters["input_shape"]) - return [input_values], sess.run( - outputs, feed_dict=dict(zip(inputs, [input_values]))) + # Build list of input values either containing 1 tensor (input) or 2 tensors + # (input, filter) based on whether filter is constant or variable input. + input_shape, filter_shape = get_tensor_shapes(parameters) + values = [create_tensor_data(np.float32, input_shape)] + if not parameters["constant_filter"]: + values.append(create_tensor_data(np.float32, filter_shape)) + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) + + make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + + +def make_split_tests(zip_path): + """Make a set of tests to do tf.split.""" + + test_parameters = [{ + "input_shape": [[1, 3, 4, 6], [2, 4, 1], [6, 4], [8]], + "num_or_size_splits": [1, 2, 3, 4, 5, [2, 2]], + "axis": [0, 1, 2, 3, -4, -3, -2, -1], + }] + + def build_graph(parameters): + input_tensor = tf.placeholder( + dtype=tf.float32, name="input", shape=parameters["input_shape"]) + out = tf.split( + input_tensor, parameters["num_or_size_splits"], parameters["axis"]) + return [input_tensor], out + + def build_inputs(parameters, sess, inputs, outputs): + values = [create_tensor_data(np.float32, parameters["input_shape"])] + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) def make_concatenation_tests(zip_path): - """Make a set of tests to do concatenatinon.""" + """Make a set of tests to do concatenation.""" test_parameters = [{ "base_shape": [[1, 3, 4, 3], [3, 4]], "num_tensors": [1, 2, 3, 4, 5, 6], - "axis": [0, 1, 2, 3], + "axis": [0, 1, 2, 3, -3, -2, -1], }] def get_shape(parameters, delta): """Return a tweaked version of 'base_shape'.""" axis = parameters["axis"] shape = parameters["base_shape"][:] + if axis < 0: + axis += len(shape) if axis < len(shape): shape[axis] += delta return shape @@ -955,32 +1104,49 @@ def make_fully_connected_tests(zip_path): "shape2": [[3, 3]], "transpose_a": [True, False], "transpose_b": [True, False], + "constant_filter": [True, False], }, { "shape1": [[4, 4], [1, 4], [4]], "shape2": [[4, 4], [4, 1], [4]], "transpose_a": [False], "transpose_b": [False], + "constant_filter": [True, False], }, { "shape1": [[40, 37]], "shape2": [[37, 40]], "transpose_a": [False], "transpose_b": [False], - + "constant_filter": [True, False], }] def build_graph(parameters): + """Build a matmul graph given `parameters`.""" input_tensor1 = tf.placeholder(dtype=tf.float32, name="input1", shape=parameters["shape1"]) - input_tensor2 = create_tensor_data(np.float32, parameters["shape2"]) + + # Get input_tensor2 either as a placeholder or constants. Also get a list of + # the input tensors that are represented as placeholders. + if parameters["constant_filter"]: + input_tensor2 = create_tensor_data(np.float32, parameters["shape2"]) + input_tensors = [input_tensor1] + else: + input_tensor2 = tf.placeholder( + dtype=tf.float32, name="input2", shape=parameters["shape2"]) + input_tensors = [input_tensor1, input_tensor2] + out = tf.matmul(input_tensor1, input_tensor2, transpose_a=parameters["transpose_a"], transpose_b=parameters["transpose_b"]) - return [input_tensor1], [out] + return input_tensors, [out] def build_inputs(parameters, sess, inputs, outputs): - input_values1 = create_tensor_data(np.float32, shape=parameters["shape1"]) - return [input_values1], sess.run( - outputs, feed_dict=dict(zip(inputs, [input_values1]))) + # Build list of input values either containing 1 tensor (input_values1) or 2 + # tensors (input_values1, input_values2) based on whether the second input + # is a constant or variable input. + values = [create_tensor_data(np.float32, shape=parameters["shape1"])] + if not parameters["constant_filter"]: + values.append(create_tensor_data(np.float32, parameters["shape2"])) + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) @@ -1048,30 +1214,50 @@ def make_local_response_norm_tests(zip_path): def make_pad_tests(zip_path): """Make a set of tests to do pad.""" - test_parameters = [{ - "dtype": [tf.int32, tf.float32], - "input_shape": [[1, 1, 2, 1], [2, 1, 1, 1]], - "paddings": [[[0, 0], [0, 1], [2, 3], [0, 0]], [[0, 1], [0, 0], [0, 0], - [2, 3]]], - }, { - "dtype": [tf.int32, tf.float32], - "input_shape": [[1, 2], [0, 1, 2]], - "paddings": [[[0, 1], [2, 3]]], - }] + # TODO(nupurgarg): Add test for tf.uint8. + test_parameters = [ + { + "dtype": [tf.int32, tf.int64, tf.float32], + "input_shape": [[1, 1, 2, 1], [2, 1, 1, 1]], + "paddings": [[[0, 0], [0, 1], [2, 3], [0, 0]], [[0, 1], [0, 0], + [0, 0], [2, 3]]], + "constant_paddings": [True, False], + }, + # Non-4D use case. + { + "dtype": [tf.int32, tf.int64, tf.float32], + "input_shape": [[1, 2], [0, 1, 2]], + "paddings": [[[0, 1], [2, 3]]], + "constant_paddings": [True, False], + }, + ] def build_graph(parameters): + """Build a pad graph given `parameters`.""" input_tensor = tf.placeholder( dtype=parameters["dtype"], name="input", shape=parameters["input_shape"]) - out = tf.pad(input_tensor, paddings=parameters["paddings"]) - return [input_tensor], [out] + + # Get paddings as either a placeholder or constants. + if parameters["constant_paddings"]: + paddings = parameters["paddings"] + input_tensors = [input_tensor] + else: + shape = [len(parameters["paddings"]), 2] + paddings = tf.placeholder(dtype=tf.int32, name="padding", shape=shape) + input_tensors = [input_tensor, paddings] + + out = tf.pad(input_tensor, paddings=paddings) + return input_tensors, [out] def build_inputs(parameters, sess, inputs, outputs): - input_values = create_tensor_data(parameters["dtype"], - parameters["input_shape"]) - return [input_values], sess.run( - outputs, feed_dict=dict(zip(inputs, [input_values]))) + values = [ + create_tensor_data(parameters["dtype"], parameters["input_shape"]) + ] + if not parameters["constant_paddings"]: + values.append(np.array(parameters["paddings"])) + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) @@ -1079,7 +1265,7 @@ def make_pad_tests(zip_path): def make_reshape_tests(zip_path): """Make a set of tests to do reshape.""" - # Alll shapes below are suitable for tensors with 420 elements. + # All shapes below are suitable for tensors with 420 elements. test_parameters = [{ "dtype": [tf.float32, tf.int32], "input_shape": [[3, 4, 5, 7], [4, 105], [21, 5, 2, 2], [420]], @@ -1202,6 +1388,78 @@ def make_space_to_depth_tests(zip_path): make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) +def make_space_to_batch_nd_tests(zip_path): + """Make a set of tests to do space_to_batch_nd.""" + + # TODO(nupurgarg): Add test for uint8. + test_parameters = [ + { + "dtype": [tf.int32, tf.int64, tf.float32], + "input_shape": [[1, 2, 2, 3], [2, 2, 4, 1]], + "block_shape": [[1, 3], [2, 2]], + "paddings": [[[0, 0], [0, 0]], [[0, 0], [2, 0]], [[1, 1], [1, 1]]], + "constant_block_shape": [True, False], + "constant_paddings": [True, False], + }, + { + "dtype": [tf.float32], + "input_shape": [[2, 3, 7, 3]], + "block_shape": [[1, 3], [2, 2]], + "paddings": [[[0, 0], [2, 0]], [[1, 0], [1, 0]]], + "constant_block_shape": [True, False], + "constant_paddings": [True, False], + }, + # Non-4D use case: 1 bath dimension, 3 spatial dimensions, 2 others. + { + "dtype": [tf.float32], + "input_shape": [[1, 4, 4, 4, 1, 1]], + "block_shape": [[2, 2, 2]], + "paddings": [[[0, 0], [0, 0], [0, 0]]], + "constant_block_shape": [True, False], + "constant_paddings": [True, False], + }, + ] + + def build_graph(parameters): + """Build a space_to_batch graph given `parameters`.""" + input_tensor = tf.placeholder( + dtype=parameters["dtype"], + name="input", + shape=parameters["input_shape"]) + input_tensors = [input_tensor] + + # Get block_shape either as a const or as a placeholder (tensor). + if parameters["constant_block_shape"]: + block_shape = parameters["block_shape"] + else: + shape = [len(parameters["block_shape"])] + block_shape = tf.placeholder(dtype=tf.int32, name="shape", shape=shape) + input_tensors.append(block_shape) + + # Get paddings either as a const or as a placeholder (tensor). + if parameters["constant_paddings"]: + paddings = parameters["paddings"] + else: + shape = [len(parameters["paddings"]), 2] + paddings = tf.placeholder(dtype=tf.int32, name="paddings", shape=shape) + input_tensors.append(paddings) + + out = tf.space_to_batch_nd(input_tensor, block_shape, paddings) + return input_tensors, [out] + + def build_inputs(parameters, sess, inputs, outputs): + values = [ + create_tensor_data(parameters["dtype"], parameters["input_shape"]) + ] + if not parameters["constant_block_shape"]: + values.append(np.array(parameters["block_shape"])) + if not parameters["constant_paddings"]: + values.append(np.array(parameters["paddings"])) + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) + + make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + + def make_batch_to_space_nd_tests(zip_path): """Make a set of tests to do batch_to_space_nd.""" @@ -1211,6 +1469,8 @@ def make_batch_to_space_nd_tests(zip_path): "input_shape": [[12, 2, 2, 1]], "block_shape": [[1, 4], [2, 2], [3, 4]], "crops": [[[0, 0], [0, 0]], [[1, 1], [1, 1]]], + "constant_block_shape": [True, False], + "constant_crops": [True, False], }, # Non-4D use case: 1 bath dimension, 3 spatial dimensions, 2 others. { @@ -1218,16 +1478,128 @@ def make_batch_to_space_nd_tests(zip_path): "input_shape": [[8, 2, 2, 2, 1, 1]], "block_shape": [[2, 2, 2]], "crops": [[[0, 0], [0, 0], [0, 0]]], + "constant_block_shape": [True, False], + "constant_crops": [True, False], }, ] def build_graph(parameters): + """Build a batch_to_space graph given `parameters`.""" input_tensor = tf.placeholder( dtype=parameters["dtype"], name="input", shape=parameters["input_shape"]) - out = tf.batch_to_space_nd(input_tensor, parameters["block_shape"], - parameters["crops"]) + input_tensors = [input_tensor] + + # Get block_shape either as a const or as a placeholder (tensor). + if parameters["constant_block_shape"]: + block_shape = parameters["block_shape"] + else: + shape = [len(parameters["block_shape"])] + block_shape = tf.placeholder(dtype=tf.int32, name="shape", shape=shape) + input_tensors.append(block_shape) + + # Get crops either as a const or as a placeholder (tensor). + if parameters["constant_crops"]: + crops = parameters["crops"] + else: + shape = [len(parameters["crops"]), 2] + crops = tf.placeholder(dtype=tf.int32, name="crops", shape=shape) + input_tensors.append(crops) + + out = tf.batch_to_space_nd(input_tensor, block_shape, crops) + return input_tensors, [out] + + def build_inputs(parameters, sess, inputs, outputs): + values = [ + create_tensor_data(parameters["dtype"], parameters["input_shape"]) + ] + if not parameters["constant_block_shape"]: + values.append(np.array(parameters["block_shape"])) + if not parameters["constant_crops"]: + values.append(np.array(parameters["crops"])) + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) + + make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + + +def make_transpose_tests(zip_path): + """Make a set of tests to do transpose.""" + + # TODO(nupurgarg): Add test for uint8. + test_parameters = [{ + "dtype": [tf.int32, tf.int64, tf.float32], + "input_shape": [[2, 2, 3]], + "perm": [[0, 1, 2], [0, 2, 1]], + "constant_perm": [True, False], + }, { + "dtype": [tf.float32], + "input_shape": [[1, 2, 3, 4]], + "perm": [[0, 1, 2, 3], [3, 0, 1, 2]], + "constant_perm": [True, False], + }, { + "dtype": [tf.float32], + "input_shape": [[1, 2, 3, 4, 5]], + "perm": [[0, 1, 2, 3, 4]], + "constant_perm": [True, False], + }] + + def build_graph(parameters): + """Build a transpose graph given `parameters`.""" + input_tensor = tf.placeholder( + dtype=parameters["dtype"], + name="input", + shape=parameters["input_shape"]) + + if parameters["constant_perm"]: + perm = parameters["perm"] + input_tensors = [input_tensor] + else: + shape = [len(parameters["perm"]), 2] + perm = tf.placeholder(dtype=tf.int32, name="perm", shape=shape) + input_tensors = [input_tensor, perm] + + out = tf.transpose(input_tensor, perm=perm) + return input_tensors, [out] + + def build_inputs(parameters, sess, inputs, outputs): + values = [ + create_tensor_data(parameters["dtype"], parameters["input_shape"]) + ] + if not parameters["constant_perm"]: + values.append(np.array(parameters["perm"])) + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) + + make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + + +def make_squeeze_tests(zip_path): + """Make a set of tests to do squeeze.""" + + test_parameters = [{ + "dtype": [tf.int32, tf.float32, tf.int64], + "input_shape": [[1, 2, 1, 3, 1, 4, 1, 1]], + "axis": [ + None, [], [0, 2], [4, 7], [-1, 0, 2, 0, 7, -6], [1], [2, 3, 2], + [-1, -2, -4, -6, -8], [0, 2, 4, 6, 7], [7, 6, 4, 2, 0], [6, 6], + [0, 1, 2, 3, 4, 5, 6, 7], [-2, -3, 1, 0, 7, -5] + ], + }, { + "dtype": [tf.int32, tf.float32, tf.int64], + "input_shape": [[1]], + "axis": [None, [], [0], [-1]], + }, { + "dtype": [tf.int32, tf.float32, tf.int64], + "input_shape": [[1, 1, 1, 1, 1]], + "axis": [None, [], [0], [3, 0], [-2, 0, 3, 2]], + }] + + def build_graph(parameters): + input_tensor = tf.placeholder( + dtype=parameters["dtype"], + name="input", + shape=parameters["input_shape"]) + out = tf.squeeze(input_tensor, axis=parameters["axis"]) return [input_tensor], [out] def build_inputs(parameters, sess, inputs, outputs): @@ -1239,6 +1611,125 @@ def make_batch_to_space_nd_tests(zip_path): make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) +def make_strided_slice_tests(zip_path): + """Make a set of tests to do strided_slice.""" + + # TODO(soroosh): add test/support for uint8. + test_parameters = [ + # 4-D + { + "dtype": [tf.float32, tf.int32, tf.int64], + "index_type": [tf.int32], + "input_shape": [[12, 2, 2, 5]], + "begin": [[0, 0, 0, 0], [1, 0, 1, 0]], + "end": [[8, 2, 2, 3], [12, 2, 2, 5]], + "strides": [None, [2, 1, 3, 1]], + "begin_mask": [None, 1, 8], + "end_mask": [None, 1, 8], + "shrink_axis_mask": [None, 1, 8, 11, 15, -1], + "constant_indices": [False, True], + }, + # + { + "dtype": [tf.float32], + "index_type": [tf.int32], + "input_shape": [[12, 2, 2, 5]], + "begin": [[0]], + "end": [[1]], + "strides": [[1]], + "begin_mask": [0], + "end_mask": [0], + "shrink_axis_mask": [1], + "constant_indices": [True], + }, + # 2-D + { + "dtype": [tf.float32, tf.int32, tf.int64], + "index_type": [tf.int32], + "input_shape": [[2, 3]], + "begin": [[0, 0], [1, 0]], + "end": [[2, 3], [2, 2]], + "strides": [None, [2, 2]], + "begin_mask": [None, 1, 2], + "end_mask": [None, 1, 2], + "shrink_axis_mask": [None, 1, 2, 3, -1], + "constant_indices": [False, True], + }, + # Negative strides + { + "dtype": [tf.float32], + "index_type": [tf.int32], + "input_shape": [[2, 3]], + "begin": [[0, -1]], + "end": [[2, -3]], + "strides": [[1, -1]], + "begin_mask": [None, 1, 2], + "end_mask": [None, 1, 2], + "shrink_axis_mask": [None, 1, 2, 3, -1], + "constant_indices": [False], + }, + ] + + def build_graph(parameters): + """Build graph for stride_slice test.""" + input_tensor = tf.placeholder( + dtype=parameters["dtype"], + name="input", + shape=parameters["input_shape"]) + if parameters["constant_indices"]: + begin = parameters["begin"] + end = parameters["end"] + strides = parameters["strides"] + tensors = [input_tensor] + else: + begin = tf.placeholder( + dtype=parameters["index_type"], + name="begin", + shape=[len(parameters["input_shape"])]) + end = tf.placeholder( + dtype=parameters["index_type"], + name="end", + shape=[len(parameters["input_shape"])]) + strides = ( + tf.placeholder( + dtype=parameters["index_type"], + name="strides", + shape=[len(parameters["input_shape"])]) + if parameters["strides"] is not None else None) + tensors = [input_tensor, begin, end] + if strides is not None: + tensors.append(strides) + out = tf.strided_slice( + input_tensor, + begin, + end, + strides, + begin_mask=parameters["begin_mask"], + end_mask=parameters["end_mask"]) + return tensors, [out] + + def build_inputs(parameters, sess, inputs, outputs): + """Build inputs for stride_slice test.""" + input_values = create_tensor_data(parameters["dtype"], + parameters["input_shape"]) + index_type = _TF_TYPE_INFO[parameters["index_type"]][0] + values = [input_values] + if not parameters["constant_indices"]: + begin_values = np.array(parameters["begin"]).astype(index_type) + end_values = np.array(parameters["end"]).astype(index_type) + stride_values = ( + np.array(parameters["strides"]).astype(index_type) + if parameters["strides"] is not None else None) + values.append(begin_values) + values.append(end_values) + if stride_values is not None: + values.append(stride_values) + + return values, sess.run(outputs, feed_dict=dict(zip(inputs, values))) + + make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + + def make_l2_pool(input_tensor, ksize, strides, padding, data_format): """Given an input perform a sequence of TensorFlow ops to produce l2pool.""" return tf.sqrt(tf.nn.avg_pool( @@ -1246,6 +1737,32 @@ def make_l2_pool(input_tensor, ksize, strides, padding, data_format): padding=padding, data_format=data_format)) +def make_topk_tests(zip_path): + """Make a set of tests to do gather.""" + + test_parameters = [{ + "input_dtype": [tf.float32, tf.int32], + "input_shape": [[10], [5, 20]], + }] + + def build_graph(parameters): + """Build the gather op testing graph.""" + input_value = tf.placeholder( + dtype=parameters["input_dtype"], + name="input", + shape=parameters["input_shape"]) + k = tf.constant(3, name="k") + out = tf.nn.top_k(input_value, k) + return [input_value], [out[1]] + + def build_inputs(parameters, sess, inputs, outputs): + input_value = create_tensor_data(parameters["input_dtype"], + parameters["input_shape"]) + return [input_value], sess.run( + outputs, feed_dict=dict(zip(inputs, [input_value]))) + + make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + # Toco binary path provided by the generate rule. bin_path = None @@ -1266,7 +1783,10 @@ def main(unused_args): dispatch = { "control_dep.zip": make_control_dep_tests, - "add.zip": make_add_tests, + "add.zip": make_binary_op_tests_func(tf.add), + "space_to_batch_nd.zip": make_space_to_batch_nd_tests, + "div.zip": make_binary_op_tests_func(tf.div), + "sub.zip": make_binary_op_tests_func(tf.subtract), "batch_to_space_nd.zip": make_batch_to_space_nd_tests, "conv.zip": make_conv_tests, "constant.zip": make_constant_tests, @@ -1278,7 +1798,7 @@ def main(unused_args): "fused_batch_norm.zip": make_fused_batch_norm_tests, "l2norm.zip": make_l2norm_tests, "local_response_norm.zip": make_local_response_norm_tests, - "mul.zip": make_mul_tests, + "mul.zip": make_binary_op_tests_func(tf.multiply), "relu.zip": make_relu_tests, "relu1.zip": make_relu1_tests, "relu6.zip": make_relu6_tests, @@ -1291,6 +1811,13 @@ def main(unused_args): "sigmoid.zip": make_sigmoid_tests, "softmax.zip": make_softmax_tests, "space_to_depth.zip": make_space_to_depth_tests, + "topk.zip": make_topk_tests, + "split.zip": make_split_tests, + "transpose.zip": make_transpose_tests, + "mean.zip": make_mean_tests, + "squeeze.zip": make_squeeze_tests, + "strided_slice.zip": make_strided_slice_tests, + "exp.zip": make_exp_tests, } out = FLAGS.zip_to_output bin_path = FLAGS.toco diff --git a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc index 4c05979e245e48ac9f389ca7a430ec761ef6764e..49766cedac8d1acd96f9b38665119e99f8bb9ac0 100644 --- a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc +++ b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc @@ -20,24 +20,25 @@ limitations under the License. #include #include #include "re2/re2.h" -#include "tensorflow/contrib/lite/builtin_op_data.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" #include "tensorflow/contrib/lite/testing/parse_testdata.h" +#include "tensorflow/contrib/lite/testing/tflite_driver.h" +#include "tensorflow/contrib/lite/testing/util.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/subprocess.h" -#include "tensorflow/core/platform/test.h" #include "tensorflow/core/util/command_line_flags.h" +namespace tflite { +namespace testing { + namespace { bool FLAGS_ignore_known_bugs = true; +// TODO(b/71769302) zip_files_dir should have a more accurate default, if +// possible +string* FLAGS_zip_files_dir = new string("./"); +string* FLAGS_unzip_binary_path = new string("/usr/bin/unzip"); } // namespace -namespace tflite { -namespace testing { - // TensorFlow system environment for file system called. tensorflow::Env* env = tensorflow::Env::Default(); @@ -46,42 +47,51 @@ tensorflow::Env* env = tensorflow::Env::Default(); // Key is a substring of the test name and value is a bug number. // TODO(ahentz): make sure we clean this list up frequently. std::map kBrokenTests = { - // Add doesn't support broadcasting. - {R"(addd.*input_shape_1=\[1,3,4,3\],input_shape_2=\[3\])", "68500195"}, - {R"(muld.*input_shape_1=\[1,3,4,3\],input_shape_2=\[3\])", "68500195"}, + // Sub and Div don't support broadcasting. + {R"(^\/diva.*input_shape_1=\[1,3,4,3\],input_shape_2=\[3\])", "68500195"}, + {R"(^\/suba.*input_shape_1=\[1,3,4,3\],input_shape_2=\[3\])", "68500195"}, // Add only supports float32. (and "constant" tests use Add) - {R"(addd.*int32)", "68808744"}, - {R"(constant.*int32)", "68808744"}, - {R"(mul.*int32)", "68808744"}, + {R"(^\/adda.*int32)", "68808744"}, + {R"(^\/constant.*int32)", "68808744"}, + {R"(^\/mul.*int32)", "68808744"}, + {R"(^\/div.*int32)", "68808744"}, + {R"(^\/sub.*int32)", "68808744"}, + + // Pad only supports 4D tensors. + {R"(^\/pad.*,input_shape=\[.,.\],paddings=\[\[.,.\],\[.,.\]\])", + "70527055"}, - // Toco or TFLite has a bug to deal with some constant functions with - // more than 1 element. - {R"(constant.*input_shape=\[(2|2,2,2,2)\])", "68721522"}, + // L2Norm only supports tensors with 4D or fewer. + {R"(^\/l2normdim=.*,epsilon=.*,input_shape=\[.,.,.,.,.*\])", "67963684"}, - // Pad only supports 4D float32 tensors. - {R"(paddtype=.*,input_shape=\[.,.\],paddings=\[\[.,.\],\[.,.\]\])", - "70527055"}, - {R"(padd.*int32)", "70527055"}, + // BatchToSpaceND doesn't support cropping. This catches test cases with + // non-const tensors as crops. + {R"(^\/batch_to_space_nd.*crops=\[\[1,1\],\[1,1\]\])", "70594634"}, - // L2Norm only supports 4D tensors. - {R"(l2normdim=.*,epsilon=.*,input_shape=\[.,.\])", "67963684"}, - {R"(l2normdim=.*,epsilon=.*,input_shape=\[.,.,.,.,.*\])", "67963684"}, + // SpaceToBatchND only supports 4D tensors. + {R"(^\/space_to_batch_nd.*input_shape=\[1,4,4,4,1,1\])", "70848787"}, // L2Norm only works for dim=-1. - {R"(l2normdim=-2,epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, - {R"(l2normdim=-2,epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, - {R"(l2normdim=2,epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, - {R"(l2normdim=2,epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, - {R"(l2normdim=0,epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, - {R"(l2normdim=0,epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, - {R"(l2normdim=1,epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, - {R"(l2normdim=1,epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, - {R"(l2normdim=\[2,3\],epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, - {R"(l2normdim=\[2,3\],epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, + {R"(^\/l2normdim=-2,epsilon=.*,input_shape=\[.,.\])", "67963812"}, + {R"(^\/l2normdim=0,epsilon=.*,input_shape=\[.,.\])", "67963812"}, + {R"(^\/l2normdim=-2,epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, + {R"(^\/l2normdim=-2,epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, + {R"(^\/l2normdim=2,epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, + {R"(^\/l2normdim=2,epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, + {R"(^\/l2normdim=0,epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, + {R"(^\/l2normdim=0,epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, + {R"(^\/l2normdim=1,epsilon=.*,input_shape=\[3,15,14,3\])", "67963812"}, + {R"(^\/l2normdim=1,epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, + {R"(^\/l2normdim=\[2,3\],epsilon=.*,input_shape=\[3,15,14,3\])", + "67963812"}, + {R"(^\/l2normdim=\[2,3\],epsilon=.*,input_shape=\[1,3,4,3\])", "67963812"}, // ResizeBilinear looks completely incompatible with Tensorflow - {R"(resize_bilinear)", "67964336"}, + {R"(^\/resize_bilinear.*dtype=tf.int32)", "72401107"}, + + // Transpose only supports 1D-4D input tensors. + {R"(^\/transpose.*input_shape=\[.,.,.,.,.\])", "71545879"}, }; // Allows test data to be unzipped into a temporary directory and makes @@ -105,8 +115,9 @@ class ZipEnvironment : public ::testing::Environment { string dir; TF_CHECK_OK(MakeTemporaryDirectory(&dir)); tensorflow::SubProcess proc; - string unzip_binary = - "/usr/bin/unzip"; + string unzip_binary = *FLAGS_unzip_binary_path; + TF_CHECK_OK(env->FileExists(unzip_binary)); + TF_CHECK_OK(env->FileExists(zip)); proc.SetProgram(unzip_binary, {"unzip", "-d", dir, zip}); proc.SetChannelAction(tensorflow::CHAN_STDOUT, tensorflow::ACTION_PIPE); proc.SetChannelAction(tensorflow::CHAN_STDERR, tensorflow::ACTION_PIPE); @@ -174,8 +185,7 @@ tensorflow::Status ReadManifest(const string& original_file, const string& dir, // Get a list of tests from a zip file `zip_file_name`. std::vector UnarchiveZipAndFindTestNames(const string& zip_file_name) { - string zip_file = ::tensorflow::testing::TensorFlowSrcRoot() + - "/contrib/lite/testing/optest/" + zip_file_name; + string zip_file = *FLAGS_zip_files_dir + "/" + zip_file_name; string decompress_tmp_dir; TF_CHECK_OK(zip_environment()->UnZip(zip_file, &decompress_tmp_dir)); std::vector stuff; @@ -187,20 +197,14 @@ class OpsTest : public ::testing::TestWithParam {}; TEST_P(OpsTest, RunStuff) { string test_path = GetParam(); - string tflite_file = test_path + ".bin"; - string tflite_examples = test_path + ".inputs"; + string tflite_test_case = test_path + "_tests.txt"; + string tflite_dir = test_path.substr(0, test_path.find_last_of("/")); string test_name = test_path.substr(test_path.find_last_of('/')); - auto model = tflite::FlatBufferModel::BuildFromFile(tflite_file.c_str()); - std::unique_ptr interpreter; - - tflite::ops::builtin::BuiltinOpResolver builtins; - ASSERT_EQ(tflite::InterpreterBuilder(*model, builtins)(&interpreter), - kTfLiteOk); - - std::vector examples; - ASSERT_EQ(tflite::testing::ParseExamples(tflite_examples.c_str(), &examples), - kTfLiteOk); + std::ifstream tflite_stream(tflite_test_case); + ASSERT_TRUE(tflite_stream.is_open()) << tflite_test_case; + tflite::testing::TfLiteDriver test_driver(/*use_nnapi=*/true); + test_driver.SetModelBaseDir(tflite_dir); string bug_number; for (const auto& p : kBrokenTests) { @@ -209,25 +213,15 @@ TEST_P(OpsTest, RunStuff) { } } - for (const auto& example : examples) { - ASSERT_EQ(interpreter->inputs().size(), example.inputs.size()); - auto result = [&]() { - TF_LITE_ENSURE_STATUS(FeedExample(interpreter.get(), example)); - TF_LITE_ENSURE_STATUS(interpreter->Invoke()); - TF_LITE_ENSURE_STATUS(CheckOutputs(interpreter.get(), example)); - return kTfLiteOk; - }(); - - if (bug_number.empty()) { - ASSERT_EQ(result, kTfLiteOk); + bool result = tflite::testing::ParseAndRunTests(&tflite_stream, &test_driver); + if (bug_number.empty()) { + EXPECT_TRUE(result) << test_driver.GetErrorMessage(); + } else { + if (FLAGS_ignore_known_bugs) { + EXPECT_FALSE(result); } else { - if (FLAGS_ignore_known_bugs) { - ASSERT_EQ(result, kTfLiteError) - << "Not failing as expected due to http://b/" << bug_number; - } else { - ASSERT_EQ(result, kTfLiteOk) - << "Possibly due to http://b/" << bug_number; - } + EXPECT_TRUE(result) << test_driver.GetErrorMessage() + << ": Possibly due to http://b/" << bug_number; } } } @@ -241,12 +235,14 @@ TEST_P(OpsTest, RunStuff) { INSTANTIATE_TESTS(add) INSTANTIATE_TESTS(avg_pool) +INSTANTIATE_TESTS(space_to_batch_nd) INSTANTIATE_TESTS(batch_to_space_nd) INSTANTIATE_TESTS(concat) INSTANTIATE_TESTS(constant) INSTANTIATE_TESTS(control_dep) INSTANTIATE_TESTS(conv) INSTANTIATE_TESTS(depthwiseconv) +INSTANTIATE_TESTS(exp) INSTANTIATE_TESTS(fully_connected) INSTANTIATE_TESTS(fused_batch_norm) INSTANTIATE_TESTS(gather) @@ -265,6 +261,13 @@ INSTANTIATE_TESTS(resize_bilinear) INSTANTIATE_TESTS(sigmoid) INSTANTIATE_TESTS(softmax) INSTANTIATE_TESTS(space_to_depth) +INSTANTIATE_TESTS(sub) +INSTANTIATE_TESTS(split) +INSTANTIATE_TESTS(div) +INSTANTIATE_TESTS(transpose) +INSTANTIATE_TESTS(mean) +INSTANTIATE_TESTS(squeeze) +INSTANTIATE_TESTS(strided_slice) } // namespace testing } // namespace tflite @@ -272,16 +275,23 @@ INSTANTIATE_TESTS(space_to_depth) int main(int argc, char** argv) { ::testing::AddGlobalTestEnvironment(tflite::testing::zip_environment()); - std::vector flags = {tensorflow::Flag( - "ignore_known_bugs", &FLAGS_ignore_known_bugs, - "If a particular model is affected by a known bug, the " - "corresponding test should expect the outputs to not match.")}; + std::vector flags = { + tensorflow::Flag( + "ignore_known_bugs", &tflite::testing::FLAGS_ignore_known_bugs, + "If a particular model is affected by a known bug, the " + "corresponding test should expect the outputs to not match."), + tensorflow::Flag("zip_files_dir", tflite::testing::FLAGS_zip_files_dir, + "Required: Location of the test zips."), + tensorflow::Flag("unzip_binary_path", + tflite::testing::FLAGS_unzip_binary_path, + "Required: Location of a suitable unzip binary.")}; bool success = tensorflow::Flags::Parse(&argc, argv, flags); if (!success || (argc == 2 && !strcmp(argv[1], "--helpfull"))) { fprintf(stderr, "%s", tensorflow::Flags::Usage(argv[0], flags).c_str()); return 1; } + ::tflite::LogToStderr(); ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); } diff --git a/tensorflow/contrib/lite/testing/join.h b/tensorflow/contrib/lite/testing/join.h new file mode 100644 index 0000000000000000000000000000000000000000..ce8c072a21c6e61e8ab8ae12ba52418e6144009a --- /dev/null +++ b/tensorflow/contrib/lite/testing/join.h @@ -0,0 +1,42 @@ +/* Copyright 2017 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_LITE_TESTING_JOIN_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_JOIN_H_ + +#include +#include +#include + +namespace tflite { +namespace testing { + +// Join a list of data separated by delimieter. +template +string Join(T* data, size_t len, const string& delimiter) { + if (len == 0 || data == nullptr) { + return ""; + } + std::stringstream result; + result << data[0]; + for (int i = 1; i < len; i++) { + result << delimiter << data[i]; + } + return result.str(); +} + +} // namespace testing +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_JOIN_H_ diff --git a/tensorflow/contrib/lite/testing/join_test.cc b/tensorflow/contrib/lite/testing/join_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..bd04528381f6d31164728a5cabbf8753e9b8d2b8 --- /dev/null +++ b/tensorflow/contrib/lite/testing/join_test.cc @@ -0,0 +1,43 @@ +/* Copyright 2017 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/lite/testing/join.h" + +#include +#include + +namespace tflite { +namespace testing { +namespace { + +TEST(JoinTest, JoinInt) { + std::vector data = {1, 2, 3}; + EXPECT_EQ(Join(data.data(), data.size(), ","), "1,2,3"); +} + +TEST(JoinTest, JoinFloat) { + float data[] = {1.0, -3, 2.3, 1e-5}; + EXPECT_EQ(Join(data, 4, " "), "1 -3 2.3 1e-05"); +} + +TEST(JoinTest, JoinNullData) { EXPECT_THAT(Join(nullptr, 3, ","), ""); } + +TEST(JoinTest, JoinZeroData) { + std::vector data; + EXPECT_THAT(Join(data.data(), 0, ","), ""); +} + +} // namespace +} // namespace testing +} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/message.h b/tensorflow/contrib/lite/testing/message.h index 78ef7e2cbe1c323753ac36f1be06a089e650aa37..e2bc4082141f0601c141a193fbea75f8f759146a 100644 --- a/tensorflow/contrib/lite/testing/message.h +++ b/tensorflow/contrib/lite/testing/message.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_MESSAGE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_MESSAGE_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TESTING_MESSAGE_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_MESSAGE_H_ #include #include @@ -79,4 +79,4 @@ class Message { } // namespace testing } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_MESSAGE_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_MESSAGE_H_ diff --git a/tensorflow/contrib/lite/testing/nnapi_example.cc b/tensorflow/contrib/lite/testing/nnapi_example.cc index 74f6cfc3de5d209671c38595434a43128966bb0e..5870782b69217f292fe60821ea8ce4ea1174c495 100644 --- a/tensorflow/contrib/lite/testing/nnapi_example.cc +++ b/tensorflow/contrib/lite/testing/nnapi_example.cc @@ -19,80 +19,35 @@ limitations under the License. // Usage: bazel run -c opt \ // tensorflow/contrib/lite/nnapi:nnapi_example -- // +#include #include #include -#include "tensorflow/contrib/lite/builtin_op_data.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" +#include +#include +#include #include "tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h" #include "tensorflow/contrib/lite/testing/parse_testdata.h" +#include "tensorflow/contrib/lite/testing/tflite_driver.h" -// TODO(aselle): FATAL leaves resources hanging. -void FATAL(const char* format, ...) { - va_list args; - va_start(args, format); - vfprintf(stderr, format, args); - va_end(args); - fflush(stderr); - exit(1); -} +string dirname(const string& s) { return s.substr(0, s.find_last_of("/")); } -#define CHECK_TFLITE_SUCCESS(x) \ - if (x != kTfLiteOk) { \ - FATAL("Aborting since tflite returned failure."); \ +bool Interpret(const char* examples_filename, bool use_nnapi) { + std::ifstream tflite_stream(examples_filename); + if (!tflite_stream.is_open()) { + fprintf(stderr, "Can't open input file."); + return false; } -void Interpret(const char* filename, const char* examples_filename, - bool use_nnapi) { - // TODO(aselle): Resize of input image should go here - // ... - // For now I am allocating all tensors. This means I am fixed size. - // So I am not using the variable size ability yet. - fprintf(stderr, "example file %s\n", examples_filename); - std::vector examples; - CHECK_TFLITE_SUCCESS( - tflite::testing::ParseExamples(examples_filename, &examples)); - - for (const tflite::testing::Example& example : examples) { - auto model = tflite::FlatBufferModel::BuildFromFile(filename); - if (!model) FATAL("Cannot read file %s\n", filename); - std::unique_ptr interpreter; - tflite::ops::builtin::BuiltinOpResolver builtins; - - CHECK_TFLITE_SUCCESS( - tflite::InterpreterBuilder(*model, builtins)(&interpreter)); + printf("Use nnapi is set to: %d\n", use_nnapi); + tflite::testing::TfLiteDriver test_driver(use_nnapi); - printf("Use nnapi is set to: %d\n", use_nnapi); - interpreter->UseNNAPI(use_nnapi); - CHECK_TFLITE_SUCCESS( - tflite::testing::FeedExample(interpreter.get(), example)); - - { - TfLiteTensor* tensor = interpreter->tensor(interpreter->outputs()[0]); - if (float* data = - interpreter->typed_tensor(interpreter->outputs()[0])) { - size_t num = tensor->bytes / sizeof(float); - for (float* p = data; p < data + num; p++) { - *p = 0; - } - } - } - interpreter->Invoke(); - - CHECK_TFLITE_SUCCESS( - tflite::testing::CheckOutputs(interpreter.get(), example)); - - printf("Result:\n"); - TfLiteTensor* tensor = interpreter->tensor(interpreter->outputs()[0]); - if (float* data = - interpreter->typed_tensor(interpreter->outputs()[0])) { - size_t num = tensor->bytes / sizeof(float); - for (float* p = data; p < data + num; p++) { - printf(" %f", *p); - } - } + test_driver.SetModelBaseDir(dirname(examples_filename)); + if (!tflite::testing::ParseAndRunTests(&tflite_stream, &test_driver)) { + fprintf(stderr, "Results from tflite don't match."); + return false; } + + return true; } int main(int argc, char* argv[]) { @@ -109,6 +64,25 @@ int main(int argc, char* argv[]) { argv[0]); return 1; } - Interpret(argv[1], argv[2], use_nnapi); + + string base_dir = dirname(argv[1]); + DIR* dir = opendir(base_dir.c_str()); + if (dir == nullptr) { + fprintf(stderr, "Can't open dir %s\n", base_dir.c_str()); + return 1; + } + while (struct dirent* ent = readdir(dir)) { + string name = ent->d_name; + if (name.rfind(".txt") == name.length() - 4) { + printf("%s: ", name.c_str()); + if (Interpret((base_dir + "/" + name).c_str(), use_nnapi)) { + printf(" %s\n", "OK"); + } else { + printf(" %s\n", "FAIL"); + } + } + } + closedir(dir); + return 0; } diff --git a/tensorflow/contrib/lite/testing/parse_testdata.cc b/tensorflow/contrib/lite/testing/parse_testdata.cc index d745ed27158cdad55bdcd97162cb3dfa9e32c112..0caef0fe2201a668b2235a98304eb353072a3c2f 100644 --- a/tensorflow/contrib/lite/testing/parse_testdata.cc +++ b/tensorflow/contrib/lite/testing/parse_testdata.cc @@ -18,6 +18,7 @@ limitations under the License. // ASCII file. #include "tensorflow/contrib/lite/testing/parse_testdata.h" +#include #include #include #include @@ -169,6 +170,11 @@ TfLiteStatus FeedExample(tflite::Interpreter* interpreter, for (size_t idx = 0; idx < example.inputs[i].flat_data.size(); idx++) { data[idx] = example.inputs[i].flat_data[idx]; } + } else if (int64_t* data = + interpreter->typed_tensor(input_index)) { + for (size_t idx = 0; idx < example.inputs[i].flat_data.size(); idx++) { + data[idx] = example.inputs[i].flat_data[idx]; + } } else { fprintf(stderr, "input[%zu] was not float or int data\n", i); return kTfLiteError; @@ -213,8 +219,22 @@ TfLiteStatus CheckOutputs(tflite::Interpreter* interpreter, int32_t computed = data[idx]; int32_t reference = example.outputs[0].flat_data[idx]; if (std::abs(computed - reference) > 0) { - fprintf(stderr, "output[%zu][%zu] did not match %d vs reference %f\n", - i, idx, data[idx], example.outputs[0].flat_data[idx]); + fprintf(stderr, "output[%zu][%zu] did not match %d vs reference %d\n", + i, idx, computed, reference); + return kTfLiteError; + } + } + fprintf(stderr, "\n"); + } else if (const int64_t* data = + interpreter->typed_tensor(output_index)) { + for (size_t idx = 0; idx < example.outputs[i].flat_data.size(); idx++) { + int64_t computed = data[idx]; + int64_t reference = example.outputs[0].flat_data[idx]; + if (std::abs(computed - reference) > 0) { + fprintf(stderr, + "output[%zu][%zu] did not match %" PRId64 + " vs reference %" PRId64 "\n", + i, idx, computed, reference); return kTfLiteError; } } diff --git a/tensorflow/contrib/lite/testing/parse_testdata.h b/tensorflow/contrib/lite/testing/parse_testdata.h index 90839fe24550b6c4a0a3a3f4115c479a71580bb0..7ebf362eb99c5f4cf6ea3654cf71e13ff1de99b3 100644 --- a/tensorflow/contrib/lite/testing/parse_testdata.h +++ b/tensorflow/contrib/lite/testing/parse_testdata.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_NNAPI_PARSE_TESTDATA_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_NNAPI_PARSE_TESTDATA_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_NNAPI_PARSE_TESTDATA_H_ +#define TENSORFLOW_CONTRIB_LITE_NNAPI_PARSE_TESTDATA_H_ #include #include "tensorflow/contrib/lite/interpreter.h" @@ -71,4 +71,4 @@ bool ParseAndRunTests(std::istream* input, TestRunner* test_runner); } // namespace testing } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_NNAPI_PARSE_TESTDATA_H_ +#endif // TENSORFLOW_CONTRIB_LITE_NNAPI_PARSE_TESTDATA_H_ diff --git a/tensorflow/contrib/lite/testing/split.h b/tensorflow/contrib/lite/testing/split.h index 24071442e8929f37443df1b98d22711b3024b87c..428cfda4f216f0ee6409a32c43a4cf91ecc11922 100644 --- a/tensorflow/contrib/lite/testing/split.h +++ b/tensorflow/contrib/lite/testing/split.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_SPLIT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_SPLIT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TESTING_SPLIT_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_SPLIT_H_ #include #include @@ -53,6 +53,15 @@ inline std::vector Split(const string& s, const string& delimiter) { return fields; } +template <> +inline std::vector Split(const string& s, const string& delimiter) { + std::vector fields; + for (const auto& p : SplitToPos(s, delimiter)) { + fields.push_back(strtoll(s.data() + p.first, nullptr, 10)); + } + return fields; +} + template <> inline std::vector Split(const string& s, const string& delimiter) { std::vector fields; @@ -74,4 +83,4 @@ inline std::vector Split(const string& s, const string& delimiter) { } // namespace testing } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_SPLIT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_SPLIT_H_ diff --git a/tensorflow/contrib/lite/testing/test_runner.h b/tensorflow/contrib/lite/testing/test_runner.h index f4b26949b57e0702ac5554afd766a6072af268a4..05770beee23275ebe210606dbfd2b33eea17612d 100644 --- a/tensorflow/contrib/lite/testing/test_runner.h +++ b/tensorflow/contrib/lite/testing/test_runner.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TEST_RUNNER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TEST_RUNNER_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TESTING_TEST_RUNNER_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_TEST_RUNNER_H_ #include #include @@ -68,6 +68,10 @@ class TestRunner { // satisfied. virtual bool CheckResults() = 0; + // Read contents of tensor into csv format. + // The given 'id' is guaranteed to be one of the ids returned by GetOutputs(). + virtual string ReadOutput(int id) = 0; + // Set the base path for loading models. void SetModelBaseDir(const string& path) { model_base_dir_ = path; @@ -121,4 +125,4 @@ class TestRunner { } // namespace testing } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TEST_RUNNER_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_TEST_RUNNER_H_ diff --git a/tensorflow/contrib/lite/testing/test_runner_test.cc b/tensorflow/contrib/lite/testing/test_runner_test.cc index f712a5347a042990ae5adb9d44325dd683193168..3f04aa20bd7de813f0acd3f5897d5ab2df6c0fd7 100644 --- a/tensorflow/contrib/lite/testing/test_runner_test.cc +++ b/tensorflow/contrib/lite/testing/test_runner_test.cc @@ -31,6 +31,7 @@ class ConcreteTestRunner : public TestRunner { void ResetTensor(int id) override {} void SetInput(int id, const string& csv_values) override {} void SetExpectation(int id, const string& csv_values) override {} + string ReadOutput(int id) override { return ""; } void Invoke() override {} bool CheckResults() override { return true; } bool CheckFloatSizes(size_t bytes, size_t values) { diff --git a/tensorflow/contrib/lite/testing/tf_driver.cc b/tensorflow/contrib/lite/testing/tf_driver.cc new file mode 100644 index 0000000000000000000000000000000000000000..2c253bb1983e5ddc5bc12858c929585d1bcee710 --- /dev/null +++ b/tensorflow/contrib/lite/testing/tf_driver.cc @@ -0,0 +1,182 @@ +/* Copyright 2017 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/lite/testing/tf_driver.h" + +#include +#include + +#include "tensorflow/contrib/lite/testing/join.h" +#include "tensorflow/contrib/lite/testing/split.h" +#include "tensorflow/core/lib/gtl/array_slice.h" + +namespace tflite { +namespace testing { + +namespace { + +tensorflow::Tensor CreateTensor(const tensorflow::DataType type, + const std::vector& dim) { + tensorflow::TensorShape shape{gtl::ArraySlice{ + reinterpret_cast(dim.data()), dim.size()}}; + return {type, shape}; +} + +template +void FillTensorWithData(tensorflow::Tensor* tensor, const string& csv_values) { + auto data = tensor->flat(); + + const auto& values = testing::Split(csv_values, ","); + for (int i = 0; i < values.size(); i++) { + data(i) = values[i]; + } +} + +template +void FillTensorWithZeros(tensorflow::Tensor* tensor) { + auto data = tensor->flat(); + for (int i = 0; i < tensor->NumElements(); i++) { + data(i) = 0; + } +} + +template +string TensorDataToCsvString(const tensorflow::Tensor& tensor) { + const auto& data = tensor.flat(); + return Join(data.data(), data.size(), ","); +} + +} // namespace + +TfDriver::TfDriver(const std::vector& input_layer, + const std::vector& input_layer_type, + const std::vector& input_layer_shape, + const std::vector& output_layer) + : input_names_(input_layer), output_names_(output_layer) { + CHECK_EQ(input_layer.size(), input_layer_type.size()); + CHECK_EQ(input_layer.size(), input_layer_shape.size()); + + input_ids_.resize(input_layer.size()); + input_tensors_.reserve(input_layer.size()); + input_types_.resize(input_layer.size()); + input_shapes_.resize(input_layer.size()); + for (int i = 0; i < input_layer.size(); i++) { + input_ids_[i] = i; + input_tensors_[input_layer[i]] = {}; + CHECK(DataTypeFromString(input_layer_type[i], &input_types_[i])); + input_shapes_[i] = Split(input_layer_shape[i], ","); + } + + output_ids_.resize(output_layer.size()); + output_tensors_.reserve(output_layer.size()); + for (int i = 0; i < output_layer.size(); i++) { + output_ids_[i] = i; + } +} + +void TfDriver::LoadModel(const string& bin_file_path) { + if (!IsValid()) return; + std::cout << std::endl << "Loading model: " << bin_file_path << std::endl; + std::ifstream model(bin_file_path); + if (model.fail()) { + Invalidate("Failed to find the model"); + return; + } + + tensorflow::GraphDef graphdef; + if (!graphdef.ParseFromIstream(&model)) { + Invalidate("Failed to parse tensorflow graphdef"); + return; + } + + tensorflow::SessionOptions options; + session_.reset(tensorflow::NewSession(options)); + auto status = session_->Create(graphdef); + if (!status.ok()) { + Invalidate("Failed to create session"); + } +} + +void TfDriver::SetInput(int id, const string& csv_values) { + if (!IsValid()) return; + + auto tensor = CreateTensor(input_types_[id], input_shapes_[id]); + switch (input_types_[id]) { + case tensorflow::DT_FLOAT: { + FillTensorWithData(&tensor, csv_values); + break; + } + case tensorflow::DT_INT32: { + FillTensorWithData(&tensor, csv_values); + break; + } + default: + fprintf(stderr, "Unsupported type %d in SetInput\n", input_types_[id]); + Invalidate("Unsupported tensor data type"); + return; + } + input_tensors_[input_names_[id]] = tensor; +} + +void TfDriver::ResetTensor(int id) { + if (!IsValid()) return; + auto tensor = input_tensors_[input_names_[id]]; + switch (input_types_[id]) { + case tensorflow::DT_FLOAT: { + FillTensorWithZeros(&tensor); + break; + } + case tensorflow::DT_INT32: { + FillTensorWithZeros(&tensor); + break; + } + default: + fprintf(stderr, "Unsupported type %d in ResetTensor\n", input_types_[id]); + Invalidate("Unsupported tensor data type"); + return; + } +} + +void TfDriver::ReshapeTensor(int id, const string& csv_values) { + input_shapes_[id] = Split(csv_values, ","); + input_tensors_[input_names_[id]] = + CreateTensor(input_types_[id], input_shapes_[id]); + ResetTensor(id); +} + +string TfDriver::ReadOutput(int id) { + if (!IsValid()) return ""; + switch (output_tensors_[id].dtype()) { + case tensorflow::DT_FLOAT: + return TensorDataToCsvString(output_tensors_[id]); + case tensorflow::DT_INT32: + return TensorDataToCsvString(output_tensors_[id]); + default: + fprintf(stderr, "Unsupported type %d in ResetTensor\n", input_types_[id]); + Invalidate("Unsupported tensor data type"); + return ""; + } +} + +void TfDriver::Invoke() { + if (!IsValid()) return; + auto status = session_->Run({input_tensors_.begin(), input_tensors_.end()}, + output_names_, {}, &output_tensors_); + if (!status.ok()) { + Invalidate("Failed to invoke interpreter"); + } +} + +} // namespace testing +} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/tf_driver.h b/tensorflow/contrib/lite/testing/tf_driver.h new file mode 100644 index 0000000000000000000000000000000000000000..b766f85c4ddee9fb7b1513c264d4159e694770ca --- /dev/null +++ b/tensorflow/contrib/lite/testing/tf_driver.h @@ -0,0 +1,75 @@ +/* Copyright 2017 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_LITE_TESTING_TF_DRIVER_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_TF_DRIVER_H_ + +#include +#include + +#include "tensorflow/contrib/lite/testing/split.h" +#include "tensorflow/contrib/lite/testing/test_runner.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/public/session.h" + +namespace tflite { +namespace testing { + +// A test runner that feeds inputs into Tensorflow and generates outputs. +class TfDriver : public TestRunner { + public: + explicit TfDriver(const std::vector& input_layer, + const std::vector& input_layer_type, + const std::vector& input_layer_shape, + const std::vector& output_layer); + ~TfDriver() override {} + + void LoadModel(const string& bin_file_path) override; + void SetInput(int id, const string& csv_values) override; + void Invoke() override; + string ReadOutput(int id) override; + + const std::vector& GetInputs() override { return input_ids_; } + const std::vector& GetOutputs() override { return output_ids_; } + void ReshapeTensor(int id, const string& csv_values) override; + // Note: ResetTensor only works for input tensor. + void ResetTensor(int id) override; + + // no-op. SetInput will overwrite existing data . + void AllocateTensors() override {} + // no-op. Tf driver is not supposed to check the results. + void SetExpectation(int id, const string& csv_values) override {} + // tf driver is not supposed to check the results. + bool CheckResults() override { return false; } + + private: + std::unique_ptr session_; + std::vector input_ids_; + std::vector input_names_; + std::vector> input_shapes_; + std::vector input_types_; + std::unordered_map input_tensors_; + + std::vector output_ids_; + std::vector output_names_; + std::vector<::tensorflow::Tensor> output_tensors_; +}; + +} // namespace testing +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_TF_DRIVER_H_ diff --git a/tensorflow/contrib/lite/testing/tf_driver_test.cc b/tensorflow/contrib/lite/testing/tf_driver_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..c0faa4676adc3e846ad398bb203b77b99a2ba360 --- /dev/null +++ b/tensorflow/contrib/lite/testing/tf_driver_test.cc @@ -0,0 +1,56 @@ +/* Copyright 2017 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/lite/testing/tf_driver.h" + +#include +#include + +namespace tflite { +namespace testing { +namespace { + +using ::testing::ElementsAre; + +TEST(TfDriverTest, SimpleTest) { + std::unique_ptr runner( + new TfDriver({"a", "b", "c", "d"}, {"float", "float", "float", "float"}, + {"1,8,8,3", "1,8,8,3", "1,8,8,3", "1,8,8,3"}, {"x", "y"})); + + runner->LoadModel( + "third_party/tensorflow/contrib/lite/testdata/multi_add.pb"); + EXPECT_TRUE(runner->IsValid()) << runner->GetErrorMessage(); + + ASSERT_THAT(runner->GetInputs(), ElementsAre(0, 1, 2, 3)); + ASSERT_THAT(runner->GetOutputs(), ElementsAre(0, 1)); + + for (int i : {0, 1, 2, 3}) { + runner->ReshapeTensor(i, "1,2,2,1"); + } + ASSERT_TRUE(runner->IsValid()); + + runner->SetInput(0, "0.1,0.2,0.3,0.4"); + runner->SetInput(1, "0.001,0.002,0.003,0.004"); + runner->SetInput(2, "0.001,0.002,0.003,0.004"); + runner->SetInput(3, "0.01,0.02,0.03,0.04"); + runner->ResetTensor(2); + runner->Invoke(); + + ASSERT_EQ(runner->ReadOutput(0), "0.101,0.202,0.303,0.404"); + ASSERT_EQ(runner->ReadOutput(1), "0.011,0.022,0.033,0.044"); +} + +} // namespace +} // namespace testing +} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/tflite_driver.cc b/tensorflow/contrib/lite/testing/tflite_driver.cc index cf9df2ec264bcff7f836a70db37afe8a5ce01c28..613223f3d4ff212cb8672494243b2d7a1d06b3db 100644 --- a/tensorflow/contrib/lite/testing/tflite_driver.cc +++ b/tensorflow/contrib/lite/testing/tflite_driver.cc @@ -31,6 +31,14 @@ float Value(const TfLitePtrUnion& data, int index) { return data.f[index]; } template <> +int32_t Value(const TfLitePtrUnion& data, int index) { + return data.i32[index]; +} +template <> +int64_t Value(const TfLitePtrUnion& data, int index) { + return data.i64[index]; +} +template <> uint8_t Value(const TfLitePtrUnion& data, int index) { return data.uint8[index]; } @@ -61,9 +69,14 @@ class TfLiteDriver::Expectation { switch (tensor.type) { case kTfLiteFloat32: return TypedCheck(verbose, tensor); + case kTfLiteInt32: + return TypedCheck(verbose, tensor); + case kTfLiteInt64: + return TypedCheck(verbose, tensor); case kTfLiteUInt8: return TypedCheck(verbose, tensor); default: + fprintf(stderr, "Unsupported type %d in Check\n", tensor.type); return false; } } @@ -71,15 +84,30 @@ class TfLiteDriver::Expectation { private: template bool TypedCheck(bool verbose, const TfLiteTensor& tensor) { + // TODO(ahentz): must find a way to configure the tolerance. + constexpr double kRelativeThreshold = 1e-2f; + constexpr double kAbsoluteThreshold = 1e-4f; + int tensor_size = tensor.bytes / sizeof(T); bool good_output = true; for (int i = 0; i < tensor_size; ++i) { - if (std::abs(Value(data_, i) - Value(tensor.data, i)) > 1e-5) { + float computed = Value(tensor.data, i); + float reference = Value(data_, i); + float diff = std::abs(computed - reference); + bool error_is_large = false; + // For very small numbers, try absolute error, otherwise go with + // relative. + if (std::abs(reference) < kRelativeThreshold) { + error_is_large = (diff > kAbsoluteThreshold); + } else { + error_is_large = (diff > kRelativeThreshold * std::abs(reference)); + } + if (error_is_large) { good_output = false; if (verbose) { - std::cerr << " index " << i << ": " << Value(data_, i) - << " != " << Value(tensor.data, i) << std::endl; + std::cerr << " index " << i << ": got " << computed + << ", but expected " << reference << std::endl; } } } @@ -95,8 +123,8 @@ TfLiteDriver::~TfLiteDriver() {} void TfLiteDriver::AllocateTensors() { if (must_allocate_tensors_) { if (interpreter_->AllocateTensors() != kTfLiteOk) { - std::cerr << "Failed to allocate tensors" << std::endl; - abort(); + Invalidate("Failed to allocate tensors"); + return; } must_allocate_tensors_ = false; } @@ -147,6 +175,18 @@ void TfLiteDriver::SetInput(int id, const string& csv_values) { SetTensorData(values, &tensor->data); break; } + case kTfLiteInt32: { + const auto& values = testing::Split(csv_values, ","); + if (!CheckSizes(tensor->bytes, values.size())) return; + SetTensorData(values, &tensor->data); + break; + } + case kTfLiteInt64: { + const auto& values = testing::Split(csv_values, ","); + if (!CheckSizes(tensor->bytes, values.size())) return; + SetTensorData(values, &tensor->data); + break; + } case kTfLiteUInt8: { const auto& values = testing::Split(csv_values, ","); if (!CheckSizes(tensor->bytes, values.size())) return; @@ -154,6 +194,7 @@ void TfLiteDriver::SetInput(int id, const string& csv_values) { break; } default: + fprintf(stderr, "Unsupported type %d in SetInput\n", tensor->type); Invalidate("Unsupported tensor data type"); return; } @@ -162,15 +203,26 @@ void TfLiteDriver::SetInput(int id, const string& csv_values) { void TfLiteDriver::SetExpectation(int id, const string& csv_values) { if (!IsValid()) return; auto* tensor = interpreter_->tensor(id); + if (expected_output_.count(id) != 0) { + fprintf(stderr, "Overriden expectation for tensor %d\n", id); + Invalidate("Overriden expectation"); + } expected_output_[id].reset(new Expectation); switch (tensor->type) { case kTfLiteFloat32: expected_output_[id]->SetData(csv_values); break; + case kTfLiteInt32: + expected_output_[id]->SetData(csv_values); + break; + case kTfLiteInt64: + expected_output_[id]->SetData(csv_values); + break; case kTfLiteUInt8: expected_output_[id]->SetData(csv_values); break; default: + fprintf(stderr, "Unsupported type %d in SetExpectation\n", tensor->type); Invalidate("Unsupported tensor data type"); return; } diff --git a/tensorflow/contrib/lite/testing/tflite_driver.h b/tensorflow/contrib/lite/testing/tflite_driver.h index 4440d4285e948c3d1622c8de5c47ff3729c5847f..02b7de1534e648734d7bc53154afa42f2ef256b4 100644 --- a/tensorflow/contrib/lite/testing/tflite_driver.h +++ b/tensorflow/contrib/lite/testing/tflite_driver.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ #include @@ -45,6 +45,7 @@ class TfLiteDriver : public TestRunner { void SetExpectation(int id, const string& csv_values) override; void Invoke() override; bool CheckResults() override; + string ReadOutput(int id) override { return "no-op"; } private: class Expectation; @@ -59,4 +60,4 @@ class TfLiteDriver : public TestRunner { } // namespace testing } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ diff --git a/tensorflow/contrib/lite/testing/tokenize.h b/tensorflow/contrib/lite/testing/tokenize.h index daccf0e84a450a0ffdf04a1eb8ff319878cfc808..7ed8eb96b7a10eecd915fe426ab3abf0e7a46ca4 100644 --- a/tensorflow/contrib/lite/testing/tokenize.h +++ b/tensorflow/contrib/lite/testing/tokenize.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TOKENIZER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TOKENIZER_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TESTING_TOKENIZER_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_TOKENIZER_H_ #include #include @@ -39,4 +39,4 @@ void Tokenize(std::istream* input, TokenProcessor* processor); } // namespace testing } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_TOKENIZER_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_TOKENIZER_H_ diff --git a/tensorflow/contrib/lite/testing/util.h b/tensorflow/contrib/lite/testing/util.h index 4d4304f022187027950f58050ececae73dedffb6..6d20aec141c7c3a3e48af290edb169c6fd7254cf 100644 --- a/tensorflow/contrib/lite/testing/util.h +++ b/tensorflow/contrib/lite/testing/util.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TESTING_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_TESTING_UTIL_H_ namespace tflite { @@ -25,4 +25,4 @@ inline void LogToStderr() { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TESTING_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TESTING_UTIL_H_ diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD index fadd2d940747236e89f0bc19d369f745978a3cc2..e2879fad327d965799d84da5c9092a12a36aa65b 100644 --- a/tensorflow/contrib/lite/toco/BUILD +++ b/tensorflow/contrib/lite/toco/BUILD @@ -160,6 +160,7 @@ cc_library( ], deps = [ # Placeholder for internal file dependency. + "@protobuf_archive//:protobuf_headers", "//tensorflow/core:framework_lite", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", @@ -171,6 +172,10 @@ cc_library( srcs = [ "graph_transformations/convert_expanddims_to_reshape.cc", "graph_transformations/convert_pure_conv_to_depthwise.cc", + "graph_transformations/convert_reorder_axes.cc", + "graph_transformations/convert_trivial_addn_to_add.cc", + "graph_transformations/convert_trivial_stack_to_reshape.cc", + "graph_transformations/convert_trivial_transpose_to_reshape.cc", "graph_transformations/create_im2col_arrays.cc", "graph_transformations/dequantize.cc", "graph_transformations/drop_fake_quant.cc", @@ -184,7 +189,10 @@ cc_library( "graph_transformations/identify_l2_normalization.cc", "graph_transformations/identify_l2_pool.cc", "graph_transformations/identify_lstm.cc", + "graph_transformations/identify_lstm_merge_inputs.cc", + "graph_transformations/identify_lstm_split_inputs.cc", "graph_transformations/identify_relu1.cc", + "graph_transformations/lstm_utils.cc", "graph_transformations/make_initial_dequantize_operator.cc", "graph_transformations/propagate_array_data_types.cc", "graph_transformations/propagate_fixed_sizes.cc", @@ -200,31 +208,42 @@ cc_library( "graph_transformations/remove_trivial_passthrough.h", "graph_transformations/remove_trivial_quantized_activation_func.cc", "graph_transformations/remove_trivial_reshape.cc", + "graph_transformations/remove_trivial_slice.cc", "graph_transformations/remove_unused_op.cc", + "graph_transformations/reorder_activation_functions.cc", "graph_transformations/resolve_batch_normalization.cc", "graph_transformations/resolve_batch_to_space_nd_attributes.cc", "graph_transformations/resolve_constant_binary.cc", "graph_transformations/resolve_constant_concatenation.cc", "graph_transformations/resolve_constant_fake_quant.cc", "graph_transformations/resolve_constant_fill.cc", - "graph_transformations/resolve_constant_tensorflow_shape.cc", + "graph_transformations/resolve_constant_range.cc", + "graph_transformations/resolve_constant_shape_or_rank.cc", + "graph_transformations/resolve_constant_stack.cc", + "graph_transformations/resolve_constant_strided_slice.cc", + "graph_transformations/resolve_constant_transpose.cc", "graph_transformations/resolve_constant_unary.cc", "graph_transformations/resolve_mean_attributes.cc", + "graph_transformations/resolve_multiply_by_zero.cc", "graph_transformations/resolve_pad_attributes.cc", "graph_transformations/resolve_reorder_axes.cc", "graph_transformations/resolve_reshape_attributes.cc", "graph_transformations/resolve_slice_attributes.cc", + "graph_transformations/resolve_space_to_batch_nd_attributes.cc", + "graph_transformations/resolve_squeeze_attributes.cc", "graph_transformations/resolve_strided_slice_attributes.cc", "graph_transformations/resolve_tensorflow_concat.cc", "graph_transformations/resolve_tensorflow_matmul.cc", "graph_transformations/resolve_tensorflow_merge.cc", - "graph_transformations/resolve_tensorflow_squeeze.cc", "graph_transformations/resolve_tensorflow_switch.cc", "graph_transformations/resolve_tensorflow_tile.cc", + "graph_transformations/resolve_transpose_attributes.cc", "graph_transformations/unfuse_activation_functions.cc", + "graph_transformations/unroll_batch_matmul.cc", ], hdrs = [ "graph_transformations/graph_transformations.h", + "graph_transformations/lstm_utils.h", ], visibility = ["//visibility:public"], deps = [ @@ -235,6 +254,7 @@ cc_library( ":tooling_util", ":types_proto_cc", "//tensorflow/core:lib", + "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], ) diff --git a/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc b/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc index 62e7282d16aa9aa02d6ebe131ead569282518753..49cc1fc2aa365925cde86ceb658ff2b354d06911 100644 --- a/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc +++ b/tensorflow/contrib/lite/toco/allocate_transient_arrays.cc @@ -148,7 +148,7 @@ std::size_t TransientArraySize(const Model& model, const string& array_name, if (!IsAllocatableTransientArray(model, array_name)) { return 0; } - const auto& array = model.arrays.at(array_name); + const auto& array = &model.GetArray(array_name); CHECK(array->has_shape()) << "Array '" << array_name << "' doesn't have a shape"; if (array->data_type == ArrayDataType::kNone) { @@ -158,9 +158,7 @@ std::size_t TransientArraySize(const Model& model, const string& array_name, LOG(FATAL) << "A RNN state array, " << array_name << ", still does not " << "have a known data type after all graph transformations have " - << "run. That's mostly a toco bug --- sorry. For now, you can " - << "work around this issue by adding manually_create:true in the " - << "--rnn_state description of this RNN state."; + << "run."; } } LOG(FATAL) << "An array, " << array_name << ", still does not " @@ -185,7 +183,7 @@ void AllocateTransientArray(const Model& model, const string& array_name, } const std::size_t size = TransientArraySize(model, array_name, transient_data_alignment); - const auto& array = model.arrays.at(array_name); + const auto& array = &model.GetArray(array_name); CHECK(!array->alloc); allocator->Allocate(size, &array->GetOrCreateAlloc()); } @@ -197,7 +195,7 @@ void DeallocateTransientArray(const Model& model, const string& array_name, if (!IsAllocatableTransientArray(model, array_name)) { return; } - const auto& array = model.arrays.at(array_name); + const auto& array = &model.GetArray(array_name); CHECK(!!array->alloc); allocator->Deallocate(*array->alloc); } @@ -231,7 +229,7 @@ void AllocateTransientArrays(Model* model, // Construct a sorted map of array names, so that other layout engines can // match exactly. std::map ordered_arrays_map; - for (const auto& pair : model->arrays) { + for (const auto& pair : model->GetArrayMap()) { ordered_arrays_map[pair.first] = pair.second.get(); } @@ -239,8 +237,8 @@ void AllocateTransientArrays(Model* model, // is a misnormer, should read 'workspace'. for (const auto& array_pair : ordered_arrays_map) { const string& array_name = array_pair.first; - const auto& array_lifespan = array_lifespans.find(array_name)->second; - if (array_lifespan.persistent) { + auto it = array_lifespans.find(array_name); + if (it != array_lifespans.end() && it->second.persistent) { AllocateTransientArray(*model, array_name, &allocator, transient_data_alignment); } @@ -282,8 +280,8 @@ void AllocateTransientArrays(Model* model, std::size_t persistent_alloc_size = 0; for (const auto& array_pair : ordered_arrays_map) { const string& array_name = array_pair.first; - const auto& array_lifespan = array_lifespans.find(array_name)->second; - if (array_lifespan.persistent) { + auto it = array_lifespans.find(array_name); + if (it != array_lifespans.end() && it->second.persistent) { persistent_alloc_size += TransientArraySize(*model, array_name, transient_data_alignment); } diff --git a/tensorflow/contrib/lite/toco/allocate_transient_arrays.h b/tensorflow/contrib/lite/toco/allocate_transient_arrays.h index 12d0d0498f5224962f2775d4e3cb7d8e360cbe46..59d8ada1e9bb985f2eaa7ff6d29bc4f1b054a070 100644 --- a/tensorflow/contrib/lite/toco/allocate_transient_arrays.h +++ b/tensorflow/contrib/lite/toco/allocate_transient_arrays.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_ALLOCATE_TRANSIENT_ARRAYS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_ALLOCATE_TRANSIENT_ARRAYS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_ALLOCATE_TRANSIENT_ARRAYS_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_ALLOCATE_TRANSIENT_ARRAYS_H_ #include "tensorflow/contrib/lite/toco/model.h" @@ -41,4 +41,4 @@ void AllocateTransientArrays(Model* model, } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_ALLOCATE_TRANSIENT_ARRAYS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_ALLOCATE_TRANSIENT_ARRAYS_H_ diff --git a/tensorflow/contrib/lite/toco/args.h b/tensorflow/contrib/lite/toco/args.h index a2f80fae9b2706d77652c89901430e7f77b90814..b97a4720a7c4e69f8b69574475d19e0522cfe86d 100644 --- a/tensorflow/contrib/lite/toco/args.h +++ b/tensorflow/contrib/lite/toco/args.h @@ -15,12 +15,15 @@ limitations under the License. // This abstracts command line arguments in toco. // Arg is a parseable type that can register a default value, be able to // parse itself, and keep track of whether it was specified. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_ARGS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_ARGS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_ARGS_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_ARGS_H_ #include #include #include +#if defined(PLATFORM_GOOGLE) +#include "strings/split.h" +#endif #include "absl/strings/numbers.h" #include "absl/strings/str_split.h" #include "tensorflow/contrib/lite/toco/toco_port.h" @@ -144,12 +147,12 @@ class Arg final { if (!TryStripPrefixString(outer_member, "{", &outer_member)) return false; if (!TryStripSuffixString(outer_member, "}", &outer_member)) return false; const std::vector inner_fields_vector = - strings::Split(outer_member, ','); + absl::StrSplit(outer_member, ','); std::unordered_map element; for (const string& member_field : inner_fields_vector) { std::vector outer_member_key_value = - strings::Split(member_field, ':'); + absl::StrSplit(member_field, ':'); if (outer_member_key_value.size() != 2) return false; string& key = outer_member_key_value[0]; string& value = outer_member_key_value[1]; @@ -205,6 +208,7 @@ struct ParsedModelFlags { Arg dump_graphviz_video = Arg(false); Arg allow_nonexistent_arrays = Arg(false); Arg allow_nonascii_arrays = Arg(false); + Arg arrays_extra_info_file; }; // Flags that describe the operation you would like to do (what conversion @@ -229,4 +233,4 @@ struct ParsedTocoFlags { }; } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_ARGS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_ARGS_H_ diff --git a/tensorflow/contrib/lite/toco/dump_graphviz.cc b/tensorflow/contrib/lite/toco/dump_graphviz.cc index 39809216c77bdadfd44aafbddc8e0979fde66a49..c726eb6d8678e2703f5acba8b3d8d740186939f5 100644 --- a/tensorflow/contrib/lite/toco/dump_graphviz.cc +++ b/tensorflow/contrib/lite/toco/dump_graphviz.cc @@ -278,8 +278,8 @@ std::vector OperatorsToDump(const Model& model) { if (last_specified) { // Return only the part of the graph between graphviz_first_array // and graphviz_last_array. - CHECK(model.arrays.count(dump_options.graphviz_first_array)); - CHECK(model.arrays.count(dump_options.graphviz_last_array)); + CHECK(model.HasArray(dump_options.graphviz_first_array)); + CHECK(model.HasArray(dump_options.graphviz_last_array)); std::unordered_set arrays_already_produced; std::vector arrays_to_produce; arrays_to_produce.push_back(dump_options.graphviz_last_array); @@ -336,7 +336,7 @@ void DumpGraphviz(const Model& model, string* output_file_contents) { op_properties.color.TextColorString().c_str()); // Add nodes and edges for all inputs of the operator. for (const auto& input : op.inputs) { - if (model.arrays.count(input) == 0) { + if (!model.HasArray(input)) { // Arrays should _always_ exist. Except, perhaps, during development. continue; } @@ -352,7 +352,7 @@ void DumpGraphviz(const Model& model, string* output_file_contents) { } // Add nodes and edges for all outputs of the operator. for (const auto& output : op.outputs) { - if (model.arrays.count(output) == 0) { + if (!model.HasArray(output)) { // Arrays should _always_ exist. Except, perhaps, during development. continue; } diff --git a/tensorflow/contrib/lite/toco/dump_graphviz.h b/tensorflow/contrib/lite/toco/dump_graphviz.h index 0fb28e3de844b123a60e36bc23c7d2add8189962..ea5a4031c39580be00130a2fd3a89c61da2acf01 100644 --- a/tensorflow/contrib/lite/toco/dump_graphviz.h +++ b/tensorflow/contrib/lite/toco/dump_graphviz.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_DUMP_GRAPHVIZ_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_DUMP_GRAPHVIZ_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_DUMP_GRAPHVIZ_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_DUMP_GRAPHVIZ_H_ #include @@ -25,4 +25,4 @@ void DumpGraphviz(const Model& model, string* output_file_contents); } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_DUMP_GRAPHVIZ_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_DUMP_GRAPHVIZ_H_ diff --git a/tensorflow/contrib/lite/toco/export_tensorflow.cc b/tensorflow/contrib/lite/toco/export_tensorflow.cc index 51d76e44a03fdf75248ce3b73998f52a00e723e4..570cc7943b2926136f6fdb21d20b8aa6acf8cd26 100644 --- a/tensorflow/contrib/lite/toco/export_tensorflow.cc +++ b/tensorflow/contrib/lite/toco/export_tensorflow.cc @@ -46,6 +46,32 @@ using tensorflow::TensorProto; namespace toco { namespace { +tensorflow::DataType GetTensorFlowDataType(ArrayDataType data_type) { + switch (data_type) { + case ArrayDataType::kBool: + return tensorflow::DT_BOOL; + case ArrayDataType::kFloat: + return tensorflow::DT_FLOAT; + case ArrayDataType::kUint8: + return tensorflow::DT_UINT8; + case ArrayDataType::kInt32: + return tensorflow::DT_INT32; + case ArrayDataType::kInt64: + return tensorflow::DT_INT64; + case ArrayDataType::kString: + return tensorflow::DT_STRING; + default: + case ArrayDataType::kNone: + LOG(FATAL) << "Unsupported data type: " << static_cast(data_type); + return tensorflow::DT_INVALID; + } +} + +tensorflow::DataType GetTensorFlowDataType(const Model& model, + const string& array_name) { + return GetTensorFlowDataType(model.GetArray(array_name).data_type); +} + // TensorFlow sometimes forbids what it calls "legacy scalars", // which are 1-D shapes where the unique shape size is 1. // See OpKernel::IsLegacyScalar and OpKernel::allow_legacy_scalars. @@ -156,8 +182,8 @@ void ConvertFloatTensorConst(const Model& model, const string& name, const_op->set_name(name); (*const_op->mutable_attr())["dtype"].set_type(DT_FLOAT); auto* tensor = (*const_op->mutable_attr())["value"].mutable_tensor(); - CHECK(model.arrays.count(name)); - const auto& input_array = *model.arrays.at(name); + CHECK(model.HasArray(name)); + const auto& input_array = model.GetArray(name); const auto& input_shape = input_array.shape(); CHECK(input_array.buffer); CHECK(input_array.buffer->type == ArrayDataType::kFloat); @@ -177,8 +203,8 @@ void ConvertFloatTensorConst(const Model& model, const string& name, const_op->set_name(name); (*const_op->mutable_attr())["dtype"].set_type(DT_FLOAT); auto* tensor = (*const_op->mutable_attr())["value"].mutable_tensor(); - CHECK(model.arrays.count(name)); - const auto& input_array = *model.arrays.at(name); + CHECK(model.HasArray(name)); + const auto& input_array = model.GetArray(name); const auto& input_shape = input_array.shape(); CHECK(input_array.buffer); CHECK(input_array.buffer->type == ArrayDataType::kFloat); @@ -193,8 +219,8 @@ void ConvertIntTensorConst(const Model& model, const string& name, if (HasAlreadyExportedConst(name, *tensorflow_graph)) { return; } - CHECK(model.arrays.count(name)); - const auto& array = *model.arrays.at(name); + CHECK(model.HasArray(name)); + const auto& array = model.GetArray(name); auto* const_op = tensorflow_graph->add_node(); const_op->set_op("Const"); const_op->set_name(name); @@ -212,6 +238,24 @@ void ConvertIntTensorConst(const Model& model, const string& name, } } +void CreateIntTensorConst(const string& name, const std::vector& data, + GraphDef* tensorflow_graph) { + if (HasAlreadyExportedConst(name, *tensorflow_graph)) { + return; + } + auto* const_op = tensorflow_graph->add_node(); + const_op->set_op("Const"); + const_op->set_name(name); + (*const_op->mutable_attr())["dtype"].set_type(DT_INT32); + auto* tensor = (*const_op->mutable_attr())["value"].mutable_tensor(); + tensor->set_dtype(DT_INT32); + for (auto index : data) { + tensor->add_int_val(index); + } + auto* shape = tensor->mutable_tensor_shape(); + shape->add_dim()->set_size(data.size()); +} + void CreateMatrixShapeTensorConst(const string& name, int rows, int cols, GraphDef* tensorflow_graph) { if (HasAlreadyExportedConst(name, *tensorflow_graph)) { @@ -324,7 +368,7 @@ void ConvertConvOperator(const Model& model, const ConvOperator& src_op, biasadd_op->add_input(conv_output); biasadd_op->add_input(src_op.inputs[2]); (*biasadd_op->mutable_attr())["T"].set_type(DT_FLOAT); - CHECK(model.arrays.count(src_op.inputs[2])); + CHECK(model.HasArray(src_op.inputs[2])); const string& bias_array_name = WalkUpToConstantArray(model, src_op.inputs[2]); const auto& bias_array = model.GetArray(bias_array_name); @@ -361,7 +405,7 @@ void ConvertDepthwiseConvOperator(const Model& model, // We need to convert that to H x W x InputDepth x Multiplier. // That's only a matter of constructing a Dims object; the actual // array layout is the same. - CHECK(model.arrays.count(src_op.inputs[1])); + CHECK(model.HasArray(src_op.inputs[1])); const string& src_weights_name = WalkUpToConstantArray(model, src_op.inputs[1]); const auto& src_weights_array = model.GetArray(src_weights_name); @@ -404,7 +448,7 @@ void ConvertDepthwiseConvOperator(const Model& model, biasadd_op->add_input(conv_output); biasadd_op->add_input(src_op.inputs[2]); (*biasadd_op->mutable_attr())["T"].set_type(DT_FLOAT); - CHECK(model.arrays.count(src_op.inputs[2])); + CHECK(model.HasArray(src_op.inputs[2])); const string& bias_name = WalkUpToConstantArray(model, src_op.inputs[2]); const auto& bias_array = model.GetArray(bias_name); // TODO(b/62904716) Bias arrays should be 1-D, and used directly. @@ -445,14 +489,23 @@ void ConvertSpaceToDepthOperator(const Model& model, void ConvertFullyConnectedOperator(const Model& model, const FullyConnectedOperator& src_op, GraphDef* tensorflow_graph) { - const string reshape_output = src_op.outputs[0] + "/reshape"; - const string reshape_shape = src_op.outputs[0] + "/reshape/shape"; + // Reshape input activations to have the shape expected by the MatMul. + const string reshape_output = + AvailableArrayName(model, src_op.outputs[0] + "/reshape"); + const string reshape_shape = + AvailableArrayName(model, reshape_output + "/shape"); + const auto& fc_weights_array = model.GetArray(src_op.inputs[1]); + const auto& fc_weights_shape = fc_weights_array.shape(); + CHECK_EQ(fc_weights_shape.dimensions_count(), 2); + CreateMatrixShapeTensorConst(reshape_shape, fc_weights_shape.dims(1), -1, + tensorflow_graph); auto* reshape_op = tensorflow_graph->add_node(); reshape_op->set_op("Reshape"); reshape_op->set_name(reshape_output); reshape_op->add_input(src_op.inputs[0]); reshape_op->add_input(reshape_shape); - (*reshape_op->mutable_attr())["T"].set_type(DT_FLOAT); + (*reshape_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); const bool has_bias = src_op.inputs.size() >= 3; string matmul_output = src_op.outputs[0]; @@ -460,40 +513,45 @@ void ConvertFullyConnectedOperator(const Model& model, matmul_output += "/matmul"; } + // Transpose the RHS input from column-major to row-major to match TensorFlow + // expectations. This is the inverse of the transpose we do during + // ResolveTensorFlowMatMul. + const string transpose_output = + AvailableArrayName(model, matmul_output + "/transpose_weights"); + const string transpose_perm = + AvailableArrayName(model, transpose_output + "/perm"); + CreateIntTensorConst(transpose_perm, {1, 0}, tensorflow_graph); + auto transpose_op = tensorflow_graph->add_node(); + transpose_op->set_op("Transpose"); + transpose_op->set_name(transpose_output); + *transpose_op->add_input() = src_op.inputs[1]; + *transpose_op->add_input() = transpose_perm; + (*transpose_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[1])); + (*transpose_op->mutable_attr())["Tperm"].set_type(DT_INT32); + auto* matmul_op = tensorflow_graph->add_node(); matmul_op->set_op("MatMul"); - matmul_op->set_name(matmul_output); *matmul_op->add_input() = reshape_output; - *matmul_op->add_input() = src_op.inputs[1]; - (*matmul_op->mutable_attr())["T"].set_type(DT_FLOAT); + *matmul_op->add_input() = transpose_op->name(); + (*matmul_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); (*matmul_op->mutable_attr())["transpose_a"].set_b(false); (*matmul_op->mutable_attr())["transpose_b"].set_b(false); - CHECK(model.arrays.count(src_op.inputs[1])); - const string& fc_weights_name = - WalkUpToConstantArray(model, src_op.inputs[1]); - const auto& fc_weights_array = *model.arrays.at(fc_weights_name); - const auto& fc_weights_shape = fc_weights_array.shape(); - CHECK_EQ(fc_weights_shape.dimensions_count(), 2); - CreateMatrixShapeTensorConst(reshape_shape, fc_weights_shape.dims(1), -1, - tensorflow_graph); - - CHECK(fc_weights_array.buffer); - CHECK(fc_weights_array.buffer->type == ArrayDataType::kFloat); - const float* fc_weights_data = - fc_weights_array.GetBuffer().data.data(); - ConvertFloatTensorConst(fc_weights_name, fc_weights_shape, fc_weights_data, - AxesOrder::kCR, AxesOrder::kRC, tensorflow_graph); + CHECK(model.HasArray(src_op.inputs[1])); + // Add the bias, if it exists. if (has_bias) { auto* biasadd_op = tensorflow_graph->add_node(); biasadd_op->set_op("BiasAdd"); biasadd_op->set_name(src_op.outputs[0]); biasadd_op->add_input(matmul_output); biasadd_op->add_input(src_op.inputs[2]); - (*biasadd_op->mutable_attr())["T"].set_type(DT_FLOAT); - CHECK(model.arrays.count(src_op.inputs[2])); - const auto& bias_array = *model.arrays.at(src_op.inputs[2]); + (*biasadd_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); + CHECK(model.HasArray(src_op.inputs[2])); + const auto& bias_array = model.GetArray(src_op.inputs[2]); // TODO(b/62904716) Bias arrays should be 1-D, and used directly. Shape bias_shape_1d = bias_array.shape(); UnextendShape(&bias_shape_1d, 1); @@ -519,6 +577,18 @@ void ConvertAddOperator(const Model& model, const AddOperator& src_op, (*add_op->mutable_attr())["T"].set_type(DT_FLOAT); } +void ConvertAddNOperator(const Model& model, const AddNOperator& src_op, + GraphDef* tensorflow_graph) { + auto* add_op = tensorflow_graph->add_node(); + add_op->set_op("AddN"); + add_op->set_name(src_op.outputs[0]); + for (const auto& input : src_op.inputs) { + *add_op->add_input() = input; + } + (*add_op->mutable_attr())["N"].set_i(src_op.inputs.size()); + (*add_op->mutable_attr())["T"].set_type(DT_FLOAT); +} + void ConvertMulOperator(const Model& model, const MulOperator& src_op, GraphDef* tensorflow_graph) { auto* add_op = tensorflow_graph->add_node(); @@ -609,7 +679,8 @@ void ConvertSoftmaxOperator(const Model& model, const SoftmaxOperator& src_op, GraphDef* tensorflow_graph) { string softmax_input; Operator* providing_op = GetOpWithOutput(model, src_op.inputs[0]); - if (providing_op->type == OperatorType::kTensorFlowReshape) { + if (providing_op != nullptr && + providing_op->type == OperatorType::kTensorFlowReshape) { softmax_input = src_op.inputs[0]; } else { // Insert a reshape operator that reduces the dimensions down to the 2 that @@ -625,7 +696,7 @@ void ConvertSoftmaxOperator(const Model& model, const SoftmaxOperator& src_op, *reshape_op->add_input() = softmax_size; (*reshape_op->mutable_attr())["T"].set_type(DT_FLOAT); - const auto& input_shape = model.arrays.at(src_op.inputs[0])->shape(); + const auto& input_shape = model.GetArray(src_op.inputs[0]).shape(); int32 flattened_size = 1; for (int i = 0; i < input_shape.dimensions_count() - 1; ++i) { flattened_size *= input_shape.dims(i); @@ -644,6 +715,45 @@ void ConvertSoftmaxOperator(const Model& model, const SoftmaxOperator& src_op, (*softmax_op->mutable_attr())["T"].set_type(DT_FLOAT); } +void ConvertLogSoftmaxOperator(const Model& model, + const LogSoftmaxOperator& src_op, + GraphDef* tensorflow_graph) { + string softmax_input; + Operator* providing_op = GetOpWithOutput(model, src_op.inputs[0]); + if (providing_op->type == OperatorType::kTensorFlowReshape) { + softmax_input = src_op.inputs[0]; + } else { + // Insert a reshape operator that reduces the dimensions down to the 2 that + // are required for TensorFlow Logits. + const string reshape_output = + src_op.outputs[0] + "/log_softmax_insert_reshape"; + const string softmax_size = src_op.outputs[0] + "/log_softmax_insert_size"; + softmax_input = reshape_output; + + auto* reshape_op = tensorflow_graph->add_node(); + reshape_op->set_op("Reshape"); + reshape_op->set_name(reshape_output); + *reshape_op->add_input() = src_op.inputs[0]; + *reshape_op->add_input() = softmax_size; + (*reshape_op->mutable_attr())["T"].set_type(DT_FLOAT); + + const auto& input_shape = model.GetArray(src_op.inputs[0]).shape(); + int32 flattened_size = 1; + for (int i = 0; i < input_shape.dimensions_count() - 1; ++i) { + flattened_size *= input_shape.dims(i); + } + const std::vector shape_data = { + flattened_size, input_shape.dims(input_shape.dimensions_count() - 1)}; + CreateReshapeShapeTensorConst(softmax_size, shape_data, tensorflow_graph); + } + + auto* log_softmax_op = tensorflow_graph->add_node(); + log_softmax_op->set_op("LogSoftmax"); + log_softmax_op->set_name(src_op.outputs[0]); + *log_softmax_op->add_input() = softmax_input; + (*log_softmax_op->mutable_attr())["T"].set_type(DT_FLOAT); +} + void ConvertL2NormalizationOperator(const L2NormalizationOperator& src_op, GraphDef* tensorflow_graph) { const string square_output = src_op.outputs[0] + "/square"; @@ -786,7 +896,8 @@ void ConvertConcatenationOperator(const Model& model, *dc_op->add_input() = input; } *dc_op->add_input() = dummy_axis; - (*dc_op->mutable_attr())["T"].set_type(DT_FLOAT); + (*dc_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); (*dc_op->mutable_attr())["Tidx"].set_type(DT_INT32); (*dc_op->mutable_attr())["N"].set_i(src_op.inputs.size()); } @@ -800,10 +911,13 @@ void ConvertTensorFlowReshapeOperator(const Model& model, CHECK_EQ(src_op.inputs.size(), 2); *reshape_op->add_input() = src_op.inputs[0]; *reshape_op->add_input() = src_op.inputs[1]; - (*reshape_op->mutable_attr())["T"].set_type(DT_FLOAT); + (*reshape_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.outputs[0])); const auto& shape_array = model.GetArray(src_op.inputs[1]); - CHECK(shape_array.data_type == ArrayDataType::kInt32); - CHECK(shape_array.buffer != nullptr); + QCHECK(shape_array.data_type == ArrayDataType::kInt32) + << "Only int32 shape is supported."; + QCHECK(shape_array.buffer != nullptr) + << "Shape inferred at runtime is not supported."; const auto& shape_data = shape_array.GetBuffer().data; CreateReshapeShapeTensorConst(src_op.inputs[1], shape_data, tensorflow_graph); } @@ -895,24 +1009,6 @@ void ConvertSplitOperator(const Model& model, tensorflow_graph); } -tensorflow::DataType GetTensorFlowDataType(const Model& model, - const string& array_name) { - auto& dtype = model.GetArray(array_name).data_type; - CHECK(dtype == ArrayDataType::kFloat || dtype == ArrayDataType::kInt32 || - dtype == ArrayDataType::kUint8 || dtype == ArrayDataType::kInt64); - if (dtype == ArrayDataType::kFloat) { - return tensorflow::DT_FLOAT; - } else if (dtype == ArrayDataType::kInt32) { - return tensorflow::DT_INT32; - } else if (dtype == ArrayDataType::kUint8) { - return tensorflow::DT_UINT8; - } else if (dtype == ArrayDataType::kInt64) { - return tensorflow::DT_INT64; - } else { - LOG(FATAL) << "Wrong data type"; - } -} - void ConvertCastOperator(const Model& model, const CastOperator& src_op, GraphDef* tensorflow_graph) { auto* cast_op = tensorflow_graph->add_node(); @@ -967,6 +1063,113 @@ void ConvertArgMaxOperator(const Model& model, const ArgMaxOperator& src_op, GetTensorFlowDataType(model, src_op.outputs[0])); } +void ConvertTransposeOperator(const Model& model, + const TransposeOperator& src_op, + GraphDef* tensorflow_graph) { + auto* transpose_op = tensorflow_graph->add_node(); + transpose_op->set_op("Transpose"); + transpose_op->set_name(src_op.outputs[0]); + CHECK_EQ(src_op.inputs.size(), 2); + *transpose_op->add_input() = src_op.inputs[0]; + *transpose_op->add_input() = src_op.inputs[1]; + (*transpose_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); + (*transpose_op->mutable_attr())["Tperm"].set_type( + GetTensorFlowDataType(model, src_op.inputs[1])); +} + +void ConvertTensorFlowShapeOperator(const Model& model, + const TensorFlowShapeOperator& src_op, + GraphDef* tensorflow_graph) { + auto* shape_op = tensorflow_graph->add_node(); + shape_op->set_op("Shape"); + shape_op->set_name(src_op.outputs[0]); + CHECK_EQ(src_op.inputs.size(), 1); + *shape_op->add_input() = src_op.inputs[0]; + (*shape_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); + (*shape_op->mutable_attr())["out_type"].set_type( + GetTensorFlowDataType(model, src_op.outputs[0])); +} + +void ConvertRankOperator(const Model& model, const RankOperator& src_op, + GraphDef* tensorflow_graph) { + auto* rank_op = tensorflow_graph->add_node(); + rank_op->set_op("Rank"); + rank_op->set_name(src_op.outputs[0]); + CHECK_EQ(src_op.inputs.size(), 1); + *rank_op->add_input() = src_op.inputs[0]; + (*rank_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); +} + +void ConvertRangeOperator(const Model& model, const RangeOperator& src_op, + GraphDef* tensorflow_graph) { + auto* range_op = tensorflow_graph->add_node(); + range_op->set_op("Range"); + range_op->set_name(src_op.outputs[0]); + CHECK_EQ(src_op.inputs.size(), 3); + *range_op->add_input() = src_op.inputs[0]; + *range_op->add_input() = src_op.inputs[1]; + *range_op->add_input() = src_op.inputs[2]; + (*range_op->mutable_attr())["Tidx"].set_type( + GetTensorFlowDataType(src_op.dtype)); +} + +void ConvertStackOperator(const Model& model, const StackOperator& src_op, + GraphDef* tensorflow_graph) { + auto* stack_op = tensorflow_graph->add_node(); + stack_op->set_op("Stack"); + stack_op->set_name(src_op.outputs[0]); + for (const auto& input : src_op.inputs) { + *stack_op->add_input() = input; + } + (*stack_op->mutable_attr())["elem_type"].set_type( + GetTensorFlowDataType(model, src_op.outputs[0])); + (*stack_op->mutable_attr())["axis"].set_i(src_op.axis); +} + +void ConvertFillOperator(const Model& model, const FillOperator& src_op, + GraphDef* tensorflow_graph) { + auto* fill_op = tensorflow_graph->add_node(); + fill_op->set_op("Fill"); + fill_op->set_name(src_op.outputs[0]); + CHECK_EQ(src_op.inputs.size(), 2); + *fill_op->add_input() = src_op.inputs[0]; + *fill_op->add_input() = src_op.inputs[1]; + (*fill_op->mutable_attr())["index_type"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); + (*fill_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[1])); +} + +void ConvertFloorDivOperator(const Model& model, const FloorDivOperator& src_op, + GraphDef* tensorflow_graph) { + auto* floor_div_op = tensorflow_graph->add_node(); + floor_div_op->set_op("FloorDiv"); + floor_div_op->set_name(src_op.outputs[0]); + CHECK_EQ(src_op.inputs.size(), 2); + *floor_div_op->add_input() = src_op.inputs[0]; + *floor_div_op->add_input() = src_op.inputs[1]; + (*floor_div_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); +} + +void ConvertExpandDimsOperator(const Model& model, + const ExpandDimsOperator& src_op, + GraphDef* tensorflow_graph) { + auto* expand_dims_op = tensorflow_graph->add_node(); + expand_dims_op->set_op("ExpandDims"); + expand_dims_op->set_name(src_op.outputs[0]); + CHECK_EQ(src_op.inputs.size(), 2); + *expand_dims_op->add_input() = src_op.inputs[0]; + *expand_dims_op->add_input() = src_op.inputs[1]; + (*expand_dims_op->mutable_attr())["T"].set_type( + GetTensorFlowDataType(model, src_op.inputs[0])); + (*expand_dims_op->mutable_attr())["Tdim"].set_type( + GetTensorFlowDataType(model, src_op.inputs[1])); +} + void ConvertResizeBilinearOperator(const Model& model, const ResizeBilinearOperator& src_op, GraphDef* tensorflow_graph) { @@ -977,6 +1180,7 @@ void ConvertResizeBilinearOperator(const Model& model, *resize_op->add_input() = src_op.inputs[0]; *resize_op->add_input() = src_op.inputs[1]; (*resize_op->mutable_attr())["T"].set_type(DT_FLOAT); + (*resize_op->mutable_attr())["align_corners"].set_b(src_op.align_corners); } namespace { @@ -1011,8 +1215,8 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op, // Op names have been chosen to match the tf.slim LSTM naming // as closely as possible. const int axis = - model.arrays.at(src_op.inputs[LstmCellOperator::PREV_ACTIV_INPUT]) - ->shape() + model.GetArray(src_op.inputs[LstmCellOperator::PREV_ACTIV_INPUT]) + .shape() .dimensions_count() - 1; // Note that DATA_INPUT may have extra size 1 dimensions, but TF concat @@ -1031,9 +1235,10 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op, // Write weights const string weights_output = base + "weights"; - CHECK(model.arrays.count(src_op.inputs[LstmCellOperator::WEIGHTS_INPUT])); - const auto& weights_array = - *model.arrays.at(src_op.inputs[LstmCellOperator::WEIGHTS_INPUT]); + CHECK(model.HasArray(src_op.inputs[LstmCellOperator::WEIGHTS_INPUT])); + const string weights_name = WalkUpToConstantArray( + model, src_op.inputs[LstmCellOperator::WEIGHTS_INPUT]); + const auto& weights_array = model.GetArray(weights_name); // Convert 4D FullyConnected weights into 2D matrix const auto& weights_shape = weights_array.shape(); CHECK_EQ(weights_shape.dimensions_count(), 2); @@ -1057,9 +1262,10 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op, // Write biases const string biases_output = base + "biases"; - CHECK(model.arrays.count(src_op.inputs[LstmCellOperator::BIASES_INPUT])); - const auto& bias_array = - *model.arrays.at(src_op.inputs[LstmCellOperator::BIASES_INPUT]); + CHECK(model.HasArray(src_op.inputs[LstmCellOperator::BIASES_INPUT])); + const string bias_name = WalkUpToConstantArray( + model, src_op.inputs[LstmCellOperator::BIASES_INPUT]); + const auto& bias_array = model.GetArray(bias_name); // TODO(b/62904716) Bias arrays should be 1-D, and used directly. Shape bias_shape_1d = bias_array.shape(); UnextendShape(&bias_shape_1d, 1); @@ -1375,6 +1581,17 @@ void ConvertTensorFlowMaximumOperator(const Model& model, (*sub_op->mutable_attr())["T"].set_type(data_type); } +void ConvertTopKV2Operator(const Model& model, const TopKV2Operator& src_op, + GraphDef* tensorflow_graph) { + auto* topk_op = tensorflow_graph->add_node(); + topk_op->set_op("TOPKV2"); + topk_op->set_name(src_op.outputs[0]); + CHECK_EQ(src_op.inputs.size(), 2); + *topk_op->add_input() = src_op.inputs[0]; + *topk_op->add_input() = src_op.inputs[1]; + (*topk_op->mutable_attr())["sorted"].set_b(true); +} + void ConvertOperator(const Model& model, const Operator& src_op, GraphDef* tensorflow_graph) { if (src_op.fused_activation_function != FusedActivationFunctionType::kNone) { @@ -1404,6 +1621,9 @@ void ConvertOperator(const Model& model, const Operator& src_op, } else if (src_op.type == OperatorType::kAdd) { ConvertAddOperator(model, static_cast(src_op), tensorflow_graph); + } else if (src_op.type == OperatorType::kAddN) { + ConvertAddNOperator(model, static_cast(src_op), + tensorflow_graph); } else if (src_op.type == OperatorType::kMul) { ConvertMulOperator(model, static_cast(src_op), tensorflow_graph); @@ -1428,6 +1648,10 @@ void ConvertOperator(const Model& model, const Operator& src_op, } else if (src_op.type == OperatorType::kSoftmax) { ConvertSoftmaxOperator(model, static_cast(src_op), tensorflow_graph); + } else if (src_op.type == OperatorType::kLogSoftmax) { + ConvertLogSoftmaxOperator(model, + static_cast(src_op), + tensorflow_graph); } else if (src_op.type == OperatorType::kLocalResponseNormalization) { ConvertLocalResponseNormalizationOperator( static_cast(src_op), @@ -1516,6 +1740,35 @@ void ConvertOperator(const Model& model, const Operator& src_op, } else if (src_op.type == OperatorType::kArgMax) { ConvertArgMaxOperator(model, static_cast(src_op), tensorflow_graph); + } else if (src_op.type == OperatorType::kTopK_V2) { + ConvertTopKV2Operator(model, static_cast(src_op), + tensorflow_graph); + } else if (src_op.type == OperatorType::kTranspose) { + ConvertTransposeOperator( + model, static_cast(src_op), tensorflow_graph); + } else if (src_op.type == OperatorType::kTensorFlowShape) { + ConvertTensorFlowShapeOperator( + model, static_cast(src_op), + tensorflow_graph); + } else if (src_op.type == OperatorType::kRank) { + ConvertRankOperator(model, static_cast(src_op), + tensorflow_graph); + } else if (src_op.type == OperatorType::kRange) { + ConvertRangeOperator(model, static_cast(src_op), + tensorflow_graph); + } else if (src_op.type == OperatorType::kStack) { + ConvertStackOperator(model, static_cast(src_op), + tensorflow_graph); + } else if (src_op.type == OperatorType::kFill) { + ConvertFillOperator(model, static_cast(src_op), + tensorflow_graph); + } else if (src_op.type == OperatorType::kFloorDiv) { + ConvertFloorDivOperator(model, static_cast(src_op), + tensorflow_graph); + } else if (src_op.type == OperatorType::kExpandDims) { + ConvertExpandDimsOperator(model, + static_cast(src_op), + tensorflow_graph); } else { LOG(FATAL) << "Unhandled operator type " << OperatorTypeName(src_op.type); } @@ -1555,7 +1808,7 @@ void AddPlaceholderForRNNState(const Model& model, const string& name, int size, (*placeholder->mutable_attr())["dtype"].set_type(DT_FLOAT); auto* shape = (*placeholder->mutable_attr())["shape"].mutable_shape(); - const auto& state_array = *model.arrays.at(name); + const auto& state_array = model.GetArray(name); if (state_array.has_shape()) { const auto& state_shape = state_array.shape(); const int kDims = state_shape.dimensions_count(); @@ -1572,7 +1825,7 @@ void ExportTensorFlowGraphDefImplementation(const Model& model, GraphDef* tensorflow_graph) { for (const auto& input_array : model.flags.input_arrays()) { AddPlaceholder(input_array.name(), - model.arrays.at(input_array.name())->data_type, + model.GetArray(input_array.name()).data_type, tensorflow_graph); } for (const auto& rnn_state : model.flags.rnn_states()) { @@ -1586,7 +1839,7 @@ void ExportTensorFlowGraphDefImplementation(const Model& model, // by the above operators export. It's important that this comes // after, as some operators need to export arrays that they reference // in a specific way, rather than in the generic way done below. - for (const auto& array_pair : model.arrays) { + for (const auto& array_pair : model.GetArrayMap()) { const string& array_name = array_pair.first; const auto& array = *array_pair.second; if (array.buffer) { @@ -1605,6 +1858,30 @@ void ExportTensorFlowGraphDefImplementation(const Model& model, } } // namespace +void EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(Model* model) { + for (const auto& array_kv : model->GetArrayMap()) { + const string& array_name = array_kv.first; + Array& array = *array_kv.second; + if (!array.buffer || !array.minmax) { + continue; + } + const string& wrapped_array_name = + AvailableArrayName(*model, array_name + "/data"); + Array& wrapped_array = model->GetOrCreateArray(wrapped_array_name); + wrapped_array.data_type = array.data_type; + wrapped_array.copy_shape(array.shape()); + wrapped_array.buffer = std::move(array.buffer); + FakeQuantOperator* fakequant_op = new FakeQuantOperator; + fakequant_op->inputs = {wrapped_array_name}; + fakequant_op->outputs = {array_name}; + fakequant_op->minmax.reset(new MinMax); + *fakequant_op->minmax = *array.minmax; + const auto& it = FindOpWithInput(*model, array_name); + model->operators.emplace(it, fakequant_op); + } + CheckInvariants(*model); +} + void ExportTensorFlowGraphDef(const Model& model, string* output_file_contents) { CHECK(output_file_contents->empty()); diff --git a/tensorflow/contrib/lite/toco/export_tensorflow.h b/tensorflow/contrib/lite/toco/export_tensorflow.h index eca97745767387a04bcd2c8deb579928edf2497c..d7310bb75f258cde25236da2a9269f18234784e4 100644 --- a/tensorflow/contrib/lite/toco/export_tensorflow.h +++ b/tensorflow/contrib/lite/toco/export_tensorflow.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_EXPORT_TENSORFLOW_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_EXPORT_TENSORFLOW_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_EXPORT_TENSORFLOW_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_EXPORT_TENSORFLOW_H_ #include #include "tensorflow/contrib/lite/toco/model.h" @@ -22,6 +22,8 @@ namespace toco { void ExportTensorFlowGraphDef(const Model& model, string* output_file_contents); +void EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(Model* model); + } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_EXPORT_TENSORFLOW_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_EXPORT_TENSORFLOW_H_ diff --git a/tensorflow/contrib/lite/toco/format_port.h b/tensorflow/contrib/lite/toco/format_port.h index 0e999001e0e35fb916b11db199dbf28572685f3d..eb81e90faf20133ed722185928f86ef45ac4f8f6 100644 --- a/tensorflow/contrib/lite/toco/format_port.h +++ b/tensorflow/contrib/lite/toco/format_port.h @@ -16,8 +16,8 @@ limitations under the License. // and util::format::AppendF. Unfortunately, type safety is not as good as a // a full C++ example. // TODO(aselle): When absl adds support for StrFormat, use that instead. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_FORMAT_PORT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_FORMAT_PORT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_FORMAT_PORT_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_FORMAT_PORT_H_ #include "tensorflow/contrib/lite/toco/toco_types.h" #include "tensorflow/core/lib/strings/stringprintf.h" @@ -74,4 +74,4 @@ inline string StringF(const char* fmt, Args&&... args) { } // namespace port } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_FORMAT_PORT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_FORMAT_PORT_H_ diff --git a/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md b/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md index 7e152f5ba887088c98055596f8245b82fbc86eaa..372c52558973f4aacc180ac44b9e95a5e9b199ef 100644 --- a/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md +++ b/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md @@ -23,7 +23,7 @@ curl https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_ bazel run --config=opt \ //tensorflow/contrib/lite/toco:toco -- \ --input_file=/tmp/mobilenet_v1_0.50_128/frozen_graph.pb \ - --output_file=/tmp/foo.lite \ + --output_file=/tmp/foo.tflite \ --input_format=TENSORFLOW_GRAPHDEF \ --output_format=TFLITE \ --inference_type=FLOAT \ @@ -101,7 +101,7 @@ direction, let us just give an example of that: ``` bazel run --config=opt \ //tensorflow/contrib/lite/toco:toco -- \ - --input_file=/tmp/foo.lite \ + --input_file=/tmp/foo.tflite \ --output_file=/tmp/foo.pb \ --input_format=TFLITE \ --output_format=TENSORFLOW_GRAPHDEF \ @@ -130,7 +130,7 @@ flatbuffer is done like this: bazel run --config=opt \ //tensorflow/contrib/lite/toco:toco -- \ --input_file=/tmp/some_quantized_graph.pb \ - --output_file=/tmp/foo.lite \ + --output_file=/tmp/foo.tflite \ --input_format=TENSORFLOW_GRAPHDEF \ --output_format=TFLITE \ --inference_type=QUANTIZED_UINT8 \ @@ -207,7 +207,7 @@ curl https://storage.googleapis.com/download.tensorflow.org/models/inception_v1_ bazel run --config=opt \ //tensorflow/contrib/lite/toco:toco -- \ --input_file=/tmp/inception_v1_2016_08_28_frozen.pb \ - --output_file=/tmp/foo.lite \ + --output_file=/tmp/foo.tflite \ --input_format=TENSORFLOW_GRAPHDEF \ --output_format=TFLITE \ --inference_type=FLOAT \ @@ -235,7 +235,7 @@ curl https://storage.googleapis.com/download.tensorflow.org/models/inception_v1_ bazel run --config=opt \ //tensorflow/contrib/lite/toco:toco -- \ --input_file=/tmp/inception_v1_2016_08_28_frozen.pb \ - --output_file=/tmp/foo.lite \ + --output_file=/tmp/foo.tflite \ --input_format=TENSORFLOW_GRAPHDEF \ --output_format=TFLITE \ --inference_type=FLOAT \ @@ -308,7 +308,7 @@ curl https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_ bazel run --config=opt \ //tensorflow/contrib/lite/toco:toco -- \ --input_file=/tmp/mobilenet_v1_0.50_128/frozen_graph.pb \ - --output_file=/tmp/foo.lite \ + --output_file=/tmp/foo.tflite \ --input_format=TENSORFLOW_GRAPHDEF \ --output_format=TFLITE \ --inference_type=FLOAT \ @@ -415,7 +415,7 @@ curl https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_ bazel run --config=opt \ //tensorflow/contrib/lite/toco:toco -- \ --input_file=/tmp/mobilenet_v1_0.50_128/frozen_graph.pb \ - --output_file=/tmp/foo.lite \ + --output_file=/tmp/foo.tflite \ --input_format=TENSORFLOW_GRAPHDEF \ --output_format=TFLITE \ --inference_type=FLOAT \ diff --git a/tensorflow/contrib/lite/toco/g3doc/cmdline_reference.md b/tensorflow/contrib/lite/toco/g3doc/cmdline_reference.md index 4776741ab9273cf3b2ef0c63a6dbfdea5475b057..5e077952235fa1aac1e12403d3d83633a617ccb7 100644 --- a/tensorflow/contrib/lite/toco/g3doc/cmdline_reference.md +++ b/tensorflow/contrib/lite/toco/g3doc/cmdline_reference.md @@ -229,7 +229,7 @@ additional information about the multiple input arrays: well-formed quantized representation of these graphs. Such graphs should be fixed, but as a temporary work-around, setting this reorder_across_fake_quant flag allows the converter to perform necessary - graph transformaitons on them, at the cost of no longer faithfully matching + graph transformations on them, at the cost of no longer faithfully matching inference and training arithmetic. ### Logging flags diff --git a/tensorflow/contrib/lite/toco/graph_transformations/convert_expanddims_to_reshape.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_expanddims_to_reshape.cc index 3bde9b0169ddfb7fc37657122e2e8eb65ccbdf6d..56f48d47de4e86ece76ceef1d09a25f50957a8dc 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/convert_expanddims_to_reshape.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_expanddims_to_reshape.cc @@ -35,7 +35,7 @@ bool ConvertExpandDimsToReshape::Run(Model* model, std::size_t op_index) { CHECK_EQ(expand_op->inputs.size(), 2); CHECK_EQ(expand_op->outputs.size(), 1); - const auto& input_array = *model->arrays[expand_op->inputs[0]]; + const auto& input_array = model->GetArray(expand_op->inputs[0]); if (!input_array.has_shape()) { // Yield until input dims have been resolved. return false; @@ -46,7 +46,7 @@ bool ConvertExpandDimsToReshape::Run(Model* model, std::size_t op_index) { return false; } - const auto& axis_array = *model->arrays[expand_op->inputs[1]]; + const auto& axis_array = model->GetArray(expand_op->inputs[1]); if (!axis_array.has_shape()) { // Yield until input axis array shape has been resolved. return false; @@ -86,7 +86,7 @@ bool ConvertExpandDimsToReshape::Run(Model* model, std::size_t op_index) { if (IsDiscardableArray(*model, axis_array_name) && CountOpsWithInput(*model, axis_array_name) == 1 && !GetOpWithOutput(*model, axis_array_name)) { - model->arrays.erase(axis_array_name); + model->EraseArray(axis_array_name); } // Replace the operator in the graph. diff --git a/tensorflow/contrib/lite/toco/graph_transformations/convert_pure_conv_to_depthwise.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_pure_conv_to_depthwise.cc index bf454c40c7b50d242d8a7e9eb6b7e579fb0da217..d38db85280d7bd935a47cda70227d383a513fbac 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/convert_pure_conv_to_depthwise.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_pure_conv_to_depthwise.cc @@ -58,7 +58,7 @@ bool ConvertPureConvToDepthwise::Run(Model* model, std::size_t op_index) { depthwiseconv_op->outputs = {conv_op->outputs[0]}; if (conv_op->outputs.size() > 1) { // delete the im2col array. - model->arrays.erase(conv_op->outputs[1]); + model->EraseArray(conv_op->outputs[1]); } depthwiseconv_op->fused_activation_function = conv_op->fused_activation_function; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/convert_reorder_axes.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_reorder_axes.cc new file mode 100644 index 0000000000000000000000000000000000000000..0d274fc687c8d42d47ddb5beb4f9c6f39b417097 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_reorder_axes.cc @@ -0,0 +1,149 @@ +/* Copyright 2017 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 +#include +#include +#include + +#include "absl/strings/str_cat.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +// Creates a Reshape operator from ReorderAxes operator. +TensorFlowReshapeOperator* CreateReshapeFromReorderAxes( + Model* model, ReorderAxesOperator* reorder_op, const Shape& input_shape) { + auto* reshape_op = new TensorFlowReshapeOperator; + + // Copy inputs and outputs to Reshape. + reshape_op->inputs.push_back(reorder_op->inputs[0]); + reshape_op->outputs = reorder_op->outputs; + + // Create reshape dimensions based on input shape. Conversion from + // ReorderAxes to Reshape requires a 4D input shape. + CHECK_EQ(input_shape.dimensions_count(), 4); + std::vector reshape_dims = {1, input_shape.dims(0), input_shape.dims(1), + input_shape.dims(3) * input_shape.dims(2)}; + + // Create a new input array for Reshape. + string reshape_array_name = + AvailableArrayName(*model, reshape_op->outputs[0]); + reshape_op->inputs.push_back(reshape_array_name); + + Array& reshape_array = model->GetOrCreateArray(reshape_array_name); + *(reshape_array.mutable_shape()->mutable_dims()) = { + 1, static_cast(reshape_dims.size())}; + reshape_array.data_type = ArrayDataType::kInt32; + auto& reshape_buffer = + reshape_array.GetMutableBuffer(); + reshape_buffer.data = reshape_dims; + + return reshape_op; +} + +// Creates a Transpose operator from ReorderAxes operator. +TransposeOperator* CreateTransposeFromReorderAxes( + Model* model, ReorderAxesOperator* reorder_op, const Shape& input_shape, + const AxesOrder& input_axes_order, const AxesOrder& output_axes_order) { + auto* transpose_op = new TransposeOperator; + + // Copy inputs and outputs to Transpose. + transpose_op->inputs.push_back(reorder_op->inputs[0]); + transpose_op->outputs = reorder_op->outputs; + + // Create permutations data based on input and output axes order. + std::vector permutations_data; + GetShuffleShape(input_axes_order, output_axes_order, &permutations_data); + + // Create a new input permutations array for Transpose. + string perm_array_name = AvailableArrayName(*model, transpose_op->outputs[0]); + transpose_op->inputs.push_back(perm_array_name); + + Array& perm_array = model->GetOrCreateArray(perm_array_name); + *(perm_array.mutable_shape()->mutable_dims()) = { + static_cast(permutations_data.size())}; + perm_array.data_type = ArrayDataType::kInt32; + auto& perm_buffer = perm_array.GetMutableBuffer(); + perm_buffer.data = permutations_data; + + return transpose_op; +} + +// Converts ReorderAxes into Transpose and Reshape which are compatible with the +// TFLite interpreter. +bool ConvertReorderAxes::Run(Model* model, std::size_t op_index) { + auto reorder_it = model->operators.begin() + op_index; + if (reorder_it->get()->type != OperatorType::kReorderAxes) return false; + + auto* reorder_op = static_cast(reorder_it->get()); + CHECK_EQ(reorder_op->inputs.size(), 1); + CHECK_EQ(reorder_op->outputs.size(), 1); + + const auto& input_array_name = reorder_op->inputs[0]; + const auto& output_array_name = reorder_op->outputs[0]; + auto& input_array = model->GetArray(input_array_name); + auto& output_array = model->GetArray(output_array_name); + + // Get input array. If kFakeQuant is the input into ReorderAxes, get the input + // array passed into kFakeQuant. kFakeQuant op is dropped when possible. + string constant_input_array_name = input_array_name; + if (!input_array.buffer) { + const auto* op_producing_input = GetOpWithOutput(*model, input_array_name); + if (op_producing_input && + op_producing_input->type == OperatorType::kFakeQuant) { + constant_input_array_name = op_producing_input->inputs[0]; + } + } + + // Yield if input array contains constants or if output array size has not + // been adjusted to reflect the permutations in ReorderAxes. ReorderAxes will + // be merged into a constant array when possible. + if (IsConstantParameterArray(*model, constant_input_array_name)) return false; + if (!output_array.has_shape()) return false; + + const auto input_axes_order = reorder_op->input_axes_order; + const auto output_axes_order = reorder_op->output_axes_order; + const Shape input_shape = input_array.shape(); + + // Creates a Reshape or Transpose operator depending on the conversion. + if (input_axes_order == AxesOrder::kHWIM && + output_axes_order == AxesOrder::k1HWO) { + // Add Reshape operator into the graph. This special case is not just a + // permutation. The input dimensions get merged into 3 dimensions while the + // order of the elements does not change. + auto* reshape_op = + CreateReshapeFromReorderAxes(model, reorder_op, input_shape); + const auto reshape_it = model->operators.emplace(reorder_it, reshape_op); + reorder_it = reshape_it + 1; + } else { + // Add Transpose operator into the graph. + auto* transpose_op = CreateTransposeFromReorderAxes( + model, reorder_op, input_shape, input_axes_order, output_axes_order); + const auto transpose_it = + model->operators.emplace(reorder_it, transpose_op); + reorder_it = transpose_it + 1; + } + + // Remove ReorderAxes operator from the graph. + CHECK_EQ(reorder_it->get(), reorder_op); + model->operators.erase(reorder_it); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_addn_to_add.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_addn_to_add.cc new file mode 100644 index 0000000000000000000000000000000000000000..dcaaddbf3b5409f0fc3ddaf32e23b1e5eefb6565 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_addn_to_add.cc @@ -0,0 +1,51 @@ +/* Copyright 2017 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/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +// This pass will convert an AddN operator with only 2 inputs into a regular Add +// operator, to which more optimizations may apply. +bool ConvertTrivialAddNToAdd::Run(Model* model, std::size_t op_index) { + auto addn_it = model->operators.begin() + op_index; + if (addn_it->get()->type != OperatorType::kAddN) { + return false; + } + AddNOperator* addn_op = static_cast(addn_it->get()); + CHECK_GE(addn_op->inputs.size(), 2); + CHECK_EQ(addn_op->outputs.size(), 1); + + // We only reduce AddN with N=2 to a regular Add. + if (addn_op->inputs.size() != 2) { + return false; + } + + // Copy inputs & outputs to regular Add. + auto* add_op = new AddOperator; + add_op->inputs.push_back(addn_op->inputs[0]); + add_op->inputs.push_back(addn_op->inputs[1]); + add_op->outputs = addn_op->outputs; + + // Replace the AddN operator in the graph. + const auto add_it = model->operators.emplace(addn_it, add_op); + addn_it = add_it + 1; + CHECK_EQ(addn_it->get(), addn_op); + model->operators.erase(addn_it); + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_stack_to_reshape.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_stack_to_reshape.cc new file mode 100644 index 0000000000000000000000000000000000000000..0615b5e6c6db910ee847188427b416fd812aa141 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_stack_to_reshape.cc @@ -0,0 +1,81 @@ +/* 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 +#include +#include +#include + +#include "absl/strings/str_cat.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool ConvertTrivialStackToReshape::Run(Model* model, std::size_t op_index) { + auto stack_it = model->operators.begin() + op_index; + if (stack_it->get()->type != OperatorType::kStack) { + return false; + } + auto* stack_op = static_cast(stack_it->get()); + if (stack_op->inputs.size() > 1) { + // Not trivial. + return false; + } + CHECK_EQ(stack_op->outputs.size(), 1); + + const auto& input_array = model->GetArray(stack_op->inputs[0]); + if (!input_array.has_shape()) { + // Yield until input dims have been resolved. + return false; + } + if (input_array.shape().dimensions_count() == 0) { + // Input array cannot be 0-D. + // (Unsure if this is TF behavior, but was required to get a test to pass.) + return false; + } + + AddMessageF("Converting trivial %s to a reshape", LogName(*stack_op)); + + // Note that we could convert to ExpandDims but toco prefers reshapes. + auto* reshape_op = new TensorFlowReshapeOperator; + reshape_op->inputs = {stack_op->inputs[0]}; + reshape_op->outputs = stack_op->outputs; + + // Create shape param. + string shape_array_name = + AvailableArrayName(*model, stack_op->outputs[0] + "_shape"); + Array& shape_array = model->GetOrCreateArray(shape_array_name); + *(shape_array.mutable_shape()->mutable_dims()) = { + 1 + input_array.shape().dimensions_count()}; + reshape_op->inputs.push_back(shape_array_name); + shape_array.data_type = ArrayDataType::kInt32; + auto& shape_buffer = shape_array.GetMutableBuffer(); + shape_buffer.data.push_back(1); + for (int dim : input_array.shape().dims()) { + shape_buffer.data.push_back(dim); + } + + // Replace the operator in the graph. + const auto reshape_it = model->operators.emplace(stack_it, reshape_op); + stack_it = reshape_it + 1; + CHECK_EQ(stack_it->get(), stack_op); + model->operators.erase(stack_it); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc b/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc new file mode 100644 index 0000000000000000000000000000000000000000..c2b166033c33b777bad88cb712adf8517be1762a --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc @@ -0,0 +1,85 @@ +/* Copyright 2017 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 + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool ConvertTrivialTransposeToReshape::Run(Model* model, std::size_t op_index) { + auto transpose_it = model->operators.begin() + op_index; + if (transpose_it->get()->type != OperatorType::kTranspose) { + return false; + } + TransposeOperator* transpose_op = + static_cast(transpose_it->get()); + + const auto& output_array = model->GetArray(transpose_op->outputs[0]); + if (!output_array.has_shape()) { + // Yield until PropagateFixedSizes has been run on this op. + return false; + } + // Note: We can assume we have error checked inputs in PropagateFixedSizes. + + // This transpose is trivial if we only have one non-unitary dimension. + std::vector const& dims = output_array.shape().dims(); + unsigned non_unitary_axis_count = 0; + for (int i = 0; i < dims.size(); i++) { + if (dims[i] != 1) { + non_unitary_axis_count++; + } + } + if (non_unitary_axis_count > 1) { + // Transpose is not trivial + return false; + } + + // This transpose is trivial. Replace it with a Reshape op. + auto* reshape_op = new TensorFlowReshapeOperator; + + // Copy input and output + reshape_op->inputs.push_back(transpose_op->inputs[0]); + reshape_op->outputs = transpose_op->outputs; + + // Create a new input array for the shape input + string perm_array_name = transpose_op->inputs[1]; + string shape_array_name = toco::AvailableArrayName(*model, perm_array_name); + Array& shape_array = model->GetOrCreateArray(shape_array_name); + *(shape_array.mutable_shape()->mutable_dims()) = { + 1, static_cast(dims.size())}; + reshape_op->inputs.push_back(shape_array_name); + shape_array.data_type = ArrayDataType::kInt32; + auto& shape_buffer = shape_array.GetMutableBuffer(); + shape_buffer.data = dims; + + // Delete perm array if unused + if (IsDiscardableArray(*model, perm_array_name) && + CountOpsWithInput(*model, perm_array_name) == 1) { + model->EraseArray(perm_array_name); + } + + // Replace the operator in the graph. + const auto reshape_it = model->operators.emplace(transpose_it, reshape_op); + transpose_it = reshape_it + 1; + CHECK_EQ(transpose_it->get(), transpose_op); + model->operators.erase(transpose_it); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/create_im2col_arrays.cc b/tensorflow/contrib/lite/toco/graph_transformations/create_im2col_arrays.cc index 1735b51e5b6ca517bad62bf55f0cc9f0c21ac440..076415ece8c1039caa32e947fe54ab3e101bec9e 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/create_im2col_arrays.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/create_im2col_arrays.cc @@ -35,7 +35,7 @@ bool CreateIm2colArrays::Run(Model* model, std::size_t op_index) { // We already have an im2col array return false; } - const auto& weights_array = *model->arrays[conv_op->inputs[1]]; + const auto& weights_array = model->GetArray(conv_op->inputs[1]); if (!weights_array.has_shape()) { // We need to yield until weights dims have been resolved, because // from the weights dims we determine whether an im2col array is diff --git a/tensorflow/contrib/lite/toco/graph_transformations/dequantize.cc b/tensorflow/contrib/lite/toco/graph_transformations/dequantize.cc index b89e3f5310cd7364294ad875cfcdf9c14660366b..498c864bde6d656c8318e981204cb42cb3a4d03f 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/dequantize.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/dequantize.cc @@ -53,7 +53,7 @@ std::vector>::iterator FindFirstOpWithInput( } void ClearArrayQuantizationParams(const string& array_name, Model* model) { - auto* array = model->arrays.at(array_name).get(); + auto* array = &model->GetArray(array_name); CHECK(array->quantization_params); for (auto& input_array : *model->flags.mutable_input_arrays()) { if (input_array.name() == array_name) { @@ -77,7 +77,7 @@ void ClearArrayQuantizationParams(const string& array_name, Model* model) { bool DequantizeArray(const string& array_name, GraphTransformation* transformation, Model* model) { - auto* array = model->arrays.at(array_name).get(); + auto* array = &model->GetArray(array_name); if (!array->quantization_params) { return false; } @@ -214,7 +214,9 @@ bool Dequantize::Run(Model* model, std::size_t op_index) { } bool changed = false; for (const string& array : arrays) { - changed |= DequantizeArray(array, this, model); + if (!model->IsOptionalArray(array)) { + changed |= DequantizeArray(array, this, model); + } } return changed; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/drop_fake_quant.cc b/tensorflow/contrib/lite/toco/graph_transformations/drop_fake_quant.cc index fea360740f4e645e1f00eaed42cbff48f430fe2a..95558ef5ece9a78825daf0203e2f6f6fee6f3cda 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/drop_fake_quant.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/drop_fake_quant.cc @@ -45,7 +45,7 @@ bool DropFakeQuant::Run(Model* model, std::size_t op_index) { // Drop min/max inputs for (int i = 1; i < fakequant_op->inputs.size(); i++) { if (CountOpsWithInput(*model, fakequant_op->inputs[i]) == 1) { - model->arrays.erase(fakequant_op->inputs[i]); + model->EraseArray(fakequant_op->inputs[i]); } } fakequant_op->inputs.resize(1); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/drop_im2col_arrays.cc b/tensorflow/contrib/lite/toco/graph_transformations/drop_im2col_arrays.cc index a3ed6663bcc80c5fc642a399b1e5c0cf3336973a..f7fd878b7e8b1c834125130ea2a778cecefd3de0 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/drop_im2col_arrays.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/drop_im2col_arrays.cc @@ -32,7 +32,7 @@ bool DropIm2colArrays::Run(Model* model, std::size_t op_index) { // Drop the im2col array. CHECK_EQ(conv_op->outputs.size(), 2); - model->arrays.erase(conv_op->outputs[1]); + model->EraseArray(conv_op->outputs[1]); conv_op->outputs.resize(1); AddMessageF("Dropped an im2col array for %s", LogName(*conv_op)); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/fuse_activation_functions.cc b/tensorflow/contrib/lite/toco/graph_transformations/fuse_activation_functions.cc index d129b5ecf2615434b8ff8387a04af9561fe617a4..ab943f72d1dd87ae9ff4bd53a807cd4923a88c38 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/fuse_activation_functions.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/fuse_activation_functions.cc @@ -68,11 +68,7 @@ bool FuseActivationFunctions::Run(Model* model, std::size_t op_index) { return false; } - // TODO(dkalenichenko): Great many ops don't support activation function - // fusing. Switch to the whilelist approach instead. - if (op->type == OperatorType::kConcatenation || - op->type == OperatorType::kSlice || - op->type == OperatorType::kTensorFlowSplit) { + if (!OperatorSupportsFusedActivation(op->type)) { AddMessageF( "Not fusing activation function because the %s op doesn't support it", LogName(*op)); @@ -90,7 +86,7 @@ bool FuseActivationFunctions::Run(Model* model, std::size_t op_index) { } else { LOG(FATAL) << "Unhandled activation function type"; } - model->arrays.erase(ac_op->inputs[0]); + model->EraseArray(ac_op->inputs[0]); op->outputs[0] = ac_op->outputs[0]; model->operators.erase(ac_it); return true; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc b/tensorflow/contrib/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc index 4619d8bbee2e52483a523277f421de5bfa155635..dcbbead517f26a227363989b5af2a4040c98ff57 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc @@ -285,13 +285,13 @@ bool FuseBinaryIntoFollowingAffine::Run(Model* model, std::size_t op_index) { AddMessageF("Fusing %s into the following %s", LogName(*binary_op), LogName(*following_op)); - model->arrays.erase(binary_op->outputs[0]); + model->EraseArray(binary_op->outputs[0]); following_op->inputs[0] = binary_op->inputs[index_of_variable_input]; const auto& old_constant_param_name = binary_op->inputs[index_of_constant_input]; CHECK(IsConstantParameterArray(*model, old_constant_param_name)); if (CountOpsWithInput(*model, old_constant_param_name) == 1) { - model->arrays.erase(old_constant_param_name); + model->EraseArray(old_constant_param_name); } model->operators.erase(binary_it); return true; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc b/tensorflow/contrib/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc index 8948653ec38f5a5a6e92cfe9e6bafdbf1aa9a962..5b57178b18d2d60e1f301a1a8b257d8057618550 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc @@ -309,7 +309,7 @@ bool FuseBinaryIntoPrecedingAffine::Run(Model* model, std::size_t op_index) { LOG(FATAL) << "should not get here"; } - model->arrays.erase(preceding_op->outputs[0]); + model->EraseArray(preceding_op->outputs[0]); preceding_op->outputs[0] = binary_op->outputs[0]; preceding_op->fused_activation_function = binary_op->fused_activation_function; @@ -317,7 +317,7 @@ bool FuseBinaryIntoPrecedingAffine::Run(Model* model, std::size_t op_index) { binary_op->inputs[index_of_constant_input]; CHECK(IsConstantParameterArray(*model, old_constant_param_name)); if (CountOpsWithInput(*model, old_constant_param_name) == 1) { - model->arrays.erase(old_constant_param_name); + model->EraseArray(old_constant_param_name); } model->operators.erase(binary_it); return true; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.cc b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.cc index 3e4c8fa2d9ae3902c626bd13c9c71f77a601c263..6961e23690a5e53643f2b2c52bb62ce395d05c95 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.cc @@ -31,13 +31,13 @@ namespace { void PrintModelStats(const string& label, const Model& model) { int quantized_arrays = 0; - for (const auto& array : model.arrays) { + for (const auto& array : model.GetArrayMap()) { if (array.second->quantization_params) { quantized_arrays++; } } LOG(INFO) << label << ": " << model.operators.size() << " operators, " - << model.arrays.size() << " arrays (" << quantized_arrays + << model.GetArrayMap().size() << " arrays (" << quantized_arrays << " quantized)"; } @@ -91,14 +91,9 @@ void DiscardUselessConnectedComponentsAndRNNBackEdges(Model* model) { } } while (found_new_useful_arrays); // Erase arrays that aren't useful, and that are discardable. - for (auto it = model->arrays.begin(); it != model->arrays.end();) { - if (useful_arrays.count(it->first) || - !IsDiscardableArray(*model, it->first)) { - ++it; - } else { - it = model->arrays.erase(it); - } - } + model->EraseArrays([&](const string& name) { + return (!useful_arrays.count(name) && IsDiscardableArray(*model, name)); + }); // Erase operators that do not produce a useful output array. for (auto it = model->operators.begin(); it != model->operators.end();) { // Only need to test the first output, as we simultaneously added all of @@ -118,8 +113,8 @@ void DiscardUselessConnectedComponentsAndRNNBackEdges(Model* model) { std::vector rnn_states_to_keep; for (const auto& rnn_state : model->flags.rnn_states()) { const bool dangling = - !model->arrays.count(rnn_state.back_edge_source_array()) || - !model->arrays.count(rnn_state.state_array()); + !model->HasArray(rnn_state.back_edge_source_array()) || + !model->HasArray(rnn_state.state_array()); if (dangling) { CHECK(rnn_state.discardable()); } else { @@ -137,6 +132,7 @@ bool GraphTransformationsPass(int increment, Model* model, CHECK(increment == 1 || increment == -1); bool changed = false; if (model->operators.empty()) { + LOG(INFO) << "Model is empty!!!"; return false; } int op_index = increment == 1 ? 0 : model->operators.size() - 1; @@ -147,23 +143,28 @@ bool GraphTransformationsPass(int increment, Model* model, CHECK(!changed_now); CHECK(transformation->Messages().empty()); changed_now = transformation->Run(model, op_index); - if (changed_now) { - DumpGraphvizVideoFrame(*model); - CHECK(!model->operators.empty()); - op_index = std::min(op_index, model->operators.size() - 1); - // Uncomment for debugging - // CheckInvariants(*model); - } const char* made_a_change_msg = changed_now ? "made a change" : "did NOT make a change"; const int log_level = changed_now ? kLogLevelModelChanged : kLogLevelModelUnchanged; + if (transformation->Messages().empty()) { + VLOG(log_level) << transformation->Name() << " " << made_a_change_msg + << " at op_index=" << op_index << "/" + << model->operators.size() - 1; + } for (const string& message : transformation->Messages()) { VLOG(log_level) << transformation->Name() << " " << made_a_change_msg << " at op_index=" << op_index << "/" << model->operators.size() - 1 << ": " << message; } transformation->ClearMessages(); + if (changed_now) { + DumpGraphvizVideoFrame(*model); + if (model->operators.empty()) return true; + op_index = std::min(op_index, model->operators.size() - 1); + // Uncomment for debugging + // CheckInvariants(*model); + } if (changed_now) { break; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h index 7349a0551a414a7e3e755c26ea73989d32e349d8..616bdac268c41d29135368d685729c961f44132b 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h +++ b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_GRAPH_TRANSFORMATIONS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_GRAPH_TRANSFORMATIONS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_GRAPH_TRANSFORMATIONS_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_GRAPH_TRANSFORMATIONS_H_ #include #include @@ -114,6 +114,10 @@ void RunGraphTransformations(Model* model, const string& message, // List of all graph transformations DECLARE_GRAPH_TRANSFORMATION(ConvertExpandDimsToReshape) DECLARE_GRAPH_TRANSFORMATION(ConvertPureConvToDepthwise) +DECLARE_GRAPH_TRANSFORMATION(ConvertTrivialAddNToAdd) +DECLARE_GRAPH_TRANSFORMATION(ConvertTrivialStackToReshape) +DECLARE_GRAPH_TRANSFORMATION(ConvertTrivialTransposeToReshape) +DECLARE_GRAPH_TRANSFORMATION(ConvertReorderAxes) DECLARE_GRAPH_TRANSFORMATION(EnsureBiasVectors) DECLARE_GRAPH_TRANSFORMATION(FuseActivationFunctions) DECLARE_GRAPH_TRANSFORMATION(FuseBinaryIntoFollowingAffine) @@ -121,6 +125,8 @@ DECLARE_GRAPH_TRANSFORMATION(FuseBinaryIntoPrecedingAffine) DECLARE_GRAPH_TRANSFORMATION(IdentifyL2Normalization) DECLARE_GRAPH_TRANSFORMATION(IdentifyL2Pool) DECLARE_GRAPH_TRANSFORMATION(IdentifyLstmCell) +DECLARE_GRAPH_TRANSFORMATION(SplitLstmCellInputs) +DECLARE_GRAPH_TRANSFORMATION(MergeLstmCellInputs) DECLARE_GRAPH_TRANSFORMATION(IdentifyRelu1) DECLARE_GRAPH_TRANSFORMATION(MakeInitialDequantizeOperator) DECLARE_GRAPH_TRANSFORMATION(PropagateArrayDataTypes) @@ -133,6 +139,7 @@ DECLARE_GRAPH_TRANSFORMATION(RemoveTensorFlowIdentity) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialBinaryOperator) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialConcatenation) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialConcatenationInput) +DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialSlice) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialQuantizedActivationFunc) DECLARE_GRAPH_TRANSFORMATION(RemoveUnusedOp) DECLARE_GRAPH_TRANSFORMATION(ResolveBatchNormalization) @@ -141,24 +148,33 @@ DECLARE_GRAPH_TRANSFORMATION(ResolveConstantUnaryOperator) DECLARE_GRAPH_TRANSFORMATION(CreateIm2colArrays) DECLARE_GRAPH_TRANSFORMATION(DropIm2colArrays) DECLARE_GRAPH_TRANSFORMATION(ReadFakeQuantMinMax) +DECLARE_GRAPH_TRANSFORMATION(ReorderActivationFunctions) DECLARE_GRAPH_TRANSFORMATION(ResolveReorderAxes) DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowConcat) DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowMatMul) DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowMerge) -DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowSqueeze) +DECLARE_GRAPH_TRANSFORMATION(ResolveSqueezeAttributes) DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowSwitch) DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowTile) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantFakeQuant) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantConcatenation) +DECLARE_GRAPH_TRANSFORMATION(ResolveConstantTranspose) DECLARE_GRAPH_TRANSFORMATION(DropFakeQuant) DECLARE_GRAPH_TRANSFORMATION(UnfuseActivationFunctions) +DECLARE_GRAPH_TRANSFORMATION(UnrollBatchMatMul) +DECLARE_GRAPH_TRANSFORMATION(ResolveSpaceToBatchNDAttributes) DECLARE_GRAPH_TRANSFORMATION(ResolveBatchToSpaceNDAttributes) DECLARE_GRAPH_TRANSFORMATION(ResolvePadAttributes) DECLARE_GRAPH_TRANSFORMATION(ResolveStridedSliceAttributes) DECLARE_GRAPH_TRANSFORMATION(ResolveSliceAttributes) DECLARE_GRAPH_TRANSFORMATION(ResolveMeanAttributes) -DECLARE_GRAPH_TRANSFORMATION(ResolveConstantTensorFlowShape) +DECLARE_GRAPH_TRANSFORMATION(ResolveTransposeAttributes) +DECLARE_GRAPH_TRANSFORMATION(ResolveConstantRange) +DECLARE_GRAPH_TRANSFORMATION(ResolveConstantShapeOrRank) +DECLARE_GRAPH_TRANSFORMATION(ResolveConstantStack) +DECLARE_GRAPH_TRANSFORMATION(ResolveConstantStridedSlice) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantFill) +DECLARE_GRAPH_TRANSFORMATION(ResolveMultiplyByZero) DECLARE_GRAPH_TRANSFORMATION(Dequantize) class ResolveReshapeAttributes : public GraphTransformation { @@ -186,4 +202,4 @@ class RemoveTrivialReshape : public GraphTransformation { } // end namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_GRAPH_TRANSFORMATIONS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_GRAPH_TRANSFORMATIONS_H_ diff --git a/tensorflow/contrib/lite/toco/graph_transformations/hardcode_min_max.cc b/tensorflow/contrib/lite/toco/graph_transformations/hardcode_min_max.cc index 9689b205cd137904504d87906cb691d0ed8235bf..1b0be858107b54f5a6ecd2a1cb87c9dbde1c06bb 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/hardcode_min_max.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/hardcode_min_max.cc @@ -177,6 +177,106 @@ bool HardcodeMinMaxForOutput(Model* model, Operator* op, double min, output_minmax.max = max; return true; } + +// Propagates MinMax from any of the listed arrays, to all others. +// If multiple of these arrays have MinMax, then these are required +// to agree with each other. +bool PropagateMinMaxAmongArrays(Model* model, + const std::vector array_names) { + string reference_array_name; + MinMax* reference_minmax = nullptr; + for (const string& array_name : array_names) { + if (model->GetArray(array_name).minmax) { + reference_array_name = array_name; + reference_minmax = model->GetArray(array_name).minmax.get(); + break; + } + } + // No MinMax info is available to propagate. + if (!reference_minmax) { + return false; + } + bool changed = false; + for (const string& array_name : array_names) { + auto& array = model->GetArray(array_name); + if (array.minmax) { + CHECK(*array.minmax == *reference_minmax) + << "Both the following arrays have minmax, and they disagree: " + << reference_array_name << " and " << array_name + << ". Expected that either only one of them would have minmax, or at " + "least that they would agree."; + } else { + array.GetOrCreateMinMax() = *reference_minmax; + changed = true; + } + } + return changed; +} + +bool HardcodeMinMaxForLstmCell(Model* model, Operator* op) { + CHECK_EQ(op->inputs.size(), LstmCellOperator::NUM_INPUTS); + CHECK_EQ(op->outputs.size(), LstmCellOperator::NUM_OUTPUTS); + + bool changed = false; + changed |= PropagateMinMaxAmongArrays( + model, {op->inputs[LstmCellOperator::PREV_STATE_INPUT], + op->outputs[LstmCellOperator::STATE_OUTPUT]}); + + auto& input_activations = + model->GetArray(op->inputs[LstmCellOperator::DATA_INPUT]); + if (!input_activations.minmax) { + auto& minmax = input_activations.GetOrCreateMinMax(); + minmax.min = -1; + minmax.max = 127. / 128.; + changed = true; + } + + auto& prev_output_activations = + model->GetArray(op->inputs[LstmCellOperator::PREV_ACTIV_INPUT]); + if (!prev_output_activations.minmax) { + auto& minmax = prev_output_activations.GetOrCreateMinMax(); + minmax.min = -1; + minmax.max = 127. / 128.; + changed = true; + } + + auto& output_concat_temp = + model->GetArray(op->outputs[LstmCellOperator::CONCAT_TEMP]); + if (!output_concat_temp.minmax) { + auto& minmax = output_concat_temp.GetOrCreateMinMax(); + minmax.min = -1; + minmax.max = 127. / 128.; + changed = true; + } + + auto& output_activations = + model->GetArray(op->outputs[LstmCellOperator::ACTIV_OUTPUT]); + if (!output_activations.minmax) { + auto& minmax = output_activations.GetOrCreateMinMax(); + minmax.min = -1; + minmax.max = 127. / 128.; + changed = true; + } + + // (This comment should morph into proper documentation for + // quantization of LSTM models. It isn't just a local implementation detail, + // the training code for LSTM models needs to be adjusted to that.) + // + // Finally, output_activations_temp holds the output of the fully-connected + // node inside the LSTM cell. For it, we hardcode a minmax of [-8, 8]. + // The rationale for that is given in a lengthy comment on the LstmCell + // quantized runtime implementation in reference_ops.h. + auto& output_activations_temp = + model->GetArray(op->outputs[LstmCellOperator::ACTIV_TEMP]); + if (!output_activations_temp.minmax) { + auto& minmax = output_activations_temp.GetOrCreateMinMax(); + minmax.min = -8; + minmax.max = 8 * 32767. / 32768.; + changed = true; + } + + return changed; +} } // namespace bool HardcodeMinMax::Run(Model* model, std::size_t op_index) { @@ -219,6 +319,16 @@ bool HardcodeMinMax::Run(Model* model, std::size_t op_index) { changed = HardcodeMinMaxForOutput(model, op, 0, 255. / 256.); break; + case OperatorType::kTanh: + // We hardcode quantization_params to: zero_point=127, scale=1/128. + // This choice of minmax is the one that is equivalent to that. + changed = HardcodeMinMaxForOutput(model, op, -127. / 128., 1.0); + break; + + case OperatorType::kLstmCell: + changed = HardcodeMinMaxForLstmCell(model, op); + break; + default: break; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_l2_normalization.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_l2_normalization.cc index 01b75e37c691d48fabf8832af04543be3f5eb3bc..419a0776a6b987a18df059d3c1d4bf4370cd24d8 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/identify_l2_normalization.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_l2_normalization.cc @@ -150,19 +150,19 @@ bool IdentifyL2Normalization::Run(Model* model, std::size_t op_index) { // Erase the subgraph that is now replaced by L2Normalization model->operators.erase(FindOperator(model, square_op)); - model->arrays.erase(sum_op->inputs[0]); + model->EraseArray(sum_op->inputs[0]); if (sum_op->inputs.size() > 1) { - model->arrays.erase(sum_op->inputs[1]); + model->EraseArray(sum_op->inputs[1]); } model->operators.erase(FindOperator(model, sum_op)); if (add_op) { - model->arrays.erase(add_op->inputs[0]); - model->arrays.erase(add_op->inputs[1]); + model->EraseArray(add_op->inputs[0]); + model->EraseArray(add_op->inputs[1]); model->operators.erase(FindOperator(model, add_op)); } - model->arrays.erase(sqrt_or_rsqrt_op->inputs[0]); + model->EraseArray(sqrt_or_rsqrt_op->inputs[0]); model->operators.erase(FindOperator(model, sqrt_or_rsqrt_op)); - model->arrays.erase(div_or_mul_op->inputs[1]); + model->EraseArray(div_or_mul_op->inputs[1]); model->operators.erase(FindOperator(model, div_or_mul_op)); return true; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_l2_pool.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_l2_pool.cc index 1865416fc2226d663dfd51a5c0a0e2129caf485c..e4d52476c649de53b3ab663f53ce7a5538dbb5ab 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/identify_l2_pool.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_l2_pool.cc @@ -92,8 +92,8 @@ bool IdentifyL2Pool::Run(Model* model, std::size_t op_index) { AddMessageF("Creating %s replacing equivalent subgraph", LogName(*l2pool_op)); // Erase intermediate arrays, keeping input to square op. - model->arrays.erase(avpool_op->inputs[0]); - model->arrays.erase(sqrt_op->inputs[0]); + model->EraseArray(avpool_op->inputs[0]); + model->EraseArray(sqrt_op->inputs[0]); // Erase three operators being replaced. model->operators.erase(FindOperator(model, square_op)); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm.cc index 082820fddcf137238867239bbc4d4eed8158e307..c363b93394f0af7bcfc37c1e8be5f98aca6667ae 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm.cc @@ -16,7 +16,6 @@ limitations under the License. #include #include -#include "absl/strings/string_view.h" #include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" #include "tensorflow/contrib/lite/toco/model.h" #include "tensorflow/contrib/lite/toco/tooling_util.h" @@ -202,23 +201,6 @@ bool MatchOperatorInputs(const Operator& op, const Model& model, return true; } -absl::string_view FindLongestCommonPrefix(absl::string_view a, - absl::string_view b) { - if (a.empty() || b.empty()) return absl::string_view(); - - const char* pa = a.data(); - const char* pb = b.data(); - size_t count = 0; - const ssize_t limit = std::min(a.size(), b.size()); - while (count < limit && *pa == *pb) { - ++pa; - ++pb; - ++count; - } - - return absl::string_view(a.data(), count); -} - } // namespace bool IdentifyLstmCell::Run(Model* model, std::size_t op_index) { diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_merge_inputs.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_merge_inputs.cc new file mode 100644 index 0000000000000000000000000000000000000000..45335fd78c99a577d535770d78acf4fcd6c04531 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_merge_inputs.cc @@ -0,0 +1,185 @@ +/* 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 +#include +#include + +#include "absl/memory/memory.h" +#include "absl/strings/string_view.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" + +namespace toco { + +bool MergeLstmCellInputs::Run(Model* model, std::size_t op_index) { + // Find lstm cell. + auto op_it = model->operators.begin() + op_index; + auto src_op = op_it->get(); + if (src_op->type != OperatorType::kLstmCell) { + return false; + } + + // Already a compact LstmCell with LstmCellOperator::NUM_INPUTS of inputs, + // do not need to merge cell inputs. + if (src_op->inputs.size() == LstmCellOperator::NUM_INPUTS) { + return false; + } + + // Identify prev_activ_input, prev_state_input as required Op inputs, + // using the rnn_states in the model flag. + string prev_activ_input; + if (!GetMatchingRnnArray(model, src_op->outputs[kOutputTensor], + &prev_activ_input)) { + return false; + } + string prev_state_input; + if (!GetMatchingRnnArray(model, src_op->outputs[kCellStateTensor], + &prev_state_input)) { + return false; + } + + // Get LstmCell's cell, input, output size. + int num_cell = model->GetArray(src_op->inputs[kInputToInputWeightsTensor]) + .shape() + .dims(0); + int num_input = model->GetArray(src_op->inputs[kInputToInputWeightsTensor]) + .shape() + .dims(1); + int num_output = + model->GetArray(src_op->inputs[kRecurrentToInputWeightsTensor]) + .shape() + .dims(1); + + // Make sure n_cell and n_output are equal as there is no projection. + CHECK_EQ(num_cell, num_output); + + // Create tensorflow_graphdef style's one big weight tensor. + const string base_name(FindLongestCommonPrefix( + src_op->outputs[kOutputTensor], src_op->outputs[kCellStateTensor])); + string merged_weights = AvailableArrayName(*model, base_name + "weights"); + auto& array = model->GetOrCreateArray(merged_weights); + array.data_type = ArrayDataType::kFloat; + int weights_dim1 = 4 * num_cell; + int weights_dim2 = num_input + num_output; + Shape shape = Shape({weights_dim1, weights_dim2}); + array.copy_shape(shape); + auto& buffer = array.GetMutableBuffer(); + buffer.data.resize(weights_dim1 * weights_dim2); + + // Merge 8 small weight tensors to 1 weight tensor. + CopyArrayToSubArray( + buffer, weights_dim2, + model->GetArray(src_op->inputs[kInputToInputWeightsTensor]), 0, 0); + CopyArrayToSubArray( + buffer, weights_dim2, + model->GetArray(src_op->inputs[kInputToCellWeightsTensor]), num_cell, 0); + CopyArrayToSubArray( + buffer, weights_dim2, + model->GetArray(src_op->inputs[kInputToForgetWeightsTensor]), + num_cell * 2, 0); + CopyArrayToSubArray( + buffer, weights_dim2, + model->GetArray(src_op->inputs[kInputToOutputWeightsTensor]), + num_cell * 3, 0); + CopyArrayToSubArray( + buffer, weights_dim2, + model->GetArray(src_op->inputs[kRecurrentToInputWeightsTensor]), 0, + num_input); + CopyArrayToSubArray( + buffer, weights_dim2, + model->GetArray(src_op->inputs[kRecurrentToCellWeightsTensor]), num_cell, + num_input); + CopyArrayToSubArray( + buffer, weights_dim2, + model->GetArray(src_op->inputs[kRecurrentToForgetWeightsTensor]), + num_cell * 2, num_input); + CopyArrayToSubArray( + buffer, weights_dim2, + model->GetArray(src_op->inputs[kRecurrentToOutputWeightsTensor]), + num_cell * 3, num_input); + + // Create tensorflow_graphdef style's one big bias tensor. + string merged_biases = AvailableArrayName(*model, base_name + "biases"); + auto& bias_array = model->GetOrCreateArray(merged_biases); + bias_array.data_type = ArrayDataType::kFloat; + bias_array.copy_shape(Shape({weights_dim1})); + auto& bias_buffer = bias_array.GetMutableBuffer(); + bias_buffer.data.resize(weights_dim1); + + // Merge 4 small bias tensors into a big one. + CopyArrayToSubArray(bias_buffer, weights_dim2, + model->GetArray(src_op->inputs[kInputGateBiasTensor]), 0, + 0); + CopyArrayToSubArray(bias_buffer, weights_dim2, + model->GetArray(src_op->inputs[kCellGateBiasTensor]), + num_cell, 0); + CopyArrayToSubArray(bias_buffer, weights_dim2, + model->GetArray(src_op->inputs[kForgetGateBiasTensor]), + num_cell * 2, 0); + CopyArrayToSubArray(bias_buffer, weights_dim2, + model->GetArray(src_op->inputs[kOutputGateBiasTensor]), + num_cell * 3, 0); + + // Emplace a new LSTM cell operator (use basic 5 inputs kernel). + auto lstm_cell_op = absl::make_unique(); + + // Compact LstmCell's 5 inputs. + lstm_cell_op->inputs.resize(LstmCellOperator::NUM_INPUTS); + lstm_cell_op->inputs[LstmCellOperator::DATA_INPUT] = + src_op->inputs[kInputTensor]; + lstm_cell_op->inputs[LstmCellOperator::WEIGHTS_INPUT] = merged_weights; + lstm_cell_op->inputs[LstmCellOperator::BIASES_INPUT] = merged_biases; + lstm_cell_op->inputs[LstmCellOperator::PREV_ACTIV_INPUT] = prev_activ_input; + lstm_cell_op->inputs[LstmCellOperator::PREV_STATE_INPUT] = prev_state_input; + + // Reorder LstmCell's 4 outputs. + lstm_cell_op->outputs.resize(LstmCellOperator::NUM_OUTPUTS); + lstm_cell_op->outputs[LstmCellOperator::ACTIV_OUTPUT] = + src_op->outputs[kOutputTensor]; + lstm_cell_op->outputs[LstmCellOperator::STATE_OUTPUT] = + src_op->outputs[kCellStateTensor]; + lstm_cell_op->outputs[LstmCellOperator::CONCAT_TEMP] = + src_op->outputs[kScratchBufferTensor]; + lstm_cell_op->outputs[LstmCellOperator::ACTIV_TEMP] = + src_op->outputs[kOutputStateTensor]; + + // Add the op into model. + model->operators.emplace(op_it, std::move(lstm_cell_op)); + AddMessageF("Creating compact LstmCell replacing previous lstm cell"); + + // Delete arrays and operators replaced by the LSTM cell operator. Order is + // important - DeleteArrayIfUnused() only succeeds if dependent operators + // have been removed first. Start at the output and work towards the input. + // Erase curr lstm op being replaced. + DeleteArrayIfUnused(src_op->inputs[kInputToInputWeightsTensor], model); + DeleteArrayIfUnused(src_op->inputs[kInputToForgetWeightsTensor], model); + DeleteArrayIfUnused(src_op->inputs[kInputToCellWeightsTensor], model); + DeleteArrayIfUnused(src_op->inputs[kInputToOutputWeightsTensor], model); + DeleteArrayIfUnused(src_op->inputs[kRecurrentToInputWeightsTensor], model); + DeleteArrayIfUnused(src_op->inputs[kRecurrentToForgetWeightsTensor], model); + DeleteArrayIfUnused(src_op->inputs[kRecurrentToCellWeightsTensor], model); + DeleteArrayIfUnused(src_op->inputs[kRecurrentToOutputWeightsTensor], model); + DeleteArrayIfUnused(src_op->inputs[kInputGateBiasTensor], model); + DeleteArrayIfUnused(src_op->inputs[kForgetGateBiasTensor], model); + DeleteArrayIfUnused(src_op->inputs[kCellGateBiasTensor], model); + DeleteArrayIfUnused(src_op->inputs[kOutputGateBiasTensor], model); + model->operators.erase(FindOp(*model, src_op)); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_split_inputs.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_split_inputs.cc new file mode 100644 index 0000000000000000000000000000000000000000..eca717680af281018b919c27068ba5d9f5699d69 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_split_inputs.cc @@ -0,0 +1,171 @@ +/* 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 +#include +#include + +#include "absl/memory/memory.h" +#include "absl/strings/string_view.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" + +namespace toco { + +bool SplitLstmCellInputs::Run(Model* model, std::size_t op_index) { + // Find lstm cell. + auto op_it = model->operators.begin() + op_index; + auto curr_op = op_it->get(); + if (curr_op->type != OperatorType::kLstmCell) { + return false; + } + + // Already an extended LstmCell with kExtendedLstmInputCount of inputs, + // do not need to split cell inputs. + if (curr_op->inputs.size() == kExtendedLstmInputCount) { + return false; + } + + // Make sure the WEIGHTS_INPUT and BIASES_INPUT are constant arrays, + // that are able to be split into smaller weight and bias tensors. + if (!IsConstantParameterArray( + *model, curr_op->inputs[LstmCellOperator::WEIGHTS_INPUT]) || + !IsConstantParameterArray( + *model, curr_op->inputs[LstmCellOperator::BIASES_INPUT])) { + return false; + } + + // Make sure propagate_fixed_sizes has defined the size of the output. + if (!model->GetArray(curr_op->outputs[LstmCellOperator::ACTIV_OUTPUT]) + .has_shape()) { + return false; + } + + // Emplace a new LstmCell operator with extended inputs (kernel/lstm.cc). + auto lstm_cell_op = absl::make_unique(); + lstm_cell_op->inputs.resize(kExtendedLstmInputCount); + int num_input = model->GetArray(curr_op->inputs[LstmCellOperator::DATA_INPUT]) + .shape() + .dims(1); + + // n_cell and n_output have the same size when there is no projection. + int num_cell = + model->GetArray(curr_op->outputs[LstmCellOperator::ACTIV_OUTPUT]) + .shape() + .dims(1); + int num_output = num_cell; + + // Data input. + lstm_cell_op->inputs[kInputTensor] = + curr_op->inputs[LstmCellOperator::ACTIV_OUTPUT]; + + // Get original weight tensor and decompose 1 tensor to 8 sub tensors. + Array& kernel = + model->GetArray(curr_op->inputs[LstmCellOperator::WEIGHTS_INPUT]); + const string base_name(FindLongestCommonPrefix( + curr_op->outputs[LstmCellOperator::ACTIV_OUTPUT], + curr_op->outputs[LstmCellOperator::STATE_OUTPUT])); + + // Input weight tensors of size {n_cell, n_input}. + CopySubArrayToArray( + model, &(lstm_cell_op->inputs[kInputToInputWeightsTensor]), + base_name + "weight_i_i", num_cell, num_input, kernel, 0, 0); + CopySubArrayToArray(model, &(lstm_cell_op->inputs[kInputToCellWeightsTensor]), + base_name + "weight_c_i", num_cell, num_input, kernel, + num_cell, 0); + CopySubArrayToArray( + model, &(lstm_cell_op->inputs[kInputToForgetWeightsTensor]), + base_name + "weight_f_i", num_cell, num_input, kernel, num_cell * 2, 0); + CopySubArrayToArray( + model, &(lstm_cell_op->inputs[kInputToOutputWeightsTensor]), + base_name + "weight_o_i", num_cell, num_input, kernel, num_cell * 3, 0); + + // Recurrent weight tensors of size {n_cell, n_output}. + CopySubArrayToArray( + model, &(lstm_cell_op->inputs[kRecurrentToInputWeightsTensor]), + base_name + "weight_i_r", num_cell, num_output, kernel, 0, num_input); + CopySubArrayToArray(model, + &(lstm_cell_op->inputs[kRecurrentToCellWeightsTensor]), + base_name + "weight_c_r", num_cell, num_output, kernel, + num_cell, num_input); + CopySubArrayToArray(model, + &(lstm_cell_op->inputs[kRecurrentToForgetWeightsTensor]), + base_name + "weight_f_r", num_cell, num_output, kernel, + num_cell * 2, num_input); + CopySubArrayToArray(model, + &(lstm_cell_op->inputs[kRecurrentToOutputWeightsTensor]), + base_name + "weight_o_r", num_cell, num_output, kernel, + num_cell * 3, num_input); + + // Peephole (optional). + CreateOptionalArray(model, &(lstm_cell_op->inputs[kCellToInputWeightsTensor]), + base_name + "peephole_c_i"); + CreateOptionalArray(model, + &(lstm_cell_op->inputs[kCellToForgetWeightsTensor]), + base_name + "peephole_c_f"); + CreateOptionalArray(model, + &(lstm_cell_op->inputs[kCellToOutputWeightsTensor]), + base_name + "peephole_c_o"); + + // Get original bias tensor and decompose 1 tensor to 4 sub tensors + Array& bias = + model->GetArray(curr_op->inputs[LstmCellOperator::BIASES_INPUT]); + CopySubArrayToArray(model, &(lstm_cell_op->inputs[kInputGateBiasTensor]), + base_name + "bias_i", num_cell, 1, bias, 0, 0); + CopySubArrayToArray(model, &(lstm_cell_op->inputs[kCellGateBiasTensor]), + base_name + "bias_c", num_cell, 1, bias, num_cell, 0); + CopySubArrayToArray(model, &(lstm_cell_op->inputs[kForgetGateBiasTensor]), + base_name + "bias_f", num_cell, 1, bias, num_cell * 2, 0); + CopySubArrayToArray(model, &(lstm_cell_op->inputs[kOutputGateBiasTensor]), + base_name + "bias_o", num_cell, 1, bias, num_cell * 3, 0); + + // Projection (optional). + CreateOptionalArray(model, &(lstm_cell_op->inputs[kProjectionWeightsTensor]), + base_name + "proj_weight"); + CreateOptionalArray(model, &(lstm_cell_op->inputs[kProjectionBiasTensor]), + base_name + "proj_bias"); + + // Reorder LstmCell's outputs. + lstm_cell_op->outputs.resize(LstmCellOperator::NUM_OUTPUTS); + lstm_cell_op->outputs[kScratchBufferTensor] = + curr_op->outputs[LstmCellOperator::CONCAT_TEMP]; + lstm_cell_op->outputs[kOutputStateTensor] = + curr_op->outputs[LstmCellOperator::ACTIV_TEMP]; + lstm_cell_op->outputs[kCellStateTensor] = + curr_op->outputs[LstmCellOperator::STATE_OUTPUT]; + lstm_cell_op->outputs[kOutputTensor] = + curr_op->outputs[LstmCellOperator::ACTIV_OUTPUT]; + + // Add the op into model. + model->operators.emplace(op_it, std::move(lstm_cell_op)); + AddMessageF("Creating extended LstmCell replacing previous lstm cell"); + + // Delete arrays and operators replaced by the LSTM cell operator. Order is + // important - DeleteArrayIfUnused() only succeeds if dependent operators + // have been removed first. Start at the output and work towards the input. + // Erase curr lstm op being replaced. + DeleteArrayIfUnused(curr_op->inputs[LstmCellOperator::WEIGHTS_INPUT], model); + DeleteArrayIfUnused(curr_op->inputs[LstmCellOperator::BIASES_INPUT], model); + DeleteArrayIfUnused(curr_op->inputs[LstmCellOperator::PREV_ACTIV_INPUT], + model); + DeleteArrayIfUnused(curr_op->inputs[LstmCellOperator::PREV_STATE_INPUT], + model); + model->operators.erase(FindOp(*model, curr_op)); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_relu1.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_relu1.cc index cfc77024e7e56038878570c9d3a462715a53ae3f..de6d8889fb4ccdb56e9639ab0dd7d093bfa4b908 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/identify_relu1.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_relu1.cc @@ -57,45 +57,60 @@ int GetSingleScalarInputIndexOfBinaryOp(Model* model, const Operator* op, } // namespace bool IdentifyRelu1::Run(Model* model, std::size_t op_index) { - const auto maximum_it = model->operators.begin() + op_index; - const auto* maximum_op = maximum_it->get(); - if (maximum_op->type != OperatorType::kTensorFlowMaximum) { + // Follow sequences of min+max and max+min. First get the leading op. + const auto op_it = model->operators.begin() + op_index; + const auto* op_0 = op_it->get(); + if (op_0->type != OperatorType::kTensorFlowMinimum && + op_0->type != OperatorType::kTensorFlowMaximum) { return false; } - CHECK_EQ(maximum_op->inputs.size(), 2); - if (maximum_op->outputs.size() != 1) { - return false; - } - int scalar_input_index = - GetSingleScalarInputIndexOfBinaryOp(model, maximum_op, -1.0f); - if (scalar_input_index == -1) { + + // Get the paired op and ensure it's the counter to the first. + const auto* op_1 = GetOpWithInput(*model, op_0->outputs[0]); + if (!op_1 || + (op_1->type != OperatorType::kTensorFlowMinimum && + op_1->type != OperatorType::kTensorFlowMaximum) || + op_0->type == op_1->type) { return false; } - const auto* minimum_op = GetOpWithInput(*model, maximum_op->outputs[0]); - if (!minimum_op || minimum_op->type != OperatorType::kTensorFlowMinimum) { + + const auto* min_op = + op_0->type == OperatorType::kTensorFlowMinimum ? op_0 : op_1; + const auto* max_op = + op_0->type == OperatorType::kTensorFlowMaximum ? op_0 : op_1; + + CHECK_EQ(min_op->inputs.size(), 2); + CHECK_EQ(max_op->inputs.size(), 2); + if (min_op->outputs.size() != 1 || max_op->outputs.size() != 1) { return false; } - if (GetSingleScalarInputIndexOfBinaryOp(model, minimum_op, 1.0f) == -1) { + + // Get the original input to the min+max pair. + int min_scalar_input_index = + GetSingleScalarInputIndexOfBinaryOp(model, min_op, 1.0f); + int max_scalar_input_index = + GetSingleScalarInputIndexOfBinaryOp(model, max_op, -1.0f); + if (min_scalar_input_index == -1 || max_scalar_input_index == -1) { return false; } - CHECK_EQ(minimum_op->inputs.size(), 2); + int op_0_scalar_input_index = + op_0 == min_op ? min_scalar_input_index : max_scalar_input_index; - // Create and emplace Relu1 node + // Create and emplace Relu1 node. auto* relu1_op = new Relu1Operator; - relu1_op->inputs = {maximum_op->inputs[!scalar_input_index]}; - relu1_op->outputs = minimum_op->outputs; - model->operators.emplace(maximum_it, relu1_op); + relu1_op->inputs = {op_0->inputs[!op_0_scalar_input_index]}; + relu1_op->outputs = op_1->outputs; + model->operators.emplace(op_it, relu1_op); AddMessageF("Creating %s replacing equivalent subgraph", LogName(*relu1_op)); - // Erase Maximum scalar input & operator - model->arrays.erase(maximum_op->inputs[scalar_input_index]); - model->operators.erase(FindOperator(model, maximum_op)); - - // Erase Minimum inputs & operator - model->arrays.erase(minimum_op->inputs[0]); - model->arrays.erase(minimum_op->inputs[1]); - model->operators.erase(FindOperator(model, minimum_op)); + // Erase op scalar inputs & operators. Note that we preserve the non-scalar + // input to the first op as that's been redirected to the relu1_op. + DeleteArrayIfUsedOnce(op_0->inputs[op_0_scalar_input_index], model); + DeleteArrayIfUsedOnce(op_1->inputs[0], model); + DeleteArrayIfUsedOnce(op_1->inputs[1], model); + model->operators.erase(FindOperator(model, op_0)); + model->operators.erase(FindOperator(model, op_1)); return true; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.cc b/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.cc new file mode 100644 index 0000000000000000000000000000000000000000..910a96058979887972b41f27b2e570e8cb5b4f4c --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.cc @@ -0,0 +1,97 @@ +/* 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/lite/toco/graph_transformations/lstm_utils.h" + +namespace toco { + +void CreateOptionalArray(Model* model, string* input_array_buffer, + const string& array_name) { + *input_array_buffer = array_name; + model->CreateOptionalArray(array_name); +} + +void CopyArrayData(const Buffer& src_buffer, + int src_stride, int src_start_idx1, int src_start_idx2, + Buffer* dst_buffer, int dst_stride, + int dst_start_idx1, int dst_start_idx2, int dim1_copy_size, + int dim2_copy_size) { + int src_offset = src_start_idx1 * src_stride + src_start_idx2; + int dst_offset = dst_start_idx1 * dst_stride + dst_start_idx2; + for (int i = 0; i < dim1_copy_size; i++) { + for (int j = 0; j < dim2_copy_size; j++) { + int idx_src = src_offset + i * src_stride + j; + int idx_dst = dst_offset + i * dst_stride + j; + dst_buffer->data[idx_dst] = src_buffer.data[idx_src]; + } + } +} + +Buffer* CreateFloatArrayBuffer(Model* model, + string* array_name, + const Shape& shape) { + *array_name = AvailableArrayName(*model, *array_name); + auto& array = model->GetOrCreateArray(*array_name); + array.data_type = ArrayDataType::kFloat; + array.copy_shape(shape); + Buffer* buffer = + &(array.GetMutableBuffer()); + buffer->data.resize(RequiredBufferSizeForShape(shape)); + return buffer; +} + +void CopySubArrayToArray(Model* model, string* array_name, + const string& tensor_name, int dim1_size, + int dim2_size, const Array& original_array, + int start_idx1, int start_idx2) { + // Determine whether it's bias or not, create shape, buffer. + bool is_bias = dim2_size == 1; + Shape shape = is_bias ? Shape({dim1_size}) : Shape({dim1_size, dim2_size}); + Buffer* buffer = + CreateFloatArrayBuffer(model, array_name, shape); + auto& orig_buffer = original_array.GetBuffer(); + + // Copy data from big tensor. + CopyArrayData(orig_buffer, is_bias ? 1 : original_array.shape().dims(1), + start_idx1, start_idx2, buffer, dim2_size, 0, 0, dim1_size, + dim2_size); +} + +void CopyArrayToSubArray(Buffer& tensor_buffer, + int tensor_stride, const Array& sub_array, + int start_idx1, int start_idx2) { + // Get tensor data. + bool is_bias = sub_array.shape().dims().size() == 1; + int dim1_copy_size = sub_array.shape().dims()[0]; + int dim2_copy_size = is_bias ? 1 : sub_array.shape().dims(1); + auto& sub_buffer = sub_array.GetBuffer(); + + // Copy data from sub tensor. + CopyArrayData(sub_buffer, dim2_copy_size, 0, 0, &tensor_buffer, + is_bias ? 1 : tensor_stride, start_idx1, start_idx2, + dim1_copy_size, dim2_copy_size); +} + +bool GetMatchingRnnArray(Model* model, const string& back_edge_source_array, + string* rnn_array) { + for (const auto& rnn_state : model->flags.rnn_states()) { + if (rnn_state.back_edge_source_array() == back_edge_source_array) { + *rnn_array = rnn_state.state_array(); + return true; + } + } + return false; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h b/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..881c2d4dc892625d4640cac867a2f49c24b638f5 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h @@ -0,0 +1,102 @@ +/* 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 +#include +#include + +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" + +namespace toco { + +// For consistency with the parameters defined in extended LstmCell's kernel +// (tensorflow/contrib/lite/kernels/lstm.cc), +// use lowercase for these constants. + +enum ExtendedLstmCellInputs { + kInputTensor = 0, + kInputToInputWeightsTensor = 1, // Optional + kInputToForgetWeightsTensor = 2, + kInputToCellWeightsTensor = 3, + kInputToOutputWeightsTensor = 4, + kRecurrentToInputWeightsTensor = 5, // Optional + kRecurrentToForgetWeightsTensor = 6, + kRecurrentToCellWeightsTensor = 7, + kRecurrentToOutputWeightsTensor = 8, + kCellToInputWeightsTensor = 9, // Optional + kCellToForgetWeightsTensor = 10, // Optional + kCellToOutputWeightsTensor = 11, // Optional + kInputGateBiasTensor = 12, // Optional + kForgetGateBiasTensor = 13, + kCellGateBiasTensor = 14, + kOutputGateBiasTensor = 15, + kProjectionWeightsTensor = 16, // Optional + kProjectionBiasTensor = 17, // Optional + kExtendedLstmInputCount = 18 +}; + +enum ExtendedLstmCellOutputs { + kScratchBufferTensor = 0, + kOutputStateTensor = 1, + kCellStateTensor = 2, + kOutputTensor = 3 +}; + +// Create optional array used for optional tensor in ExtendedLstmCell inputs. +void CreateOptionalArray(Model* model, string* input_array_buffer, + const string& array_name); + +// Create float array and get its buffer. +Buffer* CreateFloatArrayBuffer(Model* model, + string* array_name, + const Shape& shape); + +// Copy data from one array to the other one (supports 1D and 2D array), +// for 1D array, the 2nd dim's size is 1. +// Arguments: +// src_buffer: the source buffer +// src_stride: the stride of source buffer, i.e., 2nd dim's size +// src_start_idx1: the 1st dim index of start point in src matrix +// src_start_idx2: the 2nd dim index of start point in src matrix +// dst_buffer: the destination buffer +// dst_stride: the stride of destination buffer, i.e., 2nd dim's size +// dst_start_idx1: the 1st dim index of start point in dst matrix +// dst_start_idx2: the 2nd dim index of start point in dst matrix +// dim1_copy_size: 1st dim size of copy data +// dim2_copy_size: 2nd dim size of copy data +void CopyArrayData(const Buffer& src_buffer, + int src_stride, int src_start_idx1, int src_start_idx2, + Buffer* dst_buffer, int dst_stride, + int dst_start_idx1, int dst_start_idx2, int dim1_copy_size, + int dim2_copy_size); + +// Copy a subset of array data and create a smaller array, +// mostly used for spliting weights and bias for Lstm cell. +void CopySubArrayToArray(Model* model, string* array_name, + const string& tensor_name, int dim1_size, + int dim2_size, const Array& original_array, + int start_idx1, int start_idx2); + +// Copy array data to a large array's submatrix, +// mostly used for merging weights and bias for Lstm cell. +void CopyArrayToSubArray(Buffer& tensor_buffer, + int tensor_stride, const Array& sub_array, + int start_idx1, int start_idx2); + +// Get mating rnn array inputs using rnn_states flag. +bool GetMatchingRnnArray(Model* model, const string& back_edge_source_array, + string* rnn_array); + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc index 4fe127544be6f4439b184e1fcf4436eda4a53cc5..f0d107232b4517115aa3f64b39b825dbaffb83ce 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc @@ -27,7 +27,7 @@ namespace { void SetDataTypeForAllOutputs(Model* model, Operator* op, ArrayDataType data_type) { for (const auto& output : op->outputs) { - model->arrays[output]->data_type = data_type; + model->GetArray(output).data_type = data_type; } } } // namespace @@ -38,7 +38,8 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) { // If the data type of some input is unknown, we need to yield. for (const auto& input : op->inputs) { - if (model->arrays[input]->data_type == ArrayDataType::kNone) { + if (!model->IsOptionalArray(input) && + model->GetArray(input).data_type == ArrayDataType::kNone) { return false; } } @@ -46,7 +47,7 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) { // end if we changed anything, and return the correct boolean value. std::unordered_map old_output_data_types; for (const auto& output : op->outputs) { - old_output_data_types[output] = model->arrays[output]->data_type; + old_output_data_types[output] = model->GetArray(output).data_type; } // Do the actual output data types propagation. if (op->type == OperatorType::kDequantize || @@ -61,25 +62,39 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) { SetDataTypeForAllOutputs(model, op, ArrayDataType::kBool); } else if (op->type == OperatorType::kRank || op->type == OperatorType::kTensorFlowShape) { - // These operators are assumed to produce int32 outputs. + // These operators only produce int32 outputs. SetDataTypeForAllOutputs(model, op, ArrayDataType::kInt32); } else if (op->type == OperatorType::kTensorFlowSplit || op->type == OperatorType::kTensorFlowConcat || op->type == OperatorType::kFill) { // These operators produce an output with the same type as their 2nd input CHECK_GE(op->inputs.size(), 2); - const ArrayDataType data_type = model->arrays[op->inputs[1]]->data_type; + const ArrayDataType data_type = model->GetArray(op->inputs[1]).data_type; SetDataTypeForAllOutputs(model, op, data_type); } else if (op->type == OperatorType::kCast) { // Data type of the Cast op is specified. CHECK_EQ(op->outputs.size(), 1); auto* cast_op = static_cast(op); - model->arrays[op->outputs[0]]->data_type = cast_op->dst_data_type; + model->GetArray(op->outputs[0]).data_type = cast_op->dst_data_type; } else if (op->type == OperatorType::kArgMax) { // Data type of the ArgMax op is specified. CHECK_EQ(op->outputs.size(), 1); auto* argmax_op = static_cast(op); - model->arrays[op->outputs[0]]->data_type = argmax_op->output_data_type; + model->GetArray(op->outputs[0]).data_type = argmax_op->output_data_type; + } else if (op->type == OperatorType::kRange) { + auto* range_op = static_cast(op); + // Output type of the Range op can be set via an attribute + ArrayDataType data_type; + if (range_op->dtype != ArrayDataType::kNone) { + // Use the type if specified + data_type = range_op->dtype; + } else { + // Otherwise use the first input + CHECK_GE(op->inputs.size(), 1); + data_type = model->GetArray(op->inputs[0]).data_type; + } + CHECK_EQ(op->outputs.size(), 1); + SetDataTypeForAllOutputs(model, op, data_type); } else if (op->type == OperatorType::kTensorFlowUnsupported) { auto* unsupported_op = static_cast(op); if (unsupported_op->output_data_types.size() != op->outputs.size()) { @@ -88,7 +103,7 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) { for (int i = 0; i < unsupported_op->output_data_types.size(); ++i) { auto output = op->outputs[i]; auto data_type = unsupported_op->output_data_types[i]; - model->arrays[output]->data_type = data_type; + model->GetArray(output).data_type = data_type; } } else if (op->type == OperatorType::kExpandDims) { // Yield on ExpandDim until it is converted to Reshape @@ -96,12 +111,12 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) { } else { // These operators produce outputs with the same type as their 1st input CHECK_GT(op->inputs.size(), 0); - const ArrayDataType data_type = model->arrays[op->inputs[0]]->data_type; + const ArrayDataType data_type = model->GetArray(op->inputs[0]).data_type; SetDataTypeForAllOutputs(model, op, data_type); } // Return true if any output data type changed, false if none changed. for (const auto& output : op->outputs) { - if (old_output_data_types[output] != model->arrays[output]->data_type) { + if (old_output_data_types[output] != model->GetArray(output).data_type) { return true; } } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc index 89884911685d1ef9a69a17e11ec69e53f43633e2..0cf0994b43bb048616bab1abe79db1aae2223d37 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc @@ -61,23 +61,42 @@ void ComputeConvSizes(const Shape& input_shape, int output_depth, int kwidth, output_shape->ReplaceDims({batch, output_height, output_width, output_depth}); } -void ComputeBinaryOperatorOutputSize(const Shape& input_shape1, - const Shape& input_shape2, +void ComputeBinaryOperatorOutputSize(const Shape& input_shape_x, + const Shape& input_shape_y, Array* output_array) { - const int size1 = RequiredBufferSizeForShape(input_shape1); - const int size2 = RequiredBufferSizeForShape(input_shape2); - if (size1 > size2) { - output_array->copy_shape(input_shape1); - } else if (size2 > size1) { - output_array->copy_shape(input_shape2); - } else { - CHECK_EQ(size1, size2); - const int dims1 = input_shape1.dimensions_count(); - const int dims2 = input_shape2.dimensions_count(); - if (dims1 >= dims2) { - output_array->copy_shape(input_shape1); + // This matches the code in BroadcastBinaryOpShapeFn from tensorflow. + // It zips together the two input shapes and pads with 1 to make them the + // same length. For each dimension we broadcast if either dimension is 1 and + // otherwise expect them to match. + int rank_x = input_shape_x.dimensions_count(); + int rank_y = input_shape_y.dimensions_count(); + int rank_out = std::max(rank_x, rank_y); + std::vector* dims_out = output_array->mutable_shape()->mutable_dims(); + dims_out->clear(); + dims_out->reserve(rank_out); + for (int i = 0; i < rank_out; ++i) { + int dim_x = i < (rank_out - rank_x) + ? 1 + : input_shape_x.dims(i - (rank_out - rank_x)); + bool dim_y_is_one = i < (rank_out - rank_y); + int dim_y = dim_y_is_one ? 1 : input_shape_y.dims(i - (rank_out - rank_y)); + if (dim_x == -1 || dim_y == -1) { + // One or both dimensions is unknown. + QCHECK(false) << "Shapes must be specified"; + } else if (dim_x == 1 || dim_y == 1) { + // Broadcast one dimension to the other that is 1. + if (dim_x == 1 && !dim_y_is_one) { + // Broadcast dim_y to dim_x (1). + dims_out->push_back(dim_y); + } else { + // Broadcast dim_x to dim_y (1). + DCHECK_EQ(dim_y, 1); + dims_out->push_back(dim_x); + } } else { - output_array->copy_shape(input_shape2); + // Expect the dimensions to match. + CHECK_EQ(dim_x, dim_y) << "Dimensions must match"; + dims_out->push_back(dim_x); } } CHECK(output_array->has_shape()); @@ -85,7 +104,7 @@ void ComputeBinaryOperatorOutputSize(const Shape& input_shape1, int GetOutputDepthFromWeights(const Model& model, const Operator& op) { const string& weights_name = op.inputs[1]; - const auto& weights_shape = model.arrays.at(weights_name)->shape(); + const auto& weights_shape = model.GetArray(weights_name).shape(); if (op.type == OperatorType::kConv || op.type == OperatorType::kFullyConnected) { return weights_shape.dims(0); @@ -98,7 +117,7 @@ int GetOutputDepthFromWeights(const Model& model, const Operator& op) { bool EnsureBiasVectorShape(Model* model, Operator* op) { const string& weights_name = op->inputs[1]; - const auto& weights_array = *model->arrays[weights_name]; + const auto& weights_array = model->GetArray(weights_name); // Yield until weights shape has been resolved. if (!weights_array.has_shape()) { return false; @@ -107,7 +126,7 @@ bool EnsureBiasVectorShape(Model* model, Operator* op) { if (op->inputs.size() < 3) { return false; } - auto& bias_array = *model->arrays[op->inputs[2]]; + auto& bias_array = model->GetArray(op->inputs[2]); if (bias_array.has_shape()) { return true; } @@ -126,7 +145,7 @@ void ProcessConvOperator(Model* model, ConvOperator* op) { return; } - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -134,7 +153,7 @@ void ProcessConvOperator(Model* model, ConvOperator* op) { const auto& input_shape = input_array.shape(); CHECK_EQ(input_shape.dimensions_count(), 4); - const auto& weights_array = *model->arrays[op->inputs[1]]; + const auto& weights_array = model->GetArray(op->inputs[1]); // Yield until weights dims have been resolved. if (!weights_array.has_shape()) { return; @@ -156,7 +175,7 @@ void ProcessConvOperator(Model* model, ConvOperator* op) { if (op->outputs.size() == 2) { const auto& output_shape = output_array.shape(); const int input_depth = weights_shape.dims(3); - auto& im2col_array = *model->arrays[op->outputs[1]]; + auto& im2col_array = model->GetArray(op->outputs[1]); im2col_array.copy_shape(Shape{output_shape.dims(0), output_shape.dims(1), output_shape.dims(2), input_depth * kheight * kwidth}); @@ -168,7 +187,7 @@ void ProcessDepthwiseConvOperator(Model* model, DepthwiseConvOperator* op) { return; } - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -176,7 +195,7 @@ void ProcessDepthwiseConvOperator(Model* model, DepthwiseConvOperator* op) { const auto& input_shape = input_array.shape(); CHECK_EQ(input_shape.dimensions_count(), 4); - const auto& weights_array = *model->arrays[op->inputs[1]]; + const auto& weights_array = model->GetArray(op->inputs[1]); // Yield until weights dims have been resolved. if (!weights_array.has_shape()) { return; @@ -209,7 +228,7 @@ void ProcessDepthwiseConvOperator(Model* model, DepthwiseConvOperator* op) { } void ProcessDepthToSpaceOperator(Model* model, DepthToSpaceOperator* op) { - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -232,7 +251,7 @@ void ProcessDepthToSpaceOperator(Model* model, DepthToSpaceOperator* op) { } void ProcessSpaceToDepthOperator(Model* model, SpaceToDepthOperator* op) { - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -258,7 +277,7 @@ void ProcessSpaceToDepthOperator(Model* model, SpaceToDepthOperator* op) { void ProcessFillOperator(Model* model, FillOperator* op) { CHECK_EQ(op->inputs.size(), 2); CHECK_EQ(op->outputs.size(), 1); - auto& output_array = *model->arrays[op->outputs[0]]; + auto& output_array = model->GetArray(op->outputs[0]); if (output_array.has_shape()) { // We have already run return; @@ -287,7 +306,7 @@ void ProcessFullyConnectedOperator(Model* model, FullyConnectedOperator* op) { return; } - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -295,7 +314,7 @@ void ProcessFullyConnectedOperator(Model* model, FullyConnectedOperator* op) { const auto& input_shape = input_array.shape(); CHECK_GE(input_shape.dimensions_count(), 1); - const auto& weights_array = *model->arrays[op->inputs[1]]; + const auto& weights_array = model->GetArray(op->inputs[1]); // Yield until weights dims have been resolved. if (!weights_array.has_shape()) { return; @@ -315,26 +334,31 @@ void ProcessFullyConnectedOperator(Model* model, FullyConnectedOperator* op) { void ProcessTensorFlowReshapeOperator(Model* model, TensorFlowReshapeOperator* op) { - auto& output_array = *model->arrays[op->outputs[0]]; - // Bail if we already have output dims + auto& output_array = model->GetArray(op->outputs[0]); if (output_array.has_shape()) { + // We have already run return; } - const auto& input_array = *model->arrays[op->inputs[0]]; - // Yield until input dims have been resolved. + const auto& input_array = model->GetArray(op->inputs[0]); if (!input_array.has_shape()) { + // Yield until input dims have been resolved. return; } const auto& input_shape = input_array.shape(); - const string& shape_name = op->inputs[1]; - auto& shape_array = model->GetArray(shape_name); - // Yield until the shape is resolved as a constant array + auto& shape_array = model->GetArray(op->inputs[1]); + if (!shape_array.has_shape()) { + // Yield until target_shape shape been resolved. + return; + } if (!shape_array.buffer) { + // Yield until the target_shape is constant return; } - CHECK(shape_array.data_type == ArrayDataType::kInt32); + CHECK(shape_array.data_type == ArrayDataType::kInt32) + << "Reshape dims must be int32"; + // shape_data is the raw array of ints describing the shape // in the TensorFlow node. We intentionally make a copy here, rather than // modify wildcards in-place below, because in some graphs, the same shape @@ -357,23 +381,29 @@ void ProcessTensorFlowReshapeOperator(Model* model, } const int input_flat_size = RequiredBufferSizeForShape(input_shape); if (has_wildcard) { + CHECK_GE(input_flat_size, product_non_wildcard_dims) + << "Array not large enough to fill the requested dimensions for " + "Reshape op with output \"" + << op->outputs[0] << "\". Are your input shapes correct?"; shape_data[wildcard_index] = input_flat_size / product_non_wildcard_dims; } auto& output_shape = *output_array.mutable_shape(); *output_shape.mutable_dims() = shape_data; - const int output_flat_size = RequiredBufferSizeForShape(output_shape); - CHECK_EQ(output_flat_size, input_flat_size); + CHECK_EQ(input_flat_size, RequiredBufferSizeForShape(output_shape)) + << "Input cannot be reshaped to requested dimensions for Reshape op with " + "output \"" + << op->outputs[0] << "\". Are your input shapes correct?"; } void ProcessSimpleOperator(Model* model, Operator* op) { - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; } const string& output_name = op->outputs[0]; - auto& output_array = *model->arrays[output_name]; + auto& output_array = model->GetArray(output_name); if (output_array.has_shape()) { return; } @@ -383,18 +413,40 @@ void ProcessSimpleOperator(Model* model, Operator* op) { void ProcessSimpleBinaryOperator(Model* model, Operator* op) { CHECK_EQ(op->inputs.size(), 2); - const auto& input0_array = *model->arrays[op->inputs[0]]; - const auto& input1_array = *model->arrays[op->inputs[1]]; + const auto& input0_array = model->GetArray(op->inputs[0]); + const auto& input1_array = model->GetArray(op->inputs[1]); // Yield until input dims have been resolved. if (!input0_array.has_shape() || !input1_array.has_shape()) { return; } const string& output_name = op->outputs[0]; - auto& output_array = *model->arrays[output_name]; + auto& output_array = model->GetArray(output_name); ComputeBinaryOperatorOutputSize(input0_array.shape(), input1_array.shape(), &output_array); } +void ProcessAddNOperator(Model* model, Operator* op) { + // Yield until all input dims have been resolved. + // + // TODO(myenik): Since AddN does not support broadcasting, maybe we could + // actually use this to improve shape propagation by propagating the shape of + // one input to all other inputs once it is resolved instead of just the + // output, since all inputs must be the same size and shape for a well-formed + // graph. + for (const auto& input : op->inputs) { + const auto& input_array = model->GetArray(input); + if (!input_array.has_shape()) { + return; + } + } + + // AddN does not support broadcasting, all inputs must be the same shape, so + // we just take the first input shape and apply it to the output. + const auto& input0_array = model->GetArray(op->inputs[0]); + auto& output_array = model->GetArray(op->outputs[0]); + output_array.copy_shape(input0_array.shape()); +} + bool KeepDims(const Operator& op) { switch (op.type) { case OperatorType::kTensorFlowMin: @@ -413,11 +465,11 @@ bool KeepDims(const Operator& op) { void ProcessTensorFlowReductionOperator(Model* model, Operator* op) { CHECK_LE(op->inputs.size(), 2); - auto& output_array = *model->arrays[op->outputs[0]]; + auto& output_array = model->GetArray(op->outputs[0]); if (output_array.has_shape()) { return; } - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); if (!input_array.has_shape()) { return; } @@ -425,7 +477,7 @@ void ProcessTensorFlowReductionOperator(Model* model, Operator* op) { const bool keep_dims = KeepDims(*op); if (op->inputs.size() == 2) { // There is a reduction_indices input. - const auto& reduction_array = *model->arrays[op->inputs[1]]; + const auto& reduction_array = model->GetArray(op->inputs[1]); if (!reduction_array.buffer) { return; } @@ -465,11 +517,11 @@ void ProcessSliceOperator(Model* model, SliceOperator* op) { if (op->begin.empty()) return; // Yield until input dims have been resolved. - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); if (!input_array.has_shape()) return; const Shape& input_shape = input_array.shape(); - auto& output_array = *model->arrays[op->outputs[0]]; + auto& output_array = model->GetArray(op->outputs[0]); if (output_array.has_shape()) return; CHECK_EQ(input_shape.dims().size(), op->size.size()); @@ -489,7 +541,7 @@ void ProcessSliceOperator(Model* model, SliceOperator* op) { void ProcessReorderAxesOperator(Model* model, ReorderAxesOperator* op) { const string& input_name = op->inputs[0]; - const auto& input_array = *model->arrays[input_name]; + const auto& input_array = model->GetArray(input_name); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -504,20 +556,23 @@ void ProcessReorderAxesOperator(Model* model, ReorderAxesOperator* op) { void ProcessConcatenationOperator(Model* model, ConcatenationOperator* op) { // Yield until input dims have been resolved. for (const auto& input_name : op->inputs) { - auto& input_array = *model->arrays[input_name]; + auto& input_array = model->GetArray(input_name); if (!input_array.has_shape()) { return; } } auto& output_array = model->GetArray(op->outputs[0]); // Use 0 input as basis for output dimensions. - const auto& first_input_array = *model->arrays[op->inputs[0]]; + const auto& first_input_array = model->GetArray(op->inputs[0]); output_array.copy_shape(first_input_array.shape()); + // Negative axis means the count starts at the back of the dims(). + int axis = op->axis; + if (axis < 0) axis += first_input_array.shape().dims().size(); // Determine the concat size, and enfore that all inputs have // the same dimensions count. int concat_size = 0; for (const auto& input_name : op->inputs) { - auto& input_array = *model->arrays[input_name]; + auto& input_array = model->GetArray(input_name); CHECK(input_array.has_shape()); if (input_array.shape().dimensions_count() == 0) { continue; @@ -525,45 +580,114 @@ void ProcessConcatenationOperator(Model* model, ConcatenationOperator* op) { CHECK_EQ(input_array.shape().dimensions_count(), output_array.shape().dimensions_count()); const std::vector& input_dims = input_array.shape().dims(); - CHECK_LT(op->axis, input_dims.size()); - concat_size += input_dims[op->axis]; + CHECK_LT(axis, input_dims.size()); + concat_size += input_dims[axis]; } // Write out the concat_size on the output array shape. auto& output_shape = *output_array.mutable_shape(); auto& output_dims = *output_shape.mutable_dims(); - CHECK_LT(op->axis, output_shape.dimensions_count()); - output_dims[op->axis] = concat_size; + CHECK_LT(axis, output_shape.dimensions_count()); + output_dims[axis] = concat_size; +} + +void ProcessRangeOperator(Model* model, RangeOperator* op) { + CHECK_EQ(op->inputs.size(), 3); + const auto& start_array = model->GetArray(op->inputs[0]); + if (!start_array.has_shape()) { + // Yield until input dims have been resolved. + return; + } + const auto& limit_array = model->GetArray(op->inputs[1]); + if (!limit_array.has_shape()) { + return; + } + const auto& delta_array = model->GetArray(op->inputs[2]); + if (!delta_array.has_shape()) { + return; + } + + if (!IsConstantParameterArray(*model, op->inputs[0])) { + // Yield until inputs are constant. + return; + } + if (!IsConstantParameterArray(*model, op->inputs[1])) { + return; + } + if (!IsConstantParameterArray(*model, op->inputs[2])) { + return; + } + + CHECK(start_array.data_type == ArrayDataType::kInt32) + << "Range op inputs must be int32."; + CHECK(limit_array.data_type == ArrayDataType::kInt32) + << "Range op inputs must be int32."; + CHECK(delta_array.data_type == ArrayDataType::kInt32) + << "Range op inputs must be int32."; + CHECK_EQ(RequiredBufferSizeForShape(start_array.shape()), 1) + << "Range op inputs must be scalar."; + CHECK_EQ(RequiredBufferSizeForShape(limit_array.shape()), 1) + << "Range op inputs must be scalar."; + CHECK_EQ(RequiredBufferSizeForShape(delta_array.shape()), 1) + << "Range op inputs must be scalar."; + int size = floor((limit_array.GetBuffer().data[0] - + start_array.GetBuffer().data[0]) / + delta_array.GetBuffer().data[0]); + + // Only set the output shape. Contents are set by ResolveConstantRange. + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + Shape* output_shape = output_array.mutable_shape(); + output_shape->ReplaceDims({size}); } void ProcessTensorFlowSplitOperator(Model* model, TensorFlowSplitOperator* op) { CHECK_EQ(op->inputs.size(), 2); const string& input_name = op->inputs[1]; - const auto& input_array = *model->arrays[input_name]; + const auto& input_array = model->GetArray(input_name); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; } const Shape& input_shape = input_array.shape(); - // This code is slightly suspect. The TensorFlow docs say that the axis - // selection defaults to 0, but we are splitting across the final axis. - const int input_dims_count = input_shape.dimensions_count(); - const int input_depth = input_shape.dims(input_dims_count - 1); - CHECK_EQ(input_depth % op->num_split, 0); - const int split_depth = input_depth / op->num_split; + // Yield until axis is constant. + if (!IsConstantParameterArray(*model, op->inputs[0])) { + return; + } + + const auto& axis_array = model->GetArray(op->inputs[0]); + + // Yield until axis dims have been resolved. + if (!axis_array.has_shape()) { + return; + } + + CHECK(axis_array.data_type == ArrayDataType::kInt32) + << "Axis array must be int32."; + CHECK_EQ(RequiredBufferSizeForShape(axis_array.shape()), 1) + << "Axis array must be scalar."; + + int axis = axis_array.GetBuffer().data[0]; + if (axis < 0) { + axis += input_shape.dimensions_count(); + } + + const int split_dim = input_shape.dims(axis); + CHECK_EQ(split_dim % op->num_split, 0); + const int split_depth = split_dim / op->num_split; Shape output_shape = input_shape; - (*output_shape.mutable_dims())[input_dims_count - 1] = split_depth; + (*output_shape.mutable_dims())[axis] = split_depth; CHECK_EQ(op->outputs.size(), op->num_split); for (const auto& output : op->outputs) { - model->arrays[output]->copy_shape(output_shape); + model->GetArray(output).copy_shape(output_shape); } } void ProcessAveragePoolOperator(Model* model, AveragePoolOperator* op) { const string& input_name = op->inputs[0]; - const auto& input_array = *model->arrays[input_name]; + const auto& input_array = model->GetArray(input_name); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -580,7 +704,7 @@ void ProcessAveragePoolOperator(Model* model, AveragePoolOperator* op) { void ProcessMaxPoolOperator(Model* model, MaxPoolOperator* op) { const string& input_name = op->inputs[0]; - const auto& input_array = *model->arrays[input_name]; + const auto& input_array = model->GetArray(input_name); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -597,7 +721,7 @@ void ProcessMaxPoolOperator(Model* model, MaxPoolOperator* op) { void ProcessL2PoolOperator(Model* model, L2PoolOperator* op) { const string& input_name = op->inputs[0]; - const auto& input_array = *model->arrays[input_name]; + const auto& input_array = model->GetArray(input_name); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -618,33 +742,35 @@ void ProcessResizeBilinearOperator(Model* model, ResizeBilinearOperator* op) { CHECK_EQ(op->inputs.size(), 2); CHECK_EQ(op->outputs.size(), 1); - if (!model->arrays[op->inputs[0]]->has_shape() || - !model->arrays[op->inputs[1]]->has_shape()) { + if (!model->GetArray(op->inputs[0]).has_shape() || + !model->GetArray(op->inputs[1]).has_shape()) { return; } - const auto& input_data_shape = model->arrays[op->inputs[0]]->shape(); + const auto& input_data_shape = model->GetArray(op->inputs[0]).shape(); const string& output_size_name = op->inputs[1]; - const auto& output_size_array = *model->arrays[output_size_name]; + const auto& output_size_array = model->GetArray(output_size_name); CHECK(output_size_array.data_type == ArrayDataType::kInt32); CHECK(output_size_array.has_shape()); const auto& output_size_shape = output_size_array.shape(); CHECK_EQ(output_size_shape.dimensions_count(), 1); CHECK_EQ(output_size_shape.dims(0), 2); + if (!output_size_array.buffer) { + return; + } std::vector output_shape = output_size_array.GetBuffer().data; - model->arrays[op->outputs[0]]->copy_shape( - Shape({input_data_shape.dims(0), output_shape[0], output_shape[1], - input_data_shape.dims(3)})); + model->GetArray(op->outputs[0]) + .copy_shape(Shape({input_data_shape.dims(0), output_shape[0], + output_shape[1], input_data_shape.dims(3)})); } void ProcessLstmCellOperator(Model* model, LstmCellOperator* op) { - // I/O arrays should be allocated on creation of op. - QCHECK_EQ(op->inputs.size(), LstmCellOperator::NUM_INPUTS); - QCHECK_EQ(op->outputs.size(), LstmCellOperator::NUM_OUTPUTS); + // Only required for compact LstmCell with default NUM_INPUTS of inputs. + if (op->inputs.size() != LstmCellOperator::NUM_INPUTS) return; const auto& input_array = - *model->arrays[op->inputs[LstmCellOperator::DATA_INPUT]]; + model->GetArray(op->inputs[LstmCellOperator::DATA_INPUT]); // Yield until all input dims have been resolved. if (!input_array.has_shape()) { return; @@ -653,7 +779,7 @@ void ProcessLstmCellOperator(Model* model, LstmCellOperator* op) { CHECK_GE(input_shape.dimensions_count(), 2); const auto& prev_activ_array = - *model->arrays[op->inputs[LstmCellOperator::PREV_ACTIV_INPUT]]; + model->GetArray(op->inputs[LstmCellOperator::PREV_ACTIV_INPUT]); // Yield until all input dims have been resolved. if (!prev_activ_array.has_shape()) { return; @@ -662,7 +788,7 @@ void ProcessLstmCellOperator(Model* model, LstmCellOperator* op) { CHECK_GE(prev_activ_shape.dimensions_count(), 2); const auto& weights_array = - *model->arrays[op->inputs[LstmCellOperator::WEIGHTS_INPUT]]; + model->GetArray(op->inputs[LstmCellOperator::WEIGHTS_INPUT]); // Yield until weights dims have been resolved. if (!weights_array.has_shape()) { return; @@ -671,7 +797,7 @@ void ProcessLstmCellOperator(Model* model, LstmCellOperator* op) { CHECK_EQ(weights_shape.dimensions_count(), 2); const auto& bias_array = - *model->arrays[op->inputs[LstmCellOperator::BIASES_INPUT]]; + model->GetArray(op->inputs[LstmCellOperator::BIASES_INPUT]); // Yield until bias dims have been resolved. if (!bias_array.has_shape()) { return; @@ -680,7 +806,7 @@ void ProcessLstmCellOperator(Model* model, LstmCellOperator* op) { CHECK_GE(bias_shape.dimensions_count(), 1); const auto& prev_state_array = - *model->arrays[op->inputs[LstmCellOperator::PREV_STATE_INPUT]]; + model->GetArray(op->inputs[LstmCellOperator::PREV_STATE_INPUT]); // Yield until all input dims have been resolved. if (!prev_state_array.has_shape()) { return; @@ -720,21 +846,21 @@ void ProcessLstmCellOperator(Model* model, LstmCellOperator* op) { } void ProcessSpaceToBatchNDOperator(Model* model, SpaceToBatchNDOperator* op) { - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; } const auto& input_shape = input_array.shape(); + // This method only handles input dimensions of 4. if (input_shape.dimensions_count() != 4) { - // This method only handles input dimensions of 4 return; } const auto input_height = input_shape.dims(1); const auto input_width = input_shape.dims(2); - const auto& block_shape_array = *model->arrays[op->inputs[1]]; - const auto& paddings_array = *model->arrays[op->inputs[2]]; + const auto& block_shape_array = model->GetArray(op->inputs[1]); + const auto& paddings_array = model->GetArray(op->inputs[2]); const auto& block_shape_array_shape = block_shape_array.shape(); const auto& paddings_array_shape = paddings_array.shape(); QCHECK_EQ(block_shape_array_shape.dimensions_count(), 1); @@ -766,13 +892,13 @@ void ProcessSpaceToBatchNDOperator(Model* model, SpaceToBatchNDOperator* op) { int output_height = height_with_paddings / block_height; int output_width = width_with_paddings / block_width; - model->arrays[op->outputs[0]]->copy_shape( - Shape({input_shape.dims(0) * block_height * block_width, output_height, - output_width, input_shape.dims(3)})); + model->GetArray(op->outputs[0]) + .copy_shape(Shape({input_shape.dims(0) * block_height * block_width, + output_height, output_width, input_shape.dims(3)})); } void ProcessBatchToSpaceNDOperator(Model* model, BatchToSpaceNDOperator* op) { - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -782,8 +908,8 @@ void ProcessBatchToSpaceNDOperator(Model* model, BatchToSpaceNDOperator* op) { const auto input_height = input_shape.dims(1); const auto input_width = input_shape.dims(2); - const auto& block_shape_array = *model->arrays[op->inputs[1]]; - const auto& crops_array = *model->arrays[op->inputs[2]]; + const auto& block_shape_array = model->GetArray(op->inputs[1]); + const auto& crops_array = model->GetArray(op->inputs[2]); const auto& block_shape_array_shape = block_shape_array.shape(); const auto& crops_array_shape = crops_array.shape(); QCHECK_EQ(block_shape_array_shape.dimensions_count(), 1); @@ -818,15 +944,15 @@ void ProcessBatchToSpaceNDOperator(Model* model, BatchToSpaceNDOperator* op) { int output_height = input_height * block_height; int output_width = input_width * block_width; - model->arrays[op->outputs[0]]->copy_shape( - Shape({input_shape.dims(0) / (block_height * block_width), output_height, - output_width, input_shape.dims(3)})); + model->GetArray(op->outputs[0]) + .copy_shape(Shape({input_shape.dims(0) / (block_height * block_width), + output_height, output_width, input_shape.dims(3)})); } void ProcessGatherOperator(Model* model, GatherOperator* op) { - const auto& input_array = *model->arrays[op->inputs[0]]; - const auto& indices_array = *model->arrays[op->inputs[1]]; - auto& output_array = *model->arrays[op->outputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); + const auto& indices_array = model->GetArray(op->inputs[1]); + auto& output_array = model->GetArray(op->outputs[0]); // Bail if we already know the output shape. if (output_array.has_shape()) { @@ -856,11 +982,48 @@ void ProcessGatherOperator(Model* model, GatherOperator* op) { } } +void ProcessTopkV2Operator(Model* model, TopKV2Operator* op) { + const auto& input_values = model->GetArray(op->inputs[0]); + const auto& input_k = model->GetArray(op->inputs[1]); + auto& output_indexes = model->GetArray(op->outputs[0]); + auto& output_values = model->GetArray(op->outputs[1]); + + // Bail if we already know the output shape. + if (output_indexes.has_shape()) { + QCHECK(output_values.has_shape()); + return; + } + + // Yield until input dims have been resolved. + if (!input_values.has_shape()) { + return; + } + + const auto& input_values_shape = input_values.shape(); + auto output_indexes_dims = output_indexes.mutable_shape()->mutable_dims(); + auto output_values_dims = output_values.mutable_shape()->mutable_dims(); + for (int dim = 0; dim < input_values_shape.dimensions_count() - 1; dim++) { + output_indexes_dims->push_back(input_values_shape.dims(dim)); + output_values_dims->push_back(input_values_shape.dims(dim)); + } + // If the value is initialized, we can specify the last dimension, otherwise + // unknown. + if (input_k.buffer) { + const int32_t k_value = input_k.GetBuffer().data[0]; + output_indexes_dims->push_back(k_value); + output_values_dims->push_back(k_value); + + } else { + output_indexes_dims->push_back(0); + output_values_dims->push_back(0); + } +} + void ProcessPadOperator(Model* model, PadOperator* op) { CHECK_EQ(op->inputs.size(), 2); CHECK_EQ(op->outputs.size(), 1); - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) return; @@ -868,7 +1031,7 @@ void ProcessPadOperator(Model* model, PadOperator* op) { if (op->left_padding.empty()) return; CHECK_EQ(op->left_padding.size(), op->right_padding.size()); - auto& output_array = *model->arrays[op->outputs[0]]; + auto& output_array = model->GetArray(op->outputs[0]); if (output_array.has_shape()) return; Shape output_shape = input_array.shape(); @@ -882,47 +1045,179 @@ void ProcessPadOperator(Model* model, PadOperator* op) { output_array.copy_shape(output_shape); } -void ProcessStridedSliceOperator(Model* model, StridedSliceOperator* op) { - CHECK_EQ(op->inputs.size(), 4); +void ProcessRankOperator(Model* model, RankOperator* op) { + CHECK_GE(op->inputs.size(), 1); CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.has_shape()) { + // Shape already propagated + return; + } - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); + if (!input_array.has_shape()) { + // Yield until input dims have been resolved. + return; + } - // Yield until input dims have been resolved. - if (!input_array.has_shape()) return; + // Only set the output shape. Array contents are set by + // ResolveConstantShapeOrRank. + Shape* output_shape = output_array.mutable_shape(); + output_shape->ReplaceDims({}); +} - if (op->start_indices.empty()) return; - CHECK_EQ(op->start_indices.size(), op->stop_indices.size()); - CHECK_EQ(op->start_indices.size(), op->strides.size()); +void ProcessShapeOperator(Model* model, TensorFlowShapeOperator* op) { + CHECK_GE(op->inputs.size(), 1); + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.has_shape()) { + // Shape already propagated + return; + } - auto& output_array = *model->arrays[op->outputs[0]]; - if (output_array.has_shape()) return; + const auto& input_array = model->GetArray(op->inputs[0]); + if (!input_array.has_shape()) { + // Yield until input dims have been resolved. + return; + } - Shape output_shape = input_array.shape(); - std::vector& dims = *output_shape.mutable_dims(); - CHECK_EQ(op->start_indices.size(), dims.size()); + // Only set the output shape. Array contents are set by + // ResolveConstantShapeOrRank. + Shape* output_shape = output_array.mutable_shape(); + output_shape->ReplaceDims({input_array.shape().dimensions_count()}); +} - for (int i = 0; i < op->start_indices.size(); ++i) { - const int mask = 1 << i; - const int start = (op->begin_mask & mask) ? 0 : op->start_indices[i]; - const int stop = (op->end_mask & mask) ? input_array.shape().dims()[i] - : op->stop_indices[i]; - dims[i] = (stop - start) / op->strides[i]; +void ProcessStackOperator(Model* model, StackOperator* op) { + CHECK_GE(op->inputs.size(), 1); + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.has_shape()) { + // Shape already propagated + return; } - output_array.copy_shape(output_shape); + std::unique_ptr stacked_shape; + for (const auto& input : op->inputs) { + const auto& input_array = model->GetArray(input); + if (!input_array.has_shape()) { + // Yield until all input dims have been resolved. + return; + } + + Shape shape = input_array.shape(); + if (shape.dimensions_count() == 0) { + // Convert 0D scalars to 1D scalars of shape {1}. + shape.mutable_dims()->push_back(1); + } + if (!stacked_shape) { + stacked_shape.reset(new Shape(shape)); + } else { + CHECK(*stacked_shape == shape) << "All input arrays to Stack operators " + "must have the same shape. Input \"" + << input << "\" is different."; + } + } + + int axis = op->axis; + if (axis < 0) { + // Handle negative axis + axis += stacked_shape->dims().size() + 1; + } + stacked_shape->mutable_dims()->insert( + stacked_shape->mutable_dims()->begin() + axis, op->inputs.size()); + output_array.copy_shape(*stacked_shape); +} + +void ProcessStridedSliceOperator(Model* model, StridedSliceOperator* op) { + CHECK_GE(op->inputs.size(), 1); + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.has_shape()) { + // Shape already propagated + return; + } + + if (op->start_indices.empty() || op->stop_indices.empty() || + op->strides.empty()) { + // ResolveStridedSliceAttributes has not run yet. + return; + } + + const auto& input_array = model->GetArray(op->inputs[0]); + if (!input_array.has_shape()) { + // Yield until input dims have been resolved. + return; + } + + if (op->ellipsis_mask != 0) { + // Something like LOG_FIRST_N(WARNING, 10) would be prefferable to reduce + // log noise. However, the TensorFlow logging library does not appear to + // support this. + LOG(WARNING) << "Skipping StridedSlice op with output \"" << op->outputs[0] + << "\". ellipsis_mask is not supported (mask=" + << op->ellipsis_mask << ")"; + return; + } + if (op->new_axis_mask != 0) { + LOG(WARNING) << "Skipping StridedSlice op with output \"" << op->outputs[0] + << "\". new_axis_mask is not supported (mask=" + << op->new_axis_mask << ")"; + return; + } + + int dim_count = input_array.shape().dimensions_count(); + CHECK(op->start_indices.size() == dim_count) + << ": Incorrect number of start indices supplied to StridedSlice op with " + "output \"" + << op->outputs[0] << "\". Op requires " << dim_count << " start indices"; + CHECK(op->stop_indices.size() == dim_count) + << ": Incorrect number of stop indices supplied to StridedSlice op with " + "output \"" + << op->outputs[0] << "\". Op requires " << dim_count << " stop indices"; + CHECK(op->strides.size() == dim_count) + << ": Incorrect number of strides supplied to StridedSlice op with " + " output \"" + << op->outputs[0] << "\". Op requires " << dim_count << " strides"; + + // Create output shape + std::vector* dims = output_array.mutable_shape()->mutable_dims(); + + // Compute output shape + for (int i = 0; i < dim_count; ++i) { + const int mask = 1 << i; + int start = (op->begin_mask & mask) ? 0 : op->start_indices[i]; + if (start < 0) { + // handle negative indices + start += input_array.shape().dims(i); + } + int stop = (op->end_mask & mask) ? input_array.shape().dims(i) + : op->stop_indices[i]; + if (stop < 0) { + // handle negative indices + stop += input_array.shape().dims(i); + } + + int dim_size = ceil((stop - start) / static_cast(op->strides[i])); + dim_size = dim_size < 0 ? 0 : dim_size; + if (op->shrink_axis_mask & mask) { + CHECK_EQ(dim_size, 1) << "Output size for an axis must compute to 1 when " + "shrinking that axis"; + } else { + dims->push_back(dim_size); + } + } } void ProcessSqueezeOperator(Model* model, SqueezeOperator* op) { CHECK_EQ(op->inputs.size(), 1); CHECK_EQ(op->outputs.size(), 1); - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) return; - auto& output_array = *model->arrays[op->outputs[0]]; + auto& output_array = model->GetArray(op->outputs[0]); if (output_array.has_shape()) return; const std::vector& input_dims = input_array.shape().dims(); @@ -941,18 +1236,18 @@ void ProcessSqueezeOperator(Model* model, SqueezeOperator* op) { void ProcessSvdfOperator(Model* model, SvdfOperator* op) { CHECK(op->inputs.size() == 3 || op->inputs.size() == 4); - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); if (!input_array.has_shape()) return; - auto& weights_feature_array = *model->arrays[op->inputs[1]]; + auto& weights_feature_array = model->GetArray(op->inputs[1]); if (!weights_feature_array.has_shape()) return; - const auto& weights_time_array = *model->arrays[op->inputs[2]]; + const auto& weights_time_array = model->GetArray(op->inputs[2]); if (!weights_time_array.has_shape()) return; const bool has_bias = (op->inputs.size() == 4); if (has_bias) { - const auto& bias_array = *model->arrays[op->inputs[3]]; + const auto& bias_array = model->GetArray(op->inputs[3]); if (!bias_array.has_shape()) return; } @@ -968,9 +1263,49 @@ void ProcessSvdfOperator(Model* model, SvdfOperator* op) { output_array.mutable_shape()->ReplaceDims({batch_size, num_units}); } +void ProcessTransposeOperator(Model* model, TransposeOperator* op) { + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.has_shape()) { + // We have already run + return; + } + + const auto& input_array = model->GetArray(op->inputs[0]); + if (!input_array.has_shape()) { + // Yield until input dims have been resolved. + return; + } + const auto& input_shape = input_array.shape(); + + auto& perm_array = model->GetArray(op->inputs[1]); + if (!perm_array.has_shape()) { + // Yield until permutation shape been resolved. + return; + } + if (!perm_array.buffer) { + // Yield until the permutation is constant + return; + } + CHECK(perm_array.data_type == ArrayDataType::kInt32) + << "Transpose permutation input must be int32"; + + std::vector const& perm = + perm_array.GetBuffer().data; + CHECK_EQ(perm.size(), input_shape.dimensions_count()) + << "Transpose permutation input " << op->inputs[0] + << " must be same length as input dimensions"; + std::vector* output_dims = output_array.mutable_shape()->mutable_dims(); + for (int i = 0; i < perm.size(); i++) { + int axis = perm[i]; + CHECK_GE(axis, 0); + CHECK_LT(axis, input_shape.dimensions_count()); + output_dims->push_back(input_shape.dims(axis)); + } +} + void ProcessArgMaxOperator(Model* model, ArgMaxOperator* op) { CHECK_EQ(op->inputs.size(), 2); - const auto& input_array = *model->arrays[op->inputs[0]]; + const auto& input_array = model->GetArray(op->inputs[0]); // Yield until input dims have been resolved. if (!input_array.has_shape()) { return; @@ -988,7 +1323,7 @@ void ProcessArgMaxOperator(Model* model, ArgMaxOperator* op) { } output_dims.push_back(1); const string& output_name = op->outputs[0]; - auto& output_array = *model->arrays[output_name]; + auto& output_array = model->GetArray(output_name); if (output_array.has_shape()) { return; } @@ -1002,8 +1337,8 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { auto* op = it->get(); std::unordered_map> old_output_dims; for (const auto& output : op->outputs) { - if (model->arrays[output]->has_shape()) { - old_output_dims[output] = model->arrays[output]->shape().dims(); + if (model->GetArray(output).has_shape()) { + old_output_dims[output] = model->GetArray(output).shape().dims(); } } @@ -1015,6 +1350,7 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { case OperatorType::kRelu1: case OperatorType::kRelu6: case OperatorType::kSoftmax: + case OperatorType::kLogSoftmax: case OperatorType::kLogistic: case OperatorType::kTanh: case OperatorType::kLocalResponseNormalization: @@ -1028,12 +1364,15 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { case OperatorType::kTensorFlowAssert: case OperatorType::kCast: case OperatorType::kFloor: + case OperatorType::kExp: ProcessSimpleOperator(model, op); break; case OperatorType::kGather: ProcessGatherOperator(model, static_cast(op)); break; - + case OperatorType::kTopK_V2: + ProcessTopkV2Operator(model, static_cast(op)); + break; case OperatorType::kAdd: case OperatorType::kSub: case OperatorType::kMul: @@ -1048,6 +1387,9 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { case OperatorType::kTensorFlowGreaterEqual: ProcessSimpleBinaryOperator(model, op); break; + case OperatorType::kAddN: + ProcessAddNOperator(model, op); + break; case OperatorType::kConv: ProcessConvOperator(model, static_cast(op)); break; @@ -1133,13 +1475,19 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { // or else at the moment we will abort. break; case OperatorType::kExpandDims: + // Yield until ExpandDims is converted to Reshape + break; case OperatorType::kRange: + ProcessRangeOperator(model, static_cast(op)); + break; case OperatorType::kRank: + ProcessRankOperator(model, static_cast(op)); + break; case OperatorType::kTensorFlowShape: + ProcessShapeOperator(model, static_cast(op)); + break; case OperatorType::kStack: - case OperatorType::kTranspose: - // Unimplemented. Hopefully another graph transformation will drop it or - // rewrite it. + ProcessStackOperator(model, static_cast(op)); break; case OperatorType::kReorderAxes: ProcessReorderAxesOperator(model, static_cast(op)); @@ -1155,6 +1503,7 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { case OperatorType::kLstmCell: ProcessLstmCellOperator(model, static_cast(op)); break; + case OperatorType::kBatchMatMul: case OperatorType::kTensorFlowMatMul: // MatMul operators are converted to FullyConnected, after which their // shapes are propagated. @@ -1182,6 +1531,9 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { case OperatorType::kSvdf: ProcessSvdfOperator(model, static_cast(op)); break; + case OperatorType::kTranspose: + ProcessTransposeOperator(model, static_cast(op)); + break; default: // Unimplemented, another graph transformation should drop it. LOG(FATAL) << "Unhandled operator type " << OperatorTypeName(op->type); @@ -1190,10 +1542,10 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) { // Return true if any output dim changed, false if none changed. // Assumption: no transformation clears an output shape, they only add shapes. for (const auto& output : op->outputs) { - if (model->arrays[output]->has_shape() && - (old_output_dims[output] != model->arrays[output]->shape().dims())) { + if (model->GetArray(output).has_shape() && + (old_output_dims[output] != model->GetArray(output).shape().dims())) { AddMessageF("Set shape of %s to [%s]", output, - absl::StrJoin(model->arrays[output]->shape().dims(), ",")); + absl::StrJoin(model->GetArray(output).shape().dims(), ",")); return true; } } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc index 56082b965a7cbd9d61cca2e26f7d76764c0e54aa..d7f804ee432598cafe6b6c05d03219aa7d2783fa 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc @@ -41,11 +41,15 @@ bool SupportsQuantization(const Operator& op) { type == OperatorType::kConcatenation || type == OperatorType::kL2Normalization || type == OperatorType::kAdd || type == OperatorType::kAveragePool || type == OperatorType::kMaxPool || + type == OperatorType::kTensorFlowMinimum || + type == OperatorType::kTensorFlowMaximum || type == OperatorType::kLogistic || type == OperatorType::kSoftmax || + type == OperatorType::kTensorFlowSplit || type == OperatorType::kSub || type == OperatorType::kSqueeze || type == OperatorType::kPad || type == OperatorType::kTensorFlowReshape || - type == OperatorType::kMul || type == OperatorType::kSpaceToDepth || - type == OperatorType::kDepthToSpace; + type == OperatorType::kTanh || type == OperatorType::kMul || + type == OperatorType::kSpaceToDepth || + type == OperatorType::kDepthToSpace || type == OperatorType::kLstmCell; } template @@ -100,6 +104,9 @@ void QuantizeArray(GraphTransformation* transformation, Model* model, case ArrayDataType::kUint8: return QuantizeArray(transformation, model, name, quantization_params); + case ArrayDataType::kInt16: + return QuantizeArray(transformation, model, name, + quantization_params); case ArrayDataType::kInt32: return QuantizeArray(transformation, model, name, quantization_params); @@ -168,36 +175,62 @@ bool ChooseQuantizationForOperatorInput( if (array.data_type != ArrayDataType::kFloat) { return false; } + + // Quantization of bias vectors + bool is_bias_vector = false; + int activations_input_index; + int weights_input_index; if (op.type == OperatorType::kConv || op.type == OperatorType::kDepthwiseConv || op.type == OperatorType::kFullyConnected) { if (input_index == 2) { - // Quantization of bias vector. - // We need both of the mandatory inputs (input activations and weights) to - // have - // been already quantized. - const auto& input_activations = model->GetArray(op.inputs[0]); - const auto& input_weights = model->GetArray(op.inputs[1]); - if (!input_activations.quantization_params || - !input_weights.quantization_params) { - return false; - } - const auto input_activations_scale = - input_activations.quantization_params->scale; - const auto input_weights_scale = input_weights.quantization_params->scale; - quantization_params->scale = - input_activations_scale * input_weights_scale; - quantization_params->zero_point = 0; - *quantized_data_type = ArrayDataType::kInt32; - transformation->AddMessageF( - "Input array %s is a bias vector. Choosing quantization params " - "accordingly.", - input); - return true; + is_bias_vector = true; + activations_input_index = 0; + weights_input_index = 1; } } + if (op.type == OperatorType::kLstmCell) { + if (input_index == LstmCellOperator::BIASES_INPUT) { + is_bias_vector = true; + activations_input_index = LstmCellOperator::DATA_INPUT; + weights_input_index = LstmCellOperator::WEIGHTS_INPUT; + } + } + if (is_bias_vector) { + // Quantization of bias vector. + // We need both of the mandatory inputs (input activations and weights) to + // have been already quantized. + const auto& input_activations = + model->GetArray(op.inputs[activations_input_index]); + const auto& input_weights = model->GetArray(op.inputs[weights_input_index]); + if (!input_activations.quantization_params || + !input_weights.quantization_params) { + return false; + } + const auto input_activations_scale = + input_activations.quantization_params->scale; + const auto input_weights_scale = input_weights.quantization_params->scale; + quantization_params->scale = input_activations_scale * input_weights_scale; + quantization_params->zero_point = 0; + *quantized_data_type = ArrayDataType::kInt32; + transformation->AddMessageF( + "Input array %s is a bias vector. Choosing quantization params " + "accordingly.", + input); + return true; + } const MinMax& minmax = GetOrComputeMinMax(model, input); + + if (op.type == OperatorType::kLstmCell) { + if (input_index == LstmCellOperator::PREV_STATE_INPUT) { + GetQuantizationParamsFromMinMax( + model->flags, minmax, quantization_params); + *quantized_data_type = ArrayDataType::kInt16; + return true; + } + } + GetQuantizationParamsFromMinMax(model->flags, minmax, quantization_params); transformation->AddMessageF( @@ -258,6 +291,17 @@ bool ChooseHardcodedQuantizationForOperatorOutput( *quantization_params)); return true; } + if (op.type == OperatorType::kTanh) { + // Tanh has the range: [-1, 1]. + *quantized_data_type = ArrayDataType::kUint8; + quantization_params->zero_point = 128; + quantization_params->scale = 1. / 128.; + // 0 should be exactly representable, as values will typically be centered + // around 0, with many values near 0. + CHECK( + IsExactlyRepresentable(0., *quantized_data_type, *quantization_params)); + return true; + } return false; } @@ -295,6 +339,15 @@ bool ChooseQuantizationForOperatorOutput( return true; } const MinMax& minmax = GetOrComputeMinMax(model, output); + if (op.type == OperatorType::kLstmCell) { + if (output_index == LstmCellOperator::STATE_OUTPUT || + output_index == LstmCellOperator::ACTIV_TEMP) { + GetQuantizationParamsFromMinMax( + model->flags, minmax, quantization_params); + *quantized_data_type = ArrayDataType::kInt16; + return true; + } + } GetQuantizationParamsFromMinMax(model->flags, minmax, quantization_params); *quantized_data_type = ArrayDataType::kUint8; @@ -390,30 +443,52 @@ bool Quantize::Run(Model* model, std::size_t op_index) { if (ChooseQuantizationForOperatorInput(this, model, op, input_index, &quantized_data_type, &quantization_params)) { - changed = true; const auto& input = op.inputs[input_index]; if (IsConstantParameterArray(*model, input)) { QuantizeArray(this, model, input, quantized_data_type, quantization_params); + changed = true; } else { auto dequantize_it = FindOpWithOutput(*model, input); - CHECK(dequantize_it != model->operators.end()); - auto* dequantize_op = dequantize_it->get(); - CHECK(dequantize_op->type == OperatorType::kDequantize); - op.inputs[input_index] = dequantize_op->inputs[0]; - // Check if the output of that Dequantize op was not used by any - // other operator. We will then erase that Dequantize op. - if (!CountOpsWithInput(*model, dequantize_op->outputs[0])) { - // If any of the model's output_arrays was pointing to the - // Dequantize op's output, let it point to the Dequantize op's - // input instead. - for (int i = 0; i < model->flags.output_arrays_size(); i++) { - if (model->flags.output_arrays(i) == dequantize_op->outputs[0]) { - model->flags.set_output_arrays(i, dequantize_op->inputs[0]); + if (dequantize_it != model->operators.end()) { + auto* dequantize_op = dequantize_it->get(); + CHECK(dequantize_op->type == OperatorType::kDequantize); + op.inputs[input_index] = dequantize_op->inputs[0]; + // Check if the output of that Dequantize op was not used by any + // other operator. We will then erase that Dequantize op. + if (!CountOpsWithInput(*model, dequantize_op->outputs[0])) { + // If any of the model's output_arrays was pointing to the + // Dequantize op's output, let it point to the Dequantize op's + // input instead. + for (int i = 0; i < model->flags.output_arrays_size(); i++) { + if (model->flags.output_arrays(i) == dequantize_op->outputs[0]) { + model->flags.set_output_arrays(i, dequantize_op->inputs[0]); + } + } + model->EraseArray(dequantize_op->outputs[0]); + model->operators.erase(dequantize_it); + } + changed = true; + } else { + // This input array is not produced by a Dequantize op. + // We have encountered this situation in RNN graphs, whose cyclic + // nature defeats the basic assumption underlying the quantization + // algorithm implemented here. For now, when we have seen this + // happening, the array in question was a RNN state array itself, + // so let us just implement this case here, and guard that assumption + // with a CHECK. A more general fix would involve revisiting the + // design of this whole Quantization transformation. + bool is_rnn_state_array = false; + for (const auto& rnn_state : model->flags.rnn_states()) { + if (rnn_state.state_array() == input) { + is_rnn_state_array = true; + break; } } - model->arrays.erase(dequantize_op->outputs[0]); - model->operators.erase(dequantize_it); + CHECK(is_rnn_state_array); + QuantizeArray(this, model, input, quantized_data_type, + quantization_params); + changed = true; } } } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/read_fake_quant_min_max.cc b/tensorflow/contrib/lite/toco/graph_transformations/read_fake_quant_min_max.cc index 371ced388a8111c18ada32cf31a784809479291d..11f8d4b6eea836c5fe4efcbd5136e6183a59dc62 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/read_fake_quant_min_max.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/read_fake_quant_min_max.cc @@ -80,7 +80,7 @@ bool ReadFakeQuantMinMax::Run(Model* model, std::size_t op_index) { // else. for (int i = 1; i <= 2; i++) { if (CountOpsWithInput(*model, fakequant_op->inputs[i]) == 1) { - model->arrays.erase(fakequant_op->inputs[i]); + model->EraseArray(fakequant_op->inputs[i]); } } fakequant_op->inputs.resize(1); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_final_dequantize_op.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_final_dequantize_op.cc index 3992e7d1ef71edd4040e626d5848d2fd9bb3dab6..c3b2709a33d54213661ba96394b01aa2cfd1a278 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_final_dequantize_op.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_final_dequantize_op.cc @@ -51,7 +51,7 @@ bool RemoveFinalDequantizeOp::Run(Model* model, std::size_t op_index) { // Remove the node and its output array. AddMessageF("Removed final %s", LogName(*dequantize_op)); - model->arrays.erase(output); + model->EraseArray(output); model->operators.erase(dequantize_it); return true; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_binary.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_binary.cc index 6add443f2d62fd06e8c0d17e03bc78c5d74732a1..95a50c61794092b02e518d1f08d8cf4a668353a8 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_binary.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_binary.cc @@ -81,7 +81,7 @@ bool RemoveTrivialBinaryOperator::Run(Model* model, std::size_t op_index) { // Now check if the constant operand makes this binary // operator trivial. const auto& constant_input_array = - *model->arrays[binary_op->inputs[index_of_constant_input]]; + model->GetArray(binary_op->inputs[index_of_constant_input]); // For now, we only handle floats here. if (constant_input_array.data_type != ArrayDataType::kFloat) { return false; @@ -89,14 +89,14 @@ bool RemoveTrivialBinaryOperator::Run(Model* model, std::size_t op_index) { const auto& constant_input_float_data = constant_input_array.GetBuffer().data; bool is_trivial = false; - if (binary_op->type != OperatorType::kAdd) { + if (binary_op->type == OperatorType::kAdd) { is_trivial = AreAllBufferElementsEqualTo(constant_input_float_data, 0.f); - } else if (binary_op->type != OperatorType::kSub) { + } else if (binary_op->type == OperatorType::kSub) { is_trivial = index_of_constant_input == 1 && AreAllBufferElementsEqualTo(constant_input_float_data, 0.f); - } else if (binary_op->type != OperatorType::kMul) { + } else if (binary_op->type == OperatorType::kMul) { is_trivial = AreAllBufferElementsEqualTo(constant_input_float_data, 1.f); - } else if (binary_op->type != OperatorType::kDiv) { + } else if (binary_op->type == OperatorType::kDiv) { is_trivial = index_of_constant_input == 1 && AreAllBufferElementsEqualTo(constant_input_float_data, 1.f); } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_concatenation_input.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_concatenation_input.cc index 23a5c857e8b19f7edbb48f2c004d03e21008833d..936854a04fd600ea23ab5dda50370f85a311c28c 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_concatenation_input.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_concatenation_input.cc @@ -59,7 +59,7 @@ bool RemoveTrivialConcatenationInput::Run(Model* model, std::size_t op_index) { for (const string& input : trivial_inputs) { if (IsDiscardableArray(*model, input) && CountOpsWithInput(*model, input) == 1) { - model->arrays.erase(input); + model->EraseArray(input); } } concat_op->inputs = nontrivial_inputs; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc index 047389f69a1d8987b52b07478b0d3eaf46f433ba..587f171bbf823408a45083c36d52f1d38c300123 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc @@ -124,7 +124,7 @@ bool RemoveTrivialPassthroughOp(GraphTransformation* transformation, } } if (!is_referenced) { - model->arrays.erase(removal_candidate); + model->EraseArray(removal_candidate); } } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.h b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.h index a06181ca0b5f1cbb930fa4295fec3d6adf66440d..9d448c3ee9088c16b96aa7ddc84457d2cab3231a 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.h +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_REMOVE_TRIVIAL_PASSTHROUGH_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_REMOVE_TRIVIAL_PASSTHROUGH_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_REMOVE_TRIVIAL_PASSTHROUGH_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_REMOVE_TRIVIAL_PASSTHROUGH_H_ #include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" #include "tensorflow/contrib/lite/toco/model.h" @@ -54,4 +54,4 @@ bool RemoveTrivialPassthroughOp(GraphTransformation* transformation, } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_REMOVE_TRIVIAL_PASSTHROUGH_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_REMOVE_TRIVIAL_PASSTHROUGH_H_ diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_slice.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_slice.cc new file mode 100644 index 0000000000000000000000000000000000000000..0cbbcd7c814d38e32ee55e9d9271adf532d20924 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_slice.cc @@ -0,0 +1,69 @@ +/* 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 +#include +#include +#include +#include + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +namespace { + +bool IsSliceTrivial(const Model& model, const Operator& op, + RemoveTrivialSlice* transformation) { + CHECK(op.type == OperatorType::kSlice); + + // Slices are trivial if they are slicing the entire input contents. + const auto& input_array = model.GetArray(op.inputs[0]); + const auto& output_array = model.GetArray(op.outputs[0]); + if (input_array.has_shape() && output_array.has_shape()) { + if (input_array.shape() == output_array.shape()) { + transformation->AddMessageF( + "%s is trivial because its input and output shapes are equal", + LogName(op)); + return true; + } + } + + return false; +} + +} // namespace + +bool RemoveTrivialSlice::Run(Model* model, std::size_t op_index) { + const auto reshape_it = model->operators.begin() + op_index; + auto* slice_op = reshape_it->get(); + if (slice_op->type != OperatorType::kSlice) { + return false; + } + + if (!IsSliceTrivial(*model, *slice_op, this)) { + return false; + } + + AddMessageF("Removing trivial %s", LogName(*slice_op)); + + CHECK_EQ(slice_op->inputs.size(), 3); + return RemoveTrivialPassthroughOp(this, model, op_index); +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc index e6cca8acf36745d989fb731aa948f257375d7e90..aa2c293382a98b476bee783ed8e177b19d35b858 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc @@ -33,7 +33,7 @@ bool RemoveUnusedOp::Run(Model* model, std::size_t op_index) { // the model. We allow specifying an arbitrary input_array, // treating the part of the graph leading up to it as unused. for (const auto& output : op->outputs) { - CHECK(model->arrays.count(output)); + CHECK(model->HasArray(output)); // If this output is provided as the model's input array, // then we don't need this operator to produce its contents. if (IsInputArray(*model, output)) { @@ -93,7 +93,7 @@ bool RemoveUnusedOp::Run(Model* model, std::size_t op_index) { if (IsDiscardableArray(*model, input) && CountOpsWithInput(*model, input) == 1 && !GetOpWithOutput(*model, input)) { - model->arrays.erase(input); + model->EraseArray(input); } } @@ -116,7 +116,7 @@ bool RemoveUnusedOp::Run(Model* model, std::size_t op_index) { continue; } // Generic case: do delete this output array. - model->arrays.erase(output); + model->EraseArray(output); } model->operators.erase(it); return true; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/reorder_activation_functions.cc b/tensorflow/contrib/lite/toco/graph_transformations/reorder_activation_functions.cc new file mode 100644 index 0000000000000000000000000000000000000000..30a005c789bb12e880e8e4534088d99ebacba84a --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/reorder_activation_functions.cc @@ -0,0 +1,103 @@ +/* Copyright 2017 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 +#include +#include +#include + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/runtime/types.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool ReorderActivationFunctions::Run(Model* model, std::size_t op_index) { + const auto ac_it = model->operators.begin() + op_index; + std::unique_ptr& ac_op = *ac_it; + DCHECK(ac_op); + + if (ac_op->type != OperatorType::kRelu6 && + ac_op->type != OperatorType::kRelu1 && + ac_op->type != OperatorType::kRelu) { + return false; + } + + auto exchange_it = FindOpWithOutput(*model, ac_op->inputs[0]); + if (exchange_it == model->operators.end()) return false; + // Find the op producing the array passed to this activation function + std::unique_ptr& exchange_op = *exchange_it; + DCHECK(exchange_op); + + if (exchange_op->type != OperatorType::kTensorFlowReshape) { + return false; + } + + DCHECK_EQ(exchange_op->outputs[0], ac_op->inputs[0]); + const auto& exchange_op_input = exchange_op->inputs[0]; + const auto& intermediate_array = exchange_op->outputs[0]; + const auto& ac_op_output = ac_op->outputs[0]; + + int count_ops_consuming_output = + CountOpsWithInput(*model, intermediate_array); + DCHECK_GE(count_ops_consuming_output, 1); + if (count_ops_consuming_output > 1) { + AddMessageF( + "Not exchanging activation function with %s because it is consumed by " + "more than 1 other operator", + LogName(*exchange_op)); + return false; + } + + // If the ac_op was originally producing an output_array we can't reorder as + // otherwise the output array would change. It'd be nice to still be able to + // reorder but if code is relying on the fetch names instead of array indices + // this won't work. + for (int i = 0; i < model->flags.output_arrays_size(); ++i) { + if (model->flags.output_arrays(i) == ac_op->outputs[0]) { + AddMessageF( + "Not exchanging activation function with %s to preserve output array " + "name %s", + LogName(*exchange_op), ac_op->outputs[0]); + return false; + } + } + + // Rewire by changing inputs, including all consumers. + Operator* consumer = GetFirstOpWithInput(*model, ac_op_output); + while (consumer) { + for (int i = 0; i < consumer->inputs.size(); ++i) { + if (consumer->inputs[i] == ac_op_output) { + consumer->inputs[i] = intermediate_array; + } + } + consumer = GetFirstOpWithInput(*model, ac_op_output); + } + ac_op->inputs[0] = exchange_op_input; + exchange_op->inputs[0] = ac_op_output; + + // Clear shapes; this will allow shape propagation to fix the sizes for us. + model->GetOrCreateArray(ac_op->outputs[0]).clear_shape(); + model->GetOrCreateArray(exchange_op->outputs[0]).clear_shape(); + + // Finally, reorder operators. Note that this only works when there are no + // other direct descendents of the exchange_op. + ac_op.swap(exchange_op); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc index 3eb7fa3896c57ea612f21f8b4f3fa568d19420d4..fb109eb91b16e3a73005230f821c18b9ef82d2fb 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc @@ -121,9 +121,9 @@ bool ResolveBatchNormalization::Run(Model* model, std::size_t op_index) { } // Remove the old param arrays - model->arrays.erase(bn_op->inputs[1]); - model->arrays.erase(bn_op->inputs[2]); - model->arrays.erase(bn_op->inputs[3]); + model->EraseArray(bn_op->inputs[1]); + model->EraseArray(bn_op->inputs[2]); + model->EraseArray(bn_op->inputs[3]); // Remove the old operator DCHECK_EQ(bn_it->get(), bn_op); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_to_space_nd_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_to_space_nd_attributes.cc index a4f198e92f8d5e4ef9d43d84ab3aa63b6e1ac5f7..a06919e228dc2084f8943a714a0ca111d013c159 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_to_space_nd_attributes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_to_space_nd_attributes.cc @@ -37,26 +37,19 @@ bool ResolveBatchToSpaceNDAttributes::Run(Model* model, std::size_t op_index) { } CHECK_EQ(op->inputs.size(), 3); - if (!IsConstantParameterArray(*model, op->inputs[1]) or + if (!IsConstantParameterArray(*model, op->inputs[1]) || !IsConstantParameterArray(*model, op->inputs[2])) return false; - // Handling block_shape. - const auto& block_shape_array = *model->arrays[op->inputs[1]]; - if (!block_shape_array.has_shape()) return false; - const std::vector& block_shape_dims = block_shape_array.shape().dims(); - CHECK_EQ(block_shape_dims.size(), 1); - std::vector block_shape_buffer = - block_shape_array.GetBuffer().data; - for (int i = 0; i < block_shape_dims[0]; ++i) { - op->block_shape.push_back(block_shape_buffer[i]); - } - - // Handling crops. - const auto& crops_array = *model->arrays[op->inputs[2]]; + // Handle crops + const auto& crops_array = model->GetArray(op->inputs[2]); if (!crops_array.has_shape()) return false; const std::vector& crops_dims = crops_array.shape().dims(); - CHECK_EQ(crops_dims.size(), 2); + if (crops_dims.size() != 2) { + // Code only handles crops of 2 dimensions. Perhaps another transformation + // will delete this op. + return false; + } std::vector crops_buffer = crops_array.GetBuffer().data; for (int i = 0; i < crops_dims[0]; ++i) { @@ -64,6 +57,17 @@ bool ResolveBatchToSpaceNDAttributes::Run(Model* model, std::size_t op_index) { op->after_crops.push_back(crops_buffer[i * 2 + 1]); } + // Handle block_shape + const auto& block_shape_array = model->GetArray(op->inputs[1]); + if (!block_shape_array.has_shape()) return false; + const std::vector& block_shape_dims = block_shape_array.shape().dims(); + CHECK_EQ(block_shape_dims.size(), 1); + std::vector block_shape_buffer = + block_shape_array.GetBuffer().data; + for (int i = 0; i < block_shape_dims[0]; ++i) { + op->block_shape.push_back(block_shape_buffer[i]); + } + return true; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_binary.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_binary.cc index 53e1be7a05807cde305eca2a7a8901f652f986f6..5e779f6765262326bd59db886c2feed603e0102e 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_binary.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_binary.cc @@ -141,6 +141,10 @@ void EvaluateBinaryOperatorOnConstantInputs(Model* model, outval = val0 - val1; } else if (binary_op->type == OperatorType::kDiv) { outval = val0 / val1; + } else if (binary_op->type == OperatorType::kFloorDiv) { + outval = floor(val0 / val1); + } else if (binary_op->type == OperatorType::kFloorMod) { + outval = val0 - (floor(val0 / val1) * val1); } else if (binary_op->type == OperatorType::kTensorFlowMinimum) { outval = std::min(val0, val1); } else if (binary_op->type == OperatorType::kTensorFlowMaximum) { @@ -162,8 +166,9 @@ void EvaluateBinaryOperatorOnConstantInputs(Model* model, void EvaluateBinaryOperatorOnConstantInputs(Model* model, const Operator* binary_op) { - const auto inputs_data_type = model->arrays[binary_op->inputs[0]]->data_type; - const auto output_data_type = model->arrays[binary_op->outputs[0]]->data_type; + const auto inputs_data_type = model->GetArray(binary_op->inputs[0]).data_type; + const auto output_data_type = + model->GetArray(binary_op->outputs[0]).data_type; #define TOCO_HANDLE_CASE(InputsDataType, OutputDataType) \ if (inputs_data_type == InputsDataType && \ output_data_type == OutputDataType) { \ @@ -191,6 +196,8 @@ bool ResolveConstantBinaryOperator::Run(Model* model, std::size_t op_index) { binary_op->type != OperatorType::kMul && binary_op->type != OperatorType::kSub && binary_op->type != OperatorType::kDiv && + binary_op->type != OperatorType::kFloorDiv && + binary_op->type != OperatorType::kFloorMod && binary_op->type != OperatorType::kTensorFlowMinimum && binary_op->type != OperatorType::kTensorFlowMaximum && binary_op->type != OperatorType::kTensorFlowLess && @@ -208,7 +215,7 @@ bool ResolveConstantBinaryOperator::Run(Model* model, std::size_t op_index) { return false; } - auto& output_array = *model->arrays[binary_op->outputs[0]]; + auto& output_array = model->GetArray(binary_op->outputs[0]); // Yield until the output array dims have been resolved. if (!output_array.has_shape()) { return false; @@ -233,10 +240,10 @@ bool ResolveConstantBinaryOperator::Run(Model* model, std::size_t op_index) { // Remove the binary operator and its inputs if (CountOpsWithInput(*model, binary_op->inputs[0]) == 1) { - model->arrays.erase(binary_op->inputs[0]); + model->EraseArray(binary_op->inputs[0]); } if (CountOpsWithInput(*model, binary_op->inputs[1]) == 1) { - model->arrays.erase(binary_op->inputs[1]); + model->EraseArray(binary_op->inputs[1]); } AddMessageF("Resolved constant %s to the equivalent constant array", LogName(*binary_op)); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_concatenation.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_concatenation.cc index 9835f86398a37f118d3ebd5b568ffddbcd56c38b..064810b53e7c3bee4601204c9dbd976c374a6a60 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_concatenation.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_concatenation.cc @@ -73,7 +73,7 @@ void CopyTensorSegments(const std::vector& input_arrays, // Receives a series of input arrays of type Array and an integer showing the // axis on which those arrays will be concatenated. It returns the concatenated -// arrray. +// array. template void ConcatenateTensorBuffers(const std::vector& input_arrays, int concatenation_axis, @@ -189,7 +189,10 @@ bool ResolveConstantConcatenation::Run(Model* model, std::size_t op_index) { // Remove all the resolved arrays. for (const string& input_name : concat_op->inputs) { - model->arrays.erase(input_name); + // Check to prevent removal of shared tensors + if (CountOpsWithInput(*model, input_name) == 1) { + model->EraseArray(input_name); + } } // Remove concatenate operator diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_fake_quant.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_fake_quant.cc index 244adcc4c46eda9de79dd753565113bbeca970c5..944901ece77430708013ea4ca340a30511ba0174 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_fake_quant.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_fake_quant.cc @@ -50,6 +50,7 @@ bool ResolveConstantFakeQuant::Run(Model* model, std::size_t op_index) { output_array.data_type = ArrayDataType::kFloat; CHECK(!output_array.buffer); const auto& input_buffer = input_array.GetBuffer(); + output_array.GetOrCreateMinMax() = *fakequant_op->minmax; auto& output_buffer = output_array.GetMutableBuffer(); const int size = input_buffer.data.size(); output_buffer.data.resize(size); @@ -66,7 +67,7 @@ bool ResolveConstantFakeQuant::Run(Model* model, std::size_t op_index) { output_buffer.data[i] = dst_val; } if (CountOpsWithInput(*model, fakequant_op->inputs[0]) == 1) { - model->arrays.erase(fakequant_op->inputs[0]); + model->EraseArray(fakequant_op->inputs[0]); } model->operators.erase(fakequant_it); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_fill.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_fill.cc index 9da51d9147a98a935d00db04827aa7ebb12998b9..f6f95481b57f58f497b119df73d331f13d9705c0 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_fill.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_fill.cc @@ -104,11 +104,11 @@ bool ResolveConstantFill::Run(Model* model, std::size_t op_index) { // Erase input arrays if no longer used if (IsDiscardableArray(*model, op->inputs[0]) && CountOpsWithInput(*model, op->inputs[0]) == 1) { - model->arrays.erase(op->inputs[0]); + model->EraseArray(op->inputs[0]); } if (IsDiscardableArray(*model, op->inputs[1]) && CountOpsWithInput(*model, op->inputs[1]) == 1) { - model->arrays.erase(op->inputs[1]); + model->EraseArray(op->inputs[1]); } // Erase the operator diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_range.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_range.cc new file mode 100644 index 0000000000000000000000000000000000000000..1a0ba9e2bc7235720b59210cdd6affa089613077 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_range.cc @@ -0,0 +1,107 @@ +/* Copyright 2017 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/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool ResolveConstantRange::Run(Model* model, std::size_t op_index) { + const auto it = model->operators.begin() + op_index; + auto* base_op = it->get(); + if (base_op->type != OperatorType::kRange) { + return false; + } + auto* op = static_cast(base_op); + + CHECK_EQ(op->inputs.size(), 3); + const auto& start_array = model->GetArray(op->inputs[0]); + if (!start_array.has_shape()) { + // Yield until all input dims have been resolved. + return false; + } + const auto& limit_array = model->GetArray(op->inputs[1]); + if (!limit_array.has_shape()) { + // Yield until all input dims have been resolved. + return false; + } + const auto& delta_array = model->GetArray(op->inputs[2]); + if (!delta_array.has_shape()) { + // Yield until all input dims have been resolved. + return false; + } + + for (const auto& input : op->inputs) { + if (!IsConstantParameterArray(*model, input)) { + // yield if any input is mutable + return false; + } + } + + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.data_type == ArrayDataType::kNone) { + // Yield until the output type has been set by PropagateArrayDataTypes + return false; + } + + CHECK_EQ(RequiredBufferSizeForShape(start_array.shape()), 1) + << "Range op inputs must be scalar."; + CHECK_EQ(RequiredBufferSizeForShape(limit_array.shape()), 1) + << "Range op inputs must be scalar."; + CHECK_EQ(RequiredBufferSizeForShape(delta_array.shape()), 1) + << "Range op inputs must be scalar."; + + CHECK(start_array.data_type == ArrayDataType::kInt32) + << "Range op inputs must be int32."; + CHECK(limit_array.data_type == ArrayDataType::kInt32) + << "Range op inputs must be int32."; + CHECK(delta_array.data_type == ArrayDataType::kInt32) + << "Range op inputs must be int32."; + + // Compute buffer contents + int start = start_array.GetBuffer().data[0]; + int limit = limit_array.GetBuffer().data[0]; + int delta = delta_array.GetBuffer().data[0]; + auto& buffer = output_array.GetMutableBuffer(); + buffer.data.clear(); + for (int32 val = start; val < limit; val += delta) { + buffer.data.push_back(val); + } + CHECK_EQ(floor((limit - start) / delta), buffer.data.size()); + CHECK_EQ(buffer.data.size(), output_array.shape().dims()[0]); + + // Delete the input array if no longer used + if (IsDiscardableArray(*model, op->inputs[0]) && + CountOpsWithInput(*model, op->inputs[0]) == 1) { + model->EraseArray(op->inputs[0]); + } + if (IsDiscardableArray(*model, op->inputs[1]) && + CountOpsWithInput(*model, op->inputs[1]) == 1) { + model->EraseArray(op->inputs[1]); + } + if (IsDiscardableArray(*model, op->inputs[2]) && + CountOpsWithInput(*model, op->inputs[2]) == 1) { + model->EraseArray(op->inputs[2]); + } + + // Delete the operator + model->operators.erase(it); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_shape_or_rank.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_shape_or_rank.cc new file mode 100644 index 0000000000000000000000000000000000000000..9ea01acd05364224ce219bed533c999793a2a2f1 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_shape_or_rank.cc @@ -0,0 +1,72 @@ +/* Copyright 2017 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/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool ResolveConstantShapeOrRank::Run(Model* model, std::size_t op_index) { + const auto it = model->operators.begin() + op_index; + const auto* op = it->get(); + if (!(op->type == OperatorType::kTensorFlowShape || + op->type == OperatorType::kRank)) { + return false; + } + + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.data_type == ArrayDataType::kNone) { + // Yield until the output type has been resolved + return false; + } + + const auto& input_array = model->GetArray(op->inputs[0]); + if (!input_array.has_shape()) { + // Yield until the input array's shape has been resolved. + return false; + } + + if (!output_array.has_shape()) { + // Yield until the output shape has been resolved. + return false; + } + + // Compute the output + CHECK(!output_array.buffer); + auto& output_buffer = output_array.GetMutableBuffer(); + if (op->type == OperatorType::kTensorFlowShape) { + // Copy the input shape into the output buffer. + output_buffer.data = input_array.shape().dims(); + } else if (op->type == OperatorType::kRank) { + // Copy the dimension count into the output buffer. + output_buffer.data.resize(1); + output_buffer.data[0] = input_array.shape().dimensions_count(); + } + output_array.mutable_shape()->ReplaceDims( + {static_cast(output_buffer.data.size())}); + + // Delete the input array if no longer used + if (IsDiscardableArray(*model, op->inputs[0]) && + CountOpsWithInput(*model, op->inputs[0]) == 1) { + model->EraseArray(op->inputs[0]); + } + + model->operators.erase(it); + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_stack.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_stack.cc new file mode 100644 index 0000000000000000000000000000000000000000..ea0d6dc8200897db9266efbe41556dbf4c296db3 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_stack.cc @@ -0,0 +1,113 @@ +/* Copyright 2017 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 + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +namespace { + +template +void Stack(Model* model, StackOperator const& op) { + auto& output_array = model->GetArray(op.outputs[0]); + CHECK(output_array.data_type == Type); + + // Create a buffer for the output array + std::vector>& output_data = + output_array.GetMutableBuffer().data; + output_data.resize(RequiredBufferSizeForShape(output_array.shape())); + + // Stack inputs into buffer + CHECK_EQ(op.axis, 0) << "Stacking only supported along first axis"; + int dst_offset = 0; + for (int i = 0; i < op.inputs.size(); i++) { + // Append array data to output for each input array + const auto& input_array = model->GetArray(op.inputs[i]); + int input_size = RequiredBufferSizeForShape(input_array.shape()); + memcpy(&output_data[dst_offset], &input_array.GetBuffer().data[0], + input_size * sizeof(Type)); + dst_offset += input_size; + } + CHECK_EQ(dst_offset, output_data.size()); +} + +} // namespace + +bool ResolveConstantStack::Run(Model* model, std::size_t op_index) { + auto it = model->operators.begin() + op_index; + const auto* base_op = it->get(); + if (base_op->type != OperatorType::kStack) { + return false; + } + const auto* op = static_cast(base_op); + + CHECK_GE(op->inputs.size(), 1); + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.data_type == ArrayDataType::kNone) { + // Yield until the output type has been set by PropagateArrayDataTypes + return false; + } + + if (!output_array.has_shape()) { + // Yield until the output shape has been set by PropagateFixedShapes + return false; + } + + for (const auto& input : op->inputs) { + if (!IsConstantParameterArray(*model, input)) { + // Yield if any input is mutable + return false; + } + } + + CHECK(!output_array.buffer); + switch (output_array.data_type) { + case ArrayDataType::kFloat: + Stack(model, *op); + break; + case ArrayDataType::kUint8: + Stack(model, *op); + break; + case ArrayDataType::kInt32: + Stack(model, *op); + break; + case ArrayDataType::kInt64: + Stack(model, *op); + break; + default: + LOG(FATAL) << "Unsupported data type given to Stack op with output \"" + << op->outputs[0] << "\""; + break; + } + + // Erase input arrays if no longer used + for (const auto& input : op->inputs) { + if (IsDiscardableArray(*model, input) && + CountOpsWithInput(*model, input) == 1) { + model->EraseArray(input); + } + } + + // Erase the operator + model->operators.erase(it); + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_strided_slice.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_strided_slice.cc new file mode 100644 index 0000000000000000000000000000000000000000..a0cfc3d59763dc1211ed4d1ac114d371a4a7ee0b --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_strided_slice.cc @@ -0,0 +1,198 @@ +/* Copyright 2017 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 + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +namespace { + +int StartForAxis(StridedSliceOperator const& op, Shape const& input_shape, + int axis) { + int start; + if (op.begin_mask & 1 << axis) { + // If begin mask bit is set, use the first element + start = 0; + } else { + // Otherwise, use the specified element + start = op.start_indices[axis]; + if (start < 0) { + // Handle negative indices + start += input_shape.dims(axis); + } + } + return start; +} + +int StopForAxis(StridedSliceOperator const& op, Shape const& input_shape, + int axis) { + int stop; + if (op.end_mask & (1 << axis)) { + // If end mask bit set, use the last element + stop = input_shape.dims(axis); + } else { + // Otherwise, use the specified element + stop = op.stop_indices[axis]; + if (stop < 0) { + // Handle negative indices + stop += input_shape.dims(axis); + } + } + return stop; +} + +template +void StridedSlice(StridedSliceOperator const& op, Array const& input_array, + Array* output_array) { + // The TensorFlow documentation for StridedSlice is a bit ambiguous in places + // (https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/strided-slice). + // Use the source code at /third_party/tensorflow/core/util/strided_op.cc as + // "master documentation". + + CHECK(input_array.data_type == Type); + CHECK(output_array->data_type == Type); + CHECK_EQ(op.ellipsis_mask, 0); + CHECK_EQ(op.new_axis_mask, 0); + + int num_input_axes = op.start_indices.size(); + CHECK_EQ(num_input_axes, op.stop_indices.size()); + CHECK_EQ(num_input_axes, op.strides.size()); + for (int i = 0; i < op.strides.size(); i++) { + CHECK_GE(op.strides[i], 0) << "Negative strides usupported"; + } + + // Create a buffer for the output array + std::vector>& output_data = + output_array->GetMutableBuffer().data; + output_data.resize(RequiredBufferSizeForShape(output_array->shape())); + + // Initialize source coordinate + Shape const& input_shape = input_array.shape(); + Buffer const& input_buffer = input_array.GetBuffer(); + std::vector src_coord(op.start_indices.size()); + for (int axis = 0; axis < num_input_axes; axis++) { + src_coord[axis] = StartForAxis(op, input_shape, axis); + } + + // In order to handle any number (N) of dimensions, we copy elements one by + // one and treat the source coordinate as an N digit number (src_coord here). + // Each "digit" is incremented individually (by the stride). When it overflows + // (becomes greater than the stop), that digit is reset and a carry flag is + // used to increment the next digit. + int dst_offset = 0; + do { + // Copy element. + output_data[dst_offset] = input_buffer.data[Offset(input_shape, src_coord)]; + + // Compute next source input coordinates. + bool carry = true; + for (int axis = 0; axis < num_input_axes; axis++) { + // Increment this axis if we carried from the previous one + if (carry) { + src_coord[axis] += op.strides[axis]; + } + + // Check if we've overflowed. + if (src_coord[axis] >= StopForAxis(op, input_shape, axis)) { + // Reset axis and set carry + src_coord[axis] = StartForAxis(op, input_shape, axis); + carry = true; + } else { + carry = false; + } + } + // increment destination buffer offset + dst_offset++; + } while (dst_offset < output_data.size()); +} + +} // anonymous namespace + +bool ResolveConstantStridedSlice::Run(Model* model, std::size_t op_index) { + const auto it = model->operators.begin() + op_index; + const auto* base_op = it->get(); + if (base_op->type != OperatorType::kStridedSlice) { + return false; + } + + const StridedSliceOperator* op = + static_cast(base_op); + + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.data_type == ArrayDataType::kNone) { + // Yield until the output type has been set by PropagateArrayDataTypes + return false; + } + + if (!output_array.has_shape()) { + // Yield until the output shape has been set by PropagateFixedShapes + return false; + } + + if (op->start_indices.empty() || op->stop_indices.empty() || + op->strides.empty()) { + // Attributes have not resolved yet. + return false; + } + + const auto& input_array = model->GetArray(op->inputs[0]); + if (!input_array.has_shape()) { + // Yield until the value shape has been resolved. + return false; + } + if (!IsConstantParameterArray(*model, op->inputs[0])) { + // Yield until the value is constant. + return false; + } + + CHECK(!output_array.buffer); + switch (output_array.data_type) { + case ArrayDataType::kFloat: + StridedSlice(*op, input_array, &output_array); + break; + case ArrayDataType::kUint8: + StridedSlice(*op, input_array, &output_array); + break; + case ArrayDataType::kInt32: + StridedSlice(*op, input_array, &output_array); + break; + case ArrayDataType::kInt64: + StridedSlice(*op, input_array, &output_array); + break; + default: + LOG(FATAL) + << "Unsupported data type input to StridedSlice op with output \"" + << op->outputs[0] << "\""; + break; + } + + // Erase input array if no longer used + if (IsDiscardableArray(*model, op->inputs[0]) && + CountOpsWithInput(*model, op->inputs[0]) == 1) { + model->EraseArray(op->inputs[0]); + } + + // Erase the operator + model->operators.erase(it); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_tensorflow_shape.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_tensorflow_shape.cc deleted file mode 100644 index 8cc6db161987bbd834212fdfed7e1f82cac958ce..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_tensorflow_shape.cc +++ /dev/null @@ -1,62 +0,0 @@ -/* Copyright 2017 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 -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" -#include "tensorflow/contrib/lite/toco/model.h" -#include "tensorflow/contrib/lite/toco/tooling_util.h" -#include "tensorflow/core/platform/logging.h" - -namespace toco { - -bool ResolveConstantTensorFlowShape::Run(Model* model, std::size_t op_index) { - const auto tfshape_it = model->operators.begin() + op_index; - const auto* tfshape_base_op = tfshape_it->get(); - if (tfshape_base_op->type != OperatorType::kTensorFlowShape) { - return false; - } - - const auto* tfshape_op = - static_cast(tfshape_base_op); - - const auto& input_array = model->GetArray(tfshape_op->inputs[0]); - auto& output_array = model->GetArray(tfshape_op->outputs[0]); - - // Yield until the input array's shape has been resolved. - if (!input_array.has_shape()) { - return false; - } - - // Create a buffer for the output array, making it a constant array, and - // copy the input shape into the output buffer. - CHECK(!output_array.buffer); - auto& output_buffer = output_array.GetMutableBuffer(); - output_buffer.data = input_array.shape().dims(); - - // Erase the input array if no longer used - if (IsDiscardableArray(*model, tfshape_op->inputs[0]) && - CountOpsWithInput(*model, tfshape_op->inputs[0]) == 1) { - model->arrays.erase(tfshape_op->inputs[0]); - } - model->operators.erase(tfshape_it); - - return true; -} - -} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_transpose.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_transpose.cc new file mode 100644 index 0000000000000000000000000000000000000000..4f984bfde55b3457694bb411bbfdf30723c7066e --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_transpose.cc @@ -0,0 +1,180 @@ +/* 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 + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +namespace { + +// Transposes an array up to rank 4. +// This is ShuffleArrayTemplate with non-enum permutation. +template +void Transpose(Model* model, const Array& input_array, + const std::vector& perm, Array* output_array) { + const Shape& input_shape = input_array.shape(); + const std::vector>& input_data = + input_array.GetBuffer().data; + + const Shape& output_shape = output_array->shape(); + std::vector>& output_data = + output_array->GetMutableBuffer().data; + output_data.resize(RequiredBufferSizeForShape(output_shape)); + + CHECK(input_shape.dimensions_count() == output_shape.dimensions_count()); + const int dim = input_shape.dimensions_count(); + CHECK_LE(dim, 4); + CHECK(perm.size() >= dim); + for (int i = 0; i < dim; i++) { + CHECK(perm[i] >= 0 && perm[i] < dim); + CHECK(input_shape.dims(perm[i]) == output_shape.dims(i)); + } + Shape extended_input_shape = input_shape; + ExtendShape(&extended_input_shape, 4); + Shape extended_output_shape = output_shape; + ExtendShape(&extended_output_shape, 4); + std::vector extended_perm; + ExtendShuffle(perm, 4, &extended_perm); + + const std::vector& extended_input_dims = extended_input_shape.dims(); + const std::vector& extended_output_dims = extended_output_shape.dims(); + + // TODO(starka): Rework to handle different numbers of dimensions. + int input_strides[4]; + input_strides[3] = 1; + input_strides[2] = extended_input_dims[3]; + input_strides[1] = input_strides[2] * extended_input_dims[2]; + input_strides[0] = input_strides[1] * extended_input_dims[1]; + const int input_stride_0 = input_strides[extended_perm[3]]; + const int input_stride_1 = input_strides[extended_perm[2]]; + const int input_stride_2 = input_strides[extended_perm[1]]; + const int input_stride_3 = input_strides[extended_perm[0]]; + + const int output_size_0 = extended_output_dims[3]; + const int output_size_1 = extended_output_dims[2]; + const int output_size_2 = extended_output_dims[1]; + const int output_size_3 = extended_output_dims[0]; + const int output_stride_0 = 1; + const int output_stride_1 = output_size_0; + const int output_stride_2 = output_stride_1 * output_size_1; + const int output_stride_3 = output_stride_2 * output_size_2; + + for (int i3 = 0; i3 < output_size_3; i3++) { + const DataType* const input_ptr_3 = + input_data.data() + i3 * input_stride_3; + DataType* const output_ptr_3 = + output_data.data() + i3 * output_stride_3; + for (int i2 = 0; i2 < output_size_2; i2++) { + const DataType* const input_ptr_2 = + input_ptr_3 + i2 * input_stride_2; + DataType* const output_ptr_2 = output_ptr_3 + i2 * output_stride_2; + for (int i1 = 0; i1 < output_size_1; i1++) { + const DataType* input_ptr = input_ptr_2 + i1 * input_stride_1; + DataType* output_ptr = output_ptr_2 + i1 * output_stride_1; + DataType* const output_ptr_end = + output_ptr + output_size_0 * output_stride_0; + while (output_ptr != output_ptr_end) { + *output_ptr = *input_ptr; + input_ptr += input_stride_0; + output_ptr += output_stride_0; + } + } + } + } +} + +} // namespace + +bool ResolveConstantTranspose::Run(Model* model, std::size_t op_index) { + auto it = model->operators.begin() + op_index; + const auto* base_op = it->get(); + if (base_op->type != OperatorType::kTranspose) { + return false; + } + const auto* op = static_cast(base_op); + + CHECK_EQ(op->inputs.size(), 2); + CHECK_EQ(op->outputs.size(), 1); + auto& output_array = model->GetArray(op->outputs[0]); + if (output_array.data_type == ArrayDataType::kNone) { + // Yield until the output type has been set by PropagateArrayDataTypes. + return false; + } + if (!output_array.has_shape()) { + // Yield until the output shape has been set by PropagateFixedShapes. + return false; + } + + // We require constant inputs. + if (!IsConstantParameterArray(*model, op->inputs[0]) || + !IsConstantParameterArray(*model, op->inputs[1])) { + return false; + } + const Array& input_array = model->GetArray(op->inputs[0]); + + if (input_array.minmax) { + output_array.GetOrCreateMinMax() = input_array.GetMinMax(); + } + + if (op->perm.empty()) { + // Yield until perm has been populated by ResolveTransposeAttributes. + return false; + } + + // We currently only support 1-4 dimensions. + CHECK_LE(op->perm.size(), 4); + + CHECK(!output_array.buffer); + switch (output_array.data_type) { + case ArrayDataType::kFloat: + Transpose(model, input_array, op->perm, + &output_array); + break; + case ArrayDataType::kUint8: + Transpose(model, input_array, op->perm, + &output_array); + break; + case ArrayDataType::kInt32: + Transpose(model, input_array, op->perm, + &output_array); + break; + case ArrayDataType::kInt64: + Transpose(model, input_array, op->perm, + &output_array); + break; + default: + LOG(FATAL) << "Unsupported data type given to Transpose op with output \"" + << op->outputs[0] << "\""; + break; + } + + // Erase input arrays if no longer used. + for (const auto& input : op->inputs) { + if (IsDiscardableArray(*model, input) && + CountOpsWithInput(*model, input) == 1) { + model->EraseArray(input); + } + } + + // Erase the operator. + model->operators.erase(it); + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_unary.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_unary.cc index 26ff9d887b40651559ad030cd41a824679d6dd15..f227554bc505efe6a758fdd9894fee43f2500641 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_unary.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_unary.cc @@ -139,14 +139,13 @@ bool ResolveConstantUnaryOperator::Run(Model* model, std::size_t op_index) { output_buffer_size * sizeof(output_float_data[0])); } else if (unary_op->type == OperatorType::kTensorFlowSum) { // At the moment only full reduction across all dimensions is supported. - for (int i = 0; i < output_dims_count; i++) { - CHECK_EQ(output_shape.dims(i), 1); - } float sum = 0.f; for (int i = 0; i < input_buffer_size; i++) { sum += (*input_float_data)[i]; } - output_float_data[0] = sum; + for (int i = 0; i < output_buffer_size; ++i) { + output_float_data[i] = sum; + } } else if (unary_op->type == OperatorType::kTensorFlowMin) { // At the moment only full reduction across all dimensions is supported. // TODO(starka): Output should not be padded. @@ -199,7 +198,7 @@ bool ResolveConstantUnaryOperator::Run(Model* model, std::size_t op_index) { } for (const auto& input : unary_op->inputs) { if (CountOpsWithInput(*model, input) == 1) { - model->arrays.erase(input); + model->EraseArray(input); } } AddMessageF("Resolved constant %s to the equivalent constant array", diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_mean_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_mean_attributes.cc index 444f59d14bc916e306186eaa031f04493af63bb8..013b50ac9ba8a51c23b19953d987b2fbf63fcea1 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_mean_attributes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_mean_attributes.cc @@ -36,19 +36,9 @@ bool ResolveMeanAttributes::Run(Model* model, std::size_t op_index) { if (op->inputs.size() != 2) return false; if (!IsConstantParameterArray(*model, op->inputs[1])) return false; - const auto& indices_array = *model->arrays[op->inputs[1]]; + const auto& indices_array = model->GetArray(op->inputs[1]); if (!indices_array.has_shape()) return false; - - // We only support simultaneous reduction over width and height. - std::vector axis = indices_array.GetBuffer().data; - if (axis.size() != 2) { - return false; - } - if (!((axis[0] == 1 && axis[1] == 2) || (axis[0] == 2 && axis[1] == 1))) { - return false; - } - - op->axis = axis; + op->axis = indices_array.GetBuffer().data; return true; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_multiply_by_zero.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_multiply_by_zero.cc new file mode 100644 index 0000000000000000000000000000000000000000..37beb41dfc5904fc6ace79ebea2420d2ab92fbfb --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_multiply_by_zero.cc @@ -0,0 +1,152 @@ +/* Copyright 2017 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 +#include +#include +#include +#include + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" + +namespace toco { + +namespace { + +template +bool AreAllBufferElementsZero(const std::vector& buffer_data) { + for (auto x : buffer_data) { + if (x != 0) { + return false; + } + } + return true; +} + +template +void FillArrayWithZeros(Array* array) { + CHECK(array->data_type == Type); + std::vector>& data = array->GetMutableBuffer().data; + data.resize(RequiredBufferSizeForShape(array->shape())); + for (size_t i = 0; i < data.size(); i++) { + data[i] = 0; + } +} + +} // namespace + +// Removes a multiplication by array of constant zeros by making the output +// array an array of constant zeros and removing the input arrays if they are no +// longer needed. +bool ResolveMultiplyByZero::Run(Model* model, std::size_t op_index) { + const auto mul_it = model->operators.begin() + op_index; + auto* mul_op = mul_it->get(); + if (mul_op->type != OperatorType::kMul) { + return false; + } + const auto& output_array_name = mul_op->outputs[0]; + auto& output_array = model->GetArray(output_array_name); + + // Yield if the output shape is not known yet. + if (!output_array.has_shape()) { + return false; + } + + // This transformation only handles the case where one operand is all 0's and + // the other is non-constant. Other cases are handled by constant propagation + // or the trivial binary removal pass. + const bool is_input_constant[2] = { + IsConstantParameterArray(*model, mul_op->inputs[0]), + IsConstantParameterArray(*model, mul_op->inputs[1]), + }; + if (!is_input_constant[0] && !is_input_constant[1]) { + // Neither input is constant, so nothing we can resolve here. + return false; + } + if (is_input_constant[0] && is_input_constant[1]) { + // Both inputs are constants. That's a job for constants propagation, not + // for us to handle here. + return false; + } + const int index_of_constant_input = is_input_constant[0] ? 0 : 1; + const int index_of_variable_input = is_input_constant[0] ? 1 : 0; + CHECK(is_input_constant[index_of_constant_input]); + CHECK(!is_input_constant[index_of_variable_input]); + + const auto& constant_input_array = + model->GetArray(mul_op->inputs[index_of_constant_input]); + + CHECK(constant_input_array.data_type == output_array.data_type); + switch (output_array.data_type) { + case ArrayDataType::kFloat: { + const auto& constant_input_data = + constant_input_array.GetBuffer().data; + if (!AreAllBufferElementsZero>( + constant_input_data)) { + return false; + } + FillArrayWithZeros(&output_array); + } break; + case ArrayDataType::kUint8: { + const auto& constant_input_data = + constant_input_array.GetBuffer().data; + if (!AreAllBufferElementsZero>( + constant_input_data)) { + return false; + } + FillArrayWithZeros(&output_array); + } break; + case ArrayDataType::kInt32: { + const auto& constant_input_data = + constant_input_array.GetBuffer().data; + if (!AreAllBufferElementsZero>( + constant_input_data)) { + return false; + } + FillArrayWithZeros(&output_array); + } break; + case ArrayDataType::kInt64: { + const auto& constant_input_data = + constant_input_array.GetBuffer().data; + if (!AreAllBufferElementsZero>( + constant_input_data)) { + return false; + } + FillArrayWithZeros(&output_array); + } break; + default: + AddMessageF( + "Cannot resolve multiply by 0 because of unsupported data type\n"); + return false; + } + + // Erase input arrays to the multiply if no longer used + if (IsDiscardableArray(*model, mul_op->inputs[0]) && + CountOpsWithInput(*model, mul_op->inputs[0]) == 1) { + model->EraseArray(mul_op->inputs[0]); + } + if (IsDiscardableArray(*model, mul_op->inputs[1]) && + CountOpsWithInput(*model, mul_op->inputs[1]) == 1) { + model->EraseArray(mul_op->inputs[1]); + } + + // Erase the multiply operator. + model->operators.erase(mul_it); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_pad_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_pad_attributes.cc index d5f5869c625f419a825f6bd652a04eca1bce4a6f..8a8e723cf7b2d77ec199e3817464a068bf85afdd 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_pad_attributes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_pad_attributes.cc @@ -35,7 +35,7 @@ bool ResolvePadAttributes::Run(Model* model, std::size_t op_index) { CHECK_EQ(op->inputs.size(), 2); if (!IsConstantParameterArray(*model, op->inputs[1])) return false; - const auto& array = *model->arrays[op->inputs[1]]; + const auto& array = model->GetArray(op->inputs[1]); if (!array.has_shape()) return false; const std::vector& dims = array.shape().dims(); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_reorder_axes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_reorder_axes.cc index 8fa7b83bedc0da99c3a5a60f38586f712eeb3c4e..bc70db0bd8c26319fa140616de96452260a01058 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_reorder_axes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_reorder_axes.cc @@ -25,6 +25,31 @@ limitations under the License. namespace toco { +// Reorder the elements of an input_array according to the input_axes_order and +// output_axes_order. Then adjust the shapes of the input and output arrays +// accordingly. Note that input_array must have a buffer (that is, it is a +// constant array). +template +void ReorderAxes(AxesOrder input_axes_order, AxesOrder output_axes_order, + Array* input_array, Array* output_array) { + CHECK(input_array->buffer->type == DataType); + CHECK(!output_array->buffer); + auto& input_data = input_array->GetMutableBuffer().data; + std::vector reordered_data; + reordered_data.resize(RequiredBufferSizeForShape(output_array->shape())); + // TODO(b/62904716) Shapes should be used directly. + Shape input_shape = input_array->shape(); + Shape output_shape = output_array->shape(); + if (AxesCount(input_axes_order) == 2) { + UnextendShape(&input_shape, 2); + UnextendShape(&output_shape, 2); + } + ShuffleArray(input_shape, input_axes_order, output_axes_order, output_shape, + input_data.data(), reordered_data.data()); + input_data = reordered_data; + input_array->copy_shape(output_array->shape()); +} + bool ResolveReorderAxes::Run(Model* model, std::size_t op_index) { auto reorder_it = model->operators.begin() + op_index; auto* reorder_op = static_cast(reorder_it->get()); @@ -35,16 +60,7 @@ bool ResolveReorderAxes::Run(Model* model, std::size_t op_index) { const auto& output_array_name = reorder_op->outputs[0]; auto& input_array = model->GetArray(input_array_name); auto& output_array = model->GetArray(output_array_name); - string constant_input_array_name = input_array_name; if (!input_array.buffer) { - const auto* op_producing_input = GetOpWithOutput(*model, input_array_name); - if (op_producing_input && - op_producing_input->type == OperatorType::kFakeQuant) { - constant_input_array_name = op_producing_input->inputs[0]; - } - } - auto& constant_input_array = model->GetArray(constant_input_array_name); - if (!constant_input_array.buffer) { return false; } // Yield until output dims have been resolved. @@ -52,26 +68,19 @@ bool ResolveReorderAxes::Run(Model* model, std::size_t op_index) { return false; } // Reorder the input array dims and buffer data - CHECK(constant_input_array.buffer->type == ArrayDataType::kFloat); - CHECK(!output_array.buffer); - auto& input_data = - constant_input_array.GetMutableBuffer().data; - std::vector reordered_data; - reordered_data.resize(RequiredBufferSizeForShape(output_array.shape())); - const auto input_axes_order = reorder_op->input_axes_order; - const auto output_axes_order = reorder_op->output_axes_order; - // TODO(b/62904716) Shapes should be used directly. - Shape input_shape = constant_input_array.shape(); - Shape output_shape = output_array.shape(); - if (AxesCount(input_axes_order) == 2) { - UnextendShape(&input_shape, 2); - UnextendShape(&output_shape, 2); + if (input_array.buffer->type == ArrayDataType::kFloat) { + ReorderAxes(reorder_op->input_axes_order, + reorder_op->output_axes_order, + &input_array, &output_array); + } else if (input_array.buffer->type == ArrayDataType::kInt32) { + ReorderAxes(reorder_op->input_axes_order, + reorder_op->output_axes_order, + &input_array, &output_array); + } else { + LOG(FATAL) << "Cannot ReorderAxes unless input buffer is float or uint8."; } - ShuffleArray(input_shape, input_axes_order, output_axes_order, output_shape, - input_data.data(), reordered_data.data()); - input_data = reordered_data; + input_array.copy_shape(output_array.shape()); - constant_input_array.copy_shape(output_array.shape()); // Update the edges of the graph to point to the input array for (const auto& other_op : model->operators) { @@ -85,7 +94,7 @@ bool ResolveReorderAxes::Run(Model* model, std::size_t op_index) { AddMessageF("Reordered axes for array %s", input_array_name); // Remove the op and output array. - model->arrays.erase(output_array_name); + model->EraseArray(output_array_name); model->operators.erase(reorder_it); return true; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_reshape_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_reshape_attributes.cc index bed2a85bd262c49913f22e522d260c4dc6510246..2e063e35548aa5e51c3bcc94a2dfc7992180d014 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_reshape_attributes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_reshape_attributes.cc @@ -37,7 +37,7 @@ bool ResolveReshapeAttributes::Run(Model* model, std::size_t op_index) { if (!op->shape.empty()) return false; if (IsConstantParameterArray(*model, reshape_op->inputs[1])) { - const auto& constant_input_array = *model->arrays[reshape_op->inputs[1]]; + const auto& constant_input_array = model->GetArray(reshape_op->inputs[1]); op->shape = constant_input_array.GetBuffer().data; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_slice_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_slice_attributes.cc index 1d0a2ec8f6c1f532f23873062534a37e07fff72b..e760d08e5a6c2f56db6b11fee922b701d33dd1a0 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_slice_attributes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_slice_attributes.cc @@ -36,10 +36,10 @@ bool ResolveSliceAttributes::Run(Model* model, std::size_t op_index) { if (!IsConstantParameterArray(*model, op->inputs[1])) return false; if (!IsConstantParameterArray(*model, op->inputs[2])) return false; - const auto& begin_array = *model->arrays[op->inputs[1]]; + const auto& begin_array = model->GetArray(op->inputs[1]); if (!begin_array.has_shape()) return false; - const auto& size_array = *model->arrays[op->inputs[2]]; + const auto& size_array = model->GetArray(op->inputs[2]); if (!size_array.has_shape()) return false; op->begin = begin_array.GetBuffer().data; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_space_to_batch_nd_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_space_to_batch_nd_attributes.cc new file mode 100644 index 0000000000000000000000000000000000000000..dad6aceccfd201b3db07c29c99a8c6ef75bb89a1 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_space_to_batch_nd_attributes.cc @@ -0,0 +1,78 @@ +/* Copyright 2017 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 +#include +#include +#include + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool ResolveSpaceToBatchNDAttributes::Run(Model* model, std::size_t op_index) { + const auto op_it = model->operators.begin() + op_index; + if (op_it->get()->type != OperatorType::kSpaceToBatchND) return false; + + auto* op = static_cast(op_it->get()); + + // The attributes are resolved only when the 3 attributes (block_shape, + // before_paddings, after_paddings) are all constant. + if (!op->block_shape.empty()) { + return false; + } + + const int block_shape_index = 1; + const int paddings_index = 2; + + CHECK_EQ(op->inputs.size(), 3); + if (!IsConstantParameterArray(*model, op->inputs[block_shape_index]) || + !IsConstantParameterArray(*model, op->inputs[paddings_index])) + return false; + + // Handle paddings. + const auto& paddings_array = model->GetArray(op->inputs[paddings_index]); + if (!paddings_array.has_shape()) return false; + const std::vector& paddings_dims = paddings_array.shape().dims(); + if (paddings_dims.size() != 2) { + // Code only handles padding of 2 dimensions. Perhaps another transformation + // will delete this op. + return false; + } + std::vector paddings_buffer = + paddings_array.GetBuffer().data; + for (int i = 0; i < paddings_dims[0]; ++i) { + op->before_paddings.push_back(paddings_buffer[i * 2]); + op->after_paddings.push_back(paddings_buffer[i * 2 + 1]); + } + + // Handle block_shape. + const auto& block_shape_array = + model->GetArray(op->inputs[block_shape_index]); + if (!block_shape_array.has_shape()) return false; + const std::vector& block_shape_dims = block_shape_array.shape().dims(); + CHECK_EQ(block_shape_dims.size(), 1); + std::vector block_shape_buffer = + block_shape_array.GetBuffer().data; + for (int i = 0; i < block_shape_dims[0]; ++i) { + op->block_shape.push_back(block_shape_buffer[i]); + } + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_squeeze.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_squeeze_attributes.cc similarity index 86% rename from tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_squeeze.cc rename to tensorflow/contrib/lite/toco/graph_transformations/resolve_squeeze_attributes.cc index 1d3f42b5ec4cab29189c12043d12ea687d684832..dd3e73635ae0215510f0a8d1aee487da5af35700 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_squeeze.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_squeeze_attributes.cc @@ -25,15 +25,13 @@ limitations under the License. namespace toco { -bool ResolveTensorFlowSqueeze::Run(Model* model, std::size_t op_index) { - const auto squeeze_it = model->operators.begin() + op_index; - const auto* squeeze_op = squeeze_it->get(); +bool ResolveSqueezeAttributes::Run(Model* model, std::size_t op_index) { + auto* squeeze_op = model->operators[op_index].get(); if (squeeze_op->type != OperatorType::kSqueeze) { return false; } - - CHECK_EQ(squeeze_op->inputs.size(), 1); - CHECK_EQ(squeeze_op->outputs.size(), 1); + DCHECK_EQ(squeeze_op->inputs.size(), 1); + DCHECK_EQ(squeeze_op->outputs.size(), 1); // If the output is consumed by a reshape op, it's a trivial squeeze. if (CountOpsWithInput(*model, squeeze_op->outputs[0]) == 1) { @@ -47,7 +45,6 @@ bool ResolveTensorFlowSqueeze::Run(Model* model, std::size_t op_index) { return RemoveTrivialPassthroughOp(this, model, op_index); } } - return false; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_strided_slice_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_strided_slice_attributes.cc index 97946182ef07b0c3d826cafa95b2bb47fbaf0125..7e8b249b07ecca551cbb75afd8007efad0b52eaf 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_strided_slice_attributes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_strided_slice_attributes.cc @@ -12,11 +12,6 @@ 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 -#include -#include -#include - #include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" #include "tensorflow/contrib/lite/toco/model.h" #include "tensorflow/contrib/lite/toco/tooling_util.h" @@ -30,32 +25,43 @@ bool ResolveStridedSliceAttributes::Run(Model* model, std::size_t op_index) { if (slice_op->type != OperatorType::kStridedSlice) return false; auto* op = static_cast(slice_op); - if (!op->start_indices.empty()) return false; + if (!op->start_indices.empty()) { + // We have already resolved these attributes + return false; + } CHECK_EQ(op->inputs.size(), 4); - if (!IsConstantParameterArray(*model, op->inputs[1])) return false; - if (!IsConstantParameterArray(*model, op->inputs[2])) return false; - if (!IsConstantParameterArray(*model, op->inputs[3])) return false; - - const auto& start_array = *model->arrays[op->inputs[1]]; + const auto& start_array = model->GetArray(op->inputs[1]); if (!start_array.has_shape()) return false; - if (toco::RequiredBufferSizeForShape(start_array.shape()) != 4) { - // Only 4D arrays are supported for now. + if (toco::RequiredBufferSizeForShape(start_array.shape()) > 4) { + // Only 1-4D arrays are supported for now. return false; } - const auto& stop_array = *model->arrays[op->inputs[2]]; + const auto& stop_array = model->GetArray(op->inputs[2]); if (!stop_array.has_shape()) return false; - const auto& stride_array = *model->arrays[op->inputs[3]]; + const auto& stride_array = model->GetArray(op->inputs[3]); if (!stride_array.has_shape()) return false; + if (!IsConstantParameterArray(*model, op->inputs[1])) return false; + if (!IsConstantParameterArray(*model, op->inputs[2])) return false; + if (!IsConstantParameterArray(*model, op->inputs[3])) return false; + op->start_indices = start_array.GetBuffer().data; op->stop_indices = stop_array.GetBuffer().data; op->strides = stride_array.GetBuffer().data; - // TODO(dkalenichenko): Delete the extra inputs? + CHECK_GE(op->start_indices.size(), 1); + CHECK_LE(op->start_indices.size(), 4); + CHECK_EQ(op->stop_indices.size(), op->start_indices.size()); + CHECK_EQ(op->strides.size(), op->stop_indices.size()); + // Ideally, we would remove the input arrays after they have been resolved. + // However, we must then reconstitute these input arrays for all supported + // export formats. For now, leave the arrays so we don't have to modify our + // exporters. Ideally, we wouldn't have op attributes, and would work directly + // with the input arrays. return true; } } // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_concat.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_concat.cc index c6723a880ed0e51cc5828f77742a6c8eb70fa864..5c0c1e3478fa0d94104d1b76bab176b98b314c50 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_concat.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_concat.cc @@ -75,7 +75,7 @@ bool ResolveTensorFlowConcat::Run(Model* model, std::size_t op_index) { // Remove the axis array if it is not used by anything else. if (CountOpsWithInput(*model, axis_name) == 1) { - model->arrays.erase(axis_name); + model->EraseArray(axis_name); } // Remove the TensorFlowConcat op model->operators.erase(concat_it); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_matmul.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_matmul.cc index bea7487051a58344a56a3186a05d0fdceebc8727..f38203c80fcb7ab8bc1639129fd98e4e342e5cb7 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_matmul.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_matmul.cc @@ -29,7 +29,36 @@ bool ResolveTensorFlowMatMul::Run(Model* model, std::size_t op_index) { if (matmul_it->get()->type != OperatorType::kTensorFlowMatMul) { return false; } - const auto* matmul_op = matmul_it->get(); + const auto* matmul_op = + static_cast(matmul_it->get()); + + // Reorder the axes on the second input. TensorFlow uses row-major ordering + // on both inputs, however this is inefficient for the FullyConnected + // operator. We'll transpose the second input to be in column-major order now + // and let constant propagation optimize things (if possible). + auto* transpose_op = new TransposeOperator; + transpose_op->inputs = { + matmul_op->inputs[1], + CreateInt32Array( + model, + AvailableArrayName(*model, matmul_op->inputs[1] + "/transpose/perm"), + {1, 0})}; + transpose_op->outputs = { + AvailableArrayName(*model, matmul_op->inputs[1] + "/transpose")}; + model->GetOrCreateArray(transpose_op->outputs[0]); + model->operators.emplace(matmul_it, transpose_op); + + // Refresh iterator. + matmul_it = model->operators.begin(); + for (; matmul_it != model->operators.end(); ++matmul_it) { + if (matmul_it->get() == matmul_op) { + break; + } + } + DCHECK_EQ(matmul_it->get(), matmul_op); + + string input_lhs = matmul_op->inputs[0]; + string input_rhs = transpose_op->outputs[0]; // Find the op producing the array passed to this MatMul auto previous_op_it = model->operators.begin(); @@ -47,38 +76,42 @@ bool ResolveTensorFlowMatMul::Run(Model* model, std::size_t op_index) { } Operator* previous_op = (found) ? previous_op_it->get() : nullptr; - // construct the new FullyConnectedOperator + // Construct the new FullyConnectedOperator. auto* fc_op = new FullyConnectedOperator; fc_op->outputs = matmul_op->outputs; - // insert the newly constructed FullyConnectedOperator - auto fc_it = model->operators.emplace(matmul_it, fc_op); + // Insert the newly constructed FullyConnectedOperator. + model->operators.emplace(matmul_it, fc_op) + 1; - // refresh invalidated iterator - matmul_it = fc_it + 1; + // Refresh iterator. + matmul_it = model->operators.begin(); + for (; matmul_it != model->operators.end(); ++matmul_it) { + if (matmul_it->get() == matmul_op) { + break; + } + } DCHECK_EQ(matmul_it->get(), matmul_op); // The way that TensorFlow encodes FullyConnected ops is as a pair // (Reshape, MatMul), so we want to remove the Reshape op and rewrite the - // MatMul - // op as a FullyConnected. However, TensorFlow skips the Reshape ops if the - // input doesn't need reshaping, so we can't just match (Reshape, MatMul) + // MatMul op as a FullyConnected. However, TensorFlow skips the Reshape ops if + // the input doesn't need reshaping, so we can't just match (Reshape, MatMul) // pairs. if (previous_op && previous_op->type == OperatorType::kTensorFlowReshape) { AddMessageF("Combining %s and %s into %s", LogName(*previous_op), LogName(*matmul_op), LogName(*fc_op)); const auto& previous_op_output = previous_op->outputs[0]; if (CountOpsWithInput(*model, previous_op_output) == 1) { - model->arrays.erase(previous_op_output); + model->EraseArray(previous_op_output); } CHECK_EQ(previous_op->inputs.size(), 2); - fc_op->inputs = {previous_op->inputs[0], matmul_op->inputs[1]}; + input_lhs = previous_op->inputs[0]; // Only remove Reshape node if no other node uses its output. if (CountOpsWithInput(*model, previous_op_output) == 1) { const auto& previous_op_shape = previous_op->inputs[1]; if (CountOpsWithInput(*model, previous_op_shape) == 1 && !GetOpWithOutput(*model, previous_op_shape)) { - model->arrays.erase(previous_op_shape); + model->EraseArray(previous_op_shape); } model->operators.erase(previous_op_it); } @@ -95,9 +128,10 @@ bool ResolveTensorFlowMatMul::Run(Model* model, std::size_t op_index) { } else { AddMessageF("Replacing %s by a FullyConnected operator", LogName(*matmul_op)); - fc_op->inputs = {matmul_op->inputs[0], matmul_op->inputs[1]}; } + fc_op->inputs = {input_lhs, input_rhs}; + // erase the MatMul operator model->operators.erase(matmul_it); return true; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_merge.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_merge.cc index cfa5ce0716523adbfb0a76e89ce3b202f0595763..477e7f13da3d88a68547d494011cd4984936b909 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_merge.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_merge.cc @@ -55,7 +55,7 @@ bool ResolveTensorFlowMerge::Run(Model* model, std::size_t op_index) { // Remove the node and its output array. AddMessageF("Removing already-resolved %s", LogName(*merge_op)); - model->arrays.erase(merge_op->outputs[0]); + model->EraseArray(merge_op->outputs[0]); model->operators.erase(merge_it); return true; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_switch.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_switch.cc index 150cf53da3099227c5c637ee58c44512d5a41d4f..a418073441f1241a5acb1164b36f332828ea2e99 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_switch.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_switch.cc @@ -103,7 +103,7 @@ bool ResolveTensorFlowSwitch::Run(Model* model, std::size_t op_index) { // Remove the output arrays if they are now unused. for (int i = 0; i < 2; i++) { if (!GetOpWithInput(*model, switch_op->outputs[i])) { - model->arrays.erase(switch_op->outputs[i]); + model->EraseArray(switch_op->outputs[i]); } } // Remove input arrays if they are only used by the switch itself and aren't @@ -111,7 +111,7 @@ bool ResolveTensorFlowSwitch::Run(Model* model, std::size_t op_index) { for (const auto& input : switch_op->inputs) { if (CountOpsWithInput(*model, input) == 1 && !GetOpWithOutput(*model, input)) { - model->arrays.erase(input); + model->EraseArray(input); } } // Remove the switch node itself. diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_tile.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_tile.cc index 9f7e7c42a26b60c96573be6653babb78fdb5fd73..1ddf54c778cd1fae7a8fce0ecb97209274e71ac0 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_tile.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_tile.cc @@ -45,10 +45,10 @@ void RemoveTileOperator(Model* model, Operator* tile_op, Operator* binary_op, model->operators.erase(tile_it); if (!CountOpsWithInput(*model, tile_multiplier_array) && !GetOpWithOutput(*model, tile_multiplier_array)) { - model->arrays.erase(tile_multiplier_array); + model->EraseArray(tile_multiplier_array); } if (!CountOpsWithInput(*model, tile_output_array)) { - model->arrays.erase(tile_output_array); + model->EraseArray(tile_output_array); } } } // namespace diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_transpose_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_transpose_attributes.cc new file mode 100644 index 0000000000000000000000000000000000000000..a657ee00af66bd431f96c361e12d5213e203b3df --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_transpose_attributes.cc @@ -0,0 +1,53 @@ +/* Copyright 2017 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 +#include +#include +#include + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool ResolveTransposeAttributes::Run(Model* model, std::size_t op_index) { + const auto op_it = model->operators.begin() + op_index; + if (op_it->get()->type != OperatorType::kTranspose) return false; + + auto* op = static_cast(op_it->get()); + if (!op->perm.empty()) return false; + + CHECK_EQ(op->inputs.size(), 2); + if (!IsConstantParameterArray(*model, op->inputs[1])) return false; + + // Handling perm. + const auto& perm_array = model->GetArray(op->inputs[1]); + if (!perm_array.has_shape()) return false; + + const std::vector& perm_dims = perm_array.shape().dims(); + CHECK_EQ(perm_dims.size(), 1); + + std::vector perm_buffer = + perm_array.GetBuffer().data; + for (int i = 0; i < perm_dims[0]; ++i) { + op->perm.push_back(perm_buffer[i]); + } + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/tests/BUILD b/tensorflow/contrib/lite/toco/graph_transformations/tests/BUILD index 893149878293c9ef2740effe331d3b6c51b49983..2f94f9cd8a9ab24809fb3d137b5d05ab12f43003 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/tests/BUILD +++ b/tensorflow/contrib/lite/toco/graph_transformations/tests/BUILD @@ -18,6 +18,17 @@ tf_cc_test( ], ) +tf_cc_test( + name = "lstm_utils_test", + srcs = ["lstm_utils_test.cc"], + deps = [ + "//tensorflow/contrib/lite/toco:graph_transformations", + "//tensorflow/contrib/lite/toco:model", + "//tensorflow/contrib/lite/toco:tooling_util", + "@com_google_googletest//:gtest_main", + ], +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/contrib/lite/toco/graph_transformations/tests/lstm_utils_test.cc b/tensorflow/contrib/lite/toco/graph_transformations/tests/lstm_utils_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..6aae0775d3445daf7d990bcce09d335c5f686601 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/tests/lstm_utils_test.cc @@ -0,0 +1,442 @@ +/* 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 +#include + +#include +#include +#include "tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" + +namespace toco { + +namespace { + +// A gmock matcher that check that elements of a float vector match to a given +// tolerance. +std::vector> ArrayFloatNear( + const std::vector& values, float max_abs_error = 1e-5) { + std::vector> matchers; + matchers.reserve(values.size()); + for (const float& v : values) { + matchers.emplace_back(testing::FloatNear(v, max_abs_error)); + } + return matchers; +} +} // namespace + +class CopyArrayDataTest : public ::testing::Test { + public: + CopyArrayDataTest() {} + + void PrepareBuffers(Model* model, std::initializer_list src_data, + int src_dim_1, int src_dim_2, + std::initializer_list dst_data, int dst_dim_1, + int dst_dim_2) { + string src_array = "src_array"; + src_buffer_ = CreateFloatArrayBuffer( + model, &src_array, + src_dim_2 == 1 ? Shape({src_dim_1}) : Shape({src_dim_1, src_dim_2})); + PopulateBuffer(src_buffer_, src_data); + string dst_array = "dst_array"; + dst_buffer_ = CreateFloatArrayBuffer( + model, &dst_array, + dst_dim_2 == 1 ? Shape({dst_dim_1}) : Shape({dst_dim_1, dst_dim_2})); + PopulateBuffer(dst_buffer_, dst_data); + } + + Buffer* GetSrcBuffer() { return src_buffer_; } + Buffer* GetDstBuffer() { return dst_buffer_; } + + void PopulateBuffer(Buffer* buffer, + const std::vector& init_data) { + for (int i = 0; i < init_data.size(); i++) { + buffer->data[i] = init_data[i]; + } + } + void UpdateBuffer(Buffer* buffer, + std::initializer_list data) { + buffer->data.resize(data.size()); + PopulateBuffer(buffer, data); + } + + private: + Buffer* src_buffer_; + Buffer* dst_buffer_; +}; + +// Copy from 1 big 2D array to 8 smaller ones. +TEST_F(CopyArrayDataTest, CopyFromBigArrayToSmallerArrayes2D) { + // Init src_buffer, dst_buffer. + Model model; + std::initializer_list large_tf_weight_data = { + -0.320407, -0.108683, 0.406358, -0.410811, -0.285786, -0.15769, + -0.194201, 0.170866, 0.084135, 0.201878, 0.21519, -0.284458, + 0.495906, -0.073818, 0.045578, 0.149816, -0.447073, -0.453578, + 0.116766, 0.21808, 0.047326, -0.001985, 0.402193, 0.315517, + 0.38258, 0.43599, 0.11986, 0.465195, 0.33548, -0.118789, + -0.414159, 0.049269, 0.156108, 0.093459, -0.129103, -0.086274, + 0.186188, -0.324923, 0.4117, -0.344439, 0.240465, -0.343331, + -0.463082, -0.231706, -0.487465, -0.186592, -0.020756, -0.239007, + 0.364817, 0.459106, -0.171447, -0.006542, 0.204032, -0.375317, + -0.041911, 0.051664, 0.320483, 0.155899, 0.156555, -0.249823, + -0.353107, 0.031563, -0.340771, -0.052532, 0.134631, -0.257957, + -0.50141, 0.486939, -0.43853, 0.268426, -0.08754, -0.109447, + -0.502462, -0.028055, -0.121838, -0.046016, 0.105309, -0.070774, + 0.495683, -0.475088, 0.048654, -0.38582, 0.411018, -0.315606, + 0.349628, 0.21698, 0.258989, -0.097902, 0.331218, 0.034602, + 0.418069, -0.089025, -0.417513, 0.07609, 0.393821, 0.404733, + -0.055418, -0.43903, -0.447049, 0.013125, 0.278503, 0.459869, + 0.143755, -0.177335, -0.162247, -0.432371, 0.153714, -0.047403, + -0.446775, -0.418363, 0.019743, 0.042025}; + std::initializer_list tflite_lstm_input_weight = {0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0}; + PrepareBuffers(&model, large_tf_weight_data, /*src_dim_1=*/16, + /*src_dim_2=*/7, tflite_lstm_input_weight, + /*dst_dim_1=*/4, /*dst_dim_2=*/3); + + // Copy src starts at (0,0), size (4,3). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/7, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/3, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/3); + std::vector expected = {-0.320407, -0.108683, 0.406358, 0.170866, + 0.084135, 0.201878, 0.045578, 0.149816, + -0.447073, -0.001985, 0.402193, 0.315517}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy src starts at (4,0), size (4,3). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/7, /*src_start_idx1=*/4, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/3, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/3); + expected = {0.33548, -0.118789, -0.414159, -0.086274, 0.186188, -0.324923, + -0.463082, -0.231706, -0.487465, 0.459106, -0.171447, -0.006542}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy src starts at (8,0), size (4,3). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/7, /*src_start_idx1=*/8, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/3, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/3); + expected = {0.320483, 0.155899, 0.156555, -0.052532, 0.134631, -0.257957, + -0.08754, -0.109447, -0.502462, -0.070774, 0.495683, -0.475088}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy src starts at (12,0), size (4,3). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/7, /*src_start_idx1=*/12, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/3, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/3); + expected = {0.349628, 0.21698, 0.258989, -0.089025, -0.417513, 0.07609, + -0.447049, 0.013125, 0.278503, -0.432371, 0.153714, -0.047403}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // New dst_buffer with size 16. + std::initializer_list tflite_lstm_recurrent_weight = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; + PrepareBuffers(&model, large_tf_weight_data, /*src_dim_1=*/16, + /*src_dim_2=*/7, tflite_lstm_recurrent_weight, + /*dst_dim_1=*/4, /*dst_dim_2=*/4); + + // Copy src starts at (0,3), size (4,4). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/7, /*src_start_idx1=*/0, + /*src_start_idx2=*/3, GetDstBuffer(), /*dst_stride=*/4, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/4); + expected = {-0.410811, -0.285786, -0.15769, -0.194201, 0.21519, -0.284458, + 0.495906, -0.073818, -0.453578, 0.116766, 0.21808, 0.047326, + 0.38258, 0.43599, 0.11986, 0.465195}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy src starts at (4,3), size (4,4). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/7, /*src_start_idx1=*/4, + /*src_start_idx2=*/3, GetDstBuffer(), /*dst_stride=*/4, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/4); + expected = {0.049269, 0.156108, 0.093459, -0.129103, 0.4117, -0.344439, + 0.240465, -0.343331, -0.186592, -0.020756, -0.239007, 0.364817, + 0.204032, -0.375317, -0.041911, 0.051664}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy src starts at (8,3), size (4,4). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/7, /*src_start_idx1=*/8, + /*src_start_idx2=*/3, GetDstBuffer(), /*dst_stride=*/4, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/4); + expected = {-0.249823, -0.353107, 0.031563, -0.340771, -0.50141, 0.486939, + -0.43853, 0.268426, -0.028055, -0.121838, -0.046016, 0.105309, + 0.048654, -0.38582, 0.411018, -0.315606}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy src starts at (12,3), size (4,4). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/7, /*src_start_idx1=*/12, + /*src_start_idx2=*/3, GetDstBuffer(), /*dst_stride=*/4, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/4); + expected = {-0.097902, 0.331218, 0.034602, 0.418069, 0.393821, 0.404733, + -0.055418, -0.43903, 0.459869, 0.143755, -0.177335, -0.162247, + -0.446775, -0.418363, 0.019743, 0.042025}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); +} + +// Copy from 1 big 1D array to 4 small ones. +TEST_F(CopyArrayDataTest, CopyFromBigArrayToSmallerArrayes1D) { + // Init src_buffer, dst_buffer. + Model model; + std::initializer_list large_tf_bias_data = { + 0.980304, 0.419808, 0.080278, 0.728548, 0.581674, 0.672433, + 0.434190, 0.844357, 0.229587, 0.785629, 0.022065, 0.753082, + 0.422080, 0.539481, 0.878386, 0.168965}; + std::initializer_list tflite_lstm_i_bias = {0, 0, 0, 0}; + PrepareBuffers(&model, large_tf_bias_data, /*src_dim_1=*/16, + /*src_dim_2=*/1, tflite_lstm_i_bias, + /*dst_dim_1=*/4, /*dst_dim_2=*/1); + + // Copy starts at (0,), size (4,). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/1, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/1, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/1); + std::vector expected = {0.980304, 0.419808, 0.080278, 0.728548}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy starts at (4,), size (4,). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/1, /*src_start_idx1=*/4, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/1, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/1); + expected = {0.581674, 0.672433, 0.434190, 0.844357}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy starts at (8,), size (4,). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/1, /*src_start_idx1=*/8, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/1, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/1); + expected = {0.229587, 0.785629, 0.022065, 0.753082}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); + + // Copy starts at (12,), size (4,). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/1, /*src_start_idx1=*/12, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/1, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/1); + expected = {0.422080, 0.539481, 0.878386, 0.168965}; + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); +} + +// Copy from 8 small 2D arrayes to 1 big one. +TEST_F(CopyArrayDataTest, CopyFromSmallArrayesToBigArray2D) { + // Init src_buffer, dst_buffer. + Model model; + std::initializer_list large_tf_weights_data = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; + + // Copy dst starts (0, 0), size (4, 3). + std::initializer_list tflite_lstm_i2i_weight = { + -0.320407, -0.108683, 0.406358, 0.170866, 0.084135, 0.201878, + 0.045578, 0.149816, -0.447073, -0.001985, 0.402193, 0.315517}; + PrepareBuffers(&model, tflite_lstm_i2i_weight, /*src_dim_1=*/4, + /*src_dim_2=*/3, large_tf_weights_data, + /*dst_dim_1=*/16, /*dst_dim_2=*/7); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/3, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/7, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/3); + + // Copy dst starts (4, 0), size (4, 3). + std::initializer_list tflite_lstm_i2c_weight = { + 0.33548, -0.118789, -0.414159, -0.086274, 0.186188, -0.324923, + -0.463082, -0.231706, -0.487465, 0.459106, -0.171447, -0.006542}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_i2c_weight); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/3, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/7, + /*dst_start_idx1=*/4, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/3); + + // Copy dst starts (8, 0), size (4, 3). + std::initializer_list tflite_lstm_i2f_weight = { + 0.320483, 0.155899, 0.156555, -0.052532, 0.134631, -0.257957, + -0.08754, -0.109447, -0.502462, -0.070774, 0.495683, -0.475088}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_i2f_weight); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/3, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/7, + /*dst_start_idx1=*/8, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/3); + + // Copy dst starts (12, 0), size (4, 3). + std::initializer_list tflite_lstm_i2o_weight = { + 0.349628, 0.21698, 0.258989, -0.089025, -0.417513, 0.07609, + -0.447049, 0.013125, 0.278503, -0.432371, 0.153714, -0.047403}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_i2o_weight); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/3, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/7, + /*dst_start_idx1=*/12, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/3); + + // Copy dst starts (0, 3), size (4, 4). + std::initializer_list tflite_lstm_i2r_weight = { + -0.410811, -0.285786, -0.15769, -0.194201, 0.21519, -0.284458, + 0.495906, -0.073818, -0.453578, 0.116766, 0.21808, 0.047326, + 0.38258, 0.43599, 0.11986, 0.465195}; + UpdateBuffer(GetSrcBuffer(), tflite_lstm_i2r_weight); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/4, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/7, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/3, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/4); + + // Copy dst starts (4, 3), size (4, 4). + std::initializer_list tflite_lstm_c2r_weight = { + 0.049269, 0.156108, 0.093459, -0.129103, 0.4117, -0.344439, + 0.240465, -0.343331, -0.186592, -0.020756, -0.239007, 0.364817, + 0.204032, -0.375317, -0.041911, 0.051664}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_c2r_weight); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/4, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/7, + /*dst_start_idx1=*/4, /*dst_start_idx2=*/3, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/4); + + // Copy dst starts (8, 3), size (4, 4). + std::initializer_list tflite_lstm_f2r_weight = { + -0.249823, -0.353107, 0.031563, -0.340771, -0.50141, 0.486939, + -0.43853, 0.268426, -0.028055, -0.121838, -0.046016, 0.105309, + 0.048654, -0.38582, 0.411018, -0.315606}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_f2r_weight); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/4, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/7, + /*dst_start_idx1=*/8, /*dst_start_idx2=*/3, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/4); + + // Copy dst starts (12, 3), size (4, 4). + std::initializer_list tflite_lstm_o2r_weight = { + -0.097902, 0.331218, 0.034602, 0.418069, 0.393821, 0.404733, + -0.055418, -0.43903, 0.459869, 0.143755, -0.177335, -0.162247, + -0.446775, -0.418363, 0.019743, 0.042025}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_o2r_weight); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/4, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/7, + /*dst_start_idx1=*/12, /*dst_start_idx2=*/3, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/4); + + std::vector expected = { + -0.320407, -0.108683, 0.406358, -0.410811, -0.285786, -0.15769, + -0.194201, 0.170866, 0.084135, 0.201878, 0.21519, -0.284458, + 0.495906, -0.073818, 0.045578, 0.149816, -0.447073, -0.453578, + 0.116766, 0.21808, 0.047326, -0.001985, 0.402193, 0.315517, + 0.38258, 0.43599, 0.11986, 0.465195, 0.33548, -0.118789, + -0.414159, 0.049269, 0.156108, 0.093459, -0.129103, -0.086274, + 0.186188, -0.324923, 0.4117, -0.344439, 0.240465, -0.343331, + -0.463082, -0.231706, -0.487465, -0.186592, -0.020756, -0.239007, + 0.364817, 0.459106, -0.171447, -0.006542, 0.204032, -0.375317, + -0.041911, 0.051664, 0.320483, 0.155899, 0.156555, -0.249823, + -0.353107, 0.031563, -0.340771, -0.052532, 0.134631, -0.257957, + -0.50141, 0.486939, -0.43853, 0.268426, -0.08754, -0.109447, + -0.502462, -0.028055, -0.121838, -0.046016, 0.105309, -0.070774, + 0.495683, -0.475088, 0.048654, -0.38582, 0.411018, -0.315606, + 0.349628, 0.21698, 0.258989, -0.097902, 0.331218, 0.034602, + 0.418069, -0.089025, -0.417513, 0.07609, 0.393821, 0.404733, + -0.055418, -0.43903, -0.447049, 0.013125, 0.278503, 0.459869, + 0.143755, -0.177335, -0.162247, -0.432371, 0.153714, -0.047403, + -0.446775, -0.418363, 0.019743, 0.042025}; + + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); +} + +// Copy from 4 small 1D arrayes to 1 big one. +TEST_F(CopyArrayDataTest, CopyFromSmallArrayesToBigArray1D) { + // Init src_buffer, dst_buffer. + Model model; + std::initializer_list large_tf_bias_data = {0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0}; + + std::initializer_list tflite_lstm_i_bias = {0.980304, 0.419808, + 0.080278, 0.728548}; + + PrepareBuffers(&model, tflite_lstm_i_bias, /*src_dim_1=*/4, + /*src_dim_2=*/1, large_tf_bias_data, + /*dst_dim_1=*/16, /*dst_dim_2=*/1); + + // Copy starts at (0,), size (4,). + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/1, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/1, + /*dst_start_idx1=*/0, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/1); + + // Copy starts at (4,), size (4,). + std::initializer_list tflite_lstm_cell_bias = {0.581674, 0.672433, + 0.434190, 0.844357}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_cell_bias); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/1, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/1, + /*dst_start_idx1=*/4, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/1); + + // Copy starts at (8,0), size (4,). + std::initializer_list tflite_lstm_forget_bias = {0.229587, 0.785629, + 0.022065, 0.753082}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_forget_bias); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/1, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/1, + /*dst_start_idx1=*/8, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/1); + + // Copy starts at (12,), size (4,). + std::initializer_list tflite_lstm_output_bias = {0.422080, 0.539481, + 0.878386, 0.168965}; + PopulateBuffer(GetSrcBuffer(), tflite_lstm_output_bias); + CopyArrayData(*(GetSrcBuffer()), + /*src_stride=*/1, /*src_start_idx1=*/0, + /*src_start_idx2=*/0, GetDstBuffer(), /*dst_stride=*/1, + /*dst_start_idx1=*/12, /*dst_start_idx2=*/0, + /*dim1_copy_size=*/4, /*dim2_copy_size=*/1); + + std::vector expected = {0.980304, 0.419808, 0.080278, 0.728548, + 0.581674, 0.672433, 0.434190, 0.844357, + 0.229587, 0.785629, 0.022065, 0.753082, + 0.422080, 0.539481, 0.878386, 0.168965}; + + EXPECT_THAT(GetDstBuffer()->data, ElementsAreArray(ArrayFloatNear(expected))); +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc b/tensorflow/contrib/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc index a14016e8e2705a66c392118899335eb3997fa1de..3a1d175b9823f085c9b8730caba8bedd7eb87d52 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc @@ -19,7 +19,6 @@ limitations under the License. #include #include -//#include "tensorflow/contrib/lite/kernels/test_util.h" #include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" #include "tensorflow/contrib/lite/toco/model.h" #include "tensorflow/contrib/lite/toco/tooling_util.h" @@ -168,11 +167,11 @@ TEST_F(ResolveConstantConcatenationTest, ConcatAtAxis0) { GraphTransformationsSet graph_transformation_set; graph_transformation_set.Add(new toco::ResolveConstantConcatenation); - EXPECT_THAT(model.arrays.size(), 5); + EXPECT_THAT(model.GetArrayMap().size(), 5); (*graph_transformation_set.begin())->Run(&model, /*op_index=*/0); - EXPECT_THAT(model.arrays.size(), 1); + EXPECT_THAT(model.GetArrayMap().size(), 1); - auto& concatenated_array = (*model.arrays.begin()).second; + auto& concatenated_array = (*model.GetArrayMap().begin()).second; EXPECT_THAT(concatenated_array->GetBuffer().data, ElementsAreArray(ArrayFloatNear( {0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., @@ -187,11 +186,11 @@ TEST_F(ResolveConstantConcatenationTest, ConcatAtAxis1) { GraphTransformationsSet graph_transformation_set; graph_transformation_set.Add(new toco::ResolveConstantConcatenation); - EXPECT_THAT(model.arrays.size(), 5); + EXPECT_THAT(model.GetArrayMap().size(), 5); (*graph_transformation_set.begin())->Run(&model, /*op_index=*/0); - EXPECT_THAT(model.arrays.size(), 1); + EXPECT_THAT(model.GetArrayMap().size(), 1); - auto& concatenated_array = (*model.arrays.begin()).second; + auto& concatenated_array = (*model.GetArrayMap().begin()).second; EXPECT_THAT(concatenated_array->GetBuffer().data, ElementsAreArray(ArrayFloatNear( {0., 1., 2., 3., 10., 11., 12., 13., 20., 21., 22., @@ -206,11 +205,11 @@ TEST_F(ResolveConstantConcatenationTest, ConcatAtAxis2) { GraphTransformationsSet graph_transformation_set; graph_transformation_set.Add(new toco::ResolveConstantConcatenation); - EXPECT_THAT(model.arrays.size(), 5); + EXPECT_THAT(model.GetArrayMap().size(), 5); (*graph_transformation_set.begin())->Run(&model, /*op_index=*/0); - EXPECT_THAT(model.arrays.size(), 1); + EXPECT_THAT(model.GetArrayMap().size(), 1); - auto& concatenated_array = (*model.arrays.begin()).second; + auto& concatenated_array = (*model.GetArrayMap().begin()).second; EXPECT_THAT(concatenated_array->GetBuffer().data, ElementsAreArray(ArrayFloatNear( {0., 1., 10., 11., 20., 21., 30., 31., 2., 3., 12., diff --git a/tensorflow/contrib/lite/toco/graph_transformations/unfuse_activation_functions.cc b/tensorflow/contrib/lite/toco/graph_transformations/unfuse_activation_functions.cc index 4e273343df9f3e5ade8f23a2fbd868bcab72c62e..2c7046c8c77c94a89fc05a26d7d72b3661380475 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/unfuse_activation_functions.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/unfuse_activation_functions.cc @@ -63,7 +63,7 @@ bool UnfuseActivationFunctions::Run(Model* model, std::size_t op_index) { ac_op->outputs = op->outputs; const string& tmp_array_name = AvailableArrayName(*model, op->outputs[0] + "_unfused"); - CHECK(!model->arrays.count(tmp_array_name)); + CHECK(!model->HasArray(tmp_array_name)); model->GetOrCreateArray(tmp_array_name); ac_op->inputs = {tmp_array_name}; op->outputs = {tmp_array_name}; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/unroll_batch_matmul.cc b/tensorflow/contrib/lite/toco/graph_transformations/unroll_batch_matmul.cc new file mode 100644 index 0000000000000000000000000000000000000000..da81ea2ff3b4ab0bee0550874a9c4ea1044a3579 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/unroll_batch_matmul.cc @@ -0,0 +1,172 @@ +/* Copyright 2017 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 +#include +#include +#include + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +// Unrolls a BatchMatMul on the batch dimension. +// We need to slice each batch out of the inputs, matmul them individually, then +// stack them all back together at the end. +// +// This transform effectively looks like: +// result_slices = [] +// for bat in B: +// slice_a = tf.reshape(tf.slice(a, [bat, 0, 0], [1, M, N]), [M, N]) +// slice_b = tf.reshape(tf.slice(b, [bat, 0, 0], [1, M, N]), [M, N]) +// slice_c = tf.matmul(slice_a, slice_b) +// result_slices[bat] = slice_c +// result = tf.stack(result_slices) +bool UnrollBatchMatMul::Run(Model* model, std::size_t op_index) { + auto batch_op_it = model->operators.begin() + op_index; + if (batch_op_it->get()->type != OperatorType::kBatchMatMul) { + return false; + } + const auto* batch_op = + static_cast(batch_op_it->get()); + + // We must have the shape of at least one input to know our batch size. + const auto& input_array_a = model->GetArray(batch_op->inputs[0]); + const auto& input_array_b = model->GetArray(batch_op->inputs[1]); + if (!input_array_a.has_shape() || !input_array_b.has_shape()) return false; + + // We only support the rank 3 case. If you are batching on rank > 3 you'll + // have to figure that out. + CHECK_EQ(input_array_a.shape().dimensions_count(), + input_array_b.shape().dimensions_count()) + << "Input dimensions must have the same rank"; + if (input_array_a.shape().dimensions_count() == 2) { + // This is really just a MatMul. This likely means that someone hand-crafted + // a graphdef with a BatchMatMul when they really wanted a MatMul. + AddMessageF("Replacing non-batch BatchMatMul %s by a MatMul operator", + LogName(*batch_op)); + auto* matmul_op = new TensorFlowMatMulOperator; + matmul_op->inputs = batch_op->inputs; + matmul_op->outputs = batch_op->outputs; + const auto matmul_op_it = model->operators.emplace(batch_op_it, matmul_op); + batch_op_it = matmul_op_it + 1; + CHECK_EQ(batch_op_it->get(), batch_op); + model->operators.erase(batch_op_it); + return true; + } + CHECK_EQ(input_array_a.shape().dimensions_count(), 3) + << "Input arrays must have rank 3"; + + // Perform the matmul for each slice of the batch. + int batch_count = input_array_a.shape().dims(0); + AddMessageF("Unrolling BatchMatMul %s %d times", LogName(*batch_op), + batch_count); + auto tail_it = batch_op_it; + std::vector stack_inputs; + for (int batch = 0; batch < batch_count; ++batch) { + std::string batch_name = + std::string(batch_op->outputs[0]) + "_b" + std::to_string(batch); + + // tf.slice(a, ...). + auto* slice_a_op = new SliceOperator; + slice_a_op->inputs = { + batch_op->inputs[0], + CreateInt32Array(model, batch_name + "/slice_a/slice/begin", + {batch, 0, 0}), + CreateInt32Array( + model, batch_name + "/slice_a/slice/size", + {1, input_array_a.shape().dims(1), input_array_a.shape().dims(2)}), + }; + slice_a_op->outputs = {AvailableArrayName(*model, batch_name + "/slice_a")}; + auto& slice_a_op_output = model->GetOrCreateArray(slice_a_op->outputs[0]); + slice_a_op_output.data_type = input_array_a.data_type; + tail_it = model->operators.emplace(tail_it, slice_a_op) + 1; + + // Reshape to remove the first dimension ([1,M,N] -> [M,N]). + auto* slice_a_reshape_op = new TensorFlowReshapeOperator; + slice_a_reshape_op->inputs = { + slice_a_op->outputs[0], + CreateInt32Array(model, batch_name + "/slice_a/reshape/shape", + {-1, input_array_a.shape().dims(2)})}; + slice_a_reshape_op->outputs = { + AvailableArrayName(*model, batch_name + "/slice_a/reshape")}; + auto& slice_a_reshape_op_output = + model->GetOrCreateArray(slice_a_reshape_op->outputs[0]); + slice_a_reshape_op_output.data_type = input_array_a.data_type; + tail_it = model->operators.emplace(tail_it, slice_a_reshape_op) + 1; + + // tf.slice(b, ...). + auto* slice_b_op = new SliceOperator; + slice_b_op->inputs = { + batch_op->inputs[1], + CreateInt32Array(model, batch_name + "/slice_b/slice/begin", {0, 0, 0}), + CreateInt32Array( + model, batch_name + "/slice_b/slice/size", + {1, input_array_b.shape().dims(1), input_array_b.shape().dims(2)}), + }; + slice_b_op->outputs = {AvailableArrayName(*model, batch_name + "/slice_b")}; + auto& slice_b_op_output = model->GetOrCreateArray(slice_b_op->outputs[0]); + slice_b_op_output.data_type = input_array_b.data_type; + tail_it = model->operators.emplace(tail_it, slice_b_op) + 1; + + // Reshape to remove the first dimension ([1,M,N] -> [M,N]). + auto* slice_b_reshape_op = new TensorFlowReshapeOperator; + slice_b_reshape_op->inputs = { + slice_b_op->outputs[0], + CreateInt32Array(model, batch_name + "/slice_b/reshape/shape", + {-1, input_array_b.shape().dims(2)})}; + slice_b_reshape_op->outputs = { + AvailableArrayName(*model, batch_name + "/slice_b/reshape")}; + auto& slice_b_reshape_op_output = + model->GetOrCreateArray(slice_b_reshape_op->outputs[0]); + slice_b_reshape_op_output.data_type = input_array_b.data_type; + tail_it = model->operators.emplace(tail_it, slice_b_reshape_op) + 1; + + // tf.matmul(slice_a, slice_b). + auto* matmul_op = new TensorFlowMatMulOperator; + matmul_op->inputs = {slice_a_reshape_op->outputs[0], + slice_b_reshape_op->outputs[0]}; + matmul_op->outputs = {AvailableArrayName(*model, batch_name)}; + auto& matmul_op_output = model->GetOrCreateArray(matmul_op->outputs[0]); + matmul_op_output.data_type = input_array_a.data_type; + tail_it = model->operators.emplace(tail_it, matmul_op) + 1; + + // Add to stack. + stack_inputs.push_back(matmul_op->outputs[0]); + } + + // The stack that will join all the individual matmul results together. + auto* stack_op = new StackOperator; + stack_op->inputs = stack_inputs; + stack_op->outputs = {batch_op->outputs[0]}; + stack_op->axis = 0; + model->operators.emplace(tail_it, stack_op); + + // Remove the old batch matmul now that we've unrolled. + batch_op_it = model->operators.begin(); + for (; batch_op_it != model->operators.end(); ++batch_op_it) { + if (batch_op_it->get() == batch_op) { + break; + } + } + CHECK(batch_op_it != model->operators.end()); + CHECK(batch_op_it->get() == batch_op); + model->operators.erase(batch_op_it); + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc index ce38ea475414f6f29d16754e8ab614e21770fac5..9c01b67420603a0d3c0e095dafe6a3359f2514b5 100644 --- a/tensorflow/contrib/lite/toco/import_tensorflow.cc +++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc @@ -21,6 +21,7 @@ limitations under the License. #include "google/protobuf/map.h" #include "google/protobuf/text_format.h" +#include "absl/memory/memory.h" #include "absl/strings/match.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" @@ -52,6 +53,7 @@ using tensorflow::DT_BOOL; using tensorflow::DT_FLOAT; using tensorflow::DT_INT32; using tensorflow::DT_INT64; +using tensorflow::DT_QUINT8; using tensorflow::DT_STRING; using tensorflow::DT_UINT8; using tensorflow::GraphDef; @@ -172,8 +174,13 @@ void ImportFloatArray(const TensorProto& input_tensor, Array* output_array) { } auto& output_float_data = output_array->GetMutableBuffer().data; - output_float_data.resize(input_flat_size); - if (input_tensor.float_val_size()) { + output_float_data.resize(RequiredBufferSizeForShape(output_array->shape()), + 0.f); + if (input_tensor.float_val_size() == 1) { + for (int i = 0; i < input_flat_size; i++) { + output_float_data[i] = input_tensor.float_val(0); + } + } else if (input_tensor.float_val_size() == input_flat_size) { for (int i = 0; i < input_tensor.float_val_size(); i++) { output_float_data[i] = input_tensor.float_val(i); } @@ -187,6 +194,32 @@ void ImportFloatArray(const TensorProto& input_tensor, Array* output_array) { } } +void ImportQuint8Array(const TensorProto& input_tensor, Array* output_array) { + CHECK_EQ(input_tensor.dtype(), DT_QUINT8); + const auto& input_shape = input_tensor.tensor_shape(); + CHECK_LE(input_shape.dim_size(), 4); + ImportShape(input_shape.dim(), output_array->mutable_shape()); + int input_flat_size = 1; + for (int k = 0; k < input_shape.dim_size(); k++) { + input_flat_size *= input_shape.dim(k).size(); + } + auto& output_int_data = + output_array->GetMutableBuffer().data; + output_int_data.resize(RequiredBufferSizeForShape(output_array->shape()), 0); + if (input_tensor.int_val_size()) { + for (int i = 0; i < input_tensor.int_val_size(); i++) { + output_int_data[i] = input_tensor.int_val(i); + } + } else if (input_tensor.tensor_content().size() == + input_flat_size * sizeof(uint8_t)) { + toco::port::CopyToBuffer(input_tensor.tensor_content(), + reinterpret_cast(output_int_data.data())); + } else { + LOG(FATAL) << "Neither input_content nor int_val have the right " + "dimensions for this uint8 tensor."; + } +} + void ImportInt32Array(const TensorProto& input_tensor, Array* output_array) { CHECK_EQ(input_tensor.dtype(), DT_INT32); const auto& input_shape = input_tensor.tensor_shape(); @@ -198,7 +231,7 @@ void ImportInt32Array(const TensorProto& input_tensor, Array* output_array) { } auto& output_int_data = output_array->GetMutableBuffer().data; - output_int_data.resize(input_flat_size); + output_int_data.resize(RequiredBufferSizeForShape(output_array->shape()), 0); if (input_tensor.int_val_size()) { for (int i = 0; i < input_tensor.int_val_size(); i++) { output_int_data[i] = input_tensor.int_val(i); @@ -224,7 +257,7 @@ void ImportInt64Array(const TensorProto& input_tensor, Array* output_array) { } auto& output_int_data = output_array->GetMutableBuffer().data; - output_int_data.resize(input_flat_size); + output_int_data.resize(RequiredBufferSizeForShape(output_array->shape()), 0); if (input_tensor.int64_val_size()) { for (int i = 0; i < input_tensor.int64_val_size(); i++) { output_int_data[i] = input_tensor.int64_val(i); @@ -250,7 +283,7 @@ void ImportStringArray(const TensorProto& input_tensor, Array* output_array) { } auto& output_string_data = output_array->GetMutableBuffer().data; - output_string_data.resize(input_flat_size); + output_string_data.resize(RequiredBufferSizeForShape(output_array->shape())); if (input_flat_size != input_tensor.string_val_size()) { LOG(FATAL) << "Input_content string_val doesn't have the right " "dimensions for this string tensor."; @@ -277,6 +310,14 @@ int GetInputsCount(const NodeDef& node, } } +void CheckInputsCount(const NodeDef& node, + const TensorFlowImportFlags& tf_import_flags, + int expected_input_count) { + QCHECK_EQ(GetInputsCount(node, tf_import_flags), expected_input_count) + << node.op() << " node expects " << expected_input_count + << " input(s) other than control dependencies: " << node.DebugString(); +} + void ConvertConstOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { @@ -294,6 +335,10 @@ void ConvertConstOperator(const NodeDef& node, array.data_type = ArrayDataType::kInt32; ImportInt32Array(tensor, &array); break; + case DT_QUINT8: + array.data_type = ArrayDataType::kUint8; + ImportQuint8Array(tensor, &array); + break; case DT_INT64: array.data_type = ArrayDataType::kInt64; ImportInt64Array(tensor, &array); @@ -316,7 +361,7 @@ void ConvertConvOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Conv2D"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); // We only support NHWC, which is the default data_format. // So if data_format is not defined, we're all good. @@ -369,7 +414,7 @@ void ConvertDepthwiseConvOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "DepthwiseConv2dNative"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); // We only support NHWC, which is the default data_format. // So if data_format is not defined, we're all good. @@ -422,7 +467,8 @@ void ConvertDepthToSpaceOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "DepthToSpace"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); + CHECK_EQ(GetDataTypeAttr(node, "T"), DT_FLOAT); auto* op = new DepthToSpaceOperator; op->inputs.push_back(node.input(0)); @@ -436,7 +482,8 @@ void ConvertSpaceToDepthOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "SpaceToDepth"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); + CHECK_EQ(GetDataTypeAttr(node, "T"), DT_FLOAT); auto* op = new SpaceToDepthOperator; op->inputs.push_back(node.input(0)); @@ -450,7 +497,8 @@ void ConvertBiasAddOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "BiasAdd"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); + const auto& input_name = node.input(0); const auto& bias_name = node.input(1); CHECK_EQ(GetDataTypeAttr(node, "T"), DT_FLOAT); @@ -465,7 +513,7 @@ void ConvertReluOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Relu"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); const auto& input_name = node.input(0); auto* relu = new ReluOperator; relu->inputs.push_back(input_name); @@ -477,7 +525,8 @@ void ConvertRelu6Operator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Relu6"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); + const auto& input_name = node.input(0); auto* op = new Relu6Operator; op->inputs.push_back(input_name); @@ -489,7 +538,8 @@ void ConvertLogisticOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Sigmoid"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); + const auto& input_name = node.input(0); auto* op = new LogisticOperator; op->inputs.push_back(input_name); @@ -501,7 +551,8 @@ void ConvertTanhOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Tanh"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); + const auto& input_name = node.input(0); auto* op = new TanhOperator; op->inputs.push_back(input_name); @@ -513,7 +564,7 @@ void ConvertDivOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK(node.op() == "Div" || node.op() == "RealDiv"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new DivOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -532,7 +583,10 @@ void ConvertIdentityOperator(const NodeDef& node, // to be gratuitous (in the case of rajeev_lstm.pb, these are // enumerating the LSTM state arrays). We will just ignore extra // inputs beyond the first input. - CHECK_GE(node.input_size(), 1); + QCHECK_GE(node.input_size(), 1) + << node.op() + << " node expects at least 1 input other than control dependencies: " + << node.DebugString(); const auto& input_name = node.input(0); op->inputs.push_back(input_name); op->outputs.push_back(node.name()); @@ -543,7 +597,7 @@ void ConvertFakeQuantWithMinMaxArgs( const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "FakeQuantWithMinMaxArgs"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); auto* op = new FakeQuantOperator; op->inputs.push_back(node.input(0)); op->minmax.reset(new MinMax); @@ -559,7 +613,10 @@ void ConvertFakeQuantWithMinMaxVars( Model* model) { CHECK_EQ(node.op(), "FakeQuantWithMinMaxVars"); const int num_inputs = GetInputsCount(node, tf_import_flags); - CHECK(num_inputs == 3 || num_inputs == 4); + QCHECK(num_inputs == 3 || num_inputs == 4) + << "FakeQuantWithMinMaxVars node expects 3 or 4 inputs other than " + "control dependencies: " + << node.DebugString(); auto* op = new FakeQuantOperator; for (int i = 0; i < 3; i++) { op->inputs.push_back(node.input(i)); @@ -572,7 +629,7 @@ void ConvertNegOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Neg"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); auto* op = new NegOperator; op->inputs.push_back(node.input(0)); op->outputs.push_back(node.name()); @@ -583,7 +640,7 @@ void ConvertRsqrtOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Rsqrt"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); auto* op = new TensorFlowRsqrtOperator; op->inputs.push_back(node.input(0)); op->outputs.push_back(node.name()); @@ -594,7 +651,7 @@ void ConvertSqrtOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Sqrt"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); auto* op = new TensorFlowSqrtOperator; op->inputs.push_back(node.input(0)); op->outputs.push_back(node.name()); @@ -605,7 +662,7 @@ void ConvertSqueezeOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Squeeze"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); auto* op = new SqueezeOperator; op->inputs.push_back(node.input(0)); op->outputs.push_back(node.name()); @@ -622,7 +679,7 @@ void ConvertSquareOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Square"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); auto* op = new TensorFlowSquareOperator; op->inputs.push_back(node.input(0)); op->outputs.push_back(node.name()); @@ -633,7 +690,7 @@ void ConvertAddOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Add"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new AddOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -641,11 +698,24 @@ void ConvertAddOperator(const NodeDef& node, model->operators.emplace_back(op); } +void ConvertAddNOperator(const NodeDef& node, + const TensorFlowImportFlags& tf_import_flags, + Model* model) { + CHECK_EQ(node.op(), "AddN"); + const int num_inputs = GetInputsCount(node, tf_import_flags); + auto* op = new AddNOperator; + for (int i = 0; i < num_inputs; ++i) { + op->inputs.push_back(node.input(i)); + } + op->outputs.push_back(node.name()); + model->operators.emplace_back(op); +} + void ConvertMulOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Mul"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new MulOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -657,7 +727,7 @@ void ConvertSubOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Sub"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new SubOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -669,7 +739,7 @@ void ConvertSumOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Sum"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowSumOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -684,7 +754,7 @@ void ConvertTileOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Tile"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowTileOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -696,7 +766,7 @@ void ConvertSliceOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Slice"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 3); + CheckInputsCount(node, tf_import_flags, 3); auto* op = new SliceOperator; for (int i = 0; i < 3; ++i) { op->inputs.push_back(node.input(i)); @@ -709,7 +779,7 @@ void ConvertPadOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Pad"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new PadOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -721,7 +791,7 @@ void ConvertShapeOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Shape"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); auto* op = new TensorFlowShapeOperator; op->inputs.push_back(node.input(0)); op->outputs.push_back(node.name()); @@ -732,7 +802,7 @@ void ConvertSplitOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Split"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowSplitOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -749,7 +819,7 @@ void ConvertMergeOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Merge"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowMergeOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -761,7 +831,7 @@ void ConvertSwitchOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Switch"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowSwitchOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -770,11 +840,12 @@ void ConvertSwitchOperator(const NodeDef& node, op->outputs.push_back(node.name() + ":1"); model->operators.emplace_back(op); } + void ConvertSoftmaxOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Softmax"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); const auto& input_name = node.input(0); auto* softmax = new SoftmaxOperator; softmax->inputs.push_back(input_name); @@ -785,11 +856,23 @@ void ConvertSoftmaxOperator(const NodeDef& node, model->operators.emplace_back(softmax); } +void ConvertLogSoftmaxOperator(const NodeDef& node, + const TensorFlowImportFlags& tf_import_flags, + Model* model) { + CHECK_EQ(node.op(), "LogSoftmax"); + CheckInputsCount(node, tf_import_flags, 1); + const auto& input_name = node.input(0); + auto* log_softmax = new LogSoftmaxOperator; + log_softmax->inputs.push_back(input_name); + log_softmax->outputs.push_back(node.name()); + model->operators.emplace_back(log_softmax); +} + void ConvertLRNOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "LRN"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); const auto& input_name = node.input(0); auto* lrn = new LocalResponseNormalizationOperator; lrn->inputs.push_back(input_name); @@ -805,7 +888,7 @@ void ConvertMaxPoolOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "MaxPool"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); const auto& input_name = node.input(0); // We only support NHWC, which is the default data_format. // So if data_format is not defined, we're all good. @@ -847,7 +930,7 @@ void ConvertAvgPoolOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "AvgPool"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); const auto& input_name = node.input(0); // We only support NHWC, which is the default data_format. // So if data_format is not defined, we're all good. @@ -885,7 +968,7 @@ void ConvertReshapeOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Reshape"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowReshapeOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -893,49 +976,37 @@ void ConvertReshapeOperator(const NodeDef& node, model->operators.emplace_back(op); } +void ConvertBatchMatMulOperator(const NodeDef& node, + const TensorFlowImportFlags& tf_import_flags, + Model* model) { + CheckInputsCount(node, tf_import_flags, 2); + + // https://www.tensorflow.org/versions/r0.12/api_docs/python/math_ops/matrix_math_functions + CHECK(!HasAttr(node, "adj_a") || (GetBoolAttr(node, "adj_a") == false)); + CHECK(!HasAttr(node, "adj_b") || (GetBoolAttr(node, "adj_b") == false)); + + auto* batch_matmul = new BatchMatMulOperator; + batch_matmul->inputs = {node.input(0), node.input(1)}; + batch_matmul->outputs = {node.name()}; + model->operators.emplace_back(batch_matmul); +} + void ConvertMatMulOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); - if (node.op() == "MatMul") { - // Transpose flags should be easy to support, but we don't have a - // GraphDef with them to test on at the moment. - CHECK_EQ(GetBoolAttr(node, "transpose_a"), false); - CHECK_EQ(GetBoolAttr(node, "transpose_b"), false); - CHECK(!HasAttr(node, "adjoint_a") || - (GetBoolAttr(node, "adjoint_a") == false)); - CHECK(!HasAttr(node, "adjoint_b") || - (GetBoolAttr(node, "adjoint_b") == false)); - } else if (node.op() == "BatchMatMul") { - // https://www.tensorflow.org/versions/r0.12/api_docs/python/math_ops/matrix_math_functions - CHECK(!HasAttr(node, "adj_a") || (GetBoolAttr(node, "adj_a") == false)); - CHECK(!HasAttr(node, "adj_b") || (GetBoolAttr(node, "adj_b") == false)); - } else { - LOG(FATAL) << "op must be 'MatMul' or 'BatchMatMul'"; - } + CheckInputsCount(node, tf_import_flags, 2); + + // Transpose flags should be easy to support, but we don't have a + // GraphDef with them to test on at the moment. + CHECK_EQ(GetBoolAttr(node, "transpose_a"), false); + CHECK_EQ(GetBoolAttr(node, "transpose_b"), false); + CHECK(!HasAttr(node, "adjoint_a") || + (GetBoolAttr(node, "adjoint_a") == false)); + CHECK(!HasAttr(node, "adjoint_b") || + (GetBoolAttr(node, "adjoint_b") == false)); - const auto& input_name = node.input(0); - const auto& weights_name = node.input(1); - const auto& reordered_weights_name = weights_name + "_reordered"; - // Check if a ReorderAxesOperator was already created for these weights - // (that happens when multiple layers share the same weights). - const Operator* existing_reorder = - GetOpWithOutput(*model, reordered_weights_name); - if (existing_reorder) { - // Check that it is safe to rely on the _reordered naming of the output - // array! - CHECK(existing_reorder->type == OperatorType::kReorderAxes); - } else { - // Create a new ReorderAxesOperator - auto* reorder = new ReorderAxesOperator; - reorder->inputs = {weights_name}; - reorder->outputs = {reordered_weights_name}; - reorder->input_axes_order = AxesOrder::kRC; - reorder->output_axes_order = AxesOrder::kCR; - model->operators.emplace_back(reorder); - } auto* matmul = new TensorFlowMatMulOperator; - matmul->inputs = {input_name, reordered_weights_name}; + matmul->inputs = {node.input(0), node.input(1)}; matmul->outputs = {node.name()}; model->operators.emplace_back(matmul); } @@ -952,7 +1023,10 @@ void ConvertConcatOperator(const NodeDef& node, LOG(FATAL) << "Expected Concat or ConcatV2"; } const int num_inputs = GetInputsCount(node, tf_import_flags); - CHECK_GE(num_inputs, 2); + QCHECK_GE(num_inputs, 2) + << node.op() + << " node expects at least 2 inputs other than control dependencies: " + << node.DebugString(); CHECK_EQ(num_inputs, 1 + GetIntAttr(node, "N")); for (int i = 0; i < num_inputs; ++i) { op->inputs.push_back(node.input(i)); @@ -1043,7 +1117,7 @@ void ConvertMaxOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Max"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowMaxOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1058,7 +1132,7 @@ void ConvertMinOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Min"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowMinOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1073,7 +1147,7 @@ void ConvertMaximumOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Maximum"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowMaximumOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1085,7 +1159,7 @@ void ConvertMinimumOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Minimum"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TensorFlowMinimumOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1121,22 +1195,9 @@ void ConvertStridedSliceOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "StridedSlice"); - CHECK_EQ(node.input_size(), 4); - - // Only a subset of the full TF op functionality is supported now. - if ( // No 64-bit indices. - GetDataTypeAttr(node, "Index") != DT_INT32 || - // No dimensionality changes. - GetIntAttr(node, "new_axis_mask") != 0 || - GetIntAttr(node, "shrink_axis_mask") != 0 || - // No sparse indices. - GetIntAttr(node, "ellipsis_mask") != 0 || - // Only 4D tensors are supported. - GetIntAttr(node, "begin_mask") > 15 || - GetIntAttr(node, "end_mask") > 15) { - ConvertUnsupportedOperator(node, tf_import_flags, model); - return; - } + // TODO(soroosh): The 4th input (strides) should be e optional, to be + // consistent with TF. + CheckInputsCount(node, tf_import_flags, 4); auto* op = new StridedSliceOperator; for (const auto& input : node.input()) { @@ -1157,7 +1218,7 @@ void ConvertPlaceholderOperator(const NodeDef& node, Model* model) { CHECK(node.op() == "Placeholder" || node.op() == "LegacyFedInput"); if (node.op() == "Placeholder") { - CHECK_EQ(GetInputsCount(node, tf_import_flags), 0); + CheckInputsCount(node, tf_import_flags, 0); } auto& array = model->GetOrCreateArray(node.name()); if (node.attr().count("dtype")) { @@ -1192,7 +1253,7 @@ void ConvertCastOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Cast"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); const auto tf_src_dtype = GetDataTypeAttr(node, "SrcT"); const auto tf_dst_dtype = GetDataTypeAttr(node, "DstT"); auto* op = new CastOperator; @@ -1207,7 +1268,7 @@ void ConvertFloorOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Floor"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); const auto data_type = GetDataTypeAttr(node, "T"); CHECK(data_type == DT_FLOAT); auto* op = new FloorOperator; @@ -1220,7 +1281,7 @@ void ConvertGatherOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Gather"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); const auto indices_data_type = GetDataTypeAttr(node, "Tindices"); CHECK(indices_data_type == DT_INT32 || indices_data_type == DT_INT64); auto* op = new GatherOperator; @@ -1234,7 +1295,7 @@ void ConvertArgMaxOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "ArgMax"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); const auto axis_data_type = GetDataTypeAttr(node, "Tidx"); const auto output_type = GetDataTypeAttr(node, "output_type"); CHECK(axis_data_type == DT_INT64 || axis_data_type == DT_INT32); @@ -1251,8 +1312,14 @@ void ConvertResizeBilinearOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "ResizeBilinear"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new ResizeBilinearOperator; + + op->align_corners = false; + if (HasAttr(node, "align_corners")) { + op->align_corners = GetBoolAttr(node, "align_corners"); + } + op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); op->outputs.push_back(node.name()); @@ -1263,7 +1330,7 @@ void ConvertBatchNormWithGlobalNormalizationOperator( const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "BatchNormWithGlobalNormalization"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 5); + CheckInputsCount(node, tf_import_flags, 5); // TODO(ahentz): to really match tensorflow we need to add variance_epsilon // to the input, before feeding it into TensorFlowRsqrtOperator. @@ -1312,7 +1379,7 @@ void ConvertFusedBatchNormOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "FusedBatchNorm"); - CHECK_EQ(node.input_size(), 5); + CheckInputsCount(node, tf_import_flags, 5); // Declare shortcuts for the inputs. const string& gamma_input = node.input(1); @@ -1368,7 +1435,7 @@ void ConvertSpaceToBatchNDOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "SpaceToBatchND"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 3); + CheckInputsCount(node, tf_import_flags, 3); CHECK_EQ(GetDataTypeAttr(node, "Tblock_shape"), DT_INT32); CHECK_EQ(GetDataTypeAttr(node, "Tpaddings"), DT_INT32); auto* op = new SpaceToBatchNDOperator; @@ -1383,7 +1450,7 @@ void ConvertBatchToSpaceNDOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "BatchToSpaceND"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 3); + CheckInputsCount(node, tf_import_flags, 3); CHECK_EQ(GetDataTypeAttr(node, "Tblock_shape"), DT_INT32); CHECK_EQ(GetDataTypeAttr(node, "Tcrops"), DT_INT32); auto* op = new BatchToSpaceNDOperator; @@ -1394,11 +1461,22 @@ void ConvertBatchToSpaceNDOperator(const NodeDef& node, model->operators.emplace_back(op); } +void ConvertExpOperator(const NodeDef& node, + const TensorFlowImportFlags& tf_import_flags, + Model* model) { + CHECK_EQ(node.op(), "Exp"); + CheckInputsCount(node, tf_import_flags, 1); + auto* op = new ExpOperator; + op->inputs.push_back(node.input(0)); + op->outputs.push_back(node.name()); + model->operators.emplace_back(op); +} + void ConvertMeanOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Mean"); - CHECK_EQ(node.input_size(), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new MeanOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1413,7 +1491,11 @@ void ConvertSvdfOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Svdf"); - bool has_bias = (node.input_size() == 4); + const int input_size = GetInputsCount(node, tf_import_flags); + QCHECK(input_size == 3 || input_size == 4) + << "Svdf node expects 3 or 4 inputs other than control dependencies: " + << node.DebugString(); + bool has_bias = (input_size == 4); auto* op = new SvdfOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1437,7 +1519,7 @@ void ConvertTransposeConvOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Conv2DBackpropInput"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 3); + CheckInputsCount(node, tf_import_flags, 3); auto* op = new TransposeConvOperator; op->inputs.push_back(node.input(2)); op->inputs.push_back(node.input(1)); @@ -1465,7 +1547,7 @@ void ConvertExpandDimsOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "ExpandDims"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new ExpandDimsOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1477,7 +1559,7 @@ void ConvertFillOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Fill"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new FillOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1489,7 +1571,7 @@ void ConvertFloorDivOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "FloorDiv"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new FloorDivOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1500,8 +1582,8 @@ void ConvertFloorDivOperator(const NodeDef& node, void ConvertFloorModOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { - CHECK(node.op() == "FloorMod"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CHECK_EQ(node.op(), "FloorMod"); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new FloorModOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1513,7 +1595,7 @@ void ConvertRangeOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Range"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 3); + CheckInputsCount(node, tf_import_flags, 3); auto* op = new RangeOperator; if (HasAttr(node, "Tidx")) { const auto dtype = toco::GetDataTypeAttr(node, "Tidx"); @@ -1532,7 +1614,7 @@ void ConvertRankOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Rank"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 1); + CheckInputsCount(node, tf_import_flags, 1); auto* op = new RankOperator; op->inputs.push_back(node.input(0)); op->outputs.push_back(node.name()); @@ -1545,7 +1627,10 @@ void ConvertStackOperator(const NodeDef& node, CHECK((node.op() == "Stack") || (node.op() == "Pack")); auto* op = new StackOperator; const int num_inputs = GetInputsCount(node, tf_import_flags); - CHECK_GE(num_inputs, 1); + QCHECK_GE(num_inputs, 1) + << node.op() + << " node expects at least 1 input other than control dependencies: " + << node.DebugString(); CHECK_EQ(num_inputs, GetIntAttr(node, "N")); for (int i = 0; i < num_inputs; ++i) { op->inputs.push_back(node.input(i)); @@ -1560,7 +1645,7 @@ void ConvertTransposeOperator(const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, Model* model) { CHECK_EQ(node.op(), "Transpose"); - CHECK_EQ(GetInputsCount(node, tf_import_flags), 2); + CheckInputsCount(node, tf_import_flags, 2); auto* op = new TransposeOperator; op->inputs.push_back(node.input(0)); op->inputs.push_back(node.input(1)); @@ -1602,7 +1687,7 @@ void StripCaretFromArrayNames(Model* model) { output = string(absl::StripPrefix(output, "^")); } } - for (auto& array : model->arrays) { + for (auto& array : model->GetArrayMap()) { if (absl::StartsWith(array.first, "^")) { LOG(FATAL) << "What?"; } @@ -1732,6 +1817,37 @@ bool InlineAllFunctions(GraphDef* graphdef) { } return graph_modified; } + +void ConvertTopKV2Operator(const NodeDef& node, + const TensorFlowImportFlags& tf_import_flags, + Model* model) { + CHECK((node.op() == "TopK") || (node.op() == "TopKV2")); + auto op = absl::make_unique(); + op->inputs.push_back(node.input(0)); + // K can be encoded as attr (TopK) convert it to a const. + if (HasAttr(node, "k")) { + // Convert attribute into const tensor. + const string array_name = node.name() + "k"; + auto& array = model->GetOrCreateArray(array_name); + array.data_type = ArrayDataType::kInt32; + // Size of array is always 1. + array.mutable_shape()->mutable_dims()->emplace_back(1); + + auto& output_int_data = + array.GetMutableBuffer().data; + output_int_data.resize(1); + output_int_data[0] = GetIntAttr(node, "k"); + op->inputs.push_back(array_name); + + } else { + CheckInputsCount(node, tf_import_flags, 2); + op->inputs.push_back(node.input(1)); + } + // The op has two outputs. + op->outputs.push_back(node.name() + ":0"); + op->outputs.push_back(node.name() + ":1"); + model->operators.emplace_back(op.release()); +} } // namespace std::unique_ptr ImportTensorFlowGraphDef( @@ -1787,7 +1903,9 @@ std::unique_ptr ImportTensorFlowGraphDef( ConvertAvgPoolOperator(node, tf_import_flags, model); } else if (node.op() == "Reshape") { ConvertReshapeOperator(node, tf_import_flags, model); - } else if (node.op() == "MatMul" || node.op() == "BatchMatMul") { + } else if (node.op() == "BatchMatMul") { + ConvertBatchMatMulOperator(node, tf_import_flags, model); + } else if (node.op() == "MatMul") { ConvertMatMulOperator(node, tf_import_flags, model); } else if (node.op() == "Div" || node.op() == "RealDiv") { ConvertDivOperator(node, tf_import_flags, model); @@ -1810,6 +1928,8 @@ std::unique_ptr ImportTensorFlowGraphDef( ConvertSquareOperator(node, tf_import_flags, model); } else if (node.op() == "Add") { ConvertAddOperator(node, tf_import_flags, model); + } else if (node.op() == "AddN") { + ConvertAddNOperator(node, tf_import_flags, model); } else if (node.op() == "Mul") { ConvertMulOperator(node, tf_import_flags, model); } else if (node.op() == "Sub") { @@ -1824,6 +1944,8 @@ std::unique_ptr ImportTensorFlowGraphDef( ConvertLRNOperator(node, tf_import_flags, model); } else if (node.op() == "Softmax") { ConvertSoftmaxOperator(node, tf_import_flags, model); + } else if (node.op() == "LogSoftmax") { + ConvertLogSoftmaxOperator(node, tf_import_flags, model); } else if (node.op() == "All") { ConvertAllOperator(node, tf_import_flags, model); } else if (node.op() == "Assert") { @@ -1907,6 +2029,10 @@ std::unique_ptr ImportTensorFlowGraphDef( ConvertTransposeOperator(node, tf_import_flags, model); } else if (node.op() == "ArgMax") { ConvertArgMaxOperator(node, tf_import_flags, model); + } else if (node.op() == "Exp") { + ConvertExpOperator(node, tf_import_flags, model); + } else if (node.op() == "TopK" || node.op() == "TopKV2") { + ConvertTopKV2Operator(node, tf_import_flags, model); } else { ConvertUnsupportedOperator(node, tf_import_flags, model); } diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.h b/tensorflow/contrib/lite/toco/import_tensorflow.h index 312e3b8f17cfaa012bf25696937f97d396802bb2..2177872334bfec6147f865be1518e440c2c636ea 100644 --- a/tensorflow/contrib/lite/toco/import_tensorflow.h +++ b/tensorflow/contrib/lite/toco/import_tensorflow.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_IMPORT_TENSORFLOW_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_IMPORT_TENSORFLOW_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_IMPORT_TENSORFLOW_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_IMPORT_TENSORFLOW_H_ #include #include @@ -39,4 +39,4 @@ std::unique_ptr ImportTensorFlowGraphDef( } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_IMPORT_TENSORFLOW_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_IMPORT_TENSORFLOW_H_ diff --git a/tensorflow/contrib/lite/toco/model.h b/tensorflow/contrib/lite/toco/model.h index 5e7cc294967186585c78108742107be08d213852..c55bf664f8b65c6eb53ff9ae926bed11adc7b183 100644 --- a/tensorflow/contrib/lite/toco/model.h +++ b/tensorflow/contrib/lite/toco/model.h @@ -12,9 +12,10 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_H_ +#include #include #include #include @@ -32,7 +33,9 @@ enum class OperatorType { kNone, // General-purpose neural network operators. kAdd, + kAddN, kAveragePool, + kBatchMatMul, kBatchNormalization, kConv, kConcatenation, @@ -41,6 +44,7 @@ enum class OperatorType { kSpaceToDepth, kDequantize, kDiv, + kExp, kExpandDims, kFill, kFloorDiv, @@ -60,6 +64,7 @@ enum class OperatorType { kRelu1, kRelu6, kSoftmax, + kLogSoftmax, kSub, kTanh, kTransposeConv, @@ -109,6 +114,7 @@ enum class OperatorType { kTensorFlowSwitch, kTensorFlowTile, kTranspose, + kTopK_V2, // An unsupported TF operation. It's only needed to be able to represent TF // graph internally and is expected to be dropped by graph transformations. kTensorFlowUnsupported, @@ -154,12 +160,17 @@ enum class AxesOrder { // may be involved only in debug-only subgraphs that we may not be interested // in actually supporting). enum class ArrayDataType { - kNone, + kNone, // 0 kBool, kFloat, + kInt8, kUint8, + kInt16, // 5 + kUint16, kInt32, + kUint32, kInt64, + kUint64, // 10 kString }; @@ -179,18 +190,38 @@ struct DataTypeImpl { typedef float Type; }; template <> +struct DataTypeImpl { + typedef int8 Type; +}; +template <> struct DataTypeImpl { typedef uint8 Type; }; template <> +struct DataTypeImpl { + typedef int16 Type; +}; +template <> +struct DataTypeImpl { + typedef uint16 Type; +}; +template <> struct DataTypeImpl { typedef int32 Type; }; template <> +struct DataTypeImpl { + typedef uint32 Type; +}; +template <> struct DataTypeImpl { typedef int64 Type; }; template <> +struct DataTypeImpl { + typedef uint64 Type; +}; +template <> struct DataTypeImpl { typedef string Type; }; @@ -559,6 +590,16 @@ struct AddOperator : Operator { AddOperator() : Operator(OperatorType::kAdd) {} }; +// Element-wise addition operator for N inputs. +// +// Inputs: +// inputs[i]: The i-th array to add together to form the output. +// +// TensorFlow equivalent: AddN +struct AddNOperator : Operator { + AddNOperator() : Operator(OperatorType::kAddN) {} +}; + // Concatenation operator: concatenates its inputs // along the axis. // @@ -700,6 +741,19 @@ struct TensorFlowIdentityOperator : Operator { TensorFlowIdentityOperator() : Operator(OperatorType::kTensorFlowIdentity) {} }; +// Batch matrix multiplication operator. This comes from the (deprecated) +// tf.batch_matmul or a tf.matmul that has rank 3. dims(0) is the batch count +// and it can be trivially unrolled into a series of matmuls on each element. +// +// Inputs: +// inputs[0]: required: the left-hand side matrix +// inputs[1]: required: the right-hand side matrix +// +// TensorFlow equivalent: MatMul +struct BatchMatMulOperator : Operator { + BatchMatMulOperator() : Operator(OperatorType::kBatchMatMul) {} +}; + // General matrix multiplication operator. We don't want to support general // matrix multiplication at inference time, so we resolve it during tooling // to more specific operator types, namely, FullyConnected. @@ -738,6 +792,9 @@ struct PadOperator : Operator { // // Inputs: // inputs[0]: required: the input array +// inputs[1]: required: the begin array +// inputs[2]: required: the end array +// inputs[3]: optional: the strides array // // TensorFlow equivalent: StridedSlice struct StridedSliceOperator : Operator { @@ -797,6 +854,17 @@ struct TransposeConvOperator : Operator { int stride_height = 0; }; +// Given a tensor input, this operation calculates element-wise exponential +// (y = e^x). +// +// Inputs: +// inputs[0]: required: input tensor +// +// TensorFlow equivalent: Exp +struct ExpOperator : Operator { + ExpOperator() : Operator(OperatorType::kExp) {} +}; + // Given a tensor input, this operation inserts a dimension of 1 at the // dimension index axis of input's shape. The dimension index axis starts at // zero; if you specify a negative number for axis it is counted backward from @@ -957,6 +1025,7 @@ struct TensorFlowSquareOperator : Operator { // TensorFlow equivalent: Transpose struct TransposeOperator : Operator { TransposeOperator() : Operator(OperatorType::kTranspose) {} + std::vector perm; }; // Element-wise subtraction operator. @@ -1200,6 +1269,16 @@ struct SoftmaxOperator : Operator { float beta = 0.f; }; +// LogSoftmax activation function. +// +// Inputs: +// inputs[0]: required: the logits input array +// +// TensorFlow equivalent: LogSoftmax +struct LogSoftmaxOperator : Operator { + LogSoftmaxOperator() : Operator(OperatorType::kLogSoftmax) {} +}; + // Cast operator. // // Inputs: @@ -1257,6 +1336,8 @@ struct ArgMaxOperator : Operator { // TensorFlow equivalent: ResizeBilinear struct ResizeBilinearOperator : Operator { ResizeBilinearOperator() : Operator(OperatorType::kResizeBilinear) {} + + bool align_corners = false; }; // SpaceToBatchND operator. It divides spatial dimensions into a grid of @@ -1271,6 +1352,10 @@ struct ResizeBilinearOperator : Operator { // TensorFlow equivalent: SpaceToBatchND struct SpaceToBatchNDOperator : Operator { SpaceToBatchNDOperator() : Operator(OperatorType::kSpaceToBatchND) {} + + std::vector block_shape; + std::vector before_paddings; + std::vector after_paddings; }; // BatchToSpaceND operator. Rearranges data from batch into blocks of @@ -1316,6 +1401,14 @@ struct SvdfOperator : Operator { int rank; }; +// TopKV2 operator. +// +// Inputs: +// input tensor and top_k scalar. +struct TopKV2Operator : Operator { + TopKV2Operator() : Operator(OperatorType::kTopK_V2) {} +}; + // Alloc's are used for transient arrays only. An Alloc specifies which interval // of the "transient_data" workspace buffer passed to inference functions, is to // be used for the transient array at hand. The 'start' and 'end' values are @@ -1516,29 +1609,54 @@ struct Array { // Our Model struct, represents an entire model (our "top-level" struct). // Owns everything. -struct Model { +class Model { + public: + using ArrayMap = std::unordered_map>; + + bool HasArray(const string& name) const { return arrays.count(name) > 0; } Array& GetArray(const string& name) const { - DCHECK(arrays.count(name)); + DCHECK(HasArray(name)) << "Array not found: " << name; return *arrays.at(name); } Array& GetOrCreateArray(const string& name) { - if (!arrays.count(name)) { + // Make sure name is not used by an optional array + DCHECK(!optional_arrays.count(name)); + if (!HasArray(name)) { Array* ptr = new Array; arrays[name] = std::unique_ptr(ptr); } Array& result = GetArray(name); return result; } + void CreateOptionalArray(const string& name) { + DCHECK(!arrays.count(name) && !optional_arrays.count(name)); + optional_arrays.insert(name); + } + bool IsOptionalArray(const string& name) const { + return optional_arrays.count(name); + } + + // Note that this invalidates all array iterators. + void EraseArray(const string& name) { arrays.erase(name); } + void EraseArrays(std::function discardable) { + for (auto it = arrays.begin(); it != arrays.end();) { + if (discardable(it->first)) { + it = arrays.erase(it); + } else { + ++it; + } + } + } + const ArrayMap& GetArrayMap() const { return arrays; } + + // Optional arrays are used for optional tensors, + // these tensors do not have data, but with reserved names as op inputs. + std::set optional_arrays; // The list of operators. Notice how it's a list of unique_ptr's, implying // that the Model is what owns Operator's and keeps them alive. std::vector> operators; - // The associative array mapping names to Array's. - // Notice how it's a container of unique_ptr's, implying - // that the Model is what owns Array's and keeps them alive. - // The Operator's refer to these Array's by their name strings, not by their - // addresses. See Operator::inputs, Operator::outputs. - std::unordered_map> arrays; + // Generic flags, a place where we combine information passed to us via // command-line parameters (e.g. --input_width=N) with information that // we may or may not find in the input model file. @@ -1547,7 +1665,15 @@ struct Model { std::size_t transient_data_size = 0; // For code-generation only: required alignment of the transient_data buffer std::size_t transient_data_alignment = 0; + + private: + // The associative array mapping names to Array's. + // Notice how it's a container of unique_ptr's, implying + // that the Model is what owns Array's and keeps them alive. + // The Operator's refer to these Array's by their name strings, not by their + // addresses. See Operator::inputs, Operator::outputs. + std::unordered_map> arrays; }; } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_H_ diff --git a/tensorflow/contrib/lite/toco/model_cmdline_flags.cc b/tensorflow/contrib/lite/toco/model_cmdline_flags.cc index 790b3443cef1c577e19bafc5e087ca42e6fce60a..4e2dec15a534607ef9207149a2e6061069eabcb1 100644 --- a/tensorflow/contrib/lite/toco/model_cmdline_flags.cc +++ b/tensorflow/contrib/lite/toco/model_cmdline_flags.cc @@ -148,6 +148,12 @@ bool ParseModelFlagsFromCommandLineFlags( "ranging from 32 to 127. This is disallowed by default so as to " "catch common copy-and-paste issues where invisible unicode " "characters are unwittingly added to these strings."), + Flag( + "arrays_extra_info_file", parsed_flags.arrays_extra_info_file.bind(), + parsed_flags.arrays_extra_info_file.default_value(), + "Path to an optional file containing a serialized ArraysExtraInfo " + "proto allowing to pass extra information about arrays not specified " + "in the input model file, such as extra MinMax information."), }; bool asked_for_help = *argc == 2 && (!strcmp(argv[1], "--help") || !strcmp(argv[1], "-help")); @@ -327,9 +333,6 @@ void ReadModelFlagsFromCommandLineFlags( CHECK(absl::SimpleAtoi(value, &size)); CHECK_GT(size, 0); rnn_state_proto->set_size(size); - } else if (key == "manually_create") { - CHECK_EQ(absl::AsciiStrToLower(value), "true"); - rnn_state_proto->set_manually_create(true); } else { LOG(FATAL) << "Unknown key '" << key << "' in --rnn_states"; } @@ -368,6 +371,15 @@ void ReadModelFlagsFromCommandLineFlags( parsed_model_flags.allow_nonascii_arrays.value()); model_flags->set_allow_nonexistent_arrays( parsed_model_flags.allow_nonexistent_arrays.value()); + + if (parsed_model_flags.arrays_extra_info_file.specified()) { + string arrays_extra_info_file_contents; + port::file::GetContents(parsed_model_flags.arrays_extra_info_file.value(), + &arrays_extra_info_file_contents, + port::file::Defaults()); + ParseFromStringEitherTextOrBinary(arrays_extra_info_file_contents, + model_flags->mutable_arrays_extra_info()); + } } ParsedModelFlags* UncheckedGlobalParsedModelFlags(bool must_already_exist) { diff --git a/tensorflow/contrib/lite/toco/model_cmdline_flags.h b/tensorflow/contrib/lite/toco/model_cmdline_flags.h index 027d7ae1aa62b5b31b8fcebdc29d4f547507b7fe..c868d5c7d0b5a6ee81d99423414c87e4e6e7cf66 100644 --- a/tensorflow/contrib/lite/toco/model_cmdline_flags.h +++ b/tensorflow/contrib/lite/toco/model_cmdline_flags.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_CMDLINE_FLAGS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_CMDLINE_FLAGS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_CMDLINE_FLAGS_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_CMDLINE_FLAGS_H_ #include #include @@ -40,5 +40,4 @@ ParsedModelFlags* GlobalParsedModelFlags(); } // namespace toco - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_CMDLINE_FLAGS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_CMDLINE_FLAGS_H_ diff --git a/tensorflow/contrib/lite/toco/model_flags.proto b/tensorflow/contrib/lite/toco/model_flags.proto index 13fea29a07ed9ea75ebe1b9b046f2a68d814c649..e4b39b34e85e4d703c1b41cb68f8139abd1f6279 100644 --- a/tensorflow/contrib/lite/toco/model_flags.proto +++ b/tensorflow/contrib/lite/toco/model_flags.proto @@ -81,19 +81,26 @@ message RnnState { optional string state_array = 1; optional string back_edge_source_array = 2; optional bool discardable = 5; - // TODO(benoitjacob): drop the 'size' field. Should be redundant with - // --input_shapes and shapes propagation. + // size allows to specify a 1-D shape for the RNN state array. + // Will be expanded with 1's to fit the model. + // TODO(benoitjacob): should allow a generic, explicit shape. optional int32 size = 3; - // TODO(benoitjacob): manually_create is a temporary hack: - // due to discrepancies between the current toco dims tracking and - // TensorFlow shapes, for some models we need to manually create RNN state - // arrays with a specified shape. - // Maybe we should actually implement back-edges as operators of their own, - // which would remove the need for much special-casing, including here, - // we could probably consistently let PropagateFixedSizes handle state - // arrays. - // TODO(benoitjacob): should really drop manually_create now. - optional bool manually_create = 4; +} + +// An ArraysExtraInfo message stores a collection of additional Information +// about arrays in a model, complementing the information in the model itself. +// It is intentionally a separate message so that it may be serialized and +// passed separately from the model. See --arrays_extra_info_file. +// +// A typical use case is to manually specify MinMax for specific arrays in a +// model that does not already contain such MinMax information. +message ArraysExtraInfo { + message Entry { + optional string name = 1; + optional float min = 2; + optional float max = 3; + } + repeated Entry entries = 1; } // ModelFlags encodes properties of a model that, depending on the file @@ -117,7 +124,7 @@ message RnnState { // optional int32 input_dims = 11 [ default = 4]; // repeated int32 input_shape = 13; // -// Next ID to USE: 18. +// Next ID to USE: 19. message ModelFlags { // Information about the input arrays, i.e. the arrays from which input // activations will be read. @@ -160,4 +167,8 @@ message ModelFlags { // catch common copy-and-paste issues where invisible unicode // characters are unwittingly added to these strings. optional bool allow_nonascii_arrays = 17; + + // If set, this ArraysExtraInfo allows to pass extra information about arrays + // not specified in the input model file, such as extra MinMax information. + optional ArraysExtraInfo arrays_extra_info = 18; } diff --git a/tensorflow/contrib/lite/toco/runtime/common.h b/tensorflow/contrib/lite/toco/runtime/common.h index bd55544f57f9a266514e878edd8f1f7dec1cb7b7..3c6828840c4a963a4a68774ec5d559b7f80baf22 100644 --- a/tensorflow/contrib/lite/toco/runtime/common.h +++ b/tensorflow/contrib/lite/toco/runtime/common.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ #ifndef ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK #ifdef GEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK @@ -23,4 +23,4 @@ limitations under the License. #include "tensorflow/contrib/lite/kernels/internal/common.h" -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ diff --git a/tensorflow/contrib/lite/toco/runtime/types.h b/tensorflow/contrib/lite/toco/runtime/types.h index df63b2d59ea2a98f1ec9009614c18791e8822c14..f5de5a5781a5304634642680e6a3cef60e7b844b 100644 --- a/tensorflow/contrib/lite/toco/runtime/types.h +++ b/tensorflow/contrib/lite/toco/runtime/types.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ #include "tensorflow/contrib/lite/kernels/internal/common.h" #include "tensorflow/contrib/lite/kernels/internal/compatibility.h" @@ -29,4 +29,4 @@ using tflite::RequiredBufferSizeForDims; } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ diff --git a/tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_svdf_test.cc b/tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_svdf_test.cc index 664e828c19dca1117b81113f723416541f48d621..646d048496c27955aa641fd01a35d8acfbd8dd90 100644 --- a/tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_svdf_test.cc +++ b/tensorflow/contrib/lite/toco/tensorflow_graph_matching/resolve_svdf_test.cc @@ -103,11 +103,11 @@ class ResolveSvdfTest : public ::testing::Test { // Add the float vector as an attribute to the node. (*node->mutable_attr())["dtype"].set_type(tensorflow::DT_FLOAT); tensorflow::TensorProto* allocated_tensor = new tensorflow::TensorProto; - tensorflow::TensorShapeProto* allocated_tesnor_shape = + tensorflow::TensorShapeProto* allocated_tensor_shape = new tensorflow::TensorShapeProto; - auto tensor_shape_dim0 = allocated_tesnor_shape->add_dim(); + auto tensor_shape_dim0 = allocated_tensor_shape->add_dim(); tensor_shape_dim0->set_size(values.size()); - allocated_tensor->set_allocated_tensor_shape(allocated_tesnor_shape); + allocated_tensor->set_allocated_tensor_shape(allocated_tensor_shape); allocated_tensor->set_tensor_content( string(reinterpret_cast(values.data()), values.size() * sizeof(float))); @@ -122,11 +122,11 @@ class ResolveSvdfTest : public ::testing::Test { // Add the float vector as an attribute to the node. (*node->mutable_attr())["dtype"].set_type(tensorflow::DT_INT32); tensorflow::TensorProto* allocated_tensor = new tensorflow::TensorProto; - tensorflow::TensorShapeProto* allocated_tesnor_shape = + tensorflow::TensorShapeProto* allocated_tensor_shape = new tensorflow::TensorShapeProto; - auto tensor_shape_dim0 = allocated_tesnor_shape->add_dim(); + auto tensor_shape_dim0 = allocated_tensor_shape->add_dim(); tensor_shape_dim0->set_size(values.size()); - allocated_tensor->set_allocated_tensor_shape(allocated_tesnor_shape); + allocated_tensor->set_allocated_tensor_shape(allocated_tensor_shape); allocated_tensor->set_tensor_content( string(reinterpret_cast(values.data()), values.size() * sizeof(int))); diff --git a/tensorflow/contrib/lite/toco/tensorflow_util.cc b/tensorflow/contrib/lite/toco/tensorflow_util.cc index 82e2800ca2f5bb017f91b5bf43d8d3cd05e97b83..0e7e9c41a066581b14fe1b78f83d8d57b916be6c 100644 --- a/tensorflow/contrib/lite/toco/tensorflow_util.cc +++ b/tensorflow/contrib/lite/toco/tensorflow_util.cc @@ -51,7 +51,8 @@ void LogDumpGraphDef(int log_level, const string& message, BEGIN DUMP OF TENSORFLOW GRAPHDEF (%s) There are %d nodes. There are %zu different op types: -)MSG", message, tf_graph.node_size(), ops.size()); +)MSG", + message, tf_graph.node_size(), ops.size()); for (const auto& op : ops) { toco::port::AppendF(&dump, " %s\n", op); } @@ -63,7 +64,8 @@ PROTO DUMP BEGIN NODE: name = %s op = %s inputs = [ -)MSG", node.name(), node.op()); +)MSG", + node.name(), node.op()); for (const auto& input : node.input()) { toco::port::AppendF(&dump, " %s\n", input); } diff --git a/tensorflow/contrib/lite/toco/tensorflow_util.h b/tensorflow/contrib/lite/toco/tensorflow_util.h index 152b4f7a727a88f721f1a63299ea4fa709bb5d52..61f91042685288a48ba19f8c67d4c7c1960a7787 100644 --- a/tensorflow/contrib/lite/toco/tensorflow_util.h +++ b/tensorflow/contrib/lite/toco/tensorflow_util.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_UTIL_H_ #include #include @@ -29,4 +29,4 @@ void LogDumpGraphDef(int log_level, const string& message, } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TENSORFLOW_UTIL_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/BUILD b/tensorflow/contrib/lite/toco/tflite/BUILD index 332253a092aff812fb18601862c66bc0423599c2..a2b8145a67278c3ac0065f9551da6ffd1de60772 100644 --- a/tensorflow/contrib/lite/toco/tflite/BUILD +++ b/tensorflow/contrib/lite/toco/tflite/BUILD @@ -27,7 +27,7 @@ cc_library( "//tensorflow/contrib/lite/toco:model", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/memory", - "@flatbuffers//:flatbuffers", + "@flatbuffers", ], ) @@ -41,7 +41,7 @@ tf_cc_test( "//tensorflow/contrib/lite/toco:tooling_util", "//tensorflow/core:protos_all_cc", "@com_google_googletest//:gtest_main", - "@flatbuffers//:flatbuffers", + "@flatbuffers", ], ) @@ -87,7 +87,7 @@ cc_library( "//tensorflow/contrib/lite/toco:model", "//tensorflow/contrib/lite/toco:tooling_util", "@com_google_absl//absl/strings", - "@flatbuffers//:flatbuffers", + "@flatbuffers", ], ) @@ -117,7 +117,8 @@ cc_library( ":types", "//tensorflow/contrib/lite/schema:schema_fbs", "//tensorflow/contrib/lite/toco:model", - "@flatbuffers//:flatbuffers", + "//tensorflow/contrib/lite/toco:tooling_util", + "@flatbuffers", ], ) @@ -131,7 +132,7 @@ tf_cc_test( "//tensorflow/contrib/lite:schema_fbs_version", "//tensorflow/contrib/lite/schema:schema_fbs", "@com_google_googletest//:gtest_main", - "@flatbuffers//:flatbuffers", + "@flatbuffers", ], ) diff --git a/tensorflow/contrib/lite/toco/tflite/builtin_operator.h b/tensorflow/contrib/lite/toco/tflite/builtin_operator.h index 93cc79ddb64fbc46a97a47ecdc155a8aabf5c3ef..cfe7ecd9f982618dea3b3a5d02e69e3f15434bc2 100644 --- a/tensorflow/contrib/lite/toco/tflite/builtin_operator.h +++ b/tensorflow/contrib/lite/toco/tflite/builtin_operator.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_ #include "absl/memory/memory.h" #include "tensorflow/contrib/lite/toco/tflite/operator.h" @@ -71,4 +71,4 @@ class BuiltinOperator : public BaseOperator { } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_BUILTIN_OPERATOR_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/custom_operator.h b/tensorflow/contrib/lite/toco/tflite/custom_operator.h index 1a4bfac7d4f684043d2a9ce8fc2c78dd738f4b69..bd5713618ff379be42fd1b76649cfb2cf55b843d 100644 --- a/tensorflow/contrib/lite/toco/tflite/custom_operator.h +++ b/tensorflow/contrib/lite/toco/tflite/custom_operator.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_ #include "flatbuffers/flexbuffers.h" #include "absl/memory/memory.h" @@ -71,4 +71,4 @@ class CustomOperator : public BaseOperator { } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_CUSTOM_OPERATOR_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/export.cc b/tensorflow/contrib/lite/toco/tflite/export.cc index bec694a23377c7c70684000069e9c08ee446b6c0..27719599708a7eb14f72a82f8e5d76b3b8af9dc4 100644 --- a/tensorflow/contrib/lite/toco/tflite/export.cc +++ b/tensorflow/contrib/lite/toco/tflite/export.cc @@ -26,6 +26,9 @@ namespace toco { namespace tflite { +using flatbuffers::FlatBufferBuilder; +using flatbuffers::Offset; +using flatbuffers::Vector; using ::tflite::Buffer; using ::tflite::BuiltinOperator; using ::tflite::BuiltinOperator_CUSTOM; @@ -39,9 +42,6 @@ using ::tflite::Operator; using ::tflite::OperatorCode; using ::tflite::SubGraph; using ::tflite::Tensor; -using flatbuffers::FlatBufferBuilder; -using flatbuffers::Offset; -using flatbuffers::Vector; namespace { @@ -62,7 +62,7 @@ namespace details { void LoadTensorsMap(const Model& model, TensorsMap* tensors_map) { // First find a list of unique array names. std::set names; - for (const auto& array_pair : model.arrays) { + for (const auto& array_pair : model.GetArrayMap()) { names.insert(array_pair.first); } @@ -96,7 +96,7 @@ Offset>> ExportTensors( // tensors in the tensors_map. std::map> ordered_tensors; - for (const auto& array_pair : model.arrays) { + for (const auto& array_pair : model.GetArrayMap()) { const string& tensor_name = array_pair.first; const toco::Array& array = *array_pair.second; @@ -235,9 +235,10 @@ Offset>> ExportOperators( for (const auto& op : model.operators) { std::vector inputs; for (const string& input : op->inputs) { - inputs.push_back(tensors_map.at(input)); + // -1 is the ID for optional tensor in TFLite output + int id = model.IsOptionalArray(input) ? -1 : tensors_map.at(input); + inputs.push_back(id); } - std::vector outputs; for (const string& output : op->outputs) { outputs.push_back(tensors_map.at(output)); diff --git a/tensorflow/contrib/lite/toco/tflite/export.h b/tensorflow/contrib/lite/toco/tflite/export.h index 44012b7126e17d730ea248551dea2414ad0072d9..8c79cb820015e16847ce48c171e8f6e41f60c319 100644 --- a/tensorflow/contrib/lite/toco/tflite/export.h +++ b/tensorflow/contrib/lite/toco/tflite/export.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ #include "tensorflow/contrib/lite/toco/model.h" @@ -73,4 +73,4 @@ void LoadOperatorsMap(const Model& model, OperatorsMap* operators_map); } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/export_test.cc b/tensorflow/contrib/lite/toco/tflite/export_test.cc index d4c4612d62c4eb5b14898eb8846314246ecbb815..6754372330797ae30230af26a3b478c24ad44005 100644 --- a/tensorflow/contrib/lite/toco/tflite/export_test.cc +++ b/tensorflow/contrib/lite/toco/tflite/export_test.cc @@ -91,8 +91,8 @@ TEST_F(ExportTest, Export) { } } - EXPECT_THAT(names, ElementsAre("builtin:ADD", "builtin:CONV_2D", "custom:Sub", - "custom:MyCrazyOp")); + EXPECT_THAT(names, ElementsAre("builtin:ADD", "builtin:CONV_2D", + "builtin:SUB", "custom:MyCrazyOp")); std::vector indices; auto operators = (*model->subgraphs())[0]->operators(); diff --git a/tensorflow/contrib/lite/toco/tflite/import.cc b/tensorflow/contrib/lite/toco/tflite/import.cc index bbf201fd288140d990b8f739adcd9244e1196072..5b1ab514b23248cd98e66847185d0e8b9fe2d6aa 100644 --- a/tensorflow/contrib/lite/toco/tflite/import.cc +++ b/tensorflow/contrib/lite/toco/tflite/import.cc @@ -18,6 +18,7 @@ limitations under the License. #include "tensorflow/contrib/lite/schema/schema_generated.h" #include "tensorflow/contrib/lite/toco/tflite/operator.h" #include "tensorflow/contrib/lite/toco/tflite/types.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" namespace toco { @@ -119,8 +120,16 @@ void ImportOperators( auto inputs = input_op->inputs(); for (int i = 0; i < inputs->Length(); i++) { auto input_index = inputs->Get(i); - const string& input_name = tensors_table.at(input_index); - op->inputs.push_back(input_name); + // input_index == -1 indicates optional tensor. + if (input_index != -1) { + const string& input_name = tensors_table.at(input_index); + op->inputs.push_back(input_name); + } else { + const string& tensor_name = + toco::AvailableArrayName(*model, "OptionalTensor"); + model->CreateOptionalArray(tensor_name); + op->inputs.push_back(tensor_name); + } } auto outputs = input_op->outputs(); for (int i = 0; i < outputs->Length(); i++) { diff --git a/tensorflow/contrib/lite/toco/tflite/import.h b/tensorflow/contrib/lite/toco/tflite/import.h index 3c27a2843c47814ad46c8f1bbd77b7afcb324375..280677bae189fa345c2e19f6399a7b9ac7629ab5 100644 --- a/tensorflow/contrib/lite/toco/tflite/import.h +++ b/tensorflow/contrib/lite/toco/tflite/import.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_IMPORT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_IMPORT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_IMPORT_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_IMPORT_H_ #include "tensorflow/contrib/lite/schema/schema_generated.h" #include "tensorflow/contrib/lite/toco/model.h" @@ -46,4 +46,4 @@ void LoadOperatorsTable(const ::tflite::Model &input_model, } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_IMPORT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_IMPORT_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/import_test.cc b/tensorflow/contrib/lite/toco/tflite/import_test.cc index 309fa6d7f688ba1dd99a7e6eeda14d513a9e49d4..aad6e780d5eb5c3dbc880906df5053ad231ffd54 100644 --- a/tensorflow/contrib/lite/toco/tflite/import_test.cc +++ b/tensorflow/contrib/lite/toco/tflite/import_test.cc @@ -114,7 +114,7 @@ TEST_F(ImportTest, Tensors) { auto model = Import(ModelFlags(), InputModelAsString()); - ASSERT_GT(model->arrays.count("tensor_one"), 0); + ASSERT_GT(model->HasArray("tensor_one"), 0); Array& a1 = model->GetArray("tensor_one"); EXPECT_EQ(ArrayDataType::kFloat, a1.data_type); EXPECT_THAT(a1.GetBuffer().data, diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc index ede6df88ab3a48f25c93f19c3d84a3ce5afc0450..aabc7c5109ddc205a7862c3ee2253390dae25095 100644 --- a/tensorflow/contrib/lite/toco/tflite/operator.cc +++ b/tensorflow/contrib/lite/toco/tflite/operator.cc @@ -130,6 +130,63 @@ class Add : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + return ::tflite::CreateSpaceToBatchNDOptions(*builder); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override {} +}; + +class Sub : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + auto activation_function = + ActivationFunction::Serialize(op.fused_activation_function); + return ::tflite::CreateSubOptions(*builder, activation_function); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + op->fused_activation_function = + ActivationFunction::Deserialize(options.fused_activation_function()); + } +}; + +class Div : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + auto activation_function = + ActivationFunction::Serialize(op.fused_activation_function); + return ::tflite::CreateDivOptions(*builder, activation_function); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + op->fused_activation_function = + ActivationFunction::Deserialize(options.fused_activation_function()); + } +}; + class BatchToSpaceND : public BuiltinOperator WriteOptions( const TocoOperator& op, flatbuffers::FlatBufferBuilder* builder) const override { - auto block_shape = builder->CreateVector(op.block_shape); - auto before_crops = builder->CreateVector(op.before_crops); - auto after_crops = builder->CreateVector(op.after_crops); - return ::tflite::CreateBatchToSpaceNDOptions(*builder, block_shape, - before_crops, after_crops); + return ::tflite::CreateBatchToSpaceNDOptions(*builder); } void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->block_shape.insert(op->block_shape.end(), - options.block_shape()->begin(), - options.block_shape()->end()); - op->before_crops.insert(op->before_crops.end(), - options.before_crops()->begin(), - options.before_crops()->end()); - op->after_crops.insert(op->after_crops.end(), - options.after_crops()->begin(), - options.after_crops()->end()); - } + TocoOperator* op) const override {} }; class Cast : public CustomOperator { @@ -403,20 +446,11 @@ class Pad : public BuiltinOperator WriteOptions( const TocoOperator& op, flatbuffers::FlatBufferBuilder* builder) const override { - auto before_padding = builder->CreateVector(op.left_padding); - auto after_padding = builder->CreateVector(op.right_padding); - return ::tflite::CreatePadOptions(*builder, before_padding, after_padding); + return ::tflite::CreatePadOptions(*builder); } void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->left_padding.insert(op->left_padding.end(), - options.before_padding()->begin(), - options.before_padding()->end()); - op->right_padding.insert(op->right_padding.end(), - options.after_padding()->begin(), - options.after_padding()->end()); - } + TocoOperator* op) const override {} }; class Reshape @@ -475,18 +509,154 @@ class SpaceToDepth } }; -class Split : public CustomOperator { +class Transpose + : public BuiltinOperator { public: - using CustomOperator::CustomOperator; - void WriteOptions(const TocoOperator& op, - flexbuffers::Builder* fbb) const override { - fbb->Int("num_split", op.num_split); + using BuiltinOperator::BuiltinOperator; + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + return ::tflite::CreateTransposeOptions(*builder); } - void ReadOptions(const flexbuffers::Map& m, TocoOperator* op) const override { - op->num_split = m["num_split"].AsInt64(); + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override {} +}; + +class Lstm : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + // Current toco converter only supports tanh, no clip. + return ::tflite::CreateLSTMOptions(*builder, /*fused_activation_function=*/ + ::tflite::ActivationFunctionType_TANH, + /*cell_clip=*/0.0, + /*proj_clip=*/0.0); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + // Only support tanh activation, so check that tflite type is tanh. + CHECK(options.fused_activation_function() == + ::tflite::ActivationFunctionType_TANH); + } +}; + +class Mean : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + return ::tflite::CreateMeanOptions(*builder, op.keep_dims); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + op->keep_dims = options.keep_dims(); + } +}; + +class ResizeBilinear + : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + return ::tflite::CreateResizeBilinearOptions(*builder, op.align_corners); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + op->align_corners = options.align_corners(); } }; +class Squeeze + : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + auto squeeze_dims = builder->CreateVector(op.squeeze_dims); + return ::tflite::CreateSqueezeOptions(*builder, squeeze_dims); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + op->squeeze_dims.insert(op->squeeze_dims.end(), + options.squeeze_dims()->begin(), + options.squeeze_dims()->end()); + } +}; + +class Split + : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + return ::tflite::CreateSplitOptions(*builder, op.num_split); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + op->num_split = options.num_splits(); + } +}; + +class StridedSlice + : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + return ::tflite::CreateStridedSliceOptions( + *builder, op.begin_mask, op.end_mask, op.ellipsis_mask, + op.new_axis_mask, op.shrink_axis_mask); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + op->begin_mask = options.begin_mask(); + op->end_mask = options.end_mask(); + op->ellipsis_mask = options.ellipsis_mask(); + op->new_axis_mask = options.new_axis_mask(); + op->shrink_axis_mask = options.shrink_axis_mask(); + } +}; + +class TopK_V2 : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + return ::tflite::CreateTopKV2Options(*builder); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override {} +}; + class TensorFlowUnsupported : public BaseOperator { public: using BaseOperator::BaseOperator; @@ -600,8 +770,13 @@ std::vector> BuildOperatorList() { // Builtin Operators. ops.emplace_back(new Add(::tflite::BuiltinOperator_ADD, OperatorType::kAdd)); + ops.emplace_back(new Div(::tflite::BuiltinOperator_DIV, OperatorType::kDiv)); + ops.emplace_back(new Sub(::tflite::BuiltinOperator_SUB, OperatorType::kSub)); ops.emplace_back(new AveragePool(::tflite::BuiltinOperator_AVERAGE_POOL_2D, OperatorType::kAveragePool)); + ops.emplace_back( + new SpaceToBatchND(::tflite::BuiltinOperator_SPACE_TO_BATCH_ND, + OperatorType::kSpaceToBatchND)); ops.emplace_back( new BatchToSpaceND(::tflite::BuiltinOperator_BATCH_TO_SPACE_ND, OperatorType::kBatchToSpaceND)); @@ -636,24 +811,38 @@ std::vector> BuildOperatorList() { OperatorType::kSpaceToDepth)); ops.emplace_back( new Svdf(::tflite::BuiltinOperator_SVDF, OperatorType::kSvdf)); + ops.emplace_back(new Transpose(::tflite::BuiltinOperator_TRANSPOSE, + OperatorType::kTranspose)); + ops.emplace_back( + new Mean(::tflite::BuiltinOperator_MEAN, OperatorType::kMean)); + ops.emplace_back(new ResizeBilinear(::tflite::BuiltinOperator_RESIZE_BILINEAR, + OperatorType::kResizeBilinear)); + ops.emplace_back( + new Squeeze(::tflite::BuiltinOperator_SQUEEZE, OperatorType::kSqueeze)); + ops.emplace_back(new Split(::tflite::BuiltinOperator_SPLIT, + OperatorType::kTensorFlowSplit)); + ops.emplace_back(new StridedSlice(::tflite::BuiltinOperator_STRIDED_SLICE, + OperatorType::kStridedSlice)); + ops.emplace_back( + new TopK_V2(::tflite::BuiltinOperator_TOPK_V2, OperatorType::kTopK_V2)); + ops.emplace_back( + new Lstm(::tflite::BuiltinOperator_LSTM, OperatorType::kLstmCell)); // Custom Operators. ops.emplace_back(new Cast("CAST", OperatorType::kCast)); ops.emplace_back( new DepthToSpace("DEPTH_TO_SPACE", OperatorType::kDepthToSpace)); ops.emplace_back(new FakeQuant("FAKE_QUANT", OperatorType::kFakeQuant)); - ops.emplace_back(new Split("SPLIT", OperatorType::kTensorFlowSplit)); ops.emplace_back(new TensorFlowUnsupported( "TENSORFLOW_UNSUPPORTED", OperatorType::kTensorFlowUnsupported)); // There operators are supported by Toco, but not by TF Lite, and has no // attributes. + ops.emplace_back( + new SimpleOperator("ADDN", OperatorType::kAddN)); ops.emplace_back(new SimpleOperator("NEG", OperatorType::kNeg)); ops.emplace_back(new SimpleOperator( "RSQRT", OperatorType::kTensorFlowRsqrt)); - ops.emplace_back( - new SimpleOperator("DIV", OperatorType::kDiv)); - // Simple Operators. ops.emplace_back(new SimpleOperator( "DEQUANTIZE", OperatorType::kDequantize)); @@ -662,15 +851,14 @@ std::vector> BuildOperatorList() { ops.emplace_back( new SimpleOperator("RELU", OperatorType::kRelu)); ops.emplace_back( - new SimpleOperator("RELU1", OperatorType::kRelu1)); + new SimpleOperator("RELU_N1_TO_1", OperatorType::kRelu1)); ops.emplace_back( new SimpleOperator("RELU6", OperatorType::kRelu6)); - ops.emplace_back(new SimpleOperator( - "RESIZE_BILINEAR", OperatorType::kResizeBilinear)); ops.emplace_back(new SimpleOperator( "LOGISTIC", OperatorType::kLogistic)); ops.emplace_back( new SimpleOperator("TANH", OperatorType::kTanh)); + ops.emplace_back(new SimpleOperator("EXP", OperatorType::kExp)); return ops; } diff --git a/tensorflow/contrib/lite/toco/tflite/operator.h b/tensorflow/contrib/lite/toco/tflite/operator.h index 37df302d4697c78e0349bcd30e0e1adc540066bc..88af3d6ab6c6af150af83ed5c52931f9f089aa3c 100644 --- a/tensorflow/contrib/lite/toco/tflite/operator.h +++ b/tensorflow/contrib/lite/toco/tflite/operator.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_OPERATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_OPERATOR_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_OPERATOR_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_OPERATOR_H_ #include "flatbuffers/flatbuffers.h" #include "tensorflow/contrib/lite/schema/schema_generated.h" @@ -86,4 +86,4 @@ class BaseOperator { } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_OPERATOR_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_OPERATOR_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/operator_test.cc b/tensorflow/contrib/lite/toco/tflite/operator_test.cc index 735eea4ddcb3906b8fdc2fa4576b55736ad4382f..5c486f72ade9ec5f366f075fcc39274bb7b12679 100644 --- a/tensorflow/contrib/lite/toco/tflite/operator_test.cc +++ b/tensorflow/contrib/lite/toco/tflite/operator_test.cc @@ -102,12 +102,11 @@ TEST_F(OperatorTest, SimpleOperators) { OperatorType::kDequantize); CheckSimpleOperator("FLOOR", OperatorType::kFloor); CheckSimpleOperator("RELU", OperatorType::kRelu); - CheckSimpleOperator("RELU1", OperatorType::kRelu1); + CheckSimpleOperator("RELU_N1_TO_1", OperatorType::kRelu1); CheckSimpleOperator("RELU6", OperatorType::kRelu6); - CheckSimpleOperator("RESIZE_BILINEAR", - OperatorType::kResizeBilinear); CheckSimpleOperator("LOGISTIC", OperatorType::kLogistic); CheckSimpleOperator("TANH", OperatorType::kTanh); + CheckSimpleOperator("EXP", OperatorType::kExp); } TEST_F(OperatorTest, BuiltinAdd) { @@ -119,17 +118,13 @@ TEST_F(OperatorTest, BuiltinAdd) { output_toco_op->fused_activation_function); } -TEST_F(OperatorTest, BuiltinBatchToSpaceND) { - BatchToSpaceNDOperator op; - op.block_shape = {2, 2}; - op.before_crops = {1, 2}; - op.after_crops = {3, 4}; +TEST_F(OperatorTest, BuiltinMean) { + MeanOperator op; + op.keep_dims = false; - auto output_toco_op = SerializeAndDeserialize( - GetOperator("BATCH_TO_SPACE_ND", OperatorType::kBatchToSpaceND), op); - EXPECT_EQ(op.block_shape, output_toco_op->block_shape); - EXPECT_EQ(op.before_crops, output_toco_op->before_crops); - EXPECT_EQ(op.after_crops, output_toco_op->after_crops); + auto output_toco_op = + SerializeAndDeserialize(GetOperator("MEAN", OperatorType::kMean), op); + EXPECT_EQ(op.keep_dims, output_toco_op->keep_dims); } TEST_F(OperatorTest, CustomCast) { @@ -234,16 +229,6 @@ TEST_F(OperatorTest, BuiltinMaxPool) { EXPECT_EQ(op.kheight, output_toco_op->kheight); } -TEST_F(OperatorTest, BuiltinPad) { - PadOperator op; - op.left_padding = {1, 2, 3}; - op.right_padding = {1, 2, 3}; - auto output_toco_op = - SerializeAndDeserialize(GetOperator("PAD", OperatorType::kPad), op); - EXPECT_EQ(op.left_padding, output_toco_op->left_padding); - EXPECT_EQ(op.right_padding, output_toco_op->right_padding); -} - TEST_F(OperatorTest, BuiltinReshape) { TensorFlowReshapeOperator op; op.shape = {1, 2, 4, 5, 8}; @@ -345,6 +330,14 @@ TEST_F(OperatorTest, BuiltinMul) { output_toco_op->fused_activation_function); } +TEST_F(OperatorTest, ResizeBilinear) { + ResizeBilinearOperator op; + op.align_corners = true; + auto output_toco_op = SerializeAndDeserialize( + GetOperator("RESIZE_BILINEAR", OperatorType::kResizeBilinear), op); + EXPECT_EQ(op.align_corners, output_toco_op->align_corners); +} + TEST_F(OperatorTest, Svdf) { SvdfOperator op; op.fused_activation_function = FusedActivationFunctionType::kRelu; @@ -356,6 +349,44 @@ TEST_F(OperatorTest, Svdf) { EXPECT_EQ(op.rank, output_toco_op->rank); } +TEST_F(OperatorTest, Squeeze) { + SqueezeOperator op; + op.squeeze_dims = {-2, -3, 4, 1, 4}; + + auto output_toco_op = SerializeAndDeserialize( + GetOperator("SQUEEZE", OperatorType::kSqueeze), op); + EXPECT_EQ(op.squeeze_dims, output_toco_op->squeeze_dims); +} + +TEST_F(OperatorTest, StridedSlice) { + StridedSliceOperator op; + + op.begin_mask = 1; + op.end_mask = 2; + op.ellipsis_mask = 1; + op.new_axis_mask = 1; + op.shrink_axis_mask = 2; + + auto output_toco_op = SerializeAndDeserialize( + GetOperator("STRIDED_SLICE", OperatorType::kStridedSlice), op); + EXPECT_EQ(op.start_indices, output_toco_op->start_indices); + EXPECT_EQ(op.stop_indices, output_toco_op->stop_indices); + EXPECT_EQ(op.strides, output_toco_op->strides); + EXPECT_EQ(op.begin_mask, output_toco_op->begin_mask); + EXPECT_EQ(op.end_mask, output_toco_op->end_mask); + EXPECT_EQ(op.end_mask, output_toco_op->end_mask); + EXPECT_EQ(op.ellipsis_mask, output_toco_op->ellipsis_mask); + EXPECT_EQ(op.new_axis_mask, output_toco_op->new_axis_mask); + EXPECT_EQ(op.shrink_axis_mask, output_toco_op->shrink_axis_mask); +} + +TEST_F(OperatorTest, BuiltinTopKV2) { + TopKV2Operator op; + auto output_toco_op = SerializeAndDeserialize( + GetOperator("TOPK_V2", OperatorType::kTopK_V2), op); + ASSERT_NE(nullptr, output_toco_op.get()); +} + TEST_F(OperatorTest, TensorFlowUnsupported) { TensorFlowUnsupportedOperator op; op.tensorflow_op = "MyCustomUnsupportedOp"; diff --git a/tensorflow/contrib/lite/toco/tflite/simple_operator.h b/tensorflow/contrib/lite/toco/tflite/simple_operator.h index 992b98bacafecb080e792ae87a2940977482eed6..72678c82a22a7168f858747b0b1c6a2b515b6578 100644 --- a/tensorflow/contrib/lite/toco/tflite/simple_operator.h +++ b/tensorflow/contrib/lite/toco/tflite/simple_operator.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_ #include "tensorflow/contrib/lite/toco/tflite/operator.h" @@ -47,4 +47,4 @@ class SimpleOperator : public BaseOperator { } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_SIMPLE_OPERATOR_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/types.cc b/tensorflow/contrib/lite/toco/tflite/types.cc index a6fa0237bc08d7d721fa4955c802b1faea69ca75..b4c2851502a40a1ca36965d4ddd2c8a15b8fe60f 100644 --- a/tensorflow/contrib/lite/toco/tflite/types.cc +++ b/tensorflow/contrib/lite/toco/tflite/types.cc @@ -51,6 +51,8 @@ void CopyBuffer(const ::tflite::Buffer& buffer, Array* array) { return ::tflite::TensorType_FLOAT32; case ArrayDataType::kInt32: return ::tflite::TensorType_INT32; + case ArrayDataType::kInt64: + return ::tflite::TensorType_INT64; case ArrayDataType::kUint8: return ::tflite::TensorType_UINT8; case ArrayDataType::kString: @@ -68,6 +70,8 @@ ArrayDataType DataType::Deserialize(int tensor_type) { return ArrayDataType::kFloat; case ::tflite::TensorType_INT32: return ArrayDataType::kInt32; + case ::tflite::TensorType_INT64: + return ArrayDataType::kInt64; case ::tflite::TensorType_STRING: return ArrayDataType::kString; case ::tflite::TensorType_UINT8: @@ -105,6 +109,8 @@ void DataBuffer::Deserialize(const ::tflite::Tensor& tensor, return CopyBuffer(buffer, array); case ::tflite::TensorType_INT32: return CopyBuffer(buffer, array); + case ::tflite::TensorType_INT64: + return CopyBuffer(buffer, array); case ::tflite::TensorType_STRING: return CopyBuffer(buffer, array); case ::tflite::TensorType_UINT8: @@ -146,7 +152,7 @@ PaddingType Padding::Deserialize(int padding) { case FusedActivationFunctionType::kRelu6: return ::tflite::ActivationFunctionType_RELU6; case FusedActivationFunctionType::kRelu1: - return ::tflite::ActivationFunctionType_RELU1; + return ::tflite::ActivationFunctionType_RELU_N1_TO_1; default: LOG(FATAL) << "Unhandled fused activation function type."; } @@ -161,7 +167,7 @@ FusedActivationFunctionType ActivationFunction::Deserialize( return FusedActivationFunctionType::kRelu; case ::tflite::ActivationFunctionType_RELU6: return FusedActivationFunctionType::kRelu6; - case ::tflite::ActivationFunctionType_RELU1: + case ::tflite::ActivationFunctionType_RELU_N1_TO_1: return FusedActivationFunctionType::kRelu1; default: LOG(FATAL) << "Unhandled fused activation function type."; diff --git a/tensorflow/contrib/lite/toco/tflite/types.h b/tensorflow/contrib/lite/toco/tflite/types.h index f7c51405107d954fa259809b72f56af193e344fb..3923756fc94e3175a6505740a96cce8d614c3990 100644 --- a/tensorflow/contrib/lite/toco/tflite/types.h +++ b/tensorflow/contrib/lite/toco/tflite/types.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ #include "tensorflow/contrib/lite/schema/schema_generated.h" #include "tensorflow/contrib/lite/toco/model.h" @@ -55,4 +55,4 @@ struct ActivationFunction { } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/types_test.cc b/tensorflow/contrib/lite/toco/tflite/types_test.cc index 174b78f3e632fde8dc6ea0ed83ed7a67fa12c16a..a040fe135841b92a6e668f32cc5e36cf812ab15b 100644 --- a/tensorflow/contrib/lite/toco/tflite/types_test.cc +++ b/tensorflow/contrib/lite/toco/tflite/types_test.cc @@ -28,8 +28,8 @@ using flatbuffers::Vector; // These are types that exist in TF Mini but don't have a correspondence // in TF Lite. -static const ArrayDataType kUnsupportedTocoTypes[] = { - ArrayDataType::kNone, ArrayDataType::kBool, ArrayDataType::kInt64}; +static const ArrayDataType kUnsupportedTocoTypes[] = {ArrayDataType::kNone, + ArrayDataType::kBool}; // These are TF Lite types for which there is no correspondence in TF Mini. static const ::tflite::TensorType kUnsupportedTfLiteTypes[] = { @@ -70,6 +70,7 @@ TEST(DataType, SupportedTypes) { std::vector> testdata = { {ArrayDataType::kUint8, ::tflite::TensorType_UINT8}, {ArrayDataType::kInt32, ::tflite::TensorType_INT32}, + {ArrayDataType::kInt64, ::tflite::TensorType_INT64}, {ArrayDataType::kFloat, ::tflite::TensorType_FLOAT32}}; for (auto x : testdata) { EXPECT_EQ(x.second, DataType::Serialize(x.first)); @@ -172,7 +173,7 @@ TEST(ActivationFunction, All) { {FusedActivationFunctionType::kRelu6, ::tflite::ActivationFunctionType_RELU6}, {FusedActivationFunctionType::kRelu1, - ::tflite::ActivationFunctionType_RELU1}}; + ::tflite::ActivationFunctionType_RELU_N1_TO_1}}; for (auto x : testdata) { EXPECT_EQ(x.second, ActivationFunction::Serialize(x.first)); EXPECT_EQ(x.first, ActivationFunction::Deserialize(x.second)); diff --git a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc index f8281f3a5725283d472e5e1a36e4d904b4dc1c49..c5a62fdb620ee7d6b7195f6e8e2bc3cb208feb10 100644 --- a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc +++ b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc @@ -44,9 +44,11 @@ bool ParseTocoFlagsFromCommandLineFlags( "For Protobuf formats, the binary format will be used."), Flag("input_format", parsed_flags.input_format.bind(), parsed_flags.input_format.default_value(), - "Input file format. One of: tensorflow_graphdef, "), + "Input file format. One of: TENSORFLOW_GRAPHDEF, TFLITE."), Flag("output_format", parsed_flags.output_format.bind(), - parsed_flags.output_format.default_value(), "Output file format."), + parsed_flags.output_format.default_value(), + "Output file format. " + "One of TENSORFLOW_GRAPHDEF, TFLITE, GRAPHVIZ_DOT."), Flag("default_ranges_min", parsed_flags.default_ranges_min.bind(), parsed_flags.default_ranges_min.default_value(), "If defined, will be used as the default value for the min bound " @@ -58,11 +60,13 @@ bool ParseTocoFlagsFromCommandLineFlags( Flag("inference_type", parsed_flags.inference_type.bind(), parsed_flags.inference_type.default_value(), "Target data type of arrays in the output file (for input_arrays, " - "this may be overridden by inference_input_type)."), + "this may be overridden by inference_input_type). " + "One of FLOAT, QUANTIZED_UINT8."), Flag("inference_input_type", parsed_flags.inference_input_type.bind(), parsed_flags.inference_input_type.default_value(), - "Target data type of input arrays. If not specified, inference_type " - "is used."), + "Target data type of input arrays. " + "If not specified, inference_type is used. " + "One of FLOAT, QUANTIZED_UINT8."), Flag("input_type", parsed_flags.input_type.bind(), parsed_flags.input_type.default_value(), "Deprecated ambiguous flag that set both --input_data_types and " @@ -76,35 +80,31 @@ bool ParseTocoFlagsFromCommandLineFlags( Flag("drop_fake_quant", parsed_flags.drop_fake_quant.bind(), parsed_flags.drop_fake_quant.default_value(), - "Ignore and discard FakeQuant nodes. For instance, that can be used " - "to " + "Ignore and discard FakeQuant nodes. For instance, to " "generate plain float code without fake-quantization from a " - "quantized " - "graph."), + "quantized graph."), Flag( "reorder_across_fake_quant", parsed_flags.reorder_across_fake_quant.bind(), parsed_flags.reorder_across_fake_quant.default_value(), "Normally, FakeQuant nodes must be strict boundaries for graph " "transformations, in order to ensure that quantized inference has " - "the " - "exact same arithmetic behavior as quantized training --- which is " - "the " - "whole point of quantized training and of FakeQuant nodes in the " - "first " - "place. However, that entails subtle requirements on where exactly " + "the exact same arithmetic behavior as quantized training --- which " + "is the whole point of quantized training and of FakeQuant nodes in " + "the first place. " + "However, that entails subtle requirements on where exactly " "FakeQuant nodes must be placed in the graph. Some quantized graphs " "have FakeQuant nodes at unexpected locations, that prevent graph " "transformations that are necessary in order to generate inference " "code for these graphs. Such graphs should be fixed, but as a " "temporary work-around, setting this reorder_across_fake_quant flag " - "allows toco to perform necessary graph transformaitons on them, " + "allows TOCO to perform necessary graph transformaitons on them, " "at the cost of no longer faithfully matching inference and training " "arithmetic."), Flag("allow_custom_ops", parsed_flags.allow_custom_ops.bind(), parsed_flags.allow_custom_ops.default_value(), - "If true, allow TOCO to create TF Lite Custom operators for all the" - "unsupported Tensorflow ops."), + "If true, allow TOCO to create TF Lite Custom operators for all the " + "unsupported TensorFlow ops."), Flag( "drop_control_dependency", parsed_flags.drop_control_dependency.bind(), diff --git a/tensorflow/contrib/lite/toco/toco_cmdline_flags.h b/tensorflow/contrib/lite/toco/toco_cmdline_flags.h index ba35ca8d5d23f07d843ae6fa2099cc7e15b1e9a3..46eb3f57283cc52bf2877f578500f3a4a633be86 100644 --- a/tensorflow/contrib/lite/toco/toco_cmdline_flags.h +++ b/tensorflow/contrib/lite/toco/toco_cmdline_flags.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_CMDLINE_FLAGS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_CMDLINE_FLAGS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_CMDLINE_FLAGS_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_CMDLINE_FLAGS_H_ #include #include @@ -33,4 +33,4 @@ void ReadTocoFlagsFromCommandLineFlags(const ParsedTocoFlags& parsed_toco_flags, } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_CMDLINE_FLAGS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_CMDLINE_FLAGS_H_ diff --git a/tensorflow/contrib/lite/toco/toco_graphviz_dump_options.h b/tensorflow/contrib/lite/toco/toco_graphviz_dump_options.h index ae0541f62b61581e3ba183725a85fe51c54116dc..d6c3ba6543378b3e15b5fb7816f52376fe05123d 100644 --- a/tensorflow/contrib/lite/toco/toco_graphviz_dump_options.h +++ b/tensorflow/contrib/lite/toco/toco_graphviz_dump_options.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_GRAPHVIZ_DUMP_OPTIONS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_GRAPHVIZ_DUMP_OPTIONS_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_GRAPHVIZ_DUMP_OPTIONS_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_GRAPHVIZ_DUMP_OPTIONS_H_ #include @@ -31,4 +31,4 @@ struct GraphVizDumpOptions { } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_GRAPHVIZ_DUMP_OPTIONS_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_GRAPHVIZ_DUMP_OPTIONS_H_ diff --git a/tensorflow/contrib/lite/toco/toco_port.h b/tensorflow/contrib/lite/toco/toco_port.h index b5cb7a11e7c46d02d398ff937d46e52368e88098..4be3b5a0bf00ed204a1218545d9e66f7685a50d7 100644 --- a/tensorflow/contrib/lite/toco/toco_port.h +++ b/tensorflow/contrib/lite/toco/toco_port.h @@ -12,13 +12,14 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_PORT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_PORT_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_PORT_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_PORT_H_ // Portability layer for toco tool. Mainly, abstract filesystem access so we // can build and use on google internal environments and on OSX. #include +#include "google/protobuf/text_format.h" #include "tensorflow/contrib/lite/toco/format_port.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/platform.h" @@ -75,6 +76,26 @@ void CopyToBuffer(const ::Cord& src, char* dest); #endif // PLATFORM_GOOGLE void CopyToBuffer(const string& src, char* dest); } // namespace port + +inline bool ParseFromStringOverload(const std::string& in, + TFLITE_PROTO_NS::Message* proto) { + return TFLITE_PROTO_NS::TextFormat::ParseFromString(in, proto); +} + +template +bool ParseFromStringEitherTextOrBinary(const std::string& input_file_contents, + Proto* proto) { + if (proto->ParseFromString(input_file_contents)) { + return true; + } + + if (ParseFromStringOverload(input_file_contents, proto)) { + return true; + } + + return false; +} + } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_PORT_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_PORT_H_ diff --git a/tensorflow/contrib/lite/toco/toco_tooling.cc b/tensorflow/contrib/lite/toco/toco_tooling.cc index e89421b85f8b33bdcfd9ec0f0a1f6218db63dbb2..1b836fbc151db2141ad64d5370f15a43246fdd8b 100644 --- a/tensorflow/contrib/lite/toco/toco_tooling.cc +++ b/tensorflow/contrib/lite/toco/toco_tooling.cc @@ -52,45 +52,59 @@ void MakeGeneralGraphTransformationsSet( GraphTransformationsSet* transformations) { CHECK(transformations->empty()); transformations->Add(new ConvertExpandDimsToReshape); + transformations->Add(new ConvertTrivialAddNToAdd); + transformations->Add(new ConvertTrivialStackToReshape); + transformations->Add(new ConvertTrivialTransposeToReshape); + transformations->Add(new ConvertReorderAxes); transformations->Add(new ResolveReshapeAttributes); + transformations->Add(new ResolveTransposeAttributes); transformations->Add(new PropagateArrayDataTypes); transformations->Add(new PropagateFixedSizes); transformations->Add(new RemoveTensorFlowAssert); transformations->Add(new RemoveTensorFlowIdentity); transformations->Add(new RemoveTrivialConcatenation); transformations->Add(new RemoveTrivialConcatenationInput); + transformations->Add(new RemoveTrivialSlice); transformations->Add(new RemoveUnusedOp); transformations->Add(new EnsureBiasVectors); transformations->Add(new ResolveReorderAxes); + transformations->Add(new UnrollBatchMatMul); transformations->Add(new ResolveTensorFlowMatMul); transformations->Add(new FuseBinaryIntoPrecedingAffine); transformations->Add(new FuseBinaryIntoFollowingAffine); + transformations->Add(new ReorderActivationFunctions); transformations->Add(new ResolveBatchNormalization); transformations->Add(new ResolveConstantBinaryOperator); transformations->Add(new ResolveConstantFill); + transformations->Add(new ResolveConstantRange); + transformations->Add(new ResolveConstantStack); + transformations->Add(new ResolveConstantStridedSlice); + transformations->Add(new ResolveConstantTranspose); transformations->Add(new ResolveConstantUnaryOperator); transformations->Add(new ResolveTensorFlowMerge); - transformations->Add(new ResolveTensorFlowSqueeze); + transformations->Add(new ResolveSqueezeAttributes); transformations->Add(new ResolveTensorFlowSwitch); transformations->Add(new ResolveTensorFlowTile); transformations->Add(new ResolveTensorFlowConcat); + transformations->Add(new ResolveMultiplyByZero); transformations->Add(new IdentifyL2Normalization); transformations->Add(new IdentifyL2Pool); transformations->Add(new IdentifyRelu1); transformations->Add(new RemoveTrivialBinaryOperator); transformations->Add(new ReadFakeQuantMinMax); + transformations->Add(new ResolveSpaceToBatchNDAttributes); transformations->Add(new ResolveBatchToSpaceNDAttributes); transformations->Add(new ResolvePadAttributes); transformations->Add(new ResolveStridedSliceAttributes); transformations->Add(new ResolveSliceAttributes); transformations->Add(new ResolveMeanAttributes); - transformations->Add(new ResolveConstantTensorFlowShape); + transformations->Add(new ResolveConstantShapeOrRank); transformations->Add(new MakeInitialDequantizeOperator); + transformations->Add(new ResolveConstantFakeQuant); } bool SupportsQuantization(FileFormat format) { return (format == GRAPHVIZ_DOT || format == TFLITE); - ; } bool SupportsFusedActivationFunction(FileFormat format) { @@ -98,7 +112,8 @@ bool SupportsFusedActivationFunction(FileFormat format) { } bool SupportsLstmCell(FileFormat format) { - return (format == TENSORFLOW_GRAPHDEF || format == GRAPHVIZ_DOT); + return (format == TENSORFLOW_GRAPHDEF || format == GRAPHVIZ_DOT || + format == TFLITE); } bool SupportsPreallocatedWorkspace(FileFormat format) { @@ -127,7 +142,7 @@ void SetFinalDataTypeOnInputs(const TocoFlags& toco_flags, Model* model) { for (int i = 0; i < model->flags.input_arrays_size(); i++) { string const& array_name = model->flags.input_arrays(i).name(); - auto* array = model->arrays[array_name].get(); + auto* array = &model->GetArray(array_name); // Note that the notion of changing data types only applies to real-numbers // arrays (see the documentation for inference_input_type). // TODO(benoitjacob) this is assuming that uint8 arrays are quantized, @@ -174,6 +189,11 @@ std::unique_ptr Import(const TocoFlags& toco_flags, } void Transform(const TocoFlags& toco_flags, Model* model) { + // Clean up after import. + SetFinalDataTypeOnInputs(toco_flags, model); + UseArraysExtraInfo(model); + FinishBuildingRNNStates(model); + const FileFormat output_format = toco_flags.output_format(); const IODataType inference_type = toco_flags.inference_type(); @@ -185,7 +205,12 @@ void Transform(const TocoFlags& toco_flags, Model* model) { << "Quantized inference is not allowed with float inputs."; } - SetFinalDataTypeOnInputs(toco_flags, model); + // Remove unused ops before performing any other optimizations. This is to + // stop optimizations from crossing the input/output boundaries. For example + // this will stop BatchNorm fusing if the output node is in between a conv + // and BatchNorm layers. + RunGraphTransformations(model, "Removing unused ops", + {new toco::RemoveUnusedOp}); GraphTransformationsSet transformations; MakeGeneralGraphTransformationsSet(&transformations); @@ -196,9 +221,6 @@ void Transform(const TocoFlags& toco_flags, Model* model) { } else { transformations.Add(new UnfuseActivationFunctions); } - if (output_format != TENSORFLOW_GRAPHDEF) { - transformations.Add(new ResolveConstantFakeQuant); - } if (toco_flags.drop_fake_quant()) { transformations.Add(new DropFakeQuant); } else { @@ -211,13 +233,18 @@ void Transform(const TocoFlags& toco_flags, Model* model) { } } transformations.Add(new ConvertPureConvToDepthwise); - // TFLite export does not yet support fused LSTM cell. if (SupportsLstmCell(output_format)) { transformations.Add(new IdentifyLstmCell); + if (output_format == TFLITE) { + transformations.Add(new toco::SplitLstmCellInputs); + } else { + transformations.Add(new toco::MergeLstmCellInputs); + } } transformations.Add(new ResolveConstantConcatenation); RunGraphTransformations(model, "general graph transformations", transformations); + if (quantize_output) { RunGraphTransformations(model, "pre-quantization graph transformations", {new HardcodeMinMax, new DropFakeQuant}); @@ -250,6 +277,10 @@ void Transform(const TocoFlags& toco_flags, Model* model) { dequantization_transformations); } + if (output_format == TENSORFLOW_GRAPHDEF) { + EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(model); + } + LogDump(kLogLevelModelChanged, "AFTER TRANSFORMATIONS", *model); if (output_format != GRAPHVIZ_DOT && output_format != TFLITE) { diff --git a/tensorflow/contrib/lite/toco/toco_tooling.h b/tensorflow/contrib/lite/toco/toco_tooling.h index 9c5a93a21170ba773b1160eb2e1261f85cdd70e5..e731c149eef412d3048a1d5f84145ce6ff87208d 100644 --- a/tensorflow/contrib/lite/toco/toco_tooling.h +++ b/tensorflow/contrib/lite/toco/toco_tooling.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ #include #include @@ -47,4 +47,4 @@ inline void Export(const TocoFlags& toco_flags, const Model& model, } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ diff --git a/tensorflow/contrib/lite/toco/toco_types.h b/tensorflow/contrib/lite/toco/toco_types.h index ad42497ada6cb0dbda673bf3aad406c9fedfb078..d72a3bd1f382679f81061a51f35586631b571400 100644 --- a/tensorflow/contrib/lite/toco/toco_types.h +++ b/tensorflow/contrib/lite/toco/toco_types.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TYPES_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TYPES_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TYPES_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TYPES_H_ #include #include "tensorflow/core/platform/platform.h" @@ -42,4 +42,4 @@ using tensorflow::uint8; } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TYPES_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TYPES_H_ diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc index fd56f64cc9bc0a546459cbda4021a7ac30d7aa4d..dcb409c84d8d80790f3c5e41a6eb7bce1b1efd2e 100644 --- a/tensorflow/contrib/lite/toco/tooling_util.cc +++ b/tensorflow/contrib/lite/toco/tooling_util.cc @@ -25,6 +25,7 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/strings/str_replace.h" +#include "absl/strings/str_split.h" #include "tensorflow/contrib/lite/toco/dump_graphviz.h" #include "tensorflow/contrib/lite/toco/model_flags.pb.h" #include "tensorflow/contrib/lite/toco/toco_graphviz_dump_options.h" @@ -33,6 +34,24 @@ limitations under the License. namespace toco { +// Find the longest common prefix of two strings. +absl::string_view FindLongestCommonPrefix(absl::string_view a, + absl::string_view b) { + if (a.empty() || b.empty()) return absl::string_view(); + + const char* pa = a.data(); + const char* pb = b.data(); + size_t count = 0; + const size_t limit = std::min(a.size(), b.size()); + while (count < limit && *pa == *pb) { + ++pa; + ++pb; + ++count; + } + + return absl::string_view(a.data(), count); +} + string LogName(const Operator& op) { const string& opname = HelpfulOperatorTypeName(op); if (op.outputs.empty()) { @@ -92,8 +111,18 @@ int CountOpsWithInput(const Model& model, const string& array_name) { } bool DeleteArrayIfUnused(const string& array_name, Model* model) { - if (CountOpsWithInput(*model, array_name) == 0) { - model->arrays.erase(array_name); + if (IsDiscardableArray(*model, array_name) && + CountOpsWithInput(*model, array_name) == 0) { + model->EraseArray(array_name); + return true; + } + return false; +} + +bool DeleteArrayIfUsedOnce(const string& array_name, Model* model) { + if (IsDiscardableArray(*model, array_name) && + CountOpsWithInput(*model, array_name) == 1) { + model->EraseArray(array_name); return true; } return false; @@ -141,6 +170,18 @@ std::vector>::const_iterator FindOpWithInput( return model.operators.end(); } +std::vector>::iterator FindOpWithInput( + Model& model, const string& array_name) { + for (auto it = model.operators.begin(); it != model.operators.end(); ++it) { + for (auto& input : it->get()->inputs) { + if (input == array_name) { + return it; + } + } + } + return model.operators.end(); +} + std::vector>::const_iterator FindOp( const Model& model, const Operator* op) { for (auto it = model.operators.begin(); it != model.operators.end(); ++it) { @@ -197,7 +238,9 @@ const char* OperatorTypeName(OperatorType type) { case OperatorType::k##c: \ return #c; HANDLE_OPERATORTYPENAME_CASE(Add) + HANDLE_OPERATORTYPENAME_CASE(AddN) HANDLE_OPERATORTYPENAME_CASE(AveragePool) + HANDLE_OPERATORTYPENAME_CASE(BatchMatMul) HANDLE_OPERATORTYPENAME_CASE(BatchNormalization) HANDLE_OPERATORTYPENAME_CASE(Conv) HANDLE_OPERATORTYPENAME_CASE(Concatenation) @@ -219,6 +262,7 @@ const char* OperatorTypeName(OperatorType type) { HANDLE_OPERATORTYPENAME_CASE(Relu6) HANDLE_OPERATORTYPENAME_CASE(ReorderAxes) HANDLE_OPERATORTYPENAME_CASE(Softmax) + HANDLE_OPERATORTYPENAME_CASE(LogSoftmax) HANDLE_OPERATORTYPENAME_CASE(Div) HANDLE_OPERATORTYPENAME_CASE(Tanh) HANDLE_OPERATORTYPENAME_CASE(TensorFlowAll) @@ -269,7 +313,9 @@ const char* OperatorTypeName(OperatorType type) { HANDLE_OPERATORTYPENAME_CASE(Mean) HANDLE_OPERATORTYPENAME_CASE(Svdf) HANDLE_OPERATORTYPENAME_CASE(ArgMax) + HANDLE_OPERATORTYPENAME_CASE(TopK_V2) HANDLE_OPERATORTYPENAME_CASE(TensorFlowUnsupported) + HANDLE_OPERATORTYPENAME_CASE(Exp) default: LOG(FATAL) << "Unhandled op type"; #undef HANDLE_OPERATORTYPENAME_CASE @@ -285,9 +331,23 @@ string HelpfulOperatorTypeName(const Operator& op) { return OperatorTypeName(op.type); } +bool OperatorSupportsFusedActivation(OperatorType type) { + switch (type) { + case OperatorType::kConcatenation: + case OperatorType::kGather: + case OperatorType::kSlice: + case OperatorType::kSqueeze: + case OperatorType::kTensorFlowReshape: + case OperatorType::kTensorFlowSplit: + return false; + default: + return true; + } +} + void LogSummary(int log_level, const Model& model) { VLOG(log_level) << "Operators summary (" << model.operators.size() - << " operators): "; + << " operators):"; std::unordered_multiset ops_by_type; for (const auto& op : model.operators) { ops_by_type.insert(op->type); @@ -374,7 +434,7 @@ void LogArray(int log_level, const Model& model, const string& name) { } if (array.quantization_params) { VLOG(log_level) << " QuantizationParams: zero_point=" - << array.quantization_params->zero_point + << static_cast(array.quantization_params->zero_point) << ", scale=" << array.quantization_params->scale; } } @@ -404,6 +464,7 @@ void DumpGraphvizVideoFrame(const Model& model) { DumpGraphviz(model, &graphviz_dump); std::size_t hash = std::hash{}(graphviz_dump); if (!dump_hashes.count(hash)) { + LOG(INFO) << "DUMPING GRAPHVIZ VIDEO FRAME: " << dump_id; dump_hashes.insert(hash); CHECK(port::file::SetContents( port::file::JoinPath( @@ -447,7 +508,7 @@ void LogDump(int log_level, const string& message, const Model& model) { LogArray(log_level, model, input); } } - VLOG(log_level) << HelpfulOperatorTypeName(*op) << " : "; + VLOG(log_level) << HelpfulOperatorTypeName(*op) << " :"; VLOG(log_level) << " " << FormatArraysList(model, op->inputs) << " -> " << FormatArraysList(model, op->outputs); if (op->fused_activation_function != FusedActivationFunctionType::kNone) { @@ -565,14 +626,22 @@ int RequiredBufferSizeForShape(const Shape& shape) { } bool IsConstantParameterArray(const Model& model, const string& name) { - if (!model.arrays.count(name)) { + if (!model.HasArray(name)) { return false; } - return !!model.arrays.at(name)->buffer; + return !!model.GetArray(name).buffer; } namespace { +// Take an array name, which may be something like "name:3_5" and make it +// acceptable as a TF node name, say "name_3_5"; +string SanitizeNameForTFNode(const string& array_name) { + auto node_name = array_name; + std::replace(node_name.begin(), node_name.end(), ':', '_'); + return node_name; +} + void CheckInputArraysAreNotOutputArrays(const ModelFlags& model_flags) { for (const auto& input_array : model_flags.input_arrays()) { for (const string& output_array : model_flags.output_arrays()) { @@ -632,17 +701,17 @@ void CheckNonExistentIOArrays(const Model& model) { return; } for (const auto& input_array : model.flags.input_arrays()) { - CHECK(model.arrays.count(input_array.name())) + CHECK(model.HasArray(input_array.name())) << "Input array not found: " << input_array.name(); } for (const string& output_array : model.flags.output_arrays()) { - CHECK(model.arrays.count(output_array)) + CHECK(model.HasArray(output_array)) << "Output array not found: " << output_array; } for (const auto& rnn_state : model.flags.rnn_states()) { if (!rnn_state.discardable()) { - CHECK(model.arrays.count(rnn_state.state_array())); - CHECK(model.arrays.count(rnn_state.back_edge_source_array())); + CHECK(model.HasArray(rnn_state.state_array())); + CHECK(model.HasArray(rnn_state.back_edge_source_array())); } } } @@ -651,10 +720,11 @@ void CheckNonExistentIOArrays(const Model& model) { void CheckNoMissingArray(const Model& model) { for (const auto& op : model.operators) { for (const auto& input : op->inputs) { - CHECK(model.arrays.count(input)); + CHECK(model.HasArray(input) || model.optional_arrays.count(input)) + << "Input: " << input << " missing for op: " << op->outputs[0] << "."; } for (const auto& output : op->outputs) { - CHECK(model.arrays.count(output)); + CHECK(model.HasArray(output)) << "Output: " << output << " missing."; } } CheckNonExistentIOArrays(model); @@ -663,12 +733,12 @@ void CheckNoMissingArray(const Model& model) { void FixNoMissingArray(Model* model) { for (const auto& op : model->operators) { for (const auto& input : op->inputs) { - if (!model->arrays.count(input)) { + if (!model->HasArray(input)) { model->GetOrCreateArray(input); } } for (const auto& output : op->outputs) { - if (!model->arrays.count(output)) { + if (!model->HasArray(output)) { model->GetOrCreateArray(output); } } @@ -686,7 +756,7 @@ void FixNoMissingArray(Model* model) { void CheckNoOrphanedArray(const Model& model) { std::unordered_set arrays_without_known_use; - for (const auto& array : model.arrays) { + for (const auto& array : model.GetArrayMap()) { if (IsDiscardableArray(model, array.first)) { arrays_without_known_use.insert(array.first); } @@ -713,7 +783,7 @@ void CheckNoOrphanedArray(const Model& model) { void FixNoOrphanedArray(Model* model) { std::unordered_set arrays_without_known_use; - for (const auto& array : model->arrays) { + for (const auto& array : model->GetArrayMap()) { arrays_without_known_use.insert(array.first); } for (const auto& op : model->operators) { @@ -730,13 +800,16 @@ void FixNoOrphanedArray(Model* model) { } for (const auto& array : arrays_without_known_use) { if (IsDiscardableArray(*model, array)) { - model->arrays.erase(array); + model->EraseArray(array); } } } -void CheckArrayFieldsConsistent(const Model& model) { - for (const auto& array_entry : model.arrays) { +// Apply checks to arrays individually (for-each fashion). +// +// Check consistency of array fields, check name. +void CheckEachArray(const Model& model) { + for (const auto& array_entry : model.GetArrayMap()) { const auto& array = array_entry.second; if (array->has_shape()) { for (int d : array->shape().dims()) { @@ -750,16 +823,30 @@ void CheckArrayFieldsConsistent(const Model& model) { if (array->buffer) { CHECK(array->buffer->type == array->data_type); } + + // Check name. Either "name_with_suffix_8", "name_with_port:3", but not + // "name_with_both:3_8". + const string& name = array_entry.first; + auto colon_pos = name.find_first_of(":"); + if (colon_pos != string::npos) { + CHECK_EQ(name.substr(colon_pos + 1).find_first_not_of("0123456789"), + string::npos) + << "Array name must only have digits after colon"; + } + CHECK_GT(colon_pos, 0) + << "First character of array name must not be a colon."; } } void CheckOperatorOrdering(const Model& model) { std::unordered_set arrays_behind_us; - for (const auto& array_entry : model.arrays) { + for (const auto& array_entry : model.GetArrayMap()) { if (!GetOpWithOutput(model, array_entry.first)) { arrays_behind_us.insert(array_entry.first); } } + arrays_behind_us.insert(model.optional_arrays.begin(), + model.optional_arrays.end()); for (const auto& op : model.operators) { for (const auto& input : op->inputs) { if (!IsConstantParameterArray(model, input)) { @@ -778,11 +865,13 @@ void CheckOperatorOrdering(const Model& model) { void FixOperatorOrdering(Model* model) { std::unordered_set arrays_behind_us; - for (const auto& array_entry : model->arrays) { + for (const auto& array_entry : model->GetArrayMap()) { if (!GetOpWithOutput(*model, array_entry.first)) { arrays_behind_us.insert(array_entry.first); } } + arrays_behind_us.insert(model->optional_arrays.begin(), + model->optional_arrays.end()); std::vector> old_operators; std::swap(old_operators, model->operators); std::set remaining; @@ -894,7 +983,7 @@ void CheckInvariants(const Model& model) { CheckNonAsciiIOArrays(model.flags); CheckNoMissingArray(model); CheckNoOrphanedArray(model); - CheckArrayFieldsConsistent(model); + CheckEachArray(model); CheckOperatorOrdering(model); } @@ -931,7 +1020,8 @@ void CheckModelCounts(const Model& model) { if (count_type == "None") { continue; } else if (count_type == "Arrays") { - CheckCountInRange(model_check, model.arrays.size(), "count of arrays"); + CheckCountInRange(model_check, model.GetArrayMap().size(), + "count of arrays"); } else if (count_type == "Total") { CheckCountInRange(model_check, model.operators.size(), "count of all operator instances"); @@ -951,7 +1041,9 @@ void CheckModelCounts(const Model& model) { void MakeArrayDims(int num_dims, int batch, int height, int width, int depth, std::vector* out_dims) { CHECK(out_dims->empty()); - if (num_dims == 1) { + if (num_dims == 0) { + return; + } else if (num_dims == 1) { CHECK_EQ(batch, 1); *out_dims = {depth}; } else if (num_dims == 2) { @@ -983,13 +1075,10 @@ void CreateOrCheckRnnStateArray(const string& name, int size, Model* model) { if (array.has_shape()) { num_dims = array.shape().dimensions_count(); } - std::vector dims; - MakeArrayDims(num_dims, batch, 1, 1, size, &dims); - CHECK(array.data_type == ArrayDataType::kFloat || - array.data_type == ArrayDataType::kNone); - array.data_type = ArrayDataType::kFloat; - if (!array.has_shape()) { + if (!array.has_shape() && num_dims >= 0) { Shape* shape = array.mutable_shape(); + std::vector dims; + MakeArrayDims(num_dims, batch, 1, 1, size, &dims); *shape->mutable_dims() = dims; } } @@ -1009,7 +1098,7 @@ void ResolveModelFlags(const ModelFlags& model_flags, Model* model) { } } if (!dst_input_array) { - // specified_input_array from model_flags is not found in model->flags. + // Specified_input_array from model_flags is not found in model->flags. // Match a name-less specified input array when there can be no ambiguity // as there is only 1 input array. if (model->flags.input_arrays_size() == 1 && @@ -1178,9 +1267,6 @@ void ResolveModelFlags(const ModelFlags& model_flags, Model* model) { } // Creation of the RNN state arrays for (const auto& rnn_state : model->flags.rnn_states()) { - if (!rnn_state.manually_create()) { - continue; - } CreateOrCheckRnnStateArray(rnn_state.state_array(), rnn_state.size(), model); } @@ -1194,6 +1280,9 @@ void ResolveModelFlags(const ModelFlags& model_flags, Model* model) { model->flags.set_allow_nonascii_arrays(model_flags.allow_nonascii_arrays()); model->flags.set_allow_nonexistent_arrays( model_flags.allow_nonexistent_arrays()); + + CHECK(!model->flags.has_arrays_extra_info()); + *model->flags.mutable_arrays_extra_info() = model_flags.arrays_extra_info(); } void CheckIsReadyForQuantization(const Model& model) { @@ -1253,12 +1342,23 @@ int ElementSize(ArrayDataType data_type) { switch (data_type) { case ArrayDataType::kFloat: return 4; - case ArrayDataType::kInt32: - return 4; + case ArrayDataType::kInt8: + return 1; case ArrayDataType::kUint8: return 1; + case ArrayDataType::kInt16: + return 2; + case ArrayDataType::kUint16: + return 2; + case ArrayDataType::kInt32: + return 4; + case ArrayDataType::kUint32: + return 4; case ArrayDataType::kInt64: return 8; + case ArrayDataType::kUint64: + return 8; + // Usually not critical limitation because strings are only input and/or // output. case ArrayDataType::kString: @@ -1280,6 +1380,8 @@ void DropMinMax(Model* model, const string& array_name) { } bool IsAllocatableTransientArray(const Model& model, const string& array_name) { + // Optional array is not transient + if (model.IsOptionalArray(array_name)) return false; // The model's input and output arrays are externally allocated. // They are not transient arrays. if (IsInputArray(model, array_name)) { @@ -1290,7 +1392,7 @@ bool IsAllocatableTransientArray(const Model& model, const string& array_name) { return false; } } - const auto& array = model.arrays.at(array_name); + const auto& array = &model.GetArray(array_name); // An array with a constant buffer isn't a transient array. if (!!array->buffer) { return false; @@ -1303,18 +1405,23 @@ bool IsAllocatableTransientArray(const Model& model, const string& array_name) { } string AvailableArrayName(const Model& model, const string& name) { - if (!model.arrays.count(name)) { - return name; + string sanitized_name = SanitizeNameForTFNode(name); + if (!model.HasArray(sanitized_name) && + !model.IsOptionalArray(sanitized_name)) { + return sanitized_name; } const int kNumSuffixesToTry = 1000; for (int i = 0; i < kNumSuffixesToTry; i++) { - const string& name_with_suffix = toco::port::StringF("%s_%d", name, i); - if (!model.arrays.count(name_with_suffix)) { + const string& name_with_suffix = + toco::port::StringF("%s_%d", sanitized_name, i); + if (!model.HasArray(name_with_suffix) && + !model.IsOptionalArray(name_with_suffix)) { return name_with_suffix; } } - LOG(FATAL) << "Could not find an available array name starting with " << name - << ". Tried " << kNumSuffixesToTry << " suffixes, all were taken!"; + LOG(FATAL) << "Could not find an available array name starting with " + << sanitized_name << ". Tried " << kNumSuffixesToTry + << " suffixes, all were taken!"; return ""; } @@ -1327,12 +1434,12 @@ string ShapeToString(const Shape& shape) { } void PrintArrayShape(Model* model, const string& name) { - if (!model->arrays[name]->has_shape()) { + if (!model->GetArray(name).has_shape()) { LOG(INFO) << name << " has no shape"; return; } LOG(INFO) << name - << " has shape: " << ShapeToString(model->arrays[name]->shape()); + << " has shape: " << ShapeToString(model->GetArray(name).shape()); } bool IsArrayFullyConnectedWeights(const Model& model, const string& name) { @@ -1353,6 +1460,21 @@ bool IsArrayFullyConnectedWeights(const Model& model, const string& name) { return is_fc_weights; } +string CreateInt32Array(Model* model, const string& param_name, + const std::vector& value) { + auto param_array_name = AvailableArrayName(*model, param_name); + auto& param_array = model->GetOrCreateArray(param_array_name); + param_array.mutable_shape()->ReplaceDims({static_cast(value.size())}); + param_array.data_type = ArrayDataType::kInt32; + auto& param_array_data = + param_array.GetMutableBuffer().data; + param_array_data.resize(RequiredBufferSizeForShape(param_array.shape())); + for (int i = 0; i < value.size(); ++i) { + param_array_data[i] = value[i]; + } + return param_array_name; +} + bool EstimateArithmeticOpsCount(const Model& model, int64* result) { int64 total = 0; for (const auto& op : model.operators) { @@ -1388,8 +1510,19 @@ bool EstimateArithmeticOpsCount(const Model& model, int64* result) { total += RequiredBufferSizeForShape(output_array.shape()); break; } + case OperatorType::kAddN: { + const auto& output_array = model.GetArray(op->outputs[0]); + if (!output_array.has_shape()) { + return false; + } + // AddN cost is roughly the same cost as N-1 Adds. + const int num_adds = op->inputs.size() - 1; + total += num_adds * RequiredBufferSizeForShape(output_array.shape()); + break; + } case OperatorType::kLogistic: case OperatorType::kSoftmax: + case OperatorType::kLogSoftmax: case OperatorType::kTanh: { const auto& output_array = model.GetArray(op->outputs[0]); if (!output_array.has_shape()) { @@ -1455,8 +1588,6 @@ bool EstimateArithmeticOpsCount(const Model& model, int64* result) { return true; } -namespace { - void GetShuffleShape(AxesOrder input_axes_order, AxesOrder output_axes_order, std::vector* shuffle) { CHECK_EQ(AxesCount(input_axes_order), AxesCount(output_axes_order)); @@ -1491,8 +1622,6 @@ void GetShuffleShape(AxesOrder input_axes_order, AxesOrder output_axes_order, } } -// Extend shuffle is designed to match ExtendShape, which pads the shape with -// unit dimensions at the beginning. void ExtendShuffle(const std::vector& input_shuffle, int newdim, std::vector* extended_shuffle) { *extended_shuffle = input_shuffle; @@ -1507,8 +1636,6 @@ void ExtendShuffle(const std::vector& input_shuffle, int newdim, } } -} // end anonymous namespace - void ShuffleDims(const Shape& input_shape, AxesOrder input_axes_order, AxesOrder output_axes_order, Shape* output_shape) { if (input_axes_order == AxesOrder::kHWIM && @@ -1528,9 +1655,11 @@ void ShuffleDims(const Shape& input_shape, AxesOrder input_axes_order, } } -void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, - AxesOrder output_axes_order, const Shape& output_shape, - const float* input_data, float* output_data) { +template +void ShuffleArrayTemplate(const Shape& input_shape, AxesOrder input_axes_order, + AxesOrder output_axes_order, + const Shape& output_shape, const T* input_data, + T* output_data) { if (input_axes_order == AxesOrder::kHWIM && output_axes_order == AxesOrder::k1HWO) { // This special case isn't just a permutation, the IM pair of dims get @@ -1582,16 +1711,15 @@ void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, const int output_stride_3 = output_stride_2 * output_size_2; for (int i3 = 0; i3 < output_size_3; i3++) { - const float* const input_ptr_3 = input_data + i3 * input_stride_3; - float* const output_ptr_3 = output_data + i3 * output_stride_3; + const T* const input_ptr_3 = input_data + i3 * input_stride_3; + T* const output_ptr_3 = output_data + i3 * output_stride_3; for (int i2 = 0; i2 < output_size_2; i2++) { - const float* const input_ptr_2 = input_ptr_3 + i2 * input_stride_2; - float* const output_ptr_2 = output_ptr_3 + i2 * output_stride_2; + const T* const input_ptr_2 = input_ptr_3 + i2 * input_stride_2; + T* const output_ptr_2 = output_ptr_3 + i2 * output_stride_2; for (int i1 = 0; i1 < output_size_1; i1++) { - const float* input_ptr = input_ptr_2 + i1 * input_stride_1; - float* output_ptr = output_ptr_2 + i1 * output_stride_1; - float* const output_ptr_end = - output_ptr + output_size_0 * output_stride_0; + const T* input_ptr = input_ptr_2 + i1 * input_stride_1; + T* output_ptr = output_ptr_2 + i1 * output_stride_1; + T* const output_ptr_end = output_ptr + output_size_0 * output_stride_0; while (output_ptr != output_ptr_end) { *output_ptr = *input_ptr; input_ptr += input_stride_0; @@ -1602,6 +1730,20 @@ void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, } } +void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, + AxesOrder output_axes_order, const Shape& output_shape, + const uint8* input_data, uint8* output_data) { + ShuffleArrayTemplate(input_shape, input_axes_order, output_axes_order, + output_shape, input_data, output_data); +} + +void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, + AxesOrder output_axes_order, const Shape& output_shape, + const float* input_data, float* output_data) { + ShuffleArrayTemplate(input_shape, input_axes_order, output_axes_order, + output_shape, input_data, output_data); +} + int AxesCount(AxesOrder axes_order) { switch (axes_order) { case AxesOrder::kOneAxis: @@ -1651,7 +1793,7 @@ bool IsDiscardableArray(const Model& model, const string& array_name) { } void CheckFinalDataTypesSatisfied(const Model& model) { - for (const auto& array_entry : model.arrays) { + for (const auto& array_entry : model.GetArrayMap()) { const auto& array = *array_entry.second; if (array.final_data_type != ArrayDataType::kNone) { CHECK(array.final_data_type == array.data_type) @@ -1678,4 +1820,32 @@ ArrayDataType ConvertIODataTypeToArrayDataType(IODataType type) { } } +void FinishBuildingRNNStates(Model* model) { + for (const auto& rnn_state : model->flags.rnn_states()) { + if (!model->HasArray(rnn_state.back_edge_source_array()) || + !model->HasArray(rnn_state.state_array())) { + CHECK(model->HasArray(rnn_state.back_edge_source_array())); + CHECK(model->HasArray(rnn_state.state_array())); + continue; + } + const auto& src_array = model->GetArray(rnn_state.back_edge_source_array()); + auto& dst_array = model->GetArray(rnn_state.state_array()); + if (src_array.data_type == ArrayDataType::kNone && + dst_array.data_type == ArrayDataType::kNone) { + dst_array.data_type = ArrayDataType::kFloat; + } + } +} + +void UseArraysExtraInfo(Model* model) { + for (const auto& entry : model->flags.arrays_extra_info().entries()) { + QCHECK(model->HasArray(entry.name())) + << "ArraysExtraInfo refers to non-existent array name: " + << entry.name(); + auto& minmax = model->GetArray(entry.name()).GetOrCreateMinMax(); + minmax.min = entry.min(); + minmax.max = entry.max(); + } +} + } // namespace toco diff --git a/tensorflow/contrib/lite/toco/tooling_util.h b/tensorflow/contrib/lite/toco/tooling_util.h index d820d619d0de425407e88076082a3e0f8d4783a9..0aaa0f6a215288430dfdde7d5042012730c3be4c 100644 --- a/tensorflow/contrib/lite/toco/tooling_util.h +++ b/tensorflow/contrib/lite/toco/tooling_util.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_ #include #include @@ -23,7 +23,7 @@ limitations under the License. #include #include -#include "google/protobuf/text_format.h" +#include "absl/strings/string_view.h" #include "tensorflow/core/platform/logging.h" #if TOCO_SUPPORT_PORTABLE_PROTOS #include "third_party/protobuf/src/google/protobuf/text_format.h" @@ -50,6 +50,8 @@ namespace toco { constexpr int kLogLevelModelChanged = 1; constexpr int kLogLevelModelUnchanged = 2; +absl::string_view FindLongestCommonPrefix(absl::string_view a, + absl::string_view b); string LogName(const Operator& op); bool IsInputArray(const Model& model, const string& name); @@ -58,6 +60,7 @@ int CountTrueOutputs(const Model& model, const Operator& op); int CountOpsWithInput(const Model& model, const string& array_name); bool DeleteArrayIfUnused(const string& array_name, Model* model); +bool DeleteArrayIfUsedOnce(const string& array_name, Model* model); std::vector>::const_iterator FindOpWithOutput( const Model& model, const string& array_name); @@ -65,10 +68,15 @@ Operator* GetOpWithOutput(const Model& model, const string& array_name); std::vector>::iterator FindOpWithOutput( Model& model, const string& array_name); + Operator* GetOpWithOutput(const Model& model, const string& array_name); std::vector>::const_iterator FindOpWithInput( const Model& model, const string& array_name); + +std::vector>::iterator FindOpWithInput( + Model& model, const string& array_name); + Operator* GetOpWithInput(const Model& model, const string& array_name); Operator* GetFirstOpWithInput(const Model& model, const string& array_name); @@ -80,29 +88,12 @@ std::vector>::iterator FindOp(Model& model, const char* OperatorTypeName(OperatorType type); string HelpfulOperatorTypeName(const Operator& op); +bool OperatorSupportsFusedActivation(OperatorType type); + void DumpGraphvizVideoFrame(const Model& model); void LogDump(int log_level, const string& message, const Model& model); void LogSummary(int log_level, const string& message, const Model& model); -inline bool ParseFromStringOverload(const std::string& in, - TFLITE_PROTO_NS::Message* proto) { - return TFLITE_PROTO_NS::TextFormat::ParseFromString(in, proto); -} - -template -bool ParseFromStringEitherTextOrBinary(const std::string& input_file_contents, - Proto* proto) { - if (proto->ParseFromString(input_file_contents)) { - return true; - } - - if (ParseFromStringOverload(input_file_contents, proto)) { - return true; - } - - return false; -} - // TODO(b/36075966): Clean up when dims superseded by array shape. void ExtendShape(Shape* shape, int new_shape_size); @@ -270,15 +261,33 @@ void PrintArrayShape(Model* model, const string& name); void MakeArrayDims(int num_dims, int batch, int height, int width, int depth, std::vector* out_dims); +// Defines a constant int32 array with the provided values formatted for use +// as op parameters. +string CreateInt32Array(Model* model, const string& param_name, + const std::vector& value); + bool EstimateArithmeticOpsCount(const Model& model, int64* result); int AxesCount(AxesOrder axes_order); +// Returns the permutation of the dimensions based on the input axes order and +// output axes order. +void GetShuffleShape(AxesOrder input_axes_order, AxesOrder output_axes_order, + std::vector* shuffle); + +// Extend shuffle is designed to match ExtendShape, which pads the shape with +// unit dimensions at the beginning. +void ExtendShuffle(const std::vector& input_shuffle, int newdim, + std::vector* extended_shuffle); + void ShuffleDims(const Shape& input_shape, AxesOrder input_axes_order, AxesOrder output_axes_order, Shape* output_shape); void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, AxesOrder output_axes_order, const Shape& output_shape, const float* input_data, float* output_data); +void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, + AxesOrder output_axes_order, const Shape& output_shape, + const uint8* input_data, uint8* output_data); // Returns true if it may be OK for any graph transformation to ever discard // that array. The idea is that we can't ever discard arrays that are either @@ -290,6 +299,25 @@ void CheckFinalDataTypesSatisfied(const Model& model); ArrayDataType ConvertIODataTypeToArrayDataType(IODataType type); +// The process of building models varies according to the import format. +// +// (a) In some cases, such as model-proto format, the model should be fully +// specified. In these cases, no extra action should be taken by this function. +// (b) In other cases, such as TF graphdef format, the desired types of RNN +// arrays are not specified directly in the model, neither can they be inferred. +// However, we can set the types of RNN destination arrays to float. This breaks +// any cycles such as when resolution of the type of an RNN source array depends +// on the type of its destination array. +// +// This function is applied after the main import, after resolution of flags and +// after application of ArraysExtraInfo. It only defaults destination RNN arrays +// to float. If the model is subsequently quantized, it is assumed that the +// model contains sufficient information for that to be completed. If it is +// already quantized, then case (a) should hold. +void FinishBuildingRNNStates(Model* model); + +void UseArraysExtraInfo(Model* model); + } // namespace toco -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_ diff --git a/tensorflow/contrib/lite/tools/BUILD b/tensorflow/contrib/lite/tools/BUILD index 751682215bce37a8e4b8befe70b5288617053b54..999ccf2ebc009b6b7c50a9a2d1667d69a3f690e7 100644 --- a/tensorflow/contrib/lite/tools/BUILD +++ b/tensorflow/contrib/lite/tools/BUILD @@ -6,6 +6,16 @@ licenses(["notice"]) # Apache 2.0 load("//tensorflow:tensorflow.bzl", "tf_cc_binary") +py_binary( + name = "visualize", + srcs = ["visualize.py"], + data = [ + "//tensorflow/contrib/lite/schema:schema.fbs", + "@flatbuffers//:flatc", + ], + srcs_version = "PY2AND3", +) + tf_cc_binary( name = "generate_op_registrations", srcs = ["gen_op_registration_main.cc"], @@ -17,6 +27,27 @@ tf_cc_binary( ], ) +tf_cc_binary( + name = "benchmark_model", + srcs = ["benchmark_model.cc"], + linkopts = select({ + "//tensorflow:android": [ + "-pie", + "-landroid", + "-lm", + "-z defs", + "-Wl,--exclude-libs,ALL", # Exclude syms in all libs from auto export + ], + "//conditions:default": [], + }), + deps = [ + ":mutable_op_resolver", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite:string_util", + "//tensorflow/contrib/lite/kernels:builtin_ops", + ], +) + cc_library( name = "gen_op_registration", srcs = ["gen_op_registration.cc"], @@ -62,3 +93,34 @@ filegroup( ), visibility = ["//tensorflow:__subpackages__"], ) + +cc_library( + name = "verifier", + srcs = ["verifier.cc"], + hdrs = ["verifier.h"], + deps = [ + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite:schema_fbs_version", + "//tensorflow/contrib/lite:string_util", + "//tensorflow/contrib/lite/schema:schema_fbs", + "@com_google_absl//absl/base:core_headers", + ], +) + +cc_test( + name = "verifier_test", + size = "small", + srcs = ["verifier_test.cc"], + deps = [ + ":mutable_op_resolver", + ":verifier", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite:schema_fbs_version", + "//tensorflow/contrib/lite:string_util", + "//tensorflow/contrib/lite/schema:schema_fbs", + "//tensorflow/contrib/lite/testing:util", + "//tensorflow/core:framework_lite", + "@com_google_googletest//:gtest", + "@flatbuffers", + ], +) diff --git a/tensorflow/contrib/lite/tools/gen_op_registration.h b/tensorflow/contrib/lite/tools/gen_op_registration.h index 318859e23d7b404c130f003b0e249893f2ed92fe..5f2ac6ca97fde9a2fe6f4bcf20184f6ef6606f0b 100644 --- a/tensorflow/contrib/lite/tools/gen_op_registration.h +++ b/tensorflow/contrib/lite/tools/gen_op_registration.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOOLS_GEN_OP_REGISTRATION_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOOLS_GEN_OP_REGISTRATION_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOOLS_GEN_OP_REGISTRATION_H_ +#define TENSORFLOW_CONTRIB_LITE_TOOLS_GEN_OP_REGISTRATION_H_ #include "tensorflow/contrib/lite/model.h" #include "tensorflow/contrib/lite/string.h" @@ -36,4 +36,4 @@ void ReadOpsFromModel(const ::tflite::Model* model, } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOOLS_GEN_OP_REGISTRATION_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOOLS_GEN_OP_REGISTRATION_H_ diff --git a/tensorflow/contrib/lite/tools/mutable_op_resolver.h b/tensorflow/contrib/lite/tools/mutable_op_resolver.h index 906553da570720a0c4b90bbd2eebb6d8bdea6bb8..573a359c458acb6e4320c5a21cb378cdde720924 100644 --- a/tensorflow/contrib/lite/tools/mutable_op_resolver.h +++ b/tensorflow/contrib/lite/tools/mutable_op_resolver.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOOLS_MUTABLE_OP_RESOLVER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOOLS_MUTABLE_OP_RESOLVER_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_TOOLS_MUTABLE_OP_RESOLVER_H_ +#define TENSORFLOW_CONTRIB_LITE_TOOLS_MUTABLE_OP_RESOLVER_H_ #include #include "tensorflow/contrib/lite/context.h" @@ -52,4 +52,4 @@ class MutableOpResolver : public OpResolver { } // namespace tflite -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_TOOLS_MUTABLE_OP_RESOLVER_H_ +#endif // TENSORFLOW_CONTRIB_LITE_TOOLS_MUTABLE_OP_RESOLVER_H_ diff --git a/tensorflow/contrib/lite/tools/verifier.cc b/tensorflow/contrib/lite/tools/verifier.cc new file mode 100644 index 0000000000000000000000000000000000000000..59c74205f0a311ec12ff87f46622041605fb493b --- /dev/null +++ b/tensorflow/contrib/lite/tools/verifier.cc @@ -0,0 +1,234 @@ +/* Copyright 2017 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/lite/tools/verifier.h" +#include +#include "tensorflow/contrib/lite/schema/schema_generated.h" +#include "tensorflow/contrib/lite/string_util.h" +#include "tensorflow/contrib/lite/version.h" + +namespace tflite { + +namespace { + +// Reports error message when the reporter is set. +void ReportError(ErrorReporter* error_reporter, const char* format, ...) { + if (error_reporter) { + va_list args; + va_start(args, format); + error_reporter->Report(format, args); + va_end(args); + } +} + +// Returns the int32_t value pointed by ptr. +const uint32_t* GetIntPtr(const char* ptr) { + return reinterpret_cast(ptr); +} + +// Verifies flatbuffer format of the model contents and returns the in-memory +// model. +const Model* VerifyFlatbufferAndGetModel(const void* buf, size_t len) { + ::flatbuffers::Verifier verifier(static_cast(buf), len); + if (VerifyModelBuffer(verifier)) { + return ::tflite::GetModel(buf); + } else { + return nullptr; + } +} + +const uint32_t kMaxNumString = UINT_MAX / sizeof(int32_t) - 2; + +// Verifies string tensor has legit buffer contents that follow the schema +// defined in lite/string_util.h +bool VerifyStringTensorBuffer(const Buffer& buffer, + ErrorReporter* error_reporter) { + uint32_t buffer_size = buffer.data()->size(); + const char* buffer_ptr = reinterpret_cast(buffer.data()->data()); + + uint32_t num_strings = *GetIntPtr(buffer_ptr); + if (num_strings > kMaxNumString) { + ReportError(error_reporter, + "String tensor has invalid num of string set: %d", num_strings); + return false; + } + uint32_t header_offsets = + static_cast(num_strings + 2) * sizeof(int32_t); + + if (buffer_size < header_offsets) { + ReportError(error_reporter, + "String tensor buffer requires at least %d bytes, but is " + "allocated with %d bytes", + header_offsets, buffer_size); + return false; + } + + uint32_t prev_ptr = header_offsets; + uint32_t offset = sizeof(int32_t); + + if (*GetIntPtr(buffer_ptr + offset) != header_offsets) { + ReportError(error_reporter, + "String tensor buffer initial offset must be: %d", + header_offsets); + return false; + } + offset += sizeof(int32_t); + for (int i = 1; i <= num_strings; i++, offset += sizeof(int32_t)) { + int string_offset = *GetIntPtr(buffer_ptr + offset); + if (string_offset < prev_ptr || string_offset > buffer_size) { + ReportError(error_reporter, "String tensor buffer is invalid: index %d", + i); + return false; + } + } + if (*GetIntPtr(buffer_ptr + offset - sizeof(int32_t)) != buffer_size) { + ReportError(error_reporter, "String tensor buffer last offset must be %d", + buffer_size); + return false; + } + return true; +} + +// Verifies numeric tensor has legit buffer. +bool VerifyNumericTensorBuffer(const Tensor& tensor, const Buffer& buffer, + ErrorReporter* error_reporter) { + uint64_t bytes_required = 1; + for (int dim : *tensor.shape()) { + bytes_required *= dim; + if (bytes_required > UINT_MAX) { + ReportError(error_reporter, "Tensor dimension overflow"); + return false; + } + } + switch (tensor.type()) { + case TensorType_FLOAT32: + bytes_required *= sizeof(float); + break; + case TensorType_INT32: + bytes_required *= sizeof(int32_t); + break; + case TensorType_UINT8: + bytes_required *= sizeof(uint8_t); + break; + case TensorType_INT64: + bytes_required *= sizeof(int64_t); + break; + case TensorType_FLOAT16: + // FALLTHROUGH_INTENDED; + default: + ReportError(error_reporter, "Invalid tensor type: %d", tensor.type()); + return false; + } + if (bytes_required > UINT_MAX) { + ReportError(error_reporter, "Tensor dimension overflow"); + return false; + } + + if (bytes_required != buffer.data()->size()) { + ReportError( + error_reporter, + "Tensor requires %d bytes, but is allocated with %d bytes buffer", + bytes_required, buffer.data()->size()); + return false; + } + return true; + + // TODO(yichengfan): verify quantized tensors. +} + +// Verifies tensors have valid properties and legit buffer if set. +bool VerifyTensors(const Model& model, ErrorReporter* error_reporter) { + if (!model.subgraphs()) { + return true; + } + for (const auto& subgraph : *model.subgraphs()) { + if (!subgraph->tensors()) { + continue; + } + for (const auto& tensor : *subgraph->tensors()) { + if (!tensor->buffer()) { + continue; + } + if (tensor->buffer() >= model.buffers()->size()) { + ReportError(error_reporter, "Invalid tensor buffer index: %d", + tensor->buffer()); + return false; + } + auto* buffer = model.buffers()->Get(tensor->buffer()); + if (!buffer || !buffer->data()) { + ReportError(error_reporter, "Tensor buffer %d not set", + tensor->buffer()); + return false; + } + + if (tensor->type() == TensorType_STRING) { + if (!VerifyStringTensorBuffer(*buffer, error_reporter)) { + return false; + } + } else { + if (!VerifyNumericTensorBuffer(*tensor, *buffer, error_reporter)) { + return false; + } + } + } + } + return true; +} + +bool VerifyOps(const Model& model, const OpResolver& resolver, + ErrorReporter* error_reporter) { + if (!model.operator_codes()) { + return true; + } + for (const auto& opcode : *model.operator_codes()) { + if (opcode->builtin_code() == BuiltinOperator_CUSTOM) { + if (!resolver.FindOp(opcode->custom_code()->c_str())) { + ReportError(error_reporter, "Unsupported custom op: %s", + opcode->custom_code()->c_str()); + return false; + } + } else { + if (!resolver.FindOp(opcode->builtin_code())) { + ReportError(error_reporter, "Unsupported builtin op: %s", + EnumNameBuiltinOperator(opcode->builtin_code())); + return false; + } + } + } + return true; +} + +} // namespace + +bool Verify(const void* buf, size_t len, const OpResolver& resolver, + ErrorReporter* error_reporter) { + const Model* model = VerifyFlatbufferAndGetModel(buf, len); + if (model == nullptr) { + ReportError(error_reporter, "Invalid flatbuffer format"); + return false; + } + if (model->version() != TFLITE_SCHEMA_VERSION) { + ReportError(error_reporter, "Invalid model version %d", model->version()); + return false; + } + if (!VerifyTensors(*model, error_reporter)) { + return false; + } + if (!VerifyOps(*model, resolver, error_reporter)) { + return false; + } + return true; +} +} // namespace tflite diff --git a/tensorflow/contrib/lite/tools/verifier.h b/tensorflow/contrib/lite/tools/verifier.h new file mode 100644 index 0000000000000000000000000000000000000000..c2ee11215c861ed7b27696a8d786bb6e2a48e930 --- /dev/null +++ b/tensorflow/contrib/lite/tools/verifier.h @@ -0,0 +1,36 @@ +/* Copyright 2017 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_LITE_TOOLS_VERIFIER_H_ +#define TENSORFLOW_CONTRIB_LITE_TOOLS_VERIFIER_H_ + +#include + +#include "tensorflow/contrib/lite/error_reporter.h" +#include "tensorflow/contrib/lite/model.h" + +namespace tflite { + +// Verifies the integrity of a Tensorflow Lite flatbuffer model file. +// Currently, it verifies: +// * The file is following a legit flatbuffer schema. +// * The model is in supported version. +// * All ops used in the model are supported by OpResolver. +bool Verify(const void* buf, size_t len, const OpResolver& resolver, + ErrorReporter* error_reporter); + +} // namespace tflite + +#endif // TENSORFLOW_CONTRIB_LITE_TOOLS_VERIFIER_H_ diff --git a/tensorflow/contrib/lite/tools/verifier_test.cc b/tensorflow/contrib/lite/tools/verifier_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..b3e611f999b2837efbf8876bd989db44c408b8c7 --- /dev/null +++ b/tensorflow/contrib/lite/tools/verifier_test.cc @@ -0,0 +1,288 @@ +/* Copyright 2017 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 +#include + +#include "flatbuffers/flatbuffers.h" +#include "flatbuffers/util.h" +#include +#include "tensorflow/contrib/lite/allocation.h" +#include "tensorflow/contrib/lite/error_reporter.h" +#include "tensorflow/contrib/lite/schema/schema_generated.h" +#include "tensorflow/contrib/lite/testing/util.h" +#include "tensorflow/contrib/lite/tools/mutable_op_resolver.h" +#include "tensorflow/contrib/lite/tools/verifier.h" +#include "tensorflow/contrib/lite/version.h" +#include "tensorflow/core/framework/numeric_types.h" + +namespace tflite { + +using flatbuffers::FlatBufferBuilder; +using flatbuffers::Offset; +using flatbuffers::Vector; + +// Build single subgraph model. +class TfLiteFlatbufferModelBuilder { + public: + TfLiteFlatbufferModelBuilder() { + buffers_.push_back( + CreateBuffer(builder_, builder_.CreateVector(std::vector{}))); + } + + TfLiteFlatbufferModelBuilder(const std::vector& builtin_ops, + const std::vector& custom_ops) { + buffers_.push_back( + CreateBuffer(builder_, builder_.CreateVector(std::vector{}))); + + for (const auto& iter : builtin_ops) { + resolver_.AddBuiltin(iter, &fake_op_); + } + for (const auto& iter : custom_ops) { + resolver_.AddCustom(iter.data(), &fake_op_); + } + } + + void AddTensor(const std::vector& shape, tflite::TensorType type, + const std::vector& buffer, const char* name) { + int buffer_index = 0; + if (!buffer.empty()) { + buffer_index = buffers_.size(); + buffers_.push_back(CreateBuffer(builder_, builder_.CreateVector(buffer))); + } + tensors_.push_back(CreateTensorDirect(builder_, &shape, type, buffer_index, + name, /*quantization=*/0)); + } + + void AddOperator(const std::vector& inputs, + const std::vector& outputs, + tflite::BuiltinOperator builtin_op, const char* custom_op) { + operator_codes_.push_back( + CreateOperatorCodeDirect(builder_, builtin_op, custom_op)); + operators_.push_back(CreateOperator( + builder_, operator_codes_.size() - 1, builder_.CreateVector(inputs), + builder_.CreateVector(outputs), BuiltinOptions_NONE, + /*builtin_options=*/0, + /*custom_options=*/0, tflite::CustomOptionsFormat_FLEXBUFFERS)); + } + + void FinishModel(const std::vector& inputs, + const std::vector& outputs) { + auto subgraph = std::vector>({CreateSubGraph( + builder_, builder_.CreateVector(tensors_), + builder_.CreateVector(inputs), builder_.CreateVector(outputs), + builder_.CreateVector(operators_), + builder_.CreateString("test_subgraph"))}); + auto result = CreateModel( + builder_, TFLITE_SCHEMA_VERSION, builder_.CreateVector(operator_codes_), + builder_.CreateVector(subgraph), builder_.CreateString("test_model"), + builder_.CreateVector(buffers_)); + tflite::FinishModelBuffer(builder_, result); + } + + bool Verify() { + return tflite::Verify(builder_.GetBufferPointer(), builder_.GetSize(), + resolver_, DefaultErrorReporter()); + } + + private: + FlatBufferBuilder builder_; + MutableOpResolver resolver_; + TfLiteRegistration fake_op_; + std::vector> operators_; + std::vector> operator_codes_; + std::vector> tensors_; + std::vector> buffers_; +}; + +TEST(VerifyModel, TestEmptyModel) { + FlatBufferBuilder builder; + auto model = CreateModel(builder, /*version=*/TFLITE_SCHEMA_VERSION, + /*operator_codes=*/0, /*subgraphs=*/0, + /*description=*/0, /*buffers=*/0); + ::tflite::FinishModelBuffer(builder, model); + + ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize(), + MutableOpResolver{}, DefaultErrorReporter())); +} + +TEST(VerifyModel, TestSimpleModel) { + TfLiteFlatbufferModelBuilder builder({}, {"test"}); + builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "test"); + builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4, 5, 6}, "input"); + builder.AddTensor( + {2}, TensorType_STRING, + {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 19, 0, 0, 0, 'A', 'B', 'C'}, + "data"); + builder.AddTensor({2, 3}, TensorType_INT32, {}, "output"); + builder.FinishModel({0, 1}, {2}); + ASSERT_TRUE(builder.Verify()); +} + +TEST(VerifyModel, TestCorruptedData) { + std::string model = "123"; + ASSERT_FALSE(Verify(model.data(), model.size(), MutableOpResolver{}, + /*error_reporter=*/nullptr)); +} + +TEST(VerifyModel, TestUnsupportedVersion) { + FlatBufferBuilder builder; + auto model = CreateModel(builder, /*version=*/1, /*operator_codes=*/0, + /*subgraphs=*/0, /*description=*/0, /*buffers=*/0); + ::tflite::FinishModelBuffer(builder, model); + ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), + MutableOpResolver{}, DefaultErrorReporter())); +} + +TEST(VerifyModel, TestRandomModificationIsNotAllowed) { + FlatBufferBuilder builder; + auto model = CreateModel(builder, /*version=*/TFLITE_SCHEMA_VERSION, + /*operator_codes=*/0, + /*subgraphs=*/0, /*description=*/0, /*buffers=*/0); + ::tflite::FinishModelBuffer(builder, model); + + std::string model_content(reinterpret_cast(builder.GetBufferPointer()), + builder.GetSize()); + for (int i = 0; i < model_content.size(); i++) { + model_content[i] = (model_content[i] + 137) % 255; + EXPECT_FALSE(Verify(model_content.data(), model_content.size(), + MutableOpResolver{}, DefaultErrorReporter())) + << "Fail at position: " << i; + } +} + +TEST(VerifyModel, TestIntTensorShapeIsGreaterThanBuffer) { + TfLiteFlatbufferModelBuilder builder; + builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, TestIntTensorShapeIsSmallerThanBuffer) { + TfLiteFlatbufferModelBuilder builder; + builder.AddTensor({2, 1}, TensorType_UINT8, {1, 2, 3, 4}, "input"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, TestIntTensorShapeOverflow) { + TfLiteFlatbufferModelBuilder builder; + builder.AddTensor({1024, 2048, 4096}, TensorType_UINT8, {1, 2, 3, 4}, + "input"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, TensorBufferIsNotValid) { + FlatBufferBuilder builder; + std::vector shape = {2, 3}; + auto tensors = builder.CreateVector(std::vector>{ + CreateTensorDirect(builder, &shape, TensorType_INT32, /*buffer=*/2, + "input", /*quantization=*/0)}); + auto subgraph = std::vector>( + {CreateSubGraph(builder, tensors, /*inputs=*/0, /*outputs=*/0, + /*operators=*/0, builder.CreateString("Main"))}); + + auto buffers = builder.CreateVector(std::vector>{ + CreateBuffer(builder, + builder.CreateVector(std::vector{1, 2, 3, 4, 5, 6})), + }); + + auto model = CreateModel(builder, TFLITE_SCHEMA_VERSION, /*operator_codes=*/0, + builder.CreateVector(subgraph), + builder.CreateString("SmartReply"), buffers); + + ::tflite::FinishModelBuffer(builder, model); + ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), + MutableOpResolver{}, DefaultErrorReporter())); +} + +TEST(VerifyModel, StringTensorHasInvalidNumString) { + TfLiteFlatbufferModelBuilder builder; + builder.AddTensor( + {2}, TensorType_STRING, + {0x00, 0x00, 0x00, 0x20, 16, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B'}, + "input"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, StringTensorOffsetTooSmall) { + TfLiteFlatbufferModelBuilder builder; + builder.AddTensor( + {2}, TensorType_STRING, + {2, 0, 0, 0, 12, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B'}, "input"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, StringTensorOffsetOutOfRange) { + TfLiteFlatbufferModelBuilder builder; + builder.AddTensor( + {2}, TensorType_STRING, + {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 22, 0, 0, 0, 'A', 'B'}, "input"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, StringTensorIsLargerThanRequired) { + TfLiteFlatbufferModelBuilder builder; + builder.AddTensor( + {2}, TensorType_STRING, + {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B', 'C'}, + "input"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, AllOpsAreSupported) { + TfLiteFlatbufferModelBuilder builder({BuiltinOperator_ADD}, {"CustomOp"}); + builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); + builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); + builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); + builder.AddOperator({0, 1}, {2}, BuiltinOperator_ADD, nullptr); + builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "CustomOp"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, UseUnsupportedBuiltinOps) { + TfLiteFlatbufferModelBuilder builder({BuiltinOperator_SUB}, {"CustomOp"}); + builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); + builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); + builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); + builder.AddOperator({0, 1}, {2}, BuiltinOperator_ADD, nullptr); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +TEST(VerifyModel, UseUnsupportedCustomOps) { + TfLiteFlatbufferModelBuilder builder({BuiltinOperator_ADD}, {"NewOp"}); + builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); + builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); + builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); + builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "Not supported"); + builder.FinishModel({}, {}); + ASSERT_FALSE(builder.Verify()); +} + +// TODO(yichengfan): make up malicious files to test with. + +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/tools/visualize.py b/tensorflow/contrib/lite/tools/visualize.py new file mode 100644 index 0000000000000000000000000000000000000000..f571dd59da0a3f4aff264b48fba3e41f75b50404 --- /dev/null +++ b/tensorflow/contrib/lite/tools/visualize.py @@ -0,0 +1,391 @@ +#!/usr/bin/env python +# Copyright 2017 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. +# ============================================================================== +"""This tool creates an html visualization of a TensorFlow Lite graph. + +Example usage: + +python visualize.py foo.tflite foo.html +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import json +import os +import sys + +# Schema to use for flatbuffers +_SCHEMA = "third_party/tensorflow/contrib/lite/schema/schema.fbs" + +# Where the binary will be once built in for the flatc converter +_BINARY = "third_party/flatbuffers/flatc" + +# A CSS description for making the visualizer +_CSS = """ + + + + + + + + +""" + +_D3_HTML_TEMPLATE = """ + +""" + + +class OpCodeMapper(object): + """Maps an opcode index to an op name.""" + + def __init__(self, data): + self.code_to_name = {} + for idx, d in enumerate(data["operator_codes"]): + self.code_to_name[idx] = d["builtin_code"] + + def __call__(self, x): + if x not in self.code_to_name: + s = "" + else: + s = self.code_to_name[x] + return "%s (opcode=%d)" % (s, x) + + +class DataSizeMapper(object): + """For buffers, report the number of bytes.""" + + def __call__(self, x): + if x is not None: + return "%d bytes" % len(x) + else: + return "--" + + +class TensorMapper(object): + """Maps a list of tensor indices to a tooltip hoverable indicator of more.""" + + def __init__(self, subgraph_data): + self.data = subgraph_data + + def __call__(self, x): + html = "" + html += "" + for i in x: + tensor = self.data["tensors"][i] + html += str(i) + " " + html += tensor["name"] + " " + html += str(tensor["type"]) + " " + html += repr(tensor["shape"]) + "
" + html += "" + html += repr(x) + html += "" + return html + + +def GenerateGraph(subgraph_idx, g, opcode_mapper): + """Produces the HTML required to have a d3 visualization of the dag.""" + + def TensorName(idx): + return "t%d" % idx + + def OpName(idx): + return "o%d" % idx + + edges = [] + nodes = [] + first = {} + pixel_mult = 50 # TODO(aselle): multiplier for initial placement + for op_index, op in enumerate(g["operators"]): + for tensor_input_position, tensor_index in enumerate(op["inputs"]): + if tensor_index not in first: + first[tensor_index] = ( + op_index * pixel_mult, + tensor_input_position * pixel_mult - pixel_mult / 2) + edges.append({ + "source": TensorName(tensor_index), + "target": OpName(op_index) + }) + for tensor_index in op["outputs"]: + edges.append({ + "target": TensorName(tensor_index), + "source": OpName(op_index) + }) + nodes.append({ + "id": OpName(op_index), + "name": opcode_mapper(op["opcode_index"]), + "group": 2, + "x": pixel_mult, + "y": op_index * pixel_mult + }) + for tensor_index, tensor in enumerate(g["tensors"]): + initial_y = ( + first[tensor_index] if tensor_index in first else len(g["operators"])) + + nodes.append({ + "id": TensorName(tensor_index), + "name": "%s (%d)" % (tensor["name"], tensor_index), + "group": 1, + "x": 2, + "y": initial_y + }) + graph_str = json.dumps({"nodes": nodes, "edges": edges}) + + html = _D3_HTML_TEMPLATE % (graph_str, subgraph_idx) + return html + + +def GenerateTableHtml(items, keys_to_print, display_index=True): + """Given a list of object values and keys to print, make an HTML table. + + Args: + items: Items to print an array of dicts. + keys_to_print: (key, display_fn). `key` is a key in the object. i.e. + items[0][key] should exist. display_fn is the mapping function on display. + i.e. the displayed html cell will have the string returned by + `mapping_fn(items[0][key])`. + display_index: add a column which is the index of each row in `items`. + Returns: + An html table. + """ + html = "" + # Print the list of items + html += "\n" + html += "\n" + if display_index: + html += "" + for h, mapper in keys_to_print: + html += "" % h + html += "\n" + for idx, tensor in enumerate(items): + html += "\n" + if display_index: + html += "" % idx + # print tensor.keys() + for h, mapper in keys_to_print: + val = tensor[h] if h in tensor else None + val = val if mapper is None else mapper(val) + html += "\n" % val + + html += "\n" + html += "
index%s
%d%s
\n" + return html + + +def CreateHtmlFile(tflite_input, html_output): + """Given a tflite model in `tflite_input` file, produce html description.""" + + # Convert the model into a JSON flatbuffer using flatc (build if doesn't + # exist. + if not os.path.exists(tflite_input): + raise RuntimeError("Invalid filename %r" % tflite_input) + if tflite_input.endswith(".tflite") or tflite_input.endswith(".bin"): + + # Run convert + cmd = ( + _BINARY + " -t " + "--strict-json --defaults-json -o /tmp {schema} -- {input}".format( + input=tflite_input, schema=_SCHEMA)) + print(cmd) + os.system(cmd) + real_output = ("/tmp/" + os.path.splitext( + os.path.split(tflite_input)[-1])[0] + ".json") + + data = json.load(open(real_output)) + elif tflite_input.endswith(".json"): + data = json.load(open(tflite_input)) + else: + raise RuntimeError("Input file was not .tflite or .json") + html = "" + html += _CSS + html += "

TensorFlow Lite Model

" + + data["filename"] = tflite_input # Avoid special case + toplevel_stuff = [("filename", None), ("version", None), ("description", + None)] + + html += "\n" + for key, mapping in toplevel_stuff: + if not mapping: + mapping = lambda x: x + html += "\n" % (key, mapping(data[key])) + html += "
%s%s
\n" + + # Spec on what keys to display + buffer_keys_to_display = [("data", DataSizeMapper())] + operator_keys_to_display = [("builtin_code", None)] + + for subgraph_idx, g in enumerate(data["subgraphs"]): + # Subgraph local specs on what to display + html += "
" + tensor_mapper = TensorMapper(g) + opcode_mapper = OpCodeMapper(data) + op_keys_to_display = [("inputs", tensor_mapper), ("outputs", tensor_mapper), + ("builtin_options", None), ("opcode_index", + opcode_mapper)] + tensor_keys_to_display = [("name", None), ("type", None), ("shape", None), + ("buffer", None), ("quantization", None)] + + html += "

Subgraph %d

\n" % subgraph_idx + + # Inputs and outputs. + html += "

Inputs/Outputs

\n" + html += GenerateTableHtml( + [{ + "inputs": g["inputs"], + "outputs": g["outputs"] + }], [("inputs", tensor_mapper), ("outputs", tensor_mapper)], + display_index=False) + + # Print the tensors. + html += "

Tensors

\n" + html += GenerateTableHtml(g["tensors"], tensor_keys_to_display) + + # Print the ops. + html += "

Ops

\n" + html += GenerateTableHtml(g["operators"], op_keys_to_display) + + # Visual graph. + html += "\n" % ( + subgraph_idx,) + html += GenerateGraph(subgraph_idx, g, opcode_mapper) + html += "
" + + # Buffers have no data, but maybe in the future they will + html += "

Buffers

\n" + html += GenerateTableHtml(data["buffers"], buffer_keys_to_display) + + # Operator codes + html += "

Operator Codes

\n" + html += GenerateTableHtml(data["operator_codes"], operator_keys_to_display) + + html += "\n" + + open(html_output, "w").write(html) + + +def main(argv): + try: + tflite_input = argv[1] + html_output = argv[2] + except IndexError: + print("Usage: %s " % (argv[0])) + else: + CreateHtmlFile(tflite_input, html_output) + + +if __name__ == "__main__": + main(sys.argv) diff --git a/tensorflow/contrib/lite/version.h b/tensorflow/contrib/lite/version.h index a751afabe7460f0c9e88385faf1497b2c0a25d6b..efd63f4006ae661c6fdbbaa81cb02fa8947271f3 100644 --- a/tensorflow/contrib/lite/version.h +++ b/tensorflow/contrib/lite/version.h @@ -12,12 +12,12 @@ 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 THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_VERSION_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_VERSION_H_ +#ifndef TENSORFLOW_CONTRIB_LITE_VERSION_H_ +#define TENSORFLOW_CONTRIB_LITE_VERSION_H_ // The version number of the Schema. Ideally all changes will be backward // compatible. If that ever changes, we must ensure that version is the first // entry in the new tflite root so that we can see that version is not 1. #define TFLITE_SCHEMA_VERSION (3) -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_LITE_VERSION_H_ +#endif // TENSORFLOW_CONTRIB_LITE_VERSION_H_ diff --git a/tensorflow/contrib/lookup/lookup_ops.py b/tensorflow/contrib/lookup/lookup_ops.py index 66caa6a2e5d17f74706965b7ca3f7928d63ae130..a430dac4ec43ce31f0b5aaae5e7b0b51d25c9632 100644 --- a/tensorflow/contrib/lookup/lookup_ops.py +++ b/tensorflow/contrib/lookup/lookup_ops.py @@ -399,7 +399,7 @@ class MutableHashTable(LookupInterface): Raises: TypeError: when `keys` do not match the table data types. """ - if keys.dtype != self._key_dtype: + if keys.dtype.base_dtype != self._key_dtype: raise TypeError("Signature mismatch. Keys must be dtype %s, got %s." % (self._key_dtype, keys.dtype)) @@ -600,7 +600,7 @@ class MutableDenseHashTable(LookupInterface): Raises: TypeError: when `keys` do not match the table data types. """ - if keys.dtype != self._key_dtype: + if keys.dtype.base_dtype != self._key_dtype: raise TypeError("Signature mismatch. Keys must be dtype %s, got %s." % (self._key_dtype, keys.dtype)) diff --git a/tensorflow/contrib/lookup/lookup_ops_test.py b/tensorflow/contrib/lookup/lookup_ops_test.py index f0499010d476f68da541ee67d085b12e48faeaf5..f681b7b132750ef80aa56f25143418fbc4eaa1bb 100644 --- a/tensorflow/contrib/lookup/lookup_ops_test.py +++ b/tensorflow/contrib/lookup/lookup_ops_test.py @@ -187,6 +187,11 @@ class HashTableOpTest(test.TestCase): lookup.KeyValueTensorInitializer(keys, values), default_val) table.init.run() + # Ref types do not produce a lookup signature mismatch. + input_string_ref = variables.Variable("brain") + variables.global_variables_initializer().run() + self.assertEqual(0, table.lookup(input_string_ref).eval()) + input_string = constant_op.constant([1, 2, 3], dtypes.int64) with self.assertRaises(TypeError): table.lookup(input_string) @@ -629,6 +634,17 @@ class MutableHashTableOpTest(test.TestCase): table.insert(keys, values).run() self.assertAllEqual(3, table.size().eval()) + input_string_ref = variables.Variable("brain") + input_int64_ref = variables.Variable(-1, dtype=dtypes.int64) + variables.global_variables_initializer().run() + + # Ref types do not produce an insert signature mismatch. + table.insert(input_string_ref, input_int64_ref).run() + self.assertAllEqual(3, table.size().eval()) + + # Ref types do not produce a lookup signature mismatch. + self.assertEqual(-1, table.lookup(input_string_ref).eval()) + # lookup with keys of the wrong type input_string = constant_op.constant([1, 2, 3], dtypes.int64) with self.assertRaises(TypeError): @@ -1640,23 +1656,22 @@ class InitializeTableFromFileOpTest(test.TestCase): f.write("\n".join(values) + "\n") return vocabulary_file + @test_util.run_in_graph_and_eager_modes() def testInitializeStringTable(self): vocabulary_file = self._createVocabFile("one_column_1.txt") + default_value = -1 + table = lookup.HashTable( + lookup.TextFileInitializer(vocabulary_file, dtypes.string, + lookup.TextFileIndex.WHOLE_LINE, + dtypes.int64, + lookup.TextFileIndex.LINE_NUMBER), + default_value) + self.evaluate(table.init) - with self.test_session(): - default_value = -1 - table = lookup.HashTable( - lookup.TextFileInitializer(vocabulary_file, dtypes.string, - lookup.TextFileIndex.WHOLE_LINE, - dtypes.int64, - lookup.TextFileIndex.LINE_NUMBER), - default_value) - table.init.run() - - output = table.lookup(constant_op.constant(["brain", "salad", "tank"])) + output = table.lookup(constant_op.constant(["brain", "salad", "tank"])) - result = output.eval() - self.assertAllEqual([0, 1, -1], result) + result = self.evaluate(output) + self.assertAllEqual([0, 1, -1], result) def testInitializeInt64Table(self): vocabulary_file = self._createVocabFile( diff --git a/tensorflow/contrib/losses/python/losses/loss_ops.py b/tensorflow/contrib/losses/python/losses/loss_ops.py index 7c523ad49265aaf32c8d5a8ae04d3e93262a1b55..8c3a8afe7a0f6f5ad9ceae566288ba60be73d339 100644 --- a/tensorflow/contrib/losses/python/losses/loss_ops.py +++ b/tensorflow/contrib/losses/python/losses/loss_ops.py @@ -30,20 +30,13 @@ from tensorflow.python.ops import nn_ops from tensorflow.python.util.deprecation import deprecated from tensorflow.python.util.deprecation import deprecated_args -__all__ = ["absolute_difference", - "add_loss", - "cosine_distance", - "compute_weighted_loss", - "get_losses", - "get_regularization_losses", - "get_total_loss", - "hinge_loss", - "log_loss", - "mean_pairwise_squared_error", - "mean_squared_error", - "sigmoid_cross_entropy", - "softmax_cross_entropy", - "sparse_softmax_cross_entropy"] +__all__ = [ + "absolute_difference", "add_loss", "cosine_distance", + "compute_weighted_loss", "get_losses", "get_regularization_losses", + "get_total_loss", "hinge_loss", "log_loss", "mean_pairwise_squared_error", + "mean_squared_error", "sigmoid_cross_entropy", "softmax_cross_entropy", + "sparse_softmax_cross_entropy" +] def _scale_losses(losses, weights): @@ -66,8 +59,8 @@ def _scale_losses(losses, weights): # First, compute the sum of the losses over all elements: start_index = max(0, weights.get_shape().ndims) reduction_indices = list(range(start_index, losses.get_shape().ndims)) - reduced_losses = math_ops.reduce_sum(losses, - reduction_indices=reduction_indices) + reduced_losses = math_ops.reduce_sum( + losses, reduction_indices=reduction_indices) reduced_losses = math_ops.multiply(reduced_losses, weights) return math_ops.reduce_sum(reduced_losses) @@ -90,9 +83,10 @@ def _safe_div(numerator, denominator, name="value"): """ return array_ops.where( math_ops.greater(denominator, 0), - math_ops.div(numerator, array_ops.where( - math_ops.equal(denominator, 0), - array_ops.ones_like(denominator), denominator)), + math_ops.div(numerator, + array_ops.where( + math_ops.equal(denominator, 0), + array_ops.ones_like(denominator), denominator)), array_ops.zeros_like(numerator), name=name) @@ -176,14 +170,15 @@ def _num_present(losses, weights, per_batch=False): """ # If weights is a scalar, its easy to compute: if weights.get_shape().ndims == 0: - batch_size = array_ops.reshape(array_ops.slice(array_ops.shape(losses), - [0], [1]), []) - num_per_batch = math_ops.div(math_ops.to_float(array_ops.size(losses)), - math_ops.to_float(batch_size)) - num_per_batch = array_ops.where(math_ops.equal(weights, 0), - 0.0, num_per_batch) - num_per_batch = math_ops.multiply(array_ops.ones( - array_ops.reshape(batch_size, [1])), num_per_batch) + batch_size = array_ops.reshape( + array_ops.slice(array_ops.shape(losses), [0], [1]), []) + num_per_batch = math_ops.div( + math_ops.to_float(array_ops.size(losses)), + math_ops.to_float(batch_size)) + num_per_batch = array_ops.where( + math_ops.equal(weights, 0), 0.0, num_per_batch) + num_per_batch = math_ops.multiply( + array_ops.ones(array_ops.reshape(batch_size, [1])), num_per_batch) return num_per_batch if per_batch else math_ops.reduce_sum(num_per_batch) # First, count the number of nonzero weights: @@ -194,8 +189,8 @@ def _num_present(losses, weights, per_batch=False): reduction_indices=reduction_indices) # Next, determine the number of elements that weights would broadcast to: - broadcast_dims = array_ops.slice(array_ops.shape(losses), - [weights.get_shape().ndims], [-1]) + broadcast_dims = array_ops.slice( + array_ops.shape(losses), [weights.get_shape().ndims], [-1]) num_to_broadcast = math_ops.to_float(math_ops.reduce_prod(broadcast_dims)) num_per_batch = math_ops.multiply(num_nonzero_per_batch, num_to_broadcast) @@ -303,8 +298,11 @@ def absolute_difference(predictions, labels=None, weights=1.0, scope=None): @deprecated("2016-12-30", "Use tf.losses.sigmoid_cross_entropy instead. Note that the order " "of the predictions and labels arguments has been changed.") -def sigmoid_cross_entropy( - logits, multi_class_labels, weights=1.0, label_smoothing=0, scope=None): +def sigmoid_cross_entropy(logits, + multi_class_labels, + weights=1.0, + label_smoothing=0, + scope=None): """Creates a cross-entropy loss using tf.nn.sigmoid_cross_entropy_with_logits. `weights` acts as a coefficient for the loss. If a scalar is provided, @@ -340,20 +338,22 @@ def sigmoid_cross_entropy( multi_class_labels = math_ops.cast(multi_class_labels, logits.dtype) if label_smoothing > 0: - multi_class_labels = (multi_class_labels * (1 - label_smoothing) + - 0.5 * label_smoothing) + multi_class_labels = ( + multi_class_labels * (1 - label_smoothing) + 0.5 * label_smoothing) - losses = nn.sigmoid_cross_entropy_with_logits(labels=multi_class_labels, - logits=logits, - name="xentropy") + losses = nn.sigmoid_cross_entropy_with_logits( + labels=multi_class_labels, logits=logits, name="xentropy") return compute_weighted_loss(losses, weights, scope=scope) @deprecated("2016-12-30", "Use tf.losses.softmax_cross_entropy instead. Note that the order " "of the logits and labels arguments has been changed.") -def softmax_cross_entropy( - logits, onehot_labels, weights=1.0, label_smoothing=0, scope=None): +def softmax_cross_entropy(logits, + onehot_labels, + weights=1.0, + label_smoothing=0, + scope=None): """Creates a cross-entropy loss using tf.nn.softmax_cross_entropy_with_logits. `weights` acts as a coefficient for the loss. If a scalar is provided, @@ -393,9 +393,8 @@ def softmax_cross_entropy( smooth_negatives = label_smoothing / num_classes onehot_labels = onehot_labels * smooth_positives + smooth_negatives - losses = nn.softmax_cross_entropy_with_logits(labels=onehot_labels, - logits=logits, - name="xentropy") + losses = nn.softmax_cross_entropy_with_logits( + labels=onehot_labels, logits=logits, name="xentropy") return compute_weighted_loss(losses, weights, scope=scope) @@ -429,9 +428,8 @@ def sparse_softmax_cross_entropy(logits, labels, weights=1.0, scope=None): [logits, labels, weights]) as scope: labels = array_ops.reshape(labels, shape=[array_ops.shape(labels)[0]]) - losses = nn.sparse_softmax_cross_entropy_with_logits(labels=labels, - logits=logits, - name="xentropy") + losses = nn.sparse_softmax_cross_entropy_with_logits( + labels=labels, logits=logits, name="xentropy") return compute_weighted_loss(losses, weights, scope=scope) @@ -470,8 +468,7 @@ def log_loss(predictions, labels=None, weights=1.0, epsilon=1e-7, scope=None): predictions = math_ops.to_float(predictions) labels = math_ops.to_float(labels) losses = -math_ops.multiply( - labels, - math_ops.log(predictions + epsilon)) - math_ops.multiply( + labels, math_ops.log(predictions + epsilon)) - math_ops.multiply( (1 - labels), math_ops.log(1 - predictions + epsilon)) return compute_weighted_loss(losses, weights, scope=scope) @@ -490,7 +487,8 @@ def hinge_loss(logits, labels=None, scope=None): scope: The scope for the operations performed in computing the loss. Returns: - An unweighted `Tensor` of same shape as `logits` and `labels` representing the + An unweighted `Tensor` of same shape as `logits` and `labels` representing + the loss values across the batch. Raises: @@ -544,8 +542,10 @@ def mean_squared_error(predictions, labels=None, weights=1.0, scope=None): @deprecated("2016-12-30", "Use tf.losses.mean_pairwise_squared_error instead. Note that the " "order of the predictions and labels arguments has been changed.") -def mean_pairwise_squared_error( - predictions, labels=None, weights=1.0, scope=None): +def mean_pairwise_squared_error(predictions, + labels=None, + weights=1.0, + scope=None): """Adds a pairwise-errors-squared loss to the training procedure. Unlike `mean_squared_error`, which is a measure of the differences between @@ -602,31 +602,34 @@ def mean_pairwise_squared_error( reduction_indices = list(range(1, diffs.get_shape().ndims)) sum_squares_diff_per_batch = math_ops.reduce_sum( - math_ops.square(diffs), - reduction_indices=reduction_indices) + math_ops.square(diffs), reduction_indices=reduction_indices) num_present_per_batch = _num_present(diffs, weights, per_batch=True) - term1 = 2.0 * _safe_div(sum_squares_diff_per_batch, - num_present_per_batch) + term1 = 2.0 * _safe_div(sum_squares_diff_per_batch, num_present_per_batch) sum_diff = math_ops.reduce_sum(diffs, reduction_indices=reduction_indices) - term2 = 2.0 * _safe_div(math_ops.square(sum_diff), - math_ops.square(num_present_per_batch)) + term2 = 2.0 * _safe_div( + math_ops.square(sum_diff), math_ops.square(num_present_per_batch)) loss = _scale_losses(term1 - term2, weights) - mean_loss = array_ops.where(math_ops.reduce_sum(num_present_per_batch) > 0, - loss, - array_ops.zeros_like(loss), - name="value") + mean_loss = array_ops.where( + math_ops.reduce_sum(num_present_per_batch) > 0, + loss, + array_ops.zeros_like(loss), + name="value") add_loss(mean_loss) return mean_loss @deprecated("2016-12-30", "Use tf.losses.cosine_distance instead.") @deprecated_args(None, "dim is deprecated, use axis instead", "dim") -def cosine_distance( - predictions, labels=None, axis=None, weights=1.0, scope=None, dim=None): +def cosine_distance(predictions, + labels=None, + axis=None, + weights=1.0, + scope=None, + dim=None): """Adds a cosine-distance loss to the training procedure. Note that the function assumes that `predictions` and `labels` are already @@ -662,5 +665,8 @@ def cosine_distance( labels = math_ops.to_float(labels) radial_diffs = math_ops.multiply(predictions, labels) - losses = 1 - math_ops.reduce_sum(radial_diffs, reduction_indices=[axis,]) + losses = 1 - math_ops.reduce_sum( + radial_diffs, reduction_indices=[ + axis, + ]) return compute_weighted_loss(losses, weights, scope=scope) diff --git a/tensorflow/contrib/losses/python/losses/loss_ops_test.py b/tensorflow/contrib/losses/python/losses/loss_ops_test.py index 9d0f95e6f3e7fa9666a99e31578b38d52e0b6b4a..1417772e0496cb571488e5b30bd4f3fb1b591730 100644 --- a/tensorflow/contrib/losses/python/losses/loss_ops_test.py +++ b/tensorflow/contrib/losses/python/losses/loss_ops_test.py @@ -27,6 +27,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors_impl from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed +from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops @@ -274,6 +275,7 @@ class SoftmaxCrossEntropyLossTest(test.TestCase): self.assertAlmostEqual(np.average(weights) * 10.0, loss, 3) +@test_util.with_c_api class SparseSoftmaxCrossEntropyLossTest(test.TestCase): def testNoneWeightRaisesValueError(self): @@ -471,7 +473,11 @@ class SparseSoftmaxCrossEntropyLossTest(test.TestCase): labels = constant_op.constant([[0, 1], [2, 3]]) weights = constant_op.constant([1.2, 3.4, 5.6, 7.8]) - with self.assertRaises(errors_impl.InvalidArgumentError): + if ops._USE_C_API: + error_type = ValueError + else: + error_type = errors_impl.InvalidArgumentError + with self.assertRaises(error_type): loss_ops.sparse_softmax_cross_entropy( logits, labels, weights=weights).eval() diff --git a/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py b/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py index c3a57ba51bcf0a292490dfaa9e556f6e5811ed66..6842bc38eb108b46cc3eff715c9cbc74f991308b 100644 --- a/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py +++ b/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py @@ -53,12 +53,12 @@ def pairwise_distance(feature, squared=False): math_ops.reduce_sum( math_ops.square(feature), axis=[1], - keep_dims=True), + keepdims=True), math_ops.reduce_sum( math_ops.square( array_ops.transpose(feature)), axis=[0], - keep_dims=True)) - 2.0 * math_ops.matmul( + keepdims=True)) - 2.0 * math_ops.matmul( feature, array_ops.transpose(feature)) # Deal with numerical inaccuracies. Set small negatives to zero. @@ -132,10 +132,10 @@ def masked_maximum(data, mask, dim=1): masked_maximums: N-D `Tensor`. The maximized dimension is of size 1 after the operation. """ - axis_minimums = math_ops.reduce_min(data, dim, keep_dims=True) + axis_minimums = math_ops.reduce_min(data, dim, keepdims=True) masked_maximums = math_ops.reduce_max( math_ops.multiply( - data - axis_minimums, mask), dim, keep_dims=True) + axis_minimums + data - axis_minimums, mask), dim, keepdims=True) + axis_minimums return masked_maximums @@ -151,10 +151,10 @@ def masked_minimum(data, mask, dim=1): masked_minimums: N-D `Tensor`. The minimized dimension is of size 1 after the operation. """ - axis_maximums = math_ops.reduce_max(data, dim, keep_dims=True) + axis_maximums = math_ops.reduce_max(data, dim, keepdims=True) masked_minimums = math_ops.reduce_min( math_ops.multiply( - data - axis_maximums, mask), dim, keep_dims=True) + axis_maximums + data - axis_maximums, mask), dim, keepdims=True) + axis_maximums return masked_minimums @@ -203,7 +203,7 @@ def triplet_semihard_loss(labels, embeddings, margin=1.0): math_ops.greater( math_ops.reduce_sum( math_ops.cast( - mask, dtype=dtypes.float32), 1, keep_dims=True), + mask, dtype=dtypes.float32), 1, keepdims=True), 0.0), [batch_size, batch_size]) mask_final = array_ops.transpose(mask_final) @@ -290,7 +290,7 @@ def npairs_loss(labels, embeddings_anchor, embeddings_positive, labels_remapped = math_ops.to_float( math_ops.equal(labels, array_ops.transpose(labels))) - labels_remapped /= math_ops.reduce_sum(labels_remapped, 1, keep_dims=True) + labels_remapped /= math_ops.reduce_sum(labels_remapped, 1, keepdims=True) # Add the softmax loss. xent_loss = nn.softmax_cross_entropy_with_logits( @@ -395,7 +395,7 @@ def npairs_loss_multilabel(sparse_labels, embeddings_anchor, multilabel_adjacency_matrix = _build_multilabel_adjacency(sparse_labels) labels_remapped = math_ops.to_float(multilabel_adjacency_matrix) - labels_remapped /= math_ops.reduce_sum(labels_remapped, 1, keep_dims=True) + labels_remapped /= math_ops.reduce_sum(labels_remapped, 1, keepdims=True) # Add the softmax loss. xent_loss = nn.softmax_cross_entropy_with_logits( @@ -448,10 +448,10 @@ def lifted_struct_loss(labels, embeddings, margin=1.0): # Safe maximum: Temporarily shift negative distances # above zero before taking max. # this is to take the max only among negatives. - row_minimums = math_ops.reduce_min(diff, 1, keep_dims=True) + row_minimums = math_ops.reduce_min(diff, 1, keepdims=True) row_negative_maximums = math_ops.reduce_max( math_ops.multiply( - diff - row_minimums, mask), 1, keep_dims=True) + row_minimums + diff - row_minimums, mask), 1, keepdims=True) + row_minimums # Compute the loss. # Keep track of matrix of maximums where M_ij = max(m_i, m_j) @@ -470,7 +470,7 @@ def lifted_struct_loss(labels, embeddings, margin=1.0): math_ops.reduce_sum(math_ops.multiply( math_ops.exp( diff_tiled - max_elements_vect), - mask_tiled), 1, keep_dims=True), [batch_size, batch_size]) + mask_tiled), 1, keepdims=True), [batch_size, batch_size]) loss_mat = max_elements + math_ops.log( loss_exp_left + array_ops.transpose(loss_exp_left)) @@ -686,7 +686,7 @@ def _find_loss_augmented_facility_idx(pairwise_distances, labels, chosen_ids, array_ops.reshape(pairwise_distances_candidate, [1, -1]) ], 0), axis=0, - keep_dims=True), [num_candidates, -1]), + keepdims=True), [num_candidates, -1]), axis=1) nmi_scores = array_ops.zeros([num_candidates]) diff --git a/tensorflow/contrib/makefile/BUILD b/tensorflow/contrib/makefile/BUILD index a8dd59f32a7f3b27993a7ee48ee7cc07ada59a4c..701eeb44fe3f814cb3fb1cedd8618753946cc3e5 100644 --- a/tensorflow/contrib/makefile/BUILD +++ b/tensorflow/contrib/makefile/BUILD @@ -12,20 +12,3 @@ filegroup( ), visibility = ["//tensorflow:__subpackages__"], ) - -sh_test( - name = "build_all_linux", - size = "enormous", - srcs = ["build_all_linux.sh"], - data = [ - "//tensorflow:all_opensource_files", - "//third_party/eigen3:all_files", - "//third_party/fft2d:all_files", - ], - tags = [ - "manual", - "no_gpu", - "no_oss", - "notap", - ], -) diff --git a/tensorflow/contrib/makefile/Makefile b/tensorflow/contrib/makefile/Makefile index ee84b5b4c8a9e41fe07b4e9dfdc93e31f807d35d..81327407d44b4317b7aecb964a689a35aa35c163 100644 --- a/tensorflow/contrib/makefile/Makefile +++ b/tensorflow/contrib/makefile/Makefile @@ -374,12 +374,72 @@ $(MARCH_OPTION) \ ifdef ENABLE_EXPERIMENTAL_HEXNN_OPS CXXFLAGS += -DENABLE_EXPERIMENTAL_HEXNN_OPS endif - - OBJDIR := $(OBJDIR)android_$(ANDROID_ARCH)/ - LIBDIR := $(LIBDIR)android_$(ANDROID_ARCH)/ - BINDIR := $(BINDIR)android_$(ANDROID_ARCH)/ - DEPDIR := $(DEPDIR)android_$(ANDROID_ARCH)/ + ifeq ($(BUILD_FOR_TEGRA),1) + NVCC := $(JETPACK)/cuda/bin/nvcc + NVCCFLAGS := -x=cu -D__CUDACC__ -DNVCC -DANDROID_TEGRA -ccbin $(NDK_ROOT)/toolchains/$(TOOLCHAIN)/prebuilt/$(ANDROID_HOST_OS_ARCH)/bin/$(BIN_PREFIX)-g++ --std c++11 --expt-relaxed-constexpr -m64 -gencode arch=compute_53,\"code=sm_53\" -gencode arch=compute_62,\"code=sm_62\" -DEIGEN_AVOID_STL_ARRAY -DTENSORFLOW_USE_EIGEN_THREADPOOL -DLANG_CXX11 -DEIGEN_HAS_C99_MATH -DGOOGLE_CUDA=1 -DTF_EXTRA_CUDA_CAPABILITIES=5.3 + CXXFLAGS4NVCC =\ +-DIS_SLIM_BUILD \ +-DANDROID_TEGRA \ +-fno-exceptions \ +-DNDEBUG $(OPTFLAGS) \ +-march=armv8-a \ +-fPIE \ +-D__ANDROID_TYPES_FULL__ \ +--sysroot $(NDK_ROOT)/platforms/android-21/arch-arm64 + + CXXFLAGS +=\ +-DGOOGLE_CUDA=1 \ +-D__ANDROID_TYPES_FULL__ \ +-DANDROID_TEGRA \ +-DEIGEN_AVOID_STL_ARRAY \ +-DEIGEN_HAS_C99_MATH \ +-DLANG_CXX11 -DTENSORFLOW_USE_EIGEN_THREADPOOL -DTF_EXTRA_CUDA_CAPABILITIES=5.3 + + INCLUDES += \ +-Itensorflow/core/kernels \ +-I$(MAKEFILE_DIR)/downloads/cub \ +-I$(MAKEFILE_DIR)/downloads/cub/cub_archive/cub/device \ +-Ithird_party/toolchains/gpus/cuda \ +-I$(JETPACK)/cuda/include \ +-I$(JETPACK) \ +-I$(JETPACK)/cuDNN/aarch64 \ +-I$(JETPACK)/cuda/extras/CUPTI/include + + + CUDA_LIBS := \ +-ltfcuda \ +-lcudart_static \ +-lcudnn \ +-lcublas_static \ +-lcufftw_static \ +-lcusolver_static \ +-lcusparse_static \ +-lcufft \ +-lcuda \ +-lculibos \ +-lcurand_static + + OBJDIR := $(OBJDIR)android_arm64-v8a/ + LIBDIR := $(LIBDIR)android_arm64-v8a/ + BINDIR := $(BINDIR)android_arm64-v8a/ + DEPDIR := $(DEPDIR)android_arm64-v8a/ + + TEGRA_LIBS := \ +-L$(JETPACK)/cuda/targets/aarch64-linux-androideabi/lib \ +-L$(JETPACK)/cuda/targets/aarch64-linux-androideabi/lib/stubs \ +-L$(JETPACK)/cuda/targets/aarch64-linux-androideabi/lib64 \ +-L$(JETPACK)/cuda/targets/aarch64-linux-androideabi/lib64/stubs \ +-L$(JETPACK)/cuDNN/aarch64/cuda/lib64 \ +-L$(LIBDIR) + + CUDA_LIB_DEPS := $(LIBDIR)libtfcuda.a + else + OBJDIR := $(OBJDIR)android_$(ANDROID_ARCH)/ + LIBDIR := $(LIBDIR)android_$(ANDROID_ARCH)/ + BINDIR := $(BINDIR)android_$(ANDROID_ARCH)/ + DEPDIR := $(DEPDIR)android_$(ANDROID_ARCH)/ + endif # ifeq ($(BUILD_FOR_TEGRA),1) endif # ANDROID # LINT.ThenChange(//tensorflow/contrib/android/cmake/CMakeLists.txt) @@ -546,7 +606,8 @@ $(wildcard tensorflow/core/util/*/*.cc) \ tensorflow/core/util/version_info.cc # Remove duplicates (for version_info.cc) CORE_CC_ALL_SRCS := $(sort $(CORE_CC_ALL_SRCS)) -CORE_CC_EXCLUDE_SRCS := \ + +CORE_CC_EXCLUDE_SRCS_NON_GPU := \ $(wildcard tensorflow/core/*/*test.cc) \ $(wildcard tensorflow/core/*/*testutil*) \ $(wildcard tensorflow/core/*/*testlib*) \ @@ -566,25 +627,52 @@ $(wildcard tensorflow/core/lib/jpeg/*) \ $(wildcard tensorflow/core/lib/png/*) \ $(wildcard tensorflow/core/util/events_writer.*) \ $(wildcard tensorflow/core/util/reporter.*) \ -$(wildcard tensorflow/core/platform/default/cuda_libdevice_path.*) \ -$(wildcard tensorflow/core/platform/default/stream_executor.*) \ $(wildcard tensorflow/core/platform/default/test_benchmark.*) \ -$(wildcard tensorflow/core/platform/cuda.h) \ -$(wildcard tensorflow/core/platform/cuda_libdevice_path.*) \ $(wildcard tensorflow/core/platform/cloud/*) \ $(wildcard tensorflow/core/platform/google/*) \ $(wildcard tensorflow/core/platform/google/*/*) \ $(wildcard tensorflow/core/platform/jpeg.*) \ $(wildcard tensorflow/core/platform/png.*) \ $(wildcard tensorflow/core/platform/s3/*) \ -$(wildcard tensorflow/core/platform/stream_executor.*) \ $(wildcard tensorflow/core/platform/windows/*) \ -$(wildcard tensorflow/core/user_ops/*.cu.cc) \ -$(wildcard tensorflow/core/common_runtime/gpu/*) \ -$(wildcard tensorflow/core/common_runtime/gpu_device_factory.*) \ $(wildcard tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.*) \ $(wildcard tensorflow/core/grappler/inputs/file_input_yielder.*) \ -$(wildcard tensorflow/core/grappler/clusters/single_machine.*) +$(wildcard tensorflow/core/grappler/clusters/single_machine.*) \ +tensorflow/core/util/cuda_kernel_helper_test.cu.cc + +CORE_CC_EXCLUDE_SRCS := \ +$(CORE_CC_EXCLUDE_SRCS_NON_GPU) \ +$(wildcard tensorflow/core/platform/stream_executor.*) \ +$(wildcard tensorflow/core/platform/default/cuda_libdevice_path.*) \ +$(wildcard tensorflow/core/platform/cuda.h) \ +$(wildcard tensorflow/core/platform/cuda_libdevice_path.*) \ +$(wildcard tensorflow/core/user_ops/*.cu.cc) \ +$(wildcard tensorflow/core/common_runtime/gpu/*) \ +$(wildcard tensorflow/core/common_runtime/gpu_device_factory.*) + +ifeq ($(BUILD_FOR_TEGRA),1) +CORE_CC_ALL_SRCS := $(CORE_CC_ALL_SRCS) \ +tensorflow/core/kernels/concat_lib_gpu.cc \ +tensorflow/core/kernels/cuda_solvers.cc \ +tensorflow/core/kernels/cudnn_pooling_gpu.cc \ +tensorflow/core/kernels/dense_update_functor.cc \ +tensorflow/core/kernels/fractional_avg_pool_op.cc \ +tensorflow/core/kernels/fractional_max_pool_op.cc \ +tensorflow/core/kernels/fractional_pool_common.cc \ +tensorflow/core/kernels/pooling_ops_3d.cc \ +tensorflow/core/kernels/sparse_fill_empty_rows_op.cc \ +tensorflow/core/kernels/list_kernels.cc \ +$(wildcard tensorflow/core/common_runtime/gpu/*.cc) \ +$(wildcard tensorflow/stream_executor/*.cc) \ +$(wildcard tensorflow/stream_executor/*/*.cc) + +CORE_CC_EXCLUDE_SRCS := \ +$(CORE_CC_EXCLUDE_SRCS_NON_GPU) + +CUDA_CC_SRCS := $(wildcard tensorflow/core/kernels/*.cu.cc) +CUDA_CC_OBJS := $(addprefix $(OBJDIR), $(CUDA_CC_SRCS:.cc=.o)) +endif # TEGRA + # Filter out all the excluded files. TF_CC_SRCS := $(filter-out $(CORE_CC_EXCLUDE_SRCS), $(CORE_CC_ALL_SRCS)) # Add in any extra files that don't fit the patterns easily @@ -637,11 +725,23 @@ $(LIB_PATH): $(LIB_OBJS) @mkdir -p $(dir $@) $(AR) $(ARFLAGS) $(LIB_PATH) $(LIB_OBJS) -$(BENCHMARK_NAME): $(BENCHMARK_OBJS) $(LIB_PATH) +$(BENCHMARK_NAME): $(BENCHMARK_OBJS) $(LIB_PATH) $(CUDA_LIB_DEPS) @mkdir -p $(dir $@) $(CXX) $(CXXFLAGS) $(INCLUDES) \ -o $(BENCHMARK_NAME) $(BENCHMARK_OBJS) \ - $(LIBFLAGS) $(LIB_PATH) $(LDFLAGS) $(LIBS) + $(LIBFLAGS) $(TEGRA_LIBS) $(LIB_PATH) $(LDFLAGS) $(LIBS) $(CUDA_LIBS) + +# NVCC compilation rules for Tegra +ifeq ($(BUILD_FOR_TEGRA),1) +$(OBJDIR)%.cu.o: %.cu.cc + @mkdir -p $(dir $@) + @mkdir -p $(dir $(DEPDIR)$*) + $(NVCC) $(NVCCFLAGS) -Xcompiler "$(CXXFLAGS4NVCC) $(DEPFLAGS)" $(INCLUDES) -c $< -o $@ + +$(LIBDIR)libtfcuda.a: $(CUDA_CC_OBJS) + @mkdir -p $(dir $@) + $(AR) $(ARFLAGS) $@ $(CUDA_CC_OBJS) +endif # Matches on the normal hand-written TensorFlow C++ source files. $(OBJDIR)%.o: %.cc | $(PBT_GEN_FILES) @@ -730,6 +830,7 @@ clean_except_protobuf_libs: cleantarget: rm -rf $(OBJDIR) rm -rf $(BINDIR) + rm -rf $(LIBDIR) $(DEPDIR)/%.d: ; .PRECIOUS: $(DEPDIR)/%.d diff --git a/tensorflow/contrib/makefile/README.md b/tensorflow/contrib/makefile/README.md index 9345303ff11462a447ed6299b0ac3cba558ea68b..b0228c543505c3d14e41bf1dd540b027b00489e6 100644 --- a/tensorflow/contrib/makefile/README.md +++ b/tensorflow/contrib/makefile/README.md @@ -130,6 +130,105 @@ adb shell '/data/local/tmp/benchmark \ For more details, see the [benchmark documentation](../../tools/benchmark). +## CUDA support for Tegra devices running Android (Nvidia Shield TV, etc) + +With the release of TF 1.6 and JetPack for Android 3.2 (currently pending), you can now build a version of TensorFlow for compatible devices according to the following instructions which will receive the full benefits of GPU acceleration. + +#### Environment setup: + +First, download and install JetPack for Android version 3.2 or greater from [Nvidia](https://developers.nvidia.com). Note that as of the TF 1.6 release the JetPack for Android 3.2 release is still pending, and regular JetPack for L4T will not work. + +```bash +git clone https://github.com/tensorflow/tensorflow.git +cd tensorflow +JETPACK=$HOME/JetPack_Android_3.2 +TEGRA_LIBS="$JETPACK/cuDNN/aarch64/cuda/lib64/libcudnn.so $JETPACK/cuda-9.0/extras/CUPTI/lib64/libcupti.so $JETPACK/cuda/targets/aarch64-linux-androideabi/lib64/libcufft.so" +``` + +#### Building all CUDA-enabled native binaries: +This will build CUDA-enabled versions of libtensorflow_inference.so and the benchmark binary. (libtensorflow_demo.so will also be built incidentally, but it does not support CUDA) + +```bash +NDK_ROOT=$JETPACK/android-ndk-r13b +CC_PREFIX=ccache tensorflow/contrib/makefile/build_all_android.sh -s tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in -t "libtensorflow_inference.so libtensorflow_demo.so all" -a tegra +``` +(add -T on subsequent builds to skip protobuf downloading/building) + + +#### Testing the the CUDA-enabled benchmark via adb: +Build binaries first as above, then run: + +```bash +adb shell mkdir -p /data/local/tmp/lib64 +adb push $TEGRA_LIBS /data/local/tmp/lib64 +adb push tensorflow/contrib/makefile/gen/bin/android_arm64-v8a/benchmark /data/local/tmp +wget https://ci.tensorflow.org/view/Nightly/job/nightly-android/lastSuccessfulBuild/artifact/out/tensorflow_demo.apk +unzip tensorflow_demo.apk -d /tmp/tensorflow_demo +adb push /tmp/tensorflow_demo/assets/*.pb /data/local/tmp +adb shell "LD_LIBRARY_PATH=/data/local/tmp/lib64 /data/local/tmp/benchmark --graph=/data/local/tmp/tensorflow_inception_graph.pb" +``` + +#### Building the CUDA-enabled TensorFlow AAR with Bazel: +Build the native binaries first as above. Then, build the aar and package the native libs by executing the following: +```bash +mkdir -p /tmp/tf/jni/arm64-v8a +cp tensorflow/contrib/makefile/gen/lib/android_tegra/libtensorflow_*.so /tmp/tf/jni/arm64-v8a/ +cp $TEGRA_LIBS /tmp/tf/jni/arm64-v8a +bazel build //tensorflow/contrib/android:android_tensorflow_inference_java.aar +cp bazel-bin/tensorflow/contrib/android/android_tensorflow_inference_java.aar /tmp/tf/tensorflow.aar +cd /tmp/tf +chmod +w tensorflow.aar +zip -ur tensorflow.aar $(find jni -name *.so) +``` + +#### Building the CUDA-enabled TensorFlow Android demo with Bazel: +Build binaries first as above, then edit tensorflow/examples/android/BUILD and replace: +``` + srcs = [ + ":libtensorflow_demo.so", + "//tensorflow/contrib/android:libtensorflow_inference.so", + ], +``` +with: +``` +srcs = glob(["libs/arm64-v8a/*.so"]), +``` + +Then run: +```bash +# Create dir for native libs +mkdir -p tensorflow/examples/android/libs/arm64-v8a + +# Copy JetPack libs +cp $TEGRA_LIBS tensorflow/examples/android/libs/arm64-v8a + +# Copy native TensorFlow libraries +cp tensorflow/contrib/makefile/gen/lib/android_arm64-v8a/libtensorflow_*.so tensorflow/examples/android/libs/arm64-v8a/ + +# Build APK +bazel build -c opt --fat_apk_cpu=arm64-v8a tensorflow/android:tensorflow_demo + +# Install +adb install -r -f bazel-bin/tensorflow/examples/android/tensorflow_demo.apk +``` + +#### Building the CUDA-enabled Android demo with gradle/Android Studio: + +Add tensorflow/examples/android as an Android project in Android Studio as normal. + +Edit build.gradle and: +* set nativeBuildSystem = 'makefile' +* set cpuType = 'arm64-v8a' +* in "buildNativeMake", replace cpuType with 'tegra' (optional speedups like -T and ccache also work) +* set the environment "NDK_ROOT" var to $JETPACK/android-ndk-r13b + +Click "build apk" to build. + +Install: +```bash +adb install -r -f tensorflow/examples/android/gradleBuild/outputs/apk/debug/android-debug.apk +``` + ## iOS _Note: To use this library in an iOS application, see related instructions in @@ -262,6 +361,14 @@ to register ops and kernels. #### Optimization +The `build_all_ios.sh` script can take optional command-line arguments to +selectively register only for the operators used in your graph. + +```bash +tensorflow/contrib/makefile/build_all_ios.sh -a arm64 -g $HOME/graphs/inception/tensorflow_inception_graph.pb +``` +Please note this is an aggressive optimization of the operators and the resulting library may not work with other graphs but will reduce the size of the final library. + The `compile_ios_tensorflow.sh` script can take optional command-line arguments. The first argument will be passed as a C++ optimization flag and defaults to debug mode. If you are concerned about performance or are working on a release diff --git a/tensorflow/contrib/makefile/build_all_android.sh b/tensorflow/contrib/makefile/build_all_android.sh index 81cb17a311fd94aa397eb7a766cd8c668268759a..fc88f59e0948e1d3ed7cce9b809bf30ba280af12 100755 --- a/tensorflow/contrib/makefile/build_all_android.sh +++ b/tensorflow/contrib/makefile/build_all_android.sh @@ -18,7 +18,7 @@ set -e usage() { - echo "Usage: NDK_ROOT= $(basename "$0") [-Es:t:Tx:a:X]" + echo "Usage: NDK_ROOT= $(basename "$0") [-Es:t:Tx:a]" echo "-E enable experimental hexnn ops" echo "-s [sub_makefiles] sub makefiles separated by white space" echo "-t [build_target] build target for Android makefile [default=all]" @@ -26,7 +26,7 @@ usage() { echo "-x [hexagon library path] copy and hexagon libraries in the specified path" echo "-a [architecture] Architecture of target android [default=armeabi-v7a] \ (supported architecture list: \ -arm64-v8a armeabi armeabi-v7a mips mips64 x86 x86_64)" +arm64-v8a armeabi armeabi-v7a mips mips64 x86 x86_64 tegra)" exit 1 } @@ -50,6 +50,26 @@ while getopts "Es:t:Tx:a:" opt_name; do done shift $((OPTIND - 1)) +if [ "$ARCH" == "tegra" ]; then + if [[ -z "${JETPACK}" ]]; then + export JETPACK="$HOME/JetPack_Android_3.2" + fi + if [ ! -d ${JETPACK} ]; then + echo "Can't find Jetpack at ${JETPACK}" + echo "Set JETPACK= to specify a non-default Jetpack path" + exit -1 + fi + if [ ! -d ${JETPACK}/cuda ]; then + ln -s $(ls -d ${JETPACK}/cuda-*/|sort -r|head -n1) ${JETPACK}/cuda + fi + if [ ! -d ${JETPACK}/cuda ]; then + ln -s $(ls -d ${JETPACK}/cuda-*/|sort -r|head -n1) ${JETPACK}/cuda + fi + + export BUILD_FOR_TEGRA=1 + ARCH="arm64-v8a" +fi + # Make sure we're in the correct directory, at the root of the source tree. SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null && pwd)" cd "${SCRIPT_DIR}"/../../../ diff --git a/tensorflow/contrib/makefile/build_all_ios.sh b/tensorflow/contrib/makefile/build_all_ios.sh index 988e12b48287300004cc23c31cb4a20e63f72a27..2d9979183975e6a17527b40ef5ee1795ced44a7b 100755 --- a/tensorflow/contrib/makefile/build_all_ios.sh +++ b/tensorflow/contrib/makefile/build_all_ios.sh @@ -26,13 +26,16 @@ fi usage() { echo "Usage: $(basename "$0") [-a:T]" echo "-a [build_arch] build only for specified arch x86_64 [default=all]" + echo "-g [graph] optimize and selectively register ops only for this graph" echo "-T only build tensorflow (dont download other deps etc)" exit 1 } -while getopts "a:T" opt_name; do +DEFAULT_ARCH="i386 x86_64 armv7 armv7s arm64" +while getopts "a:g:T" opt_name; do case "$opt_name" in a) BUILD_ARCH="${OPTARG}";; + g) OPTIMIZE_FOR_GRAPH="${OPTARG}";; T) ONLY_MAKE_TENSORFLOW="true";; *) usage;; esac @@ -42,7 +45,8 @@ shift $((OPTIND - 1)) # Make sure we're in the correct directory, at the root of the source tree. SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd ${SCRIPT_DIR}/../../../ +TOP_SRCDIR="${SCRIPT_DIR}/../../../" +cd ${TOP_SRCDIR} source "${SCRIPT_DIR}/build_helper.subr" JOB_COUNT="${JOB_COUNT:-$(get_job_count)}" @@ -56,6 +60,32 @@ if [[ -n MACOSX_DEPLOYMENT_TARGET ]]; then export MACOSX_DEPLOYMENT_TARGET=$(sw_vers -productVersion) fi +PRNT_SLCTV_BIN="${TOP_SRCDIR}bazel-bin/tensorflow/python/tools/print_selective_registration_header" + +if [[ ! -z "${OPTIMIZE_FOR_GRAPH}" ]]; then + echo "Request to optimize for graph: ${OPTIMIZE_FOR_GRAPH}" + #Request to trim the OPs by selectively registering + if [ ! -f ${PRNT_SLCTV_BIN} ]; then + #Build bazel build tensorflow/python/tools:print_selective_registration_header + echo "${PRNT_SLCTV_BIN} not found. Trying to build it" + cd ${TOP_SRCDIR} + bazel build --copt="-DUSE_GEMM_FOR_CONV" tensorflow/python/tools:print_selective_registration_header + if [ ! -f ${PRNT_SLCTV_BIN} ]; then + echo "Building print_selective_registration_header failed" + echo "You may want to build TensorFlow with: " + echo "./configure" + echo "bazel build --copt="-DUSE_GEMM_FOR_CONV" tensorflow/python/tools:print_selective_registration_header" + echo "and then run this script again" + exit 1 + fi + else + echo "${PRNT_SLCTV_BIN} found. Using it" + ${PRNT_SLCTV_BIN} --graphs=${OPTIMIZE_FOR_GRAPH} > ${TOP_SRCDIR}/tensorflow/core/framework/ops_to_register.h + + fi + +fi + if [[ "${ONLY_MAKE_TENSORFLOW}" != "true" ]]; then # Remove any old files first. make -f tensorflow/contrib/makefile/Makefile clean @@ -64,8 +94,13 @@ if [[ "${ONLY_MAKE_TENSORFLOW}" != "true" ]]; then # Pull down the required versions of the frameworks we need. tensorflow/contrib/makefile/download_dependencies.sh - # Compile protobuf for the target iOS device architectures. - tensorflow/contrib/makefile/compile_ios_protobuf.sh + if [[ -z "${BUILD_ARCH}" ]]; then + # Compile protobuf for the target iOS device architectures. + tensorflow/contrib/makefile/compile_ios_protobuf.sh + else + # Compile protobuf for the target iOS device architectures. + tensorflow/contrib/makefile/compile_ios_protobuf.sh -a ${BUILD_ARCH} + fi fi # Compile nsync for the target iOS device architectures. @@ -80,13 +115,24 @@ else fi export HOST_NSYNC_LIB TARGET_NSYNC_LIB -if [[ -z "${BUILD_ARCH}" ]]; then - # build the ios tensorflow libraries. - tensorflow/contrib/makefile/compile_ios_tensorflow.sh -f "-O3" -h $HOST_NSYNC_LIB -n $TARGET_NSYNC_LIB -else +TF_CC_FLAGS="-O3" +TF_SCRIPT_FLAGS="-h ${HOST_NSYNC_LIB} -n ${TARGET_NSYNC_LIB}" + +if [[ ! -z "${OPTIMIZE_FOR_GRAPH}" ]]; then + # arch specified so build just that + TF_CC_FLAGS="${TF_CC_FLAGS} -DANDROID_TYPES=__ANDROID_TYPES_FULL__ -DSELECTIVE_REGISTRATION -DSUPPORT_SELECTIVE_REGISTRATION" + # The Makefile checks the env var to decide which ANDROID_TYPES to build + export ANDROID_TYPES="-D__ANDROID_TYPES_FULL__" +fi + +if [[ ! -z "${BUILD_ARCH}" ]]; then # arch specified so build just that - tensorflow/contrib/makefile/compile_ios_tensorflow.sh -f "-O3" -a "${BUILD_ARCH}" -h $HOST_NSYNC_LIB -n $TARGET_NSYNC_LIB + TF_SCRIPT_FLAGS="${TF_SCRIPT_FLAGS} -a ${BUILD_ARCH}" fi +# build the ios tensorflow libraries. +echo "Building TensorFlow with flags: ${TF_SCRIPT_FLAGS} -f ${TF_CC_FLAGS}" +tensorflow/contrib/makefile/compile_ios_tensorflow.sh ${TF_SCRIPT_FLAGS} -f "${TF_CC_FLAGS}" + # Creates a static universal library in # tensorflow/contrib/makefile/gen/lib/libtensorflow-core.a diff --git a/tensorflow/contrib/makefile/download_dependencies.sh b/tensorflow/contrib/makefile/download_dependencies.sh index b61044130897cf0dddc37e460b4e1618c3a7e2e9..4ae18b2cef28335a90bbc967529c0cf76b0a5da2 100755 --- a/tensorflow/contrib/makefile/download_dependencies.sh +++ b/tensorflow/contrib/makefile/download_dependencies.sh @@ -34,6 +34,7 @@ PROTOBUF_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/protobuf/. RE2_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/re2/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1)" FFT2D_URL="$(grep -o 'http.*fft\.tgz' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" ABSL_URL="$(grep -o 'https://github.com/abseil/abseil-cpp/.*tar.gz' "${BZL_FILE_PATH}" | head -n1)" +CUB_URL="$(grep -o 'https.*cub/archive.*zip' "${BZL_FILE_PATH}" | grep -v bazel-mirror | head -n1)" # TODO(petewarden): Some new code in Eigen triggers a clang bug with iOS arm64, # so work around it by patching the source. @@ -62,12 +63,17 @@ download_and_extract() { elif [[ "${url}" == *zip ]]; then tempdir=$(mktemp -d) tempdir2=$(mktemp -d) - wget -P ${tempdir} ${url} - unzip ${tempdir}/* -d ${tempdir2} + if [[ "$OSTYPE" == "darwin"* ]]; then + # macOS (AKA darwin) doesn't have wget. + (cd "${tempdir}"; curl --remote-name --silent --location "${url}") + else + wget -P "${tempdir}" "${url}" + fi + unzip "${tempdir}"/* -d "${tempdir2}" # unzip has no strip components, so unzip to a temp dir, and move the files # we want from the tempdir to destination. - cp -R ${tempdir2}/*/* ${dir}/ - rm -rf ${tempdir2} ${tempdir} + cp -R "${tempdir2}"/*/* "${dir}"/ + rm -rf "${tempdir2}" "${tempdir}" fi # Delete any potential BUILD files, which would interfere with Bazel builds. @@ -82,6 +88,7 @@ download_and_extract "${PROTOBUF_URL}" "${DOWNLOADS_DIR}/protobuf" download_and_extract "${RE2_URL}" "${DOWNLOADS_DIR}/re2" download_and_extract "${FFT2D_URL}" "${DOWNLOADS_DIR}/fft2d" download_and_extract "${ABSL_URL}" "${DOWNLOADS_DIR}/absl" +download_and_extract "${CUB_URL}" "${DOWNLOADS_DIR}/cub/external/cub_archive" replace_by_sed 's#static uint32x4_t p4ui_CONJ_XOR = vld1q_u32( conj_XOR_DATA );#static uint32x4_t p4ui_CONJ_XOR; // = vld1q_u32( conj_XOR_DATA ); - Removed by script#' \ "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h" diff --git a/tensorflow/contrib/makefile/samples/build_and_run_inception_hexagon.sh b/tensorflow/contrib/makefile/samples/build_and_run_inception_hexagon.sh index 861bb885c7031b996b48dbc50887cfce55c638f3..421ddd210fd5b1ac6487918d5797eab5953316df 100755 --- a/tensorflow/contrib/makefile/samples/build_and_run_inception_hexagon.sh +++ b/tensorflow/contrib/makefile/samples/build_and_run_inception_hexagon.sh @@ -36,7 +36,7 @@ while getopts "bc:Eps" opt_name; do b) BUILD_ONLY="true";; c) TEST_COUNT="${OPTARG}";; E) ENABLE_EXPERIMENTAL_HEXNN_OPS="true";; - p) USE_PREBUILT_HEXAOGON_BINARIES="true";; + p) USE_PREBUILT_HEXAGON_BINARIES="true";; s) SKIP_DOWNLOAD_IF_EXIST="true";; *) usage;; esac @@ -49,7 +49,7 @@ if [[ -z "${NDK_ROOT}" ]]; then exit 1 fi -if [[ "${USE_PREBUILT_HEXAOGON_BINARIES}" != "true" && +if [[ "${USE_PREBUILT_HEXAGON_BINARIES}" != "true" && -z "${QUALCOMM_SDK}" ]]; then echo "QUALCOMM_SDK is empty" 1>&2 usage @@ -76,13 +76,15 @@ GEN_LIBS_DIR="${GEN_DIR}/libs" GEN_DOWNLOAD_DIR="${GEN_DIR}/downloads" URL_BASE="https://storage.googleapis.com/download.tensorflow.org" +ARCH="armeabi-v7a" + source "${SCRIPT_DIR}/../build_helper.subr" rm -rf "${GEN_DIR}" mkdir -p "${GEN_LIBS_DIR}" mkdir -p "${GEN_DOWNLOAD_DIR}" -if [[ "${USE_PREBUILT_HEXAOGON_BINARIES}" == "true" ]]; then +if [[ "${USE_PREBUILT_HEXAGON_BINARIES}" == "true" ]]; then echo "Download prebuilt hexagon binaries" if [[ "${BUILD_ONLY}" != "true" ]]; then CONTROLLER_PUSH_DEST="/data/local/tmp" @@ -219,7 +221,7 @@ if [[ "${BUILD_ONLY}" != "true" ]]; then adb push "${GEN_LIBS_DIR}/libhexagon_nn_skel.so" "/vendor/lib/rfsa/adsp" adb push -p \ - "${TF_ROOT_DIR}/tensorflow/contrib/makefile/gen/bin/hexagon_graph_execution" \ + "${TF_ROOT_DIR}/tensorflow/contrib/makefile/gen/bin/android_${ARCH}/hexagon_graph_execution" \ "/data/local/tmp/" adb wait-for-device adb shell chmod "${ANDROID_EXEC_FILE_MODE}" \ diff --git a/tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in b/tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in index 26c1ad4947363e98d9bb8e400f40290fb87b2e4e..3081084ee76e41de801f49a67c1fec07f4ff03b9 100644 --- a/tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in +++ b/tensorflow/contrib/makefile/sub_makefiles/android/Makefile.in @@ -48,13 +48,13 @@ INFERENCE_OBJS := $(addprefix $(OBJDIR), $(INFERENCE_SRCS:.cc=.o)) INFERENCE_SO_NAME := libtensorflow_inference.so INFERENCE_SO_PATH := $(LIBDIR)$(INFERENCE_SO_NAME) -$(INFERENCE_SO_PATH): $(LIB_OBJS) $(INFERENCE_OBJS) +$(INFERENCE_SO_PATH): $(LIB_OBJS) $(INFERENCE_OBJS) $(CUDA_LIB_DEPS) @mkdir -p $(dir $@) $(CXX) $(CXXFLAGS) $(INCLUDES) \ - -o $@ $(INFERENCE_OBJS) $(LIB_OBJS) \ + -o $@ $(INFERENCE_OBJS) $(LIB_OBJS) $(TEGRA_LIBS) \ $(LIBFLAGS) $(LDFLAGS) \ -shared -Wl,-soname,$(INFERENCE_SO_NAME) \ - $(LIBS) + $(LIBS) $(CUDA_LIBS) $(INFERENCE_SO_NAME): $(INFERENCE_SO_PATH) diff --git a/tensorflow/contrib/makefile/tf_op_files.txt b/tensorflow/contrib/makefile/tf_op_files.txt index 5f275663986f9d480659880ab601eeb5c41037be..5a812af4e95fe7a05b9c2634b0cc1d860fb7f619 100644 --- a/tensorflow/contrib/makefile/tf_op_files.txt +++ b/tensorflow/contrib/makefile/tf_op_files.txt @@ -91,6 +91,7 @@ tensorflow/core/kernels/reduction_ops_max.cc tensorflow/core/kernels/reduction_ops_common.cc tensorflow/core/kernels/reduction_ops_any.cc tensorflow/core/kernels/reduction_ops_all.cc +tensorflow/core/kernels/roll_op.cc tensorflow/core/kernels/queue_ops.cc tensorflow/core/kernels/queue_base.cc tensorflow/core/kernels/pooling_ops_common.cc @@ -270,6 +271,7 @@ tensorflow/core/ops/parsing_ops.cc tensorflow/core/ops/no_op.cc tensorflow/core/ops/nn_ops.cc tensorflow/core/ops/nn_grad.cc +tensorflow/core/ops/manip_ops.cc tensorflow/core/ops/math_ops.cc tensorflow/core/ops/math_grad.cc tensorflow/core/ops/logging_ops.cc @@ -291,3 +293,4 @@ tensorflow/core/kernels/batchtospace_op.cc tensorflow/core/kernels/warn_about_ints.cc tensorflow/core/kernels/segment_reduction_ops.cc tensorflow/core/kernels/batch_util.cc +tensorflow/core/ops/audio_ops.cc diff --git a/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc b/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc index 39c0d5af45b4a81fa4dde0b5deac14a3af372cbb..974fb537499c5ea4591a0a128f53d2dea67b9e57 100644 --- a/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc +++ b/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc @@ -80,9 +80,9 @@ REGISTER_KERNEL_BUILDER(Name("BytesLimit").Device(DEVICE_GPU).HostMemory("out"), BytesLimitOp); #ifdef TENSORFLOW_USE_SYCL -REGISTER_KERNEL_BUILDER(Name("BytesLimit").Device(DEVICE_SYCL).HostMemory("out"), - BytesLimitOp); -#endif // TENSORFLOW_USE_SYCL +REGISTER_KERNEL_BUILDER( + Name("BytesLimit").Device(DEVICE_SYCL).HostMemory("out"), BytesLimitOp); +#endif // TENSORFLOW_USE_SYCL // Op that measures the peak memory in bytes. class MaxBytesInUseOp : public MemoryStatsOp { @@ -107,6 +107,6 @@ REGISTER_KERNEL_BUILDER( REGISTER_KERNEL_BUILDER( Name("MaxBytesInUse").Device(DEVICE_SYCL).HostMemory("out"), MaxBytesInUseOp); -#endif // TENSORFLOW_USE_SYCL +#endif // TENSORFLOW_USE_SYCL } // namespace tensorflow diff --git a/tensorflow/contrib/meta_graph_transform/meta_graph_transform.py b/tensorflow/contrib/meta_graph_transform/meta_graph_transform.py index 2932ae1c8df32cd936cff932b061571c513fda79..ff88b4fa841673fc52b9f6fdc5ca43d30c44bbfd 100644 --- a/tensorflow/contrib/meta_graph_transform/meta_graph_transform.py +++ b/tensorflow/contrib/meta_graph_transform/meta_graph_transform.py @@ -171,7 +171,14 @@ def _clean_save_and_restore(graph_def, op, removed_op_names): shape_op_value_tensor.tensor_shape.dim[0].size = len(shapes) op.attr['dtypes'].list.type[:] = dtypes + if not name_op.attr['_output_shapes'].list.shape: + name_op.attr['_output_shapes'].list.shape.add() + name_op.attr['_output_shapes'].list.shape[0].dim.add() name_op.attr['_output_shapes'].list.shape[0].dim[0].size = len(names) + + if not shape_op.attr['_output_shapes'].list.shape: + shape_op.attr['_output_shapes'].list.shape.add() + shape_op.attr['_output_shapes'].list.shape[0].dim.add() shape_op.attr['_output_shapes'].list.shape[0].dim[0].size = len(shapes) diff --git a/tensorflow/contrib/metrics/BUILD b/tensorflow/contrib/metrics/BUILD index 9de664c822bf7a9abf7b8082f444c61dfa45f499..e90c525113348532a3ebdadde7e712bf2d98cee9 100644 --- a/tensorflow/contrib/metrics/BUILD +++ b/tensorflow/contrib/metrics/BUILD @@ -43,6 +43,7 @@ py_library( "//tensorflow/python:util", "//tensorflow/python:variable_scope", "//tensorflow/python:weights_broadcast_ops", + "//tensorflow/python/ops/distributions", ], ) diff --git a/tensorflow/contrib/metrics/__init__.py b/tensorflow/contrib/metrics/__init__.py index 27dad5379a2e56b91960a1f2274610e4f2568dbc..de02dc8f457364450929776035829d86035d706b 100644 --- a/tensorflow/contrib/metrics/__init__.py +++ b/tensorflow/contrib/metrics/__init__.py @@ -16,6 +16,7 @@ See the @{$python/contrib.metrics} guide. +@@auc_with_confidence_intervals @@streaming_accuracy @@streaming_mean @@streaming_recall @@ -66,6 +67,7 @@ See the @{$python/contrib.metrics} guide. @@set_intersection @@set_size @@set_union +@@cohen_kappa @@count @@precision_recall_at_equal_thresholds @@recall_at_precision @@ -82,6 +84,8 @@ from tensorflow.contrib.metrics.python.ops.confusion_matrix_ops import confusion from tensorflow.contrib.metrics.python.ops.histogram_ops import auc_using_histogram from tensorflow.contrib.metrics.python.ops.metric_ops import aggregate_metric_map from tensorflow.contrib.metrics.python.ops.metric_ops import aggregate_metrics +from tensorflow.contrib.metrics.python.ops.metric_ops import auc_with_confidence_intervals +from tensorflow.contrib.metrics.python.ops.metric_ops import cohen_kappa from tensorflow.contrib.metrics.python.ops.metric_ops import count from tensorflow.contrib.metrics.python.ops.metric_ops import precision_recall_at_equal_thresholds from tensorflow.contrib.metrics.python.ops.metric_ops import recall_at_precision diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops.py b/tensorflow/contrib/metrics/python/ops/metric_ops.py index 2f2798563481cc0c53360944f967e6b31991057d..31e274c5fd7c670458b1b40a4f58c668a23776c7 100644 --- a/tensorflow/contrib/metrics/python/ops/metric_ops.py +++ b/tensorflow/contrib/metrics/python/ops/metric_ops.py @@ -24,10 +24,12 @@ from __future__ import print_function import collections as collections_lib +from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops +from tensorflow.python.ops import confusion_matrix from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics @@ -36,6 +38,7 @@ from tensorflow.python.ops import nn from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import weights_broadcast_ops +from tensorflow.python.ops.distributions.normal import Normal from tensorflow.python.util.deprecation import deprecated # Epsilon constant used to represent extremely small quantity. @@ -337,9 +340,9 @@ def streaming_mean_tensor(values, name=name) -@deprecated( - None, 'Please switch to tf.metrics.accuracy. Note that the order of the ' - 'labels and predictions arguments has been switched.') +@deprecated(None, + 'Please switch to tf.metrics.accuracy. Note that the order of the ' + 'labels and predictions arguments has been switched.') def streaming_accuracy(predictions, labels, weights=None, @@ -737,7 +740,7 @@ def _streaming_confusion_matrix_at_thresholds(predictions, else: for include in includes: if include not in all_includes: - raise ValueError('Invaild key: %s.' % include) + raise ValueError('Invalid key: %s.' % include) predictions, labels, weights = metrics_impl._remove_squeezable_dimensions( # pylint: disable=protected-access predictions, labels, weights) @@ -934,8 +937,9 @@ def streaming_curve_points(labels=None, if curve != 'ROC' and curve != 'PR': raise ValueError('curve must be either ROC or PR, %s unknown' % (curve)) kepsilon = _EPSILON # to account for floating point imprecisions - thresholds = [(i + 1) * 1.0 / (num_thresholds - 1) - for i in range(num_thresholds - 2)] + thresholds = [ + (i + 1) * 1.0 / (num_thresholds - 1) for i in range(num_thresholds - 2) + ] thresholds = [0.0 - kepsilon] + thresholds + [1.0 + kepsilon] values, update_ops = _streaming_confusion_matrix_at_thresholds( @@ -971,9 +975,8 @@ def streaming_curve_points(labels=None, return points, update_op -@deprecated( - None, 'Please switch to tf.metrics.auc. Note that the order of the ' - 'labels and predictions arguments has been switched.') +@deprecated(None, 'Please switch to tf.metrics.auc. Note that the order of the ' + 'labels and predictions arguments has been switched.') def streaming_auc(predictions, labels, weights=None, @@ -1103,8 +1106,7 @@ def _compute_dynamic_auc(labels, predictions, curve='ROC'): # For conformance, set precision to 1 when the number of positive # classifications is 0. y_axis_values = array_ops.where( - math_ops.greater(splits, 0), - math_ops.truediv(true_positives, splits), + math_ops.greater(splits, 0), math_ops.truediv(true_positives, splits), array_ops.ones_like(true_positives, dtype=dtypes.float64)) # Calculate trapezoid areas. @@ -1117,9 +1119,8 @@ def _compute_dynamic_auc(labels, predictions, curve='ROC'): # exception seems excessive) so we return 0, otherwise we finish computing. return control_flow_ops.cond( math_ops.logical_or( - math_ops.equal(total_positive, 0), - math_ops.equal(total_positive, size) - ), + math_ops.equal(total_positive, 0), math_ops.equal( + total_positive, size)), true_fn=lambda: array_ops.constant(0, dtypes.float64), false_fn=continue_computing_dynamic_auc) @@ -1183,10 +1184,10 @@ def streaming_dynamic_auc(labels, array_ops.ones_like(labels, dtypes.int64), message='labels must be 0 or 1, at least one is >1') ]): - preds_accum, update_preds = streaming_concat(predictions, - name='concat_preds') - labels_accum, update_labels = streaming_concat(labels, - name='concat_labels') + preds_accum, update_preds = streaming_concat( + predictions, name='concat_preds') + labels_accum, update_labels = streaming_concat( + labels, name='concat_labels') update_op = control_flow_ops.group(update_labels, update_preds) auc = _compute_dynamic_auc(labels_accum, preds_accum, curve=curve) if updates_collections: @@ -1196,6 +1197,295 @@ def streaming_dynamic_auc(labels, return auc, update_op +def _compute_placement_auc(labels, predictions, weights, alpha, + logit_transformation, is_valid): + """Computes the AUC and asymptotic normally distributed confidence interval. + + The calculations are achieved using the fact that AUC = P(Y_1>Y_0) and the + concept of placement values for each labeled group, as presented by Delong and + Delong (1988). The actual algorithm used is a more computationally efficient + approach presented by Sun and Xu (2014). This could be slow for large batches, + but has the advantage of not having its results degrade depending on the + distribution of predictions. + + Args: + labels: A `Tensor` of ground truth labels with the same shape as + `predictions` with values of 0 or 1 and type `int64`. + predictions: A 1-D `Tensor` of predictions whose values are `float64`. + weights: `Tensor` whose rank is either 0, or the same rank as `labels`. + alpha: Confidence interval level desired. + logit_transformation: A boolean value indicating whether the estimate should + be logit transformed prior to calculating the confidence interval. Doing + so enforces the restriction that the AUC should never be outside the + interval [0,1]. + is_valid: A bool tensor describing whether the input is valid. + + Returns: + A 1-D `Tensor` containing the area-under-curve, lower, and upper confidence + interval values. + """ + # Disable the invalid-name checker so that we can capitalize the name. + # pylint: disable=invalid-name + AucData = collections_lib.namedtuple('AucData', ['auc', 'lower', 'upper']) + # pylint: enable=invalid-name + + # If all the labels are the same or if number of observations are too few, + # AUC isn't well-defined + size = array_ops.size(predictions, out_type=dtypes.int32) + + # Count the total number of positive and negative labels in the input. + total_0 = math_ops.reduce_sum( + math_ops.cast(1 - labels, weights.dtype) * weights) + total_1 = math_ops.reduce_sum( + math_ops.cast(labels, weights.dtype) * weights) + + # Sort the predictions ascending, as well as + # (i) the corresponding labels and + # (ii) the corresponding weights. + ordered_predictions, indices = nn.top_k(predictions, k=size, sorted=True) + ordered_predictions = array_ops.reverse( + ordered_predictions, axis=array_ops.zeros(1, dtypes.int32)) + indices = array_ops.reverse(indices, axis=array_ops.zeros(1, dtypes.int32)) + ordered_labels = array_ops.gather(labels, indices) + ordered_weights = array_ops.gather(weights, indices) + + # We now compute values required for computing placement values. + + # We generate a list of indices (segmented_indices) of increasing order. An + # index is assigned for each unique prediction float value. Prediction + # values that are the same share the same index. + _, segmented_indices = array_ops.unique(ordered_predictions) + + # We create 2 tensors of weights. weights_for_true is non-zero for true + # labels. weights_for_false is non-zero for false labels. + float_labels_for_true = math_ops.cast(ordered_labels, dtypes.float32) + float_labels_for_false = 1.0 - float_labels_for_true + weights_for_true = ordered_weights * float_labels_for_true + weights_for_false = ordered_weights * float_labels_for_false + + # For each set of weights with the same segmented indices, we add up the + # weight values. Note that for each label, we deliberately rely on weights + # for the opposite label. + weight_totals_for_true = math_ops.segment_sum(weights_for_false, + segmented_indices) + weight_totals_for_false = math_ops.segment_sum(weights_for_true, + segmented_indices) + + # These cumulative sums of weights importantly exclude the current weight + # sums. + cum_weight_totals_for_true = math_ops.cumsum(weight_totals_for_true, + exclusive=True) + cum_weight_totals_for_false = math_ops.cumsum(weight_totals_for_false, + exclusive=True) + + # Compute placement values using the formula. Values with the same segmented + # indices and labels share the same placement values. + placements_for_true = ( + (cum_weight_totals_for_true + weight_totals_for_true / 2.0) / + (math_ops.reduce_sum(weight_totals_for_true) + _EPSILON)) + placements_for_false = ( + (cum_weight_totals_for_false + weight_totals_for_false / 2.0) / + (math_ops.reduce_sum(weight_totals_for_false) + _EPSILON)) + + # We expand the tensors of placement values (for each label) so that their + # shapes match that of predictions. + placements_for_true = array_ops.gather(placements_for_true, segmented_indices) + placements_for_false = array_ops.gather(placements_for_false, + segmented_indices) + + # Select placement values based on the label for each index. + placement_values = ( + placements_for_true * float_labels_for_true + + placements_for_false * float_labels_for_false) + + # Split placement values by labeled groups. + placement_values_0 = placement_values * math_ops.cast( + 1 - ordered_labels, weights.dtype) + weights_0 = ordered_weights * math_ops.cast( + 1 - ordered_labels, weights.dtype) + placement_values_1 = placement_values * math_ops.cast( + ordered_labels, weights.dtype) + weights_1 = ordered_weights * math_ops.cast( + ordered_labels, weights.dtype) + + # Calculate AUC using placement values + auc_0 = (math_ops.reduce_sum(weights_0 * (1. - placement_values_0)) / + (total_0 + _EPSILON)) + auc_1 = (math_ops.reduce_sum(weights_1 * (placement_values_1)) / + (total_1 + _EPSILON)) + auc = array_ops.where(math_ops.less(total_0, total_1), auc_1, auc_0) + + # Calculate variance and standard error using the placement values. + var_0 = ( + math_ops.reduce_sum( + weights_0 * math_ops.square(1. - placement_values_0 - auc_0)) / + (total_0 - 1. + _EPSILON)) + var_1 = ( + math_ops.reduce_sum( + weights_1 * math_ops.square(placement_values_1 - auc_1)) / + (total_1 - 1. + _EPSILON)) + auc_std_err = math_ops.sqrt( + (var_0 / (total_0 + _EPSILON)) + (var_1 / (total_1 + _EPSILON))) + + # Calculate asymptotic normal confidence intervals + std_norm_dist = Normal(loc=0., scale=1.) + z_value = std_norm_dist.quantile((1.0 - alpha) / 2.0) + if logit_transformation: + estimate = math_ops.log(auc / (1. - auc + _EPSILON)) + std_err = auc_std_err / (auc * (1. - auc + _EPSILON)) + transformed_auc_lower = estimate + (z_value * std_err) + transformed_auc_upper = estimate - (z_value * std_err) + def inverse_logit_transformation(x): + exp_negative = math_ops.exp(math_ops.negative(x)) + return 1. / (1. + exp_negative + _EPSILON) + + auc_lower = inverse_logit_transformation(transformed_auc_lower) + auc_upper = inverse_logit_transformation(transformed_auc_upper) + else: + estimate = auc + std_err = auc_std_err + auc_lower = estimate + (z_value * std_err) + auc_upper = estimate - (z_value * std_err) + + ## If estimate is 1 or 0, no variance is present so CI = 1 + ## n.b. This can be misleading, since number obs can just be too low. + lower = array_ops.where( + math_ops.logical_or( + math_ops.equal(auc, array_ops.ones_like(auc)), + math_ops.equal(auc, array_ops.zeros_like(auc))), + auc, auc_lower) + upper = array_ops.where( + math_ops.logical_or( + math_ops.equal(auc, array_ops.ones_like(auc)), + math_ops.equal(auc, array_ops.zeros_like(auc))), + auc, auc_upper) + + # If all the labels are the same, AUC isn't well-defined (but raising an + # exception seems excessive) so we return 0, otherwise we finish computing. + trivial_value = array_ops.constant(0.0) + + return AucData(*control_flow_ops.cond( + is_valid, lambda: [auc, lower, upper], lambda: [trivial_value]*3)) + + +def auc_with_confidence_intervals(labels, + predictions, + weights=None, + alpha=0.95, + logit_transformation=True, + metrics_collections=(), + updates_collections=(), + name=None): + """Computes the AUC and asymptotic normally distributed confidence interval. + + USAGE NOTE: this approach requires storing all of the predictions and labels + for a single evaluation in memory, so it may not be usable when the evaluation + batch size and/or the number of evaluation steps is very large. + + Computes the area under the ROC curve and its confidence interval using + placement values. This has the advantage of being resilient to the + distribution of predictions by aggregating across batches, accumulating labels + and predictions and performing the final calculation using all of the + concatenated values. + + Args: + labels: A `Tensor` of ground truth labels with the same shape as `labels` + and with values of 0 or 1 whose values are castable to `int64`. + predictions: A `Tensor` of predictions whose values are castable to + `float64`. Will be flattened into a 1-D `Tensor`. + weights: Optional `Tensor` whose rank is either 0, or the same rank as + `labels`. + alpha: Confidence interval level desired. + logit_transformation: A boolean value indicating whether the estimate should + be logit transformed prior to calculating the confidence interval. Doing + so enforces the restriction that the AUC should never be outside the + interval [0,1]. + metrics_collections: An optional iterable of collections that `auc` should + be added to. + updates_collections: An optional iterable of collections that `update_op` + should be added to. + name: An optional name for the variable_scope that contains the metric + variables. + + Returns: + auc: A 1-D `Tensor` containing the current area-under-curve, lower, and + upper confidence interval values. + update_op: An operation that concatenates the input labels and predictions + to the accumulated values. + + Raises: + ValueError: If `labels`, `predictions`, and `weights` have mismatched shapes + or if `alpha` isn't in the range (0,1). + """ + if not (alpha > 0 and alpha < 1): + raise ValueError('alpha must be between 0 and 1; currently %.02f' % alpha) + + if weights is None: + weights = array_ops.ones_like(predictions) + + with variable_scope.variable_scope( + name, + default_name='auc_with_confidence_intervals', + values=[labels, predictions, weights]): + + predictions, labels, weights = metrics_impl._remove_squeezable_dimensions( # pylint: disable=protected-access + predictions=predictions, + labels=labels, + weights=weights) + + total_weight = math_ops.reduce_sum(weights) + + weights = array_ops.reshape(weights, [-1]) + predictions = array_ops.reshape( + math_ops.cast(predictions, dtypes.float64), [-1]) + labels = array_ops.reshape(math_ops.cast(labels, dtypes.int64), [-1]) + + with ops.control_dependencies([ + check_ops.assert_greater_equal( + labels, + array_ops.zeros_like(labels, dtypes.int64), + message='labels must be 0 or 1, at least one is <0'), + check_ops.assert_less_equal( + labels, + array_ops.ones_like(labels, dtypes.int64), + message='labels must be 0 or 1, at least one is >1'), + ]): + preds_accum, update_preds = streaming_concat( + predictions, name='concat_preds') + labels_accum, update_labels = streaming_concat(labels, + name='concat_labels') + weights_accum, update_weights = streaming_concat( + weights, name='concat_weights') + update_op_for_valid_case = control_flow_ops.group( + update_labels, update_preds, update_weights) + + # Only perform updates if this case is valid. + all_labels_positive_or_0 = math_ops.logical_and( + math_ops.equal(math_ops.reduce_min(labels), 0), + math_ops.equal(math_ops.reduce_max(labels), 1)) + sums_of_weights_at_least_1 = math_ops.greater_equal(total_weight, 1.0) + is_valid = math_ops.logical_and(all_labels_positive_or_0, + sums_of_weights_at_least_1) + + update_op = control_flow_ops.cond( + sums_of_weights_at_least_1, + lambda: update_op_for_valid_case, control_flow_ops.no_op) + + auc = _compute_placement_auc( + labels_accum, + preds_accum, + weights_accum, + alpha=alpha, + logit_transformation=logit_transformation, + is_valid=is_valid) + + if updates_collections: + ops.add_to_collections(updates_collections, update_op) + if metrics_collections: + ops.add_to_collections(metrics_collections, auc) + return auc, update_op + + def precision_recall_at_equal_thresholds(labels, predictions, weights=None, @@ -1226,7 +1516,7 @@ def precision_recall_at_equal_thresholds(labels, predictions: A floating point `Tensor` of arbitrary shape and whose values are in the range `[0, 1]`. weights: Optional; If provided, a `Tensor` that has the same dtype as, - and broadcastable to, `predictions`. This tensor is multplied by counts. + and broadcastable to, `predictions`. This tensor is multiplied by counts. num_thresholds: Optional; Number of thresholds, evenly distributed in `[0, 1]`. Should be `>= 2`. Defaults to 201. Note that the number of bins is 1 less than `num_thresholds`. Using an even `num_thresholds` value @@ -1569,9 +1859,9 @@ def streaming_precision_at_thresholds(predictions, name=name) -@deprecated( - None, 'Please switch to tf.metrics.recall_at_thresholds. Note that the ' - 'order of the labels and predictions arguments has been switched.') +@deprecated(None, + 'Please switch to tf.metrics.recall_at_thresholds. Note that the ' + 'order of the labels and predictions arguments has been switched.') def streaming_recall_at_thresholds(predictions, labels, thresholds, @@ -3297,9 +3587,143 @@ def count(values, return count_, update_op +def cohen_kappa(labels, + predictions_idx, + num_classes, + weights=None, + metrics_collections=None, + updates_collections=None, + name=None): + """Calculates Cohen's kappa. + + [Cohen's kappa](https://en.wikipedia.org/wiki/Cohen's_kappa) is a statistic + that measures inter-annotator agreement. + + The `cohen_kappa` function calculates the confusion matrix, and creates three + local variables to compute the Cohen's kappa: `po`, `pe_row`, and `pe_col`, + which refer to the diagonal part, rows and columns totals of the confusion + matrix, respectively. This value is ultimately returned as `kappa`, an + idempotent operation that is calculated by + + pe = (pe_row * pe_col) / N + k = (sum(po) - sum(pe)) / (N - sum(pe)) + + For estimation of the metric over a stream of data, the function creates an + `update_op` operation that updates these variables and returns the + `kappa`. `update_op` weights each prediction by the corresponding value in + `weights`. + + 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 `weights` is `None`, weights default to 1. Use weights of 0 to mask values. + + NOTE: Equivalent to `sklearn.metrics.cohen_kappa_score`, but the method + doesn't support weighted matrix yet. + + Args: + labels: 1-D `Tensor` of real labels for the classification task. Must be + one of the following types: int16, int32, int64. + predictions_idx: 1-D `Tensor` of predicted class indices for a given + classification. Must have the same type as `labels`. + num_classes: The possible number of labels. + weights: Optional `Tensor` whose shape matches `predictions`. + metrics_collections: An optional list of collections that `kappa` should + be added to. + updates_collections: An optional list of collections that `update_op` should + be added to. + name: An optional variable_scope name. + + Returns: + kappa: Scalar float `Tensor` representing the current Cohen's kappa. + update_op: `Operation` that increments `po`, `pe_row` and `pe_col` + variables appropriately and whose value matches `kappa`. + + Raises: + ValueError: If `num_classes` is less than 2, or `predictions` and `labels` + have mismatched shapes, or if `weights` is not `None` and its shape + doesn't match `predictions`, or if either `metrics_collections` or + `updates_collections` are not a list or tuple. + RuntimeError: If eager execution is enabled. + """ + if context.in_eager_mode(): + raise RuntimeError('tf.contrib.metrics.cohen_kappa is not supported' + 'when eager execution is enabled.') + if num_classes < 2: + raise ValueError('`num_classes` must be >= 2.' + 'Found: {}'.format(num_classes)) + with variable_scope.variable_scope(name, 'cohen_kappa', + (labels, predictions_idx, weights)): + # Convert 2-dim (num, 1) to 1-dim (num,) + labels.get_shape().with_rank_at_most(2) + if labels.get_shape().ndims == 2: + labels = array_ops.squeeze(labels, axis=[-1]) + predictions_idx, labels, weights = ( + metrics_impl._remove_squeezable_dimensions( # pylint: disable=protected-access + predictions=predictions_idx, + labels=labels, + weights=weights)) + predictions_idx.get_shape().assert_is_compatible_with(labels.get_shape()) + + stat_dtype = ( + dtypes.int64 + if weights is None or weights.dtype.is_integer else dtypes.float32) + po = metrics_impl.metric_variable((num_classes,), stat_dtype, name='po') + pe_row = metrics_impl.metric_variable( + (num_classes,), stat_dtype, name='pe_row') + pe_col = metrics_impl.metric_variable( + (num_classes,), stat_dtype, name='pe_col') + + # Table of the counts of agreement: + counts_in_table = confusion_matrix.confusion_matrix( + labels, + predictions_idx, + num_classes=num_classes, + weights=weights, + dtype=stat_dtype, + name='counts_in_table') + + po_t = array_ops.diag_part(counts_in_table) + pe_row_t = math_ops.reduce_sum(counts_in_table, axis=0) + pe_col_t = math_ops.reduce_sum(counts_in_table, axis=1) + update_po = state_ops.assign_add(po, po_t) + update_pe_row = state_ops.assign_add(pe_row, pe_row_t) + update_pe_col = state_ops.assign_add(pe_col, pe_col_t) + + def _calculate_k(po, pe_row, pe_col, name): + po_sum = math_ops.reduce_sum(po) + total = math_ops.reduce_sum(pe_row) + pe_sum = math_ops.reduce_sum( + metrics_impl._safe_div( # pylint: disable=protected-access + pe_row * pe_col, total, None)) + po_sum, pe_sum, total = (math_ops.to_double(po_sum), + math_ops.to_double(pe_sum), + math_ops.to_double(total)) + # kappa = (po - pe) / (N - pe) + k = metrics_impl._safe_scalar_div( # pylint: disable=protected-access + po_sum - pe_sum, + total - pe_sum, + name=name) + return k + + kappa = _calculate_k(po, pe_row, pe_col, name='value') + update_op = _calculate_k( + update_po, update_pe_row, update_pe_col, name='update_op') + + if metrics_collections: + ops.add_to_collections(metrics_collections, kappa) + + if updates_collections: + ops.add_to_collections(updates_collections, update_op) + + return kappa, update_op + + __all__ = [ + 'auc_with_confidence_intervals', 'aggregate_metric_map', 'aggregate_metrics', + 'cohen_kappa', 'count', 'precision_recall_at_equal_thresholds', 'recall_at_precision', diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops_test.py b/tensorflow/contrib/metrics/python/ops/metric_ops_test.py index f05ae394e6b46809f9f3f963733076f1a3933059..b387f26c0195432fb972dac450d2919bdaa702a1 100644 --- a/tensorflow/contrib/metrics/python/ops/metric_ops_test.py +++ b/tensorflow/contrib/metrics/python/ops/metric_ops_test.py @@ -46,8 +46,7 @@ def _enqueue_vector(sess, queue, values, shape=None): shape = (1, len(values)) dtype = queue.dtypes[0] sess.run( - queue.enqueue(constant_op.constant( - values, dtype=dtype, shape=shape))) + queue.enqueue(constant_op.constant(values, dtype=dtype, shape=shape))) def _binary_2d_label_to_sparse_value(labels): @@ -79,8 +78,8 @@ def _binary_2d_label_to_sparse_value(labels): batch += 1 shape = [len(labels), len(labels[0])] return sparse_tensor.SparseTensorValue( - np.array(indices, np.int64), - np.array(values, np.int64), np.array(shape, np.int64)) + np.array(indices, np.int64), np.array(values, np.int64), + np.array(shape, np.int64)) def _binary_2d_label_to_sparse(labels): @@ -125,8 +124,8 @@ def _binary_3d_label_to_sparse_value(labels): assert label == 0 shape = [len(labels), len(labels[0]), len(labels[0][0])] return sparse_tensor.SparseTensorValue( - np.array(indices, np.int64), - np.array(values, np.int64), np.array(shape, np.int64)) + np.array(indices, np.int64), np.array(values, np.int64), + np.array(shape, np.int64)) def _binary_3d_label_to_sparse(labels): @@ -669,20 +668,18 @@ class StreamingTruePositivesTest(test.TestCase): for expand_predictions in [True, False]: for expand_labels in [True, False]: for dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): - predictions = math_ops.cast(constant_op.constant( - ((1, 0, 1, 0), - (0, 1, 1, 1), - (0, 0, 0, 0))), dtype=dtype) + predictions = math_ops.cast( + constant_op.constant(((1, 0, 1, 0), (0, 1, 1, 1), (0, 0, 0, 0))), + dtype=dtype) if expand_predictions: predictions = array_ops.expand_dims(predictions, 2) - labels = math_ops.cast(constant_op.constant( - ((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))), dtype=dtype) + labels = math_ops.cast( + constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))), + dtype=dtype) if expand_labels: labels = array_ops.expand_dims(labels, 2) - tp, tp_update_op = metrics.streaming_true_positives(predictions, - labels) + tp, tp_update_op = metrics.streaming_true_positives( + predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -692,14 +689,12 @@ class StreamingTruePositivesTest(test.TestCase): def testWeighted(self): for dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): - predictions = math_ops.cast(constant_op.constant( - ((1, 0, 1, 0), - (0, 1, 1, 1), - (0, 0, 0, 0))), dtype=dtype) - labels = math_ops.cast(constant_op.constant( - ((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))), dtype=dtype) + predictions = math_ops.cast( + constant_op.constant(((1, 0, 1, 0), (0, 1, 1, 1), (0, 0, 0, 0))), + dtype=dtype) + labels = math_ops.cast( + constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))), + dtype=dtype) tp, tp_update_op = metrics.streaming_true_positives( predictions, labels, weights=37.0) @@ -717,28 +712,25 @@ class StreamingFalseNegativesTest(test.TestCase): ops.reset_default_graph() def testVars(self): - metrics.streaming_false_negatives((0, 1, 0), - (0, 1, 1)) + metrics.streaming_false_negatives((0, 1, 0), (0, 1, 1)) _assert_metric_variables(self, ('false_negatives/count:0',)) def testUnweighted(self): for expand_predictions in [True, False]: for expand_labels in [True, False]: for dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): - predictions = math_ops.cast(constant_op.constant( - ((1, 0, 1, 0), - (0, 1, 1, 1), - (0, 0, 0, 0))), dtype=dtype) + predictions = math_ops.cast( + constant_op.constant(((1, 0, 1, 0), (0, 1, 1, 1), (0, 0, 0, 0))), + dtype=dtype) if expand_predictions: predictions = array_ops.expand_dims(predictions, 2) - labels = math_ops.cast(constant_op.constant( - ((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))), dtype=dtype) + labels = math_ops.cast( + constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))), + dtype=dtype) if expand_labels: labels = array_ops.expand_dims(labels, 2) - fn, fn_update_op = metrics.streaming_false_negatives(predictions, - labels) + fn, fn_update_op = metrics.streaming_false_negatives( + predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -748,14 +740,12 @@ class StreamingFalseNegativesTest(test.TestCase): def testWeighted(self): for dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): - predictions = math_ops.cast(constant_op.constant( - ((1, 0, 1, 0), - (0, 1, 1, 1), - (0, 0, 0, 0))), dtype=dtype) - labels = math_ops.cast(constant_op.constant( - ((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))), dtype=dtype) + predictions = math_ops.cast( + constant_op.constant(((1, 0, 1, 0), (0, 1, 1, 1), (0, 0, 0, 0))), + dtype=dtype) + labels = math_ops.cast( + constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))), + dtype=dtype) fn, fn_update_op = metrics.streaming_false_negatives( predictions, labels, weights=((3.0,), (5.0,), (7.0,))) @@ -773,28 +763,25 @@ class StreamingFalsePositivesTest(test.TestCase): ops.reset_default_graph() def testVars(self): - metrics.streaming_false_positives((0, 1, 0), - (0, 1, 1)) + metrics.streaming_false_positives((0, 1, 0), (0, 1, 1)) _assert_metric_variables(self, ('false_positives/count:0',)) def testUnweighted(self): for expand_predictions in [True, False]: for expand_labels in [True, False]: for dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): - predictions = math_ops.cast(constant_op.constant( - ((1, 0, 1, 0), - (0, 1, 1, 1), - (0, 0, 0, 0))), dtype=dtype) + predictions = math_ops.cast( + constant_op.constant(((1, 0, 1, 0), (0, 1, 1, 1), (0, 0, 0, 0))), + dtype=dtype) if expand_predictions: predictions = array_ops.expand_dims(predictions, 2) - labels = math_ops.cast(constant_op.constant( - ((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))), dtype=dtype) + labels = math_ops.cast( + constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))), + dtype=dtype) if expand_labels: labels = array_ops.expand_dims(labels, 2) - fp, fp_update_op = metrics.streaming_false_positives(predictions, - labels) + fp, fp_update_op = metrics.streaming_false_positives( + predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -804,20 +791,17 @@ class StreamingFalsePositivesTest(test.TestCase): def testWeighted(self): for dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): - predictions = math_ops.cast(constant_op.constant( - ((1, 0, 1, 0), - (0, 1, 1, 1), - (0, 0, 0, 0))), dtype=dtype) - labels = math_ops.cast(constant_op.constant( - ((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))), dtype=dtype) + predictions = math_ops.cast( + constant_op.constant(((1, 0, 1, 0), (0, 1, 1, 1), (0, 0, 0, 0))), + dtype=dtype) + labels = math_ops.cast( + constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))), + dtype=dtype) fp, fp_update_op = metrics.streaming_false_positives( predictions, labels, - weights=((1.0, 2.0, 3.0, 5.0), - (7.0, 11.0, 13.0, 17.0), - (19.0, 23.0, 29.0, 31.0))) + weights=((1.0, 2.0, 3.0, 5.0), (7.0, 11.0, 13.0, 17.0), (19.0, 23.0, + 29.0, 31.0))) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -833,28 +817,25 @@ class StreamingTrueNegativesTest(test.TestCase): ops.reset_default_graph() def testVars(self): - metrics.streaming_true_negatives((0, 1, 0), - (0, 1, 1)) + metrics.streaming_true_negatives((0, 1, 0), (0, 1, 1)) _assert_metric_variables(self, ('true_negatives/count:0',)) def testUnweighted(self): for expand_predictions in [True, False]: for expand_labels in [True, False]: for dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): - predictions = math_ops.cast(constant_op.constant( - ((1, 0, 1, 0), - (0, 1, 1, 1), - (0, 0, 0, 0))), dtype=dtype) + predictions = math_ops.cast( + constant_op.constant(((1, 0, 1, 0), (0, 1, 1, 1), (0, 0, 0, 0))), + dtype=dtype) if expand_predictions: predictions = array_ops.expand_dims(predictions, 2) - labels = math_ops.cast(constant_op.constant( - ((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))), dtype=dtype) + labels = math_ops.cast( + constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))), + dtype=dtype) if expand_labels: labels = array_ops.expand_dims(labels, 2) - tn, tn_update_op = metrics.streaming_true_negatives(predictions, - labels) + tn, tn_update_op = metrics.streaming_true_negatives( + predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -864,14 +845,12 @@ class StreamingTrueNegativesTest(test.TestCase): def testWeighted(self): for dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): - predictions = math_ops.cast(constant_op.constant( - ((1, 0, 1, 0), - (0, 1, 1, 1), - (0, 0, 0, 0))), dtype=dtype) - labels = math_ops.cast(constant_op.constant( - ((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))), dtype=dtype) + predictions = math_ops.cast( + constant_op.constant(((1, 0, 1, 0), (0, 1, 1, 1), (0, 0, 0, 0))), + dtype=dtype) + labels = math_ops.cast( + constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))), + dtype=dtype) tn, tn_update_op = metrics.streaming_true_negatives( predictions, labels, weights=((0.0, 2.0, 3.0, 5.0),)) @@ -894,12 +873,9 @@ class StreamingTruePositivesAtThresholdsTest(test.TestCase): _assert_metric_variables(self, ('true_positives:0',)) def testUnweighted(self): - predictions = constant_op.constant(((0.9, 0.2, 0.8, 0.1), - (0.2, 0.9, 0.7, 0.6), - (0.1, 0.2, 0.4, 0.3))) - labels = constant_op.constant(((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))) + predictions = constant_op.constant( + ((0.9, 0.2, 0.8, 0.1), (0.2, 0.9, 0.7, 0.6), (0.1, 0.2, 0.4, 0.3))) + labels = constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))) tp, tp_update_op = metrics.streaming_true_positives_at_thresholds( predictions, labels, thresholds=(0.15, 0.5, 0.85)) @@ -910,12 +886,9 @@ class StreamingTruePositivesAtThresholdsTest(test.TestCase): self.assertAllEqual((3, 1, 0), tp.eval()) def testWeighted(self): - predictions = constant_op.constant(((0.9, 0.2, 0.8, 0.1), - (0.2, 0.9, 0.7, 0.6), - (0.1, 0.2, 0.4, 0.3))) - labels = constant_op.constant(((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))) + predictions = constant_op.constant( + ((0.9, 0.2, 0.8, 0.1), (0.2, 0.9, 0.7, 0.6), (0.1, 0.2, 0.4, 0.3))) + labels = constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))) tp, tp_update_op = metrics.streaming_true_positives_at_thresholds( predictions, labels, weights=37.0, thresholds=(0.15, 0.5, 0.85)) @@ -937,16 +910,14 @@ class StreamingFalseNegativesAtThresholdsTest(test.TestCase): (0.0, 1.0, 0.0), (0, 1, 1), thresholds=( 0.15, 0.5, - 0.85,)) + 0.85, + )) _assert_metric_variables(self, ('false_negatives:0',)) def testUnweighted(self): - predictions = constant_op.constant(((0.9, 0.2, 0.8, 0.1), - (0.2, 0.9, 0.7, 0.6), - (0.1, 0.2, 0.4, 0.3))) - labels = constant_op.constant(((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))) + predictions = constant_op.constant( + ((0.9, 0.2, 0.8, 0.1), (0.2, 0.9, 0.7, 0.6), (0.1, 0.2, 0.4, 0.3))) + labels = constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))) fn, fn_update_op = metrics.streaming_false_negatives_at_thresholds( predictions, labels, thresholds=(0.15, 0.5, 0.85)) @@ -957,12 +928,9 @@ class StreamingFalseNegativesAtThresholdsTest(test.TestCase): self.assertAllEqual((0, 2, 3), fn.eval()) def testWeighted(self): - predictions = constant_op.constant(((0.9, 0.2, 0.8, 0.1), - (0.2, 0.9, 0.7, 0.6), - (0.1, 0.2, 0.4, 0.3))) - labels = constant_op.constant(((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))) + predictions = constant_op.constant( + ((0.9, 0.2, 0.8, 0.1), (0.2, 0.9, 0.7, 0.6), (0.1, 0.2, 0.4, 0.3))) + labels = constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))) fn, fn_update_op = metrics.streaming_false_negatives_at_thresholds( predictions, labels, @@ -988,12 +956,9 @@ class StreamingFalsePositivesAtThresholdsTest(test.TestCase): _assert_metric_variables(self, ('false_positives:0',)) def testUnweighted(self): - predictions = constant_op.constant(((0.9, 0.2, 0.8, 0.1), - (0.2, 0.9, 0.7, 0.6), - (0.1, 0.2, 0.4, 0.3))) - labels = constant_op.constant(((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))) + predictions = constant_op.constant( + ((0.9, 0.2, 0.8, 0.1), (0.2, 0.9, 0.7, 0.6), (0.1, 0.2, 0.4, 0.3))) + labels = constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))) fp, fp_update_op = metrics.streaming_false_positives_at_thresholds( predictions, labels, thresholds=(0.15, 0.5, 0.85)) @@ -1004,18 +969,14 @@ class StreamingFalsePositivesAtThresholdsTest(test.TestCase): self.assertAllEqual((7, 4, 2), fp.eval()) def testWeighted(self): - predictions = constant_op.constant(((0.9, 0.2, 0.8, 0.1), - (0.2, 0.9, 0.7, 0.6), - (0.1, 0.2, 0.4, 0.3))) - labels = constant_op.constant(((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))) + predictions = constant_op.constant( + ((0.9, 0.2, 0.8, 0.1), (0.2, 0.9, 0.7, 0.6), (0.1, 0.2, 0.4, 0.3))) + labels = constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))) fp, fp_update_op = metrics.streaming_false_positives_at_thresholds( predictions, labels, - weights=((1.0, 2.0, 3.0, 5.0), - (7.0, 11.0, 13.0, 17.0), - (19.0, 23.0, 29.0, 31.0)), + weights=((1.0, 2.0, 3.0, 5.0), (7.0, 11.0, 13.0, 17.0), (19.0, 23.0, + 29.0, 31.0)), thresholds=(0.15, 0.5, 0.85)) with self.test_session() as sess: @@ -1037,12 +998,9 @@ class StreamingTrueNegativesAtThresholdsTest(test.TestCase): _assert_metric_variables(self, ('true_negatives:0',)) def testUnweighted(self): - predictions = constant_op.constant(((0.9, 0.2, 0.8, 0.1), - (0.2, 0.9, 0.7, 0.6), - (0.1, 0.2, 0.4, 0.3))) - labels = constant_op.constant(((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))) + predictions = constant_op.constant( + ((0.9, 0.2, 0.8, 0.1), (0.2, 0.9, 0.7, 0.6), (0.1, 0.2, 0.4, 0.3))) + labels = constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))) tn, tn_update_op = metrics.streaming_true_negatives_at_thresholds( predictions, labels, thresholds=(0.15, 0.5, 0.85)) @@ -1053,12 +1011,9 @@ class StreamingTrueNegativesAtThresholdsTest(test.TestCase): self.assertAllEqual((2, 5, 7), tn.eval()) def testWeighted(self): - predictions = constant_op.constant(((0.9, 0.2, 0.8, 0.1), - (0.2, 0.9, 0.7, 0.6), - (0.1, 0.2, 0.4, 0.3))) - labels = constant_op.constant(((0, 1, 1, 0), - (1, 0, 0, 0), - (0, 0, 0, 0))) + predictions = constant_op.constant( + ((0.9, 0.2, 0.8, 0.1), (0.2, 0.9, 0.7, 0.6), (0.1, 0.2, 0.4, 0.3))) + labels = constant_op.constant(((0, 1, 1, 0), (1, 0, 0, 0), (0, 0, 0, 0))) tn, tn_update_op = metrics.streaming_true_negatives_at_thresholds( predictions, labels, @@ -1393,8 +1348,7 @@ class StreamingFPRTest(test.TestCase): (10, 3), maxval=1, dtype=dtypes_lib.int64, seed=1) labels = random_ops.random_uniform( (10, 3), maxval=2, dtype=dtypes_lib.int64, seed=2) - fpr, update_op = metrics.streaming_false_positive_rate( - predictions, labels) + fpr, update_op = metrics.streaming_false_positive_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1413,8 +1367,7 @@ class StreamingFPRTest(test.TestCase): predictions = constant_op.constant(np_inputs) labels = constant_op.constant(np_inputs) - fpr, update_op = metrics.streaming_false_positive_rate( - predictions, labels) + fpr, update_op = metrics.streaming_false_positive_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1424,8 +1377,7 @@ class StreamingFPRTest(test.TestCase): def testSomeCorrect(self): predictions = constant_op.constant([1, 0, 1, 0], shape=(1, 4)) labels = constant_op.constant([0, 1, 1, 0], shape=(1, 4)) - fpr, update_op = metrics.streaming_false_positive_rate( - predictions, labels) + fpr, update_op = metrics.streaming_false_positive_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1467,8 +1419,7 @@ class StreamingFPRTest(test.TestCase): predictions = constant_op.constant(np_inputs) labels = constant_op.constant(1 - np_inputs) - fpr, update_op = metrics.streaming_false_positive_rate( - predictions, labels) + fpr, update_op = metrics.streaming_false_positive_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1478,8 +1429,7 @@ class StreamingFPRTest(test.TestCase): def testZeroFalsePositivesAndTrueNegativesGivesZeroFPR(self): predictions = array_ops.ones((1, 4)) labels = array_ops.ones((1, 4)) - fpr, update_op = metrics.streaming_false_positive_rate( - predictions, labels) + fpr, update_op = metrics.streaming_false_positive_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1521,8 +1471,7 @@ class StreamingFNRTest(test.TestCase): (10, 3), maxval=1, dtype=dtypes_lib.int64, seed=1) labels = random_ops.random_uniform( (10, 3), maxval=2, dtype=dtypes_lib.int64, seed=2) - fnr, update_op = metrics.streaming_false_negative_rate( - predictions, labels) + fnr, update_op = metrics.streaming_false_negative_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1541,8 +1490,7 @@ class StreamingFNRTest(test.TestCase): predictions = constant_op.constant(np_inputs) labels = constant_op.constant(np_inputs) - fnr, update_op = metrics.streaming_false_negative_rate( - predictions, labels) + fnr, update_op = metrics.streaming_false_negative_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1552,8 +1500,7 @@ class StreamingFNRTest(test.TestCase): def testSomeCorrect(self): predictions = constant_op.constant([1, 0, 1, 0], shape=(1, 4)) labels = constant_op.constant([0, 1, 1, 0], shape=(1, 4)) - fnr, update_op = metrics.streaming_false_negative_rate( - predictions, labels) + fnr, update_op = metrics.streaming_false_negative_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1595,8 +1542,7 @@ class StreamingFNRTest(test.TestCase): predictions = constant_op.constant(np_inputs) labels = constant_op.constant(1 - np_inputs) - fnr, update_op = metrics.streaming_false_negative_rate( - predictions, labels) + fnr, update_op = metrics.streaming_false_negative_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1606,8 +1552,7 @@ class StreamingFNRTest(test.TestCase): def testZeroFalseNegativesAndTruePositivesGivesZeroFNR(self): predictions = array_ops.zeros((1, 4)) labels = array_ops.zeros((1, 4)) - fnr, update_op = metrics.streaming_false_negative_rate( - predictions, labels) + fnr, update_op = metrics.streaming_false_negative_rate(predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -1857,9 +1802,9 @@ class StreamingAUCTest(test.TestCase): auc, update_op = metrics.streaming_auc(predictions, labels, curve='PR') sess.run(variables.local_variables_initializer()) - self.assertAlmostEqual(0.79166, sess.run(update_op), delta=1e-3) + self.assertAlmostEqual(0.54166603, sess.run(update_op), delta=1e-3) - self.assertAlmostEqual(0.79166, auc.eval(), delta=1e-3) + self.assertAlmostEqual(0.54166603, auc.eval(), delta=1e-3) def testAnotherAUCPRSpecialCase(self): with self.test_session() as sess: @@ -1871,9 +1816,9 @@ class StreamingAUCTest(test.TestCase): auc, update_op = metrics.streaming_auc(predictions, labels, curve='PR') sess.run(variables.local_variables_initializer()) - self.assertAlmostEqual(0.610317, sess.run(update_op), delta=1e-3) + self.assertAlmostEqual(0.44365042, sess.run(update_op), delta=1e-3) - self.assertAlmostEqual(0.610317, auc.eval(), delta=1e-3) + self.assertAlmostEqual(0.44365042, auc.eval(), delta=1e-3) def testThirdAUCPRSpecialCase(self): with self.test_session() as sess: @@ -1885,9 +1830,9 @@ class StreamingAUCTest(test.TestCase): auc, update_op = metrics.streaming_auc(predictions, labels, curve='PR') sess.run(variables.local_variables_initializer()) - self.assertAlmostEqual(0.90277, sess.run(update_op), delta=1e-3) + self.assertAlmostEqual(0.73611039, sess.run(update_op), delta=1e-3) - self.assertAlmostEqual(0.90277, auc.eval(), delta=1e-3) + self.assertAlmostEqual(0.73611039, auc.eval(), delta=1e-3) def testAllIncorrect(self): inputs = np.random.randint(0, 2, size=(100, 1)) @@ -1920,9 +1865,9 @@ class StreamingAUCTest(test.TestCase): auc, update_op = metrics.streaming_auc(predictions, labels, curve='PR') sess.run(variables.local_variables_initializer()) - self.assertAlmostEqual(1, sess.run(update_op), 6) + self.assertAlmostEqual(0.49999976, sess.run(update_op), 6) - self.assertAlmostEqual(1, auc.eval(), 6) + self.assertAlmostEqual(0.49999976, auc.eval(), 6) def testWithMultipleUpdates(self): num_samples = 1000 @@ -1944,16 +1889,17 @@ class StreamingAUCTest(test.TestCase): enqueue_ops[i].append(x_queue.enqueue(x_batches[i, :])) return x_queue.dequeue() - for weights in (None, np.ones(num_samples), np.random.exponential( - scale=1.0, size=num_samples)): + for weights in (None, np.ones(num_samples), + np.random.exponential(scale=1.0, size=num_samples)): expected_auc = _np_auc(predictions, labels, weights) with self.test_session() as sess: enqueue_ops = [[] for i in range(num_batches)] tf_predictions = _enqueue_as_batches(predictions, enqueue_ops) tf_labels = _enqueue_as_batches(labels, enqueue_ops) - tf_weights = (_enqueue_as_batches(weights, enqueue_ops) if - weights is not None else None) + tf_weights = ( + _enqueue_as_batches(weights, enqueue_ops) + if weights is not None else None) for i in range(num_batches): sess.run(enqueue_ops[i]) @@ -1985,17 +1931,18 @@ class StreamingDynamicAUCTest(test.TestCase): def testUnknownCurve(self): with self.assertRaisesRegexp( ValueError, 'curve must be either ROC or PR, TEST_CURVE unknown'): - metrics.streaming_dynamic_auc(labels=array_ops.ones((10, 1)), - predictions=array_ops.ones((10, 1)), - curve='TEST_CURVE') + metrics.streaming_dynamic_auc( + labels=array_ops.ones((10, 1)), + predictions=array_ops.ones((10, 1)), + curve='TEST_CURVE') def testVars(self): metrics.streaming_dynamic_auc( labels=array_ops.ones((10, 1)), predictions=array_ops.ones((10, 1))) - _assert_metric_variables(self, ['dynamic_auc/concat_labels/array:0', - 'dynamic_auc/concat_labels/size:0', - 'dynamic_auc/concat_preds/array:0', - 'dynamic_auc/concat_preds/size:0']) + _assert_metric_variables(self, [ + 'dynamic_auc/concat_labels/array:0', 'dynamic_auc/concat_labels/size:0', + 'dynamic_auc/concat_preds/array:0', 'dynamic_auc/concat_preds/size:0' + ]) def testMetricsCollection(self): my_collection_name = '__metrics__' @@ -2049,8 +1996,8 @@ class StreamingDynamicAUCTest(test.TestCase): def testNonZeroOnePredictions(self): with self.test_session() as sess: - predictions = constant_op.constant([2.5, -2.5, 2.5, -2.5], - dtype=dtypes_lib.float32) + predictions = constant_op.constant( + [2.5, -2.5, 2.5, -2.5], dtype=dtypes_lib.float32) labels = constant_op.constant([1, 0, 1, 0]) auc, update_op = metrics.streaming_dynamic_auc(labels, predictions) sess.run(variables.local_variables_initializer()) @@ -2122,9 +2069,10 @@ class StreamingDynamicAUCTest(test.TestCase): num_batches = 100 labels = np.array([]) predictions = np.array([]) - tf_labels = variables.Variable(array_ops.ones(batch_size, dtypes_lib.int32), - collections=[ops.GraphKeys.LOCAL_VARIABLES], - dtype=dtypes_lib.int32) + tf_labels = variables.Variable( + array_ops.ones(batch_size, dtypes_lib.int32), + collections=[ops.GraphKeys.LOCAL_VARIABLES], + dtype=dtypes_lib.int32) tf_predictions = variables.Variable( array_ops.ones(batch_size), collections=[ops.GraphKeys.LOCAL_VARIABLES], @@ -2180,7 +2128,7 @@ class StreamingDynamicAUCTest(test.TestCase): self.assertAlmostEqual(0.90277, auc.eval(), delta=1e-5) -class StreamingPrecisionRecallAtEqualThresholdsTest(test.TestCase): +class AucWithConfidenceIntervalsTest(test.TestCase): def setUp(self): np.random.seed(1) @@ -2192,7 +2140,7 @@ class StreamingPrecisionRecallAtEqualThresholdsTest(test.TestCase): Args: expected_dict: A dictionary with keys that are the names of properties of PrecisionRecallData and whose values are lists of floats. - gotten_result: A PrecisionRecallData object. + gotten_result: A AucWithConfidenceIntervalData object. """ gotten_dict = {k: t.eval() for k, t in gotten_result._asdict().items()} self.assertItemsEqual( @@ -2201,6 +2149,204 @@ class StreamingPrecisionRecallAtEqualThresholdsTest(test.TestCase): for key, expected_values in expected_dict.items(): self.assertAllClose(expected_values, gotten_dict[key]) + def _testCase(self, predictions, labels, expected_result, weights=None): + """Performs a test given a certain scenario of labels, predictions, weights. + + Args: + predictions: The predictions tensor. Of type float32. + labels: The labels tensor. Of type bool. + expected_result: The expected result (dict) that maps to tensors. + weights: Optional weights tensor. + """ + with self.test_session() as sess: + predictions_tensor = constant_op.constant( + predictions, dtype=dtypes_lib.float32) + labels_tensor = constant_op.constant(labels, dtype=dtypes_lib.int64) + weights_tensor = None + if weights: + weights_tensor = constant_op.constant(weights, dtype=dtypes_lib.float32) + gotten_result, update_op = ( + metric_ops.auc_with_confidence_intervals( + labels=labels_tensor, + predictions=predictions_tensor, + weights=weights_tensor)) + + sess.run(variables.local_variables_initializer()) + sess.run(update_op) + + self._testResultsEqual(expected_result, gotten_result) + + def testAucAllCorrect(self): + self._testCase( + predictions=[0., 0.2, 0.3, 0.3, 0.4, 0.5, 0.6, 0.6, 0.8, 1.0], + labels=[0, 0, 1, 0, 0, 1, 0, 1, 1, 0], + expected_result={ + 'auc': 0.66666667, + 'lower': 0.27826795, + 'upper': 0.91208512, + }) + + def testAucUnorderedInput(self): + self._testCase( + predictions=[1.0, 0.6, 0., 0.3, 0.4, 0.2, 0.5, 0.3, 0.6, 0.8], + labels=[0, 1, 0, 1, 0, 0, 1, 0, 0, 1], + expected_result={ + 'auc': 0.66666667, + 'lower': 0.27826795, + 'upper': 0.91208512, + }) + + def testAucWithWeights(self): + self._testCase( + predictions=[0., 0.2, 0.3, 0.3, 0.4, 0.5, 0.6, 0.6, 0.8, 1.0], + labels=[0, 0, 1, 0, 0, 1, 0, 1, 1, 0], + weights=[0.5, 0.6, 1.2, 1.5, 2.0, 2.0, 1.5, 1.2, 0.6, 0.5], + expected_result={ + 'auc': 0.65151515, + 'lower': 0.28918604, + 'upper': 0.89573906, + }) + + def testAucEqualOne(self): + self._testCase( + predictions=[0, 0.2, 0.3, 0.3, 0.4, 0.5, 0.6, 0.6, 0.8, 1.0], + labels=[0, 0, 0, 0, 0, 1, 1, 1, 1, 1], + expected_result={ + 'auc': 1.0, + 'lower': 1.0, + 'upper': 1.0, + }) + + def testAucEqualZero(self): + self._testCase( + predictions=[0, 0.2, 0.3, 0.3, 0.4, 0.5, 0.6, 0.6, 0.8, 1.0], + labels=[1, 1, 1, 1, 1, 0, 0, 0, 0, 0], + expected_result={ + 'auc': 0.0, + 'lower': 0.0, + 'upper': 0.0, + }) + + def testNonZeroOnePredictions(self): + self._testCase( + predictions=[2.5, -2.5, .5, -.5, 1], + labels=[1, 0, 1, 0, 0], + expected_result={ + 'auc': 0.83333333, + 'lower': 0.15229267, + 'upper': 0.99286517, + }) + + def testAllLabelsOnes(self): + self._testCase( + predictions=[1., 1., 1., 1., 1.], + labels=[1, 1, 1, 1, 1], + expected_result={ + 'auc': 0., + 'lower': 0., + 'upper': 0., + }) + + def testAllLabelsZeros(self): + self._testCase( + predictions=[0., 0., 0., 0., 0.], + labels=[0, 0, 0, 0, 0], + expected_result={ + 'auc': 0., + 'lower': 0., + 'upper': 0., + }) + + def testWeightSumLessThanOneAll(self): + self._testCase( + predictions=[1., 1., 0., 1., 0., 0.], + labels=[1, 1, 1, 0, 0, 0], + weights=[0.1, 0.1, 0.1, 0.1, 0.1, 0.1], + expected_result={ + 'auc': 0., + 'lower': 0., + 'upper': 0., + }) + + def testWithMultipleUpdates(self): + batch_size = 50 + num_batches = 100 + labels = np.array([]) + predictions = np.array([]) + tf_labels = variables.Variable(array_ops.ones(batch_size, dtypes_lib.int32), + collections=[ops.GraphKeys.LOCAL_VARIABLES], + dtype=dtypes_lib.int32) + tf_predictions = variables.Variable( + array_ops.ones(batch_size), + collections=[ops.GraphKeys.LOCAL_VARIABLES], + dtype=dtypes_lib.float32) + auc, update_op = metrics.auc_with_confidence_intervals(tf_labels, + tf_predictions) + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + for _ in xrange(num_batches): + new_labels = np.random.randint(0, 2, size=batch_size) + noise = np.random.normal(0.0, scale=0.2, size=batch_size) + new_predictions = 0.4 + 0.2 * new_labels + noise + labels = np.concatenate([labels, new_labels]) + predictions = np.concatenate([predictions, new_predictions]) + sess.run(tf_labels.assign(new_labels)) + sess.run(tf_predictions.assign(new_predictions)) + sess.run(update_op) + expected_auc = _np_auc(predictions, labels) + self.assertAllClose(expected_auc, auc.auc.eval()) + + def testExceptionOnFloatLabels(self): + with self.test_session() as sess: + predictions = constant_op.constant([1, 0.5, 0, 1, 0], dtypes_lib.float32) + labels = constant_op.constant([0.7, 0, 1, 0, 1]) + _, update_op = metrics.auc_with_confidence_intervals(labels, predictions) + sess.run(variables.local_variables_initializer()) + self.assertRaises(TypeError, sess.run(update_op)) + + def testExceptionOnGreaterThanOneLabel(self): + with self.test_session() as sess: + predictions = constant_op.constant([1, 0.5, 0, 1, 0], dtypes_lib.float32) + labels = constant_op.constant([2, 1, 0, 1, 0]) + _, update_op = metrics.auc_with_confidence_intervals(labels, predictions) + sess.run(variables.local_variables_initializer()) + with self.assertRaisesRegexp( + errors_impl.InvalidArgumentError, + '.*labels must be 0 or 1, at least one is >1.*'): + sess.run(update_op) + + def testExceptionOnNegativeLabel(self): + with self.test_session() as sess: + predictions = constant_op.constant([1, 0.5, 0, 1, 0], dtypes_lib.float32) + labels = constant_op.constant([1, 0, -1, 1, 0]) + _, update_op = metrics.auc_with_confidence_intervals(labels, predictions) + sess.run(variables.local_variables_initializer()) + with self.assertRaisesRegexp( + errors_impl.InvalidArgumentError, + '.*labels must be 0 or 1, at least one is <0.*'): + sess.run(update_op) + + +class StreamingPrecisionRecallAtEqualThresholdsTest(test.TestCase): + + def setUp(self): + np.random.seed(1) + ops.reset_default_graph() + + def _testResultsEqual(self, expected_dict, gotten_result): + """Tests that 2 results (dicts) represent the same data. + + Args: + expected_dict: A dictionary with keys that are the names of properties + of PrecisionRecallData and whose values are lists of floats. + gotten_result: A PrecisionRecallData object. + """ + gotten_dict = {k: t.eval() for k, t in gotten_result._asdict().items()} + self.assertItemsEqual(list(expected_dict.keys()), list(gotten_dict.keys())) + + for key, expected_values in expected_dict.items(): + self.assertAllClose(expected_values, gotten_dict[key]) + def _testCase(self, predictions, labels, expected_result, weights=None): """Performs a test given a certain scenario of labels, predictions, weights. @@ -2261,60 +2407,65 @@ class StreamingPrecisionRecallAtEqualThresholdsTest(test.TestCase): sess.run(update_op) # Then verify idempotency. - initial_result = {k: value.eval().tolist() for k, value in - result._asdict().items()} + initial_result = { + k: value.eval().tolist() + for k, value in result._asdict().items() + } for _ in range(3): self._testResultsEqual(initial_result, result) def testAllTruePositives(self): - self._testCase([[1]], [[True]], { - 'tp': [1, 1, 1], - 'fp': [0, 0, 0], - 'tn': [0, 0, 0], - 'fn': [0, 0, 0], - 'precision': [1.0, 1.0, 1.0], - 'recall': [1.0, 1.0, 1.0], - 'thresholds': [0.0, 0.5, 1.0], - }) + self._testCase( + [[1]], [[True]], { + 'tp': [1, 1, 1], + 'fp': [0, 0, 0], + 'tn': [0, 0, 0], + 'fn': [0, 0, 0], + 'precision': [1.0, 1.0, 1.0], + 'recall': [1.0, 1.0, 1.0], + 'thresholds': [0.0, 0.5, 1.0], + }) def testAllTrueNegatives(self): - self._testCase([[0]], [[False]], { - 'tp': [0, 0, 0], - 'fp': [1, 0, 0], - 'tn': [0, 1, 1], - 'fn': [0, 0, 0], - 'precision': [0.0, 0.0, 0.0], - 'recall': [0.0, 0.0, 0.0], - 'thresholds': [0.0, 0.5, 1.0], - }) + self._testCase( + [[0]], [[False]], { + 'tp': [0, 0, 0], + 'fp': [1, 0, 0], + 'tn': [0, 1, 1], + 'fn': [0, 0, 0], + 'precision': [0.0, 0.0, 0.0], + 'recall': [0.0, 0.0, 0.0], + 'thresholds': [0.0, 0.5, 1.0], + }) def testAllFalsePositives(self): - self._testCase([[1]], [[False]], { - 'tp': [0, 0, 0], - 'fp': [1, 1, 1], - 'tn': [0, 0, 0], - 'fn': [0, 0, 0], - 'precision': [0.0, 0.0, 0.0], - 'recall': [0.0, 0.0, 0.0], - 'thresholds': [0.0, 0.5, 1.0], - }) + self._testCase( + [[1]], [[False]], { + 'tp': [0, 0, 0], + 'fp': [1, 1, 1], + 'tn': [0, 0, 0], + 'fn': [0, 0, 0], + 'precision': [0.0, 0.0, 0.0], + 'recall': [0.0, 0.0, 0.0], + 'thresholds': [0.0, 0.5, 1.0], + }) def testAllFalseNegatives(self): - self._testCase([[0]], [[True]], { - 'tp': [1, 0, 0], - 'fp': [0, 0, 0], - 'tn': [0, 0, 0], - 'fn': [0, 1, 1], - 'precision': [1.0, 0.0, 0.0], - 'recall': [1.0, 0.0, 0.0], - 'thresholds': [0.0, 0.5, 1.0], - }) + self._testCase( + [[0]], [[True]], { + 'tp': [1, 0, 0], + 'fp': [0, 0, 0], + 'tn': [0, 0, 0], + 'fn': [0, 1, 1], + 'precision': [1.0, 0.0, 0.0], + 'recall': [1.0, 0.0, 0.0], + 'thresholds': [0.0, 0.5, 1.0], + }) def testManyValues(self): self._testCase( [[0.2, 0.3, 0.4, 0.6, 0.7, 0.8]], - [[True, False, False, True, True, True]], - { + [[True, False, False, True, True, True]], { 'tp': [4, 3, 0], 'fp': [2, 0, 0], 'tn': [0, 2, 2], @@ -2327,8 +2478,7 @@ class StreamingPrecisionRecallAtEqualThresholdsTest(test.TestCase): def testManyValuesWithWeights(self): self._testCase( [[0.2, 0.3, 0.4, 0.6, 0.7, 0.8]], - [[True, False, False, True, True, True]], - { + [[True, False, False, True, True, True]], { 'tp': [1.5, 1.5, 0.0], 'fp': [2.5, 0.0, 0.0], 'tn': [0.0, 2.5, 2.5], @@ -2644,11 +2794,10 @@ class StreamingPrecisionRecallThresholdsTest(test.TestCase): labels = random_ops.random_uniform( (10, 3), maxval=2, dtype=dtypes_lib.int64, seed=2) thresholds = [0, 0.5, 1.0] - prec, prec_op = metrics.streaming_precision_at_thresholds(predictions, - labels, - thresholds) - rec, rec_op = metrics.streaming_recall_at_thresholds(predictions, labels, - thresholds) + prec, prec_op = metrics.streaming_precision_at_thresholds( + predictions, labels, thresholds) + rec, rec_op = metrics.streaming_recall_at_thresholds( + predictions, labels, thresholds) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -2672,11 +2821,10 @@ class StreamingPrecisionRecallThresholdsTest(test.TestCase): predictions = constant_op.constant(inputs, dtype=dtypes_lib.float32) labels = constant_op.constant(inputs) thresholds = [0.5] - prec, prec_op = metrics.streaming_precision_at_thresholds(predictions, - labels, - thresholds) - rec, rec_op = metrics.streaming_recall_at_thresholds(predictions, labels, - thresholds) + prec, prec_op = metrics.streaming_precision_at_thresholds( + predictions, labels, thresholds) + rec, rec_op = metrics.streaming_recall_at_thresholds( + predictions, labels, thresholds) sess.run(variables.local_variables_initializer()) sess.run([prec_op, rec_op]) @@ -2690,11 +2838,10 @@ class StreamingPrecisionRecallThresholdsTest(test.TestCase): [1, 0, 1, 0], shape=(1, 4), dtype=dtypes_lib.float32) labels = constant_op.constant([0, 1, 1, 0], shape=(1, 4)) thresholds = [0.5] - prec, prec_op = metrics.streaming_precision_at_thresholds(predictions, - labels, - thresholds) - rec, rec_op = metrics.streaming_recall_at_thresholds(predictions, labels, - thresholds) + prec, prec_op = metrics.streaming_precision_at_thresholds( + predictions, labels, thresholds) + rec, rec_op = metrics.streaming_recall_at_thresholds( + predictions, labels, thresholds) sess.run(variables.local_variables_initializer()) sess.run([prec_op, rec_op]) @@ -2709,11 +2856,10 @@ class StreamingPrecisionRecallThresholdsTest(test.TestCase): predictions = constant_op.constant(inputs, dtype=dtypes_lib.float32) labels = constant_op.constant(1 - inputs, dtype=dtypes_lib.float32) thresholds = [0.5] - prec, prec_op = metrics.streaming_precision_at_thresholds(predictions, - labels, - thresholds) - rec, rec_op = metrics.streaming_recall_at_thresholds(predictions, labels, - thresholds) + prec, prec_op = metrics.streaming_precision_at_thresholds( + predictions, labels, thresholds) + rec, rec_op = metrics.streaming_recall_at_thresholds( + predictions, labels, thresholds) sess.run(variables.local_variables_initializer()) sess.run([prec_op, rec_op]) @@ -2779,11 +2925,10 @@ class StreamingPrecisionRecallThresholdsTest(test.TestCase): [1, 0, 1, 0], shape=(1, 4), dtype=dtypes_lib.float32) labels = constant_op.constant([0, 1, 1, 1], shape=(1, 4)) thresholds = [-1.0, 2.0] # lower/higher than any values - prec, prec_op = metrics.streaming_precision_at_thresholds(predictions, - labels, - thresholds) - rec, rec_op = metrics.streaming_recall_at_thresholds(predictions, labels, - thresholds) + prec, prec_op = metrics.streaming_precision_at_thresholds( + predictions, labels, thresholds) + rec, rec_op = metrics.streaming_recall_at_thresholds( + predictions, labels, thresholds) prec_low = prec[0] prec_high = prec[1] @@ -2803,11 +2948,10 @@ class StreamingPrecisionRecallThresholdsTest(test.TestCase): predictions = array_ops.zeros([4], dtype=dtypes_lib.float32) labels = array_ops.zeros([4]) thresholds = [0.5] - prec, prec_op = metrics.streaming_precision_at_thresholds(predictions, - labels, - thresholds) - rec, rec_op = metrics.streaming_recall_at_thresholds(predictions, labels, - thresholds) + prec, prec_op = metrics.streaming_precision_at_thresholds( + predictions, labels, thresholds) + rec, rec_op = metrics.streaming_recall_at_thresholds( + predictions, labels, thresholds) sess.run(variables.local_variables_initializer()) sess.run([prec_op, rec_op]) @@ -2872,12 +3016,10 @@ class StreamingPrecisionRecallThresholdsTest(test.TestCase): tf_predictions = predictions_queue.dequeue() tf_labels = labels_queue.dequeue() - prec, prec_op = metrics.streaming_precision_at_thresholds(tf_predictions, - tf_labels, - thresholds) - rec, rec_op = metrics.streaming_recall_at_thresholds(tf_predictions, - tf_labels, - thresholds) + prec, prec_op = metrics.streaming_precision_at_thresholds( + tf_predictions, tf_labels, thresholds) + rec, rec_op = metrics.streaming_recall_at_thresholds( + tf_predictions, tf_labels, thresholds) sess.run(variables.local_variables_initializer()) for _ in range(int(num_samples / batch_size)): @@ -2921,8 +3063,7 @@ class StreamingFPRThresholdsTest(test.TestCase): labels=array_ops.ones((10, 1)), thresholds=[0, 0.5, 1.0], updates_collections=[my_collection_name]) - self.assertListEqual( - ops.get_collection(my_collection_name), [update_op]) + self.assertListEqual(ops.get_collection(my_collection_name), [update_op]) def testValueTensorIsIdempotent(self): predictions = random_ops.random_uniform( @@ -3271,8 +3412,7 @@ class StreamingFNRThresholdsTest(test.TestCase): labels=array_ops.ones((10, 1)), thresholds=[0, 0.5, 1.0], updates_collections=[my_collection_name]) - self.assertListEqual( - ops.get_collection(my_collection_name), [update_op]) + self.assertListEqual(ops.get_collection(my_collection_name), [update_op]) def testValueTensorIsIdempotent(self): predictions = random_ops.random_uniform( @@ -3492,8 +3632,7 @@ class StreamingRecallAtKTest(test.TestCase): def testVars(self): metrics.streaming_recall_at_k( predictions=array_ops.ones((self._batch_size, self._num_classes)), - labels=array_ops.ones( - (self._batch_size,), dtype=dtypes_lib.int32), + labels=array_ops.ones((self._batch_size,), dtype=dtypes_lib.int32), k=1) _assert_metric_variables(self, ('recall_at_1/count:0', 'recall_at_1/total:0')) @@ -3502,8 +3641,7 @@ class StreamingRecallAtKTest(test.TestCase): my_collection_name = '__metrics__' mean, _ = metrics.streaming_recall_at_k( predictions=array_ops.ones((self._batch_size, self._num_classes)), - labels=array_ops.ones( - (self._batch_size,), dtype=dtypes_lib.int32), + labels=array_ops.ones((self._batch_size,), dtype=dtypes_lib.int32), k=1, metrics_collections=[my_collection_name]) self.assertListEqual(ops.get_collection(my_collection_name), [mean]) @@ -3512,8 +3650,7 @@ class StreamingRecallAtKTest(test.TestCase): my_collection_name = '__updates__' _, update_op = metrics.streaming_recall_at_k( predictions=array_ops.ones((self._batch_size, self._num_classes)), - labels=array_ops.ones( - (self._batch_size,), dtype=dtypes_lib.int32), + labels=array_ops.ones((self._batch_size,), dtype=dtypes_lib.int32), k=1, updates_collections=[my_collection_name]) self.assertListEqual(ops.get_collection(my_collection_name), [update_op]) @@ -3715,9 +3852,17 @@ class StreamingSparsePrecisionTest(test.TestCase): # top_k_predictions has rank < 2. top_k_predictions = [9, 4, 6, 2, 0] sp_labels = sparse_tensor.SparseTensorValue( - indices=np.array([[0,], [1,], [2,]], np.int64), + indices=np.array([[ + 0, + ], [ + 1, + ], [ + 2, + ]], np.int64), values=np.array([2, 7, 8], np.int64), - dense_shape=np.array([10,], np.int64)) + dense_shape=np.array([ + 10, + ], np.int64)) with self.assertRaises(ValueError): precision, _ = metrics.streaming_sparse_precision_at_top_k( @@ -3774,8 +3919,9 @@ class StreamingSparsePrecisionTest(test.TestCase): # average of the 2 examples. labels = np.array([labels_ex1, labels_ex2], dtype=np.int64) predictions = (predictions_ex1, predictions_ex2) - streaming_precision = [(ex1 + ex2) / 2 - for ex1, ex2 in zip(precision_ex1, precision_ex2)] + streaming_precision = [ + (ex1 + ex2) / 2 for ex1, ex2 in zip(precision_ex1, precision_ex2) + ] streaming_average_precision = [ (ex1 + ex2) / 2 for ex1, ex2 in zip(avg_precision_ex1, avg_precision_ex2) @@ -3835,29 +3981,29 @@ class StreamingSparsePrecisionTest(test.TestCase): (predictions_top_k_ex1[:k],), labels, expected=avg_precision_ex1[i]) def test_average_precision_at_top_k_static_shape_check(self): - predictions_top_k = array_ops.placeholder(shape=(2, None), - dtype=dtypes_lib.int64) + predictions_top_k = array_ops.placeholder( + shape=(2, None), dtype=dtypes_lib.int64) labels = np.array(((1,), (2,)), dtype=np.int64) # Fails due to non-static predictions_idx shape. with self.assertRaises(ValueError): - metric_ops.streaming_sparse_average_precision_at_top_k(predictions_top_k, - labels) + metric_ops.streaming_sparse_average_precision_at_top_k( + predictions_top_k, labels) predictions_top_k = (2, 1) # Fails since rank of predictions_idx is less than one. with self.assertRaises(ValueError): - metric_ops.streaming_sparse_average_precision_at_top_k(predictions_top_k, - labels) + metric_ops.streaming_sparse_average_precision_at_top_k( + predictions_top_k, labels) predictions_top_k = ((2,), (1,)) # Valid static shape. - metric_ops.streaming_sparse_average_precision_at_top_k(predictions_top_k, - labels) + metric_ops.streaming_sparse_average_precision_at_top_k( + predictions_top_k, labels) def test_one_label_at_k1_nan(self): predictions = [[0.1, 0.3, 0.2, 0.4], [0.1, 0.2, 0.3, 0.4]] top_k_predictions = [[3], [3]] - sparse_labels = _binary_2d_label_to_sparse_value( - [[0, 0, 0, 1], [0, 0, 1, 0]]) + sparse_labels = _binary_2d_label_to_sparse_value([[0, 0, 0, 1], + [0, 0, 1, 0]]) dense_labels = np.array([[3], [2]], dtype=np.int64) for labels in (sparse_labels, dense_labels): @@ -3871,8 +4017,8 @@ class StreamingSparsePrecisionTest(test.TestCase): def test_one_label_at_k1(self): predictions = [[0.1, 0.3, 0.2, 0.4], [0.1, 0.2, 0.3, 0.4]] top_k_predictions = [[3], [3]] - sparse_labels = _binary_2d_label_to_sparse_value( - [[0, 0, 0, 1], [0, 0, 1, 0]]) + sparse_labels = _binary_2d_label_to_sparse_value([[0, 0, 0, 1], + [0, 0, 1, 0]]) dense_labels = np.array([[3], [2]], dtype=np.int64) for labels in (sparse_labels, dense_labels): @@ -3971,8 +4117,8 @@ class StreamingSparsePrecisionTest(test.TestCase): [5, 7, 2, 9, 6], ] sp_labels = sparse_tensor.SparseTensorValue( - indices=[[0, 0], [0, 1], [0, 2], [0, 3], [1, 0], [1, 1], [1, 2], - [1, 3]], + indices=[[0, 0], [0, 1], [0, 2], [0, 3], [1, 0], [1, 1], [1, 2], [1, + 3]], # values -1 and 10 are outside the [0, n_classes) range and are ignored. values=np.array([2, 7, -1, 8, 1, 2, 5, 10], np.int64), dense_shape=[2, 4]) @@ -4324,8 +4470,8 @@ class StreamingSparseRecallTest(test.TestCase): def test_one_label_at_k1_nan(self): predictions = [[0.1, 0.3, 0.2, 0.4], [0.1, 0.2, 0.3, 0.4]] top_k_predictions = [[3], [3]] - sparse_labels = _binary_2d_label_to_sparse_value( - [[0, 0, 0, 1], [0, 0, 1, 0]]) + sparse_labels = _binary_2d_label_to_sparse_value([[0, 0, 0, 1], + [0, 0, 1, 0]]) dense_labels = np.array([[3], [2]], dtype=np.int64) # Classes 0,1 have 0 labels, 0 predictions, classes -1 and 4 are out of @@ -4340,8 +4486,8 @@ class StreamingSparseRecallTest(test.TestCase): def test_one_label_at_k1_no_predictions(self): predictions = [[0.1, 0.3, 0.2, 0.4], [0.1, 0.2, 0.3, 0.4]] top_k_predictions = [[3], [3]] - sparse_labels = _binary_2d_label_to_sparse_value( - [[0, 0, 0, 1], [0, 0, 1, 0]]) + sparse_labels = _binary_2d_label_to_sparse_value([[0, 0, 0, 1], + [0, 0, 1, 0]]) dense_labels = np.array([[3], [2]], dtype=np.int64) for labels in (sparse_labels, dense_labels): @@ -4354,8 +4500,8 @@ class StreamingSparseRecallTest(test.TestCase): def test_one_label_at_k1(self): predictions = [[0.1, 0.3, 0.2, 0.4], [0.1, 0.2, 0.3, 0.4]] top_k_predictions = [[3], [3]] - sparse_labels = _binary_2d_label_to_sparse_value( - [[0, 0, 0, 1], [0, 0, 1, 0]]) + sparse_labels = _binary_2d_label_to_sparse_value([[0, 0, 0, 1], + [0, 0, 1, 0]]) dense_labels = np.array([[3], [2]], dtype=np.int64) for labels in (sparse_labels, dense_labels): @@ -4374,8 +4520,8 @@ class StreamingSparseRecallTest(test.TestCase): def test_one_label_at_k1_weighted(self): predictions = [[0.1, 0.3, 0.2, 0.4], [0.1, 0.2, 0.3, 0.4]] top_k_predictions = [[3], [3]] - sparse_labels = _binary_2d_label_to_sparse_value( - [[0, 0, 0, 1], [0, 0, 1, 0]]) + sparse_labels = _binary_2d_label_to_sparse_value([[0, 0, 0, 1], + [0, 0, 1, 0]]) dense_labels = np.array([[3], [2]], dtype=np.int64) for labels in (sparse_labels, dense_labels): @@ -4647,8 +4793,8 @@ class StreamingSparseRecallTest(test.TestCase): [5, 7, 2, 9, 6], ] sp_labels = sparse_tensor.SparseTensorValue( - indices=[[0, 0], [0, 1], [0, 2], [0, 3], [1, 0], [1, 1], [1, 2], - [1, 3]], + indices=[[0, 0], [0, 1], [0, 2], [0, 3], [1, 0], [1, 1], [1, 2], [1, + 3]], # values -1 and 10 are outside the [0, n_classes) range. values=np.array([2, 7, -1, 8, 1, 2, 5, 10], np.int64), dense_shape=[2, 4]) @@ -4661,10 +4807,7 @@ class StreamingSparseRecallTest(test.TestCase): expected=2.0 / 2, class_id=2) self._test_sparse_recall_at_top_k( - sp_labels, - top_k_predictions, - expected=2.0 / 2, - class_id=2) + sp_labels, top_k_predictions, expected=2.0 / 2, class_id=2) # Class 5: 1 label, incorrect. self._test_streaming_sparse_recall_at_k( @@ -4674,10 +4817,7 @@ class StreamingSparseRecallTest(test.TestCase): expected=1.0 / 1, class_id=5) self._test_sparse_recall_at_top_k( - sp_labels, - top_k_predictions, - expected=1.0 / 1, - class_id=5) + sp_labels, top_k_predictions, expected=1.0 / 1, class_id=5) # Class 7: 1 label, incorrect. self._test_streaming_sparse_recall_at_k( @@ -4687,10 +4827,7 @@ class StreamingSparseRecallTest(test.TestCase): expected=0.0 / 1, class_id=7) self._test_sparse_recall_at_top_k( - sp_labels, - top_k_predictions, - expected=0.0 / 1, - class_id=7) + sp_labels, top_k_predictions, expected=0.0 / 1, class_id=7) # All classes: 8 labels, 3 correct. self._test_streaming_sparse_recall_at_k( @@ -4740,10 +4877,8 @@ class StreamingSparseRecallTest(test.TestCase): [9, 4, 6, 2, 0], ]] sparse_labels = _binary_3d_label_to_sparse_value( - [[[0, 0, 1, 0, 0, 0, 0, 1, 1, 0], - [0, 1, 1, 0, 0, 1, 0, 0, 0, 0]], - [[0, 1, 1, 0, 0, 1, 0, 0, 0, 0], - [0, 0, 1, 0, 0, 0, 0, 1, 1, 0]]]) + [[[0, 0, 1, 0, 0, 0, 0, 1, 1, 0], [0, 1, 1, 0, 0, 1, 0, 0, 0, 0]], + [[0, 1, 1, 0, 0, 1, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 1, 1, 0]]]) dense_labels = np.array( [[[2, 7, 8], [1, 2, 5]], [ [1, 2, 5], @@ -4771,10 +4906,8 @@ class StreamingSparseRecallTest(test.TestCase): [9, 4, 6, 2, 0], ]] labels = _binary_3d_label_to_sparse_value( - [[[0, 0, 1, 0, 0, 0, 0, 1, 1, 0], - [0, 1, 1, 0, 0, 1, 0, 0, 0, 0]], - [[0, 1, 1, 0, 0, 1, 0, 1, 0, 0], - [0, 0, 1, 0, 0, 0, 0, 0, 1, 0]]]) + [[[0, 0, 1, 0, 0, 0, 0, 1, 1, 0], [0, 1, 1, 0, 0, 1, 0, 0, 0, 0]], + [[0, 1, 1, 0, 0, 1, 0, 1, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0, 1, 0]]]) # Class 2: 4 labels, all correct. self._test_streaming_sparse_recall_at_k( @@ -4813,10 +4946,8 @@ class StreamingSparseRecallTest(test.TestCase): [9, 4, 6, 2, 0], ]] labels = _binary_3d_label_to_sparse_value( - [[[0, 0, 1, 0, 0, 0, 0, 1, 1, 0], - [0, 1, 1, 0, 0, 1, 0, 0, 0, 0]], - [[0, 1, 1, 0, 0, 1, 0, 1, 0, 0], - [0, 0, 1, 0, 0, 0, 0, 0, 1, 0]]]) + [[[0, 0, 1, 0, 0, 0, 0, 1, 1, 0], [0, 1, 1, 0, 0, 1, 0, 0, 0, 0]], + [[0, 1, 1, 0, 0, 1, 0, 1, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0, 1, 0]]]) for class_id in xrange(10): self._test_streaming_sparse_recall_at_k( @@ -4867,10 +4998,8 @@ class StreamingSparseRecallTest(test.TestCase): [9, 4, 6, 2, 0], ]] labels = _binary_3d_label_to_sparse_value( - [[[0, 0, 1, 0, 0, 0, 0, 1, 1, 0], - [0, 1, 1, 0, 0, 1, 0, 0, 0, 0]], - [[0, 1, 1, 0, 0, 1, 0, 1, 0, 0], - [0, 0, 1, 0, 0, 0, 0, 0, 1, 0]]]) + [[[0, 0, 1, 0, 0, 0, 0, 1, 1, 0], [0, 1, 1, 0, 0, 1, 0, 0, 0, 0]], + [[0, 1, 1, 0, 0, 1, 0, 1, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0, 1, 0]]]) # Class 2: 2 labels, both correct. self._test_streaming_sparse_recall_at_k( @@ -4963,10 +5092,8 @@ class StreamingSparseRecallTest(test.TestCase): weights=[[0, 1], [0, 1]]) def test_sparse_tensor_value(self): - predictions = [[0.1, 0.3, 0.2, 0.4], - [0.1, 0.2, 0.3, 0.4]] - labels = [[0, 0, 1, 0], - [0, 0, 0, 1]] + predictions = [[0.1, 0.3, 0.2, 0.4], [0.1, 0.2, 0.3, 0.4]] + labels = [[0, 0, 1, 0], [0, 0, 0, 1]] expected_recall = 0.5 with self.test_session(): _, recall = metrics.streaming_sparse_recall_at_k( @@ -5009,8 +5136,8 @@ class StreamingMeanAbsoluteErrorTest(test.TestCase): def testValueTensorIsIdempotent(self): predictions = random_ops.random_normal((10, 3), seed=1) labels = random_ops.random_normal((10, 3), seed=2) - error, update_op = metrics.streaming_mean_absolute_error(predictions, - labels) + error, update_op = metrics.streaming_mean_absolute_error( + predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -5031,8 +5158,8 @@ class StreamingMeanAbsoluteErrorTest(test.TestCase): [1, 3, 2, 3], shape=(1, 4), dtype=dtypes_lib.float32) weights = constant_op.constant([0, 1, 0, 1], shape=(1, 4)) - error, update_op = metrics.streaming_mean_absolute_error(predictions, - labels, weights) + error, update_op = metrics.streaming_mean_absolute_error( + predictions, labels, weights) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -5075,8 +5202,8 @@ class StreamingMeanRelativeErrorTest(test.TestCase): predictions = random_ops.random_normal((10, 3), seed=1) labels = random_ops.random_normal((10, 3), seed=2) normalizer = random_ops.random_normal((10, 3), seed=3) - error, update_op = metrics.streaming_mean_relative_error(predictions, - labels, normalizer) + error, update_op = metrics.streaming_mean_relative_error( + predictions, labels, normalizer) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -5200,8 +5327,8 @@ class StreamingMeanSquaredErrorTest(test.TestCase): [1, 3, 2, 3], shape=(1, 4), dtype=dtypes_lib.float32) weights = constant_op.constant([0, 1, 0, 1], shape=(1, 4)) - error, update_op = metrics.streaming_mean_squared_error(predictions, labels, - weights) + error, update_op = metrics.streaming_mean_squared_error( + predictions, labels, weights) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -5224,8 +5351,8 @@ class StreamingMeanSquaredErrorTest(test.TestCase): _enqueue_vector(sess, labels_queue, [2, 4, 6]) labels = labels_queue.dequeue() - error, update_op = metrics.streaming_mean_squared_error(predictions, - labels) + error, update_op = metrics.streaming_mean_squared_error( + predictions, labels) sess.run(variables.local_variables_initializer()) sess.run(update_op) @@ -5292,10 +5419,10 @@ class StreamingMeanSquaredErrorTest(test.TestCase): _enqueue_vector(sess, labels_queue, [2, 4, 6]) labels = labels_queue.dequeue() - mae, ma_update_op = metrics.streaming_mean_absolute_error(predictions, - labels) - mse, ms_update_op = metrics.streaming_mean_squared_error(predictions, - labels) + mae, ma_update_op = metrics.streaming_mean_absolute_error( + predictions, labels) + mse, ms_update_op = metrics.streaming_mean_squared_error( + predictions, labels) sess.run(variables.local_variables_initializer()) sess.run([ma_update_op, ms_update_op]) @@ -5336,8 +5463,8 @@ class StreamingRootMeanSquaredErrorTest(test.TestCase): def testValueTensorIsIdempotent(self): predictions = random_ops.random_normal((10, 3), seed=1) labels = random_ops.random_normal((10, 3), seed=2) - error, update_op = metrics.streaming_root_mean_squared_error(predictions, - labels) + error, update_op = metrics.streaming_root_mean_squared_error( + predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -5357,8 +5484,8 @@ class StreamingRootMeanSquaredErrorTest(test.TestCase): 0.0, shape=(1, 3), dtype=dtypes_lib.float32) labels = constant_op.constant(0.0, shape=(1, 3), dtype=dtypes_lib.float32) - rmse, update_op = metrics.streaming_root_mean_squared_error(predictions, - labels) + rmse, update_op = metrics.streaming_root_mean_squared_error( + predictions, labels) sess.run(variables.local_variables_initializer()) self.assertEqual(0, sess.run(update_op)) @@ -5372,8 +5499,8 @@ class StreamingRootMeanSquaredErrorTest(test.TestCase): labels = constant_op.constant( [1, 3, 2], shape=(1, 3), dtype=dtypes_lib.float32) - rmse, update_op = metrics.streaming_root_mean_squared_error(predictions, - labels) + rmse, update_op = metrics.streaming_root_mean_squared_error( + predictions, labels) sess.run(variables.local_variables_initializer()) self.assertAlmostEqual(math.sqrt(6), update_op.eval(), 5) @@ -5387,9 +5514,8 @@ class StreamingRootMeanSquaredErrorTest(test.TestCase): [1, 3, 2, 3], shape=(1, 4), dtype=dtypes_lib.float32) weights = constant_op.constant([0, 1, 0, 1], shape=(1, 4)) - rmse, update_op = metrics.streaming_root_mean_squared_error(predictions, - labels, - weights) + rmse, update_op = metrics.streaming_root_mean_squared_error( + predictions, labels, weights) sess.run(variables.local_variables_initializer()) self.assertAlmostEqual(math.sqrt(13), sess.run(update_op)) @@ -5404,8 +5530,8 @@ class StreamingCovarianceTest(test.TestCase): def testVars(self): metrics.streaming_covariance( - predictions=math_ops.to_float(math_ops.range(10)) + array_ops.ones( - [10, 10]), + predictions=math_ops.to_float(math_ops.range(10)) + + array_ops.ones([10, 10]), labels=math_ops.to_float(math_ops.range(10)) + array_ops.ones([10, 10])) _assert_metric_variables(self, ( 'covariance/comoment:0', @@ -5417,8 +5543,8 @@ class StreamingCovarianceTest(test.TestCase): def testMetricsCollection(self): my_collection_name = '__metrics__' cov, _ = metrics.streaming_covariance( - predictions=math_ops.to_float(math_ops.range(10)) + array_ops.ones( - [10, 10]), + predictions=math_ops.to_float(math_ops.range(10)) + + array_ops.ones([10, 10]), labels=math_ops.to_float(math_ops.range(10)) + array_ops.ones([10, 10]), metrics_collections=[my_collection_name]) self.assertListEqual(ops.get_collection(my_collection_name), [cov]) @@ -5426,8 +5552,8 @@ class StreamingCovarianceTest(test.TestCase): def testUpdatesCollection(self): my_collection_name = '__updates__' _, update_op = metrics.streaming_covariance( - predictions=math_ops.to_float(math_ops.range(10)) + array_ops.ones( - [10, 10]), + predictions=math_ops.to_float(math_ops.range(10)) + + array_ops.ones([10, 10]), labels=math_ops.to_float(math_ops.range(10)) + array_ops.ones([10, 10]), updates_collections=[my_collection_name]) self.assertListEqual(ops.get_collection(my_collection_name), [update_op]) @@ -5487,9 +5613,8 @@ class StreamingCovarianceTest(test.TestCase): cov, update_op = metrics.streaming_covariance( predictions, labels, weights=weights) - expected_cov = np.cov([2, 4, 6, 8], - [1, 3, 2, 7], - fweights=[0, 1, 3, 1])[0, 1] + expected_cov = np.cov( + [2, 4, 6, 8], [1, 3, 2, 7], fweights=[0, 1, 3, 1])[0, 1] sess.run(variables.local_variables_initializer()) self.assertAlmostEqual(expected_cov, sess.run(update_op)) self.assertAlmostEqual(expected_cov, cov.eval()) @@ -5514,17 +5639,18 @@ class StreamingCovarianceTest(test.TestCase): predictions_t: predictions[stride * i:stride * (i + 1)], labels_t: labels[stride * i:stride * (i + 1)] } - self.assertEqual(np.isnan(prev_expected_cov), - np.isnan(sess.run(cov, feed_dict=feed_dict))) + self.assertEqual( + np.isnan(prev_expected_cov), + np.isnan(sess.run(cov, feed_dict=feed_dict))) if not np.isnan(prev_expected_cov): - self.assertAlmostEqual( - prev_expected_cov, sess.run(cov, feed_dict=feed_dict), 5) + self.assertAlmostEqual(prev_expected_cov, + sess.run(cov, feed_dict=feed_dict), 5) expected_cov = np.cov(predictions[:stride * (i + 1)], labels[:stride * (i + 1)])[0, 1] - self.assertAlmostEqual( - expected_cov, sess.run(update_op, feed_dict=feed_dict), 5) - self.assertAlmostEqual( - expected_cov, sess.run(cov, feed_dict=feed_dict), 5) + self.assertAlmostEqual(expected_cov, + sess.run(update_op, feed_dict=feed_dict), 5) + self.assertAlmostEqual(expected_cov, sess.run(cov, feed_dict=feed_dict), + 5) prev_expected_cov = expected_cov def testMultiUpdateWithErrorAndWeights(self): @@ -5552,18 +5678,20 @@ class StreamingCovarianceTest(test.TestCase): labels_t: labels[stride * i:stride * (i + 1)], weights_t: weights[stride * i:stride * (i + 1)] } - self.assertEqual(np.isnan(prev_expected_cov), - np.isnan(sess.run(cov, feed_dict=feed_dict))) + self.assertEqual( + np.isnan(prev_expected_cov), + np.isnan(sess.run(cov, feed_dict=feed_dict))) if not np.isnan(prev_expected_cov): - self.assertAlmostEqual( - prev_expected_cov, sess.run(cov, feed_dict=feed_dict), 5) - expected_cov = np.cov(predictions[:stride * (i + 1)], - labels[:stride * (i + 1)], - fweights=weights[:stride * (i + 1)])[0, 1] - self.assertAlmostEqual( - expected_cov, sess.run(update_op, feed_dict=feed_dict), 5) - self.assertAlmostEqual( - expected_cov, sess.run(cov, feed_dict=feed_dict), 5) + self.assertAlmostEqual(prev_expected_cov, + sess.run(cov, feed_dict=feed_dict), 5) + expected_cov = np.cov( + predictions[:stride * (i + 1)], + labels[:stride * (i + 1)], + fweights=weights[:stride * (i + 1)])[0, 1] + self.assertAlmostEqual(expected_cov, + sess.run(update_op, feed_dict=feed_dict), 5) + self.assertAlmostEqual(expected_cov, sess.run(cov, feed_dict=feed_dict), + 5) prev_expected_cov = expected_cov @@ -5574,8 +5702,8 @@ class StreamingPearsonRTest(test.TestCase): def testVars(self): metrics.streaming_pearson_correlation( - predictions=math_ops.to_float(math_ops.range(10)) + array_ops.ones( - [10, 10]), + predictions=math_ops.to_float(math_ops.range(10)) + + array_ops.ones([10, 10]), labels=math_ops.to_float(math_ops.range(10)) + array_ops.ones([10, 10])) _assert_metric_variables(self, ( 'pearson_r/covariance/comoment:0', @@ -5595,8 +5723,8 @@ class StreamingPearsonRTest(test.TestCase): def testMetricsCollection(self): my_collection_name = '__metrics__' pearson_r, _ = metrics.streaming_pearson_correlation( - predictions=math_ops.to_float(math_ops.range(10)) + array_ops.ones( - [10, 10]), + predictions=math_ops.to_float(math_ops.range(10)) + + array_ops.ones([10, 10]), labels=math_ops.to_float(math_ops.range(10)) + array_ops.ones([10, 10]), metrics_collections=[my_collection_name]) self.assertListEqual(ops.get_collection(my_collection_name), [pearson_r]) @@ -5604,8 +5732,8 @@ class StreamingPearsonRTest(test.TestCase): def testUpdatesCollection(self): my_collection_name = '__updates__' _, update_op = metrics.streaming_pearson_correlation( - predictions=math_ops.to_float(math_ops.range(10)) + array_ops.ones( - [10, 10]), + predictions=math_ops.to_float(math_ops.range(10)) + + array_ops.ones([10, 10]), labels=math_ops.to_float(math_ops.range(10)) + array_ops.ones([10, 10]), updates_collections=[my_collection_name]) self.assertListEqual(ops.get_collection(my_collection_name), [update_op]) @@ -5613,8 +5741,8 @@ class StreamingPearsonRTest(test.TestCase): def testValueTensorIsIdempotent(self): labels = random_ops.random_normal((10, 3), seed=2) predictions = labels * 0.5 + random_ops.random_normal((10, 3), seed=1) * 0.5 - pearson_r, update_op = metrics.streaming_pearson_correlation(predictions, - labels) + pearson_r, update_op = metrics.streaming_pearson_correlation( + predictions, labels) with self.test_session() as sess: sess.run(variables.local_variables_initializer()) @@ -5633,8 +5761,8 @@ class StreamingPearsonRTest(test.TestCase): predictions = math_ops.to_float(math_ops.range(10)) labels = math_ops.to_float(math_ops.range(10)) - pearson_r, update_op = metrics.streaming_pearson_correlation(predictions, - labels) + pearson_r, update_op = metrics.streaming_pearson_correlation( + predictions, labels) expected_r = np.corrcoef(np.arange(10), np.arange(10))[0, 1] sess.run(variables.local_variables_initializer()) @@ -5648,8 +5776,8 @@ class StreamingPearsonRTest(test.TestCase): labels = constant_op.constant( [1, 3, 2], shape=(1, 3), dtype=dtypes_lib.float32) - pearson_r, update_op = metrics.streaming_pearson_correlation(predictions, - labels) + pearson_r, update_op = metrics.streaming_pearson_correlation( + predictions, labels) expected_r = np.corrcoef([2, 4, 6], [1, 3, 2])[0, 1] sess.run(variables.local_variables_initializer()) @@ -5698,17 +5826,18 @@ class StreamingPearsonRTest(test.TestCase): predictions_t: predictions[stride * i:stride * (i + 1)], labels_t: labels[stride * i:stride * (i + 1)] } - self.assertEqual(np.isnan(prev_expected_r), - np.isnan(sess.run(pearson_r, feed_dict=feed_dict))) + self.assertEqual( + np.isnan(prev_expected_r), + np.isnan(sess.run(pearson_r, feed_dict=feed_dict))) if not np.isnan(prev_expected_r): - self.assertAlmostEqual( - prev_expected_r, sess.run(pearson_r, feed_dict=feed_dict), 5) + self.assertAlmostEqual(prev_expected_r, + sess.run(pearson_r, feed_dict=feed_dict), 5) expected_r = np.corrcoef(predictions[:stride * (i + 1)], labels[:stride * (i + 1)])[0, 1] - self.assertAlmostEqual( - expected_r, sess.run(update_op, feed_dict=feed_dict), 5) - self.assertAlmostEqual( - expected_r, sess.run(pearson_r, feed_dict=feed_dict), 5) + self.assertAlmostEqual(expected_r, + sess.run(update_op, feed_dict=feed_dict), 5) + self.assertAlmostEqual(expected_r, + sess.run(pearson_r, feed_dict=feed_dict), 5) prev_expected_r = expected_r def testMultiUpdateWithErrorAndWeights(self): @@ -5736,19 +5865,21 @@ class StreamingPearsonRTest(test.TestCase): labels_t: labels[stride * i:stride * (i + 1)], weights_t: weights[stride * i:stride * (i + 1)] } - self.assertEqual(np.isnan(prev_expected_r), - np.isnan(sess.run(pearson_r, feed_dict=feed_dict))) + self.assertEqual( + np.isnan(prev_expected_r), + np.isnan(sess.run(pearson_r, feed_dict=feed_dict))) if not np.isnan(prev_expected_r): - self.assertAlmostEqual( - prev_expected_r, sess.run(pearson_r, feed_dict=feed_dict), 5) - cmat = np.cov(predictions[:stride * (i + 1)], - labels[:stride * (i + 1)], - fweights=weights[:stride * (i + 1)]) + self.assertAlmostEqual(prev_expected_r, + sess.run(pearson_r, feed_dict=feed_dict), 5) + cmat = np.cov( + predictions[:stride * (i + 1)], + labels[:stride * (i + 1)], + fweights=weights[:stride * (i + 1)]) expected_r = cmat[0, 1] / np.sqrt(cmat[0, 0] * cmat[1, 1]) - self.assertAlmostEqual( - expected_r, sess.run(update_op, feed_dict=feed_dict), 5) - self.assertAlmostEqual( - expected_r, sess.run(pearson_r, feed_dict=feed_dict), 5) + self.assertAlmostEqual(expected_r, + sess.run(update_op, feed_dict=feed_dict), 5) + self.assertAlmostEqual(expected_r, + sess.run(pearson_r, feed_dict=feed_dict), 5) prev_expected_r = expected_r def testMultiUpdateWithErrorAndSingletonBatches(self): @@ -5758,7 +5889,7 @@ class StreamingPearsonRTest(test.TestCase): predictions = np.random.randn(n) labels = 0.5 * predictions + np.random.randn(n) stride = 10 - weights = (np.arange(n).reshape(n//stride, stride) % stride == 0) + weights = (np.arange(n).reshape(n // stride, stride) % stride == 0) for row in weights: np.random.shuffle(row) # Now, weights is one-hot by row - one item per batch has non-zero weight. @@ -5778,19 +5909,20 @@ class StreamingPearsonRTest(test.TestCase): labels_t: labels[stride * i:stride * (i + 1)], weights_t: weights[stride * i:stride * (i + 1)] } - cmat = np.cov(predictions[:stride * (i + 1)], - labels[:stride * (i + 1)], - fweights=weights[:stride * (i + 1)]) + cmat = np.cov( + predictions[:stride * (i + 1)], + labels[:stride * (i + 1)], + fweights=weights[:stride * (i + 1)]) expected_r = cmat[0, 1] / np.sqrt(cmat[0, 0] * cmat[1, 1]) actual_r = sess.run(update_op, feed_dict=feed_dict) self.assertEqual(np.isnan(expected_r), np.isnan(actual_r)) - self.assertEqual(np.isnan(expected_r), - np.isnan(sess.run(pearson_r, feed_dict=feed_dict))) + self.assertEqual( + np.isnan(expected_r), + np.isnan(sess.run(pearson_r, feed_dict=feed_dict))) if not np.isnan(expected_r): - self.assertAlmostEqual( - expected_r, actual_r, 5) - self.assertAlmostEqual( - expected_r, sess.run(pearson_r, feed_dict=feed_dict), 5) + self.assertAlmostEqual(expected_r, actual_r, 5) + self.assertAlmostEqual(expected_r, + sess.run(pearson_r, feed_dict=feed_dict), 5) class StreamingMeanCosineDistanceTest(test.TestCase): @@ -6191,20 +6323,14 @@ class StreamingMeanIOUTest(test.TestCase): self.assertAlmostEqual(desired_output, miou.eval()) def testUpdateOpEvalIsAccumulatedConfusionMatrix(self): - predictions = array_ops.concat( - [ - constant_op.constant( - 0, shape=[5]), constant_op.constant( - 1, shape=[5]) - ], - 0) - labels = array_ops.concat( - [ - constant_op.constant( - 0, shape=[3]), constant_op.constant( - 1, shape=[7]) - ], - 0) + predictions = array_ops.concat([ + constant_op.constant(0, shape=[5]), + constant_op.constant(1, shape=[5]) + ], 0) + labels = array_ops.concat([ + constant_op.constant(0, shape=[3]), + constant_op.constant(1, shape=[7]) + ], 0) num_classes = 2 with self.test_session() as sess: miou, update_op = metrics.streaming_mean_iou(predictions, labels, @@ -6238,29 +6364,20 @@ class StreamingMeanIOUTest(test.TestCase): self.assertEqual(0., miou.eval()) def testResultsWithSomeMissing(self): - predictions = array_ops.concat( - [ - constant_op.constant( - 0, shape=[5]), constant_op.constant( - 1, shape=[5]) - ], - 0) - labels = array_ops.concat( - [ - constant_op.constant( - 0, shape=[3]), constant_op.constant( - 1, shape=[7]) - ], - 0) + predictions = array_ops.concat([ + constant_op.constant(0, shape=[5]), + constant_op.constant(1, shape=[5]) + ], 0) + labels = array_ops.concat([ + constant_op.constant(0, shape=[3]), + constant_op.constant(1, shape=[7]) + ], 0) num_classes = 2 - weights = array_ops.concat( - [ - constant_op.constant( - 0, shape=[1]), constant_op.constant( - 1, shape=[8]), constant_op.constant( - 0, shape=[1]) - ], - 0) + weights = array_ops.concat([ + constant_op.constant(0, shape=[1]), + constant_op.constant(1, shape=[8]), + constant_op.constant(0, shape=[1]) + ], 0) with self.test_session() as sess: miou, update_op = metrics.streaming_mean_iou( predictions, labels, num_classes, weights=weights) @@ -6270,56 +6387,45 @@ class StreamingMeanIOUTest(test.TestCase): self.assertAlmostEqual(desired_miou, miou.eval()) def testMissingClassInLabels(self): - labels = constant_op.constant([ - [[0, 0, 1, 1, 0, 0], - [1, 0, 0, 0, 0, 1]], - [[1, 1, 1, 1, 1, 1], - [0, 0, 0, 0, 0, 0]]]) - predictions = constant_op.constant([ - [[0, 0, 2, 1, 1, 0], - [0, 1, 2, 2, 0, 1]], - [[0, 0, 2, 1, 1, 1], - [1, 1, 2, 0, 0, 0]]]) + labels = constant_op.constant([[[0, 0, 1, 1, 0, 0], [1, 0, 0, 0, 0, 1]], + [[1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0]]]) + predictions = constant_op.constant( + [[[0, 0, 2, 1, 1, 0], [0, 1, 2, 2, 0, 1]], [[0, 0, 2, 1, 1, 1], + [1, 1, 2, 0, 0, 0]]]) num_classes = 3 with self.test_session() as sess: - miou, update_op = metrics.streaming_mean_iou( - predictions, labels, num_classes) + miou, update_op = metrics.streaming_mean_iou(predictions, labels, + num_classes) sess.run(variables.local_variables_initializer()) self.assertAllEqual([[7, 4, 3], [3, 5, 2], [0, 0, 0]], update_op.eval()) - self.assertAlmostEqual( - 1 / 3 * (7 / (7 + 3 + 7) + 5 / (5 + 4 + 5) + 0 / (0 + 5 + 0)), - miou.eval()) + self.assertAlmostEqual(1 / 3 * (7 / (7 + 3 + 7) + 5 / (5 + 4 + 5) + 0 / + (0 + 5 + 0)), miou.eval()) def testMissingClassOverallSmall(self): labels = constant_op.constant([0]) predictions = constant_op.constant([0]) num_classes = 2 with self.test_session() as sess: - miou, update_op = metrics.streaming_mean_iou( - predictions, labels, num_classes) + miou, update_op = metrics.streaming_mean_iou(predictions, labels, + num_classes) sess.run(variables.local_variables_initializer()) self.assertAllEqual([[1, 0], [0, 0]], update_op.eval()) self.assertAlmostEqual(1, miou.eval()) def testMissingClassOverallLarge(self): - labels = constant_op.constant([ - [[0, 0, 1, 1, 0, 0], - [1, 0, 0, 0, 0, 1]], - [[1, 1, 1, 1, 1, 1], - [0, 0, 0, 0, 0, 0]]]) - predictions = constant_op.constant([ - [[0, 0, 1, 1, 0, 0], - [1, 1, 0, 0, 1, 1]], - [[0, 0, 0, 1, 1, 1], - [1, 1, 1, 0, 0, 0]]]) + labels = constant_op.constant([[[0, 0, 1, 1, 0, 0], [1, 0, 0, 0, 0, 1]], + [[1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0]]]) + predictions = constant_op.constant( + [[[0, 0, 1, 1, 0, 0], [1, 1, 0, 0, 1, 1]], [[0, 0, 0, 1, 1, 1], + [1, 1, 1, 0, 0, 0]]]) num_classes = 3 with self.test_session() as sess: - miou, update_op = metrics.streaming_mean_iou( - predictions, labels, num_classes) + miou, update_op = metrics.streaming_mean_iou(predictions, labels, + num_classes) sess.run(variables.local_variables_initializer()) self.assertAllEqual([[9, 5, 0], [3, 7, 0], [0, 0, 0]], update_op.eval()) - self.assertAlmostEqual( - 1 / 2 * (9 / (9 + 3 + 5) + 7 / (7 + 5 + 3)), miou.eval()) + self.assertAlmostEqual(1 / 2 * (9 / (9 + 3 + 5) + 7 / (7 + 5 + 3)), + miou.eval()) class StreamingConcatTest(test.TestCase): @@ -6660,5 +6766,211 @@ class CountTest(test.TestCase): self.assertAlmostEqual(4.1, result.eval(), 5) +class CohenKappaTest(test.TestCase): + + def _confusion_matrix_to_samples(self, confusion_matrix): + x, y = confusion_matrix.shape + pairs = [] + for label in range(x): + for feature in range(y): + pairs += [label, feature] * confusion_matrix[label, feature] + pairs = np.array(pairs).reshape((-1, 2)) + return pairs[:, 0], pairs[:, 1] + + def setUp(self): + np.random.seed(1) + ops.reset_default_graph() + + def testVars(self): + metrics.cohen_kappa( + predictions_idx=array_ops.ones((10, 1)), + labels=array_ops.ones((10, 1)), + num_classes=2) + _assert_metric_variables(self, ( + 'cohen_kappa/po:0', + 'cohen_kappa/pe_row:0', + 'cohen_kappa/pe_col:0', + )) + + def testMetricsCollection(self): + my_collection_name = '__metrics__' + kappa, _ = metrics.cohen_kappa( + predictions_idx=array_ops.ones((10, 1)), + labels=array_ops.ones((10, 1)), + num_classes=2, + metrics_collections=[my_collection_name]) + self.assertListEqual(ops.get_collection(my_collection_name), [kappa]) + + def testUpdatesCollection(self): + my_collection_name = '__updates__' + _, update_op = metrics.cohen_kappa( + predictions_idx=array_ops.ones((10, 1)), + labels=array_ops.ones((10, 1)), + num_classes=2, + updates_collections=[my_collection_name]) + self.assertListEqual(ops.get_collection(my_collection_name), [update_op]) + + def testValueTensorIsIdempotent(self): + predictions = random_ops.random_uniform( + (10, 1), maxval=3, dtype=dtypes_lib.int64, seed=1) + labels = random_ops.random_uniform( + (10, 1), maxval=3, dtype=dtypes_lib.int64, seed=2) + kappa, update_op = metrics.cohen_kappa(labels, predictions, 3) + + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + + # Run several updates. + for _ in range(10): + sess.run(update_op) + + # Then verify idempotency. + initial_kappa = kappa.eval() + for _ in range(10): + self.assertAlmostEqual(initial_kappa, kappa.eval(), 5) + + def testBasic(self): + confusion_matrix = np.array([[9, 3, 1], [4, 8, 2], [2, 1, 6]]) + # overall total = 36 + # po = [9, 8, 6], sum(po) = 23 + # pe_row = [15, 12, 9], pe_col = [13, 14, 9], so pe = [5.42, 4.67, 2.25] + # finally, kappa = (sum(po) - sum(pe)) / (N - sum(pe)) + # = (23 - 12.34) / (36 - 12.34) + # = 0.45 + # see: http://psych.unl.edu/psycrs/handcomp/hckappa.PDF + expect = 0.45 + labels, predictions = self._confusion_matrix_to_samples(confusion_matrix) + + dtypes = [dtypes_lib.int16, dtypes_lib.int32, dtypes_lib.int64] + shapes = [ + (len(labels,)), # 1-dim + (len(labels), 1) + ] # 2-dim + weights = [None, np.ones_like(labels)] + + for dtype in dtypes: + for shape in shapes: + for weight in weights: + with self.test_session() as sess: + predictions_tensor = constant_op.constant( + np.reshape(predictions, shape), dtype=dtype) + labels_tensor = constant_op.constant( + np.reshape(labels, shape), dtype=dtype) + kappa, update_op = metrics.cohen_kappa( + labels_tensor, predictions_tensor, 3, weights=weight) + + sess.run(variables.local_variables_initializer()) + self.assertAlmostEqual(expect, sess.run(update_op), 2) + self.assertAlmostEqual(expect, kappa.eval(), 2) + + def testAllCorrect(self): + inputs = np.arange(0, 100) % 4 + # confusion matrix + # [[25, 0, 0], + # [0, 25, 0], + # [0, 0, 25]] + # Calculated by v0.19: sklearn.metrics.cohen_kappa_score(inputs, inputs) + expect = 1.0 + + with self.test_session() as sess: + predictions = constant_op.constant(inputs, dtype=dtypes_lib.float32) + labels = constant_op.constant(inputs) + kappa, update_op = metrics.cohen_kappa(labels, predictions, 4) + + sess.run(variables.local_variables_initializer()) + self.assertAlmostEqual(expect, sess.run(update_op), 5) + self.assertAlmostEqual(expect, kappa.eval(), 5) + + def testAllIncorrect(self): + labels = np.arange(0, 100) % 4 + predictions = (labels + 1) % 4 + # confusion matrix + # [[0, 25, 0], + # [0, 0, 25], + # [25, 0, 0]] + # Calculated by v0.19: sklearn.metrics.cohen_kappa_score( + # labels, predictions) + expect = -0.333333333333 + + with self.test_session() as sess: + predictions = constant_op.constant(predictions, dtype=dtypes_lib.float32) + labels = constant_op.constant(labels) + kappa, update_op = metrics.cohen_kappa(labels, predictions, 4) + + sess.run(variables.local_variables_initializer()) + self.assertAlmostEqual(expect, sess.run(update_op), 5) + self.assertAlmostEqual(expect, kappa.eval(), 5) + + def testWeighted(self): + confusion_matrix = np.array([[9, 3, 1], [4, 8, 2], [2, 1, 6]]) + labels, predictions = self._confusion_matrix_to_samples(confusion_matrix) + num_samples = np.sum(confusion_matrix, dtype=np.int32) + weights = (np.arange(0, num_samples) % 5) / 5.0 + # Calculated by v0.19: sklearn.metrics.cohen_kappa_score( + # labels, predictions, sample_weight=weights) + expect = 0.453466583385 + + with self.test_session() as sess: + predictions = constant_op.constant(predictions, dtype=dtypes_lib.float32) + labels = constant_op.constant(labels) + kappa, update_op = metrics.cohen_kappa( + labels, predictions, 4, weights=weights) + + sess.run(variables.local_variables_initializer()) + self.assertAlmostEqual(expect, sess.run(update_op), 5) + self.assertAlmostEqual(expect, kappa.eval(), 5) + + def testWithMultipleUpdates(self): + confusion_matrix = np.array([[90, 30, 10, 20], [40, 80, 20, 30], + [20, 10, 60, 35], [15, 25, 30, 25]]) + labels, predictions = self._confusion_matrix_to_samples(confusion_matrix) + num_samples = np.sum(confusion_matrix, dtype=np.int32) + weights = (np.arange(0, num_samples) % 5) / 5.0 + num_classes = confusion_matrix.shape[0] + + batch_size = num_samples // 10 + predictions_t = array_ops.placeholder( + dtypes_lib.float32, shape=(batch_size,)) + labels_t = array_ops.placeholder(dtypes_lib.int32, shape=(batch_size,)) + weights_t = array_ops.placeholder(dtypes_lib.float32, shape=(batch_size,)) + kappa, update_op = metrics.cohen_kappa( + labels_t, predictions_t, num_classes, weights=weights_t) + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + + for idx in range(0, num_samples, batch_size): + batch_start, batch_end = idx, idx + batch_size + sess.run( + update_op, + feed_dict={ + labels_t: labels[batch_start:batch_end], + predictions_t: predictions[batch_start:batch_end], + weights_t: weights[batch_start:batch_end] + }) + # Calculated by v0.19: sklearn.metrics.cohen_kappa_score( + # labels_np, predictions_np, + # sample_weight=weights_np) + expect = 0.289965397924 + self.assertAlmostEqual(expect, kappa.eval(), 5) + + def testInvalidNumClasses(self): + predictions = array_ops.placeholder(dtypes_lib.float32, shape=(4, 1)) + labels = array_ops.placeholder(dtypes_lib.int32, shape=(4, 1)) + with self.assertRaisesRegexp(ValueError, 'num_classes'): + metrics.cohen_kappa(labels, predictions, 1) + + def testInvalidDimension(self): + predictions = array_ops.placeholder(dtypes_lib.float32, shape=(4, 1)) + invalid_labels = array_ops.placeholder(dtypes_lib.int32, shape=(4, 2)) + with self.assertRaises(ValueError): + metrics.cohen_kappa(invalid_labels, predictions, 3) + + invalid_predictions = array_ops.placeholder( + dtypes_lib.float32, shape=(4, 2)) + labels = array_ops.placeholder(dtypes_lib.int32, shape=(4, 1)) + with self.assertRaises(ValueError): + metrics.cohen_kappa(labels, invalid_predictions, 3) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/model_pruning/examples/cifar10/cifar10_input.py b/tensorflow/contrib/model_pruning/examples/cifar10/cifar10_input.py index d07fece4bc668612d517e8dcaab1a35451a0238e..6a3b535eb447dd80f8e39d1d005f8f1d4f503549 100644 --- a/tensorflow/contrib/model_pruning/examples/cifar10/cifar10_input.py +++ b/tensorflow/contrib/model_pruning/examples/cifar10/cifar10_input.py @@ -58,6 +58,7 @@ def read_cifar10(filename_queue): class CIFAR10Record(object): pass + result = CIFAR10Record() # Dimensions of the images in the CIFAR-10 dataset. @@ -147,8 +148,9 @@ def distorted_inputs(data_dir, batch_size): images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size. labels: Labels. 1D tensor of [batch_size] size. """ - filenames = [os.path.join(data_dir, 'data_batch_%d.bin' % i) - for i in xrange(1, 6)] + filenames = [ + os.path.join(data_dir, 'data_batch_%d.bin' % i) for i in xrange(1, 6) + ] for f in filenames: if not tf.gfile.Exists(f): raise ValueError('Failed to find file: ' + f) @@ -174,10 +176,9 @@ def distorted_inputs(data_dir, batch_size): # Because these operations are not commutative, consider randomizing # the order their operation. - distorted_image = tf.image.random_brightness(distorted_image, - max_delta=63) - distorted_image = tf.image.random_contrast(distorted_image, - lower=0.2, upper=1.8) + distorted_image = tf.image.random_brightness(distorted_image, max_delta=63) + distorted_image = tf.image.random_contrast( + distorted_image, lower=0.2, upper=1.8) # Subtract off the mean and divide by the variance of the pixels. float_image = tf.image.per_image_standardization(distorted_image) @@ -188,15 +189,18 @@ def distorted_inputs(data_dir, batch_size): # Ensure that the random shuffling has good mixing properties. min_fraction_of_examples_in_queue = 0.4 - min_queue_examples = int(NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * - min_fraction_of_examples_in_queue) - print ('Filling queue with %d CIFAR images before starting to train. ' - 'This will take a few minutes.' % min_queue_examples) + min_queue_examples = int( + NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * min_fraction_of_examples_in_queue) + print('Filling queue with %d CIFAR images before starting to train. ' + 'This will take a few minutes.' % min_queue_examples) # Generate a batch of images and labels by building up a queue of examples. - return _generate_image_and_label_batch(float_image, read_input.label, - min_queue_examples, batch_size, - shuffle=True) + return _generate_image_and_label_batch( + float_image, + read_input.label, + min_queue_examples, + batch_size, + shuffle=True) def inputs(eval_data, data_dir, batch_size): @@ -212,8 +216,9 @@ def inputs(eval_data, data_dir, batch_size): labels: Labels. 1D tensor of [batch_size] size. """ if not eval_data: - filenames = [os.path.join(data_dir, 'data_batch_%d.bin' % i) - for i in xrange(1, 6)] + filenames = [ + os.path.join(data_dir, 'data_batch_%d.bin' % i) for i in xrange(1, 6) + ] num_examples_per_epoch = NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN else: filenames = [os.path.join(data_dir, 'test_batch.bin')] @@ -235,8 +240,8 @@ def inputs(eval_data, data_dir, batch_size): # Image processing for evaluation. # Crop the central [height, width] of the image. - resized_image = tf.image.resize_image_with_crop_or_pad(reshaped_image, - width, height) + resized_image = tf.image.resize_image_with_crop_or_pad( + reshaped_image, width, height) # Subtract off the mean and divide by the variance of the pixels. float_image = tf.image.per_image_standardization(resized_image) @@ -247,10 +252,13 @@ def inputs(eval_data, data_dir, batch_size): # Ensure that the random shuffling has good mixing properties. min_fraction_of_examples_in_queue = 0.4 - min_queue_examples = int(num_examples_per_epoch * - min_fraction_of_examples_in_queue) + min_queue_examples = int( + num_examples_per_epoch * min_fraction_of_examples_in_queue) # Generate a batch of images and labels by building up a queue of examples. - return _generate_image_and_label_batch(float_image, read_input.label, - min_queue_examples, batch_size, - shuffle=False) + return _generate_image_and_label_batch( + float_image, + read_input.label, + min_queue_examples, + batch_size, + shuffle=False) diff --git a/tensorflow/contrib/model_pruning/examples/cifar10/cifar10_pruning.py b/tensorflow/contrib/model_pruning/examples/cifar10/cifar10_pruning.py index 0d1de869f6ef91791a235cfe545b3b3a9b734e72..660f0168b10aa1e5b320cb476b051918804d2bde 100644 --- a/tensorflow/contrib/model_pruning/examples/cifar10/cifar10_pruning.py +++ b/tensorflow/contrib/model_pruning/examples/cifar10/cifar10_pruning.py @@ -48,16 +48,16 @@ from tensorflow.contrib.model_pruning.python import pruning # Global constants describing the CIFAR-10 data set. IMAGE_SIZE = cifar10_input.IMAGE_SIZE NUM_CLASSES = cifar10_input.NUM_CLASSES -NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = cifar10_input.NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN +NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = cifar10_input.NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN # pylint: disable=line-too-long NUM_EXAMPLES_PER_EPOCH_FOR_EVAL = cifar10_input.NUM_EXAMPLES_PER_EPOCH_FOR_EVAL BATCH_SIZE = 128 DATA_DIR = '/tmp/cifar10_data' # Constants describing the training process. -MOVING_AVERAGE_DECAY = 0.9999 # The decay to use for the moving average. -NUM_EPOCHS_PER_DECAY = 350.0 # Epochs after which learning rate decays. +MOVING_AVERAGE_DECAY = 0.9999 # The decay to use for the moving average. +NUM_EPOCHS_PER_DECAY = 350.0 # Epochs after which learning rate decays. LEARNING_RATE_DECAY_FACTOR = 0.1 # Learning rate decay factor. -INITIAL_LEARNING_RATE = 0.1 # Initial learning rate. +INITIAL_LEARNING_RATE = 0.1 # Initial learning rate. # If a model is trained with multiple GPUs, prefix all Op names with tower_name # to differentiate the operations. Note that this prefix is removed from the @@ -82,8 +82,7 @@ def _activation_summary(x): # session. This helps the clarity of presentation on tensorboard. tensor_name = re.sub('%s_[0-9]*/' % TOWER_NAME, '', x.op.name) tf.summary.histogram(tensor_name + '/activations', x) - tf.summary.scalar(tensor_name + '/sparsity', - tf.nn.zero_fraction(x)) + tf.summary.scalar(tensor_name + '/sparsity', tf.nn.zero_fraction(x)) def _variable_on_cpu(name, shape, initializer): @@ -120,10 +119,9 @@ def _variable_with_weight_decay(name, shape, stddev, wd): Variable Tensor """ dtype = tf.float32 - var = _variable_on_cpu( - name, - shape, - tf.truncated_normal_initializer(stddev=stddev, dtype=dtype)) + var = _variable_on_cpu(name, shape, + tf.truncated_normal_initializer( + stddev=stddev, dtype=dtype)) if wd is not None: weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss') tf.add_to_collection('losses', weight_decay) @@ -188,10 +186,8 @@ def inference(images): # Note that the masks are applied only to the weight tensors # conv1 with tf.variable_scope('conv1') as scope: - kernel = _variable_with_weight_decay('weights', - shape=[5, 5, 3, 64], - stddev=5e-2, - wd=0.0) + kernel = _variable_with_weight_decay( + 'weights', shape=[5, 5, 3, 64], stddev=5e-2, wd=0.0) conv = tf.nn.conv2d( images, pruning.apply_mask(kernel, scope), [1, 1, 1, 1], padding='SAME') @@ -201,18 +197,20 @@ def inference(images): _activation_summary(conv1) # pool1 - pool1 = tf.nn.max_pool(conv1, ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1], - padding='SAME', name='pool1') + pool1 = tf.nn.max_pool( + conv1, + ksize=[1, 3, 3, 1], + strides=[1, 2, 2, 1], + padding='SAME', + name='pool1') # norm1 - norm1 = tf.nn.lrn(pool1, 4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, - name='norm1') + norm1 = tf.nn.lrn( + pool1, 4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, name='norm1') # conv2 with tf.variable_scope('conv2') as scope: - kernel = _variable_with_weight_decay('weights', - shape=[5, 5, 64, 64], - stddev=5e-2, - wd=0.0) + kernel = _variable_with_weight_decay( + 'weights', shape=[5, 5, 64, 64], stddev=5e-2, wd=0.0) conv = tf.nn.conv2d( norm1, pruning.apply_mask(kernel, scope), [1, 1, 1, 1], padding='SAME') biases = _variable_on_cpu('biases', [64], tf.constant_initializer(0.1)) @@ -221,19 +219,23 @@ def inference(images): _activation_summary(conv2) # norm2 - norm2 = tf.nn.lrn(conv2, 4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, - name='norm2') + norm2 = tf.nn.lrn( + conv2, 4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, name='norm2') # pool2 - pool2 = tf.nn.max_pool(norm2, ksize=[1, 3, 3, 1], - strides=[1, 2, 2, 1], padding='SAME', name='pool2') + pool2 = tf.nn.max_pool( + norm2, + ksize=[1, 3, 3, 1], + strides=[1, 2, 2, 1], + padding='SAME', + name='pool2') # local3 with tf.variable_scope('local3') as scope: # Move everything into depth so we can perform a single matrix multiply. reshape = tf.reshape(pool2, [BATCH_SIZE, -1]) dim = reshape.get_shape()[1].value - weights = _variable_with_weight_decay('weights', shape=[dim, 384], - stddev=0.04, wd=0.004) + weights = _variable_with_weight_decay( + 'weights', shape=[dim, 384], stddev=0.04, wd=0.004) biases = _variable_on_cpu('biases', [384], tf.constant_initializer(0.1)) local3 = tf.nn.relu( tf.matmul(reshape, pruning.apply_mask(weights, scope)) + biases, @@ -242,8 +244,8 @@ def inference(images): # local4 with tf.variable_scope('local4') as scope: - weights = _variable_with_weight_decay('weights', shape=[384, 192], - stddev=0.04, wd=0.004) + weights = _variable_with_weight_decay( + 'weights', shape=[384, 192], stddev=0.04, wd=0.004) biases = _variable_on_cpu('biases', [192], tf.constant_initializer(0.1)) local4 = tf.nn.relu( tf.matmul(local3, pruning.apply_mask(weights, scope)) + biases, @@ -255,8 +257,8 @@ def inference(images): # tf.nn.sparse_softmax_cross_entropy_with_logits accepts the unscaled logits # and performs the softmax internally for efficiency. with tf.variable_scope('softmax_linear') as scope: - weights = _variable_with_weight_decay('weights', [192, NUM_CLASSES], - stddev=1/192.0, wd=0.0) + weights = _variable_with_weight_decay( + 'weights', [192, NUM_CLASSES], stddev=1 / 192.0, wd=0.0) biases = _variable_on_cpu('biases', [NUM_CLASSES], tf.constant_initializer(0.0)) softmax_linear = tf.add( @@ -337,11 +339,12 @@ def train(total_loss, global_step): decay_steps = int(num_batches_per_epoch * NUM_EPOCHS_PER_DECAY) # Decay the learning rate exponentially based on the number of steps. - lr = tf.train.exponential_decay(INITIAL_LEARNING_RATE, - global_step, - decay_steps, - LEARNING_RATE_DECAY_FACTOR, - staircase=True) + lr = tf.train.exponential_decay( + INITIAL_LEARNING_RATE, + global_step, + decay_steps, + LEARNING_RATE_DECAY_FACTOR, + staircase=True) tf.summary.scalar('learning_rate', lr) # Generate moving averages of all losses and associated summaries. @@ -365,8 +368,8 @@ def train(total_loss, global_step): tf.summary.histogram(var.op.name + '/gradients', grad) # Track the moving averages of all trainable variables. - variable_averages = tf.train.ExponentialMovingAverage( - MOVING_AVERAGE_DECAY, global_step) + variable_averages = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY, + global_step) variables_averages_op = variable_averages.apply(tf.trainable_variables()) with tf.control_dependencies([apply_gradient_op, variables_averages_op]): @@ -383,10 +386,13 @@ def maybe_download_and_extract(): filename = DATA_URL.split('/')[-1] filepath = os.path.join(dest_directory, filename) if not os.path.exists(filepath): + def _progress(count, block_size, total_size): - sys.stdout.write('\r>> Downloading %s %.1f%%' % (filename, - float(count * block_size) / float(total_size) * 100.0)) + sys.stdout.write('\r>> Downloading %s %.1f%%' % + (filename, + float(count * block_size) / float(total_size) * 100.0)) sys.stdout.flush() + filepath, _ = urllib.request.urlretrieve(DATA_URL, filepath, _progress) print() statinfo = os.stat(filepath) diff --git a/tensorflow/contrib/model_pruning/python/layers/core_layers.py b/tensorflow/contrib/model_pruning/python/layers/core_layers.py index 95dfd8f4213a8729f5954eb0626f28ecc9265bbb..764ab620bc2227ff5e8e3f473d689e0e133e83d4 100644 --- a/tensorflow/contrib/model_pruning/python/layers/core_layers.py +++ b/tensorflow/contrib/model_pruning/python/layers/core_layers.py @@ -210,7 +210,7 @@ class _MaskedConv(base.Layer): return self.activation(outputs) return outputs - def _compute_output_shape(self, input_shape): + def compute_output_shape(self, input_shape): input_shape = tensor_shape.TensorShape(input_shape).as_list() if self.data_format == 'channels_last': space = input_shape[1:-1] @@ -467,7 +467,7 @@ class MaskedFullyConnected(base.Layer): return self.activation(outputs) # pylint: disable=not-callable return outputs - def _compute_output_shape(self, input_shape): + def compute_output_shape(self, input_shape): input_shape = tensor_shape.TensorShape(input_shape) input_shape = input_shape.with_rank_at_least(2) if input_shape[-1].value is None: diff --git a/tensorflow/contrib/model_pruning/python/layers/layers.py b/tensorflow/contrib/model_pruning/python/layers/layers.py index dfebb9a6794056dd43b0699ccbcc5797f2f172f7..988748ad75bdf72f1da3f4e1c6e85aabb04a5954 100644 --- a/tensorflow/contrib/model_pruning/python/layers/layers.py +++ b/tensorflow/contrib/model_pruning/python/layers/layers.py @@ -21,7 +21,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np import six from tensorflow.contrib.framework.python.ops import add_arg_scope diff --git a/tensorflow/contrib/mpi/BUILD b/tensorflow/contrib/mpi/BUILD index d9d55faf50b7f5043bfd0ed3b3d9ca5c404c7627..23f90cf77ef0bde34f3938688aa6ca2f6e9bbc53 100644 --- a/tensorflow/contrib/mpi/BUILD +++ b/tensorflow/contrib/mpi/BUILD @@ -71,6 +71,8 @@ cc_library( "//tensorflow/core:protos_cc", "//tensorflow/core:worker_proto_cc", "//tensorflow/core/distributed_runtime:base_rendezvous_mgr", + "//tensorflow/core/distributed_runtime:recent_request_ids", + "//tensorflow/core/distributed_runtime:request_id", "//tensorflow/core/distributed_runtime:session_mgr", "//tensorflow/core/distributed_runtime:tensor_coding", "//tensorflow/core/distributed_runtime:worker_env", diff --git a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc index 1a2563d20fdc33d3c5e4a85561b61d04d3eeabff..6a7f5efecdb4062874a09df227d139ad20d59f3f 100644 --- a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc +++ b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc @@ -24,17 +24,19 @@ limitations under the License. #include #include -#include "tensorflow/core/distributed_runtime/tensor_coding.h" #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/gpu/gpu_util.h" #include "tensorflow/core/distributed_runtime/session_mgr.h" +#include "tensorflow/core/distributed_runtime/tensor_coding.h" namespace tensorflow { MPIRendezvousMgr::MPIRendezvousMgr(const WorkerEnv* env) - : BaseRendezvousMgr(env), worker_env_2(env), use_optimal_transfer_(false) { - + : BaseRendezvousMgr(env), + worker_env_2(env), + use_optimal_transfer_(false), + recv_tensor_recent_request_ids_(100000) { const char* mpienv = getenv("MPI_OPTIMAL_PATH"); if (mpienv && mpienv[0] == '1') { LOG(INFO) << "MPI Optimal copy path enabled (Requires CUDA-Aware MPI when " @@ -60,7 +62,6 @@ BaseRemoteRendezvous* MPIRendezvousMgr::Create(int64 step_id, void MPIRemoteRendezvous::RecvFromRemoteAsync( const Rendezvous::ParsedKey& parsed, const Rendezvous::Args& recv_args, DoneCallback done) { - Status s = Status::OK(); MPIRequestTensorCall* rendezvous_call = new MPIRequestTensorCall(); @@ -101,37 +102,37 @@ void MPIRemoteRendezvous::RecvFromRemoteAsync( // Create the function which is called when the Tensor is send by remote const int64 temp1 = step_id_; rendezvous_call->recv_call_ = - [this, parsed, recv_args, done, dst, temp1, rendezvous_call]( - MPIRecvTensorResponse mpi_response) { - Status s; - Device* dst_device; - if (s.ok()) { - s = env_->device_mgr->LookupDevice(parsed.dst_device, &dst_device); - CHECK(s.ok()) << "Device lookup failed"; - } - - VLOG(3) << "MPI Received tensor " << parsed.FullKey() - << " @ step: " << temp1 - << " single-send: " << mpi_response.singlesend(); - - Tensor val; - if (mpi_response.singlesend()) { - dst_device->MakeTensorFromProto(mpi_response.response().tensor(), - recv_args.alloc_attrs, &val); - } else { - TensorResponse tr; - tr.InitAlloc(dst_device, recv_args.alloc_attrs); - tr.InitPartial(mpi_response.response()); - const size_t nBytes = tr.tensor().TotalBytes(); - void* data = const_cast(DMAHelper::base(&tr.tensor())); - MPI_Status status; - MPI_CHECK(MPI_Recv(data, static_cast(nBytes), MPI_BYTE, dst, - TAG_SENDTENSOR2, MPI_COMM_WORLD, &status)); - val = std::move(tr.tensor()); - } - - done(s, Args(), recv_args, val, mpi_response.response().is_dead()); - }; + [this, parsed, recv_args, done, dst, temp1, + rendezvous_call](MPIRecvTensorResponse mpi_response) { + Status s; + Device* dst_device; + if (s.ok()) { + s = env_->device_mgr->LookupDevice(parsed.dst_device, &dst_device); + CHECK(s.ok()) << "Device lookup failed"; + } + + VLOG(3) << "MPI Received tensor " << parsed.FullKey() + << " @ step: " << temp1 + << " single-send: " << mpi_response.singlesend(); + + Tensor val; + if (mpi_response.singlesend()) { + dst_device->MakeTensorFromProto(mpi_response.response().tensor(), + recv_args.alloc_attrs, &val); + } else { + TensorResponse tr; + tr.InitAlloc(dst_device, recv_args.alloc_attrs); + tr.InitPartial(mpi_response.response()); + const size_t nBytes = tr.tensor().TotalBytes(); + void* data = const_cast(DMAHelper::base(&tr.tensor())); + MPI_Status status; + MPI_CHECK(MPI_Recv(data, static_cast(nBytes), MPI_BYTE, dst, + TAG_SENDTENSOR2, MPI_COMM_WORLD, &status)); + val = std::move(tr.tensor()); + } + + done(s, Args(), recv_args, val, mpi_response.response().is_dead()); + }; MPIRendezvousMgr* mgr = reinterpret_cast(this->rendezvous_mgr_); @@ -149,15 +150,19 @@ MPIRemoteRendezvous::~MPIRemoteRendezvous() {} */ void MPIRendezvousMgr::AddRequest(RecvTensorRequest request, const int mpi_dst) { + TF_CHECK_OK(recv_tensor_recent_request_ids_.TrackUnique( + request.request_id(), "RecvTensor (MPIRendezvousMgr)", request)); const int64 step_id = request.step_id(); const std::string& key = request.rendezvous_key(); Rendezvous::ParsedKey parsed; TF_CHECK_OK(Rendezvous::ParseKey(key, &parsed)); MPIRecvTensorCallBack send_cb = [this, mpi_dst, parsed]( - const Status& status, const Rendezvous::Args& send_args, - const Rendezvous::Args& recv_args, const Tensor& val, bool is_dead, - MPISendTensorCall* mpi_send_call) { + const Status& status, + const Rendezvous::Args& send_args, + const Rendezvous::Args& recv_args, + const Tensor& val, bool is_dead, + MPISendTensorCall* mpi_send_call) { // TODO(jbedorf) this should be a loop over max size CHECK(mpi_send_call->mRes_.ByteSize() < INT_MAX) << "Buffer too large for single transfer"; @@ -190,74 +195,78 @@ void MPIRendezvousMgr::AddRequest(RecvTensorRequest request, }; // Wrapper around the read callback to place the callback on our queue - Rendezvous::DoneCallback done_cb = [this, parsed, step_id, send_cb]( - const Status& status, const Rendezvous::Args& send_args, - const Rendezvous::Args& recv_args, const Tensor& val, bool is_dead) { - if (!status.ok()) { - CHECK(status.ok()) << "RecvLocalAsync was not ok, key: " - << parsed.FullKey() << " step: " << step_id - << " error message: " << status.error_message(); - return; - } - - VLOG(3) << "MPI Sending tensor " << parsed.FullKey() - << " @ step: " << step_id << std::endl; - - auto mpi_send_call = new MPISendTensorCall(); - mpi_send_call->Init(parsed, step_id, is_dead); - - Device* src_dev = nullptr; - Status s = this->worker_env_2->device_mgr->LookupDevice(parsed.src_device, - &src_dev); - CHECK(s.ok()) << "src device not found"; - - // Control if shape and data should be send together or if we can optimize - // it in two different transfers, thereby reducing memory copies - bool doOptimalTransfer = true; - if (!DataTypeCanUseMemcpy(val.dtype())) doOptimalTransfer = false; - if (val.TotalBytes() < 1024) doOptimalTransfer = false; - - doOptimalTransfer = doOptimalTransfer && use_optimal_transfer_; - - if (doOptimalTransfer) { - // First send the Tensor description and in a follow up transfer the data - mpi_send_call->mRes_.mutable_response()->mutable_tensor()->set_dtype( - val.dtype()); - val.shape().AsProto(mpi_send_call->mRes_.mutable_response() - ->mutable_tensor() - ->mutable_tensor_shape()); - mpi_send_call->mRes_.set_singlesend(false); - } else { - // Send the Tensor description and data in a single transfer - if (src_dev->tensorflow_gpu_device_info() && - (!send_args.alloc_attrs.on_host())) { - Notification n; - GPUUtil::SetProtoFromGPU( - val, src_dev, send_args.device_context, - mpi_send_call->mRes_.mutable_response()->mutable_tensor(), is_dead, - [&n, &s](const Status& s_) { - s = s_; - n.Notify(); - }); - n.WaitForNotification(); - } else { - val.AsProtoTensorContent( - mpi_send_call->mRes_.mutable_response()->mutable_tensor()); - } - } - - std::function res = std::bind( - send_cb, status, send_args, recv_args, val, is_dead, mpi_send_call); - - SendQueueEntry req(parsed.FullKey().ToString().c_str(), std::move(res)); - - this->QueueSendRequest(req); - - // Wait for the notification that indicates the tensor has been - // successfully transmitted to the remote process. Only needed if we - // have not parsed the tensor to proto - if (doOptimalTransfer) mpi_send_call->n_.WaitForNotification(); - }; // done_cb + Rendezvous::DoneCallback done_cb = + [this, parsed, step_id, send_cb]( + const Status& status, const Rendezvous::Args& send_args, + const Rendezvous::Args& recv_args, const Tensor& val, bool is_dead) { + if (!status.ok()) { + CHECK(status.ok()) + << "RecvLocalAsync was not ok, key: " << parsed.FullKey() + << " step: " << step_id + << " error message: " << status.error_message(); + return; + } + + VLOG(3) << "MPI Sending tensor " << parsed.FullKey() + << " @ step: " << step_id << std::endl; + + auto mpi_send_call = new MPISendTensorCall(); + mpi_send_call->Init(parsed, step_id, is_dead); + + Device* src_dev = nullptr; + Status s = this->worker_env_2->device_mgr->LookupDevice( + parsed.src_device, &src_dev); + CHECK(s.ok()) << "src device not found"; + + // Control if shape and data should be send together or if we can + // optimize it in two different transfers, thereby reducing memory + // copies + bool doOptimalTransfer = true; + if (!DataTypeCanUseMemcpy(val.dtype())) doOptimalTransfer = false; + if (val.TotalBytes() < 1024) doOptimalTransfer = false; + + doOptimalTransfer = doOptimalTransfer && use_optimal_transfer_; + + if (doOptimalTransfer) { + // First send the Tensor description and in a follow up transfer the + // data + mpi_send_call->mRes_.mutable_response()->mutable_tensor()->set_dtype( + val.dtype()); + val.shape().AsProto(mpi_send_call->mRes_.mutable_response() + ->mutable_tensor() + ->mutable_tensor_shape()); + mpi_send_call->mRes_.set_singlesend(false); + } else { + // Send the Tensor description and data in a single transfer + if (src_dev->tensorflow_gpu_device_info() && + (!send_args.alloc_attrs.on_host())) { + Notification n; + GPUUtil::SetProtoFromGPU( + val, src_dev, send_args.device_context, + mpi_send_call->mRes_.mutable_response()->mutable_tensor(), + is_dead, [&n, &s](const Status& s_) { + s = s_; + n.Notify(); + }); + n.WaitForNotification(); + } else { + val.AsProtoTensorContent( + mpi_send_call->mRes_.mutable_response()->mutable_tensor()); + } + } + + std::function res = std::bind( + send_cb, status, send_args, recv_args, val, is_dead, mpi_send_call); + + SendQueueEntry req(parsed.FullKey().ToString().c_str(), std::move(res)); + + this->QueueSendRequest(req); + + // Wait for the notification that indicates the tensor has been + // successfully transmitted to the remote process. Only needed if we + // have not parsed the tensor to proto + if (doOptimalTransfer) mpi_send_call->n_.WaitForNotification(); + }; // done_cb worker_env_2->compute_pool->Schedule([this, step_id, parsed, done_cb]() { this->RecvLocalAsync(step_id, parsed, done_cb); @@ -289,9 +298,8 @@ void MPIRendezvousMgr::MPIBackgroundThread() { } // Remove sends that have been completed - active_sends.remove_if([](std::unique_ptr& i) { - return i->IsFinished(); - }); + active_sends.remove_if( + [](std::unique_ptr& i) { return i->IsFinished(); }); // send a Tensor request RequestQueueEntry req; diff --git a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.h b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.h index b15748d63c9fdbc5134069b63fd998e46c499e16..5596601ddb9846c0e4f5be4bf33114fc19c0a59d 100644 --- a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.h +++ b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.h @@ -18,22 +18,24 @@ limitations under the License. #ifdef TENSORFLOW_USE_MPI -#include -#include #include -#include -#include #include +#include +#include +#include +#include #include #include #include #include +#include "tensorflow/contrib/mpi/mpi_msg.pb.h" #include "tensorflow/contrib/mpi/mpi_utils.h" #include "tensorflow/core/distributed_runtime/base_rendezvous_mgr.h" +#include "tensorflow/core/distributed_runtime/recent_request_ids.h" +#include "tensorflow/core/distributed_runtime/request_id.h" #include "tensorflow/core/distributed_runtime/worker_env.h" -#include "tensorflow/contrib/mpi/mpi_msg.pb.h" #include "tensorflow/core/protobuf/worker.pb.h" #define TAG_REQTENSOR 1010 @@ -104,6 +106,7 @@ class MPIRequestTensorCall { void Init(const Rendezvous::ParsedKey& parsed, const int64 step_id) { req_.set_step_id(step_id); req_.set_rendezvous_key(parsed.FullKey().data(), parsed.FullKey().size()); + req_.set_request_id(GetUniqueRequestId()); request_buffer_size_ = req_.ByteSize(); // request_buffer_ = new char[request_buffer_size_]; // req_.SerializeToArray(request_buffer_, request_buffer_size_); @@ -158,7 +161,8 @@ class MPIRendezvousMgr : public BaseRendezvousMgr { private: typedef std::function MPIRecvTensorCallBack; + const Tensor&, const bool, MPISendTensorCall*)> + MPIRecvTensorCallBack; typedef std::pair> RequestQueueEntry; typedef std::pair> @@ -177,6 +181,8 @@ class MPIRendezvousMgr : public BaseRendezvousMgr { std::map> recv_tensor_map_ GUARDED_BY(mrq_); + RecentRequestIds recv_tensor_recent_request_ids_; + void AddRequest(RecvTensorRequest, const int); void MPIBackgroundThread(); diff --git a/tensorflow/contrib/mpi/mpi_server_lib.cc b/tensorflow/contrib/mpi/mpi_server_lib.cc index d585c0565eb234655e7a1bbc92df5741e18c8f33..a31fa9ce0b3110d875689d74a41ca9f9cc85f532 100644 --- a/tensorflow/contrib/mpi/mpi_server_lib.cc +++ b/tensorflow/contrib/mpi/mpi_server_lib.cc @@ -22,8 +22,8 @@ limitations under the License. #include "grpc/support/alloc.h" -#include "tensorflow/core/distributed_runtime/server_lib.h" #include "tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.h" +#include "tensorflow/core/distributed_runtime/server_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/env.h" diff --git a/tensorflow/contrib/mpi/mpi_utils.h b/tensorflow/contrib/mpi/mpi_utils.h index 45e21f2b25ab4897641ffec776eb1b3c32ab9a2e..fa297c28cb47d43ba927ab941854bd472d90b465 100644 --- a/tensorflow/contrib/mpi/mpi_utils.h +++ b/tensorflow/contrib/mpi/mpi_utils.h @@ -18,8 +18,8 @@ limitations under the License. #ifdef TENSORFLOW_USE_MPI -#include #include +#include #include #include "tensorflow/core/lib/strings/str_util.h" diff --git a/tensorflow/contrib/mpi_collectives/BUILD b/tensorflow/contrib/mpi_collectives/BUILD index 11c5d6e776d6adbf7c439012027752e2235883ab..9f9802b8fe12356c0da82ebb2b48b565cf3f7319 100644 --- a/tensorflow/contrib/mpi_collectives/BUILD +++ b/tensorflow/contrib/mpi_collectives/BUILD @@ -6,20 +6,9 @@ package(default_visibility = [ licenses(["notice"]) # Apache 2.0 -filegroup( - name = "all_files", - srcs = glob( - ["**/*"], - exclude = [ - "**/METADATA", - "**/OWNERS", - ], - ), - visibility = ["//tensorflow:__subpackages__"], -) - load( "//tensorflow/core:platform/default/build_config.bzl", + "tf_additional_mpi_lib_defines", "tf_proto_library_cc", ) @@ -33,26 +22,98 @@ tf_proto_library_cc( ], ) -load("//tensorflow:tensorflow.bzl", "tf_custom_op_library") -load("//tensorflow:tensorflow.bzl", "tf_py_test") +cc_library( + name = "mpi_defines", + defines = tf_additional_mpi_lib_defines(), +) + +load( + "//tensorflow:tensorflow.bzl", + "tf_custom_op_py_library", + "tf_custom_op_library", + "tf_gen_op_wrapper_py", + "tf_gen_op_libs", + "tf_kernel_library", + "tf_py_test", +) tf_custom_op_library( - name = "mpi_collectives.so", + name = "python/ops/_mpi_ops.so", srcs = [ - "mpi_ops.cc", - "ring.cc", - "ring.h", + "kernels/mpi_ops.cc", + "kernels/ring.cc", + "kernels/ring.h", + "ops/mpi_ops.cc", ], gpu_srcs = [ - "ring.cu.cc", - "ring.h", + "kernels/ring.cu.cc", + "kernels/ring.h", ], deps = [ + ":mpi_defines", ":mpi_message_proto_cc", "//third_party/mpi", ], ) +tf_kernel_library( + name = "mpi_ops_kernels", + srcs = [ + "kernels/mpi_ops.cc", + "kernels/ring.cc", + ], + hdrs = [ + "kernels/ring.h", + ], + gpu_srcs = [ + "kernels/ring.cu.cc", + ], + deps = [ + ":mpi_defines", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:gpu_headers_lib", + "//tensorflow/core:lib", + "//tensorflow/core:proto_text", + "//tensorflow/core:stream_executor", + ], + # TODO: Include? alwayslink = 1, +) + +tf_gen_op_libs( + op_lib_names = ["mpi_ops"], +) + +tf_gen_op_wrapper_py( + name = "mpi_ops", + deps = [":mpi_ops_op_lib"], +) + +tf_custom_op_py_library( + name = "mpi_collectives_py", + srcs = [ + "__init__.py", + "python/ops/mpi_ops.py", + ], + dso = [ + ":python/ops/_mpi_ops.so", + ], + kernels = [ + ":mpi_ops_kernels", + ":mpi_ops_op_lib", + ], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + ":mpi_ops", + "//tensorflow/contrib/util:util_py", + "//tensorflow/python:device", + "//tensorflow/python:framework_ops", + "//tensorflow/python:platform", + "//tensorflow/python:util", + ], +) + tf_py_test( name = "mpi_ops_test", srcs = ["mpi_ops_test.py"], @@ -61,20 +122,19 @@ tf_py_test( "//tensorflow/python:platform", ], data = [ - ":mpi_collectives.so", + ":python/ops/_mpi_ops.so", ], tags = ["manual"], ) -py_library( - name = "mpi_ops_py", - srcs = [ - "__init__.py", - "mpi_ops.py", - ], - data = [ - ":mpi_collectives.so", - ], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], ) diff --git a/tensorflow/contrib/mpi_collectives/__init__.py b/tensorflow/contrib/mpi_collectives/__init__.py index 9ed16a6f078a506b60fd14f4356ff65a0a692203..52029cbc36a3bb77ea38d3973a75bfd37e93bfa4 100644 --- a/tensorflow/contrib/mpi_collectives/__init__.py +++ b/tensorflow/contrib/mpi_collectives/__init__.py @@ -37,7 +37,7 @@ for detecting the running MPI configuration. Example: ```python -from tensorflow.contrib import mpi +import tensorflow.contrib.mpi_collectives as mpi # Use `mpi.Session` instead of `tf.Session` with mpi.Session() as session: @@ -48,8 +48,10 @@ with mpi.Session() as session: print("MPI Size:", session.run(mpi.size())) ``` -@@rank +@@init @@size +@@rank +@@local_rank ### Ring Allreduce and Allgather @@ -123,12 +125,12 @@ from __future__ import print_function import tensorflow as tf -from tensorflow.contrib.mpi_collectives.mpi_ops import size -from tensorflow.contrib.mpi_collectives.mpi_ops import rank -from tensorflow.contrib.mpi_collectives.mpi_ops import local_rank -from tensorflow.contrib.mpi_collectives.mpi_ops import allgather -from tensorflow.contrib.mpi_collectives.mpi_ops import _allreduce -from tensorflow.contrib.mpi_collectives.mpi_ops import init +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import init +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import size +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import rank +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import local_rank +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import allgather +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import _allreduce def allreduce(tensor, average=True): diff --git a/tensorflow/contrib/mpi_collectives/mpi_ops.cc b/tensorflow/contrib/mpi_collectives/kernels/mpi_ops.cc similarity index 93% rename from tensorflow/contrib/mpi_collectives/mpi_ops.cc rename to tensorflow/contrib/mpi_collectives/kernels/mpi_ops.cc index a051ab0004626d034071112bb37671137ca5a3f0..8dca90a1e34d6a234c2b1479ca5594e88afcc194 100644 --- a/tensorflow/contrib/mpi_collectives/mpi_ops.cc +++ b/tensorflow/contrib/mpi_collectives/kernels/mpi_ops.cc @@ -21,7 +21,6 @@ limitations under the License. #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/framework/shape_inference.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/platform/mutex.h" @@ -36,8 +35,8 @@ limitations under the License. #define OMPI_SKIP_MPICXX #include "third_party/mpi/mpi.h" +#include "tensorflow/contrib/mpi_collectives/kernels/ring.h" #include "tensorflow/contrib/mpi_collectives/mpi_message.pb.h" -#include "tensorflow/contrib/mpi_collectives/ring.h" /* * MPI Allreduce and Allgather Ops for TensorFlow. @@ -81,7 +80,7 @@ using GPUDevice = Eigen::GpuDevice; namespace tensorflow { namespace contrib { -namespace mpi { +namespace mpi_collectives { // Make sure template specializations are generated in the ring.cu.cc and the // ring.cc file, not in this file. @@ -877,14 +876,6 @@ REGISTER_KERNEL_BUILDER(Name("MPIInit").Device(DEVICE_GPU), MPIInitOp); #endif -REGISTER_OP("MPIInit").Doc(R"doc( -Initialize MPI for the current process. - -If this is run on a GPU, then that GPU must be used for all future MPI -operations. If it is run on CPU, then all future MPI operations must also -run on CPU. -)doc"); - // Op to get the current MPI Size. template class MPISizeOp : public OpKernel { @@ -911,21 +902,6 @@ REGISTER_KERNEL_BUILDER(Name("MPISize").Device(DEVICE_GPU).HostMemory("size"), MPISizeOp); #endif -REGISTER_OP("MPISize") - .Output("size: int32") - .SetShapeFn([](shape_inference::InferenceContext* c) { - c->set_output(0, c->Scalar()); - return Status::OK(); - }) - .Doc(R"doc( -Returns the number of running MPI processes. - -More precisely, returns the number of MPI processes in the group associated -with the MPI_COMM_WORLD communicator. - -size: Size of the MPI group. -)doc"); - // Op to get the current MPI Rank. template class MPIRankOp : public OpKernel { @@ -952,21 +928,6 @@ REGISTER_KERNEL_BUILDER(Name("MPIRank").Device(DEVICE_GPU).HostMemory("rank"), MPIRankOp); #endif -REGISTER_OP("MPIRank") - .Output("rank: int32") - .SetShapeFn([](shape_inference::InferenceContext* c) { - c->set_output(0, c->Scalar()); - return Status::OK(); - }) - .Doc(R"doc( -Returns the index of the current process in the MPI group. - -More precisely, returns the rank of the calling process in the MPI_COMM_WORLD -communicator. - -rank: Rank of the calling process. -)doc"); - // Op to get the current local MPI Rank. template class MPILocalRankOp : public OpKernel { @@ -994,21 +955,6 @@ REGISTER_KERNEL_BUILDER( MPILocalRankOp); #endif -REGISTER_OP("MPILocalRank") - .Output("rank: int32") - .SetShapeFn([](shape_inference::InferenceContext* c) { - c->set_output(0, c->Scalar()); - return Status::OK(); - }) - .Doc(R"doc( -Returns the index of the current process in the node it is on. - -More precisely, returns the rank of the calling process in communicator that -only spans the MPI processes running on that node. - -rank: Rank of the calling process on the node it is on. -)doc"); - template class MPIAllreduceOp : public AsyncOpKernel { public: @@ -1083,28 +1029,6 @@ REGISTER_KERNEL_BUILDER(Name("MPIAllreduce").Device(DEVICE_GPU), MPIAllreduceOp); #endif -REGISTER_OP("MPIAllreduce") - .Attr("T: {int32, int64, float32}") - .Input("tensor: T") - .Output("sum: T") - .SetShapeFn([](shape_inference::InferenceContext* c) { - c->set_output(0, c->input(0)); - return Status::OK(); - }) - .Doc(R"doc( -Perform an MPI Allreduce on a tensor. All other processes that do a reduction -on a tensor with the same name must have the same dimension for that tensor. -Tensors are reduced with other tensors that have the same node name for the -allreduce. - -Arguments - tensor: A tensor to reduce. - -Output - sum: A tensor with the same shape as `tensor`, summed across all - MPI processes. -)doc"); - template class MPIAllgatherOp : public AsyncOpKernel { public: @@ -1192,34 +1116,6 @@ class MPIAllgatherOp : public AsyncOpKernel { } }; -REGISTER_OP("MPIAllgather") - .Attr("T: {int32, int64, float32}") - .Attr("S: {int64}") - .Input("tensor: T") - .Input("sizes: S") - .Output("gathered: T") - .SetShapeFn([](shape_inference::InferenceContext* c) { - shape_inference::ShapeHandle output; - TF_RETURN_IF_ERROR( - c->ReplaceDim(c->input(0), 0, c->UnknownDim(), &output)); - c->set_output(0, output); - return Status::OK(); - }) - .Doc(R"doc( -Perform an MPI Allgather on a tensor. All other processes that do a gather on a -tensor with the same name must have the same rank for that tensor, and have the -same dimension on all but the first dimension. - -Arguments - tensor: A tensor to gather. - sizes: A tensor containing the first-dimension sizes of tensors to be - gathered from other ranks - -Output - gathered: A tensor with the same shape as `tensor` except for the first - dimension, which is the sum of dimensions in `sizes`. -)doc"); - REGISTER_KERNEL_BUILDER( Name("MPIAllgather").Device(DEVICE_CPU).HostMemory("sizes"), MPIAllgatherOp); @@ -1229,7 +1125,7 @@ REGISTER_KERNEL_BUILDER( MPIAllgatherOp); #endif -} // namespace mpi +} // namespace mpi_collectives } // namespace contrib } // namespace tensorflow diff --git a/tensorflow/contrib/mpi_collectives/ring.cc b/tensorflow/contrib/mpi_collectives/kernels/ring.cc similarity index 96% rename from tensorflow/contrib/mpi_collectives/ring.cc rename to tensorflow/contrib/mpi_collectives/kernels/ring.cc index d93233eb210b80df10fd9c2c7975ce77112d18a2..8970ceb1a206ff2f9d6e18f7d19e313b8a036042 100644 --- a/tensorflow/contrib/mpi_collectives/ring.cc +++ b/tensorflow/contrib/mpi_collectives/kernels/ring.cc @@ -17,11 +17,11 @@ limitations under the License. #define EIGEN_USE_THREADS -#include "tensorflow/contrib/mpi_collectives/ring.h" +#include "tensorflow/contrib/mpi_collectives/kernels/ring.h" namespace tensorflow { namespace contrib { -namespace mpi { +namespace mpi_collectives { using CPUDevice = Eigen::ThreadPoolDevice; @@ -73,7 +73,7 @@ GENERATE_ACCUMULATE(long long); GENERATE_ACCUMULATE(float); #undef GENERATE_ACCUMULATE -} // namespace mpi +} // namespace mpi_collectives } // namespace contrib } // namespace tensorflow diff --git a/tensorflow/contrib/mpi_collectives/ring.cu.cc b/tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc similarity index 97% rename from tensorflow/contrib/mpi_collectives/ring.cu.cc rename to tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc index 2f3eef366a9a3c10e59cd5298fc1626e1094dff8..b04abde4694199d827a1738850bded9bf696d56c 100644 --- a/tensorflow/contrib/mpi_collectives/ring.cu.cc +++ b/tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc @@ -19,11 +19,11 @@ limitations under the License. #define EIGEN_USE_GPU -#include "tensorflow/contrib/mpi_collectives/ring.h" +#include "tensorflow/contrib/mpi_collectives/kernels/ring.h" namespace tensorflow { namespace contrib { -namespace mpi { +namespace mpi_collectives { using CPUDevice = Eigen::ThreadPoolDevice; @@ -109,7 +109,7 @@ GENERATE_ACCUMULATE(long long); GENERATE_ACCUMULATE(float); #undef GENERATE_ACCUMULATE -} // namespace mpi +} // namespace mpi_collectives } // namespace contrib } // namespace tensorflow #endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/mpi_collectives/ring.h b/tensorflow/contrib/mpi_collectives/kernels/ring.h similarity index 99% rename from tensorflow/contrib/mpi_collectives/ring.h rename to tensorflow/contrib/mpi_collectives/kernels/ring.h index cae57ce60eb09509af69f8ccab9eacedea361548..1d56d588bc49eda542303ae6ebb19602352ae01d 100644 --- a/tensorflow/contrib/mpi_collectives/ring.h +++ b/tensorflow/contrib/mpi_collectives/kernels/ring.h @@ -37,7 +37,7 @@ limitations under the License. namespace tensorflow { namespace contrib { -namespace mpi { +namespace mpi_collectives { using CPUDevice = Eigen::ThreadPoolDevice; using GPUDevice = Eigen::GpuDevice; @@ -317,7 +317,7 @@ Status RingAllgather(OpKernelContext* context, const Tensor* input, return Status::OK(); } -} // namespace mpi +} // namespace mpi_collectives } // namespace contrib } // namespace tensorflow diff --git a/tensorflow/contrib/mpi_collectives/mpi_message.proto b/tensorflow/contrib/mpi_collectives/mpi_message.proto index 7fa5e203010465766b8ab9562cac010de51a7bbc..afbce981ae1bdd5ae143ba5c45a4d9790a52fafc 100644 --- a/tensorflow/contrib/mpi_collectives/mpi_message.proto +++ b/tensorflow/contrib/mpi_collectives/mpi_message.proto @@ -15,7 +15,7 @@ limitations under the License. syntax = "proto3"; -package tensorflow.contrib.mpi; +package tensorflow.contrib.mpi_collectives; import "tensorflow/core/framework/tensor_shape.proto"; import "tensorflow/core/framework/types.proto"; diff --git a/tensorflow/contrib/mpi_collectives/ops/mpi_ops.cc b/tensorflow/contrib/mpi_collectives/ops/mpi_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..18e6bb61cffc6471412cb4c5141655839d7ddb3a --- /dev/null +++ b/tensorflow/contrib/mpi_collectives/ops/mpi_ops.cc @@ -0,0 +1,132 @@ +/* Copyright 2016 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. +==============================================================================*/ + +#ifdef TENSORFLOW_USE_MPI + +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" + +namespace tensorflow { +namespace contrib { +namespace mpi_collectives { + +REGISTER_OP("MPIInit").Doc(R"doc( +Initialize MPI for the current process. + +If this is run on a GPU, then that GPU must be used for all future MPI +operations. If it is run on CPU, then all future MPI operations must also +run on CPU. +)doc"); + +REGISTER_OP("MPISize") + .Output("size: int32") + .SetShapeFn([](shape_inference::InferenceContext* c) { + c->set_output(0, c->Scalar()); + return Status::OK(); + }) + .Doc(R"doc( +Returns the number of running MPI processes. + +More precisely, returns the number of MPI processes in the group associated +with the MPI_COMM_WORLD communicator. + +size: Size of the MPI group. +)doc"); + +REGISTER_OP("MPIRank") + .Output("rank: int32") + .SetShapeFn([](shape_inference::InferenceContext* c) { + c->set_output(0, c->Scalar()); + return Status::OK(); + }) + .Doc(R"doc( +Returns the index of the current process in the MPI group. + +More precisely, returns the rank of the calling process in the MPI_COMM_WORLD +communicator. + +rank: Rank of the calling process. +)doc"); + +REGISTER_OP("MPILocalRank") + .Output("rank: int32") + .SetShapeFn([](shape_inference::InferenceContext* c) { + c->set_output(0, c->Scalar()); + return Status::OK(); + }) + .Doc(R"doc( +Returns the index of the current process in the node it is on. + +More precisely, returns the rank of the calling process in communicator that +only spans the MPI processes running on that node. + +rank: Rank of the calling process on the node it is on. +)doc"); + +REGISTER_OP("MPIAllreduce") + .Attr("T: {int32, int64, float32}") + .Input("tensor: T") + .Output("sum: T") + .SetShapeFn([](shape_inference::InferenceContext* c) { + c->set_output(0, c->input(0)); + return Status::OK(); + }) + .Doc(R"doc( +Perform an MPI Allreduce on a tensor. All other processes that do a reduction +on a tensor with the same name must have the same dimension for that tensor. +Tensors are reduced with other tensors that have the same node name for the +allreduce. + +Arguments + tensor: A tensor to reduce. + +Output + sum: A tensor with the same shape as `tensor`, summed across all + MPI processes. +)doc"); + +REGISTER_OP("MPIAllgather") + .Attr("T: {int32, int64, float32}") + .Attr("S: {int64}") + .Input("tensor: T") + .Input("sizes: S") + .Output("gathered: T") + .SetShapeFn([](shape_inference::InferenceContext* c) { + shape_inference::ShapeHandle output; + TF_RETURN_IF_ERROR( + c->ReplaceDim(c->input(0), 0, c->UnknownDim(), &output)); + c->set_output(0, output); + return Status::OK(); + }) + .Doc(R"doc( +Perform an MPI Allgather on a tensor. All other processes that do a gather on a +tensor with the same name must have the same rank for that tensor, and have the +same dimension on all but the first dimension. + +Arguments + tensor: A tensor to gather. + sizes: A tensor containing the first-dimension sizes of tensors to be + gathered from other ranks + +Output + gathered: A tensor with the same shape as `tensor` except for the first + dimension, which is the sum of dimensions in `sizes`. +)doc"); + +} // namespace mpi_collectives +} // namespace contrib +} // namespace tensorflow + +#endif // TENSORFLOW_USE_MPI diff --git a/tensorflow/contrib/mpi_collectives/mpi_ops.py b/tensorflow/contrib/mpi_collectives/python/ops/mpi_ops.py similarity index 67% rename from tensorflow/contrib/mpi_collectives/mpi_ops.py rename to tensorflow/contrib/mpi_collectives/python/ops/mpi_ops.py index 81567cc688ac8666c3755d5f84162a6dff869107..2fbefef0d36f6a1507827427ebbafe5e81e35ea3 100644 --- a/tensorflow/contrib/mpi_collectives/mpi_ops.py +++ b/tensorflow/contrib/mpi_collectives/python/ops/mpi_ops.py @@ -20,43 +20,13 @@ from __future__ import print_function import tensorflow as tf -from tensorflow.python.framework import errors -from tensorflow.python.framework import load_library +from tensorflow.contrib.mpi_collectives.ops import gen_mpi_ops +from tensorflow.contrib.util import loader from tensorflow.python.framework import ops from tensorflow.python.platform import resource_loader -from tensorflow.python.platform import tf_logging as logging - -def _load_library(name, op_list=None): - """Loads a .so file containing the specified operators. - - Args: - name: The name of the .so file to load. - op_list: A list of names of operators that the library should have. If None - then the .so file's contents will not be verified. - - Raises: - NameError if one of the required ops is missing. - """ - try: - filename = resource_loader.get_path_to_datafile(name) - library = load_library.load_op_library(filename) - for expected_op in (op_list or []): - for lib_op in library.OP_LIST.op: - if lib_op.name == expected_op: - break - else: - raise NameError( - 'Could not find operator %s in dynamic library %s' % - (expected_op, name)) - return library - except errors.NotFoundError: - logging.warning('%s file could not be loaded.', name) - - -MPI_LIB = _load_library('mpi_collectives.so', ['MPISize', 'MPIRank', - 'MPILocalRank', 'MPIAllgather', - 'MPIAllreduce']) +_mpi_ops_so = loader.load_op_library( + resource_loader.get_path_to_datafile('_mpi_ops.so')) def size(name=None): @@ -68,7 +38,7 @@ def size(name=None): Returns: An integer scalar containing the number of MPI processes. """ - return MPI_LIB.mpi_size(name=name) + return gen_mpi_ops.mpi_size(name=name) ops.NotDifferentiable('MPISize') @@ -83,7 +53,7 @@ def rank(name=None): Returns: An integer scalar with the MPI rank of the calling process. """ - return MPI_LIB.mpi_rank(name=name) + return gen_mpi_ops.mpi_rank(name=name) ops.NotDifferentiable('MPIRank') @@ -95,7 +65,7 @@ def init(name=None): All future MPI ops must be run on the same device that the `init` op was run on. """ - return MPI_LIB.mpi_init(name=name) + return gen_mpi_ops.mpi_init(name=name) ops.NotDifferentiable('MPIInit') @@ -112,7 +82,7 @@ def local_rank(name=None): Returns: An integer scalar with the local MPI rank of the calling process. """ - return MPI_LIB.mpi_local_rank(name=name) + return gen_mpi_ops.mpi_local_rank(name=name) ops.NotDifferentiable('MPILocalRank') @@ -129,7 +99,7 @@ def _allreduce(tensor, name=None): A tensor of the same shape and type as `tensor`, summed across all processes. """ - return MPI_LIB.mpi_allreduce(tensor, name=name) + return gen_mpi_ops.mpi_allreduce(tensor, name=name) ops.NotDifferentiable('MPIAllreduce') @@ -151,15 +121,14 @@ def allgather(tensor, name=None): """ # Specify that first allgather is to collect the tensor gather sizes, # indicated by passing in a scalar (0-D tensor) of value 0 - sizes_flag = tf.constant(0, dtype=tf.int64, name="size_flag_const") - my_size = tf.slice(tf.shape(tensor, out_type=tf.int64), [0], [1], name="size_slice") + sizes_flag = tf.constant(0, dtype=tf.int64, name='size_flag_const') + my_size = tf.slice( + tf.shape(tensor, out_type=tf.int64), [0], [1], name='size_slice') if name is None: - name = "allgather" - sizing_name = "{}_sizing".format(name) - sizes = MPI_LIB.mpi_allgather(my_size, sizes_flag, name=sizing_name) - return MPI_LIB.mpi_allgather(tensor, sizes, name=name) + name = 'allgather' + sizing_name = '{}_sizing'.format(name) + sizes = gen_mpi_ops.mpi_allgather(my_size, sizes_flag, name=sizing_name) + return gen_mpi_ops.mpi_allgather(tensor, sizes, name=name) ops.NotDifferentiable('MPIAllgather') - - diff --git a/tensorflow/contrib/nccl/BUILD b/tensorflow/contrib/nccl/BUILD index df9dbb457ace32ab804f7fc736a23f5b08bd077a..5ac96007df7ee08b1e32aacd28f83768859810a9 100644 --- a/tensorflow/contrib/nccl/BUILD +++ b/tensorflow/contrib/nccl/BUILD @@ -23,15 +23,17 @@ load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") tf_custom_op_library( name = "python/ops/_nccl_ops.so", srcs = [ + "ops/nccl_ops.cc", + ], + gpu_srcs = [ "kernels/nccl_manager.cc", "kernels/nccl_manager.h", "kernels/nccl_ops.cc", - "ops/nccl_ops.cc", ], - deps = [ - "//tensorflow/core:gpu_headers_lib", + deps = if_cuda([ "@nccl_archive//:nccl", - ], + "//tensorflow/core:gpu_headers_lib", + ]), ) tf_cuda_cc_test( @@ -52,17 +54,14 @@ tf_cuda_cc_test( "no_oss", "notap", ], - deps = if_cuda( + deps = [ - "@nccl_archive//:nccl", "//tensorflow/core:cuda", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + "@nccl_archive//:nccl", ], - [], - ) + [ - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core:testlib", - ], ) tf_kernel_library( @@ -103,11 +102,8 @@ tf_custom_op_py_library( "__init__.py", "python/ops/nccl_ops.py", ], - dso = [ - ":python/ops/_nccl_ops.so", - ], - kernels = [ - ":nccl_kernels", + dso = [":python/ops/_nccl_ops.so"], + kernels = if_cuda([":nccl_kernels"]) + [ ":nccl_ops_op_lib", ], srcs_version = "PY2AND3", diff --git a/tensorflow/contrib/nccl/kernels/nccl_manager.h b/tensorflow/contrib/nccl/kernels/nccl_manager.h index cb1719c3be6a5c042db6e258d68663e70bfbfa15..bb219e0edc8a2c4ba0ce0583cbe4018a4fa3a1d1 100644 --- a/tensorflow/contrib/nccl/kernels/nccl_manager.h +++ b/tensorflow/contrib/nccl/kernels/nccl_manager.h @@ -12,8 +12,8 @@ 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 THIRD_PARTY_TENSORFLOW_CORE_KERNELS_NCCL_COMMUNICATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CORE_KERNELS_NCCL_COMMUNICATOR_H_ +#ifndef TENSORFLOW_CORE_KERNELS_NCCL_COMMUNICATOR_H_ +#define TENSORFLOW_CORE_KERNELS_NCCL_COMMUNICATOR_H_ #ifdef GOOGLE_CUDA @@ -136,4 +136,4 @@ class NcclManager { #endif // GOOGLE_CUDA -#endif // THIRD_PARTY_TENSORFLOW_CORE_KERNELS_NCCL_COMMUNICATOR_H_ +#endif // TENSORFLOW_CORE_KERNELS_NCCL_COMMUNICATOR_H_ diff --git a/tensorflow/contrib/nccl/kernels/nccl_manager_test.cc b/tensorflow/contrib/nccl/kernels/nccl_manager_test.cc index 658baf18d3c706d3b7fbba1ec9d02a1f0cda638e..985b2bae2566c38dfb2c71a899e4b03bbb8fa55d 100644 --- a/tensorflow/contrib/nccl/kernels/nccl_manager_test.cc +++ b/tensorflow/contrib/nccl/kernels/nccl_manager_test.cc @@ -16,6 +16,7 @@ limitations under the License. #ifdef GOOGLE_CUDA #include +#include #include #include "tensorflow/contrib/nccl/kernels/nccl_manager.h" @@ -239,7 +240,8 @@ TEST_F(NcclManagerTest, MultipleCallers) { SE_ASSERT_OK(stream->BlockHostUntilDone()); } - std::random_shuffle(case_and_device_num.begin(), case_and_device_num.end()); + std::shuffle(case_and_device_num.begin(), case_and_device_num.end(), + std::mt19937(std::random_device()())); mutex mu; // guards case_and_device_num. std::unique_ptr pool( diff --git a/tensorflow/contrib/nccl/python/ops/nccl_ops_test.py b/tensorflow/contrib/nccl/python/ops/nccl_ops_test.py index bad0abd44cc507c6ebbe4481f80b8cafd8480322..98fe394c5b38294700617591992d3207b0a4706b 100644 --- a/tensorflow/contrib/nccl/python/ops/nccl_ops_test.py +++ b/tensorflow/contrib/nccl/python/ops/nccl_ops_test.py @@ -77,10 +77,6 @@ class NcclTestCase(test.TestCase): # same communicator across multiple sessions. with self.test_session(use_gpu=True) as sess: - # Check GPU availability *after* creating test session, see b/68975239. - if not test.is_gpu_available(): - return # Test requires access to a GPU - for devices in device_sets: shape = (3, 4) random = (np.random.random_sample(shape) - .5) * 1024 @@ -100,6 +96,11 @@ class NcclTestCase(test.TestCase): result_tensors = [array_ops.identity(t) for t in reduce_tensors] + # Check GPU availability *after* creating session, see b/68975239. + if not test.is_gpu_available(): + # If no GPU is available, only test graph construction. + continue + # Test execution and results. for t in sess.run(result_tensors): self.assertAllClose(t, np_ans) @@ -114,6 +115,7 @@ class NcclTestCase(test.TestCase): numpy_fn: A function taking two tensors and returning the gradient of the reduction of the two. """ + def _Gradient(tensors, devices): inputs = [array_ops.placeholder(t.dtype, t.shape) for t in tensors] reduce_tensors = nccl_reduce(inputs, devices) @@ -164,12 +166,17 @@ class BroadcastTest(NcclTestCase): (['/device:GPU:0', '/device:GPU:0'],)) def testBroadcastToCpuError(self): - # Broadcasts to CPU is not supported. - with self.assertRaisesRegexp( - errors.NotFoundError, - "No registered '_NcclBroadcastRecv' OpKernel for CPU devices"): + try: + # Broadcasts to CPU is not supported. self._Test(_NcclBroadcast, lambda x, y: x, (['/device:GPU:0', '/device:CPU:0'],)) + except errors.NotFoundError as e: + self.assertRegexpMatches( + str(e), "No registered '_NcclBroadcastRecv' OpKernel for CPU devices") + else: + # Session isn't executed when no GPU is available. + if test.is_gpu_available(): + self.fail("Didn't raise NotFoundError trying to broadcast to CPU") class CombinedTest(NcclTestCase): diff --git a/tensorflow/contrib/ndlstm/BUILD b/tensorflow/contrib/ndlstm/BUILD deleted file mode 100644 index 8403f841884d4640ce8156ff4db46868dbe1788c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/BUILD +++ /dev/null @@ -1,92 +0,0 @@ -# Description: -# Contains classes implementing 1D and 2D LSTMs for image and signal -# processing problems. - -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - -package(default_visibility = ["//tensorflow:__subpackages__"]) - -load("//tensorflow:tensorflow.bzl", "tf_py_test") - -py_library( - name = "ndlstm", - srcs = [ - "__init__.py", - "python/__init__.py", - "python/lstm1d.py", - "python/lstm2d.py", - "python/misc.py", - ], - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/contrib/framework:framework_py", - "//tensorflow/contrib/layers:layers_py", - "//tensorflow/contrib/rnn:rnn_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:framework", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:math_ops", - "//tensorflow/python:nn_ops", - "//tensorflow/python:ops", - "//tensorflow/python:platform", - "//tensorflow/python:random_ops", - "//tensorflow/python:rnn", - "//tensorflow/python:rnn_cell", - "//tensorflow/python:sparse_ops", - "//tensorflow/python:training", - "//tensorflow/python:variable_scope", - ], -) - -tf_py_test( - name = "lstm1d_test", - srcs = ["python/lstm1d_test.py"], - additional_deps = [ - ":ndlstm", - "//third_party/py/numpy", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:gradients", - "//tensorflow/python:variables", - ], -) - -tf_py_test( - name = "lstm2d_test", - srcs = ["python/lstm2d_test.py"], - additional_deps = [ - ":ndlstm", - "//third_party/py/numpy", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:variables", - ], -) - -tf_py_test( - name = "misc_test", - srcs = ["python/misc_test.py"], - additional_deps = [ - ":ndlstm", - "//third_party/py/numpy", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:variables", - ], -) - -filegroup( - name = "all_files", - srcs = glob( - ["**/*"], - exclude = [ - "**/METADATA", - "**/OWNERS", - ], - ), - visibility = ["//tensorflow:__subpackages__"], -) diff --git a/tensorflow/contrib/ndlstm/README.md b/tensorflow/contrib/ndlstm/README.md deleted file mode 100644 index 7ccb57f1b34a24af7d776f7dbb12a2a00bb5ca30..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/README.md +++ /dev/null @@ -1,31 +0,0 @@ -Library of multidimensional LSTM models and related code. - -# 2D LSTM code - -The 2D LSTM layers take tensors of the form (batch_size, height, width, -depth), compatible with convolutional layers, as inputs. The library -transposes and reshapes these tensors in a way that allows batches of -images to be processed by LSTMs. - -The library currently provides: - - - a separable 2D LSTM layer - - a simple 2D convolutional layer that can be swapped out against 2D LSTM - - layers to reduce images to sequences and images to final state vectors - - layers for sequence classification, pixel-wise classification - -# Other Dimensions - -There is 1D LSTM code in `lstm1d.py`. This code implements 1D LSTM versions -suitable as a basis for higher dimensional LSTMs. It is intended for constant -batch size and uses a different layout. Although the code is perfectly fine for -1D use, you may find other 1D LSTM implementations to be more convenient if you -are interested in sequence problems. - -# Upcoming Changes - - - PyramidLSTM - - support for 3D and 4D - - optional use of native fused LSTM op - - easy-to-use command line drivers and examples - - operators for patch-wise processing diff --git a/tensorflow/contrib/ndlstm/python/lstm1d.py b/tensorflow/contrib/ndlstm/python/lstm1d.py deleted file mode 100644 index d3c3531f405a74d89ce736dae0134939e189f7ae..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/python/lstm1d.py +++ /dev/null @@ -1,192 +0,0 @@ -# Copyright 2016 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. -# ============================================================================== -"""LSTM layers for sequences.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from six.moves import xrange # pylint: disable=redefined-builtin -from tensorflow.contrib.framework.python.ops import variables -from tensorflow.python.framework import constant_op -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nn_ops -from tensorflow.python.ops import random_ops -from tensorflow.python.ops import rnn -from tensorflow.python.ops import rnn_cell -from tensorflow.python.ops import variable_scope - - -def _shape(tensor): - return tensor.get_shape().as_list() - - -def ndlstm_base_unrolled(inputs, noutput, scope=None, reverse=False): - """Run an LSTM, either forward or backward. - - This is a 1D LSTM implementation using unrolling and the TensorFlow - LSTM op. - - Args: - inputs: input sequence (length, batch_size, ninput) - noutput: depth of output - scope: optional scope name - reverse: run LSTM in reverse - - Returns: - Output sequence (length, batch_size, noutput) - - """ - with variable_scope.variable_scope(scope, "SeqLstmUnrolled", [inputs]): - length, batch_size, _ = _shape(inputs) - lstm_cell = rnn_cell.BasicLSTMCell(noutput, state_is_tuple=False) - state = array_ops.zeros([batch_size, lstm_cell.state_size]) - output_u = [] - inputs_u = array_ops.unstack(inputs) - if reverse: - inputs_u = list(reversed(inputs_u)) - for i in xrange(length): - if i > 0: - variable_scope.get_variable_scope().reuse_variables() - output, state = lstm_cell(inputs_u[i], state) - output_u += [output] - if reverse: - output_u = list(reversed(output_u)) - outputs = array_ops.stack(output_u) - return outputs - - -def ndlstm_base_dynamic(inputs, noutput, scope=None, reverse=False): - """Run an LSTM, either forward or backward. - - This is a 1D LSTM implementation using dynamic_rnn and - the TensorFlow LSTM op. - - Args: - inputs: input sequence (length, batch_size, ninput) - noutput: depth of output - scope: optional scope name - reverse: run LSTM in reverse - - Returns: - Output sequence (length, batch_size, noutput) - """ - with variable_scope.variable_scope(scope, "SeqLstm", [inputs]): - # TODO(tmb) make batch size, sequence_length dynamic - # example: sequence_length = tf.shape(inputs)[0] - _, batch_size, _ = _shape(inputs) - lstm_cell = rnn_cell.BasicLSTMCell(noutput, state_is_tuple=False) - state = array_ops.zeros([batch_size, lstm_cell.state_size]) - sequence_length = int(inputs.get_shape()[0]) - sequence_lengths = math_ops.to_int64( - array_ops.fill([batch_size], sequence_length)) - if reverse: - inputs = array_ops.reverse_v2(inputs, [0]) - outputs, _ = rnn.dynamic_rnn( - lstm_cell, inputs, sequence_lengths, state, time_major=True) - if reverse: - outputs = array_ops.reverse_v2(outputs, [0]) - return outputs - - -def ndlstm_base(inputs, noutput, scope=None, reverse=False, dynamic=True): - """Implements a 1D LSTM, either forward or backward. - - This is a base case for multidimensional LSTM implementations, which - tend to be used differently from sequence-to-sequence - implementations. For general 1D sequence to sequence - transformations, you may want to consider another implementation - from TF slim. - - Args: - inputs: input sequence (length, batch_size, ninput) - noutput: depth of output - scope: optional scope name - reverse: run LSTM in reverse - dynamic: use dynamic_rnn - - Returns: - Output sequence (length, batch_size, noutput) - - """ - # TODO(tmb) maybe add option for other LSTM implementations, like - # slim.rnn.basic_lstm_cell - if dynamic: - return ndlstm_base_dynamic(inputs, noutput, scope=scope, reverse=reverse) - else: - return ndlstm_base_unrolled(inputs, noutput, scope=scope, reverse=reverse) - - -def sequence_to_final(inputs, noutput, scope=None, name=None, reverse=False): - """Run an LSTM across all steps and returns only the final state. - - Args: - inputs: (length, batch_size, depth) tensor - noutput: size of output vector - scope: optional scope name - name: optional name for output tensor - reverse: run in reverse - - Returns: - Batch of size (batch_size, noutput). - """ - with variable_scope.variable_scope(scope, "SequenceToFinal", [inputs]): - length, batch_size, _ = _shape(inputs) - lstm = rnn_cell.BasicLSTMCell(noutput, state_is_tuple=False) - state = array_ops.zeros([batch_size, lstm.state_size]) - inputs_u = array_ops.unstack(inputs) - if reverse: - inputs_u = list(reversed(inputs_u)) - for i in xrange(length): - if i > 0: - variable_scope.get_variable_scope().reuse_variables() - output, state = lstm(inputs_u[i], state) - outputs = array_ops.reshape(output, [batch_size, noutput], name=name) - return outputs - - -def sequence_softmax(inputs, noutput, scope=None, name=None, linear_name=None): - """Run a softmax layer over all the time steps of an input sequence. - - Args: - inputs: (length, batch_size, depth) tensor - noutput: output depth - scope: optional scope name - name: optional name for output tensor - linear_name: name for linear (pre-softmax) output - - Returns: - A tensor of size (length, batch_size, noutput). - - """ - length, _, ninputs = _shape(inputs) - inputs_u = array_ops.unstack(inputs) - output_u = [] - with variable_scope.variable_scope(scope, "SequenceSoftmax", [inputs]): - initial_w = random_ops.truncated_normal([0 + ninputs, noutput], stddev=0.1) - initial_b = constant_op.constant(0.1, shape=[noutput]) - w = variables.model_variable("weights", initializer=initial_w) - b = variables.model_variable("biases", initializer=initial_b) - for i in xrange(length): - with variable_scope.variable_scope(scope, "SequenceSoftmaxStep", - [inputs_u[i]]): - # TODO(tmb) consider using slim.fully_connected(..., - # activation_fn=tf.nn.softmax) - linear = nn_ops.xw_plus_b(inputs_u[i], w, b, name=linear_name) - output = nn_ops.softmax(linear) - output_u += [output] - outputs = array_ops.stack(output_u, name=name) - return outputs diff --git a/tensorflow/contrib/ndlstm/python/lstm1d_test.py b/tensorflow/contrib/ndlstm/python/lstm1d_test.py deleted file mode 100644 index 49b15cc814cc54aaea7c67c4e509e5aa144e063e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/python/lstm1d_test.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright 2016 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 1D LSTM.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -from tensorflow.contrib.ndlstm.python import lstm1d as lstm1d_lib -from tensorflow.python.framework import constant_op -from tensorflow.python.ops import gradient_checker -from tensorflow.python.ops import gradients_impl -from tensorflow.python.ops import variables -from tensorflow.python.platform import test - -lstm1d = lstm1d_lib - - -def _rand(*size): - return np.random.uniform(size=size).astype("f") - - -class Lstm1DTest(test.TestCase): - - def testSequenceToSequenceDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(17, 1, 5)) - outputs = lstm1d.ndlstm_base(inputs, 8) - variables.global_variables_initializer().run() - names = [v.name for v in variables.trainable_variables()] - self.assertEqual(len(names), 2) - result = outputs.eval() - self.assertEqual(tuple(result.shape), (17, 1, 8)) - - def testSequenceToSequenceGradient(self): - with self.test_session(): - size = (17, 1, 15) - output_size = (17, 1, 8) - inputs = constant_op.constant(_rand(*size)) - outputs = lstm1d.ndlstm_base(inputs, 8, dynamic=False) - variables.global_variables_initializer().run() - gradients = gradients_impl.gradients(outputs, inputs) - if 1: # pylint: disable=using-constant-test - gradients = gradients_impl.gradients(outputs, inputs)[0].eval() - self.assertEqual(gradients.shape, size) - else: - # TODO(tmb) tf.test.compute_gradient error is currently broken - # with dynamic_rnn. Enable this test case eventually. - err = gradient_checker.compute_gradient_error( - inputs, size, outputs, output_size, delta=1e-4) - self.assert_(not np.isnan(err)) - self.assert_(err < 0.1) - - def testSequenceToSequenceGradientReverse(self): - with self.test_session(): - size = (17, 1, 15) - output_size = (17, 1, 8) - inputs = constant_op.constant(_rand(*size)) - outputs = lstm1d.ndlstm_base(inputs, 8, reverse=1, dynamic=False) - variables.global_variables_initializer().run() - if 1: # pylint: disable=using-constant-test - gradients = gradients_impl.gradients(outputs, inputs)[0].eval() - self.assertEqual(gradients.shape, size) - else: - # TODO(tmb) tf.test.compute_gradient error is currently broken - # with dynamic_rnn. Enable this test case eventually. - err = gradient_checker.compute_gradient_error( - inputs, size, outputs, output_size, delta=1e-4) - self.assert_(not np.isnan(err)) - self.assert_(err < 0.1) - - def testSequenceToFinalDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(17, 6, 5)) - outputs = lstm1d.sequence_to_final(inputs, 8) - variables.global_variables_initializer().run() - names = [v.name for v in variables.trainable_variables()] - self.assertEqual(len(names), 2) - result = outputs.eval() - self.assertEqual(tuple(result.shape), (6, 8)) - - def testSequenceSoftmaxDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(17, 1, 5)) - outputs = lstm1d.sequence_softmax(inputs, 8) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (17, 1, 8)) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/ndlstm/python/lstm2d.py b/tensorflow/contrib/ndlstm/python/lstm2d.py deleted file mode 100644 index ebbb4ccf11b219e86578d05e99a7a02ebe08271e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/python/lstm2d.py +++ /dev/null @@ -1,213 +0,0 @@ -# Copyright 2016 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. -# ============================================================================== -"""A small library of functions dealing with LSTMs applied to images. - -Tensors in this library generally have the shape (num_images, height, width, -depth). -""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.ndlstm.python import lstm1d -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import variable_scope - - -def _shape(tensor): - """Get the shape of a tensor as an int list.""" - return tensor.get_shape().as_list() - - -def images_to_sequence(tensor): - """Convert a batch of images into a batch of sequences. - - Args: - tensor: a (num_images, height, width, depth) tensor - - Returns: - (width, num_images*height, depth) sequence tensor - """ - - num_image_batches, height, width, depth = _shape(tensor) - transposed = array_ops.transpose(tensor, [2, 0, 1, 3]) - return array_ops.reshape(transposed, - [width, num_image_batches * height, depth]) - - -def sequence_to_images(tensor, num_image_batches): - """Convert a batch of sequences into a batch of images. - - Args: - tensor: (num_steps, num_batches, depth) sequence tensor - num_image_batches: the number of image batches - - Returns: - (num_images, height, width, depth) tensor - """ - - width, num_batches, depth = _shape(tensor) - height = num_batches // num_image_batches - reshaped = array_ops.reshape(tensor, - [width, num_image_batches, height, depth]) - return array_ops.transpose(reshaped, [1, 2, 0, 3]) - - -def horizontal_lstm(images, num_filters_out, scope=None): - """Run an LSTM bidirectionally over all the rows of each image. - - Args: - images: (num_images, height, width, depth) tensor - num_filters_out: output depth - scope: optional scope name - - Returns: - (num_images, height, width, num_filters_out) tensor, where - num_steps is width and new num_batches is num_image_batches * height - """ - with variable_scope.variable_scope(scope, "HorizontalLstm", [images]): - batch_size, _, _, _ = _shape(images) - sequence = images_to_sequence(images) - with variable_scope.variable_scope("lr"): - hidden_sequence_lr = lstm1d.ndlstm_base(sequence, num_filters_out // 2) - with variable_scope.variable_scope("rl"): - hidden_sequence_rl = (lstm1d.ndlstm_base( - sequence, num_filters_out - num_filters_out // 2, reverse=1)) - output_sequence = array_ops.concat([hidden_sequence_lr, hidden_sequence_rl], - 2) - output = sequence_to_images(output_sequence, batch_size) - return output - - -def get_blocks(images, kernel_size): - """Split images in blocks - - Args: - images: (num_images, height, width, depth) tensor - kernel_size: A list of length 2 holding the [kernel_height, kernel_width] of - of the pooling. Can be an int if both values are the same. - - Returns: - (num_images, height/kernel_height, width/kernel_width, - depth*kernel_height*kernel_width) tensor - """ - with variable_scope.variable_scope("image_blocks"): - batch_size, height, width, chanels = _shape(images) - - if height % kernel_size[0] != 0: - offset = array_ops.zeros([batch_size, - kernel_size[0] - (height % kernel_size[0]), - width, - chanels]) - images = array_ops.concat([images, offset], 1) - batch_size, height, width, chanels = _shape(images) - if width % kernel_size[1] != 0: - offset = array_ops.zeros([batch_size, - height, - kernel_size[1] - (width % kernel_size[1]), - chanels]) - images = array_ops.concat([images, offset], 2) - batch_size, height, width, chanels = _shape(images) - - h, w = int(height / kernel_size[0]), int(width / kernel_size[1]) - features = kernel_size[1] * kernel_size[0] * chanels - - lines = array_ops.split(images, h, axis=1) - line_blocks = [] - for line in lines: - line = array_ops.transpose(line, [0, 2, 3, 1]) - line = array_ops.reshape(line, [batch_size, w, features]) - line_blocks.append(line) - - return array_ops.stack(line_blocks, axis=1) - - -def separable_lstm(images, num_filters_out, - kernel_size=None, nhidden=None, scope=None): - """Run bidirectional LSTMs first horizontally then vertically. - - Args: - images: (num_images, height, width, depth) tensor - num_filters_out: output layer depth - kernel_size: A list of length 2 holding the [kernel_height, kernel_width] of - of the pooling. Can be an int if both values are the same. Set to None for - not using blocks - nhidden: hidden layer depth - scope: optional scope name - - Returns: - (num_images, height/kernel_height, width/kernel_width, - num_filters_out) tensor - """ - with variable_scope.variable_scope(scope, "SeparableLstm", [images]): - if nhidden is None: - nhidden = num_filters_out - if kernel_size is not None: - images = get_blocks(images, kernel_size) - hidden = horizontal_lstm(images, nhidden) - with variable_scope.variable_scope("vertical"): - transposed = array_ops.transpose(hidden, [0, 2, 1, 3]) - output_transposed = horizontal_lstm(transposed, num_filters_out) - output = array_ops.transpose(output_transposed, [0, 2, 1, 3]) - return output - - -def reduce_to_sequence(images, num_filters_out, scope=None): - """Reduce an image to a sequence by scanning an LSTM vertically. - - Args: - images: (num_images, height, width, depth) tensor - num_filters_out: output layer depth - scope: optional scope name - - Returns: - A (width, num_images, num_filters_out) sequence. - """ - with variable_scope.variable_scope(scope, "ReduceToSequence", [images]): - batch_size, height, width, depth = _shape(images) - transposed = array_ops.transpose(images, [1, 0, 2, 3]) - reshaped = array_ops.reshape(transposed, - [height, batch_size * width, depth]) - reduced = lstm1d.sequence_to_final(reshaped, num_filters_out) - output = array_ops.reshape(reduced, [batch_size, width, num_filters_out]) - return output - - -def reduce_to_final(images, num_filters_out, nhidden=None, scope=None): - """Reduce an image to a final state by running two LSTMs. - - Args: - images: (num_images, height, width, depth) tensor - num_filters_out: output layer depth - nhidden: hidden layer depth (defaults to num_filters_out) - scope: optional scope name - - Returns: - A (num_images, num_filters_out) batch. - """ - with variable_scope.variable_scope(scope, "ReduceToFinal", [images]): - nhidden = nhidden or num_filters_out - batch_size, height, width, depth = _shape(images) - transposed = array_ops.transpose(images, [1, 0, 2, 3]) - reshaped = array_ops.reshape(transposed, - [height, batch_size * width, depth]) - with variable_scope.variable_scope("reduce1"): - reduced = lstm1d.sequence_to_final(reshaped, nhidden) - transposed_hidden = array_ops.reshape(reduced, - [batch_size, width, nhidden]) - hidden = array_ops.transpose(transposed_hidden, [1, 0, 2]) - with variable_scope.variable_scope("reduce2"): - output = lstm1d.sequence_to_final(hidden, num_filters_out) - return output diff --git a/tensorflow/contrib/ndlstm/python/lstm2d_test.py b/tensorflow/contrib/ndlstm/python/lstm2d_test.py deleted file mode 100644 index f1b37d701b868438dcbac4e713ccc2136dacd983..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/python/lstm2d_test.py +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright 2016 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 2D LSTMs.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -from tensorflow.contrib.ndlstm.python import lstm2d as lstm2d_lib -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import test_util -from tensorflow.python.ops import variables -from tensorflow.python.platform import test - -lstm2d = lstm2d_lib - - -def _rand(*size): - return np.random.uniform(size=size).astype("f") - - -class Lstm2DTest(test_util.TensorFlowTestCase): - - def testImagesToSequenceDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(2, 7, 11, 5)) - outputs = lstm2d.images_to_sequence(inputs) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (11, 14, 5)) - - def testSequenceToImagesDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(11, 14, 5)) - outputs = lstm2d.sequence_to_images(inputs, 2) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (2, 7, 11, 5)) - - def testImagesAndSequenceDims(self): - with self.test_session(): - size = (2, 7, 11, 5) - inputs = constant_op.constant(_rand(*size)) - sequence = lstm2d.images_to_sequence(inputs) - outputs = lstm2d.sequence_to_images(sequence, size[0]) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), size) - - def testSeparableLstmDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(2, 7, 11, 5)) - outputs = lstm2d.separable_lstm(inputs, 8) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (2, 7, 11, 8)) - - def testSeparableLstmDimsBlocks(self): - with self.test_session(): - inputs = constant_op.constant(_rand(2, 7, 11, 5)) - outputs = lstm2d.separable_lstm(inputs, 8, kernel_size=[2, 2]) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (2, 4, 6, 8)) - - def testReduceToSequenceDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(2, 7, 11, 5)) - outputs = lstm2d.reduce_to_sequence(inputs, 8) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (2, 11, 8)) - - def testReduceToFinalDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(2, 7, 11, 5)) - outputs = lstm2d.reduce_to_final(inputs, 8, 12) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (2, 8)) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/ndlstm/python/misc.py b/tensorflow/contrib/ndlstm/python/misc.py deleted file mode 100644 index 38eeff84ca4e5afbe45d6c9e0c52af9ae86de24f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/python/misc.py +++ /dev/null @@ -1,99 +0,0 @@ -# Copyright 2016 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. -# ============================================================================== -"""Miscellaneous functions useful for nD-LSTM models. - -Some of these functions duplicate functionality in tfslim with -slightly different interfaces. - -Tensors in this library generally have the shape (num_images, height, width, -depth). -""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.layers.python.layers import layers -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import sparse_ops - - -def _shape(tensor): - """Get the shape of a tensor as an int list.""" - return tensor.get_shape().as_list() - - -def pixels_as_vector(images, scope=None): - """Reduce images to vectors by combining all pixels.""" - with ops.name_scope(scope, "PixelsAsVector", [images]): - batch_size, height, width, depth = _shape(images) - return array_ops.reshape(images, [batch_size, height * width * depth]) - - -def pool_as_vector(images, scope=None): - """Reduce images to vectors by averaging all pixels.""" - with ops.name_scope(scope, "PoolAsVector", [images]): - return math_ops.reduce_mean(images, [1, 2]) - - -def one_hot_planes(labels, num_classes, scope=None): - """Compute 1-hot encodings for planes. - - Given a label, this computes a label image that contains - 1 at all pixels in the plane corresponding to the target - class and 0 in all other planes. - - Args: - labels: (batch_size,) tensor - num_classes: number of classes - scope: optional scope name - - Returns: - Tensor of shape (batch_size, 1, 1, num_classes) with a 1-hot encoding. - """ - with ops.name_scope(scope, "OneHotPlanes", [labels]): - batch_size, = _shape(labels) - batched = layers.one_hot_encoding(labels, num_classes) - return array_ops.reshape(batched, [batch_size, 1, 1, num_classes]) - - -def one_hot_mask(labels, num_classes, scope=None): - """Compute 1-hot encodings for masks. - - Given a label image, this computes the one hot encoding at - each pixel. - - Args: - labels: (batch_size, width, height, 1) tensor containing labels. - num_classes: number of classes - scope: optional scope name - - Returns: - Tensor of shape (batch_size, width, height, num_classes) with - a 1-hot encoding. - """ - with ops.name_scope(scope, "OneHotMask", [labels]): - height, width, depth = _shape(labels) - assert depth == 1 - sparse_labels = math_ops.to_int32(array_ops.reshape(labels, [-1, 1])) - sparse_size, _ = _shape(sparse_labels) - indices = array_ops.reshape(math_ops.range(0, sparse_size, 1), [-1, 1]) - concated = array_ops.concat([indices, sparse_labels], 1) - dense_result = sparse_ops.sparse_to_dense(concated, - [sparse_size, num_classes], 1.0, - 0.0) - result = array_ops.reshape(dense_result, [height, width, num_classes]) - return result diff --git a/tensorflow/contrib/ndlstm/python/misc_test.py b/tensorflow/contrib/ndlstm/python/misc_test.py deleted file mode 100644 index fac9023da3b23b89a5494358c6e7ad82c12f9bdf..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ndlstm/python/misc_test.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright 2016 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. -# ============================================================================== -"""Miscellaneous tests.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -from tensorflow.contrib.ndlstm.python import misc as misc_lib -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import test_util -from tensorflow.python.ops import variables -from tensorflow.python.platform import test - -misc = misc_lib - - -def _rand(*size): - return np.random.uniform(size=size).astype("f") - - -class LstmMiscTest(test_util.TensorFlowTestCase): - - def testPixelsAsVectorDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(2, 7, 11, 5)) - outputs = misc.pixels_as_vector(inputs) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (2, 7 * 11 * 5)) - - def testPoolAsVectorDims(self): - with self.test_session(): - inputs = constant_op.constant(_rand(2, 7, 11, 5)) - outputs = misc.pool_as_vector(inputs) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (2, 5)) - - def testOneHotPlanes(self): - with self.test_session(): - inputs = constant_op.constant([0, 1, 3]) - outputs = misc.one_hot_planes(inputs, 4) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (3, 1, 1, 4)) - target = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]]) - self.assertAllClose(result.reshape(-1), target.reshape(-1)) - - def testOneHotMask(self): - with self.test_session(): - data = np.array([[0, 1, 2], [2, 0, 1]]).reshape(2, 3, 1) - inputs = constant_op.constant(data) - outputs = misc.one_hot_mask(inputs, 3) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (2, 3, 3)) - target = np.array([[[1, 0, 0], [0, 1, 0]], [[0, 1, 0], [0, 0, 1]], - [[0, 0, 1], [1, 0, 0]]]).transpose(1, 2, 0) - self.assertAllClose(result.reshape(-1), target.reshape(-1)) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/nearest_neighbor/kernels/heap.h b/tensorflow/contrib/nearest_neighbor/kernels/heap.h index 6e33a574e25d39a13a256383cbc9848fdb8b788f..a2dbb8052bfa1634d27c8b38a9bb6ca27fae42a2 100644 --- a/tensorflow/contrib/nearest_neighbor/kernels/heap.h +++ b/tensorflow/contrib/nearest_neighbor/kernels/heap.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HEAP_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HEAP_H_ +#ifndef TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HEAP_H_ +#define TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HEAP_H_ #include #include @@ -56,7 +56,7 @@ class HeapBase { // This method adds an element at the end of the internal array without // "heapifying" the array afterwards. This is useful for setting up a heap - // where a single call to heapify at the end of the inital insertion + // where a single call to heapify at the end of the initial insertion // operations suffices. void InsertUnsorted(const KeyType& key, const DataType& data) { if (v_.size() == static_cast(num_elements_)) { @@ -205,4 +205,4 @@ class AugmentedHeap : public HeapBase { } // namespace nearest_neighbor } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HEAP_H_ +#endif // TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HEAP_H_ diff --git a/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.cc b/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.cc index 2b412fac9a621f01bd21c6b4391da3c462dd78b3..13db6f62f525b6318687e3bf4b6499eee2c61ea8 100644 --- a/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.cc +++ b/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.cc @@ -75,7 +75,8 @@ class HyperplaneLSHProbesOp : public OpKernel { num_hyperplanes_per_table, ".")); OP_REQUIRES(context, num_hyperplanes_per_table <= 30, InvalidArgument("Need num_hyperplanes_per_table <= 30, got ", - num_hyperplanes_per_table, ". " + num_hyperplanes_per_table, + ". " "If you need more hyperplanes, change this Op" " to work for larger integer types (int64).")); @@ -88,12 +89,13 @@ class HyperplaneLSHProbesOp : public OpKernel { InvalidArgument("num_probes must be at least 1.")); int expected_num_hyperplanes = num_tables * num_hyperplanes_per_table; - OP_REQUIRES( - context, products_tensor.dim_size(1) == expected_num_hyperplanes, - InvalidArgument("Expected number of hyperplanes is ", - expected_num_hyperplanes, " but received ", - products_tensor.dim_size(1), " inner products per " - "point.")); + OP_REQUIRES(context, + products_tensor.dim_size(1) == expected_num_hyperplanes, + InvalidArgument("Expected number of hyperplanes is ", + expected_num_hyperplanes, " but received ", + products_tensor.dim_size(1), + " inner products per " + "point.")); auto products_eigen_tensor = products_tensor.matrix(); ConstMatrixMap products_matrix(products_eigen_tensor.data(), @@ -116,13 +118,11 @@ class HyperplaneLSHProbesOp : public OpKernel { // lschmidt's workstation. int64 cost_per_unit = 21 * num_hyperplanes_per_table * num_tables; if (num_probes > num_tables) { - cost_per_unit += 110 * num_hyperplanes_per_table - * (num_probes - num_tables); + cost_per_unit += + 110 * num_hyperplanes_per_table * (num_probes - num_tables); } context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor( - batch_size, - cost_per_unit, - [&](int64 start, int64 end) { + batch_size, cost_per_unit, [&](int64 start, int64 end) { HyperplaneMultiprobe multiprobe( num_hyperplanes_per_table, num_tables); diff --git a/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.h b/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.h index 1670e2f83b3afa10ca76b765bf97cc1c08038fba..c53205e1a4089c8bb5159621662496b798acf242 100644 --- a/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.h +++ b/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HYPERPLANE_LSH_PROBES_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HYPERPLANE_LSH_PROBES_H_ +#ifndef TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HYPERPLANE_LSH_PROBES_H_ +#define TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HYPERPLANE_LSH_PROBES_H_ #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" @@ -232,4 +232,4 @@ class HyperplaneMultiprobe { } // namespace nearest_neighbor } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HYPERPLANE_LSH_PROBES_H_ +#endif // TENSORFLOW_CONTRIB_NEAREST_NEIGHBOR_KERNELS_HYPERPLANE_LSH_PROBES_H_ diff --git a/tensorflow/contrib/nn/BUILD b/tensorflow/contrib/nn/BUILD index 56a24ac77f0b9a87b6e4db48cddacdf35f4855d0..5543eb6c6e3785978e9c878f309b9bd0863b0b0a 100644 --- a/tensorflow/contrib/nn/BUILD +++ b/tensorflow/contrib/nn/BUILD @@ -17,6 +17,7 @@ py_library( "python/ops/__init__.py", "python/ops/alpha_dropout.py", "python/ops/cross_entropy.py", + "python/ops/fwd_gradients.py", "python/ops/sampling_ops.py", "python/ops/scaled_softplus.py", ], @@ -28,6 +29,7 @@ py_library( "//tensorflow/python:embedding_ops", "//tensorflow/python:framework_ops", "//tensorflow/python:function", + "//tensorflow/python:gradients", "//tensorflow/python:math_ops", "//tensorflow/python:nn", "//tensorflow/python:nn_ops", @@ -55,6 +57,19 @@ py_test( ], ) +py_test( + name = "fwd_gradients_test", + size = "small", + srcs = ["python/ops/fwd_gradients_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":nn_py", + "//tensorflow/python:client_testlib", + "//tensorflow/python:constant_op", + "//tensorflow/python:math_ops", + ], +) + py_test( name = "sampling_ops_test", size = "small", diff --git a/tensorflow/contrib/nn/python/ops/alpha_dropout.py b/tensorflow/contrib/nn/python/ops/alpha_dropout.py index d7b61a584478f701726248a41c4992382189223d..2f92d05ba81f30a91f68f3c3ec51b6695d3d0371 100644 --- a/tensorflow/contrib/nn/python/ops/alpha_dropout.py +++ b/tensorflow/contrib/nn/python/ops/alpha_dropout.py @@ -18,7 +18,6 @@ from __future__ import print_function import numbers -from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import tensor_util @@ -26,7 +25,6 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nn_impl def alpha_dropout(x, keep_prob, noise_shape=None, seed=None, name=None): # pylint: disable=invalid-name diff --git a/tensorflow/contrib/nn/python/ops/alpha_dropout_test.py b/tensorflow/contrib/nn/python/ops/alpha_dropout_test.py index 2ff978ab89727c0ba2a8654013466838732377e4..54a98e6f142b7ba58c9418a8ac88269d38944aab 100644 --- a/tensorflow/contrib/nn/python/ops/alpha_dropout_test.py +++ b/tensorflow/contrib/nn/python/ops/alpha_dropout_test.py @@ -21,7 +21,6 @@ from __future__ import print_function from tensorflow.contrib.nn.python.ops.alpha_dropout import alpha_dropout from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import nn_impl diff --git a/tensorflow/contrib/nn/python/ops/fwd_gradients.py b/tensorflow/contrib/nn/python/ops/fwd_gradients.py new file mode 100644 index 0000000000000000000000000000000000000000..922497779b1d6ce426df9d7bb8fb343eea48502b --- /dev/null +++ b/tensorflow/contrib/nn/python/ops/fwd_gradients.py @@ -0,0 +1,76 @@ +# Copyright 2017 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. +# ============================================================================== +"""Forward-mode derivatives.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops.gradients_impl import gradients + + +def fwd_gradients(ys, xs, grad_xs=None, assert_unused=False): + """Computes forward-mode derivatives. + + This is accomplished in pure-python using tensorflow's existing (reverse-mode) + gradients. There is additional overhead on graph construction, but runtime + performance should be equal to a manual implementation [citation needed]. + + See https://j-towns.github.io/2017/06/12/A-new-trick.html and + https://github.com/HIPS/autograd/pull/175 for the original discussion of this + method, and https://github.com/renmengye/tensorflow-forward-ad for a "direct" + implementation. + + Args: + ys: A list of tensors. + xs: A list of tensors. + grad_xs: An optional list of tensors. If provided, must have the same length + and shapes compatible with xs. + assert_unused: Add assertions that intermediate values are not computed. + Returns: + A list of tensors of the same shapes as ys. The directional derivatives of + ys with respect to xs in the direction grad_xs. Leaving grad_xs unspecified + is equivalent to passing in 1s for each x in xs. + """ + # This version of forward-mode autodiff is based on code by Tim Cooijmans + # and handles list arguments and certain special cases such as when the + # ys doesn't depend on one or more of the xs, and when tf.IndexedSlices are + # generated by the first tf.gradients call. + + us = [array_ops.zeros_like(y) + float('nan') for y in ys] + + dydxs = gradients(ys, xs, grad_ys=us) + + # deal with strange types that tf.gradients returns but can't deal with + dydxs = [ops.convert_to_tensor(dydx) if isinstance(dydx, ops.IndexedSlices) + else dydx for dydx in dydxs] + + if assert_unused: + with ops.control_dependencies(dydxs): + assert_unused = control_flow_ops.Assert(False, [1], name='fwd_gradients') + with ops.control_dependencies([assert_unused]): + dydxs = array_ops.identity_n(dydxs) + + dydxs = [array_ops.zeros_like(x) if dydx is None else dydx + for x, dydx in zip(xs, dydxs)] + for x, dydx in zip(xs, dydxs): + dydx.set_shape(x.shape) + + dysdx = gradients(dydxs, us, grad_ys=grad_xs) + + return dysdx diff --git a/tensorflow/contrib/nn/python/ops/fwd_gradients_test.py b/tensorflow/contrib/nn/python/ops/fwd_gradients_test.py new file mode 100644 index 0000000000000000000000000000000000000000..56062c3cab32d727dd22a78d1f60c823a2f86a79 --- /dev/null +++ b/tensorflow/contrib/nn/python/ops/fwd_gradients_test.py @@ -0,0 +1,52 @@ +# Copyright 2017 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 forward_ad.py.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.nn.python.ops import fwd_gradients +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test + + +class ForwardAdTest(test.TestCase): + + def testSquare(self): + x = constant_op.constant(1.) + y = math_ops.square(x) + grad_x = 3. + + dydx_tf = fwd_gradients.fwd_gradients([y], [x], [grad_x])[0] + dydx_py = 2. * grad_x + + with self.test_session() as sess: + self.assertAllClose(sess.run(dydx_tf), dydx_py, 1e-6) + + def testGather(self): + x = constant_op.constant([1., 2., 3.]) + y = array_ops.gather(x, [0, 1]) + y.set_shape([2]) + dydx = fwd_gradients.fwd_gradients([y], [x], assert_unused=True) + + with self.test_session() as sess: + sess.run(dydx) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/opt/BUILD b/tensorflow/contrib/opt/BUILD index 9c961f2b9c828f7406516860b7e3fd3dc343d993..827279bd476f9666a972f43ad557fde6d0b6c59a 100644 --- a/tensorflow/contrib/opt/BUILD +++ b/tensorflow/contrib/opt/BUILD @@ -19,6 +19,7 @@ py_library( "python/training/elastic_average_optimizer.py", "python/training/external_optimizer.py", "python/training/lazy_adam_optimizer.py", + "python/training/model_average_optimizer.py", "python/training/moving_average_optimizer.py", "python/training/multitask_optimizer_wrapper.py", "python/training/nadam_optimizer.py", @@ -193,6 +194,27 @@ tf_py_test( ], ) +tf_py_test( + name = "model_average_optimizer_test", + srcs = ["python/training/model_average_optimizer_test.py"], + additional_deps = [ + ":opt_py", + "//tensorflow/python:client", + "//tensorflow/python:client_testlib", + "//tensorflow/python:array_ops", + "//tensorflow/python:variables", + "//tensorflow/python:framework", + "//tensorflow/python:platform", + "//tensorflow/python:training", + "//tensorflow/python:ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//third_party/py/numpy", + ], + tags = [ + "notap", # This test launches local server. + ], +) + py_test( name = "sign_decay_test", srcs = ["python/training/sign_decay_test.py"], diff --git a/tensorflow/contrib/opt/__init__.py b/tensorflow/contrib/opt/__init__.py index 90d2f924629800ccf26c160edd22c13b817f4584..6c1bb1adc096f5b8e6945ea1492727d16cf29e65 100644 --- a/tensorflow/contrib/opt/__init__.py +++ b/tensorflow/contrib/opt/__init__.py @@ -29,6 +29,7 @@ from tensorflow.contrib.opt.python.training.nadam_optimizer import * from tensorflow.contrib.opt.python.training.powersign import * from tensorflow.contrib.opt.python.training.variable_clipping_optimizer import * from tensorflow.contrib.opt.python.training.elastic_average_optimizer import * +from tensorflow.contrib.opt.python.training.model_average_optimizer import * # pylint: enable=wildcard-import from tensorflow.python.util.all_util import remove_undocumented @@ -48,7 +49,9 @@ _allowed_symbols = [ 'MultitaskOptimizerWrapper', 'clip_gradients_by_global_norm', 'ElasticAverageOptimizer', - 'ElasticAverageCustomGetter' + 'ElasticAverageCustomGetter', + 'ModelAverageOptimizer', + 'ModelAverageCustomGetter' ] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py b/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py index 6132cba1f5aecbafd8ca820ecda39355dd768847..5763593b81497f5d6945ff1e5d000042d295c093 100644 --- a/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py +++ b/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Wrapper optimizer for Elastic Average SGD """ from __future__ import absolute_import from __future__ import division @@ -78,23 +77,24 @@ class ElasticAverageCustomGetter(object): def __call__(self, getter, name, trainable, collections, *args, **kwargs): if trainable: with ops.device(self._worker_device): - local_var = getter(name, trainable=True, - collections=[ops.GraphKeys.LOCAL_VARIABLES], - *args, **kwargs) + local_var = getter( + name, + trainable=True, + collections=[ops.GraphKeys.LOCAL_VARIABLES], + *args, + **kwargs) global_center_variable = variable_scope.variable( - name='%s/%s' % - (GLOBAL_VARIABLE_NAME, - name), - initial_value=local_var.initialized_value(), - trainable=False, - collections=[ops.GraphKeys.GLOBAL_VARIABLES]) + name='%s/%s' % (GLOBAL_VARIABLE_NAME, name), + initial_value=local_var.initialized_value(), + trainable=False, + collections=[ops.GraphKeys.GLOBAL_VARIABLES]) with ops.device(self._worker_device): local_center_variable = variable_scope.variable( - name='%s/%s' % (LOCAL_VARIABLE_NAME, name), - initial_value=local_var.initialized_value(), - trainable=False, - collections=[ops.GraphKeys.LOCAL_VARIABLES]) + name='%s/%s' % (LOCAL_VARIABLE_NAME, name), + initial_value=local_var.initialized_value(), + trainable=False, + collections=[ops.GraphKeys.LOCAL_VARIABLES]) self._local_map[local_var] = local_center_variable self._global_map[local_var] = global_center_variable @@ -117,16 +117,15 @@ class ElasticAverageOptimizer(optimizer.Optimizer): # Default value as paper described BETA = 0.9 - def __init__( - self, - opt, - num_worker, - ea_custom_getter, - communication_period=10, - moving_rate=None, - rho=None, - use_locking=True, - name="ElasticAverageOptimizer"): + def __init__(self, + opt, + num_worker, + ea_custom_getter, + communication_period=10, + moving_rate=None, + rho=None, + use_locking=True, + name='ElasticAverageOptimizer'): """Construct a new gradient descent optimizer. Args: @@ -151,7 +150,7 @@ class ElasticAverageOptimizer(optimizer.Optimizer): self._global_map = ea_custom_getter._global_map if moving_rate is None: - self._moving_rate = BETA / communication_period / num_worker + self._moving_rate = self.BETA / communication_period / num_worker else: self._moving_rate = moving_rate if rho is None: @@ -160,13 +159,15 @@ class ElasticAverageOptimizer(optimizer.Optimizer): self._rho = rho self._local_step = variable_scope.get_variable( - initializer=0, - trainable=False, - collections=[ops.GraphKeys.LOCAL_VARIABLES], - name="local_step") + initializer=0, + trainable=False, + collections=[ops.GraphKeys.LOCAL_VARIABLES], + name='local_step') self._opt._prepare() - def compute_gradients(self, loss, var_list=None, + def compute_gradients(self, + loss, + var_list=None, gate_gradients=optimizer.Optimizer.GATE_OP, aggregation_method=None, colocate_gradients_with_ops=False, @@ -204,16 +205,18 @@ class ElasticAverageOptimizer(optimizer.Optimizer): if not var_list: var_list = variables.trainable_variables() - elastic_difference = [math_ops.subtract(v, lv) for v, lv in zip( - variables.trainable_variables(), - [self._local_map[var] for var in var_list])] + elastic_difference = [ + math_ops.subtract(v, lv) + for v, lv in zip(variables.trainable_variables(), + [self._local_map[var] for var in var_list]) + ] distance_loss = self._rho * math_ops.add_n( - [gen_nn_ops.l2_loss(ed) for ed in elastic_difference]) + [gen_nn_ops.l2_loss(ed) for ed in elastic_difference]) total_loss = loss + distance_loss - return self._opt.compute_gradients(total_loss, var_list, - gate_gradients, aggregation_method, + return self._opt.compute_gradients(total_loss, var_list, gate_gradients, + aggregation_method, colocate_gradients_with_ops, grad_loss) def apply_gradients(self, grads_and_vars, global_step=None, name=None): @@ -241,7 +244,7 @@ class ElasticAverageOptimizer(optimizer.Optimizer): apply_updates = self._opt.apply_gradients(grads_and_vars) with ops.control_dependencies([apply_updates]): local_update = state_ops.assign_add( - self._local_step, 1, name='local_step_update').op + self._local_step, 1, name='local_step_update').op # update global variables. def _Update_global_variables(): @@ -259,12 +262,16 @@ class ElasticAverageOptimizer(optimizer.Optimizer): differences.append(math_ops.subtract(v, lv)) for lvar, diff in zip(local_vars, differences): with ops.device(lvar.device): - update_ops.append(state_ops.assign_sub(lvar, math_ops.multiply( - self._moving_rate, diff))) + update_ops.append( + state_ops.assign_sub(lvar, + math_ops.multiply(self._moving_rate, + diff))) for var, diff in zip(global_center_vars, differences): with ops.device(var.device): - update_ops.append(state_ops.assign_add(var, math_ops.multiply( - self._moving_rate, diff))) + update_ops.append( + state_ops.assign_add(var, + math_ops.multiply(self._moving_rate, + diff))) if global_step: with ops.colocate_with(global_step): update_ops.append(state_ops.assign_add(global_step, 1)) @@ -272,10 +279,10 @@ class ElasticAverageOptimizer(optimizer.Optimizer): return variable_update with ops.control_dependencies([local_update]): - condition = math_ops.equal(math_ops.mod( - self._local_step, self._period), 0) + condition = math_ops.equal( + math_ops.mod(self._local_step, self._period), 0) conditional_update = control_flow_ops.cond( - condition, _Update_global_variables, control_flow_ops.no_op) + condition, _Update_global_variables, control_flow_ops.no_op) return conditional_update def get_init_op(self, task_index): @@ -285,10 +292,12 @@ class ElasticAverageOptimizer(optimizer.Optimizer): def _Add_sync_queues_and_barrier(enqueue_after_list): """Adds ops to enqueu on all worker queues""" sync_queues = [ - data_flow_ops.FIFOQueue(self._num_worker, [dtypes.bool], shapes=[[]], - shared_name='%s%s' % ( - 'variable_init_sync_queue', i)) for i in - range(self._num_worker)] + data_flow_ops.FIFOQueue( + self._num_worker, [dtypes.bool], + shapes=[[]], + shared_name='%s%s' % ('variable_init_sync_queue', i)) + for i in range(self._num_worker) + ] queue_ops = [] # For each other worker, add an entry in a queue token = constant_op.constant(False) @@ -299,7 +308,7 @@ class ElasticAverageOptimizer(optimizer.Optimizer): else: queue_ops.append(q.enqueue(token)) queue_ops.append( - sync_queues[task_index].dequeue_many(len(sync_queues) - 1)) + sync_queues[task_index].dequeue_many(len(sync_queues) - 1)) return control_flow_ops.group(*queue_ops) init_ops = [] @@ -307,11 +316,10 @@ class ElasticAverageOptimizer(optimizer.Optimizer): global_center_vars = [self._global_map[var] for var in local_vars] local_center_vars = [self._local_map[var] for var in local_vars] if not (local_vars and global_center_vars and local_center_vars): - raise ValueError( - 'The lists of local_variables, global_center_variables, ' - 'local_center_variables should not be empty ') - for lvar, gc_var, lc_var in zip( - local_vars, global_center_vars, local_center_vars): + raise ValueError('The lists of local_variables, global_center_variables, ' + 'local_center_variables should not be empty ') + for lvar, gc_var, lc_var in zip(local_vars, global_center_vars, + local_center_vars): init_ops.append(state_ops.assign(lvar, gc_var)) init_ops.append(state_ops.assign(lc_var, gc_var)) @@ -325,6 +333,7 @@ class ElasticAverageOptimizer(optimizer.Optimizer): class _ElasticAverageOptimizerHook(session_run_hook.SessionRunHook): + def __init__(self, ea_optimizer, is_chief, task_index): """Creates hook to handle ElasticAverageOptimizer initialization ops. diff --git a/tensorflow/contrib/opt/python/training/elastic_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/elastic_average_optimizer_test.py index 446e91018d477d75116f6b78a2443ed79ed3b3ef..37539b959959b5cf1f7b2c8e8d2b6b05191565ad 100644 --- a/tensorflow/contrib/opt/python/training/elastic_average_optimizer_test.py +++ b/tensorflow/contrib/opt/python/training/elastic_average_optimizer_test.py @@ -38,20 +38,20 @@ def create_local_cluster(num_workers, num_ps, protocol="grpc"): worker_ports = [portpicker.pick_unused_port() for _ in range(num_workers)] ps_ports = [portpicker.pick_unused_port() for _ in range(num_ps)] cluster_dict = { - "worker": ["localhost:%s" % port for port in worker_ports], - "ps": ["localhost:%s" % port for port in ps_ports] + "worker": ["localhost:%s" % port for port in worker_ports], + "ps": ["localhost:%s" % port for port in ps_ports] } cs = server_lib.ClusterSpec(cluster_dict) workers = [ - server_lib.Server( - cs, job_name="worker", protocol=protocol, task_index=ix, start=True) - for ix in range(num_workers) + server_lib.Server( + cs, job_name="worker", protocol=protocol, task_index=ix, start=True) + for ix in range(num_workers) ] ps_servers = [ - server_lib.Server( - cs, job_name="ps", protocol=protocol, task_index=ix, start=True) - for ix in range(num_ps) + server_lib.Server( + cs, job_name="ps", protocol=protocol, task_index=ix, start=True) + for ix in range(num_ps) ] return cluster_dict, workers, ps_servers @@ -68,15 +68,14 @@ def _get_workers(num_workers, period, workers, moving_rate): is_chief = (worker_id == 0) with graph.as_default(): worker_device = "/job:worker/task:%d/cpu:0" % (worker_id) - ea_coustom = ElasticAverageCustomGetter( - worker_device=worker_device) - with variable_scope.variable_scope('', - custom_getter=ea_coustom), ops.device( - device_setter.replica_device_setter(worker_device=worker_device, - ps_device="/job:ps/task:0/cpu:0", - ps_tasks=1)): - global_step = variables.Variable(0, name='global_step', - trainable=False) + ea_coustom = ElasticAverageCustomGetter(worker_device=worker_device) + with variable_scope.variable_scope( + "", custom_getter=ea_coustom), ops.device( + device_setter.replica_device_setter( + worker_device=worker_device, + ps_device="/job:ps/task:0/cpu:0", + ps_tasks=1)): + global_step = variables.Variable(0, name="global_step", trainable=False) var_0 = variable_scope.get_variable(initializer=0.0, name="v0") var_1 = variable_scope.get_variable(initializer=1.0, name="v1") @@ -86,21 +85,19 @@ def _get_workers(num_workers, period, workers, moving_rate): sgd_opt = gradient_descent.GradientDescentOptimizer(1.0) opt = ElasticAverageOptimizer( - opt=sgd_opt, - num_worker=num_workers, - moving_rate=moving_rate, - communication_period=period, - ea_custom_getter=ea_coustom - ) + opt=sgd_opt, + num_worker=num_workers, + moving_rate=moving_rate, + communication_period=period, + ea_custom_getter=ea_coustom) train_op = [ - opt.apply_gradients( - ([grads_0, var_0], - [grads_1, var_1]), global_step) + opt.apply_gradients(([grads_0, var_0], [grads_1, var_1]), + global_step) ] easgd_hook = opt.make_session_run_hook(is_chief, worker_id) # Creates MonitoredSession - sess = training.MonitoredTrainingSession(workers[worker_id].target, - hooks=[easgd_hook]) + sess = training.MonitoredTrainingSession( + workers[worker_id].target, hooks=[easgd_hook]) sessions.append(sess) graphs.append(graph) @@ -110,6 +107,7 @@ def _get_workers(num_workers, period, workers, moving_rate): class ElasticAverageOptimizerTest(test.TestCase): + def _run(self, train_op, sess): sess.run(train_op) @@ -117,15 +115,14 @@ class ElasticAverageOptimizerTest(test.TestCase): num_workers = 1 communication_period = 2 num_ps = 1 - cluster, workers, _ = create_local_cluster(num_workers=num_workers, - num_ps=num_ps) + cluster, workers, _ = create_local_cluster( + num_workers=num_workers, num_ps=num_ps) - sessions, graphs, train_ops = _get_workers(num_workers, - communication_period, - workers, 1.0) + sessions, graphs, train_ops = _get_workers( + num_workers, communication_period, workers, 1.0) - var_0 = graphs[0].get_tensor_by_name('v0:0') - var_1 = graphs[0].get_tensor_by_name('v1:0') + var_0 = graphs[0].get_tensor_by_name("v0:0") + var_1 = graphs[0].get_tensor_by_name("v1:0") global_step = training_util.get_global_step(graphs[0]) var_0_g = graphs[0].get_tensor_by_name(GLOBAL_VARIABLE_NAME + "/v0:0") var_1_g = graphs[0].get_tensor_by_name(GLOBAL_VARIABLE_NAME + "/v1:0") @@ -166,18 +163,17 @@ class ElasticAverageOptimizerTest(test.TestCase): num_workers = 2 communication_period = 1 num_ps = 2 - cluster, workers, _ = create_local_cluster(num_workers=num_workers, - num_ps=num_ps) + cluster, workers, _ = create_local_cluster( + num_workers=num_workers, num_ps=num_ps) - sessions, graphs, train_ops = _get_workers(num_workers, - communication_period, - workers, 0.5) + sessions, graphs, train_ops = _get_workers( + num_workers, communication_period, workers, 0.5) - var_0 = graphs[0].get_tensor_by_name('v0:0') - var_1 = graphs[0].get_tensor_by_name('v1:0') + var_0 = graphs[0].get_tensor_by_name("v0:0") + var_1 = graphs[0].get_tensor_by_name("v1:0") - var_0_1 = graphs[1].get_tensor_by_name('v0:0') - var_1_1 = graphs[1].get_tensor_by_name('v1:0') + var_0_1 = graphs[1].get_tensor_by_name("v0:0") + var_1_1 = graphs[1].get_tensor_by_name("v1:0") var_0_g = graphs[0].get_tensor_by_name(GLOBAL_VARIABLE_NAME + "/v0:0") var_1_g = graphs[0].get_tensor_by_name(GLOBAL_VARIABLE_NAME + "/v1:0") @@ -201,25 +197,24 @@ class ElasticAverageOptimizerTest(test.TestCase): def testPS2TasksWithClusterSpecClass(self): cluster_spec = server_lib.ClusterSpec({ - "ps": ["ps0:2222", "ps1:2222"], - "worker": ["worker0:2222", "worker1:2222", "worker2:2222"] + "ps": ["ps0:2222", "ps1:2222"], + "worker": ["worker0:2222", "worker1:2222", "worker2:2222"] }) - ea_coustom = ElasticAverageCustomGetter( - worker_device="/job:worker/task:0") + ea_coustom = ElasticAverageCustomGetter(worker_device="/job:worker/task:0") from tensorflow.python.training import device_setter with ops.device( device_setter.replica_device_setter(cluster=cluster_spec, worker_device="/job:worker/task:0", ps_device="/job:ps")), \ - variable_scope.variable_scope('', custom_getter=ea_coustom): + variable_scope.variable_scope("", custom_getter=ea_coustom): v = variable_scope.get_variable(initializer=[1, 2], name="v") - w = variable_scope.get_variable(initializer=[2, 1], name='w') - v_g, w_g = ea_coustom._global_map[v],ea_coustom._global_map[w] + w = variable_scope.get_variable(initializer=[2, 1], name="w") + v_g, w_g = ea_coustom._global_map[v], ea_coustom._global_map[w] self.assertDeviceEqual("/job:worker/task:0", v.device) self.assertDeviceEqual("job:ps/task:0", v_g.device) self.assertDeviceEqual("/job:worker/task:0", w.device) self.assertDeviceEqual("job:ps/task:1", w_g.device) -if __name__ == '__main__': +if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/opt/python/training/external_optimizer.py b/tensorflow/contrib/opt/python/training/external_optimizer.py index f243317f1df2ec8d93d44ad534f3fa58527f3217..82ebca7f20306e5658c8321716e39f9c7f8b8970 100644 --- a/tensorflow/contrib/opt/python/training/external_optimizer.py +++ b/tensorflow/contrib/opt/python/training/external_optimizer.py @@ -397,10 +397,6 @@ class ScipyOptimizerInterface(ExternalOptimizerInterface): 'automatically and cannot be injected manually'.format(kwarg)) minimize_kwargs.update(optimizer_kwargs) - if method == 'SLSQP': - # SLSQP doesn't support step callbacks. Obviate associated warning - # message. - del minimize_kwargs['callback'] import scipy.optimize # pylint: disable=g-import-not-at-top result = scipy.optimize.minimize(*minimize_args, **minimize_kwargs) diff --git a/tensorflow/contrib/opt/python/training/external_optimizer_test.py b/tensorflow/contrib/opt/python/training/external_optimizer_test.py index 0f597d0a246a53892d72939edd1499a86c01017d..953586ee70cd4137295dd254bfb2d37cab0bcfe4 100644 --- a/tensorflow/contrib/opt/python/training/external_optimizer_test.py +++ b/tensorflow/contrib/opt/python/training/external_optimizer_test.py @@ -299,6 +299,45 @@ class ScipyOptimizerInterfaceTest(TestCase): method = optimizer.optimizer_kwargs.get('method') self.assertEqual('SLSQP', method) + def test_callbacks(self): + vector_val = np.array([7., -2.], dtype=np.float32) + vector = variables.Variable(vector_val, 'vector') + + minimum_location_val = np.arange(2) + minimum_location = constant_op.constant( + minimum_location_val, dtype=dtypes.float32) + + loss = math_ops.reduce_sum(math_ops.square(vector - minimum_location)) / 2. + loss_val_first = ((vector_val - minimum_location_val)**2).sum() / 2. + + optimizer = external_optimizer.ScipyOptimizerInterface(loss, method='SLSQP') + + with self.test_session() as sess: + sess.run(variables.global_variables_initializer()) + + initial_vector_val = sess.run(vector) + + extra_fetches = [loss] + + step_callback = test.mock.Mock() + loss_callback = test.mock.Mock() + + optimizer.minimize( + sess, + fetches=extra_fetches, + loss_callback=loss_callback, + step_callback=step_callback) + + loss_val_last = sess.run(loss) + + call_first = test.mock.call(loss_val_first) + call_last = test.mock.call(loss_val_last) + loss_calls = [call_first, call_last] + loss_callback.assert_has_calls(loss_calls, any_order=True) + + args, _ = step_callback.call_args + self.assertAllClose(minimum_location_val, args[0]) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py b/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py index 4c3fec067287e8edefcc4e36ca9fa91f5657013b..aeca900bc8ff4c4cc26da490ce43dfec70fd9f11 100644 --- a/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py +++ b/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py @@ -47,8 +47,9 @@ class LazyAdamOptimizer(adam.AdamOptimizer): """ def _apply_sparse(self, grad, var): - beta1_power = math_ops.cast(self._beta1_power, var.dtype.base_dtype) - beta2_power = math_ops.cast(self._beta2_power, var.dtype.base_dtype) + beta1_power, beta2_power = self._get_beta_accumulators() + beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype) + beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype) lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype) beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype) beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype) diff --git a/tensorflow/contrib/opt/python/training/model_average_optimizer.py b/tensorflow/contrib/opt/python/training/model_average_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..a7c97a1da2baf29914337094c6153447c997af08 --- /dev/null +++ b/tensorflow/contrib/opt/python/training/model_average_optimizer.py @@ -0,0 +1,308 @@ +# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Wrapper optimizer for Model Average.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import data_flow_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables +from tensorflow.python.training import optimizer +from tensorflow.python.training import session_run_hook + +GLOBAL_VARIABLE_NAME = "global_center_variable" + + +class ModelAverageCustomGetter(object): + """Custom_getter class is used to do. + + 1. Change trainable variables to local collection and place them at worker + device + 2. Generate global variables + Notice that the class should be used with tf.replica_device_setter, + so that the global center variables and global step variable can be placed + at ps device. Besides, use 'tf.get_variable' instead of 'tf.Variable' to + use this custom getter. + + For example, + ma_custom_getter = ModelAverageCustomGetter(worker_device) + with tf.device( + tf.train.replica_device_setter( + worker_device=worker_device, + ps_device="/job:ps/cpu:0", + cluster=cluster)), + tf.variable_scope('',custom_getter=ma_custom_getter): + hid_w = tf.get_variable( + initializer=tf.truncated_normal( + [IMAGE_PIXELS * IMAGE_PIXELS, FLAGS.hidden_units], + stddev=1.0 / IMAGE_PIXELS), + name="hid_w") + hid_b = tf.get_variable(initializer=tf.zeros([FLAGS.hidden_units]), + name="hid_b") + """ + + def __init__(self, worker_device): + """Create a new `ElasticAverageCustomGetter`. + + Args: + worker_device: String. Name of the `worker` job. + """ + self._worker_device = worker_device + self._local_2_global = {} + + def __call__(self, getter, name, trainable, collections, *args, **kwargs): + if trainable: + with ops.device(self._worker_device): + local_var = getter( + name, + trainable=True, + collections=[ops.GraphKeys.LOCAL_VARIABLES], + *args, + **kwargs) + + global_variable = variable_scope.variable( + name="%s/%s" % (GLOBAL_VARIABLE_NAME, name), + initial_value=local_var.initialized_value(), + trainable=False, + collections=[ops.GraphKeys.GLOBAL_VARIABLES]) + + self._local_2_global[local_var] = global_variable + return local_var + else: + return getter(name, trainable, collections, *args, **kwargs) + + +class ModelAverageOptimizer(optimizer.Optimizer): + """Wrapper optimizer that implements the Model Average algorithm. + + This is a sync optimizer. During the training, each worker will update + the local variables and maintains its own local_step, which starts from 0 + and is incremented by 1 after each update of local variables. Whenever the + interval_steps divides the local step, the local variables from all the + workers will be averaged and assigned to global center variables. Then the + local variables will be assigned by global center variables. + """ + + def __init__(self, + opt, + num_worker, + is_chief, + ma_custom_getter, + interval_steps=100, + use_locking=True, + name="ModelAverageOptimizer"): + """Construct a new model average optimizer. + + Args: + opt: The actual optimizer that will be used to update local variables + num_worker: The number of workers + is_chief: whether chief worker + ma_custom_getter: ModelAverageCustomGetter + interval_steps: An int point value to controls the frequency of the + average of local variables + use_locking: If True use locks for update operations + name: string. Optional name of the returned operation + """ + super(ModelAverageOptimizer, self).__init__(use_locking, name) + self._opt = opt + self._num_worker = num_worker + self._is_chief = is_chief + self._local_2_global = ma_custom_getter._local_2_global # pylint:disable=protected-access + self._interval_steps = interval_steps + self._accumulator_list = [] + self._chief_init_op = None + + self._local_step = variable_scope.get_variable( + initializer=0, + trainable=False, + collections=[ops.GraphKeys.LOCAL_VARIABLES], + name="local_step") + + self._opt._prepare() # pylint:disable=protected-access + + def compute_gradients(self, *args, **kwargs): + """Compute gradients of "loss" for the variables in "var_list". + + This simply wraps the compute_gradients() from the real optimizer. + + Args: + *args: Arguments for compute_gradients(). + **kwargs: Keyword arguments for compute_gradients(). + + Returns: + A list of (gradient, variable) pairs. + """ + return self._opt.compute_gradients(*args, **kwargs) + + def _local_vars_update(self, var_list): + """Get the update ops for the local variables in "var_list". + + Args: + var_list: Optional list or tuple of 'tf.Variable' to update + + Returns: + An update op + + Raises: + ValueError: if var_list is empty. + """ + if not var_list: + raise ValueError("The list of local_variables should not be empty") + update_ops = [] + global_center_vars = [self._local_2_global[var] for var in var_list] + for lvar, gvar in zip(var_list, global_center_vars): + with ops.device(lvar.device): + update_ops.append(state_ops.assign(lvar, gvar.read_value())) + return control_flow_ops.group(*(update_ops)) + + def apply_gradients(self, grads_and_vars, global_step=None, name=None): + """Apply gradients to variables. + + This contains most of the synchronization implementation and also wraps the + apply_gradients() from the real optimizer. The chief work updates global + variables. + + Args: + grads_and_vars: List of (gradient, variable) pairs as returned by + compute_gradients(). + global_step: Optional Variable to increment by one after the + variables have been updated. + name: Optional name for the returned operation. Default to the + name passed to the Optimizer constructor. + + Returns: + A conditional 'Operation' that update both local and global variables or + just local variables + + Raises: + ValueError: If the grads_and_vars is empty. + ValueError: If global step is not provided, the staleness cannot be + checked. + """ + + # update local variables + if not grads_and_vars: + raise ValueError("Must supply at least one variable") + if global_step is None: + raise ValueError("Global step is required") + + apply_updates = self._opt.apply_gradients(grads_and_vars) + with ops.control_dependencies([apply_updates]): + local_update = state_ops.assign_add( + self._local_step, 1, name="local_step_update").op + + # update global variables. + def _update_global_variables(): # pylint: disable=missing-docstring + local_vars = [v for g, v in grads_and_vars if g is not None] + global_vars = [self._local_2_global[v] for v in local_vars] + # sync queue + with ops.colocate_with(global_step): + sync_queue = data_flow_ops.FIFOQueue( + -1, [dtypes.bool], shapes=[[]], shared_name="sync_queue") + train_ops = [] + aggregated_vars = [] + with ops.name_scope(None, self._name + "/global"): + for var, gvar in zip(local_vars, global_vars): + # pylint: disable=protected-access + with ops.device(gvar.device): + if isinstance(var._ref(), ops.Tensor): + var_accum = data_flow_ops.ConditionalAccumulator( + var.dtype, + shape=var.get_shape(), + shared_name=gvar.name + "/var_accum") + train_ops.append( + var_accum.apply_grad(var._ref(), local_step=global_step)) + aggregated_vars.append(var_accum.take_grad(self._num_worker)) + else: + raise ValueError("Unknown local variable type!") + self._accumulator_list.append((var_accum, gvar.device)) + # chief worker updates global vars and enqueues tokens to the sync queue + if self._is_chief: + update_ops = [] + with ops.control_dependencies(train_ops): + for avg_var, gvar in zip(aggregated_vars, global_vars): + with ops.device(gvar.device): + update_ops.append(state_ops.assign(gvar, avg_var)) + with ops.device(global_step.device): + update_ops.append(state_ops.assign_add(global_step, 1)) + with ops.control_dependencies(update_ops), ops.device( + global_step.device): + tokens = array_ops.fill([self._num_worker - 1], + constant_op.constant(False)) + sync_op = sync_queue.enqueue_many(tokens) + else: + with ops.control_dependencies(train_ops), ops.device( + global_step.device): + sync_op = sync_queue.dequeue() + + with ops.control_dependencies([sync_op]): + local_update_op = self._local_vars_update(local_vars) + return local_update_op + + with ops.control_dependencies([local_update]): + condition = math_ops.equal( + math_ops.mod(self._local_step, self._interval_steps), 0) + conditional_update = control_flow_ops.cond( + condition, _update_global_variables, control_flow_ops.no_op) + + chief_init_ops = [] + for accum, dev in self._accumulator_list: + with ops.device(dev): + chief_init_ops.append( + accum.set_global_step(global_step, name="SetGlobalStep")) + self._chief_init_op = control_flow_ops.group(*(chief_init_ops)) + + return conditional_update + + def get_init_op(self): + """Returns the op. + + This method lets all the local variables equal to the global + variables before the training begins. + """ + return self._local_vars_update(variables.trainable_variables()) + + def make_session_run_hook(self): + """Creates a hook to handle ModelAverage ops such as initialization.""" + return _ModelAverageOptimizerHook(self, self._is_chief) + + +class _ModelAverageOptimizerHook(session_run_hook.SessionRunHook): # pylint: disable=missing-docstring + + def __init__(self, ma_optimizer, is_chief): + """Creates hook to handle ModelAverageOptimizer initialization ops. + + Args: + ma_optimizer: `ModelAverageOptimizer` which this hook will initialize. + is_chief: `Bool`, whether is this a chief replica or not. + """ + self._ma_optimizer = ma_optimizer + self._is_chief = is_chief + + def begin(self): + self._local_init_op = variables.local_variables_initializer() + self._global_init_op = None + if self._is_chief: + self._global_init_op = variables.global_variables_initializer() + self._chief_init_op = self._ma_optimizer._chief_init_op # pylint: disable=protected-access + self._variable_init_op = self._ma_optimizer.get_init_op() diff --git a/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6cca0a8a009456f266245fd9a638bfab371c9b34 --- /dev/null +++ b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py @@ -0,0 +1,198 @@ +# Copyright 2017 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 ModelAverageOptimizer.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import portpicker + +from tensorflow.contrib.opt.python.training import model_average_optimizer +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables +from tensorflow.python.platform import test +from tensorflow.python.training import device_setter +from tensorflow.python.training import gradient_descent +from tensorflow.python.training import server_lib +from tensorflow.python.training import training +from tensorflow.python.training import training_util + + +def create_local_cluster(num_workers, num_ps, protocol="grpc"): + """Create local GRPC servers and return them.""" + worker_ports = [portpicker.pick_unused_port() for _ in range(num_workers)] + ps_ports = [portpicker.pick_unused_port() for _ in range(num_ps)] + cluster_dict = { + "worker": ["localhost:%s" % port for port in worker_ports], + "ps": ["localhost:%s" % port for port in ps_ports] + } + cs = server_lib.ClusterSpec(cluster_dict) + + workers = [ + server_lib.Server( + cs, job_name="worker", protocol=protocol, task_index=ix, start=True) + for ix in range(num_workers) + ] + ps_servers = [ + server_lib.Server( + cs, job_name="ps", protocol=protocol, task_index=ix, start=True) + for ix in range(num_ps) + ] + + return cluster_dict, workers, ps_servers + + +# Creates the workers and return their sessions, graphs, train_ops. +# Cheif worker will update at last +def _get_workers(num_workers, steps, workers): + sessions = [] + graphs = [] + train_ops = [] + for worker_id in range(num_workers): + graph = ops.Graph() + is_chief = (worker_id == 0) + with graph.as_default(): + worker_device = "/job:worker/task:%d/cpu:0" % (worker_id) + ma_coustom = model_average_optimizer.ModelAverageCustomGetter( + worker_device=worker_device) + with variable_scope.variable_scope( + "", custom_getter=ma_coustom), ops.device( + device_setter.replica_device_setter( + worker_device=worker_device, + ps_device="/job:ps/task:0/cpu:0", + ps_tasks=1)): + + global_step = variables.Variable(0, name="global_step", trainable=False) + var_0 = variable_scope.get_variable(initializer=0.0, name="v0") + var_1 = variable_scope.get_variable(initializer=1.0, name="v1") + + with ops.device("/job:worker/task:" + str(worker_id)): + if worker_id == 0: + grads_0 = constant_op.constant(-1.0) + grads_1 = constant_op.constant(-1.0) + else: + grads_0 = constant_op.constant(-2.0) + grads_1 = constant_op.constant(-2.0) + sgd_opt = gradient_descent.GradientDescentOptimizer(1.0) + opt = model_average_optimizer.ModelAverageOptimizer( + opt=sgd_opt, + num_worker=num_workers, + ma_custom_getter=ma_coustom, + is_chief=is_chief, + interval_steps=steps) + train_op = [ + opt.apply_gradients([[grads_0, var_0], [grads_1, var_1]], + global_step) + ] + easgd_hook = opt.make_session_run_hook() + # Creates MonitoredSession + sess = training.MonitoredTrainingSession( + workers[worker_id].target, hooks=[easgd_hook]) + + sessions.append(sess) + graphs.append(graph) + train_ops.append(train_op) + return sessions, graphs, train_ops + + +class ModelAverageOptimizerTest(test.TestCase): + def _run(self, train_op, sess): + sess.run(train_op) + + def test1Workers2Period(self): + num_workers = 2 + steps = 2 + num_ps = 1 + _, workers, _ = create_local_cluster( + num_workers=num_workers, num_ps=num_ps) + + sessions, graphs, train_ops = _get_workers(num_workers, steps, workers) + + var_0 = graphs[0].get_tensor_by_name("v0:0") + var_1 = graphs[0].get_tensor_by_name("v1:0") + global_step = training_util.get_global_step(graphs[0]) + global_var_0 = graphs[0].get_tensor_by_name( + model_average_optimizer.GLOBAL_VARIABLE_NAME + "/v0:0") + global_var_1 = graphs[0].get_tensor_by_name( + model_average_optimizer.GLOBAL_VARIABLE_NAME + "/v1:0") + + # Verify the initialized value. + self.assertAllEqual(0.0, sessions[0].run(var_0)) + self.assertAllEqual(1.0, sessions[0].run(var_1)) + self.assertAllEqual(0.0, sessions[0].run(global_var_0)) + self.assertAllEqual(1.0, sessions[0].run(global_var_1)) + self.assertAllEqual(0, sessions[0].run(global_step)) + + sessions[0].run(train_ops[0]) + sessions[1].run(train_ops[1]) + + self.assertAllEqual(1.0, sessions[0].run(var_0)) + self.assertAllEqual(2.0, sessions[0].run(var_1)) + self.assertAllEqual(0.0, sessions[0].run(global_var_0)) + self.assertAllEqual(1.0, sessions[0].run(global_var_1)) + self.assertAllEqual(0, sessions[0].run(global_step)) + + # iteration 2, global varibale update + thread_0 = self.checkedThread( + target=self._run, args=(train_ops[0], sessions[0])) + thread_1 = self.checkedThread( + target=self._run, args=(train_ops[1], sessions[1])) + thread_0.start() + thread_1.start() + thread_0.join() + thread_1.join() + + self.assertAllEqual(3.0, sessions[0].run(var_0)) + self.assertAllEqual(4.0, sessions[0].run(var_1)) + self.assertAllEqual(3.0, sessions[0].run(global_var_0)) + self.assertAllEqual(4.0, sessions[0].run(global_var_1)) + self.assertAllEqual(1, sessions[0].run(global_step)) + + # iteration 3 + sessions[0].run(train_ops[0]) + + self.assertAllEqual(4.0, sessions[0].run(var_0)) + self.assertAllEqual(5.0, sessions[0].run(var_1)) + self.assertAllEqual(3.0, sessions[0].run(global_var_0)) + self.assertAllEqual(4.0, sessions[0].run(global_var_1)) + self.assertAllEqual(1, sessions[0].run(global_step)) + + def testPS2TasksWithClusterSpecClass(self): + cluster_spec = server_lib.ClusterSpec({ + "ps": ["ps0:2222", "ps1:2222"], + "worker": ["worker0:2222", "worker1:2222", "worker2:2222"] + }) + worker_device = "/job:worker/task:0" + ma_coustom = model_average_optimizer.ModelAverageCustomGetter( + worker_device=worker_device) + from tensorflow.python.training import device_setter + with ops.device( + device_setter.replica_device_setter(cluster=cluster_spec, + worker_device=worker_device, + ps_device="/job:ps")), \ + variable_scope.variable_scope("", custom_getter=ma_coustom): + v = variable_scope.get_variable(initializer=[1, 2], name="v") + w = variable_scope.get_variable(initializer=[2, 1], name="w") + v_g, w_g = ma_coustom._local_2_global[v], ma_coustom._local_2_global[w] + self.assertDeviceEqual("/job:worker/task:0", v.device) + self.assertDeviceEqual("job:ps/task:0", v_g.device) + self.assertDeviceEqual("/job:worker/task:0", w.device) + self.assertDeviceEqual("job:ps/task:1", w_g.device) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/opt/python/training/moving_average_optimizer.py b/tensorflow/contrib/opt/python/training/moving_average_optimizer.py index d68ad23d65500cc2348459cdc53030c2ea08373a..9ce50bfe1054072b315adecb87f1ba729dfe0d83 100644 --- a/tensorflow/contrib/opt/python/training/moving_average_optimizer.py +++ b/tensorflow/contrib/opt/python/training/moving_average_optimizer.py @@ -83,7 +83,7 @@ class MovingAverageOptimizer(optimizer.Optimizer): self._optimizer = opt self._ema = moving_averages.ExponentialMovingAverage( average_decay, num_updates=num_updates) - self._variable_map = None + self._swapped_variable_name_map = None self._sequential_update = sequential_update def compute_gradients(self, *args, **kwargs): @@ -93,7 +93,7 @@ class MovingAverageOptimizer(optimizer.Optimizer): train_op = self._optimizer.apply_gradients( grads_and_vars, global_step=global_step, name=name) var_list = [x[1] for x in grads_and_vars if x[0] is not None] - self._variable_map = {} + self._swapped_variable_name_map = {} if self._sequential_update: with ops.control_dependencies([train_op]): ma_op = self._ema.apply(var_list) @@ -102,9 +102,9 @@ class MovingAverageOptimizer(optimizer.Optimizer): for v in var_list: v_avg = self._ema.average(v) - self._variable_map[v.op.name] = v_avg - self._variable_map[v_avg.op.name] = v - return control_flow_ops.group(train_op, ma_op, name="train_with_avg") + self._swapped_variable_name_map[v.op.name] = v_avg.op.name + self._swapped_variable_name_map[v_avg.op.name] = v.op.name + return control_flow_ops.group(train_op, ma_op, name='train_with_avg') def swapping_saver(self, var_list=None, name='swapping_saver', **kwargs): """Create a saver swapping moving averages and variables. @@ -129,22 +129,45 @@ class MovingAverageOptimizer(optimizer.Optimizer): Raises: RuntimeError: If apply_gradients or minimize has not been called before. + ValueError: If var_list is provided and contains some variables but not + their moving average counterpart. """ - if self._variable_map is None: + if self._swapped_variable_name_map is None: raise RuntimeError('Must call apply_gradients or minimize before ' 'creating the swapping_saver') if var_list is None: var_list = variables.global_variables() if not isinstance(var_list, dict): var_list = saver.BaseSaverBuilder.OpListToDict(var_list) + + # OpListToDict converts variables to tensors. We make sure we can get + # the unique variable name for normal and resource vaiables. + def get_v_name(tensor): + if tensor.op.type == 'ReadVariableOp': + return tensor.op.inputs[0].op.name + else: + return tensor.op.name + + v_name_to_tensor = {} + for tensor in six.itervalues(var_list): + v_name = get_v_name(tensor) + v_name_to_tensor[v_name] = tensor + # Now swap variables and moving averages swapped_var_list = {} - for k, v in six.iteritems(var_list): - v_swap = self._variable_map.get(v.op.name, None) - if v_swap: - swapped_var_list[k] = v_swap - else: - swapped_var_list[k] = v + for k, tensor in six.iteritems(var_list): + v_name = get_v_name(tensor) + swapped_v_name = self._swapped_variable_name_map.get(v_name, None) + tensor_to_save = tensor + if swapped_v_name is not None: + if swapped_v_name in v_name_to_tensor: + tensor_to_save = v_name_to_tensor[swapped_v_name] + else: + raise ValueError( + ('Variable to swap %s is not part of variables to save. ' + 'This breaks MovingAverageOptimizer.') % swapped_v_name) + swapped_var_list[k] = tensor_to_save + # Build the swapping saver. return saver.Saver(swapped_var_list, name=name, **kwargs) diff --git a/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py index 60929add198f2e69b5acc2eb5516dafc82b1f3ba..85e3e8d3791f2331ed249c0b7f67a3dbde4fca08 100644 --- a/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py +++ b/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py @@ -24,6 +24,10 @@ import six from tensorflow.contrib.opt.python.training import moving_average_optimizer from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import test from tensorflow.python.training import gradient_descent @@ -33,13 +37,26 @@ from tensorflow.python.training import saver class MovingAverageOptimizerTest(test.TestCase): def testRun(self): + self._helpTestRun(use_resource=False) + + def testRunUseResource(self): + # Test that MovingAverageOptimizer works with resource variables. + self._helpTestRun(use_resource=True) + + def _helpTestRun(self, use_resource=False): for sequential_update in [True, False]: for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.test_session() as sess: + with self.test_session(graph=ops.Graph()) as sess: orig_val0 = [1.0, 2.0] orig_val1 = [3.0, 4.0] - var0 = variables.Variable(orig_val0, name='var0', dtype=dtype) - var1 = variables.Variable(orig_val1, name='var1', dtype=dtype) + var0 = variable_scope.get_variable( + 'var0', + initializer=constant_op.constant(orig_val0, dtype=dtype), + use_resource=use_resource) + var1 = variable_scope.get_variable( + 'var1', + initializer=constant_op.constant(orig_val1, dtype=dtype), + use_resource=use_resource) grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) @@ -52,22 +69,63 @@ class MovingAverageOptimizerTest(test.TestCase): save_path = os.path.join(save_dir, 'model') update = opt.apply_gradients( list(six.moves.zip([grads0, grads1], [var0, var1]))) + global_vars = variables.global_variables() + ema_var0 = [ + v for v in global_vars + if v.op.name == 'var0/ExponentialMovingAverage' + ][0] + ema_var1 = [ + v for v in global_vars + if v.op.name == 'var1/ExponentialMovingAverage' + ][0] + perturb = control_flow_ops.group([ + state_ops.assign_add(var0, [1.0, 1.0]), + state_ops.assign_add(var1, [2.0, 2.0]), + state_ops.assign_add(ema_var0, [3.0, 3.0]), + state_ops.assign_add(ema_var1, [4.0, 4.0]) + ]) + + # Test taht saver with missing ema variables will fail. + with self.assertRaisesRegexp(ValueError, r'Variable to swap'): + opt.swapping_saver(var_list=[var0]) + train_saver = opt.swapping_saver() + train_saver_subset = opt.swapping_saver(var_list=[var0, ema_var0]) inference_saver = saver.Saver() variables.global_variables_initializer().run() # Step 1. update.run() - val0 = var0.eval() - val1 = var1.eval() self.assertAllCloseAccordingToType([0.8, 1.8], var0.eval()) self.assertAllCloseAccordingToType([2.98, 3.98], var1.eval()) + if sequential_update: + self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) + self.assertAllCloseAccordingToType([2.99, 3.99], ema_var1.eval()) # Test that the swapping saver save/restore operation is identity. train_saver.save(sess, save_path) train_saver.restore(sess, save_path) - val0 = var0.eval() - val1 = var1.eval() self.assertAllCloseAccordingToType([0.8, 1.8], var0.eval()) self.assertAllCloseAccordingToType([2.98, 3.98], var1.eval()) + if sequential_update: + self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) + self.assertAllCloseAccordingToType([2.99, 3.99], ema_var1.eval()) + # Test that the subset saver saves the EMA variable as well. + if sequential_update: + subset_save_path = save_path + '_subset' + train_saver_subset.save(sess, subset_save_path) + perturb.run() + self.assertAllCloseAccordingToType([1.8, 2.8], var0.eval()) + self.assertAllCloseAccordingToType([3.9, 4.9], ema_var0.eval()) + self.assertAllCloseAccordingToType([4.98, 5.98], var1.eval()) + self.assertAllCloseAccordingToType([6.99, 7.99], ema_var1.eval()) + # Restoring should only restore var0 and ema_var0. + train_saver_subset.restore(sess, subset_save_path) + self.assertAllCloseAccordingToType([0.8, 1.8], var0.eval()) + self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) + self.assertAllCloseAccordingToType([4.98, 5.98], var1.eval()) + self.assertAllCloseAccordingToType([6.99, 7.99], ema_var1.eval()) + # Restore back to previou state. + train_saver.restore(sess, save_path) + # If updates are parallel, this is not always true after the 1st step. if sequential_update: # Test that the normal saver will have the averaged variables. diff --git a/tensorflow/contrib/opt/python/training/nadam_optimizer.py b/tensorflow/contrib/opt/python/training/nadam_optimizer.py index a4421ecfe6b0af9759c6aaa51d644f1211965b6a..44a8890cb107440b79cf8fbbdfcfda503b1c910f 100644 --- a/tensorflow/contrib/opt/python/training/nadam_optimizer.py +++ b/tensorflow/contrib/opt/python/training/nadam_optimizer.py @@ -34,12 +34,13 @@ class NadamOptimizer(adam.AdamOptimizer): def _apply_dense(self, grad, var): m = self.get_slot(var, "m") v = self.get_slot(var, "v") + beta1_power, beta2_power = self._get_beta_accumulators() return training_ops.apply_adam( var, m, v, - math_ops.cast(self._beta1_power, var.dtype.base_dtype), - math_ops.cast(self._beta2_power, var.dtype.base_dtype), + math_ops.cast(beta1_power, var.dtype.base_dtype), + math_ops.cast(beta2_power, var.dtype.base_dtype), math_ops.cast(self._lr_t, var.dtype.base_dtype), math_ops.cast(self._beta1_t, var.dtype.base_dtype), math_ops.cast(self._beta2_t, var.dtype.base_dtype), @@ -51,12 +52,13 @@ class NadamOptimizer(adam.AdamOptimizer): def _resource_apply_dense(self, grad, var): m = self.get_slot(var, "m") v = self.get_slot(var, "v") + beta1_power, beta2_power = self._get_beta_accumulators() return training_ops.resource_apply_adam( var.handle, m.handle, v.handle, - math_ops.cast(self._beta1_power, grad.dtype.base_dtype), - math_ops.cast(self._beta2_power, grad.dtype.base_dtype), + math_ops.cast(beta1_power, grad.dtype.base_dtype), + math_ops.cast(beta2_power, grad.dtype.base_dtype), math_ops.cast(self._lr_t, grad.dtype.base_dtype), math_ops.cast(self._beta1_t, grad.dtype.base_dtype), math_ops.cast(self._beta2_t, grad.dtype.base_dtype), @@ -66,8 +68,9 @@ class NadamOptimizer(adam.AdamOptimizer): use_nesterov=True) def _apply_sparse_shared(self, grad, var, indices, scatter_add): - beta1_power = math_ops.cast(self._beta1_power, var.dtype.base_dtype) - beta2_power = math_ops.cast(self._beta2_power, var.dtype.base_dtype) + beta1_power, beta2_power = self._get_beta_accumulators() + beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype) + beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype) lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype) beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype) beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype) diff --git a/tensorflow/contrib/opt/python/training/nadam_optimizer_test.py b/tensorflow/contrib/opt/python/training/nadam_optimizer_test.py index b0a257d264f83ae0a54cdc0e9265d6e7098b7b56..825c08a09a05894df1656a9bb6981f1862195244 100644 --- a/tensorflow/contrib/opt/python/training/nadam_optimizer_test.py +++ b/tensorflow/contrib/opt/python/training/nadam_optimizer_test.py @@ -21,12 +21,9 @@ from __future__ import print_function import numpy as np from tensorflow.contrib.opt.python.training import nadam_optimizer -from tensorflow.python.client import session from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import variables from tensorflow.python.platform import test diff --git a/tensorflow/contrib/periodic_resample/BUILD b/tensorflow/contrib/periodic_resample/BUILD index 71582f9c9a01eb221666e2c71c4a2edb18e7cb98..bd9078ae76ee27ec26c09d1aa2012f871cbdf5e9 100644 --- a/tensorflow/contrib/periodic_resample/BUILD +++ b/tensorflow/contrib/periodic_resample/BUILD @@ -6,6 +6,7 @@ exports_files(["LICENSE"]) load( "//tensorflow:tensorflow.bzl", + "py_test", "tf_gen_op_libs", "tf_custom_op_library", "tf_custom_op_py_library", @@ -64,11 +65,28 @@ py_library( "python/__init__.py", ], srcs_version = "PY2AND3", + tags = [ + "notap", + ], deps = [ ":periodic_resample_op_py", ], ) +py_test( + name = "periodic_resample_op_test", + srcs = ["python/kernel_tests/periodic_resample_op_test.py"], + srcs_version = "PY2AND3", + tags = [ + "notap", + ], + deps = [ + ":init_py", + "//tensorflow/contrib/util:util_py", + "//tensorflow/python:framework_test_lib", + ], +) + # py_library( # name = "periodic_resample_op_py", # srcs = ["python/ops/periodic_resample_op.py"], diff --git a/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.cc b/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.cc index 9cee405cef25f54fd064f8002265c42016c4fa50..e18923c8aae74c66ce78f98eb5e615e99463af74 100644 --- a/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.cc +++ b/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.cc @@ -14,13 +14,12 @@ // limitations under the License. // ============================================================================= -#include "tensorflow/core/framework/register_types.h" #include "tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h" +#include "tensorflow/core/framework/register_types.h" namespace tensorflow { -REGISTER_KERNEL_BUILDER(Name("PeriodicResample") - .Device(DEVICE_CPU), +REGISTER_KERNEL_BUILDER(Name("PeriodicResample").Device(DEVICE_CPU), PeriodicResampleOp); } // namespace tensorflow diff --git a/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h b/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h index bef21f7a5c8a27011f95eb7fae8451ca944d3cde..3ab588c45881c8f93b4c1bcdf7ccde39086a1ed7 100644 --- a/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h +++ b/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h @@ -100,6 +100,8 @@ template = input_tensor_shape.dim_size(i), + tensorflow::errors::InvalidArgument( + "periodic_resample expects the size of non-adjustable " + "dimensions be at least as large as size of input tensor." + " Dimension ", + i, " input tensor has size ", input_tensor_shape.dim_size(i), + ", desired shape has size ", desired_shape[i], ".")); + // target_dimensions[i] = desired_shape(i); target_dimensions[i] = desired_shape[i]; new_sliced_size *= target_dimensions[i]; diff --git a/tensorflow/contrib/periodic_resample/ops/array_ops.cc b/tensorflow/contrib/periodic_resample/ops/array_ops.cc index c90fc06c7fb9d79e8fd7a937e786a34947d8c1cb..82bd79695646e3673c2c78ad99dd2bd200fc2fbf 100644 --- a/tensorflow/contrib/periodic_resample/ops/array_ops.cc +++ b/tensorflow/contrib/periodic_resample/ops/array_ops.cc @@ -34,26 +34,40 @@ This function implements a slightly more generic version of the subpixel convolutions found in this [paper](https://arxiv.org/abs/1609.05158). The formula for computing the elements in the `output` tensor is as follows: + `T` = `values` tensor of rank `R` + `S` = desired `shape` of output tensor (vector of length `R`) + `P` = `output` tensor of rank `R` - \((T_1,\ldots,T_R)\) = shape(`T`) - \([S_1,\ldots,S_q,\ldots,S_R]\) = elements of vector `S` - A single element in `S` is left unspecified (denoted \(S_q=-1\)). - Let \(f_i\) denote the (possibly non-integer) factor that relates the original - dimension to the desired dimensions, \(S_i=f_i T_i\), for \(i\neq q\) where - \(f_i>0\). + \\((T_1,\\ldots,T_R)\\) = shape(`T`) + + \\([S_1,\\ldots,S_q,\\ldots,S_R]\\) = elements of vector `S` + + A single element in `S` is left unspecified (denoted \\(S_q=-1\\)). + + Let \\(f_i\\) denote the (possibly non-integer) factor that relates the original + dimension to the desired dimensions, \\(S_i=f_i T_i\\), for \\(i\\neq q\\) where + \\(f_i>0\\). + Define the following: - \(g_i=\lceil f_i\rceil\) - \(t=\prod_i T_i\) - \(s=\prod_{i\neq q} S_i\) - \(S_q\) can then be defined as by \(S_q=\lfloor t/s\rfloor\). + + \\(g_i=\\lceil f_i\\rceil\\) + + \\(t=\\prod_i T_i\\) + + \\(s=\\prod_{i\\neq q} S_i\\) + + \\(S_q\\) can then be defined by \\(S_q=\\lfloor t/s\\rfloor\\). The elements of the resulting tensor are defined as - \(P_{s_1,\ldots,s_R}=T_{h_1,\ldots,h_q,\ldots,h_R}\). - The \(h_i\) (\(i\neq q\)) are defined by \(h_i=\lfloor s_i/g_i\rfloor\). - \(h_q=S_q\sum_{j\neq q}^{q-1}G_j \mathrm{mod}(s_j,g_j) + s_q\), where - \(G_j=\prod_{i}^{j-1}g_i\) (\(G_0=1\)). + + \\(P_{s_1,\\ldots,s_R}=T_{h_1,\\ldots,h_q,\\ldots,h_R}\\). + + The \\(h_i\\) (\\(i\\neq q\\)) are defined by \\(h_i=\\lfloor s_i/g_i\\rfloor\\). + + \\(h_q=S_q\\sum_{j\\neq q}^{q-1}G_j \\mathrm{mod}(s_j,g_j) + s_q\\), where + \\(G_j=\\prod_{i}^{j-1}g_i\\) (\\(G_0=1\\)). One drawback of this method is that whenever the output dimensions are slightly less than integer multiples of the input dimensions, many of the tensor elements diff --git a/tensorflow/contrib/periodic_resample/python/kernel_tests/periodic_resample_op_test.py b/tensorflow/contrib/periodic_resample/python/kernel_tests/periodic_resample_op_test.py index 1d727870f652f3606218928983ea18e990d0afe6..a25de55e18b223db2b724aafb54b18d8f48a5baa 100644 --- a/tensorflow/contrib/periodic_resample/python/kernel_tests/periodic_resample_op_test.py +++ b/tensorflow/contrib/periodic_resample/python/kernel_tests/periodic_resample_op_test.py @@ -19,8 +19,9 @@ from __future__ import division from __future__ import print_function import numpy -import tensorflow + from tensorflow.contrib.periodic_resample import periodic_resample +from tensorflow.python.framework import errors_impl from tensorflow.python.framework import test_util from tensorflow.python.ops import variables from tensorflow.python.platform import googletest @@ -52,12 +53,11 @@ class PeriodicResampleTest(test_util.TensorFlowTestCase): def testPeriodicResampleBasic3D(self): - input_tensor = numpy.arange(2*2*4).reshape((2, 2, 4)) + input_tensor = numpy.arange(2 * 2 * 4).reshape((2, 2, 4)) desired_shape = numpy.array([4, 4, None]) - output_tensor = numpy.array([[[0], [2], [4], [6]], - [[1], [3], [5], [7]], - [[8], [10], [12], [14]], - [[9], [11], [13], [15]]]) + output_tensor = numpy.array([[[0], [2], [4], [6]], [[1], [3], [5], [7]], + [[8], [10], [12], [14]], [[9], [11], [13], + [15]]]) # NOTE: output_tensor != input_tensor.reshape((4, 4, -1)) with self.test_session(): @@ -71,24 +71,18 @@ class PeriodicResampleTest(test_util.TensorFlowTestCase): def testPeriodicResampleBasic4D(self): - input_tensor = numpy.arange(2*2*2*8).reshape((2, 2, 2, 8)) + input_tensor = numpy.arange(2 * 2 * 2 * 8).reshape((2, 2, 2, 8)) desired_shape = numpy.array([4, 4, 4, None]) - output_tensor = numpy.array([[[[0], [4], [8], [12]], - [[2], [6], [10], [14]], - [[16], [20], [24], [28]], - [[18], [22], [26], [30]]], - [[[1], [5], [9], [13]], - [[3], [7], [11], [15]], - [[17], [21], [25], [29]], - [[19], [23], [27], [31]]], - [[[32], [36], [40], [44]], - [[34], [38], [42], [46]], - [[48], [52], [56], [60]], - [[50], [54], [58], [62]]], - [[[33], [37], [41], [45]], - [[35], [39], [43], [47]], - [[49], [53], [57], [61]], - [[51], [55], [59], [63]]]]) + output_tensor = numpy.array( + [[[[0], [4], [8], [12]], [[2], [6], [10], [14]], + [[16], [20], [24], [28]], [[18], [22], [26], [30]]], + [[[1], [5], [9], [13]], [[3], [7], [11], [15]], [[17], [21], [25], + [29]], + [[19], [23], [27], + [31]]], [[[32], [36], [40], [44]], [[34], [38], [42], [46]], + [[48], [52], [56], [60]], [[50], [54], [58], [62]]], + [[[33], [37], [41], [45]], [[35], [39], [43], [47]], + [[49], [53], [57], [61]], [[51], [55], [59], [63]]]]) # NOTE: output_tensor != input_tensor.reshape((4, 4, 4, -1)) with self.test_session(): @@ -96,6 +90,19 @@ class PeriodicResampleTest(test_util.TensorFlowTestCase): result = periodic_resample(input_tensor, desired_shape).eval() self.assertAllEqual(result, output_tensor) + def testPeriodicResampleErrors(self): + input_tensor = numpy.zeros(shape=[1, 2, 2, 4]) + with self.test_session(): + variables.global_variables_initializer().run() + with self.assertRaisesWithPredicateMatch( + errors_impl.InvalidArgumentError, + 'Dimension 3 input tensor has size 4, desired shape has size 1'): + periodic_resample(input_tensor, [None, 4, 4, 1]).eval() + with self.assertRaisesWithPredicateMatch( + errors_impl.InvalidArgumentError, + '4, to be the same as the length of the desired shape, 3'): + periodic_resample(input_tensor, [None, 4, 4]).eval() + -if __name__ == "__main__": +if __name__ == '__main__': googletest.main() diff --git a/tensorflow/contrib/periodic_resample/python/ops/periodic_resample_op.py b/tensorflow/contrib/periodic_resample/python/ops/periodic_resample_op.py index 6a09f70f442131da7da2a4e98a238f21c3ccb6ec..348623d8f8d0c2ed60f559eca281343722038100 100644 --- a/tensorflow/contrib/periodic_resample/python/ops/periodic_resample_op.py +++ b/tensorflow/contrib/periodic_resample/python/ops/periodic_resample_op.py @@ -18,13 +18,14 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function - +# pylint: disable=unused-import from tensorflow.contrib.periodic_resample.python.ops import gen_periodic_resample_op from tensorflow.contrib.periodic_resample.python.ops.gen_periodic_resample_op import periodic_resample from tensorflow.contrib.util import loader from tensorflow.python.platform import resource_loader +# pylint: enable=unused-import _periodic_resample_op = loader.load_op_library( resource_loader.get_path_to_datafile('_periodic_resample_op.so')) diff --git a/tensorflow/contrib/pi_examples/camera/camera.cc b/tensorflow/contrib/pi_examples/camera/camera.cc index cb20661662922a0a160272fcd02a9fe6f9daf6fb..8110185ea8d4684fbc6b4ce54a9b192a020a97f8 100644 --- a/tensorflow/contrib/pi_examples/camera/camera.cc +++ b/tensorflow/contrib/pi_examples/camera/camera.cc @@ -17,16 +17,16 @@ limitations under the License. #include #include -#include #include #include #include #include #include #include -#include -#include #include +#include +#include +#include #include #include "tensorflow/core/framework/graph.pb.h" @@ -46,10 +46,10 @@ limitations under the License. // These are all common classes it's handy to reference with no namespace. using tensorflow::Flag; -using tensorflow::Tensor; +using tensorflow::int32; using tensorflow::Status; using tensorflow::string; -using tensorflow::int32; +using tensorflow::Tensor; // Used to store the memory-mapped buffers we use for capture. struct CameraBuffer { diff --git a/tensorflow/contrib/pi_examples/label_image/label_image.cc b/tensorflow/contrib/pi_examples/label_image/label_image.cc index 0b18045789f3a87ceb228033407d6b696bdb33f6..c6935a093f728353caeeb79a9ed85c957d87f066 100644 --- a/tensorflow/contrib/pi_examples/label_image/label_image.cc +++ b/tensorflow/contrib/pi_examples/label_image/label_image.cc @@ -23,9 +23,9 @@ limitations under the License. // // Full build instructions are at tensorflow/contrib/pi_examples/README.md. -#include #include #include +#include #include #include @@ -46,10 +46,10 @@ limitations under the License. // These are all common classes it's handy to reference with no namespace. using tensorflow::Flag; -using tensorflow::Tensor; +using tensorflow::int32; using tensorflow::Status; using tensorflow::string; -using tensorflow::int32; +using tensorflow::Tensor; // Takes a file name, and loads a list of labels from it, one per line, and // returns a vector of the strings. It pads with empty strings so the length @@ -77,23 +77,22 @@ Status ReadLabelsFile(string file_name, std::vector* result, // Error handling for JPEG decoding. void CatchError(j_common_ptr cinfo) { (*cinfo->err->output_message)(cinfo); - jmp_buf *jpeg_jmpbuf = reinterpret_cast(cinfo->client_data); + jmp_buf* jpeg_jmpbuf = reinterpret_cast(cinfo->client_data); jpeg_destroy(cinfo); longjmp(*jpeg_jmpbuf, 1); } // Decompresses a JPEG file from disk. Status LoadJpegFile(string file_name, std::vector* data, - int* width, int* height, int* channels) { + int* width, int* height, int* channels) { struct jpeg_decompress_struct cinfo; - FILE * infile; + FILE* infile; JSAMPARRAY buffer; int row_stride; if ((infile = fopen(file_name.c_str(), "rb")) == NULL) { LOG(ERROR) << "Can't open " << file_name; - return tensorflow::errors::NotFound("JPEG file ", file_name, - " not found"); + return tensorflow::errors::NotFound("JPEG file ", file_name, " not found"); } struct jpeg_error_mgr jerr; @@ -116,10 +115,11 @@ Status LoadJpegFile(string file_name, std::vector* data, data->resize((*height) * (*width) * (*channels)); row_stride = cinfo.output_width * cinfo.output_components; - buffer = (*cinfo.mem->alloc_sarray) - ((j_common_ptr) &cinfo, JPOOL_IMAGE, row_stride, 1); + buffer = (*cinfo.mem->alloc_sarray)((j_common_ptr)&cinfo, JPOOL_IMAGE, + row_stride, 1); while (cinfo.output_scanline < cinfo.output_height) { - tensorflow::uint8* row_address = &((*data)[cinfo.output_scanline * row_stride]); + tensorflow::uint8* row_address = + &((*data)[cinfo.output_scanline * row_stride]); jpeg_read_scanlines(&cinfo, buffer, 1); memcpy(row_address, buffer[0], row_stride); } @@ -141,24 +141,25 @@ Status ReadTensorFromImageFile(string file_name, const int wanted_height, int image_height; int image_channels; TF_RETURN_IF_ERROR(LoadJpegFile(file_name, &image_data, &image_width, - &image_height, &image_channels)); - LOG(INFO) << "Loaded JPEG: " << image_width << "x" << image_height - << "x" << image_channels; + &image_height, &image_channels)); + LOG(INFO) << "Loaded JPEG: " << image_width << "x" << image_height << "x" + << image_channels; const int wanted_channels = 3; if (image_channels < wanted_channels) { - return tensorflow::errors::FailedPrecondition("Image needs to have at least ", - wanted_channels, " but only has ", - image_channels); + return tensorflow::errors::FailedPrecondition( + "Image needs to have at least ", wanted_channels, " but only has ", + image_channels); } - // In these loops, we convert the eight-bit data in the image into float, resize - // it using bilinear filtering, and scale it numerically to the float range that - // the model expects (given by input_mean and input_std). + // In these loops, we convert the eight-bit data in the image into float, + // resize it using bilinear filtering, and scale it numerically to the float + // range that the model expects (given by input_mean and input_std). tensorflow::Tensor image_tensor( - tensorflow::DT_FLOAT, tensorflow::TensorShape( - {1, wanted_height, wanted_width, wanted_channels})); + tensorflow::DT_FLOAT, + tensorflow::TensorShape( + {1, wanted_height, wanted_width, wanted_channels})); auto image_tensor_mapped = image_tensor.tensor(); tensorflow::uint8* in = image_data.data(); - float *out = image_tensor_mapped.data(); + float* out = image_tensor_mapped.data(); const size_t image_rowlen = image_width * image_channels; const float width_scale = static_cast(image_width) / wanted_width; const float height_scale = static_cast(image_height) / wanted_height; @@ -166,35 +167,37 @@ Status ReadTensorFromImageFile(string file_name, const int wanted_height, const float in_y = y * height_scale; const int top_y_index = static_cast(floorf(in_y)); const int bottom_y_index = - std::min(static_cast(ceilf(in_y)), (image_height - 1)); + std::min(static_cast(ceilf(in_y)), (image_height - 1)); const float y_lerp = in_y - top_y_index; tensorflow::uint8* in_top_row = in + (top_y_index * image_rowlen); tensorflow::uint8* in_bottom_row = in + (bottom_y_index * image_rowlen); - float *out_row = out + (y * wanted_width * wanted_channels); + float* out_row = out + (y * wanted_width * wanted_channels); for (int x = 0; x < wanted_width; ++x) { const float in_x = x * width_scale; const int left_x_index = static_cast(floorf(in_x)); const int right_x_index = - std::min(static_cast(ceilf(in_x)), (image_width - 1)); + std::min(static_cast(ceilf(in_x)), (image_width - 1)); tensorflow::uint8* in_top_left_pixel = - in_top_row + (left_x_index * wanted_channels); + in_top_row + (left_x_index * wanted_channels); tensorflow::uint8* in_top_right_pixel = - in_top_row + (right_x_index * wanted_channels); + in_top_row + (right_x_index * wanted_channels); tensorflow::uint8* in_bottom_left_pixel = - in_bottom_row + (left_x_index * wanted_channels); + in_bottom_row + (left_x_index * wanted_channels); tensorflow::uint8* in_bottom_right_pixel = - in_bottom_row + (right_x_index * wanted_channels); + in_bottom_row + (right_x_index * wanted_channels); const float x_lerp = in_x - left_x_index; - float *out_pixel = out_row + (x * wanted_channels); + float* out_pixel = out_row + (x * wanted_channels); for (int c = 0; c < wanted_channels; ++c) { - const float top_left((in_top_left_pixel[c] - input_mean) / input_std); - const float top_right((in_top_right_pixel[c] - input_mean) / input_std); - const float bottom_left((in_bottom_left_pixel[c] - input_mean) / input_std); - const float bottom_right((in_bottom_right_pixel[c] - input_mean) / input_std); - const float top = top_left + (top_right - top_left) * x_lerp; - const float bottom = - bottom_left + (bottom_right - bottom_left) * x_lerp; - out_pixel[c] = top + (bottom - top) * y_lerp; + const float top_left((in_top_left_pixel[c] - input_mean) / input_std); + const float top_right((in_top_right_pixel[c] - input_mean) / input_std); + const float bottom_left((in_bottom_left_pixel[c] - input_mean) / + input_std); + const float bottom_right((in_bottom_right_pixel[c] - input_mean) / + input_std); + const float top = top_left + (top_right - top_left) * x_lerp; + const float bottom = + bottom_left + (bottom_right - bottom_left) * x_lerp; + out_pixel[c] = top + (bottom - top) * y_lerp; } } } @@ -233,10 +236,10 @@ Status GetTopLabels(const std::vector& outputs, int how_many_labels, scores.push_back(std::pair({i, unsorted_scores_flat(i)})); } std::sort(scores.begin(), scores.end(), - [](const std::pair &left, - const std::pair &right) { - return left.second > right.second; - }); + [](const std::pair& left, + const std::pair& right) { + return left.second > right.second; + }); scores.resize(how_many_labels); Tensor sorted_indices(tensorflow::DT_INT32, {scores.size()}); Tensor sorted_scores(tensorflow::DT_FLOAT, {scores.size()}); diff --git a/tensorflow/contrib/predictor/BUILD b/tensorflow/contrib/predictor/BUILD index d7c3d6c3be5c38bc89e0f65549306c0e411fb97a..a80f060b91df3b6d5e2ca9ff63c721382f0cbb0a 100644 --- a/tensorflow/contrib/predictor/BUILD +++ b/tensorflow/contrib/predictor/BUILD @@ -144,6 +144,7 @@ py_test( tags = ["no_pip"], deps = [ ":predictor_factories", + ":testing_common", ], ) diff --git a/tensorflow/contrib/predictor/core_estimator_predictor.py b/tensorflow/contrib/predictor/core_estimator_predictor.py index bd5174aef81a84488c896d259de83d0714745fee..d78d94c2699b14c80e7decee2181d190a6d91f99 100644 --- a/tensorflow/contrib/predictor/core_estimator_predictor.py +++ b/tensorflow/contrib/predictor/core_estimator_predictor.py @@ -68,10 +68,10 @@ class CoreEstimatorPredictor(predictor.Predictor): serving_input_receiver = serving_input_receiver_fn() signature_def = _get_signature_def( serving_input_receiver, estimator, output_key) - checkpoint_path = estimator.model_dir + checkpoint_dir = estimator.model_dir self._session = monitored_session.MonitoredSession( session_creator=monitored_session.ChiefSessionCreator( - checkpoint_filename_with_path=checkpoint_path)) + checkpoint_dir=checkpoint_dir)) feed_tensor_info = signature_def.inputs self._feed_tensors = {k: self._graph.get_tensor_by_name(v.name) diff --git a/tensorflow/contrib/predictor/predictor_factories.py b/tensorflow/contrib/predictor/predictor_factories.py index 9485187c5d54737120f94c8e0ec0c4a57bc1ef62..04b5d5bdf158dc6a478d7a24b538c75d1dca8d45 100644 --- a/tensorflow/contrib/predictor/predictor_factories.py +++ b/tensorflow/contrib/predictor/predictor_factories.py @@ -21,6 +21,8 @@ from __future__ import print_function from tensorflow.contrib.predictor import contrib_estimator_predictor from tensorflow.contrib.predictor import core_estimator_predictor from tensorflow.contrib.predictor import saved_model_predictor + +from tensorflow.contrib.learn.python.learn.estimators import estimator as contrib_estimator from tensorflow.python.estimator import estimator as core_estimator @@ -85,7 +87,7 @@ def from_estimator(estimator, TypeError: if `estimator` is a contrib `Estimator` instead of a core `Estimator`. """ - if isinstance(estimator, estimator.Estimator): + if isinstance(estimator, contrib_estimator.Estimator): raise TypeError('Espected estimator to be of type ' 'tf.python.estimator.Estimator, but got type ' 'tf.contrib.learn.Estimator. You likely want to call ' diff --git a/tensorflow/contrib/predictor/predictor_factories_test.py b/tensorflow/contrib/predictor/predictor_factories_test.py index 60ffeec6536e785f8830b37715222491c3a8a6fc..578d9424b25dd38f1d77a267d1fdf1ff9ff2da88 100644 --- a/tensorflow/contrib/predictor/predictor_factories_test.py +++ b/tensorflow/contrib/predictor/predictor_factories_test.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib.predictor import predictor_factories +from tensorflow.contrib.predictor import testing_common from tensorflow.python.platform import test MODEL_DIR_NAME = 'contrib/predictor/test_export_dir' @@ -46,6 +47,29 @@ class PredictorFactoriesTest(test.TestCase): with self.assertRaisesRegexp(RuntimeError, bad_tags_regex): predictor_factories.from_saved_model(self._export_dir, tags='bad_tag') + def testFromContribEstimator(self): + estimator = testing_common.get_arithmetic_estimator(core=False) + input_fn = testing_common.get_arithmetic_input_fn(core=False) + predictor_factories.from_contrib_estimator( + estimator, input_fn, output_alternative_key='sum') + + def testFromContribEstimatorWithCoreEstimatorRaises(self): + estimator = testing_common.get_arithmetic_estimator(core=True) + input_fn = testing_common.get_arithmetic_input_fn(core=True) + with self.assertRaises(TypeError): + predictor_factories.from_contrib_estimator(estimator, input_fn) + + def testFromCoreEstimator(self): + estimator = testing_common.get_arithmetic_estimator(core=True) + input_fn = testing_common.get_arithmetic_input_fn(core=True) + predictor_factories.from_estimator(estimator, input_fn) + + def testFromCoreEstimatorWithContribEstimatorRaises(self): + estimator = testing_common.get_arithmetic_estimator(core=False) + input_fn = testing_common.get_arithmetic_input_fn(core=False) + with self.assertRaises(TypeError): + predictor_factories.from_estimator(estimator, input_fn) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/py2tf/BUILD b/tensorflow/contrib/py2tf/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..d91220f6ddb859ff52d4e5853948cb667981009b --- /dev/null +++ b/tensorflow/contrib/py2tf/BUILD @@ -0,0 +1,31 @@ +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "py_test") + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +py_library( + name = "py2tf", + srcs = [ + "__init__.py", + ], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/contrib/py2tf/impl", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/contrib/py2tf/utils", + "@gast_archive//:gast", + "@six_archive//:six", + ], +) diff --git a/tensorflow/contrib/py2tf/__init__.py b/tensorflow/contrib/py2tf/__init__.py index f7fd8725e0a57737aa8294f7d1389060706697fe..379fa7fd5c2a22b5b16a21cca8c2ea8afdcaeefa 100644 --- a/tensorflow/contrib/py2tf/__init__.py +++ b/tensorflow/contrib/py2tf/__init__.py @@ -21,9 +21,16 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib.py2tf import utils +from tensorflow.contrib.py2tf.impl.api import convert +from tensorflow.contrib.py2tf.impl.api import graph_ready +from tensorflow.contrib.py2tf.impl.api import to_code +from tensorflow.contrib.py2tf.impl.api import to_graph +from tensorflow.contrib.py2tf.pyct.transformer import PyFlowParseError from tensorflow.python.util.all_util import remove_undocumented - -_allowed_symbols = [] +_allowed_symbols = [ + 'to_graph', 'to_code', 'convert', 'graph_ready', 'utils', 'PyFlowParseError' +] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/py2tf/converters/BUILD b/tensorflow/contrib/py2tf/converters/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..93c751b28dae3aa480aed839029bd37a2f47056b --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/BUILD @@ -0,0 +1,173 @@ +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "py_test") + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +py_library( + name = "converters", + srcs = [ + "asserts.py", + "break_statements.py", + "builtin_functions.py", + "call_trees.py", + "continue_statements.py", + "control_flow.py", + "decorators.py", + "for_loops.py", + "list_comprehension.py", + "logical_expressions.py", + "side_effect_guards.py", + ], + srcs_version = "PY2AND3", + visibility = ["//tensorflow:__subpackages__"], + deps = [ + "@gast_archive//:gast", + ], +) + +py_library( + name = "test_lib", + srcs = [ + "converter_test_base.py", + ], + srcs_version = "PY2AND3", + visibility = ["//tensorflow:__subpackages__"], + deps = [ + ":converters", + "//tensorflow/contrib/py2tf/pyct/static_analysis", + "//tensorflow/contrib/py2tf/utils", + "@gast_archive//:gast", + "@six_archive//:six", + ], +) + +py_test( + name = "asserts_test", + srcs = ["asserts_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "break_statements_test", + srcs = ["break_statements_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "builtin_functions_test", + srcs = ["builtin_functions_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "call_trees_test", + srcs = ["call_trees_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "continue_statements_test", + srcs = ["continue_statements_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "control_flow_test", + srcs = ["control_flow_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "decorators_test", + srcs = ["decorators_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "for_loops_test", + srcs = ["for_loops_test.py"], + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "list_comprehension_test", + srcs = ["list_comprehension_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "logical_expressions_test", + srcs = ["logical_expressions_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "side_effect_guards_test", + srcs = ["side_effect_guards_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":test_lib", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) diff --git a/tensorflow/contrib/py2tf/converters/__init__.py b/tensorflow/contrib/py2tf/converters/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ca10896ee5c6c23d9b20ff23add9945de68e5bf9 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/__init__.py @@ -0,0 +1,22 @@ +# Copyright 2016 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. +# ============================================================================== +"""Code converters used by Py2TF.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# TODO(mdan): Define a base transformer class that can recognize skip_processing +# TODO(mdan): All converters are incomplete, especially those that change blocks diff --git a/tensorflow/contrib/py2tf/converters/asserts.py b/tensorflow/contrib/py2tf/converters/asserts.py new file mode 100644 index 0000000000000000000000000000000000000000..5b9b8e772bed82df2429fd6cb94dbf7b565e22b3 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/asserts.py @@ -0,0 +1,53 @@ +# Copyright 2016 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. +# ============================================================================== +"""Converts Assert statements to their corresponding TF calls.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer + + +class AssertsTransformer(transformer.Base): + """Transforms Print nodes to Call so they can be handled as functions.""" + + # pylint:disable=invalid-name + + def visit_Assert(self, node): + self.generic_visit(node) + + # Note: The lone tf.Assert call will be wrapped with control_dependencies + # by side_effect_guards. + template = """ + tf.Assert(test, [msg]) + """ + + if node.msg is None: + return templates.replace( + template, test=node.test, msg=gast.Str('Assertion error')) + elif isinstance(node.msg, gast.Str): + return templates.replace(template, test=node.test, msg=node.msg) + else: + raise NotImplementedError('Can only convert string messages for now.') + + # pylint:enable=invalid-name + + +def transform(node, context): + return AssertsTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/asserts_test.py b/tensorflow/contrib/py2tf/converters/asserts_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6611f2777a93a7e819c8becfa06a09b27f4e6aaf --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/asserts_test.py @@ -0,0 +1,42 @@ +# Copyright 2017 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 asserts module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.converters import asserts +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.python.platform import test + + +class AssertsTest(converter_test_base.TestCase): + + def test_transform(self): + + def test_fn(a): + assert a > 0 + + node = self.parse_and_analyze(test_fn, {}) + node = asserts.transform(node, self.ctx) + + self.assertTrue(isinstance(node.body[0].body[0].value, gast.Call)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/break_statements.py b/tensorflow/contrib/py2tf/converters/break_statements.py new file mode 100644 index 0000000000000000000000000000000000000000..bfb709c5e32c6f19dc0fd109df61ece925d701a3 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/break_statements.py @@ -0,0 +1,120 @@ +# Copyright 2017 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. +# ============================================================================== +"""Canonicalizes break statements by de-sugaring into a control boolean.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno + + +class BreakCanonicalizationTransformer(transformer.Base): + """Canonicalizes continue statements into additional conditionals.""" + + def __init__(self, context): + super(BreakCanonicalizationTransformer, self).__init__(context) + # This is a stack structure, to correctly process nested loops. + self.break_uses = [] + + def _create_break_check(self): + template = """ + (not var_name) + """ + expr, = templates.replace(template, var_name=self.break_uses[-1][1]) + return expr.value + + def _create_break_trigger(self): + template = """ + var_name = True + """ + block = templates.replace(template, var_name=self.break_uses[-1][1]) + block.append(gast.Continue()) + return block + + def _create_break_init(self): + template = """ + var_name = False + """ + assign, = templates.replace(template, var_name=self.break_uses[-1][1]) + return assign + + # TODO(mdan): Surely the transformer supports this better? + def _manual_visit_list(self, block): + new_block = [] + for n in block: + new_n = self.visit(n) + if isinstance(new_n, list): + new_block.extend(new_n) + else: + new_block.append(new_n) + return new_block + + def visit_While(self, node): + self.generic_visit(node.test) + scope = anno.getanno(node, NodeAnno.BODY_SCOPE) + + break_var = self.context.namer.new_symbol('break_requested', + scope.referenced) + self.break_uses.append([False, break_var]) + node.body = self._manual_visit_list(node.body) + if self.break_uses[-1][0]: + node.test = gast.BoolOp(gast.And(), [ + node.test, + gast.UnaryOp(gast.Not(), gast.Name(break_var, gast.Load(), None)) + ]) + final_nodes = [self._create_break_init(), node] + else: + final_nodes = node + self.break_uses.pop() + + for n in node.orelse: + self.generic_visit(n) + return final_nodes + + def visit_For(self, node): + self.generic_visit(node.target) + self.generic_visit(node.iter) + scope = anno.getanno(node, NodeAnno.BODY_SCOPE) + + break_var = self.context.namer.new_symbol('break_requested', + scope.referenced) + self.break_uses.append([False, break_var]) + node.body = self._manual_visit_list(node.body) + if self.break_uses[-1][0]: + anno.setanno(node, 'extra_cond', + gast.UnaryOp(gast.Not(), + gast.Name(break_var, gast.Load(), None))) + final_nodes = [self._create_break_init(), node] + else: + final_nodes = node + self.break_uses.pop() + + for n in node.orelse: + self.generic_visit(n) + return final_nodes + + def visit_Break(self, node): + self.break_uses[-1][0] = True + return self._create_break_trigger() + + +def transform(node, context): + return BreakCanonicalizationTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/break_statements_test.py b/tensorflow/contrib/py2tf/converters/break_statements_test.py new file mode 100644 index 0000000000000000000000000000000000000000..095fcdff07d44ecc6b9bb7f8d3e2c7c43df72a02 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/break_statements_test.py @@ -0,0 +1,112 @@ +# Copyright 2017 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 break_statements module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.converters import break_statements +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.python.platform import test + + +class BreakCanonicalizationTest(converter_test_base.TestCase): + + def test_basic_break(self): + + def test_fn(x): + v = [] + while x > 0: + x -= 1 + if x % 2 == 0: + break + v.append(x) + return v + + node = self.parse_and_analyze(test_fn, {}) + node = break_statements.transform(node, self.ctx) + + with self.compiled(node) as result: + self.assertEqual(test_fn(0), result.test_fn(0)) + self.assertEqual(test_fn(1), result.test_fn(1)) + self.assertEqual(test_fn(2), result.test_fn(2)) + self.assertEqual(test_fn(3), result.test_fn(3)) + self.assertEqual(test_fn(4), result.test_fn(4)) + + def test_basic_break_for_loop(self): + + def test_fn(a): + v = [] + for x in a: + x -= 1 + if x % 2 == 0: + break + v.append(x) + return v + + # The break is incompletely canonicalized for for loops. Everything is + # in place except for the condition verification. + def test_equiv_fn(a): + v = [] + for x in a: + x -= 1 + if x % 2 == 0: + continue + v.append(x) + return v + + node = self.parse_and_analyze(test_fn, {}) + node = break_statements.transform(node, self.ctx) + + with self.compiled(node) as result: + # The break is incompletely canonicalized. Everything is in place, but + # the loop does not break. + self.assertEqual(test_equiv_fn([]), result.test_fn([])) + self.assertEqual(test_equiv_fn([1]), result.test_fn([1])) + self.assertEqual(test_equiv_fn([2]), result.test_fn([2])) + self.assertEqual( + test_equiv_fn([1, 2, 3, 4]), result.test_fn([1, 2, 3, 4])) + + def test_continue_deeply_nested(self): + + def test_fn(x): + v = [] + u = [] + w = [] + while x > 0: + x -= 1 + if x % 2 == 0: + if x % 3 != 0: + u.append(x) + else: + w.append(x) + continue + v.append(x) + return v, u, w + + node = self.parse_and_analyze(test_fn, {}) + node = break_statements.transform(node, self.ctx) + + with self.compiled(node) as result: + self.assertEqual(test_fn(0), result.test_fn(0)) + self.assertEqual(test_fn(1), result.test_fn(1)) + self.assertEqual(test_fn(2), result.test_fn(2)) + self.assertEqual(test_fn(3), result.test_fn(3)) + self.assertEqual(test_fn(4), result.test_fn(4)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/builtin_functions.py b/tensorflow/contrib/py2tf/converters/builtin_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..2eb00f90575920ac948e799b0e97a9cfccb42fad --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/builtin_functions.py @@ -0,0 +1,76 @@ +# Copyright 2016 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. +# ============================================================================== +"""Handles builtins and other special functions.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer + + +class BuiltinFunctionTransformer(transformer.Base): + """Handles builtin functions. + + This transformer only covers functions that are translated into a + TF equivalent, like `len`. + """ + + def __init__(self, context): + super(BuiltinFunctionTransformer, self).__init__(context) + + # pylint:disable=invalid-name + + def _convert_len(self, node): + template = """ + tf.shape(args)[0] + """ + return templates.replace(template, args=node.args)[0].value + + def _convert_print(self, node): + template = """ + py2tf_utils.call_print(args) + """ + return templates.replace(template, args=node.args)[0].value + + def visit_Call(self, node): + self.generic_visit(node) + # TODO(mdan): This won't work if the function was hidden. + if isinstance(node.func, gast.Name) and node.func.id == 'len': + return self._convert_len(node) + if isinstance(node.func, gast.Name) and node.func.id == 'print': + return self._convert_print(node) + return node + + def visit_Print(self, node): + self.generic_visit(node) + args = node.values + # Following is the case when calling print(a, b) + if len(args) == 1 and isinstance(args[0], gast.Tuple): + args = args[0].elts + template = """ + fname(args) + """ + function_call = templates.replace(template, fname='print', args=args)[0] + return self.visit(function_call) + + # pylint:enable=invalid-name + + +def transform(node, context): + return BuiltinFunctionTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/builtin_functions_test.py b/tensorflow/contrib/py2tf/converters/builtin_functions_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b279ff77ef10b96586d3d68585adb0d5424afb90 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/builtin_functions_test.py @@ -0,0 +1,115 @@ +# Copyright 2017 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 builtin_functions module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import sys + +import six + +from tensorflow.contrib.py2tf.converters import builtin_functions +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import logging_ops +from tensorflow.python.ops import script_ops +from tensorflow.python.platform import test + + +class BuiltinFunctionsTest(converter_test_base.TestCase): + + def test_len(self): + + def test_fn(a): + return len(a) + + node = self.parse_and_analyze(test_fn, {'len': len}) + node = builtin_functions.transform(node, self.ctx) + + with self.compiled(node, array_ops.shape) as result: + with self.test_session() as sess: + self.assertEqual(3, + sess.run( + result.test_fn(constant_op.constant([0, 0, 0])))) + + def test_print_with_op(self): + + def test_fn(a): + print(a) + + node = self.parse_and_analyze(test_fn, {'print': print}) + node = builtin_functions.transform(node, self.ctx) + + # Note: it's relevant not to include script_ops.py_func here, to verify + # that tf.Print is used. + with self.compiled(node, logging_ops.Print) as result: + with self.test_session() as sess: + try: + out_capturer = six.StringIO() + sys.stdout = out_capturer + result.test_fn('a') + sess.run(sess.graph.get_operations()) + self.assertEqual(out_capturer.getvalue(), 'a\n') + finally: + sys.stdout = sys.__stdout__ + + def test_print_with_op_multiple_values(self): + + def test_fn(a, b): + print(a, b) + + node = self.parse_and_analyze(test_fn, {'print': print}) + node = builtin_functions.transform(node, self.ctx) + + # Note: it's relevant not to include script_ops.py_func here, to verify + # that tf.Print is used. + with self.compiled(node, logging_ops.Print) as result: + with self.test_session() as sess: + try: + out_capturer = six.StringIO() + sys.stdout = out_capturer + result.test_fn('a', 1) + sess.run(sess.graph.get_operations()) + self.assertEqual(out_capturer.getvalue(), 'a 1\n') + finally: + sys.stdout = sys.__stdout__ + + def test_print_with_py_func(self): + + def test_fn(a, b, c): + print(a, b, c) + + node = self.parse_and_analyze(test_fn, {'print': print}) + node = builtin_functions.transform(node, self.ctx) + + # Note: it's relevant not to include logging_ops.Print here, to verify + # that py_func is used. + with self.compiled(node, script_ops.py_func) as result: + with self.test_session() as sess: + try: + out_capturer = six.StringIO() + sys.stdout = out_capturer + result.test_fn('a', 1, [2, 3]) + sess.run(sess.graph.get_operations()) + self.assertEqual(out_capturer.getvalue(), 'a 1 [2, 3]\n') + finally: + sys.stdout = sys.__stdout__ + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/call_trees.py b/tensorflow/contrib/py2tf/converters/call_trees.py new file mode 100644 index 0000000000000000000000000000000000000000..1050ba654c63bb52c1c5e71c981a6a0baa3fc987 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/call_trees.py @@ -0,0 +1,275 @@ +# Copyright 2016 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. +# ============================================================================== +"""Handles function calls, by generating compiled function names and calls. + +Note: this transformer does not rename the top level object being converted; +that is the caller's responsibility. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import types + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.python.util import tf_inspect + + +class FunctionNamer(object): + """Describes the interface for CallTreeTransformer's namer.""" + + def compiled_function_name(self, + original_fqn, + live_entity=None, + owner_type=None): + """Generate the name corresponding to the compiled version of a function. + + Args: + original_fqn: string or tuple(string) + live_entity: Callable, the actual target function, if known. + owner_type: Optional object. If present, it indicates that the function is + a member of the given type. + Returns: + string, bool + """ + raise NotImplementedError() + + def compiled_class_name(self, original_fqn, live_entity=None): + """Generate the name corresponding to the compiled version of a class. + + Args: + original_fqn: string or tuple(string) + live_entity: The actual target class, if known. + Returns: + string + """ + raise NotImplementedError() + + +class CallTreeTransformer(transformer.Base): + """Transforms the call tree by renaming transformed symbols.""" + + def __init__(self, context, uncompiled_modules, nocompile_decorators): + super(CallTreeTransformer, self).__init__(context) + self.uncompiled_modules = uncompiled_modules + self.nocompile_decorators = nocompile_decorators + + # pylint:disable=invalid-name + + def _resolve_name(self, node): + if isinstance(node, gast.Call): + return self._resolve_name(node.func) + if isinstance(node, gast.Name): + return self.context.namespace.get(node.id) + if isinstance(node, gast.Attribute): + parent = self._resolve_name(node.value) + if parent is not None: + return getattr(parent, node.attr) + return None + raise ValueError(node) + + def _try_resolve_target(self, node): + """Works for methods of objects of known type.""" + if anno.hasanno(node, 'live_val'): + return anno.getanno(node, 'live_val') + if isinstance(node, gast.Attribute) and anno.hasanno(node, 'type'): + owner_type = anno.getanno(node, 'type') + if hasattr(owner_type, node.attr): + return getattr(owner_type, node.attr) + else: + raise ValueError('Type "%s" has not attribute "%s". Is it dynamic?' % + (owner_type, node.attr)) + return None + + def _should_compile(self, node, fqn): + for i in range(1, len(fqn)): + if fqn[:i] in self.uncompiled_modules: + return False + + # Check for local decorations + if anno.hasanno(node, 'graph_ready'): + return False + + # The decorators themselves are not to be converted. + # If present, the decorators should appear as static functions. + target_entity = self._try_resolve_target(node.func) + if target_entity is not None: + # This attribute is set by the decorator itself. + # TODO(mdan): This may not play nicely with other wrapping decorators. + if hasattr(target_entity, '__pyct_is_compile_decorator'): + return False + + if target_entity in self.nocompile_decorators: + return False + + # Inspect the target function decorators. If any include a @convert + # or @graph_ready annotation, then they must be called as they are. + # TODO(mdan): This may be quite heavy. + # To parse and re-analize each function for every call site could be quite + # wasteful. Maybe we could cache the parsed AST? + try: + target_node, _ = parser.parse_entity(target_entity) + target_node = target_node.body[0] + except TypeError: + # Functions whose source we cannot access are compilable (e.g. wrapped + # to py_func). + return True + + for dec in target_node.decorator_list: + decorator_fn = self._resolve_name(dec) + if (decorator_fn is not None and + decorator_fn in self.nocompile_decorators): + return False + + return True + + def _determine_function_owner(self, m): + # TODO(mdan): The parent type should be known at analysis. Use that instead. + if hasattr(m, 'im_class'): # Python 2 + return m.im_class + if hasattr(m, '__qualname__'): # Python 3 + # Object attributes: should be bound to "self". + if hasattr(m, '__self__'): + return type(m.__self__) + + # Class attributes: should have the owner name in their namespace. + qn = m.__qualname__.split('.') + if len(qn) < 2: + return None + owner_name, func_name = qn[-2:] + if func_name != m.__name__: + raise ValueError('Inconsistent names detected ' + '(__qualname__[1] = "%s", __name__ = "%s") for %s.' % + (func_name, m.__name__, m)) + if owner_name == '': + return None + if owner_name not in self.context.namespace: + raise ValueError( + 'Could not resolve name "%s" while analyzing %s. Namespace:\n%s' % + (owner_name, m, self.context.namespace)) + return self.context.namespace[owner_name] + return None + + def _rename_compilable_function(self, node): + assert anno.hasanno(node.func, 'live_val') + assert anno.hasanno(node.func, 'fqn') + target_entity = anno.getanno(node.func, 'live_val') + target_fqn = anno.getanno(node.func, 'fqn') + + if not self._should_compile(node, target_fqn): + return node + + if anno.hasanno(node, 'is_constructor'): + new_name = self.context.namer.compiled_class_name( + target_fqn, live_entity=target_entity) + do_rename = True + else: + owner_type = self._determine_function_owner(target_entity) + new_name, do_rename = self.context.namer.compiled_function_name( + target_fqn, live_entity=target_entity, owner_type=owner_type) + + if do_rename: + if target_entity is not None: + if tf_inspect.ismethod(target_entity): + # The renaming process will transform it into a regular function. + # TODO(mdan): Is this complete? How does it work with nested members? + node.args = [node.func.value] + node.args + node.func = templates.replace('func_name', func_name=new_name)[0] + return node + + def _wrap_to_py_func_no_return(self, node): + # TODO(mdan): Properly handle varargs, kwargs, etc. + template = """ + py2tf_utils.wrap_py_func(func, None, (original_args,), True) + """ + return templates.replace(template, func=node.func, original_args=node.args) + + def _function_is_compilable(self, target_entity): + # TODO(mdan): This is just a placeholder. Implement. + return not isinstance(target_entity, types.BuiltinFunctionType) + + def visit_Expr(self, node): + if isinstance(node.value, gast.Call): + if anno.hasanno(node.value.func, 'live_val'): + target_entity = anno.getanno(node.value.func, 'live_val') + if not self._function_is_compilable(target_entity): + if anno.hasanno(node.value.func, 'fqn'): + target_fqn = anno.getanno(node.value.func, 'fqn') + if not self._should_compile(node.value, target_fqn): + return node + node = self._wrap_to_py_func_no_return(node.value) + return node + # Only the case of py_func with no return value is special. + # Everything else is processed by visit_Call. + self.visit(node.value) + else: + self.generic_visit(node) + return node + + def visit_Call(self, node): + # If the function is wrapped by one of the marker decorators, + # consider it graph ready. + if anno.hasanno(node.func, 'live_val'): + target_entity = anno.getanno(node.func, 'live_val') + if target_entity in self.nocompile_decorators: + if len(node.args) < 1: + raise ValueError( + 'Found call to decorator function "%s", but it had no arguments. ' + 'A decorator needs at least an argument.') + anno.setanno(node.args[0], 'graph_ready', True) + + self.generic_visit(node) + if anno.hasanno(node.func, 'live_val'): + target_entity = anno.getanno(node.func, 'live_val') + if self._function_is_compilable(target_entity): + node = self._rename_compilable_function(node) + else: + raise NotImplementedError('py_func with return values') + else: + if self.context.recursive: + raise NotImplementedError('Could not resolve target function.') + else: + # TODO(mdan): Double check. Is this reachable code? + pass + return node + + # pylint:enable=invalid-name + + +def transform(node, context, uncompiled_modules, nocompile_decorators): + """Transform function call to the compiled counterparts. + + Args: + node: AST to transform. + context: An EntityContext object. + uncompiled_modules: set of string tuples, each tuple represents the fully + qualified name of a package containing functions that will not be + compiled. + nocompile_decorators: A tuple containing decorators to be stripped from + functions during conversion. + Returns: + A tuple (node, new_names): + node: The transformed AST + new_names: set(string), containing any newly-generated names + """ + t = CallTreeTransformer(context, uncompiled_modules, nocompile_decorators) + node = t.visit(node) + return node diff --git a/tensorflow/contrib/py2tf/converters/call_trees_test.py b/tensorflow/contrib/py2tf/converters/call_trees_test.py new file mode 100644 index 0000000000000000000000000000000000000000..777648dc0b31863227262fbf931aba680bb4ed98 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/call_trees_test.py @@ -0,0 +1,118 @@ +# Copyright 2017 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 call_trees module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.converters import call_trees +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test + + +class CallTreesTest(converter_test_base.TestCase): + + def test_basic(self): + + def test_fn_1(_): + raise ValueError('This should not be called in the compiled verison.') + + def renamed_test_fn_1(a): + return a + 1 + + def test_fn_2(a): + return test_fn_1(a) + 1 + + node = self.parse_and_analyze(test_fn_2, {'test_fn_1': test_fn_1}) + node = call_trees.transform(node, self.ctx, (), ()) + + with self.compiled(node) as result: + # Only test_fn_2 is transformed, so we'll insert renamed_test_fn_1 + # manually. + result.renamed_test_fn_1 = renamed_test_fn_1 + self.assertEquals(3, result.test_fn_2(1)) + + def test_simple_methods(self): + + class TestClass(object): + + def test_fn_1(self, a): + return a + 1 + + def test_fn_2(self, a): + return self.test_fn_1(a) + 1 + + node = self.parse_and_analyze( + TestClass.test_fn_2, {'TestClass': TestClass}, + arg_types={'self': (TestClass.__name__, TestClass)}) + node = call_trees.transform(node, self.ctx, (), ()) + + with self.compiled(node) as result: + tc = TestClass() + self.assertEquals(3, result.test_fn_2(tc, 1)) + + def test_py_func_wrap_no_retval(self): + + def test_fn(a): + setattr(a, 'foo', 'bar') + + node = self.parse_and_analyze(test_fn, {'setattr': setattr}) + node = call_trees.transform(node, self.ctx, (), ()) + + with self.compiled(node) as result: + with self.test_session() as sess: + # The function has no return value, so we do some tricks to grab the + # generated py_func node and ensure its effect only happens at graph + # execution. + + class Dummy(object): + pass + + a = Dummy() + result.test_fn(a) + self.assertFalse(hasattr(a, 'foo')) + sess.run(sess.graph.get_operations()[0]) + self.assertEquals('bar', a.foo) + + def test_uncompiled_modules(self): + + def test_fn(a): + a = math_ops.multiply(a, constant_op.constant(2)) + a = math_ops.add(a, constant_op.constant(1)) + return a + + node = self.parse_and_analyze(test_fn, { + 'math_ops': math_ops, + 'constant_op': constant_op + }) + node = call_trees.transform(node, self.ctx, + set(((math_ops.__name__,), + (constant_op.__name__,))), ()) + + with self.compiled(node) as result: + result.math_ops = math_ops + result.constant_op = constant_op + with self.test_session() as sess: + # Not renamed, because the converter doesn't rename the definition + # itself (the caller is responsible for that). + result_tensor = result.test_fn(constant_op.constant(1)) + self.assertEquals(3, sess.run(result_tensor)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/continue_statements.py b/tensorflow/contrib/py2tf/converters/continue_statements.py new file mode 100644 index 0000000000000000000000000000000000000000..4069a678b118b56b59d2e5491bb80cf52efd8143 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/continue_statements.py @@ -0,0 +1,127 @@ +# Copyright 2017 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. +# ============================================================================== +"""Canonicalizes continue statements by de-sugaring into a control boolean.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno + + +class ContinueCanonicalizationTransformer(transformer.Base): + """Canonicalizes continue statements into additional conditionals.""" + + def __init__(self, context): + super(ContinueCanonicalizationTransformer, self).__init__(context) + # This is a stack structure, to correctly process nested loops. + self.continuation_uses = [] + + def _create_continuation_check(self): + template = """ + if not var_name: + pass + """ + cond, = templates.replace(template, var_name=self.continuation_uses[-1][1]) + cond.body = [] + return cond + + def _create_continuation_trigger(self): + template = """ + var_name = True + """ + assign, = templates.replace( + template, var_name=self.continuation_uses[-1][1]) + return assign + + def _create_continuation_init(self): + template = """ + var_name = False + """ + assign, = templates.replace( + template, var_name=self.continuation_uses[-1][1]) + return assign + + def _visit_and_reindent_if_necessary(self, nodes): + reorganized_nodes = [] + current_dest = reorganized_nodes + continue_used_in_block = False + for i, n in enumerate(nodes): + # TODO(mdan): This could be optimized if control structures are simple. + self.continuation_uses[-1][0] = False + n = self.visit(n) + current_dest.append(n) + if self.continuation_uses[-1][0]: + continue_used_in_block = True + if i < len(nodes) - 1: # Last statement in block needs no protection. + cond = self._create_continuation_check() + current_dest.append(cond) + current_dest = cond.body + self.continuation_uses[-1][0] = continue_used_in_block + return reorganized_nodes + + def _process_loop_block(self, block, scope): + cont_var = self.context.namer.new_symbol('cont_requested', scope.referenced) + self.continuation_uses.append([False, cont_var]) + block = self._visit_and_reindent_if_necessary(block) + if self.continuation_uses[-1][0]: + block.insert(0, self._create_continuation_init()) + self.continuation_uses.pop() + return block + + def visit_While(self, node): + self.generic_visit(node.test) + node.body = self._process_loop_block(node.body, + anno.getanno(node, + NodeAnno.BODY_SCOPE)) + for n in node.orelse: + self.generic_visit(n) + return node + + def visit_For(self, node): + self.generic_visit(node.target) + self.generic_visit(node.iter) + node.body = self._process_loop_block(node.body, + anno.getanno(node, + NodeAnno.BODY_SCOPE)) + for n in node.orelse: + self.generic_visit(n) + return node + + def visit_If(self, node): + if self.continuation_uses: + self.generic_visit(node.test) + node.body = self._visit_and_reindent_if_necessary(node.body) + continue_used_in_body = self.continuation_uses[-1][0] + node.orelse = self._visit_and_reindent_if_necessary(node.orelse) + self.continuation_uses[-1][0] = ( + continue_used_in_body or self.continuation_uses[-1][0]) + else: + node = self.generic_visit(node) + return node + + def visit_Continue(self, node): + self.continuation_uses[-1][0] = True + return self._create_continuation_trigger() + + def visit_Break(self, node): + assert False, 'break statement should be desugared at this point' + + +def transform(node, namer): + return ContinueCanonicalizationTransformer(namer).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/continue_statements_test.py b/tensorflow/contrib/py2tf/converters/continue_statements_test.py new file mode 100644 index 0000000000000000000000000000000000000000..a598dcd1aed29478b7e3fe27e3c1b20010247dd9 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/continue_statements_test.py @@ -0,0 +1,98 @@ +# Copyright 2017 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 continue_statements module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.converters import continue_statements +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.python.platform import test + + +class ContinueCanonicalizationTest(converter_test_base.TestCase): + + def test_basic_continue(self): + + def test_fn(x): + v = [] + while x > 0: + x -= 1 + if x % 2 == 0: + continue + v.append(x) + return v + + node = self.parse_and_analyze(test_fn, {}) + node = continue_statements.transform(node, self.ctx) + + with self.compiled(node) as result: + self.assertEqual(test_fn(0), result.test_fn(0)) + self.assertEqual(test_fn(1), result.test_fn(1)) + self.assertEqual(test_fn(2), result.test_fn(2)) + self.assertEqual(test_fn(3), result.test_fn(3)) + self.assertEqual(test_fn(4), result.test_fn(4)) + + def test_basic_continue_for_loop(self): + + def test_fn(a): + v = [] + for x in a: + x -= 1 + if x % 2 == 0: + continue + v.append(x) + return v + + node = self.parse_and_analyze(test_fn, {}) + node = continue_statements.transform(node, self.ctx) + + with self.compiled(node) as result: + self.assertEqual(test_fn([]), result.test_fn([])) + self.assertEqual(test_fn([1]), result.test_fn([1])) + self.assertEqual(test_fn([2]), result.test_fn([2])) + self.assertEqual(test_fn([1, 2, 3]), result.test_fn([1, 2, 3])) + + def test_continue_deeply_nested(self): + + def test_fn(x): + v = [] + u = [] + w = [] + while x > 0: + x -= 1 + if x % 2 == 0: + if x % 3 != 0: + u.append(x) + else: + w.append(x) + continue + v.append(x) + return v, u, w + + node = self.parse_and_analyze(test_fn, {}) + node = continue_statements.transform(node, self.ctx) + + with self.compiled(node) as result: + self.assertEqual(test_fn(0), result.test_fn(0)) + self.assertEqual(test_fn(1), result.test_fn(1)) + self.assertEqual(test_fn(2), result.test_fn(2)) + self.assertEqual(test_fn(3), result.test_fn(3)) + self.assertEqual(test_fn(4), result.test_fn(4)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/control_flow.py b/tensorflow/contrib/py2tf/converters/control_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..d53e3e4fd6d87004cbe55bd430346ad263e898ea --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/control_flow.py @@ -0,0 +1,221 @@ +# Copyright 2016 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. +# ============================================================================== +"""Handles control flow statements: while, if.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import ast_util +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno + + +class SymbolNamer(object): + """Describes the interface for ControlFlowTransformer's namer.""" + + def new_symbol(self, name_root, reserved_locals): + """Generate a new unique symbol. + + Args: + name_root: String, used as stem in the new name. + reserved_locals: Set(string), additional local symbols that are reserved + and which should not be used. + Returns: + String. + """ + raise NotImplementedError() + + +class ControlFlowTransformer(transformer.Base): + """Transforms control flow structures like loops an conditionals.""" + + def __init__(self, context): + super(ControlFlowTransformer, self).__init__(context) + + # pylint:disable=invalid-name + + def visit_For(self, node): + assert False, 'for statement should have been canonicalized at this point' + + def _create_cond_branch(self, body_name, aliased_orig_names, + aliased_new_names, body, returns): + if aliased_orig_names: + template = """ + def body_name(): + aliased_new_names, = aliased_orig_names, + body + return (returns,) + """ + return templates.replace( + template, + body_name=body_name, + body=body, + aliased_orig_names=aliased_orig_names, + aliased_new_names=aliased_new_names, + returns=returns) + else: + template = """ + def body_name(): + body + return (returns,) + """ + return templates.replace( + template, body_name=body_name, body=body, returns=returns) + + def _create_cond_expr(self, results, test, body_name, orelse_name): + if results is not None: + template = """ + results = py2tf_utils.run_cond(test, body_name, orelse_name) + """ + return templates.replace( + template, + test=test, + results=results, + body_name=body_name, + orelse_name=orelse_name) + else: + template = """ + py2tf_utils.run_cond(test, body_name, orelse_name) + """ + return templates.replace( + template, test=test, body_name=body_name, orelse_name=orelse_name) + + def visit_If(self, node): + self.generic_visit(node) + + body_scope = anno.getanno(node, NodeAnno.BODY_SCOPE) + orelse_scope = anno.getanno(node, NodeAnno.ORELSE_SCOPE) + + if body_scope.created - orelse_scope.created: + raise ValueError( + 'The if branch creates new symbols that the else branch does not.') + if orelse_scope.created - body_scope.created: + raise ValueError( + 'The else branch creates new symbols that the if branch does not.') + + modified = tuple(body_scope.modified | orelse_scope.modified) + all_referenced = body_scope.referenced | orelse_scope.referenced + + # Alias the closure variables inside the conditional functions + # to avoid errors caused by the local variables created in the branch + # functions. + need_alias = ( + (body_scope.modified | orelse_scope.modified) - + (body_scope.created | orelse_scope.created)) + aliased_orig_names = tuple(need_alias) + aliased_new_names = tuple( + self.context.namer.new_symbol(s.ssf(), all_referenced) + for s in aliased_orig_names) + alias_map = dict(zip(aliased_orig_names, aliased_new_names)) + node_body = ast_util.rename_symbols(node.body, alias_map) + node_orelse = ast_util.rename_symbols(node.orelse, alias_map) + + if not modified: + # When the cond would return no value, we leave the cond called without + # results. That in turn should trigger the side effect guards. The + # branch functions will return a dummy value that ensures cond + # actually has some return value as well. + results = None + elif len(modified) == 1: + results = modified[0] + else: + results = gast.Tuple([s.ast() for s in modified], None) + + body_name = self.context.namer.new_symbol('if_true', all_referenced) + orelse_name = self.context.namer.new_symbol('if_false', all_referenced) + if modified: + body_returns = tuple( + alias_map[s] if s in aliased_orig_names else s for s in modified) + else: + body_returns = templates.replace('tf.ones(())')[0].value + + body_def = self._create_cond_branch( + body_name, + aliased_orig_names=tuple(aliased_orig_names), + aliased_new_names=tuple(aliased_new_names), + body=node_body, + returns=body_returns) + orelse_def = self._create_cond_branch( + orelse_name, + aliased_orig_names=tuple(aliased_orig_names), + aliased_new_names=tuple(aliased_new_names), + body=node_orelse, + returns=body_returns) + cond_expr = self._create_cond_expr(results, node.test, body_name, + orelse_name) + + return body_def + orelse_def + cond_expr + + def visit_While(self, node): + self.generic_visit(node) + + body_scope = anno.getanno(node, NodeAnno.BODY_SCOPE) + body_closure = body_scope.modified - body_scope.created + all_referenced = body_scope.referenced + + state = list(body_closure) + state_ssf = [ + self.context.namer.new_symbol(s.ssf(), all_referenced) for s in state + ] + ssf_map = { + name: ssf + for name, ssf in zip(state, state_ssf) + if str(name) != ssf + } + + if len(state) == 1: + state = state[0] + state_ssf = state_ssf[0] + state_ast_tuple = state + else: + state_ast_tuple = gast.Tuple([n.ast() for n in state], None) + + node_body = ast_util.rename_symbols(node.body, ssf_map) + test = ast_util.rename_symbols(node.test, ssf_map) + + template = """ + def test_name(state_ssf): + return test + def body_name(state_ssf): + body + return state_ssf, + state_ast_tuple = py2tf_utils.run_while(test_name, body_name, [state]) + """ + node = templates.replace( + template, + state=state, + state_ssf=state_ssf, + state_ast_tuple=state_ast_tuple, + test_name=self.context.namer.new_symbol('loop_test', + body_scope.referenced), + test=test, + body_name=self.context.namer.new_symbol('loop_body', + body_scope.referenced), + body=node_body) + + return node + + # pylint:enable=invalid-name + + +def transform(node, context): + t = ControlFlowTransformer(context) + node = t.visit(node) + return node diff --git a/tensorflow/contrib/py2tf/converters/control_flow_test.py b/tensorflow/contrib/py2tf/converters/control_flow_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b785b284a7fb7a0257551326c88b44a341b295ba --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/control_flow_test.py @@ -0,0 +1,99 @@ +# Copyright 2017 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 control_flow module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.converters import control_flow +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.platform import test + + +class ControlFlowTest(converter_test_base.TestCase): + + def test_simple_while(self): + + def test_fn(n): + i = 0 + s = 0 + while i < n: + s += i + i += 1 + return s, i, n + + node = self.parse_and_analyze(test_fn, {}) + node = control_flow.transform(node, self.ctx) + + with self.compiled(node, control_flow_ops.while_loop) as result: + with self.test_session() as sess: + self.assertEqual((10, 5, 5), + sess.run(result.test_fn(constant_op.constant(5)))) + + def test_while_single_var(self): + + def test_fn(n): + while n > 0: + n -= 1 + return n + + node = self.parse_and_analyze(test_fn, {}) + node = control_flow.transform(node, self.ctx) + + with self.compiled(node, control_flow_ops.while_loop) as result: + with self.test_session() as sess: + self.assertEqual(0, sess.run(result.test_fn(constant_op.constant(5)))) + + def test_simple_if(self): + + def test_fn(n): + a = 0 + b = 0 + if n > 0: + a = -n + else: + b = 2 * n + return a, b + + node = self.parse_and_analyze(test_fn, {}) + node = control_flow.transform(node, self.ctx) + + with self.compiled(node, control_flow_ops.cond) as result: + with self.test_session() as sess: + self.assertEqual((-1, 0), + sess.run(result.test_fn(constant_op.constant(1)))) + self.assertEqual((0, -2), + sess.run(result.test_fn(constant_op.constant(-1)))) + + def test_if_single_var(self): + + def test_fn(n): + if n > 0: + n = -n + return n + + node = self.parse_and_analyze(test_fn, {}) + node = control_flow.transform(node, self.ctx) + + with self.compiled(node, control_flow_ops.cond) as result: + with self.test_session() as sess: + self.assertEqual(-1, sess.run(result.test_fn(constant_op.constant(1)))) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/converter_test_base.py b/tensorflow/contrib/py2tf/converters/converter_test_base.py new file mode 100644 index 0000000000000000000000000000000000000000..67747183dd323a799a04943ce4c7fe8c4093d002 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/converter_test_base.py @@ -0,0 +1,103 @@ +# Copyright 2017 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. +# ============================================================================== +"""Base class for tests in this module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import contextlib +import imp + +from tensorflow.contrib.py2tf import utils +from tensorflow.contrib.py2tf.pyct import compiler +from tensorflow.contrib.py2tf.pyct import context +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import qual_names +from tensorflow.contrib.py2tf.pyct.static_analysis import activity +from tensorflow.contrib.py2tf.pyct.static_analysis import live_values +from tensorflow.contrib.py2tf.pyct.static_analysis import type_info +from tensorflow.python.platform import test + + +class FakeNamer(object): + + def new_symbol(self, name_root, used): + i = 0 + while True: + name = '%s%d' % (name_root, i) + if name not in used: + return name + i += 1 + + def compiled_function_name(self, + original_fqn, + live_entity=None, + owner_type=None): + del live_entity + if owner_type is not None: + return None, False + return ('renamed_%s' % '_'.join(original_fqn)), True + + +class TestCase(test.TestCase): + """Base class for unit tests in this module. Contains relevant utilities.""" + + @contextlib.contextmanager + def compiled(self, node, *symbols): + source = '' + try: + result, source = compiler.ast_to_object(node) + result.tf = self.make_fake_tf(*symbols) + result.py2tf_utils = utils + yield result + except Exception: # pylint:disable=broad-except + print('Offending compiled code:\n%s' % source) + raise + + def make_fake_tf(self, *symbols): + fake_tf = imp.new_module('fake_tf') + for s in symbols: + setattr(fake_tf, s.__name__, s) + return fake_tf + + def attach_namespace(self, module, **ns): + for k, v in ns.items(): + setattr(module, k, v) + + def parse_and_analyze(self, + test_fn, + namespace, + namer=None, + arg_types=None, + include_type_analysis=True, + recursive=True): + node, source = parser.parse_entity(test_fn) + ctx = context.EntityContext( + namer=namer or FakeNamer(), + source_code=source, + source_file=None, + namespace=namespace, + arg_values=None, + arg_types=arg_types, + recursive=recursive) + node = qual_names.resolve(node) + node = activity.resolve(node, ctx) + node = live_values.resolve(node, ctx, {}) + if include_type_analysis: + node = type_info.resolve(node, ctx) + node = live_values.resolve(node, ctx, {}) + self.ctx = ctx + return node diff --git a/tensorflow/contrib/py2tf/converters/decorators.py b/tensorflow/contrib/py2tf/converters/decorators.py new file mode 100644 index 0000000000000000000000000000000000000000..3f620c1cd2d9b75f82410754a7e812e13eabe3ae --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/decorators.py @@ -0,0 +1,62 @@ +# Copyright 2016 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. +# ============================================================================== +"""Handles decorators. + +Note: this module only deals with functions whose decorators are still recorded +in the AST. This does not always happen. See the unit test for an example. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import pretty_printer + + +class DecoratorsTransformer(gast.NodeTransformer): + """Converts or removes decorators.""" + + def __init__(self, remove_decorators): + self.remove_decorators = remove_decorators + + # pylint:disable=invalid-name + + def visit_FunctionDef(self, node): + self.generic_visit(node) + kept_decorators = [] + for dec in node.decorator_list: + if isinstance(dec, gast.Call): + dec_func = dec.func + else: + dec_func = dec + if not anno.hasanno(dec_func, 'live_val'): + raise ValueError( + 'Could not resolve decorator: %s' % pretty_printer.fmt(dec_func)) + dec_value = anno.getanno(dec_func, 'live_val') + if dec_value not in self.remove_decorators: + kept_decorators.append(dec) + node.decorator_list = kept_decorators + return node + + # pylint:enable=invalid-name + + +def transform(node, remove_decorators): + transformer = DecoratorsTransformer(remove_decorators) + node = transformer.visit(node) + return node diff --git a/tensorflow/contrib/py2tf/converters/decorators_test.py b/tensorflow/contrib/py2tf/converters/decorators_test.py new file mode 100644 index 0000000000000000000000000000000000000000..402fa0dda28e696f70d0354ca4abf3a6c83506d9 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/decorators_test.py @@ -0,0 +1,102 @@ +# Copyright 2017 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 decorators module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import textwrap + +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.contrib.py2tf.converters import decorators +from tensorflow.contrib.py2tf.pyct import compiler +from tensorflow.python.platform import test +from tensorflow.python.util import tf_inspect + + +class DecoratorsTest(converter_test_base.TestCase): + + def test_function_decorator(self): + + def function_decorator(): + + def decorator(f): + return lambda a: f(a) + 1 + + return decorator + + # The Python parser does capture decorators into the AST. + # However, the interpreter desugars them on load, and refering to the + # decorated function at runtime usually loses any trace of the decorator. + # Below is an example when that doesn't happen. + def static_wrapper(): + + @function_decorator() + def test_fn(a): # pylint:disable=unused-variable + return a + + node = self.parse_and_analyze(static_wrapper, + {'function_decorator': function_decorator}) + node = node.body[0].body[0] + + node = decorators.transform(node, remove_decorators=()) + # Since the decorator is not removed, we need to include its source + # code. We cannot do it after the fact because decorators are executed + # on load. + result, _ = compiler.ast_to_object( + node, + source_prefix=textwrap.dedent(tf_inspect.getsource(function_decorator))) + self.assertEqual(2, result.test_fn(1)) + + node = decorators.transform(node, remove_decorators=(function_decorator,)) + with self.compiled(node) as result: + self.assertEqual(1, result.test_fn(1)) + + def test_simple_decorator(self): + + def simple_decorator(f): + return lambda a: f(a) + 1 + + # The Python parser does capture decorators into the AST. + # However, the interpreter desugars them upon load, and refering to the + # decorated function at runtime usually loses any trace of the decorator. + # Below is an example when that doesn't happen. + def static_wrapper(): + + @simple_decorator + def test_fn(a): # pylint:disable=unused-variable + return a + + node = self.parse_and_analyze(static_wrapper, + {'simple_decorator': simple_decorator}) + node = node.body[0].body[0] + + node = decorators.transform(node, remove_decorators=()) + # Since the decorator is not removed, we need to include its source + # code. We cannot do it after the fact because decorators are executed + # on load. + result, _ = compiler.ast_to_object( + node, + source_prefix=textwrap.dedent(tf_inspect.getsource(simple_decorator))) + self.assertEqual(2, result.test_fn(1)) + + node = decorators.transform(node, remove_decorators=(simple_decorator,)) + with self.compiled(node) as result: + self.assertEqual(1, result.test_fn(1)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/for_loops.py b/tensorflow/contrib/py2tf/converters/for_loops.py new file mode 100644 index 0000000000000000000000000000000000000000..935dade0ed30975dd29c8ffe5be875993936d241 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/for_loops.py @@ -0,0 +1,86 @@ +# Copyright 2017 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. +# ============================================================================== +"""Canonicalizes for loops into while loops. + +This canonicalizer uses the len function on its argument. That should be +converted to a tf.shape separately. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno + + +class ForLoopCanonicalizationTransformer(transformer.Base): + """Canonicalizes for loops (e.g. into while loops).""" + + def __init__(self, context): + super(ForLoopCanonicalizationTransformer, self).__init__(context) + + def visit_For(self, node): + self.generic_visit(node) + body_scope = anno.getanno(node, NodeAnno.BODY_SCOPE) + + if anno.hasanno(node, 'extra_cond'): + template = """ + i = 0 + n = len(loop_iter) + while i < n and extra_cond: + # TODO(mdan): Use TensorListFromTensor(loop_iter) here. + target = loop_iter[i] + body + i += 1 + """ + return templates.replace( + template, + loop_iter=node.iter, + target=node.target, + body=node.body, + i=self.context.namer.new_symbol('i', body_scope.referenced), + n=self.context.namer.new_symbol('n', body_scope.referenced), + extra_cond=anno.getanno(node, 'extra_cond')) + else: + template = """ + i = 0 + n = len(loop_iter) + while i < n: + # TODO(mdan): Use TensorListFromTensor(loop_iter) here. + target = loop_iter[i] + body # pylint:disable=pointless-statement + i += 1 + """ + repl = templates.replace( + template, + loop_iter=node.iter, + target=node.target, + body=node.body, + i=self.context.namer.new_symbol('i', body_scope.referenced), + n=self.context.namer.new_symbol('n', body_scope.referenced)) + return repl + + def visit_Continue(self, node): + assert False, 'continue statement should be desugared at this point' + + def visit_Break(self, node): + assert False, 'break statement should be desugared at this point' + + +def transform(node, context): + return ForLoopCanonicalizationTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/for_loops_test.py b/tensorflow/contrib/py2tf/converters/for_loops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..70a367d3b517e528b67f260d607431d324d2ab7d --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/for_loops_test.py @@ -0,0 +1,47 @@ +# Copyright 2017 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 for_loops module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.contrib.py2tf.converters import for_loops +from tensorflow.python.platform import test + + +class ControlFlowTest(converter_test_base.TestCase): + + def test_basic_for(self): + + def test_fn(l): + s = 0 + for e in l: + s += e + return s + + node = self.parse_and_analyze(test_fn, {}) + node = for_loops.transform(node, self.ctx) + + with self.compiled(node) as result: + l = [1, 2, 3] + self.assertEqual(test_fn(l), result.test_fn(l)) + l = [] + self.assertEqual(test_fn(l), result.test_fn(l)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/list_comprehension.py b/tensorflow/contrib/py2tf/converters/list_comprehension.py new file mode 100644 index 0000000000000000000000000000000000000000..e8744831100e4852919b5cd1253b74acea4d790d --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/list_comprehension.py @@ -0,0 +1,80 @@ +# Copyright 2016 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. +# ============================================================================== +"""Canonicalizing list comprehensions into for and if statements. + +e.g. +result = [x * x for x in xs] + +becomes + +result = [] +for x in xs: + elt = x * x + result.append(elt) +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer + + +class ListCompCanonicalizationTransformer(transformer.Base): + """NodeTransformer to canonicalize list comprehensions.""" + + def __init__(self, context): + super(ListCompCanonicalizationTransformer, self).__init__(context) + + def make_update_list_node(self, list_, elt): + return templates.replace('list_.append(elt)', list_=list_, elt=elt)[0] + + def instantiate_list_node(self): + return parser.parse_str('[]').body[0].value + + def visit_Assign(self, node): + if not isinstance(node.value, gast.ListComp): + return node + if len(node.targets) > 1: + raise ValueError('Only support single assignment.') + return self.canonicalize_listcomp(node.targets[0], node.value) + + def canonicalize_listcomp(self, result_node, list_comp_node): + + make_list = templates.replace( + 'list_ = create_list', + list_=result_node, + create_list=self.instantiate_list_node()) + loop_body = self.make_update_list_node(result_node, list_comp_node.elt) + + for gen in reversed(list_comp_node.generators): + for gen_if in reversed(gen.ifs): + loop_body = templates.replace( + 'if test: loop_body', test=gen_if, loop_body=loop_body) + loop_body = templates.replace( + 'for target in iter_: loop_body', + iter_=gen.iter, + target=gen.target, + loop_body=loop_body) + + return make_list + loop_body + + +def transform(node, context): + return ListCompCanonicalizationTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/list_comprehension_test.py b/tensorflow/contrib/py2tf/converters/list_comprehension_test.py new file mode 100644 index 0000000000000000000000000000000000000000..025fac11e41e6771fbb9b80ff3da70dc3ceec73e --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/list_comprehension_test.py @@ -0,0 +1,75 @@ +# Copyright 2017 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 list_comprehension module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.contrib.py2tf.converters import list_comprehension +from tensorflow.python.platform import test + + +class ListCompTest(converter_test_base.TestCase): + + def test_basic(self): + + def test_fn(l): + s = [e * e for e in l] + return s + + node = self.parse_and_analyze(test_fn, {}) + node = list_comprehension.transform(node, self.ctx) + + with self.compiled(node) as result: + l = [1, 2, 3] + self.assertEqual(test_fn(l), result.test_fn(l)) + l = [] + self.assertEqual(test_fn(l), result.test_fn(l)) + + def test_multiple_generators(self): + + def test_fn(l): + s = [e * e for sublist in l for e in sublist] + return s + + node = self.parse_and_analyze(test_fn, {}) + node = list_comprehension.transform(node, self.ctx) + + with self.compiled(node) as result: + l = [[1], [2], [3]] + self.assertEqual(test_fn(l), result.test_fn(l)) + l = [] + self.assertEqual(test_fn(l), result.test_fn(l)) + + def test_conds(self): + + def test_fn(l): + s = [e * e for e in l if e > 1] + return s + + node = self.parse_and_analyze(test_fn, {}) + node = list_comprehension.transform(node, self.ctx) + + with self.compiled(node) as result: + l = [1, 2, 3] + self.assertEqual(test_fn(l), result.test_fn(l)) + l = [] + self.assertEqual(test_fn(l), result.test_fn(l)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/logical_expressions.py b/tensorflow/contrib/py2tf/converters/logical_expressions.py new file mode 100644 index 0000000000000000000000000000000000000000..df980d41c9c57e325bee9a1fa870d9c95f46ea41 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/logical_expressions.py @@ -0,0 +1,74 @@ +# Copyright 2016 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. +# ============================================================================== +"""Converter for logical expressions. + +e.g. `a and b -> tf.logical_and(a, b)`. This is not done automatically in TF. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import parser + + +class LogicalExpressionTransformer(gast.NodeTransformer): + """Converts logical expressions to corresponding TF calls.""" + + def __init__(self): + # TODO(mdan): Look into replacing with bitwise operators instead. + self.op_mapping = { + gast.And: 'tf.logical_and', + gast.Or: 'tf.logical_or', + gast.Not: 'tf.logical_not', + gast.Eq: 'tf.equal', + } + + def visit_Compare(self, node): + node = self.generic_visit(node) + if len(node.ops) > 1: + raise NotImplementedError() + cmp_type = type(node.ops[0]) + if cmp_type in self.op_mapping: + tf_function = parser.parse_str(self.op_mapping[cmp_type]).body[0].value + return gast.Call( + func=tf_function, args=[node.left, node.comparators[0]], keywords=[]) + return node + + def visit_UnaryOp(self, node): + node = self.generic_visit(node) + if isinstance(node.op, gast.Not): + tf_function = parser.parse_str(self.op_mapping[type( + node.op)]).body[0].value + node = gast.Call(func=tf_function, args=[node.operand], keywords=[]) + return node + + def visit_BoolOp(self, node): + # TODO(mdan): A normalizer may be useful here. Use ANF? + node = self.generic_visit(node) + tf_function = parser.parse_str(self.op_mapping[type(node.op)]).body[0].value + left = node.values[0] + for i in range(1, len(node.values)): + left = gast.Call( + func=tf_function, args=[left, node.values[i]], keywords=[]) + return left + + +def transform(node): + transformer = LogicalExpressionTransformer() + node = transformer.visit(node) + return node diff --git a/tensorflow/contrib/py2tf/converters/logical_expressions_test.py b/tensorflow/contrib/py2tf/converters/logical_expressions_test.py new file mode 100644 index 0000000000000000000000000000000000000000..a28326c517d468230f35e45f0fbfe5257d769895 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/logical_expressions_test.py @@ -0,0 +1,57 @@ +# Copyright 2017 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 logical_expressions module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.contrib.py2tf.converters import logical_expressions +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test + + +class GradientsFunctionTest(converter_test_base.TestCase): + + def test_equals(self): + + def test_fn(a, b): + return a == b + + node = self.parse_and_analyze(test_fn, {}) + node = logical_expressions.transform(node) + + with self.compiled(node, math_ops.equal) as result: + with self.test_session() as sess: + self.assertTrue(sess.run(result.test_fn(1, 1))) + self.assertFalse(sess.run(result.test_fn(1, 2))) + + def test_bool_ops(self): + + def test_fn(a, b, c): + return (a or b) and (a or b or c) + + node = self.parse_and_analyze(test_fn, {}) + node = logical_expressions.transform(node) + + with self.compiled(node, math_ops.logical_or, + math_ops.logical_and) as result: + with self.test_session() as sess: + self.assertTrue(sess.run(result.test_fn(True, False, True))) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/converters/side_effect_guards.py b/tensorflow/contrib/py2tf/converters/side_effect_guards.py new file mode 100644 index 0000000000000000000000000000000000000000..30976b3ec6db5a6607023ac804d9d54cfb296190 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/side_effect_guards.py @@ -0,0 +1,190 @@ +# Copyright 2016 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. +# ============================================================================== +"""Adds guards against function calls with side effects. + +Only standalone calls are guarded. + +WARNING: This mechanism is incomplete. Particularly, it only guards the +arguments passed to functions, and does not account for indirectly modified +state. + +Example: + y = tf.layers.dense(x) # Creates TF variable 'foo' + loss = loss(y) + opt.minimize(loss) # indirectly affects 'foo' + z = tf.get_variable('foo') # Indirectly affects `loss` and 'foo' + # Here, `loss` can be guarded. But `z` cannot. + +# TODO(mdan): We should probably define a safe mode where we guard everything. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import ast_util +from tensorflow.contrib.py2tf.pyct import qual_names +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno + + +class SymbolNamer(object): + """Describes the interface for SideEffectGuardTransformer's namer.""" + + def new_symbol(self, name_root, reserved_locals): + """Generate a new unique function_name. + + Args: + name_root: String, used as stem in the new name. + reserved_locals: Set(string), additional local symbols that are reserved. + Returns: + String. + """ + raise NotImplementedError() + + +class SideEffectGuardTransformer(transformer.Base): + """Adds control dependencies to functions with side effects.""" + + def __init__(self, context): + super(SideEffectGuardTransformer, self).__init__(context) + + # pylint:disable=invalid-name + + def _visit_and_reindent(self, nodes): + new_nodes = [] + current_dest = new_nodes + alias_map = {} + reindent_requested = False + for n in nodes: + n = self.visit(n) + # NOTE: the order in which these statements execute is important; in + # particular, watch out for ending up with cycles in the AST. + if alias_map: + n = ast_util.rename_symbols(n, alias_map) + if isinstance(n, (list, tuple)): + current_dest.extend(n) + else: + current_dest.append(n) + if anno.hasanno(n, anno.Basic.INDENT_BLOCK_REMAINDER): + reindent_requested = True + new_dest, new_alias_map = anno.getanno( + n, anno.Basic.INDENT_BLOCK_REMAINDER) + anno.delanno(n, anno.Basic.INDENT_BLOCK_REMAINDER) + new_alias_map.update(alias_map) + alias_map = new_alias_map + current_dest = new_dest + if reindent_requested and not current_dest: + # TODO(mdan): There may still be something that could be done. + raise ValueError('Unable to insert statement into the computation flow: ' + 'it is not followed by any computation which ' + 'the statement could gate.') + return new_nodes + + def visit_FunctionDef(self, node): + node.body = self._visit_and_reindent(node.body) + return node + + def visit_With(self, node): + node.body = self._visit_and_reindent(node.body) + return node + + def visit_If(self, node): + node.body = self._visit_and_reindent(node.body) + node.orelse = self._visit_and_reindent(node.orelse) + return node + + def visit_While(self, node): + node.body = self._visit_and_reindent(node.body) + node.orelse = self._visit_and_reindent(node.orelse) + return node + + def visit_Expr(self, node): + self.generic_visit(node) + if isinstance(node.value, gast.Call): + # Patterns of single function calls, like: + # opt.minimize(loss) + # or: + # tf.py_func(...) + + # First, attempt to gate future evaluation of args. If that's not + # possible, gate all remaining statements (and that may fail too, see + # _visit_and_reindent. + args_scope = anno.getanno(node.value, NodeAnno.ARGS_SCOPE) + # NOTE: We can't guard object attributes because they may not be writable. + # In addition, avoid renaming well-known names. + # TODO(mdan): Move these names into config. + unguarded_names = (qual_names.QN('self'), qual_names.QN('tf')) + guarded_args = tuple(s for s in args_scope.used + if not s.is_composite() and s not in unguarded_names) + + # TODO(mdan): Include all arguments which depended on guarded_args too. + # For example, the following will still cause a race: + # tf.assign(a, a + 1) + # b = a + 1 + # tf.assign(a, a + 1) # Control deps here should include `b` + # c = b + 1 + # Or maybe we should just raise an "unsafe assign" error? + + if guarded_args: + # The aliases may need new names to avoid incorrectly making them local. + # TODO(mdan): This is brutal. It will even rename modules - any fix? + need_alias = tuple( + s for s in guarded_args if s not in args_scope.parent.modified) + aliased_new_names = tuple( + qual_names.QN( + self.context.namer.new_symbol( + s.ssf(), args_scope.parent.referenced)) for s in need_alias) + alias_map = dict(zip(need_alias, aliased_new_names)) + if len(guarded_args) == 1: + s, = guarded_args + aliased_guarded_args = alias_map.get(s, s) + else: + aliased_guarded_args = gast.Tuple( + [alias_map.get(s, s).ast() for s in guarded_args], None) + + template = """ + with py2tf_utils.control_dependency_on_returns(call): + aliased_guarded_args = py2tf_utils.alias_tensors(guarded_args) + """ + control_deps_guard = templates.replace( + template, + call=node.value, + aliased_guarded_args=aliased_guarded_args, + guarded_args=guarded_args)[-1] + else: + alias_map = {} + + template = """ + with py2tf_utils.control_dependency_on_returns(call): + pass + """ + control_deps_guard = templates.replace(template, call=node.value)[-1] + control_deps_guard.body = [] + + node = control_deps_guard + anno.setanno(node, anno.Basic.INDENT_BLOCK_REMAINDER, + (node.body, alias_map)) + return node + + # pylint:enable=invalid-name + + +def transform(node, context): + return SideEffectGuardTransformer(context).visit(node) diff --git a/tensorflow/contrib/py2tf/converters/side_effect_guards_test.py b/tensorflow/contrib/py2tf/converters/side_effect_guards_test.py new file mode 100644 index 0000000000000000000000000000000000000000..463db2e770213ba9636d2537b095a77dece5d8f6 --- /dev/null +++ b/tensorflow/contrib/py2tf/converters/side_effect_guards_test.py @@ -0,0 +1,165 @@ +# Copyright 2017 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 side_effect_guards module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.converters import converter_test_base +from tensorflow.contrib.py2tf.converters import side_effect_guards +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import errors_impl +from tensorflow.python.framework import ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +class SideEffectGuardsTest(converter_test_base.TestCase): + + def test_side_effect_on_return_only_variable(self): + + tf = None + + def test_fn(a): + tf.assign(a, a + 1) + return a + + node = self.parse_and_analyze(test_fn, {}) + node = side_effect_guards.transform(node, self.ctx) + + with self.compiled(node, state_ops.assign) as result: + self.assertEqual(len(node.body[0].body), 1) + with self.test_session() as sess: + v = variables.Variable(2) + sess.run(v.initializer) + # NOTE: We don't expect the assignment to execute in this case, because + # variables cannot be reliably guarded. + self.assertEqual(2, sess.run(result.test_fn(v))) + + def test_side_effect_on_used_variable(self): + + tf = None + + def test_fn(a): + tf.assign(a, a + 1) + return a + 1 + + node = self.parse_and_analyze(test_fn, {}) + node = side_effect_guards.transform(node, self.ctx) + + with self.compiled(node, state_ops.assign) as result: + self.assertEqual(len(node.body[0].body), 1) + with self.test_session() as sess: + v = variables.Variable(2) + sess.run(v.initializer) + # NOTE: Unlike test_side_effect_on_return_only_variable, the variable + # was used in the local scope and so we could catch the assign's side + # effect. + self.assertEqual(4, sess.run(result.test_fn(v))) + + def test_side_effect_on_tensor(self): + + tf = None + + def test_fn(a): + tf.Assert(a > 0, ['expected in throw']) + return a + + node = self.parse_and_analyze(test_fn, {}) + node = side_effect_guards.transform(node, self.ctx) + + with self.compiled(node, control_flow_ops.Assert) as result: + self.assertEqual(len(node.body[0].body), 1) + with self.test_session() as sess: + # NOTE: In this case we can also capture the side effect because the + # argument is a tensor ans we can wrap it inside an identity. + with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, + 'expected in throw'): + sess.run(result.test_fn(constant_op.constant(-1))) + + def test_multiline_block(self): + + tf = None + + def test_fn(a): + tf.assign(a, a + 1) + b = a + 1 + tf.assign(a, b + 1) + c = b + 1 + d = c + 1 + return d + + node = self.parse_and_analyze(test_fn, {}) + node = side_effect_guards.transform(node, self.ctx) + + with self.compiled(node, state_ops.assign) as result: + self.assertEqual(len(node.body[0].body), 1) + with self.test_session() as sess: + v = variables.Variable(2) + sess.run(v.initializer) + self.assertEqual(6, sess.run(result.test_fn(v))) + + def test_multiline_nested_block(self): + + tf = None + + def test_fn(a): + with tf.name_scope('foo'): + tf.assign(a, a + 1) + b = a + 1 + c = b + 1 + d = c + 1 + return d + + node = self.parse_and_analyze(test_fn, {}) + node = side_effect_guards.transform(node, self.ctx) + + with self.compiled(node, state_ops.assign, ops.name_scope) as result: + self.assertEqual(len(node.body[0].body[0].body), 1) + with self.test_session() as sess: + v = variables.Variable(2) + sess.run(v.initializer) + self.assertEqual(6, sess.run(result.test_fn(v))) + + def test_multiline_block_unsafe(self): + + tf = None + + def test_fn(a): + tf.assign(a, a + 1) + b = a + 1 + tf.assign(a, a + 1) + c = b + 1 + d = c + 1 + return d + + node = self.parse_and_analyze(test_fn, {}) + node = side_effect_guards.transform(node, self.ctx) + + with self.compiled(node, state_ops.assign) as result: + self.assertEqual(len(node.body[0].body), 1) + with self.test_session() as sess: + v = variables.Variable(2) + sess.run(v.initializer) + # NOTE: This intentionally highlights the flakiness. The test should be + # tightened down once that is solved. + self.assertTrue(sess.run(result.test_fn(v)) in (6, 7)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/impl/BUILD b/tensorflow/contrib/py2tf/impl/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..90ffabbc9bf4524ec2ebf54b6dd847bd8768a486 --- /dev/null +++ b/tensorflow/contrib/py2tf/impl/BUILD @@ -0,0 +1,67 @@ +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "py_test") + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +py_library( + name = "impl", + srcs = [ + "api.py", + "config.py", + "conversion.py", + "naming.py", + ], + srcs_version = "PY2AND3", + visibility = ["//tensorflow:__subpackages__"], + deps = [ + "//tensorflow/contrib/py2tf/converters", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/contrib/py2tf/pyct/static_analysis", + "//tensorflow/contrib/py2tf/utils", + "@gast_archive//:gast", + "@six_archive//:six", + ], +) + +py_test( + name = "api_test", + srcs = ["api_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":impl", + "//tensorflow/contrib/py2tf/utils", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "conversion_test", + srcs = ["conversion_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":impl", + "//tensorflow/python:client_testlib", + "@gast_archive//:gast", + ], +) + +py_test( + name = "naming_test", + srcs = ["naming_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":impl", + "//tensorflow/python:client_testlib", + ], +) diff --git a/tensorflow/contrib/py2tf/impl/api.py b/tensorflow/contrib/py2tf/impl/api.py new file mode 100644 index 0000000000000000000000000000000000000000..8ae1c701698ae9a4efbde45222ff6c3db6e92521 --- /dev/null +++ b/tensorflow/contrib/py2tf/impl/api.py @@ -0,0 +1,233 @@ +# Copyright 2016 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. +# ============================================================================== +"""Public API.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from functools import wraps + +import gast +import six + +from tensorflow.contrib.py2tf.impl import config +from tensorflow.contrib.py2tf.impl import conversion +from tensorflow.contrib.py2tf.pyct import compiler +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util import tf_inspect + +# TODO(mdan): Properly document the type hints. +# TODO(mdan): Reduce the type hint information to (module, type). +# (currently we require (module + class name, type)) + + +def graph_ready(f): + """No-op decorator that explicitly marks a function as graph-ready. + + Graph-ready functions are assumed to not need any conversion. + + Args: + f: Any callable. + Returns: + f itself. + """ + setattr(f, '__pyct_is_compile_decorator', True) + return f + + +def convert_inline(f, *args, **kwargs): + """Shorthand to convert and call a function. + + For example, the following two statements are equivalent: + + @convert() + def foo(): + ... + foo(bar) + + def foo(): + ... + convert_inline(foo, bar) + + Args: + f: Function to convert. Only this call will be converted. + *args: Passed through to f. + **kwargs: Passed through to f, with the following exceptions: + * arg_value_hints: A dict mapping parameter names to objects that can + hint at the type of those parameters. + + Returns: + The result of the converted f applied to args and kwargs. + """ + if 'arg_value_hints' in kwargs: + arg_value_hints = kwargs['arg_value_hints'] + del kwargs['arg_value_hints'] + else: + arg_value_hints = None + if tf_inspect.ismethod(f): + # When converting methods, the result is still an unbound function. + args = (f.__self__,) + args + return convert(arg_value_hints)(f)(*args, **kwargs) + + +def convert(recursive=False, verbose=False, arg_types=None): + """Decorator that compiles a function to graph mode. + + The decorator is dynamic - invoking compilation whenever the decorated + function is called. This means the parameter values are known at compilation. + + Args: + recursive: Whether to recusrively convert any functions that the decorator + function may call. + verbose: Whether to output the compiled code in the logs. + arg_types: See to_graph. + + Returns: + A decorator that compiles the given function to graph mode. + + Raises: + ValueError: If any of the arguments are illegal. + """ + if arg_types is None: + arg_types = {} + + def decorator(f): + """Decorator implementation.""" + + @wraps(f) + def wrapper(*args, **kwargs): + """Wrapper that calls the compiled version of the wrapped function.""" + partial_types = () + arg_values = {} + arg_names = tf_inspect.getargspec(f)[0] + for name, arg in zip(arg_names, args): + arg_values[name] = arg + arg_class = arg.__class__ + # If arg_value_hints specifies any name, use that instead. + if name not in arg_types: + arg_types[name] = (arg_class.__name__, arg_class) + if name == 'self' and tf_inspect.isclass(arg_class): + # Annotated methods need to specify that their owner type is partial, + # otherwise other members they call will not be converted. + partial_types = (arg_class,) + wrapped = to_graph( + f, + recursive=recursive, + verbose=verbose, + arg_values=arg_values, + arg_types=arg_types, + partial_types=partial_types) + return wrapped(*args, **kwargs) + + # Sometimes the decorator is just desugared, making it impossible to detect. + # This attribute makes detection easier. + setattr(wrapper, '__pyct_is_compile_decorator', True) + return wrapper + + return decorator + + +def to_graph(e, + recursive=True, + verbose=False, + arg_values=None, + arg_types=None, + partial_types=None): + """Compile a Python entity into equivalent TensorFlow code. + + Currently supported entities: + * functions + * classes + + Classes are handled by converting all their methods into a new class. + + Args: + e: A Python entity. + recursive: Whether to recusrively convert any functions that the decorator + function may call. + verbose: Whether to output the compiled code in the logs. + arg_values: A dict containing value hints for symbols like function + parameters. + arg_types: A dict containing type hints for symbols like function + parameters. + partial_types: A set of types (e.g. classes) that will not be converted + entirely. Calls to member functions for these types will be renamed + independently. + + Returns: + A function with a signature identical to `o`, but which when executed it + creates TF a graph that has the same functionality as the original entity. + """ + conversion_map = conversion.ConversionMap( + recursive=recursive, + nocompile_decorators=(convert, graph_ready, convert_inline), + partial_types=partial_types) + _, name = conversion.entity_to_graph(e, conversion_map, arg_values, arg_types) + + module = gast.Module([]) + for import_line in config.COMPILED_IMPORT_STATEMENTS: + module.body.append(parser.parse_str(import_line)) + for dep in conversion_map.dependency_cache.values(): + module.body.append(dep) + compiled_node, compiled_src = compiler.ast_to_object(module) + + # The compiled code should see everything the entry function saw. + # TODO(mdan): This might not work well if the call tree spans modules? + if tf_inspect.isfunction(e): + compiled_node.__dict__.update(six.get_function_globals(e)) + compiled_fn = getattr(compiled_node, name) + + if verbose: + logging.info('Compiled output of %s:\n\n%s\n', e, compiled_src) + + return compiled_fn + + +def to_code(e, + recursive=True, + arg_values=None, + arg_types=None, + partial_types=None, + indentation=' '): + """Return the equivalent of an entity in TensorFlow code. + + See `to_graph` for more details. + + Args: + e: A Python entity. + recursive: See to_graph. + arg_values: See to_graph. + arg_types: See to_graph. + partial_types: See to_graph. + indentation: String, when to use for each level of indentation. + + Returns: + String. + """ + conversion_map = conversion.ConversionMap( + recursive=recursive, + nocompile_decorators=(convert, graph_ready, convert_inline), + partial_types=partial_types) + conversion.entity_to_graph(e, conversion_map, arg_values, arg_types) + + imports = '\n'.join(config.COMPILED_IMPORT_STATEMENTS) + code = '\n'.join( + compiler.ast_to_source(dep, indentation) + for dep in reversed(tuple( + six.itervalues(conversion_map.dependency_cache)))) + + return imports + '\n\n' + code diff --git a/tensorflow/contrib/py2tf/impl/api_test.py b/tensorflow/contrib/py2tf/impl/api_test.py new file mode 100644 index 0000000000000000000000000000000000000000..02cd8ed2d0ffee8ef2d31ea65902d2b493df9d64 --- /dev/null +++ b/tensorflow/contrib/py2tf/impl/api_test.py @@ -0,0 +1,193 @@ +# Copyright 2017 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 api module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.impl import api +from tensorflow.contrib.py2tf.impl import config +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test + + +class ApiTest(test.TestCase): + + def setUp(self): + config.DEFAULT_UNCOMPILED_MODULES.add((math_ops.__name__,)) + config.COMPILED_IMPORT_STATEMENTS = ( + 'from tensorflow.python.ops ' + 'import control_flow_ops as tf', + 'from tensorflow.contrib.py2tf import utils as ' + 'py2tf_utils') + + def test_decorator_recurses(self): + + class TestClass(object): + + def called_member(self, a): + if a < 0: + a = -a + return a + + @api.convert(recursive=True) + def test_method(self, x, s, a): + while math_ops.reduce_sum(x) > s: + x //= self.called_member(a) + return x + + tc = TestClass() + with self.test_session() as sess: + x = tc.test_method( + constant_op.constant([2, 4]), constant_op.constant(1), + constant_op.constant(-2)) + self.assertListEqual([0, 1], sess.run(x).tolist()) + + def test_decorator_does_not_recurse(self): + + class TestClass(object): + + def called_member(self, a): + return math_ops.negative(a) + + @api.convert(recursive=False) + def test_method(self, x, s, a): + while math_ops.reduce_sum(x) > s: + x //= self.called_member(a) + return x + + tc = TestClass() + with self.test_session() as sess: + x = tc.test_method( + constant_op.constant([2, 4]), constant_op.constant(1), + constant_op.constant(-2)) + self.assertListEqual([0, 1], sess.run(x).tolist()) + + def test_decorator_calls_converted(self): + + class TestClass(object): + + @api.graph_ready + def called_member(self, a): + return math_ops.negative(a) + + @api.convert(recursive=True) + def test_method(self, x, s, a): + while math_ops.reduce_sum(x) > s: + x //= self.called_member(a) + return x + + tc = TestClass() + with self.test_session() as sess: + x = tc.test_method( + constant_op.constant([2, 4]), constant_op.constant(1), + constant_op.constant(-2)) + self.assertListEqual([0, 1], sess.run(x).tolist()) + + def test_decorator_calls_decorated(self): + + class TestClass(object): + + @api.convert() + def called_member(self, a): + if a < 0: + a = -a + return a + + @api.convert(recursive=True) + def test_method(self, x, s, a): + while math_ops.reduce_sum(x) > s: + x //= self.called_member(a) + return x + + tc = TestClass() + with self.test_session() as sess: + x = tc.test_method( + constant_op.constant([2, 4]), constant_op.constant(1), + constant_op.constant(-2)) + self.assertListEqual([0, 1], sess.run(x).tolist()) + + def test_convert_call_site_decorator(self): + + class TestClass(object): + + def called_member(self, a): + if a < 0: + a = -a + return a + + @api.convert(recursive=True) + def test_method(self, x, s, a): + while math_ops.reduce_sum(x) > s: + x //= api.convert_inline(self.called_member, a) + return x + + tc = TestClass() + with self.test_session() as sess: + x = tc.test_method( + constant_op.constant([2, 4]), constant_op.constant(1), + constant_op.constant(-2)) + self.assertListEqual([0, 1], sess.run(x).tolist()) + + def test_graph_ready_call_site_decorator(self): + + class TestClass(object): + + def called_member(self, a): + return math_ops.negative(a) + + @api.convert(recursive=True) + def test_method(self, x, s, a): + while math_ops.reduce_sum(x) > s: + x //= api.graph_ready(self.called_member(a)) + return x + + tc = TestClass() + with self.test_session() as sess: + x = tc.test_method( + constant_op.constant([2, 4]), constant_op.constant(1), + constant_op.constant(-2)) + self.assertListEqual([0, 1], sess.run(x).tolist()) + + def test_to_graph_basic(self): + def test_fn(x, s): + while math_ops.reduce_sum(x) > s: + x //= 2 + return x + + compiled_fn = api.to_graph(test_fn) + + with self.test_session() as sess: + x = compiled_fn(constant_op.constant([4, 8]), 4) + self.assertListEqual([1, 2], sess.run(x).tolist()) + + def test_to_code_basic(self): + def test_fn(x, s): + while math_ops.reduce_sum(x) > s: + x /= 2 + return x + + compiled_code = api.to_code(test_fn) + + # Just check for some key words and that it is parseable Python code. + self.assertRegexpMatches(compiled_code, 'py2tf_utils\\.run_while') + self.assertIsNotNone(parser.parse_str(compiled_code)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/impl/config.py b/tensorflow/contrib/py2tf/impl/config.py new file mode 100644 index 0000000000000000000000000000000000000000..7c3ecefff0f8858d5505ff30e1270b2fd42c9ad8 --- /dev/null +++ b/tensorflow/contrib/py2tf/impl/config.py @@ -0,0 +1,45 @@ +# Copyright 2016 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. +# ============================================================================== +"""Global configuration.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf import utils + + +PYTHON_LITERALS = { + 'None': None, + 'False': False, + 'True': True, + 'float': float, +} + +DEFAULT_UNCOMPILED_MODULES = set(( + ('tensorflow',), + (utils.__name__,), +)) + +NO_SIDE_EFFECT_CONSTRUCTORS = set(('tensorflow',)) + +# TODO(mdan): Also allow controlling the generated names (for testability). +# TODO(mdan): Verify that these names are not hidden by generated code. +# TODO(mdan): Make sure copybara renames the reference below. +COMPILED_IMPORT_STATEMENTS = ( + 'from __future__ import print_function', + 'import tensorflow as tf', + 'from tensorflow.contrib.py2tf import utils as ' + 'py2tf_utils') diff --git a/tensorflow/contrib/py2tf/impl/conversion.py b/tensorflow/contrib/py2tf/impl/conversion.py new file mode 100644 index 0000000000000000000000000000000000000000..3d5624b187ed47e9eed8afbb2e101e1098f81c15 --- /dev/null +++ b/tensorflow/contrib/py2tf/impl/conversion.py @@ -0,0 +1,282 @@ +# Copyright 2016 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. +# ============================================================================== +"""High level conversion support.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast +import six + +from tensorflow.contrib.py2tf import utils +from tensorflow.contrib.py2tf.converters import asserts +from tensorflow.contrib.py2tf.converters import break_statements +from tensorflow.contrib.py2tf.converters import builtin_functions +from tensorflow.contrib.py2tf.converters import call_trees +from tensorflow.contrib.py2tf.converters import continue_statements +from tensorflow.contrib.py2tf.converters import control_flow +from tensorflow.contrib.py2tf.converters import decorators +from tensorflow.contrib.py2tf.converters import for_loops +from tensorflow.contrib.py2tf.converters import logical_expressions +from tensorflow.contrib.py2tf.converters import side_effect_guards +from tensorflow.contrib.py2tf.impl import config +from tensorflow.contrib.py2tf.impl import naming +from tensorflow.contrib.py2tf.pyct import context +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import qual_names +from tensorflow.contrib.py2tf.pyct.static_analysis import activity +from tensorflow.contrib.py2tf.pyct.static_analysis import live_values +from tensorflow.contrib.py2tf.pyct.static_analysis import type_info +from tensorflow.python.util import tf_inspect + + +# TODO(mdan): Might we not need any renaming at all? + + +class ConversionMap(object): + """ConversionMaps keep track of converting function hierarchies. + + Attributes: + recursive: Whether to recusrively convert any functions that the decorator + function may call. + nocompile_decorators: tuple of decorator functions that toggle compilation + off. + dependency_cache: dict[object]: ast; maps original entities to their + converted AST + name_map: dict[string]: string; maps original entities to the name of + their converted counterparts + """ + + # TODO(mdan): Rename to ConversionContext, and pull in additional flags. + + def __init__(self, recursive, nocompile_decorators, partial_types): + self.recursive = recursive + self.nocompile_decorators = nocompile_decorators + self.partial_types = partial_types if partial_types else () + self.dependency_cache = {} + self.name_map = {} + + def new_namer(self, namespace): + return naming.Namer(namespace, self.recursive, self.name_map, + self.partial_types) + + def update_name_map(self, namer): + for o, name in namer.renamed_calls.items(): + if o in self.name_map: + if self.name_map[o] != name: + raise ValueError( + 'Calls to %s were converted using multiple names (%s). This is ' + 'possible when an entity with one of these names already ' + 'existed. To fix, avoid using any of these names.') + else: + self.name_map[o] = name + + def add_to_cache(self, original_entity, converted_ast): + self.dependency_cache[original_entity] = converted_ast + + +def entity_to_graph(o, conversion_map, arg_values, arg_types): + """Compile a Python entity into equivalent TensorFlow. + + The function will also recursively compile all the entities that `o` + references, updating `dependency_cache`. + + This function is reentrant, and relies on dependency_cache to avoid + generating duplicate code. + + Args: + o: A Python entity. + conversion_map: A ConversionMap object. + arg_values: A dict containing value hints for symbols like function + parameters. + arg_types: A dict containing type hints for symbols like function + parameters. + + Returns: + A tuple (ast, new_name): + * ast: An AST representing an entity with interface equivalent to `o`, + but which when executed it creates TF a graph. + * new_name: The symbol name under which the new entity can be found. + + Raises: + ValueError: if the entity type is not supported. + """ + if tf_inspect.isclass(o): + node, new_name = class_to_graph(o, conversion_map) + elif tf_inspect.isfunction(o): + node, new_name = function_to_graph(o, conversion_map, arg_values, arg_types) + elif tf_inspect.ismethod(o): + node, new_name = function_to_graph(o, conversion_map, arg_values, arg_types) + else: + raise ValueError( + 'Entity "%s" has unsupported type "%s". Only functions and classes are ' + 'supported for now.' % (o, type(o))) + + conversion_map.add_to_cache(o, node) + if conversion_map.recursive: + for obj in conversion_map.name_map.keys(): + if obj not in conversion_map.dependency_cache: + if (hasattr(obj, 'im_class') and + getattr(obj, 'im_class') not in conversion_map.partial_types): + # Class members are converted with their objects, unless they're + # only converted partially. + continue + entity_to_graph(obj, conversion_map, {}, {}) + + return node, new_name + + +def class_to_graph(c, conversion_map): + """Specialization of `entity_to_graph` for classes.""" + converted_members = {} + members = tf_inspect.getmembers(c, predicate=tf_inspect.ismethod) + if not members: + raise ValueError('Cannot convert %s: it has no member methods.') + + class_globals = None + for _, m in members: + node, _ = function_to_graph( + m, + conversion_map=conversion_map, + arg_values={}, + arg_types={'self': (c.__name__, c)}, + owner_type=c) + # TODO(mdan): Do not assume all members have the same view of globals. + if class_globals is None: + class_globals = six.get_function_globals(m) + converted_members[m] = node + namer = conversion_map.new_namer(class_globals) + class_name = namer.compiled_class_name(c.__name__, c) + node = gast.ClassDef( + class_name, + bases=[], + keywords=[], + body=converted_members.values(), + decorator_list=[]) + + return node, class_name + + +def function_to_graph(f, conversion_map, arg_values, arg_types, + owner_type=None): + """Specialization of `entity_to_graph` for callable functions.""" + node, source = parser.parse_entity(f) + node = node.body[0] + namespace = six.get_function_globals(f) + + # This is needed for non-global functions. + closure = six.get_function_closure(f) + if closure: + for e in closure: + if callable(e.cell_contents): + fn = e.cell_contents + namespace[fn.__name__] = fn + + # Manually add the utils namespace which may be used from generated code. + if 'py2tf_util' not in namespace: + namespace['py2tf_utils'] = utils + elif namespace['py2tf_utils'] != utils: + raise ValueError( + 'The module name py2tf_utils is reserved and may not be used.') + + namer = conversion_map.new_namer(namespace) + ctx = context.EntityContext( + namer=namer, + source_code=source, + source_file='', + namespace=namespace, + arg_values=arg_values, + arg_types=arg_types, + recursive=conversion_map.recursive) + node = node_to_graph(node, ctx, conversion_map.nocompile_decorators) + + # TODO(mdan): This somewhat duplicates the call rename logic in call_treest.py + new_name, did_rename = namer.compiled_function_name(f.__name__, f, owner_type) + if not did_rename: + new_name = f.__name__ + if node.name != f.__name__: + raise NotImplementedError('Strange corner case. Send us offending code!') + + node.name = new_name + conversion_map.update_name_map(namer) + return node, new_name + + +def _static_analysis_pass(node, ctx): + node = qual_names.resolve(node) + node = activity.resolve(node, ctx, None) + node = live_values.resolve(node, ctx, config.PYTHON_LITERALS) + node = type_info.resolve(node, ctx) + return node + + +def node_to_graph(node, ctx, nocompile_decorators): + """Convert Python code to equivalent TF graph mode code. + + Args: + node: A Python AST node representing the code to convert. + ctx: An EntityContext object. + nocompile_decorators: A tuple containing decorators to be stripped from + functions during conversion. + + Returns: + A tuple (node, deps): + * node: A Python ast node, representing the converted code. + * deps: A set of strings, the fully qualified names of entity + dependencies that this node has. + """ + # TODO(mdan): Verify arguments for correctness. + + # TODO(mdan): Factor out common elements. + # These include: + # * code move between blocks + # * visiting blocks in transformers + + # Certain steps, especially canonicalization, insert new symbols into the + # tree, which must be accounted. Although less efficient, it is most robust + # to re-run the analysis. + + node = _static_analysis_pass(node, ctx) + # Past this point, line numbers are no longer accurate so we ignore the + # source. + # TODO(mdan): Is it feasible to reconstruct intermediate source code? + ctx.source_code = None + node = decorators.transform(node, nocompile_decorators) + node = break_statements.transform(node, ctx) + node = asserts.transform(node, ctx) + + # Note: sequencing continue canonicalization before for loop one avoids + # dealing with the extra loop increment operation that the for + # canonicalization creates. + node = continue_statements.transform(node, ctx) + ctx.namespace['len'] = len + + node = _static_analysis_pass(node, ctx) + node = for_loops.transform(node, ctx) + # for_loops may insert new global references. + node = builtin_functions.transform(node, ctx) + + node = _static_analysis_pass(node, ctx) + node = call_trees.transform(node, ctx, config.DEFAULT_UNCOMPILED_MODULES, + nocompile_decorators) + node = control_flow.transform(node, ctx) + + # control_flow may create new symbols and change scopes. + node = _static_analysis_pass(node, ctx) + node = logical_expressions.transform(node) + node = side_effect_guards.transform(node, ctx) + + return node diff --git a/tensorflow/contrib/py2tf/impl/conversion_test.py b/tensorflow/contrib/py2tf/impl/conversion_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3888958f19b9fa13b759924c5188722e500e30a1 --- /dev/null +++ b/tensorflow/contrib/py2tf/impl/conversion_test.py @@ -0,0 +1,64 @@ +# Copyright 2017 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 conversion module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.impl import conversion +from tensorflow.python.platform import test + + +class ConversionTest(test.TestCase): + + def test_entity_to_graph_unsupported_types(self): + with self.assertRaises(ValueError): + conversion_map = conversion.ConversionMap(True, (), ()) + conversion.entity_to_graph('dummy', conversion_map, None, None) + + def test_entity_to_graph_callable(self): + + def f(a): + return a + + conversion_map = conversion.ConversionMap(True, (), ()) + ast, new_name = conversion.entity_to_graph(f, conversion_map, None, None) + self.assertTrue(isinstance(ast, gast.FunctionDef), ast) + self.assertEqual('tf__f', new_name) + + def test_entity_to_graph_call_tree(self): + + def g(a): + return a + + def f(a): + return g(a) + + conversion_map = conversion.ConversionMap(True, (), ()) + conversion.entity_to_graph(f, conversion_map, None, None) + + self.assertTrue(f in conversion_map.dependency_cache) + self.assertTrue(g in conversion_map.dependency_cache) + self.assertEqual('tf__f', conversion_map.dependency_cache[f].name) + self.assertEqual( + 'tf__g', conversion_map.dependency_cache[f].body[0].value.func.id) + self.assertEqual('tf__g', conversion_map.dependency_cache[g].name) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/impl/naming.py b/tensorflow/contrib/py2tf/impl/naming.py new file mode 100644 index 0000000000000000000000000000000000000000..51326091de13715c32d0a79279f1d3274e48ad10 --- /dev/null +++ b/tensorflow/contrib/py2tf/impl/naming.py @@ -0,0 +1,132 @@ +# Copyright 2017 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. +# ============================================================================== +"""Symbol naming utilities.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.pyct import qual_names + + +class Namer(object): + """Implementation of the namer interfaces required by various converters. + + This implementation performs additional tasks like keeping track of the + function calls that have been encountered and replaced with calls to their + corresponding compiled counterparts. + + Interfaces currently implemented: + * call_trees.FunctionNamer + * control_flow.SymbolNamer + * side_effect_guards.SymbolNamer + """ + + def __init__(self, global_namespace, recursive, name_map, partial_types): + self.global_namespace = global_namespace + self.recursive = recursive + self.partial_types = partial_types + + self.renamed_calls = {} + if name_map is not None: + self.renamed_calls.update(name_map) + + self.generated_names = set() + + def compiled_class_name(self, original_fqn, live_entity=None): + """See call_trees.FunctionNamer.compiled_class_name.""" + if live_entity is not None and live_entity in self.renamed_calls: + return self.renamed_calls[live_entity] + + if isinstance(original_fqn, tuple): + original_name = '__'.join(original_fqn) + else: + original_name = original_fqn + + new_name_root = 'Tf%s' % original_name + new_name = new_name_root + n = 0 + while new_name in self.global_namespace: + n += 1 + new_name = '%s_%d' % (new_name_root, n) + + if live_entity is not None: + self.renamed_calls[live_entity] = new_name + self.generated_names.add(new_name) + if live_entity is not None: + self.renamed_calls[live_entity] = new_name + return new_name + + def compiled_function_name(self, + original_fqn, + live_entity=None, + owner_type=None): + """See call_trees.FunctionNamer.compiled_function_name.""" + + if not self.recursive: + return None, False + + if owner_type is not None and owner_type not in self.partial_types: + # Members are not renamed when part of an entire converted class. + return None, False + + if isinstance(original_fqn, tuple): + original_name = '__'.join(original_fqn) + else: + original_name = original_fqn + + if live_entity is not None and live_entity in self.renamed_calls: + return self.renamed_calls[live_entity], True + + new_name_root = 'tf__%s' % original_name + new_name = new_name_root + n = 0 + while new_name in self.global_namespace: + n += 1 + new_name = '%s_%d' % (new_name_root, n) + + if live_entity is not None: + self.renamed_calls[live_entity] = new_name + self.generated_names.add(new_name) + + return new_name, True + + def new_symbol(self, name_root, reserved_locals): + """See control_flow.SymbolNamer.new_symbol.""" + # reserved_locals may contain QNs. + all_reserved_locals = set() + for s in reserved_locals: + if isinstance(s, qual_names.QN): + all_reserved_locals.update(s.qn) + elif isinstance(s, str): + all_reserved_locals.add(s) + else: + raise ValueError('Unexpected symbol type "%s"' % type(s)) + + pieces = name_root.split('_') + if pieces[-1].isdigit(): + name_root = '_'.join(pieces[:-1]) + n = int(pieces[-1]) + else: + n = 0 + new_name = name_root + + while (new_name in self.global_namespace or + new_name in all_reserved_locals or new_name in self.generated_names): + n += 1 + new_name = '%s_%d' % (name_root, n) + + self.generated_names.add(new_name) + return new_name diff --git a/tensorflow/contrib/py2tf/impl/naming_test.py b/tensorflow/contrib/py2tf/impl/naming_test.py new file mode 100644 index 0000000000000000000000000000000000000000..beb4e54937bbb91b19157c9b9e3c528353206c62 --- /dev/null +++ b/tensorflow/contrib/py2tf/impl/naming_test.py @@ -0,0 +1,77 @@ +# Copyright 2017 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 naming module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.impl import naming +from tensorflow.python.platform import test + + +class NamerTest(test.TestCase): + + def test_compiled_function_name_tracks_names(self): + def bar(): + pass + + namer = naming.Namer({}, True, None, ()) + self.assertEqual(('tf__foo', True), namer.compiled_function_name('foo')) + self.assertEqual(('tf__bar', True), namer.compiled_function_name( + 'bar', bar)) + self.assertEqual({bar: 'tf__bar'}, namer.renamed_calls) + self.assertItemsEqual(('tf__bar', 'tf__foo'), namer.generated_names) + + def test_compiled_function_name_consistent(self): + def foo(): + pass + + namer = naming.Namer({}, True, None, ()) + self.assertEqual(('tf__foo', True), namer.compiled_function_name( + 'foo', foo)) + self.assertEqual(('tf__foo', True), namer.compiled_function_name( + 'foo', foo)) + + def test_compiled_function_name_avoids_global_conflicts(self): + def foo(): + pass + + namer = naming.Namer({'tf__foo': 1}, True, None, ()) + self.assertEqual(('tf__foo_1', True), + namer.compiled_function_name('foo', foo)) + + def test_new_symbol_tracks_names(self): + namer = naming.Namer({}, True, None, ()) + self.assertEqual('temp', namer.new_symbol('temp', set())) + self.assertItemsEqual(('temp',), namer.generated_names) + + def test_new_symbol_avoids_duplicates(self): + namer = naming.Namer({}, True, None, ()) + self.assertEqual('temp', namer.new_symbol('temp', set())) + self.assertEqual('temp_1', namer.new_symbol('temp', set())) + self.assertItemsEqual(('temp', 'temp_1'), namer.generated_names) + + def test_new_symbol_avoids_conflicts(self): + namer = naming.Namer({'temp': 1}, True, None, ()) + # temp is reserved in the global namespace + self.assertEqual('temp_1', namer.new_symbol('temp', set())) + # temp_2 is reserved in the local namespace + self.assertEqual('temp_3', namer.new_symbol('temp', set(('temp_2',)))) + self.assertItemsEqual(('temp_1', 'temp_3'), namer.generated_names) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/BUILD b/tensorflow/contrib/py2tf/pyct/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..e3c0da4b10f9ffbee1b2a906b64d4762f41d97b4 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/BUILD @@ -0,0 +1,114 @@ +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + +load("//tensorflow:tensorflow.bzl", "py_test") + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +py_library( + name = "pyct", + srcs = [ + "__init__.py", + "anno.py", + "ast_util.py", + "compiler.py", + "context.py", + "parser.py", + "pretty_printer.py", + "qual_names.py", + "templates.py", + "transformer.py", + ], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + "@astor_archive//:astor", + "@gast_archive//:gast", + "@six_archive//:six", + "@termcolor_archive//:termcolor", + ], +) + +py_test( + name = "anno_test", + srcs = ["anno_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "ast_util_test", + srcs = ["ast_util_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + "@gast_archive//:gast", + ], +) + +py_test( + name = "compiler_test", + srcs = ["compiler_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + "@gast_archive//:gast", + ], +) + +py_test( + name = "parser_test", + srcs = ["parser_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "pretty_printer_test", + srcs = ["pretty_printer_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "qual_names_test", + srcs = ["qual_names_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "templates_test", + srcs = ["templates_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + "@gast_archive//:gast", + ], +) diff --git a/tensorflow/contrib/ndlstm/__init__.py b/tensorflow/contrib/py2tf/pyct/__init__.py similarity index 94% rename from tensorflow/contrib/ndlstm/__init__.py rename to tensorflow/contrib/py2tf/pyct/__init__.py index 52e83069cb0c68b510da46149248369dce376647..d787e56bbecbd1d891fdf41207256c4c5096224f 100644 --- a/tensorflow/contrib/ndlstm/__init__.py +++ b/tensorflow/contrib/py2tf/pyct/__init__.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +"""Python source code transformation library.""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/py2tf/pyct/anno.py b/tensorflow/contrib/py2tf/pyct/anno.py new file mode 100644 index 0000000000000000000000000000000000000000..7a0528b6d0b65b6604930b7a13d8493af9d61f02 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/anno.py @@ -0,0 +1,72 @@ +# Copyright 2017 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. +# ============================================================================== +"""Handling annotations on AST nodes. + +Adapted from Tangent. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from enum import Enum + + +class NoValue(Enum): + + def __repr__(self): + return self.name + + +class Basic(NoValue): + """Container for annotation keys. + + The enum values are used strictly for documentation purposes. + """ + + QN = 'Qualified name, as it appeared in the code.' + SKIP_PROCESSING = ( + 'This node should be preserved as is and not processed any further.') + INDENT_BLOCK_REMAINDER = ( + 'When a node is annotated with this, the remainder of the block should ' + 'be indented below it. The annotation contains a tuple ' + '(new_body, name_map), where `new_body` is the new indented block and ' + '`name_map` allows renaming symbols.') + + +def getanno(node, key, field_name='___pyct_anno'): + return getattr(node, field_name)[key] + + +def hasanno(node, key, field_name='___pyct_anno'): + return hasattr(node, field_name) and key in getattr(node, field_name) + + +def setanno(node, key, value, field_name='___pyct_anno'): + annotations = getattr(node, field_name, {}) + setattr(node, field_name, annotations) + annotations[key] = value + + # So that the annotations survive gast_to_ast() and ast_to_gast() + if field_name not in node._fields: + node._fields += (field_name,) + + +def delanno(node, key, field_name='___pyct_anno'): + annotations = getattr(node, field_name) + del annotations[key] + if not annotations: + delattr(node, field_name) + node._fields = tuple(f for f in node._fields if f != field_name) diff --git a/tensorflow/contrib/py2tf/pyct/anno_test.py b/tensorflow/contrib/py2tf/pyct/anno_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ff40bfe1f50ae731648afdf509c26c3a70d3f6cb --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/anno_test.py @@ -0,0 +1,47 @@ +# Copyright 2017 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 anno module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.python.platform import test + + +class AnnoTest(test.TestCase): + + def test_basic(self): + node = ast.Name() + + self.assertFalse(anno.hasanno(node, 'foo')) + with self.assertRaises(AttributeError): + anno.getanno(node, 'foo') + + anno.setanno(node, 'foo', 3) + self.assertTrue(anno.hasanno(node, 'foo')) + self.assertEqual(3, anno.getanno(node, 'foo')) + + anno.delanno(node, 'foo') + self.assertFalse(anno.hasanno(node, 'foo')) + with self.assertRaises(AttributeError): + anno.getanno(node, 'foo') + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/ast_util.py b/tensorflow/contrib/py2tf/pyct/ast_util.py new file mode 100644 index 0000000000000000000000000000000000000000..f916775b9cf3cec960ec2896c334f1d737862205 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/ast_util.py @@ -0,0 +1,96 @@ +# Copyright 2017 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. +# ============================================================================== +"""Copy an AST tree, discarding annotations.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ast + +import gast + +from tensorflow.contrib.py2tf.pyct import anno + + +class CleanCopier(gast.NodeVisitor): + """Copy AST nodes. + + The copied nodes will ignore almost all fields that prefixed by '__'. + Exceptions make some annotations. + """ + + # TODO(mdan): Parametrize which annotations get carried over. + + def generic_visit(self, node): + new_fields = {} + for f in node._fields: + if f.startswith('__'): + continue + if not hasattr(node, f): + continue + v = getattr(node, f) + if isinstance(v, list): + v = [self.generic_visit(n) for n in v] + elif isinstance(v, tuple): + v = tuple(self.generic_visit(n) for n in v) + elif isinstance(v, (gast.AST, ast.AST)): + v = self.generic_visit(v) + else: + # Assume everything else is a value type. + pass + new_fields[f] = v + new_node = type(node)(**new_fields) + if anno.hasanno(node, anno.Basic.SKIP_PROCESSING): + anno.setanno(new_node, anno.Basic.SKIP_PROCESSING, True) + return new_node + + +def copy_clean(node): + copier = CleanCopier() + if isinstance(node, list): + return [copier.visit(n) for n in node] + elif isinstance(node, tuple): + return tuple(copier.visit(n) for n in node) + else: + return copier.visit(node) + + +class SymbolRenamer(gast.NodeTransformer): + """Transformer that can rename symbols to a simple names.""" + + def __init__(self, name_map): + self.name_map = name_map + + def _process(self, node): + qn = anno.getanno(node, anno.Basic.QN) + if qn in self.name_map: + return gast.Name(str(self.name_map[qn]), node.ctx, None) + return self.generic_visit(node) + + def visit_Name(self, node): + return self._process(node) + + def visit_Attribute(self, node): + return self._process(node) + + +def rename_symbols(node, name_map): + renamer = SymbolRenamer(name_map) + if isinstance(node, list): + return [renamer.visit(n) for n in node] + elif isinstance(node, tuple): + return tuple(renamer.visit(n) for n in node) + return renamer.visit(node) diff --git a/tensorflow/contrib/py2tf/pyct/ast_util_test.py b/tensorflow/contrib/py2tf/pyct/ast_util_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e0b00c178168f96e656c57cc75a76e6da8af1d8a --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/ast_util_test.py @@ -0,0 +1,79 @@ +# Copyright 2017 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 ast_util module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ast + +from tensorflow.contrib.py2tf.pyct import ast_util +from tensorflow.contrib.py2tf.pyct import qual_names +from tensorflow.python.platform import test + + +class AstUtilTest(test.TestCase): + + def test_rename_symbols(self): + node = ast.Tuple([ + ast.Name('a', ast.Load()), + ast.Name('b', ast.Load()), + ast.Attribute(ast.Name('b', None), 'c', ast.Store()), + ast.Attribute( + ast.Attribute(ast.Name('b', None), 'c', ast.Load()), 'd', + None) + ], None) + node = qual_names.resolve(node) + node = ast_util.rename_symbols( + node, + { + qual_names.QN('a'): qual_names.QN('renamed_a'), + qual_names.QN('b.c'): qual_names.QN('renamed_b_c'), + }) + + self.assertEqual(node.elts[0].id, 'renamed_a') + self.assertTrue(isinstance(node.elts[0].ctx, ast.Load)) + self.assertEqual(node.elts[1].id, 'b') + self.assertEqual(node.elts[2].id, 'renamed_b_c') + self.assertTrue(isinstance(node.elts[2].ctx, ast.Store)) + self.assertEqual(node.elts[3].value.id, 'renamed_b_c') + self.assertTrue(isinstance(node.elts[3].value.ctx, ast.Load)) + + def test_copy_clean(self): + ret = ast.Return( + ast.BinOp( + op=ast.Add(), + left=ast.Name(id='a', ctx=ast.Load()), + right=ast.Num(1))) + setattr(ret, '__foo', 'bar') + node = ast.FunctionDef( + name='f', + args=ast.arguments( + args=[ast.Name(id='a', ctx=ast.Param())], + vararg=None, + kwarg=None, + defaults=[]), + body=[ret], + decorator_list=[], + returns=None) + new_node = ast_util.copy_clean(node) + self.assertFalse(node is new_node) + self.assertFalse(ret is new_node.body[0]) + self.assertFalse(hasattr(new_node.body[0], '__foo')) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/compiler.py b/tensorflow/contrib/py2tf/pyct/compiler.py new file mode 100644 index 0000000000000000000000000000000000000000..51cf6930e8bcb3728ee55bf5d4781f01a5ef73bd --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/compiler.py @@ -0,0 +1,72 @@ +# Copyright 2017 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. +# ============================================================================== +"""Converting AST to code. + +Adapted from Tangent. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# TODO(mdan): Use six for compatibility here. +import atexit +import imp +import os +import tempfile + +import astor +import gast + + +def ast_to_source(node, indentation): + """Return the source code of given AST.""" + if isinstance(node, gast.AST): + node = gast.gast_to_ast(node) + generator = astor.codegen.SourceGenerator(indentation, False, + astor.string_repr.pretty_string) + generator.visit(node) + generator.result.append('\n') + return astor.source_repr.pretty_source(generator.result).lstrip() + + +def ast_to_object( + node, indentation=' ', source_prefix=None, delete_on_exit=True): + """Return the Python objects represented by given AST. + + Compiling the AST code this way ensures that the source code is readable by + e.g. `pdb` or `inspect`. + + Args: + node: The code to compile, as an AST object. + indentation: The string to use for indentation. + source_prefix: Optional string to print as-is into the source file. + delete_on_exit: Whether to delete the temporary file used for compilation + on exit. + + Returns: + A module object containing the compiled source code. + """ + source = ast_to_source(node, indentation) + + with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f: + module_name = os.path.basename(f.name[:-3]) + if source_prefix: + f.write(source_prefix) + f.write('\n') + f.write(source) + if delete_on_exit: + atexit.register(lambda: os.remove(f.name)) + return imp.load_source(module_name, f.name), source diff --git a/tensorflow/contrib/py2tf/pyct/compiler_test.py b/tensorflow/contrib/py2tf/pyct/compiler_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c1f84238efa7dd6fc0748748a2cb4f074572b4c6 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/compiler_test.py @@ -0,0 +1,91 @@ +# Copyright 2017 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 compiler module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import textwrap + +import gast + +from tensorflow.contrib.py2tf.pyct import compiler +from tensorflow.python.platform import test + + +class CompilerTest(test.TestCase): + + def test_ast_to_source(self): + node = gast.If( + test=gast.Num(1), + body=[ + gast.Assign( + targets=[gast.Name('a', gast.Store(), None)], + value=gast.Name('b', gast.Load(), None)) + ], + orelse=[ + gast.Assign( + targets=[gast.Name('a', gast.Store(), None)], + value=gast.Str('c')) + ]) + + self.assertEqual( + textwrap.dedent(""" + if 1: + a = b + else: + a = 'c' + """).strip(), + compiler.ast_to_source(node, indentation=' ').strip()) + + def test_ast_to_object(self): + node = gast.FunctionDef( + name='f', + args=gast.arguments( + args=[gast.Name('a', gast.Param(), None)], + vararg=None, + kwonlyargs=[], + kwarg=None, + defaults=[], + kw_defaults=[]), + body=[ + gast.Return( + gast.BinOp( + op=gast.Add(), + left=gast.Name('a', gast.Load(), None), + right=gast.Num(1))) + ], + decorator_list=[], + returns=None) + + module, source = compiler.ast_to_object(node) + + expected_source = """ + def f(a): + return a + 1 + """ + self.assertEqual( + textwrap.dedent(expected_source).strip(), + source.strip()) + self.assertEqual(2, module.f(1)) + with open(module.__file__, 'r') as temp_output: + self.assertEqual( + textwrap.dedent(expected_source).strip(), + temp_output.read().strip()) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/context.py b/tensorflow/contrib/py2tf/pyct/context.py new file mode 100644 index 0000000000000000000000000000000000000000..fef74ebefa290369c7310af6d7e4faeef44d9aee --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/context.py @@ -0,0 +1,43 @@ +# Copyright 2016 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. +# ============================================================================== +"""Conversion context containers.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + + +class EntityContext(object): + """Contains information about an entity, like source code. + + Attributes: + namer: Namer that matches the contract of all converters. + source_code: The entity's source code. + source_file: The entity's source file. + namespace: Dict[str->*], containing symbols visible to the entity + (excluding parameters). + arg_values: Dict[str->*], containing parameter values, if known. + arg_types: Dict[str->*], containing parameter types, if known. + """ + + def __init__(self, namer, source_code, source_file, namespace, arg_values, + arg_types, recursive): + self.namer = namer + self.source_code = source_code + self.source_file = source_file + self.namespace = namespace + self.arg_values = {} if arg_values is None else arg_values + self.arg_types = {} if arg_types is None else arg_types + self.recursive = recursive diff --git a/tensorflow/contrib/quantize/python/copy_graph.py b/tensorflow/contrib/py2tf/pyct/parser.py similarity index 65% rename from tensorflow/contrib/quantize/python/copy_graph.py rename to tensorflow/contrib/py2tf/pyct/parser.py index 0376fcba82b99feabdba3b683f9db9a32db51efb..dc7df883b349becd860bb0dbceab22cb39c750b5 100644 --- a/tensorflow/contrib/quantize/python/copy_graph.py +++ b/tensorflow/contrib/py2tf/pyct/parser.py @@ -12,21 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Utility to copy a tf.Graph.""" +"""Converting code to AST. + +Adapted from Tangent. +""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.framework import ops -from tensorflow.python.training import saver as saver_lib +import textwrap + +import gast + +from tensorflow.python.util import tf_inspect + + +def parse_entity(entity): + """Return the AST of given entity.""" + source = tf_inspect.getsource(entity) + source = textwrap.dedent(source) + return parse_str(source), source -def CopyGraph(graph): - """Return a copy of graph.""" - meta_graph = saver_lib.export_meta_graph( - graph=graph, collection_list=graph.get_all_collection_keys()) - graph_copy = ops.Graph() - with graph_copy.as_default(): - _ = saver_lib.import_meta_graph(meta_graph) - return graph_copy +def parse_str(src): + """Return the AST of given piece of code.""" + return gast.parse(src) diff --git a/tensorflow/contrib/py2tf/pyct/parser_test.py b/tensorflow/contrib/py2tf/pyct/parser_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f35dfa04c70dc191078248c32f9a04d28133129a --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/parser_test.py @@ -0,0 +1,47 @@ +# Copyright 2017 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 parser module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import textwrap + +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.python.platform import test + + +def f(x): + return x + 1 + + +class ParserTest(test.TestCase): + + def test_parse_entity(self): + mod, _ = parser.parse_entity(f) + self.assertEqual('f', mod.body[0].name) + + def test_parse_str(self): + mod = parser.parse_str( + textwrap.dedent(""" + def f(x): + return x + 1 + """)) + self.assertEqual('f', mod.body[0].name) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/pretty_printer.py b/tensorflow/contrib/py2tf/pyct/pretty_printer.py new file mode 100644 index 0000000000000000000000000000000000000000..bacc1e4a7774ec5b84495255042392fe089150d5 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/pretty_printer.py @@ -0,0 +1,113 @@ +# Copyright 2017 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. +# ============================================================================== +"""Print an AST tree in a form more readable than ast.dump.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast +import termcolor + + +class PrettyPrinter(gast.NodeVisitor): + """Print AST nodes.""" + + def __init__(self, color): + self.indent_lvl = 0 + self.result = '' + self.color = color + + def _color(self, string, color, attrs=None): + if self.color: + return termcolor.colored(string, color, attrs=attrs) + return string + + def _type(self, node): + return self._color(node.__class__.__name__, None, ['bold']) + + def _field(self, name): + return self._color(name, 'blue') + + def _value(self, name): + return self._color(name, 'magenta') + + def _warning(self, name): + return self._color(name, 'red') + + def _indent(self): + return self._color('| ' * self.indent_lvl, None, ['dark']) + + def _print(self, s): + self.result += s + self.result += '\n' + + def generic_visit(self, node, name=None): + if node._fields: + cont = ':' + else: + cont = '()' + + if name: + self._print('%s%s=%s%s' % (self._indent(), self._field(name), + self._type(node), cont)) + else: + self._print('%s%s%s' % (self._indent(), self._type(node), cont)) + + self.indent_lvl += 1 + for f in node._fields: + if not hasattr(node, f): + self._print('%s%s' % (self._indent(), self._warning('%s=' % f))) + continue + v = getattr(node, f) + if isinstance(v, list): + if v: + self._print('%s%s=[' % (self._indent(), self._field(f))) + self.indent_lvl += 1 + for n in v: + self.generic_visit(n) + self.indent_lvl -= 1 + self._print('%s]' % (self._indent())) + else: + self._print('%s%s=[]' % (self._indent(), self._field(f))) + elif isinstance(v, tuple): + if v: + self._print('%s%s=(' % (self._indent(), self._field(f))) + self.indent_lvl += 1 + for n in v: + self.generic_visit(n) + self.indent_lvl -= 1 + self._print('%s)' % (self._indent())) + else: + self._print('%s%s=()' % (self._indent(), self._field(f))) + elif isinstance(v, gast.AST): + self.generic_visit(v, f) + elif isinstance(v, str): + self._print('%s%s=%s' % (self._indent(), self._field(f), + self._value('"%s"' % v))) + else: + self._print('%s%s=%s' % (self._indent(), self._field(f), + self._value(v))) + self.indent_lvl -= 1 + + +def fmt(node, color=True): + printer = PrettyPrinter(color) + if isinstance(node, (list, tuple)): + for n in node: + printer.visit(n) + else: + printer.visit(node) + return printer.result diff --git a/tensorflow/contrib/py2tf/pyct/pretty_printer_test.py b/tensorflow/contrib/py2tf/pyct/pretty_printer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..81e3f47b80b6cb3bb7ba9f4a1787d03df4151a99 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/pretty_printer_test.py @@ -0,0 +1,52 @@ +# Copyright 2017 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 pretty_printer module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ast + +from tensorflow.contrib.py2tf.pyct import pretty_printer +from tensorflow.python.platform import test + + +class PrettyPrinterTest(test.TestCase): + + def test_format(self): + node = ast.FunctionDef( + name='f', + args=ast.arguments( + args=[ast.Name(id='a', ctx=ast.Param())], + vararg=None, + kwarg=None, + defaults=[]), + body=[ + ast.Return( + ast.BinOp( + op=ast.Add(), + left=ast.Name(id='a', ctx=ast.Load()), + right=ast.Num(1))) + ], + decorator_list=[], + returns=None) + # Just checking for functionality, the color control characters make it + # difficult to inspect the result. + self.assertIsNotNone(pretty_printer.fmt(node)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/qual_names.py b/tensorflow/contrib/py2tf/pyct/qual_names.py new file mode 100644 index 0000000000000000000000000000000000000000..8717ee6cff198ff31f6cbdb7213e5a8dd3df1149 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/qual_names.py @@ -0,0 +1,104 @@ +# Copyright 2017 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. +# ============================================================================== +"""Utilities for manipulating qualified names. + +A qualified name is a uniform way to refer to simple (e.g. 'foo') and composite +(e.g. 'foo.bar') syntactic symbols. + +This is *not* related to the __qualname__ attribute used by inspect, which +refers to scopes. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import anno + + +class QN(object): + """Represents a qualified name.""" + + def __init__(self, base, attr=None): + if attr: + if not isinstance(base, QN): + raise ValueError('For attribute QNs, base must be a QN.') + self._parent = base + self.qn = base.qn + (attr,) + else: + if isinstance(base, QN): + if base.is_composite(): + self._parent = base.parent + else: + self._parent = None + self.qn = base.qn + else: + self._parent = None + self.qn = tuple(base.split('.')) + + def is_composite(self): + return len(self.qn) > 1 + + @property + def parent(self): + if self._parent is None: + raise ValueError('Cannot get parent of simple name "%s".' % self.qn[0]) + return self._parent + + def __hash__(self): + return hash(self.qn) + + def __eq__(self, other): + return self.qn == other.qn + + def __str__(self): + return '.'.join(self.qn) + + def __repr__(self): + return str(self) + + def ssf(self): + """Simple symbol form.""" + return '_'.join(self.qn) + + def ast(self): + # The caller must adjust the context appropriately. + if self.is_composite(): + return gast.Attribute(self.parent.ast(), self.qn[-1], None) + return gast.Name(self.qn[0], None, None) + + +class QnResolver(gast.NodeTransformer): + """Annotates nodes with QN information. + + Note: Not using NodeAnnos to avoid circular dependencies. + """ + + def visit_Name(self, node): + self.generic_visit(node) + anno.setanno(node, anno.Basic.QN, QN(node.id)) + return node + + def visit_Attribute(self, node): + self.generic_visit(node) + anno.setanno(node, anno.Basic.QN, + QN(anno.getanno(node.value, anno.Basic.QN), node.attr)) + return node + + +def resolve(node): + return QnResolver().visit(node) diff --git a/tensorflow/contrib/py2tf/pyct/qual_names_test.py b/tensorflow/contrib/py2tf/pyct/qual_names_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1b1eee2deca18bb0540c17d6ee85d421602aa2b7 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/qual_names_test.py @@ -0,0 +1,108 @@ +# Copyright 2017 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 qual_names module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import textwrap + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import qual_names +from tensorflow.python.platform import test + + +class QNTest(test.TestCase): + + def test_basic(self): + a = qual_names.QN('a') + self.assertEqual(a.qn, ('a',)) + self.assertEqual(str(a), 'a') + self.assertEqual(a.ssf(), 'a') + self.assertEqual(a.ast().id, 'a') + self.assertFalse(a.is_composite()) + with self.assertRaises(ValueError): + _ = a.parent + + a_b = qual_names.QN(a, 'b') + self.assertEqual(a_b.qn, ('a', 'b')) + self.assertEqual(str(a_b), 'a.b') + self.assertEqual(a_b.ssf(), 'a_b') + self.assertEqual(a_b.ast().value.id, 'a') + self.assertEqual(a_b.ast().attr, 'b') + self.assertTrue(a_b.is_composite()) + self.assertEqual(a_b.parent.qn, ('a',)) + + a2 = qual_names.QN(a) + self.assertEqual(a2.qn, ('a',)) + with self.assertRaises(ValueError): + _ = a.parent + + a_b2 = qual_names.QN(a_b) + self.assertEqual(a_b2.qn, ('a', 'b')) + self.assertEqual(a_b2.parent.qn, ('a',)) + + self.assertTrue(a2 == a) + self.assertFalse(a2 is a) + + self.assertTrue(a_b.parent == a) + self.assertTrue(a_b2.parent == a) + + self.assertTrue(a_b2 == a_b) + self.assertFalse(a_b2 is a_b) + self.assertFalse(a_b2 == a) + + with self.assertRaises(ValueError): + qual_names.QN('a', 'b') + + def test_hashable(self): + d = {qual_names.QN('a'): 'a', qual_names.QN('b'): 'b'} + + self.assertEqual(d[qual_names.QN('a')], 'a') + self.assertEqual(d[qual_names.QN('b')], 'b') + self.assertTrue(qual_names.QN('c') not in d) + + +class QNResolverTest(test.TestCase): + + def assertQNStringIs(self, node, qn_str): + self.assertEqual(str(anno.getanno(node, anno.Basic.QN)), qn_str) + + def test_resolve(self): + samples = """ + a + a.b + (c, d.e) + [f, (g.h.i)] + j(k, l) + """ + nodes = qual_names.resolve(parser.parse_str(textwrap.dedent(samples))) + nodes = tuple(n.value for n in nodes.body) + + self.assertQNStringIs(nodes[0], 'a') + self.assertQNStringIs(nodes[1], 'a.b') + self.assertQNStringIs(nodes[2].elts[0], 'c') + self.assertQNStringIs(nodes[2].elts[1], 'd.e') + self.assertQNStringIs(nodes[3].elts[0], 'f') + self.assertQNStringIs(nodes[3].elts[1], 'g.h.i') + self.assertQNStringIs(nodes[4].func, 'j') + self.assertQNStringIs(nodes[4].args[0], 'k') + self.assertQNStringIs(nodes[4].args[1], 'l') + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/BUILD b/tensorflow/contrib/py2tf/pyct/static_analysis/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..fbfce18c60cca4b105e7de3c3ea7b9c3438f6b2a --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/BUILD @@ -0,0 +1,65 @@ +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "py_test") + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +py_library( + name = "static_analysis", + srcs = [ + "activity.py", + "annos.py", + "live_values.py", + "type_info.py", + ], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/contrib/py2tf/pyct", + "@gast_archive//:gast", + ], +) + +py_test( + name = "activity_test", + srcs = ["activity_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":static_analysis", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + "@gast_archive//:gast", + ], +) + +py_test( + name = "live_values_test", + srcs = ["live_values_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":static_analysis", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "type_info_test", + srcs = ["type_info_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":static_analysis", + "//tensorflow/contrib/py2tf/pyct", + "//tensorflow/python:client_testlib", + ], +) diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/__init__.py b/tensorflow/contrib/py2tf/pyct/static_analysis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c325e19f28376da3be6db4b00b9f664eac047af2 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2017 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. +# ============================================================================== +"""Static information resolution. + +This module contains utilities to help annotate AST nodes with as much runtime +information as can be possibly extracted without actually executing the code, +under that assumption that the context in which the code will run is known. + +Note: It's a fair bet that this analysis cannot be reused across contexts +without re-running it. In most cases, the context usually means referenced +modules, which should be static enough to allow reuse, but that is not being +reliably verified. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/activity.py b/tensorflow/contrib/py2tf/pyct/static_analysis/activity.py new file mode 100644 index 0000000000000000000000000000000000000000..1c93e1603113d48176af7a97a0f37321e6f67586 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/activity.py @@ -0,0 +1,270 @@ +# Copyright 2017 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. +# ============================================================================== +"""Activity analysis.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import copy + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno + +# TODO(mdan): Add support for PY3 (e.g. Param vs arg). + + +class Scope(object): + """Encloses local symbol definition and usage information. + + This can track for instance whether a symbol is modified in the current scope. + Note that scopes do not necessarily align with Python's scopes. For example, + the body of an if statement may be considered a separate scope. + + Attributes: + modified: identifiers modified in this scope + created: identifiers created in this scope + used: identifiers referenced in this scope + """ + + def __init__(self, parent, isolated=True): + """Create a new scope. + + Args: + parent: A Scope or None. + isolated: Whether the scope is isolated, that is, whether variables + created in this scope should be visible to the parent scope. + """ + self.isolated = isolated + self.parent = parent + self.modified = set() + self.created = set() + self.used = set() + self.params = set() + self.returned = set() + + # TODO(mdan): Rename to `locals` + @property + def referenced(self): + if not self.isolated and self.parent is not None: + return self.used | self.parent.referenced + return self.used + + def __repr__(self): + return 'Scope{r=%s, c=%s, w=%s}' % (tuple(self.used), tuple(self.created), + tuple(self.modified)) + + def copy_from(self, other): + self.modified = copy.copy(other.modified) + self.created = copy.copy(other.created) + self.used = copy.copy(other.used) + self.params = copy.copy(other.params) + self.returned = copy.copy(other.returned) + + def merge_from(self, other): + self.modified |= other.modified + self.created |= other.created + self.used |= other.used + self.params |= other.params + self.returned |= other.returned + + def has(self, name): + if name in self.modified or name in self.params: + return True + elif self.parent is not None: + return self.parent.has(name) + return False + + def is_modified_since_entry(self, name): + if name in self.modified: + return True + elif self.parent is not None and not self.isolated: + return self.parent.is_modified_since_entry(name) + return False + + def is_param(self, name): + if name in self.params: + return True + elif self.parent is not None and not self.isolated: + return self.parent.is_param(name) + return False + + def mark_read(self, name): + self.used.add(name) + if self.parent is not None and name not in self.created: + self.parent.mark_read(name) + + def mark_param(self, name): + self.params.add(name) + + def mark_creation(self, name): + if name.is_composite(): + parent = name.parent + if self.has(parent): + # This is considered mutation of the parent, not creation. + # TODO(mdan): Is that really so? + return + else: + raise ValueError('Unknown symbol "%s".' % parent) + self.created.add(name) + + def mark_write(self, name): + self.modified.add(name) + if self.isolated: + self.mark_creation(name) + else: + if self.parent is None: + self.mark_creation(name) + else: + if not self.parent.has(name): + self.mark_creation(name) + self.parent.mark_write(name) + + def mark_returned(self, name): + self.returned.add(name) + if not self.isolated and self.parent is not None: + self.parent.mark_returned(name) + + +class ActivityAnalizer(transformer.Base): + """Annotates nodes with local scope information. See Scope.""" + + def __init__(self, context, parent_scope): + super(ActivityAnalizer, self).__init__(context) + self.scope = Scope(parent_scope) + self._in_return_statement = False + + def _track_symbol(self, node): + qn = anno.getanno(node, anno.Basic.QN) + + if isinstance(node.ctx, gast.Store): + self.scope.mark_write(qn) + elif isinstance(node.ctx, gast.Load): + self.scope.mark_read(qn) + elif isinstance(node.ctx, gast.Param): + # Param contexts appear in function defs, so they have the meaning of + # defining a variable. + # TODO(mdan): This bay be incorrect with nested functions. + # For nested functions, we'll have to add the notion of hiding args from + # the parent scope, not writing to them. + self.scope.mark_creation(qn) + self.scope.mark_param(qn) + else: + raise ValueError('Unknown context %s for node %s.' % (type(node.ctx), qn)) + + anno.setanno(node, NodeAnno.IS_LOCAL, self.scope.has(qn)) + anno.setanno(node, NodeAnno.IS_MODIFIED_SINCE_ENTRY, + self.scope.is_modified_since_entry(qn)) + anno.setanno(node, NodeAnno.IS_PARAM, self.scope.is_param(qn)) + + if self._in_return_statement: + self.scope.mark_returned(qn) + + def visit_Name(self, node): + self.generic_visit(node) + self._track_symbol(node) + return node + + def visit_Attribute(self, node): + self.generic_visit(node) + self._track_symbol(node) + return node + + def visit_Print(self, node): + current_scope = self.scope + args_scope = Scope(current_scope) + self.scope = args_scope + for n in node.values: + self.visit(n) + anno.setanno(node, NodeAnno.ARGS_SCOPE, args_scope) + self.scope = current_scope + return node + + def visit_Call(self, node): + current_scope = self.scope + args_scope = Scope(current_scope, isolated=False) + self.scope = args_scope + for n in node.args: + self.visit(n) + # TODO(mdan): Account starargs, kwargs + for n in node.keywords: + self.visit(n) + anno.setanno(node, NodeAnno.ARGS_SCOPE, args_scope) + self.scope = current_scope + self.visit(node.func) + return node + + def _process_block_node(self, node, block, scope_name): + current_scope = self.scope + block_scope = Scope(current_scope, isolated=False) + self.scope = block_scope + for n in block: + self.visit(n) + anno.setanno(node, scope_name, block_scope) + self.scope = current_scope + return node + + def _process_parallel_blocks(self, parent, children): + # Because the scopes are not isolated, processing any child block + # modifies the parent state causing the other child blocks to be + # processed incorrectly. So we need to checkpoint the parent scope so that + # each child sees the same context. + before_parent = Scope(None) + before_parent.copy_from(self.scope) + after_children = [] + for child, scope_name in children: + self.scope.copy_from(before_parent) + parent = self._process_block_node(parent, child, scope_name) + after_child = Scope(None) + after_child.copy_from(self.scope) + after_children.append(after_child) + for after_child in after_children: + self.scope.merge_from(after_child) + return parent + + def visit_If(self, node): + self.visit(node.test) + node = self._process_parallel_blocks(node, + ((node.body, NodeAnno.BODY_SCOPE), + (node.orelse, NodeAnno.ORELSE_SCOPE))) + return node + + def visit_For(self, node): + self.visit(node.target) + self.visit(node.iter) + node = self._process_parallel_blocks(node, + ((node.body, NodeAnno.BODY_SCOPE), + (node.orelse, NodeAnno.ORELSE_SCOPE))) + return node + + def visit_While(self, node): + self.visit(node.test) + node = self._process_parallel_blocks(node, + ((node.body, NodeAnno.BODY_SCOPE), + (node.orelse, NodeAnno.ORELSE_SCOPE))) + return node + + def visit_Return(self, node): + self._in_return_statement = True + node = self.generic_visit(node) + self._in_return_statement = False + return node + + +def resolve(node, context, parent_scope=None): + return ActivityAnalizer(context, parent_scope).visit(node) diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/activity_test.py b/tensorflow/contrib/py2tf/pyct/static_analysis/activity_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e1eb954a5efef4d6a00ac492e7c85394d54e28c9 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/activity_test.py @@ -0,0 +1,271 @@ +# Copyright 2017 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 activity module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import context +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import qual_names +from tensorflow.contrib.py2tf.pyct.qual_names import QN +from tensorflow.contrib.py2tf.pyct.static_analysis import activity +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno +from tensorflow.python.platform import test + + +class ScopeTest(test.TestCase): + + def test_basic(self): + scope = activity.Scope(None) + self.assertFalse(scope.has(QN('foo'))) + + scope.mark_read(QN('foo')) + self.assertFalse(scope.has(QN('foo'))) + + scope.mark_write(QN('foo')) + self.assertTrue(scope.has(QN('foo'))) + + scope.mark_read(QN('bar')) + self.assertFalse(scope.has(QN('bar'))) + + def test_copy(self): + scope = activity.Scope(None) + scope.mark_write(QN('foo')) + + other = activity.Scope(None) + other.copy_from(scope) + + self.assertTrue(QN('foo') in other.created) + + scope.mark_write(QN('bar')) + scope.copy_from(other) + + self.assertFalse(QN('bar') in scope.created) + + scope.mark_write(QN('bar')) + scope.merge_from(other) + + self.assertTrue(QN('bar') in scope.created) + self.assertFalse(QN('bar') in other.created) + + def test_nesting(self): + scope = activity.Scope(None) + scope.mark_write(QN('foo')) + scope.mark_read(QN('bar')) + + child = activity.Scope(scope) + self.assertTrue(child.has(QN('foo'))) + self.assertTrue(scope.has(QN('foo'))) + + child.mark_write(QN('bar')) + self.assertTrue(child.has(QN('bar'))) + self.assertFalse(scope.has(QN('bar'))) + + def test_referenced(self): + scope = activity.Scope(None) + scope.mark_read(QN('a')) + + child = activity.Scope(scope) + child.mark_read(QN('b')) + + child2 = activity.Scope(child, isolated=False) + child2.mark_read(QN('c')) + + self.assertTrue(QN('c') in child2.referenced) + self.assertTrue(QN('b') in child2.referenced) + self.assertFalse(QN('a') in child2.referenced) + + self.assertTrue(QN('c') in child.referenced) + self.assertTrue(QN('b') in child.referenced) + self.assertFalse(QN('a') in child.referenced) + + +class ActivityAnalizerTest(test.TestCase): + + def _parse_and_analyze(self, test_fn): + node, source = parser.parse_entity(test_fn) + ctx = context.EntityContext( + namer=None, + source_code=source, + source_file=None, + namespace={}, + arg_values=None, + arg_types=None, + recursive=True) + node = qual_names.resolve(node) + node = activity.resolve(node, ctx) + return node + + def test_local_markers(self): + + def test_fn(a): # pylint:disable=unused-argument + b = c # pylint:disable=undefined-variable + while b > 0: + b -= 1 + return b + + node = self._parse_and_analyze(test_fn) + self.assertFalse( + anno.getanno(node.body[0].body[0].value, + NodeAnno.IS_LOCAL)) # c in b = c + self.assertTrue( + anno.getanno(node.body[0].body[1].test.left, + NodeAnno.IS_LOCAL)) # b in b > 0 + self.assertTrue( + anno.getanno(node.body[0].body[2].value, + NodeAnno.IS_LOCAL)) # b in return b + + def assertScopeIs(self, scope, used, modified, created): + self.assertItemsEqual(used, tuple(str(s) for s in scope.used)) + self.assertItemsEqual(modified, tuple(str(s) for s in scope.modified)) + self.assertItemsEqual(created, tuple(str(s) for s in scope.created)) + + def test_print_statement(self): + + def test_fn(a): + b = 0 + c = 1 + print(a, b) + return c + + node = self._parse_and_analyze(test_fn) + print_node = node.body[0].body[2] + if isinstance(print_node, gast.Print): + # Python 2 + print_args_scope = anno.getanno(print_node, NodeAnno.ARGS_SCOPE) + else: + # Python 3 + assert isinstance(print_node, gast.Expr) + # The call node should be the one being annotated. + print_node = print_node.value + print_args_scope = anno.getanno(print_node, NodeAnno.ARGS_SCOPE) + # We basically need to detect which variables are captured by the call + # arguments. + self.assertScopeIs(print_args_scope, ('a', 'b'), (), ()) + + def test_call(self): + + def test_fn(a): + b = 0 + c = 1 + foo(a, b) # pylint:disable=undefined-variable + return c + + node = self._parse_and_analyze(test_fn) + call_node = node.body[0].body[2].value + # We basically need to detect which variables are captured by the call + # arguments. + self.assertScopeIs( + anno.getanno(call_node, NodeAnno.ARGS_SCOPE), ('a', 'b'), (), ()) + + def test_while(self): + + def test_fn(a): + b = a + while b > 0: + c = b + b -= 1 + return b, c + + node = self._parse_and_analyze(test_fn) + while_node = node.body[0].body[1] + self.assertScopeIs( + anno.getanno(while_node, NodeAnno.BODY_SCOPE), ('b',), ('b', 'c'), + ('c',)) + self.assertScopeIs( + anno.getanno(while_node, NodeAnno.BODY_SCOPE).parent, ('a', 'b', 'c'), + ('b', 'c'), ('a', 'b', 'c')) + + def test_for(self): + + def test_fn(a): + b = a + for _ in a: + c = b + b -= 1 + return b, c + + node = self._parse_and_analyze(test_fn) + for_node = node.body[0].body[1] + self.assertScopeIs( + anno.getanno(for_node, NodeAnno.BODY_SCOPE), ('b',), ('b', 'c'), ('c',)) + self.assertScopeIs( + anno.getanno(for_node, NodeAnno.BODY_SCOPE).parent, ('a', 'b', 'c'), + ('b', 'c', '_'), ('a', 'b', 'c', '_')) + + def test_if(self): + + def test_fn(x): + if x > 0: + x = -x + y = 2 * x + z = -y + else: + x = 2 * x + y = -x + u = -y + return z, u + + node = self._parse_and_analyze(test_fn) + if_node = node.body[0].body[0] + self.assertScopeIs( + anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('x', 'y'), ('x', 'y', 'z'), + ('y', 'z')) + # TODO(mdan): Double check: is it ok to not mark a local symbol as not read? + self.assertScopeIs( + anno.getanno(if_node, NodeAnno.BODY_SCOPE).parent, ('x', 'z', 'u'), + ('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u')) + self.assertScopeIs( + anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), ('x', 'y'), + ('x', 'y', 'u'), ('y', 'u')) + self.assertScopeIs( + anno.getanno(if_node, NodeAnno.ORELSE_SCOPE).parent, ('x', 'z', 'u'), + ('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u')) + + def test_call_with_composite_names(self): + + def foo(*_): + pass + + def test_fn(a): + foo(a.b, a.c) + if a > 0: + a.b = 2 + else: + d = 2 + d.e = a.c + f = d.e + 1 + a.c = f + + node = self._parse_and_analyze(test_fn) + call_node = node.body[0].body[0].value + self.assertScopeIs( + anno.getanno(call_node, NodeAnno.ARGS_SCOPE), ('a', 'a.b', 'a.c'), (), + ()) + if_node = node.body[0].body[1] + self.assertScopeIs( + anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('a',), ('a.b',), ()) + self.assertScopeIs( + anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), + ('a', 'a.c', 'd', 'd.e', 'f'), ('a.c', 'd', 'd.e', 'f'), ('d', 'f')) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/annos.py b/tensorflow/contrib/py2tf/pyct/static_analysis/annos.py new file mode 100644 index 0000000000000000000000000000000000000000..2d8e49442364fdd4a4752c8a83a5f3b76117fe57 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/annos.py @@ -0,0 +1,50 @@ +# Copyright 2017 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. +# ============================================================================== +"""Annotations used by the static analizer.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from enum import Enum + + +class NoValue(Enum): + + def __repr__(self): + return self.name + + +class NodeAnno(NoValue): + """Additionnal annotations used by the static analyzer. + + These are in addition to the basic annotations declared in anno.py. + """ + + # Symbols + + IS_LOCAL = 'Symbol is local to the function scope being analized.' + IS_PARAM = 'Symbol is a parameter to the function being analized.' + IS_MODIFIED_SINCE_ENTRY = ( + 'Symbol has been explicitly replaced in the current function scope.') + + # Scopes + ARGS_SCOPE = 'The scope for the argument list of a function call.' + BODY_SCOPE = ( + 'The scope for the main body of a statement (True branch for if ' + 'statements, main body for loops).') + ORELSE_SCOPE = ( + 'The scope for the orelse body of a statement (False branch for if ' + 'statements, orelse body for loops).') diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/live_values.py b/tensorflow/contrib/py2tf/pyct/static_analysis/live_values.py new file mode 100644 index 0000000000000000000000000000000000000000..9c0a9a9e74eccb3d22840032e8f0c2b81e051e7e --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/live_values.py @@ -0,0 +1,112 @@ +# Copyright 2017 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. +# ============================================================================== +"""Live value resolution. + +Live values are extracted from the known execution context. + +Requires activity analysis annotations. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.contrib.py2tf.pyct.static_analysis.annos import NodeAnno + + +class LiveValueResolver(transformer.Base): + """Annotates nodes with live values.""" + + def __init__(self, context, literals): + super(LiveValueResolver, self).__init__(context) + self.literals = literals + + def visit_ClassDef(self, node): + self.generic_visit(node) + anno.setanno(node, 'live_val', self.context.namespace[node.name]) + return node + + def visit_Name(self, node): + self.generic_visit(node) + if isinstance(node.ctx, gast.Load): + assert anno.hasanno(node, NodeAnno.IS_LOCAL), node + symbol_is_local = anno.getanno(node, NodeAnno.IS_LOCAL) + assert anno.hasanno(node, NodeAnno.IS_MODIFIED_SINCE_ENTRY), node + symbol_is_modified = anno.getanno(node, NodeAnno.IS_MODIFIED_SINCE_ENTRY) + assert anno.hasanno(node, NodeAnno.IS_PARAM), node + symbol_is_param = anno.getanno(node, NodeAnno.IS_PARAM) + + if not symbol_is_local and not symbol_is_param: + if node.id in self.literals: + anno.setanno(node, 'live_val', self.literals[node.id]) + # TODO(mdan): Could live values have FQNs? i.e. 'a'.join() + elif node.id in self.context.namespace: + obj = self.context.namespace[node.id] + anno.setanno(node, 'live_val', obj) + anno.setanno(node, 'fqn', (obj.__name__,)) + else: + pass + # TODO(mdan): Should we raise an error here? + # Can encounter this when: + # * a symbol truly lacks reference + # * a symbol is new, like the new name of a function we just renamed. + else: + pass + # TODO(mdan): Attempt to trace its value through the local chain. + # TODO(mdan): Use type annotations as fallback. + + if not symbol_is_modified: + if node.id in self.context.arg_values: + obj = self.context.arg_values[node.id] + anno.setanno(node, 'live_val', obj) + anno.setanno(node, 'fqn', (obj.__class__.__name__,)) + return node + + def visit_Attribute(self, node): + self.generic_visit(node) + if anno.hasanno(node.value, 'live_val'): + assert anno.hasanno(node.value, 'fqn') + parent_object = anno.getanno(node.value, 'live_val') + if not hasattr(parent_object, node.attr): + raise AttributeError('%s has no attribute %s' % (parent_object, + node.attr)) + anno.setanno(node, 'live_val', getattr(parent_object, node.attr)) + anno.setanno(node, 'fqn', anno.getanno(node.value, 'fqn') + (node.attr,)) + # TODO(mdan): Investigate the role built-in annotations can play here. + elif anno.hasanno(node.value, 'type'): + parent_type = anno.getanno(node.value, 'type') + if hasattr(parent_type, node.attr): + # This should hold for static members like methods. + # This would not hold for dynamic members like function attributes. + # For the dynamic case, we simply leave the node without an annotation, + # and let downstream consumers figure out what to do. + anno.setanno(node, 'live_val', getattr(parent_type, node.attr)) + anno.setanno(node, 'fqn', + anno.getanno(node.value, 'type_fqn') + (node.attr,)) + elif isinstance(node.value, gast.Name): + stem_name = node.value + # All nonlocal symbols should be fully resolved. + assert anno.hasanno(stem_name, NodeAnno.IS_LOCAL), stem_name + # TODO(mdan): Figure out what to do when calling attribute on local object + # Maybe just leave as-is? + return node + + +def resolve(node, context, literals): + return LiveValueResolver(context, literals).visit(node) diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/live_values_test.py b/tensorflow/contrib/py2tf/pyct/static_analysis/live_values_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9f64689401e3594a77fbdd7b6f02880bd6e90492 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/live_values_test.py @@ -0,0 +1,109 @@ +# Copyright 2017 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 live_values module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import context +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import qual_names +from tensorflow.contrib.py2tf.pyct.static_analysis import activity +from tensorflow.contrib.py2tf.pyct.static_analysis import live_values +from tensorflow.contrib.py2tf.pyct.static_analysis import type_info +from tensorflow.python.framework import constant_op +from tensorflow.python.platform import test + + +class LiveValuesResolverTest(test.TestCase): + + def _parse_and_analyze(self, + test_fn, + namespace, + literals=None, + arg_types=None): + literals = literals or {} + arg_types = arg_types or {} + node, source = parser.parse_entity(test_fn) + ctx = context.EntityContext( + namer=None, + source_code=source, + source_file=None, + namespace=namespace, + arg_values=None, + arg_types=arg_types, + recursive=True) + node = qual_names.resolve(node) + node = activity.resolve(node, ctx) + node = live_values.resolve(node, ctx, literals) + node = type_info.resolve(node, ctx) + node = live_values.resolve(node, ctx, literals) + return node + + def test_literals(self): + + def test_fn(): + return Foo # pylint: disable=undefined-variable + + node = self._parse_and_analyze(test_fn, {}, {'Foo': 'bar'}) + retval_node = node.body[0].body[0].value + self.assertEquals('bar', anno.getanno(retval_node, 'live_val')) + + def test_namespace(self): + + def foo(): + return 'bar' + + def test_fn(): + return foo() + + node = self._parse_and_analyze(test_fn, {'foo': foo}) + func_node = node.body[0].body[0].value.func + self.assertEquals(foo, anno.getanno(func_node, 'live_val')) + self.assertEquals(('foo',), anno.getanno(func_node, 'fqn')) + + def test_attribute_names(self): + + def test_fn(): + return constant_op.constant(0) + + node = self._parse_and_analyze(test_fn, {'constant_op': constant_op}) + func_node = node.body[0].body[0].value.func + self.assertEquals(constant_op.constant, anno.getanno(func_node, 'live_val')) + self.assertEquals((constant_op.__name__, 'constant'), + anno.getanno(func_node, 'fqn')) + + def test_attributes_with_type_hints(self): + + class TestClass(object): + + def member(self): + pass + + def test_fn(self): + return self.member() + + node = self._parse_and_analyze( + TestClass.test_fn, {'constant_op': constant_op}, + arg_types={'self': (TestClass.__name__, TestClass)}) + func_node = node.body[0].body[0].value.func + self.assertEquals(TestClass.member, anno.getanno(func_node, 'live_val')) + self.assertEquals(('TestClass', 'member'), anno.getanno(func_node, 'fqn')) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/type_info.py b/tensorflow/contrib/py2tf/pyct/static_analysis/type_info.py new file mode 100644 index 0000000000000000000000000000000000000000..8203bda0f9a792a5b24b9abb25d8f39b61625748 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/type_info.py @@ -0,0 +1,186 @@ +# Copyright 2017 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. +# ============================================================================== +"""Type resolution. + +Requires annotations generated by LiveValuesResolver. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import transformer +from tensorflow.python.util import tf_inspect + + +class Scope(object): + """Encloses symbol value references. + + Attributes: + values: A dict mapping string to gast.Node, containing the value that was + most recently assigned to the symbol. + """ + + def __init__(self, parent): + """Create a new scope. + + Args: + parent: A Scope or None. + """ + self.parent = parent + self.values = {} + + def __repr__(self): + return 'Scope[%s]' % self.values.keys() + + def copy(self): + s = Scope(self.parent) + s.values = self.values.copy() + return s + + def setval(self, name, value): + self.values[name] = value + + def hasval(self, name): + return (name in self.values or + (self.parent is not None and self.parent.hasval(name))) + + def getval(self, name): + if name in self.values: + return self.values[name] + if self.parent is not None: + return self.parent.getval(name) + raise KeyError(name) + + +class TypeInfoResolver(transformer.Base): + """Annotates symbols with type information where possible. + + Nodes currently annotated: + * Call (helps detect class constructors) + * Attribute (helps resolve object methods) + """ + + def __init__(self, context): + super(TypeInfoResolver, self).__init__(context) + self.scope = Scope(None) + self.function_level = 0 + + def visit_FunctionDef(self, node): + self.scope = Scope(self.scope) + self.function_level += 1 + self.generic_visit(node) + self.function_level -= 1 + self.scope = self.scope.parent + return node + + def _visit_block(self, block): + self.scope = Scope(self.scope) + for i, n in enumerate(block): + block[i] = self.generic_visit(n) + self.scope = self.scope.parent + return block + + def visit_For(self, node): + self.generic_visit(node.target) + self.generic_visit(node.iter) + node.body = self._visit_block(node.body) + node.orelse = self._visit_block(node.orelse) + return node + + def visit_While(self, node): + self.generic_visit(node.test) + node.body = self._visit_block(node.body) + node.orelse = self._visit_block(node.orelse) + return node + + def visit_If(self, node): + self.generic_visit(node.test) + node.body = self._visit_block(node.body) + node.orelse = self._visit_block(node.orelse) + return node + + def _process_function_arg(self, arg_name): + str_name = str(arg_name) + if self.function_level == 1 and str_name in self.context.arg_types: + # Forge a node to hold the type information, so that method calls on + # it can resolve the type. + type_holder = arg_name.ast() + type_string, type_obj = self.context.arg_types[str_name] + anno.setanno(type_holder, 'type', type_obj) + anno.setanno(type_holder, 'type_fqn', tuple(type_string.split('.'))) + self.scope.setval(arg_name, type_holder) + + def visit_arg(self, node): + self._process_function_arg(anno.getanno(node.arg, anno.Basic.QN)) + return node + + def visit_Name(self, node): + self.generic_visit(node) + qn = anno.getanno(node, anno.Basic.QN) + if isinstance(node.ctx, gast.Param): + self._process_function_arg(qn) + elif isinstance(node.ctx, gast.Load) and self.scope.hasval(qn): + # E.g. if we had + # a = b + # then for future references to `a` we should have traced_source = `b` + traced_source = self.scope.getval(qn) + if anno.hasanno(traced_source, 'type'): + anno.setanno(node, 'type', anno.getanno(traced_source, 'type')) + anno.setanno(node, 'type_fqn', anno.getanno(traced_source, 'type_fqn')) + return node + + def _process_variable_assignment(self, source, targets): + if isinstance(source, gast.Call): + func = source.func + if anno.hasanno(func, 'live_val'): + func_obj = anno.getanno(func, 'live_val') + if tf_inspect.isclass(func_obj): + anno.setanno(source, 'is_constructor', True) + anno.setanno(source, 'type', func_obj) + anno.setanno(source, 'type_fqn', anno.getanno(func, 'fqn')) + # TODO(mdan): Raise an error if constructor has side effects. + # We can have a whitelist of no-side-effects constructors. + # We can also step inside the constructor and further analyze. + + for t in targets: + if isinstance(t, gast.Tuple): + for i, e in enumerate(t.elts): + self.scope.setval( + anno.getanno(e, anno.Basic.QN), + gast.Subscript(source, gast.Index(i), ctx=gast.Store())) + elif isinstance(t, (gast.Name, gast.Attribute)): + self.scope.setval(anno.getanno(t, anno.Basic.QN), source) + else: + raise ValueError('Dont know how to handle assignment to %s' % t) + + def visit_With(self, node): + for wi in node.items: + if wi.optional_vars is not None: + self._process_variable_assignment(wi.context_expr, (wi.optional_vars,)) + self.generic_visit(node) + return node + + def visit_Assign(self, node): + self.generic_visit(node) + self._process_variable_assignment(node.value, node.targets) + return node + + +def resolve(node, context): + return TypeInfoResolver(context).visit(node) diff --git a/tensorflow/contrib/py2tf/pyct/static_analysis/type_info_test.py b/tensorflow/contrib/py2tf/pyct/static_analysis/type_info_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3659f949db9910534870d8dd9e42fd4ee8297253 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/static_analysis/type_info_test.py @@ -0,0 +1,179 @@ +# Copyright 2017 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 type_info module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import context +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import qual_names +from tensorflow.contrib.py2tf.pyct.static_analysis import activity +from tensorflow.contrib.py2tf.pyct.static_analysis import live_values +from tensorflow.contrib.py2tf.pyct.static_analysis import type_info +from tensorflow.python.client import session +from tensorflow.python.platform import test +from tensorflow.python.training import training + + +class ScopeTest(test.TestCase): + + def test_basic(self): + scope = type_info.Scope(None) + self.assertFalse(scope.hasval('foo')) + + scope.setval('foo', 'bar') + self.assertTrue(scope.hasval('foo')) + + self.assertFalse(scope.hasval('baz')) + + def test_nesting(self): + scope = type_info.Scope(None) + scope.setval('foo', '') + + child = type_info.Scope(scope) + self.assertTrue(child.hasval('foo')) + self.assertTrue(scope.hasval('foo')) + + child.setval('bar', '') + self.assertTrue(child.hasval('bar')) + self.assertFalse(scope.hasval('bar')) + + +class TypeInfoResolverTest(test.TestCase): + + def _parse_and_analyze(self, test_fn, namespace, arg_types=None): + node, source = parser.parse_entity(test_fn) + ctx = context.EntityContext( + namer=None, + source_code=source, + source_file=None, + namespace=namespace, + arg_values=None, + arg_types=arg_types, + recursive=True) + node = qual_names.resolve(node) + node = activity.resolve(node, ctx) + node = live_values.resolve(node, ctx, {}) + node = type_info.resolve(node, ctx) + node = live_values.resolve(node, ctx, {}) + return node + + def test_constructor_detection(self): + + def test_fn(): + opt = training.GradientDescentOptimizer(0.1) + return opt + + node = self._parse_and_analyze(test_fn, {'training': training}) + call_node = node.body[0].body[0].value + self.assertEquals(training.GradientDescentOptimizer, + anno.getanno(call_node, 'type')) + self.assertEquals((training.__name__, 'GradientDescentOptimizer'), + anno.getanno(call_node, 'type_fqn')) + + def test_class_members_of_detected_constructor(self): + + def test_fn(): + opt = training.GradientDescentOptimizer(0.1) + opt.minimize(0) + + node = self._parse_and_analyze(test_fn, {'training': training}) + method_call = node.body[0].body[1].value.func + self.assertEquals(training.GradientDescentOptimizer.minimize, + anno.getanno(method_call, 'live_val')) + + def test_class_members_in_with_stmt(self): + + def test_fn(x): + with session.Session() as sess: + sess.run(x) + + node = self._parse_and_analyze(test_fn, {'session': session}) + constructor_call = node.body[0].body[0].items[0].context_expr + self.assertEquals(session.Session, anno.getanno(constructor_call, 'type')) + self.assertEquals((session.__name__, 'Session'), + anno.getanno(constructor_call, 'type_fqn')) + + method_call = node.body[0].body[0].body[0].value.func + self.assertEquals(session.Session.run, anno.getanno(method_call, + 'live_val')) + + def test_constructor_data_dependent(self): + + def test_fn(x): + if x > 0: + opt = training.GradientDescentOptimizer(0.1) + else: + opt = training.GradientDescentOptimizer(0.01) + opt.minimize(0) + + node = self._parse_and_analyze(test_fn, {'training': training}) + method_call = node.body[0].body[1].value.func + self.assertFalse(anno.hasanno(method_call, 'live_val')) + + def test_parameter_class_members(self): + + def test_fn(opt): + opt.minimize(0) + + node = self._parse_and_analyze(test_fn, {}) + method_call = node.body[0].body[0].value.func + self.assertFalse(anno.hasanno(method_call, 'live_val')) + + def test_parameter_class_members_with_value_hints(self): + + def test_fn(opt): + opt.minimize(0) + + node = self._parse_and_analyze( + test_fn, {'training': training}, + arg_types={ + 'opt': (training.GradientDescentOptimizer.__name__, + training.GradientDescentOptimizer) + }) + + method_call = node.body[0].body[0].value.func + self.assertEquals(training.GradientDescentOptimizer.minimize, + anno.getanno(method_call, 'live_val')) + + def test_function_variables(self): + + def bar(): + pass + + def test_fn(): + foo = bar + foo() + + node = self._parse_and_analyze(test_fn, {'bar': bar}) + method_call = node.body[0].body[1].value.func + self.assertFalse(anno.hasanno(method_call, 'live_val')) + + def test_nested_members(self): + + def test_fn(): + foo = training.GradientDescentOptimizer(0.1) + foo.bar.baz() + + node = self._parse_and_analyze(test_fn, {'training': training}) + method_call = node.body[0].body[1].value.func + self.assertFalse(anno.hasanno(method_call, 'live_val')) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/templates.py b/tensorflow/contrib/py2tf/pyct/templates.py new file mode 100644 index 0000000000000000000000000000000000000000..6ee6c0c5ceb70d87779ee313670135cadc5214b5 --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/templates.py @@ -0,0 +1,156 @@ +# Copyright 2017 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. +# ============================================================================== +"""AST conversion templates. + +Adapted from Tangent. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ast +import textwrap + +import gast + +from tensorflow.contrib.py2tf.pyct import ast_util +from tensorflow.contrib.py2tf.pyct import parser +from tensorflow.contrib.py2tf.pyct import qual_names + + +class ReplaceTransformer(gast.NodeTransformer): + """Replace AST nodes.""" + + def __init__(self, replacements): + """Create a new ReplaceTransformer. + + Args: + replacements: A mapping from placeholder names to (lists of) AST nodes + that these placeholders will be replaced by. + """ + self.replacements = replacements + self.in_replacements = False + + # TODO(mdan): Make a more detailed pass and clean up if needed. + + def visit_Expr(self, node): + if (isinstance(node.value, gast.Name) and + node.value.id in self.replacements): + return self.visit(node.value) + self.generic_visit(node) + return node + + def visit_FunctionDef(self, node): + node = self.generic_visit(node) + if node.name in self.replacements: + repl = self.replacements[node.name] + if not isinstance(repl, (gast.Name, ast.Name)): + raise ValueError( + 'A function name can only be replaced by a Name node. Found: %s' % + repl) + node.name = repl.id + return node + + def _set_inner_child_context(self, node, ctx): + if isinstance(node, gast.Attribute): + self._set_inner_child_context(node.value, ctx) + node.ctx = gast.Load() + elif isinstance(node, gast.Tuple): + for e in node.elts: + self._set_inner_child_context(e, ctx) + node.ctx = ctx + elif isinstance(node, gast.Name): + node.ctx = ctx + elif isinstance(node, (gast.Str, gast.Num)): + pass + else: + raise ValueError('unexpected node type "%s"' % node) + + def visit_Name(self, node): + if node.id not in self.replacements: + return node + + new_nodes = ast_util.copy_clean(self.replacements[node.id]) + if isinstance(new_nodes, gast.AST): + new_nodes = [new_nodes] + + # Preserve the target context. + for n in new_nodes: + if isinstance(n, gast.Tuple): + for e in n.elts: + self._set_inner_child_context(e, node.ctx) + if isinstance(n, gast.Attribute): + # For attributes, the inner Name node receives the context, while the + # outer ones have it set to Load. + self._set_inner_child_context(n, node.ctx) + else: + n.ctx = node.ctx + + if len(new_nodes) == 1: + new_nodes, = new_nodes + + return new_nodes + + +def _convert_to_ast(n): + """Convert from a known data type to AST.""" + if isinstance(n, str): + # Note: the node will receive the ctx value from the template, see + # ReplaceTransformer.visit_Name. + return gast.Name(id=n, ctx=None, annotation=None) + if isinstance(n, qual_names.QN): + return n.ast() + if isinstance(n, list): + return [_convert_to_ast(e) for e in n] + if isinstance(n, tuple): + return tuple(_convert_to_ast(e) for e in n) + return n + + +def replace(template, **replacements): + """Replace placeholders in a Python template. + + AST Name and Tuple nodes always receive the context that inferred from + the template. However, when replacing more complex nodes (that can potentially + contain Name children), then the caller is responsible for setting the + appropriate context. + + Args: + template: A string representing Python code. Any symbol name can be used + that appears in the template code can be used as placeholder. + **replacements: A mapping from placeholder names to (lists of) AST nodes + that these placeholders will be replaced by. String values are also + supported as a shorthand for AST Name nodes with the respective ID. + + Returns: + An AST node or list of AST nodes with the replacements made. If the + template was a function, a list will be returned. If the template was a + node, the same node will be returned. If the template was a string, an + AST node will be returned (a `Module` node in the case of a multi-line + string, an `Expr` node otherwise). + + Raises: + ValueError: if the arguments are incorrect. + """ + if not isinstance(template, str): + raise ValueError('Expected string template, got %s' % type(template)) + tree = parser.parse_str(textwrap.dedent(template)) + for k in replacements: + replacements[k] = _convert_to_ast(replacements[k]) + results = ReplaceTransformer(replacements).visit(tree).body + if isinstance(results, list): + return [qual_names.resolve(r) for r in results] + return qual_names.resolve(results) diff --git a/tensorflow/contrib/py2tf/pyct/templates_test.py b/tensorflow/contrib/py2tf/pyct/templates_test.py new file mode 100644 index 0000000000000000000000000000000000000000..8ccfde8573724741b0bbe4eacb3c54beb381ee7e --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/templates_test.py @@ -0,0 +1,84 @@ +# Copyright 2017 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 templates module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gast + +from tensorflow.contrib.py2tf.pyct import compiler +from tensorflow.contrib.py2tf.pyct import templates +from tensorflow.python.platform import test + + +class TemplatesTest(test.TestCase): + + def test_replace_tuple(self): + template = """ + def test_fn(a, c): + return b, + """ + + node = templates.replace(template, b=('a', 'c'))[0] + result, _ = compiler.ast_to_object(node) + + self.assertEquals((2, 3), result.test_fn(2, 3)) + + def test_replace_variable(self): + template = """ + def test_fn(a): + a += 1 + a = 2 * a + 1 + return b + """ + + node = templates.replace(template, a='b')[0] + result, _ = compiler.ast_to_object(node) + self.assertEquals(7, result.test_fn(2)) + + def test_replace_function_name(self): + template = """ + def fname(a): + a += 1 + a = 2 * a + 1 + return a + """ + + node = templates.replace(template, fname='test_fn')[0] + result, _ = compiler.ast_to_object(node) + self.assertEquals(7, result.test_fn(2)) + + def test_code_block(self): + template = """ + def test_fn(a): + block + return a + """ + + node = templates.replace( + template, + block=[ + gast.Assign([ + gast.Name('a', None, None) + ], gast.BinOp(gast.Name('a', None, None), gast.Add(), gast.Num(1))), + ] * 2)[0] + result, _ = compiler.ast_to_object(node) + self.assertEquals(3, result.test_fn(1)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/pyct/transformer.py b/tensorflow/contrib/py2tf/pyct/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..877d52af016af720424c8a56257fec9ab64611cb --- /dev/null +++ b/tensorflow/contrib/py2tf/pyct/transformer.py @@ -0,0 +1,69 @@ +# Copyright 2017 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. +# ============================================================================== +"""A node transformer that includes utilities for SCT.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import sys + +import gast +import six + +from tensorflow.contrib.py2tf.pyct import anno +from tensorflow.contrib.py2tf.pyct import pretty_printer + + +class PyFlowParseError(SyntaxError): + pass + + +class Base(gast.NodeTransformer): + """Base class for specialized transformers.""" + + def __init__(self, context): + """Initialize the transformer. Subclasses should call this. + + Args: + context: An EntityContext. + """ + self._lineno = 0 + self._col_offset = 0 + self.context = context + + def visit(self, node): + source_code = self.context.source_code + source_file = self.context.source_file + try: + if source_code and hasattr(node, 'lineno'): + self._lineno = node.lineno + self._col_offset = node.col_offset + if anno.hasanno(node, anno.Basic.SKIP_PROCESSING): + return node + return super(Base, self).visit(node) + except (ValueError, AttributeError, KeyError, NotImplementedError, + AssertionError) as e: + msg = '%s: %s\nOccurred at node:\n%s' % ( + e.__class__.__name__, str(e), pretty_printer.fmt(node, color=False)) + if source_code: + line = source_code.splitlines()[self._lineno - 1] + else: + line = '' + six.reraise(PyFlowParseError, + PyFlowParseError( + msg, + (source_file, self._lineno, self._col_offset + 1, line)), + sys.exc_info()[2]) diff --git a/tensorflow/contrib/py2tf/utils/BUILD b/tensorflow/contrib/py2tf/utils/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..c2fdd40707775783140390e4b5c0186c9c3e562e --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/BUILD @@ -0,0 +1,108 @@ +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + +load("//tensorflow:tensorflow.bzl", "py_test") + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +py_library( + name = "utils", + srcs = [ + "__init__.py", + "context_managers.py", + "misc.py", + "multiple_dispatch.py", + "printing.py", + "py_func.py", + "tensor_list.py", + "type_check.py", + ], + srcs_version = "PY2AND3", + visibility = ["//tensorflow:__subpackages__"], + deps = [ + "//tensorflow/python:script_ops", + "@six_archive//:six", + ], +) + +py_test( + name = "context_managers_test", + srcs = ["context_managers_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":utils", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "misc_test", + srcs = ["misc_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":utils", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "multiple_dispatch_test", + srcs = ["multiple_dispatch_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":utils", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "py_func_test", + srcs = ["py_func_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":utils", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "printing_test", + srcs = ["printing_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":utils", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "type_check_test", + srcs = ["type_check_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":utils", + "//tensorflow/python:client_testlib", + ], +) + +py_test( + name = "tensor_list_test", + srcs = ["tensor_list_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":utils", + "//tensorflow/python:client_testlib", + "//tensorflow/python:list_ops", + ], +) diff --git a/tensorflow/contrib/py2tf/utils/__init__.py b/tensorflow/contrib/py2tf/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0a1b993fd366e1317e5f7e01fe849d86c93b8fc2 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2016 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. +# ============================================================================== +"""Utility module that contains APIs usable in the generated code.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.utils.context_managers import control_dependency_on_returns +from tensorflow.contrib.py2tf.utils.misc import alias_tensors +from tensorflow.contrib.py2tf.utils.multiple_dispatch import run_cond +from tensorflow.contrib.py2tf.utils.multiple_dispatch import run_while +from tensorflow.contrib.py2tf.utils.printing import call_print +from tensorflow.contrib.py2tf.utils.py_func import wrap_py_func +from tensorflow.contrib.py2tf.utils.type_check import is_tensor diff --git a/tensorflow/contrib/py2tf/utils/context_managers.py b/tensorflow/contrib/py2tf/utils/context_managers.py new file mode 100644 index 0000000000000000000000000000000000000000..38d9e11fe9069722b9023fee848bf53e1f72de6a --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/context_managers.py @@ -0,0 +1,42 @@ +# Copyright 2017 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. +# ============================================================================== +"""Various context managers.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import contextlib + +from tensorflow.python.framework import ops + + +def control_dependency_on_returns(return_value): + """Create a TF control dependency on the return values of a function. + + If the function had no return value, a no-op context is returned. + + Args: + return_value: The return value to set as control dependency. + + Returns: + A context manager. + """ + if return_value is None: + return contextlib.contextmanager(lambda: (yield))() + # TODO(mdan): Filter to tensor objects. + if not isinstance(return_value, (list, tuple)): + return_value = (return_value,) + return ops.control_dependencies(return_value) diff --git a/tensorflow/contrib/py2tf/utils/context_managers_test.py b/tensorflow/contrib/py2tf/utils/context_managers_test.py new file mode 100644 index 0000000000000000000000000000000000000000..633ba93540e696889a6b2b71b40b999da39d48ff --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/context_managers_test.py @@ -0,0 +1,42 @@ +# Copyright 2017 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 context_managers module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.utils import context_managers +from tensorflow.python.framework import constant_op +from tensorflow.python.platform import test + + +class ContextManagersTest(test.TestCase): + + def test_control_dependency_on_returns(self): + # Just dry run them. + with context_managers.control_dependency_on_returns(None): + pass + with context_managers.control_dependency_on_returns( + constant_op.constant(1)): + pass + with context_managers.control_dependency_on_returns( + [constant_op.constant(1), + constant_op.constant(2)]): + pass + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/utils/misc.py b/tensorflow/contrib/py2tf/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..1b06caf0bdeb6f4a079e33f2e887d2dca017adc2 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/misc.py @@ -0,0 +1,50 @@ +# Copyright 2017 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. +# ============================================================================== +"""Miscellaneous utilities that don't fit anywhere else.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops + + +def alias_tensors(*args): + """Wrap any Tensor arguments with an identity op. + + Any other argument, including Variables, is returned unchanged. + + Args: + *args: Any arguments. Must contain at least one element. + + Returns: + Same as *args, with Tensor instances replaced as described. + + Raises: + ValueError: If args doesn't meet the requirements. + """ + + def alias_if_tensor(a): + return array_ops.identity(a) if isinstance(a, ops.Tensor) else a + + # TODO(mdan): Recurse into containers? + # TODO(mdan): Anything we can do about variables? Fake a scope reuse? + if len(args) > 1: + return (alias_if_tensor(a) for a in args) + elif len(args) == 1: + return alias_if_tensor(args[0]) + + raise ValueError('at least one argument required') diff --git a/tensorflow/contrib/py2tf/utils/misc_test.py b/tensorflow/contrib/py2tf/utils/misc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..bfcb304c838df69e9e3961907362c7939c065117 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/misc_test.py @@ -0,0 +1,54 @@ +# Copyright 2017 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 misc module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.utils import misc +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +class ContextManagersTest(test.TestCase): + + def test_alias_single_tensor(self): + a = constant_op.constant(1) + + new_a = misc.alias_tensors(a) + self.assertFalse(new_a is a) + with self.test_session() as sess: + self.assertEqual(1, sess.run(new_a)) + + def test_alias_tensors(self): + a = constant_op.constant(1) + v = variables.Variable(2) + s = 'a' + l = [1, 2, 3] + + new_a, new_v, new_s, new_l = misc.alias_tensors(a, v, s, l) + + self.assertFalse(new_a is a) + self.assertTrue(new_v is v) + self.assertTrue(new_s is s) + self.assertTrue(new_l is l) + with self.test_session() as sess: + self.assertEqual(1, sess.run(new_a)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/utils/multiple_dispatch.py b/tensorflow/contrib/py2tf/utils/multiple_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..a855fdc075941915035d1e3380846ff912803494 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/multiple_dispatch.py @@ -0,0 +1,90 @@ +# Copyright 2016 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. +# ============================================================================== +"""Utilities for type-dependent behavior used in py2tf-generated code.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import six + +from tensorflow.contrib.py2tf.utils.type_check import is_tensor +from tensorflow.python.ops import control_flow_ops + + +def run_cond(condition, true_fn, false_fn): + """Type-dependent functional conditional. + + Args: + condition: A Tensor or Python bool. + true_fn: A Python callable implementing the true branch of the conditional. + false_fn: A Python callable implementing the false branch of the + conditional. + + Returns: + result: The result of calling the appropriate branch. If condition is a + Tensor, tf.cond will be used. Otherwise, a standard Python if statement will + be ran. + """ + if is_tensor(condition): + return control_flow_ops.cond(condition, true_fn, false_fn) + else: + return py_cond(condition, true_fn, false_fn) + + +def py_cond(condition, true_fn, false_fn): + if condition: + return true_fn() + else: + return false_fn() + + +def run_while(cond_fn, body_fn, init_args): + """Type-dependent functional while loop. + + Args: + cond_fn: A Python callable implementing the stop conditions of the loop. + body_fn: A Python callable implementing the body of the loop. + init_args: The initial values of the arguments that will be passed to both + cond_fn and body_fn. + + Returns: + result: A list of values with the same shape and type as init_args. If any + of the init_args, or any variables closed-over in cond_fn are Tensors, + tf.while_loop will be used, otherwise a Python while loop will be ran. + + Raises: + ValueError: if init_args is not a tuple or list with one or more elements. + """ + if not isinstance(init_args, (tuple, list)) or not init_args: + raise ValueError( + 'init_args must be a non-empty list or tuple, found %s' % init_args) + + # TODO(alexbw): statically determine all active variables in cond_fn, + # and pass them directly + closure_vars = tuple( + [c.cell_contents for c in six.get_function_closure(cond_fn) or []]) + possibly_tensors = tuple(init_args) + closure_vars + if is_tensor(*possibly_tensors): + return control_flow_ops.while_loop(cond_fn, body_fn, init_args) + else: + return py_while_loop(cond_fn, body_fn, init_args) + + +def py_while_loop(cond_fn, body_fn, init_args): + state = init_args + while cond_fn(*state): + state = body_fn(*state) + return state diff --git a/tensorflow/contrib/py2tf/utils/multiple_dispatch_test.py b/tensorflow/contrib/py2tf/utils/multiple_dispatch_test.py new file mode 100644 index 0000000000000000000000000000000000000000..5bb4d4086b002211eebb86783bb7212c707a1418 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/multiple_dispatch_test.py @@ -0,0 +1,69 @@ +# Copyright 2017 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 multiple_dispatch.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +from tensorflow.contrib.py2tf.utils import multiple_dispatch +from tensorflow.python.client.session import Session +from tensorflow.python.framework.constant_op import constant +from tensorflow.python.platform import test + + +class MultipleDispatchTest(test.TestCase): + + def test_run_cond_python(self): + true_fn = lambda: 2.0 + false_fn = lambda: 3.0 + self.assertEqual(multiple_dispatch.run_cond(True, true_fn, false_fn), 2.0) + self.assertEqual(multiple_dispatch.run_cond(False, true_fn, false_fn), 3.0) + + def test_run_cond_tf(self): + + true_fn = lambda: constant([2.0]) + false_fn = lambda: constant([3.0]) + with Session() as sess: + out = multiple_dispatch.run_cond(constant(True), true_fn, false_fn) + self.assertEqual(sess.run(out), 2.0) + out = multiple_dispatch.run_cond(constant(False), true_fn, false_fn) + self.assertEqual(sess.run(out), 3.0) + + def test_run_while_python(self): + cond_fn = lambda x, t, s: x > t + body_fn = lambda x, t, s: (x * s, t, s) + + x, _, _ = multiple_dispatch.run_while(cond_fn, body_fn, [3.0, 1.0, 0.5]) + self.assertEqual(x, 0.75) + + x, _, _ = multiple_dispatch.run_while(cond_fn, body_fn, [3.0, 4.0, 0.5]) + self.assertEqual(x, 3.0) + + def test_run_while_tf(self): + cond_fn = lambda x, t, s: x > t + body_fn = lambda x, t, s: (x * s, t, s) + + with Session() as sess: + x, _, _ = multiple_dispatch.run_while(cond_fn, body_fn, + [constant(3.0), 1.0, 0.5]) + self.assertEqual(sess.run(x), 0.75) + + x, _, _ = multiple_dispatch.run_while(cond_fn, body_fn, + [constant(3.0), 4.0, 0.5]) + self.assertEqual(sess.run(x), 3.0) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/utils/printing.py b/tensorflow/contrib/py2tf/utils/printing.py new file mode 100644 index 0000000000000000000000000000000000000000..95a62bd80b5f4854e6a062df18d882f7bd495555 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/printing.py @@ -0,0 +1,47 @@ +# Copyright 2017 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. +# ============================================================================== +"""TensorFlow printing support utilities.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.utils import py_func +from tensorflow.python.ops import logging_ops + + +def is_tf_print_compatible(value): + # TODO(mdan): Enable once we can reliably test this. + # This is currently disabled because we can't capture the output of + # op kernels from Python. + del value + return False + + +def call_print(*values): + """Compiled counterpart of the print builtin. + + The function attempts to use tf.Print if all the values are compatible. + Otherwise, it will fall back to py_func. + + Args: + *values: values to print + Returns: + A dummy value indicating the print completed. If tf. + """ + + if all(map(is_tf_print_compatible, values)): + return logging_ops.Print(1, values) + return py_func.wrap_py_func(print, None, values, use_dummy_return=True) diff --git a/tensorflow/contrib/py2tf/utils/printing_test.py b/tensorflow/contrib/py2tf/utils/printing_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2070deb304d8df2433fb9a95ae36d48415578482 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/printing_test.py @@ -0,0 +1,53 @@ +# Copyright 2017 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 printing module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import sys + +import six + +from tensorflow.contrib.py2tf.utils import printing +from tensorflow.python.platform import test + + +class ContextManagersTest(test.TestCase): + + def test_call_print_tf(self): + try: + out_capturer = six.StringIO() + sys.stdout = out_capturer + with self.test_session() as sess: + sess.run(printing.call_print('test message', 1)) + self.assertEqual(out_capturer.getvalue(), 'test message 1\n') + finally: + sys.stdout = sys.__stdout__ + + def test_call_print_py_func(self): + try: + out_capturer = six.StringIO() + sys.stdout = out_capturer + with self.test_session() as sess: + sess.run(printing.call_print('test message', [1, 2])) + self.assertEqual(out_capturer.getvalue(), 'test message [1, 2]\n') + finally: + sys.stdout = sys.__stdout__ + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/utils/py_func.py b/tensorflow/contrib/py2tf/utils/py_func.py new file mode 100644 index 0000000000000000000000000000000000000000..838872d092a3ab07e965180eff4fec7ff6c4ccf9 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/py_func.py @@ -0,0 +1,69 @@ +# Copyright 2017 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. +# ============================================================================== +"""Pyfunc creation utilities.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import tensor_util +from tensorflow.python.ops import script_ops + + +def wrap_py_func(f, return_dtypes, arguments, use_dummy_return=False): + """Helper that wraps a callable to py_func. + + The helper passes tensor arguments through the py_func interface. Non-tensor + arguments are allowed, and will be passed to f directly. Note that non-tensor + arguments are captured by f will not update every time the wrapper is + called (this is consistent with its argument list, which only includes + the tensor arguments). In general, it's safest not to reuse this wrapper. + + Args: + f: Callable + return_dtypes: DType, tuple, list or None, the data type for each of f's + return value. None if f has no return values or use_dummy_return is + True. + arguments: Arguments for f + use_dummy_return: If True, the function will return a dummy value of 1 + and discard its actual return value. + Returns: + The return values of f converted to tensor. + Raises: + ValueError: if the arguments are incorrect. + """ + + if return_dtypes and use_dummy_return: + raise ValueError('if use_dummy_return is True, return_dtypes must be empty') + + n = len(arguments) + arg_is_tensor = tuple(map(tensor_util.is_tensor, arguments)) + index_in_tensor_list = [0] * n + i = 0 + for j in range(n): + index_in_tensor_list[j] = i + if arg_is_tensor[j]: + i += 1 + + def f_wrapper(*tensor_args): + f_args = tuple(tensor_args[index_in_tensor_list[i]] + if arg_is_tensor[i] else arguments[i] for i in range(n)) + retval = f(*f_args) + return 1 if use_dummy_return else retval + + return script_ops.py_func( + f_wrapper, tuple(arguments[i] for i in range(n) if arg_is_tensor[i]), + dtypes.int64 if use_dummy_return else return_dtypes) diff --git a/tensorflow/contrib/py2tf/utils/py_func_test.py b/tensorflow/contrib/py2tf/utils/py_func_test.py new file mode 100644 index 0000000000000000000000000000000000000000..776b5309c6f027bb2008aa83d48e4155e817ed97 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/py_func_test.py @@ -0,0 +1,91 @@ +# Copyright 2017 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 wrap_py_func module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.utils import py_func +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.platform import test + + +class PyFuncTest(test.TestCase): + + def test_wrap_py_func_simple(self): + + def test_fn(a, b, c): + return a + b + c + + with self.test_session() as sess: + tensor_1 = constant_op.constant(1) + self.assertEqual(3, + sess.run( + py_func.wrap_py_func(test_fn, dtypes.int64, + (1, tensor_1, 1)))) + self.assertEqual(3, + sess.run( + py_func.wrap_py_func(test_fn, dtypes.int64, + (1, 1, 1)))) + self.assertEqual(3, + sess.run( + py_func.wrap_py_func(test_fn, dtypes.int64, + (tensor_1, 1, tensor_1)))) + + def test_wrap_py_func_complex_args(self): + + class TestClass(object): + + def __init__(self): + self.foo = 5 + + def test_fn(a, b): + return a * b.foo + + with self.test_session() as sess: + self.assertEqual(35, + sess.run( + py_func.wrap_py_func(test_fn, dtypes.int64, + (7, TestClass())))) + self.assertEqual( + 35, + sess.run( + py_func.wrap_py_func(test_fn, dtypes.int64, + (constant_op.constant(7), TestClass())))) + + def test_wrap_py_func_dummy_return(self): + + side_counter = [0] + + def test_fn(_): + side_counter[0] += 1 + + with self.test_session() as sess: + self.assertEqual(1, + sess.run( + py_func.wrap_py_func(test_fn, None, (5,), True))) + self.assertEqual([1], side_counter) + self.assertEqual(1, + sess.run( + py_func.wrap_py_func(test_fn, None, + (constant_op.constant(5),), + True))) + self.assertEqual([2], side_counter) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/utils/tensor_list.py b/tensorflow/contrib/py2tf/utils/tensor_list.py new file mode 100644 index 0000000000000000000000000000000000000000..b6ff49e2a0eff384f10903e12212ab929e267804 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/tensor_list.py @@ -0,0 +1,49 @@ +# Copyright 2016 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. +# ============================================================================== +"""A typed list in Python.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.ops import list_ops + + +class TensorList(object): + """Tensor list wrapper API-compatible with Python built-in list.""" + + def __init__(self, shape, dtype): + self.dtype = dtype + self.shape = shape + self.clear() + + def append(self, value): + self.list_ = list_ops.tensor_list_push_back(self.list_, value) + + def pop(self): + self.list_, value = list_ops.tensor_list_pop_back(self.list_, self.dtype) + return value + + def clear(self): + self.list_ = list_ops.empty_tensor_list(self.shape, self.dtype) + + def count(self): + return list_ops.tensor_list_length(self.list_) + + def __getitem__(self, key): + return list_ops.tensor_list_get_item(self.list_, key, self.dtype) + + def __setitem__(self, key, value): + self.list_ = list_ops.tensor_list_set_item(self.list_, key, value) diff --git a/tensorflow/contrib/py2tf/utils/tensor_list_test.py b/tensorflow/contrib/py2tf/utils/tensor_list_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b5e554a162674e08da21785dcbe193c54647f128 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/tensor_list_test.py @@ -0,0 +1,89 @@ +# Copyright 2017 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 PyFlow list.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.py2tf.utils import tensor_list as tl +from tensorflow.python.client.session import Session +from tensorflow.python.eager import context +from tensorflow.python.framework import ops +from tensorflow.python.framework.constant_op import constant +from tensorflow.python.platform import test + + +class TensorListTest(test.TestCase): + + def test_list_append_python(self): + with context.eager_mode(): + a = constant(3.0) + l = tl.TensorList(a.shape, a.dtype) + l.append(a) + self.assertEqual(l.count().numpy(), 1) + l.append(a) + self.assertEqual(l.count().numpy(), 2) + _ = l.pop() + self.assertEqual(l.count().numpy(), 1) + a2 = l.pop() + self.assertEqual(l.count().numpy(), 0) + self.assertEqual(a.numpy(), a2.numpy()) + + def test_list_index_python(self): + with context.eager_mode(): + a = constant(3.0) + b = constant(2.0) + l = tl.TensorList(a.shape, a.dtype) + l.append(a) + self.assertEqual(l[0].numpy(), a.numpy()) + l[0] = ops.convert_to_tensor(b) + self.assertEqual(l[0].numpy(), b.numpy()) + + def test_list_append_tf(self): + a = constant(3.0) + l = tl.TensorList(a.shape, a.dtype) + l.append(a) + c1 = l.count() + l.append(a) + c2 = l.count() + _ = l.pop() + c3 = l.count() + a2 = l.pop() + c4 = l.count() + with Session() as sess: + c1, c2, c3, c4, a, a2 = sess.run([c1, c2, c3, c4, a, a2]) + self.assertEqual(c1, 1) + self.assertEqual(c2, 2) + self.assertEqual(c3, 1) + self.assertEqual(c4, 0) + self.assertEqual(a, a2) + + def test_list_index_tf(self): + a = constant(3.0) + b = constant(2.0) + l = tl.TensorList(a.shape, a.dtype) + l.append(a) + l0 = l[0] + l[0] = b + l1 = l[0] + with self.test_session() as sess: + l0, l1, a, b = sess.run([l0, l1, a, b]) + self.assertEqual(l0, a) + self.assertEqual(l1, b) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/py2tf/utils/type_check.py b/tensorflow/contrib/py2tf/utils/type_check.py new file mode 100644 index 0000000000000000000000000000000000000000..9ca2dec872c8a9ca7bedaa8603f70e3214a3e24a --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/type_check.py @@ -0,0 +1,33 @@ +# Copyright 2016 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. +# ============================================================================== +"""Utilities used in py2tf-generated code.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import tensor_util + + +def is_tensor(*args): + """Check if all arguments are tensors. + + Args: + *args: Python objects that may or may not be tensors. + + Returns: + True if all *args are TensorFlow types, False if one or more are not. + """ + return any([tensor_util.is_tensor(a) for a in args]) diff --git a/tensorflow/contrib/py2tf/utils/type_check_test.py b/tensorflow/contrib/py2tf/utils/type_check_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7d0428e9cccecdc67511e236bc00655a055aea29 --- /dev/null +++ b/tensorflow/contrib/py2tf/utils/type_check_test.py @@ -0,0 +1,43 @@ +# Copyright 2017 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 type_check.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy + +from tensorflow.contrib.py2tf.utils import type_check +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import test_util +from tensorflow.python.platform import test + + +class TypeCheckTest(test.TestCase): + + def test_checks(self): + self.assertTrue(type_check.is_tensor(constant_op.constant([1, 2, 3]))) + self.assertTrue( + type_check.is_tensor(test_util.variables.Variable([1, 2, 3]))) + self.assertTrue( + type_check.is_tensor( + test_util.array_ops.placeholder(test_util.dtypes.float32))) + self.assertFalse(type_check.is_tensor(3)) + self.assertFalse(type_check.is_tensor(numpy.eye(3))) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/quantize/BUILD b/tensorflow/contrib/quantize/BUILD index 3c5b34a0a6adb2f4e340a8e378c1eb51a2e2b534..aec9f47ccb20349c08bbe2fd813ee24a807f9fe3 100644 --- a/tensorflow/contrib/quantize/BUILD +++ b/tensorflow/contrib/quantize/BUILD @@ -13,6 +13,20 @@ py_library( deps = [], ) +py_test( + name = "common_test", + size = "small", + srcs = ["python/common_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":common", + "//tensorflow/python:framework_ops", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:platform_test", + "//tensorflow/python:session", + ], +) + py_library( name = "graph_matcher", srcs = [ @@ -75,11 +89,18 @@ py_library( ":graph_matcher", ":input_to_ops", "//tensorflow/contrib/graph_editor:graph_editor_py", + "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", + "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", + "//tensorflow/python:layers", "//tensorflow/python:math_ops", "//tensorflow/python:nn", "//tensorflow/python:nn_ops", + "//tensorflow/python:ops", + "//tensorflow/python:training", + "//tensorflow/python:util", + "//tensorflow/python:variables", ], ) @@ -88,7 +109,6 @@ py_test( srcs = ["python/fold_batch_norms_test.py"], srcs_version = "PY2AND3", deps = [ - ":copy_graph", ":fold_batch_norms", "//tensorflow/contrib/layers:layers_py", "//tensorflow/python:array_ops", @@ -103,31 +123,7 @@ py_test( "//tensorflow/python:random_ops", "//tensorflow/python:random_seed", "//tensorflow/python:session", - "//tensorflow/python:variables", - ], -) - -py_library( - name = "copy_graph", - srcs = ["python/copy_graph.py"], - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/python:framework_ops", "//tensorflow/python:training", - ], -) - -py_test( - name = "copy_graph_test", - size = "small", - srcs = ["python/copy_graph_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":copy_graph", - "//tensorflow/python:constant_op", - "//tensorflow/python:framework_ops", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", "//tensorflow/python:variables", ], ) @@ -158,7 +154,6 @@ py_test( "//tensorflow/python:array_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", - "//tensorflow/python:framework_test_lib", "//tensorflow/python:platform_test", "//tensorflow/python:session", "//tensorflow/python:variables", @@ -170,7 +165,7 @@ py_library( srcs = ["python/quantize.py"], srcs_version = "PY2AND3", deps = [ - ":common", + ":graph_matcher", ":input_to_ops", ":quant_ops", "//tensorflow/contrib/graph_editor:graph_editor_py", @@ -217,7 +212,6 @@ py_test( "//tensorflow/python:math_ops", "//tensorflow/python:nn_ops", "//tensorflow/python:platform_test", - "//tensorflow/python:training", ], ) @@ -229,12 +223,9 @@ py_library( ], srcs_version = "PY2AND3", deps = [ - ":copy_graph", ":fold_batch_norms", ":quantize", - "//tensorflow/python:framework_ops", "//tensorflow/python:util", - "//tensorflow/python:variables", ], ) @@ -247,13 +238,11 @@ py_test( ":quantize_graph", "//tensorflow/contrib/layers:layers_py", "//tensorflow/python:array_ops", - "//tensorflow/python:constant_op", "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", "//tensorflow/python:init_ops", "//tensorflow/python:nn_ops", "//tensorflow/python:platform_test", - "//tensorflow/python:variables", ], ) diff --git a/tensorflow/contrib/quantize/__init__.py b/tensorflow/contrib/quantize/__init__.py index 5d4e4575c935e0a888c6e5e4d0db640d93e1bd49..933200e60749e62094040672793953c1c79de6cf 100644 --- a/tensorflow/contrib/quantize/__init__.py +++ b/tensorflow/contrib/quantize/__init__.py @@ -27,6 +27,8 @@ from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ "create_eval_graph", "create_training_graph", + "experimental_create_eval_graph", + "experimental_create_training_graph", ] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/quantize/python/common.py b/tensorflow/contrib/quantize/python/common.py index d0b0674c31239ee903f5ab7ef9ae0262bb20d189..3a1fa61e43986af1a1315d5a9e6f010e802ea157 100644 --- a/tensorflow/contrib/quantize/python/common.py +++ b/tensorflow/contrib/quantize/python/common.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Constants used across this package.""" +"""Common utilities used across this package.""" from __future__ import absolute_import from __future__ import division @@ -21,6 +21,12 @@ from __future__ import print_function import collections import re +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.ops import init_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import variable_scope + # Skip all operations that are backprop related or export summaries. SKIPPED_PREFIXES = ( 'gradients/', 'RMSProp/', 'Adagrad/', 'Const_', 'HistogramSummary', @@ -86,3 +92,31 @@ def _GetOperationByNameDontThrow(graph, name): return graph.get_operation_by_name(name) except KeyError: return None + + +def CreateOrGetQuantizationStep(): + """Returns a Tensor of the number of steps the quantized graph has run. + + Returns: + Quantization step Tensor. + """ + quantization_step_name = 'fake_quantization_step' + quantization_step_tensor_name = quantization_step_name + '/AssignAdd:0' + g = ops.get_default_graph() + try: + return g.get_tensor_by_name(quantization_step_tensor_name) + except KeyError: + # Create in proper graph and base name_scope. + with g.name_scope(None): + quantization_step_tensor = variable_scope.get_variable( + quantization_step_name, + shape=[], + dtype=dtypes.int64, + initializer=init_ops.zeros_initializer(), + trainable=False, + collections=[ops.GraphKeys.GLOBAL_VARIABLES]) + with g.name_scope(quantization_step_tensor.op.name + '/'): + # We return the incremented variable tensor. Since this is used in conds + # for quant_delay and freeze_bn_delay, it will run once per graph + # execution. + return state_ops.assign_add(quantization_step_tensor, 1) diff --git a/tensorflow/contrib/quantize/python/common_test.py b/tensorflow/contrib/quantize/python/common_test.py new file mode 100644 index 0000000000000000000000000000000000000000..d6237fe5e38d905bf262d7be3746b9ee6046da47 --- /dev/null +++ b/tensorflow/contrib/quantize/python/common_test.py @@ -0,0 +1,59 @@ +# 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 common utilities in this package.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.quantize.python import common +from tensorflow.python.client import session +from tensorflow.python.framework import ops +from tensorflow.python.framework import test_util +from tensorflow.python.ops import variables +from tensorflow.python.platform import googletest + + +class CommonTest(test_util.TensorFlowTestCase): + + def testCreateOrGetQuantizationStep(self): + g = ops.Graph() + with session.Session(graph=g) as sess: + quantization_step_tensor = common.CreateOrGetQuantizationStep() + + # Check that operations are added to the graph. + num_nodes = len(g.get_operations()) + self.assertGreater(num_nodes, 0) + + # Check that getting the quantization step doesn't change the graph. + get_quantization_step_tensor = common.CreateOrGetQuantizationStep() + self.assertEqual(quantization_step_tensor, get_quantization_step_tensor) + self.assertEqual(num_nodes, len(g.get_operations())) + + # Ensure that running the graph increments the quantization step. + sess.run(variables.global_variables_initializer()) + step_val = sess.run(quantization_step_tensor) + self.assertEqual(step_val, 1) + + # Ensure that even running a graph that depends on the quantization step + # multiple times only executes it once. + a = quantization_step_tensor + 1 + b = a + quantization_step_tensor + _, step_val = sess.run([b, quantization_step_tensor]) + self.assertEqual(step_val, 2) + + +if __name__ == '__main__': + googletest.main() diff --git a/tensorflow/contrib/quantize/python/copy_graph_test.py b/tensorflow/contrib/quantize/python/copy_graph_test.py deleted file mode 100644 index 7ff9ad9f8412d7076bf12d6cf10772244444013f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/quantize/python/copy_graph_test.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright 2017 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 copy_graph.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.quantize.python import copy_graph -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util -from tensorflow.python.ops import variables -from tensorflow.python.platform import googletest - - -class CopyGraphTest(test_util.TensorFlowTestCase): - - def _CompareNodeInGraph(self, node, graph): - graph_node = graph.get_operation_by_name(node.name) - self.assertEqual(str(node.node_def), str(graph_node.node_def)) - - def testCopyGraph(self): - graph = ops.Graph() - with graph.as_default(): - a = constant_op.constant(1.0) - b = variables.Variable(2.0) - c = a + b - graph_copy = copy_graph.CopyGraph(graph) - # Ensure that the three original nodes are in the new graph. - # import_meta_graph also adds a saver node to the graph which we don't care - # about in this specific use case. - for tensor in [a, b, c]: - self._CompareNodeInGraph(tensor.op, graph_copy) - # Test that the graph collections are the same. - for key in graph.get_all_collection_keys(): - self.assertEqual( - len(graph.get_collection(key)), - len(graph_copy.get_collection(key)), 'Collection %s differs.') - - -if __name__ == '__main__': - googletest.main() diff --git a/tensorflow/contrib/quantize/python/fold_batch_norms.py b/tensorflow/contrib/quantize/python/fold_batch_norms.py index aa605e6caadf4d1e69a4a331b1e580797e4fdef8..75d9eb0e58d96e4bb2946684febd250e2e1a6b4a 100644 --- a/tensorflow/contrib/quantize/python/fold_batch_norms.py +++ b/tensorflow/contrib/quantize/python/fold_batch_norms.py @@ -26,6 +26,7 @@ from tensorflow.contrib.quantize.python import input_to_ops from tensorflow.core.framework import attr_value_pb2 from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.layers import utils from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn @@ -33,7 +34,7 @@ from tensorflow.python.ops import nn_ops from tensorflow.python.util import compat -def FoldBatchNorms(graph): +def FoldBatchNorms(graph, is_training, freeze_batch_norm_delay=None): """Finds batch norm layers and folds them into preceding layers. Folding only affects the following layers: Conv2D, fully connected, depthwise @@ -41,15 +42,22 @@ def FoldBatchNorms(graph): Args: graph: Graph to walk and modify. - + is_training: Bool, true if training. + freeze_batch_norm_delay: How many steps to wait before freezing moving mean + and variance and using them for batch normalization. This value is used + only when is_training is True. Raises: ValueError: When batch norm folding fails. """ - _FoldFusedBatchNorms(graph) - _FoldUnfusedBatchNorms(graph) + _FoldFusedBatchNorms( + graph, is_training, freeze_batch_norm_delay=freeze_batch_norm_delay) + _FoldUnfusedBatchNorms( + graph, + is_training=is_training, + freeze_batch_norm_delay=freeze_batch_norm_delay) -def _FoldFusedBatchNorms(graph): +def _FoldFusedBatchNorms(graph, is_training, freeze_batch_norm_delay): """Finds fused batch norm layers and folds them into preceding layers. Folding only affects the following layers: Conv2D, fully connected, depthwise @@ -57,6 +65,9 @@ def _FoldFusedBatchNorms(graph): Args: graph: Graph to walk and modify. + is_training: Bool, true if training. + freeze_batch_norm_delay: How many steps to wait before freezing moving mean + and variance and using them for batch normalization. Raises: ValueError: When batch norm folding fails. @@ -67,8 +78,7 @@ def _FoldFusedBatchNorms(graph): # `bn_op`. The '/' (i.e. `sep`) ensures that we reuse the existing scope # named `scope`. Otherwise, TF creates a unique scope whose name starts with # `scope`. - with graph.as_default(), graph.name_scope(scope + sep), ops.device( - match.bn_op.device): + with graph.as_default(), graph.name_scope(scope + sep): with graph.name_scope(scope + sep + 'BatchNorm_Fold' + sep): # new weights = old weights * gamma / sqrt(variance + epsilon) # new biases = -mean * gamma / sqrt(variance + epsilon) + beta @@ -79,9 +89,18 @@ def _FoldFusedBatchNorms(graph): match.mean_tensor * multiplier_tensor, name='bias') + correction_scale, correction_recip, correction_offset = None, None, None + if is_training: + correction_scale, correction_recip, correction_offset = ( + _ComputeBatchNormCorrections( + context='', + match=match, + freeze_batch_norm_delay=freeze_batch_norm_delay, + fused_batch_norm=True)) # The shape of depthwise weights is different, so we need to reshape the # multiplier_tensor to ensure that the scaled_weight_tensor has the # expected shape. + weights = match.weight_tensor if match.layer_op.type == 'DepthwiseConv2dNative': new_shape = [ match.weight_tensor.get_shape().as_list()[2], @@ -90,15 +109,25 @@ def _FoldFusedBatchNorms(graph): multiplier_tensor = array_ops.reshape( multiplier_tensor, new_shape, name='scale_reshape') - # TODO(suharshs): This naming of the following ops needs to carefully - # follow the naming expected by quantize.py. Generalize the quantize code - # to not require these delicate naming conventions. - scaled_weight_tensor = math_ops.multiply( - match.weight_tensor, multiplier_tensor, name='mul_fold') + if correction_scale is not None: + correction_scale = array_ops.reshape( + correction_scale, new_shape, name='correction_reshape') + + if correction_scale is not None: + weights = math_ops.multiply( + correction_scale, weights, name='correction_mult') + scaled_weight_tensor = math_ops.multiply( + weights, multiplier_tensor, name='mul_fold') new_layer_tensor = _CloneWithNewOperands( match.layer_op, match.input_tensor, scaled_weight_tensor) + if correction_recip is not None: + new_layer_tensor = math_ops.multiply( + correction_recip, new_layer_tensor, name='post_conv_mul') + new_layer_tensor = math_ops.add(new_layer_tensor, (correction_offset), + 'correction_add') + bias_add_tensor = math_ops.add( new_layer_tensor, bias_tensor, name='add_fold') @@ -109,46 +138,6 @@ def _FoldFusedBatchNorms(graph): 'Unexpected inputs to op: %s' % match.output_tensor.name) -def _CloneWithNewOperands(layer_op, input_tensor, weight_tensor): - """Clones layer_op with input_tensor and weight_tensor as new inputs.""" - new_layer_name = layer_op.name.split('/')[-1] + '_Fold' - if layer_op.type == 'Conv2D': - return nn_ops.conv2d( - input_tensor, - weight_tensor, - strides=layer_op.get_attr('strides'), - padding=layer_op.get_attr('padding'), - use_cudnn_on_gpu=layer_op.get_attr('use_cudnn_on_gpu'), - data_format=layer_op.get_attr('data_format'), - name=new_layer_name) - elif layer_op.type == 'MatMul': - return math_ops.matmul( - input_tensor, - weight_tensor, - transpose_a=layer_op.get_attr('transpose_a'), - transpose_b=layer_op.get_attr('transpose_b'), - name=new_layer_name) - elif layer_op.type == 'DepthwiseConv2dNative': - return nn.depthwise_conv2d( - input_tensor, - weight_tensor, - strides=layer_op.get_attr('strides'), - padding=layer_op.get_attr('padding'), - name=new_layer_name) - else: - raise ValueError('Cannot handle operation of type: %s' % layer_op.type) - - -@ops.RegisterGradient('FoldFusedBatchNormGrad') -def _FoldFusedBatchNormGrad(op, unused_grad_y, grad_mean, grad_var, unused_1, - unused_2): - x = op.inputs[0] - n = x.get_shape().num_elements() / grad_mean.get_shape().num_elements() - dmean_dx = grad_mean / n - dvar_dx = 2 * grad_var * (x - op.outputs[1]) / (n - 1) - return (dmean_dx + dvar_dx), None, None, None, None - - def _FindFusedBatchNorms(graph): """Finds all ops and tensors related to found FusedBatchNorms. @@ -165,37 +154,59 @@ def _FindFusedBatchNorms(graph): mean_pattern = graph_matcher.OpTypePattern('*') variance_pattern = graph_matcher.OpTypePattern('*') - conv_pattern = graph_matcher.OpTypePattern( - 'Conv2D|DepthwiseConv2dNative', inputs=[input_pattern, weight_pattern]) + moving_average_pattern = graph_matcher.OpTypePattern('*') + bn_decay_pattern = graph_matcher.OpTypePattern('*') + layer_pattern = graph_matcher.OpTypePattern( + 'Conv2D|DepthwiseConv2dNative|MatMul', + inputs=[input_pattern, weight_pattern]) # MatMul has a Reshape between it and FusedBatchNorm. - matmul_pattern = graph_matcher.OpTypePattern( - 'MatMul', inputs=[input_pattern, weight_pattern]) matmul_reshape_pattern = graph_matcher.OpTypePattern( - 'Reshape', inputs=[matmul_pattern, + 'Reshape', inputs=[layer_pattern, graph_matcher.OpTypePattern('*')]) - conv_batch_norm_pattern = graph_matcher.OpTypePattern( + batch_norm_pattern = graph_matcher.OpTypePattern( 'FusedBatchNorm', inputs=[ - conv_pattern, gamma_pattern, beta_pattern, mean_pattern, - variance_pattern - ]) - matmul_batch_norm_pattern = graph_matcher.OpTypePattern( - 'FusedBatchNorm', - inputs=[ - matmul_reshape_pattern, gamma_pattern, beta_pattern, mean_pattern, - variance_pattern + graph_matcher.OneofPattern([matmul_reshape_pattern, layer_pattern]), + gamma_pattern, beta_pattern, mean_pattern, variance_pattern ]) matmul_bn_output_reshape_pattern = graph_matcher.OpTypePattern( - 'Reshape', - inputs=[matmul_batch_norm_pattern, - graph_matcher.OpTypePattern('*')]) + 'Reshape', inputs=[batch_norm_pattern, + graph_matcher.OpTypePattern('*')]) + + bn_matcher = graph_matcher.GraphMatcher( + graph_matcher.OneofPattern( + [matmul_bn_output_reshape_pattern, batch_norm_pattern])) - conv_matcher = graph_matcher.GraphMatcher(conv_batch_norm_pattern) - matmul_matcher = graph_matcher.GraphMatcher(matmul_bn_output_reshape_pattern) + moving_average_sub_pattern = graph_matcher.OpTypePattern( + 'Sub', inputs=[moving_average_pattern, batch_norm_pattern]) + moving_average_mul_pattern = graph_matcher.OpTypePattern( + 'Mul', inputs=[moving_average_sub_pattern, bn_decay_pattern]) + + moving_avg_mul_matcher = graph_matcher.GraphMatcher( + moving_average_mul_pattern) + + for match_result in bn_matcher.match_graph(graph): + moving_mean_tensor = None + moving_variance_tensor = None + bn_decay_mean_tensor = None + bn_decay_var_tensor = None + layer_op = match_result.get_op(layer_pattern) + layer_tensor = match_result.get_tensor(layer_pattern) + bn_op = match_result.get_op(batch_norm_pattern) + batch_epsilon_tensor = bn_op.get_attr('epsilon') + + # In the MatMul case, the output of batch norm is reshaped back into a + # 2D tensor, so the output_tensor is the output of the Reshape op. + output_tensor = bn_op.outputs[0] + if layer_op.type == 'MatMul': + output_reshape_op = match_result.get_op(matmul_bn_output_reshape_pattern) + # If the matcher didn't match matmul_bn_output_reshape, there will be + # another match for this 'MatMul' later, so we can skip this one. + if output_reshape_op is None: + continue + output_tensor = output_reshape_op.outputs[0] - def _GetCommonTensors(match_result, bn_op, bn_input_tensor): - """Gets tensors needed for FusedBatchNormMatch from match_result.""" input_tensor = match_result.get_tensor(input_pattern) weight_tensor = match_result.get_tensor(weight_pattern) gamma_tensor = match_result.get_tensor(gamma_pattern) @@ -222,48 +233,30 @@ def _FindFusedBatchNorms(graph): # calculation, the variance is corrected by the term N/N-1 (Bessel's # correction). The variance tensor read from FuseBatchNorm has bessel's # correction applied, so we undo it here. - n = math_ops.cast( - array_ops.size(bn_input_tensor) / array_ops.size(mean_tensor), - dtypes.float32) - variance_tensor = bn_op.outputs[2] * (n - 1) / n + scope, sep, _ = bn_op.name.rpartition('/') + g = ops.get_default_graph() + with g.as_default(), g.name_scope(scope + sep): + n = math_ops.cast( + array_ops.size(layer_tensor) / array_ops.size(mean_tensor), + dtypes.float32) + variance_tensor = math_ops.multiply( + bn_op.outputs[2], (n - 1) / n, name='Undo_Bessel_Correction') + # TODO(suharshs): Find a way to get rid of this inner match. + for mul_match_result in moving_avg_mul_matcher.match_graph(graph): + sub_op = mul_match_result.get_op(moving_average_sub_pattern) + if sub_op.inputs[1].name == bn_op.outputs[1].name: + # During training: Batch Mean is bn_op.outputs[1] + moving_mean_tensor = sub_op.inputs[0] + bn_decay_mean_tensor = mul_match_result.get_tensor(bn_decay_pattern) + if sub_op.inputs[1].name == bn_op.outputs[2].name: + # During training: Batch Var is bn_op.outputs[2] + moving_variance_tensor = sub_op.inputs[0] + bn_decay_var_tensor = mul_match_result.get_tensor(bn_decay_pattern) else: mean_tensor = match_result.get_tensor(mean_pattern) variance_tensor = match_result.get_tensor(variance_pattern) - return (input_tensor, weight_tensor, gamma_tensor, beta_tensor, mean_tensor, - variance_tensor) - - for match_result in conv_matcher.match_graph(graph): - layer_op = match_result.get_op(conv_pattern) - layer_tensor = match_result.get_tensor(conv_pattern) - bn_op = match_result.get_op(conv_batch_norm_pattern) - # In the case of convolution the output_tensor is the output of bn_op. - output_tensor = bn_op.outputs[0] - - (input_tensor, weight_tensor, gamma_tensor, beta_tensor, mean_tensor, - variance_tensor) = _GetCommonTensors(match_result, bn_op, layer_tensor) - yield _FusedBatchNormMatch( - layer_op=layer_op, - bn_op=bn_op, - output_tensor=output_tensor, - input_tensor=input_tensor, - weight_tensor=weight_tensor, - gamma_tensor=gamma_tensor, - beta_tensor=beta_tensor, - mean_tensor=mean_tensor, - variance_tensor=variance_tensor) - for match_result in matmul_matcher.match_graph(graph): - layer_op = match_result.get_op(matmul_pattern) - layer_tensor = match_result.get_tensor(matmul_pattern) - bn_op = match_result.get_op(matmul_batch_norm_pattern) - # In the MatMul case, the output of batch norm is reshaped back into a - # 2D tensor, so the output_tensor is the output of the Reshape op. - output_reshape_op = match_result.get_op(matmul_bn_output_reshape_pattern) - output_tensor = output_reshape_op.outputs[0] - - (input_tensor, weight_tensor, gamma_tensor, beta_tensor, mean_tensor, - variance_tensor) = _GetCommonTensors(match_result, bn_op, layer_tensor) - yield _FusedBatchNormMatch( + yield _BatchNormMatch( layer_op=layer_op, bn_op=bn_op, output_tensor=output_tensor, @@ -272,63 +265,156 @@ def _FindFusedBatchNorms(graph): gamma_tensor=gamma_tensor, beta_tensor=beta_tensor, mean_tensor=mean_tensor, - variance_tensor=variance_tensor) - - -class _FusedBatchNormMatch(object): - """Contains all information related to a found FusedBatchNorm.""" - - def __init__(self, layer_op, bn_op, output_tensor, input_tensor, - weight_tensor, gamma_tensor, beta_tensor, mean_tensor, - variance_tensor): - self._layer_op = layer_op - self._bn_op = bn_op - self._output_tensor = output_tensor - self._input_tensor = input_tensor - self._weight_tensor = weight_tensor - self._gamma_tensor = gamma_tensor - self._beta_tensor = beta_tensor - self._mean_tensor = mean_tensor - self._variance_tensor = variance_tensor + variance_tensor=variance_tensor, + moving_mean_tensor=moving_mean_tensor, + moving_variance_tensor=moving_variance_tensor, + bn_decay_mean_tensor=bn_decay_mean_tensor, + bn_decay_var_tensor=bn_decay_var_tensor, + batch_epsilon_tensor=batch_epsilon_tensor) + + +def _ComputeBatchNormCorrections(context, match, freeze_batch_norm_delay, + fused_batch_norm): + """Computes batch norm correction params. + + Before batch normalization is frozen: + We use batch statistics for batch norm. + correction_scale = sigma_b/sigma_mv + correction_recip = 1/correction_scale + correction_offset = 0 + + After batch normalization is frozen: + correction_scale = sigma_b/sigma_mv + correction_recip = 1 + correction_offset = gamma*(mu_b/sigma_b-mu_mv/sigma_mv). + + Batch norm is frozen if global_step > bn_freeze_delay. + The corrections ensure that: + a) The weights are quantized after scaling by gamma/sigma_mv. This enables + smoother training as the scaling on the weights changes slowly, rather than + jump across mini-batches + b) Changing the values of the corrections allows for one to switch between + using batch statistics to using moving mean and average, without requiring + changes to batch_norm - @property - def layer_op(self): - return self._layer_op - - @property - def bn_op(self): - return self._bn_op - @property - def output_tensor(self): - return self._output_tensor + Args: + context: The scope under which we look for batch norm params + match: Object containg required batch norm tensors for correction + computation. + freeze_batch_norm_delay: Delay in steps at which computation switches + from regular batch norm to frozen mean and variance. + fused_batch_norm: Bool, true if fused batch norm is used. - @property - def input_tensor(self): - return self._input_tensor + Returns: + A tuple of correction_scale, correction_recip, correction_offset + """ - @property - def weight_tensor(self): - return self._weight_tensor + g = ops.get_default_graph() + with g.name_scope(context + '/batch_norm_correction'): + recip_sigma_mv = math_ops.rsqrt( + match.moving_variance_tensor + match.batch_epsilon_tensor) + recip_sigma = math_ops.rsqrt( + match.variance_tensor + match.batch_epsilon_tensor) + correction_scale = math_ops.divide( + recip_sigma_mv, recip_sigma, name='scale_compute') + correction_scale = array_ops.identity( + correction_scale, name='correction_scale') + correction_recip = math_ops.reciprocal( + correction_scale, name='reciprocal_compute') + correction_offset = math_ops.multiply( + match.gamma_tensor, + match.mean_tensor * recip_sigma - + match.moving_mean_tensor * recip_sigma_mv, + name='offset_compute') + + if freeze_batch_norm_delay is not None: + use_mv_avg = math_ops.greater_equal( + common.CreateOrGetQuantizationStep(), + freeze_batch_norm_delay, + name='use_moving_average') + else: + use_mv_avg = False + + bn_decay_zero = 0.0 + bn_decay_mean_consumers = list(match.bn_decay_mean_tensor.consumers()) + bn_decay_var_consumers = list(match.bn_decay_mean_tensor.consumers()) + + bn_decay_mean_out = utils.smart_cond( + use_mv_avg, + lambda: bn_decay_zero, + lambda: match.bn_decay_mean_tensor, + name='freeze_moving_mean') + graph_editor.reroute_ts( + [bn_decay_mean_out], [match.bn_decay_mean_tensor], + can_modify=bn_decay_mean_consumers) + + if fused_batch_norm is False: + bn_decay_var_consumers = list(match.bn_decay_var_tensor.consumers()) + bn_decay_var_out = utils.smart_cond( + use_mv_avg, + lambda: bn_decay_zero, + lambda: match.bn_decay_var_tensor, + name='freeze_moving_var') + graph_editor.reroute_ts( + [bn_decay_var_out], [match.bn_decay_var_tensor], + can_modify=bn_decay_var_consumers) + + correction_recip = utils.smart_cond( + use_mv_avg, + lambda: array_ops.ones(correction_scale.shape), + lambda: correction_recip, + name='correction_recip') + + correction_offset = utils.smart_cond( + use_mv_avg, + lambda: correction_offset, + lambda: array_ops.zeros(correction_offset.shape), + name='correction_offset') + return correction_scale, correction_recip, correction_offset - @property - def gamma_tensor(self): - return self._gamma_tensor - @property - def beta_tensor(self): - return self._beta_tensor +def _CloneWithNewOperands(layer_op, input_tensor, weight_tensor): + """Clones layer_op with input_tensor and weight_tensor as new inputs.""" + new_layer_name = layer_op.name.split('/')[-1] + '_Fold' + if layer_op.type == 'Conv2D': + return nn_ops.conv2d( + input_tensor, + weight_tensor, + strides=layer_op.get_attr('strides'), + padding=layer_op.get_attr('padding'), + use_cudnn_on_gpu=layer_op.get_attr('use_cudnn_on_gpu'), + data_format=layer_op.get_attr('data_format'), + name=new_layer_name) + elif layer_op.type == 'MatMul': + return math_ops.matmul( + input_tensor, + weight_tensor, + transpose_a=layer_op.get_attr('transpose_a'), + transpose_b=layer_op.get_attr('transpose_b'), + name=new_layer_name) + elif layer_op.type == 'DepthwiseConv2dNative': + return nn.depthwise_conv2d( + input_tensor, + weight_tensor, + strides=layer_op.get_attr('strides'), + padding=layer_op.get_attr('padding'), + name=new_layer_name) + else: + raise ValueError('Cannot handle operation of type: %s' % layer_op.type) - @property - def mean_tensor(self): - return self._mean_tensor - @property - def variance_tensor(self): - return self._variance_tensor +@ops.RegisterGradient('FoldFusedBatchNormGrad') +def _FoldFusedBatchNormGrad(op, unused_grad_y, grad_mean, grad_var, unused_1, + unused_2): + x = op.inputs[0] + n = x.get_shape().num_elements() / grad_mean.get_shape().num_elements() + dmean_dx = grad_mean / n + dvar_dx = 2 * grad_var * (x - op.outputs[1]) / (n - 1) + return (dmean_dx + dvar_dx), None, None, None, None -def _FoldUnfusedBatchNorms(graph): +def _FoldUnfusedBatchNorms(graph, is_training, freeze_batch_norm_delay): """Finds unfused batch norm layers and folds them into preceding layers. Folding only affects the following layers: Conv2D, fully connected, depthwise @@ -336,6 +422,9 @@ def _FoldUnfusedBatchNorms(graph): Args: graph: Graph to walk and modify. + is_training: Bool, True if training. + freeze_batch_norm_delay: How many steps to wait before freezing moving mean + and variance and using them for batch normalization. Raises: ValueError: When batch norm folding fails. @@ -346,7 +435,12 @@ def _FoldUnfusedBatchNorms(graph): has_scaling = _HasScaling(graph, input_to_ops_map, bn) # The mangling code intimately depends on BatchNorm node's internals. - original_op, folded_op = _CreateFoldedOp(graph, bn, has_scaling=has_scaling) + original_op, folded_op = _CreateFoldedOp( + graph, + bn, + has_scaling=has_scaling, + freeze_batch_norm_delay=freeze_batch_norm_delay, + is_training=is_training) activation = common.GetEndpointActivationOp(graph, bn) if activation: @@ -368,46 +462,84 @@ def _FoldUnfusedBatchNorms(graph): raise ValueError('Unexpected inputs to op: %s' % add_bypass.name) -def _HasScaling(graph, input_to_ops_map, bn): - r"""Checks if batch norm has scaling enabled. - - Difference between batch norm with scaling and without is that with scaling: - - Rsqrt -> mul -> mul_1 - \-> mul_2 - - where - mul multiplies gamma by inverse square root of EMA of batch variance, - mul_1 multiplies output of mul with output from the base operation - (convolution, FC or depthwise convolution), - mul_2 multiplies output of mul with EMA of batch mean, - and without scaling: - - Rsqrt -> mul - \-> mul_1 - - where - mul multiplies the inverse square root of EMA of batch variance with output - from the base operation, - mul_1 multiplies inverse square root of EMA of batch variance with EMA - of batch mean. +def _GetBatchNormParams(graph, context, has_scaling): + """Extracts relevant tensors for folding batch norms. Args: graph: Graph to inspect. - input_to_ops_map: InputToOps object containing mapping from tensor's name - to ops that take it as input. - bn: Batch norm layer prefix string. + context: The scope under which we look for batch norm params + has_scaling: Bool that specifies if scaling is done as part of batch norm. Returns: - A boolean indicating whether this batch norm layer has scaling enabled. + _BatchNormMatch containing all required batch norm parameters. """ - rsqrt_op = graph.get_operation_by_name(bn + '/BatchNorm/batchnorm/Rsqrt') - rsqrt_consumers = input_to_ops_map.ConsumerOperations(rsqrt_op) - - return sum(1 for op in rsqrt_consumers if op.type == 'Mul') == 1 - - -def _CreateFoldedOp(graph, context, has_scaling): + gamma_tensor = None + batch_mean_tensor = None + batch_variance_tensor = None + moving_mean_tensor = None + moving_variance_tensor = None + batch_epsilon_tensor = None + bn_decay_mean_tensor = None + bn_decay_var_tensor = None + + split_context = context.split('/') + base_context = split_context[-1] + + oplist = graph.get_operations() + op_suffix_gamma = base_context + '/BatchNorm/gamma' + op_suffix_mean = base_context + '/BatchNorm/moments/Squeeze' + op_suffix_variance = base_context + '/BatchNorm/moments/Squeeze_1' + op_suffix_moving_variance = base_context + '/BatchNorm/moving_variance/read' + op_suffix_moving_mean = base_context + '/BatchNorm/moving_mean/read' + op_suffix_epsilon = base_context + '/BatchNorm/batchnorm/add/y' + op_suffix_bn_decay_mean = base_context + '/BatchNorm/AssignMovingAvg/decay' + op_suffix_bn_decay_var = base_context + '/BatchNorm/AssignMovingAvg_1/decay' + + # Parse through list of ops to find relevant ops + for op in oplist: + if op.name.endswith(op_suffix_mean): + # This is an efficient way to check for two things: + # Is batch norm present and is it training mode? + # Batch statistics are computed only during batch norm in training + batch_mean_tensor = graph.get_tensor_by_name(op.name + ':0') + if op.name.endswith(op_suffix_variance): + batch_variance_tensor = graph.get_tensor_by_name(op.name + ':0') + if op.name.endswith(op_suffix_moving_mean): + moving_mean_tensor = graph.get_tensor_by_name(op.name + ':0') + if op.name.endswith(op_suffix_moving_variance): + moving_variance_tensor = graph.get_tensor_by_name(op.name + ':0') + if op.name.endswith(op_suffix_epsilon): + batch_epsilon_tensor = graph.get_tensor_by_name(op.name + ':0') + if op.name.endswith(op_suffix_bn_decay_mean): + bn_decay_mean_tensor = graph.get_tensor_by_name(op.name + ':0') + if op.name.endswith(op_suffix_bn_decay_var): + bn_decay_var_tensor = graph.get_tensor_by_name(op.name + ':0') + if has_scaling: + if op.name.endswith(op_suffix_gamma): + gamma_tensor = graph.get_tensor_by_name(op.name + ':0') + + if not has_scaling: + gamma_tensor = array_ops.ones(batch_mean_tensor.shape) + + return _BatchNormMatch( + layer_op=None, + bn_op=None, + output_tensor=None, + input_tensor=None, + weight_tensor=None, + gamma_tensor=gamma_tensor, + beta_tensor=None, + mean_tensor=batch_mean_tensor, + variance_tensor=batch_variance_tensor, + moving_mean_tensor=moving_mean_tensor, + moving_variance_tensor=moving_variance_tensor, + bn_decay_mean_tensor=bn_decay_mean_tensor, + bn_decay_var_tensor=bn_decay_var_tensor, + batch_epsilon_tensor=batch_epsilon_tensor) + + +def _CreateFoldedOp(graph, context, has_scaling, freeze_batch_norm_delay, + is_training): """Folds in batch norm layer into preceding convolution or FC layer. Creates 3 new nodes, connects their inputs and adds them to the graph: @@ -417,17 +549,20 @@ def _CreateFoldedOp(graph, context, has_scaling): Args: graph: Graph to modify. context: String, batch norm context, i.e. node into which BatchNorm is - nested. + nested. has_scaling: Whether the batch norm has scaling enabled. + freeze_batch_norm_delay: How many steps to wait before freezing moving mean + and variance and using them for batch normalization. + is_training: Bool, true if training. Raises: ValueError: When operation type is not supported, or input and output tensor - shapes mismatch for created operations: mul_fold, add_fold. + shapes mismatch for created operations: mul_fold, add_fold. Returns: A pair of Operations, the first is the original consumer node of the batch - norm (../BatchNorm/batchnorm/add_1), the second is the consumer node of - the folded graph (add_fold). + norm (../BatchNorm/batchnorm/add_1), the second is the consumer node of + the folded graph (add_fold). """ mul_scale_name = 'mul_1' if has_scaling else 'mul' mul_scale = graph.get_operation_by_name(context + @@ -435,19 +570,43 @@ def _CreateFoldedOp(graph, context, has_scaling): mul_scale_name) op_below = mul_scale.inputs[0].op weights = op_below.inputs[1] - + match = _GetBatchNormParams( + graph=graph, context=context, has_scaling=has_scaling) + correction_scale, correction_recip, correction_offset = None, None, None + if is_training: + correction_scale, correction_recip, correction_offset = ( + _ComputeBatchNormCorrections( + context=context, + match=match, + freeze_batch_norm_delay=freeze_batch_norm_delay, + fused_batch_norm=False)) # Special handling for weights of depthwise convolution. if op_below.type == 'DepthwiseConv2dNative': - new_shape = [weights.get_shape().as_list()[2], - weights.get_shape().as_list()[3]] + new_shape = [ + weights.get_shape().as_list()[2], + weights.get_shape().as_list()[3] + ] scale_name = 'mul' if has_scaling else 'Rsqrt' - scale = graph.get_operation_by_name(context + '/BatchNorm/batchnorm/' + - scale_name) + scale = graph.get_operation_by_name( + context + '/BatchNorm/batchnorm/' + scale_name) scale = array_ops.reshape(scale.outputs[0], new_shape, context + '/scale_reshape') - mul_fold = _CloneOp(mul_scale, context + '/mul_fold', - [(0, weights), (1, scale)]) + + if correction_scale is not None: + correction_scale = array_ops.reshape(correction_scale, new_shape, + context + '/correction_reshape') + with ops.device(mul_scale.device): + weights = math_ops.multiply(correction_scale, weights, + context + '/correction_mult') + + mul_fold = _CloneOp(mul_scale, context + '/mul_fold', [(0, weights), + (1, scale)]) elif op_below.type in ['Conv2D', 'MatMul']: + + if correction_scale is not None: + with ops.device(mul_scale.device): + weights = math_ops.multiply(correction_scale, weights, + context + '/correction_mult') mul_fold = _CloneOp(mul_scale, context + '/mul_fold', [(0, weights)]) else: raise ValueError('Cannot handle operation of type: %s' % op_below.op) @@ -456,10 +615,17 @@ def _CreateFoldedOp(graph, context, has_scaling): conv_or_fc_folded = _CloneOp(op_below, op_below.name + '_Fold', [(1, mul_fold.outputs[0])]) - add_shift = graph.get_operation_by_name(context + - '/BatchNorm/batchnorm/add_1') - add_fold = _CloneOp(add_shift, context + '/add_fold', - [(0, conv_or_fc_folded.outputs[0])]) + add_shift = graph.get_operation_by_name( + context + '/BatchNorm/batchnorm/add_1') + + corrected_output = conv_or_fc_folded.outputs[0] + if correction_offset is not None: + with ops.device(conv_or_fc_folded.device): + corrected_output = math_ops.multiply(correction_recip, corrected_output, + context + '/post_conv_mul') + corrected_output = math_ops.add(corrected_output, (correction_offset), + context + '/correction_add') + add_fold = _CloneOp(add_shift, context + '/add_fold', [(0, corrected_output)]) _AssertShapesMatch('add_fold', add_fold.inputs[0], add_fold.outputs[0]) return add_shift, add_fold @@ -471,7 +637,7 @@ def _CloneOp(op, new_name, new_inputs): op: Operation to modify. new_name: String, a new name to set on cloned op. new_inputs: A list of tuples (idx, tensor), each input with corresponding - index will be replaced by the given Tensor in the cloned op. + index will be replaced by the given Tensor in the cloned op. Returns: Operation, the cloned op. @@ -603,3 +769,121 @@ def _AssertShapesMatch(op_name, in_tensor, out_tensor): if not in_shape.is_compatible_with(out_shape): raise ValueError('%s should not change tensor shape: input %s, ' 'output %s' % (op_name, in_shape, out_shape)) + + +def _HasScaling(graph, input_to_ops_map, bn): + r"""Checks if batch norm has scaling enabled. + + Difference between batch norm with scaling and without is that with scaling: + + Rsqrt -> mul -> mul_1 + \-> mul_2 + + where + mul multiplies gamma by inverse square root of EMA of batch variance, + mul_1 multiplies output of mul with output from the base operation + (convolution, FC or depthwise convolution), + mul_2 multiplies output of mul with EMA of batch mean, + and without scaling: + + Rsqrt -> mul + \-> mul_1 + + where + mul multiplies the inverse square root of EMA of batch variance with output + from the base operation, + mul_1 multiplies inverse square root of EMA of batch variance with EMA + of batch mean. + + Args: + graph: Graph to inspect. + input_to_ops_map: InputToOps object containing mapping from tensor's name + to ops that take it as input. + bn: Batch norm layer prefix string. + + Returns: + A boolean indicating whether this batch norm layer has scaling enabled. + """ + rsqrt_op = graph.get_operation_by_name(bn + '/BatchNorm/batchnorm/Rsqrt') + rsqrt_consumers = input_to_ops_map.ConsumerOperations(rsqrt_op) + + return sum(1 for op in rsqrt_consumers if op.type == 'Mul') == 1 + + +class _BatchNormMatch(object): + """Contains all information related to a found Fused/UnfusedBatchNorm.""" + + def __init__(self, layer_op, bn_op, output_tensor, input_tensor, + weight_tensor, gamma_tensor, beta_tensor, mean_tensor, + variance_tensor, moving_mean_tensor, moving_variance_tensor, + bn_decay_mean_tensor, bn_decay_var_tensor, batch_epsilon_tensor): + self._layer_op = layer_op + self._bn_op = bn_op + self._output_tensor = output_tensor + self._input_tensor = input_tensor + self._weight_tensor = weight_tensor + self._gamma_tensor = gamma_tensor + self._beta_tensor = beta_tensor + self._mean_tensor = mean_tensor + self._variance_tensor = variance_tensor + self._moving_mean_tensor = moving_mean_tensor + self._moving_variance_tensor = moving_variance_tensor + self._bn_decay_mean_tensor = bn_decay_mean_tensor + self._bn_decay_var_tensor = bn_decay_var_tensor + self._batch_epsilon_tensor = batch_epsilon_tensor + + @property + def layer_op(self): + return self._layer_op + + @property + def bn_op(self): + return self._bn_op + + @property + def output_tensor(self): + return self._output_tensor + + @property + def input_tensor(self): + return self._input_tensor + + @property + def weight_tensor(self): + return self._weight_tensor + + @property + def gamma_tensor(self): + return self._gamma_tensor + + @property + def beta_tensor(self): + return self._beta_tensor + + @property + def mean_tensor(self): + return self._mean_tensor + + @property + def variance_tensor(self): + return self._variance_tensor + + @property + def moving_mean_tensor(self): + return self._moving_mean_tensor + + @property + def moving_variance_tensor(self): + return self._moving_variance_tensor + + @property + def batch_epsilon_tensor(self): + return self._batch_epsilon_tensor + + @property + def bn_decay_mean_tensor(self): + return self._bn_decay_mean_tensor + + @property + def bn_decay_var_tensor(self): + return self._bn_decay_var_tensor diff --git a/tensorflow/contrib/quantize/python/fold_batch_norms_test.py b/tensorflow/contrib/quantize/python/fold_batch_norms_test.py index ecf321ff573181c7a2e325770a8dde223bf0c021..c90a18ab0357f1bcbc5d8ccd48edf894d7baf5f9 100644 --- a/tensorflow/contrib/quantize/python/fold_batch_norms_test.py +++ b/tensorflow/contrib/quantize/python/fold_batch_norms_test.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib.layers.python.layers import layers -from tensorflow.contrib.quantize.python import copy_graph from tensorflow.contrib.quantize.python import fold_batch_norms from tensorflow.python.client import session from tensorflow.python.framework import dtypes @@ -34,6 +33,7 @@ from tensorflow.python.ops import nn_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import variables from tensorflow.python.platform import googletest +from tensorflow.python.training import saver as saver_lib batch_norm = layers.batch_norm conv2d = layers.conv2d @@ -46,26 +46,27 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): def _RunTestOverParameters(self, test_fn): parameters_list = [ - # (relu, relu_op_name, with_bypass, has_scaling, fused_batch_norm) - (nn_ops.relu6, 'Relu6', False, False, False), - (nn_ops.relu, 'Relu', False, False, False), - (nn_ops.relu6, 'Relu6', True, False, False), - (nn_ops.relu, 'Relu', True, False, False), - (nn_ops.relu6, 'Relu6', False, True, False), - (nn_ops.relu, 'Relu', False, True, False), - (nn_ops.relu6, 'Relu6', True, True, False), - (nn_ops.relu, 'Relu', True, True, False), + # (relu, relu_op_name, with_bypass, has_scaling, fused_batch_norm, + # freeze_batch_norm_delay) + (nn_ops.relu6, 'Relu6', False, False, False, 100), + (nn_ops.relu, 'Relu', False, False, False, None), + (nn_ops.relu6, 'Relu6', True, False, False, 100), + (nn_ops.relu, 'Relu', True, False, False, None), + (nn_ops.relu6, 'Relu6', False, True, False, 100), + (nn_ops.relu, 'Relu', False, True, False, None), + (nn_ops.relu6, 'Relu6', True, True, False, 100), + (nn_ops.relu, 'Relu', True, True, False, None), # Fused batch norm always has scaling enabled. - (nn_ops.relu6, 'Relu6', False, True, True), - (nn_ops.relu, 'Relu', False, True, True), - (nn_ops.relu6, 'Relu6', True, True, True), - (nn_ops.relu, 'Relu', True, True, True), + (nn_ops.relu6, 'Relu6', False, True, True, None), + (nn_ops.relu, 'Relu', False, True, True, 100), + (nn_ops.relu6, 'Relu6', True, True, True, None), + (nn_ops.relu, 'Relu', True, True, True, 100), ] for params in parameters_list: - test_fn(params[0], params[1], params[2], params[3], params[4]) + test_fn(params[0], params[1], params[2], params[3], params[4], params[5]) def _TestFoldConv2d(self, relu, relu_op_name, with_bypass, has_scaling, - fused_batch_norm): + fused_batch_norm, freeze_batch_norm_delay): """Tests folding cases: inputs -> Conv2d with batch norm -> Relu*. Args: @@ -75,6 +76,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): inputs to just before Relu*. has_scaling: Bool, when true the batch norm has scaling. fused_batch_norm: Bool, when true the batch norm is fused. + freeze_batch_norm_delay: None or the number of steps after which training + switches to using frozen mean and variance """ g = ops.Graph() with g.as_default(): @@ -99,12 +102,13 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): node = math_ops.add(inputs, node, name='test/Add') relu(node, name='test/' + relu_op_name) - fold_batch_norms.FoldBatchNorms(g) + fold_batch_norms.FoldBatchNorms( + g, is_training=True, freeze_batch_norm_delay=freeze_batch_norm_delay) folded_mul = g.get_operation_by_name(scope + '/mul_fold') self.assertEqual(folded_mul.type, 'Mul') self._AssertInputOpsAre(folded_mul, [ - scope + '/weights/read', + scope + '/correction_mult', self._BatchNormMultiplierName(scope, has_scaling, fused_batch_norm) ]) self._AssertOutputGoesToOps(folded_mul, g, [scope + '/Conv2D_Fold']) @@ -113,12 +117,12 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): self.assertEqual(folded_conv.type, 'Conv2D') self._AssertInputOpsAre(folded_conv, [scope + '/mul_fold', inputs.op.name]) - self._AssertOutputGoesToOps(folded_conv, g, [scope + '/add_fold']) + self._AssertOutputGoesToOps(folded_conv, g, [scope + '/post_conv_mul']) folded_add = g.get_operation_by_name(scope + '/add_fold') self.assertEqual(folded_add.type, 'Add') self._AssertInputOpsAre(folded_add, [ - scope + '/Conv2D_Fold', + scope + '/correction_add', self._BathNormBiasName(scope, fused_batch_norm) ]) output_op_names = ['test/Add' if with_bypass else 'test/' + relu_op_name] @@ -128,7 +132,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): self._RunTestOverParameters(self._TestFoldConv2d) def _TestFoldConv2dUnknownShape(self, relu, relu_op_name, with_bypass, - has_scaling, fused_batch_norm): + has_scaling, fused_batch_norm, + freeze_batch_norm_delay): """Tests folding cases: inputs -> Conv2d with batch norm -> Relu*. Tests that folding works even with an input shape where some dimensions are @@ -141,6 +146,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): inputs to just before Relu*. has_scaling: Bool, when true the batch norm has scaling. fused_batch_norm: Bool, when true the batch norm is fused. + freeze_batch_norm_delay: None or the number of steps after which training + switches to using frozen mean and variance """ g = ops.Graph() with g.as_default(): @@ -164,12 +171,13 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): node = math_ops.add(inputs, node, name='test/Add') relu(node, name='test/' + relu_op_name) - fold_batch_norms.FoldBatchNorms(g) + fold_batch_norms.FoldBatchNorms( + g, is_training=True, freeze_batch_norm_delay=freeze_batch_norm_delay) folded_mul = g.get_operation_by_name(scope + '/mul_fold') self.assertEqual(folded_mul.type, 'Mul') self._AssertInputOpsAre(folded_mul, [ - scope + '/weights/read', + scope + '/correction_mult', self._BatchNormMultiplierName(scope, has_scaling, fused_batch_norm) ]) self._AssertOutputGoesToOps(folded_mul, g, [scope + '/Conv2D_Fold']) @@ -177,12 +185,12 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): folded_conv = g.get_operation_by_name(scope + '/Conv2D_Fold') self.assertEqual(folded_conv.type, 'Conv2D') self._AssertInputOpsAre(folded_conv, [scope + '/mul_fold', inputs.op.name]) - self._AssertOutputGoesToOps(folded_conv, g, [scope + '/add_fold']) + self._AssertOutputGoesToOps(folded_conv, g, [scope + '/post_conv_mul']) folded_add = g.get_operation_by_name(scope + '/add_fold') self.assertEqual(folded_add.type, 'Add') self._AssertInputOpsAre(folded_add, [ - scope + '/Conv2D_Fold', + scope + '/correction_add', self._BathNormBiasName(scope, fused_batch_norm) ]) output_op_names = ['test/Add' if with_bypass else 'test/' + relu_op_name] @@ -192,7 +200,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): self._RunTestOverParameters(self._TestFoldConv2dUnknownShape) def _TestFoldFullyConnectedLayer(self, relu, relu_op_name, with_bypass, - has_scaling, fused_batch_norm): + has_scaling, fused_batch_norm, + freeze_batch_norm_delay): """Tests folding cases: inputs -> FC with batch norm -> Relu*. Args: @@ -202,6 +211,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): inputs to just before Relu*. has_scaling: Bool, when true the batch norm has scaling. fused_batch_norm: Bool, when true the batch norm is fused. + freeze_batch_norm_delay: None or the number of steps after which training + switches to using frozen mean and variance """ g = ops.Graph() with g.as_default(): @@ -223,12 +234,13 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): node = math_ops.add(inputs, node, name='test/Add') relu(node, name='test/' + relu_op_name) - fold_batch_norms.FoldBatchNorms(g) + fold_batch_norms.FoldBatchNorms( + g, is_training=True, freeze_batch_norm_delay=freeze_batch_norm_delay) folded_mul = g.get_operation_by_name(scope + '/mul_fold') self.assertEqual(folded_mul.type, 'Mul') self._AssertInputOpsAre(folded_mul, [ - scope + '/weights/read', + scope + '/correction_mult', self._BatchNormMultiplierName(scope, has_scaling, fused_batch_norm) ]) self._AssertOutputGoesToOps(folded_mul, g, [scope + '/MatMul_Fold']) @@ -237,12 +249,12 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): self.assertEqual(folded_conv.type, 'MatMul') self._AssertInputOpsAre(folded_conv, [scope + '/mul_fold', inputs.op.name]) - self._AssertOutputGoesToOps(folded_conv, g, [scope + '/add_fold']) + self._AssertOutputGoesToOps(folded_conv, g, [scope + '/post_conv_mul']) folded_add = g.get_operation_by_name(scope + '/add_fold') self.assertEqual(folded_add.type, 'Add') self._AssertInputOpsAre(folded_add, [ - scope + '/MatMul_Fold', + scope + '/correction_add', self._BathNormBiasName(scope, fused_batch_norm) ]) output_op_names = ['test/Add' if with_bypass else 'test/' + relu_op_name] @@ -252,7 +264,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): self._RunTestOverParameters(self._TestFoldFullyConnectedLayer) def _TestFoldDepthwiseConv2d(self, relu, relu_op_name, with_bypass, - has_scaling, fused_batch_norm): + has_scaling, fused_batch_norm, + freeze_batch_norm_delay): """Tests folding: inputs -> DepthwiseConv2d with batch norm -> Relu*. Args: @@ -262,6 +275,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): inputs to just before Relu*. has_scaling: Bool, when true the batch norm has scaling. fused_batch_norm: Bool, when true the batch norm is fused. + freeze_batch_norm_delay: None or the number of steps after which training + switches to using frozen mean and variance """ g = ops.Graph() with g.as_default(): @@ -286,7 +301,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): node = math_ops.add(inputs, node, name='test/Add') relu(node, name='test/' + relu_op_name) - fold_batch_norms.FoldBatchNorms(g) + fold_batch_norms.FoldBatchNorms( + g, is_training=True, freeze_batch_norm_delay=freeze_batch_norm_delay) folded_mul = g.get_operation_by_name(scope + '/mul_fold') self.assertEqual(folded_mul.type, 'Mul') @@ -295,8 +311,7 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): else: scale_reshape_op_name = scope + '/scale_reshape' self._AssertInputOpsAre(folded_mul, - [scope + '/depthwise_weights/read', - scale_reshape_op_name]) + [scope + '/correction_mult', scale_reshape_op_name]) self._AssertOutputGoesToOps(folded_mul, g, [scope + '/depthwise_Fold']) scale_reshape = g.get_operation_by_name(scale_reshape_op_name) @@ -311,12 +326,12 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): self.assertEqual(folded_conv.type, 'DepthwiseConv2dNative') self._AssertInputOpsAre(folded_conv, [scope + '/mul_fold', inputs.op.name]) - self._AssertOutputGoesToOps(folded_conv, g, [scope + '/add_fold']) + self._AssertOutputGoesToOps(folded_conv, g, [scope + '/post_conv_mul']) folded_add = g.get_operation_by_name(scope + '/add_fold') self.assertEqual(folded_add.type, 'Add') self._AssertInputOpsAre(folded_add, [ - scope + '/depthwise_Fold', + scope + '/correction_add', self._BathNormBiasName(scope, fused_batch_norm) ]) output_op_names = ['test/Add' if with_bypass else 'test/' + relu_op_name] @@ -326,7 +341,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): self._RunTestOverParameters(self._TestFoldDepthwiseConv2d) def _TestCompareFoldAndUnfolded(self, relu, relu_op_name, with_bypass, - has_scaling, fused_batch_norm): + has_scaling, fused_batch_norm, + freeze_batch_norm_delay): """Tests that running folded and unfolded BN returns the same results. Args: @@ -336,6 +352,8 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): inputs to just before Relu*. has_scaling: Bool, when true the batch norm has scaling. fused_batch_norm: Bool, when true the batch norm is fused. + freeze_batch_norm_delay: None or the number of steps after which training + switches to using frozen mean and variance """ random_seed.set_random_seed(1234) unfolded_g = ops.Graph() @@ -361,11 +379,12 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): if with_bypass: node = math_ops.add(inputs, node, name='test/Add') relu_node = relu(node, name='test/' + relu_op_name) - - folded_g = copy_graph.CopyGraph(unfolded_g) + folded_g = self._CopyGraph(unfolded_g) with folded_g.as_default(): - fold_batch_norms.FoldBatchNorms(folded_g) - + fold_batch_norms.FoldBatchNorms( + folded_g, + is_training=True, + freeze_batch_norm_delay=freeze_batch_norm_delay) with session.Session(graph=unfolded_g) as sess: sess.run(variables.global_variables_initializer()) grad_node = gradients.gradients(relu_node, inputs) @@ -443,5 +462,15 @@ class FoldBatchNormsTest(test_util.TensorFlowTestCase): out_op = graph.get_operation_by_name(out_op_name) self.assertIn(op.outputs[0].name, [str(t.name) for t in out_op.inputs]) + def _CopyGraph(self, graph): + """Return a copy of graph.""" + meta_graph = saver_lib.export_meta_graph( + graph=graph, collection_list=graph.get_all_collection_keys()) + graph_copy = ops.Graph() + with graph_copy.as_default(): + _ = saver_lib.import_meta_graph(meta_graph) + return graph_copy + + if __name__ == '__main__': googletest.main() diff --git a/tensorflow/contrib/quantize/python/graph_matcher.py b/tensorflow/contrib/quantize/python/graph_matcher.py index e3581cc55905a0af7d0464bc0ec673d3ed7f0363..b458f039df0523b5b8b07cff7d14643154124b95 100644 --- a/tensorflow/contrib/quantize/python/graph_matcher.py +++ b/tensorflow/contrib/quantize/python/graph_matcher.py @@ -18,8 +18,19 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import abc -class OpTypePattern(object): + +class Pattern(object): + """The parent class of all patterns (e.g. OpTypePattern and OneofPattern).""" + + @abc.abstractmethod + def match(self, op, tensor): + """Returns the result of matching op/tensor against this pattern.""" + raise NotImplementedError('Method "match" not implemented.') + + +class OpTypePattern(Pattern): """A tree pattern that matches TF expressions with certain op types.""" def __init__(self, op_type, name=None, inputs=None): @@ -34,7 +45,7 @@ class OpTypePattern(object): similar TF op types. name: Optional string. The name of the pattern that can be looked up in MatchResult. - inputs: Optional list of `OpTypePattern`s or strings that specify the + inputs: Optional list of `Pattern`s or strings that specify the patterns for the inputs of a matching op. If None, this pattern accepts any inputs of a matching op. """ @@ -43,22 +54,51 @@ class OpTypePattern(object): if inputs is None: inputs = [] self._inputs = [ - input_pattern if isinstance(input_pattern, OpTypePattern) else - OpTypePattern(input_pattern) for input_pattern in inputs + input_pattern + if isinstance(input_pattern, Pattern) else OpTypePattern(input_pattern) + for input_pattern in inputs ] - @property - def op_type(self): - return self._op_type - - @property - def inputs(self): - return self._inputs - @property def name(self): return self._name + def match(self, op, tensor): + if self._op_type != '*': + if op.type not in self._op_type.split('|'): + return None + + match_result = MatchResult() + match_result.add(self, op, tensor) + + if not self._inputs: + # If pattern.inputs is empty, skips the rest and accepts all the inputs. + return match_result + + if len(op.inputs) != len(self._inputs): + return None + + for input_tensor, input_pattern in zip(op.inputs, self._inputs): + input_match_result = input_pattern.match(input_tensor.op, input_tensor) + if input_match_result is None: + return None + match_result.merge_from(input_match_result) + return match_result + + +class OneofPattern(Pattern): + """Matches one of the given sub-patterns.""" + + def __init__(self, sub_patterns): + self._sub_patterns = sub_patterns + + def match(self, op, tensor): + for sub_pattern in self._sub_patterns: + match_result = sub_pattern.match(op, tensor) + if match_result is not None: + return match_result + return None + class MatchResult(object): r"""Encapsulates the result of a match done by GraphMatcher. @@ -102,16 +142,36 @@ class MatchResult(object): return pattern_or_name if isinstance(pattern_or_name, str): + if pattern_or_name not in self._name_to_pattern: + return None return self._name_to_pattern[pattern_or_name] raise ValueError('pattern_or_name has type %s. Expect OpTypePattern or str.' % type(pattern_or_name)) + def _get_op_tensor(self, pattern_or_name): + pattern = self._to_pattern(pattern_or_name) + if pattern is None: + return None + + if pattern not in self._pattern_to_op_tensor: + return None + + return self._pattern_to_op_tensor[pattern] + def get_op(self, pattern_or_name): - return self._pattern_to_op_tensor[self._to_pattern(pattern_or_name)][0] + op_tensor = self._get_op_tensor(pattern_or_name) + return op_tensor[0] if op_tensor else None def get_tensor(self, pattern_or_name): - return self._pattern_to_op_tensor[self._to_pattern(pattern_or_name)][1] + op_tensor = self._get_op_tensor(pattern_or_name) + return op_tensor[1] if op_tensor else None + + def merge_from(self, other_match_result): + # pylint: disable=protected-access + self._pattern_to_op_tensor.update(other_match_result._pattern_to_op_tensor) + self._name_to_pattern.update(other_match_result._name_to_pattern) + # pylint: enable=protected-access class GraphMatcher(object): @@ -121,7 +181,7 @@ class GraphMatcher(object): """Initializes a GraphMatcher. Args: - pattern: The `OpTypePattern` against which `GraphMatcher` matches + pattern: The `Pattern` against which `GraphMatcher` matches subgraphs. """ self._pattern = pattern @@ -133,7 +193,7 @@ class GraphMatcher(object): with key `pattern`. Args: - pattern: An `OpTypePattern`. + pattern: An `Pattern`. op: A `tf.Operation` to match against the pattern. tensor: the output `tf.Tensor` of `op` that is used by the matching op of `pattern`'s parent. Can be None if `pattern` is already the root of the @@ -142,20 +202,11 @@ class GraphMatcher(object): Returns: True if an TF expression rooted at `op` matches `pattern`. """ - if pattern.op_type != '*': - if op.type not in pattern.op_type.split('|'): - return False - - self._match_result.add(pattern, op, tensor) - - if not pattern.inputs: - # If pattern.inputs is empty, skips the rest and accepts all the inputs. - return True - - return len(op.inputs) == len(pattern.inputs) and all([ - self._match_pattern(input_pattern, input_tensor.op, input_tensor) - for input_tensor, input_pattern in zip(op.inputs, pattern.inputs) - ]) + match_result = pattern.match(op, tensor) + if match_result is None: + return False + self._match_result.merge_from(match_result) + return True def match_op(self, op): """Matches `op` against `self._pattern`. diff --git a/tensorflow/contrib/quantize/python/graph_matcher_test.py b/tensorflow/contrib/quantize/python/graph_matcher_test.py index e1572865e423e569ee3b280036c0e02b71b70648..6d587572181c125faa02d36fb54933cff24f11c6 100644 --- a/tensorflow/contrib/quantize/python/graph_matcher_test.py +++ b/tensorflow/contrib/quantize/python/graph_matcher_test.py @@ -105,7 +105,7 @@ class GraphMatcherTest(test_util.TensorFlowTestCase): self.assertEqual(match_result.get_op(y1_pattern), y1.op) self.assertEqual(match_result.get_tensor(y1_pattern), y1) - def test_oneof_pattern(self): + def test_oneof_type_pattern(self): # - + # / \ / \ # x y z @@ -125,6 +125,44 @@ class GraphMatcherTest(test_util.TensorFlowTestCase): for match_result in matcher.match_graph(g) ], [plus.op, minus.op]) + def test_oneof_pattern(self): + reshape_pattern = graph_matcher.OpTypePattern('Reshape') + transpose_pattern = graph_matcher.OneofPattern([ + graph_matcher.OpTypePattern( + 'Transpose', + name='transpose', + inputs=[ + graph_matcher.OpTypePattern( + 'Slice', name='slice', inputs=[reshape_pattern, '*', '*']), + '*' + ]), + graph_matcher.OpTypePattern( + 'Transpose', name='transpose', inputs=[reshape_pattern, '*']) + ]) + + matcher = graph_matcher.GraphMatcher(transpose_pattern) + + g = ops.Graph() + with g.as_default(): + inputs = array_ops.placeholder(dtypes.float32, shape=[6]) + reshape = array_ops.reshape(inputs, [2, 3]) + transpose = array_ops.transpose(reshape) + [match_result] = list(matcher.match_graph(g)) + self.assertEqual(match_result.get_tensor(reshape_pattern), reshape) + self.assertEqual(match_result.get_tensor('slice'), None) + self.assertEqual(match_result.get_op('transpose'), transpose.op) + + g = ops.Graph() + with g.as_default(): + inputs = array_ops.placeholder(dtypes.float32, shape=[6]) + reshape = array_ops.reshape(inputs, [2, 3]) + slicing = array_ops.slice(reshape, [0, 0], [-1, -1]) + transpose = array_ops.transpose(slicing) + [match_result] = list(matcher.match_graph(g)) + self.assertEqual(match_result.get_tensor(reshape_pattern), reshape) + self.assertEqual(match_result.get_tensor('slice'), slicing) + self.assertEqual(match_result.get_op('transpose'), transpose.op) + if __name__ == '__main__': googletest.main() diff --git a/tensorflow/contrib/quantize/python/quant_ops.py b/tensorflow/contrib/quantize/python/quant_ops.py index f80d427ff0a6573ecd6562c443182797b5d22527..0a8e35080cb08f71dc28e33c6138a12656e5a5ea 100644 --- a/tensorflow/contrib/quantize/python/quant_ops.py +++ b/tensorflow/contrib/quantize/python/quant_ops.py @@ -53,7 +53,7 @@ def LastValueQuantize(inputs, init_max=6.0, updates_collection=ops.GraphKeys.UPDATE_OPS, vars_collection=ops.GraphKeys.MOVING_AVERAGE_VARIABLES, - scope=None, + name_prefix='LastValueQuant', reuse=None, is_training=True, num_bits=8, @@ -73,7 +73,7 @@ def LastValueQuantize(inputs, computation. vars_collection: (Optional) collection where to store variables for quantization interval ends. - scope: Optional scope for variable_scope. + name_prefix: name_prefix for created nodes. reuse: whether or not the layer and its variables should be reused. To be able to reuse the layer scope must be given. is_training: Whether the op is applied to a training or eval graph. @@ -84,13 +84,13 @@ def LastValueQuantize(inputs, a tensor containing quantized values. """ with variable_scope.variable_scope( - scope, 'LastValueQuantize', values=[inputs], reuse=reuse): + None, default_name=name_prefix, values=[inputs], reuse=reuse): input_shape = inputs.get_shape() input_dim = len(input_shape) if per_channel: # Only support quantizing 1-, 2- and 4-dimensional tensors. assert input_dim in [1, 2, 4], ('Expected 1D, 2D or 4D input, was: %s in ' - ' scope: %s' % (input_shape, scope)) + ' scope: %s' % (input_shape, name_prefix)) min_max_shape = [input_shape[-1]] else: min_max_shape = [] @@ -165,7 +165,7 @@ def MovingAvgQuantize(inputs, ema_decay=0.999, updates_collection=ops.GraphKeys.UPDATE_OPS, vars_collection=ops.GraphKeys.MOVING_AVERAGE_VARIABLES, - scope=None, + name_prefix='MovingAvgQuantize', reuse=None, is_training=True, num_bits=8, @@ -186,7 +186,7 @@ def MovingAvgQuantize(inputs, computation. vars_collection: (Optional) collection where to store variables for quantization interval ends. - scope: Optional scope for variable_scope. + name_prefix: name_prefix for created nodes. reuse: whether or not the layer and its variables should be reused. To be able to reuse the layer scope must be given. is_training: Whether the op is applied to a training or eval graph. @@ -197,13 +197,13 @@ def MovingAvgQuantize(inputs, a tensor containing quantized values. """ with variable_scope.variable_scope( - scope, 'MovingAvgQuantize', values=[inputs], reuse=reuse): + None, default_name=name_prefix, values=[inputs], reuse=reuse): input_shape = inputs.get_shape() input_dim = len(input_shape) if per_channel: # Only support quantizing 1-, 2- and 4-dimensional tensors. assert input_dim in [1, 2, 4], ('Expected 1D, 2D or 4D input, was: %s in ' - ' scope: %s' % (input_shape, scope)) + ' scope: %s' % (input_shape, name_prefix)) min_max_shape = [input_shape[-1]] else: min_max_shape = [] diff --git a/tensorflow/contrib/quantize/python/quantize.py b/tensorflow/contrib/quantize/python/quantize.py index 50a2b4c91c9e7a2681f6041646a023a4225fb0c5..7a3f92f503a5d6f2b0fab2a499f8e8758809d0ed 100644 --- a/tensorflow/contrib/quantize/python/quantize.py +++ b/tensorflow/contrib/quantize/python/quantize.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Logic to update a Tensorflow model graph with quantization operations.""" +"""Logic to update a TensorFlow model graph with quantization operations.""" from __future__ import absolute_import from __future__ import division @@ -21,37 +21,37 @@ from __future__ import print_function import re from tensorflow.contrib import graph_editor from tensorflow.contrib.quantize.python import common +from tensorflow.contrib.quantize.python import graph_matcher from tensorflow.contrib.quantize.python import input_to_ops from tensorflow.contrib.quantize.python import quant_ops from tensorflow.python.framework import ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops -from tensorflow.python.training import training_util -# Operation types used to select operations of interest. +# Quantizable operation types that are supported by the quantization rewrite. _QUANTIZABLE_TYPES = {'Conv2D', 'MatMul', 'DepthwiseConv2dNative'} -# Custom key for storing and retrieving update ops used by quantizing nodes. -_UPDATE_QUANT_OPS = 'update_quant_ops' +# Activations that are supported by the quantization rewrite. +_ACTIVATION_TYPES = {'Relu', 'Relu6', 'Identity'} + +# Weight types that are supported by the quantization rewrite. +# TODO(suharshs): Add support for ResourceVariable. +_WEIGHT_TYPES = {'Variable', 'VariableV2'} def Quantize(graph, + is_training, weight_bits=8, - weight_narrow_range=False, activation_bits=8, ema_decay=0.999, quant_delay=None, - vars_collection=ops.GraphKeys.MOVING_AVERAGE_VARIABLES, - is_training=True, - quantize_folded_weights_use_ema=False): + vars_collection=ops.GraphKeys.MOVING_AVERAGE_VARIABLES): """Updates graph with quantization operations. Args: graph: Graph to modify. + is_training: Whether quantizing training graph or eval graph. weight_bits: Number of bits to use for quantizing weights. - weight_narrow_range: Whether to use a more efficient narrow range for - weights quantization. With weight_narrow_range true, the range is - [1; 2^weight_bits - 1], with it false [0; 2^weight_bits - 1]. activation_bits: Number of bits to use for quantizing activations. ema_decay: (Optional) Float, EMA decay parameter. EMA is used to update quantization intervals for quantizing activations (see here about EMA: @@ -61,346 +61,280 @@ def Quantize(graph, training. vars_collection: (Optional) Collection where to store the variables for quantization interval ends. - is_training: (Optional) Whether quantizing training graph or eval graph. - quantize_folded_weights_use_ema: (Optional, default False) Whether to - quantize weights after batchnorm-folding with exponential average - quantization. Raises: ValueError: When quantization fails. """ - context = _QuantizeContext(graph, weight_bits, weight_narrow_range, - activation_bits, ema_decay, quant_delay, - vars_collection, is_training, - quantize_folded_weights_use_ema) - - graph_ops = graph.get_operations() - - # Filter out backprop and summary related operations, leave only interesting - # op types. - def _IsInterestingOpWithWeights(op): - return (op.type in _QUANTIZABLE_TYPES and - not op.name.startswith(common.SKIPPED_PREFIXES)) - - for op in (op for op in graph_ops if _IsInterestingOpWithWeights(op)): - if op.name.endswith('/depthwise'): - # Separable convolution may consist of 2 convolution nodes. If so, skip - # .../depthwise and only quantize the top one. - separable_conv = context.GetOperationByNameDontThrow( - op.name[:-len('/depthwise')]) - if separable_conv and separable_conv.type == 'Conv2D': - continue - # Quantize add ops that come after Conv2D or DepthwiseConv2dNative. - if op.type in ['Conv2D', 'DepthwiseConv2dNative']: - add_context_re = re.search(r'^(.*)/[^/]+/', op.name) - if add_context_re is not None: - context.add_contexts.add(add_context_re.group(1)) - if not op.name.endswith('_Fold'): - folded_op = context.GetOperationByNameDontThrow(op.name + '_Fold') - # Do nothing if found, it will be quantized when it is iterated over. - if not folded_op: - context.QuantizeOpWithWeights(op, folded=False) - else: - context.QuantizeOpWithWeights(op, folded=True) - - context.QuantizeAddContexts() - - # Once all quantization ops have been inserted in the graph, collect update - # ops for their variables and modify the TF Slim update barrier (see - # https://www.tensorflow.org/code/tensorflow/contrib/slim/python/slim/learning.py) - # to depend on them. - try: - update_barrier = graph.get_operation_by_name('update_barrier') - except KeyError: - # In evaluation graph, this barrier may not exist. - return None - update_quant_ops = graph.get_collection_ref(_UPDATE_QUANT_OPS) - graph_editor.add_control_inputs(update_barrier, update_quant_ops) - - -class _QuantizeContext(object): - """Context holds references needed for quantization.""" - - def __init__(self, - graph, - weight_bits, - weight_narrow_range, - activation_bits, - ema_decay=0.999, - quant_delay=None, - vars_collection=ops.GraphKeys.MOVING_AVERAGE_VARIABLES, - is_training=True, - quantize_folded_weights_use_ema=False): - """Initializes context to hold references needed for quantization. - - Args: - graph: Graph to modify. - weight_bits: Number of bits to use for quantizing weights. - weight_narrow_range: Whether to use a more efficient narrow range for - weights quantization. With weight_narrow_range true, the range is - [1; 2^weight_bits - 1], with it false [0; 2^weight_bits - 1]. - activation_bits: Number of bits to use for quantizing activations. - ema_decay: (Optional) Float, EMA decay parameter. - quant_delay: (Optional, default None) Int, count of global steps for which - to delay quantization. This helps weights stabilize at the start of - training. - vars_collection: (Optional) Collection where to store the variables for - quantization interval ends. - is_training: (Optional) Whether quantizing training or eval graph. - quantize_folded_weights_use_ema: (Optional, default False) Whether to - quantize weights after batchnorm-folding with exponential average - quantization. - """ - self.graph = graph - self.weight_bits = weight_bits - self.weight_narrow_range = weight_narrow_range - self.activation_bits = activation_bits - self.ema_decay = ema_decay - self.quant_delay = quant_delay - self.vars_collection = vars_collection - self.is_training = is_training - self.quantize_folded_weights_use_ema = quantize_folded_weights_use_ema - self.input_to_ops_map = input_to_ops.InputToOps(graph) - self.add_contexts = set() - - def QuantizeAddContexts(self): - """Quantizes all add ops in self.add_contexts.""" - # Loop through sorted self.add_contexts so that op creation is - # deterministic. This is needed when using multiple worker replicas so that - # the ops can be initialized consistently. - for add_context in sorted(self.add_contexts): - add_op = self.GetOperationByNamesDontThrow([ - add_context + '/Add', add_context + '/add']) - if add_op is not None: - self._InsertQuantOp( - add_context, - add_op, - self.input_to_ops_map.ConsumerOperations(add_op), - name='add_quant', - moving_avg=True, - bits=self.activation_bits, - narrow_range=False) - - def QuantizeOpWithWeights(self, op, folded): - """Quantizes around the specific operation with or without batch norm. - - Args: - op: Operation to quantize. - folded: Operation has been folded and needs special handling if True. - Raises: - ValueError: When quantization fails. - """ - # Op name component before the last slash will be used as context. - context = re.search(r'^(.*)/([^/]+)', op.name).group(1) - - # Quantize weights. - if folded: - producer_op = self.graph.get_operation_by_name(context + '/mul_fold') - else: - try: - input_idx = next(i for i, v in enumerate(op.inputs) - if '/weights/' in v.name or - '/depthwise_weights' in v.name) - except StopIteration: - raise ValueError('No inputs to quantize for op: %s' % op) - producer_op = op.inputs[input_idx].op - - # If batch norm is used, the folded weights depend on the batch std, hence - # it is sensible to use EMA during training to smooth out the noise. This is - # controlled by the flag quantize_folded_weights_use_ema. Its default is - # False for backward compatibility. - # If there is no batch norm, weights do not depend on the batch and using - # the latest value of min and max is more efficient. - weight_use_ema = folded and self.quantize_folded_weights_use_ema - self._InsertQuantOp( - context, - producer_op, [op], - name='weights_quant', - moving_avg=weight_use_ema, - delay_requested=weight_use_ema, - bits=self.weight_bits, - narrow_range=self.weight_narrow_range) - - # Important: do not quantize biases here. During inference they are - # quantized to 32 bits, which is much finer than 8 bit quantization and - # depends on weight and input activation ranges. - - # Find activation and (optionally) Add operations to quantize. - activation_op, add_op, add_context = self._GetReluAndAddOperations(context, - op) - if add_op: - original_context = context - context = add_context - - # Quantize activation outputs. - consumer_ops = self.input_to_ops_map.ConsumerOperations(activation_op) - self._InsertQuantOp( + input_to_ops_map = input_to_ops.InputToOps(graph) + for layer_match in _FindLayersToQuantize(graph): + # Quantize the weights. + context = _GetContextFromOp(layer_match.layer_op) + _InsertQuantOp( context, - activation_op, + 'weights_quant', + layer_match.weight_tensor.op, [layer_match.layer_op], + is_training, + moving_avg=False, + ema_decay=ema_decay, + quant_delay=quant_delay, + narrow_range=True, + vars_collection=vars_collection, + bits=weight_bits) + + # Quantize the activations. + consumer_ops = input_to_ops_map.ConsumerOperations( + layer_match.activation_op) + add_context = context + if layer_match.bypass_op: + add_context = re.search(r'^(.*)/([^/]+)', context).group(1) + _InsertQuantOp( + add_context, + 'act_quant', + layer_match.activation_op, consumer_ops, - name='act_quant', + is_training, moving_avg=True, - init_min=0.0, - bits=self.activation_bits, - narrow_range=False) - - # When a bypass connection was found, also quantize Add op input. - if add_op: - def _QuantizeAddInput(add_input): - if folded: - return add_input.op.name.endswith('/add_fold') - else: - return add_input.op.name.startswith(original_context + '/') - - for add_input in add_op.inputs: - if _QuantizeAddInput(add_input): - self._InsertQuantOp( - original_context, - add_input.op, [add_op], - name='conv_quant', - moving_avg=True, - bits=self.activation_bits, - narrow_range=False) - - def _GetReluAndAddOperations(self, context, op): - """Looks up a Relu* and Add operations in given context. - - Args: - context: Context where to look for operations. - op: Operation to quantize. - - Returns: - A triplet (Operation, Operation, string), the first element is an end - point operation, the second is Add operation (optional), the third element - is string context where the Add operation was found (optional). - - Raises: - ValueError: When operations cannot be found. - """ - activation_op = common.GetEndpointActivationOp(self.graph, context) - if activation_op: - return activation_op, None, None - - if '/' in context: - # If no activation op is there, look for them one level up. - add_context = re.search(r'^(.*)/([^/]+)', context).group(1) - activation_op = common.GetEndpointActivationOp(self.graph, add_context) - if not activation_op: - # Still no Relu, can happen on the top layer, just find the next node up, - # make sure it is BiasAdd. - consumers = [c for outp in op.outputs for c in outp.consumers()] - if len(consumers) != 1 or consumers[0].type != 'BiasAdd': - raise ValueError('Failed to quantize op: %s, %s' % (op.name, op.type)) - return consumers[0], None, None - if add_context: - add_op = self.GetOperationByNamesDontThrow([ - add_context + '/Add', add_context + '/add']) - return activation_op, add_op, add_context - else: - raise ValueError('Failed to quantize op: %s, %s' % (op.name, op.type)) - - def GetOperationByNameDontThrow(self, name): - """Returns an Operation with the given name. - - Args: - name: Name of Operation to return. - - Returns: - The Operation with the given name. None if the name does not correspond to - any operation in the graph. - """ - try: - return self.graph.get_operation_by_name(name) - except KeyError: - return None - - def GetOperationByNamesDontThrow(self, names): - """Returns an Operation with one of the given names. - - Args: - names: Names of Operation to return. - - Returns: - The Operation with one of the given names. None if none of the names - corresponds to any operation in the graph. - """ - for name in names: - op = self.GetOperationByNameDontThrow(name) - if op is not None: - return op - return None - - def _InsertQuantOp( - self, - context, - producer, - consumers, - name, - moving_avg=True, - init_min=-6.0, - init_max=6.0, - delay_requested=True, - bits=8, - narrow_range=False,): - """Inserts a quant op between a producer op and (multiple) consumer ops. - - Args: - context: Context where producer and consumer operations are nested. - producer: Producer operation of the pairs where quantization will be - inserted. - consumers: Consumer operations of the pairs. - name: Name for the new quantization op within the context. - moving_avg: Specifies whether to use exponential moving average or just - the last value seen. - init_min: Starting minimum value for the new quantization op. - init_max: Starting maximum value for the new quantization op. - delay_requested: If true, implement quantization delay where needed. - False value explicitly disables delay quantization everywhere. - bits: Number of bits to use for quantization, must be between 2 and 8. - narrow_range: Whether to use the narrow quantization range - [1; 2^bits - 1] or wide range [0; 2^bits - 1]. - Raises: - ValueError: When producer operation is not directly connected to the - consumer operation. - """ - scope = context + '/' + name - inputs = producer.outputs[0] - if moving_avg: - quant = (quant_ops.MovingAvgQuantize( - inputs, - init_min=init_min, - init_max=init_max, - ema_decay=self.ema_decay, - is_training=self.is_training, - num_bits=bits, - narrow_range=narrow_range, - updates_collection=_UPDATE_QUANT_OPS, - vars_collection=self.vars_collection, - scope=scope)) - else: - quant = (quant_ops.LastValueQuantize( - inputs, - init_min=init_min, - init_max=init_max, - is_training=self.is_training, - num_bits=bits, - narrow_range=narrow_range, - updates_collection=_UPDATE_QUANT_OPS, - vars_collection=self.vars_collection, - scope=scope)) - - if delay_requested and self.quant_delay and self.quant_delay > 0: - activate_quant = math_ops.greater_equal( - training_util.get_or_create_global_step(), - self.quant_delay, - name=scope + '/activate_quant') - quant = control_flow_ops.cond( - activate_quant, - lambda: quant, - lambda: inputs, - name=scope + '/delayed_quant') - - nodes_modified_count = graph_editor.reroute_ts( - [quant], [inputs], can_modify=consumers) - if nodes_modified_count != len(consumers): - raise ValueError('Some inputs not quantized for ops: [%s]' % - ', '.join([consumer.name for consumer in consumers])) + ema_decay=ema_decay, + quant_delay=quant_delay, + vars_collection=vars_collection, + bits=activation_bits, + init_min=0.0) + + # Quantize the inputs and output to the bypass (if it exists). The input to + # the bypass is the bias add, and the output is the activation. + if layer_match.bypass_op is not None: + _InsertQuantOp( + context, + 'conv_quant', + layer_match.bias_add_op, [layer_match.bypass_op], + is_training, + moving_avg=True, + ema_decay=ema_decay, + quant_delay=quant_delay, + vars_collection=vars_collection, + bits=activation_bits) + _InsertQuantOp( + add_context, + 'add_quant', + layer_match.bypass_op, + input_to_ops_map.ConsumerOperations(layer_match.bypass_op), + is_training, + moving_avg=True, + ema_decay=ema_decay, + quant_delay=quant_delay, + vars_collection=vars_collection, + bits=activation_bits) + + +def _FindLayersToQuantize(graph): + """Matches layers in graph to quantize. + + Args: + graph: Graph to perform match on. + + Yields: + _LayerMatches. + """ + input_pattern = graph_matcher.OpTypePattern('*') + weight_var_pattern = graph_matcher.OpTypePattern('|'.join(_WEIGHT_TYPES)) + weight_pattern = graph_matcher.OpTypePattern( + 'Identity', inputs=[weight_var_pattern]) + + folded_weight_pattern = graph_matcher.OpTypePattern('Mul') + + # The weights inputs to the layer operation can either be from the Variable or + # the folded weight (Mul). + layer_pattern = graph_matcher.OpTypePattern( + '|'.join(_QUANTIZABLE_TYPES), + inputs=[ + input_pattern, + graph_matcher.OneofPattern([weight_pattern, folded_weight_pattern]) + ]) + + folded_bias_mul_pattern = graph_matcher.OpTypePattern( + 'Mul', inputs=[graph_matcher.OpTypePattern('*'), layer_pattern]) + post_layer_op_correction_pattern = graph_matcher.OpTypePattern( + 'Add', inputs=[folded_bias_mul_pattern, + graph_matcher.OpTypePattern('*')]) + folded_bias_add_pattern = graph_matcher.OpTypePattern( + 'Add', + inputs=[ + post_layer_op_correction_pattern, + graph_matcher.OpTypePattern('*') + ]) + + bias_add_pattern = graph_matcher.OpTypePattern( + 'Add|BiasAdd', inputs=[layer_pattern, '*']) + + # The bias can come from the bias add or the folded bias add. + bypass_pattern_a = graph_matcher.OpTypePattern( + 'Add', + inputs=[ + graph_matcher.OneofPattern( + [bias_add_pattern, folded_bias_add_pattern]), '*' + ]) + bypass_pattern_b = graph_matcher.OpTypePattern( + 'Add', + inputs=[ + '*', + graph_matcher.OneofPattern( + [bias_add_pattern, folded_bias_add_pattern]) + ]) + + # The input to the activation can come from bias add, fold bias add or the + # bypasses. + activation_pattern = graph_matcher.OpTypePattern( + '|'.join(_ACTIVATION_TYPES), + inputs=[ + graph_matcher.OneofPattern([ + bias_add_pattern, folded_bias_add_pattern, bypass_pattern_a, + bypass_pattern_b + ]) + ]) + + layer_matcher = graph_matcher.GraphMatcher(activation_pattern) + for match_result in layer_matcher.match_graph(graph): + layer_op = match_result.get_op(layer_pattern) + weight_tensor = match_result.get_tensor(weight_pattern) + if weight_tensor is None: + weight_tensor = match_result.get_tensor(folded_weight_pattern) + activation_op = match_result.get_op(activation_pattern) + bias_add_op = match_result.get_op(bias_add_pattern) + if bias_add_op is None: + bias_add_op = match_result.get_op(folded_bias_add_pattern) + bypass_op = match_result.get_op(bypass_pattern_a) + if bypass_op is None: + bypass_op = match_result.get_op(bypass_pattern_b) + yield _LayerMatch(layer_op, weight_tensor, activation_op, bypass_op, + bias_add_op) + + +class _LayerMatch(object): + """Contains all information related to a matched Layer.""" + + def __init__(self, layer_op, weight_tensor, activation_op, bypass_op, + bias_add_op): + self._layer_op = layer_op + self._weight_tensor = weight_tensor + self._activation_op = activation_op + self._bypass_op = bypass_op + self._bias_add_op = bias_add_op + + @property + def layer_op(self): + return self._layer_op + + @property + def weight_tensor(self): + return self._weight_tensor + + @property + def activation_op(self): + return self._activation_op + + @property + def bypass_op(self): + return self._bypass_op + + @property + def bias_add_op(self): + return self._bias_add_op + + +def _InsertQuantOp(context, + name, + producer, + consumers, + is_training, + moving_avg=True, + init_min=-6.0, + init_max=6.0, + bits=8, + ema_decay=0.999, + quant_delay=None, + vars_collection=ops.GraphKeys.MOVING_AVERAGE_VARIABLES, + narrow_range=False): + """Inserts a quant op between a producer op and (multiple) consumer ops. + + Args: + context: Context w,here producer and consumer operations are nested. + name: Name for the new quantization op within the context. + producer: Producer operation of the pairs where quantization will be + inserted. + consumers: Consumer operations of the pairs. + is_training: Whether quantizing training graph or eval graph. + moving_avg: Specifies whether to use exponential moving average or just + the last value seen. + init_min: Starting minimum value for the new quantization op. + init_max: Starting maximum value for the new quantization op. + bits: Number of bits to use for quantization, must be between 2 and 8. + ema_decay: (Optional) Float, EMA decay parameter. EMA is used to update + quantization intervals for quantizing activations (see here about EMA: + https://en.wikipedia.org/wiki/Moving_average#Exponential_moving_average). + quant_delay: (Optional, default None) Int, count of global steps for which + to delay quantization. This helps weights stabilize at the start of + training. + vars_collection: (Optional) Collection where to store the variables for + quantization interval ends. + narrow_range: Whether to use the narrow quantization range + [1; 2^bits - 1] or wide range [0; 2^bits - 1]. + Raises: + ValueError: When producer operation is not directly connected to the + consumer operation. + """ + name_prefix = _AddContextToName(context, name) + inputs = producer.outputs[0] + if moving_avg: + quant = ( + quant_ops.MovingAvgQuantize( + inputs, + init_min=init_min, + init_max=init_max, + ema_decay=ema_decay, + is_training=is_training, + num_bits=bits, + narrow_range=narrow_range, + vars_collection=vars_collection, + name_prefix=name_prefix)) + else: + quant = ( + quant_ops.LastValueQuantize( + inputs, + init_min=init_min, + init_max=init_max, + is_training=is_training, + num_bits=bits, + narrow_range=narrow_range, + vars_collection=vars_collection, + name_prefix=name_prefix)) + + if quant_delay and quant_delay > 0: + activate_quant = math_ops.greater_equal( + common.CreateOrGetQuantizationStep(), + quant_delay, + name=name_prefix + '/activate_quant') + quant = control_flow_ops.cond( + activate_quant, + lambda: quant, + lambda: inputs, + name=name_prefix + '/delayed_quant') + + nodes_modified_count = graph_editor.reroute_ts( + [quant], [inputs], can_modify=consumers) + if nodes_modified_count != len(consumers): + raise ValueError('Some inputs not quantized for ops: [%s]' % ', '.join( + [consumer.name for consumer in consumers])) + + +def _GetContextFromOp(op): + """Gets the root context name from the op name.""" + context_re = re.search(r'^(.*)/([^/]+)', op.name) + if context_re: + return context_re.group(1) + return '' + + +def _AddContextToName(context, name): + """Adds the context to the name if it exists.""" + if not context: + return name + return context + '/' + name diff --git a/tensorflow/contrib/quantize/python/quantize_graph.py b/tensorflow/contrib/quantize/python/quantize_graph.py index d647bb94e849c713c2aca93c53f372bae5857c43..5a3a74cec4864ad3808d485849334c81f569d300 100644 --- a/tensorflow/contrib/quantize/python/quantize_graph.py +++ b/tensorflow/contrib/quantize/python/quantize_graph.py @@ -18,113 +18,198 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.quantize.python import copy_graph from tensorflow.contrib.quantize.python import fold_batch_norms from tensorflow.contrib.quantize.python import quantize from tensorflow.python.framework import ops -from tensorflow.python.ops import variables -def _create_graph(input_graph, - is_training, - elements=None, - device_name_or_function=None): - """Returns a transformed training input_graph for simulated quantization. +def _create_graph(input_graph=None, + is_training=True, + weight_bits=8, + activation_bits=8, + quant_delay=None, + freeze_bn_delay=None): + """Rewrites an input_graph in place for simulated quantization. - The forward pass has fake quantization ops inserted to simulate the error - introduced by quantization. + The graph has fake quantization ops inserted to simulate the error + introduced by quantization. Since the graph is transformed in place, + the expected behavior of previously held references to nodes and tensors may + change. Args: - input_graph: The tf.Graph to be transformed. + input_graph: The tf.Graph to be transformed, if None then defaults to the + default graph. is_training: Whether quantizing training or eval graph. - elements: (Optional) List of Tensors and Operations in input_graph whose - corresponding elements in the new graph will be returned. - device_name_or_function: (Optional) The device name or function to use. - - Returns: - g is new tf.Graph that is rewritten for simulated quantization. - l is a list of Tensors/Operations in g corresponding to the provided input - elements, if elements is not None. + weight_bits: Number of bits to use for quantizing weights. + activation_bits: Number of bits to use for quantizing activations. + quant_delay: Number of steps after which weights and activations are + quantized during training. + freeze_bn_delay: Number of steps after which moving mean and variance are + frozen and used instead of batch statistics during training. + freeze_bn_delay should be greater than quant_delay and should correspond + to the number of steps when training has almost converged Raises: ValueError: If elements contains an element that isn't a tf.Tensor or - tf.Operation. + tf.Operation. """ - # TODO(suharshs): Describe the process in more detail in the doc string. - g = copy_graph.CopyGraph(input_graph) - with g.as_default(): - with ops.device(device_name_or_function): - fold_batch_norms.FoldBatchNorms(g) - quantize.Quantize(g, is_training=is_training) - if elements is None: - return g - - return_elements = [] - for element in elements: - if isinstance(element, (ops.Tensor, variables.Variable)): - return_elements.append(g.get_tensor_by_name(element.name)) - elif isinstance(element, ops.Operation): - return_elements.append(g.get_operation_by_name(element.name)) - else: - raise ValueError( - 'elements must consist of Tensor or Operation objects, got: ', - str(element)) - return g, return_elements - - -def create_training_graph(input_graph, - elements=None, - device_name_or_function=None): - """Returns a transformed training input_graph for simulated quantization. - - The forward pass has fake quantization ops inserted to simulate the error - introduced by quantization. + + if input_graph is None: + input_graph = ops.get_default_graph() + with input_graph.as_default(): + fold_batch_norms.FoldBatchNorms( + input_graph, + freeze_batch_norm_delay=freeze_bn_delay, + is_training=is_training) + quantize.Quantize( + input_graph, + is_training, + quant_delay=quant_delay, + weight_bits=weight_bits, + activation_bits=activation_bits) + + +def create_training_graph(input_graph=None, quant_delay=0): + """Rewrites a training input_graph in place for simulated quantization. + + The graph has fake quantization ops inserted to simulate the error + introduced by quantization. Since the graph is transformed in place, + the expected behavior of previously held references to nodes and tensors may + change. + + The default value of quant_delay is suitable for finetuning an already trained + floating point model (recommended). + If one wants to train a quantized model from scratch, quant_delay should be + set to the number of steps it take the floating point model to converge. + Quantization will be activated at this point and effectively finetune the + model. If quant_delay is not provided when training from scratch, training can + often fail. Args: input_graph: The tf.Graph to be transformed. - elements: (Optional) List of Tensors and Operations in input_graph whose - corresponding elements in the new graph will be returned. - device_name_or_function: (Optional) The device name or function to use. + quant_delay: Number of steps after which weights and activations are + quantized during training. + + Raises: + ValueError: If elements contains an element that isn't a tf.Tensor or + tf.Operation. + """ + # TODO(raghuramank) Need to have freeze_bn_delay be a function of batch size + # Currently the values below are hardcoded for mobilenetV1 on imagenet + # Please use the experimental API if you need to tune these values. + if quant_delay == 0: + # Corresponds to case of restoring from a floating point checkpoint + # In this case, we can freeze the moving mean and variance early on and + # switch to using them during training. Therefore, freeze_bn_delay is set to + # 2e5. + freeze_bn_delay = int(2e5) + else: + # If training from scratch, set freeze_bn_delay to 100 epochs after quant + # delay. With a batch size of 64, this corresponds to 20000*100=2M steps. + freeze_bn_delay = quant_delay + int(2e6) + + _create_graph( + input_graph=input_graph, + is_training=True, + quant_delay=quant_delay, + freeze_bn_delay=freeze_bn_delay) + + +def create_eval_graph(input_graph=None): + """Rewrites an eval input_graph in place for simulated quantization. + + The graph has fake quantization ops inserted to simulate the error + introduced by quantization. Since the graph is transformed in place, + the expected behavior of previously held references to nodes and tensors may + change. + + Args: + input_graph: The tf.Graph to be transformed, if None then defaults to the + default graph. - Returns: - g is new tf.Graph that is rewritten for simulated quantization. - l is a list of Tensors/Operations in g corresponding to the provided input - elements, if elements is not None. + Raises: + ValueError: If elements contains an element that isn't a tf.Tensor or + tf.Operation. + """ + _create_graph(input_graph=input_graph, is_training=False) + + +def experimental_create_training_graph(input_graph=None, + weight_bits=8, + activation_bits=8, + quant_delay=0, + freeze_bn_delay=int(2e5)): + """Rewrites a training input_graph in place for simulated quantization. + + This function has additional experimental options not (yet) available to + create_training_graph. The resulting behavior may be undefined. + + The graph has fake quantization ops inserted to simulate the error + introduced by quantization. Since the graph is transformed in place, + the expected behavior of previously held references to nodes and tensors may + change. + + The default value of quant_delay is suitable for finetuning an already trained + floating point model (recommended). + If one wants to train a quantized model from scratch, quant_delay should be + set to the number of steps it take the floating point model to converge. + Quantization will be activated at this point and effectively finetune the + model. If quant_delay is not provided when training from scratch, training can + often fail. + + Args: + input_graph: The tf.Graph to be transformed,if None then defaults to the + default graph. + weight_bits: Number of bits to use for quantizing weights. + activation_bits: Number of bits to use for quantizing activations. + quant_delay: Number of steps after which weights and activations are + quantized during training. + freeze_bn_delay: Number of steps after which moving mean and variance are + frozen and used instead of batch statistics during training. + freeze_bn_delay should be greater than quant_delay and should correspond + to when training has almost converged Raises: ValueError: If elements contains an element that isn't a tf.Tensor or tf.Operation. """ - return _create_graph( + + _create_graph( input_graph=input_graph, is_training=True, - elements=elements, - device_name_or_function=device_name_or_function) + weight_bits=weight_bits, + activation_bits=activation_bits, + quant_delay=quant_delay, + freeze_bn_delay=freeze_bn_delay) + +def experimental_create_eval_graph(input_graph=None, + weight_bits=8, + activation_bits=8): + """Rewrites an eval input_graph in place for simulated quantization. -def create_eval_graph(input_graph, elements=None, device_name_or_function=None): - """Returns a transformed eval input_graph for simulated quantization. + This function has additional experimental options not (yet) available to + create_eval_graph. The resulting behavior may be undefined. - The forward pass has fake quantization ops inserted to simulate the error - introduced by quantization. + The graph has fake quantization ops inserted to simulate the error + introduced by quantization. Since the graph is transformed in place, + the expected behavior of previously held references to nodes and tensors may + change. Args: - input_graph: The tf.Graph to be transformed. - elements: (Optional) List of Tensors and Operations in input_graph whose - corresponding elements in the new graph will be returned. - device_name_or_function: (Optional) The device name or function to use. + input_graph: The tf.Graph to be transformed, if None then defaults to the + default graph. + weight_bits: Number of bits to use for quantizing weights. + activation_bits: Number of bits to use for quantizing activations. + - Returns: - g is new tf.Graph that is rewritten for simulated quantization. - l is a list of Tensors/Operations in g corresponding to the provided input - elements, if elements is not None. Raises: ValueError: If elements contains an element that isn't a tf.Tensor or - tf.Operation. + tf.Operation. """ - return _create_graph( + _create_graph( input_graph=input_graph, is_training=False, - elements=elements, - device_name_or_function=device_name_or_function) + weight_bits=weight_bits, + activation_bits=activation_bits) diff --git a/tensorflow/contrib/quantize/python/quantize_graph_test.py b/tensorflow/contrib/quantize/python/quantize_graph_test.py index 3407ace3914fe2de2506a2952ea5d1bf19028bb9..6b9289ef5f4b847172e1f093a1e4b5b2d3bdab57 100644 --- a/tensorflow/contrib/quantize/python/quantize_graph_test.py +++ b/tensorflow/contrib/quantize/python/quantize_graph_test.py @@ -20,121 +20,211 @@ from __future__ import print_function from tensorflow.contrib.layers.python.layers import layers from tensorflow.contrib.quantize.python import quantize_graph -from tensorflow.python.framework import constant_op from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import nn_ops -from tensorflow.python.ops import variables from tensorflow.python.platform import googletest class QuantizeGraphTest(test_util.TensorFlowTestCase): - # We have a lot of other tests that test the details of the rewrite, here we # just the specific features of the quantize_graph API. - def testReturnedElementsTraining(self): - self._TestReturnElements(True) - - def testReturnedElementsEval(self): - self._TestReturnElements(False) - def _TestReturnElements(self, is_training): - graph = ops.Graph() - with graph.as_default(): - a = constant_op.constant(1.0) - b = variables.Variable(2.0) - c = a + b - elements = [a, b, c.op] - if is_training: - q_graph, returned_elements = quantize_graph.create_training_graph( - graph, elements=elements) - else: - q_graph, returned_elements = quantize_graph.create_eval_graph( - graph, elements=elements) - # Make sure q_graph is different from graph. - self.assertTrue(graph != q_graph) - # Check that the returned elements are part of the new graph. - for returned_element in returned_elements: - self.assertEqual(q_graph, returned_element.graph) - # Check that the elements match with the one from the input graph. - for element, returned_element in zip(elements, returned_elements): - self.assertEqual(element.name, returned_element.name) - - def testNoReturnElementsTraining(self): - self._TestNoReturnElements(True) - - def testNoReturnElementsEval(self): - self._TestNoReturnElements(False) - - def _TestNoReturnElements(self, is_training): - graph = ops.Graph() - with graph.as_default(): - a = constant_op.constant(1.0) - b = variables.Variable(2.0) - _ = a + b - if is_training: - q_graph = quantize_graph.create_training_graph(graph) - else: - q_graph = quantize_graph.create_eval_graph(graph) - # Check that quantize_graph didn't return a tuple when elements isn't - # provided. - self.assertTrue(isinstance(q_graph, ops.Graph)) - # Make sure q_graph is different from graph. - self.assertTrue(graph != q_graph) - - def testDeviceNameTraining(self): - self._TestDeviceName(True) - - def testDeviceNameEval(self): - self._TestDeviceName(False) - - def _TestDeviceName(self, is_training): + def _RunTestOverAllRewrites(self, test_fn): + rewrite_fns = [ + quantize_graph.create_training_graph, + quantize_graph.create_eval_graph, + quantize_graph.experimental_create_training_graph, + quantize_graph.experimental_create_eval_graph, + ] + for fn in rewrite_fns: + test_fn(fn) + + def _RunTestOverTrainingRewrites(self, test_fn): + rewrite_fns = [ + quantize_graph.create_training_graph, + quantize_graph.experimental_create_training_graph, + ] + for fn in rewrite_fns: + test_fn(fn) + + def _RunTestOverEvalRewrites(self, test_fn): + rewrite_fns = [ + quantize_graph.create_eval_graph, + quantize_graph.experimental_create_eval_graph, + ] + for fn in rewrite_fns: + test_fn(fn) + + def _RunTestOverExperimentalRewrites(self, test_fn): + rewrite_fns = [ + quantize_graph.experimental_create_training_graph, + quantize_graph.experimental_create_eval_graph, + ] + for fn in rewrite_fns: + test_fn(fn) + + def testRewrite(self): + self._RunTestOverAllRewrites(self._TestRewrite) + + def _TestRewrite(self, rewrite_fn): graph = ops.Graph() with graph.as_default(): - batch_size, height, width, depth = 5, 128, 128, 3 - inputs = array_ops.zeros((batch_size, height, width, depth)) - conv = layers.conv2d( - inputs, - 32, [5, 5], - stride=2, - padding='SAME', - weights_initializer=self._WeightInit(0.09), - activation_fn=None, - scope='test') - _ = nn_ops.relu6(conv) - - device_name = '/job:oink/task:0/device:CPU:0' - if is_training: - q_graph = quantize_graph.create_training_graph( - graph, device_name_or_function=device_name) - else: - q_graph = quantize_graph.create_eval_graph( - graph, device_name_or_function=device_name) + self._ConvLayer() orig_variable_names = set( [v.name for v in graph.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)]) - q_variables = q_graph.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) - # Ensure that variables were added. - self.assertTrue(len(orig_variable_names) < len(q_variables)) - # All added variables should have the specified device name. - for var in q_variables: - if var.name not in orig_variable_names: - self.assertEqual(var.device, device_name) - - def _WeightInit(self, stddev): - """Returns truncated normal variable initializer. - Function is defined purely to shorten the name so that it stops wrapping. + rewrite_fn(graph) - Args: - stddev: Standard deviation of normal variable. + q_variables = graph.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + # Ensure that variables were added. + self.assertTrue(len(orig_variable_names) < len(q_variables)) - Returns: - An initialized that initialzes with a truncated normal variable. - """ - return init_ops.truncated_normal_initializer(stddev=stddev) + def testDefaultGraph(self): + self._RunTestOverAllRewrites(self._TestRewrite) + + def _TestDefaultGraph(self, rewrite_fn): + # Tests that the default graph is correctly used when no args are provided + # to rewrite_fn. + with ops.Graph().as_default() as g: + self._ConvLayer() + orig_variable_names = set( + [v.name for v in g.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)]) + rewrite_fn() + + q_variables = g.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + # Ensure that variables were added. + self.assertTrue(len(orig_variable_names) < len(q_variables)) + + def testQuantDelay(self): + self._RunTestOverTrainingRewrites(self._TestQuantDelay) + + def _TestQuantDelay(self, rewrite_fn): + with ops.Graph().as_default() as g: + self._ConvLayer() + quant_delay = 100 + rewrite_fn(quant_delay=quant_delay) + + quant_delay_found = False + for op in g.get_operations(): + # Check to see if the quant_delay is correctly set. + if 'activate_quant' in op.name and op.type == 'Const': + quant_delay_found = True + const_value = str(op.get_attr('value')) + self.assertTrue(('int64_val: %i' % quant_delay) in const_value) + self.assertTrue(quant_delay_found) + + def testWeightBits(self): + self._RunTestOverExperimentalRewrites(self._TestWeightBits) + + def _TestWeightBits(self, rewrite_fn): + with ops.Graph().as_default() as g: + self._ConvLayer() + weight_bits = 4 + rewrite_fn(weight_bits=weight_bits) + + weights_quant_found = False + for op in g.get_operations(): + # Check to see if FakeQuant operations for weights have the right bits + # set. + if 'weights_quant' in op.name and op.type == 'FakeQuantWithMinMaxVars': + weights_quant_found = True + self.assertEqual(op.get_attr('num_bits'), weight_bits) + self.assertTrue(weights_quant_found) + + def testActivationBits(self): + self._RunTestOverExperimentalRewrites(self._TestActivationBits) + + def _TestActivationBits(self, rewrite_fn): + with ops.Graph().as_default() as g: + self._ConvLayer() + activation_bits = 4 + rewrite_fn(activation_bits=activation_bits) + + act_quant_found = False + for op in g.get_operations(): + # Check to see if FakeQuant operations for activations have the right bits + # set. + act_quant_names = ['act_quant', 'conv_quant', 'add_quant'] + if any(s in op.name + for s in act_quant_names) and op.type == 'FakeQuantWithMinMaxVars': + act_quant_found = True + self.assertEqual(op.get_attr('num_bits'), activation_bits) + self.assertTrue(act_quant_found) + + def testTrainingQuantization(self): + self._RunTestOverTrainingRewrites(self._TestTrainingQuantization) + + def _TestTrainingQuantization(self, rewrite_fn): + with ops.Graph().as_default() as g: + self._ConvLayer() + rewrite_fn() + + # Ensure that FakeQuant and variable update nodes were found. + quant_found = False + assign_min_last_found = False + assign_min_ema_found = False + assign_max_last_found = False + assign_max_ema_found = False + for op in g.get_operations(): + # Check that FakeQuant operations were added. + if op.type == 'FakeQuantWithMinMaxVars': + quant_found = True + # Check that update operations for the added min max variables exist in + # the graph. + if 'AssignMinLast' in op.name: + assign_min_last_found = True + elif 'AssignMinEma' in op.name: + assign_min_ema_found = True + elif 'AssignMaxLast' in op.name: + assign_max_last_found = True + elif 'AssignMaxEma' in op.name: + assign_max_ema_found = True + self.assertTrue(assign_min_last_found) + self.assertTrue(assign_min_ema_found) + self.assertTrue(assign_max_last_found) + self.assertTrue(assign_max_ema_found) + self.assertTrue(quant_found) + + def testEvalQuantization(self): + self._RunTestOverEvalRewrites(self._TestEvalQuantization) + + def _TestEvalQuantization(self, rewrite_fn): + with ops.Graph().as_default() as g: + self._ConvLayer() + rewrite_fn() + + # Ensure that FakeQuant and variable update nodes were found. + quant_found = False + for op in g.get_operations(): + # Check that FakeQuant operations were added. + if op.type == 'FakeQuantWithMinMaxVars': + quant_found = True + # Check that update operations for the added min max variables don't + # exist in the graph. + update_names = [ + 'AssignMinLast', 'AssignMinEma', 'AssignMaxLast', 'AssignMaxEma' + ] + self.assertFalse(any(s in op.name for s in update_names)) + self.assertTrue(quant_found) + + def _ConvLayer(self): + """Add a basic convolution layer to the default graph.""" + batch_size, height, width, depth = 5, 128, 128, 3 + inputs = array_ops.zeros((batch_size, height, width, depth)) + weight_init = init_ops.truncated_normal_initializer + conv = layers.conv2d( + inputs, + 32, [5, 5], + stride=2, + padding='SAME', + weights_initializer=weight_init(0.09), + activation_fn=None, + scope='test') + _ = nn_ops.relu6(conv) if __name__ == '__main__': diff --git a/tensorflow/contrib/quantize/python/quantize_parameterized_test.py b/tensorflow/contrib/quantize/python/quantize_parameterized_test.py index 57dab03f162629f84adf1d15521b05f4014c4a80..639a7454a92aebd7289c59498cebff82cc003f75 100644 --- a/tensorflow/contrib/quantize/python/quantize_parameterized_test.py +++ b/tensorflow/contrib/quantize/python/quantize_parameterized_test.py @@ -29,7 +29,6 @@ from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops from tensorflow.python.platform import googletest -from tensorflow.python.training import training batch_norm = layers.batch_norm conv2d = layers.conv2d @@ -73,8 +72,6 @@ class QuantizeTest(test_util.TensorFlowTestCase): """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - batch_size, height, width, depth = 5, 128, 128, 3 inputs = array_ops.zeros((batch_size, height, width, depth)) stride = 1 if with_bypass else 2 @@ -91,7 +88,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - quantize.Quantize(graph, quant_delay=delay) + quantize.Quantize(graph, True, quant_delay=delay) quantization_node_name = 'FakeQuantWithMinMaxVars' weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' + quantization_node_name) @@ -101,7 +98,11 @@ class QuantizeTest(test_util.TensorFlowTestCase): scope + '/weights_quant/AssignMaxLast', scope + '/weights/read' ] self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + '/Conv2D' + if delay and delay > 0: + output_op_name = scope + '/weights_quant/delayed_quant/Switch_1' + else: + output_op_name = scope + '/Conv2D' + self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) if with_bypass: @@ -148,8 +149,6 @@ class QuantizeTest(test_util.TensorFlowTestCase): """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - batch_size, depth = 5, 256 inputs = array_ops.zeros((batch_size, depth)) out_depth = 256 if with_bypass else 128 @@ -165,7 +164,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - quantize.Quantize(graph, quant_delay=delay) + quantize.Quantize(graph, True, quant_delay=delay) quantization_node_name = 'FakeQuantWithMinMaxVars' weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' + @@ -176,7 +175,10 @@ class QuantizeTest(test_util.TensorFlowTestCase): scope + '/weights_quant/AssignMaxLast', scope + '/weights/read' ] self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + '/MatMul' + if delay and delay > 0: + output_op_name = scope + '/weights_quant/delayed_quant/Switch_1' + else: + output_op_name = scope + '/MatMul' self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) if with_bypass: @@ -222,8 +224,6 @@ class QuantizeTest(test_util.TensorFlowTestCase): """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - batch_size, height, width, depth = 5, 128, 128, 3 inputs = array_ops.zeros((batch_size, height, width, depth)) stride = 1 if with_bypass else 2 @@ -240,7 +240,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - quantize.Quantize(graph, quant_delay=delay) + quantize.Quantize(graph, True, quant_delay=delay) quantization_node_name = 'FakeQuantWithMinMaxVars' weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' + @@ -252,7 +252,10 @@ class QuantizeTest(test_util.TensorFlowTestCase): scope + '/depthwise_weights/read' ] self._AssertInputOpsAre(weights_quant, expected_inputs) - output_op_name = scope + '/depthwise' + if delay and delay > 0: + output_op_name = scope + '/weights_quant/delayed_quant/Switch_1' + else: + output_op_name = scope + '/depthwise' self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) if with_bypass: @@ -316,40 +319,11 @@ class QuantizeTest(test_util.TensorFlowTestCase): for params in parameters_list: test_fn(params[0], params[1], params[2], params[3], params[4]) - def _TestQuantize_Conv2dWithBatchNorm(self, activation, activation_op_name, - with_bypass, delay, fused_batch_norm): - """Tests quantization: inputs -> Conv2d with batch norm -> Activation. - - Args: - activation: Callable that returns an Operation, a factory method for the - Activation. - activation_op_name: String, name of the Activation operation. - with_bypass: Bool, when true there is an extra connection added from - inputs to just before Activation. - delay: Int (optional), delay in number of steps until quantization starts. - fused_batch_norm: Bool, when true use FusedBatchNorm. - """ - self._testQuantize_Conv2dWithBatchNorm( - activation, - activation_op_name, - with_bypass, - delay, - fused_batch_norm, - use_ema=True) - self._testQuantize_Conv2dWithBatchNorm( - activation, - activation_op_name, - with_bypass, - delay, - fused_batch_norm, - use_ema=False) - def testQuantize_Conv2dWithBatchNorm(self): self._RunBatchNormTestOverParameters(self._TestQuantize_Conv2dWithBatchNorm) - def _testQuantize_Conv2dWithBatchNorm(self, activation, activation_op_name, - with_bypass, delay, fused_batch_norm, - use_ema): + def _TestQuantize_Conv2dWithBatchNorm(self, activation, activation_op_name, + with_bypass, delay, fused_batch_norm): """Tests quantization: inputs -> Conv2d with batch norm -> Activation. Args: @@ -360,12 +334,9 @@ class QuantizeTest(test_util.TensorFlowTestCase): inputs to just before Activation. delay: Int (optional), delay in number of steps until quantization starts. fused_batch_norm: Bool, when true use FusedBatchNorm. - use_ema: Bool, when true uses EMA quantization for BN folded weights. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - batch_size, height, width, depth = 5, 128, 128, 3 inputs = array_ops.zeros((batch_size, height, width, depth)) stride = 1 if with_bypass else 2 @@ -392,25 +363,21 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - fold_batch_norms.FoldBatchNorms(graph) + fold_batch_norms.FoldBatchNorms(graph, is_training=True) - quantize.Quantize( - graph, quant_delay=delay, quantize_folded_weights_use_ema=use_ema) + quantize.Quantize(graph, True, quant_delay=delay) quantization_node_name = 'FakeQuantWithMinMaxVars' weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' + quantization_node_name) self.assertEqual(weights_quant.type, quantization_node_name) expected_inputs = [ - scope + '/weights_quant/' + ('AssignMinEma' - if use_ema else 'AssignMinLast'), - scope + '/weights_quant/' + ('AssignMaxEma' - if use_ema else 'AssignMaxLast'), - scope + '/mul_fold' + scope + '/weights_quant/' + 'AssignMinLast', + scope + '/weights_quant/' + 'AssignMaxLast', scope + '/mul_fold' ] self._AssertInputOpsAre(weights_quant, expected_inputs) output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1' - if (delay and use_ema) else '/Conv2D_Fold') + if delay else '/Conv2D_Fold') self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) if with_bypass: @@ -438,40 +405,11 @@ class QuantizeTest(test_util.TensorFlowTestCase): if delay else 'control_dependency') self._AssertOutputGoesToOps(act_quant, graph, [output_op_name]) - def _TestQuantize_FCWithBatchNorm(self, activation, activation_op_name, - with_bypass, delay, fused_batch_norm): - """Tests quantization: inputs -> FC with batch norm -> Activation. - - Args: - activation: Callable that returns an Operation, a factory method for the - Activation. - activation_op_name: String, name of the Activation operation. - with_bypass: Bool, when true there is an extra connection added from - inputs to just before Activation. - delay: Int (optional), delay in number of steps until quantization starts. - fused_batch_norm: Bool, when true use FusedBatchNorm. - """ - self._testQuantize_FCWithBatchNorm( - activation, - activation_op_name, - with_bypass, - delay, - fused_batch_norm, - use_ema=True) - self._testQuantize_FCWithBatchNorm( - activation, - activation_op_name, - with_bypass, - delay, - fused_batch_norm, - use_ema=False) - def testQuantize_FCWithBatchNorm(self): self._RunBatchNormTestOverParameters(self._TestQuantize_FCWithBatchNorm) - def _testQuantize_FCWithBatchNorm(self, activation, activation_op_name, - with_bypass, delay, fused_batch_norm, - use_ema): + def _TestQuantize_FCWithBatchNorm(self, activation, activation_op_name, + with_bypass, delay, fused_batch_norm): """Tests quantization: inputs -> FC with batch norm -> Activation. Args: @@ -482,12 +420,9 @@ class QuantizeTest(test_util.TensorFlowTestCase): inputs to just before Activation. delay: Int (optional), delay in number of steps until quantization starts. fused_batch_norm: Bool, when true use FusedBatchNorm. - use_ema: Bool, when true uses EMA quantization for BN folded weights. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - batch_size, depth = 5, 256 inputs = array_ops.zeros((batch_size, depth)) out_depth = 256 if with_bypass else 128 @@ -511,25 +446,21 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - fold_batch_norms.FoldBatchNorms(graph) + fold_batch_norms.FoldBatchNorms(graph, is_training=True) - quantize.Quantize( - graph, quant_delay=delay, quantize_folded_weights_use_ema=use_ema) + quantize.Quantize(graph, True, quant_delay=delay) quantization_node_name = 'FakeQuantWithMinMaxVars' weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' + quantization_node_name) self.assertEqual(weights_quant.type, quantization_node_name) expected_inputs = [ - scope + '/weights_quant/' + ('AssignMinEma' - if use_ema else 'AssignMinLast'), - scope + '/weights_quant/' + ('AssignMaxEma' - if use_ema else 'AssignMaxLast'), - scope + '/mul_fold' + scope + '/weights_quant/' + 'AssignMinLast', + scope + '/weights_quant/' + 'AssignMaxLast', scope + '/mul_fold' ] self._AssertInputOpsAre(weights_quant, expected_inputs) output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1' - if delay and use_ema else '/MatMul_Fold') + if delay else '/MatMul_Fold') self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) if with_bypass: @@ -557,42 +488,13 @@ class QuantizeTest(test_util.TensorFlowTestCase): if delay else 'control_dependency') self._AssertOutputGoesToOps(act_quant, graph, [output_op_name]) - def _TestQuantize_DepthwiseConv2dWithBatchNorm( - self, activation, activation_op_name, with_bypass, delay, - fused_batch_norm): - """Tests quantization: inputs -> DWConv2d with batch norm -> Activation. - - Args: - activation: Callable that returns an Operation, a factory method for the - Activation. - activation_op_name: String, name of the Activation operation. - with_bypass: Bool, when true there is an extra connection added from - inputs to just before Activation. - delay: Int (optional), delay in number of steps until quantization starts. - fused_batch_norm: Bool, when true use FusedBatchNorm. - """ - self._testQuantize_DepthwiseConv2dWithBatchNorm( - activation, - activation_op_name, - with_bypass, - delay, - fused_batch_norm, - use_ema=True) - self._testQuantize_DepthwiseConv2dWithBatchNorm( - activation, - activation_op_name, - with_bypass, - delay, - fused_batch_norm, - use_ema=False) - def testQuantize_DepthwiseConv2dWithBatchNorm(self): self._RunBatchNormTestOverParameters( self._TestQuantize_DepthwiseConv2dWithBatchNorm) - def _testQuantize_DepthwiseConv2dWithBatchNorm( + def _TestQuantize_DepthwiseConv2dWithBatchNorm( self, activation, activation_op_name, with_bypass, delay, - fused_batch_norm, use_ema): + fused_batch_norm): """Tests quantization: inputs -> DWConv2d with batch norm -> Activation. Args: @@ -603,12 +505,9 @@ class QuantizeTest(test_util.TensorFlowTestCase): inputs to just before Activation. delay: Int (optional), delay in number of steps until quantization starts. fused_batch_norm: Bool, when true use FusedBatchNorm. - use_ema: Bool, when true uses EMA quantization for BN folded weights. """ graph = ops.Graph() with graph.as_default(): - training.create_global_step(graph) - batch_size, height, width, depth = 5, 128, 128, 3 inputs = array_ops.zeros((batch_size, height, width, depth)) stride = 1 if with_bypass else 2 @@ -635,24 +534,20 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - fold_batch_norms.FoldBatchNorms(graph) + fold_batch_norms.FoldBatchNorms(graph, is_training=True) - quantize.Quantize( - graph, quant_delay=delay, quantize_folded_weights_use_ema=use_ema) + quantize.Quantize(graph, True, quant_delay=delay) quantization_node_name = 'FakeQuantWithMinMaxVars' weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' + quantization_node_name) self.assertEqual(weights_quant.type, quantization_node_name) expected_inputs = [ - scope + '/weights_quant/' + ('AssignMinEma' - if use_ema else 'AssignMinLast'), - scope + '/weights_quant/' + ('AssignMaxEma' - if use_ema else 'AssignMaxLast'), - scope + '/mul_fold' + scope + '/weights_quant/' + 'AssignMinLast', + scope + '/weights_quant/' + 'AssignMaxLast', scope + '/mul_fold' ] self._AssertInputOpsAre(weights_quant, expected_inputs) output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1' - if delay and use_ema else '/depthwise_Fold') + if delay else '/depthwise_Fold') self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name]) if with_bypass: diff --git a/tensorflow/contrib/quantize/python/quantize_test.py b/tensorflow/contrib/quantize/python/quantize_test.py index 1e4dd7cf67dbfbd16386fd740c7dcc83e05ad82a..bb7be0809421b64a019e73f00aac6c58524222e8 100644 --- a/tensorflow/contrib/quantize/python/quantize_test.py +++ b/tensorflow/contrib/quantize/python/quantize_test.py @@ -35,7 +35,15 @@ separable_conv2d = layers.separable_conv2d class QuantizeTest(test_util.TensorFlowTestCase): + def _RunTestOverParameters(self, test_fn): + params = [True, False] + for is_training in params: + test_fn(is_training) + def testInsertQuantOpFailsWhenOpsNotConnected(self): + pass + + def _TestInsertQuantOpFailsWhenOpsNotConnected(self, is_training): graph = ops.Graph() with graph.as_default(): batch_size, height, width, depth = 5, 128, 128, 3 @@ -45,17 +53,18 @@ class QuantizeTest(test_util.TensorFlowTestCase): activation_fn=None, scope='test') relu = nn_ops.relu6(inputs) - context = quantize._QuantizeContext(graph=graph, weight_bits=8, - weight_narrow_range=True, - activation_bits=8) # Inserting a quantization op between two unconnected ops should fail with # ValueError. with self.assertRaises(ValueError) as err: - context._InsertQuantOp('test', conv.op, [relu.op], 'FailingQuantOp') + quantize._InsertQuantOp('test', is_training, conv.op, [relu.op], + 'FailingQuantOp') self.assertEqual( str(err.exception), 'Some inputs not quantized for ops: [Relu6]') def testInsertQuantOpForAddAfterConv2d(self): + self._RunTestOverParameters(self._TestInsertQuantOpForAddAfterConv2d) + + def _TestInsertQuantOpForAddAfterConv2d(self, is_training): graph = ops.Graph() with graph.as_default(): batch_size, height, width, depth = 5, 128, 128, 3 @@ -70,8 +79,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - quantize.Quantize(graph=graph, weight_bits=8, weight_narrow_range=True, - activation_bits=8) + quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8) quantization_node_name = 'FakeQuantWithMinMaxVars' add_quant = graph.get_operation_by_name('test/add_quant/' + @@ -79,6 +87,10 @@ class QuantizeTest(test_util.TensorFlowTestCase): self.assertEqual(add_quant.type, quantization_node_name) def testInsertQuantOpForAddAfterSeparableConv2d(self): + self._RunTestOverParameters( + self._TestInsertQuantOpForAddAfterSeparableConv2d) + + def _TestInsertQuantOpForAddAfterSeparableConv2d(self, is_training): graph = ops.Graph() with graph.as_default(): batch_size, height, width, depth = 5, 128, 128, 3 @@ -94,8 +106,7 @@ class QuantizeTest(test_util.TensorFlowTestCase): with ops.control_dependencies([update_barrier]): array_ops.identity(node, name='control_dependency') - quantize.Quantize(graph=graph, weight_bits=8, weight_narrow_range=True, - activation_bits=8) + quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8) quantization_node_name = 'FakeQuantWithMinMaxVars' add_quant = graph.get_operation_by_name('test/add_quant/' + diff --git a/tensorflow/contrib/receptive_field/BUILD b/tensorflow/contrib/receptive_field/BUILD index d16b2908a0285e04ef5d3ede2050bf24c508228d..e975aeaea7ee78f8e912be8ab1be61b9acc7b418 100644 --- a/tensorflow/contrib/receptive_field/BUILD +++ b/tensorflow/contrib/receptive_field/BUILD @@ -15,7 +15,6 @@ load("//tensorflow:tensorflow.bzl", "py_test") py_library( name = "receptive_field_pip", deps = [ - ":graph_compute_order_py", ":receptive_field_py", ], ) @@ -23,28 +22,75 @@ py_library( py_library( name = "graph_compute_order_py", srcs = [ - "__init__.py", "python/util/graph_compute_order.py", ], srcs_version = "PY2AND3", + deps = [ + ":parse_layer_parameters_py", + "//tensorflow/python:platform", + ], +) + +py_library( + name = "parse_layer_parameters_py", + srcs = [ + "python/util/parse_layer_parameters.py", + ], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/contrib/util:util_py", + "//tensorflow/python:platform", + ], ) py_library( name = "receptive_field_py", srcs = [ - "__init__.py", + "python/util/parse_layer_parameters.py", "python/util/receptive_field.py", + "receptive_field_api.py", ], srcs_version = "PY2AND3", deps = [ ":graph_compute_order_py", - "//tensorflow/contrib/util:util_py", "//tensorflow/python:framework_ops", "//tensorflow/python:platform", "//third_party/py/numpy", ], ) +py_test( + name = "graph_compute_order_test", + srcs = ["python/util/graph_compute_order_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":graph_compute_order_py", + ":receptive_field_py", + "//tensorflow/contrib/slim", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:nn", + ], +) + +py_test( + name = "parse_layer_parameters_test", + srcs = ["python/util/parse_layer_parameters_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":graph_compute_order_py", + ":parse_layer_parameters_py", + "//tensorflow/contrib/slim", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:nn", + ], +) + py_test( name = "receptive_field_test", srcs = ["python/util/receptive_field_test.py"], diff --git a/tensorflow/contrib/receptive_field/README.md b/tensorflow/contrib/receptive_field/README.md index dfe53cdf14299b055fa4cdc13008d88234e93aa8..3ff85faf611afad71b6e6203453bbe97c56f9242 100644 --- a/tensorflow/contrib/receptive_field/README.md +++ b/tensorflow/contrib/receptive_field/README.md @@ -17,7 +17,6 @@ For example, if your model is constructed using the function ```python import tensorflow as tf -from tensorflow.contrib import receptive_field # Construct graph. g = tf.Graph() @@ -27,7 +26,7 @@ with g.as_default(): # Compute receptive field parameters. rf_x, rf_y, eff_stride_x, eff_stride_y, eff_pad_x, eff_pad_y = \ - receptive_field.compute_receptive_field_from_graph_def( \ + tf.contrib.receptive_field.compute_receptive_field_from_graph_def( \ g.as_graph_def(), 'input_image', 'my_output_endpoint') ``` @@ -47,7 +46,6 @@ You can then compute the receptive field parameters for Inception-Resnet-v2 as: ```python from nets import inception import tensorflow as tf -from tensorflow.contrib import receptive_field # Construct graph. g = tf.Graph() @@ -57,7 +55,7 @@ with g.as_default(): # Compute receptive field parameters. rf_x, rf_y, eff_stride_x, eff_stride_y, eff_pad_x, eff_pad_y = \ - receptive_field.compute_receptive_field_from_graph_def( \ + tf.contrib.receptive_field.compute_receptive_field_from_graph_def( \ g.as_graph_def(), 'input_image', 'InceptionResnetV2/Conv2d_7b_1x1/Relu') ``` diff --git a/tensorflow/contrib/receptive_field/python/util/examples/compute_rf.py b/tensorflow/contrib/receptive_field/python/util/examples/compute_rf.py index 1cf978b90a3661a075130790d82a499da4d8a0cc..d6fdd12bbe37fb0e0cb12f1d0adc3fce29b19e8a 100644 --- a/tensorflow/contrib/receptive_field/python/util/examples/compute_rf.py +++ b/tensorflow/contrib/receptive_field/python/util/examples/compute_rf.py @@ -26,7 +26,7 @@ import sys from google.protobuf import text_format -from tensorflow.contrib import receptive_field +from tensorflow.contrib.receptive_field import receptive_field_api as receptive_field from tensorflow.core.framework import graph_pb2 from tensorflow.python.platform import app from tensorflow.python.platform import gfile diff --git a/tensorflow/contrib/receptive_field/python/util/examples/rf_benchmark.py b/tensorflow/contrib/receptive_field/python/util/examples/rf_benchmark.py index 94228dfa61b1de617f131611173fda7c3917d250..a298b4d49038468299b58140758c69675368e855 100644 --- a/tensorflow/contrib/receptive_field/python/util/examples/rf_benchmark.py +++ b/tensorflow/contrib/receptive_field/python/util/examples/rf_benchmark.py @@ -28,19 +28,19 @@ import argparse import csv import sys -from nets import alexnet -from nets import inception -from nets import mobilenet_v1 -from nets import resnet_v1 -from nets import resnet_v2 -from nets import vgg from tensorflow.contrib import framework -from tensorflow.contrib import receptive_field from tensorflow.contrib import slim +from tensorflow.contrib.receptive_field import receptive_field_api as receptive_field from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.platform import app +from nets import alexnet +from nets import inception +from nets import mobilenet_v1 +from nets import resnet_v1 +from nets import resnet_v2 +from nets import vgg cmd_args = None @@ -245,7 +245,8 @@ def _model_rf(graphdef, end_points, desired_end_point_keys, model_type='resnet_v1_50', - csv_writer=None): + csv_writer=None, + input_resolution=None): """Computes receptive field information for a given CNN model. The information will be printed to stdout. If the RF parameters are the same @@ -261,45 +262,93 @@ def _model_rf(graphdef, information will be computed. model_type: Type of model to be used, used only for printing purposes. csv_writer: A CSV writer for RF parameters, which is used if it is not None. + input_resolution: Input resolution to use when computing RF + parameters. This is important for the case where padding can only be + defined if the input resolution is known, which may happen if using SAME + padding. This is assumed the resolution for both height and width. If + None, we consider the resolution is unknown. """ for desired_end_point_key in desired_end_point_keys: print('- %s:' % desired_end_point_key) output_node_with_colon = end_points[desired_end_point_key].name pos = output_node_with_colon.rfind(':') output_node = output_node_with_colon[:pos] - (receptive_field_x, receptive_field_y, effective_stride_x, - effective_stride_y, effective_padding_x, effective_padding_y - ) = receptive_field.compute_receptive_field_from_graph_def( - graphdef, _INPUT_NODE, output_node) - # If values are the same in horizontal/vertical directions, just report one - # of them. Otherwise, report both. - if (receptive_field_x == receptive_field_y) and ( - effective_stride_x == effective_stride_y) and ( - effective_padding_x == effective_padding_y): - print('Receptive field size = %5s, effective stride = %5s, effective ' - 'padding = %5s' % (str(receptive_field_x), str(effective_stride_x), - str(effective_padding_x))) - else: - print('Receptive field size: horizontal = %5s, vertical = %5s. ' - 'Effective stride: horizontal = %5s, vertical = %5s. Effective ' - 'padding: horizontal = %5s, vertical = %5s' % - (str(receptive_field_x), str(receptive_field_y), - str(effective_stride_x), str(effective_stride_y), - str(effective_padding_x), str(effective_padding_y))) - if csv_writer is not None: - csv_writer.writerow({ - 'CNN': model_type, - 'end_point': desired_end_point_key, - 'RF size hor': str(receptive_field_x), - 'RF size ver': str(receptive_field_y), - 'effective stride hor': str(effective_stride_x), - 'effective stride ver': str(effective_stride_y), - 'effective padding hor': str(effective_padding_x), - 'effective padding ver': str(effective_padding_y) - }) - - -def _process_model_rf(model_type='resnet_v1_50', csv_writer=None, arg_sc=None): + try: + (receptive_field_x, receptive_field_y, effective_stride_x, + effective_stride_y, effective_padding_x, effective_padding_y + ) = receptive_field.compute_receptive_field_from_graph_def( + graphdef, _INPUT_NODE, output_node, input_resolution=input_resolution) + # If values are the same in horizontal/vertical directions, just report + # one of them. Otherwise, report both. + if (receptive_field_x == receptive_field_y) and ( + effective_stride_x == effective_stride_y) and ( + effective_padding_x == effective_padding_y): + print('Receptive field size = %5s, effective stride = %5s, effective ' + 'padding = %5s' % (str(receptive_field_x), + str(effective_stride_x), + str(effective_padding_x))) + else: + print('Receptive field size: horizontal = %5s, vertical = %5s. ' + 'Effective stride: horizontal = %5s, vertical = %5s. Effective ' + 'padding: horizontal = %5s, vertical = %5s' % + (str(receptive_field_x), str(receptive_field_y), + str(effective_stride_x), str(effective_stride_y), + str(effective_padding_x), str(effective_padding_y))) + if csv_writer is not None: + csv_writer.writerow({ + 'CNN': + model_type, + 'input resolution': + str(input_resolution[0]) + if input_resolution is not None else 'None', + 'end_point': + desired_end_point_key, + 'RF size hor': + str(receptive_field_x), + 'RF size ver': + str(receptive_field_y), + 'effective stride hor': + str(effective_stride_x), + 'effective stride ver': + str(effective_stride_y), + 'effective padding hor': + str(effective_padding_x), + 'effective padding ver': + str(effective_padding_y) + }) + except ValueError as e: + print('---->ERROR: Computing RF parameters for model %s with final end ' + 'point %s and input resolution %s did not work' % + (model_type, desired_end_point_key, input_resolution)) + print('---->The returned error is: %s' % e) + if csv_writer is not None: + csv_writer.writerow({ + 'CNN': + model_type, + 'input resolution': + str(input_resolution[0]) + if input_resolution is not None else 'None', + 'end_point': + desired_end_point_key, + 'RF size hor': + 'None', + 'RF size ver': + 'None', + 'effective stride hor': + 'None', + 'effective stride ver': + 'None', + 'effective padding hor': + 'None', + 'effective padding ver': + 'None' + }) + + +def _process_model_rf(model_type='resnet_v1_50', + csv_writer=None, + arg_sc=None, + input_resolutions=None): """Contructs model graph and desired end-points, and compute RF. The computed RF parameters are printed to stdout by the _model_rf function. @@ -308,13 +357,30 @@ def _process_model_rf(model_type='resnet_v1_50', csv_writer=None, arg_sc=None): model_type: Type of model to be used. csv_writer: A CSV writer for RF parameters, which is used if it is not None. arg_sc: Optional arg scope to use in constructing the graph. + input_resolutions: List of 1D input resolutions to use when computing RF + parameters. This is important for the case where padding can only be + defined if the input resolution is known, which may happen if using SAME + padding. The entries in the list are assumed the resolution for both + height and width. If one of the elements in the list is None, we consider + it to mean that the resolution is unknown. If the list itself is None, + we use the default list [None, 224, 321]. """ - print('********************%s' % model_type) - graphdef, end_points = _model_graph_def(model_type, arg_sc) - desired_end_point_keys = _get_desired_end_point_keys(model_type) - _model_rf(graphdef, end_points, desired_end_point_keys, model_type, - csv_writer) + # Process default value for this list. + if input_resolutions is None: + input_resolutions = [None, 224, 321] + + for n in input_resolutions: + print('********************%s, input resolution = %s' % (model_type, n)) + graphdef, end_points = _model_graph_def(model_type, arg_sc) + desired_end_point_keys = _get_desired_end_point_keys(model_type) + _model_rf( + graphdef, + end_points, + desired_end_point_keys, + model_type, + csv_writer, + input_resolution=[n, n] if n is not None else None) def _resnet_rf(csv_writer=None): @@ -421,7 +487,7 @@ def main(unused_argv): if cmd_args.csv_path: csv_file = open(cmd_args.csv_path, 'w') field_names = [ - 'CNN', 'end_point', 'RF size hor', 'RF size ver', + 'CNN', 'input resolution', 'end_point', 'RF size hor', 'RF size ver', 'effective stride hor', 'effective stride ver', 'effective padding hor', 'effective padding ver' ] diff --git a/tensorflow/contrib/receptive_field/python/util/examples/write_inception_resnet_v2_graph.py b/tensorflow/contrib/receptive_field/python/util/examples/write_inception_resnet_v2_graph.py index 793ae163d807fdda62c2025cb8176b96832cb61a..a494883396614bcee04886b5c05d0393df958580 100644 --- a/tensorflow/contrib/receptive_field/python/util/examples/write_inception_resnet_v2_graph.py +++ b/tensorflow/contrib/receptive_field/python/util/examples/write_inception_resnet_v2_graph.py @@ -22,12 +22,12 @@ from __future__ import print_function import argparse import sys -from nets import inception from tensorflow.python.framework import dtypes from tensorflow.python.framework import graph_io from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.platform import app +from nets import inception cmd_args = None diff --git a/tensorflow/contrib/receptive_field/python/util/graph_compute_order.py b/tensorflow/contrib/receptive_field/python/util/graph_compute_order.py index 8af4be16d6c17286287713a1fb6f5017355e3b32..0388079f20dee0a6b249d568e2c51d1407d7466f 100644 --- a/tensorflow/contrib/receptive_field/python/util/graph_compute_order.py +++ b/tensorflow/contrib/receptive_field/python/util/graph_compute_order.py @@ -20,69 +20,173 @@ from __future__ import division from __future__ import print_function import collections +import math +from tensorflow.contrib.receptive_field.python.util import parse_layer_parameters +from tensorflow.python.platform import tf_logging as logging -class GraphDefHelper(object): - """Helper class to collect node names and definitions. +def parse_graph_nodes(graph_def): + """Helper function to parse GraphDef's nodes. - Example: - b = GraphDefHelper(graph_def) - # Prints node that produces given output. - print b.output_of['conv/foo/bar'] + It returns a dict mapping from node name to NodeDef. + + Args: + graph_def: A GraphDef object. + + Returns: + name_to_node: Dict keyed by node name, each entry containing the node's + NodeDef. """ + name_to_node = {} + for node_def in graph_def.node: + name_to_node[node_def.name] = node_def + return name_to_node - def __init__(self, gd): - self.output_of = {} - for each in gd.node: - self.output_of[each.name] = each +# Named tuple used to collect information from each node in a computation graph. +_node_info = collections.namedtuple( + 'NodeInfo', field_names=['order', 'node', 'input_size', 'output_size']) -# pylint: disable=invalid-name -_NodeEntry = collections.namedtuple('NodeEntry', field_names=['order', 'node']) +def _compute_output_resolution(input_spatial_resolution, kernel_size, stride, + total_padding): + """Computes output resolution, given input resolution and layer parameters. -def _get_computed_nodes(g, output, seen): - """Traverses the graph in topological order. + Note that this computation is done only over one dimension (eg, x or y). + If any of the inputs is None, returns None. + + Args: + input_spatial_resolution: Input spatial resolution (int). + kernel_size: Kernel size (int). + stride: Stride (int). + total_padding: Total padding to be applied (int). + Returns: + output_resolution: Output dimension (int) or None. + """ + if (input_spatial_resolution is None) or (kernel_size is None) or ( + stride is None) or (total_padding is None): + return None + return int( + math.ceil(( + input_spatial_resolution + total_padding - kernel_size + 1) / stride)) + + +def _get_computed_nodes(name_to_node, + current, + node_info, + input_node_name='', + input_node_size=None): + """Traverses the graph recursively to compute its topological order. + + Optionally, the function may also compute the input and output feature map + resolutions at each node. In this case, input_node_name and input_node_size + must be set. Note that if a node's op type is unknown, the input and output + resolutions are ignored and set to None. Args: - g: GraphDefHelper object. - output: current node. - seen: map of nodes we've already traversed. + name_to_node: Dict keyed by node name, each entry containing the node's + NodeDef. + current: Current node name. + node_info: Map of nodes we've already traversed, containing their _node_info + information. + input_node_name: Name of node with fixed input resolution (optional). + input_node_size: Fixed input resolution to use (optional). Returns: - order in topological sort for 'output'. + order: Order in topological sort for 'current'. + input_size: Tensor spatial resolution at input of current node. + output_size: Tensor spatial resolution at output of current node. """ - if output in seen: - return seen[output].order - node_def = g.output_of.get(output, None) - if node_def is None: - seen[output] = _NodeEntry(0, None) - return 0 - - r = 0 + if current in node_info: + return (node_info[current].order, node_info[current].input_size, + node_info[current].output_size) + + node_def = name_to_node[current] + + if current == input_node_name: + order = 0 + input_size = None + output_size = input_node_size + node_info[current] = _node_info(order, node_def, input_size, output_size) + return (order, input_size, output_size) + + input_size = None + output_size = None + + order = 0 + number_inputs = 0 for each in node_def.input: # Parses name of input node. if each.startswith('^'): - each = each[1:] + # The character '^' denotes a control dependency, so this input node can + # be safely ignored. + continue each = each.split(':')[0] # Recursively computes ordering. - new_v = _get_computed_nodes(g, each, seen) - r = max(r, new_v + 1) - - seen[output] = _NodeEntry(r, node_def) - - return seen[output].order - - -def get_compute_order(graph_def): - """Computes order of computation for a given graph. + (parent_order, _, parent_output_size) = _get_computed_nodes( + name_to_node, each, node_info, input_node_name, input_node_size) + order = max(order, parent_order + 1) + if number_inputs == 0: + # For all the types of nodes we consider, the first input corresponds to + # the feature map. + input_size = parent_output_size + number_inputs += 1 + + # Figure out output size for this layer. + logging.vlog(3, 'input_size = %s', input_size) + if input_size is None: + output_size = None + else: + (kernel_size_x, kernel_size_y, stride_x, stride_y, _, _, total_padding_x, + total_padding_y) = ( + parse_layer_parameters.get_layer_params( + node_def, name_to_node, input_size, force=True)) + logging.vlog(3, 'kernel_size_x = %s, kernel_size_y = %s, ' + 'stride_x = %s, stride_y = %s, ' + 'total_padding_x = %s, total_padding_y = %s' % + (kernel_size_x, kernel_size_y, stride_x, stride_y, + total_padding_x, total_padding_y)) + output_size = [None] * 2 + output_size[0] = _compute_output_resolution(input_size[0], kernel_size_x, + stride_x, total_padding_x) + output_size[1] = _compute_output_resolution(input_size[1], kernel_size_y, + stride_y, total_padding_y) + + logging.vlog(3, 'output_size = %s', output_size) + node_info[current] = _node_info(order, node_def, input_size, output_size) + + return order, input_size, output_size + + +def get_compute_order(graph_def, input_node_name='', input_node_size=None): + """Computes order of computation for a given CNN graph. + + Optionally, the function may also compute the input and output feature map + resolutions at each node. In this case, input_node_name and input_node_size + must be set. Note that if a node's op type is unknown, the input and output + resolutions are ignored and set to None. Args: graph_def: GraphDef object. + input_node_name: Name of node with fixed input resolution (optional). This + is usually the node name for the input image in a CNN. + input_node_size: 2D list of integers, fixed input resolution to use + (optional). This is usually the input resolution used for the input image + in a CNN (common examples are: [224, 224], [299, 299], [321, 321]). Returns: - map: name -> {order, node} + node_info: Default dict keyed by node name, mapping to a named tuple with + the following fields: + - order: Integer denoting topological order; + - node: NodeDef for the given node; + - input_size: 2D list of integers, denoting the input spatial resolution + to the node; + - output_size: 2D list of integers, denoting the output spatial resolution + of the node. + name_to_node: Dict keyed by node name, each entry containing the node's + NodeDef. """ - helper = GraphDefHelper(graph_def) - seen = collections.defaultdict(_NodeEntry) + name_to_node = parse_graph_nodes(graph_def) + node_info = collections.defaultdict(_node_info) for each in graph_def.node: - _get_computed_nodes(helper, each.name, seen) - return seen + _get_computed_nodes(name_to_node, each.name, node_info, input_node_name, + input_node_size) + return node_info, name_to_node diff --git a/tensorflow/contrib/receptive_field/python/util/graph_compute_order_test.py b/tensorflow/contrib/receptive_field/python/util/graph_compute_order_test.py new file mode 100644 index 0000000000000000000000000000000000000000..94c992ad2155d9c80cceee260a90c478a54371fb --- /dev/null +++ b/tensorflow/contrib/receptive_field/python/util/graph_compute_order_test.py @@ -0,0 +1,152 @@ +# Copyright 2017 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 graph_compute_order module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib import slim +from tensorflow.contrib.receptive_field import receptive_field_api as receptive_field +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_math_ops +from tensorflow.python.ops import nn +from tensorflow.python.platform import test + + +def create_test_network(): + """Convolutional neural network for test. + + Returns: + g: Tensorflow graph object (Graph proto). + """ + g = ops.Graph() + with g.as_default(): + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') + # Left branch before first addition. + l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='VALID') + # Right branch before first addition. + l2_pad = array_ops.pad(x, [[0, 0], [1, 0], [1, 0], [0, 0]], name='L2_pad') + l2 = slim.conv2d(l2_pad, 1, [3, 3], stride=2, scope='L2', padding='VALID') + l3 = slim.max_pool2d(l2, [3, 3], stride=2, scope='L3', padding='SAME') + # First addition. + l4 = nn.relu(l1 + l3, name='L4_relu') + # Left branch after first addition. + l5 = slim.conv2d(l4, 1, [1, 1], stride=2, scope='L5', padding='SAME') + # Right branch after first addition. + l6 = slim.conv2d(l4, 1, [3, 3], stride=2, scope='L6', padding='SAME') + # Final addition. + gen_math_ops.add(l5, l6, name='L7_add') + + return g + + +class GraphComputeOrderTest(test.TestCase): + + def check_topological_sort_and_sizes(self, + node_info, + expected_input_sizes=None, + expected_output_sizes=None): + """Helper function to check topological sorting and sizes are correct. + + The arguments expected_input_sizes and expected_output_sizes are used to + check that the sizes are correct, if they are given. + + Args: + node_info: Default dict keyed by node name, mapping to a named tuple with + the following keys: {order, node, input_size, output_size}. + expected_input_sizes: Dict mapping node names to expected input sizes + (optional). + expected_output_sizes: Dict mapping node names to expected output sizes + (optional). + """ + # Loop over nodes in sorted order, collecting those that were already seen. + # These will be used to make sure that the graph is topologically sorted. + # At the same time, we construct dicts from node name to input/output size, + # which will be used to check those. + already_seen_nodes = [] + input_sizes = {} + output_sizes = {} + for _, (_, node, input_size, output_size) in sorted( + node_info.items(), key=lambda x: x[1].order): + for inp_name in node.input: + # Since the graph is topologically sorted, the inputs to the current + # node must have been seen beforehand. + self.assertIn(inp_name, already_seen_nodes) + input_sizes[node.name] = input_size + output_sizes[node.name] = output_size + already_seen_nodes.append(node.name) + + # Check input sizes, if desired. + if expected_input_sizes is not None: + for k, v in expected_input_sizes.items(): + self.assertIn(k, input_sizes) + self.assertEqual(input_sizes[k], v) + + # Check output sizes, if desired. + if expected_output_sizes is not None: + for k, v in expected_output_sizes.items(): + self.assertIn(k, output_sizes) + self.assertEqual(output_sizes[k], v) + + def testGraphOrderIsCorrect(self): + """Tests that the order and sizes of create_test_network() are correct.""" + + graph_def = create_test_network().as_graph_def() + + # Case 1: Input node name/size are not given. + node_info, _ = receptive_field.get_compute_order(graph_def) + self.check_topological_sort_and_sizes(node_info) + + # Case 2: Input node name is given, but not size. + node_info, _ = receptive_field.get_compute_order( + graph_def, input_node_name='input_image') + self.check_topological_sort_and_sizes(node_info) + + # Case 3: Input node name and size (224) are given. + node_info, _ = receptive_field.get_compute_order( + graph_def, input_node_name='input_image', input_node_size=[224, 224]) + expected_input_sizes = { + 'input_image': None, + 'L1/Conv2D': [224, 224], + 'L2_pad': [224, 224], + 'L2/Conv2D': [225, 225], + 'L3/MaxPool': [112, 112], + 'L4_relu': [56, 56], + 'L5/Conv2D': [56, 56], + 'L6/Conv2D': [56, 56], + 'L7_add': [28, 28], + } + expected_output_sizes = { + 'input_image': [224, 224], + 'L1/Conv2D': [56, 56], + 'L2_pad': [225, 225], + 'L2/Conv2D': [112, 112], + 'L3/MaxPool': [56, 56], + 'L4_relu': [56, 56], + 'L5/Conv2D': [28, 28], + 'L6/Conv2D': [28, 28], + 'L7_add': [28, 28], + } + self.check_topological_sort_and_sizes(node_info, expected_input_sizes, + expected_output_sizes) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters.py b/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters.py new file mode 100644 index 0000000000000000000000000000000000000000..bc383a803496380aaba4d0248d2b7f93253b2b50 --- /dev/null +++ b/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters.py @@ -0,0 +1,326 @@ +# Copyright 2017 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. +# ============================================================================== +"""Functions to parse RF-related parameters from TF layers.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import math +from tensorflow.contrib.util import make_ndarray +from tensorflow.python.platform import tf_logging as logging + +# White-listed layer operations, which do not affect the receptive field +# computation. +_UNCHANGED_RF_LAYER_OPS = [ + "Add", "BiasAdd", "Cast", "Ceil", "ConcatV2", "Const", "Floor", + "FusedBatchNorm", "Identity", "Log", "Mul", "Pow", "RealDiv", "Relu", + "Relu6", "Round", "Rsqrt", "Softplus", "Sub", "VariableV2" +] + +# Different ways in which padding modes may be spelled. +_VALID_PADDING = ["VALID", b"VALID"] +_SAME_PADDING = ["SAME", b"SAME"] + + +def _stride_size(node, name_to_node): + """Computes stride size given a TF node. + + Args: + node: Tensorflow node (NodeDef proto). + name_to_node: For MaxPoolV2, mapping from variable name Tensorflow node. + + Returns: + stride_x: Stride size for horizontal direction (integer). + stride_y: Stride size for vertical direction (integer). + + Raises: + ValueError: If stride input cannot be found in `name_to_node`. + """ + if node.op == "MaxPoolV2": + strides_input_name = node.input[2] + if not strides_input_name.endswith("/strides"): + raise ValueError("Strides name does not end with '/strides'") + strides_node = name_to_node[strides_input_name] + value = strides_node.attr["value"] + t = make_ndarray(value.tensor) + stride_y = t[1] + stride_x = t[2] + else: + strides_attr = node.attr["strides"] + logging.vlog(4, "strides_attr = %s", strides_attr) + stride_y = strides_attr.list.i[1] + stride_x = strides_attr.list.i[2] + return stride_x, stride_y + + +def _conv_kernel_size(node, name_to_node): + """Computes kernel size given a TF convolution or pooling node. + + Args: + node: Tensorflow node (NodeDef proto). + name_to_node: Dict keyed by node name, each entry containing the node's + NodeDef. + + Returns: + kernel_size_x: Kernel size for horizontal direction (integer). + kernel_size_y: Kernel size for vertical direction (integer). + + Raises: + ValueError: If the weight layer node is invalid. + """ + weights_layer_read_name = node.input[1] + if not weights_layer_read_name.endswith("/read"): + raise ValueError( + "Weight layer's name input to conv layer does not end with '/read'") + weights_layer_param_name = weights_layer_read_name[:-5] + weights_node = name_to_node[weights_layer_param_name] + if weights_node.op != "VariableV2": + raise ValueError("Weight layer is not of type VariableV2") + shape = weights_node.attr["shape"] + logging.vlog(4, "weight shape = %s", shape) + kernel_size_y = shape.shape.dim[0].size + kernel_size_x = shape.shape.dim[1].size + return kernel_size_x, kernel_size_y + + +def _padding_size_conv_pool(node, kernel_size, stride, input_resolution=None): + """Computes padding size given a TF convolution or pooling node. + + Args: + node: Tensorflow node (NodeDef proto). + kernel_size: Kernel size of node (integer). + stride: Stride size of node (integer). + input_resolution: Input resolution to assume, if not None (integer). + + Returns: + total_padding: Total padding size (integer). + padding: Padding size, applied to the left or top (integer). + + Raises: + ValueError: If padding is invalid. + """ + # In this case, we need to carefully consider the different TF padding modes. + # The padding depends on kernel size, and may depend on input size. If it + # depends on input size and input_resolution is None, we raise an exception. + padding_attr = node.attr["padding"] + logging.vlog(4, "padding_attr = %s", padding_attr) + if padding_attr.s in _VALID_PADDING: + total_padding = 0 + padding = 0 + elif padding_attr.s in _SAME_PADDING: + if input_resolution is None: + # In this case, we do not know the input resolution, so we can only know + # the padding in some special cases. + if kernel_size == 1: + total_padding = 0 + padding = 0 + elif stride == 1: + total_padding = kernel_size - 1 + padding = int(math.floor(float(total_padding) / 2)) + elif stride == 2 and kernel_size % 2 == 0: + # In this case, we can be sure of the left/top padding, but not of the + # total padding. + total_padding = None + padding = int(math.floor((float(kernel_size) - 1) / 2)) + else: + total_padding = None + padding = None + logging.warning( + "Padding depends on input size, which means that the effective " + "padding may be different depending on the input image " + "dimensionality. In this case, alignment check will be skipped. If" + " you know the input resolution, please set it.") + else: + # First, compute total_padding based on documentation. + if input_resolution % stride == 0: + total_padding = int(max(float(kernel_size - stride), 0.0)) + else: + total_padding = int( + max(float(kernel_size - (input_resolution % stride)), 0.0)) + # Then, compute left/top padding. + padding = int(math.floor(float(total_padding) / 2)) + + else: + raise ValueError("Invalid padding operation %s" % padding_attr.s) + return total_padding, padding + + +def _pool_kernel_size(node, name_to_node): + """Computes kernel size given a TF pooling node. + + Args: + node: Tensorflow node (NodeDef proto). + name_to_node: For MaxPoolV2, mapping from node name to NodeDef. + + Returns: + kernel_size_x: Kernel size for horizontal direction (integer). + kernel_size_y: Kernel size for vertical direction (integer). + + Raises: + ValueError: If pooling is invalid. + """ + if node.op == "MaxPoolV2": + ksize_input_name = node.input[1] + if not ksize_input_name.endswith("/ksize"): + raise ValueError("Kernel size name does not end with '/ksize'") + ksize_node = name_to_node[ksize_input_name] + value = ksize_node.attr["value"] + t = make_ndarray(value.tensor) + kernel_size_y = t[1] + kernel_size_x = t[2] + if t[0] != 1: + raise ValueError("pool ksize for first dim is not 1") + if t[3] != 1: + raise ValueError("pool ksize for last dim is not 1") + else: + ksize = node.attr["ksize"] + kernel_size_y = ksize.list.i[1] + kernel_size_x = ksize.list.i[2] + if ksize.list.i[0] != 1: + raise ValueError("pool ksize for first dim is not 1") + if ksize.list.i[3] != 1: + raise ValueError("pool ksize for last dim is not 1") + return kernel_size_x, kernel_size_y + + +def _padding_size_pad_layer(node, name_to_node): + """Computes padding size given a TF padding node. + + Args: + node: Tensorflow node (NodeDef proto). + name_to_node: Dict keyed by node name, each entry containing the node's + NodeDef. + + Returns: + total_padding_x: Total padding size for horizontal direction (integer). + padding_x: Padding size for horizontal direction, left side (integer). + total_padding_y: Total padding size for vertical direction (integer). + padding_y: Padding size for vertical direction, top side (integer). + + Raises: + ValueError: If padding layer is invalid. + """ + paddings_layer_name = node.input[1] + if not paddings_layer_name.endswith("/paddings"): + raise ValueError("Padding layer name does not end with '/paddings'") + paddings_node = name_to_node[paddings_layer_name] + if paddings_node.op != "Const": + raise ValueError("Padding op is not Const") + value = paddings_node.attr["value"] + t = make_ndarray(value.tensor) + padding_y = t[1][0] + padding_x = t[2][0] + total_padding_y = padding_y + t[1][1] + total_padding_x = padding_x + t[2][1] + if (t[0][0] != 0) or (t[0][1] != 0): + raise ValueError("padding is not zero for first tensor dim") + if (t[3][0] != 0) or (t[3][1] != 0): + raise ValueError("padding is not zero for last tensor dim") + return total_padding_x, padding_x, total_padding_y, padding_y + + +def get_layer_params(node, name_to_node, input_resolution=None, force=False): + """Gets layer parameters relevant for RF computation. + + Currently, only these nodes are supported: + - Conv2D + - DepthwiseConv2dNative + - Pad + - MaxPool + - AvgPool + - all nodes listed in _UNCHANGED_RF_LAYER_OPS + + Args: + node: Tensorflow node (NodeDef proto). + name_to_node: Dict keyed by node name, each entry containing the node's + NodeDef. + input_resolution: List with 2 dimensions, denoting the height/width of the + input feature map to this layer. If set to None, then the padding may be + undefined (in tensorflow, SAME padding depends on input spatial + resolution). + force: If True, the function does not raise a ValueError if the layer op is + unknown. Instead, in this case it sets each of the returned parameters to + None. + + Returns: + kernel_size_x: Kernel size for horizontal direction (integer). + kernel_size_y: Kernel size for vertical direction (integer). + stride_x: Stride size for horizontal direction (integer). + stride_y: Stride size for vertical direction (integer). + padding_x: Padding size for horizontal direction, left side (integer). + padding_y: Padding size for vertical direction, top side (integer). + total_padding_x: Total padding size for horizontal direction (integer). + total_padding_y: Total padding size for vertical direction (integer). + + Raises: + ValueError: If layer op is unknown and force is False. + """ + logging.vlog(3, "node.name = %s", node.name) + logging.vlog(3, "node.op = %s", node.op) + logging.vlog(4, "node = %s", node) + if node.op == "Conv2D" or node.op == "DepthwiseConv2dNative": + stride_x, stride_y = _stride_size(node, name_to_node) + kernel_size_x, kernel_size_y = _conv_kernel_size(node, name_to_node) + # Compute the padding for this node separately for each direction. + total_padding_x, padding_x = _padding_size_conv_pool( + node, kernel_size_x, stride_x, input_resolution[1] + if input_resolution is not None else None) + total_padding_y, padding_y = _padding_size_conv_pool( + node, kernel_size_y, stride_y, input_resolution[0] + if input_resolution is not None else None) + elif node.op == "Pad": + # Kernel and stride are simply 1 in this case. + kernel_size_x = 1 + kernel_size_y = 1 + stride_x = 1 + stride_y = 1 + total_padding_x, padding_x, total_padding_y, padding_y = ( + _padding_size_pad_layer(node, name_to_node)) + elif node.op == "MaxPool" or node.op == "MaxPoolV2" or node.op == "AvgPool": + stride_x, stride_y = _stride_size(node, name_to_node) + kernel_size_x, kernel_size_y = _pool_kernel_size(node, name_to_node) + # Compute the padding for this node separately for each direction. + total_padding_x, padding_x = _padding_size_conv_pool( + node, kernel_size_x, stride_x, input_resolution[1] + if input_resolution is not None else None) + total_padding_y, padding_y = _padding_size_conv_pool( + node, kernel_size_y, stride_y, input_resolution[0] + if input_resolution is not None else None) + elif node.op in _UNCHANGED_RF_LAYER_OPS: + # These nodes do not modify the RF parameters. + kernel_size_x = 1 + kernel_size_y = 1 + stride_x = 1 + stride_y = 1 + total_padding_x = 0 + padding_x = 0 + total_padding_y = 0 + padding_y = 0 + else: + if force: + kernel_size_x = None + kernel_size_y = None + stride_x = None + stride_y = None + total_padding_x = None + padding_x = None + total_padding_y = None + padding_y = None + else: + raise ValueError("Unknown layer for operation '%s': %s" % (node.name, + node.op)) + return (kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, + padding_y, total_padding_x, total_padding_y) diff --git a/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters_test.py b/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters_test.py new file mode 100644 index 0000000000000000000000000000000000000000..369758a28402f7c9e55cca9c6f9ffa9182c91140 --- /dev/null +++ b/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters_test.py @@ -0,0 +1,149 @@ +# Copyright 2017 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 parse_layer_parameters module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib import slim +from tensorflow.contrib.receptive_field.python.util import graph_compute_order +from tensorflow.contrib.receptive_field.python.util import parse_layer_parameters +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_math_ops +from tensorflow.python.ops import nn +from tensorflow.python.platform import test + + +def create_test_network(): + """Convolutional neural network for test. + + Returns: + name_to_node: Dict keyed by node name, each entry containing the node's + NodeDef. + """ + g = ops.Graph() + with g.as_default(): + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') + # Left branch before first addition. + l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='VALID') + # Right branch before first addition. + l2_pad = array_ops.pad(x, [[0, 0], [1, 0], [1, 0], [0, 0]], name='L2_pad') + l2 = slim.conv2d(l2_pad, 1, [3, 3], stride=2, scope='L2', padding='VALID') + l3 = slim.max_pool2d(l2, [3, 3], stride=2, scope='L3', padding='SAME') + # First addition. + l4 = nn.relu(l1 + l3, name='L4_relu') + # Left branch after first addition. + l5 = slim.conv2d(l4, 1, [1, 1], stride=2, scope='L5', padding='SAME') + # Right branch after first addition. + l6 = slim.conv2d(l4, 1, [3, 3], stride=2, scope='L6', padding='SAME') + # Final addition. + gen_math_ops.add(l5, l6, name='L7_add') + + name_to_node = graph_compute_order.parse_graph_nodes(g.as_graph_def()) + return name_to_node + + +class ParseLayerParametersTest(test.TestCase): + + def testParametersAreParsedCorrectly(self): + """Checks parameters from create_test_network() are parsed correctly.""" + name_to_node = create_test_network() + + # L1. + l1_node_name = 'L1/Conv2D' + l1_params = parse_layer_parameters.get_layer_params( + name_to_node[l1_node_name], name_to_node) + expected_l1_params = (1, 1, 4, 4, 0, 0, 0, 0) + self.assertEqual(l1_params, expected_l1_params) + + # L2 padding. + l2_pad_name = 'L2_pad' + l2_pad_params = parse_layer_parameters.get_layer_params( + name_to_node[l2_pad_name], name_to_node) + expected_l2_pad_params = (1, 1, 1, 1, 1, 1, 1, 1) + self.assertEqual(l2_pad_params, expected_l2_pad_params) + + # L2. + l2_node_name = 'L2/Conv2D' + l2_params = parse_layer_parameters.get_layer_params( + name_to_node[l2_node_name], name_to_node) + expected_l2_params = (3, 3, 2, 2, 0, 0, 0, 0) + self.assertEqual(l2_params, expected_l2_params) + + # L3. + l3_node_name = 'L3/MaxPool' + # - Without knowing input size. + l3_params = parse_layer_parameters.get_layer_params( + name_to_node[l3_node_name], name_to_node) + expected_l3_params = (3, 3, 2, 2, None, None, None, None) + self.assertEqual(l3_params, expected_l3_params) + # - Input size is even. + l3_even_params = parse_layer_parameters.get_layer_params( + name_to_node[l3_node_name], name_to_node, input_resolution=[4, 4]) + expected_l3_even_params = (3, 3, 2, 2, 0, 0, 1, 1) + self.assertEqual(l3_even_params, expected_l3_even_params) + # - Input size is odd. + l3_odd_params = parse_layer_parameters.get_layer_params( + name_to_node[l3_node_name], name_to_node, input_resolution=[5, 5]) + expected_l3_odd_params = (3, 3, 2, 2, 1, 1, 2, 2) + self.assertEqual(l3_odd_params, expected_l3_odd_params) + + # L4. + l4_node_name = 'L4_relu' + l4_params = parse_layer_parameters.get_layer_params( + name_to_node[l4_node_name], name_to_node) + expected_l4_params = (1, 1, 1, 1, 0, 0, 0, 0) + self.assertEqual(l4_params, expected_l4_params) + + # L5. + l5_node_name = 'L5/Conv2D' + l5_params = parse_layer_parameters.get_layer_params( + name_to_node[l5_node_name], name_to_node) + expected_l5_params = (1, 1, 2, 2, 0, 0, 0, 0) + self.assertEqual(l5_params, expected_l5_params) + + # L6. + l6_node_name = 'L6/Conv2D' + # - Without knowing input size. + l6_params = parse_layer_parameters.get_layer_params( + name_to_node[l6_node_name], name_to_node) + expected_l6_params = (3, 3, 2, 2, None, None, None, None) + self.assertEqual(l6_params, expected_l6_params) + # - Input size is even. + l6_even_params = parse_layer_parameters.get_layer_params( + name_to_node[l6_node_name], name_to_node, input_resolution=[4, 4]) + expected_l6_even_params = (3, 3, 2, 2, 0, 0, 1, 1) + self.assertEqual(l6_even_params, expected_l6_even_params) + # - Input size is odd. + l6_odd_params = parse_layer_parameters.get_layer_params( + name_to_node[l6_node_name], name_to_node, input_resolution=[5, 5]) + expected_l6_odd_params = (3, 3, 2, 2, 1, 1, 2, 2) + self.assertEqual(l6_odd_params, expected_l6_odd_params) + + # L7. + l7_node_name = 'L7_add' + l7_params = parse_layer_parameters.get_layer_params( + name_to_node[l7_node_name], name_to_node) + expected_l7_params = (1, 1, 1, 1, 0, 0, 0, 0) + self.assertEqual(l7_params, expected_l7_params) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/receptive_field/python/util/receptive_field.py b/tensorflow/contrib/receptive_field/python/util/receptive_field.py index 8b34465d21d14508c24056b588f2533d8fea6a1d..b9bd2f09761ab10a62d37e8e2580b93b9b8a4453 100644 --- a/tensorflow/contrib/receptive_field/python/util/receptive_field.py +++ b/tensorflow/contrib/receptive_field/python/util/receptive_field.py @@ -23,242 +23,11 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import math +import numpy as np from tensorflow.contrib.receptive_field.python.util import graph_compute_order -from tensorflow.contrib.util import make_ndarray -from tensorflow.python.platform import tf_logging as logging +from tensorflow.contrib.receptive_field.python.util import parse_layer_parameters from tensorflow.python.framework import ops as framework_ops -import numpy as np - -# White-listed layer operations, which do not affect the receptive field -# computation. -_UNCHANGED_RF_LAYER_OPS = [ - 'Add', 'BiasAdd', 'Ceil', 'ConcatV2', 'Const', 'Floor', 'Identity', 'Log', - 'Mul', 'Pow', 'RealDiv', 'Relu', 'Round', 'Rsqrt', 'Softplus', 'Sub', - 'VariableV2'] - -# Different ways in which padding modes may be spelled. -_VALID_PADDING = ["VALID", b"VALID"] -_SAME_PADDING = ["SAME", b"SAME"] - - -def _stride_size(node): - """Computes stride size given a TF node. - - Args: - node: Tensorflow node (NodeDef proto). - - Returns: - stride_x: Stride size for horizontal direction (integer). - stride_y: Stride size for vertical direction (integer). - """ - strides_attr = node.attr["strides"] - logging.vlog(4, "strides_attr = %s", strides_attr) - stride_y = strides_attr.list.i[1] - stride_x = strides_attr.list.i[2] - return stride_x, stride_y - - -def _conv_kernel_size(node, name_to_order_node): - """Computes kernel size given a TF convolution or pooling node. - - Args: - node: Tensorflow node (NodeDef proto). - name_to_order_node: Map from name to {order, node}. Output of - graph_compute_order.get_compute_order(). - - Returns: - kernel_size_x: Kernel size for horizontal direction (integer). - kernel_size_y: Kernel size for vertical direction (integer). - - Raises: - ValueError: If the weight layer node is invalid. - """ - weights_layer_read_name = node.input[1] - if not weights_layer_read_name.endswith("/read"): - raise ValueError( - "Weight layer's name input to conv layer does not end with '/read'") - weights_layer_param_name = weights_layer_read_name[:-5] - weights_node = name_to_order_node[weights_layer_param_name].node - if weights_node.op != "VariableV2": - raise ValueError("Weight layer is not of type VariableV2") - shape = weights_node.attr["shape"] - logging.vlog(4, "weight shape = %s", shape) - kernel_size_y = shape.shape.dim[0].size - kernel_size_x = shape.shape.dim[1].size - return kernel_size_x, kernel_size_y - - -def _padding_size_conv_pool(node, kernel_size, stride): - """Computes padding size given a TF convolution or pooling node. - - Args: - node: Tensorflow node (NodeDef proto). - kernel_size: Kernel size of node (integer). - stride: Stride size of node (integer). - - Returns: - padding: Padding size (integer). - - Raises: - ValueError: If padding is invalid. - """ - # In this case, we need to carefully consider the different TF padding modes. - # The padding depends on kernel size, and may depend on input size. If it - # depends on input size, we raise an exception. - padding_attr = node.attr["padding"] - logging.vlog(4, "padding_attr = %s", padding_attr) - if padding_attr.s in _VALID_PADDING: - padding = 0 - elif padding_attr.s in _SAME_PADDING: - if kernel_size == 1: - padding = 0 - elif stride == 1: - padding = int(math.floor((float(kernel_size) - 1) / 2)) - elif stride == 2 and kernel_size % 2 == 0: - padding = int(math.floor((float(kernel_size) - 1) / 2)) - else: - padding = None - logging.warning( - "Padding depends on input size, which means that the effective " - "padding may be different depending on the input image " - "dimensionality. In this case, alignment check will be skipped.") - else: - raise ValueError("Invalid padding operation %s" % padding_attr.s) - return padding - - -def _pool_kernel_size(node): - """Computes kernel size given a TF pooling node. - - Args: - node: Tensorflow node (NodeDef proto). - - Returns: - kernel_size_x: Kernel size for horizontal direction (integer). - kernel_size_y: Kernel size for vertical direction (integer). - - Raises: - ValueError: If pooling is invalid. - """ - ksize = node.attr["ksize"] - kernel_size_y = ksize.list.i[1] - kernel_size_x = ksize.list.i[2] - if ksize.list.i[0] != 1: - raise ValueError("pool ksize for first dim is not 1") - if ksize.list.i[3] != 1: - raise ValueError("pool ksize for last dim is not 1") - return kernel_size_x, kernel_size_y - - -def _padding_size_pad_layer(node, name_to_order_node): - """Computes padding size given a TF padding node. - - Args: - node: Tensorflow node (NodeDef proto). - name_to_order_node: Map from name to {order, node}. Output of - graph_compute_order.get_compute_order(). - - Returns: - padding_x: Padding size for horizontal direction (integer). - padding_y: Padding size for vertical direction (integer). - - Raises: - ValueError: If padding layer is invalid. - """ - paddings_layer_name = node.input[1] - if not paddings_layer_name.endswith("/paddings"): - raise ValueError("Padding layer name does not end with '/paddings'") - paddings_node = name_to_order_node[paddings_layer_name].node - if paddings_node.op != "Const": - raise ValueError("Padding op is not Const") - value = paddings_node.attr["value"] - t = make_ndarray(value.tensor) - padding_y = t[1][0] - padding_x = t[2][0] - if t[0][0] != 0: - raise ValueError("padding is not zero for first tensor dim") - if t[3][0] != 0: - raise ValueError("padding is not zero for last tensor dim") - return padding_x, padding_y - - -def _get_layer_params(node, name_to_order_node): - """Gets layer parameters relevant for RF computation. - - Currently, only these nodes are supported: - - Conv2D - - DepthwiseConv2dNative - - Pad - - MaxPool - - AvgPool - - all nodes listed in _UNCHANGED_RF_LAYER_OPS - - Args: - node: Tensorflow node (NodeDef proto). - name_to_order_node: Map from name to {order, node}. Output of - graph_compute_order.get_compute_order(). - - Returns: - kernel_size_x: Kernel size for horizontal direction (integer). - kernel_size_y: Kernel size for vertical direction (integer). - stride_x: Stride size for horizontal direction (integer). - stride_y: Stride size for vertical direction (integer). - padding_x: Padding size for horizontal direction (integer). - padding_y: Padding size for vertical direction (integer). - - Raises: - ValueError: If layer op is unknown. - """ - logging.vlog(3, "node.op = %s", node.op) - logging.vlog(4, "node = %s", node) - if node.op == "Conv2D" or node.op == "DepthwiseConv2dNative": - stride_x, stride_y = _stride_size(node) - kernel_size_x, kernel_size_y = _conv_kernel_size(node, name_to_order_node) - # Compute the padding for this node separately for each direction. - padding_x = _padding_size_conv_pool(node, kernel_size_x, stride_x) - padding_y = _padding_size_conv_pool(node, kernel_size_y, stride_y) - elif node.op == "Pad": - # Kernel and stride are simply 1 in this case. - kernel_size_x = 1 - kernel_size_y = 1 - stride_x = 1 - stride_y = 1 - padding_x, padding_y = _padding_size_pad_layer(node, name_to_order_node) - elif node.op == "MaxPool" or node.op == "AvgPool": - stride_x, stride_y = _stride_size(node) - kernel_size_x, kernel_size_y = _pool_kernel_size(node) - # Compute the padding for this node separately for each direction. - padding_x = _padding_size_conv_pool(node, kernel_size_x, stride_x) - padding_y = _padding_size_conv_pool(node, kernel_size_y, stride_y) - elif node.op in _UNCHANGED_RF_LAYER_OPS: - # These nodes do not modify the RF parameters. - kernel_size_x = 1 - kernel_size_y = 1 - stride_x = 1 - stride_y = 1 - padding_x = 0 - padding_y = 0 - else: - raise ValueError("Unknown layer for operation '%s': %s" % - (node.name, node.op)) - return kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, padding_y - - -def _reverse_sort_by_order(name_to_order_node): - """Sorts map of name_to_order_node nodes in reverse order. - - The output is such that the nodes in name_to_order_node are sorted in - descending order of the "order" field. - - Args: - name_to_order_node: Map from name to {order, node}. Output of - graph_compute_order.get_compute_order(). - - Returns: - sorted_name_to_order_node: Sorted version of the input, in descending order. - """ - return sorted(name_to_order_node.items(), key=lambda x: -x[1].order) +from tensorflow.python.platform import tf_logging as logging def _get_rf_size_node_input(stride, kernel_size, rf_size_output): @@ -307,23 +76,22 @@ def _get_effective_padding_node_input(stride, padding, return stride * effective_padding_output + padding -class ReceptiveField: - """ - Receptive field of a convolutional neural network. +class ReceptiveField(object): + """Receptive field of a convolutional neural network. Args: size: Receptive field size. stride: Effective stride. padding: Effective padding. """ + def __init__(self, size, stride, padding): self.size = np.asarray(size) self.stride = np.asarray(stride) self.padding = np.asarray(padding) def compute_input_center_coordinates(self, y, axis=None): - """ - Computes the center of the receptive field that generated a feature. + """Computes the center of the receptive field that generated a feature. Args: y: An array of feature coordinates with shape `(..., d)`, where `d` is the @@ -350,12 +118,11 @@ class ReceptiveField: raise ValueError("Dimensionality of the feature coordinates `y` (%d) " "does not match dimensionality of `axis` (%d)" % (y.shape[-1], len(axis))) - return - self.padding[axis] + y * self.stride[axis] + \ - (self.size[axis] - 1) / 2 + return -self.padding[axis] + y * self.stride[axis] + ( + self.size[axis] - 1) / 2 def compute_feature_coordinates(self, x, axis=None): - """ - Computes the position of a feature given the center of a receptive field. + """Computes the position of a feature given the center of a receptive field. Args: x: An array of input center coordinates with shape `(..., d)`, where `d` @@ -381,15 +148,18 @@ class ReceptiveField: raise ValueError("Dimensionality of the input center coordinates `x` " "(%d) does not match dimensionality of `axis` (%d)" % (x.shape[-1], len(axis))) - return (x + self.padding[axis] + (1 - self.size[axis]) / 2) / \ - self.stride[axis] + return (x + self.padding[axis] + + (1 - self.size[axis]) / 2) / self.stride[axis] def __iter__(self): return iter(np.concatenate([self.size, self.stride, self.padding])) -def compute_receptive_field_from_graph_def(graph_def, input_node, output_node, - stop_propagation=None): +def compute_receptive_field_from_graph_def(graph_def, + input_node, + output_node, + stop_propagation=None, + input_resolution=None): """Computes receptive field (RF) parameters from a Graph or GraphDef object. The algorithm stops the calculation of the receptive field whenever it @@ -402,8 +172,14 @@ def compute_receptive_field_from_graph_def(graph_def, input_node, output_node, graph_def: Graph or GraphDef object. input_node: Name of the input node or Tensor object from graph. output_node: Name of the output node or Tensor object from graph. - stop_propagation: List of operation or scope names for which to stop the + stop_propagation: List of operations or scope names for which to stop the propagation of the receptive field. + input_resolution: 2D list. If the input resolution to the model is fixed and + known, this may be set. This is helpful for cases where the RF parameters + vary depending on the input resolution (this happens since SAME padding in + tensorflow depends on input resolution in general). If this is None, it is + assumed that the input resolution is unknown, so some RF parameters may be + unknown (depending on the model architecture). Returns: rf_size_x: Receptive field size of network in the horizontal direction, with @@ -437,11 +213,13 @@ def compute_receptive_field_from_graph_def(graph_def, input_node, output_node, stop_propagation = stop_propagation or [] # Computes order of computation for a given graph. - name_to_order_node = graph_compute_order.get_compute_order( - graph_def=graph_def) + node_info, name_to_node = graph_compute_order.get_compute_order( + graph_def=graph_def, + input_node_name=input_node, + input_node_size=input_resolution) # Sort in reverse topological order. - order = _reverse_sort_by_order(name_to_order_node) + ordered_node_info = sorted(node_info.items(), key=lambda x: -x[1].order) # Dictionaries to keep track of receptive field, effective stride and # effective padding of different nodes. @@ -470,7 +248,7 @@ def compute_receptive_field_from_graph_def(graph_def, input_node, output_node, # alignment checks are skipped, and the effective padding is None. undefined_padding = False - for _, (o, node) in order: + for _, (o, node, _, _) in ordered_node_info: if node: logging.vlog(3, "%10d %-100s %-20s" % (o, node.name[:90], node.op)) else: @@ -496,13 +274,14 @@ def compute_receptive_field_from_graph_def(graph_def, input_node, output_node, continue # Get params for this layer. - kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, padding_y = ( - _get_layer_params(node, name_to_order_node)) + (kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, + padding_y, _, _) = parse_layer_parameters.get_layer_params( + node, name_to_node, node_info[node.name].input_size) logging.vlog(3, "kernel_size_x = %s, kernel_size_y = %s, " "stride_x = %s, stride_y = %s, " - "padding_x = %s, padding_y = %s" % + "padding_x = %s, padding_y = %s, input size = %s" % (kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, - padding_y)) + padding_y, node_info[node.name].input_size)) if padding_x is None or padding_y is None: undefined_padding = True @@ -524,72 +303,93 @@ def compute_receptive_field_from_graph_def(graph_def, input_node, output_node, else: effective_padding_input_x = None effective_padding_input_y = None + logging.vlog( + 4, "rf_size_input_x = %s, rf_size_input_y = %s, " + "effective_stride_input_x = %s, effective_stride_input_y = %s, " + "effective_padding_input_x = %s, effective_padding_input_y = %s" % + (rf_size_input_x, rf_size_input_y, effective_stride_input_x, + effective_stride_input_y, effective_padding_input_x, + effective_padding_input_y)) # Loop over this node's inputs and potentially propagate information down. for inp_name in node.input: # Stop the propagation of the receptive field. if any(inp_name.startswith(stop) for stop in stop_propagation): - logging.vlog(3, "Skipping explicitly ignored node %s.", node.name) + logging.vlog(3, "Skipping explicitly ignored node %s.", inp_name) continue + logging.vlog(4, "inp_name = %s", inp_name) - inp_node = name_to_order_node[inp_name].node + if inp_name.startswith("^"): + # The character "^" denotes a control dependency, so this input node + # can be safely ignored. + continue + + inp_node = name_to_node[inp_name] logging.vlog(4, "inp_node = \n%s", inp_node) - if inp_node.name in rf_sizes_x: - assert inp_node.name in rf_sizes_y, ( - "Node %s is in rf_sizes_x, but " - "not in rf_sizes_y" % inp_node.name) + if inp_name in rf_sizes_x: + assert inp_name in rf_sizes_y, ("Node %s is in rf_sizes_x, but " + "not in rf_sizes_y" % inp_name) + logging.vlog( + 4, "rf_sizes_x[inp_name] = %s," + " rf_sizes_y[inp_name] = %s, " + "effective_strides_x[inp_name] = %s," + " effective_strides_y[inp_name] = %s, " + "effective_paddings_x[inp_name] = %s," + " effective_paddings_y[inp_name] = %s" % + (rf_sizes_x[inp_name], rf_sizes_y[inp_name], + effective_strides_x[inp_name], effective_strides_y[inp_name], + effective_paddings_x[inp_name], effective_paddings_y[inp_name])) # This node was already discovered through a previous path, so we need # to make sure that graph is aligned. This alignment check is skipped # if the padding is not defined, since in this case alignment cannot # be checked. if not undefined_padding: - if effective_strides_x[inp_node.name] != effective_stride_input_x: + if effective_strides_x[inp_name] != effective_stride_input_x: raise ValueError( "Graph is not aligned since effective stride from different " "paths is different in horizontal direction") - if effective_strides_y[inp_node.name] != effective_stride_input_y: + if effective_strides_y[inp_name] != effective_stride_input_y: raise ValueError( "Graph is not aligned since effective stride from different " "paths is different in vertical direction") - if (rf_sizes_x[inp_node.name] - 1 - ) / 2 - effective_paddings_x[inp_node.name] != ( + if (rf_sizes_x[inp_name] - 1 + ) / 2 - effective_paddings_x[inp_name] != ( rf_size_input_x - 1) / 2 - effective_padding_input_x: raise ValueError( "Graph is not aligned since center shift from different " "paths is different in horizontal direction") - if (rf_sizes_y[inp_node.name] - 1 - ) / 2 - effective_paddings_y[inp_node.name] != ( + if (rf_sizes_y[inp_name] - 1 + ) / 2 - effective_paddings_y[inp_name] != ( rf_size_input_y - 1) / 2 - effective_padding_input_y: raise ValueError( "Graph is not aligned since center shift from different " "paths is different in vertical direction") # Keep track of path with largest RF, for both directions. - if rf_sizes_x[inp_node.name] < rf_size_input_x: - rf_sizes_x[inp_node.name] = rf_size_input_x - effective_strides_x[inp_node.name] = effective_stride_input_x - effective_paddings_x[inp_node.name] = effective_padding_input_x - if rf_sizes_y[inp_node.name] < rf_size_input_y: - rf_sizes_y[inp_node.name] = rf_size_input_y - effective_strides_y[inp_node.name] = effective_stride_input_y - effective_paddings_y[inp_node.name] = effective_padding_input_y + if rf_sizes_x[inp_name] < rf_size_input_x: + rf_sizes_x[inp_name] = rf_size_input_x + effective_strides_x[inp_name] = effective_stride_input_x + effective_paddings_x[inp_name] = effective_padding_input_x + if rf_sizes_y[inp_name] < rf_size_input_y: + rf_sizes_y[inp_name] = rf_size_input_y + effective_strides_y[inp_name] = effective_stride_input_y + effective_paddings_y[inp_name] = effective_padding_input_y else: - assert inp_node.name not in rf_sizes_y, ( - "Node %s is in rf_sizes_y, but " - "not in rf_sizes_x" % inp_node.name) + assert inp_name not in rf_sizes_y, ("Node %s is in rf_sizes_y, but " + "not in rf_sizes_x" % inp_name) # In this case, it is the first time we encounter this node. So we # propagate the RF parameters. - rf_sizes_x[inp_node.name] = rf_size_input_x - rf_sizes_y[inp_node.name] = rf_size_input_y - effective_strides_x[inp_node.name] = effective_stride_input_x - effective_strides_y[inp_node.name] = effective_stride_input_y - effective_paddings_x[inp_node.name] = effective_padding_input_x - effective_paddings_y[inp_node.name] = effective_padding_input_y + rf_sizes_x[inp_name] = rf_size_input_x + rf_sizes_y[inp_name] = rf_size_input_y + effective_strides_x[inp_name] = effective_stride_input_x + effective_strides_y[inp_name] = effective_stride_input_y + effective_paddings_x[inp_name] = effective_padding_input_x + effective_paddings_y[inp_name] = effective_padding_input_y if not found_output_node: raise ValueError("Output node was not found") if input_node not in rf_sizes_x: raise ValueError("Input node was not found") return ReceptiveField( - (rf_sizes_x[input_node], rf_sizes_y[input_node]), - (effective_strides_x[input_node], effective_strides_y[input_node]), - (effective_paddings_x[input_node], effective_paddings_y[input_node])) + (rf_sizes_x[input_node], rf_sizes_y[input_node]), + (effective_strides_x[input_node], effective_strides_y[input_node]), + (effective_paddings_x[input_node], effective_paddings_y[input_node])) diff --git a/tensorflow/contrib/receptive_field/python/util/receptive_field_test.py b/tensorflow/contrib/receptive_field/python/util/receptive_field_test.py index 8d7d5440f630a3a78749e04a5eb058d637c258fc..cf55da27236d17c709cbde689831ad68da9a8a7b 100644 --- a/tensorflow/contrib/receptive_field/python/util/receptive_field_test.py +++ b/tensorflow/contrib/receptive_field/python/util/receptive_field_test.py @@ -18,16 +18,21 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numpy as np + from tensorflow.contrib import slim -from tensorflow.contrib.receptive_field.python.util import receptive_field +from tensorflow.contrib.receptive_field import receptive_field_api as receptive_field from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import nn from tensorflow.python.platform import test -import numpy as np +# TODO(andrearaujo): Rename the create_test_network_* functions in order to have +# more descriptive names. def create_test_network_1(): """Aligned network for test. @@ -39,8 +44,9 @@ def create_test_network_1(): """ g = ops.Graph() with g.as_default(): - # An 8x8 test image. - x = array_ops.placeholder(dtypes.float32, (1, 8, 8, 1), name='input_image') + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') # Left branch. l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='VALID') # Right branch. @@ -66,8 +72,9 @@ def create_test_network_2(): """ g = ops.Graph() with g.as_default(): - # An 8x8 test image. - x = array_ops.placeholder(dtypes.float32, (1, 8, 8, 1), name='input_image') + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') # Left branch. l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='VALID') # Right branch. @@ -90,8 +97,9 @@ def create_test_network_3(): """ g = ops.Graph() with g.as_default(): - # An 8x8 test image. - x = array_ops.placeholder(dtypes.float32, (1, 8, 8, 1), name='input_image') + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') # Left branch. l1_pad = array_ops.pad(x, [[0, 0], [2, 1], [2, 1], [0, 0]]) l1 = slim.conv2d(l1_pad, 1, [5, 5], stride=2, scope='L1', padding='VALID') @@ -117,8 +125,9 @@ def create_test_network_4(): """ g = ops.Graph() with g.as_default(): - # An 8x8 test image. - x = array_ops.placeholder(dtypes.float32, (1, 8, 8, 1), name='input_image') + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') # Left branch. l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='VALID') # Right branch. @@ -141,8 +150,9 @@ def create_test_network_5(): """ g = ops.Graph() with g.as_default(): - # An 8x8 test image. - x = array_ops.placeholder(dtypes.float32, (1, 8, 8, 1), name='input_image') + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') # Two convolutional layers, where the first one has non-square kernel. l1 = slim.conv2d(x, 1, [3, 5], stride=2, scope='L1', padding='VALID') l2 = slim.conv2d(l1, 1, [3, 1], stride=2, scope='L2', padding='VALID') @@ -162,8 +172,9 @@ def create_test_network_6(): """ g = ops.Graph() with g.as_default(): - # An 8x8 test image. - x = array_ops.placeholder(dtypes.float32, (1, 8, 8, 1), name='input_image') + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') # Left branch. l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='VALID') # Right branch. @@ -176,7 +187,102 @@ def create_test_network_6(): return g -class RfUtilsTest(test.TestCase): +def create_test_network_7(): + """Aligned network for test, with a control dependency. + + The graph is similar to create_test_network_1(), except that it includes an + assert operation on the left branch. + + Returns: + g: Tensorflow graph object (Graph proto). + """ + g = ops.Graph() + with g.as_default(): + # An 8x8 test image. + x = array_ops.placeholder(dtypes.float32, (1, 8, 8, 1), name='input_image') + # Left branch. + l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='VALID') + l1_shape = array_ops.shape(l1) + assert_op = control_flow_ops.Assert( + gen_math_ops.equal(l1_shape[1], 2), [l1_shape], summarize=4) + # Right branch. + l2_pad = array_ops.pad(x, [[0, 0], [1, 0], [1, 0], [0, 0]]) + l2 = slim.conv2d(l2_pad, 1, [3, 3], stride=2, scope='L2', padding='VALID') + l3 = slim.conv2d(l2, 1, [1, 1], stride=2, scope='L3', padding='VALID') + # Addition. + with ops.control_dependencies([assert_op]): + nn.relu(l1 + l3, name='output') + return g + + +def create_test_network_8(): + """Aligned network for test, including an intermediate addition. + + The graph is similar to create_test_network_1(), except that it includes a few + more layers on top. The added layers compose two different branches whose + receptive fields are different. This makes this test case more challenging; in + particular, this test fails if a naive DFS-like algorithm is used for RF + computation. + + Returns: + g: Tensorflow graph object (Graph proto). + """ + g = ops.Graph() + with g.as_default(): + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') + # Left branch before first addition. + l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='VALID') + # Right branch before first addition. + l2_pad = array_ops.pad(x, [[0, 0], [1, 0], [1, 0], [0, 0]]) + l2 = slim.conv2d(l2_pad, 1, [3, 3], stride=2, scope='L2', padding='VALID') + l3 = slim.conv2d(l2, 1, [1, 1], stride=2, scope='L3', padding='VALID') + # First addition. + l4 = nn.relu(l1 + l3) + # Left branch after first addition. + l5 = slim.conv2d(l4, 1, [1, 1], stride=2, scope='L5', padding='VALID') + # Right branch after first addition. + l6_pad = array_ops.pad(l4, [[0, 0], [1, 0], [1, 0], [0, 0]]) + l6 = slim.conv2d(l6_pad, 1, [3, 3], stride=2, scope='L6', padding='VALID') + # Final addition. + nn.relu(l5 + l6, name='output') + + return g + + +def create_test_network_9(): + """Aligned network for test, including an intermediate addition. + + The graph is the same as create_test_network_8(), except that VALID padding is + changed to SAME. + + Returns: + g: Tensorflow graph object (Graph proto). + """ + g = ops.Graph() + with g.as_default(): + # An input test image with unknown spatial resolution. + x = array_ops.placeholder( + dtypes.float32, (None, None, None, 1), name='input_image') + # Left branch before first addition. + l1 = slim.conv2d(x, 1, [1, 1], stride=4, scope='L1', padding='SAME') + # Right branch before first addition. + l2 = slim.conv2d(x, 1, [3, 3], stride=2, scope='L2', padding='SAME') + l3 = slim.conv2d(l2, 1, [1, 1], stride=2, scope='L3', padding='SAME') + # First addition. + l4 = nn.relu(l1 + l3) + # Left branch after first addition. + l5 = slim.conv2d(l4, 1, [1, 1], stride=2, scope='L5', padding='SAME') + # Right branch after first addition. + l6 = slim.conv2d(l4, 1, [3, 3], stride=2, scope='L6', padding='SAME') + # Final addition. + nn.relu(l5 + l6, name='output') + + return g + + +class ReceptiveFieldTest(test.TestCase): def testComputeRFFromGraphDefAligned(self): graph_def = create_test_network_1().as_graph_def() @@ -216,7 +322,7 @@ class RfUtilsTest(test.TestCase): receptive_field.compute_receptive_field_from_graph_def( graph_def, input_node, output_node) - def testComputeRFFromGraphDefUnaligned2(self): + def testComputeRFFromGraphDefUndefinedPadding(self): graph_def = create_test_network_4().as_graph_def() input_node = 'input_image' output_node = 'output' @@ -231,6 +337,29 @@ class RfUtilsTest(test.TestCase): self.assertEqual(effective_padding_x, None) self.assertEqual(effective_padding_y, None) + def testComputeRFFromGraphDefFixedInputDim(self): + graph_def = create_test_network_4().as_graph_def() + input_node = 'input_image' + output_node = 'output' + (receptive_field_x, receptive_field_y, effective_stride_x, + effective_stride_y, effective_padding_x, effective_padding_y) = ( + receptive_field.compute_receptive_field_from_graph_def( + graph_def, input_node, output_node, input_resolution=[9, 9])) + self.assertEqual(receptive_field_x, 3) + self.assertEqual(receptive_field_y, 3) + self.assertEqual(effective_stride_x, 4) + self.assertEqual(effective_stride_y, 4) + self.assertEqual(effective_padding_x, 1) + self.assertEqual(effective_padding_y, 1) + + def testComputeRFFromGraphDefUnalignedFixedInputDim(self): + graph_def = create_test_network_4().as_graph_def() + input_node = 'input_image' + output_node = 'output' + with self.assertRaises(ValueError): + receptive_field.compute_receptive_field_from_graph_def( + graph_def, input_node, output_node, input_resolution=[8, 8]) + def testComputeRFFromGraphDefNonSquareRF(self): graph_def = create_test_network_5().as_graph_def() input_node = 'input_image' @@ -269,7 +398,7 @@ class RfUtilsTest(test.TestCase): input_node = 'input_image' output_node = 'output' rf = receptive_field.compute_receptive_field_from_graph_def( - graph_def, input_node, output_node) + graph_def, input_node, output_node) x = np.random.randint(0, 100, (50, 2)) y = rf.compute_feature_coordinates(x) @@ -277,5 +406,52 @@ class RfUtilsTest(test.TestCase): self.assertAllEqual(x, x2) + def testComputeRFFromGraphDefAlignedWithControlDependencies(self): + graph_def = create_test_network_7().as_graph_def() + input_node = 'input_image' + output_node = 'output' + (receptive_field_x, receptive_field_y, effective_stride_x, + effective_stride_y, effective_padding_x, effective_padding_y) = ( + receptive_field.compute_receptive_field_from_graph_def( + graph_def, input_node, output_node)) + self.assertEqual(receptive_field_x, 3) + self.assertEqual(receptive_field_y, 3) + self.assertEqual(effective_stride_x, 4) + self.assertEqual(effective_stride_y, 4) + self.assertEqual(effective_padding_x, 1) + self.assertEqual(effective_padding_y, 1) + + def testComputeRFFromGraphDefWithIntermediateAddNode(self): + graph_def = create_test_network_8().as_graph_def() + input_node = 'input_image' + output_node = 'output' + (receptive_field_x, receptive_field_y, effective_stride_x, + effective_stride_y, effective_padding_x, effective_padding_y) = ( + receptive_field.compute_receptive_field_from_graph_def( + graph_def, input_node, output_node)) + self.assertEqual(receptive_field_x, 11) + self.assertEqual(receptive_field_y, 11) + self.assertEqual(effective_stride_x, 8) + self.assertEqual(effective_stride_y, 8) + self.assertEqual(effective_padding_x, 5) + self.assertEqual(effective_padding_y, 5) + + def testComputeRFFromGraphDefWithIntermediateAddNodeSamePaddingFixedInputDim( + self): + graph_def = create_test_network_9().as_graph_def() + input_node = 'input_image' + output_node = 'output' + (receptive_field_x, receptive_field_y, effective_stride_x, + effective_stride_y, effective_padding_x, effective_padding_y) = ( + receptive_field.compute_receptive_field_from_graph_def( + graph_def, input_node, output_node, input_resolution=[17, 17])) + self.assertEqual(receptive_field_x, 11) + self.assertEqual(receptive_field_y, 11) + self.assertEqual(effective_stride_x, 8) + self.assertEqual(effective_stride_y, 8) + self.assertEqual(effective_padding_x, 5) + self.assertEqual(effective_padding_y, 5) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/receptive_field/__init__.py b/tensorflow/contrib/receptive_field/receptive_field_api.py similarity index 89% rename from tensorflow/contrib/receptive_field/__init__.py rename to tensorflow/contrib/receptive_field/receptive_field_api.py index 10745a6a53d5b3ef9521b2313ddc28799ee8b886..4d81b4292df5f696b761b2977fec078abc28569f 100644 --- a/tensorflow/contrib/receptive_field/__init__.py +++ b/tensorflow/contrib/receptive_field/receptive_field_api.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Module to compute receptive field parameters for CNN tensorflow models.""" +"""Module that declares the functions in tf.contrib.receptive_field's API.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function @@ -21,3 +21,7 @@ from __future__ import print_function from tensorflow.contrib.receptive_field.python.util.graph_compute_order import get_compute_order from tensorflow.contrib.receptive_field.python.util.receptive_field import compute_receptive_field_from_graph_def # pylint: enable=unused-import + +del absolute_import +del division +del print_function diff --git a/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.cc b/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.cc index c33804906fc21cf2573b79091a76ab1ea86f5966..2def4f3f176b8d4d26c2c94168e9698f14649d94 100644 --- a/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.cc +++ b/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.cc @@ -15,8 +15,8 @@ limitations under the License. #define EIGEN_USE_THREADS -#include #include "tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.h" +#include #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" diff --git a/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.h b/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.h index fc3a2da9b398b16df223d60e2e913f952fa24434..d8c0a0631d38e55ef9653e0e88e90604ec0f0329 100644 --- a/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.h +++ b/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.h @@ -13,13 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CORE_KERNELS_PARTIAL_REDUCTION_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CORE_KERNELS_PARTIAL_REDUCTION_OPS_H_ +#ifndef TENSORFLOW_CORE_KERNELS_PARTIAL_REDUCTION_OPS_H_ +#define TENSORFLOW_CORE_KERNELS_PARTIAL_REDUCTION_OPS_H_ +#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_types.h" -#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #define Sum(a, b) ((a) + (b)) #define Prod(a, b) ((a) * (b)) @@ -58,11 +58,11 @@ inline T negative_infinity() { } // namespace reduce_functions -#define CALL_ALL_REDUCEOPS(func, ...) \ - func(Sum, functor::reduce_functions::zero, ##__VA_ARGS__) \ - func(Prod, functor::reduce_functions::one, ##__VA_ARGS__) \ - func(Max, functor::reduce_functions::negative_infinity, ##__VA_ARGS__) \ - func(Min, functor::reduce_functions::infinity, ##__VA_ARGS__) +#define CALL_ALL_REDUCEOPS(func, ...) \ + func(Sum, functor::reduce_functions::zero, ##__VA_ARGS__) \ + func(Prod, functor::reduce_functions::one, ##__VA_ARGS__) func( \ + Max, functor::reduce_functions::negative_infinity, ##__VA_ARGS__) \ + func(Min, functor::reduce_functions::infinity, ##__VA_ARGS__) #define ReduceSliceFunctorReduceop(reduceop, dummy) \ template \ @@ -81,4 +81,4 @@ CALL_ALL_REDUCEOPS(ReduceSliceFunctorReduceop) } // namespace functor } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CORE_KERNELS_PARTIAL_REDUCTION_OPS_H_ +#endif // TENSORFLOW_CORE_KERNELS_PARTIAL_REDUCTION_OPS_H_ diff --git a/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops_gpu.cu.cc b/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops_gpu.cu.cc index 8e6870fadd428ae8a1937a5c0cb43b6763f6be28..9f2be03d718364058da6b63add8752c046798c5b 100644 --- a/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops_gpu.cu.cc +++ b/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops_gpu.cu.cc @@ -17,10 +17,10 @@ limitations under the License. #define EIGEN_USE_GPU +#include "tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" -#include "tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops.h" #include "tensorflow/core/util/cuda_kernel_helper.h" namespace tensorflow { @@ -34,9 +34,9 @@ namespace functor { __global__ void ReduceSliceDeviceKernel##reduceop( \ Cuda3DLaunchConfig config, Index indices_width, Index bound, \ const T begin, const Index *indices, const T *input, T *out) { \ - CUDA_AXIS_KERNEL_LOOP(x, config.virtual_thread_count, x) { \ - CUDA_AXIS_KERNEL_LOOP(y, config.virtual_thread_count, y) { \ - CUDA_AXIS_KERNEL_LOOP(z, config.virtual_thread_count, z) { \ + CUDA_AXIS_KERNEL_LOOP(x, config.virtual_thread_count.x, X) { \ + CUDA_AXIS_KERNEL_LOOP(y, config.virtual_thread_count.y, Y) { \ + CUDA_AXIS_KERNEL_LOOP(z, config.virtual_thread_count.z, Z) { \ Index outidx = x * config.virtual_thread_count.y * \ config.virtual_thread_count.z + \ y * config.virtual_thread_count.z + z; \ @@ -68,8 +68,9 @@ namespace functor { if (sizex * sizey * sizez == 0) { \ return; \ } \ - Cuda3DLaunchConfig config = GetCuda3DLaunchConfig(sizex, sizey, sizez, d,\ - ReduceSliceDeviceKernel##reduceop, 0, 0); \ + Cuda3DLaunchConfig config = GetCuda3DLaunchConfig( \ + sizex, sizey, sizez, d, ReduceSliceDeviceKernel##reduceop, \ + 0, 0); \ \ ReduceSliceDeviceKernel##reduceop \ <<>>( \ diff --git a/tensorflow/contrib/reduce_slice_ops/ops/reduce_slice_ops.cc b/tensorflow/contrib/reduce_slice_ops/ops/reduce_slice_ops.cc index b8b56c0e229563a4e9bc930512c9fe49bd636e31..92879ab5356623dfa82fce8dff8db4d3036ae46c 100644 --- a/tensorflow/contrib/reduce_slice_ops/ops/reduce_slice_ops.cc +++ b/tensorflow/contrib/reduce_slice_ops/ops/reduce_slice_ops.cc @@ -87,9 +87,9 @@ and 'indices' is [[0,1] [1,1] [0,2]], -the the output will be [[ 1, 2, 3] - [ 0, 0, 0] - [41,52,63]]. +the output will be [[ 1, 2, 3] + [ 0, 0, 0] + [41,52,63]]. ``` The data must be at least rank 1. The indices must be of shape (?,2) where the @@ -132,9 +132,9 @@ and 'indices' is [[0,1] [1,1] [0,2]], -the the output will be [[ 1, 2, 3] - [ 1, 1, 1] - [40,100,180]]. +the output will be [[ 1, 2, 3] + [ 1, 1, 1] + [40,100,180]]. ``` The data must be at least rank 1. The indices can be of shape (?,2) where the @@ -189,9 +189,9 @@ and 'indices' is [[0,1] [1,1] [0,2]], -the the output will be [[ 1, 20, 3] - [ -BIG_VALUE, -BIG_VALUE, -BIG_VALUE] - [ 400, 20, 60]]. +the output will be [[ 1, 20, 3] + [ -BIG_VALUE, -BIG_VALUE, -BIG_VALUE] + [ 400, 20, 60]]. ``` The data must be at least rank 1. The indices can be of shape (?,2) where the @@ -246,9 +246,9 @@ and 'indices' is [[0,1] [1,1] [0,2]], -the the output will be [[ 1, 20, 3] - [ +BIG_VALUE, +BIG_VALUE, +BIG_VALUE] - [ 1, 5, 3]]. +the output will be [[ 1, 20, 3] + [ +BIG_VALUE, +BIG_VALUE, +BIG_VALUE] + [ 1, 5, 3]]. ``` The data must be at least rank 1. The indices can be of shape (?,2) where the diff --git a/tensorflow/contrib/reduce_slice_ops/python/kernel_tests/reduce_slice_ops_test.py b/tensorflow/contrib/reduce_slice_ops/python/kernel_tests/reduce_slice_ops_test.py index 60a193db4c7f084d3262a69e2b8c5df66273e138..468886da20021646089bd1d222da1ebd4b5c7822 100644 --- a/tensorflow/contrib/reduce_slice_ops/python/kernel_tests/reduce_slice_ops_test.py +++ b/tensorflow/contrib/reduce_slice_ops/python/kernel_tests/reduce_slice_ops_test.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function import numpy as np -import unittest from tensorflow.contrib.reduce_slice_ops.python.ops import reduce_slice_ops from tensorflow.python.framework.test_util import TensorFlowTestCase diff --git a/tensorflow/contrib/resampler/kernels/resampler_ops.cc b/tensorflow/contrib/resampler/kernels/resampler_ops.cc index e02c1b6a2bd9daf9e1f81059f7c1f92106cebc8f..63c72836d793a3df4e96a0134f3a1534c288c8c8 100644 --- a/tensorflow/contrib/resampler/kernels/resampler_ops.cc +++ b/tensorflow/contrib/resampler/kernels/resampler_ops.cc @@ -36,17 +36,12 @@ using GPUDevice = Eigen::GpuDevice; namespace functor { template -struct Resampler2DFunctor{ - void operator ()(::tensorflow::OpKernelContext* ctx, - const CPUDevice& d, - const T* __restrict__ data, - const T* __restrict__ warp, - T* __restrict__ output, - const int batch_size, - const int data_height, - const int data_width, - const int data_channels, - const int num_sampling_points){ +struct Resampler2DFunctor { + void operator()(::tensorflow::OpKernelContext* ctx, const CPUDevice& d, + const T* __restrict__ data, const T* __restrict__ warp, + T* __restrict__ output, const int batch_size, + const int data_height, const int data_width, + const int data_channels, const int num_sampling_points) { const int warp_batch_stride = num_sampling_points * 2; const int data_batch_stride = data_height * data_width * data_channels; const int output_batch_stride = num_sampling_points * data_channels; @@ -59,24 +54,19 @@ struct Resampler2DFunctor{ // The functions take care of performing the relevant pointer // arithmetics abstracting away the low level details in the // main loop over samples. Note that data is stored in NHWC format. - auto set_output = [&](const int sample_id, - const int channel, + auto set_output = [&](const int sample_id, const int channel, const T value) { - output[batch_id * output_batch_stride + - sample_id * data_channels + + output[batch_id * output_batch_stride + sample_id * data_channels + channel] = value; }; - auto get_data_point = [&](const int x, - const int y, - const int chan) { + auto get_data_point = [&](const int x, const int y, const int chan) { const bool point_is_in_range = (x >= 0 && y >= 0 && x <= data_width - 1 && y <= data_height - 1); return point_is_in_range - ? data[batch_id * data_batch_stride + - data_channels * (y * data_width + x) + - chan] - : zero; + ? data[batch_id * data_batch_stride + + data_channels * (y * data_width + x) + chan] + : zero; }; for (int sample_id = 0; sample_id < num_sampling_points; ++sample_id) { @@ -89,8 +79,7 @@ struct Resampler2DFunctor{ // The effect is that the sampled signal smoothly goes to 0 outside // the original input domain, rather than presenting a jump // discontinuity at the image boundaries. - if (x > static_cast(-1.0) && - y > static_cast(-1.0) && + if (x > static_cast(-1.0) && y > static_cast(-1.0) && x < static_cast(data_width) && y < static_cast(data_height)) { // Precompute floor (f) and ceil (c) values for x and y. @@ -103,12 +92,10 @@ struct Resampler2DFunctor{ for (int chan = 0; chan < data_channels; ++chan) { const T img_fxfy = dx * dy * get_data_point(fx, fy, chan); - const T img_cxcy = (one - dx) * (one - dy) * - get_data_point(cx, cy, chan); - const T img_fxcy = dx * (one - dy) * - get_data_point(fx, cy, chan); - const T img_cxfy = (one - dx) * dy * - get_data_point(cx, fy, chan); + const T img_cxcy = + (one - dx) * (one - dy) * get_data_point(cx, cy, chan); + const T img_fxcy = dx * (one - dy) * get_data_point(fx, cy, chan); + const T img_cxfy = (one - dx) * dy * get_data_point(cx, fy, chan); set_output(sample_id, chan, img_fxfy + img_cxcy + img_fxcy + img_cxfy); } @@ -125,8 +112,8 @@ struct Resampler2DFunctor{ // estimate of the cost of each work unit is needed to correctly shard the // workload. Shard assumes each cost unit is 1ns, minimum cost per shard // being 10us. - const int64 cost = static_cast(num_sampling_points) * - data_channels * 1000; + const int64 cost = + static_cast(num_sampling_points) * data_channels * 1000; auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads()); ::tensorflow::Shard(worker_threads.num_threads, worker_threads.workers, batch_size, cost, resample_batches); @@ -138,8 +125,8 @@ struct Resampler2DFunctor{ template class ResamplerOp : public ::tensorflow::OpKernel { public: - explicit ResamplerOp(::tensorflow::OpKernelConstruction* context) : - ::tensorflow::OpKernel(context) {} + explicit ResamplerOp(::tensorflow::OpKernelConstruction* context) + : ::tensorflow::OpKernel(context) {} void Compute(::tensorflow::OpKernelContext* ctx) override { const ::tensorflow::Tensor& data = ctx->input(0); @@ -158,16 +145,17 @@ class ResamplerOp : public ::tensorflow::OpKernel { ::tensorflow::errors::InvalidArgument( "warp should be at least a matrix, got shape ", warp_shape.DebugString())); - OP_REQUIRES(ctx, warp_shape.dim_size(warp_shape.dims()-1) == 2, + OP_REQUIRES(ctx, warp_shape.dim_size(warp_shape.dims() - 1) == 2, ::tensorflow::errors::Unimplemented( "Only bilinear interpolation is supported, warping " "coordinates must be 2D; warp shape last entry should be " - "2, but shape vector is: ", warp_shape.DebugString())); + "2, but shape vector is: ", + warp_shape.DebugString())); OP_REQUIRES(ctx, data_shape.dim_size(0) == warp_shape.dim_size(0), ::tensorflow::errors::InvalidArgument( "Batch size of data and warp tensor must be the same, but " - "input shapes are: ", data_shape.DebugString(), ", ", - warp_shape.DebugString())); + "input shapes are: ", + data_shape.DebugString(), ", ", warp_shape.DebugString())); const int batch_size = data_shape.dim_size(0); const int data_height = data_shape.dim_size(1); const int data_width = data_shape.dim_size(2); @@ -180,16 +168,10 @@ class ResamplerOp : public ::tensorflow::OpKernel { // Execute kernel only for nonempty output; otherwise Eigen crashes on GPU. if (num_sampling_points > 0) { - functor::Resampler2DFunctor()(ctx, - ctx->eigen_device(), - data.flat().data(), - warp.flat().data(), - output->flat().data(), - batch_size, - data_height, - data_width, - data_channels, - num_sampling_points); + functor::Resampler2DFunctor()( + ctx, ctx->eigen_device(), data.flat().data(), + warp.flat().data(), output->flat().data(), batch_size, + data_height, data_width, data_channels, num_sampling_points); } } @@ -197,12 +179,9 @@ class ResamplerOp : public ::tensorflow::OpKernel { TF_DISALLOW_COPY_AND_ASSIGN(ResamplerOp); }; - -#define REGISTER(TYPE) \ - REGISTER_KERNEL_BUILDER( \ - Name("Resampler") \ - .Device(DEVICE_CPU) \ - .TypeConstraint("T"), \ +#define REGISTER(TYPE) \ + REGISTER_KERNEL_BUILDER( \ + Name("Resampler").Device(DEVICE_CPU).TypeConstraint("T"), \ ResamplerOp); TF_CALL_half(REGISTER); @@ -211,40 +190,32 @@ TF_CALL_double(REGISTER); #undef REGISTER #if GOOGLE_CUDA -#define REGISTER(TYPE) \ - REGISTER_KERNEL_BUILDER(Name("Resampler") \ - .Device(DEVICE_GPU) \ - .TypeConstraint("T"), \ - ResamplerOp) +#define REGISTER(TYPE) \ + REGISTER_KERNEL_BUILDER( \ + Name("Resampler").Device(DEVICE_GPU).TypeConstraint("T"), \ + ResamplerOp) TF_CALL_float(REGISTER); TF_CALL_double(REGISTER); #undef REGISTER #endif // GOOGLE_CUDA - namespace functor { template -struct ResamplerGrad2DFunctor{ - void operator ()(::tensorflow::OpKernelContext* ctx, - const CPUDevice& d, - const T* __restrict__ data, - const T* __restrict__ warp, - const T* __restrict__ grad_output, - T* __restrict__ grad_data, - T* __restrict__ grad_warp, - const int batch_size, - const int data_height, - const int data_width, - const int data_channels, - const int num_sampling_points){ +struct ResamplerGrad2DFunctor { + void operator()(::tensorflow::OpKernelContext* ctx, const CPUDevice& d, + const T* __restrict__ data, const T* __restrict__ warp, + const T* __restrict__ grad_output, T* __restrict__ grad_data, + T* __restrict__ grad_warp, const int batch_size, + const int data_height, const int data_width, + const int data_channels, const int num_sampling_points) { // Set gradients to 0, because the kernel incrementally updates the // tensor entries by adding partial contributions. - const int resampler_output_size = batch_size * num_sampling_points * - data_channels; + const int resampler_output_size = + batch_size * num_sampling_points * data_channels; const int grad_warp_size = resampler_output_size / data_channels * 2; - const int grad_data_size = data_height * data_width * data_channels * - batch_size; + const int grad_data_size = + data_height * data_width * data_channels * batch_size; memset(grad_data, 0, sizeof(T) * grad_data_size); memset(grad_warp, 0, sizeof(T) * grad_warp_size); @@ -260,35 +231,29 @@ struct ResamplerGrad2DFunctor{ // The functions take care of performing the relevant pointer // arithmetics abstracting away the low level details in the // main loop over samples. Note that data is stored in NHWC format. - auto get_data_point = [&](const int x, - const int y, - const int chan) { + auto get_data_point = [&](const int x, const int y, const int chan) { const bool point_is_in_range = - (x >= 0 && y >= 0 && x <= data_width - 1 && y <= data_height - 1); + (x >= 0 && y >= 0 && x <= data_width - 1 && y <= data_height - 1); return point_is_in_range - ? data[batch_id * data_batch_stride + - data_channels * (y * data_width + x) + - chan] - : zero; + ? data[batch_id * data_batch_stride + + data_channels * (y * data_width + x) + chan] + : zero; }; auto update_grad_data = [&](const int x, const int y, const int chan, const T value) { const bool point_is_in_range = (x >= 0 && y >= 0 && x <= data_width - 1 && y <= data_height - 1); - if (point_is_in_range){ + if (point_is_in_range) { grad_data[batch_id * data_batch_stride + - data_channels * (y * data_width + x) + - chan] += value; + data_channels * (y * data_width + x) + chan] += value; } }; - auto update_grad_warp = [&](const int sample_id, - const int channel, + auto update_grad_warp = [&](const int sample_id, const int channel, const T value) { - grad_warp[batch_id * warp_batch_stride + - sample_id * 2 + - channel] += value; + grad_warp[batch_id * warp_batch_stride + sample_id * 2 + channel] += + value; }; for (int sample_id = 0; sample_id < num_sampling_points; ++sample_id) { @@ -301,8 +266,7 @@ struct ResamplerGrad2DFunctor{ // The effect is that the sampled signal smoothly goes to 0 outside // the original input domain, rather than presenting a jump // discontinuity at the image boundaries. - if (x > static_cast(-1.0) && - y > static_cast(-1.0) && + if (x > static_cast(-1.0) && y > static_cast(-1.0) && x < static_cast(data_width) && y < static_cast(data_height)) { // Precompute floor (f) and ceil (c) values for x and y. @@ -316,27 +280,25 @@ struct ResamplerGrad2DFunctor{ for (int chan = 0; chan < data_channels; ++chan) { const T grad_output_value = grad_output[batch_id * output_batch_stride + - sample_id * data_channels + - chan]; + sample_id * data_channels + chan]; const T img_fxfy = get_data_point(fx, fy, chan); const T img_cxcy = get_data_point(cx, cy, chan); const T img_fxcy = get_data_point(fx, cy, chan); const T img_cxfy = get_data_point(cx, fy, chan); // Update partial gradients wrt relevant warp field entries - update_grad_warp(sample_id, 0, - grad_output_value * - ((one - dy) * (img_cxcy - img_fxcy) + - dy * (img_cxfy - img_fxfy))); + update_grad_warp( + sample_id, 0, + grad_output_value * ((one - dy) * (img_cxcy - img_fxcy) + + dy * (img_cxfy - img_fxfy))); - update_grad_warp(sample_id, 1, - grad_output_value * - ((one - dx) * (img_cxcy - img_cxfy) + - dx * (img_fxcy - img_fxfy))); + update_grad_warp( + sample_id, 1, + grad_output_value * ((one - dx) * (img_cxcy - img_cxfy) + + dx * (img_fxcy - img_fxfy))); // Update partial gradients wrt sampled data - update_grad_data(fx, fy, chan, - grad_output_value * dx * dy); + update_grad_data(fx, fy, chan, grad_output_value * dx * dy); update_grad_data(cx, cy, chan, grad_output_value * (one - dx) * (one - dy)); update_grad_data(fx, cy, chan, @@ -355,8 +317,8 @@ struct ResamplerGrad2DFunctor{ // being 10us. // TODO(fviola): Check out if there is a better way of doing this. auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads()); - const int64 cost = static_cast(num_sampling_points) * - data_channels * 1000; + const int64 cost = + static_cast(num_sampling_points) * data_channels * 1000; ::tensorflow::Shard(worker_threads.num_threads, worker_threads.workers, batch_size, cost, update_grads_for_batches); } @@ -364,12 +326,11 @@ struct ResamplerGrad2DFunctor{ } // namespace functor - template class ResamplerGradOp : public ::tensorflow::OpKernel { public: - explicit ResamplerGradOp(::tensorflow::OpKernelConstruction* context) : - ::tensorflow::OpKernel(context) {} + explicit ResamplerGradOp(::tensorflow::OpKernelConstruction* context) + : ::tensorflow::OpKernel(context) {} void Compute(::tensorflow::OpKernelContext* ctx) override { const ::tensorflow::Tensor& data = ctx->input(0); @@ -383,7 +344,7 @@ class ResamplerGradOp : public ::tensorflow::OpKernel { "tensor must be a batch of 2d data; data shape should have " "4 entries corresponding to [batch_size, data_height, " "data_width, data_channels], but is: ", - data_shape.DebugString())); + data_shape.DebugString())); const int batch_size = data_shape.dim_size(0); const int data_height = data_shape.dim_size(1); const int data_width = data_shape.dim_size(2); @@ -394,7 +355,7 @@ class ResamplerGradOp : public ::tensorflow::OpKernel { ::tensorflow::errors::InvalidArgument( "warp should be at least a matrix, got shape ", warp_shape.DebugString())); - OP_REQUIRES(ctx, warp_shape.dim_size(warp_shape.dims()-1) == 2, + OP_REQUIRES(ctx, warp_shape.dim_size(warp_shape.dims() - 1) == 2, ::tensorflow::errors::Unimplemented( "Only bilinear interpolation is supported, warping " "coordinates must be 2D; warp shape last entry should be " @@ -417,18 +378,11 @@ class ResamplerGradOp : public ::tensorflow::OpKernel { OP_REQUIRES_OK(ctx, ctx->allocate_output(1, warp.shape(), &grad_warp)); // Execute kernel only for nonempty output; otherwise Eigen crashes on GPU. if (num_sampling_points > 0) { - functor::ResamplerGrad2DFunctor()(ctx, - ctx->eigen_device(), - data.flat().data(), - warp.flat().data(), - grad_output.flat().data(), - grad_data->flat().data(), - grad_warp->flat().data(), - batch_size, - data_height, - data_width, - data_channels, - num_sampling_points); + functor::ResamplerGrad2DFunctor()( + ctx, ctx->eigen_device(), data.flat().data(), + warp.flat().data(), grad_output.flat().data(), + grad_data->flat().data(), grad_warp->flat().data(), batch_size, + data_height, data_width, data_channels, num_sampling_points); } } @@ -436,11 +390,9 @@ class ResamplerGradOp : public ::tensorflow::OpKernel { TF_DISALLOW_COPY_AND_ASSIGN(ResamplerGradOp); }; -#define REGISTER(TYPE) \ - REGISTER_KERNEL_BUILDER( \ - Name("ResamplerGrad") \ - .Device(DEVICE_CPU) \ - .TypeConstraint("T"), \ +#define REGISTER(TYPE) \ + REGISTER_KERNEL_BUILDER( \ + Name("ResamplerGrad").Device(DEVICE_CPU).TypeConstraint("T"), \ ResamplerGradOp); TF_CALL_half(REGISTER); @@ -449,11 +401,10 @@ TF_CALL_double(REGISTER); #undef REGISTER #if GOOGLE_CUDA -#define REGISTER(TYPE) \ - REGISTER_KERNEL_BUILDER(Name("ResamplerGrad") \ - .Device(DEVICE_GPU) \ - .TypeConstraint("T"), \ - ResamplerGradOp) +#define REGISTER(TYPE) \ + REGISTER_KERNEL_BUILDER( \ + Name("ResamplerGrad").Device(DEVICE_GPU).TypeConstraint("T"), \ + ResamplerGradOp) // Disable half and double precision since atomicAdds are not supported // TF_CALL_half(REGISTER); // TF_CALL_double(REGISTER); diff --git a/tensorflow/contrib/resampler/kernels/resampler_ops.h b/tensorflow/contrib/resampler/kernels/resampler_ops.h index 8258ecaf5d3ba67094194c5cb12ca6d4d6efc85f..7fe3b9c0df71f51e07d38ea15a672d79fdc70453 100644 --- a/tensorflow/contrib/resampler/kernels/resampler_ops.h +++ b/tensorflow/contrib/resampler/kernels/resampler_ops.h @@ -13,8 +13,8 @@ // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_RESAMPLER_KERNELS_RESAMPLER_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_RESAMPLER_KERNELS_RESAMPLER_OPS_H_ +#ifndef TENSORFLOW_CONTRIB_RESAMPLER_KERNELS_RESAMPLER_OPS_H_ +#define TENSORFLOW_CONTRIB_RESAMPLER_KERNELS_RESAMPLER_OPS_H_ #if PLATFORM_WINDOWS #define __restrict__ __restrict @@ -29,40 +29,26 @@ namespace functor { // Helper functor for the Resampler Op in 2D template -struct Resampler2DFunctor{ - void operator ()(::tensorflow::OpKernelContext* ctx, - const Device& d, - const T* __restrict__ data, - const T* __restrict__ warp, - T* __restrict__ output, - const int batch_size, - const int data_height, - const int data_width, - const int data_channels, - const int num_sampling_points); +struct Resampler2DFunctor { + void operator()(::tensorflow::OpKernelContext* ctx, const Device& d, + const T* __restrict__ data, const T* __restrict__ warp, + T* __restrict__ output, const int batch_size, + const int data_height, const int data_width, + const int data_channels, const int num_sampling_points); }; - // Helper functor for the Resampler Gradient Op in 2D template -struct ResamplerGrad2DFunctor{ - void operator ()(::tensorflow::OpKernelContext* ctx, - const Device& d, - const T* __restrict__ data, - const T* __restrict__ warp, - const T* __restrict__ grad_output, - T* __restrict__ grad_data, - T* __restrict__ grad_warp, - const int batch_size, - const int data_height, - const int data_width, - const int data_channels, - const int num_sampling_points); +struct ResamplerGrad2DFunctor { + void operator()(::tensorflow::OpKernelContext* ctx, const Device& d, + const T* __restrict__ data, const T* __restrict__ warp, + const T* __restrict__ grad_output, T* __restrict__ grad_data, + T* __restrict__ grad_warp, const int batch_size, + const int data_height, const int data_width, + const int data_channels, const int num_sampling_points); }; - } // namespace functor } // namespace tensorflow - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_RESAMPLER_KERNELS_RESAMPLER_OPS_H_ +#endif // TENSORFLOW_CONTRIB_RESAMPLER_KERNELS_RESAMPLER_OPS_H_ diff --git a/tensorflow/contrib/resampler/kernels/resampler_ops_gpu.cu.cc b/tensorflow/contrib/resampler/kernels/resampler_ops_gpu.cu.cc index 636847a212f27c738032128e3f3f653ec32f851b..3c07051f685c74b6e45fb782c80871f38dffbbf4 100644 --- a/tensorflow/contrib/resampler/kernels/resampler_ops_gpu.cu.cc +++ b/tensorflow/contrib/resampler/kernels/resampler_ops_gpu.cu.cc @@ -31,18 +31,15 @@ using GPUDevice = Eigen::GpuDevice; namespace { -#define GET_DATA_POINT(x, y) \ - data[batch_id * data_batch_stride + \ - data_channels * (y * data_width + x) + \ +#define GET_DATA_POINT(x, y) \ + data[batch_id * data_batch_stride + data_channels * (y * data_width + x) + \ chan] template __global__ void Resampler2DKernel(const T* __restrict__ data, const T* __restrict__ warp, - T* __restrict__ output, - const int batch_size, - const int data_height, - const int data_width, + T* __restrict__ output, const int batch_size, + const int data_height, const int data_width, const int data_channels, const int num_sampling_points) { const int output_data_size = batch_size * num_sampling_points * data_channels; @@ -75,10 +72,8 @@ __global__ void Resampler2DKernel(const T* __restrict__ data, // The effect is that the sampled signal smoothly goes to 0 outside // the original input domain, rather than presenting a jump // discontinuity at the image boundaries. - if (x > static_cast(-1.0) && - y > static_cast(-1.0) && - x < static_cast(data_width) && - y < static_cast(data_height)) { + if (x > static_cast(-1.0) && y > static_cast(-1.0) && + x < static_cast(data_width) && y < static_cast(data_height)) { // Precompute floor (f) and ceil (c) values for x and y. const int fx = std::floor(static_cast(x)); const int fy = std::floor(static_cast(y)); @@ -87,21 +82,20 @@ __global__ void Resampler2DKernel(const T* __restrict__ data, const T dx = static_cast(cx) - x; const T dy = static_cast(cy) - y; - const T img_fxfy = (fx >= 0 && fy >= 0) - ? dx * dy * GET_DATA_POINT(fx, fy) - : zero; + const T img_fxfy = + (fx >= 0 && fy >= 0) ? dx * dy * GET_DATA_POINT(fx, fy) : zero; const T img_cxcy = (cx <= data_width - 1 && cy <= data_height - 1) - ? (one - dx) * (one - dy) * GET_DATA_POINT(cx, cy) - : zero; + ? (one - dx) * (one - dy) * GET_DATA_POINT(cx, cy) + : zero; const T img_fxcy = (fx >= 0 && cy <= data_height - 1) - ? dx * (one - dy) * GET_DATA_POINT(fx, cy) - : zero; + ? dx * (one - dy) * GET_DATA_POINT(fx, cy) + : zero; const T img_cxfy = (cx <= data_width - 1 && fy >= 0) - ? (one - dx) * dy * GET_DATA_POINT(cx, fy) - : zero; + ? (one - dx) * dy * GET_DATA_POINT(cx, fy) + : zero; output[out_index] = img_fxfy + img_cxcy + img_fxcy + img_cxfy; } else { @@ -115,24 +109,20 @@ __global__ void Resampler2DKernel(const T* __restrict__ data, namespace functor { template -struct Resampler2DFunctor{ - void operator ()(::tensorflow::OpKernelContext* ctx, - const GPUDevice& d, - const T* __restrict__ data, - const T* __restrict__ warp, - T* __restrict__ output, - const int batch_size, - const int data_height, - const int data_width, - const int data_channels, - const int num_sampling_points) { - const int output_data_size = batch_size * num_sampling_points * data_channels; - ::tensorflow::CudaLaunchConfig config = - ::tensorflow::GetCudaLaunchConfig(output_data_size, d); - Resampler2DKernel - <<>>( - data, warp, output, batch_size, data_height, data_width, - data_channels, num_sampling_points); +struct Resampler2DFunctor { + void operator()(::tensorflow::OpKernelContext* ctx, const GPUDevice& d, + const T* __restrict__ data, const T* __restrict__ warp, + T* __restrict__ output, const int batch_size, + const int data_height, const int data_width, + const int data_channels, const int num_sampling_points) { + const int output_data_size = + batch_size * num_sampling_points * data_channels; + ::tensorflow::CudaLaunchConfig config = + ::tensorflow::GetCudaLaunchConfig(output_data_size, d); + Resampler2DKernel + <<>>( + data, warp, output, batch_size, data_height, data_width, + data_channels, num_sampling_points); } }; @@ -145,26 +135,20 @@ template struct Resampler2DFunctor; namespace { -#define UPDATE_GRAD_DATA_POINT(x, y, v) \ - atomicAdd(grad_data + (batch_id * data_batch_stride + \ - data_channels * (y * data_width + x) + \ - chan), \ +#define UPDATE_GRAD_DATA_POINT(x, y, v) \ + atomicAdd(grad_data + (batch_id * data_batch_stride + \ + data_channels * (y * data_width + x) + chan), \ v) - template -__global__ void ResamplerGrad2DKernel(const T* __restrict__ data, - const T* __restrict__ warp, - const T* __restrict__ grad_output, - T* __restrict__ grad_data, - T* __restrict__ grad_warp, - const int batch_size, - const int data_height, - const int data_width, - const int data_channels, - const int num_sampling_points) { - const int resampler_output_size = batch_size * num_sampling_points * - data_channels; +__global__ void ResamplerGrad2DKernel( + const T* __restrict__ data, const T* __restrict__ warp, + const T* __restrict__ grad_output, T* __restrict__ grad_data, + T* __restrict__ grad_warp, const int batch_size, const int data_height, + const int data_width, const int data_channels, + const int num_sampling_points) { + const int resampler_output_size = + batch_size * num_sampling_points * data_channels; CUDA_1D_KERNEL_LOOP(index, resampler_output_size) { const int out_index = index; @@ -199,10 +183,8 @@ __global__ void ResamplerGrad2DKernel(const T* __restrict__ data, // The effect is that the sampled signal smoothly goes to 0 outside // the original input domain, rather than presenting a jump // discontinuity at the image boundaries. - if (x > static_cast(-1.0) && - y > static_cast(-1.0) && - x < static_cast(data_width) && - y < static_cast(data_height)) { + if (x > static_cast(-1.0) && y > static_cast(-1.0) && + x < static_cast(data_width) && y < static_cast(data_height)) { // Precompute floor (f) and ceil (c) values for x and y. const int fx = std::floor(static_cast(x)); const int fy = std::floor(static_cast(y)); @@ -211,21 +193,17 @@ __global__ void ResamplerGrad2DKernel(const T* __restrict__ data, const T dx = static_cast(cx) - x; const T dy = static_cast(cy) - y; - const T img_fxfy = (fx >= 0 && fy >= 0) - ? GET_DATA_POINT(fx, fy) - : zero; + const T img_fxfy = (fx >= 0 && fy >= 0) ? GET_DATA_POINT(fx, fy) : zero; const T img_cxcy = (cx <= data_width - 1 && cy <= data_height - 1) - ? GET_DATA_POINT(cx, cy) - : zero; + ? GET_DATA_POINT(cx, cy) + : zero; - const T img_fxcy = (fx >= 0 && cy <= data_height - 1) - ? GET_DATA_POINT(fx, cy) - : zero; + const T img_fxcy = + (fx >= 0 && cy <= data_height - 1) ? GET_DATA_POINT(fx, cy) : zero; - const T img_cxfy = (cx <= data_width - 1 && fy >= 0) - ? GET_DATA_POINT(cx, fy) - : zero; + const T img_cxfy = + (cx <= data_width - 1 && fy >= 0) ? GET_DATA_POINT(cx, fy) : zero; // Update partial gradients wrt relevant warp field entries atomicAdd(grad_warp + warp_id_x, @@ -241,7 +219,7 @@ __global__ void ResamplerGrad2DKernel(const T* __restrict__ data, } if (cx <= data_width - 1 && cy <= data_height - 1) { UPDATE_GRAD_DATA_POINT(cx, cy, - grad_output_value * (one - dx) * (one - dy)); + grad_output_value * (one - dx) * (one - dy)); } if (fx >= 0 && cy <= data_height - 1) { UPDATE_GRAD_DATA_POINT(fx, cy, grad_output_value * dx * (one - dy)); @@ -261,43 +239,37 @@ __global__ void ResamplerGrad2DKernel(const T* __restrict__ data, namespace functor { template -struct ResamplerGrad2DFunctor{ - void operator ()(::tensorflow::OpKernelContext* ctx, - const GPUDevice& d, - const T* __restrict__ data, - const T* __restrict__ warp, - const T* __restrict__ grad_output, - T* __restrict__ grad_data, - T* __restrict__ grad_warp, - const int batch_size, - const int data_height, - const int data_width, - const int data_channels, - const int num_sampling_points) { - // Set gradients to 0, because the kernel incrementally updates the - // tensor entries by adding partial contributions. - const int grad_warp_size = batch_size * num_sampling_points * 2; - const int grad_data_size = batch_size * data_height * data_width * - data_channels; - - ::tensorflow::CudaLaunchConfig config = - ::tensorflow::GetCudaLaunchConfig(grad_warp_size, d); - ::tensorflow::SetZero - <<>>( - grad_warp_size, grad_warp); - - config = ::tensorflow::GetCudaLaunchConfig(grad_data_size, d); - ::tensorflow::SetZero - <<>>( - grad_data_size, grad_data); - - const int resampler_output_size = batch_size * num_sampling_points * - data_channels; - config = ::tensorflow::GetCudaLaunchConfig(resampler_output_size, d); - ResamplerGrad2DKernel - <<>>( - data, warp, grad_output, grad_data, grad_warp, batch_size, - data_height, data_width, data_channels, num_sampling_points); +struct ResamplerGrad2DFunctor { + void operator()(::tensorflow::OpKernelContext* ctx, const GPUDevice& d, + const T* __restrict__ data, const T* __restrict__ warp, + const T* __restrict__ grad_output, T* __restrict__ grad_data, + T* __restrict__ grad_warp, const int batch_size, + const int data_height, const int data_width, + const int data_channels, const int num_sampling_points) { + // Set gradients to 0, because the kernel incrementally updates the + // tensor entries by adding partial contributions. + const int grad_warp_size = batch_size * num_sampling_points * 2; + const int grad_data_size = + batch_size * data_height * data_width * data_channels; + + ::tensorflow::CudaLaunchConfig config = + ::tensorflow::GetCudaLaunchConfig(grad_warp_size, d); + ::tensorflow:: + SetZero<<>>( + grad_warp_size, grad_warp); + + config = ::tensorflow::GetCudaLaunchConfig(grad_data_size, d); + ::tensorflow:: + SetZero<<>>( + grad_data_size, grad_data); + + const int resampler_output_size = + batch_size * num_sampling_points * data_channels; + config = ::tensorflow::GetCudaLaunchConfig(resampler_output_size, d); + ResamplerGrad2DKernel + <<>>( + data, warp, grad_output, grad_data, grad_warp, batch_size, + data_height, data_width, data_channels, num_sampling_points); } }; diff --git a/tensorflow/contrib/rnn/__init__.py b/tensorflow/contrib/rnn/__init__.py index c568c6760fd67b1902b0c1e6dc1aa439cb63de9b..67f31785b57fddef67733c18c3b744322532c28c 100644 --- a/tensorflow/contrib/rnn/__init__.py +++ b/tensorflow/contrib/rnn/__init__.py @@ -18,6 +18,7 @@ See @{$python/contrib.rnn} guide. @@RNNCell +@@LayerRNNCell @@BasicRNNCell @@BasicLSTMCell @@GRUCell @@ -68,6 +69,10 @@ See @{$python/contrib.rnn} guide. @@static_bidirectional_rnn @@stack_bidirectional_dynamic_rnn @@stack_bidirectional_rnn + + +@@transpose_batch_time +@@best_effort_input_batch_size """ from __future__ import absolute_import @@ -85,6 +90,8 @@ from tensorflow.contrib.rnn.python.ops.lstm_ops import * from tensorflow.contrib.rnn.python.ops.rnn import * from tensorflow.contrib.rnn.python.ops.rnn_cell import * +from tensorflow.python.ops.rnn import _best_effort_input_batch_size as best_effort_input_batch_size +from tensorflow.python.ops.rnn import _transpose_batch_time as transpose_batch_time from tensorflow.python.ops.rnn import static_bidirectional_rnn from tensorflow.python.ops.rnn import static_rnn from tensorflow.python.ops.rnn import static_state_saving_rnn diff --git a/tensorflow/contrib/rnn/kernels/blas_gemm.cc b/tensorflow/contrib/rnn/kernels/blas_gemm.cc index e62501e9b100484a7be3cc6ae0fc25905c0d0724..03006dab323a7c6dc83d9a17c035ef705f7b0366 100644 --- a/tensorflow/contrib/rnn/kernels/blas_gemm.cc +++ b/tensorflow/contrib/rnn/kernels/blas_gemm.cc @@ -36,11 +36,10 @@ perftools::gputools::DeviceMemory AsDeviceMemory(const T* cuda_memory) { namespace functor { template -void TensorCuBlasGemm::operator()(OpKernelContext* ctx, - bool transa, bool transb, uint64 m, - uint64 n, uint64 k, T alpha, const T* a, - int lda, const T* b, int ldb, T beta, T* c, - int ldc) { +void TensorCuBlasGemm::operator()(OpKernelContext* ctx, bool transa, + bool transb, uint64 m, uint64 n, uint64 k, + T alpha, const T* a, int lda, const T* b, + int ldb, T beta, T* c, int ldc) { #if GOOGLE_CUDA perftools::gputools::blas::Transpose trans[] = { perftools::gputools::blas::Transpose::kNoTranspose, diff --git a/tensorflow/contrib/rnn/kernels/blas_gemm.h b/tensorflow/contrib/rnn/kernels/blas_gemm.h index e33eceadff17fc3811f98fc29b3cb916b6a79766..a52c934233af3dc63e1a60d70fac6a9eba6a655b 100644 --- a/tensorflow/contrib/rnn/kernels/blas_gemm.h +++ b/tensorflow/contrib/rnn/kernels/blas_gemm.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_BLAS_GEMM_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_BLAS_GEMM_H_ +#ifndef TENSORFLOW_CONTRIB_RNN_KERNELS_BLAS_GEMM_H_ +#define TENSORFLOW_CONTRIB_RNN_KERNELS_BLAS_GEMM_H_ #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/framework/tensor_types.h" @@ -74,4 +74,4 @@ struct TensorBlasGemm { } // namespace functor } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_BLAS_GEMM_H_ +#endif // TENSORFLOW_CONTRIB_RNN_KERNELS_BLAS_GEMM_H_ diff --git a/tensorflow/contrib/rnn/kernels/gru_ops.cc b/tensorflow/contrib/rnn/kernels/gru_ops.cc index 0796f82b214620dd71d154fb8f8ec953dbcbb9ec..bd3d898fb09da0f490050c85b1e585502d8ecb2c 100644 --- a/tensorflow/contrib/rnn/kernels/gru_ops.cc +++ b/tensorflow/contrib/rnn/kernels/gru_ops.cc @@ -15,8 +15,8 @@ limitations under the License. #define EIGEN_USE_THREADS -#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/contrib/rnn/kernels/gru_ops.h" +#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/framework/op_kernel.h" namespace tensorflow { @@ -61,9 +61,9 @@ class GRUCellBlockOp : public OpKernel { h_prev_tensor->dim_size(0), " vs. ", batch_size)); OP_REQUIRES(ctx, h_prev_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("h_prev.dims(1) != cell_size: ", - h_prev_tensor->dim_size(1), " vs. ", - cell_size)); + errors::InvalidArgument( + "h_prev.dims(1) != cell_size: ", h_prev_tensor->dim_size(1), + " vs. ", cell_size)); // Shape of 'w_ru' must be [input_size+cell_size, 2*cell_size] OP_REQUIRES(ctx, w_ru_tensor->dim_size(0) == input_size + cell_size, @@ -82,10 +82,10 @@ class GRUCellBlockOp : public OpKernel { "w_c.dim_size(0) != input_size + cell_size: ", w_c_tensor->dim_size(0), " vs. ", input_size + cell_size)); - OP_REQUIRES( - ctx, w_c_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("w_c.dim_size(1) != cell_size: ", - w_c_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, w_c_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "w_c.dim_size(1) != cell_size: ", w_c_tensor->dim_size(1), + " vs. ", cell_size)); // Shape of 'b_ru' must be [2*cell_size] OP_REQUIRES(ctx, b_ru_tensor->dim_size(0) == cell_size * 2, @@ -97,10 +97,10 @@ class GRUCellBlockOp : public OpKernel { errors::InvalidArgument("Rank of b_ru must be 1", b_ru_tensor->dims(), " vs. 1", 1)); // Shape of 'b_c' must be [cell_size] - OP_REQUIRES( - ctx, b_c_tensor->dim_size(0) == cell_size, - errors::InvalidArgument("b_c.dim_size(0) != cell_size: ", - b_c_tensor->dim_size(0), " vs. ", cell_size)); + OP_REQUIRES(ctx, b_c_tensor->dim_size(0) == cell_size, + errors::InvalidArgument( + "b_c.dim_size(0) != cell_size: ", b_c_tensor->dim_size(0), + " vs. ", cell_size)); OP_REQUIRES(ctx, b_c_tensor->dims() == 1, errors::InvalidArgument("Rank of b_c must be 1", b_c_tensor->dims(), " vs. 1")); @@ -216,9 +216,9 @@ class GRUBlockCellGradOp : public OpKernel { h_prev_tensor->dim_size(0), " vs. ", batch_size)); OP_REQUIRES(ctx, h_prev_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("h_prev.dims(1) != cell_size: ", - h_prev_tensor->dim_size(1), " vs. ", - cell_size)); + errors::InvalidArgument( + "h_prev.dims(1) != cell_size: ", h_prev_tensor->dim_size(1), + " vs. ", cell_size)); // Shape of 'w_ru' must be [input_size+cell_size, 2*cell_size] OP_REQUIRES(ctx, w_ru_tensor->dim_size(0) == input_size + cell_size, @@ -237,10 +237,10 @@ class GRUBlockCellGradOp : public OpKernel { "w_c.dim_size(0) != input_size + cell_size: ", w_c_tensor->dim_size(0), " vs. ", input_size + cell_size)); - OP_REQUIRES( - ctx, w_c_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("w_c.dim_size(1) != cell_size: ", - w_c_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, w_c_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "w_c.dim_size(1) != cell_size: ", w_c_tensor->dim_size(1), + " vs. ", cell_size)); // Shape of 'b_ru' must be [2*cell_size] OP_REQUIRES(ctx, b_ru_tensor->dim_size(0) == cell_size * 2, @@ -253,54 +253,54 @@ class GRUBlockCellGradOp : public OpKernel { b_ru_tensor->dims(), " vs. 1")); // Shape of 'b_c' must be [cell_size] - OP_REQUIRES( - ctx, b_c_tensor->dim_size(0) == cell_size, - errors::InvalidArgument("b_c.dim_size(0) != cell_size: ", - b_c_tensor->dim_size(0), " vs. ", cell_size)); + OP_REQUIRES(ctx, b_c_tensor->dim_size(0) == cell_size, + errors::InvalidArgument( + "b_c.dim_size(0) != cell_size: ", b_c_tensor->dim_size(0), + " vs. ", cell_size)); OP_REQUIRES(ctx, b_c_tensor->dims() == 1, errors::InvalidArgument("Rank of b_c must be 1 ", b_c_tensor->dims(), " vs. 1")); // Shape of 'r' must be [batch_size, cell_size] - OP_REQUIRES( - ctx, r_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("r.dims(0) != batch_size: ", - r_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, r_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("r.dims(1) != cell_size: ", - r_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, r_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "r.dims(0) != batch_size: ", r_tensor->dim_size(0), " vs. ", + batch_size)); + OP_REQUIRES(ctx, r_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "r.dims(1) != cell_size: ", r_tensor->dim_size(1), " vs. ", + cell_size)); // Shape of 'u' must be [batch_size, cell_size] - OP_REQUIRES( - ctx, u_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("u.dims(0) != batch_size: ", - u_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, u_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("u.dims(1) != cell_size: ", - u_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, u_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "u.dims(0) != batch_size: ", u_tensor->dim_size(0), " vs. ", + batch_size)); + OP_REQUIRES(ctx, u_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "u.dims(1) != cell_size: ", u_tensor->dim_size(1), " vs. ", + cell_size)); // Shape of 'c' must be [batch_size, cell_size] - OP_REQUIRES( - ctx, c_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("c.dims(0) != batch_size: ", - c_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, c_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("c.dims(1) != cell_size: ", - c_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, c_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "c.dims(0) != batch_size: ", c_tensor->dim_size(0), " vs. ", + batch_size)); + OP_REQUIRES(ctx, c_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "c.dims(1) != cell_size: ", c_tensor->dim_size(1), " vs. ", + cell_size)); // Shape of 'd_h' must be [batch_size, cell_size] - OP_REQUIRES( - ctx, d_h_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("d_h.dims(0) != batch_size: ", - d_h_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, d_h_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("d_h.dims(1) != cell_size: ", - d_h_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, d_h_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "d_h.dims(0) != batch_size: ", d_h_tensor->dim_size(0), + " vs. ", batch_size)); + OP_REQUIRES(ctx, d_h_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "d_h.dims(1) != cell_size: ", d_h_tensor->dim_size(1), + " vs. ", cell_size)); // Create output tensors. Tensor* d_x_tensor = nullptr; diff --git a/tensorflow/contrib/rnn/kernels/gru_ops.h b/tensorflow/contrib/rnn/kernels/gru_ops.h index 06a566506296dd658a01bb3038407f77a32cde84..3e2cb39e64bb3f0b22ea66c5601af36c5fb9b0fd 100644 --- a/tensorflow/contrib/rnn/kernels/gru_ops.h +++ b/tensorflow/contrib/rnn/kernels/gru_ops.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_GRU_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_GRU_OPS_H_ +#ifndef TENSORFLOW_CONTRIB_RNN_KERNELS_GRU_OPS_H_ +#define TENSORFLOW_CONTRIB_RNN_KERNELS_GRU_OPS_H_ #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/contrib/rnn/kernels/blas_gemm.h" @@ -181,4 +181,4 @@ struct GRUBlockCellBprop : public GRUCell { } // namespace functor } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_GRU_OPS_H_ +#endif // TENSORFLOW_CONTRIB_RNN_KERNELS_GRU_OPS_H_ diff --git a/tensorflow/contrib/rnn/kernels/lstm_ops.cc b/tensorflow/contrib/rnn/kernels/lstm_ops.cc index 941a457fd3ada312b981fb23c769ff9ecea9ff13..5e7cf0ce84d332bd24088cd78995f7843813328b 100644 --- a/tensorflow/contrib/rnn/kernels/lstm_ops.cc +++ b/tensorflow/contrib/rnn/kernels/lstm_ops.cc @@ -281,23 +281,23 @@ class LSTMBlockCellOp : public OpKernel { h_prev_tensor->dim_size(0), " vs. ", batch_size)); OP_REQUIRES(ctx, h_prev_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("h_prev.dims(1) != cell_size: ", - h_prev_tensor->dim_size(1), " vs. ", - cell_size)); + errors::InvalidArgument( + "h_prev.dims(1) != cell_size: ", h_prev_tensor->dim_size(1), + " vs. ", cell_size)); OP_REQUIRES(ctx, w_tensor->dim_size(0) == input_size + cell_size, errors::InvalidArgument( "w.dim_size(0) != input_size + cell_size: ", w_tensor->dim_size(0), " vs. ", input_size + cell_size)); - OP_REQUIRES( - ctx, w_tensor->dim_size(1) == cell_size * 4, - errors::InvalidArgument("w.dim_size(1) != cell_size * 4: ", - w_tensor->dim_size(1), " vs. ", cell_size * 4)); + OP_REQUIRES(ctx, w_tensor->dim_size(1) == cell_size * 4, + errors::InvalidArgument( + "w.dim_size(1) != cell_size * 4: ", w_tensor->dim_size(1), + " vs. ", cell_size * 4)); - OP_REQUIRES( - ctx, b_tensor->dim_size(0) == cell_size * 4, - errors::InvalidArgument("b.dim_size(0) != cell_size * 4: ", - b_tensor->dim_size(0), " vs. ", cell_size * 4)); + OP_REQUIRES(ctx, b_tensor->dim_size(0) == cell_size * 4, + errors::InvalidArgument( + "b.dim_size(0) != cell_size * 4: ", b_tensor->dim_size(0), + " vs. ", cell_size * 4)); // Allocate our output tensors. Tensor* i_tensor = nullptr; @@ -484,77 +484,77 @@ class LSTMBlockCellGradOp : public OpKernel { h_prev_tensor->dim_size(0), " vs. ", batch_size)); OP_REQUIRES(ctx, h_prev_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("h_prev.dims(1) != cell_size: ", - h_prev_tensor->dim_size(1), " vs. ", - cell_size)); + errors::InvalidArgument( + "h_prev.dims(1) != cell_size: ", h_prev_tensor->dim_size(1), + " vs. ", cell_size)); OP_REQUIRES(ctx, w_tensor->dim_size(0) == input_size + cell_size, errors::InvalidArgument( "w.dim_size(0) != input_size + cell_size: ", w_tensor->dim_size(0), " vs. ", input_size + cell_size)); - OP_REQUIRES( - ctx, w_tensor->dim_size(1) == cell_size * 4, - errors::InvalidArgument("w.dim_size(1) != cell_size * 4: ", - w_tensor->dim_size(1), " vs. ", cell_size * 4)); + OP_REQUIRES(ctx, w_tensor->dim_size(1) == cell_size * 4, + errors::InvalidArgument( + "w.dim_size(1) != cell_size * 4: ", w_tensor->dim_size(1), + " vs. ", cell_size * 4)); - OP_REQUIRES( - ctx, b_tensor->dim_size(0) == cell_size * 4, - errors::InvalidArgument("b.dim_size(0) != cell_size * 4: ", - b_tensor->dim_size(0), " vs. ", cell_size * 4)); + OP_REQUIRES(ctx, b_tensor->dim_size(0) == cell_size * 4, + errors::InvalidArgument( + "b.dim_size(0) != cell_size * 4: ", b_tensor->dim_size(0), + " vs. ", cell_size * 4)); - OP_REQUIRES( - ctx, i_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("i.dim_size(0) != batch_size: ", - i_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, i_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("i.dim_size(1) != cell_size: ", - i_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, i_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "i.dim_size(0) != batch_size: ", i_tensor->dim_size(0), + " vs. ", batch_size)); + OP_REQUIRES(ctx, i_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "i.dim_size(1) != cell_size: ", i_tensor->dim_size(1), + " vs. ", cell_size)); - OP_REQUIRES( - ctx, cs_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("cs.dim_size(0) != batch_size: ", - cs_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, cs_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("cs.dim_size(1) != cell_size: ", - cs_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, cs_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "cs.dim_size(0) != batch_size: ", cs_tensor->dim_size(0), + " vs. ", batch_size)); + OP_REQUIRES(ctx, cs_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "cs.dim_size(1) != cell_size: ", cs_tensor->dim_size(1), + " vs. ", cell_size)); - OP_REQUIRES( - ctx, f_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("f.dim_size(0) != batch_size: ", - f_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, f_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("i.dim_size(1) != cell_size: ", - f_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, f_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "f.dim_size(0) != batch_size: ", f_tensor->dim_size(0), + " vs. ", batch_size)); + OP_REQUIRES(ctx, f_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "i.dim_size(1) != cell_size: ", f_tensor->dim_size(1), + " vs. ", cell_size)); - OP_REQUIRES( - ctx, o_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("o.dim_size(0) != batch_size: ", - o_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, o_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("o.dim_size(1) != cell_size: ", - o_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, o_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "o.dim_size(0) != batch_size: ", o_tensor->dim_size(0), + " vs. ", batch_size)); + OP_REQUIRES(ctx, o_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "o.dim_size(1) != cell_size: ", o_tensor->dim_size(1), + " vs. ", cell_size)); - OP_REQUIRES( - ctx, ci_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("ci.dim_size(0) != batch_size: ", - ci_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, ci_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("ci.dim_size(1) != cell_size: ", - ci_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, ci_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "ci.dim_size(0) != batch_size: ", ci_tensor->dim_size(0), + " vs. ", batch_size)); + OP_REQUIRES(ctx, ci_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "ci.dim_size(1) != cell_size: ", ci_tensor->dim_size(1), + " vs. ", cell_size)); - OP_REQUIRES( - ctx, co_tensor->dim_size(0) == batch_size, - errors::InvalidArgument("co.dim_size(0) != batch_size: ", - co_tensor->dim_size(0), " vs. ", batch_size)); - OP_REQUIRES( - ctx, co_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("co.dim_size(1) != cell_size: ", - co_tensor->dim_size(1), " vs. ", cell_size)); + OP_REQUIRES(ctx, co_tensor->dim_size(0) == batch_size, + errors::InvalidArgument( + "co.dim_size(0) != batch_size: ", co_tensor->dim_size(0), + " vs. ", batch_size)); + OP_REQUIRES(ctx, co_tensor->dim_size(1) == cell_size, + errors::InvalidArgument( + "co.dim_size(1) != cell_size: ", co_tensor->dim_size(1), + " vs. ", cell_size)); OP_REQUIRES(ctx, cs_grad_tensor->dim_size(0) == batch_size, errors::InvalidArgument( @@ -860,9 +860,9 @@ class BlockLSTMOp : public OpKernel { h_prev_tensor->dim_size(0), " vs. ", batch_size)); OP_REQUIRES(ctx, h_prev_tensor->dim_size(1) == cell_size, - errors::InvalidArgument("h_prev.dims(1) != cell_size: ", - h_prev_tensor->dim_size(1), " vs. ", - cell_size)); + errors::InvalidArgument( + "h_prev.dims(1) != cell_size: ", h_prev_tensor->dim_size(1), + " vs. ", cell_size)); const Tensor* w_tensor = nullptr; OP_REQUIRES_OK(ctx, ctx->input("w", &w_tensor)); @@ -872,46 +872,46 @@ class BlockLSTMOp : public OpKernel { errors::InvalidArgument( "w.dim_size(0) != input_size + cell_size: ", w_tensor->dim_size(0), " vs. ", input_size + cell_size)); - OP_REQUIRES( - ctx, w_tensor->dim_size(1) == cell_size * 4, - errors::InvalidArgument("w.dim_size(1) != cell_size * 4: ", - w_tensor->dim_size(1), " vs. ", cell_size * 4)); + OP_REQUIRES(ctx, w_tensor->dim_size(1) == cell_size * 4, + errors::InvalidArgument( + "w.dim_size(1) != cell_size * 4: ", w_tensor->dim_size(1), + " vs. ", cell_size * 4)); const Tensor* wci_tensor = nullptr; OP_REQUIRES_OK(ctx, ctx->input("wci", &wci_tensor)); OP_REQUIRES(ctx, wci_tensor->dims() == 1, errors::InvalidArgument("wci must be 1D")); - OP_REQUIRES( - ctx, wci_tensor->dim_size(0) == cell_size, - errors::InvalidArgument("wci.dim_size(0) != cell_size: ", - wci_tensor->dim_size(0), " vs. ", cell_size)); + OP_REQUIRES(ctx, wci_tensor->dim_size(0) == cell_size, + errors::InvalidArgument( + "wci.dim_size(0) != cell_size: ", wci_tensor->dim_size(0), + " vs. ", cell_size)); const Tensor* wcf_tensor = nullptr; OP_REQUIRES_OK(ctx, ctx->input("wcf", &wcf_tensor)); OP_REQUIRES(ctx, wcf_tensor->dims() == 1, errors::InvalidArgument("wcf must be 1D")); - OP_REQUIRES( - ctx, wcf_tensor->dim_size(0) == cell_size, - errors::InvalidArgument("wcf.dim_size(0) != cell_size: ", - wcf_tensor->dim_size(0), " vs. ", cell_size)); + OP_REQUIRES(ctx, wcf_tensor->dim_size(0) == cell_size, + errors::InvalidArgument( + "wcf.dim_size(0) != cell_size: ", wcf_tensor->dim_size(0), + " vs. ", cell_size)); const Tensor* wco_tensor = nullptr; OP_REQUIRES_OK(ctx, ctx->input("wco", &wco_tensor)); OP_REQUIRES(ctx, wco_tensor->dims() == 1, errors::InvalidArgument("wco must be 1D")); - OP_REQUIRES( - ctx, wco_tensor->dim_size(0) == cell_size, - errors::InvalidArgument("wco.dim_size(0) != cell_size: ", - wco_tensor->dim_size(0), " vs. ", cell_size)); + OP_REQUIRES(ctx, wco_tensor->dim_size(0) == cell_size, + errors::InvalidArgument( + "wco.dim_size(0) != cell_size: ", wco_tensor->dim_size(0), + " vs. ", cell_size)); const Tensor* b_tensor = nullptr; OP_REQUIRES_OK(ctx, ctx->input("b", &b_tensor)); OP_REQUIRES(ctx, b_tensor->dims() == 1, errors::InvalidArgument("b must be 1D")); - OP_REQUIRES( - ctx, b_tensor->dim_size(0) == cell_size * 4, - errors::InvalidArgument("b.dim_size(0) != cell_size * 4: ", - b_tensor->dim_size(0), " vs. ", cell_size * 4)); + OP_REQUIRES(ctx, b_tensor->dim_size(0) == cell_size * 4, + errors::InvalidArgument( + "b.dim_size(0) != cell_size * 4: ", b_tensor->dim_size(0), + " vs. ", cell_size * 4)); TensorShape batch_cell_shape({timelen, batch_size, cell_size}); Tensor* i_out; @@ -1065,9 +1065,9 @@ class BlockLSTMGradOp : public OpKernel { OP_REQUIRES_OK(ctx, ctx->input("w", &w_tensor)); const int64 cell_size = w_tensor->dim_size(1) / 4; OP_REQUIRES(ctx, input_size + cell_size == w_tensor->dim_size(0), - errors::InvalidArgument("w matrix rows don't match: ", - input_size + cell_size, " vs. ", - w_tensor->dim_size(0))); + errors::InvalidArgument( + "w matrix rows don't match: ", input_size + cell_size, + " vs. ", w_tensor->dim_size(0))); const Tensor* wci_tensor = nullptr; OP_REQUIRES_OK(ctx, ctx->input("wci", &wci_tensor)); @@ -1193,7 +1193,6 @@ class BlockLSTMGradOp : public OpKernel { OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum::v(), batch_cell_shape, &h_grad_tensor)); - const Device& device = ctx->eigen_device(); functor::TensorZero()(device, cs_grad_tensor.flat()); diff --git a/tensorflow/contrib/rnn/kernels/lstm_ops.h b/tensorflow/contrib/rnn/kernels/lstm_ops.h index 1906581b16b2e76243320bc67c8ac831323fb8e7..d23cedc234b8c0e1a784346f28164ae79b8cbf89 100644 --- a/tensorflow/contrib/rnn/kernels/lstm_ops.h +++ b/tensorflow/contrib/rnn/kernels/lstm_ops.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_LSTM_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_LSTM_OPS_H_ +#ifndef TENSORFLOW_CONTRIB_RNN_KERNELS_LSTM_OPS_H_ +#define TENSORFLOW_CONTRIB_RNN_KERNELS_LSTM_OPS_H_ #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/contrib/rnn/kernels/blas_gemm.h" @@ -92,7 +92,6 @@ struct TensorZeroPadding { } }; - struct LSTMBlockCell { LSTMBlockCell(const int batch_size, const int input_size, const int cell_size) : batch_size_(batch_size), @@ -291,4 +290,4 @@ struct BlockLSTMBprop : public LSTMBlockCell { } // namespace functor } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_RNN_KERNELS_LSTM_OPS_H_ +#endif // TENSORFLOW_CONTRIB_RNN_KERNELS_LSTM_OPS_H_ diff --git a/tensorflow/contrib/rnn/ops/lstm_ops_test.cc b/tensorflow/contrib/rnn/ops/lstm_ops_test.cc index 544cd163c50062093acf7f5e942f67606936c0e3..68184b643e5e7a04ffecb804703051638514b7b2 100644 --- a/tensorflow/contrib/rnn/ops/lstm_ops_test.cc +++ b/tensorflow/contrib/rnn/ops/lstm_ops_test.cc @@ -149,8 +149,9 @@ TEST_F(LSTMOpsTest, BlockLSTMGrad_ShapeFn) { INFER_ERROR("must be rank 1", op, "?;?;?;?;?;?;?;?;[1,?]" + suffix); // Output with all input knowns makes known rank outputs. - INFER_OK(op, JoinedCopies("?", 18), "[?,?,?];" + JoinedCopies("[?,?]", 3) + - ";" + JoinedCopies("[?]", 4)); + INFER_OK( + op, JoinedCopies("?", 18), + "[?,?,?];" + JoinedCopies("[?,?]", 3) + ";" + JoinedCopies("[?]", 4)); // Output with copies input shapes to output. string input = strings::StrCat("?;[?,?,?];", JoinedCopies("[?,?]", 3), ";", diff --git a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py index b5d81b7caac5186b34548a06c67ba48afab0a1a5..0e62b315b61cb3ceeb5cfd33bf5102a71abef83b 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py @@ -39,9 +39,6 @@ from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables as variables_lib from tensorflow.python.platform import test -from tensorflow.python.framework import test_util -from tensorflow.contrib.rnn.python.ops import rnn_cell as contrib_rnn_cell - # pylint: enable=protected-access Linear = core_rnn_cell._Linear # pylint: disable=invalid-name @@ -84,19 +81,22 @@ class RNNCellTest(test.TestCase): ], [v.name for v in cell.trainable_variables]) self.assertFalse(cell.non_trainable_variables) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g], {x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]])}) + res = sess.run([g], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) self.assertEqual(res[0].shape, (1, 2)) def testBasicRNNCellNotTrainable(self): with self.test_session() as sess: + def not_trainable_getter(getter, *args, **kwargs): kwargs["trainable"] = False return getter(*args, **kwargs) with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5), + "root", + initializer=init_ops.constant_initializer(0.5), custom_getter=not_trainable_getter): x = array_ops.zeros([1, 2]) m = array_ops.zeros([1, 2]) @@ -108,9 +108,10 @@ class RNNCellTest(test.TestCase): "root/basic_rnn_cell/%s:0" % rnn_cell_impl._BIAS_VARIABLE_NAME ], [v.name for v in cell.non_trainable_variables]) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g], {x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]])}) + res = sess.run([g], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) self.assertEqual(res[0].shape, (1, 2)) def testGRUCell(self): @@ -121,9 +122,10 @@ class RNNCellTest(test.TestCase): m = array_ops.zeros([1, 2]) g, _ = rnn_cell_impl.GRUCell(2)(x, m) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g], {x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]])}) + res = sess.run([g], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) # Smoke test self.assertAllClose(res[0], [[0.175991, 0.175991]]) with variable_scope.variable_scope( @@ -133,10 +135,10 @@ class RNNCellTest(test.TestCase): m = array_ops.zeros([1, 2]) g, _ = rnn_cell_impl.GRUCell(2)(x, m) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g], - {x.name: np.array([[1., 1., 1.]]), - m.name: np.array([[0.1, 0.1]])}) + res = sess.run([g], { + x.name: np.array([[1., 1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) # Smoke test self.assertAllClose(res[0], [[0.156736, 0.156736]]) @@ -148,11 +150,27 @@ class RNNCellTest(test.TestCase): m = array_ops.zeros([1, 2]) g, _ = contrib_rnn_cell.SRUCell(2)(x, m) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g], {x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]])}) + res = sess.run([g], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) + # Smoke test + self.assertAllClose(res[0], [[0.509682, 0.509682]]) + + def testSRUCellWithDiffSize(self): + with self.test_session() as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 3]) + m = array_ops.zeros([1, 2]) + g, _ = contrib_rnn_cell.SRUCell(2)(x, m) + sess.run([variables_lib.global_variables_initializer()]) + res = sess.run([g], { + x.name: np.array([[1., 1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) # Smoke test - self.assertAllClose(res[0], [[0.509682, 0.509682]]) + self.assertAllClose(res[0], [[0.55255556, 0.55255556]]) def testBasicLSTMCell(self): for dtype in [dtypes.float16, dtypes.float32]: @@ -164,8 +182,7 @@ class RNNCellTest(test.TestCase): m = array_ops.zeros([1, 8], dtype=dtype) cell = rnn_cell_impl.MultiRNNCell( [ - rnn_cell_impl.BasicLSTMCell( - 2, state_is_tuple=False) + rnn_cell_impl.BasicLSTMCell(2, state_is_tuple=False) for _ in range(2) ], state_is_tuple=False) @@ -183,22 +200,21 @@ class RNNCellTest(test.TestCase): "root/multi_rnn_cell/cell_1/basic_lstm_cell/%s:0" % rnn_cell_impl._BIAS_VARIABLE_NAME ] - self.assertEqual( - expected_variable_names, - [v.name for v in cell.trainable_variables]) + self.assertEqual(expected_variable_names, + [v.name for v in cell.trainable_variables]) self.assertFalse(cell.non_trainable_variables) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g, out_m], - {x.name: np.array([[1., 1.]]), - m.name: 0.1 * np.ones([1, 8])}) + res = sess.run([g, out_m], { + x.name: np.array([[1., 1.]]), + m.name: 0.1 * np.ones([1, 8]) + }) self.assertEqual(len(res), 2) variables = variables_lib.global_variables() self.assertEqual(expected_variable_names, [v.name for v in variables]) # The numbers in results were not calculated, this is just a # smoke test. - self.assertAllClose( - res[0], np.array([[0.240, 0.240]], dtype=np_dtype), 1e-2) + self.assertAllClose(res[0], np.array( + [[0.240, 0.240]], dtype=np_dtype), 1e-2) expected_mem = np.array( [[0.689, 0.689, 0.448, 0.448, 0.398, 0.398, 0.240, 0.240]], dtype=np_dtype) @@ -208,13 +224,13 @@ class RNNCellTest(test.TestCase): # Test BasicLSTMCell with input_size != num_units. x = array_ops.zeros([1, 3], dtype=dtype) m = array_ops.zeros([1, 4], dtype=dtype) - g, out_m = rnn_cell_impl.BasicLSTMCell( - 2, state_is_tuple=False)(x, m) + g, out_m = rnn_cell_impl.BasicLSTMCell(2, state_is_tuple=False)(x, m) sess.run([variables_lib.global_variables_initializer()]) res = sess.run( - [g, out_m], - {x.name: np.array([[1., 1., 1.]], dtype=np_dtype), - m.name: 0.1 * np.ones([1, 4], dtype=np_dtype)}) + [g, out_m], { + x.name: np.array([[1., 1., 1.]], dtype=np_dtype), + m.name: 0.1 * np.ones([1, 4], dtype=np_dtype) + }) self.assertEqual(len(res), 2) def testBasicLSTMCellDimension0Error(self): @@ -232,9 +248,11 @@ class RNNCellTest(test.TestCase): g, out_m = rnn_cell_impl.BasicLSTMCell( num_units, state_is_tuple=False)(x, m) sess.run([variables_lib.global_variables_initializer()]) - sess.run([g, out_m], - {x.name: 1 * np.ones([batch_size, input_size]), - m.name: 0.1 * np.ones([batch_size - 1, state_size])}) + sess.run( + [g, out_m], { + x.name: 1 * np.ones([batch_size, input_size]), + m.name: 0.1 * np.ones([batch_size - 1, state_size]) + }) def testBasicLSTMCellStateSizeError(self): """Tests that state_size must be num_units * 2.""" @@ -251,9 +269,11 @@ class RNNCellTest(test.TestCase): g, out_m = rnn_cell_impl.BasicLSTMCell( num_units, state_is_tuple=False)(x, m) sess.run([variables_lib.global_variables_initializer()]) - sess.run([g, out_m], - {x.name: 1 * np.ones([batch_size, input_size]), - m.name: 0.1 * np.ones([batch_size, state_size])}) + sess.run( + [g, out_m], { + x.name: 1 * np.ones([batch_size, input_size]), + m.name: 0.1 * np.ones([batch_size, state_size]) + }) def testBasicLSTMCellStateTupleType(self): with self.test_session(): @@ -301,11 +321,12 @@ class RNNCellTest(test.TestCase): state_is_tuple=True) g, (out_m0, out_m1) = cell(x, (m0, m1)) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g, out_m0, out_m1], { - x.name: np.array([[1., 1.]]), - m0.name: 0.1 * np.ones([1, 4]), - m1.name: 0.1 * np.ones([1, 4]) - }) + res = sess.run( + [g, out_m0, out_m1], { + x.name: np.array([[1., 1.]]), + m0.name: 0.1 * np.ones([1, 4]), + m1.name: 0.1 * np.ones([1, 4]) + }) self.assertEqual(len(res), 3) # The numbers in results were not calculated, this is just a smoke test. # Note, however, these values should match the original @@ -336,10 +357,11 @@ class RNNCellTest(test.TestCase): state_is_tuple=False) output, state = cell(x, m) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([output, state], { - x.name: np.array([[1., 1.], [2., 2.], [3., 3.]]), - m.name: 0.1 * np.ones((batch_size, state_size)) - }) + res = sess.run( + [output, state], { + x.name: np.array([[1., 1.], [2., 2.], [3., 3.]]), + m.name: 0.1 * np.ones((batch_size, state_size)) + }) self.assertEqual(len(res), 2) # The numbers in results were not calculated, this is mostly just a # smoke test. @@ -442,10 +464,10 @@ class RNNCellTest(test.TestCase): rnn_cell_impl.GRUCell(3), num_proj=3) g, new_m = cell(x, m) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g, new_m], - {x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1, 0.1]])}) + res = sess.run([g, new_m], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1, 0.1]]) + }) self.assertEqual(res[1].shape, (1, 3)) # The numbers in results were not calculated, this is just a smoke test. self.assertAllClose(res[0], [[0.154605, 0.154605, 0.154605]]) @@ -479,9 +501,11 @@ class RNNCellTest(test.TestCase): base_cell = rnn_cell_impl.GRUCell(3) g, m_new = base_cell(x, m) variable_scope.get_variable_scope().reuse_variables() + def residual_with_slice_fn(inp, out): inp_sliced = array_ops.slice(inp, [0, 0], [-1, 3]) return inp_sliced + out + g_res, m_new_res = rnn_cell_impl.ResidualWrapper( base_cell, residual_with_slice_fn)(x, m) sess.run([variables_lib.global_variables_initializer()]) @@ -551,10 +575,10 @@ class RNNCellTest(test.TestCase): self.assertEqual(embedding_cell.output_size, 2) g, new_m = embedding_cell(x, m) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g, new_m], - {x.name: np.array([[1]]), - m.name: np.array([[0.1, 0.1]])}) + res = sess.run([g, new_m], { + x.name: np.array([[1]]), + m.name: np.array([[0.1, 0.1]]) + }) self.assertEqual(res[1].shape, (1, 2)) # The numbers in results were not calculated, this is just a smoke test. self.assertAllClose(res[0], [[0.17139, 0.17139]]) @@ -584,8 +608,8 @@ class RNNCellTest(test.TestCase): x = array_ops.zeros([1, 2]) m = array_ops.zeros([1, 4]) _, ml = rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.GRUCell(2) - for _ in range(2)], state_is_tuple=False)(x, m) + [rnn_cell_impl.GRUCell(2) for _ in range(2)], + state_is_tuple=False)(x, m) sess.run([variables_lib.global_variables_initializer()]) res = sess.run(ml, { x.name: np.array([[1., 1.]]), @@ -605,19 +629,20 @@ class RNNCellTest(test.TestCase): # Test incorrectness of state with self.assertRaisesRegexp(ValueError, "Expected state .* a tuple"): rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.GRUCell(2) - for _ in range(2)], state_is_tuple=True)(x, m_bad) + [rnn_cell_impl.GRUCell(2) for _ in range(2)], + state_is_tuple=True)(x, m_bad) _, ml = rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.GRUCell(2) - for _ in range(2)], state_is_tuple=True)(x, m_good) + [rnn_cell_impl.GRUCell(2) for _ in range(2)], + state_is_tuple=True)(x, m_good) sess.run([variables_lib.global_variables_initializer()]) - res = sess.run(ml, { - x.name: np.array([[1., 1.]]), - m_good[0].name: np.array([[0.1, 0.1]]), - m_good[1].name: np.array([[0.1, 0.1]]) - }) + res = sess.run( + ml, { + x.name: np.array([[1., 1.]]), + m_good[0].name: np.array([[0.1, 0.1]]), + m_good[1].name: np.array([[0.1, 0.1]]) + }) # The numbers in results were not calculated, this is just a # smoke test. However, these numbers should match those of @@ -628,8 +653,11 @@ class RNNCellTest(test.TestCase): class DropoutWrapperTest(test.TestCase): - def _testDropoutWrapper(self, batch_size=None, time_steps=None, - parallel_iterations=None, **kwargs): + def _testDropoutWrapper(self, + batch_size=None, + time_steps=None, + parallel_iterations=None, + **kwargs): with self.test_session() as sess: with variable_scope.variable_scope( "root", initializer=init_ops.constant_initializer(0.5)): @@ -640,14 +668,14 @@ class DropoutWrapperTest(test.TestCase): x = constant_op.constant( [[[2., 2., 2.]], [[1., 1., 1.]]], dtype=dtypes.float32) m = rnn_cell_impl.LSTMStateTuple( - *[constant_op.constant([[0.1, 0.1, 0.1]], dtype=dtypes.float32) - ] * 2) + *[constant_op.constant([[0.1, 0.1, 0.1]], dtype=dtypes.float32 + )] * 2) else: x = constant_op.constant( np.random.randn(time_steps, batch_size, 3).astype(np.float32)) m = rnn_cell_impl.LSTMStateTuple(*[ - constant_op.constant( - [[0.1, 0.1, 0.1]] * batch_size, dtype=dtypes.float32) + constant_op. + constant([[0.1, 0.1, 0.1]] * batch_size, dtype=dtypes.float32) ] * 2) outputs, final_state = rnn.dynamic_rnn( cell=rnn_cell_impl.DropoutWrapper( @@ -663,13 +691,19 @@ class DropoutWrapperTest(test.TestCase): self.assertEqual(res[1].h.shape, (batch_size, 3)) return res + def testWrappedCellProperty(self): + cell = rnn_cell_impl.BasicRNNCell(10) + wrapper = rnn_cell_impl.DropoutWrapper(cell) + # Github issue 15810 + self.assertEqual(wrapper.wrapped_cell, cell) + def testDropoutWrapperKeepAllConstantInput(self): keep = array_ops.ones([]) res = self._testDropoutWrapper( input_keep_prob=keep, output_keep_prob=keep, state_keep_prob=keep) true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], - [[0.895509, 0.895509, 0.895509]]], dtype=np.float32) + [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], + dtype=np.float32) true_full_final_c = np.array( [[1.949385, 1.949385, 1.949385]], dtype=np.float32) self.assertAllClose(true_full_output, res[0]) @@ -681,8 +715,8 @@ class DropoutWrapperTest(test.TestCase): res = self._testDropoutWrapper( input_keep_prob=keep, output_keep_prob=keep, state_keep_prob=keep) true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], - [[0.895509, 0.895509, 0.895509]]], dtype=np.float32) + [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], + dtype=np.float32) true_full_final_c = np.array( [[1.949385, 1.949385, 1.949385]], dtype=np.float32) self.assertAllClose(true_full_output, res[0]) @@ -697,16 +731,20 @@ class DropoutWrapperTest(test.TestCase): ## consistent across both calls. Otherwise the seed may not end ## up being munged consistently across both graphs. res_standard_1 = self._testDropoutWrapper( - input_keep_prob=keep_some, output_keep_prob=keep_some, - state_keep_prob=keep_some, seed=10, + input_keep_prob=keep_some, + output_keep_prob=keep_some, + state_keep_prob=keep_some, + seed=10, parallel_iterations=1) # Clear away the graph and the test session (which keeps variables around) ops.reset_default_graph() self._ClearCachedSession() random_seed.set_random_seed(2) res_standard_2 = self._testDropoutWrapper( - input_keep_prob=keep_some, output_keep_prob=keep_some, - state_keep_prob=keep_some, seed=10, + input_keep_prob=keep_some, + output_keep_prob=keep_some, + state_keep_prob=keep_some, + seed=10, parallel_iterations=1) self.assertAllClose(res_standard_1[0], res_standard_2[0]) self.assertAllClose(res_standard_1[1].c, res_standard_2[1].c) @@ -716,11 +754,12 @@ class DropoutWrapperTest(test.TestCase): keep_all = variable_scope.get_variable("all", initializer=1.0) keep_none = variable_scope.get_variable("none", initializer=1e-10) res = self._testDropoutWrapper( - input_keep_prob=keep_all, output_keep_prob=keep_none, + input_keep_prob=keep_all, + output_keep_prob=keep_none, state_keep_prob=keep_all) true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], - [[0.895509, 0.895509, 0.895509]]], dtype=np.float32) + [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], + dtype=np.float32) true_full_final_c = np.array( [[1.949385, 1.949385, 1.949385]], dtype=np.float32) self.assertAllClose(np.zeros(res[0].shape), res[0]) @@ -733,13 +772,13 @@ class DropoutWrapperTest(test.TestCase): # Even though we dropout state, by default DropoutWrapper never # drops out the memory ("c") term of an LSTMStateTuple. res = self._testDropoutWrapper( - input_keep_prob=keep_all, output_keep_prob=keep_all, + input_keep_prob=keep_all, + output_keep_prob=keep_all, state_keep_prob=keep_none) - true_c_state = np.array( - [[1.713925, 1.713925, 1.713925]], dtype=np.float32) + true_c_state = np.array([[1.713925, 1.713925, 1.713925]], dtype=np.float32) true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], - [[0.895509, 0.895509, 0.895509]]], dtype=np.float32) + [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], + dtype=np.float32) self.assertAllClose(true_full_output[0], res[0][0]) # Second output is modified by zero input state self.assertGreater(np.linalg.norm(true_full_output[1] - res[0][1]), 1e-4) @@ -752,13 +791,14 @@ class DropoutWrapperTest(test.TestCase): keep_all = variable_scope.get_variable("all", initializer=1.0) keep_none = variable_scope.get_variable("none", initializer=1e-10) true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], - [[0.895509, 0.895509, 0.895509]]], dtype=np.float32) + [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], + dtype=np.float32) true_full_final_c = np.array( [[1.949385, 1.949385, 1.949385]], dtype=np.float32) # All outputs are different because inputs are zeroed out res = self._testDropoutWrapper( - input_keep_prob=keep_none, output_keep_prob=keep_all, + input_keep_prob=keep_none, + output_keep_prob=keep_all, state_keep_prob=keep_all) self.assertGreater(np.linalg.norm(res[0] - true_full_output), 1e-4) self.assertGreater(np.linalg.norm(res[1].h - true_full_output[1]), 1e-4) @@ -768,9 +808,13 @@ class DropoutWrapperTest(test.TestCase): keep_some = 0.8 keep_all = variable_scope.get_variable("all", initializer=1.0) res = self._testDropoutWrapper( - input_keep_prob=keep_all, output_keep_prob=keep_some, - state_keep_prob=keep_all, variational_recurrent=True, - input_size=3, batch_size=5, time_steps=7) + input_keep_prob=keep_all, + output_keep_prob=keep_some, + state_keep_prob=keep_all, + variational_recurrent=True, + input_size=3, + batch_size=5, + time_steps=7) # Ensure the same dropout pattern for all time steps output_mask = np.abs(res[0]) > 1e-6 for m in output_mask[1:]: @@ -779,9 +823,13 @@ class DropoutWrapperTest(test.TestCase): def testDropoutWrapperRecurrentStateInputAndOutput(self): keep_some = 0.9 res = self._testDropoutWrapper( - input_keep_prob=keep_some, output_keep_prob=keep_some, - state_keep_prob=keep_some, variational_recurrent=True, - input_size=3, batch_size=5, time_steps=7) + input_keep_prob=keep_some, + output_keep_prob=keep_some, + state_keep_prob=keep_some, + variational_recurrent=True, + input_size=3, + batch_size=5, + time_steps=7) # Smoke test for the state/input masks. output_mask = np.abs(res[0]) > 1e-6 @@ -805,17 +853,27 @@ class DropoutWrapperTest(test.TestCase): random_seed.set_random_seed(2347) np.random.seed(23487) res0 = self._testDropoutWrapper( - input_keep_prob=keep_some, output_keep_prob=keep_some, - state_keep_prob=keep_some, variational_recurrent=True, - input_size=3, batch_size=5, time_steps=7, seed=-234987) + input_keep_prob=keep_some, + output_keep_prob=keep_some, + state_keep_prob=keep_some, + variational_recurrent=True, + input_size=3, + batch_size=5, + time_steps=7, + seed=-234987) ops.reset_default_graph() self._ClearCachedSession() random_seed.set_random_seed(2347) np.random.seed(23487) res1 = self._testDropoutWrapper( - input_keep_prob=keep_some, output_keep_prob=keep_some, - state_keep_prob=keep_some, variational_recurrent=True, - input_size=3, batch_size=5, time_steps=7, seed=-234987) + input_keep_prob=keep_some, + output_keep_prob=keep_some, + state_keep_prob=keep_some, + variational_recurrent=True, + input_size=3, + batch_size=5, + time_steps=7, + seed=-234987) output_mask = np.abs(res0[0]) > 1e-6 for time_step in output_mask: @@ -852,9 +910,10 @@ class SlimRNNCellTest(test.TestCase): g, _ = rnn_cell_impl._SlimRNNCell(my_cell)(x, m) # pylint: enable=protected-access sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g], {x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]])}) + res = sess.run([g], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) self.assertEqual(res[0].shape, (1, 2)) def testBasicRNNCellMatch(self): diff --git a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py index 0258d7202df20a536ae4240a532249b6b5e7e641..57521c6a9ba0b2d66639017b09c541e270276323 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py @@ -45,6 +45,7 @@ from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging from tensorflow.python.util import nest + class Plus1RNNCell(rnn_lib.RNNCell): """RNN Cell generating (output, new_state) = (input + 1, state + 1).""" @@ -160,8 +161,7 @@ class RNNTest(test.TestCase): input_size = 5 max_length = 8 # unrolled up to this length inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] outputs, state = rnn.static_rnn(cell, inputs, dtype=dtypes.float32) self.assertEqual(len(outputs), len(inputs)) @@ -178,10 +178,9 @@ class RNNTest(test.TestCase): self.assertAllClose(v, input_value + 1.0) # Final state - self.assertAllClose( - values[-1], - max_length * np.ones( - (batch_size, input_size), dtype=np.float32)) + self.assertAllClose(values[-1], + max_length * np.ones( + (batch_size, input_size), dtype=np.float32)) def testDropout(self): cell = Plus1RNNCell() @@ -191,8 +190,7 @@ class RNNTest(test.TestCase): input_size = 5 max_length = 8 inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] with variable_scope.variable_scope("share_scope"): outputs, state = rnn.static_rnn(cell, inputs, dtype=dtypes.float32) @@ -207,8 +205,10 @@ class RNNTest(test.TestCase): with self.test_session(use_gpu=True) as sess: input_value = np.random.randn(batch_size, input_size) values = sess.run(outputs + [state], feed_dict={inputs[0]: input_value}) - full_dropout_values = sess.run(dropped_outputs, - feed_dict={inputs[0]: input_value}) + full_dropout_values = sess.run( + dropped_outputs, feed_dict={ + inputs[0]: input_value + }) for v in values[:-1]: self.assertAllClose(v, input_value + 1.0) @@ -222,8 +222,7 @@ class RNNTest(test.TestCase): input_size = 5 max_length = 8 inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] with variable_scope.variable_scope("drop_scope"): dynamic_outputs, dynamic_state = rnn.static_rnn( @@ -234,12 +233,16 @@ class RNNTest(test.TestCase): input_value = np.random.randn(batch_size, input_size) dynamic_values = sess.run( dynamic_outputs, - feed_dict={inputs[0]: input_value, - sequence_length: [2, 3]}) + feed_dict={ + inputs[0]: input_value, + sequence_length: [2, 3] + }) dynamic_state_value = sess.run( [dynamic_state], - feed_dict={inputs[0]: input_value, - sequence_length: [2, 3]}) + feed_dict={ + inputs[0]: input_value, + sequence_length: [2, 3] + }) # outputs are fully calculated for t = 0, 1 for v in dynamic_values[:2]: @@ -289,8 +292,7 @@ class RNNTest(test.TestCase): input_size = 5 max_length = 8 # unrolled up to this length inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] return rnn.static_rnn(cell, inputs, dtype=dtypes.float32, scope=scope) @@ -316,8 +318,7 @@ class LSTMTest(test.TestCase): cell = rnn_cell.LSTMCell( num_units, initializer=initializer, state_is_tuple=False) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] outputs, _ = rnn.static_rnn(cell, inputs, dtype=dtypes.float32) self.assertEqual(len(outputs), len(inputs)) @@ -343,8 +344,7 @@ class LSTMTest(test.TestCase): initializer=initializer, state_is_tuple=False) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] outputs, _ = rnn.static_rnn(cell, inputs, dtype=dtypes.float32) self.assertEqual(len(outputs), len(inputs)) @@ -374,8 +374,7 @@ class LSTMTest(test.TestCase): initializer=initializer, state_is_tuple=False) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] with variable_scope.variable_scope("share_scope"): outputs, state = rnn.static_state_saving_rnn( @@ -388,7 +387,9 @@ class LSTMTest(test.TestCase): input_value = np.random.randn(batch_size, input_size) (last_state_value, saved_state_value) = sess.run( [state, state_saver.saved_state["save_lstm"]], - feed_dict={inputs[0]: input_value}) + feed_dict={ + inputs[0]: input_value + }) self.assertAllEqual(last_state_value, saved_state_value) def testNoProjNoShardingTupleStateSaver(self): @@ -406,8 +407,7 @@ class LSTMTest(test.TestCase): initializer=initializer, state_is_tuple=True) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] with variable_scope.variable_scope("share_scope"): outputs, state = rnn.static_state_saving_rnn( @@ -420,7 +420,9 @@ class LSTMTest(test.TestCase): input_value = np.random.randn(batch_size, input_size) last_and_saved_states = sess.run( state + (state_saver.saved_state["c"], state_saver.saved_state["m"]), - feed_dict={inputs[0]: input_value}) + feed_dict={ + inputs[0]: input_value + }) self.assertEqual(4, len(last_and_saved_states)) self.assertAllEqual(last_and_saved_states[:2], last_and_saved_states[2:]) @@ -432,16 +434,17 @@ class LSTMTest(test.TestCase): with self.test_session(graph=ops_lib.Graph()) as sess: initializer = init_ops.random_uniform_initializer( -0.01, 0.01, seed=self._seed) - state_saver = TestStateSaver(batch_size, { - "c0": num_units, - "m0": num_units, - "c1": num_units + 1, - "m1": num_units + 1, - "c2": num_units + 2, - "m2": num_units + 2, - "c3": num_units + 3, - "m3": num_units + 3 - }) + state_saver = TestStateSaver( + batch_size, { + "c0": num_units, + "m0": num_units, + "c1": num_units + 1, + "m1": num_units + 1, + "c2": num_units + 2, + "m2": num_units + 2, + "c3": num_units + 3, + "m3": num_units + 3 + }) def _cell(i): return rnn_cell.LSTMCell( @@ -459,8 +462,7 @@ class LSTMTest(test.TestCase): self.assertEqual(len(cell.state_size[i]), 2) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] state_names = (("c0", "m0"), ("c1", "m1"), ("c2", "m2"), ("c3", "m3")) @@ -475,10 +477,15 @@ class LSTMTest(test.TestCase): variables_lib.global_variables_initializer().run() input_value = np.random.randn(batch_size, input_size) - last_states = sess.run(list(nest.flatten(state)), - feed_dict={inputs[0]: input_value}) - saved_states = sess.run(list(state_saver.saved_state.values()), - feed_dict={inputs[0]: input_value}) + last_states = sess.run( + list(nest.flatten(state)), feed_dict={ + inputs[0]: input_value + }) + saved_states = sess.run( + list(state_saver.saved_state.values()), + feed_dict={ + inputs[0]: input_value + }) self.assertEqual(8, len(last_states)) self.assertEqual(8, len(saved_states)) flat_state_names = nest.flatten(state_names) @@ -499,8 +506,7 @@ class LSTMTest(test.TestCase): initializer = init_ops.random_uniform_initializer( -0.01, 0.01, seed=self._seed) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None, input_size)) + array_ops.placeholder(dtypes.float32, shape=(None, input_size)) ] cell = rnn_cell.LSTMCell( num_units, @@ -526,8 +532,7 @@ class LSTMTest(test.TestCase): initializer = init_ops.random_uniform_initializer( -0.01, 0.01, seed=self._seed) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None, input_size)) + array_ops.placeholder(dtypes.float32, shape=(None, input_size)) ] cell_notuple = rnn_cell.LSTMCell( num_units, @@ -569,14 +574,20 @@ class LSTMTest(test.TestCase): variables_lib.global_variables_initializer().run() input_value = np.random.randn(batch_size, input_size) - outputs_notuple_v = sess.run(outputs_notuple, - feed_dict={inputs[0]: input_value}) - outputs_tuple_v = sess.run(outputs_tuple, - feed_dict={inputs[0]: input_value}) + outputs_notuple_v = sess.run( + outputs_notuple, feed_dict={ + inputs[0]: input_value + }) + outputs_tuple_v = sess.run( + outputs_tuple, feed_dict={ + inputs[0]: input_value + }) self.assertAllEqual(outputs_notuple_v, outputs_tuple_v) - (state_notuple_v,) = sess.run((state_notuple,), - feed_dict={inputs[0]: input_value}) + (state_notuple_v,) = sess.run( + (state_notuple,), feed_dict={ + inputs[0]: input_value + }) state_tuple_v = sess.run(state_tuple, feed_dict={inputs[0]: input_value}) self.assertAllEqual(state_notuple_v, np.hstack(state_tuple_v)) @@ -593,8 +604,7 @@ class LSTMTest(test.TestCase): -0.01, 0.01, seed=self._seed) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None, input_size)) + array_ops.placeholder(dtypes.float32, shape=(None, input_size)) ] cell = rnn_cell.LSTMCell( @@ -625,8 +635,7 @@ class LSTMTest(test.TestCase): with self.test_session(use_gpu=True, graph=ops_lib.Graph()) as sess: initializer = init_ops.random_uniform_initializer(-1, 1, seed=self._seed) inputs = max_length * [ - array_ops.placeholder( - dtypes.float64, shape=(None, input_size)) + array_ops.placeholder(dtypes.float64, shape=(None, input_size)) ] cell = rnn_cell.LSTMCell( @@ -661,8 +670,7 @@ class LSTMTest(test.TestCase): max_length = 8 with self.test_session(use_gpu=True, graph=ops_lib.Graph()) as sess: inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None, input_size)) + array_ops.placeholder(dtypes.float32, shape=(None, input_size)) ] initializer = init_ops.constant_initializer(0.001) @@ -721,8 +729,7 @@ class LSTMTest(test.TestCase): initializer = init_ops.random_uniform_initializer( -0.01, 0.01, seed=self._seed) inputs = max_length * [ - array_ops.placeholder( - dtypes.float64, shape=(None, input_size)) + array_ops.placeholder(dtypes.float64, shape=(None, input_size)) ] cell = rnn_cell.LSTMCell( @@ -743,16 +750,21 @@ class LSTMTest(test.TestCase): self.assertEqual(len(outputs), len(inputs)) - variables_lib.global_variables_initializer().run( - feed_dict={sequence_length: [2, 3]}) + variables_lib.global_variables_initializer().run(feed_dict={ + sequence_length: [2, 3] + }) input_value = np.asarray( np.random.randn(batch_size, input_size), dtype=np.float64) values = sess.run( - outputs, feed_dict={inputs[0]: input_value, - sequence_length: [2, 3]}) + outputs, feed_dict={ + inputs[0]: input_value, + sequence_length: [2, 3] + }) state_value = sess.run( - [state], feed_dict={inputs[0]: input_value, - sequence_length: [2, 3]}) + [state], feed_dict={ + inputs[0]: input_value, + sequence_length: [2, 3] + }) self.assertEqual(values[0].dtype, input_value.dtype) self.assertEqual(state_value[0].dtype, input_value.dtype) @@ -767,8 +779,7 @@ class LSTMTest(test.TestCase): initializer_d = init_ops.random_uniform_initializer( -1, 1, seed=self._seed + 1) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None, input_size)) + array_ops.placeholder(dtypes.float32, shape=(None, input_size)) ] cell = rnn_cell.LSTMCell( num_units, @@ -792,8 +803,10 @@ class LSTMTest(test.TestCase): variables_lib.global_variables_initializer().run() input_value = np.random.randn(batch_size, input_size) - output_values = sess.run(outputs0 + outputs1 + outputs2, - feed_dict={inputs[0]: input_value}) + output_values = sess.run( + outputs0 + outputs1 + outputs2, feed_dict={ + inputs[0]: input_value + }) outputs0_values = output_values[:max_length] outputs1_values = output_values[max_length:2 * max_length] outputs2_values = output_values[2 * max_length:] @@ -814,8 +827,7 @@ class LSTMTest(test.TestCase): with self.test_session(graph=ops_lib.Graph()) as sess: initializer = init_ops.random_uniform_initializer(-1, 1, seed=self._seed) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None, input_size)) + array_ops.placeholder(dtypes.float32, shape=(None, input_size)) ] cell = rnn_cell.LSTMCell( num_units, @@ -833,8 +845,10 @@ class LSTMTest(test.TestCase): variables_lib.global_variables_initializer().run() input_value = np.random.randn(batch_size, input_size) - output_values = sess.run(outputs0 + outputs1, - feed_dict={inputs[0]: input_value}) + output_values = sess.run( + outputs0 + outputs1, feed_dict={ + inputs[0]: input_value + }) outputs0_values = output_values[:max_length] outputs1_values = output_values[max_length:] self.assertEqual(len(outputs0_values), len(outputs1_values)) @@ -861,8 +875,7 @@ class LSTMTest(test.TestCase): -0.01, 0.01, seed=self._seed) if in_graph_mode: inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None, input_size)) + array_ops.placeholder(dtypes.float32, shape=(None, input_size)) ] else: inputs = max_length * [ @@ -939,8 +952,7 @@ class LSTMTest(test.TestCase): -0.01, 0.01, seed=self._seed) if in_graph_mode: inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None, input_size)) + array_ops.placeholder(dtypes.float32, shape=(None, input_size)) ] else: inputs = max_length * [ @@ -1100,8 +1112,8 @@ class LSTMTest(test.TestCase): # Test gradients to inputs and variables w.r.t. outputs & final state static_grad_values = sess.run(static_gradients, feed_dict=feeds) - static_individual_grad_values = sess.run(static_individual_gradients, - feed_dict=feeds) + static_individual_grad_values = sess.run( + static_individual_gradients, feed_dict=feeds) static_individual_var_grad_values = sess.run( static_individual_variable_gradients, feed_dict=feeds) @@ -1148,8 +1160,10 @@ class LSTMTest(test.TestCase): # Generate gradients of several individual outputs w.r.t. inputs dynamic_individual_gradients = nest.flatten([ gradients_impl.gradients(y, [concat_inputs]) - for y in - [split_outputs_dynamic[0], split_outputs_dynamic[-1], state_dynamic] + for y in [ + split_outputs_dynamic[0], split_outputs_dynamic[-1], + state_dynamic + ] ]) # Generate gradients of individual variables w.r.t. inputs @@ -1159,8 +1173,10 @@ class LSTMTest(test.TestCase): "Count of trainable variables: %d" % len(trainable_variables)) dynamic_individual_variable_gradients = nest.flatten([ gradients_impl.gradients(y, trainable_variables) - for y in - [split_outputs_dynamic[0], split_outputs_dynamic[-1], state_dynamic] + for y in [ + split_outputs_dynamic[0], split_outputs_dynamic[-1], + state_dynamic + ] ]) # Test forward pass @@ -1170,8 +1186,8 @@ class LSTMTest(test.TestCase): # Test gradients to inputs and variables w.r.t. outputs & final state dynamic_grad_values = sess.run(dynamic_gradients, feed_dict=feeds) - dynamic_individual_grad_values = sess.run(dynamic_individual_gradients, - feed_dict=feeds) + dynamic_individual_grad_values = sess.run( + dynamic_individual_gradients, feed_dict=feeds) dynamic_individual_var_grad_values = sess.run( dynamic_individual_variable_gradients, feed_dict=feeds) @@ -1207,8 +1223,8 @@ class LSTMTest(test.TestCase): for i, (a, b) in enumerate( zip(static_individual_var_grad_values, dynamic_individual_var_grad_values)): - tf_logging.info("Comparing individual variable gradients iteration %d" % - i) + tf_logging.info( + "Comparing individual variable gradients iteration %d" % i) self.assertAllEqual(a, b) @test_util.run_in_graph_and_eager_modes() @@ -1223,10 +1239,7 @@ class BidirectionalRNNTest(test.TestCase): self._seed = 23489 np.random.seed(self._seed) - def _createBidirectionalRNN(self, - use_shape, - use_sequence_length, - scope=None): + def _createBidirectionalRNN(self, use_shape, use_sequence_length, scope=None): num_units = 3 input_size = 5 batch_size = 2 @@ -1270,8 +1283,10 @@ class BidirectionalRNNTest(test.TestCase): # Run with pre-specified sequence length of 2, 3 out, s_fw, s_bw = sess.run( [outputs, state_fw, state_bw], - feed_dict={inputs[0]: input_value, - sequence_length: [2, 3]}) + feed_dict={ + inputs[0]: input_value, + sequence_length: [2, 3] + }) # Since the forward and backward LSTM cells were initialized with the # same parameters, the forward and backward output has to be the same, @@ -1312,8 +1327,10 @@ class BidirectionalRNNTest(test.TestCase): input_value, inputs, outputs, state_fw, state_bw, _ = ( self._createBidirectionalRNN(use_shape, False)) variables_lib.global_variables_initializer().run() - out, s_fw, s_bw = sess.run([outputs, state_fw, state_bw], - feed_dict={inputs[0]: input_value}) + out, s_fw, s_bw = sess.run( + [outputs, state_fw, state_bw], feed_dict={ + inputs[0]: input_value + }) # Since the forward and backward LSTM cells were initialized with the # same parameters, the forward and backward output has to be the same, @@ -1396,13 +1413,11 @@ class BidirectionalRNNTest(test.TestCase): use_time_major, use_sequence_length): with self.test_session(use_gpu=True, graph=ops_lib.Graph()) as sess: input_value, inputs, outputs, state_fw, state_bw, sequence_length = ( - self._createBidirectionalDynamicRNN(use_shape, - use_state_tuple, use_time_major, - use_sequence_length)) + self._createBidirectionalDynamicRNN( + use_shape, use_state_tuple, use_time_major, use_sequence_length)) variables_lib.global_variables_initializer().run() # Run with pre-specified sequence length of 2, 3 - feed_dict = ( - {sequence_length: [2, 3]} if use_sequence_length else {}) + feed_dict = ({sequence_length: [2, 3]} if use_sequence_length else {}) feed_dict.update({inputs[0]: input_value}) if use_state_tuple: out, c_fw, m_fw, c_bw, m_bw = sess.run( @@ -1538,8 +1553,7 @@ class MultiDimensionalLSTMTest(test.TestCase): sequence_length = [4, 6] with self.test_session(graph=ops_lib.Graph()) as sess: inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(None,) + input_size) + array_ops.placeholder(dtypes.float32, shape=(None,) + input_size) ] inputs_using_dim = max_length * [ array_ops.placeholder( @@ -1585,14 +1599,22 @@ class MultiDimensionalLSTMTest(test.TestCase): input_total_size = (batch_size,) + input_size input_value = np.random.randn(*input_total_size) - outputs_static_v = sess.run(outputs_static, - feed_dict={inputs[0]: input_value}) - outputs_dynamic_v = sess.run(outputs_dynamic, - feed_dict={inputs[0]: input_value}) - outputs_bid_v = sess.run(outputs_bid, - feed_dict={inputs_using_dim[0]: input_value}) - outputs_sav_v = sess.run(outputs_sav, - feed_dict={inputs_using_dim[0]: input_value}) + outputs_static_v = sess.run( + outputs_static, feed_dict={ + inputs[0]: input_value + }) + outputs_dynamic_v = sess.run( + outputs_dynamic, feed_dict={ + inputs[0]: input_value + }) + outputs_bid_v = sess.run( + outputs_bid, feed_dict={ + inputs_using_dim[0]: input_value + }) + outputs_sav_v = sess.run( + outputs_sav, feed_dict={ + inputs_using_dim[0]: input_value + }) self.assertAllEqual(outputs_static_v, outputs_dynamic_v) self.assertAllEqual(outputs_static_v, outputs_sav_v) @@ -1602,16 +1624,26 @@ class MultiDimensionalLSTMTest(test.TestCase): outputs_bid_array = np.array(outputs_bid_v) self.assertAllEqual(outputs_static_array_double, outputs_bid_array) - state_static_v = sess.run(state_static, - feed_dict={inputs[0]: input_value}) - state_dynamic_v = sess.run(state_dynamic, - feed_dict={inputs[0]: input_value}) - state_bid_fw_v = sess.run(state_fw, - feed_dict={inputs_using_dim[0]: input_value}) - state_bid_bw_v = sess.run(state_bw, - feed_dict={inputs_using_dim[0]: input_value}) - state_sav_v = sess.run(state_sav, - feed_dict={inputs_using_dim[0]: input_value}) + state_static_v = sess.run( + state_static, feed_dict={ + inputs[0]: input_value + }) + state_dynamic_v = sess.run( + state_dynamic, feed_dict={ + inputs[0]: input_value + }) + state_bid_fw_v = sess.run( + state_fw, feed_dict={ + inputs_using_dim[0]: input_value + }) + state_bid_bw_v = sess.run( + state_bw, feed_dict={ + inputs_using_dim[0]: input_value + }) + state_sav_v = sess.run( + state_sav, feed_dict={ + inputs_using_dim[0]: input_value + }) self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_dynamic_v)) self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_sav_v)) self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_bid_fw_v)) @@ -1633,16 +1665,17 @@ class NestedLSTMTest(test.TestCase): with self.test_session(graph=ops_lib.Graph()) as sess: state_saver = TestStateSaver(batch_size, state_size) single_input = (array_ops.placeholder( - dtypes.float32, shape=(None, input_size)), array_ops.placeholder( - dtypes.float32, shape=(None, input_size))) + dtypes.float32, shape=(None, input_size)), + array_ops.placeholder( + dtypes.float32, shape=(None, input_size))) inputs = max_length * [single_input] inputs_c = (array_ops.stack([input_[0] for input_ in inputs]), array_ops.stack([input_[1] for input_ in inputs])) - single_input_using_dim = ( - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)), - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size))) + single_input_using_dim = (array_ops.placeholder( + dtypes.float32, shape=(batch_size, input_size)), + array_ops.placeholder( + dtypes.float32, + shape=(batch_size, input_size))) inputs_using_dim = max_length * [single_input_using_dim] # Create a cell for the whole test. This is fine because the cell has no @@ -1688,14 +1721,22 @@ class NestedLSTMTest(test.TestCase): input_total_size = (batch_size, input_size) input_value = (np.random.randn(*input_total_size), np.random.randn(*input_total_size)) - outputs_dynamic_v = sess.run(outputs_dynamic, - feed_dict={single_input: input_value}) - outputs_static_v = sess.run(outputs_static, - feed_dict={single_input: input_value}) - outputs_sav_v = sess.run(outputs_sav, - feed_dict={single_input_using_dim: input_value}) - outputs_bid_v = sess.run(outputs_bid, - feed_dict={single_input_using_dim: input_value}) + outputs_dynamic_v = sess.run( + outputs_dynamic, feed_dict={ + single_input: input_value + }) + outputs_static_v = sess.run( + outputs_static, feed_dict={ + single_input: input_value + }) + outputs_sav_v = sess.run( + outputs_sav, feed_dict={ + single_input_using_dim: input_value + }) + outputs_bid_v = sess.run( + outputs_bid, feed_dict={ + single_input_using_dim: input_value + }) self.assertAllEqual(outputs_static_v, np.transpose(outputs_dynamic_v, (1, 0, 2, 3))) @@ -1706,16 +1747,26 @@ class NestedLSTMTest(test.TestCase): outputs_bid_array = np.array(outputs_bid_v) self.assertAllEqual(outputs_static_array_double, outputs_bid_array) - state_dynamic_v = sess.run(state_dynamic, - feed_dict={single_input: input_value}) - state_static_v = sess.run(state_static, - feed_dict={single_input: input_value}) - state_bid_fw_v = sess.run(state_fw, - feed_dict={single_input_using_dim: input_value}) - state_bid_bw_v = sess.run(state_bw, - feed_dict={single_input_using_dim: input_value}) - state_sav_v = sess.run(state_sav, - feed_dict={single_input_using_dim: input_value}) + state_dynamic_v = sess.run( + state_dynamic, feed_dict={ + single_input: input_value + }) + state_static_v = sess.run( + state_static, feed_dict={ + single_input: input_value + }) + state_bid_fw_v = sess.run( + state_fw, feed_dict={ + single_input_using_dim: input_value + }) + state_bid_bw_v = sess.run( + state_bw, feed_dict={ + single_input_using_dim: input_value + }) + state_sav_v = sess.run( + state_sav, feed_dict={ + single_input_using_dim: input_value + }) self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_dynamic_v)) self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_sav_v)) self.assertAllEqual(np.hstack(state_static_v), np.hstack(state_bid_fw_v)) @@ -1764,8 +1815,7 @@ class StateSaverRNNTest(test.TestCase): initializer=initializer, state_is_tuple=False) inputs = max_length * [ - array_ops.placeholder( - dtypes.float32, shape=(batch_size, input_size)) + array_ops.placeholder(dtypes.float32, shape=(batch_size, input_size)) ] return rnn.static_state_saving_rnn( cell, @@ -1931,8 +1981,10 @@ class RawRNNTest(test.TestCase): (outputs_val, outputs_dynamic_rnn_val, final_state_val, final_state_dynamic_rnn_val) = sess.run( [outputs, outputs_dynamic_rnn, final_state, final_state_dynamic_rnn], - feed_dict={inputs: rand_input, - sequence_length: rand_seq_len}) + feed_dict={ + inputs: rand_input, + sequence_length: rand_seq_len + }) self.assertAllClose(outputs_dynamic_rnn_val, outputs_val) self.assertAllClose(final_state_dynamic_rnn_val, final_state_val) @@ -1945,12 +1997,16 @@ class RawRNNTest(test.TestCase): self.assertEqual(len(gradients), len(gradients_dynamic_rnn)) gradients_val = sess.run( gradients, - feed_dict={inputs: rand_input, - sequence_length: rand_seq_len}) + feed_dict={ + inputs: rand_input, + sequence_length: rand_seq_len + }) gradients_dynamic_rnn_val = sess.run( gradients_dynamic_rnn, - feed_dict={inputs: rand_input, - sequence_length: rand_seq_len}) + feed_dict={ + inputs: rand_input, + sequence_length: rand_seq_len + }) self.assertEqual(len(gradients_val), len(gradients_dynamic_rnn_val)) input_gradients_val = gradients_val[0] input_gradients_dynamic_rnn_val = gradients_dynamic_rnn_val[0] @@ -2067,14 +2123,13 @@ class RawRNNTest(test.TestCase): def loop_fn(time_, cell_output, cell_state, _): if cell_output is None: - emit_output = (array_ops.zeros( - [2, 3], dtype=dtypes.int32), array_ops.zeros( - [unknown_dim], dtype=dtypes.int64)) + emit_output = (array_ops.zeros([2, 3], dtype=dtypes.int32), + array_ops.zeros([unknown_dim], dtype=dtypes.int64)) next_state = cell.zero_state(batch_size, dtypes.float32) else: - emit_output = (array_ops.ones( - [batch_size, 2, 3], dtype=dtypes.int32), array_ops.ones( - [batch_size, unknown_dim], dtype=dtypes.int64)) + emit_output = (array_ops.ones([batch_size, 2, 3], dtype=dtypes.int32), + array_ops.ones( + [batch_size, unknown_dim], dtype=dtypes.int64)) next_state = cell_state elements_finished = array_ops.tile([time_ >= max_time], [batch_size]) finished = math_ops.reduce_all(elements_finished) @@ -2193,8 +2248,8 @@ class TensorArrayOnCorrectDeviceTest(test.TestCase): cell = rnn_cell.LSTMCell(num_units, use_peepholes=True) gpu_cell = DeviceWrapperCell(cell, cell_device) - inputs = np.random.randn(batch_size, time_steps, - input_size).astype(np.float32) + inputs = np.random.randn(batch_size, time_steps, input_size).astype( + np.float32) sequence_length = np.random.randint(0, time_steps, size=batch_size) if input_device is not None: @@ -2262,8 +2317,7 @@ class TensorArrayOnCorrectDeviceTest(test.TestCase): gpu_dev = test.gpu_device_name() run_metadata = self._execute_rnn_on( - rnn_device="/cpu:0", cell_device="/cpu:0", - input_device=gpu_dev) + rnn_device="/cpu:0", cell_device="/cpu:0", input_device=gpu_dev) cpu_stats, gpu_stats = self._retrieve_cpu_gpu_stats(run_metadata) def _assert_in(op_str, in_stats, out_stats): @@ -2278,8 +2332,7 @@ class TensorArrayOnCorrectDeviceTest(test.TestCase): return # Test requires access to a GPU gpu_dev = test.gpu_device_name() - run_metadata = self._execute_rnn_on( - input_device=gpu_dev) + run_metadata = self._execute_rnn_on(input_device=gpu_dev) cpu_stats, gpu_stats = self._retrieve_cpu_gpu_stats(run_metadata) def _assert_in(op_str, in_stats, out_stats): diff --git a/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py b/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py index 46823fa3643c5b4a3d857fa38d1a70792d97ca40..7b883ebc5d7756f1bdf445f900500a4b89e6cffd 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py @@ -53,14 +53,12 @@ class RNNCellTest(test.TestCase): batch_size = 3 input_size = 4 expected_output = np.array( - [[0.121753, 0.121753], - [0.103349, 0.103349], - [0.100178, 0.100178]], + [[0.121753, 0.121753], [0.103349, 0.103349], [0.100178, 0.100178]], dtype=np.float32) expected_state = np.array( - [[0.137523, 0.137523, 0.121753, 0.121753], - [0.105450, 0.105450, 0.103349, 0.103349], - [0.100742, 0.100742, 0.100178, 0.100178]], + [[0.137523, 0.137523, 0.121753, 0.121753], [ + 0.105450, 0.105450, 0.103349, 0.103349 + ], [0.100742, 0.100742, 0.100178, 0.100178]], dtype=np.float32) with variable_scope.variable_scope( "root", initializer=init_ops.constant_initializer(0.5)): @@ -69,14 +67,14 @@ class RNNCellTest(test.TestCase): output, state = contrib_rnn_cell.CoupledInputForgetGateLSTMCell( num_units=num_units, forget_bias=1.0, state_is_tuple=False)(x, m) sess.run([variables.global_variables_initializer()]) - res = sess.run([output, state], { - x.name: - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]]), - m.name: - 0.1 * np.ones((batch_size, state_size)) - }) + res = sess.run( + [output, state], { + x.name: + np.array([[1., 1., 1., 1.], [2., 2., 2., 2.], + [3., 3., 3., 3.]]), + m.name: + 0.1 * np.ones((batch_size, state_size)) + }) # This is a smoke test: Only making sure expected values didn't change. self.assertEqual(len(res), 2) self.assertAllClose(res[0], expected_output) @@ -101,14 +99,14 @@ class RNNCellTest(test.TestCase): frequency_skip=frequency_skip, forget_bias=1.0)(x, m) sess.run([variables.global_variables_initializer()]) - res = sess.run([output, state], { - x.name: - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]]), - m.name: - 0.1 * np.ones((batch_size, int(state_size * (num_shifts)))) - }) + res = sess.run( + [output, state], { + x.name: + np.array([[1., 1., 1., 1.], [2., 2., 2., 2.], + [3., 3., 3., 3.]]), + m.name: + 0.1 * np.ones((batch_size, int(state_size * (num_shifts)))) + }) self.assertEqual(len(res), 2) # The numbers in results were not calculated, this is mostly just a # smoke test. @@ -141,17 +139,14 @@ class RNNCellTest(test.TestCase): state_is_tuple=True) inputs = constant_op.constant( np.array( - [[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]], + [[1., 1., 1., 1.], [2., 2., 2., 2.], [3., 3., 3., 3.]], dtype=np.float32), dtype=dtypes.float32) state_value = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) - init_state = cell.state_tuple_type( - *([state_value, state_value] * num_shifts)) + init_state = cell.state_tuple_type(*( + [state_value, state_value] * num_shifts)) output, state = cell(inputs, init_state) sess.run([variables.global_variables_initializer()]) res = sess.run([output, state]) @@ -198,11 +193,10 @@ class RNNCellTest(test.TestCase): dtype=np.float32), dtype=dtypes.float32) state_value = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) - init_state = cell.state_tuple_type( - *([state_value, state_value] * total_blocks)) + init_state = cell.state_tuple_type(*( + [state_value, state_value] * total_blocks)) output, state = cell(inputs, init_state) sess.run([variables.global_variables_initializer()]) res = sess.run([output, state]) @@ -230,20 +224,28 @@ class RNNCellTest(test.TestCase): frequency_skip = 1 num_shifts = int((input_size - feature_size) / frequency_skip + 1) expected_output = np.array( - [[0.416383, 0.416383, 0.403238, 0.403238, 0.524020, 0.524020, - 0.565425, 0.565425, 0.557865, 0.557865, 0.609699, 0.609699], - [0.627331, 0.627331, 0.622393, 0.622393, 0.688342, 0.688342, - 0.708078, 0.708078, 0.694245, 0.694245, 0.715171, 0.715171], - [0.711050, 0.711050, 0.709197, 0.709197, 0.736533, 0.736533, - 0.744264, 0.744264, 0.737390, 0.737390, 0.745250, 0.745250]], + [[ + 0.416383, 0.416383, 0.403238, 0.403238, 0.524020, 0.524020, + 0.565425, 0.565425, 0.557865, 0.557865, 0.609699, 0.609699 + ], [ + 0.627331, 0.627331, 0.622393, 0.622393, 0.688342, 0.688342, + 0.708078, 0.708078, 0.694245, 0.694245, 0.715171, 0.715171 + ], [ + 0.711050, 0.711050, 0.709197, 0.709197, 0.736533, 0.736533, + 0.744264, 0.744264, 0.737390, 0.737390, 0.745250, 0.745250 + ]], dtype=np.float32) expected_state = np.array( - [[0.625556, 0.625556, 0.416383, 0.416383, 0.759134, 0.759134, - 0.524020, 0.524020, 0.798795, 0.798795, 0.557865, 0.557865], - [0.875488, 0.875488, 0.627331, 0.627331, 0.936432, 0.936432, - 0.688342, 0.688342, 0.941961, 0.941961, 0.694245, 0.694245], - [0.957327, 0.957327, 0.711050, 0.711050, 0.979522, 0.979522, - 0.736533, 0.736533, 0.980245, 0.980245, 0.737390, 0.737390]], + [[ + 0.625556, 0.625556, 0.416383, 0.416383, 0.759134, 0.759134, + 0.524020, 0.524020, 0.798795, 0.798795, 0.557865, 0.557865 + ], [ + 0.875488, 0.875488, 0.627331, 0.627331, 0.936432, 0.936432, + 0.688342, 0.688342, 0.941961, 0.941961, 0.694245, 0.694245 + ], [ + 0.957327, 0.957327, 0.711050, 0.711050, 0.979522, 0.979522, + 0.736533, 0.736533, 0.980245, 0.980245, 0.737390, 0.737390 + ]], dtype=np.float32) for state_is_tuple in [False, True]: with self.test_session() as sess: @@ -259,18 +261,16 @@ class RNNCellTest(test.TestCase): couple_input_forget_gates=True, state_is_tuple=state_is_tuple) inputs = constant_op.constant( - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]], - dtype=np.float32), + np.array( + [[1., 1., 1., 1.], [2., 2., 2., 2.], [3., 3., 3., 3.]], + dtype=np.float32), dtype=dtypes.float32) if state_is_tuple: state_value = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) - init_state = cell.state_tuple_type( - *([state_value, state_value] * num_shifts)) + init_state = cell.state_tuple_type(*( + [state_value, state_value] * num_shifts)) else: init_state = constant_op.constant( 0.1 * np.ones( @@ -302,32 +302,40 @@ class RNNCellTest(test.TestCase): frequency_skip = 1 num_shifts = int((input_size - feature_size) / frequency_skip + 1) expected_output = np.array( - [[0.464130, 0.464130, 0.419165, 0.419165, 0.593283, 0.593283, - 0.738350, 0.738350, 0.661638, 0.661638, 0.866774, 0.866774, - 0.520789, 0.520789, 0.476968, 0.476968, 0.604341, 0.604341, - 0.760207, 0.760207, 0.635773, 0.635773, 0.850218, 0.850218], - [0.669636, 0.669636, 0.628966, 0.628966, 0.736057, 0.736057, - 0.895927, 0.895927, 0.755559, 0.755559, 0.954359, 0.954359, - 0.692621, 0.692621, 0.652363, 0.652363, 0.737517, 0.737517, - 0.899558, 0.899558, 0.745984, 0.745984, 0.946840, 0.946840], - [0.751109, 0.751109, 0.711716, 0.711716, 0.778357, 0.778357, - 0.940779, 0.940779, 0.784530, 0.784530, 0.980604, 0.980604, - 0.759940, 0.759940, 0.720652, 0.720652, 0.778552, 0.778552, - 0.941606, 0.941606, 0.781035, 0.781035, 0.977731, 0.977731]], + [[ + 0.464130, 0.464130, 0.419165, 0.419165, 0.593283, 0.593283, + 0.738350, 0.738350, 0.661638, 0.661638, 0.866774, 0.866774, + 0.520789, 0.520789, 0.476968, 0.476968, 0.604341, 0.604341, + 0.760207, 0.760207, 0.635773, 0.635773, 0.850218, 0.850218 + ], [ + 0.669636, 0.669636, 0.628966, 0.628966, 0.736057, 0.736057, + 0.895927, 0.895927, 0.755559, 0.755559, 0.954359, 0.954359, + 0.692621, 0.692621, 0.652363, 0.652363, 0.737517, 0.737517, + 0.899558, 0.899558, 0.745984, 0.745984, 0.946840, 0.946840 + ], [ + 0.751109, 0.751109, 0.711716, 0.711716, 0.778357, 0.778357, + 0.940779, 0.940779, 0.784530, 0.784530, 0.980604, 0.980604, + 0.759940, 0.759940, 0.720652, 0.720652, 0.778552, 0.778552, + 0.941606, 0.941606, 0.781035, 0.781035, 0.977731, 0.977731 + ]], dtype=np.float32) expected_state = np.array( - [[0.710660, 0.710660, 0.464130, 0.464130, 0.877293, 0.877293, - 0.593283, 0.593283, 0.958505, 0.958505, 0.661638, 0.661638, - 0.785405, 0.785405, 0.520789, 0.520789, 0.890836, 0.890836, - 0.604341, 0.604341, 0.928512, 0.928512, 0.635773, 0.635773], - [0.967579, 0.967579, 0.669636, 0.669636, 1.038811, 1.038811, - 0.736057, 0.736057, 1.058201, 1.058201, 0.755559, 0.755559, - 0.993088, 0.993088, 0.692621, 0.692621, 1.040288, 1.040288, - 0.737517, 0.737517, 1.048773, 1.048773, 0.745984, 0.745984], - [1.053842, 1.053842, 0.751109, 0.751109, 1.079919, 1.079919, - 0.778357, 0.778357, 1.085620, 1.085620, 0.784530, 0.784530, - 1.062455, 1.062455, 0.759940, 0.759940, 1.080101, 1.080101, - 0.778552, 0.778552, 1.082402, 1.082402, 0.781035, 0.781035]], + [[ + 0.710660, 0.710660, 0.464130, 0.464130, 0.877293, 0.877293, + 0.593283, 0.593283, 0.958505, 0.958505, 0.661638, 0.661638, + 0.785405, 0.785405, 0.520789, 0.520789, 0.890836, 0.890836, + 0.604341, 0.604341, 0.928512, 0.928512, 0.635773, 0.635773 + ], [ + 0.967579, 0.967579, 0.669636, 0.669636, 1.038811, 1.038811, + 0.736057, 0.736057, 1.058201, 1.058201, 0.755559, 0.755559, + 0.993088, 0.993088, 0.692621, 0.692621, 1.040288, 1.040288, + 0.737517, 0.737517, 1.048773, 1.048773, 0.745984, 0.745984 + ], [ + 1.053842, 1.053842, 0.751109, 0.751109, 1.079919, 1.079919, + 0.778357, 0.778357, 1.085620, 1.085620, 0.784530, 0.784530, + 1.062455, 1.062455, 0.759940, 0.759940, 1.080101, 1.080101, + 0.778552, 0.778552, 1.082402, 1.082402, 0.781035, 0.781035 + ]], dtype=np.float32) with variable_scope.variable_scope( "root", initializer=init_ops.constant_initializer(0.5)): @@ -339,17 +347,16 @@ class RNNCellTest(test.TestCase): forget_bias=1.0, num_frequency_blocks=[num_shifts]) inputs = constant_op.constant( - np.array([[1.0, 1.1, 1.2, 1.3], - [2.0, 2.1, 2.2, 2.3], - [3.0, 3.1, 3.2, 3.3]], - dtype=np.float32), + np.array( + [[1.0, 1.1, 1.2, 1.3], [2.0, 2.1, 2.2, 2.3], + [3.0, 3.1, 3.2, 3.3]], + dtype=np.float32), dtype=dtypes.float32) state_value = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) - init_state = cell.state_tuple_type( - *([state_value, state_value] * num_shifts * 2)) + init_state = cell.state_tuple_type(*( + [state_value, state_value] * num_shifts * 2)) output, state = cell(inputs, init_state) sess.run([variables.global_variables_initializer()]) res = sess.run([output, state]) @@ -375,32 +382,40 @@ class RNNCellTest(test.TestCase): frequency_skip = 1 num_shifts = int((input_size - feature_size) / frequency_skip + 1) expected_output = np.array( - [[0.464130, 0.464130, 0.419165, 0.419165, 0.593283, 0.593283, - 0.738350, 0.738350, 0.661638, 0.661638, 0.866774, 0.866774, - 0.322645, 0.322645, 0.276068, 0.276068, 0.584654, 0.584654, - 0.690292, 0.690292, 0.640446, 0.640446, 0.840071, 0.840071], - [0.669636, 0.669636, 0.628966, 0.628966, 0.736057, 0.736057, - 0.895927, 0.895927, 0.755559, 0.755559, 0.954359, 0.954359, - 0.493625, 0.493625, 0.449236, 0.449236, 0.730828, 0.730828, - 0.865996, 0.865996, 0.749429, 0.749429, 0.944958, 0.944958], - [0.751109, 0.751109, 0.711716, 0.711716, 0.778357, 0.778357, - 0.940779, 0.940779, 0.784530, 0.784530, 0.980604, 0.980604, - 0.608587, 0.608587, 0.566683, 0.566683, 0.777345, 0.777345, - 0.925820, 0.925820, 0.782597, 0.782597, 0.976858, 0.976858]], + [[ + 0.464130, 0.464130, 0.419165, 0.419165, 0.593283, 0.593283, + 0.738350, 0.738350, 0.661638, 0.661638, 0.866774, 0.866774, + 0.322645, 0.322645, 0.276068, 0.276068, 0.584654, 0.584654, + 0.690292, 0.690292, 0.640446, 0.640446, 0.840071, 0.840071 + ], [ + 0.669636, 0.669636, 0.628966, 0.628966, 0.736057, 0.736057, + 0.895927, 0.895927, 0.755559, 0.755559, 0.954359, 0.954359, + 0.493625, 0.493625, 0.449236, 0.449236, 0.730828, 0.730828, + 0.865996, 0.865996, 0.749429, 0.749429, 0.944958, 0.944958 + ], [ + 0.751109, 0.751109, 0.711716, 0.711716, 0.778357, 0.778357, + 0.940779, 0.940779, 0.784530, 0.784530, 0.980604, 0.980604, + 0.608587, 0.608587, 0.566683, 0.566683, 0.777345, 0.777345, + 0.925820, 0.925820, 0.782597, 0.782597, 0.976858, 0.976858 + ]], dtype=np.float32) expected_state = np.array( - [[0.710660, 0.710660, 0.464130, 0.464130, 0.877293, 0.877293, - 0.593283, 0.593283, 0.958505, 0.958505, 0.661638, 0.661638, - 0.516575, 0.516575, 0.322645, 0.322645, 0.866628, 0.866628, - 0.584654, 0.584654, 0.934002, 0.934002, 0.640446, 0.640446], - [0.967579, 0.967579, 0.669636, 0.669636, 1.038811, 1.038811, - 0.736057, 0.736057, 1.058201, 1.058201, 0.755559, 0.755559, - 0.749836, 0.749836, 0.493625, 0.493625, 1.033488, 1.033488, - 0.730828, 0.730828, 1.052186, 1.052186, 0.749429, 0.749429], - [1.053842, 1.053842, 0.751109, 0.751109, 1.079919, 1.079919, - 0.778357, 0.778357, 1.085620, 1.085620, 0.784530, 0.784530, - 0.895999, 0.895999, 0.608587, 0.608587, 1.078978, 1.078978, - 0.777345, 0.777345, 1.083843, 1.083843, 0.782597, 0.782597]], + [[ + 0.710660, 0.710660, 0.464130, 0.464130, 0.877293, 0.877293, + 0.593283, 0.593283, 0.958505, 0.958505, 0.661638, 0.661638, + 0.516575, 0.516575, 0.322645, 0.322645, 0.866628, 0.866628, + 0.584654, 0.584654, 0.934002, 0.934002, 0.640446, 0.640446 + ], [ + 0.967579, 0.967579, 0.669636, 0.669636, 1.038811, 1.038811, + 0.736057, 0.736057, 1.058201, 1.058201, 0.755559, 0.755559, + 0.749836, 0.749836, 0.493625, 0.493625, 1.033488, 1.033488, + 0.730828, 0.730828, 1.052186, 1.052186, 0.749429, 0.749429 + ], [ + 1.053842, 1.053842, 0.751109, 0.751109, 1.079919, 1.079919, + 0.778357, 0.778357, 1.085620, 1.085620, 0.784530, 0.784530, + 0.895999, 0.895999, 0.608587, 0.608587, 1.078978, 1.078978, + 0.777345, 0.777345, 1.083843, 1.083843, 0.782597, 0.782597 + ]], dtype=np.float32) with variable_scope.variable_scope( "root", initializer=init_ops.constant_initializer(0.5)): @@ -413,17 +428,16 @@ class RNNCellTest(test.TestCase): num_frequency_blocks=[num_shifts], backward_slice_offset=1) inputs = constant_op.constant( - np.array([[1.0, 1.1, 1.2, 1.3], - [2.0, 2.1, 2.2, 2.3], - [3.0, 3.1, 3.2, 3.3]], - dtype=np.float32), + np.array( + [[1.0, 1.1, 1.2, 1.3], [2.0, 2.1, 2.2, 2.3], + [3.0, 3.1, 3.2, 3.3]], + dtype=np.float32), dtype=dtypes.float32) state_value = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) - init_state = cell.state_tuple_type( - *([state_value, state_value] * num_shifts * 2)) + init_state = cell.state_tuple_type(*( + [state_value, state_value] * num_shifts * 2)) output, state = cell(inputs, init_state) sess.run([variables.global_variables_initializer()]) res = sess.run([output, state]) @@ -474,8 +488,8 @@ class RNNCellTest(test.TestCase): for state_is_tuple in [False, True]: with ops.Graph().as_default(): with self.test_session() as sess: - with variable_scope.variable_scope("state_is_tuple_" + str( - state_is_tuple)): + with variable_scope.variable_scope( + "state_is_tuple_" + str(state_is_tuple)): lstm_cell = rnn_cell.BasicLSTMCell( num_units, state_is_tuple=state_is_tuple) cell = contrib_rnn_cell.AttentionCellWrapper( @@ -525,16 +539,15 @@ class RNNCellTest(test.TestCase): for state_is_tuple in [False, True]: with ops.Graph().as_default(): with self.test_session() as sess: - with variable_scope.variable_scope("state_is_tuple_" + str( - state_is_tuple)): + with variable_scope.variable_scope( + "state_is_tuple_" + str(state_is_tuple)): lstm_cell = rnn_cell.BasicLSTMCell( num_units, state_is_tuple=state_is_tuple) cell = contrib_rnn_cell.AttentionCellWrapper( lstm_cell, attn_length, state_is_tuple=state_is_tuple) if state_is_tuple: zeros = constant_op.constant( - 0.1 * np.ones( - [batch_size, num_units], dtype=np.float32), + 0.1 * np.ones([batch_size, num_units], dtype=np.float32), dtype=dtypes.float32) attn_state_zeros = constant_op.constant( 0.1 * np.ones( @@ -579,22 +592,25 @@ class RNNCellTest(test.TestCase): [1.018088, 0.378983, -0.572179, 0.268591]], dtype=np.float32) expected_state = np.array( - [[0.74946702, 0.34681597, 0.26474735, 1.06485605, 0.38465962, - 0.11420801, 0.10272158, 0.30925757, 0.63899988, 0.7181077, - 0.47534478, 0.33715725, 0.58086717, 0.49446869, 0.7641536, - 0.12814975, 0.92231739, 0.89857256, 0.21889746, 0.38442063, - 0.53481543, 0.8876909, 0.45823169, 0.5905602, 0.78038228, - 0.56501579, 0.03971386, 0.09870267, 0.8074435, 0.66821432, - 0.99211812, 0.12295902, 1.14606023, 0.34370938, -0.79251152, - 0.51843399], - [0.5179342, 0.48682183, -0.25426468, 0.96810579, 0.28809637, - 0.13607743, -0.11446252, 0.26792109, 0.78047138, 0.63460857, - 0.49122369, 0.52007174, 0.73000264, 0.66986895, 0.73576689, - 0.86301267, 0.87887371, 0.35185754, 0.93417215, 0.64732957, - 0.63173044, 0.66627824, 0.53644657, 0.20477486, 0.98458421, - 0.38277245, 0.03746676, 0.92510188, 0.57714164, 0.84932971, - 0.36127412, 0.12125921, 1.1362772, 0.34361625, -0.78150457, - 0.70582712]], + [[ + 0.74946702, 0.34681597, 0.26474735, 1.06485605, 0.38465962, + 0.11420801, 0.10272158, 0.30925757, 0.63899988, 0.7181077, + 0.47534478, 0.33715725, 0.58086717, 0.49446869, 0.7641536, + 0.12814975, 0.92231739, 0.89857256, 0.21889746, 0.38442063, + 0.53481543, 0.8876909, 0.45823169, 0.5905602, 0.78038228, + 0.56501579, 0.03971386, 0.09870267, 0.8074435, 0.66821432, + 0.99211812, 0.12295902, 1.14606023, 0.34370938, -0.79251152, + 0.51843399 + ], [ + 0.5179342, 0.48682183, -0.25426468, 0.96810579, 0.28809637, + 0.13607743, -0.11446252, 0.26792109, 0.78047138, 0.63460857, + 0.49122369, 0.52007174, 0.73000264, 0.66986895, 0.73576689, + 0.86301267, 0.87887371, 0.35185754, 0.93417215, 0.64732957, + 0.63173044, 0.66627824, 0.53644657, 0.20477486, 0.98458421, + 0.38277245, 0.03746676, 0.92510188, 0.57714164, 0.84932971, + 0.36127412, 0.12125921, 1.1362772, 0.34361625, -0.78150457, + 0.70582712 + ]], dtype=np.float32) seed = 12345 random_seed.set_random_seed(seed) @@ -602,7 +618,8 @@ class RNNCellTest(test.TestCase): for state_is_tuple in [False, True]: with session.Session() as sess: with variable_scope.variable_scope( - "state_is_tuple", reuse=state_is_tuple, + "state_is_tuple", + reuse=state_is_tuple, initializer=init_ops.glorot_uniform_initializer()): lstm_cell = rnn_cell.BasicLSTMCell( num_units, state_is_tuple=state_is_tuple) @@ -646,36 +663,31 @@ class RNNCellTest(test.TestCase): def testNASCell(self): num_units = 6 batch_size = 3 - expected_output = np.array([[0.576751, 0.576751, 0.576751, 0.576751, - 0.576751, 0.576751], - [0.618936, 0.618936, 0.618936, 0.618936, - 0.618936, 0.618936], - [0.627393, 0.627393, 0.627393, 0.627393, - 0.627393, 0.627393]]) - expected_state = np.array([[0.71579772, 0.71579772, 0.71579772, 0.71579772, - 0.71579772, 0.71579772, 0.57675087, 0.57675087, - 0.57675087, 0.57675087, 0.57675087, 0.57675087], - [0.78041625, 0.78041625, 0.78041625, 0.78041625, - 0.78041625, 0.78041625, 0.6189357, 0.6189357, - 0.61893570, 0.6189357, 0.6189357, 0.6189357], - [0.79457647, 0.79457647, 0.79457647, 0.79457647, - 0.79457653, 0.79457653, 0.62739348, 0.62739348, - 0.62739348, 0.62739348, 0.62739348, 0.62739348] - ]) + expected_output = np.array( + [[0.576751, 0.576751, 0.576751, 0.576751, 0.576751, 0.576751], + [0.618936, 0.618936, 0.618936, 0.618936, 0.618936, 0.618936], + [0.627393, 0.627393, 0.627393, 0.627393, 0.627393, 0.627393]]) + expected_state = np.array([[ + 0.71579772, 0.71579772, 0.71579772, 0.71579772, 0.71579772, 0.71579772, + 0.57675087, 0.57675087, 0.57675087, 0.57675087, 0.57675087, 0.57675087 + ], [ + 0.78041625, 0.78041625, 0.78041625, 0.78041625, 0.78041625, 0.78041625, + 0.6189357, 0.6189357, 0.61893570, 0.6189357, 0.6189357, 0.6189357 + ], [ + 0.79457647, 0.79457647, 0.79457647, 0.79457647, 0.79457653, 0.79457653, + 0.62739348, 0.62739348, 0.62739348, 0.62739348, 0.62739348, 0.62739348 + ]]) with self.test_session() as sess: with variable_scope.variable_scope( - "nas_test", - initializer=init_ops.constant_initializer(0.5)): + "nas_test", initializer=init_ops.constant_initializer(0.5)): cell = contrib_rnn_cell.NASCell(num_units=num_units) inputs = constant_op.constant( - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]], - dtype=np.float32), + np.array( + [[1., 1., 1., 1.], [2., 2., 2., 2.], [3., 3., 3., 3.]], + dtype=np.float32), dtype=dtypes.float32) state_value = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) init_state = rnn_cell.LSTMStateTuple(state_value, state_value) output, state = cell(inputs, init_state) @@ -699,39 +711,34 @@ class RNNCellTest(test.TestCase): num_units = 6 batch_size = 3 num_proj = 5 - expected_output = np.array([[1.697418, 1.697418, 1.697418, 1.697418, - 1.697418], - [1.840037, 1.840037, 1.840037, 1.840037, - 1.840037], - [1.873985, 1.873985, 1.873985, 1.873985, - 1.873985]]) - expected_state = np.array([[0.69855207, 0.69855207, 0.69855207, 0.69855207, - 0.69855207, 0.69855207, 1.69741797, 1.69741797, - 1.69741797, 1.69741797, 1.69741797], - [0.77073824, 0.77073824, 0.77073824, 0.77073824, - 0.77073824, 0.77073824, 1.84003687, 1.84003687, - 1.84003687, 1.84003687, 1.84003687], - [0.78973997, 0.78973997, 0.78973997, 0.78973997, - 0.78973997, 0.78973997, 1.87398517, 1.87398517, - 1.87398517, 1.87398517, 1.87398517]]) + expected_output = np.array( + [[1.697418, 1.697418, 1.697418, 1.697418, + 1.697418], [1.840037, 1.840037, 1.840037, 1.840037, 1.840037], + [1.873985, 1.873985, 1.873985, 1.873985, 1.873985]]) + expected_state = np.array([[ + 0.69855207, 0.69855207, 0.69855207, 0.69855207, 0.69855207, 0.69855207, + 1.69741797, 1.69741797, 1.69741797, 1.69741797, 1.69741797 + ], [ + 0.77073824, 0.77073824, 0.77073824, 0.77073824, 0.77073824, 0.77073824, + 1.84003687, 1.84003687, 1.84003687, 1.84003687, 1.84003687 + ], [ + 0.78973997, 0.78973997, 0.78973997, 0.78973997, 0.78973997, 0.78973997, + 1.87398517, 1.87398517, 1.87398517, 1.87398517, 1.87398517 + ]]) with self.test_session() as sess: with variable_scope.variable_scope( - "nas_proj_test", - initializer=init_ops.constant_initializer(0.5)): + "nas_proj_test", initializer=init_ops.constant_initializer(0.5)): cell = contrib_rnn_cell.NASCell(num_units=num_units, num_proj=num_proj) inputs = constant_op.constant( - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]], - dtype=np.float32), + np.array( + [[1., 1., 1., 1.], [2., 2., 2., 2.], [3., 3., 3., 3.]], + dtype=np.float32), dtype=dtypes.float32) state_value_c = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) state_value_h = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_proj), dtype=np.float32), + 0.1 * np.ones((batch_size, num_proj), dtype=np.float32), dtype=dtypes.float32) init_state = rnn_cell.LSTMStateTuple(state_value_c, state_value_h) output, state = cell(inputs, init_state) @@ -755,24 +762,20 @@ class RNNCellTest(test.TestCase): num_units = 2 batch_size = 3 expected_state_and_output = np.array( - [[0.13752282, 0.13752282], - [0.10545051, 0.10545051], + [[0.13752282, 0.13752282], [0.10545051, 0.10545051], [0.10074195, 0.10074195]], dtype=np.float32) with self.test_session() as sess: with variable_scope.variable_scope( - "ugrnn_cell_test", - initializer=init_ops.constant_initializer(0.5)): + "ugrnn_cell_test", initializer=init_ops.constant_initializer(0.5)): cell = contrib_rnn_cell.UGRNNCell(num_units=num_units) inputs = constant_op.constant( - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]], - dtype=np.float32), + np.array( + [[1., 1., 1., 1.], [2., 2., 2., 2.], [3., 3., 3., 3.]], + dtype=np.float32), dtype=dtypes.float32) init_state = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) output, state = cell(inputs, init_state) sess.run([variables.global_variables_initializer()]) @@ -786,13 +789,11 @@ class RNNCellTest(test.TestCase): num_units = 2 batch_size = 3 expected_state = np.array( - [[0.13752282, 0.13752282], - [0.10545051, 0.10545051], + [[0.13752282, 0.13752282], [0.10545051, 0.10545051], [0.10074195, 0.10074195]], dtype=np.float32) expected_output = np.array( - [[2.00431061, 2.00431061], - [4.00060606, 4.00060606], + [[2.00431061, 2.00431061], [4.00060606, 4.00060606], [6.00008249, 6.00008249]], dtype=np.float32) with self.test_session() as sess: @@ -802,14 +803,12 @@ class RNNCellTest(test.TestCase): cell = contrib_rnn_cell.IntersectionRNNCell( num_units=num_units, num_in_proj=num_units) inputs = constant_op.constant( - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]], - dtype=np.float32), + np.array( + [[1., 1., 1., 1.], [2., 2., 2., 2.], [3., 3., 3., 3.]], + dtype=np.float32), dtype=dtypes.float32) init_state = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) output, state = cell(inputs, init_state) sess.run([variables.global_variables_initializer()]) @@ -824,19 +823,17 @@ class RNNCellTest(test.TestCase): batch_size = 3 cell = contrib_rnn_cell.IntersectionRNNCell(num_units=num_units) inputs = constant_op.constant( - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]], - dtype=np.float32), + np.array( + [[1., 1., 1., 1.], [2., 2., 2., 2.], [3., 3., 3., 3.]], + dtype=np.float32), dtype=dtypes.float32) init_state = constant_op.constant( - 0.1 * np.ones( - (batch_size, num_units), dtype=np.float32), + 0.1 * np.ones((batch_size, num_units), dtype=np.float32), dtype=dtypes.float32) - with self.assertRaisesRegexp( - ValueError, "Must have input size == output size for " - "Intersection RNN. To fix, num_in_proj should " - "be set to num_units at cell init."): + with self.assertRaisesRegexp(ValueError, + "Must have input size == output size for " + "Intersection RNN. To fix, num_in_proj should " + "be set to num_units at cell init."): cell(inputs, init_state) def testPhasedLSTMCell(self): @@ -845,12 +842,12 @@ class RNNCellTest(test.TestCase): batch_size = 3 input_size = 4 expected_state_c = np.array( - [[0.00072015, 0.00036633], [0.00083481, 0.00047266], - [0.00085111, 0.00053054]], + [[6.450831e-04, 4.697885e-04], [9.862894e-05, 7.212213e-04], + [4.401947e-04, 9.143004e-04]], dtype=np.float32) expected_state_h = np.array( - [[0.0005159, 0.00026243], [0.00062958, 0.00035646], - [0.00064732, 0.00040351]], + [[4.621217e-04, 3.365449e-04], [7.438179e-05, 5.439147e-04], + [3.347936e-04, 6.953785e-04]], dtype=np.float32) with variable_scope.variable_scope( "root", initializer=init_ops.constant_initializer(0.5)): @@ -862,14 +859,14 @@ class RNNCellTest(test.TestCase): output, state = contrib_rnn_cell.PhasedLSTMCell(num_units=num_units)( (t, x), state0) sess.run([variables.global_variables_initializer()]) - res = sess.run([output, state], { - t.name: - np.array([[1.], [2.], [3.]]), - x.name: - np.array([[1., 1., 1., 1.], - [2., 2., 2., 2.], - [3., 3., 3., 3.]]), - }) + res = sess.run( + [output, state], { + t.name: + np.array([[1.], [2.], [3.]]), + x.name: + np.array([[1., 1., 1., 1.], [2., 2., 2., 2.], + [3., 3., 3., 3.]]), + }) # This is a smoke test, making sure expected values are unchanged. self.assertEqual(len(res), 2) self.assertAllClose(res[0], res[1].h) @@ -878,36 +875,32 @@ class RNNCellTest(test.TestCase): def testConv1DLSTMCell(self): with self.test_session() as sess: - shape = [2,1] + shape = [2, 1] filter_size = [3] num_features = 1 batch_size = 2 expected_state_c = np.array( - [[[1.4375670191], [1.4375670191]], - [[2.7542609292], [2.7542609292]]], + [[[1.4375670191], [1.4375670191]], [[2.7542609292], [2.7542609292]]], dtype=np.float32) expected_state_h = np.array( - [[[0.6529865603], [0.6529865603]], - [[0.8736877431], [0.8736877431]]], + [[[0.6529865603], [0.6529865603]], [[0.8736877431], [0.8736877431]]], dtype=np.float32) with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(1.0/2.0)): + "root", initializer=init_ops.constant_initializer(1.0 / 2.0)): x = array_ops.placeholder(dtypes.float32, [None, None, 1]) - cell = contrib_rnn_cell.Conv1DLSTMCell(input_shape=shape, - kernel_shape=filter_size, - output_channels=num_features) + cell = contrib_rnn_cell.Conv1DLSTMCell( + input_shape=shape, + kernel_shape=filter_size, + output_channels=num_features) hidden = cell.zero_state(array_ops.shape(x)[0], dtypes.float32) output, state = cell(x, hidden) sess.run([variables.global_variables_initializer()]) - res = sess.run([output, state], { - hidden[0].name: - np.array([[[1.],[1.]], - [[2.],[2.]]]), - x.name: - np.array([[[1.],[1.]], - [[2.],[2.]]]), - }) + res = sess.run( + [output, state], { + hidden[0].name: np.array([[[1.], [1.]], [[2.], [2.]]]), + x.name: np.array([[[1.], [1.]], [[2.], [2.]]]), + }) # This is a smoke test, making sure expected values are unchanged. self.assertEqual(len(res), 2) self.assertAllClose(res[0], res[1].h) @@ -916,44 +909,40 @@ class RNNCellTest(test.TestCase): def testConv2DLSTMCell(self): with self.test_session() as sess: - shape = [2,2,1] - filter_size = [3,3] + shape = [2, 2, 1] + filter_size = [3, 3] num_features = 1 batch_size = 2 expected_state_c = np.array( - [[[[1.4375670191], [1.4375670191]], - [[1.4375670191], [1.4375670191]]], - [[[2.7542609292], [2.7542609292]], - [[2.7542609292], [2.7542609292]]]], + [[[[1.4375670191], [1.4375670191]], [[1.4375670191], [1.4375670191]]], + [[[2.7542609292], [2.7542609292]], [[2.7542609292], [2.7542609292]] + ]], dtype=np.float32) expected_state_h = np.array( - [[[[0.6529865603], [0.6529865603]], - [[0.6529865603], [0.6529865603]]], - [[[0.8736877431], [0.8736877431]], - [[0.8736877431], [0.8736877431]]]], + [[[[0.6529865603], [0.6529865603]], [[0.6529865603], [0.6529865603]]], + [[[0.8736877431], [0.8736877431]], [[0.8736877431], [0.8736877431]] + ]], dtype=np.float32) with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(1.0/4.0)): + "root", initializer=init_ops.constant_initializer(1.0 / 4.0)): x = array_ops.placeholder(dtypes.float32, [None, None, None, 1]) - cell = contrib_rnn_cell.Conv2DLSTMCell(input_shape=shape, - kernel_shape=filter_size, - output_channels=num_features) + cell = contrib_rnn_cell.Conv2DLSTMCell( + input_shape=shape, + kernel_shape=filter_size, + output_channels=num_features) hidden = cell.zero_state(array_ops.shape(x)[0], dtypes.float32) output, state = cell(x, hidden) sess.run([variables.global_variables_initializer()]) - res = sess.run([output, state], { - hidden[0].name: - np.array([[[[1.],[1.]], - [[1.],[1.]]], - [[[2.],[2.]], - [[2.],[2.]]]]), - x.name: - np.array([[[[1.],[1.]], - [[1.],[1.]]], - [[[2.],[2.]], - [[2.],[2.]]]]), - }) + res = sess.run( + [output, state], { + hidden[0].name: + np.array([[[[1.], [1.]], [[1.], [1.]]], [[[2.], [2.]], + [[2.], [2.]]]]), + x.name: + np.array([[[[1.], [1.]], [[1.], [1.]]], [[[2.], [2.]], + [[2.], [2.]]]]), + }) # This is a smoke test, making sure expected values are unchanged. self.assertEqual(len(res), 2) self.assertAllClose(res[0], res[1].h) @@ -962,36 +951,33 @@ class RNNCellTest(test.TestCase): def testConv3DLSTMCell(self): with self.test_session() as sess: - shape = [2,2,2,1] - filter_size = [3,3,3] + shape = [2, 2, 2, 1] + filter_size = [3, 3, 3] num_features = 1 batch_size = 2 expected_state_c = np.array( - [[[[[1.4375670191], [1.4375670191]], - [[1.4375670191], [1.4375670191]]], - [[[1.4375670191], [1.4375670191]], - [[1.4375670191], [1.4375670191]]]], - [[[[2.7542609292], [2.7542609292]], - [[2.7542609292], [2.7542609292]]], - [[[2.7542609292], [2.7542609292]], - [[2.7542609292], [2.7542609292]]]]], + [[[[[1.4375670191], [1.4375670191]], [[1.4375670191], [1.4375670191]] + ], [[[1.4375670191], [1.4375670191]], [[1.4375670191], + [1.4375670191]]]], + [[[[2.7542609292], [2.7542609292]], [[2.7542609292], [2.7542609292]] + ], [[[2.7542609292], [2.7542609292]], [[2.7542609292], + [2.7542609292]]]]], dtype=np.float32) expected_state_h = np.array( - [[[[[0.6529865603], [0.6529865603]], - [[0.6529865603], [0.6529865603]]], - [[[0.6529865603], [0.6529865603]], - [[0.6529865603], [0.6529865603]]]], - [[[[0.8736877431], [0.8736877431]], - [[0.8736877431], [0.8736877431]]], - [[[0.8736877431], [0.8736877431]], - [[0.8736877431], [0.8736877431]]]]], + [[[[[0.6529865603], [0.6529865603]], [[0.6529865603], [0.6529865603]] + ], [[[0.6529865603], [0.6529865603]], [[0.6529865603], + [0.6529865603]]]], + [[[[0.8736877431], [0.8736877431]], [[0.8736877431], [0.8736877431]] + ], [[[0.8736877431], [0.8736877431]], [[0.8736877431], + [0.8736877431]]]]], dtype=np.float32) with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(1.0/8.0)): + "root", initializer=init_ops.constant_initializer(1.0 / 8.0)): x = array_ops.placeholder(dtypes.float32, [None, None, None, None, 1]) - cell = contrib_rnn_cell.Conv3DLSTMCell(input_shape=shape, - kernel_shape=filter_size, - output_channels=num_features) + cell = contrib_rnn_cell.Conv3DLSTMCell( + input_shape=shape, + kernel_shape=filter_size, + output_channels=num_features) hidden = cell.zero_state(array_ops.shape(x)[0], dtypes.float32) output, state = cell(x, hidden) @@ -1054,8 +1040,8 @@ class RNNCellTest(test.TestCase): num_units=num_units, number_of_groups=number_of_groups) cell = rnn_cell.LSTMCell(num_units=num_units) self.assertTrue(isinstance(gcell.state_size, tuple)) - zero_state = gcell.zero_state(batch_size=batch_size, - dtype=dtypes.float32) + zero_state = gcell.zero_state( + batch_size=batch_size, dtype=dtypes.float32) gh, gs = gcell(x, zero_state) h, g = cell(x, zero_state) @@ -1078,16 +1064,16 @@ class RNNCellTest(test.TestCase): glstm_input = array_ops.ones([batch_size, num_units]) gcell = contrib_rnn_cell.GLSTMCell( num_units=num_units, number_of_groups=number_of_groups) - gcell_zero_state = gcell.zero_state(batch_size=batch_size, - dtype=dtypes.float32) + gcell_zero_state = gcell.zero_state( + batch_size=batch_size, dtype=dtypes.float32) gh, gs = gcell(glstm_input, gcell_zero_state) # input for LSTM cell simulating single G-LSTM group lstm_input = array_ops.ones([batch_size, num_units / number_of_groups]) # note division by number_of_groups. This cell one simulates G-LSTM group cell = rnn_cell.LSTMCell(num_units=int(num_units / number_of_groups)) - cell_zero_state = cell.zero_state(batch_size=batch_size, - dtype=dtypes.float32) + cell_zero_state = cell.zero_state( + batch_size=batch_size, dtype=dtypes.float32) h, g = cell(lstm_input, cell_zero_state) sess.run([variables.global_variables_initializer()]) @@ -1097,6 +1083,7 @@ class RNNCellTest(test.TestCase): self.assertAllClose(gh_res[:, int(num_units / number_of_groups):], h_res, 1e-5) + class LayerNormBasicLSTMCellTest(test.TestCase): # NOTE: all the values in the current test case have been calculated. @@ -1117,13 +1104,14 @@ class LayerNormBasicLSTMCellTest(test.TestCase): cell = rnn_cell.MultiRNNCell([single_cell() for _ in range(2)]) g, out_m = cell(x, state) sess.run([variables.global_variables_initializer()]) - res = sess.run([g, out_m], { - x.name: np.array([[1., 1.]]), - c0.name: 0.1 * np.asarray([[0, 1]]), - h0.name: 0.1 * np.asarray([[2, 3]]), - c1.name: 0.1 * np.asarray([[4, 5]]), - h1.name: 0.1 * np.asarray([[6, 7]]), - }) + res = sess.run( + [g, out_m], { + x.name: np.array([[1., 1.]]), + c0.name: 0.1 * np.asarray([[0, 1]]), + h0.name: 0.1 * np.asarray([[2, 3]]), + c1.name: 0.1 * np.asarray([[4, 5]]), + h1.name: 0.1 * np.asarray([[6, 7]]), + }) expected_h = np.array([[-0.38079708, 0.38079708]]) expected_state0_c = np.array([[-1.0, 1.0]]) @@ -1153,11 +1141,12 @@ class LayerNormBasicLSTMCellTest(test.TestCase): cell = contrib_rnn_cell.LayerNormBasicLSTMCell(2) g, out_m = cell(x, state) sess.run([variables.global_variables_initializer()]) - res = sess.run([g, out_m], { - x.name: np.array([[1., 1., 1.]]), - c.name: 0.1 * np.asarray([[0, 1]]), - h.name: 0.1 * np.asarray([[2, 3]]), - }) + res = sess.run( + [g, out_m], { + x.name: np.array([[1., 1., 1.]]), + c.name: 0.1 * np.asarray([[0, 1]]), + h.name: 0.1 * np.asarray([[2, 3]]), + }) expected_h = np.array([[-0.38079708, 0.38079708]]) expected_c = np.array([[-1.0, 1.0]]) @@ -1166,7 +1155,6 @@ class LayerNormBasicLSTMCellTest(test.TestCase): self.assertAllClose(res[1].c, expected_c, 1e-5) self.assertAllClose(res[1].h, expected_h, 1e-5) - def testBasicLSTMCellWithoutNorm(self): """Tests that BasicLSTMCell with layer_norm=False.""" with self.test_session() as sess: @@ -1184,19 +1172,20 @@ class LayerNormBasicLSTMCellTest(test.TestCase): cell = rnn_cell.MultiRNNCell([single_cell() for _ in range(2)]) g, out_m = cell(x, state) sess.run([variables.global_variables_initializer()]) - res = sess.run([g, out_m], { - x.name: np.array([[1., 1.]]), - c0.name: 0.1 * np.asarray([[0, 1]]), - h0.name: 0.1 * np.asarray([[2, 3]]), - c1.name: 0.1 * np.asarray([[4, 5]]), - h1.name: 0.1 * np.asarray([[6, 7]]), - }) + res = sess.run( + [g, out_m], { + x.name: np.array([[1., 1.]]), + c0.name: 0.1 * np.asarray([[0, 1]]), + h0.name: 0.1 * np.asarray([[2, 3]]), + c1.name: 0.1 * np.asarray([[4, 5]]), + h1.name: 0.1 * np.asarray([[6, 7]]), + }) - expected_h = np.array([[ 0.70230919, 0.72581059]]) - expected_state0_c = np.array([[ 0.8020075, 0.89599884]]) - expected_state0_h = np.array([[ 0.56668288, 0.60858738]]) - expected_state1_c = np.array([[ 1.17500675, 1.26892781]]) - expected_state1_h = np.array([[ 0.70230919, 0.72581059]]) + expected_h = np.array([[0.70230919, 0.72581059]]) + expected_state0_c = np.array([[0.8020075, 0.89599884]]) + expected_state0_h = np.array([[0.56668288, 0.60858738]]) + expected_state1_c = np.array([[1.17500675, 1.26892781]]) + expected_state1_h = np.array([[0.70230919, 0.72581059]]) actual_h = res[0] actual_state0_c = res[1][0].c @@ -1213,21 +1202,22 @@ class LayerNormBasicLSTMCellTest(test.TestCase): with variable_scope.variable_scope( "other", initializer=init_ops.constant_initializer(0.5)) as vs: x = array_ops.zeros( - [1, 3]) # Test BasicLSTMCell with input_size != num_units. + [1, 3]) # Test BasicLSTMCell with input_size != num_units. c = array_ops.zeros([1, 2]) h = array_ops.zeros([1, 2]) state = rnn_cell.LSTMStateTuple(c, h) cell = contrib_rnn_cell.LayerNormBasicLSTMCell(2, layer_norm=False) g, out_m = cell(x, state) sess.run([variables.global_variables_initializer()]) - res = sess.run([g, out_m], { - x.name: np.array([[1., 1., 1.]]), - c.name: 0.1 * np.asarray([[0, 1]]), - h.name: 0.1 * np.asarray([[2, 3]]), - }) - - expected_h = np.array([[ 0.64121795, 0.68166804]]) - expected_c = np.array([[ 0.88477188, 0.98103917]]) + res = sess.run( + [g, out_m], { + x.name: np.array([[1., 1., 1.]]), + c.name: 0.1 * np.asarray([[0, 1]]), + h.name: 0.1 * np.asarray([[2, 3]]), + }) + + expected_h = np.array([[0.64121795, 0.68166804]]) + expected_c = np.array([[0.88477188, 0.98103917]]) self.assertEqual(len(res), 2) self.assertAllClose(res[0], expected_h, 1e-5) self.assertAllClose(res[1].c, expected_c, 1e-5) @@ -1248,13 +1238,14 @@ class LayerNormBasicLSTMCellTest(test.TestCase): [contrib_rnn_cell.LayerNormBasicLSTMCell(2) for _ in range(2)]) h, (s0, s1) = cell(x, (state0, state1)) sess.run([variables.global_variables_initializer()]) - res = sess.run([h, s0, s1], { - x.name: np.array([[1., 1.]]), - c0.name: 0.1 * np.asarray([[0, 1]]), - h0.name: 0.1 * np.asarray([[2, 3]]), - c1.name: 0.1 * np.asarray([[4, 5]]), - h1.name: 0.1 * np.asarray([[6, 7]]), - }) + res = sess.run( + [h, s0, s1], { + x.name: np.array([[1., 1.]]), + c0.name: 0.1 * np.asarray([[0, 1]]), + h0.name: 0.1 * np.asarray([[2, 3]]), + c1.name: 0.1 * np.asarray([[4, 5]]), + h1.name: 0.1 * np.asarray([[6, 7]]), + }) expected_h = np.array([[-0.38079708, 0.38079708]]) expected_h0 = np.array([[-0.38079708, 0.38079708]]) @@ -1328,7 +1319,7 @@ class LayerNormBasicLSTMCellTest(test.TestCase): h_low = 0.761552567265 h_high = 0.995008519604 num_units = 5 - allowed_low = [2, 3] + allowed_low = [1, 2, 3] with self.test_session() as sess: with variable_scope.variable_scope( @@ -1342,11 +1333,12 @@ class LayerNormBasicLSTMCellTest(test.TestCase): g, s = cell(x, state) sess.run([variables.global_variables_initializer()]) - res = sess.run([g, s], { - x.name: np.ones([1, 5]), - c.name: np.ones([1, 5]), - h.name: np.ones([1, 5]), - }) + res = sess.run( + [g, s], { + x.name: np.ones([1, 5]), + c.name: np.ones([1, 5]), + h.name: np.ones([1, 5]), + }) # Since the returned tensors are of size [1,n] # get the first component right now. @@ -1372,35 +1364,35 @@ class LayerNormBasicLSTMCellTest(test.TestCase): self.assertIn(dropped_count, allowed_low) -def _create_multi_lstm_cell_ops(batch_size, num_units, input_depth, - num_layers, max_time, compiled): +def _create_multi_lstm_cell_ops(batch_size, num_units, input_depth, num_layers, + max_time, compiled): with variable_scope.variable_scope( "root", initializer=init_ops.random_uniform_initializer(-0.1, 0.1, seed=2)): inputs = variable_scope.get_variable( - "inputs", initializer=random_ops.random_uniform( + "inputs", + initializer=random_ops.random_uniform( (max_time, batch_size, input_depth), seed=1)) maybe_xla = lambda c: contrib_rnn_cell.CompiledWrapper(c) if compiled else c cell = rnn_cell.MultiRNNCell( [maybe_xla(rnn_cell.LSTMCell(num_units)) for _ in range(num_layers)]) - initial_state = cell.zero_state( - batch_size=batch_size, dtype=dtypes.float32) + initial_state = cell.zero_state(batch_size=batch_size, dtype=dtypes.float32) outputs, final_state = rnn.dynamic_rnn( - cell=cell, inputs=inputs, initial_state=initial_state, - time_major=True) + cell=cell, inputs=inputs, initial_state=initial_state, time_major=True) flat_final_state = nest.flatten(final_state) trainable_variables = variables.trainable_variables() outputs_grad = gradients_impl.gradients( - [outputs], - trainable_variables + [inputs] + nest.flatten(initial_state)) + [outputs], trainable_variables + [inputs] + nest.flatten(initial_state)) final_state_grad = gradients_impl.gradients( flat_final_state, trainable_variables + [inputs] + nest.flatten(initial_state)) - return {"outputs": outputs, - "final_state": flat_final_state, - "outputs_grad": outputs_grad, - "final_state_grad": final_state_grad} + return { + "outputs": outputs, + "final_state": flat_final_state, + "outputs_grad": outputs_grad, + "final_state_grad": final_state_grad + } class CompiledWrapperTest(test.TestCase): @@ -1418,8 +1410,10 @@ class CompiledWrapperTest(test.TestCase): random_seed.set_random_seed(1234) with self.test_session(graph=ops.Graph()) as sess: xla_ops = _create_multi_lstm_cell_ops( - batch_size=batch_size, num_units=num_units, - input_depth=input_depth, num_layers=num_layers, + batch_size=batch_size, + num_units=num_units, + input_depth=input_depth, + num_layers=num_layers, max_time=max_time, compiled=True) sess.run([variables.global_variables_initializer()]) @@ -1428,8 +1422,10 @@ class CompiledWrapperTest(test.TestCase): random_seed.set_random_seed(1234) with self.test_session(graph=ops.Graph()) as sess: non_xla_ops = _create_multi_lstm_cell_ops( - batch_size=batch_size, num_units=num_units, - input_depth=input_depth, num_layers=num_layers, + batch_size=batch_size, + num_units=num_units, + input_depth=input_depth, + num_layers=num_layers, max_time=max_time, compiled=False) sess.run([variables.global_variables_initializer()]) @@ -1438,16 +1434,16 @@ class CompiledWrapperTest(test.TestCase): self.assertAllClose( non_xla_results["outputs"], xla_results["outputs"], atol=atol) - for xla_value, non_xla_value in zip( - xla_results["final_state"], non_xla_results["final_state"]): + for xla_value, non_xla_value in zip(xla_results["final_state"], + non_xla_results["final_state"]): self.assertAllClose(xla_value, non_xla_value, atol=atol) - for xla_g, non_xla_g in zip( - xla_results["outputs_grad"], non_xla_results["outputs_grad"]): + for xla_g, non_xla_g in zip(xla_results["outputs_grad"], + non_xla_results["outputs_grad"]): self.assertAllClose(xla_g, non_xla_g, atol=atol) - for xla_g, non_xla_g in zip( - xla_results["final_state_grad"], non_xla_results["final_state_grad"]): + for xla_g, non_xla_g in zip(xla_results["final_state_grad"], + non_xla_results["final_state_grad"]): self.assertAllClose(xla_g, non_xla_g, atol=atol) def testMultiRNNCellWithStateTuple(self): @@ -1461,19 +1457,20 @@ class CompiledWrapperTest(test.TestCase): # Test incorrectness of state with self.assertRaisesRegexp(ValueError, "Expected state .* a tuple"): rnn_cell.MultiRNNCell( - [rnn_cell.GRUCell(2) - for _ in range(2)], state_is_tuple=True)(x, m_bad) + [rnn_cell.GRUCell(2) for _ in range(2)], + state_is_tuple=True)(x, m_bad) _, ml = rnn_cell.MultiRNNCell( - [rnn_cell.GRUCell(2) - for _ in range(2)], state_is_tuple=True)(x, m_good) + [rnn_cell.GRUCell(2) for _ in range(2)], + state_is_tuple=True)(x, m_good) sess.run([variables.global_variables_initializer()]) - res = sess.run(ml, { - x.name: np.array([[1., 1.]]), - m_good[0].name: np.array([[0.1, 0.1]]), - m_good[1].name: np.array([[0.1, 0.1]]) - }) + res = sess.run( + ml, { + x.name: np.array([[1., 1.]]), + m_good[0].name: np.array([[0.1, 0.1]]), + m_good[1].name: np.array([[0.1, 0.1]]) + }) # The numbers in results were not calculated, this is just a # smoke test. However, these numbers should match those of @@ -1488,24 +1485,20 @@ class BenchmarkLSTMCellXLA(test.Benchmark): num_layers = 3 max_time = 50 print("benchmarkDynamicRNNWithMultiLSTMCell") - print("\t" + - "\t".join(["inter_th", "intra_th", - "batch_size", "num_units", "input_depth", "device", - "compiled", "wall_time"])) + print("\t" + "\t".join([ + "inter_th", "intra_th", "batch_size", "num_units", "input_depth", + "device", "compiled", "wall_time" + ])) warmup_run = True - for (threads, - device, - num_units, - batch_size, - input_depth, - compiled) in itertools.product( - [{"inter": 0, "intra": 0}, {"inter": 1, "intra": 4}], - ["cpu", "gpu"], - [32, 512], - [1, 32, 256], - [32, 512], - [False, True]): + for (threads, device, num_units, batch_size, input_depth, + compiled) in itertools.product([{ + "inter": 0, + "intra": 0 + }, { + "inter": 1, + "intra": 4 + }], ["cpu", "gpu"], [32, 512], [1, 32, 256], [32, 512], [False, True]): if threads["inter"] != 0: # We only care about testing inter/intra op limitations on # CPU with small batch size, to mimic embedded devices. @@ -1521,31 +1514,131 @@ class BenchmarkLSTMCellXLA(test.Benchmark): with session.Session(config=config, graph=ops.Graph()) as sess: with ops.device("/%s:0" % device): ops_dict = _create_multi_lstm_cell_ops( - batch_size=batch_size, num_units=num_units, - input_depth=input_depth, num_layers=num_layers, + batch_size=batch_size, + num_units=num_units, + input_depth=input_depth, + num_layers=num_layers, max_time=max_time, compiled=compiled) sess.run([variables.global_variables_initializer()]) all_ops = nest.flatten(ops_dict.values()) all_ops_group = control_flow_ops.group(*all_ops) - name_suffix = ( - "inter_th_%d_intra_th_%d_bs_%d_units_%d_inputdepth_%d" - "_device_%s_xla_%s" % ( - threads["inter"], threads["intra"], - batch_size, num_units, input_depth, device, compiled)) + name_suffix = ("inter_th_%d_intra_th_%d_bs_%d_units_%d_inputdepth_%d" + "_device_%s_xla_%s" % + (threads["inter"], threads["intra"], batch_size, + num_units, input_depth, device, compiled)) if warmup_run: self.run_op_benchmark( sess, all_ops_group, min_iters=30, name="ignore_warmup") warmup_run = False benchmark_results = self.run_op_benchmark( - sess, all_ops_group, min_iters=50, + sess, + all_ops_group, + min_iters=50, name="benchmarkDynamicRNNWithMultiLSTMCell_%s" % name_suffix) - print("\t" + - "\t".join(["%s" % x for x in [ - threads["inter"], threads["intra"], - batch_size, num_units, input_depth, device, compiled, - benchmark_results["wall_time"]]])) + print("\t" + "\t".join([ + "%s" % x + for x in [ + threads["inter"], threads["intra"], batch_size, num_units, + input_depth, device, compiled, benchmark_results["wall_time"] + ] + ])) + + +class WeightNormLSTMCellTest(test.TestCase): + """Compared cell output with pre-calculated values.""" + + def _cell_output(self, cell): + """Calculate cell output""" + + with self.test_session() as sess: + init = init_ops.constant_initializer(0.5) + 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]) + + state0 = rnn_cell.LSTMStateTuple(c0, h0) + + 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]]), + }) + + actual_state_c = res[1].c + actual_state_h = res[1].h + + return actual_state_c, actual_state_h + + def testBasicCell(self): + """Tests cell w/o peepholes and w/o normalisation""" + + def cell(): + return contrib_rnn_cell.WeightNormLSTMCell(2, + norm=False, + use_peepholes=False) + + actual_c, actual_h = self._cell_output(cell) + + expected_c = np.array([[0.65937078, 0.74983585]]) + expected_h = np.array([[0.44923624, 0.49362513]]) + + self.assertAllClose(expected_c, actual_c, 1e-5) + self.assertAllClose(expected_h, actual_h, 1e-5) + + def testNonbasicCell(self): + """Tests cell with peepholes and w/o normalisation""" + + def cell(): + return contrib_rnn_cell.WeightNormLSTMCell(2, + norm=False, + use_peepholes=True) + + actual_c, actual_h = self._cell_output(cell) + + expected_c = np.array([[0.65937084, 0.7574988]]) + expected_h = np.array([[0.4792085, 0.53470564]]) + + self.assertAllClose(expected_c, actual_c, 1e-5) + self.assertAllClose(expected_h, actual_h, 1e-5) + + + def testBasicCellWithNorm(self): + """Tests cell w/o peepholes and with normalisation""" + + def cell(): + return contrib_rnn_cell.WeightNormLSTMCell(2, + norm=True, + use_peepholes=False) + + actual_c, actual_h = self._cell_output(cell) + + expected_c = np.array([[0.50125383, 0.58805949]]) + expected_h = np.array([[0.32770363, 0.37397948]]) + + self.assertAllClose(expected_c, actual_c, 1e-5) + self.assertAllClose(expected_h, actual_h, 1e-5) + + def testNonBasicCellWithNorm(self): + """Tests cell with peepholes and with normalisation""" + + def cell(): + return contrib_rnn_cell.WeightNormLSTMCell(2, + norm=True, + use_peepholes=True) + + actual_c, actual_h = self._cell_output(cell) + + expected_c = np.array([[0.50125383, 0.59587258]]) + expected_h = np.array([[0.35041603, 0.40873795]]) + self.assertAllClose(expected_c, actual_c, 1e-5) + self.assertAllClose(expected_h, actual_h, 1e-5) if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/rnn/python/ops/gru_ops.py b/tensorflow/contrib/rnn/python/ops/gru_ops.py index 4c964ec201f153d6c8293d3bf93bc231ff8f751d..81ca12317be484ba420b7bbfac822e91d6d38bff 100644 --- a/tensorflow/contrib/rnn/python/ops/gru_ops.py +++ b/tensorflow/contrib/rnn/python/ops/gru_ops.py @@ -32,7 +32,7 @@ from tensorflow.python.util.deprecation import deprecated_args _gru_ops_so = loader.load_op_library( resource_loader.get_path_to_datafile("_gru_ops.so")) -LayerRNNCell = rnn_cell_impl._LayerRNNCell # pylint: disable=invalid-name,protected-access +LayerRNNCell = rnn_cell_impl.LayerRNNCell # pylint: disable=invalid-name @ops.RegisterGradient("GRUBlockCell") diff --git a/tensorflow/contrib/rnn/python/ops/lstm_ops.py b/tensorflow/contrib/rnn/python/ops/lstm_ops.py index 04f342cd18271425068b2b02c2937236c900c5e2..4eb4fbcd92f0d7cb3bee712862c8950a1971b632 100644 --- a/tensorflow/contrib/rnn/python/ops/lstm_ops.py +++ b/tensorflow/contrib/rnn/python/ops/lstm_ops.py @@ -34,7 +34,7 @@ from tensorflow.python.platform import resource_loader _lstm_ops_so = loader.load_op_library( resource_loader.get_path_to_datafile("_lstm_ops.so")) -LayerRNNCell = rnn_cell_impl._LayerRNNCell # pylint: disable=invalid-name,protected-access +LayerRNNCell = rnn_cell_impl.LayerRNNCell # pylint: disable=invalid-name # pylint: disable=invalid-name @@ -572,9 +572,8 @@ class LSTMBlockWrapper(base_layer.Layer): def _gather_states(self, data, indices, batch_size): """Produce `out`, s.t. out(i, j) = data(indices(i), i, j).""" - mod_indices = indices * batch_size + math_ops.range(batch_size) - return array_ops.gather( - array_ops.reshape(data, [-1, self.num_units]), mod_indices) + return array_ops.gather_nd( + data, array_ops.stack([indices, math_ops.range(batch_size)], axis=1)) class LSTMBlockFusedCell(LSTMBlockWrapper): diff --git a/tensorflow/contrib/rnn/python/ops/rnn_cell.py b/tensorflow/contrib/rnn/python/ops/rnn_cell.py index e4667828cdaad627143efcb823eee39aec24fab7..a6c2d9cdbb2b6f61d59960f708000e945c6115e9 100644 --- a/tensorflow/contrib/rnn/python/ops/rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/rnn_cell.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Module for constructing RNN Cells.""" from __future__ import absolute_import from __future__ import division @@ -33,11 +32,12 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import clip_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn_impl # pylint: disable=unused-import from tensorflow.python.ops import nn_ops +from tensorflow.python.ops import partitioned_variables # pylint: disable=unused-import from tensorflow.python.ops import random_ops from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.ops import variable_scope as vs -from tensorflow.python.ops import partitioned_variables from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import nest @@ -55,16 +55,15 @@ def _get_concat_variable(name, shape, dtype, num_shards): return value concat_variable = array_ops.concat(sharded_variable, 0, name=concat_name) - ops.add_to_collection(ops.GraphKeys.CONCATENATED_VARIABLES, - concat_variable) + ops.add_to_collection(ops.GraphKeys.CONCATENATED_VARIABLES, concat_variable) return concat_variable def _get_sharded_variable(name, shape, dtype, num_shards): """Get a list of sharded variables with the given dtype.""" if num_shards > shape[0]: - raise ValueError("Too many shards: shape=%s, num_shards=%d" % - (shape, num_shards)) + raise ValueError("Too many shards: shape=%s, num_shards=%d" % (shape, + num_shards)) unit_shard_size = int(math.floor(shape[0] / num_shards)) remaining_rows = shape[0] - unit_shard_size * num_shards @@ -73,8 +72,9 @@ def _get_sharded_variable(name, shape, dtype, num_shards): current_size = unit_shard_size if i < remaining_rows: current_size += 1 - shards.append(vs.get_variable(name + "_%d" % i, [current_size] + shape[1:], - dtype=dtype)) + shards.append( + vs.get_variable( + name + "_%d" % i, [current_size] + shape[1:], dtype=dtype)) return shards @@ -176,9 +176,8 @@ class CoupledInputForgetGateLSTMCell(rnn_cell_impl.RNNCell): """ super(CoupledInputForgetGateLSTMCell, self).__init__(_reuse=reuse) if not state_is_tuple: - logging.warn( - "%s: Using a concatenated state is slower and will soon be " - "deprecated. Use state_is_tuple=True.", self) + logging.warn("%s: Using a concatenated state is slower and will soon be " + "deprecated. Use state_is_tuple=True.", self) self._num_units = num_units self._use_peepholes = use_peepholes self._initializer = initializer @@ -195,12 +194,14 @@ class CoupledInputForgetGateLSTMCell(rnn_cell_impl.RNNCell): self._norm_shift = norm_shift if num_proj: - self._state_size = (rnn_cell_impl.LSTMStateTuple(num_units, num_proj) - if state_is_tuple else num_units + num_proj) + self._state_size = ( + rnn_cell_impl.LSTMStateTuple(num_units, num_proj) + if state_is_tuple else num_units + num_proj) self._output_size = num_proj else: - self._state_size = (rnn_cell_impl.LSTMStateTuple(num_units, num_units) - if state_is_tuple else 2 * num_units) + self._state_size = ( + rnn_cell_impl.LSTMStateTuple(num_units, num_units) + if state_is_tuple else 2 * num_units) self._output_size = num_units @property @@ -250,8 +251,8 @@ class CoupledInputForgetGateLSTMCell(rnn_cell_impl.RNNCell): if input_size.value is None: raise ValueError("Could not infer input size from inputs.get_shape()[-1]") concat_w = _get_concat_variable( - "W", [input_size.value + num_proj, 3 * self._num_units], - dtype, self._num_unit_shards) + "W", [input_size.value + num_proj, 3 * self._num_units], dtype, + self._num_unit_shards) b = vs.get_variable( "B", @@ -298,9 +299,9 @@ class CoupledInputForgetGateLSTMCell(rnn_cell_impl.RNNCell): m = sigmoid(o) * self._activation(c) if self._num_proj is not None: - concat_w_proj = _get_concat_variable( - "W_P", [self._num_units, self._num_proj], - dtype, self._num_proj_shards) + concat_w_proj = _get_concat_variable("W_P", + [self._num_units, self._num_proj], + dtype, self._num_proj_shards) m = math_ops.matmul(m, concat_w_proj) if self._proj_clip is not None: @@ -308,8 +309,9 @@ class CoupledInputForgetGateLSTMCell(rnn_cell_impl.RNNCell): m = clip_ops.clip_by_value(m, -self._proj_clip, self._proj_clip) # pylint: enable=invalid-unary-operand-type - new_state = (rnn_cell_impl.LSTMStateTuple(c, m) - if self._state_is_tuple else array_ops.concat([c, m], 1)) + new_state = ( + rnn_cell_impl.LSTMStateTuple(c, m) + if self._state_is_tuple else array_ops.concat([c, m], 1)) return m, new_state @@ -325,10 +327,15 @@ class TimeFreqLSTMCell(rnn_cell_impl.RNNCell): It uses peep-hole connections and optional cell clipping. """ - def __init__(self, num_units, use_peepholes=False, - cell_clip=None, initializer=None, - num_unit_shards=1, forget_bias=1.0, - feature_size=None, frequency_skip=None, + def __init__(self, + num_units, + use_peepholes=False, + cell_clip=None, + initializer=None, + num_unit_shards=1, + forget_bias=1.0, + feature_size=None, + frequency_skip=1, reuse=None): """Initialize the parameters for an LSTM cell. @@ -398,7 +405,7 @@ class TimeFreqLSTMCell(rnn_cell_impl.RNNCell): actual_input_size = freq_inputs[0].get_shape().as_list()[1] concat_w = _get_concat_variable( - "W", [actual_input_size + 2*self._num_units, 4 * self._num_units], + "W", [actual_input_size + 2 * self._num_units, 4 * self._num_units], dtype, self._num_unit_shards) b = vs.get_variable( @@ -417,23 +424,24 @@ class TimeFreqLSTMCell(rnn_cell_impl.RNNCell): "W_O_diag", shape=[self._num_units], dtype=dtype) # initialize the first freq state to be zero - m_prev_freq = array_ops.zeros([int(inputs.get_shape()[0]), - self._num_units], dtype) + m_prev_freq = array_ops.zeros( + [inputs.shape[0].value or inputs.get_shape()[0], self._num_units], + dtype) for fq in range(len(freq_inputs)): - c_prev = array_ops.slice(state, [0, 2*fq*self._num_units], + c_prev = array_ops.slice(state, [0, 2 * fq * self._num_units], [-1, self._num_units]) - m_prev = array_ops.slice(state, [0, (2*fq+1)*self._num_units], + m_prev = array_ops.slice(state, [0, (2 * fq + 1) * self._num_units], [-1, self._num_units]) # i = input_gate, j = new_input, f = forget_gate, o = output_gate - cell_inputs = array_ops.concat([freq_inputs[fq], m_prev, m_prev_freq], - 1) + cell_inputs = array_ops.concat([freq_inputs[fq], m_prev, m_prev_freq], 1) lstm_matrix = nn_ops.bias_add(math_ops.matmul(cell_inputs, concat_w), b) i, j, f, o = array_ops.split( value=lstm_matrix, num_or_size_splits=4, axis=1) if self._use_peepholes: - c = (sigmoid(f + self._forget_bias + w_f_diag * c_prev) * c_prev + - sigmoid(i + w_i_diag * c_prev) * tanh(j)) + c = ( + sigmoid(f + self._forget_bias + w_f_diag * c_prev) * c_prev + + sigmoid(i + w_i_diag * c_prev) * tanh(j)) else: c = (sigmoid(f + self._forget_bias) * c_prev + sigmoid(i) * tanh(j)) @@ -471,11 +479,11 @@ class TimeFreqLSTMCell(rnn_cell_impl.RNNCell): input_size = input_feat.get_shape().with_rank(2)[-1].value if input_size is None: raise ValueError("Cannot infer input_size from static shape inference.") - num_feats = int((input_size - self._feature_size) / ( - self._frequency_skip)) + 1 + num_feats = int( + (input_size - self._feature_size) / (self._frequency_skip)) + 1 freq_inputs = [] for f in range(num_feats): - cur_input = array_ops.slice(input_feat, [0, f*self._frequency_skip], + cur_input = array_ops.slice(input_feat, [0, f * self._frequency_skip], [-1, self._feature_size]) freq_inputs.append(cur_input) return freq_inputs @@ -497,11 +505,16 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): The code uses optional peephole connections, shared_weights and cell clipping. """ - def __init__(self, num_units, use_peepholes=False, + def __init__(self, + num_units, + use_peepholes=False, share_time_frequency_weights=False, - cell_clip=None, initializer=None, - num_unit_shards=1, forget_bias=1.0, - feature_size=None, frequency_skip=None, + cell_clip=None, + initializer=None, + num_unit_shards=1, + forget_bias=1.0, + feature_size=None, + frequency_skip=None, num_frequency_blocks=None, start_freqindex_list=None, end_freqindex_list=None, @@ -579,10 +592,10 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): for freq_index in range(self._num_frequency_blocks[block_index]): name_prefix = "state_f%02d_b%02d" % (freq_index, block_index) state_names += ("%s_c, %s_m," % (name_prefix, name_prefix)) - self._state_tuple_type = collections.namedtuple( - "GridLSTMStateTuple", state_names.strip(",")) - self._state_size = self._state_tuple_type( - *([num_units, num_units] * self._total_blocks)) + self._state_tuple_type = collections.namedtuple("GridLSTMStateTuple", + state_names.strip(",")) + self._state_size = self._state_tuple_type(*( + [num_units, num_units] * self._total_blocks)) else: self._state_tuple_type = None self._state_size = num_units * self._total_blocks * 2 @@ -625,7 +638,10 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): state_out_lst = [] for block in range(len(freq_inputs)): m_out_lst_current, state_out_lst_current = self._compute( - freq_inputs[block], block, state, batch_size, + freq_inputs[block], + block, + state, + batch_size, state_is_tuple=self._state_is_tuple) m_out_lst.extend(m_out_lst_current) state_out_lst.extend(state_out_lst_current) @@ -636,7 +652,11 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): m_out = array_ops.concat(m_out_lst, 1) return m_out, state_out - def _compute(self, freq_inputs, block, state, batch_size, + def _compute(self, + freq_inputs, + block, + state, + batch_size, state_prefix="state", state_is_tuple=True): """Run the actual computation of one step LSTM. @@ -665,8 +685,8 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): actual_input_size = freq_inputs[0].get_shape().as_list()[1] concat_w_f = _get_concat_variable( - "W_f_%d" % block, [actual_input_size + 2 * self._num_units, - num_gates * self._num_units], + "W_f_%d" % block, + [actual_input_size + 2 * self._num_units, num_gates * self._num_units], dtype, self._num_unit_shards) b_f = vs.get_variable( "B_f_%d" % block, @@ -674,10 +694,9 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): initializer=init_ops.zeros_initializer(), dtype=dtype) if not self._share_time_frequency_weights: - concat_w_t = _get_concat_variable( - "W_t_%d" % block, [actual_input_size + 2 * self._num_units, - num_gates * self._num_units], - dtype, self._num_unit_shards) + concat_w_t = _get_concat_variable("W_t_%d" % block, [ + actual_input_size + 2 * self._num_units, num_gates * self._num_units + ], dtype, self._num_unit_shards) b_t = vs.get_variable( "B_t_%d" % block, shape=[num_gates * self._num_units], @@ -690,7 +709,7 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): w_f_diag_freqf = vs.get_variable( "W_F_diag_freqf_%d" % block, shape=[self._num_units], dtype=dtype) w_f_diag_freqt = vs.get_variable( - "W_F_diag_freqt_%d"% block, shape=[self._num_units], dtype=dtype) + "W_F_diag_freqt_%d" % block, shape=[self._num_units], dtype=dtype) w_i_diag_freqf = vs.get_variable( "W_I_diag_freqf_%d" % block, shape=[self._num_units], dtype=dtype) w_i_diag_freqt = vs.get_variable( @@ -724,8 +743,7 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): m_prev_time = getattr(state, name_prefix + "_m") else: c_prev_time = array_ops.slice( - state, [0, 2 * freq_index * self._num_units], - [-1, self._num_units]) + state, [0, 2 * freq_index * self._num_units], [-1, self._num_units]) m_prev_time = array_ops.slice( state, [0, (2 * freq_index + 1) * self._num_units], [-1, self._num_units]) @@ -735,8 +753,8 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): [freq_inputs[freq_index], m_prev_time, m_prev_freq], 1) # F-LSTM - lstm_matrix_freq = nn_ops.bias_add(math_ops.matmul(cell_inputs, - concat_w_f), b_f) + lstm_matrix_freq = nn_ops.bias_add( + math_ops.matmul(cell_inputs, concat_w_f), b_f) if self._couple_input_forget_gates: i_freq, j_freq, o_freq = array_ops.split( value=lstm_matrix_freq, num_or_size_splits=num_gates, axis=1) @@ -751,8 +769,8 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): f_time = f_freq o_time = o_freq else: - lstm_matrix_time = nn_ops.bias_add(math_ops.matmul(cell_inputs, - concat_w_t), b_t) + lstm_matrix_time = nn_ops.bias_add( + math_ops.matmul(cell_inputs, concat_w_t), b_t) if self._couple_input_forget_gates: i_time, j_time, o_time = array_ops.split( value=lstm_matrix_time, num_or_size_splits=num_gates, axis=1) @@ -764,8 +782,7 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): # F-LSTM c_freq # input gate activations if self._use_peepholes: - i_freq_g = sigmoid(i_freq + - w_i_diag_freqf * c_prev_freq + + i_freq_g = sigmoid(i_freq + w_i_diag_freqf * c_prev_freq + w_i_diag_freqt * c_prev_time) else: i_freq_g = sigmoid(i_freq) @@ -774,9 +791,8 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): f_freq_g = 1.0 - i_freq_g else: if self._use_peepholes: - f_freq_g = sigmoid(f_freq + self._forget_bias + - w_f_diag_freqf * c_prev_freq + - w_f_diag_freqt * c_prev_time) + f_freq_g = sigmoid(f_freq + self._forget_bias + w_f_diag_freqf * + c_prev_freq + w_f_diag_freqt * c_prev_time) else: f_freq_g = sigmoid(f_freq + self._forget_bias) # cell state @@ -791,12 +807,10 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): # input gate activations if self._use_peepholes: if self._share_time_frequency_weights: - i_time_g = sigmoid(i_time + - w_i_diag_freqf * c_prev_freq + + i_time_g = sigmoid(i_time + w_i_diag_freqf * c_prev_freq + w_i_diag_freqt * c_prev_time) else: - i_time_g = sigmoid(i_time + - w_i_diag_timef * c_prev_freq + + i_time_g = sigmoid(i_time + w_i_diag_timef * c_prev_freq + w_i_diag_timet * c_prev_time) else: i_time_g = sigmoid(i_time) @@ -806,13 +820,11 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): else: if self._use_peepholes: if self._share_time_frequency_weights: - f_time_g = sigmoid(f_time + self._forget_bias + - w_f_diag_freqf * c_prev_freq + - w_f_diag_freqt * c_prev_time) + f_time_g = sigmoid(f_time + self._forget_bias + w_f_diag_freqf * + c_prev_freq + w_f_diag_freqt * c_prev_time) else: - f_time_g = sigmoid(f_time + self._forget_bias + - w_f_diag_timef * c_prev_freq + - w_f_diag_timet * c_prev_time) + f_time_g = sigmoid(f_time + self._forget_bias + w_f_diag_timef * + c_prev_freq + w_f_diag_timet * c_prev_time) else: f_time_g = sigmoid(f_time + self._forget_bias) # cell state @@ -825,8 +837,7 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): # F-LSTM m_freq if self._use_peepholes: - m_freq = sigmoid(o_freq + - w_o_diag_freqf * c_freq + + m_freq = sigmoid(o_freq + w_o_diag_freqf * c_freq + w_o_diag_freqt * c_time) * tanh(c_freq) else: m_freq = sigmoid(o_freq) * tanh(c_freq) @@ -834,12 +845,10 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): # T-LSTM m_time if self._use_peepholes: if self._share_time_frequency_weights: - m_time = sigmoid(o_time + - w_o_diag_freqf * c_freq + + m_time = sigmoid(o_time + w_o_diag_freqf * c_freq + w_o_diag_freqt * c_time) * tanh(c_time) else: - m_time = sigmoid(o_time + - w_o_diag_timef * c_freq + + m_time = sigmoid(o_time + w_o_diag_timef * c_freq + w_o_diag_timet * c_time) * tanh(c_time) else: m_time = sigmoid(o_time) * tanh(c_time) @@ -878,16 +887,18 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): raise ValueError("Cannot infer input_size from static shape inference.") if slice_offset > 0: # Padding to the end - inputs = array_ops.pad( - input_feat, array_ops.constant([0, 0, 0, slice_offset], shape=[2, 2], - dtype=dtypes.int32), - "CONSTANT") + inputs = array_ops.pad(input_feat, + array_ops.constant( + [0, 0, 0, slice_offset], + shape=[2, 2], + dtype=dtypes.int32), "CONSTANT") elif slice_offset < 0: # Padding to the front - inputs = array_ops.pad( - input_feat, array_ops.constant([0, 0, -slice_offset, 0], shape=[2, 2], - dtype=dtypes.int32), - "CONSTANT") + inputs = array_ops.pad(input_feat, + array_ops.constant( + [0, 0, -slice_offset, 0], + shape=[2, 2], + dtype=dtypes.int32), "CONSTANT") slice_offset = 0 else: inputs = input_feat @@ -897,13 +908,13 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): raise ValueError("Length of num_frequency_blocks" " is not 1, but instead is %d", len(self._num_frequency_blocks)) - num_feats = int((input_size - self._feature_size) / ( - self._frequency_skip)) + 1 + num_feats = int( + (input_size - self._feature_size) / (self._frequency_skip)) + 1 if num_feats != self._num_frequency_blocks[0]: raise ValueError( "Invalid num_frequency_blocks, requires %d but gets %d, please" - " check the input size and filter config are correct." % ( - self._num_frequency_blocks[0], num_feats)) + " check the input size and filter config are correct." % + (self._num_frequency_blocks[0], num_feats)) block_inputs = [] for f in range(num_feats): cur_input = array_ops.slice( @@ -926,18 +937,18 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): start_index = self._start_freqindex_list[b] end_index = self._end_freqindex_list[b] cur_size = end_index - start_index - block_feats = int((cur_size - self._feature_size) / ( - self._frequency_skip)) + 1 + block_feats = int( + (cur_size - self._feature_size) / (self._frequency_skip)) + 1 if block_feats != self._num_frequency_blocks[b]: raise ValueError( "Invalid num_frequency_blocks, requires %d but gets %d, please" - " check the input size and filter config are correct." % ( - self._num_frequency_blocks[b], block_feats)) + " check the input size and filter config are correct." % + (self._num_frequency_blocks[b], block_feats)) block_inputs = [] for f in range(block_feats): cur_input = array_ops.slice( - inputs, [0, start_index + slice_offset + f * - self._frequency_skip], + inputs, + [0, start_index + slice_offset + f * self._frequency_skip], [-1, self._feature_size]) block_inputs.append(cur_input) freq_inputs.append(block_inputs) @@ -953,11 +964,16 @@ class BidirectionalGridLSTMCell(GridLSTMCell): The current implementation uses different weights for the two directions. """ - def __init__(self, num_units, use_peepholes=False, + def __init__(self, + num_units, + use_peepholes=False, share_time_frequency_weights=False, - cell_clip=None, initializer=None, - num_unit_shards=1, forget_bias=1.0, - feature_size=None, frequency_skip=None, + cell_clip=None, + initializer=None, + num_unit_shards=1, + forget_bias=1.0, + feature_size=None, + frequency_skip=None, num_frequency_blocks=None, start_freqindex_list=None, end_freqindex_list=None, @@ -1016,8 +1032,8 @@ class BidirectionalGridLSTMCell(GridLSTMCell): state_names += ("%s_c, %s_m," % (name_prefix, name_prefix)) self._state_tuple_type = collections.namedtuple( "BidirectionalGridLSTMStateTuple", state_names.strip(",")) - self._state_size = self._state_tuple_type( - *([num_units, num_units] * self._total_blocks * 2)) + self._state_size = self._state_tuple_type(*( + [num_units, num_units] * self._total_blocks * 2)) self._output_size = 2 * num_units * self._total_blocks * 2 def call(self, inputs, state): @@ -1051,8 +1067,12 @@ class BidirectionalGridLSTMCell(GridLSTMCell): fwd_state_out_lst = [] for block in range(len(fwd_inputs)): fwd_m_out_lst_current, fwd_state_out_lst_current = self._compute( - fwd_inputs[block], block, state, batch_size, - state_prefix="fwd_state", state_is_tuple=True) + fwd_inputs[block], + block, + state, + batch_size, + state_prefix="fwd_state", + state_is_tuple=True) fwd_m_out_lst.extend(fwd_m_out_lst_current) fwd_state_out_lst.extend(fwd_state_out_lst_current) # Backward processing @@ -1063,8 +1083,12 @@ class BidirectionalGridLSTMCell(GridLSTMCell): # Reverse the blocks bwd_inputs_reverse = bwd_inputs[block][::-1] bwd_m_out_lst_current, bwd_state_out_lst_current = self._compute( - bwd_inputs_reverse, block, state, batch_size, - state_prefix="bwd_state", state_is_tuple=True) + bwd_inputs_reverse, + block, + state, + batch_size, + state_prefix="bwd_state", + state_is_tuple=True) bwd_m_out_lst.extend(bwd_m_out_lst_current) bwd_state_out_lst.extend(bwd_state_out_lst_current) state_out = self._state_tuple_type(*(fwd_state_out_lst + bwd_state_out_lst)) @@ -1075,6 +1099,7 @@ class BidirectionalGridLSTMCell(GridLSTMCell): # pylint: disable=protected-access _Linear = core_rnn_cell._Linear # pylint: disable=invalid-name + # pylint: enable=protected-access @@ -1084,8 +1109,14 @@ class AttentionCellWrapper(rnn_cell_impl.RNNCell): Implementation based on https://arxiv.org/abs/1409.0473. """ - def __init__(self, cell, attn_length, attn_size=None, attn_vec_size=None, - input_size=None, state_is_tuple=True, reuse=None): + def __init__(self, + cell, + attn_length, + attn_size=None, + attn_vec_size=None, + input_size=None, + state_is_tuple=True, + reuse=None): """Create a cell with attention. Args: @@ -1115,16 +1146,15 @@ class AttentionCellWrapper(rnn_cell_impl.RNNCell): if not rnn_cell_impl._like_rnncell(cell): # pylint: disable=protected-access raise TypeError("The parameter cell is not RNNCell.") if nest.is_sequence(cell.state_size) and not state_is_tuple: - raise ValueError("Cell returns tuple of states, but the flag " - "state_is_tuple is not set. State size is: %s" - % str(cell.state_size)) + raise ValueError( + "Cell returns tuple of states, but the flag " + "state_is_tuple is not set. State size is: %s" % str(cell.state_size)) if attn_length <= 0: - raise ValueError("attn_length should be greater than zero, got %s" - % str(attn_length)) + raise ValueError( + "attn_length should be greater than zero, got %s" % str(attn_length)) if not state_is_tuple: - logging.warn( - "%s: Using a concatenated state is slower and will soon be " - "deprecated. Use state_is_tuple=True.", self) + logging.warn("%s: Using a concatenated state is slower and will soon be " + "deprecated. Use state_is_tuple=True.", self) if attn_size is None: attn_size = cell.output_size if attn_vec_size is None: @@ -1160,8 +1190,8 @@ class AttentionCellWrapper(rnn_cell_impl.RNNCell): else: states = state state = array_ops.slice(states, [0, 0], [-1, self._cell.state_size]) - attns = array_ops.slice( - states, [0, self._cell.state_size], [-1, self._attn_size]) + attns = array_ops.slice(states, [0, self._cell.state_size], + [-1, self._attn_size]) attn_states = array_ops.slice( states, [0, self._cell.state_size + self._attn_size], [-1, self._attn_size * self._attn_length]) @@ -1199,8 +1229,8 @@ class AttentionCellWrapper(rnn_cell_impl.RNNCell): tanh = math_ops.tanh with vs.variable_scope("attention"): - k = vs.get_variable( - "attn_w", [1, 1, self._attn_size, self._attn_vec_size]) + k = vs.get_variable("attn_w", + [1, 1, self._attn_size, self._attn_vec_size]) v = vs.get_variable("attn_v", [self._attn_vec_size]) hidden = array_ops.reshape(attn_states, [-1, self._attn_length, 1, self._attn_size]) @@ -1227,7 +1257,8 @@ class HighwayWrapper(rnn_cell_impl.RNNCell): https://arxiv.org/abs/1505.00387 """ - def __init__(self, cell, + def __init__(self, + cell, couple_carry_transform_gates=True, carry_bias_init=1.0): """Constructs a `HighwayWrapper` for `cell`. @@ -1259,8 +1290,7 @@ class HighwayWrapper(rnn_cell_impl.RNNCell): carry_weight = vs.get_variable("carry_w", [input_size, input_size]) carry_bias = vs.get_variable( "carry_b", [input_size], - initializer=init_ops.constant_initializer( - self._carry_bias_init)) + initializer=init_ops.constant_initializer(self._carry_bias_init)) carry = math_ops.sigmoid(nn_ops.xw_plus_b(inp, carry_weight, carry_bias)) if self._couple_carry_transform_gates: transform = 1 - carry @@ -1269,11 +1299,9 @@ class HighwayWrapper(rnn_cell_impl.RNNCell): [input_size, input_size]) transform_bias = vs.get_variable( "transform_b", [input_size], - initializer=init_ops.constant_initializer( - -self._carry_bias_init)) - transform = math_ops.sigmoid(nn_ops.xw_plus_b(inp, - transform_weight, - transform_bias)) + initializer=init_ops.constant_initializer(-self._carry_bias_init)) + transform = math_ops.sigmoid( + nn_ops.xw_plus_b(inp, transform_weight, transform_bias)) return inp * carry + out * transform def __call__(self, inputs, state, scope=None): @@ -1293,9 +1321,11 @@ class HighwayWrapper(rnn_cell_impl.RNNCell): """ outputs, new_state = self._cell(inputs, state, scope=scope) nest.assert_same_structure(inputs, outputs) + # Ensure shapes match def assert_shape_match(inp, out): inp.get_shape().assert_is_compatible_with(out.get_shape()) + nest.map_structure(assert_shape_match, inputs, outputs) res_outputs = nest.map_structure(self._highway, inputs, outputs) return (res_outputs, new_state) @@ -1321,10 +1351,16 @@ class LayerNormBasicLSTMCell(rnn_cell_impl.RNNCell): Stanislau Semeniuta, Aliaksei Severyn, Erhardt Barth. """ - def __init__(self, num_units, forget_bias=1.0, - input_size=None, activation=math_ops.tanh, - layer_norm=True, norm_gain=1.0, norm_shift=0.0, - dropout_keep_prob=1.0, dropout_prob_seed=None, + def __init__(self, + num_units, + forget_bias=1.0, + input_size=None, + activation=math_ops.tanh, + layer_norm=True, + norm_gain=1.0, + norm_shift=0.0, + dropout_keep_prob=1.0, + dropout_prob_seed=None, reuse=None): """Initializes the basic LSTM cell. @@ -1409,8 +1445,8 @@ class LayerNormBasicLSTMCell(rnn_cell_impl.RNNCell): if (not isinstance(self._keep_prob, float)) or self._keep_prob < 1: g = nn_ops.dropout(g, self._keep_prob, seed=self._seed) - new_c = (c * math_ops.sigmoid(f + self._forget_bias) - + math_ops.sigmoid(i) * g) + new_c = ( + c * math_ops.sigmoid(f + self._forget_bias) + math_ops.sigmoid(i) * g) if self._layer_norm: new_c = self._norm(new_c, "state", dtype=dtype) new_h = self._activation(new_c) * math_ops.sigmoid(o) @@ -1432,8 +1468,7 @@ class NASCell(rnn_cell_impl.RNNCell): The class uses an optional projection layer. """ - def __init__(self, num_units, num_proj=None, - use_biases=False, reuse=None): + def __init__(self, num_units, num_proj=None, use_biases=False, reuse=None): """Initialize the parameters for a NAS cell. Args: @@ -1503,12 +1538,10 @@ class NASCell(rnn_cell_impl.RNNCell): raise ValueError("Could not infer input size from inputs.get_shape()[-1]") # Variables for the NAS cell. W_m is all matrices multiplying the # hiddenstate and W_inputs is all matrices multiplying the inputs. - concat_w_m = vs.get_variable( - "recurrent_kernel", [num_proj, 8 * self._num_units], - dtype) + concat_w_m = vs.get_variable("recurrent_kernel", + [num_proj, 8 * self._num_units], dtype) concat_w_inputs = vs.get_variable( - "kernel", [input_size.value, 8 * self._num_units], - dtype) + "kernel", [input_size.value, 8 * self._num_units], dtype) m_matrix = math_ops.matmul(m_prev, concat_w_m) inputs_matrix = math_ops.matmul(inputs, concat_w_inputs) @@ -1523,10 +1556,10 @@ class NASCell(rnn_cell_impl.RNNCell): # The NAS cell branches into 8 different splits for both the hiddenstate # and the input - m_matrix_splits = array_ops.split(axis=1, num_or_size_splits=8, - value=m_matrix) - inputs_matrix_splits = array_ops.split(axis=1, num_or_size_splits=8, - value=inputs_matrix) + m_matrix_splits = array_ops.split( + axis=1, num_or_size_splits=8, value=m_matrix) + inputs_matrix_splits = array_ops.split( + axis=1, num_or_size_splits=8, value=inputs_matrix) # First layer layer1_0 = sigmoid(inputs_matrix_splits[0] + m_matrix_splits[0]) @@ -1558,9 +1591,8 @@ class NASCell(rnn_cell_impl.RNNCell): # Projection layer if specified if self._num_proj is not None: - concat_w_proj = vs.get_variable( - "projection_weights", [self._num_units, self._num_proj], - dtype) + concat_w_proj = vs.get_variable("projection_weights", + [self._num_units, self._num_proj], dtype) new_m = math_ops.matmul(new_m, concat_w_proj) new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_m) @@ -1583,8 +1615,12 @@ class UGRNNCell(rnn_cell_impl.RNNCell): "Capacity and Trainability in Recurrent Neural Networks" Proc. ICLR 2017. """ - def __init__(self, num_units, initializer=None, forget_bias=1.0, - activation=math_ops.tanh, reuse=None): + def __init__(self, + num_units, + initializer=None, + forget_bias=1.0, + activation=math_ops.tanh, + reuse=None): """Initialize the parameters for an UGRNN cell. Args: @@ -1639,8 +1675,8 @@ class UGRNNCell(rnn_cell_impl.RNNCell): if input_size.value is None: raise ValueError("Could not infer input size from inputs.get_shape()[-1]") - with vs.variable_scope(vs.get_variable_scope(), - initializer=self._initializer): + with vs.variable_scope( + vs.get_variable_scope(), initializer=self._initializer): cell_inputs = array_ops.concat([inputs, state], 1) if self._linear is None: self._linear = _Linear(cell_inputs, 2 * self._num_units, True) @@ -1680,9 +1716,13 @@ class IntersectionRNNCell(rnn_cell_impl.RNNCell): RNNs so it may not achieve best performance with depth 1. """ - def __init__(self, num_units, num_in_proj=None, - initializer=None, forget_bias=1.0, - y_activation=nn_ops.relu, reuse=None): + def __init__(self, + num_units, + num_in_proj=None, + initializer=None, + forget_bias=1.0, + y_activation=nn_ops.relu, + reuse=None): """Initialize the parameters for an +RNN cell. Args: @@ -1746,8 +1786,8 @@ class IntersectionRNNCell(rnn_cell_impl.RNNCell): if input_size.value is None: raise ValueError("Could not infer input size from inputs.get_shape()[-1]") - with vs.variable_scope(vs.get_variable_scope(), - initializer=self._initializer): + with vs.variable_scope( + vs.get_variable_scope(), initializer=self._initializer): # read-in projections (should be used for first layer in deep +RNN # to transform size of inputs from I --> N) if input_size.value != self._num_units: @@ -1764,13 +1804,13 @@ class IntersectionRNNCell(rnn_cell_impl.RNNCell): n_dim = i_dim = self._num_units cell_inputs = array_ops.concat([inputs, state], 1) if self._linear2 is None: - self._linear2 = _Linear(cell_inputs, 2*n_dim + 2*i_dim, True) + self._linear2 = _Linear(cell_inputs, 2 * n_dim + 2 * i_dim, True) rnn_matrix = self._linear2(cell_inputs) - gh_act = rnn_matrix[:, :n_dim] # b x n - h_act = rnn_matrix[:, n_dim:2*n_dim] # b x n - gy_act = rnn_matrix[:, 2*n_dim:2*n_dim+i_dim] # b x i - y_act = rnn_matrix[:, 2*n_dim+i_dim:2*n_dim+2*i_dim] # b x i + gh_act = rnn_matrix[:, :n_dim] # b x n + h_act = rnn_matrix[:, n_dim:2 * n_dim] # b x n + gy_act = rnn_matrix[:, 2 * n_dim:2 * n_dim + i_dim] # b x i + y_act = rnn_matrix[:, 2 * n_dim + i_dim:2 * n_dim + 2 * i_dim] # b x i h = tanh(h_act) y = self._y_activation(y_act) @@ -1816,6 +1856,7 @@ class CompiledWrapper(rnn_cell_impl.RNNCell): if self._compile_stateful: compile_ops = True else: + def compile_ops(node_def): global _REGISTERED_OPS if _REGISTERED_OPS is None: @@ -1826,10 +1867,7 @@ class CompiledWrapper(rnn_cell_impl.RNNCell): return self._cell(inputs, state, scope=scope) -def _random_exp_initializer(minval, - maxval, - seed=None, - dtype=dtypes.float32): +def _random_exp_initializer(minval, maxval, seed=None, dtype=dtypes.float32): """Returns an exponential distribution initializer. Args: @@ -1848,10 +1886,7 @@ def _random_exp_initializer(minval, del partition_info # Unused. return math_ops.exp( random_ops.random_uniform( - shape, - math_ops.log(minval), - math_ops.log(maxval), - dtype, + shape, math_ops.log(minval), math_ops.log(maxval), dtype, seed=seed)) return _initializer @@ -1955,8 +1990,7 @@ class PhasedLSTMCell(rnn_cell_impl.RNNCell): if self._linear1 is None: self._linear1 = _Linear(in_mask_gates, 2 * self._num_units, True) - mask_gates = math_ops.sigmoid( - self._linear1(in_mask_gates)) + mask_gates = math_ops.sigmoid(self._linear1(in_mask_gates)) [input_gate, forget_gate] = array_ops.split( axis=1, num_or_size_splits=2, value=mask_gates) @@ -1980,12 +2014,12 @@ class PhasedLSTMCell(rnn_cell_impl.RNNCell): period = vs.get_variable( "period", [self._num_units], - initializer=_random_exp_initializer( - self._period_init_min, self._period_init_max)) + initializer=_random_exp_initializer(self._period_init_min, + self._period_init_max)) phase = vs.get_variable( "phase", [self._num_units], - initializer=init_ops.random_uniform_initializer( - 0., period.initial_value)) + initializer=init_ops.random_uniform_initializer(0., + period.initial_value)) ratio_on = vs.get_variable( "ratio_on", [self._num_units], initializer=init_ops.constant_initializer(self._ratio_on), @@ -2007,6 +2041,7 @@ class PhasedLSTMCell(rnn_cell_impl.RNNCell): return new_h, new_state + class ConvLSTMCell(rnn_cell_impl.RNNCell): """Convolutional LSTM recurrent network cell. @@ -2040,7 +2075,7 @@ class ConvLSTMCell(rnn_cell_impl.RNNCell): """ super(ConvLSTMCell, self).__init__(name=name) - if conv_ndims != len(input_shape)-1: + if conv_ndims != len(input_shape) - 1: raise ValueError("Invalid input_shape {} for conv_ndims={}.".format( input_shape, conv_ndims)) @@ -2059,8 +2094,8 @@ class ConvLSTMCell(rnn_cell_impl.RNNCell): state_size = tensor_shape.TensorShape( self._input_shape[:-1] + [self._output_channels]) self._state_size = rnn_cell_impl.LSTMStateTuple(state_size, state_size) - self._output_size = tensor_shape.TensorShape(self._input_shape[:-1] - + [self._total_output_channels]) + self._output_size = tensor_shape.TensorShape( + self._input_shape[:-1] + [self._total_output_channels]) @property def output_size(self): @@ -2072,13 +2107,10 @@ class ConvLSTMCell(rnn_cell_impl.RNNCell): def call(self, inputs, state, scope=None): cell, hidden = state - new_hidden = _conv([inputs, hidden], - self._kernel_shape, - 4*self._output_channels, - self._use_bias) - gates = array_ops.split(value=new_hidden, - num_or_size_splits=4, - axis=self._conv_ndims+1) + new_hidden = _conv([inputs, hidden], self._kernel_shape, + 4 * self._output_channels, self._use_bias) + gates = array_ops.split( + value=new_hidden, num_or_size_splits=4, axis=self._conv_ndims + 1) input_gate, new_input, forget_gate, output_gate = gates new_cell = math_ops.sigmoid(forget_gate + self._forget_bias) * cell @@ -2090,29 +2122,35 @@ class ConvLSTMCell(rnn_cell_impl.RNNCell): new_state = rnn_cell_impl.LSTMStateTuple(new_cell, output) return output, new_state + class Conv1DLSTMCell(ConvLSTMCell): """1D Convolutional LSTM recurrent network cell. https://arxiv.org/pdf/1506.04214v1.pdf """ + def __init__(self, name="conv_1d_lstm_cell", **kwargs): """Construct Conv1DLSTM. See `ConvLSTMCell` for more details.""" super(Conv1DLSTMCell, self).__init__(conv_ndims=1, **kwargs) + class Conv2DLSTMCell(ConvLSTMCell): """2D Convolutional LSTM recurrent network cell. https://arxiv.org/pdf/1506.04214v1.pdf """ + def __init__(self, name="conv_2d_lstm_cell", **kwargs): """Construct Conv2DLSTM. See `ConvLSTMCell` for more details.""" super(Conv2DLSTMCell, self).__init__(conv_ndims=2, **kwargs) + class Conv3DLSTMCell(ConvLSTMCell): """3D Convolutional LSTM recurrent network cell. https://arxiv.org/pdf/1506.04214v1.pdf """ + def __init__(self, name="conv_3d_lstm_cell", **kwargs): """Construct Conv3DLSTM. See `ConvLSTMCell` for more details.""" super(Conv3DLSTMCell, self).__init__(conv_ndims=3, **kwargs) @@ -2137,7 +2175,7 @@ def _conv(args, filter_size, num_features, bias, bias_start=0.0): shapes = [a.get_shape().as_list() for a in args] shape_length = len(shapes[0]) for shape in shapes: - if len(shape) not in [3,4,5]: + if len(shape) not in [3, 4, 5]: raise ValueError("Conv Linear expects 3D, 4D " "or 5D arguments: %s" % str(shapes)) if len(shape) != len(shapes[0]): @@ -2148,40 +2186,36 @@ def _conv(args, filter_size, num_features, bias, bias_start=0.0): dtype = [a.dtype for a in args][0] # determine correct conv operation - if shape_length == 3: + if shape_length == 3: conv_op = nn_ops.conv1d strides = 1 elif shape_length == 4: conv_op = nn_ops.conv2d - strides = shape_length*[1] + strides = shape_length * [1] elif shape_length == 5: conv_op = nn_ops.conv3d - strides = shape_length*[1] + strides = shape_length * [1] # Now the computation. kernel = vs.get_variable( - "kernel", - filter_size + [total_arg_size_depth, num_features], - dtype=dtype) + "kernel", filter_size + [total_arg_size_depth, num_features], dtype=dtype) if len(args) == 1: - res = conv_op(args[0], - kernel, - strides, - padding='SAME') + res = conv_op(args[0], kernel, strides, padding="SAME") else: - res = conv_op(array_ops.concat(axis=shape_length-1, values=args), - kernel, - strides, - padding='SAME') + res = conv_op( + array_ops.concat(axis=shape_length - 1, values=args), + kernel, + strides, + padding="SAME") if not bias: return res bias_term = vs.get_variable( "biases", [num_features], dtype=dtype, - initializer=init_ops.constant_initializer( - bias_start, dtype=dtype)) + initializer=init_ops.constant_initializer(bias_start, dtype=dtype)) return res + bias_term + class GLSTMCell(rnn_cell_impl.RNNCell): """Group LSTM cell (G-LSTM). @@ -2193,8 +2227,13 @@ class GLSTMCell(rnn_cell_impl.RNNCell): "Factorization Tricks for LSTM Networks", ICLR 2017 workshop. """ - def __init__(self, num_units, initializer=None, num_proj=None, - number_of_groups=1, forget_bias=1.0, activation=math_ops.tanh, + def __init__(self, + num_units, + initializer=None, + num_proj=None, + number_of_groups=1, + forget_bias=1.0, + activation=math_ops.tanh, reuse=None): """Initialize the parameters of G-LSTM cell. @@ -2231,11 +2270,15 @@ class GLSTMCell(rnn_cell_impl.RNNCell): if self._num_proj: if self._num_proj % self._number_of_groups != 0: raise ValueError("num_proj must be divisible by number_of_groups") - self._group_shape = [int(self._num_proj / self._number_of_groups), - int(self._num_units / self._number_of_groups)] + self._group_shape = [ + int(self._num_proj / self._number_of_groups), + int(self._num_units / self._number_of_groups) + ] else: - self._group_shape = [int(self._num_units / self._number_of_groups), - int(self._num_units / self._number_of_groups)] + self._group_shape = [ + int(self._num_units / self._number_of_groups), + int(self._num_units / self._number_of_groups) + ] if num_proj: self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_proj) @@ -2243,7 +2286,7 @@ class GLSTMCell(rnn_cell_impl.RNNCell): else: self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_units) self._output_size = num_units - self._linear1 = None + self._linear1 = [None] * number_of_groups self._linear2 = None @property @@ -2267,10 +2310,11 @@ class GLSTMCell(rnn_cell_impl.RNNCell): subset of inputs corresponding to group "group_id", a Tensor, 2D, [batch x num_units/number_of_groups] """ - return array_ops.slice(input_=inputs, - begin=[0, group_id * group_size], - size=[self._batch_size, group_size], - name=("GLSTM_group%d_input_generation" % group_id)) + return array_ops.slice( + input_=inputs, + begin=[0, group_id * group_size], + size=[self._batch_size, group_size], + name=("GLSTM_group%d_input_generation" % group_id)) def call(self, inputs, state): """Run one step of G-LSTM. @@ -2309,13 +2353,18 @@ class GLSTMCell(rnn_cell_impl.RNNCell): for group_id in range(self._number_of_groups): with vs.variable_scope("group%d" % group_id): x_g_id = array_ops.concat( - [self._get_input_for_group(inputs, group_id, - self._group_shape[0]), - self._get_input_for_group(m_prev, group_id, - self._group_shape[0])], axis=1) - if self._linear1 is None: - self._linear1 = _Linear(x_g_id, 4 * self._group_shape[1], False) - R_k = self._linear1(x_g_id) # pylint: disable=invalid-name + [ + self._get_input_for_group(inputs, group_id, + self._group_shape[0]), + self._get_input_for_group(m_prev, group_id, + self._group_shape[0]) + ], + axis=1) + linear = self._linear1[group_id] + if linear is None: + linear = _Linear(x_g_id, 4 * self._group_shape[1], False) + self._linear1[group_id] = linear + R_k = linear(x_g_id) # pylint: disable=invalid-name i_k, j_k, f_k, o_k = array_ops.split(R_k, 4, 1) i_parts.append(i_k) @@ -2323,34 +2372,35 @@ class GLSTMCell(rnn_cell_impl.RNNCell): f_parts.append(f_k) o_parts.append(o_k) - bi = vs.get_variable(name="bias_i", - shape=[self._num_units], - dtype=dtype, - initializer= - init_ops.constant_initializer(0.0, dtype=dtype)) - bj = vs.get_variable(name="bias_j", - shape=[self._num_units], - dtype=dtype, - initializer= - init_ops.constant_initializer(0.0, dtype=dtype)) - bf = vs.get_variable(name="bias_f", - shape=[self._num_units], - dtype=dtype, - initializer= - init_ops.constant_initializer(0.0, dtype=dtype)) - bo = vs.get_variable(name="bias_o", - shape=[self._num_units], - dtype=dtype, - initializer= - init_ops.constant_initializer(0.0, dtype=dtype)) + bi = vs.get_variable( + name="bias_i", + shape=[self._num_units], + dtype=dtype, + initializer=init_ops.constant_initializer(0.0, dtype=dtype)) + bj = vs.get_variable( + name="bias_j", + shape=[self._num_units], + dtype=dtype, + initializer=init_ops.constant_initializer(0.0, dtype=dtype)) + bf = vs.get_variable( + name="bias_f", + shape=[self._num_units], + dtype=dtype, + initializer=init_ops.constant_initializer(0.0, dtype=dtype)) + bo = vs.get_variable( + name="bias_o", + shape=[self._num_units], + dtype=dtype, + initializer=init_ops.constant_initializer(0.0, dtype=dtype)) i = nn_ops.bias_add(array_ops.concat(i_parts, axis=1), bi) j = nn_ops.bias_add(array_ops.concat(j_parts, axis=1), bj) f = nn_ops.bias_add(array_ops.concat(f_parts, axis=1), bf) o = nn_ops.bias_add(array_ops.concat(o_parts, axis=1), bo) - c = (math_ops.sigmoid(f + self._forget_bias) * c_prev + - math_ops.sigmoid(i) * math_ops.tanh(j)) + c = ( + math_ops.sigmoid(f + self._forget_bias) * c_prev + + math_ops.sigmoid(i) * math_ops.tanh(j)) m = math_ops.sigmoid(o) * self._activation(c) if self._num_proj is not None: @@ -2633,12 +2683,14 @@ class LayerNormLSTMCell(rnn_cell_impl.RNNCell): return m, new_state -class SRUCell(rnn_cell_impl._LayerRNNCell): +class SRUCell(rnn_cell_impl.LayerRNNCell): """SRU, Simple Recurrent Unit + Implementation based on Training RNNs as Fast as CNNs (cf. https://arxiv.org/abs/1709.02755). - This variation of RNN cell is characterized by the simplified data dependence + This variation of RNN cell is characterized by the simplified data + dependence between hidden states of two consecutive time steps. Traditionally, hidden states from a cell at time step t-1 needs to be multiplied with a matrix W_hh before being fed into the ensuing cell at time step t. @@ -2656,8 +2708,8 @@ class SRUCell(rnn_cell_impl._LayerRNNCell): will share weights, but to avoid mistakes we require reuse=True in such cases. """ - def __init__(self, num_units, - activation=None, reuse=None, name=None): + + def __init__(self, num_units, activation=None, reuse=None, name=None): super(SRUCell, self).__init__(_reuse=reuse, name=name) self._num_units = num_units self._activation = activation or math_ops.tanh @@ -2675,30 +2727,14 @@ class SRUCell(rnn_cell_impl._LayerRNNCell): def build(self, inputs_shape): if inputs_shape[1].value is None: - raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s" - % inputs_shape) + raise ValueError( + "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape) input_depth = inputs_shape[1].value - # Here the contributor believes that the following constraints - # are implied. The reasoning is explained here with reference to - # the paper https://arxiv.org/pdf/1709.02755.pdf upon which this - # implementation is based. - # In section 2.1 Equation 5, specifically: - # h_t = r_t \odot g(c_t) + (1 - r_t) \odot x_t - # the pointwise operation between r_t and x_t means they have - # the same shape (since we are implementing an RNN cell, braodcasting - # does not happen to input of a single timestep); by the same - # reasons, x_t has the same shape as h_t, essentially mandating that - # input_depth = unit_num. - if input_depth != self._num_units: - raise ValueError("SRU requires input_depth == num_units, got " - "input_depth = %s, num_units = %s" % (input_depth, - self._num_units)) - self._kernel = self.add_variable( rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - shape=[input_depth, 3 * self._num_units]) + shape=[input_depth, 4 * self._num_units]) self._bias = self.add_variable( rnn_cell_impl._BIAS_VARIABLE_NAME, @@ -2711,15 +2747,276 @@ class SRUCell(rnn_cell_impl._LayerRNNCell): """Simple recurrent unit (SRU) with num_units cells.""" U = math_ops.matmul(inputs, self._kernel) - x_bar, f_intermediate, r_intermediate = array_ops.split(value=U, - num_or_size_splits=3, - axis=1) + x_bar, f_intermediate, r_intermediate, x_tx = array_ops.split( + value=U, num_or_size_splits=4, axis=1) - f_r = math_ops.sigmoid(nn_ops.bias_add(array_ops.concat( - [f_intermediate, r_intermediate], 1), self._bias)) + f_r = math_ops.sigmoid( + nn_ops.bias_add( + array_ops.concat([f_intermediate, r_intermediate], 1), self._bias)) f, r = array_ops.split(value=f_r, num_or_size_splits=2, axis=1) c = f * state + (1.0 - f) * x_bar - h = r * self._activation(c) + (1.0 - r) * inputs + h = r * self._activation(c) + (1.0 - r) * x_tx return h, c + + +class WeightNormLSTMCell(rnn_cell_impl.RNNCell): + """Weight normalized LSTM Cell. Adapted from `rnn_cell_impl.LSTMCell`. + + The weight-norm implementation is based on: + https://arxiv.org/abs/1602.07868 + Tim Salimans, Diederik P. Kingma. + Weight Normalization: A Simple Reparameterization to Accelerate + Training of Deep Neural Networks + + The default LSTM implementation based on: + http://www.bioinf.jku.at/publications/older/2604.pdf + S. Hochreiter and J. Schmidhuber. + "Long Short-Term Memory". Neural Computation, 9(8):1735-1780, 1997. + + The class uses optional peephole connections, optional cell clipping + and an optional projection layer. + + The optional peephole implementation is based on: + https://research.google.com/pubs/archive/43905.pdf + Hasim Sak, Andrew Senior, and Francoise Beaufays. + "Long short-term memory recurrent neural network architectures for + large scale acoustic modeling." INTERSPEECH, 2014. + """ + + def __init__(self, + num_units, + norm=True, + use_peepholes=False, + cell_clip=None, + initializer=None, + num_proj=None, + proj_clip=None, + forget_bias=1, + activation=None, + reuse=None): + """Initialize the parameters of a weight-normalized LSTM cell. + + Args: + num_units: int, The number of units in the LSTM cell + norm: If `True`, apply normalization to the weight matrices. If False, + the result is identical to that obtained from `rnn_cell_impl.LSTMCell` + use_peepholes: bool, set `True` to enable diagonal/peephole connections. + cell_clip: (optional) A float value, if provided the cell state is clipped + by this value prior to the cell output activation. + initializer: (optional) The initializer to use for the weight matrices. + num_proj: (optional) int, The output dimensionality for the projection + matrices. If None, no projection is performed. + proj_clip: (optional) A float value. If `num_proj > 0` and `proj_clip` is + provided, then the projected values are clipped elementwise to within + `[-proj_clip, proj_clip]`. + forget_bias: Biases of the forget gate are initialized by default to 1 + in order to reduce the scale of forgetting at the beginning of + the training. + activation: Activation function of the inner states. Default: `tanh`. + reuse: (optional) Python boolean describing whether to reuse variables + in an existing scope. If not `True`, and the existing scope already has + the given variables, an error is raised. + """ + super(WeightNormLSTMCell, self).__init__(_reuse=reuse) + + self._scope = "wn_lstm_cell" + self._num_units = num_units + self._norm = norm + self._initializer = initializer + self._use_peepholes = use_peepholes + self._cell_clip = cell_clip + self._num_proj = num_proj + self._proj_clip = proj_clip + self._activation = activation or math_ops.tanh + self._forget_bias = forget_bias + + self._weights_variable_name = "kernel" + self._bias_variable_name = "bias" + + if num_proj: + self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_proj) + self._output_size = num_proj + else: + self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_units) + self._output_size = num_units + + @property + def state_size(self): + return self._state_size + + @property + def output_size(self): + return self._output_size + + def _normalize(self, weight, name): + """Apply weight normalization. + + Args: + weight: a 2D tensor with known number of columns. + name: string, variable name for the normalizer. + Returns: + A tensor with the same shape as `weight`. + """ + + output_size = weight.get_shape().as_list()[1] + g = vs.get_variable(name, [output_size], dtype=weight.dtype) + return nn_impl.l2_normalize(weight, dim=0) * g + + def _linear(self, + args, + output_size, + norm, + bias, + bias_initializer=None, + kernel_initializer=None): + """Linear map: sum_i(args[i] * W[i]), where W[i] is a variable. + + Args: + args: a 2D Tensor or a list of 2D, batch x n, Tensors. + output_size: int, second dimension of W[i]. + bias: boolean, whether to add a bias term or not. + bias_initializer: starting value to initialize the bias + (default is all zeros). + kernel_initializer: starting value to initialize the weight. + + Returns: + A 2D Tensor with shape [batch x output_size] equal to + sum_i(args[i] * W[i]), where W[i]s are newly created matrices. + + Raises: + ValueError: if some of the arguments has unspecified or wrong shape. + """ + if args is None or (nest.is_sequence(args) and not args): + raise ValueError("`args` must be specified") + if not nest.is_sequence(args): + args = [args] + + # Calculate the total size of arguments on dimension 1. + total_arg_size = 0 + shapes = [a.get_shape() for a in args] + for shape in shapes: + if shape.ndims != 2: + raise ValueError("linear is expecting 2D arguments: %s" % shapes) + if shape[1].value is None: + raise ValueError("linear expects shape[1] to be provided for shape %s, " + "but saw %s" % (shape, shape[1])) + else: + total_arg_size += shape[1].value + + dtype = [a.dtype for a in args][0] + + # Now the computation. + scope = vs.get_variable_scope() + with vs.variable_scope(scope) as outer_scope: + weights = vs.get_variable( + self._weights_variable_name, [total_arg_size, output_size], + dtype=dtype, + initializer=kernel_initializer) + if norm: + wn = [] + st = 0 + with ops.control_dependencies(None): + for i in range(len(args)): + en = st + shapes[i][1].value + wn.append( + self._normalize(weights[st:en, :], name="norm_{}".format(i))) + st = en + + weights = array_ops.concat(wn, axis=0) + + if len(args) == 1: + res = math_ops.matmul(args[0], weights) + else: + res = math_ops.matmul(array_ops.concat(args, 1), weights) + if not bias: + return res + + with vs.variable_scope(outer_scope) as inner_scope: + inner_scope.set_partitioner(None) + if bias_initializer is None: + bias_initializer = init_ops.constant_initializer(0.0, dtype=dtype) + + biases = vs.get_variable( + self._bias_variable_name, [output_size], + dtype=dtype, + initializer=bias_initializer) + + return nn_ops.bias_add(res, biases) + + def call(self, inputs, state): + """Run one step of LSTM. + + Args: + inputs: input Tensor, 2D, batch x num_units. + state: A tuple of state Tensors, both `2-D`, with column sizes + `c_state` and `m_state`. + + Returns: + A tuple containing: + + - A `2-D, [batch x output_dim]`, Tensor representing the output of the + LSTM after reading `inputs` when previous state was `state`. + Here output_dim is: + num_proj if num_proj was set, + num_units otherwise. + - Tensor(s) representing the new state of LSTM after reading `inputs` when + the previous state was `state`. Same type and shape(s) as `state`. + + Raises: + ValueError: If input size cannot be inferred from inputs via + static shape inference. + """ + dtype = inputs.dtype + num_units = self._num_units + sigmoid = math_ops.sigmoid + c, h = state + + input_size = inputs.get_shape().with_rank(2)[1] + if input_size.value is None: + raise ValueError("Could not infer input size from inputs.get_shape()[-1]") + + with vs.variable_scope(self._scope, initializer=self._initializer): + + concat = self._linear( + [inputs, h], 4 * num_units, norm=self._norm, bias=True) + + # i = input_gate, j = new_input, f = forget_gate, o = output_gate + i, j, f, o = array_ops.split(value=concat, num_or_size_splits=4, axis=1) + + if self._use_peepholes: + w_f_diag = vs.get_variable("w_f_diag", shape=[num_units], dtype=dtype) + w_i_diag = vs.get_variable("w_i_diag", shape=[num_units], dtype=dtype) + w_o_diag = vs.get_variable("w_o_diag", shape=[num_units], dtype=dtype) + + new_c = ( + c * sigmoid(f + self._forget_bias + w_f_diag * c) + + sigmoid(i + w_i_diag * c) * self._activation(j)) + else: + new_c = ( + c * sigmoid(f + self._forget_bias) + + sigmoid(i) * self._activation(j)) + + if self._cell_clip is not None: + # pylint: disable=invalid-unary-operand-type + new_c = clip_ops.clip_by_value(new_c, -self._cell_clip, self._cell_clip) + # pylint: enable=invalid-unary-operand-type + if self._use_peepholes: + new_h = sigmoid(o + w_o_diag * new_c) * self._activation(new_c) + else: + new_h = sigmoid(o) * self._activation(new_c) + + if self._num_proj is not None: + with vs.variable_scope("projection"): + new_h = self._linear( + new_h, self._num_proj, norm=self._norm, bias=False) + + if self._proj_clip is not None: + # pylint: disable=invalid-unary-operand-type + new_h = clip_ops.clip_by_value(new_h, -self._proj_clip, + self._proj_clip) + # pylint: enable=invalid-unary-operand-type + + new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_h) + return new_h, new_state diff --git a/tensorflow/contrib/rnn/python/tools/checkpoint_convert.py b/tensorflow/contrib/rnn/python/tools/checkpoint_convert.py index 5536a01328676e5fe01251fefdaaecb0f9569918..460e172a6d949804319b8833e34b6590f5fcf93b 100644 --- a/tensorflow/contrib/rnn/python/tools/checkpoint_convert.py +++ b/tensorflow/contrib/rnn/python/tools/checkpoint_convert.py @@ -128,10 +128,8 @@ RNN_NAME_REPLACEMENTS = collections.OrderedDict([ 'attention_cell_wrapper/attention/bias'), ############################################################################ # contrib/legacy_seq2seq/python/ops/seq2seq.py - ('attention_decoder/weights', - 'attention_decoder/kernel'), - ('attention_decoder/biases', - 'attention_decoder/bias'), + ('attention_decoder/weights', 'attention_decoder/kernel'), + ('attention_decoder/biases', 'attention_decoder/bias'), ('attention_decoder/Attention_0/weights', 'attention_decoder/Attention_0/kernel'), ('attention_decoder/Attention_0/biases', @@ -140,6 +138,19 @@ RNN_NAME_REPLACEMENTS = collections.OrderedDict([ 'attention_decoder/AttnOutputProjection/kernel'), ('attention_decoder/AttnOutputProjection/biases', 'attention_decoder/AttnOutputProjection/bias'), + # contrib/legacy_seq2seq/python/ops/seq2seq.py before cl/140060366 + ('attention_decoder/Attention_0/Linear/Bias', + 'attention_decoder/Attention_0/bias'), + ('attention_decoder/Attention_0/Linear/Matrix', + 'attention_decoder/Attention_0/kernel'), + ('attention_decoder/AttnOutputProjection/Linear/Bias', + 'attention_decoder/AttnOutputProjection/bias'), + ('attention_decoder/AttnOutputProjection/Linear/Matrix', + 'attention_decoder/AttnOutputProjection/kernel'), + ('attention_decoder/LSTMCell/B', 'attention_decoder/lstm_cell/bias'), + ('attention_decoder/LSTMCell/W_0', 'attention_decoder/lstm_cell/kernel'), + ('attention_decoder/Linear/Bias', 'attention_decoder/bias'), + ('attention_decoder/Linear/Matrix', 'attention_decoder/kernel') ]) _RNN_SHARDED_NAME_REPLACEMENTS = collections.OrderedDict([ diff --git a/tensorflow/contrib/rnn/python/tools/checkpoint_convert_test.py b/tensorflow/contrib/rnn/python/tools/checkpoint_convert_test.py index a9e79494639418c22b7380b5b78092052fbf305d..b4785ee395a2452d9595d81c3bdb88711a8fe66a 100644 --- a/tensorflow/contrib/rnn/python/tools/checkpoint_convert_test.py +++ b/tensorflow/contrib/rnn/python/tools/checkpoint_convert_test.py @@ -67,7 +67,7 @@ class CheckpointConvertTest(test.TestCase): self._old_ckpt_path, self._new_ckpt_path) self.assertTrue(glob.glob(self._new_ckpt_path + "*")) self.assertItemsEqual( - ["a"] + list(checkpoint_convert.RNN_NAME_REPLACEMENTS.values()), + set(checkpoint_convert.RNN_NAME_REPLACEMENTS.values()).union(["a"]), new_var_map.keys()) self.assertEqual(checkpoint_convert.RNN_NAME_REPLACEMENTS, conversion_map) diff --git a/tensorflow/contrib/saved_model/cc/saved_model/signature_def_utils.h b/tensorflow/contrib/saved_model/cc/saved_model/signature_def_utils.h index c0df224bc8cffcb485db38dea270600c71070dff..b732cdd41e5c39793c17fa920c115e2bbe96f5de 100644 --- a/tensorflow/contrib/saved_model/cc/saved_model/signature_def_utils.h +++ b/tensorflow/contrib/saved_model/cc/saved_model/signature_def_utils.h @@ -15,8 +15,8 @@ limitations under the License. // Helpers for working with the SignatureDefs of TensorFlow SavedModels. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_SAVED_MODEL_CC_SAVED_MODEL_SIGNATURE_DEF_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_SAVED_MODEL_CC_SAVED_MODEL_SIGNATURE_DEF_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_SAVED_MODEL_CC_SAVED_MODEL_SIGNATURE_DEF_UTILS_H_ +#define TENSORFLOW_CONTRIB_SAVED_MODEL_CC_SAVED_MODEL_SIGNATURE_DEF_UTILS_H_ #include #include @@ -66,4 +66,4 @@ Status FindOutputTensorNameByKey(const SignatureDef& signature_def, } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_SAVED_MODEL_CC_SAVED_MODEL_SIGNATURE_DEF_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_SAVED_MODEL_CC_SAVED_MODEL_SIGNATURE_DEF_UTILS_H_ diff --git a/tensorflow/contrib/seq2seq/kernels/beam_search_ops.cc b/tensorflow/contrib/seq2seq/kernels/beam_search_ops.cc index 64973ccccdc962757a727d7183bd70e94edcfd1b..dfa12e873a6aca806031c48d6f92e0432d0ea6e0 100644 --- a/tensorflow/contrib/seq2seq/kernels/beam_search_ops.cc +++ b/tensorflow/contrib/seq2seq/kernels/beam_search_ops.cc @@ -80,12 +80,12 @@ class GatherTreeOp : public OpKernel { max_sequence_lengths.shape().DebugString())); Tensor* beams; OP_REQUIRES_OK(ctx, ctx->allocate_output(0, step_ids_shape, &beams)); - typename TTypes::ConstTensor step_ids_t = step_ids.tensor(); - typename TTypes::ConstTensor parent_ids_t = parent_ids.tensor(); + typename TTypes::ConstTensor step_ids_t(step_ids.tensor()); + typename TTypes::ConstTensor parent_ids_t(parent_ids.tensor()); typename TTypes::ConstVec max_seq_lens_t = max_sequence_lengths.vec(); - typename TTypes::ConstScalar end_token_t = end_token.scalar(); - typename TTypes::Tensor beams_t = beams->tensor(); + typename TTypes::ConstScalar end_token_t(end_token.scalar()); + typename TTypes::Tensor beams_t(beams->tensor()); const T end_token_value = end_token_t(); functor::GatherTree()(ctx, device, step_ids_t, parent_ids_t, max_seq_lens_t, end_token_value, beams_t); diff --git a/tensorflow/contrib/seq2seq/kernels/beam_search_ops.h b/tensorflow/contrib/seq2seq/kernels/beam_search_ops.h index 693b02dc437afdf14c38e4224c5469bb3e569540..34da8c82cdab9b6f82af328c49a365ae1cb951ed 100644 --- a/tensorflow/contrib/seq2seq/kernels/beam_search_ops.h +++ b/tensorflow/contrib/seq2seq/kernels/beam_search_ops.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_SEQ2SEQ_KERNELS_BEAM_SEARCH_OPS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_SEQ2SEQ_KERNELS_BEAM_SEARCH_OPS_H_ +#ifndef TENSORFLOW_CONTRIB_SEQ2SEQ_KERNELS_BEAM_SEARCH_OPS_H_ +#define TENSORFLOW_CONTRIB_SEQ2SEQ_KERNELS_BEAM_SEARCH_OPS_H_ #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/framework/tensor_types.h" @@ -38,4 +38,4 @@ struct GatherTree { } // namespace functor } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_SEQ2SEQ_KERNELS_BEAM_SEARCH_OPS_H_ +#endif // TENSORFLOW_CONTRIB_SEQ2SEQ_KERNELS_BEAM_SEARCH_OPS_H_ diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py index e5d591788fa6350bb59458d12e9bb01c910ba37d..b427dff88b2d586ccf8c512bb498cdaf879ac781 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py @@ -69,7 +69,7 @@ class AttentionWrapperTest(test.TestCase): def assertAllCloseOrEqual(self, x, y, **kwargs): if isinstance(x, np.ndarray) or isinstance(x, float): return super(AttentionWrapperTest, self).assertAllClose( - x, y, atol=1e-4, **kwargs) + x, y, atol=1e-3, **kwargs) else: self.assertAllEqual(x, y, **kwargs) @@ -80,6 +80,28 @@ class AttentionWrapperTest(test.TestCase): self.assertEqual(state.time, None) self.assertEqual(new_state.time, 1) + def testAttentionWrapperStateShapePropgation(self): + batch_size = 5 + max_time = 5 + num_units = 5 + + memory = random_ops.random_uniform( + [batch_size, max_time, num_units], seed=1) + mechanism = wrapper.LuongAttention(num_units, memory) + cell = wrapper.AttentionWrapper(rnn_cell.LSTMCell(num_units), mechanism) + + # Create zero state with static batch size. + static_state = cell.zero_state(batch_size, dtypes.float32) + # Create zero state without static batch size. + state = cell.zero_state(array_ops.shape(memory)[0], dtypes.float32) + + state = static_state.clone( + cell_state=state.cell_state, attention=state.attention) + + self.assertEqual(state.cell_state.c.shape, static_state.cell_state.c.shape) + self.assertEqual(state.cell_state.h.shape, static_state.cell_state.h.shape) + self.assertEqual(state.attention.shape, static_state.attention.shape) + def _testWithAttention(self, create_attention_mechanism, expected_final_output, @@ -276,7 +298,7 @@ class AttentionWrapperTest(test.TestCase): rnn_output=ResultSummary( shape=(5, 3, 6), dtype=dtype('float32'), mean=-0.00597103), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=1.4)) + shape=(5, 3), dtype=dtype('int32'), mean=1.6)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -305,7 +327,7 @@ class AttentionWrapperTest(test.TestCase): rnn_output=ResultSummary( shape=(5, 3, 6), dtype=dtype('float32'), mean=-0.0052615386), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=1.4666666666666666)) + shape=(5, 3), dtype=dtype('int32'), mean=1.3333333333)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -336,7 +358,7 @@ class AttentionWrapperTest(test.TestCase): rnn_output=ResultSummary( shape=(5, 3, 6), dtype=dtype('float32'), mean=-0.0052615386), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=1.4666666666666666)) + shape=(5, 3), dtype=dtype('int32'), mean=1.3333333333333333)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -578,7 +600,7 @@ class AttentionWrapperTest(test.TestCase): rnn_output=ResultSummary( shape=(5, 3, 6), dtype=dtype('float32'), mean=-0.0025896581), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=1.8666666666666667)) + shape=(5, 3), dtype=dtype('int32'), mean=1.6)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -594,7 +616,7 @@ class AttentionWrapperTest(test.TestCase): shape=(5, 8), dtype=dtype('float32'), mean=0.028698336), alignment_history=()) expected_final_alignment_history = ResultSummary( - shape=(3, 5, 8), dtype=dtype('float32'), mean=0.046009291) + shape=(3, 5, 8), dtype=dtype('float32'), mean=0.04865776002407074) self._testWithAttention( create_attention_mechanism, @@ -761,9 +783,9 @@ class AttentionWrapperTest(test.TestCase): expected_final_output = BasicDecoderOutput( rnn_output=ResultSummary( - shape=(5, 3, 20), dtype=dtype('float32'), mean=0.11691988), + shape=(5, 3, 20), dtype=dtype('float32'), mean=0.11798714846372604), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=7.2666666666666666)) + shape=(5, 3), dtype=dtype('int32'), mean=7.933333333333334)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -771,7 +793,7 @@ class AttentionWrapperTest(test.TestCase): h=ResultSummary( shape=(5, 9), dtype=dtype('float32'), mean=-0.0018835809)), attention=ResultSummary( - shape=(5, 20), dtype=dtype('float32'), mean=0.11680689), + shape=(5, 20), dtype=dtype('float32'), mean=0.11798714846372604), time=3, alignments=( ResultSummary(shape=(5, 8), dtype=dtype('float32'), mean=0.125), diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/beam_search_decoder_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/beam_search_decoder_test.py index d2beac5f31460ec1c0d978a9f6fcd0e0f09cb9b4..926554031775202d7f7d9018cf6ae4efb34fe96b 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/beam_search_decoder_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/beam_search_decoder_test.py @@ -46,20 +46,18 @@ class TestGatherTree(test.TestCase): # create (batch_size, max_time, beam_width) matrix and transpose it predicted_ids = np.array( - [[[1, 2, 3], [4, 5, 6], [7, 8, 9]], - [[2, 3, 4], [5, 6, 7], [8, 9, 10]]], + [[[1, 2, 3], [4, 5, 6], [7, 8, 9]], [[2, 3, 4], [5, 6, 7], [8, 9, 10]]], dtype=np.int32).transpose([1, 0, 2]) parent_ids = np.array( - [[[0, 0, 0], [0, 1, 1], [2, 1, 2]], - [[0, 0, 0], [1, 2, 0], [2, 1, 1]]], + [[[0, 0, 0], [0, 1, 1], [2, 1, 2]], [[0, 0, 0], [1, 2, 0], [2, 1, 1]]], dtype=np.int32).transpose([1, 0, 2]) # sequence_lengths is shaped (batch_size = 3) max_sequence_lengths = [3, 3] - expected_result = np.array( - [[[2, 2, 2], [6, 5, 6], [7, 8, 9]], - [[2, 4, 4], [7, 6, 6], [8, 9, 10]]]).transpose([1, 0, 2]) + expected_result = np.array([[[2, 2, 2], [6, 5, 6], [7, 8, 9]], + [[2, 4, 4], [7, 6, 6], + [8, 9, 10]]]).transpose([1, 0, 2]) res = beam_search_ops.gather_tree( predicted_ids, @@ -157,8 +155,8 @@ class TestBeamStep(test.TestCase): self.assertAllEqual(outputs_.predicted_ids, [[3, 3, 2], [2, 2, 1]]) self.assertAllEqual(outputs_.parent_ids, [[1, 0, 0], [2, 1, 0]]) self.assertAllEqual(next_state_.lengths, [[3, 3, 3], [3, 3, 3]]) - self.assertAllEqual(next_state_.finished, [[False, False, False], - [False, False, False]]) + self.assertAllEqual(next_state_.finished, + [[False, False, False], [False, False, False]]) expected_log_probs = [] expected_log_probs.append(state_.log_probs[0][[1, 0, 0]]) @@ -212,8 +210,8 @@ class TestBeamStep(test.TestCase): self.assertAllEqual(outputs_.parent_ids, [[1, 0, 0], [1, 2, 0]]) self.assertAllEqual(outputs_.predicted_ids, [[0, 3, 2], [2, 0, 1]]) self.assertAllEqual(next_state_.lengths, [[1, 3, 3], [3, 1, 3]]) - self.assertAllEqual(next_state_.finished, [[True, False, False], - [False, True, False]]) + self.assertAllEqual(next_state_.finished, + [[True, False, False], [False, True, False]]) expected_log_probs = [] expected_log_probs.append(state_.log_probs[0][[1, 0, 0]]) @@ -225,6 +223,100 @@ class TestBeamStep(test.TestCase): self.assertAllEqual(next_state_.log_probs, expected_log_probs) +class TestLargeBeamStep(test.TestCase): + """Tests large beam step. + + Tests a single step of beam search in such case that beam size is larger than + vocabulary size. + """ + + def setUp(self): + super(TestLargeBeamStep, self).setUp() + self.batch_size = 2 + self.beam_width = 8 + self.vocab_size = 5 + self.end_token = 0 + self.length_penalty_weight = 0.6 + + def test_step(self): + + def get_probs(): + """this simulates the initialize method in BeamSearchDecoder.""" + log_prob_mask = array_ops.one_hot( + array_ops.zeros([self.batch_size], dtype=dtypes.int32), + depth=self.beam_width, + on_value=True, + off_value=False, + dtype=dtypes.bool) + + log_prob_zeros = array_ops.zeros( + [self.batch_size, self.beam_width], dtype=dtypes.float32) + log_prob_neg_inf = array_ops.ones( + [self.batch_size, self.beam_width], dtype=dtypes.float32) * -np.Inf + + log_probs = array_ops.where(log_prob_mask, log_prob_zeros, + log_prob_neg_inf) + return log_probs + + log_probs = get_probs() + dummy_cell_state = array_ops.zeros([self.batch_size, self.beam_width]) + + # pylint: disable=invalid-name + _finished = array_ops.one_hot( + array_ops.zeros([self.batch_size], dtype=dtypes.int32), + depth=self.beam_width, + on_value=False, + off_value=True, + dtype=dtypes.bool) + _lengths = np.zeros([self.batch_size, self.beam_width], dtype=np.int64) + _lengths[:, 0] = 2 + _lengths = constant_op.constant(_lengths, dtype=dtypes.int64) + + beam_state = beam_search_decoder.BeamSearchDecoderState( + cell_state=dummy_cell_state, + log_probs=log_probs, + lengths=_lengths, + finished=_finished) + + logits_ = np.full([self.batch_size, self.beam_width, self.vocab_size], + 0.0001) + logits_[0, 0, 2] = 1.9 + logits_[0, 0, 3] = 2.1 + logits_[0, 1, 3] = 3.1 + logits_[0, 1, 4] = 0.9 + logits_[1, 0, 1] = 0.5 + logits_[1, 1, 2] = 2.7 + logits_[1, 2, 2] = 10.0 + logits_[1, 2, 3] = 0.2 + logits = constant_op.constant(logits_, dtype=dtypes.float32) + log_probs = nn_ops.log_softmax(logits) + + outputs, next_beam_state = beam_search_decoder._beam_search_step( + time=2, + logits=logits, + next_cell_state=dummy_cell_state, + beam_state=beam_state, + batch_size=ops.convert_to_tensor(self.batch_size), + beam_width=self.beam_width, + end_token=self.end_token, + length_penalty_weight=self.length_penalty_weight) + + with self.test_session() as sess: + outputs_, next_state_, _, _ = sess.run( + [outputs, next_beam_state, beam_state, log_probs]) + + self.assertEqual(outputs_.predicted_ids[0, 0], 3) + self.assertEqual(outputs_.predicted_ids[0, 1], 2) + self.assertEqual(outputs_.predicted_ids[1, 0], 1) + neg_inf = -np.Inf + self.assertAllEqual( + next_state_.log_probs[:, -3:], + [[neg_inf, neg_inf, neg_inf], [neg_inf, neg_inf, neg_inf]]) + self.assertEqual((next_state_.log_probs[:, :-3] > neg_inf).all(), True) + self.assertEqual((next_state_.lengths[:, :-3] > 0).all(), True) + self.assertAllEqual(next_state_.lengths[:, -3:], [[0, 0, 0], [0, 0, 0]]) + + class BeamSearchDecoderTest(test.TestCase): def _testDynamicDecodeRNN(self, time_major, has_attention): @@ -250,8 +342,8 @@ class BeamSearchDecoderTest(test.TestCase): initial_state = cell.zero_state(batch_size, dtypes.float32) if has_attention: inputs = array_ops.placeholder_with_default( - np.random.randn(batch_size, decoder_max_time, - input_depth).astype(np.float32), + np.random.randn(batch_size, decoder_max_time, input_depth).astype( + np.float32), shape=(None, None, input_depth)) tiled_inputs = beam_search_decoder.tile_batch( inputs, multiplier=beam_width) @@ -271,8 +363,7 @@ class BeamSearchDecoderTest(test.TestCase): cell_state = cell.zero_state( dtype=dtypes.float32, batch_size=batch_size_tensor * beam_width) if has_attention: - cell_state = cell_state.clone( - cell_state=initial_state) + cell_state = cell_state.clone(cell_state=initial_state) bsd = beam_search_decoder.BeamSearchDecoder( cell=cell, embedding=embedding, diff --git a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py index 36bfc5685d51350acb86c270626cce5cec9647f8..0a53fd66dbe4d28ea102773b9c5bae50b9d18e9c 100644 --- a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py +++ b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py @@ -24,6 +24,7 @@ import math import numpy as np +from tensorflow.contrib.framework.python.framework import tensor_util from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape @@ -330,7 +331,7 @@ def _luong_score(query, keys, scale): # batched matmul on: # [batch_size, 1, depth] . [batch_size, depth, max_time] # resulting in an output shape of: - # [batch_time, 1, max_time]. + # [batch_size, 1, max_time]. # we then squeeze out the center singleton dimension. score = math_ops.matmul(query, keys, transpose_b=True) score = array_ops.squeeze(score, [1]) @@ -923,8 +924,7 @@ class LuongMonotonicAttention(_BaseMonotonicAttentionMechanism): _monotonic_probability_fn, sigmoid_noise=sigmoid_noise, mode=mode, seed=sigmoid_noise_seed) super(LuongMonotonicAttention, self).__init__( - query_layer=layers_core.Dense( - num_units, name="query_layer", use_bias=False, dtype=dtype), + query_layer=None, memory_layer=layers_core.Dense( num_units, name="memory_layer", use_bias=False, dtype=dtype), memory=memory, @@ -989,6 +989,10 @@ class AttentionWrapperState( def clone(self, **kwargs): """Clone this object, overriding components provided by kwargs. + The new state fields' shape must match original state fields' shape. This + will be validated, and original fields' shape will be propagated to new + fields. + Example: ```python @@ -1004,7 +1008,16 @@ class AttentionWrapperState( A new `AttentionWrapperState` whose properties are the same as this one, except any overridden properties as provided in `kwargs`. """ - return super(AttentionWrapperState, self)._replace(**kwargs) + def with_same_shape(old, new): + """Check and set new tensor's shape.""" + if isinstance(old, ops.Tensor) and isinstance(new, ops.Tensor): + return tensor_util.with_same_shape(old, new) + return new + + return nest.map_structure( + with_same_shape, + self, + super(AttentionWrapperState, self)._replace(**kwargs)) def hardmax(logits, name=None): diff --git a/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py b/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py index c7c4182f0d9a17dacebc1cda693cda6eaaf8451f..ed226239b860e2250072a28a5538b816642ec54b 100644 --- a/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py @@ -90,7 +90,7 @@ class BasicDecoder(decoder.Decoder): output_shape_with_unknown_batch = nest.map_structure( lambda s: tensor_shape.TensorShape([None]).concatenate(s), size) - layer_output_shape = self._output_layer._compute_output_shape( # pylint: disable=protected-access + layer_output_shape = self._output_layer.compute_output_shape( output_shape_with_unknown_batch) return nest.map_structure(lambda s: s[1:], layer_output_shape) diff --git a/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py b/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py index 5be0c92243da10af438be97fab982515266be1de..d6184d61095f727f9dcab56fe59e2601868c1624 100644 --- a/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function import collections - import numpy as np from tensorflow.contrib.seq2seq.python.ops import beam_search_ops @@ -38,7 +37,6 @@ from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.ops import tensor_array_ops from tensorflow.python.util import nest - __all__ = [ "BeamSearchDecoderOutput", "BeamSearchDecoderState", @@ -49,8 +47,8 @@ __all__ = [ class BeamSearchDecoderState( - collections.namedtuple("BeamSearchDecoderState", ("cell_state", "log_probs", - "finished", "lengths"))): + collections.namedtuple("BeamSearchDecoderState", + ("cell_state", "log_probs", "finished", "lengths"))): pass @@ -67,7 +65,8 @@ class FinalBeamSearchDecoderOutput( Args: predicted_ids: The final prediction. A tensor of shape - `[T, batch_size, beam_width]`. + `[batch_size, T, beam_width]` (or `[T, batch_size, beam_width]` if + `output_time_major` is True). Beams are ordered from best to worst. beam_search_decoder_output: An instance of `BeamSearchDecoderOutput` that describes the state of the beam search. """ @@ -85,11 +84,12 @@ def _tile_batch(t, multiplier): tiled_static_batch_size = ( t.shape[0].value * multiplier if t.shape[0].value is not None else None) tiled = array_ops.tile(array_ops.expand_dims(t, 1), tiling) - tiled = array_ops.reshape( - tiled, array_ops.concat(([shape_t[0] * multiplier], shape_t[1:]), 0)) + tiled = array_ops.reshape(tiled, + array_ops.concat( + ([shape_t[0] * multiplier], shape_t[1:]), 0)) tiled.set_shape( - tensor_shape.TensorShape( - [tiled_static_batch_size]).concatenate(t.shape[1:])) + tensor_shape.TensorShape([tiled_static_batch_size]).concatenate( + t.shape[1:])) return tiled @@ -197,8 +197,8 @@ class BeamSearchDecoder(decoder.Decoder): """ if not rnn_cell_impl._like_rnncell(cell): # pylint: disable=protected-access raise TypeError("cell must be an RNNCell, received: %s" % type(cell)) - if (output_layer is not None - and not isinstance(output_layer, layers_base.Layer)): + if (output_layer is not None and + not isinstance(output_layer, layers_base.Layer)): raise TypeError( "output_layer must be a Layer, received: %s" % type(output_layer)) self._cell = cell @@ -223,13 +223,17 @@ class BeamSearchDecoder(decoder.Decoder): self._beam_width = beam_width self._length_penalty_weight = length_penalty_weight self._initial_cell_state = nest.map_structure( - self._maybe_split_batch_beams, - initial_state, self._cell.state_size) + self._maybe_split_batch_beams, initial_state, self._cell.state_size) self._start_tokens = array_ops.tile( array_ops.expand_dims(self._start_tokens, 1), [1, self._beam_width]) self._start_inputs = self._embedding_fn(self._start_tokens) - self._finished = array_ops.zeros( - [self._batch_size, self._beam_width], dtype=dtypes.bool) + + self._finished = array_ops.one_hot( + array_ops.zeros([self._batch_size], dtype=dtypes.int32), + depth=self._beam_width, + on_value=False, + off_value=True, + dtype=dtypes.bool) @property def batch_size(self): @@ -247,9 +251,8 @@ class BeamSearchDecoder(decoder.Decoder): # dimensions to get the output size of the rnn with the layer # applied to the top. output_shape_with_unknown_batch = nest.map_structure( - lambda s: tensor_shape.TensorShape([None]).concatenate(s), - size) - layer_output_shape = self._output_layer._compute_output_shape( # pylint: disable=protected-access + lambda s: tensor_shape.TensorShape([None]).concatenate(s), size) + layer_output_shape = self._output_layer.compute_output_shape( output_shape_with_unknown_batch) return nest.map_structure(lambda s: s[1:], layer_output_shape) @@ -297,11 +300,16 @@ class BeamSearchDecoder(decoder.Decoder): """ finished, start_inputs = self._finished, self._start_inputs + log_probs = array_ops.one_hot( # shape(batch_sz, beam_sz) + array_ops.zeros([self._batch_size], dtype=dtypes.int32), + depth=self._beam_width, + on_value=0.0, + off_value=-np.Inf, + dtype=nest.flatten(self._initial_cell_state)[0].dtype) + initial_state = BeamSearchDecoderState( cell_state=self._initial_cell_state, - log_probs=array_ops.zeros( - [self._batch_size, self._beam_width], - dtype=nest.flatten(self._initial_cell_state)[0].dtype), + log_probs=log_probs, finished=finished, lengths=array_ops.zeros( [self._batch_size, self._beam_width], dtype=dtypes.int64)) @@ -358,11 +366,12 @@ class BeamSearchDecoder(decoder.Decoder): t_shape = array_ops.shape(t) static_batch_size = tensor_util.constant_value(self._batch_size) batch_size_beam_width = ( - None if static_batch_size is None - else static_batch_size * self._beam_width) + None + if static_batch_size is None else static_batch_size * self._beam_width) reshaped_t = array_ops.reshape( - t, array_ops.concat( - ([self._batch_size * self._beam_width], t_shape[2:]), 0)) + t, + array_ops.concat(([self._batch_size * self._beam_width], t_shape[2:]), + 0)) reshaped_t.set_shape( (tensor_shape.TensorShape([batch_size_beam_width]).concatenate(s))) return reshaped_t @@ -391,8 +400,9 @@ class BeamSearchDecoder(decoder.Decoder): s = tensor_shape.TensorShape(s) t_shape = array_ops.shape(t) reshaped_t = array_ops.reshape( - t, array_ops.concat( - ([self._batch_size, self._beam_width], t_shape[1:]), 0)) + t, + array_ops.concat(([self._batch_size, self._beam_width], t_shape[1:]), + 0)) static_batch_size = tensor_util.constant_value(self._batch_size) expected_reshaped_shape = tensor_shape.TensorShape( [static_batch_size, self._beam_width]).concatenate(s) @@ -402,8 +412,8 @@ class BeamSearchDecoder(decoder.Decoder): "We expected it to have shape " "(batch_size, beam_width, depth) == %s. Perhaps you " "forgot to create a zero_state with " - "batch_size=encoder_batch_size * beam_width?" - % (reshaped_t.shape, expected_reshaped_shape)) + "batch_size=encoder_batch_size * beam_width?" % + (reshaped_t.shape, expected_reshaped_shape)) reshaped_t.set_shape(expected_reshaped_shape) return reshaped_t @@ -475,15 +485,13 @@ class BeamSearchDecoder(decoder.Decoder): cell_state = state.cell_state inputs = nest.map_structure( lambda inp: self._merge_batch_beams(inp, s=inp.shape[2:]), inputs) - cell_state = nest.map_structure( - self._maybe_merge_batch_beams, - cell_state, self._cell.state_size) + cell_state = nest.map_structure(self._maybe_merge_batch_beams, cell_state, + self._cell.state_size) cell_outputs, next_cell_state = self._cell(inputs, cell_state) cell_outputs = nest.map_structure( lambda out: self._split_batch_beams(out, out.shape[1:]), cell_outputs) next_cell_state = nest.map_structure( - self._maybe_split_batch_beams, - next_cell_state, self._cell.state_size) + self._maybe_split_batch_beams, next_cell_state, self._cell.state_size) if self._output_layer is not None: cell_outputs = self._output_layer(cell_outputs) @@ -546,7 +554,8 @@ def _beam_search_step(time, logits, next_cell_state, beam_state, batch_size, lengths_to_add = array_ops.one_hot( indices=array_ops.fill([batch_size, beam_width], end_token), depth=vocab_size, - on_value=np.int64(0), off_value=np.int64(1), + on_value=np.int64(0), + off_value=np.int64(1), dtype=dtypes.int64) add_mask = math_ops.to_int64(math_ops.logical_not(previously_finished)) lengths_to_add *= array_ops.expand_dims(add_mask, 2) @@ -562,18 +571,11 @@ def _beam_search_step(time, logits, next_cell_state, beam_state, batch_size, time = ops.convert_to_tensor(time, name="time") # During the first time step we only consider the initial beam scores_shape = array_ops.shape(scores) - scores_flat = control_flow_ops.cond( - time > 0, - lambda: array_ops.reshape(scores, [batch_size, -1]), - lambda: scores[:, 0]) - num_available_beam = control_flow_ops.cond( - time > 0, lambda: math_ops.reduce_prod(scores_shape[1:]), - lambda: math_ops.reduce_prod(scores_shape[2:])) + scores_flat = array_ops.reshape(scores, [batch_size, -1]) # Pick the next beams according to the specified successors function - next_beam_size = math_ops.minimum( - ops.convert_to_tensor(beam_width, dtype=dtypes.int32, name="beam_width"), - num_available_beam) + next_beam_size = ops.convert_to_tensor( + beam_width, dtype=dtypes.int32, name="beam_width") next_beam_scores, word_indices = nn_ops.top_k(scores_flat, k=next_beam_size) next_beam_scores.set_shape([static_batch_size, beam_width]) @@ -592,11 +594,11 @@ def _beam_search_step(time, logits, next_cell_state, beam_state, batch_size, # name="next_beam_word_ids") # would be a lot cleaner but for reasons unclear, that hides the results of # the op which prevents capturing it with tfdbg debug ops. - raw_next_word_ids = math_ops.mod(word_indices, vocab_size, - name="next_beam_word_ids") + raw_next_word_ids = math_ops.mod( + word_indices, vocab_size, name="next_beam_word_ids") next_word_ids = math_ops.to_int32(raw_next_word_ids) - next_beam_ids = math_ops.to_int32(word_indices / vocab_size, - name="next_beam_parent_ids") + next_beam_ids = math_ops.to_int32( + word_indices / vocab_size, name="next_beam_parent_ids") # Append new ids to current predictions previously_finished = _tensor_gather_helper( @@ -605,9 +607,10 @@ def _beam_search_step(time, logits, next_cell_state, beam_state, batch_size, batch_size=batch_size, range_size=beam_width, gather_shape=[-1]) - next_finished = math_ops.logical_or(previously_finished, - math_ops.equal(next_word_ids, end_token), - name="next_beam_finished") + next_finished = math_ops.logical_or( + previously_finished, + math_ops.equal(next_word_ids, end_token), + name="next_beam_finished") # Calculate the length of the next predictions. # 1. Finished beams remain unchanged. @@ -768,8 +771,12 @@ def _maybe_tensor_gather_helper(gather_indices, gather_from, batch_size, return gather_from -def _tensor_gather_helper(gather_indices, gather_from, batch_size, - range_size, gather_shape, name=None): +def _tensor_gather_helper(gather_indices, + gather_from, + batch_size, + range_size, + gather_shape, + name=None): """Helper for gathering the right indices from the tensor. This works by reshaping gather_from to gather_shape (e.g. [-1]) and then @@ -800,9 +807,9 @@ def _tensor_gather_helper(gather_indices, gather_from, batch_size, array_ops.reshape(gather_from, gather_shape), gather_indices) final_shape = array_ops.shape(gather_from)[:1 + len(gather_shape)] static_batch_size = tensor_util.constant_value(batch_size) - final_static_shape = (tensor_shape.TensorShape([static_batch_size]) - .concatenate( - gather_from.shape[1:1 + len(gather_shape)])) + final_static_shape = ( + tensor_shape.TensorShape([static_batch_size]).concatenate( + gather_from.shape[1:1 + len(gather_shape)])) output = array_ops.reshape(output, final_shape, name="output") output.set_shape(final_static_shape) return output diff --git a/tensorflow/contrib/seq2seq/python/ops/helper.py b/tensorflow/contrib/seq2seq/python/ops/helper.py index ef3722ee41bb0b49e5f81d4d6514e2f40d2ad9f1..3245cc5e72154289ea3ba000b9a30586a7ad03a9 100644 --- a/tensorflow/contrib/seq2seq/python/ops/helper.py +++ b/tensorflow/contrib/seq2seq/python/ops/helper.py @@ -184,6 +184,7 @@ class TrainingHelper(Helper): """ with ops.name_scope(name, "TrainingHelper", [inputs, sequence_length]): inputs = ops.convert_to_tensor(inputs, name="inputs") + self._inputs = inputs if not time_major: inputs = nest.map_structure(_transpose_batch_time, inputs) @@ -200,6 +201,14 @@ class TrainingHelper(Helper): self._batch_size = array_ops.size(sequence_length) + @property + def inputs(self): + return self._inputs + + @property + def sequence_length(self): + return self._sequence_length + @property def batch_size(self): return self._batch_size diff --git a/tensorflow/contrib/session_bundle/bundle_shim.cc b/tensorflow/contrib/session_bundle/bundle_shim.cc index a367ea059c9a2017e94c1541e42d6296665cc466..4fc36d85edf4ac2d48769d209f0b78d6d29d9a62 100644 --- a/tensorflow/contrib/session_bundle/bundle_shim.cc +++ b/tensorflow/contrib/session_bundle/bundle_shim.cc @@ -371,9 +371,15 @@ Status LoadSessionBundleOrSavedModelBundle( return LoadSavedModelFromLegacySessionBundlePath( session_options, run_options, export_dir, saved_model_bundle); } - return Status(error::Code::NOT_FOUND, - "Session bundle or SavedModel bundle not found at specified " - "export location"); + return Status( + error::Code::NOT_FOUND, + strings::StrCat( + "Specified file path does not appear to contain a:\n" + "- Session bundle (should have a file called `export.meta`)\n" + "- or, SavedModel bundle (should have a file called " + "`saved_model.pb`)\n" + "Specified file path: ", + export_dir)); } } // namespace serving diff --git a/tensorflow/contrib/session_bundle/bundle_shim.h b/tensorflow/contrib/session_bundle/bundle_shim.h index e24efa0de14824044591b954b8465ebeebc10dd5..4628b6ab1b1164addef6aaf930a0dbe7091cd16d 100644 --- a/tensorflow/contrib/session_bundle/bundle_shim.h +++ b/tensorflow/contrib/session_bundle/bundle_shim.h @@ -15,8 +15,8 @@ limitations under the License. // Shim for systems that need to load both SessionBundle and // SavedModelBundle interchangeably during migration to SavedModel. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_BUNDLE_SHIM_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_BUNDLE_SHIM_H_ +#ifndef TENSORFLOW_CONTRIB_SESSION_BUNDLE_BUNDLE_SHIM_H_ +#define TENSORFLOW_CONTRIB_SESSION_BUNDLE_BUNDLE_SHIM_H_ #include @@ -67,4 +67,4 @@ Status LoadSessionBundleOrSavedModelBundle( } // namespace serving } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_BUNDLE_SHIM_H_ +#endif // TENSORFLOW_CONTRIB_SESSION_BUNDLE_BUNDLE_SHIM_H_ diff --git a/tensorflow/contrib/session_bundle/bundle_shim.py b/tensorflow/contrib/session_bundle/bundle_shim.py index 062c9cc68046c59ffd04190dad0fa69f5f9dfa0a..1db97020a2a81f4d034543e722a6cb7ba823f44a 100644 --- a/tensorflow/contrib/session_bundle/bundle_shim.py +++ b/tensorflow/contrib/session_bundle/bundle_shim.py @@ -82,7 +82,8 @@ def _convert_default_signature_to_signature_def(signatures): """ default_signature = signatures.default_signature signature_def = meta_graph_pb2.SignatureDef() - if default_signature.WhichOneof("type") == "regression_signature": + if (default_signature.WhichOneof("type") == + legacy_constants.REGRESSION_SIGNATURE): regression_signature = default_signature.regression_signature signature_def.method_name = signature_constants.REGRESS_METHOD_NAME _add_input_to_signature_def(regression_signature.input.tensor_name, @@ -91,7 +92,8 @@ def _convert_default_signature_to_signature_def(signatures): _add_output_to_signature_def(regression_signature.output.tensor_name, signature_constants.REGRESS_OUTPUTS, signature_def) - elif default_signature.WhichOneof("type") == "classification_signature": + elif (default_signature.WhichOneof("type") == + legacy_constants.CLASSIFICATION_SIGNATURE): classification_signature = default_signature.classification_signature signature_def.method_name = signature_constants.CLASSIFY_METHOD_NAME _add_input_to_signature_def(classification_signature.input.tensor_name, @@ -132,8 +134,9 @@ def _convert_named_signatures_to_signature_def(signatures): signature_constants.PREDICT_OUTPUTS] # TODO(pdudnik): what if there are other signatures? Mimic cr/140900781 once # it is submitted. - if (input_signature.WhichOneof("type") != "generic_signature" or - output_signature.WhichOneof("type") != "generic_signature"): + if (input_signature.WhichOneof("type") != legacy_constants.GENERIC_SIGNATURE + or output_signature.WhichOneof("type") != + legacy_constants.GENERIC_SIGNATURE): raise RuntimeError("Named input and output signatures can only be " "up-converted if they are generic signature. " "Input signature type is %s, output signature type is " diff --git a/tensorflow/contrib/session_bundle/bundle_shim_test.cc b/tensorflow/contrib/session_bundle/bundle_shim_test.cc index 72f32a0f5554e4dd3e7cbf498a57ee6bfba57211..9a1dd9303f43591888dc49984d81c4a0c6af9846 100644 --- a/tensorflow/contrib/session_bundle/bundle_shim_test.cc +++ b/tensorflow/contrib/session_bundle/bundle_shim_test.cc @@ -493,17 +493,15 @@ TEST(BundleShimTest, DefaultAndNamedSignatureWithPredict) { ASSERT_FALSE( actual_signature_def_predict->second.inputs().find("foo-input") == actual_signature_def_predict->second.inputs().end()); - EXPECT_EQ("foo-input", - actual_signature_def_predict->second.inputs() - .find("foo-input") - ->second.name()); + EXPECT_EQ("foo-input", actual_signature_def_predict->second.inputs() + .find("foo-input") + ->second.name()); ASSERT_FALSE( actual_signature_def_predict->second.outputs().find("foo-output") == actual_signature_def_predict->second.outputs().end()); - EXPECT_EQ("foo-output", - actual_signature_def_predict->second.outputs() - .find("foo-output") - ->second.name()); + EXPECT_EQ("foo-output", actual_signature_def_predict->second.outputs() + .find("foo-output") + ->second.name()); EXPECT_EQ(kPredictMethodName, actual_signature_def_predict->second.method_name()); } diff --git a/tensorflow/contrib/session_bundle/constants.py b/tensorflow/contrib/session_bundle/constants.py index 6ced73241afdda047b8feacb26fedd72363b6240..e833baee791f97df5829ee289bcaf17c31a17deb 100644 --- a/tensorflow/contrib/session_bundle/constants.py +++ b/tensorflow/contrib/session_bundle/constants.py @@ -32,3 +32,6 @@ INIT_OP_KEY = "serving_init_op" SIGNATURES_KEY = "serving_signatures" ASSETS_KEY = "serving_assets" GRAPH_KEY = "serving_graph" +REGRESSION_SIGNATURE = "regression_signature" +CLASSIFICATION_SIGNATURE = "classification_signature" +GENERIC_SIGNATURE = "generic_signature" diff --git a/tensorflow/contrib/session_bundle/exporter.py b/tensorflow/contrib/session_bundle/exporter.py index f6f663aae766b783b85139f57a93e10f553e6bd1..08983337fccc138d40eb959cecc5bf9e47cf6cac 100644 --- a/tensorflow/contrib/session_bundle/exporter.py +++ b/tensorflow/contrib/session_bundle/exporter.py @@ -281,11 +281,12 @@ class Exporter(object): tmp_export_dir = compat.as_text(export_dir) + "-tmp" gfile.MakeDirs(tmp_export_dir) - self._saver.save(sess, - os.path.join( - compat.as_text(tmp_export_dir), - compat.as_text(constants.EXPORT_BASE_NAME)), - meta_graph_suffix=constants.EXPORT_SUFFIX_NAME) + self._saver.save( + sess, + os.path.join( + compat.as_text(tmp_export_dir), + compat.as_text(constants.EXPORT_BASE_NAME)), + meta_graph_suffix=constants.EXPORT_SUFFIX_NAME) # Run the asset callback. if self._assets_callback and self._assets_to_copy: @@ -301,12 +302,12 @@ class Exporter(object): if exports_to_keep: # create a simple parser that pulls the export_version from the directory. def parser(path): - if os.name == 'nt': - match = re.match("^" + export_dir_base.replace('\\','/') + "/(\\d{8})$", - path.path.replace('\\','/')) + if os.name == "nt": + match = re.match( + "^" + export_dir_base.replace("\\", "/") + "/(\\d{8})$", + path.path.replace("\\", "/")) else: - match = re.match("^" + export_dir_base + "/(\\d{8})$", - path.path) + match = re.match("^" + export_dir_base + "/(\\d{8})$", path.path) if not match: return None return path._replace(export_version=int(match.group(1))) diff --git a/tensorflow/contrib/session_bundle/gc.py b/tensorflow/contrib/session_bundle/gc.py index 249c23c88f3043403e322b73b6c9df97e932a92a..514cc0f652c8d174bdb9bff2b2cf1ea38fdd7b1f 100644 --- a/tensorflow/contrib/session_bundle/gc.py +++ b/tensorflow/contrib/session_bundle/gc.py @@ -70,7 +70,6 @@ import heapq import math import os -from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.python.platform import gfile from tensorflow.python.util.deprecation import deprecated diff --git a/tensorflow/contrib/session_bundle/session_bundle.h b/tensorflow/contrib/session_bundle/session_bundle.h index 2ff258411d1928cea7da4f637ffe94f144b2b60a..b2be46efa6d1e7ceb1fb66a7148735b86cc68dd3 100644 --- a/tensorflow/contrib/session_bundle/session_bundle.h +++ b/tensorflow/contrib/session_bundle/session_bundle.h @@ -15,8 +15,8 @@ limitations under the License. // Low-level functionality for setting up a inference Session. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_SESSION_BUNDLE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_SESSION_BUNDLE_H_ +#ifndef TENSORFLOW_CONTRIB_SESSION_BUNDLE_SESSION_BUNDLE_H_ +#define TENSORFLOW_CONTRIB_SESSION_BUNDLE_SESSION_BUNDLE_H_ #include @@ -82,4 +82,4 @@ bool IsPossibleExportDirectory(const StringPiece export_dir); } // namespace serving } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_SESSION_BUNDLE_H_ +#endif // TENSORFLOW_CONTRIB_SESSION_BUNDLE_SESSION_BUNDLE_H_ diff --git a/tensorflow/contrib/session_bundle/signature.cc b/tensorflow/contrib/session_bundle/signature.cc index 7133875ad53e77625bbe799f4f886c074a08f1bd..ed70a5b91b231067e8e69951ef7010406e6b22cf 100644 --- a/tensorflow/contrib/session_bundle/signature.cc +++ b/tensorflow/contrib/session_bundle/signature.cc @@ -38,9 +38,9 @@ namespace { Status BatchSizesMatch(const Tensor& input, const Tensor& output) { // Ensure the number of outputs match the number of inputs. if (input.dim_size(0) != output.dim_size(0)) { - return errors::Internal( - strings::StrCat("Input batch size did not match output batch size: ", - input.dim_size(0), " vs. ", output.dim_size(0))); + return errors::Internal(strings::StrCat( + "Input batch size did not match output batch size: ", input.dim_size(0), + " vs. ", output.dim_size(0))); } return Status::OK(); } @@ -100,8 +100,8 @@ Status GetNamedClassificationSignature( const auto& it = signatures.named_signatures().find(name); if (it == signatures.named_signatures().end()) { return errors::NotFound( - strings::StrCat("Missing signature named \"", name, "\" in: ", - DebugStringIfAvailable(signatures))); + strings::StrCat("Missing signature named \"", name, + "\" in: ", DebugStringIfAvailable(signatures))); } if (!it->second.has_classification_signature()) { return errors::FailedPrecondition( @@ -232,8 +232,8 @@ Status GetNamedSignature(const string& name, const auto& it = signatures.named_signatures().find(name); if (it == signatures.named_signatures().end()) { return errors::NotFound( - strings::StrCat("Missing signature named \"", name, "\" in: ", - DebugStringIfAvailable(signatures))); + strings::StrCat("Missing signature named \"", name, + "\" in: ", DebugStringIfAvailable(signatures))); } *signature = it->second; return Status::OK(); diff --git a/tensorflow/contrib/session_bundle/signature.h b/tensorflow/contrib/session_bundle/signature.h index 0049bea00822db85c606b9e6d00ae4db83804bab..4ef1277cec413a6fcfb54721520279d024f18bc1 100644 --- a/tensorflow/contrib/session_bundle/signature.h +++ b/tensorflow/contrib/session_bundle/signature.h @@ -15,8 +15,8 @@ limitations under the License. // Helpers for working with TensorFlow exports and their signatures. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_SIGNATURE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_SIGNATURE_H_ +#ifndef TENSORFLOW_CONTRIB_SESSION_BUNDLE_SIGNATURE_H_ +#define TENSORFLOW_CONTRIB_SESSION_BUNDLE_SIGNATURE_H_ #include #include @@ -121,4 +121,4 @@ Status BindGenericNames(const GenericSignature& signature, } // namespace serving } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_SIGNATURE_H_ +#endif // TENSORFLOW_CONTRIB_SESSION_BUNDLE_SIGNATURE_H_ diff --git a/tensorflow/contrib/session_bundle/test_util.h b/tensorflow/contrib/session_bundle/test_util.h index dd0fc8d1c0c47c444ac6fe807435fa671f3939f0..f0d41ce5a4b901db80a7a01475dc0917e966dc89 100644 --- a/tensorflow/contrib/session_bundle/test_util.h +++ b/tensorflow/contrib/session_bundle/test_util.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_TEST_UTIL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_TEST_UTIL_H_ +#ifndef TENSORFLOW_CONTRIB_SESSION_BUNDLE_TEST_UTIL_H_ +#define TENSORFLOW_CONTRIB_SESSION_BUNDLE_TEST_UTIL_H_ #include @@ -35,4 +35,4 @@ string TestSrcDirPath(const string& relative_path); } // namespace serving } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_SESSION_BUNDLE_TEST_UTIL_H_ +#endif // TENSORFLOW_CONTRIB_SESSION_BUNDLE_TEST_UTIL_H_ diff --git a/tensorflow/contrib/signal/python/kernel_tests/mel_ops_test.py b/tensorflow/contrib/signal/python/kernel_tests/mel_ops_test.py index b861476b67fc360f383465145ccd1cc620de5a99..35c4b5bec172858b39dd4628a37e164efe87bdbf 100644 --- a/tensorflow/contrib/signal/python/kernel_tests/mel_ops_test.py +++ b/tensorflow/contrib/signal/python/kernel_tests/mel_ops_test.py @@ -158,6 +158,9 @@ class LinearToMelTest(test.TestCase): with self.assertRaises(ValueError): mel_ops.linear_to_mel_weight_matrix(lower_edge_hertz=100, upper_edge_hertz=10) + with self.assertRaises(ValueError): + mel_ops.linear_to_mel_weight_matrix(upper_edge_hertz=1000, + sample_rate=800) with self.assertRaises(ValueError): mel_ops.linear_to_mel_weight_matrix(dtype=dtypes.int32) diff --git a/tensorflow/contrib/signal/python/kernel_tests/mfcc_ops_test.py b/tensorflow/contrib/signal/python/kernel_tests/mfcc_ops_test.py index c04f1cf5bad358a14a1827df05a129339502c86f..e7743bdcba180929007d17bdf3b143c64643aacc 100644 --- a/tensorflow/contrib/signal/python/kernel_tests/mfcc_ops_test.py +++ b/tensorflow/contrib/signal/python/kernel_tests/mfcc_ops_test.py @@ -20,6 +20,7 @@ from __future__ import print_function from tensorflow.contrib.signal.python.ops import mfcc_ops from tensorflow.python.framework import dtypes +from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import spectral_ops_test_util @@ -49,6 +50,14 @@ class MFCCTest(test.TestCase): signal = random_ops.random_normal((2, 3, 5)) mfcc_ops.mfccs_from_log_mel_spectrograms(signal).eval() + def test_unknown_shape(self): + """A test that the op runs when shape and rank are unknown.""" + with spectral_ops_test_util.fft_kernel_label_map(): + with self.test_session(use_gpu=True): + signal = array_ops.placeholder_with_default( + random_ops.random_normal((2, 3, 5)), tensor_shape.TensorShape(None)) + self.assertIsNone(signal.shape.ndims) + mfcc_ops.mfccs_from_log_mel_spectrograms(signal).eval() if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/signal/python/ops/mel_ops.py b/tensorflow/contrib/signal/python/ops/mel_ops.py index 2ad07027aa73f1ae7f44fd684a18fc67400a4f90..d1a36548d95cf44d2bf7e6108141aeb00853db04 100644 --- a/tensorflow/contrib/signal/python/ops/mel_ops.py +++ b/tensorflow/contrib/signal/python/ops/mel_ops.py @@ -80,6 +80,10 @@ def _validate_arguments(num_mel_bins, num_spectrogram_bins, sample_rate, if lower_edge_hertz >= upper_edge_hertz: raise ValueError('lower_edge_hertz %.1f >= upper_edge_hertz %.1f' % (lower_edge_hertz, upper_edge_hertz)) + if upper_edge_hertz > sample_rate / 2: + raise ValueError('upper_edge_hertz must not be larger than the Nyquist ' + 'frequency (sample_rate / 2). Got: %s for sample_rate: %s' + % (upper_edge_hertz, sample_rate)) if not dtype.is_floating: raise ValueError('dtype must be a floating point type. Got: %s' % dtype) @@ -138,8 +142,8 @@ def linear_to_mel_weight_matrix(num_mel_bins=20, Raises: ValueError: If num_mel_bins/num_spectrogram_bins/sample_rate are not - positive, lower_edge_hertz is negative, or frequency edges are incorrectly - ordered. + positive, lower_edge_hertz is negative, frequency edges are incorrectly + ordered, or upper_edge_hertz is larger than the Nyquist frequency. [mel]: https://en.wikipedia.org/wiki/Mel_scale """ diff --git a/tensorflow/contrib/signal/python/ops/mfcc_ops.py b/tensorflow/contrib/signal/python/ops/mfcc_ops.py index 7bc7b57cd4f1033a8bda0845ccd8e777e0213d6b..4e842f7f10ae07448cc07e5f636ae80a820e656f 100644 --- a/tensorflow/contrib/signal/python/ops/mfcc_ops.py +++ b/tensorflow/contrib/signal/python/ops/mfcc_ops.py @@ -50,7 +50,7 @@ def mfccs_from_log_mel_spectrograms(log_mel_spectrograms, name=None): # A 1024-point STFT with frames of 64 ms and 75% overlap. stfts = tf.contrib.signal.stft(pcm, frame_length=1024, frame_step=256, fft_length=1024) - spectrograms = tf.abs(stft) + spectrograms = tf.abs(stfts) # Warp the linear scale spectrograms into the mel-scale. num_spectrogram_bins = stfts.shape[-1].value @@ -105,4 +105,4 @@ def mfccs_from_log_mel_spectrograms(log_mel_spectrograms, name=None): num_mel_bins = array_ops.shape(log_mel_spectrograms)[-1] dct2 = spectral_ops.dct(log_mel_spectrograms) - return dct2 * math_ops.rsqrt(num_mel_bins * 2.0) + return dct2 * math_ops.rsqrt(math_ops.to_float(num_mel_bins) * 2.0) diff --git a/tensorflow/contrib/signal/python/ops/spectral_ops.py b/tensorflow/contrib/signal/python/ops/spectral_ops.py index bca2e01d7bbefb18fd69a0eba27e3afb8f636724..a8b5deff6ca3a4a756d31b904e577f08f6155fd7 100644 --- a/tensorflow/contrib/signal/python/ops/spectral_ops.py +++ b/tensorflow/contrib/signal/python/ops/spectral_ops.py @@ -144,7 +144,7 @@ def inverse_stft_window_fn(frame_step, overlaps = -(-frame_length // frame_step) # Ceiling division. denom = array_ops.pad(denom, [(0, overlaps * frame_step - frame_length)]) denom = array_ops.reshape(denom, [overlaps, frame_step]) - denom = math_ops.reduce_sum(denom, 0, keep_dims=True) + denom = math_ops.reduce_sum(denom, 0, keepdims=True) denom = array_ops.tile(denom, [overlaps, 1]) denom = array_ops.reshape(denom, [overlaps * frame_step]) diff --git a/tensorflow/contrib/slim/README.md b/tensorflow/contrib/slim/README.md index dc92ae0c859394f44ba83d814adbef7d324a9ada..2d9df8f27ee98431f51fd39c168325b8f625dce9 100644 --- a/tensorflow/contrib/slim/README.md +++ b/tensorflow/contrib/slim/README.md @@ -145,7 +145,7 @@ regular_variables_and_model_variables = slim.get_variables() How does this work? When you create a model variable via TF-Slim's layers or directly via the `slim.model_variable` function, TF-Slim adds the variable to -a the `tf.GraphKeys.MODEL_VARIABLES` collection. What if you have your own +the `tf.GraphKeys.MODEL_VARIABLES` collection. What if you have your own custom layers or variable creation routine but still want TF-Slim to manage or be aware of your model variables? TF-Slim provides a convenience function for adding the model variable to its collection: @@ -676,7 +676,7 @@ file were implicitly obtained from each provided variable's `var.op.name`. This works well when the variable names in the checkpoint file match those in the graph. However, sometimes, we want to restore a model from a checkpoint -whose variables have different names those in the current graph. In this case, +whose variables have different names to those in the current graph. In this case, we must provide the `Saver` a dictionary that maps from each checkpoint variable name to each graph variable. Consider the following example where the checkpoint variables names are obtained via a simple function: diff --git a/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py b/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py index 0544404e9e252cca6d3650b805b91be25d705eea..b3b61e1dfe5671a7fbbee20b0c577ee5fad0fb9b 100644 --- a/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py +++ b/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py @@ -349,7 +349,8 @@ class Image(ItemHandler): shape=None, channels=3, dtype=dtypes.uint8, - repeated=False): + repeated=False, + dct_method=''): """Initializes the image. Args: @@ -368,6 +369,11 @@ class Image(ItemHandler): tf.decode_raw, repeated: if False, decodes a single image. If True, decodes a variable number of image strings from a 1D tensor of strings. + dct_method: An optional string. Defaults to empty string. It only takes + effect when image format is jpeg, used to specify a hint about the + algorithm used for jpeg decompression. Currently valid values + are ['INTEGER_FAST', 'INTEGER_ACCURATE']. The hint may be ignored, for + example, the jpeg library does not have that specific option. """ if not image_key: image_key = 'image/encoded' @@ -381,6 +387,7 @@ class Image(ItemHandler): self._channels = channels self._dtype = dtype self._repeated = repeated + self._dct_method = dct_method def tensors_to_item(self, keys_to_tensors): """See base class.""" @@ -406,9 +413,25 @@ class Image(ItemHandler): A tensor that represents decoded image of self._shape, or (?, ?, self._channels) if self._shape is not specified. """ + def decode_image(): - """Decodes a png or jpg based on the headers.""" - return image_ops.decode_image(image_buffer, self._channels) + """Decodes a image based on the headers.""" + return image_ops.decode_image(image_buffer, channels=self._channels) + + def decode_jpeg(): + """Decodes a jpeg image with specified '_dct_method'.""" + return image_ops.decode_jpeg( + image_buffer, channels=self._channels, dct_method=self._dct_method) + + def check_jpeg(): + """Checks if an image is jpeg.""" + # For jpeg, we directly use image_ops.decode_jpeg rather than decode_image + # in order to feed the jpeg specify parameter 'dct_method'. + return control_flow_ops.cond( + image_ops.is_jpeg(image_buffer), + decode_jpeg, + decode_image, + name='cond_jpeg') def decode_raw(): """Decodes a raw image.""" @@ -420,7 +443,7 @@ class Image(ItemHandler): math_ops.equal(image_format, 'RAW')): decode_raw, } image = control_flow_ops.case( - pred_fn_pairs, default=decode_image, exclusive=True) + pred_fn_pairs, default=check_jpeg, exclusive=True) image.set_shape([None, None, self._channels]) if self._shape is not None: diff --git a/tensorflow/contrib/slim/python/slim/evaluation_test.py b/tensorflow/contrib/slim/python/slim/evaluation_test.py index 870f504d10362ed5226951adefc3ba9a934900c1..c24bd048512daaae116e732ac437f7c9b6f6d7fc 100644 --- a/tensorflow/contrib/slim/python/slim/evaluation_test.py +++ b/tensorflow/contrib/slim/python/slim/evaluation_test.py @@ -42,7 +42,7 @@ from tensorflow.python.platform import flags from tensorflow.python.platform import gfile from tensorflow.python.platform import test from tensorflow.python.summary import summary_iterator -from tensorflow.python.training import input +from tensorflow.python.training import input # pylint: disable=redefined-builtin from tensorflow.python.training import saver as saver_lib from tensorflow.python.training import session_run_hook diff --git a/tensorflow/contrib/slim/python/slim/learning.py b/tensorflow/contrib/slim/python/slim/learning.py index def00b76184ba4e1fc630cd83d8e055448100562..6a200de1ea172b4ccb38c0f5d889566ccaeef893 100644 --- a/tensorflow/contrib/slim/python/slim/learning.py +++ b/tensorflow/contrib/slim/python/slim/learning.py @@ -738,6 +738,7 @@ def train(train_op, if summary_writer is not None: train_step_kwargs['summary_writer'] = sv.summary_writer + total_loss = None should_retry = True while should_retry: try: @@ -753,9 +754,10 @@ def train(train_op, if logdir: sv.start_standard_services(sess) elif startup_delay_steps > 0: + # (use sys.maxsize because sys.maxint doesn't exist in Python 3) _wait_for_step(sess, global_step, min(startup_delay_steps, number_of_steps or - sys.maxint)) + sys.maxsize)) threads = sv.start_queue_runners(sess) logging.info('Starting Queues.') if is_chief and sync_optimizer is not None: @@ -769,10 +771,10 @@ def train(train_op, logging.info('Stopping Training.') sv.request_stop() break - except errors.OutOfRangeError: + except errors.OutOfRangeError as e: # OutOfRangeError is thrown when epoch limit per # tf.train.limit_epochs is reached. - logging.info('Caught OutOfRangeError. Stopping Training.') + logging.info('Caught OutOfRangeError. Stopping Training. %s', e) if logdir and sv.is_chief: logging.info('Finished training! Saving model to disk.') sv.saver.save(sess, sv.save_path, global_step=sv.global_step) diff --git a/tensorflow/contrib/slim/python/slim/learning_test.py b/tensorflow/contrib/slim/python/slim/learning_test.py index 4e816f9b11be2986d042f336bdc320ff47d8cc49..831c6e427ae78932bec09cea935f05a87723f1a3 100644 --- a/tensorflow/contrib/slim/python/slim/learning_test.py +++ b/tensorflow/contrib/slim/python/slim/learning_test.py @@ -197,9 +197,7 @@ class MultiplyGradientsTest(test.TestCase): gradient = constant_op.constant(self._grad_vec, dtype=dtypes.float32) variable = variables_lib.Variable(array_ops.zeros_like(gradient)) multiplier_flag = variables_lib.Variable(True) - tensor_multiplier = array_ops.where(multiplier_flag, - self._multiplier, - 1.0) + tensor_multiplier = array_ops.where(multiplier_flag, self._multiplier, 1.0) grad_to_var = (gradient, variable) gradient_multipliers = {variable: tensor_multiplier} @@ -212,11 +210,8 @@ class MultiplyGradientsTest(test.TestCase): sess.run(multiplier_flag.assign(False)) gradient_false_flag = sess.run(grad_to_var[0]) np_testing.assert_almost_equal(gradient_true_flag, - self._multiplied_grad_vec, - 5) - np_testing.assert_almost_equal(gradient_false_flag, - self._grad_vec, - 5) + self._multiplied_grad_vec, 5) + np_testing.assert_almost_equal(gradient_false_flag, self._grad_vec, 5) def LogisticClassifier(inputs): @@ -502,6 +497,7 @@ class TrainTest(test.TestCase): purpose. """ dump_root = tempfile.mkdtemp() + def dumping_wrapper(sess): # pylint: disable=invalid-name return dumping_wrapper_lib.DumpingDebugWrapperSession(sess, dump_root) @@ -519,16 +515,13 @@ class TrainTest(test.TestCase): train_op = learning.create_train_op(total_loss, optimizer) loss = learning.train( - train_op, - None, - number_of_steps=1, - session_wrapper=dumping_wrapper) + train_op, None, number_of_steps=1, session_wrapper=dumping_wrapper) self.assertIsNotNone(loss) run_root = glob.glob(os.path.join(dump_root, 'run_*'))[-1] dump = debug_data.DebugDumpDir(run_root) - self.assertAllEqual( - 0, dump.get_tensors('global_step', 0, 'DebugIdentity')[0]) + self.assertAllEqual(0, + dump.get_tensors('global_step', 0, 'DebugIdentity')[0]) def testTrainWithTrace(self): logdir = os.path.join( @@ -961,8 +954,8 @@ class TrainTest(test.TestCase): self.assertGreater(losses[0], losses[1]) def testTrainWithEpochLimit(self): - logdir = os.path.join(tempfile.mkdtemp(prefix=self.get_temp_dir()), - 'tmp_logs') + logdir = os.path.join( + tempfile.mkdtemp(prefix=self.get_temp_dir()), 'tmp_logs') with ops.Graph().as_default(): random_seed.set_random_seed(0) tf_inputs = constant_op.constant(self._inputs, dtype=dtypes.float32) @@ -982,7 +975,8 @@ class TrainTest(test.TestCase): self.assertIsNotNone(loss) self.assertLess(loss, .015) self.assertTrue(os.path.isfile('{}/model.ckpt-300.index'.format(logdir))) - self.assertTrue(os.path.isfile('{}/model.ckpt-300.data-00000-of-00001'.format(logdir))) + self.assertTrue( + os.path.isfile('{}/model.ckpt-300.data-00000-of-00001'.format(logdir))) if __name__ == '__main__': diff --git a/tensorflow/contrib/slim/python/slim/nets/inception_v3.py b/tensorflow/contrib/slim/python/slim/nets/inception_v3.py index 432e1f79f12ed5536887738339de845f6343e5ba..afe261e43a9f144992318086f958e21d50286d11 100644 --- a/tensorflow/contrib/slim/python/slim/nets/inception_v3.py +++ b/tensorflow/contrib/slim/python/slim/nets/inception_v3.py @@ -548,7 +548,10 @@ def inception_v3(inputs, parameters or computation cost of the model. prediction_fn: a function to get predictions out of logits. spatial_squeeze: if True, logits is of shape is [B, C], if false logits is - of shape [B, 1, 1, C], where B is batch_size and C is number of classes. + of shape [B, 1, 1, C], where B is batch_size and C is number of classes. + To use this parameter, the input images must be smaller + than 300x300 pixels, in which case the output logit layer + does not contain spatial information and can be removed. reuse: whether or not the network and its variables should be reused. To be able to reuse 'scope' must be given. scope: Optional variable_scope. @@ -677,6 +680,9 @@ def _reduced_kernel_size_for_small_input(input_tensor, kernel_size): def inception_v3_arg_scope(weight_decay=0.00004, batch_norm_var_collection='moving_vars', + batch_norm_decay=0.9997, + batch_norm_epsilon=0.001, + updates_collections=ops.GraphKeys.UPDATE_OPS, use_fused_batchnorm=True): """Defines the default InceptionV3 arg scope. @@ -684,6 +690,9 @@ def inception_v3_arg_scope(weight_decay=0.00004, weight_decay: The weight decay to use for regularizing the model. batch_norm_var_collection: The name of the collection for the batch norm variables. + batch_norm_decay: Decay for batch norm moving average + batch_norm_epsilon: Small float added to variance to avoid division by zero + updates_collections: Collections for the update ops of the layer use_fused_batchnorm: Enable fused batchnorm. Returns: @@ -691,11 +700,11 @@ def inception_v3_arg_scope(weight_decay=0.00004, """ batch_norm_params = { # Decay for the moving averages. - 'decay': 0.9997, + 'decay': batch_norm_decay, # epsilon to prevent 0s in variance. - 'epsilon': 0.001, + 'epsilon': batch_norm_epsilon, # collection containing update_ops. - 'updates_collections': ops.GraphKeys.UPDATE_OPS, + 'updates_collections': updates_collections, # Use fused batch norm if possible. 'fused': use_fused_batchnorm, # collection containing the moving mean and moving variance. diff --git a/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py b/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py index 930df2414bc907703c2670ffd92134727a28e856..a1282847bef981717d7fdb1474adbbaaae4621c0 100644 --- a/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py +++ b/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py @@ -45,32 +45,67 @@ def _get_linear_equations_tests(dtype_, use_static_shape_, shape_): low=-1.0, high=1.0, size=np.prod(shape_)).reshape(shape_).astype(dtype_) # Make a selfadjoint, positive definite. a_np = np.dot(a_np.T, a_np) + # jacobi preconditioner + jacobi_np = np.zeros_like(a_np) + jacobi_np[range(a_np.shape[0]), range(a_np.shape[1])] = ( + 1.0 / a_np.diagonal()) rhs_np = np.random.uniform( low=-1.0, high=1.0, size=shape_[0]).astype(dtype_) + x_np = np.zeros_like(rhs_np) tol = 1e-6 if dtype_ == np.float64 else 1e-3 max_iter = 20 with self.test_session() as sess: if use_static_shape_: a = constant_op.constant(a_np) rhs = constant_op.constant(rhs_np) + x = constant_op.constant(x_np) + jacobi = constant_op.constant(jacobi_np) else: a = array_ops.placeholder(dtype_) rhs = array_ops.placeholder(dtype_) + x = array_ops.placeholder(dtype_) + jacobi = array_ops.placeholder(dtype_) operator = util.create_operator(a) - cg_graph = linear_equations.conjugate_gradient( - operator, rhs, tol=tol, max_iter=max_iter) - if use_static_shape_: - cg_val = sess.run(cg_graph) - else: - cg_val = sess.run(cg_graph, feed_dict={a: a_np, rhs: rhs_np}) - norm_r0 = np.linalg.norm(rhs_np) - norm_r = np.sqrt(cg_val.gamma) - self.assertLessEqual(norm_r, tol * norm_r0) - # Validate that we get an equally small residual norm with numpy - # using the computed solution. - r_np = rhs_np - np.dot(a_np, cg_val.x) - norm_r_np = np.linalg.norm(r_np) - self.assertLessEqual(norm_r_np, tol * norm_r0) + preconditioners = [ + None, util.identity_operator(a), + util.create_operator(jacobi) + ] + cg_results = [] + for preconditioner in preconditioners: + cg_graph = linear_equations.conjugate_gradient( + operator, + rhs, + preconditioner=preconditioner, + x=x, + tol=tol, + max_iter=max_iter) + if use_static_shape_: + cg_val = sess.run(cg_graph) + else: + cg_val = sess.run( + cg_graph, + feed_dict={ + a: a_np, + rhs: rhs_np, + x: x_np, + jacobi: jacobi_np + }) + norm_r0 = np.linalg.norm(rhs_np) + norm_r = np.linalg.norm(cg_val.r) + self.assertLessEqual(norm_r, tol * norm_r0) + # Validate that we get an equally small residual norm with numpy + # using the computed solution. + r_np = rhs_np - np.dot(a_np, cg_val.x) + norm_r_np = np.linalg.norm(r_np) + self.assertLessEqual(norm_r_np, tol * norm_r0) + cg_results.append(cg_val) + # Validate that we get same results using identity_preconditioner + # and None + self.assertEqual(cg_results[0].i, cg_results[1].i) + self.assertAlmostEqual(cg_results[0].gamma, cg_results[1].gamma) + self.assertAllClose(cg_results[0].r, cg_results[1].r, rtol=tol) + self.assertAllClose(cg_results[0].x, cg_results[1].x, rtol=tol) + self.assertAllClose(cg_results[0].p, cg_results[1].p, rtol=tol) return [test_conjugate_gradient] diff --git a/tensorflow/contrib/solvers/python/kernel_tests/util_test.py b/tensorflow/contrib/solvers/python/kernel_tests/util_test.py index 1566984b27fdab4c2a8c91bd16f587747e69e9e5..5d7534657bff27f7169e6a97bf4b03d4f6a35bc9 100644 --- a/tensorflow/contrib/solvers/python/kernel_tests/util_test.py +++ b/tensorflow/contrib/solvers/python/kernel_tests/util_test.py @@ -63,6 +63,43 @@ class UtilTest(test.TestCase): def testCreateOperatorUnknownShape(self): self._testCreateOperator(False) + def _testIdentityOperator(self, use_static_shape_): + for dtype in np.float32, np.float64: + a_np = np.array([[1., 2.], [3., 4.], [5., 6.]], dtype=dtype) + x_np = np.array([[2.], [-3.]], dtype=dtype) + y_np = np.array([[2], [-3.], [5.]], dtype=dtype) + with self.test_session() as sess: + if use_static_shape_: + a = constant_op.constant(a_np, dtype=dtype) + x = constant_op.constant(x_np, dtype=dtype) + y = constant_op.constant(y_np, dtype=dtype) + else: + a = array_ops.placeholder(dtype) + x = array_ops.placeholder(dtype) + y = array_ops.placeholder(dtype) + id_op = util.identity_operator(a) + ax = id_op.apply(x) + aty = id_op.apply_adjoint(y) + op_shape = ops.convert_to_tensor(id_op.shape) + if use_static_shape_: + op_shape_val, ax_val, aty_val = sess.run([op_shape, ax, aty]) + else: + op_shape_val, ax_val, aty_val = sess.run( + [op_shape, ax, aty], feed_dict={ + a: a_np, + x: x_np, + y: y_np + }) + self.assertAllEqual(op_shape_val, [3, 2]) + self.assertAllClose(ax_val, x_np) + self.assertAllClose(aty_val, y_np) + + def testIdentityOperator(self): + self._testIdentityOperator(True) + + def testIdentityOperatorUnknownShape(self): + self._testIdentityOperator(False) + def testL2Norm(self): with self.test_session(): x_np = np.array([[2], [-3.], [5.]]) diff --git a/tensorflow/contrib/solvers/python/ops/linear_equations.py b/tensorflow/contrib/solvers/python/ops/linear_equations.py index 8cba56eba6b3046b8efbbbbf130705255e1c13bb..d791d467639b572e7831c1d1a582aa15585649b6 100644 --- a/tensorflow/contrib/solvers/python/ops/linear_equations.py +++ b/tensorflow/contrib/solvers/python/ops/linear_equations.py @@ -26,11 +26,15 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops +from tensorflow.python.ops import linalg_ops def conjugate_gradient(operator, rhs, + preconditioner=None, + x=None, tol=1e-4, max_iter=20, name="conjugate_gradient"): @@ -55,6 +59,15 @@ def conjugate_gradient(operator, vector with the result of applying the operator to `x`, i.e. if `operator` represents matrix `A`, `apply` should return `A * x`. rhs: A rank-1 `Tensor` of shape `[N]` containing the right-hand size vector. + preconditioner: An object representing a linear operator, see `operator` + for detail. The preconditioner should approximate the inverse of `A`. + An efficient preconditioner could dramatically improve the rate of + convergence. If `preconditioner` represents matrix `M`(`M` approximates + `A^{-1}`), the algorithm uses `preconditioner.apply(x)` to estimate + `A^{-1}x`. For this to be useful, the cost of applying `M` should be + much lower than computing `A^{-1}` directly. + x: A rank-1 `Tensor` of shape `[N]` containing the initial guess for the + solution. tol: A float scalar convergence tolerance. max_iter: An integer giving the maximum number of iterations. name: A name scope for the operation. @@ -65,35 +78,49 @@ def conjugate_gradient(operator, - x: A rank-1 `Tensor` of shape `[N]` containing the computed solution. - r: A rank-1 `Tensor` of shape `[M]` containing the residual vector. - p: A rank-1 `Tensor` of shape `[N]`. `A`-conjugate basis vector. - - gamma: \\(||r||_2^2\\) + - gamma: \\(r \dot M \dot r\\), equivalent to \\(||r||_2^2\\) when + `preconditioner=None`. """ # ephemeral class holding CG state. cg_state = collections.namedtuple("CGState", ["i", "x", "r", "p", "gamma"]) def stopping_criterion(i, state): - return math_ops.logical_and(i < max_iter, state.gamma > tol) + return math_ops.logical_and(i < max_iter, linalg_ops.norm(state.r) > tol) - # TODO(rmlarsen): add preconditioning - def cg_step(i, state): + def cg_step(i, state): # pylint: disable=missing-docstring z = operator.apply(state.p) alpha = state.gamma / util.dot(state.p, z) x = state.x + alpha * state.p r = state.r - alpha * z - gamma = util.l2norm_squared(r) - beta = gamma / state.gamma - p = r + beta * state.p + if preconditioner is None: + gamma = util.dot(r, r) + beta = gamma / state.gamma + p = r + beta * state.p + else: + q = preconditioner.apply(r) + gamma = util.dot(r, q) + beta = gamma / state.gamma + p = q + beta * state.p return i + 1, cg_state(i + 1, x, r, p, gamma) with ops.name_scope(name): n = operator.shape[1:] rhs = array_ops.expand_dims(rhs, -1) - gamma0 = util.l2norm_squared(rhs) - tol = tol * tol * gamma0 - x = array_ops.expand_dims( - array_ops.zeros( - n, dtype=rhs.dtype.base_dtype), -1) + if x is None: + x = array_ops.expand_dims( + array_ops.zeros(n, dtype=rhs.dtype.base_dtype), -1) + r0 = rhs + else: + x = array_ops.expand_dims(x, -1) + r0 = rhs - operator.apply(x) + if preconditioner is None: + p0 = r0 + else: + p0 = preconditioner.apply(r0) + gamma0 = util.dot(r0, p0) + tol *= linalg_ops.norm(r0) i = constant_op.constant(0, dtype=dtypes.int32) - state = cg_state(i=i, x=x, r=rhs, p=rhs, gamma=gamma0) + state = cg_state(i=i, x=x, r=r0, p=p0, gamma=gamma0) _, state = control_flow_ops.while_loop(stopping_criterion, cg_step, [i, state]) return cg_state( diff --git a/tensorflow/contrib/solvers/python/ops/util.py b/tensorflow/contrib/solvers/python/ops/util.py index 777e0c185d6c9fffab6a7fe6e6ae4c133c62ad1a..96947e8eea1006bcd03cf09cd13cd1266695cc2e 100644 --- a/tensorflow/contrib/solvers/python/ops/util.py +++ b/tensorflow/contrib/solvers/python/ops/util.py @@ -45,6 +45,23 @@ def create_operator(matrix): apply_adjoint=lambda v: math_ops.matmul(matrix, v, adjoint_a=True)) +def identity_operator(matrix): + """Creates a linear operator from a rank-2 identity tensor.""" + + linear_operator = collections.namedtuple( + "LinearOperator", ["shape", "dtype", "apply", "apply_adjoint"]) + shape = matrix.get_shape() + if shape.is_fully_defined(): + shape = shape.as_list() + else: + shape = array_ops.shape(matrix) + return linear_operator( + shape=shape, + dtype=matrix.dtype, + apply=lambda v: v, + apply_adjoint=lambda v: v) + + # TODO(rmlarsen): Measure if we should just call matmul. def dot(x, y): return math_ops.reduce_sum(math_ops.conj(x) * y) diff --git a/tensorflow/contrib/sparsemax/python/kernel_tests/sparsemax_loss_test.py b/tensorflow/contrib/sparsemax/python/kernel_tests/sparsemax_loss_test.py index c8b4e472c99e0bf081a7222a7976b1fbbb680825..360e7dbe75f595ff61fb83379089294371203813 100644 --- a/tensorflow/contrib/sparsemax/python/kernel_tests/sparsemax_loss_test.py +++ b/tensorflow/contrib/sparsemax/python/kernel_tests/sparsemax_loss_test.py @@ -105,8 +105,8 @@ class SparsemaxLossTest(test.TestCase): tf_loss_op, tf_loss_out = self._tf_sparsemax_loss(z, q, dtype, use_gpu) np_loss = self._np_sparsemax_loss(z, q).astype(dtype) - self.assertAllCloseAccordingToType(np_loss, tf_loss_out, - half_atol=1e-2, half_rtol=5e-3) + self.assertAllCloseAccordingToType( + np_loss, tf_loss_out, half_atol=1e-2, half_rtol=5e-3) self.assertShapeEqual(np_loss, tf_loss_op) def _test_constant_add(self, dtype, random, use_gpu): @@ -116,17 +116,17 @@ class SparsemaxLossTest(test.TestCase): q = np.zeros((test_obs, 10)) q[np.arange(0, test_obs), np.random.randint(0, 10, size=test_obs)] = 1 - _, tf_loss_zpc = self._tf_sparsemax_loss( - z + c, q, dtype, use_gpu - ) + _, tf_loss_zpc = self._tf_sparsemax_loss(z + c, q, dtype, use_gpu) - _, tf_loss_z = self._tf_sparsemax_loss( - z, q, dtype, use_gpu - ) + _, tf_loss_z = self._tf_sparsemax_loss(z, q, dtype, use_gpu) - self.assertAllCloseAccordingToType(tf_loss_zpc, tf_loss_z, - float_atol=5e-6, float_rtol=5e-6, - half_atol=1e-2, half_rtol=1e-2) + self.assertAllCloseAccordingToType( + tf_loss_zpc, + tf_loss_z, + float_atol=5e-6, + float_rtol=5e-6, + half_atol=1e-2, + half_rtol=1e-2) def _test_sparsemax_loss_positive(self, dtype, random, use_gpu): """check sparsemax-loss proposition 4""" @@ -170,10 +170,7 @@ class SparsemaxLossTest(test.TestCase): with self.test_session(use_gpu=use_gpu): err = gradient_checker.compute_gradient_error( - logits, z.shape, - loss_op, (test_obs, ), - x_init_value=z, delta=1e-9 - ) + logits, z.shape, loss_op, (test_obs,), x_init_value=z, delta=1e-9) self.assertLess(err, 1e-4) @@ -192,8 +189,8 @@ class SparsemaxLossTest(test.TestCase): tf_grad = loss_grad_op.eval() np_grad = self._np_sparsemax_loss_grad(z, q).astype(dtype) - self.assertAllCloseAccordingToType(np_grad, tf_grad, - half_atol=1e-2, half_rtol=5e-3) + self.assertAllCloseAccordingToType( + np_grad, tf_grad, half_atol=1e-2, half_rtol=5e-3) self.assertShapeEqual(np_grad, loss_grad_op) def _test_dtype(self, dtype): @@ -220,5 +217,6 @@ class SparsemaxLossTest(test.TestCase): def testDouble(self): self._test_dtype('float64') -if __name__ == "__main__": + +if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/sparsemax/python/kernel_tests/sparsemax_test.py b/tensorflow/contrib/sparsemax/python/kernel_tests/sparsemax_test.py index 82d36ee9cb21fb822e6df0c3632c49a4fd616825..259e62bd864fba3cc7d9aa387e02c8319438d658 100644 --- a/tensorflow/contrib/sparsemax/python/kernel_tests/sparsemax_test.py +++ b/tensorflow/contrib/sparsemax/python/kernel_tests/sparsemax_test.py @@ -83,8 +83,8 @@ class SparsemaxTest(test.TestCase): tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype, use_gpu) p_sparemax = self._np_sparsemax(z).astype(dtype) - self.assertAllCloseAccordingToType(p_sparemax, tf_sparsemax_out, - half_atol=5e-3) + self.assertAllCloseAccordingToType( + p_sparemax, tf_sparsemax_out, half_atol=5e-3) self.assertShapeEqual(p_sparemax, tf_sparsemax_op) def _test_sparsemax_of_zero(self, dtype, random, use_gpu): @@ -111,9 +111,8 @@ class SparsemaxTest(test.TestCase): p_expected = np.zeros((test_obs, 10), dtype=dtype) p_expected[np.arange(0, test_obs), z_sort_arg[:, 0]] = 1 - tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax( - (1 / epsilon) * z, dtype, use_gpu - ) + tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax((1 / epsilon) * z, + dtype, use_gpu) self.assertAllCloseAccordingToType(p_expected, tf_sparsemax_out) self.assertShapeEqual(p_expected, tf_sparsemax_op) @@ -123,16 +122,12 @@ class SparsemaxTest(test.TestCase): z = random.uniform(low=-3, high=3, size=(test_obs, 10)).astype(dtype) c = random.uniform(low=-3, high=3, size=(test_obs, 1)).astype(dtype) - _, tf_sparsemax_zpc = self._tf_sparsemax( - z + c, dtype, use_gpu - ) + _, tf_sparsemax_zpc = self._tf_sparsemax(z + c, dtype, use_gpu) - _, tf_sparsemax_z = self._tf_sparsemax( - z, dtype, use_gpu - ) + _, tf_sparsemax_z = self._tf_sparsemax(z, dtype, use_gpu) - self.assertAllCloseAccordingToType(tf_sparsemax_zpc, tf_sparsemax_z, - half_atol=5e-3) + self.assertAllCloseAccordingToType( + tf_sparsemax_zpc, tf_sparsemax_z, half_atol=5e-3) def _test_permutation(self, dtype, random, use_gpu): """check sparsemax proposition 3""" @@ -143,12 +138,11 @@ class SparsemaxTest(test.TestCase): per = random.permutation(10) tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax( - z[i, per].reshape(1, -1), dtype, use_gpu - ) + z[i, per].reshape(1, -1), dtype, use_gpu) p_expected = p[i, per].reshape(1, -1) - self.assertAllCloseAccordingToType(p_expected, tf_sparsemax_out, - half_atol=5e-3) + self.assertAllCloseAccordingToType( + p_expected, tf_sparsemax_out, half_atol=5e-3) self.assertShapeEqual(p_expected, tf_sparsemax_op) def _test_diffrence(self, dtype, random, use_gpu): @@ -166,18 +160,14 @@ class SparsemaxTest(test.TestCase): continue self.assertTrue( - 0 <= p[val, j] - p[val, i] <= z[val, j] - z[val, i] + etol, - "0 <= %.10f <= %.10f" % ( - p[val, j] - p[val, i], z[val, j] - z[val, i] + etol - ) - ) + 0 <= p[val, j] - p[val, i] <= z[val, j] - z[val, i] + etol, + '0 <= %.10f <= %.10f' % (p[val, j] - p[val, i], + z[val, j] - z[val, i] + etol)) def _test_two_dimentional(self, dtype, random, use_gpu): """check two dimentation sparsemax case""" t = np.linspace(-2, 2, test_obs, dtype=dtype) - z = np.vstack([ - t, np.zeros(test_obs, dtype=dtype) - ]).T + z = np.vstack([t, np.zeros(test_obs, dtype=dtype)]).T tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype, use_gpu) @@ -196,10 +186,7 @@ class SparsemaxTest(test.TestCase): with self.test_session(use_gpu=use_gpu): err = gradient_checker.compute_gradient_error( - logits, z.shape, - sparsemax_op, z.shape, - x_init_value=z, delta=1e-9 - ) + logits, z.shape, sparsemax_op, z.shape, x_init_value=z, delta=1e-9) self.assertLess(err, 1e-4) @@ -248,5 +235,6 @@ class SparsemaxTest(test.TestCase): def testDouble(self): self._test_dtype('float64') -if __name__ == "__main__": + +if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/sparsemax/python/ops/sparsemax.py b/tensorflow/contrib/sparsemax/python/ops/sparsemax.py index 73a5cf1e9287ea4e4350d88165744cf12db954bb..890ca20f4cabd65146e803e54e554a5c97e72427 100644 --- a/tensorflow/contrib/sparsemax/python/ops/sparsemax.py +++ b/tensorflow/contrib/sparsemax/python/ops/sparsemax.py @@ -23,7 +23,6 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn -from tensorflow.python.platform import resource_loader __all__ = ["sparsemax"] diff --git a/tensorflow/contrib/sparsemax/python/ops/sparsemax_loss.py b/tensorflow/contrib/sparsemax/python/ops/sparsemax_loss.py index ba18f89e16c76a6ef3cb05df0c13f62eace6bbb1..582d1e6136df4d3ad3c8108ae9607d5fef519145 100644 --- a/tensorflow/contrib/sparsemax/python/ops/sparsemax_loss.py +++ b/tensorflow/contrib/sparsemax/python/ops/sparsemax_loss.py @@ -18,8 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.util import loader -from tensorflow.python.platform import resource_loader from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops diff --git a/tensorflow/contrib/specs/BUILD b/tensorflow/contrib/specs/BUILD index 4b688690aef513dd683817b0b5c2ba4cb50f73d9..084953a0a226cde46ebd9d2031d20cb839180ca8 100644 --- a/tensorflow/contrib/specs/BUILD +++ b/tensorflow/contrib/specs/BUILD @@ -23,7 +23,6 @@ py_library( srcs_version = "PY2AND3", deps = [ "//tensorflow/contrib/layers:layers_py", - "//tensorflow/contrib/ndlstm", "//tensorflow/python:array_ops", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:logging_ops", diff --git a/tensorflow/contrib/specs/README.md b/tensorflow/contrib/specs/README.md index b764e6e714ea907cd4474a07843bda300a8e4c8b..bcf34e601f1ffe3ab7a8c0d2ad573da4c8c977e9 100644 --- a/tensorflow/contrib/specs/README.md +++ b/tensorflow/contrib/specs/README.md @@ -59,17 +59,6 @@ Reshaping: - `Squeeze` = tf.squeeze - `Expand` = tf.expand_dims -Multidimensional LSTM: - -These are intended as alternatives to 2D convolutions. For sequence models, -there will be other modeling primitives. - - - `Lstm2` = Fun(lstm2d.separable_lstm) # 2D-to-2D - - `Lstm2to1` = Fun(lstm2d.reduce_to_sequence) # 2D-to-1D - - `Lstm2to0` = Fun(lstm2d.reduce_to_final) # 2D-to-vector - - `Clstm2(n, m)` is a `Cl(n, [3,3])` followed by `Lstm2(m)` - - `Dws(n)` is a depthwise convolution `Cs(n, [1, 1])` - Other: - `Id` = identity diff --git a/tensorflow/contrib/specs/python/__init__.py b/tensorflow/contrib/specs/python/__init__.py index 52db61e421a52f4106ab1e2a4d7ee5c100b6b4bc..b6cc754023859f8d3668545dd5c2fd1d1581ecf5 100644 --- a/tensorflow/contrib/specs/python/__init__.py +++ b/tensorflow/contrib/specs/python/__init__.py @@ -18,10 +18,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -# pylint: disable=wildcard-import,g-importing-member +# pylint: disable=wildcard-import,g-importing-member,redefined-builtin from tensorflow.contrib.specs.python.params_ops import * from tensorflow.contrib.specs.python.specs import * from tensorflow.contrib.specs.python.specs_lib import * from tensorflow.contrib.specs.python.specs_ops import * from tensorflow.contrib.specs.python.summaries import * -# pylint: enable=wildcard-import +# pylint: enable=wildcard-import,redefined-builtin diff --git a/tensorflow/contrib/specs/python/specs_ops.py b/tensorflow/contrib/specs/python/specs_ops.py index a6bd4d16c284a8b1a370005a7c55d3b74b4fbf95..49b989b8d0fc83a3793263a2b59a98a8fe292c6a 100644 --- a/tensorflow/contrib/specs/python/specs_ops.py +++ b/tensorflow/contrib/specs/python/specs_ops.py @@ -23,8 +23,6 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib.layers.python.layers import layers -from tensorflow.contrib.ndlstm.python import lstm1d -from tensorflow.contrib.ndlstm.python import lstm2d from tensorflow.contrib.specs.python import specs_lib from tensorflow.python.ops import array_ops from tensorflow.python.ops import logging_ops @@ -122,17 +120,6 @@ Sig = Fun(math_ops.sigmoid) Tanh = Fun(math_ops.tanh) Smax = Fun(nn_ops.softmax) -# 2D LSTM - -Lstm2 = Fun(lstm2d.separable_lstm) -Lstm2to1 = Fun(lstm2d.reduce_to_sequence) # 2D to 1D -Lstm2to0 = Fun(lstm2d.reduce_to_final) # 2D to depth-only - - -def Clstm2(n, *args, **kw): - """2D LSTM with 3x3 pre-convolution.""" - return Cl(n, [3, 3]) | Lstm2(*args, **kw) - def Dws(n): """Depth-wise convolution + sigmoid (used after LSTM).""" @@ -143,13 +130,6 @@ def Dwm(n): """Depth-wise convolution + softmax (used after LSTM).""" return Cm(n, [1, 1]) - -# 1D LSTM - -Lstm1 = Fun(lstm1d.ndlstm_base) -Lstm1to0 = Fun(lstm1d.sequence_to_final) # 1D to depth-only -Ssm = Fun(lstm1d.sequence_softmax) - # Sharing of Variables diff --git a/tensorflow/contrib/specs/python/specs_test.py b/tensorflow/contrib/specs/python/specs_test.py index 41782a9fc9ada3d8a1ff30847971aea18f0ca1c7..9a4ad36793542a83105ad0dc1ef7c0624a6c1f99 100644 --- a/tensorflow/contrib/specs/python/specs_test.py +++ b/tensorflow/contrib/specs/python/specs_test.py @@ -149,36 +149,6 @@ class SpecsTest(test.TestCase): self.assertEqual(tuple(result.shape), (10, 20)) self.assertEqual(summaries.tf_spec_structure(spec, inputs), "_ sig sig") - def testLstm2(self): - with self.test_session(): - inputs = constant_op.constant(_rand(1, 64, 64, 5)) - spec = "net = Lstm2(15)" - outputs = specs.create_net(spec, inputs) - self.assertEqual(outputs.get_shape().as_list(), [1, 64, 64, 15]) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (1, 64, 64, 15)) - - def testLstm2to1(self): - with self.test_session(): - inputs = constant_op.constant(_rand(1, 64, 64, 5)) - spec = "net = Lstm2to1(15)" - outputs = specs.create_net(spec, inputs) - self.assertEqual(outputs.get_shape().as_list(), [1, 64, 15]) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (1, 64, 15)) - - def testLstm2to0(self): - with self.test_session(): - inputs = constant_op.constant(_rand(1, 64, 64, 5)) - spec = "net = Lstm2to0(15)" - outputs = specs.create_net(spec, inputs) - self.assertEqual(outputs.get_shape().as_list(), [1, 15]) - variables.global_variables_initializer().run() - result = outputs.eval() - self.assertEqual(tuple(result.shape), (1, 15)) - def testKeywordRestriction(self): with self.test_session(): inputs = constant_op.constant(_rand(10, 20)) diff --git a/tensorflow/contrib/summary/BUILD b/tensorflow/contrib/summary/BUILD index 5ee5f1ae763db0ede9df464a08a9f1c7341b7cab..b58c83fdaf574fb349fac57c922f1178b7d13b66 100644 --- a/tensorflow/contrib/summary/BUILD +++ b/tensorflow/contrib/summary/BUILD @@ -112,5 +112,6 @@ py_library( "//tensorflow/core:protos_all_py", "//tensorflow/python:lib", "//tensorflow/python:platform", + "@org_sqlite//:python", ], ) diff --git a/tensorflow/contrib/summary/summary.py b/tensorflow/contrib/summary/summary.py index 7d3b8b7437a9ff5aaa0834db79bca8883cd679c8..2d6d7ea6a3eff2562ba8def4117e3aa6f818b6fd 100644 --- a/tensorflow/contrib/summary/summary.py +++ b/tensorflow/contrib/summary/summary.py @@ -18,6 +18,42 @@ The operations in this package are safe to use with eager execution turned on or off. It has a more flexible API that allows summaries to be written directly from ops to places other than event log files, rather than propagating protos from @{tf.summary.merge_all} to @{tf.summary.FileWriter}. + +To use with eager execution enabled, write your code as follows: + +global_step = tf.train.get_or_create_global_step() +summary_writer = tf.contrib.summary.create_file_writer( + train_dir, flush_millis=10000) +with summary_writer.as_default(), tf.contrib.summary.always_record_summaries(): + # model code goes here + # and in it call + tf.contrib.summary.scalar("loss", my_loss) + # In this case every call to tf.contrib.summary.scalar will generate a record + # ... + +To use it with graph execution, write your code as follows: + +global_step = tf.train.get_or_create_global_step() +summary_writer = tf.contrib.summary.create_file_writer( + train_dir, flush_millis=10000) +with summary_writer.as_default(), tf.contrib.summary.always_record_summaries(): + # model definition code goes here + # and in it call + tf.contrib.summary.scalar("loss", my_loss) + # In this case every call to tf.contrib.summary.scalar will generate an op, + # note the need to run tf.contrib.summary.all_summary_ops() to make sure these + # ops get executed. + # ... + train_op = .... + +with tf.Session(...) as sess: + tf.global_variables_initializer().run() + tf.contrib.summary.initialize(graph=tf.get_default_graph()) + # ... + while not_done_training: + sess.run([train_op, tf.contrib.summary.all_summary_ops()]) + # ... + """ from __future__ import absolute_import diff --git a/tensorflow/contrib/summary/summary_ops.py b/tensorflow/contrib/summary/summary_ops.py index ee661dfdc11451bb72bc2741b0b54ebf5c1e6543..b6249fc92f712b21197c2167fb5d1c4af1f48ca5 100644 --- a/tensorflow/contrib/summary/summary_ops.py +++ b/tensorflow/contrib/summary/summary_ops.py @@ -110,7 +110,7 @@ class SummaryWriter(object): def __init__(self, resource): self._resource = resource - if context.in_eager_mode(): + if context.in_eager_mode() and self._resource is not None: self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0") @@ -154,10 +154,12 @@ def initialize( to @{tf.get_default_session}. Raises: - RuntimeError: If in eager mode, or if the current thread has no - default @{tf.contrib.summary.SummaryWriter}. + RuntimeError: If the current thread has no default + @{tf.contrib.summary.SummaryWriter}. ValueError: If session wasn't passed and no default session. """ + if context.in_eager_mode(): + return if context.context().summary_writer_resource is None: raise RuntimeError("No default tf.contrib.summary.SummaryWriter found") if session is None: @@ -202,7 +204,7 @@ def create_file_writer(logdir, if flush_millis is None: flush_millis = constant_op.constant(2 * 60 * 1000) if filename_suffix is None: - filename_suffix = constant_op.constant("") + filename_suffix = constant_op.constant(".v2") return _make_summary_writer( name, gen_summary_ops.create_summary_file_writer, @@ -292,13 +294,9 @@ def all_summary_ops(): Returns: The summary ops. - - Raises: - RuntimeError: If in Eager mode. """ if context.in_eager_mode(): - raise RuntimeError( - "tf.contrib.summary.all_summary_ops is only supported in graph mode.") + return None return ops.get_collection(ops.GraphKeys._SUMMARY_COLLECTION) # pylint: disable=protected-access diff --git a/tensorflow/contrib/summary/summary_ops_test.py b/tensorflow/contrib/summary/summary_ops_test.py index 4ef03434b76ee04ce1bb0bd09c27a46db115bab3..bb7215f879411e91a1c47b87f5caede63fffea74 100644 --- a/tensorflow/contrib/summary/summary_ops_test.py +++ b/tensorflow/contrib/summary/summary_ops_test.py @@ -18,16 +18,17 @@ from __future__ import print_function import tempfile +import numpy as np import six from tensorflow.contrib.summary import summary_ops from tensorflow.contrib.summary import summary_test_util from tensorflow.core.framework import graph_pb2 from tensorflow.core.framework import node_def_pb2 +from tensorflow.core.framework import types_pb2 from tensorflow.python.eager import function from tensorflow.python.eager import test from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import state_ops @@ -37,14 +38,25 @@ from tensorflow.python.training import training_util get_all = summary_test_util.get_all get_one = summary_test_util.get_one +_NUMPY_NUMERIC_TYPES = { + types_pb2.DT_HALF: np.float16, + types_pb2.DT_FLOAT: np.float32, + types_pb2.DT_DOUBLE: np.float64, + types_pb2.DT_INT8: np.int8, + types_pb2.DT_INT16: np.int16, + types_pb2.DT_INT32: np.int32, + types_pb2.DT_INT64: np.int64, + types_pb2.DT_UINT8: np.uint8, + types_pb2.DT_UINT16: np.uint16, + types_pb2.DT_UINT32: np.uint32, + types_pb2.DT_UINT64: np.uint64, + types_pb2.DT_COMPLEX64: np.complex64, + types_pb2.DT_COMPLEX128: np.complex128, + types_pb2.DT_BOOL: np.bool_, +} -class TargetTest(test_util.TensorFlowTestCase): - def testInvalidDirectory(self): - logdir = '/tmp/apath/that/doesnt/exist' - self.assertFalse(gfile.Exists(logdir)) - with self.assertRaises(errors.NotFoundError): - summary_ops.create_file_writer(logdir, max_queue=0, name='t0') +class TargetTest(test_util.TensorFlowTestCase): def testShouldRecordSummary(self): self.assertFalse(summary_ops.should_record_summaries()) @@ -154,8 +166,9 @@ class DbTest(summary_test_util.SummaryDbTest): with writer.as_default(): self.assertEqual(5, adder(int64(2), int64(3)).numpy()) - six.assertCountEqual(self, [1, 1, 1], - get_all(self.db, 'SELECT step FROM Tensors')) + six.assertCountEqual( + self, [1, 1, 1], + get_all(self.db, 'SELECT step FROM Tensors WHERE dtype IS NOT NULL')) six.assertCountEqual(self, ['x', 'y', 'sum'], get_all(self.db, 'SELECT tag_name FROM Tags')) x_id = get_one(self.db, 'SELECT tag_id FROM Tags WHERE tag_name = "x"') @@ -166,8 +179,9 @@ class DbTest(summary_test_util.SummaryDbTest): with writer.as_default(): self.assertEqual(9, adder(int64(4), int64(5)).numpy()) - six.assertCountEqual(self, [1, 1, 1, 2, 2, 2], - get_all(self.db, 'SELECT step FROM Tensors')) + six.assertCountEqual( + self, [1, 1, 1, 2, 2, 2], + get_all(self.db, 'SELECT step FROM Tensors WHERE dtype IS NOT NULL')) six.assertCountEqual(self, [x_id, y_id, sum_id], get_all(self.db, 'SELECT tag_id FROM Tags')) self.assertEqual(2, get_tensor(self.db, x_id, 1)) @@ -212,9 +226,15 @@ class DbTest(summary_test_util.SummaryDbTest): def get_tensor(db, tag_id, step): - return get_one( - db, 'SELECT tensor FROM Tensors WHERE tag_id = ? AND step = ?', tag_id, - step) + cursor = db.execute( + 'SELECT dtype, shape, data FROM Tensors WHERE series = ? AND step = ?', + (tag_id, step)) + dtype, shape, data = cursor.fetchone() + assert dtype in _NUMPY_NUMERIC_TYPES + buf = np.frombuffer(data, dtype=_NUMPY_NUMERIC_TYPES[dtype]) + if not shape: + return buf[0] + return buf.reshape([int(i) for i in shape.split(',')]) def int64(x): diff --git a/tensorflow/contrib/summary/summary_test_internal.py b/tensorflow/contrib/summary/summary_test_internal.py new file mode 100644 index 0000000000000000000000000000000000000000..d0d3384735fb1eb1a048c7aa6da0037ee9fc6936 --- /dev/null +++ b/tensorflow/contrib/summary/summary_test_internal.py @@ -0,0 +1,60 @@ +# Copyright 2017 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. +# ============================================================================== +"""Internal helpers for tests in this directory.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +import os + +import sqlite3 + +from tensorflow.contrib.summary import summary_ops +from tensorflow.python.framework import test_util + + +class SummaryDbTest(test_util.TensorFlowTestCase): + """Helper for summary database testing.""" + + def setUp(self): + super(SummaryDbTest, self).setUp() + self.db_path = os.path.join(self.get_temp_dir(), 'DbTest.sqlite') + if os.path.exists(self.db_path): + os.unlink(self.db_path) + self.db = sqlite3.connect(self.db_path) + self.create_db_writer = functools.partial( + summary_ops.create_db_writer, + db_uri=self.db_path, + experiment_name='experiment', + run_name='run', + user_name='user') + + def tearDown(self): + self.db.close() + super(SummaryDbTest, self).tearDown() + + +def get_one(db, q, *p): + return db.execute(q, p).fetchone()[0] + + +def get_all(db, q, *p): + return unroll(db.execute(q, p).fetchall()) + + +def unroll(list_of_tuples): + return sum(list_of_tuples, ()) diff --git a/tensorflow/contrib/summary/summary_test_util.py b/tensorflow/contrib/summary/summary_test_util.py index bda57e6a0ca8e1ddb979a80de276911c7738f0aa..8506c4be9c4ca8305b62da17c7246e6e18313bd3 100644 --- a/tensorflow/contrib/summary/summary_test_util.py +++ b/tensorflow/contrib/summary/summary_test_util.py @@ -21,6 +21,7 @@ from __future__ import print_function import functools import os + import sqlite3 from tensorflow.contrib.summary import summary_ops diff --git a/tensorflow/contrib/tensor_forest/BUILD b/tensorflow/contrib/tensor_forest/BUILD index dc6cc674e6206a215b3295b1ea75cbb9ebfb61db..1e4cc3f0952ef74a1c89b7ed2d8c357fa8847ad5 100644 --- a/tensorflow/contrib/tensor_forest/BUILD +++ b/tensorflow/contrib/tensor_forest/BUILD @@ -497,6 +497,7 @@ py_library( ":tensor_forest_v4_ops_py", "//tensorflow/contrib/decision_trees/proto:generic_tree_model_py", "//tensorflow/contrib/framework:framework_py", + "//tensorflow/contrib/tensor_forest/proto:fertile_stats_proto_py", "//tensorflow/contrib/tensor_forest/proto:tensor_forest_params_proto_py", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", @@ -530,7 +531,6 @@ py_library( srcs_version = "PY2AND3", deps = [ ":client_lib", - "//tensorflow/contrib/framework:framework_py", "//tensorflow/contrib/layers:layers_py", "//tensorflow/contrib/learn", "//tensorflow/python:array_ops", diff --git a/tensorflow/contrib/tensor_forest/client/random_forest.py b/tensorflow/contrib/tensor_forest/client/random_forest.py index 807c8398439b9d225c974d2a1cdc9b845df4d26e..4abcc20ed334e706c8ae59e2127dfd6f4e152361 100644 --- a/tensorflow/contrib/tensor_forest/client/random_forest.py +++ b/tensorflow/contrib/tensor_forest/client/random_forest.py @@ -17,9 +17,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib import framework as contrib_framework from tensorflow.contrib import layers - +from tensorflow.contrib.learn.python.learn.estimators import constants from tensorflow.contrib.learn.python.learn.estimators import estimator from tensorflow.contrib.learn.python.learn.estimators import head as head_lib from tensorflow.contrib.learn.python.learn.estimators import model_fn as model_fn_lib @@ -44,8 +43,8 @@ from tensorflow.python.training import training_util KEYS_NAME = 'keys' LOSS_NAME = 'rf_training_loss' TREE_PATHS_PREDICTION_KEY = 'tree_paths' -VARIANCE_PREDICTION_KEY = 'regression_variance' - +VARIANCE_PREDICTION_KEY = 'prediction_variance' +ALL_SERVING_KEY = 'tensorforest_all' EPSILON = 0.000001 @@ -135,7 +134,8 @@ def get_model_fn(params, trainer_id=0, report_feature_importances=False, local_eval=False, - head_scope=None): + head_scope=None, + include_all_in_serving=False): """Return a model function given a way to construct a graph builder.""" if model_head is None: model_head = get_default_head(params, weights_name) @@ -190,7 +190,7 @@ def get_model_fn(params, features, labels, input_weights=weights, num_trainers=num_trainers, trainer_id=trainer_id), - state_ops.assign_add(contrib_framework.get_global_step(), 1)) + state_ops.assign_add(training_util.get_global_step(), 1)) # Put weights back in if weights is not None: @@ -238,9 +238,14 @@ def get_model_fn(params, if params.inference_tree_paths: model_ops.predictions[TREE_PATHS_PREDICTION_KEY] = tree_paths - if params.regression: - model_ops.predictions[VARIANCE_PREDICTION_KEY] = regression_variance - + model_ops.predictions[VARIANCE_PREDICTION_KEY] = regression_variance + if include_all_in_serving: + # In order to serve the variance we need to add the prediction dict + # to output_alternatives dict. + if not model_ops.output_alternatives: + model_ops.output_alternatives = {} + model_ops.output_alternatives[ALL_SERVING_KEY] = ( + constants.ProblemType.UNSPECIFIED, model_ops.predictions) return model_ops return _model_fn @@ -295,7 +300,8 @@ class TensorForestEstimator(estimator.Estimator): report_feature_importances=False, local_eval=False, version=None, - head=None): + head=None, + include_all_in_serving=False): """Initializes a TensorForestEstimator instance. Args: @@ -341,6 +347,23 @@ class TensorForestEstimator(estimator.Estimator): version: Unused. head: A heads_lib.Head object that calculates losses and such. If None, one will be automatically created based on params. + include_all_in_serving: if True, allow preparation of the complete + prediction dict including the variance to be exported for serving with + the Servo lib; and it also requires calling export_savedmodel with + default_output_alternative_key=ALL_SERVING_KEY, i.e. + estimator.export_savedmodel(export_dir_base=your_export_dir, + serving_input_fn=your_export_input_fn, + default_output_alternative_key=ALL_SERVING_KEY) + if False, resort to default behavior, i.e. export scores and + probabilities but no variances. In this case + default_output_alternative_key should be None while calling + export_savedmodel(). + Note, that due to backward compatibility we cannot always set + include_all_in_serving to True because in this case calling + export_saved_model() without + default_output_alternative_key=ALL_SERVING_KEY (legacy behavior) the + saved_model_export_utils.get_output_alternatives() would raise + ValueError. Returns: A `TensorForestEstimator` instance. @@ -359,7 +382,9 @@ class TensorForestEstimator(estimator.Estimator): num_trainers=num_trainers, trainer_id=trainer_id, report_feature_importances=report_feature_importances, - local_eval=local_eval), + local_eval=local_eval, + include_all_in_serving=include_all_in_serving, + ), model_dir=model_dir, config=config, feature_engineering_fn=feature_engineering_fn) diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/hard_routing_function_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/hard_routing_function_op.cc index 76cfb4c9ca02269f9fee61c767acc6cb4a0b4ca7..cf0db788a419f64ed891df8aa097fa8826f6de91 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/hard_routing_function_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/hard_routing_function_op.cc @@ -99,18 +99,17 @@ class HardRoutingFunction : public OpKernel { const Tensor& tree_biases_tensor = context->input(2); if (input_data.shape().dim_size(0) > 0) { - OP_REQUIRES(context, input_data.shape().dims() == 2, - errors::InvalidArgument( - "input_data should be two-dimensional")); + OP_REQUIRES( + context, input_data.shape().dims() == 2, + errors::InvalidArgument("input_data should be two-dimensional")); } // Check tensor bounds. if (!CheckTensorBounds(context, input_data)) return; - const int32 num_data = static_cast( - input_data.shape().dim_size(0)); - const int32 num_features = static_cast( - input_data.shape().dim_size(1)); + const int32 num_data = static_cast(input_data.shape().dim_size(0)); + const int32 num_features = + static_cast(input_data.shape().dim_size(1)); Tensor* output_probability = nullptr; TensorShape output_probability_shape; @@ -125,9 +124,8 @@ class HardRoutingFunction : public OpKernel { OP_REQUIRES_OK(context, context->allocate_output(0, output_probability_shape, &output_probability)); - OP_REQUIRES_OK(context, - context->allocate_output(1, output_path_shape, - &output_path)); + OP_REQUIRES_OK( + context, context->allocate_output(1, output_path_shape, &output_path)); auto out_probability = output_probability->tensor(); auto out_path = output_path->tensor(); @@ -144,12 +142,11 @@ class HardRoutingFunction : public OpKernel { out_probability(i, 0) = 1.0; out_path(i, 0) = 0; for (int j = 0; j < tree_depth_ - 1; j++) { - float left_prob = LeftProbability(point, - tree_parameters_tensor.Slice(j, j+1), - tree_biases(j), - num_features); + float left_prob = + LeftProbability(point, tree_parameters_tensor.Slice(j, j + 1), + tree_biases(j), num_features); - int32 left_child = 2*node + 1; + int32 left_child = 2 * node + 1; int32 right_child = left_child + 1; float dot_product = 0.0; diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/k_feature_gradient_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/k_feature_gradient_op.cc index 28f50f1a32eb1827a242d527cd42c58487877959..f64155fa55af22d57c6619d8a39da0455dc0de65 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/k_feature_gradient_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/k_feature_gradient_op.cc @@ -85,12 +85,9 @@ REGISTER_OP("KFeatureGradient") class KFeatureGradient : public OpKernel { public: - explicit KFeatureGradient(OpKernelConstruction* context) - : OpKernel(context) { - OP_REQUIRES_OK(context, context->GetAttr("layer_num", - &layer_num_)); - OP_REQUIRES_OK(context, context->GetAttr("random_seed", - &random_seed_)); + explicit KFeatureGradient(OpKernelConstruction* context) : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("layer_num", &layer_num_)); + OP_REQUIRES_OK(context, context->GetAttr("random_seed", &random_seed_)); } void Compute(OpKernelContext* context) override { @@ -101,14 +98,14 @@ class KFeatureGradient : public OpKernel { const Tensor& routing_tensor = context->input(3); // Extract dimensions from input tensors. - const int32 num_data = static_cast( - input_data_tensor.shape().dim_size(0)); - const int32 num_features = static_cast( - input_data_tensor.shape().dim_size(1)); - const int32 num_nodes = static_cast( - tree_parameters_tensor.shape().dim_size(0)); - const int32 num_features_per_node = static_cast( - tree_parameters_tensor.shape().dim_size(1)); + const int32 num_data = + static_cast(input_data_tensor.shape().dim_size(0)); + const int32 num_features = + static_cast(input_data_tensor.shape().dim_size(1)); + const int32 num_nodes = + static_cast(tree_parameters_tensor.shape().dim_size(0)); + const int32 num_features_per_node = + static_cast(tree_parameters_tensor.shape().dim_size(1)); // Construct output tensors. Tensor* out_routes = nullptr; @@ -127,12 +124,12 @@ class KFeatureGradient : public OpKernel { out_weights_shape.AddDim(num_nodes); out_weights_shape.AddDim(num_features_per_node); - OP_REQUIRES_OK(context, context->allocate_output( - 0, out_routes_shape, &out_routes)); - OP_REQUIRES_OK(context, context->allocate_output( - 1, out_data_shape, &out_data)); - OP_REQUIRES_OK(context, context->allocate_output( - 2, out_weights_shape, &out_weights)); + OP_REQUIRES_OK(context, + context->allocate_output(0, out_routes_shape, &out_routes)); + OP_REQUIRES_OK(context, + context->allocate_output(1, out_data_shape, &out_data)); + OP_REQUIRES_OK( + context, context->allocate_output(2, out_weights_shape, &out_weights)); tensorforest::Initialize(*out_data, 0.0f); @@ -148,18 +145,13 @@ class KFeatureGradient : public OpKernel { std::vector feature_set; for (int i = 0; i < num_data; i++) { - const Tensor point = input_data_tensor.Slice(i, i+1); + const Tensor point = input_data_tensor.Slice(i, i + 1); feature_set.clear(); // Traverse the tree from the bottom up. for (int j = num_nodes - 1; j >= 0; j--) { - tensorforest::GetFeatureSet( - layer_num_, - j, - random_seed_, - num_features, - num_features_per_node, - &feature_set); + tensorforest::GetFeatureSet(layer_num_, j, random_seed_, num_features, + num_features_per_node, &feature_set); // Compute routing gradient. // j is a leaf node. @@ -170,12 +162,8 @@ class KFeatureGradient : public OpKernel { int32 right_child = left_child + 1; float left_prob = LeftProbabilityK( - point, - feature_set, - tree_parameters_tensor.Slice(j, j+1), - tree_biases(j), - num_features, - num_features_per_node); + point, feature_set, tree_parameters_tensor.Slice(j, j + 1), + tree_biases(j), num_features, num_features_per_node); float right_prob = 1.0f - left_prob; diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/k_feature_routing_function_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/k_feature_routing_function_op.cc index 9bc42eb61fae013de3e4ea73aaf371cdaa4ccf9a..e7cafb144da84865ad2b4ea0c33866ddb89119a5 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/k_feature_routing_function_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/k_feature_routing_function_op.cc @@ -43,7 +43,6 @@ using shape_inference::ShapeHandle; using tensorforest::CheckTensorBounds; using tensorforest::LeftProbabilityK; - // The term 'routing function' is synonymous with 'the probability // that an instance is routed to each leaf node.' It is defined in // 'Deep Neural Decision Forests' by Kontschieder et al. @@ -96,10 +95,8 @@ class KFeatureRoutingFunction : public OpKernel { OP_REQUIRES_OK(context, context->GetAttr("max_nodes", &max_nodes_)); OP_REQUIRES_OK(context, context->GetAttr("num_features_per_node", &num_features_per_node_)); - OP_REQUIRES_OK(context, context->GetAttr("layer_num", - &layer_num_)); - OP_REQUIRES_OK(context, context->GetAttr("random_seed", - &random_seed_)); + OP_REQUIRES_OK(context, context->GetAttr("layer_num", &layer_num_)); + OP_REQUIRES_OK(context, context->GetAttr("random_seed", &random_seed_)); } void Compute(OpKernelContext* context) override { @@ -108,27 +105,25 @@ class KFeatureRoutingFunction : public OpKernel { const Tensor& tree_biases_tensor = context->input(2); if (input_data.shape().dim_size(0) > 0) { - OP_REQUIRES(context, input_data.shape().dims() == 2, - errors::InvalidArgument( - "input_data should be two-dimensional")); + OP_REQUIRES( + context, input_data.shape().dims() == 2, + errors::InvalidArgument("input_data should be two-dimensional")); } // Check tensor bounds. if (!CheckTensorBounds(context, input_data)) return; - const int32 num_data = static_cast( - input_data.shape().dim_size(0)); - const int32 num_features = static_cast( - input_data.shape().dim_size(1)); + const int32 num_data = static_cast(input_data.shape().dim_size(0)); + const int32 num_features = + static_cast(input_data.shape().dim_size(1)); Tensor* output_probabilities = nullptr; TensorShape output_shape; output_shape.AddDim(num_data); output_shape.AddDim(max_nodes_); - OP_REQUIRES_OK(context, - context->allocate_output(0, output_shape, - &output_probabilities)); + OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, + &output_probabilities)); auto out_probs = output_probabilities->tensor(); const auto tree_biases = tree_biases_tensor.tensor(); @@ -136,30 +131,22 @@ class KFeatureRoutingFunction : public OpKernel { // Iteratively compute the probability of reaching each leaf. std::vector feature_set; for (int i = 0; i < num_data; i++) { - const Tensor point = input_data.Slice(i, i+1); + const Tensor point = input_data.Slice(i, i + 1); out_probs(i, 0) = 1.0f; for (int j = 0; j < max_nodes_ / 2; j++) { feature_set.clear(); - tensorforest::GetFeatureSet( - layer_num_, - i, - random_seed_, - num_features, - num_features_per_node_, - &feature_set); - - int32 left_child = 2*j + 1; + tensorforest::GetFeatureSet(layer_num_, i, random_seed_, num_features, + num_features_per_node_, &feature_set); + + int32 left_child = 2 * j + 1; int32 right_child = left_child + 1; float prob = out_probs(i, j); - float left_prob = LeftProbabilityK(point, - feature_set, - tree_parameters_tensor.Slice(j, j+1), - tree_biases(j), - num_features, - num_features_per_node_); + float left_prob = LeftProbabilityK( + point, feature_set, tree_parameters_tensor.Slice(j, j + 1), + tree_biases(j), num_features, num_features_per_node_); out_probs(i, left_child) = prob * left_prob; out_probs(i, right_child) = prob * (1.0f - left_prob); diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/routing_function_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/routing_function_op.cc index 4027e732b3f52585c2149c3cdc71535664f04ed4..0c2eaabe8f3e1e1377a8d5c5308aaec00030a20f 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/routing_function_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/routing_function_op.cc @@ -90,46 +90,43 @@ class RoutingFunction : public OpKernel { const Tensor& tree_biases_tensor = context->input(2); if (input_data.shape().dim_size(0) > 0) { - OP_REQUIRES(context, input_data.shape().dims() == 2, - errors::InvalidArgument( - "input_data should be two-dimensional")); + OP_REQUIRES( + context, input_data.shape().dims() == 2, + errors::InvalidArgument("input_data should be two-dimensional")); } // Check tensor bounds. if (!CheckTensorBounds(context, input_data)) return; - const int32 num_data = static_cast( - input_data.shape().dim_size(0)); - const int32 num_features = static_cast( - input_data.shape().dim_size(1)); + const int32 num_data = static_cast(input_data.shape().dim_size(0)); + const int32 num_features = + static_cast(input_data.shape().dim_size(1)); Tensor* output_probabilities = nullptr; TensorShape output_shape; output_shape.AddDim(num_data); output_shape.AddDim(max_nodes_); - OP_REQUIRES_OK(context, - context->allocate_output(0, output_shape, - &output_probabilities)); + OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, + &output_probabilities)); auto out_probs = output_probabilities->tensor(); const auto tree_biases = tree_biases_tensor.tensor(); // Iteratively compute the probability of reaching each leaf. for (int i = 0; i < num_data; i++) { - const Tensor point = input_data.Slice(i, i+1); + const Tensor point = input_data.Slice(i, i + 1); out_probs(i, 0) = 1.0; for (int j = 0; j < max_nodes_ / 2; j++) { - int32 left_child = 2*j + 1; + int32 left_child = 2 * j + 1; int32 right_child = left_child + 1; float prob = out_probs(i, j); - float left_prob = LeftProbability(point, - tree_parameters_tensor.Slice(j, j+1), - tree_biases(j), - num_features); + float left_prob = + LeftProbability(point, tree_parameters_tensor.Slice(j, j + 1), + tree_biases(j), num_features); out_probs(i, left_child) = prob * left_prob; out_probs(i, right_child) = prob * (1.0 - left_prob); diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_function_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_function_op.cc index 66aa293dc1cb93b82f06d838ad7b0f9c09761585..c9df09bfda44e665ed013da383e1e9a2c665c454 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_function_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_function_op.cc @@ -96,10 +96,9 @@ class StochasticHardRoutingFunction : public OpKernel { explicit StochasticHardRoutingFunction(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr("tree_depth", &tree_depth_)); - OP_REQUIRES_OK(context, context->GetAttr("random_seed", - &random_seed_)); + OP_REQUIRES_OK(context, context->GetAttr("random_seed", &random_seed_)); single_rand_ = std::unique_ptr( - new random::PhiloxRandom(random_seed_)); + new random::PhiloxRandom(random_seed_)); rng_ = std::unique_ptr( new random::SimplePhilox(single_rand_.get())); } @@ -111,20 +110,19 @@ class StochasticHardRoutingFunction : public OpKernel { const Tensor& tree_biases_tensor = context->input(2); if (input_data.shape().dim_size(0) > 0) { - OP_REQUIRES(context, input_data.shape().dims() == 2, - errors::InvalidArgument( - "input_data should be two-dimensional")); + OP_REQUIRES( + context, input_data.shape().dims() == 2, + errors::InvalidArgument("input_data should be two-dimensional")); } // Check tensor bounds. if (!CheckTensorBounds(context, input_data)) return; - const int32 num_data = static_cast( - input_data.shape().dim_size(0)); - const int32 num_features = static_cast( - input_data.shape().dim_size(1)); - const int32 num_nodes = static_cast( - tree_parameters_tensor.shape().dim_size(0)); + const int32 num_data = static_cast(input_data.shape().dim_size(0)); + const int32 num_features = + static_cast(input_data.shape().dim_size(1)); + const int32 num_nodes = + static_cast(tree_parameters_tensor.shape().dim_size(0)); Tensor* output_probability = nullptr; TensorShape output_probability_shape; @@ -139,9 +137,8 @@ class StochasticHardRoutingFunction : public OpKernel { OP_REQUIRES_OK(context, context->allocate_output(0, output_probability_shape, &output_probability)); - OP_REQUIRES_OK(context, - context->allocate_output(1, output_path_shape, - &output_path)); + OP_REQUIRES_OK( + context, context->allocate_output(1, output_path_shape, &output_path)); auto out_probability = output_probability->tensor(); auto out_path = output_path->tensor(); @@ -150,19 +147,18 @@ class StochasticHardRoutingFunction : public OpKernel { // Stochastically traverse the tree to a leaf. for (int i = 0; i < num_data; i++) { - const Tensor point = input_data.Slice(i, i+1); + const Tensor point = input_data.Slice(i, i + 1); int32 node = 0; out_probability(i, 0) = 1.0; out_path(i, 0) = 0; for (int j = 0; j < tree_depth_ - 1; j++) { - int32 left_child = 2*node + 1; + int32 left_child = 2 * node + 1; int32 right_child = left_child + 1; - float left_prob = LeftProbability(point, - tree_parameters_tensor.Slice(j, j+1), - tree_biases(j), - num_features); + float left_prob = + LeftProbability(point, tree_parameters_tensor.Slice(j, j + 1), + tree_biases(j), num_features); if (left_prob < rng_->RandFloat()) { CHECK_LT(i, num_data); diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_gradient_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_gradient_op.cc index 0b5afe464f4b9608af0feca584aaa799f5980f46..b0d8b832b5437db7a4b3026e80ae99d0391d7f7a 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_gradient_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_gradient_op.cc @@ -149,14 +149,14 @@ class StochasticHardRoutingGradient : public OpKernel { TensorShape output_bias_shape; output_bias_shape.AddDim(num_data); - OP_REQUIRES_OK(context, context->allocate_output( - 0, output_routing_shape, &output_routing)); - OP_REQUIRES_OK(context, context->allocate_output( - 1, output_data_shape, &output_data)); - OP_REQUIRES_OK(context, context->allocate_output( - 2, output_parameters_shape, &output_parameters)); - OP_REQUIRES_OK(context, context->allocate_output( - 3, output_bias_shape, &output_bias)); + OP_REQUIRES_OK(context, context->allocate_output(0, output_routing_shape, + &output_routing)); + OP_REQUIRES_OK( + context, context->allocate_output(1, output_data_shape, &output_data)); + OP_REQUIRES_OK(context, context->allocate_output(2, output_parameters_shape, + &output_parameters)); + OP_REQUIRES_OK( + context, context->allocate_output(3, output_bias_shape, &output_bias)); tensorforest::Initialize(*output_routing, 0.0); tensorforest::Initialize(*output_data, 0.0); @@ -178,7 +178,7 @@ class StochasticHardRoutingGradient : public OpKernel { const Tensor point = input_data.Slice(i, i + 1); // Traverses the tree from the bottom up. - for (int j = tree_depth_-1; j > -1; j--) { + for (int j = tree_depth_ - 1; j > -1; j--) { int32 node = path(i, j); CHECK_LT(node, num_nodes); diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/unpack_path_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/unpack_path_op.cc index cacad03e274c3279eb3706e71e1bcdf8433ca1ef..25825a78a1498490009fe4ff6bbfc67493727037 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/unpack_path_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/unpack_path_op.cc @@ -64,8 +64,7 @@ REGISTER_OP("UnpackPath") class UnpackPath : public OpKernel { public: - explicit UnpackPath(OpKernelConstruction* context) - : OpKernel(context) {} + explicit UnpackPath(OpKernelConstruction* context) : OpKernel(context) {} void Compute(OpKernelContext* context) override { VLOG(1) << "unpack start"; @@ -73,8 +72,8 @@ class UnpackPath : public OpKernel { const Tensor& path_values_tensor = context->input(1); const int32 num_data = static_cast(path_tensor.shape().dim_size(0)); - const int32 tree_depth = static_cast( - path_tensor.shape().dim_size(1)); + const int32 tree_depth = + static_cast(path_tensor.shape().dim_size(1)); const int32 num_nodes = MathUtil::IPow(2, tree_depth) - 1; @@ -107,7 +106,6 @@ class UnpackPath : public OpKernel { } }; -REGISTER_KERNEL_BUILDER(Name("UnpackPath").Device(DEVICE_CPU), - UnpackPath); +REGISTER_KERNEL_BUILDER(Name("UnpackPath").Device(DEVICE_CPU), UnpackPath); } // namespace tensorflow diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/utils.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/utils.cc index c091a73c4e48a47bdccea3ec99371faab9c586c2..34388fe1aab72895a805141ec66a71ecf0f42ba4 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/utils.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/utils.cc @@ -25,9 +25,7 @@ namespace tensorforest { using tensorflow::Tensor; -float LeftProbability(const Tensor& point, - const Tensor& weight, - float bias, +float LeftProbability(const Tensor& point, const Tensor& weight, float bias, int num_features) { const auto p = point.unaligned_flat(); const auto w = weight.unaligned_flat(); @@ -41,11 +39,8 @@ float LeftProbability(const Tensor& point, return 1.0 / (1.0 + exp(-dot_product + bias)); } -float LeftProbabilityK(const Tensor& point, - std::vector feature_set, - const Tensor& weight, - float bias, - int num_features, +float LeftProbabilityK(const Tensor& point, std::vector feature_set, + const Tensor& weight, float bias, int num_features, int k) { const auto p = point.unaligned_flat(); const auto w = weight.unaligned_flat(); diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/utils.h b/tensorflow/contrib/tensor_forest/hybrid/core/ops/utils.h index c5902184f95ea8f97be4a10d1101a38333359d44..69a0143a4e319157a4526ca80fbb3f6472902b31 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/utils.h +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/utils.h @@ -24,16 +24,11 @@ namespace tensorflow { namespace tensorforest { // Returns the probability that the point falls to the left. -float LeftProbability(const Tensor& point, - const Tensor& weight, - float bias, +float LeftProbability(const Tensor& point, const Tensor& weight, float bias, int num_features); -float LeftProbabilityK(const Tensor& point, - std::vector feature_set, - const Tensor& weight, - float bias, - int num_features, +float LeftProbabilityK(const Tensor& point, std::vector feature_set, + const Tensor& weight, float bias, int num_features, int k); // Returns a random set of num_features_to_pick features in the @@ -49,5 +44,3 @@ void GetFeatureSet(int32 tree_num, int32 node_num, int32 random_seed, } // namespace tensorflow #endif // LEARNING_LIB_TENSOR_FOREST_HYBRID_CORE_OPS_UTILS_H_ - - diff --git a/tensorflow/contrib/tensor_forest/kernels/data_spec.h b/tensorflow/contrib/tensor_forest/kernels/data_spec.h index 05590d6992e2fd7eeee8d242561229ab53bb16de..0a3abe56dfc4f611ac8ed0815e4c74a639d2477e 100644 --- a/tensorflow/contrib/tensor_forest/kernels/data_spec.h +++ b/tensorflow/contrib/tensor_forest/kernels/data_spec.h @@ -15,8 +15,8 @@ // This is a surrogate for using a proto, since it doesn't seem to be possible // to use protos in a dynamically-loaded/shared-linkage library, which is // what is used for custom ops in tensorflow/contrib. -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_DATA_SPEC_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_DATA_SPEC_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_DATA_SPEC_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_DATA_SPEC_H_ #include #include "tensorflow/core/lib/strings/numbers.h" @@ -138,4 +138,4 @@ class TensorForestDataSpec { } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_DATA_SPEC_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_DATA_SPEC_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/reinterpret_string_to_float_op.cc b/tensorflow/contrib/tensor_forest/kernels/reinterpret_string_to_float_op.cc index 47b49a379c4b7a17d35b52c1403f67c2f07aeeaf..b21a9179777c21f65435e136aa6082e27fb3b78c 100644 --- a/tensorflow/contrib/tensor_forest/kernels/reinterpret_string_to_float_op.cc +++ b/tensorflow/contrib/tensor_forest/kernels/reinterpret_string_to_float_op.cc @@ -30,15 +30,13 @@ namespace tensorflow { using tensorforest::CheckTensorBounds; - float Convert(const string& in) { const std::size_t intval = std::hash()(in); return static_cast(intval); } - -void Evaluate(const Tensor& input_data, Tensor output_data, - int32 start, int32 end) { +void Evaluate(const Tensor& input_data, Tensor output_data, int32 start, + int32 end) { auto out_data = output_data.unaligned_flat(); const auto in_data = input_data.unaligned_flat(); @@ -59,9 +57,8 @@ class ReinterpretStringToFloat : public OpKernel { if (!CheckTensorBounds(context, input_data)) return; Tensor* output_data = nullptr; - OP_REQUIRES_OK(context, - context->allocate_output(0, input_data.shape(), - &output_data)); + OP_REQUIRES_OK( + context, context->allocate_output(0, input_data.shape(), &output_data)); // Evaluate input data in parallel. const int32 num_data = static_cast(input_data.NumElements()); @@ -73,8 +70,8 @@ class ReinterpretStringToFloat : public OpKernel { auto work = [&input_data, output_data, num_data](int64 start, int64 end) { CHECK(start <= end); CHECK(end <= num_data); - Evaluate(input_data, *output_data, - static_cast(start), static_cast(end)); + Evaluate(input_data, *output_data, static_cast(start), + static_cast(end)); }; Shard(num_threads, worker_threads->workers, num_data, 100, work); } diff --git a/tensorflow/contrib/tensor_forest/kernels/scatter_add_ndim_op.cc b/tensorflow/contrib/tensor_forest/kernels/scatter_add_ndim_op.cc index dd2a98b08cdb486c98c161390a3a1f81d31e1f4b..60740c2be3703141805c7eae0ac384edf934ab3d 100644 --- a/tensorflow/contrib/tensor_forest/kernels/scatter_add_ndim_op.cc +++ b/tensorflow/contrib/tensor_forest/kernels/scatter_add_ndim_op.cc @@ -22,7 +22,6 @@ #include "tensorflow/core/framework/shape_inference.h" #include "tensorflow/core/platform/logging.h" - namespace tensorflow { using tensorforest::CheckTensorBounds; @@ -38,20 +37,19 @@ class ScatterAddNdim : public OpKernel { if (indices_tensor.shape().dim_size(0) > 0) { OP_REQUIRES(context, indices_tensor.shape().dims() == 2, - errors::InvalidArgument( - "indices should be two-dimensional")); + errors::InvalidArgument("indices should be two-dimensional")); const int32 delta_dims = deltas_tensor.shape().dims(); OP_REQUIRES( context, indices_tensor.shape().dim_size(1) + delta_dims == - input_tensor.shape().dims() + 1, + input_tensor.shape().dims() + 1, errors::InvalidArgument( "Number of indices dimensions should be the same as input " "rank.")); OP_REQUIRES( context, indices_tensor.shape().dim_size(0) == - deltas_tensor.shape().dim_size(0), + deltas_tensor.shape().dim_size(0), errors::InvalidArgument( "Number of updates should be same as number of indices.")); } else { @@ -68,8 +66,8 @@ class ScatterAddNdim : public OpKernel { const auto indices = indices_tensor.tensor(); const auto deltas = deltas_tensor.unaligned_flat(); - const int32 num_dims = static_cast( - indices_tensor.shape().dim_size(1)); + const int32 num_dims = + static_cast(indices_tensor.shape().dim_size(1)); // Figure out if indices don't specify a complete position in the // input tensor. @@ -80,10 +78,9 @@ class ScatterAddNdim : public OpKernel { // Calculate index multipliers. std::vector multipliers; - OP_REQUIRES( - context, input.size() < std::numeric_limits::max(), - errors::InvalidArgument( - "Input must contain less than 2^31 total elements")); + OP_REQUIRES(context, input.size() < std::numeric_limits::max(), + errors::InvalidArgument( + "Input must contain less than 2^31 total elements")); int32 last_size = static_cast(input.size()); for (int32 j = 0; j < num_dims; j++) { diff --git a/tensorflow/contrib/tensor_forest/kernels/tree_utils.cc b/tensorflow/contrib/tensor_forest/kernels/tree_utils.cc index 94e12cea5a072f0746e642196d55f3a3b13a16c3..44997ec5d6d5fdb9aab52ab7a50f46a731bfda66 100644 --- a/tensorflow/contrib/tensor_forest/kernels/tree_utils.cc +++ b/tensorflow/contrib/tensor_forest/kernels/tree_utils.cc @@ -65,8 +65,8 @@ void GetTwoBest(int max, const std::function& score_fn, float ClassificationSplitScore( const Eigen::Tensor& splits, - const Eigen::Tensor& rights, - int32 num_classes, int i) { + const Eigen::Tensor& rights, int32 num_classes, + int i) { Eigen::array offsets; // Class counts are stored with the total in [0], so the length of each // count vector is num_classes + 1. @@ -74,7 +74,7 @@ float ClassificationSplitScore( Eigen::array extents; extents[0] = num_classes; return WeightedGiniImpurity(splits.slice(offsets, extents)) + - WeightedGiniImpurity(rights.slice(offsets, extents)); + WeightedGiniImpurity(rights.slice(offsets, extents)); } void GetTwoBestClassification(const Tensor& total_counts, @@ -90,29 +90,28 @@ void GetTwoBestClassification(const Tensor& total_counts, // in seg faults, so we have to go with flat views of these tensors. However, // it is still pretty efficient because we put off evaluation until the // score is actually returned. - const auto tc = total_counts.Slice( - accumulator, accumulator + 1).unaligned_flat(); + const auto tc = + total_counts.Slice(accumulator, accumulator + 1).unaligned_flat(); // TODO(gilberth): See if we can delay evaluation here by templating the // arguments to ClassificationSplitScore. - const Eigen::Tensor splits = split_counts.Slice( - accumulator, accumulator + 1).unaligned_flat(); + const Eigen::Tensor splits = + split_counts.Slice(accumulator, accumulator + 1).unaligned_flat(); Eigen::array bcast; bcast[0] = num_splits; const Eigen::Tensor rights = tc.broadcast(bcast) - splits; - std::function score_fn = std::bind( - ClassificationSplitScore, splits, rights, num_classes, - std::placeholders::_1); + std::function score_fn = + std::bind(ClassificationSplitScore, splits, rights, num_classes, + std::placeholders::_1); GetTwoBest(num_splits, score_fn, best_score, best_index, second_best_score, second_best_index); } -int32 BestFeatureClassification( - const Tensor& total_counts, const Tensor& split_counts, - int32 accumulator) { +int32 BestFeatureClassification(const Tensor& total_counts, + const Tensor& split_counts, int32 accumulator) { float best_score; float second_best_score; int best_feature_index; @@ -130,8 +129,7 @@ float RegressionSplitScore( const Eigen::Tensor& splits_square, const Eigen::Tensor& right_sums, const Eigen::Tensor& right_squares, - int32 accumulator, - int32 num_regression_dims, int i) { + int32 accumulator, int32 num_regression_dims, int i) { Eigen::array offsets = {i * num_regression_dims + 1}; Eigen::array extents = {num_regression_dims - 1}; float left_count = splits_count_accessor(accumulator, i, 0); @@ -141,15 +139,15 @@ float RegressionSplitScore( // Guard against divide-by-zero. if (left_count > 0) { - score += WeightedVariance( - splits_sum.slice(offsets, extents), - splits_square.slice(offsets, extents), left_count); + score += + WeightedVariance(splits_sum.slice(offsets, extents), + splits_square.slice(offsets, extents), left_count); } if (right_count > 0) { - score += WeightedVariance(right_sums.slice(offsets, extents), - right_squares.slice(offsets, extents), - right_count); + score += + WeightedVariance(right_sums.slice(offsets, extents), + right_squares.slice(offsets, extents), right_count); } return score; } @@ -159,20 +157,20 @@ void GetTwoBestRegression(const Tensor& total_sums, const Tensor& total_squares, int32 accumulator, float* best_score, int* best_index, float* second_best_score, int* second_best_index) { const int32 num_splits = static_cast(split_sums.shape().dim_size(1)); - const int32 num_regression_dims = static_cast( - split_sums.shape().dim_size(2)); + const int32 num_regression_dims = + static_cast(split_sums.shape().dim_size(2)); // Ideally, Eigen::Tensor::chip would be best to use here but it results // in seg faults, so we have to go with flat views of these tensors. However, // it is still pretty efficient because we put off evaluation until the // score is actually returned. - const auto tc_sum = total_sums.Slice( - accumulator, accumulator + 1).unaligned_flat(); - const auto tc_square = total_squares.Slice( - accumulator, accumulator + 1).unaligned_flat(); - const auto splits_sum = split_sums.Slice( - accumulator, accumulator + 1).unaligned_flat(); - const auto splits_square = split_squares.Slice( - accumulator, accumulator + 1).unaligned_flat(); + const auto tc_sum = + total_sums.Slice(accumulator, accumulator + 1).unaligned_flat(); + const auto tc_square = + total_squares.Slice(accumulator, accumulator + 1).unaligned_flat(); + const auto splits_sum = + split_sums.Slice(accumulator, accumulator + 1).unaligned_flat(); + const auto splits_square = + split_squares.Slice(accumulator, accumulator + 1).unaligned_flat(); // Eigen is infuriating to work with, usually resulting in all kinds of // unhelpful compiler errors when trying something that seems sane. This // helps us do a simple thing like access the first element (the counts) @@ -193,10 +191,10 @@ void GetTwoBestRegression(const Tensor& total_sums, const Tensor& total_squares, best_score, best_index, second_best_score, second_best_index); } -int32 BestFeatureRegression( - const Tensor& total_sums, const Tensor& total_squares, - const Tensor& split_sums, const Tensor& split_squares, - int32 accumulator) { +int32 BestFeatureRegression(const Tensor& total_sums, + const Tensor& total_squares, + const Tensor& split_sums, + const Tensor& split_squares, int32 accumulator) { float best_score; float second_best_score; int best_feature_index; @@ -207,10 +205,11 @@ int32 BestFeatureRegression( return best_feature_index; } -bool BestSplitDominatesRegression( - const Tensor& total_sums, const Tensor& total_squares, - const Tensor& split_sums, const Tensor& split_squares, - int32 accumulator) { +bool BestSplitDominatesRegression(const Tensor& total_sums, + const Tensor& total_squares, + const Tensor& split_sums, + const Tensor& split_squares, + int32 accumulator) { // TODO(thomaswc): Implement this, probably as part of v3. return false; } @@ -599,7 +598,6 @@ bool Decide(float value, float bias, DataColumnTypes type) { } } - void GetParentWeightedMean(float leaf_sum, const float* leaf_data, float parent_sum, const float* parent_data, float valid_leaf_threshold, int num_outputs, diff --git a/tensorflow/contrib/tensor_forest/kernels/tree_utils.h b/tensorflow/contrib/tensor_forest/kernels/tree_utils.h index 35f9fb7eaf4d73e98293f1ea6a4b45b71212a92c..edbac6700677633cbd4d41f7040b4859ca599c4a 100644 --- a/tensorflow/contrib/tensor_forest/kernels/tree_utils.h +++ b/tensorflow/contrib/tensor_forest/kernels/tree_utils.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_TREE_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_TREE_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_TREE_UTILS_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_TREE_UTILS_H_ #include @@ -45,13 +45,10 @@ const int32 LEAF_NODE = -1; const int32 FREE_NODE = -2; // Used to indicate column types, e.g. categorical vs. float -enum DataColumnTypes { - kDataFloat = 0, - kDataCategorical = 1 -}; +enum DataColumnTypes { kDataFloat = 0, kDataCategorical = 1 }; // Calculates the sum of a tensor. -template +template T Sum(Tensor counts) { Eigen::Tensor count_sum = counts.unaligned_flat().sum(); @@ -97,7 +94,7 @@ float WeightedGiniImpurity(const T& counts) { return RawWeightedGiniImpurity(smoothed); } -template +template float WeightedVariance(const T1& sums, const T2& squares, float count) { const auto e_x = sums / count; const auto e_x2 = squares / count; @@ -120,10 +117,11 @@ int32 BestFeatureRegression(const Tensor& total_sums, // Returns true if the best split's variance is sufficiently smaller than // that of the next best split. -bool BestSplitDominatesRegression( - const Tensor& total_sums, const Tensor& total_squares, - const Tensor& split_sums, const Tensor& split_squares, - int32 accumulator); +bool BestSplitDominatesRegression(const Tensor& total_sums, + const Tensor& total_squares, + const Tensor& split_sums, + const Tensor& split_squares, + int32 accumulator); // Performs booststrap_samples bootstrap samples of the best split's class // counts and the second best splits's class counts, and returns true if at @@ -178,10 +176,8 @@ bool DecideNode(const GetFeatureFnType& get_dense, // isn't present in sparse_input_indices. sparse_input_indices is assumed // to be sorted. template -float FindSparseValue( - const T1& sparse_input_indices, - const T2& sparse_input_values, - int32 i, int32 j) { +float FindSparseValue(const T1& sparse_input_indices, + const T2& sparse_input_values, int32 i, int32 j) { int32 low = 0; int32 high = sparse_input_values.dimension(0); while (low < high) { @@ -273,7 +269,6 @@ int32 GetNumSparseFeatures(const T1& indices, int32 input_index, // categorical data, it is value != bias. bool Decide(float value, float bias, DataColumnTypes type = kDataFloat); - // Returns true if all the splits are initialized. Since they get initialized // in order, we can simply infer this from the last split. // This should only be called for a single allocator's candidate features @@ -307,4 +302,4 @@ void GetParentWeightedMean(float leaf_sum, const float* leaf_data, } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_TREE_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_CORE_OPS_TREE_UTILS_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/tree_utils_test.cc b/tensorflow/contrib/tensor_forest/kernels/tree_utils_test.cc index 7485a695dfba93fd3f57c19096b205b10e2fa8b5..08553545502c21eb8f2d68bfd342f8ba7c081adb 100644 --- a/tensorflow/contrib/tensor_forest/kernels/tree_utils_test.cc +++ b/tensorflow/contrib/tensor_forest/kernels/tree_utils_test.cc @@ -44,11 +44,13 @@ TEST(TestWeightedVariance, Basic) { Tensor squares = test::AsTensor({29, 12}, {2}); EXPECT_FLOAT_EQ(WeightedVariance(sums.unaligned_flat(), - squares.unaligned_flat(), 3), 2.0); + squares.unaligned_flat(), 3), + 2.0); Tensor zero = test::AsTensor({0}, {1}); EXPECT_FLOAT_EQ(WeightedVariance(zero.unaligned_flat(), - zero.unaligned_flat(), 1), 0); + zero.unaligned_flat(), 1), + 0); } TEST(TestInitialize, Basic) { @@ -94,17 +96,16 @@ TEST(BestFeatureClassification, Basic) { const int32 num_accumulators = 4; const int32 num_splits = 3; const int32 num_classes = 4; - Tensor totals = test::AsTensor({1, 5, 6, 7, - 0, 0, 0, 0, - 30, 10, 10, 10, // this one - -1, -1, -1, -1}, - {num_accumulators, num_classes}); - Tensor splits = test::AsTensor( - {1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 30, 10, 10, 10, 10, 0, 0, 10, 19, 5, 6, 8, // this one - -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, - {num_accumulators, num_splits, num_classes}); + Tensor totals = test::AsTensor( + {1, 5, 6, 7, 0, 0, 0, 0, 30, 10, 10, 10, // this one + -1, -1, -1, -1}, + {num_accumulators, num_classes}); + Tensor splits = + test::AsTensor({1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 30, 10, + 10, 10, 10, 0, 0, 10, 19, 5, 6, 8, // this one + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, + {num_accumulators, num_splits, num_classes}); EXPECT_EQ(BestFeatureClassification(totals, splits, 2), 1); } @@ -114,17 +115,16 @@ TEST(BestFeatureClassification, NoWinner) { const int32 num_splits = 3; const int32 num_classes = 4; // When counts are all the same, the most reasonable thing to do is pick 0. - Tensor totals = test::AsTensor({1, 5, 6, 7, - 0, 0, 0, 0, - 18, 6, 6, 6, // this one - -1, -1, -1, -1}, - {num_accumulators, num_classes}); - Tensor splits = test::AsTensor( - {1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 9, 3, 3, 3, 9, 3, 3, 3, 9, 3, 3, 3, // this one - -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, - {num_accumulators, num_splits, num_classes}); + Tensor totals = + test::AsTensor({1, 5, 6, 7, 0, 0, 0, 0, 18, 6, 6, 6, // this one + -1, -1, -1, -1}, + {num_accumulators, num_classes}); + Tensor splits = + test::AsTensor({1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9, 3, + 3, 3, 9, 3, 3, 3, 9, 3, 3, 3, // this one + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, + {num_accumulators, num_splits, num_classes}); EXPECT_EQ(BestFeatureClassification(totals, splits, 2), 0); } @@ -133,36 +133,34 @@ TEST(BestFeatureRegression, Basic) { const int32 num_accumulators = 4; const int32 num_splits = 3; const int32 num_classes = 4; - Tensor total_sums = test::AsTensor( - {1, 5, 6, 7, - 0, 0, 0, 0, - 10, 8, 6, 9, // this one - -1, -1, -1, -1}, - {num_accumulators, num_classes}); + Tensor total_sums = + test::AsTensor({1, 5, 6, 7, 0, 0, 0, 0, 10, 8, 6, 9, // this one + -1, -1, -1, -1}, + {num_accumulators, num_classes}); Tensor total_squares = test::AsTensor( - {1, 5, 6, 7, - 0, 0, 0, 0, - 100, 50, 40, 45, // this one + {1, 5, 6, 7, 0, 0, 0, 0, 100, 50, 40, 45, // this one -1, -1, -1, -1}, {num_accumulators, num_classes}); - Tensor split_sums = test::AsTensor( - {1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 10, 8, 6, 9, 9, 8, 5, 9, 0, 0, 0, 0, // this one - -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, - {num_accumulators, num_splits, num_classes}); + Tensor split_sums = + test::AsTensor({1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 8, + 6, 9, 9, 8, 5, 9, 0, 0, 0, 0, // this one + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, + {num_accumulators, num_splits, num_classes}); // lower the variance by lowering one of the squares just a little. - Tensor split_squares = test::AsTensor( - {1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 100, 50, 40, 45, 100, 50, 40, 43, 0, 0, 0, 0, // this one - -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, - {num_accumulators, num_splits, num_classes}); + Tensor split_squares = + test::AsTensor( + {1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 100, 50, 40, 45, 100, 50, 40, 43, 0, 0, 0, 0, // this one + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, + {num_accumulators, num_splits, num_classes}); EXPECT_EQ(BestFeatureRegression(total_sums, total_squares, split_sums, - split_squares, 2), 1); + split_squares, 2), + 1); } TEST(BestFeatureRegression, NoWinner) { @@ -170,37 +168,33 @@ TEST(BestFeatureRegression, NoWinner) { const int32 num_splits = 3; const int32 num_classes = 4; // when counts are all the same, the most reasonable thing to do is pick 0. - Tensor total_sums = test::AsTensor( - {1, 5, 6, 7, - 0, 0, 0, 0, - 10, 8, 6, 9, // this one - -1, -1, -1, -1}, - {num_accumulators, num_classes}); + Tensor total_sums = + test::AsTensor({1, 5, 6, 7, 0, 0, 0, 0, 10, 8, 6, 9, // this one + -1, -1, -1, -1}, + {num_accumulators, num_classes}); Tensor total_squares = test::AsTensor( - {1, 5, 6, 7, - 0, 0, 0, 0, - 100, 50, 40, 45, // this one + {1, 5, 6, 7, 0, 0, 0, 0, 100, 50, 40, 45, // this one -1, -1, -1, -1}, {num_accumulators, num_classes}); - Tensor split_sums = test::AsTensor( - {1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 10, 8, 6, 9, 10, 8, 6, 9, 10, 8, 6, 9, // this one - -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, - {num_accumulators, num_splits, num_classes}); + Tensor split_sums = + test::AsTensor({1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 8, + 6, 9, 10, 8, 6, 9, 10, 8, 6, 9, // this one + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, + {num_accumulators, num_splits, num_classes}); Tensor split_squares = test::AsTensor( - {1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 100, 50, 40, 45, 100, 50, 40, 45, 100, 50, 40, 45, // this one - -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, + {1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 100, 50, 40, 45, 100, 50, 40, 45, 100, 50, 40, 45, // this one + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}, {num_accumulators, num_splits, num_classes}); EXPECT_EQ(BestFeatureRegression(total_sums, total_squares, split_sums, - split_squares, 2), 0); + split_squares, 2), + 0); } } // namespace tensorforest } // namespace tensorflow - diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/BUILD b/tensorflow/contrib/tensor_forest/kernels/v4/BUILD index b7876e1df6c77d1fa3fa31abb37fc16d06540f15..794b76d8583c3608d540d34a5aaf1d1a799f35e3 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/BUILD +++ b/tensorflow/contrib/tensor_forest/kernels/v4/BUILD @@ -302,6 +302,7 @@ cc_library( "//tensorflow/contrib/tensor_forest/proto:fertile_stats_proto_cc", ], [ + "//third_party/eigen3", "//tensorflow/contrib/decision_trees/proto:generic_tree_model_cc_headers_only", "//tensorflow/contrib/tensor_forest/proto:fertile_stats_proto_cc_headers_only", ], @@ -322,6 +323,7 @@ cc_library( srcs = ["params.cc"], hdrs = ["params.h"], deps = [ + "//third_party/eigen3", "//tensorflow/core:framework_headers_lib", ] + if_static( [ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.cc b/tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.cc index 81e2a1b2a1b720574210e376fa786923367794a6..f4a7058ddb8bfdd6393a9369006aabc29d058d3b 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.cc @@ -14,8 +14,8 @@ // ============================================================================= #include "tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.h" -#include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/platform/env.h" namespace tensorflow { @@ -58,8 +58,7 @@ CandidateGraphRunner::CandidateGraphRunner( // Features don't change, store them in a tensor. const auto& oblique = split.inequality_left_child_test().oblique(); const int32 feat_size = oblique.features_size(); - features_.reset( - new Tensor(tensorflow::DT_INT32, TensorShape({feat_size}))); + features_.reset(new Tensor(tensorflow::DT_INT32, TensorShape({feat_size}))); auto feat = features_->flat(); int i = 0; for (const auto& id : oblique.features()) { @@ -67,10 +66,10 @@ CandidateGraphRunner::CandidateGraphRunner( } } -void CandidateGraphRunner::RunOp( - const string& name, const TensorNameValueList& inputs, - const std::vector& output_tensor_names, - std::vector* outputs) { +void CandidateGraphRunner::RunOp(const string& name, + const TensorNameValueList& inputs, + const std::vector& output_tensor_names, + std::vector* outputs) { std::vector op_name; if (name != kNoOp) { op_name.push_back(name); diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.h b/tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.h index 4bd1f06c72945f73e50301c337692e0b510d3693..2e7368dc12c74b9dc44b72394668bf2de71f2f90 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/candidate_graph_runner.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_CANDIDATE_GRAPH_RUNNER_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_CANDIDATE_GRAPH_RUNNER_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_CANDIDATE_GRAPH_RUNNER_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_CANDIDATE_GRAPH_RUNNER_H_ #include #include @@ -70,4 +70,4 @@ class CandidateGraphRunner { } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_CANDIDATE_GRAPH_RUNNER_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_CANDIDATE_GRAPH_RUNNER_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h b/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h index bf88216d663cc9b69746a93379124bf1d9a30df9..328af28725af016e90b30ae2d303ffba15c81c1f 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_TREE_RESOURCE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_TREE_RESOURCE_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_TREE_RESOURCE_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_TREE_RESOURCE_H_ #include "tensorflow/contrib/decision_trees/proto/generic_tree_model.pb.h" #include "tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h" @@ -26,7 +26,6 @@ namespace tensorflow { namespace tensorforest { - // Keep a tree ensemble in memory for efficient evaluation and mutation. class DecisionTreeResource : public ResourceBase { public: @@ -35,15 +34,12 @@ class DecisionTreeResource : public ResourceBase { string DebugString() override { return strings::StrCat("DecisionTree[size=", - decision_tree_->decision_tree().nodes_size(), - "]"); + decision_tree_->decision_tree().nodes_size(), "]"); } void MaybeInitialize(); - const decision_trees::Model& decision_tree() const { - return *decision_tree_; - } + const decision_trees::Model& decision_tree() const { return *decision_tree_; } decision_trees::Model* mutable_decision_tree() { return decision_tree_.get(); @@ -59,9 +55,7 @@ class DecisionTreeResource : public ResourceBase { // Resets the resource and frees the proto. // Caller needs to hold the mutex lock while calling this. - void Reset() { - decision_tree_.reset(new decision_trees::Model()); - } + void Reset() { decision_tree_.reset(new decision_trees::Model()); } mutex* get_mutex() { return &mu_; } @@ -84,8 +78,7 @@ class DecisionTreeResource : public ResourceBase { std::vector> node_evaluators_; }; - } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_TREE_RESOURCE_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_TREE_RESOURCE_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h b/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h index 3f03c2d05bb1090fa75f4b6e7ad4f00caaea61a4..bf2b2aaa3c8f433ab4fc145217857112f7a0a579 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_NODE_EVALUATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_NODE_EVALUATOR_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_NODE_EVALUATOR_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_NODE_EVALUATOR_H_ #include "tensorflow/contrib/decision_trees/proto/generic_tree_model.pb.h" #include "tensorflow/contrib/decision_trees/proto/generic_tree_model_extensions.pb.h" @@ -22,7 +22,6 @@ namespace tensorflow { namespace tensorforest { - // Base class for evaluators of decision nodes that effectively copy proto // contents into C++ structures for faster execution. class DecisionNodeEvaluator { @@ -104,4 +103,4 @@ struct CandidateEvalatorCollection { } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_NODE_EVALUATOR_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_DECISION_NODE_EVALUATOR_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator_test.cc b/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator_test.cc index 5c49b87443e7b1f4ef532256ae2efdc9fa985d8a..af5cf72a3c0bea0eef45c3446acf52ff389c6751 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator_test.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator_test.cc @@ -20,11 +20,11 @@ namespace tensorflow { namespace { +using tensorflow::decision_trees::InequalityTest; +using tensorflow::decision_trees::MatchingValuesTest; using tensorflow::tensorforest::InequalityDecisionNodeEvaluator; using tensorflow::tensorforest::MatchingValuesDecisionNodeEvaluator; using tensorflow::tensorforest::ObliqueInequalityDecisionNodeEvaluator; -using tensorflow::decision_trees::InequalityTest; -using tensorflow::decision_trees::MatchingValuesTest; TEST(InequalityDecisionNodeEvaluatorTest, TestLessOrEqual) { InequalityTest test; @@ -124,4 +124,3 @@ TEST(ObliqueDecisionNodeEvaluatorTest, Basic) { } // namespace } // namespace tensorflow - diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/fertile-stats-resource.h b/tensorflow/contrib/tensor_forest/kernels/v4/fertile-stats-resource.h index dacf033d99018d47787b644b12d3181780df7113..eea0be27caf0a022ba7acaacd359c75a2df4eedb 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/fertile-stats-resource.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/fertile-stats-resource.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_FERTILE_STATS_RESOURCE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_FERTILE_STATS_RESOURCE_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_FERTILE_STATS_RESOURCE_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_FERTILE_STATS_RESOURCE_H_ #include @@ -40,9 +40,7 @@ class FertileStatsResource : public ResourceBase { model_op_ = LeafModelOperatorFactory::CreateLeafModelOperator(params_); } - string DebugString() override { - return "FertileStats"; - } + string DebugString() override { return "FertileStats"; } void ExtractFromProto(const FertileStats& stats); @@ -50,8 +48,7 @@ class FertileStatsResource : public ResourceBase { // Resets the resource and frees the proto. // Caller needs to hold the mutex lock while calling this. - void Reset() { - } + void Reset() {} // Reset the stats for a node, but leave the leaf_stats intact. void ResetSplitStats(int32 node_id, int32 depth) { @@ -84,7 +81,6 @@ class FertileStatsResource : public ResourceBase { // was found. bool BestSplit(int32 node_id, SplitCandidate* best, int32* depth); - private: mutex mu_; std::shared_ptr model_op_; @@ -94,8 +90,7 @@ class FertileStatsResource : public ResourceBase { void AllocateNode(int32 node_id, int32 depth); }; - } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_FERTILE_STATS_RESOURCE_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_FERTILE_STATS_RESOURCE_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/graph_collection_operator.h b/tensorflow/contrib/tensor_forest/kernels/v4/graph_collection_operator.h index 2ae3a79b3dd69b3fd3d31a055589b2edc63afa3c..4ae48179afc8452e6a3ec61dede16b9941482bcc 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/graph_collection_operator.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/graph_collection_operator.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GRAPH_COLLECTION_OPERATOR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GRAPH_COLLECTION_OPERATOR_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GRAPH_COLLECTION_OPERATOR_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GRAPH_COLLECTION_OPERATOR_H_ #include #include "tensorflow/contrib/decision_trees/proto/generic_tree_model.pb.h" @@ -78,4 +78,4 @@ class GraphRunnerSplitCollectionOperator : public SplitCollectionOperator { } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GRAPH_COLLECTION_OPERATOR_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GRAPH_COLLECTION_OPERATOR_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.cc b/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.cc index 3ce630e3a9691b87ad291a9f29616f741953dd84..da600d34eacdf27514709240723e5bb730cfe7f0 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.cc @@ -20,7 +20,6 @@ #include "tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.h" #include "tensorflow/core/lib/random/distribution_sampler.h" - namespace tensorflow { namespace tensorforest { @@ -454,14 +453,14 @@ void DenseClassificationGrowStats::PackToProto(FertileSlot* slot) const { class_stats->add_value()->set_float_value(total_counts_[i]); } - for (int split_num = 0; split_num < num_splits(); ++split_num) { + for (int split_num = 0; split_num < num_splits(); ++split_num) { auto* cand = slot->add_candidates(); *cand->mutable_split() = splits_[split_num]; auto* left_stats = cand->mutable_left_stats() ->mutable_classification() ->mutable_dense_counts(); for (int i = 0; i < num_outputs_; ++i) { - left_stats->add_value()->set_float_value(left_count(split_num, i)); + left_stats->add_value()->set_float_value(left_count(split_num, i)); } } } @@ -546,7 +545,7 @@ void SparseClassificationGrowStats::PackToProto(FertileSlot* slot) const { (*class_stats)[entry.first] = val; } - for (int split_num = 0; split_num < num_splits(); ++split_num) { + for (int split_num = 0; split_num < num_splits(); ++split_num) { auto* cand = slot->add_candidates(); *cand->mutable_split() = splits_[split_num]; auto* left_stats = cand->mutable_left_stats() @@ -561,8 +560,8 @@ void SparseClassificationGrowStats::PackToProto(FertileSlot* slot) const { } } -float SparseClassificationGrowStats::GiniScore( - int split, float* left_sum, float* right_sum) const { +float SparseClassificationGrowStats::GiniScore(int split, float* left_sum, + float* right_sum) const { float left_square = 0, right_square = 0; *left_sum = 0; *right_sum = 0; @@ -844,12 +843,11 @@ void LeastSquaresRegressionGrowStats::PackToProto(FertileSlot* slot) const { total_squares->add_value()->set_float_value(total_sum_squares_[i]); } - for (int split_num = 0; split_num < num_splits(); ++split_num) { + for (int split_num = 0; split_num < num_splits(); ++split_num) { auto* cand = slot->add_candidates(); *cand->mutable_split() = splits_[split_num]; - auto* sums = cand->mutable_left_stats() - ->mutable_regression() - ->mutable_mean_output(); + auto* sums = + cand->mutable_left_stats()->mutable_regression()->mutable_mean_output(); auto* squares = cand->mutable_left_stats() ->mutable_regression() ->mutable_mean_output_squares(); @@ -891,20 +889,17 @@ float LeastSquaresRegressionGrowStats::SplitVariance(int split) const { float total_variance = 0; for (int i = 0; i < params_.num_outputs(); ++i) { // Left side - const float le_x = - left_sum(split, i) / left_counts_[split]; + const float le_x = left_sum(split, i) / left_counts_[split]; - const float le_x2 = - left_square(split, i) / left_counts_[split]; + const float le_x2 = left_square(split, i) / left_counts_[split]; total_variance += le_x2 - le_x * le_x; // Right side const float re_x = (total_sum_[i] - left_sum(split, i)) / (weight_sum_ - left_counts_[split]); - const float re_x2 = - (total_sum_squares_[i] - left_square(split, i)) / - (weight_sum_ - left_counts_[split]); + const float re_x2 = (total_sum_squares_[i] - left_square(split, i)) / + (weight_sum_ - left_counts_[split]); total_variance += re_x2 - re_x * re_x; } return total_variance; @@ -937,8 +932,7 @@ bool LeastSquaresRegressionGrowStats::BestSplit(SplitCandidate* best) const { left->set_weight_sum(left_counts_[best_index]); auto* left_output_sum = left_reg_stats->mutable_mean_output(); for (int i = 0; i < num_outputs; ++i) { - left_output_sum->add_value()->set_float_value( - left_sum(best_index, i)); + left_output_sum->add_value()->set_float_value(left_sum(best_index, i)); } // Right @@ -947,8 +941,8 @@ bool LeastSquaresRegressionGrowStats::BestSplit(SplitCandidate* best) const { right->set_weight_sum(weight_sum_ - left_counts_[best_index]); auto* right_output_sum = right_reg_stats->mutable_mean_output(); for (int i = 0; i < num_outputs; ++i) { - right_output_sum->add_value()->set_float_value( - total_sum_[i] - left_sum(best_index, i)); + right_output_sum->add_value()->set_float_value(total_sum_[i] - + left_sum(best_index, i)); } return true; } diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.h b/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.h index 3e41ab50b9d78943db8ee58aab85a8c7541e2320..04e6b0a735320dd024e326a94ef910593a326245 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GROW_STATS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GROW_STATS_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GROW_STATS_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GROW_STATS_H_ #include #include @@ -73,21 +73,15 @@ class GrowStats { const InputTarget* target, int example) {} void RemoveSplit(int split_num); - int num_splits() const { - return splits_.size(); - } + int num_splits() const { return splits_.size(); } - float weight_sum() const { - return weight_sum_; - } + float weight_sum() const { return weight_sum_; } virtual bool IsInitialized() const { return weight_sum_ > 0 || splits_.size() == num_splits_to_consider_; } - int32 depth() const { - return depth_; - } + int32 depth() const { return depth_; } protected: GrowStats(const TensorForestParams& params, int32 depth); @@ -206,8 +200,8 @@ class ClassificationStats : public GrowStats { virtual float left_count(int split, int class_num) const = 0; virtual float right_count(int split, int class_num) const = 0; - virtual void ClassificationAddLeftExample( - int split, int64 int_label, float weight) = 0; + virtual void ClassificationAddLeftExample(int split, int64 int_label, + float weight) = 0; virtual void ClassificationAddRightExample(int split, int64 int_label, float weight) { // Does nothing by default, but sub-classes can override. @@ -316,7 +310,7 @@ class DenseClassificationGrowStats : public ClassificationStats { 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 { @@ -375,15 +369,13 @@ class SparseClassificationGrowStats : public ClassificationStats { SparseClassificationGrowStats(const TensorForestParams& params, int32 depth) : ClassificationStats(params, depth) {} - void Initialize() override { - Clear(); - } + void Initialize() override { Clear(); } void ExtractFromProto(const FertileSlot& slot) override; 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 { @@ -562,9 +554,9 @@ class LeastSquaresRegressionGrowStats : public GrowStats { } void RemoveSplitStats(int split_num) override { left_sums_.erase(left_sums_.begin() + num_outputs_ * split_num, - left_sums_.begin() + num_outputs_ * (split_num + 1)); + left_sums_.begin() + num_outputs_ * (split_num + 1)); left_squares_.erase(left_squares_.begin() + num_outputs_ * split_num, - left_squares_.begin() + num_outputs_ * (split_num + 1)); + left_squares_.begin() + num_outputs_ * (split_num + 1)); left_counts_.erase(left_counts_.begin() + split_num, left_counts_.begin() + (split_num + 1)); } @@ -605,8 +597,7 @@ class LeastSquaresRegressionGrowStats : public GrowStats { std::vector left_counts_; }; - } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GROW_STATS_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_GROW_STATS_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats_test.cc b/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats_test.cc index ceb58d2ead5c2f148c96d9cb9532a73688593d33..26e989928e00de1b2ae1646abf216adfbec2be4f 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats_test.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/grow_stats_test.cc @@ -24,21 +24,21 @@ namespace tensorflow { namespace { -using tensorflow::tensorforest::GrowStats; -using tensorflow::tensorforest::TestableInputTarget; -using tensorflow::tensorforest::FertileSlot; +using tensorflow::decision_trees::BinaryNode; +using tensorflow::decision_trees::FeatureId; +using tensorflow::decision_trees::InequalityTest; using tensorflow::tensorforest::DenseClassificationGrowStats; -using tensorflow::tensorforest::SparseClassificationGrowStats; +using tensorflow::tensorforest::FertileSlot; using tensorflow::tensorforest::FixedSizeClassStats; using tensorflow::tensorforest::FixedSizeSparseClassificationGrowStats; +using tensorflow::tensorforest::GrowStats; using tensorflow::tensorforest::LeastSquaresRegressionGrowStats; -using tensorflow::tensorforest::TensorForestParams; +using tensorflow::tensorforest::SparseClassificationGrowStats; using tensorflow::tensorforest::SPLIT_FINISH_BASIC; using tensorflow::tensorforest::SPLIT_FINISH_DOMINATE_HOEFFDING; using tensorflow::tensorforest::SPLIT_PRUNE_HOEFFDING; -using tensorflow::decision_trees::BinaryNode; -using tensorflow::decision_trees::InequalityTest; -using tensorflow::decision_trees::FeatureId; +using tensorflow::tensorforest::TensorForestParams; +using tensorflow::tensorforest::TestableInputTarget; BinaryNode MakeSplit(const string& feat, float val) { BinaryNode split; @@ -52,8 +52,7 @@ BinaryNode MakeSplit(const string& feat, float val) { return split; } -void RunBatch(GrowStats* stats, - const TestableInputTarget* target) { +void RunBatch(GrowStats* stats, const TestableInputTarget* target) { std::unique_ptr dataset( new tensorflow::tensorforest::TestableDataSet( {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, 2)); @@ -102,18 +101,10 @@ class TestableRunningStats : public DenseClassificationGrowStats { TestableRunningStats(const TensorForestParams& params, int32 depth) : DenseClassificationGrowStats(params, depth) {} - float test_left_sum(int split) { - return get_left_gini()->sum(split); - } - float test_left_square(int split) { - return get_left_gini()->square(split); - } - float test_right_sum(int split) { - return get_right_gini()->sum(split); - } - float test_right_square(int split) { - return get_right_gini()->square(split); - } + float test_left_sum(int split) { return get_left_gini()->sum(split); } + float test_left_square(int split) { return get_left_gini()->square(split); } + float test_right_sum(int split) { return get_right_gini()->sum(split); } + float test_right_square(int split) { return get_right_gini()->square(split); } }; TEST(GrowStatsDenseClassificationTest, BasicRunningStats) { @@ -166,9 +157,7 @@ class TestableFinishEarly : public DenseClassificationGrowStats { int num_times_called_; protected: - void CheckFinishEarlyHoeffding() override { - ++num_times_called_; - } + void CheckFinishEarlyHoeffding() override { ++num_times_called_; } }; TEST(GrowStatsDenseClassificationTest, TestFinishEarly) { @@ -212,7 +201,6 @@ TEST(GrowStatsDenseClassificationTest, TestFinishEarly) { ASSERT_EQ(stat->num_times_called_, 9); } - TEST(GrowStatsDenseClassificationTest, TestCheckPruneHoeffding) { TensorForestParams params; params.set_num_outputs(2); @@ -224,7 +212,8 @@ TEST(GrowStatsDenseClassificationTest, TestCheckPruneHoeffding) { finish->set_type(SPLIT_FINISH_BASIC); finish->mutable_check_every_steps()->set_constant_value(100); params.mutable_pruning_type()->set_type(SPLIT_PRUNE_HOEFFDING); - params.mutable_pruning_type()->mutable_prune_every_samples() + params.mutable_pruning_type() + ->mutable_prune_every_samples() ->set_constant_value(1); // On each iteration, we add two examples, one of class 0 and one @@ -234,8 +223,8 @@ TEST(GrowStatsDenseClassificationTest, TestCheckPruneHoeffding) { std::vector weights = {1, 1}; TestableInputTarget target(labels, weights, 1); std::unique_ptr dataset( - new tensorflow::tensorforest::TestableDataSet( - {-1.0, -1.0, 1.0, -1.0}, 2)); + new tensorflow::tensorforest::TestableDataSet({-1.0, -1.0, 1.0, -1.0}, + 2)); DenseClassificationGrowStats stats(params, 1); stats.Initialize(); diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/input_data.cc b/tensorflow/contrib/tensor_forest/kernels/v4/input_data.cc index 14cb19d36f33e478728aba3e28b7bea11b691d34..d43884481afbbbc988d6eb80e01e49663df6914b 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/input_data.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/input_data.cc @@ -21,8 +21,6 @@ namespace tensorflow { namespace tensorforest { namespace { -const int32 SPARSE_DEFAULT = 0; - bool DecideInequalityTest(const decision_trees::InequalityTest& test, float value) { float bias = test.threshold().float_value(); @@ -111,10 +109,10 @@ void TensorDataSet::set_input_tensors(const Tensor& dense, dense_data_.reset(new DenseStorageType(dense.tensor())); } if (sparse_indices.shape().dims() == 2) { - sparse_indices_.reset(new SparseIndicesStorageType( - sparse_indices.tensor())); - sparse_values_.reset(new SparseValuesStorageType( - sparse_values.tensor())); + sparse_indices_.reset( + new SparseIndicesStorageType(sparse_indices.tensor())); + sparse_values_.reset( + new SparseValuesStorageType(sparse_values.tensor())); sparse_batch_size_ = sparse_shape.tensor()(0); } original_dense_tensor_ = dense; diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/input_data.h b/tensorflow/contrib/tensor_forest/kernels/v4/input_data.h index e3d4edbf8a512a027e4b67916d1f2ad3f347a18b..c544a8c75e9bfe8fe6bbea8913e7be17d868bfef 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/input_data.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/input_data.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_DATA_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_DATA_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_DATA_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_DATA_H_ #include #include #include "google/protobuf/any.pb.h" @@ -93,9 +93,7 @@ class TensorDataSet { // an int32 you can avoid the atoi32. virtual float GetExampleValue(int example, int32 feature_id) const; - int num_features() { - return available_features_.size(); - } + int num_features() { return available_features_.size(); } const Tensor& original_tensor() const { return original_dense_tensor_; } @@ -123,4 +121,4 @@ class TensorDataSet { } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_DATA_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_DATA_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/input_target.h b/tensorflow/contrib/tensor_forest/kernels/v4/input_target.h index 0309ec1de9aec1044eb87e01cafc40c26ba3de14..d4402b6055a36d38042a0e6cfa07b532ec11c093 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/input_target.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/input_target.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_TARGET_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_TARGET_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_TARGET_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_TARGET_H_ #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" @@ -79,9 +79,7 @@ class TensorInputTarget : public StoredInputTarget { return (*target_)(example_index * num_targets_ + target_index); } - const Tensor& original_tensor() const { - return original_tensor_; - } + const Tensor& original_tensor() const { return original_tensor_; } protected: Tensor original_tensor_; @@ -89,4 +87,4 @@ class TensorInputTarget : public StoredInputTarget { } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_TARGET_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_INPUT_TARGET_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators.cc b/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators.cc index d43c068e462ff78b114fb29bd8cf0ee0c6080fcd..83614a25314117ef9ba29b4dcf6ebee8f7f3e226 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators.cc @@ -160,6 +160,5 @@ void RegressionLeafModelOperator::ExportModel( } } - } // namespace tensorforest } // namespace tensorflow diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators.h b/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators.h index 946a648f22ff4175782c42cc70c59440e6ac0e17..cc4ec8dc9e330784bbcfeb54fa92e0a2db9449a8 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_LEAF_MODEL_OPERATORS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_LEAF_MODEL_OPERATORS_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_LEAF_MODEL_OPERATORS_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_LEAF_MODEL_OPERATORS_H_ #include "tensorflow/contrib/decision_trees/proto/generic_tree_model.pb.h" #include "tensorflow/contrib/tensor_forest/kernels/v4/input_target.h" @@ -146,4 +146,4 @@ class LeafModelOperatorFactory { } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_LEAF_MODEL_OPERATORS_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_LEAF_MODEL_OPERATORS_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators_test.cc b/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators_test.cc index ffd92c01f9a59719e6bb2458c2f28253c364a2e8..ab4191809b6a7400114acf85991c74acfac55505 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators_test.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/leaf_model_operators_test.cc @@ -26,19 +26,19 @@ namespace { using tensorflow::decision_trees::Leaf; using tensorflow::tensorforest::DenseClassificationLeafModelOperator; using tensorflow::tensorforest::LeafModelOperator; -using tensorflow::tensorforest::SparseClassificationLeafModelOperator; -using tensorflow::tensorforest::SparseOrDenseClassificationLeafModelOperator; using tensorflow::tensorforest::LeafStat; using tensorflow::tensorforest::RegressionLeafModelOperator; -using tensorflow::tensorforest::TestableInputTarget; +using tensorflow::tensorforest::SparseClassificationLeafModelOperator; +using tensorflow::tensorforest::SparseOrDenseClassificationLeafModelOperator; using tensorflow::tensorforest::TensorForestParams; +using tensorflow::tensorforest::TestableInputTarget; const int32 kNumClasses = 3; constexpr char kRegressionStatProto[] = - "weight_sum: 3 " - "regression { " - "mean_output { " + "weight_sum: 3 " + "regression { " + "mean_output { " "value { " " float_value: 27 " "} " @@ -48,8 +48,8 @@ constexpr char kRegressionStatProto[] = "value { " " float_value: 10 " "} " - "} " - "mean_output_squares { " + "} " + "mean_output_squares { " "value {" " float_value: 245" "}" @@ -59,8 +59,8 @@ constexpr char kRegressionStatProto[] = "value {" " float_value: 46" "}" - "}" -"}"; + "}" + "}"; void TestClassificationNormalUse(const std::unique_ptr& op) { Leaf l; @@ -83,7 +83,6 @@ void TestClassificationNormalUse(const std::unique_ptr& op) { EXPECT_FLOAT_EQ(op->GetOutputValue(l, 1), 3.4); } - TEST(DenseLeafModelOperatorsTest, NormalUse) { TensorForestParams params; params.set_num_outputs(kNumClasses); @@ -182,7 +181,7 @@ TEST(SparseLeafModelOperatorsTest, InitWithExisting) { std::unique_ptr leaf(new Leaf); - op->ExportModel( *stat, leaf.get()); + op->ExportModel(*stat, leaf.get()); // Make sure it was initialized correctly. EXPECT_FLOAT_EQ(op->GetOutputValue(*leaf, 0), 1.1); @@ -194,7 +193,6 @@ TEST(SparseLeafModelOperatorsTest, InitWithExisting) { EXPECT_EQ(leaf->sparse_vector().sparse_value().size(), kNumClasses); } - TEST(RegressionLeafModelOperatorsTest, NormalUse) { TensorForestParams params; params.set_num_outputs(kNumClasses); diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/params.h b/tensorflow/contrib/tensor_forest/kernels/v4/params.h index 97a9d8d096311faaae774e9e4b2e45f28ed7fa29..7583e3d0402a3a1d07f3696727b285747dc887de 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/params.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/params.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_PARAMS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_PARAMS_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_PARAMS_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_PARAMS_H_ #include "tensorflow/contrib/tensor_forest/proto/tensor_forest_params.pb.h" #include "tensorflow/core/platform/types.h" @@ -24,9 +24,7 @@ namespace tensorforest { // Return the value of the given depth-dependent parameter given a leaf's depth. float ResolveParam(const DepthDependentParam& param, int32 depth); - } // namespace tensorforest } // namespace tensorflow - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_PARAMS_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_PARAMS_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/params_test.cc b/tensorflow/contrib/tensor_forest/kernels/v4/params_test.cc index 801881af1368dc33f00b356d12bea07ae3161ef6..4010a71006d58df0bec6d3686a9c47433b46fdd4 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/params_test.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/params_test.cc @@ -71,5 +71,3 @@ TEST(ParamsTest, TestThreshold) { } } // namespace - - diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/split_collection_operators.cc b/tensorflow/contrib/tensor_forest/kernels/v4/split_collection_operators.cc index cdb1d80a4bbd47d1481ecde2348bef500bd125f1..b7b60d0ab8c2670cec8b029d1f42c5edd3690afe 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/split_collection_operators.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/split_collection_operators.cc @@ -52,8 +52,8 @@ std::unique_ptr SplitCollectionOperator::CreateGrowStats( new SparseClassificationGrowStats(params_, depth)); case STATS_LEAST_SQUARES_REGRESSION: - return std::unique_ptr(new LeastSquaresRegressionGrowStats( - params_, depth)); + return std::unique_ptr( + new LeastSquaresRegressionGrowStats(params_, depth)); case STATS_FIXED_SIZE_SPARSE_GINI: return std::unique_ptr( @@ -136,8 +136,7 @@ void SplitCollectionOperator::CreateAndInitializeCandidateWithExample( stats_.at(node_id)->AddSplit(split, input_data, target, example); } -bool SplitCollectionOperator::BestSplit(int32 node_id, - SplitCandidate* best, +bool SplitCollectionOperator::BestSplit(int32 node_id, SplitCandidate* best, int32* depth) const { auto* slot = stats_.at(node_id).get(); *depth = slot->depth(); diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/split_collection_operators.h b/tensorflow/contrib/tensor_forest/kernels/v4/split_collection_operators.h index 6c21c0bd3443347bdb0102727b15b26754a0ed53..c606ff98c67f411a5817f0282238fdaf3be03642 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/split_collection_operators.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/split_collection_operators.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_SPLIT_COLLECTION_OPERATORS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_SPLIT_COLLECTION_OPERATORS_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_SPLIT_COLLECTION_OPERATORS_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_SPLIT_COLLECTION_OPERATORS_H_ #include #include "tensorflow/contrib/decision_trees/proto/generic_tree_model.pb.h" @@ -71,9 +71,7 @@ class SplitCollectionOperator { } // Perform any necessary cleanup for any tracked state for the slot. - virtual void ClearSlot(int32 node_id) { - stats_.erase(node_id); - } + virtual void ClearSlot(int32 node_id) { stats_.erase(node_id); } // Return true if slot is fully initialized. virtual bool IsInitialized(int32 node_id) const; @@ -128,6 +126,4 @@ class AnyCollectionCreator : public CollectionCreator { } // namespace tensorforest } // namespace tensorflow - - -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_SPLIT_COLLECTION_OPERATORS_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_SPLIT_COLLECTION_OPERATORS_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.cc b/tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.cc index 0bec198e97e8215d2cfdb9ada5355dd5b0d2d97b..c749fbe69e17769c2f2b69bcf541eb0eb8b9e7e8 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.cc +++ b/tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.cc @@ -32,9 +32,9 @@ namespace tensorforest { // smoothed_sum = stats.sum() + #_classes float GiniImpurity(const LeafStat& stats, int32 num_classes) { const float smoothed_sum = num_classes + stats.weight_sum(); - return 1.0 - ( - (stats.classification().gini().square() - + 2 * stats.weight_sum() + num_classes) / (smoothed_sum * smoothed_sum)); + return 1.0 - ((stats.classification().gini().square() + + 2 * stats.weight_sum() + num_classes) / + (smoothed_sum * smoothed_sum)); } float WeightedGiniImpurity(const LeafStat& stats, int32 num_classes) { @@ -46,21 +46,20 @@ void UpdateGini(LeafStat* stats, float old_val, float weight) { // Equivalent to stats->square() - old_val * old_val + new_val * new_val, // (for new_val = old_val + weight), but more numerically stable. stats->mutable_classification()->mutable_gini()->set_square( - stats->classification().gini().square() - + weight * weight + 2 * old_val * weight); + stats->classification().gini().square() + weight * weight + + 2 * old_val * weight); } - float Variance(const LeafStat& stats, int output) { if (stats.weight_sum() == 0) { return 0; } const float e_x = - stats.regression().mean_output().value(output).float_value() - / stats.weight_sum(); + stats.regression().mean_output().value(output).float_value() / + stats.weight_sum(); const auto e_x2 = - stats.regression().mean_output_squares().value(output).float_value() - / stats.weight_sum(); + stats.regression().mean_output_squares().value(output).float_value() / + stats.weight_sum(); return e_x2 - e_x * e_x; } @@ -75,8 +74,7 @@ float TotalVariance(const LeafStat& stats) { float SmoothedGini(float sum, float square, int num_classes) { // See comments for GiniImpurity above. const float smoothed_sum = num_classes + sum; - return 1.0 - - (square + 2 * sum + num_classes) / (smoothed_sum * smoothed_sum); + return 1.0 - (square + 2 * sum + num_classes) / (smoothed_sum * smoothed_sum); } float WeightedSmoothedGini(float sum, float square, int num_classes) { diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.h b/tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.h index 8e002d0414f48a1f409952f56c57b4e37815bca0..e6140065bbf12f2eb92c28e4affb3327f86af5d3 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/stat_utils.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_STAT_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_STAT_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_STAT_UTILS_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_STAT_UTILS_H_ #include "tensorflow/contrib/tensor_forest/proto/fertile_stats.pb.h" #include "tensorflow/core/platform/types.h" @@ -47,4 +47,4 @@ float WeightedSmoothedGini(float sum, float square, int num_classes); } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_STAT_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_STAT_UTILS_H_ diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/test_utils.h b/tensorflow/contrib/tensor_forest/kernels/v4/test_utils.h index b6e543b96fd5a00f78555eaf8558f0a95d0a6713..38deb3e3cd816aae5fe66f26cd4b934316d03ce4 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/test_utils.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/test_utils.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================= -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_TEST_UTILS_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_TEST_UTILS_H_ +#ifndef TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_TEST_UTILS_H_ +#define TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_TEST_UTILS_H_ #include "tensorflow/contrib/tensor_forest/kernels/v4/input_data.h" #include "tensorflow/contrib/tensor_forest/kernels/v4/input_target.h" @@ -27,9 +27,7 @@ class TestableInputTarget : public StoredInputTarget> { : StoredInputTarget(new std::vector(t), new std::vector(w), num_t) {} - int NumItems() const { - return target_->size(); - } + int NumItems() const { return target_->size(); } int32 GetTargetAsClassIndex(int example_index, int target_index) const override { @@ -51,7 +49,6 @@ class TestableInputTarget : public StoredInputTarget> { } }; - class TestableDataSet : public TensorDataSet { public: TestableDataSet(const std::vector& data, int num_features) @@ -71,4 +68,4 @@ class TestableDataSet : public TensorDataSet { } // namespace tensorforest } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_TEST_UTILS_H_ +#endif // TENSORFLOW_CONTRIB_TENSOR_FOREST_KERNELS_V4_TEST_UTILS_H_ diff --git a/tensorflow/contrib/tensor_forest/python/tensor_forest.py b/tensorflow/contrib/tensor_forest/python/tensor_forest.py index eb938763f12efd9281bec4321384acd4617cdfcf..3650b5d52fe8a1b87a239d41ecfa3de677fffc72 100644 --- a/tensorflow/contrib/tensor_forest/python/tensor_forest.py +++ b/tensorflow/contrib/tensor_forest/python/tensor_forest.py @@ -478,8 +478,7 @@ class RandomForestGraphs(object): **inference_args: Keyword arguments to pass through to each tree. Returns: - A tuple of (probabilities, tree_paths, variance), where variance - is the variance over all the trees for regression problems only. + A tuple of (probabilities, tree_paths, variance). Raises: NotImplementedError: If trying to use feature bagging with sparse @@ -513,13 +512,12 @@ class RandomForestGraphs(object): self.params.num_trees, name='probabilities') tree_paths = array_ops.stack(paths, axis=1) - regression_variance = None - if self.params.regression: - expected_squares = math_ops.div( - math_ops.reduce_sum(all_predict * all_predict, 1), - self.params.num_trees) - regression_variance = math_ops.maximum( - 0., expected_squares - average_values * average_values) + + expected_squares = math_ops.div( + math_ops.reduce_sum(all_predict * all_predict, 1), + self.params.num_trees) + regression_variance = math_ops.maximum( + 0., expected_squares - average_values * average_values) return average_values, tree_paths, regression_variance def average_size(self): diff --git a/tensorflow/contrib/tensor_forest/python/tensor_forest_test.py b/tensorflow/contrib/tensor_forest/python/tensor_forest_test.py index 113dfb85d3bf671e0a9448e0cb0fbfd7f3ea04e7..bbe627b15773fafe83a0700da696f429876c0968 100644 --- a/tensorflow/contrib/tensor_forest/python/tensor_forest_test.py +++ b/tensorflow/contrib/tensor_forest/python/tensor_forest_test.py @@ -108,7 +108,7 @@ class TensorForestTest(test_util.TensorFlowTestCase): probs, paths, var = graph_builder.inference_graph(input_data) self.assertTrue(isinstance(probs, ops.Tensor)) self.assertTrue(isinstance(paths, ops.Tensor)) - self.assertIsNone(var) + self.assertTrue(isinstance(var, ops.Tensor)) def testTrainingConstructionClassificationSparse(self): input_data = sparse_tensor.SparseTensor( diff --git a/tensorflow/contrib/tensorboard/db/BUILD b/tensorflow/contrib/tensorboard/db/BUILD index 9d3d60c24d72e28cf449cd196e34e53d5450d85f..4175d8adb58a85728519042a9870e8c4590232ba 100644 --- a/tensorflow/contrib/tensorboard/db/BUILD +++ b/tensorflow/contrib/tensorboard/db/BUILD @@ -5,12 +5,18 @@ package(default_visibility = ["//tensorflow:internal"]) licenses(["notice"]) # Apache 2.0 -load("//tensorflow:tensorflow.bzl", "tf_cc_test") +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_binary", + "tf_cc_test", + "tf_copts", +) cc_library( name = "schema", srcs = ["schema.cc"], hdrs = ["schema.h"], + copts = tf_copts(), deps = [ "//tensorflow/core:lib", "//tensorflow/core/lib/db:sqlite", @@ -19,6 +25,7 @@ cc_library( tf_cc_test( name = "schema_test", + size = "small", srcs = ["schema_test.cc"], deps = [ ":schema", @@ -31,8 +38,10 @@ cc_library( name = "summary_db_writer", srcs = ["summary_db_writer.cc"], hdrs = ["summary_db_writer.h"], + copts = tf_copts(), deps = [ - ":schema", + ":summary_converter", + "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", @@ -46,6 +55,7 @@ tf_cc_test( size = "small", srcs = ["summary_db_writer_test.cc"], deps = [ + ":schema", ":summary_db_writer", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", @@ -55,6 +65,77 @@ tf_cc_test( ], ) +cc_library( + name = "summary_file_writer", + srcs = ["summary_file_writer.cc"], + hdrs = ["summary_file_writer.h"], + copts = tf_copts(), + deps = [ + ":summary_converter", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:ptr_util", + "//tensorflow/core/kernels:summary_interface", + ], +) + +tf_cc_test( + name = "summary_file_writer_test", + size = "medium", # file i/o + timeout = "short", + srcs = ["summary_file_writer_test.cc"], + deps = [ + ":summary_file_writer", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +cc_library( + name = "summary_converter", + srcs = ["summary_converter.cc"], + hdrs = ["summary_converter.h"], + copts = tf_copts(), + visibility = ["//visibility:private"], + deps = [ + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + ], +) + +tf_cc_binary( + name = "loader", + srcs = ["loader.cc"], + linkstatic = 1, + deps = [ + ":schema", + ":summary_db_writer", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core/lib/db:sqlite", + ], +) + +tf_cc_binary( + name = "vacuum", + srcs = ["vacuum.cc"], + linkstatic = 1, + deps = [ + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core/lib/db:sqlite", + ], +) + filegroup( name = "all_files", srcs = glob(["*"]), diff --git a/tensorflow/contrib/tensorboard/db/loader.cc b/tensorflow/contrib/tensorboard/db/loader.cc new file mode 100644 index 0000000000000000000000000000000000000000..4d7337a53d025f29ae5f85151b7f60d2cca6f771 --- /dev/null +++ b/tensorflow/contrib/tensorboard/db/loader.cc @@ -0,0 +1,124 @@ +/* Copyright 2017 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 +#include + +#include "tensorflow/contrib/tensorboard/db/schema.h" +#include "tensorflow/contrib/tensorboard/db/summary_db_writer.h" +#include "tensorflow/core/lib/db/sqlite.h" +#include "tensorflow/core/lib/io/record_reader.h" +#include "tensorflow/core/platform/init_main.h" +#include "tensorflow/core/util/command_line_flags.h" +#include "tensorflow/core/util/event.pb.h" + +namespace tensorflow { +namespace { + +template +string AddCommas(T n) { + static_assert(std::is_integral::value, "is_integral"); + string s = strings::StrCat(n); + if (s.size() > 3) { + int extra = s.size() / 3 - (s.size() % 3 == 0 ? 1 : 0); + s.append(extra, 'X'); + int c = 0; + for (int i = s.size() - 1; i > 0; --i) { + s[i] = s[i - extra]; + if (++c % 3 == 0) { + s[--i] = ','; + --extra; + } + } + } + return s; +} + +int main(int argc, char* argv[]) { + string path; + string events; + string experiment_name; + string run_name; + string user_name; + std::vector flag_list = { + Flag("db", &path, "Path of SQLite DB file"), + Flag("events", &events, "TensorFlow record proto event log file"), + Flag("experiment_name", &experiment_name, "The DB experiment_name value"), + Flag("run_name", &run_name, "The DB run_name value"), + Flag("user_name", &user_name, "The DB user_name value"), + }; + string usage = Flags::Usage(argv[0], flag_list); + bool parse_result = Flags::Parse(&argc, argv, flag_list); + if (!parse_result || path.empty()) { + std::cerr << "The loader tool imports tf.Event record files, created by\n" + << "SummaryFileWriter, into the sorts of SQLite database files\n" + << "created by SummaryDbWriter.\n\n" + << "In addition to the flags below, the environment variables\n" + << "defined by core/lib/db/sqlite.cc can also be set.\n\n" + << usage; + return -1; + } + port::InitMain(argv[0], &argc, &argv); + Env* env = Env::Default(); + + LOG(INFO) << "Opening SQLite file: " << path; + Sqlite* db; + TF_CHECK_OK(Sqlite::Open( + path, SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_NOMUTEX, + &db)); + core::ScopedUnref unref_db(db); + + LOG(INFO) << "Initializing TensorBoard schema"; + TF_CHECK_OK(SetupTensorboardSqliteDb(db)); + + LOG(INFO) << "Creating SummaryDbWriter"; + SummaryWriterInterface* db_writer; + TF_CHECK_OK(CreateSummaryDbWriter(db, experiment_name, run_name, user_name, + env, &db_writer)); + core::ScopedUnref unref(db_writer); + + LOG(INFO) << "Loading TF event log: " << events; + std::unique_ptr file; + TF_CHECK_OK(env->NewRandomAccessFile(events, &file)); + io::RecordReader reader(file.get()); + + uint64 start = env->NowMicros(); + uint64 records = 0; + uint64 offset = 0; + string record; + while (true) { + std::unique_ptr event = std::unique_ptr(new Event); + Status s = reader.ReadRecord(&offset, &record); + if (s.code() == error::OUT_OF_RANGE) break; + TF_CHECK_OK(s); + if (!ParseProtoUnlimited(event.get(), record)) { + LOG(FATAL) << "Corrupt tf.Event record" + << " offset=" << (offset - record.size()) + << " size=" << static_cast(record.size()); + } + TF_CHECK_OK(db_writer->WriteEvent(std::move(event))); + ++records; + } + uint64 elapsed = env->NowMicros() - start; + LOG(INFO) << "Loaded " << AddCommas(offset) << " bytes with " + << AddCommas(records) << " records at " + << AddCommas(offset / (elapsed / 1000000)) << " bps"; + + return 0; +} + +} // namespace +} // namespace tensorflow + +int main(int argc, char* argv[]) { return tensorflow::main(argc, argv); } diff --git a/tensorflow/contrib/tensorboard/db/schema.cc b/tensorflow/contrib/tensorboard/db/schema.cc index fd024d692c3feddea2e5cbd29380686e8a0e9839..3c7bc87e4a2dbeadef2b9589d58c845204049123 100644 --- a/tensorflow/contrib/tensorboard/db/schema.cc +++ b/tensorflow/contrib/tensorboard/db/schema.cc @@ -14,486 +14,430 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/tensorboard/db/schema.h" +#include "tensorflow/core/lib/core/errors.h" + namespace tensorflow { namespace { -class SqliteSchema { - public: - explicit SqliteSchema(std::shared_ptr db) : db_(std::move(db)) {} - - /// \brief Creates Ids table. - /// - /// This table must be used to randomly allocate Permanent IDs for - /// all top-level tables, in order to maintain an invariant where - /// foo_id != bar_id for all IDs of any two tables. - /// - /// A row should only be deleted from this table if it can be - /// guaranteed that it exists absolutely nowhere else in the entire - /// system. - /// - /// Fields: - /// id: An ID that was allocated globally. This must be in the - /// range [1,2**47). 0 is assigned the same meaning as NULL and - /// shouldn't be stored; 2**63-1 is reserved for statically - /// allocating space in a page to UPDATE later; and all other - /// int64 values are reserved for future use. - Status CreateIdsTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Ids ( - id INTEGER PRIMARY KEY - ) - )sql"); - } - - /// \brief Creates Descriptions table. - /// - /// This table allows TensorBoard to associate Markdown text with any - /// object in the database that has a Permanent ID. - /// - /// Fields: - /// id: The Permanent ID of the associated object. This is also the - /// SQLite rowid. - /// description: Arbitrary Markdown text. - Status CreateDescriptionsTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Descriptions ( - id INTEGER PRIMARY KEY, - description TEXT - ) - )sql"); - } - - /// \brief Creates Tensors table. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// tag_id: ID of associated Tag. - /// computed_time: Float UNIX timestamp with microsecond precision. - /// In the old summaries system that uses FileWriter, this is the - /// wall time around when tf.Session.run finished. In the new - /// summaries system, it is the wall time of when the tensor was - /// computed. On systems with monotonic clocks, it is calculated - /// by adding the monotonic run duration to Run.started_time. - /// This field is not indexed because, in practice, it should be - /// ordered the same or nearly the same as TensorIndex, so local - /// insertion sort might be more suitable. - /// step: User-supplied number, ordering this tensor in Tag. - /// If NULL then the Tag must have only one Tensor. - /// tensor: Can be an INTEGER (DT_INT64), FLOAT (DT_DOUBLE), or - /// BLOB. The structure of a BLOB is currently undefined, but in - /// essence it is a Snappy tf.TensorProto that spills over into - /// TensorChunks. - Status CreateTensorsTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Tensors ( - rowid INTEGER PRIMARY KEY, - tag_id INTEGER NOT NULL, - computed_time REAL, - step INTEGER, - tensor BLOB - ) - )sql"); - } - - /// \brief Creates TensorChunks table. - /// - /// This table can be used to split up a tensor across many rows, - /// which has the advantage of not slowing down table scans on the - /// main table, allowing asynchronous fetching, minimizing copying, - /// and preventing large buffers from being allocated. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// tag_id: ID of associated Tag. - /// step: Same as corresponding Tensors.step. - /// sequence: 1-indexed sequence number for ordering chunks. Please - /// note that the 0th index is Tensors.tensor. - /// chunk: Bytes of next chunk in tensor. - Status CreateTensorChunksTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS TensorChunks ( - rowid INTEGER PRIMARY KEY, - tag_id INTEGER NOT NULL, - step INTEGER, - sequence INTEGER, - chunk BLOB - ) - )sql"); - } - - /// \brief Creates Tags table. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// tag_id: The Permanent ID of the Tag. - /// run_id: Optional ID of associated Run. - /// tag_name: The tag field in summary.proto, unique across Run. - /// inserted_time: Float UNIX timestamp with µs precision. This is - /// always the wall time of when the row was inserted into the - /// DB. It may be used as a hint for an archival job. - /// display_name: Optional for GUI and defaults to tag_name. - /// plugin_name: Arbitrary TensorBoard plugin name for dispatch. - /// plugin_data: Arbitrary data that plugin wants. - Status CreateTagsTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Tags ( - rowid INTEGER PRIMARY KEY, - run_id INTEGER, - tag_id INTEGER NOT NULL, - tag_name TEXT, - inserted_time DOUBLE, - display_name TEXT, - plugin_name TEXT, - plugin_data BLOB - ) - )sql"); - } - - /// \brief Creates Runs table. - /// - /// This table stores information about Runs. Each row usually - /// represents a single attempt at training or testing a TensorFlow - /// model, with a given set of hyper-parameters, whose summaries are - /// written out to a single event logs directory with a monotonic step - /// counter. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// run_id: The Permanent ID of the Run. This has a 1:1 mapping - /// with a SummaryWriter instance. If two writers spawn for a - /// given (user_name, run_name, run_name) then each should - /// allocate its own run_id and whichever writer puts it in the - /// database last wins. The Tags / Tensors associated with the - /// previous invocations will then enter limbo, where they may be - /// accessible for certain operations, but should be garbage - /// collected eventually. - /// experiment_id: Optional ID of associated Experiment. - /// run_name: User-supplied string, unique across Experiment. - /// inserted_time: Float UNIX timestamp with µs precision. This is - /// always the time the row was inserted into the database. It - /// does not change. - /// started_time: Float UNIX timestamp with µs precision. In the - /// old summaries system that uses FileWriter, this is - /// approximated as the first tf.Event.wall_time. In the new - /// summaries system, it is the wall time of when summary writing - /// started, from the perspective of whichever machine talks to - /// the database. This field will be mutated if the run is - /// restarted. - /// finished_time: Float UNIX timestamp with µs precision of when - /// SummaryWriter resource that created this run was destroyed. - /// Once this value becomes non-NULL a Run and its Tags and - /// Tensors should be regarded as immutable. - /// graph_id: ID of associated Graphs row. - Status CreateRunsTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Runs ( - rowid INTEGER PRIMARY KEY, - experiment_id INTEGER, - run_id INTEGER NOT NULL, - run_name TEXT, - inserted_time REAL, - started_time REAL, - finished_time REAL, - graph_id INTEGER - ) - )sql"); - } - - /// \brief Creates Experiments table. - /// - /// This table stores information about experiments, which are sets of - /// runs. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// user_id: Optional ID of associated User. - /// experiment_id: The Permanent ID of the Experiment. - /// experiment_name: User-supplied string, unique across User. - /// inserted_time: Float UNIX timestamp with µs precision. This is - /// always the time the row was inserted into the database. It - /// does not change. - /// started_time: Float UNIX timestamp with µs precision. This is - /// the MIN(experiment.started_time, run.started_time) of each - /// Run added to the database, including Runs which have since - /// been overwritten. - Status CreateExperimentsTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Experiments ( - rowid INTEGER PRIMARY KEY, - user_id INTEGER, - experiment_id INTEGER NOT NULL, - experiment_name TEXT, - inserted_time REAL, - started_time REAL - ) - )sql"); - } - - /// \brief Creates Users table. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// user_id: The Permanent ID of the User. - /// user_name: Unique user name. - /// email: Optional unique email address. - /// inserted_time: Float UNIX timestamp with µs precision. This is - /// always the time the row was inserted into the database. It - /// does not change. - Status CreateUsersTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Users ( - rowid INTEGER PRIMARY KEY, - user_id INTEGER NOT NULL, - user_name TEXT, - email TEXT, - inserted_time REAL - ) - )sql"); - } - - /// \brief Creates Graphs table. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// graph_id: The Permanent ID of the Graph. - /// inserted_time: Float UNIX timestamp with µs precision. This is - /// always the wall time of when the row was inserted into the - /// DB. It may be used as a hint for an archival job. - /// node_def: Contains Snappy tf.GraphDef proto. All fields will be - /// cleared except those not expressed in SQL. - Status CreateGraphsTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Graphs ( - rowid INTEGER PRIMARY KEY, - graph_id INTEGER NOT NULL, - inserted_time REAL, - graph_def BLOB - ) - )sql"); - } - - /// \brief Creates Nodes table. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// graph_id: The Permanent ID of the associated Graph. - /// node_id: ID for this node. This is more like a 0-index within - /// the Graph. Please note indexes are allowed to be removed. - /// node_name: Unique name for this Node within Graph. This is - /// copied from the proto so it can be indexed. This is allowed - /// to be NULL to save space on the index, in which case the - /// node_def.name proto field must not be cleared. - /// op: Copied from tf.NodeDef proto. - /// device: Copied from tf.NodeDef proto. - /// node_def: Contains Snappy tf.NodeDef proto. All fields will be - /// cleared except those not expressed in SQL. - Status CreateNodesTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS Nodes ( - rowid INTEGER PRIMARY KEY, - graph_id INTEGER NOT NULL, - node_id INTEGER NOT NULL, - node_name TEXT, - op TEXT, - device TEXT, - node_def BLOB - ) - )sql"); - } - - /// \brief Creates NodeInputs table. - /// - /// Fields: - /// rowid: Ephemeral b-tree ID dictating locality. - /// graph_id: The Permanent ID of the associated Graph. - /// node_id: Index of Node in question. This can be considered the - /// 'to' vertex. - /// idx: Used for ordering inputs on a given Node. - /// input_node_id: Nodes.node_id of the corresponding input node. - /// This can be considered the 'from' vertex. - /// is_control: If non-zero, indicates this input is a controlled - /// dependency, which means this isn't an edge through which - /// tensors flow. NULL means 0. - Status CreateNodeInputsTable() { - return Run(R"sql( - CREATE TABLE IF NOT EXISTS NodeInputs ( - rowid INTEGER PRIMARY KEY, - graph_id INTEGER NOT NULL, - node_id INTEGER NOT NULL, - idx INTEGER NOT NULL, - input_node_id INTEGER NOT NULL, - is_control INTEGER - ) - )sql"); - } - - /// \brief Uniquely indexes (tag_id, step) on Tensors table. - Status CreateTensorIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS TensorIndex - ON Tensors (tag_id, step) - )sql"); - } - - /// \brief Uniquely indexes (tag_id, step, sequence) on TensorChunks table. - Status CreateTensorChunkIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS TensorChunkIndex - ON TensorChunks (tag_id, step, sequence) - )sql"); - } - - /// \brief Uniquely indexes tag_id on Tags table. - Status CreateTagIdIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS TagIdIndex - ON Tags (tag_id) - )sql"); - } - - /// \brief Uniquely indexes run_id on Runs table. - Status CreateRunIdIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS RunIdIndex - ON Runs (run_id) - )sql"); - } - - /// \brief Uniquely indexes experiment_id on Experiments table. - Status CreateExperimentIdIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS ExperimentIdIndex - ON Experiments (experiment_id) - )sql"); - } - - /// \brief Uniquely indexes user_id on Users table. - Status CreateUserIdIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS UserIdIndex - ON Users (user_id) - )sql"); - } - - /// \brief Uniquely indexes graph_id on Graphs table. - Status CreateGraphIdIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS GraphIdIndex - ON Graphs (graph_id) - )sql"); - } - - /// \brief Uniquely indexes (graph_id, node_id) on Nodes table. - Status CreateNodeIdIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS NodeIdIndex - ON Nodes (graph_id, node_id) - )sql"); - } - - /// \brief Uniquely indexes (graph_id, node_id, idx) on NodeInputs table. - Status CreateNodeInputsIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS NodeInputsIndex - ON NodeInputs (graph_id, node_id, idx) - )sql"); - } - - /// \brief Uniquely indexes (run_id, tag_name) on Tags table. - Status CreateTagNameIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS TagNameIndex - ON Tags (run_id, tag_name) - WHERE tag_name IS NOT NULL - )sql"); - } - - /// \brief Uniquely indexes (experiment_id, run_name) on Runs table. - Status CreateRunNameIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS RunNameIndex - ON Runs (experiment_id, run_name) - WHERE run_name IS NOT NULL - )sql"); - } - - /// \brief Uniquely indexes (user_id, experiment_name) on Experiments table. - Status CreateExperimentNameIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS ExperimentNameIndex - ON Experiments (user_id, experiment_name) - WHERE experiment_name IS NOT NULL - )sql"); - } - - /// \brief Uniquely indexes user_name on Users table. - Status CreateUserNameIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS UserNameIndex - ON Users (user_name) - WHERE user_name IS NOT NULL - )sql"); - } - - /// \brief Uniquely indexes email on Users table. - Status CreateUserEmailIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS UserEmailIndex - ON Users (email) - WHERE email IS NOT NULL - )sql"); - } - - /// \brief Uniquely indexes (graph_id, node_name) on Nodes table. - Status CreateNodeNameIndex() { - return Run(R"sql( - CREATE UNIQUE INDEX IF NOT EXISTS NodeNameIndex - ON Nodes (graph_id, node_name) - WHERE node_name IS NOT NULL - )sql"); - } - - Status Run(const char* sql) { - auto stmt = db_->Prepare(sql); - TF_RETURN_WITH_CONTEXT_IF_ERROR(stmt.StepAndReset(), sql); - return Status::OK(); - } - - private: - std::shared_ptr db_; -}; +Status Run(Sqlite* db, const char* sql) { + SqliteStatement stmt; + TF_RETURN_IF_ERROR(db->Prepare(sql, &stmt)); + return stmt.StepAndReset(); +} } // namespace -Status SetupTensorboardSqliteDb(std::shared_ptr db) { - SqliteSchema s(std::move(db)); - TF_RETURN_IF_ERROR(s.CreateIdsTable()); - TF_RETURN_IF_ERROR(s.CreateDescriptionsTable()); - TF_RETURN_IF_ERROR(s.CreateTensorsTable()); - TF_RETURN_IF_ERROR(s.CreateTensorChunksTable()); - TF_RETURN_IF_ERROR(s.CreateTagsTable()); - TF_RETURN_IF_ERROR(s.CreateRunsTable()); - TF_RETURN_IF_ERROR(s.CreateExperimentsTable()); - TF_RETURN_IF_ERROR(s.CreateUsersTable()); - TF_RETURN_IF_ERROR(s.CreateGraphsTable()); - TF_RETURN_IF_ERROR(s.CreateNodeInputsTable()); - TF_RETURN_IF_ERROR(s.CreateNodesTable()); - TF_RETURN_IF_ERROR(s.CreateTensorIndex()); - TF_RETURN_IF_ERROR(s.CreateTensorChunkIndex()); - TF_RETURN_IF_ERROR(s.CreateTagIdIndex()); - TF_RETURN_IF_ERROR(s.CreateRunIdIndex()); - TF_RETURN_IF_ERROR(s.CreateExperimentIdIndex()); - TF_RETURN_IF_ERROR(s.CreateUserIdIndex()); - TF_RETURN_IF_ERROR(s.CreateGraphIdIndex()); - TF_RETURN_IF_ERROR(s.CreateNodeIdIndex()); - TF_RETURN_IF_ERROR(s.CreateNodeInputsIndex()); - TF_RETURN_IF_ERROR(s.CreateTagNameIndex()); - TF_RETURN_IF_ERROR(s.CreateRunNameIndex()); - TF_RETURN_IF_ERROR(s.CreateExperimentNameIndex()); - TF_RETURN_IF_ERROR(s.CreateUserNameIndex()); - TF_RETURN_IF_ERROR(s.CreateUserEmailIndex()); - TF_RETURN_IF_ERROR(s.CreateNodeNameIndex()); - return Status::OK(); +Status SetupTensorboardSqliteDb(Sqlite* db) { + // Note: GCC raw strings macros are broken. + // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=55971 + TF_RETURN_IF_ERROR( + db->PrepareOrDie(strings::StrCat("PRAGMA application_id=", + kTensorboardSqliteApplicationId)) + .StepAndReset()); + db->PrepareOrDie("PRAGMA user_version=0").StepAndResetOrDie(); + Status s; + + // Ids identify resources. + // + // This table can be used to efficiently generate Permanent IDs in + // conjunction with a random number generator. Unlike rowids these + // IDs safe to use in URLs and unique across tables. + // + // Within any given system, there can't be any foo_id == bar_id for + // all rows of any two (Foos, Bars) tables. A row should only be + // deleted from this table if there's a very high level of confidence + // it exists nowhere else in the system. + // + // Fields: + // id: The system-wide ID. This must be in the range [1,2**47). 0 + // is assigned the same meaning as NULL and shouldn't be stored + // and all other int64 values are reserved for future use. Please + // note that id is also the rowid. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Ids ( + id INTEGER PRIMARY KEY + ) + )sql")); + + // Descriptions are Markdown text that can be associated with any + // resource that has a Permanent ID. + // + // Fields: + // id: The foo_id of the associated row in Foos. + // description: Arbitrary NUL-terminated Markdown text. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Descriptions ( + id INTEGER PRIMARY KEY, + description TEXT + ) + )sql")); + + // Tensors are 0..n-dimensional numbers or strings. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // series: The Permanent ID of a different resource, e.g. tag_id. A + // tensor will be vacuumed if no series == foo_id exists for all + // rows of all Foos. When series is NULL this tensor may serve + // undefined purposes. This field should be set on placeholders. + // step: Arbitrary number to uniquely order tensors within series. + // The meaning of step is undefined when series is NULL. This may + // be set on placeholders to prepopulate index pages. + // computed_time: Float UNIX timestamp with microsecond precision. + // In the old summaries system that uses FileWriter, this is the + // wall time around when tf.Session.run finished. In the new + // summaries system, it is the wall time of when the tensor was + // computed. On systems with monotonic clocks, it is calculated + // by adding the monotonic run duration to Run.started_time. + // dtype: The tensorflow::DataType ID. For example, DT_INT64 is 9. + // When NULL or 0 this must be treated as a placeholder row that + // does not officially exist. + // shape: A comma-delimited list of int64 >=0 values representing + // length of each dimension in the tensor. This must be a valid + // shape. That means no -1 values and, in the case of numeric + // tensors, length(data) == product(shape) * sizeof(dtype). Empty + // means this is a scalar a.k.a. 0-dimensional tensor. + // data: Little-endian raw tensor memory. If dtype is DT_STRING and + // shape is empty, the nullness of this field indicates whether or + // not it contains the tensor contents; otherwise TensorStrings + // must be queried. If dtype is NULL then ZEROBLOB can be used on + // this field to reserve row space to be updated later. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Tensors ( + rowid INTEGER PRIMARY KEY, + series INTEGER, + step INTEGER, + dtype INTEGER, + computed_time REAL, + shape TEXT, + data BLOB + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS + TensorSeriesStepIndex + ON + Tensors (series, step) + WHERE + series IS NOT NULL + AND step IS NOT NULL + )sql")); + + // TensorStrings are the flat contents of 1..n dimensional DT_STRING + // Tensors. + // + // The number of rows associated with a Tensor must be equal to the + // product of its Tensors.shape. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // tensor_rowid: References Tensors.rowid. + // idx: Index in flattened tensor, starting at 0. + // data: The string value at a particular index. NUL characters are + // permitted. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS TensorStrings ( + rowid INTEGER PRIMARY KEY, + tensor_rowid INTEGER NOT NULL, + idx INTEGER NOT NULL, + data BLOB + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS TensorStringIndex + ON TensorStrings (tensor_rowid, idx) + )sql")); + + // Tags are series of Tensors. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // tag_id: The Permanent ID of the Tag. + // run_id: Optional ID of associated Run. + // inserted_time: Float UNIX timestamp with µs precision. This is + // always the wall time of when the row was inserted into the + // DB. It may be used as a hint for an archival job. + // tag_name: The tag field in summary.proto, unique across Run. + // display_name: Optional for GUI and defaults to tag_name. + // plugin_name: Arbitrary TensorBoard plugin name for dispatch. + // plugin_data: Arbitrary data that plugin wants. + // + // TODO(jart): Maybe there should be a Plugins table? + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Tags ( + rowid INTEGER PRIMARY KEY, + run_id INTEGER, + tag_id INTEGER NOT NULL, + inserted_time DOUBLE, + tag_name TEXT, + display_name TEXT, + plugin_name TEXT, + plugin_data BLOB + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS TagIdIndex + ON Tags (tag_id) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS + TagRunNameIndex + ON + Tags (run_id, tag_name) + WHERE + run_id IS NOT NULL + AND tag_name IS NOT NULL + )sql")); + + // Runs are groups of Tags. + // + // Each Run usually represents a single attempt at training or testing + // a TensorFlow model, with a given set of hyper-parameters, whose + // summaries are written out to a single event logs directory with a + // monotonic step counter. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // run_id: The Permanent ID of the Run. This has a 1:1 mapping + // with a SummaryWriter instance. If two writers spawn for a + // given (user_name, run_name, run_name) then each should + // allocate its own run_id and whichever writer puts it in the + // database last wins. The Tags / Tensors associated with the + // previous invocations will then enter limbo, where they may be + // accessible for certain operations, but should be garbage + // collected eventually. + // run_name: User-supplied string, unique across Experiment. + // experiment_id: Optional ID of associated Experiment. + // inserted_time: Float UNIX timestamp with µs precision. This is + // always the time the row was inserted into the database. It + // does not change. + // started_time: Float UNIX timestamp with µs precision. In the + // old summaries system that uses FileWriter, this is + // approximated as the first tf.Event.wall_time. In the new + // summaries system, it is the wall time of when summary writing + // started, from the perspective of whichever machine talks to + // the database. This field will be mutated if the run is + // restarted. + // finished_time: Float UNIX timestamp with µs precision of when + // SummaryWriter resource that created this run was destroyed. + // Once this value becomes non-NULL a Run and its Tags and + // Tensors should be regarded as immutable. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Runs ( + rowid INTEGER PRIMARY KEY, + experiment_id INTEGER, + run_id INTEGER NOT NULL, + inserted_time REAL, + started_time REAL, + finished_time REAL, + run_name TEXT + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS RunIdIndex + ON Runs (run_id) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS RunNameIndex + ON Runs (experiment_id, run_name) + WHERE run_name IS NOT NULL + )sql")); + + // Experiments are groups of Runs. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // user_id: Optional ID of associated User. + // experiment_id: The Permanent ID of the Experiment. + // experiment_name: User-supplied string, unique across User. + // inserted_time: Float UNIX timestamp with µs precision. This is + // always the time the row was inserted into the database. It + // does not change. + // started_time: Float UNIX timestamp with µs precision. This is + // the MIN(experiment.started_time, run.started_time) of each + // Run added to the database, including Runs which have since + // been overwritten. + // is_watching: A boolean indicating if someone is actively + // looking at this Experiment in the TensorBoard GUI. Tensor + // writers that do reservoir sampling can query this value to + // decide if they want the "keep last" behavior. This improves + // the performance of long running training while allowing low + // latency feedback in TensorBoard. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Experiments ( + rowid INTEGER PRIMARY KEY, + user_id INTEGER, + experiment_id INTEGER NOT NULL, + inserted_time REAL, + started_time REAL, + is_watching INTEGER, + experiment_name TEXT + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS ExperimentIdIndex + ON Experiments (experiment_id) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS ExperimentNameIndex + ON Experiments (user_id, experiment_name) + WHERE experiment_name IS NOT NULL + )sql")); + + // Users are people who love TensorBoard. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // user_id: The Permanent ID of the User. + // user_name: Unique user name. + // email: Optional unique email address. + // inserted_time: Float UNIX timestamp with µs precision. This is + // always the time the row was inserted into the database. It + // does not change. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Users ( + rowid INTEGER PRIMARY KEY, + user_id INTEGER NOT NULL, + inserted_time REAL, + user_name TEXT, + email TEXT + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS UserIdIndex + ON Users (user_id) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS UserNameIndex + ON Users (user_name) + WHERE user_name IS NOT NULL + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS UserEmailIndex + ON Users (email) + WHERE email IS NOT NULL + )sql")); + + // Graphs define how Tensors flowed in Runs. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // run_id: The Permanent ID of the associated Run. Only one Graph + // can be associated with a Run. + // graph_id: The Permanent ID of the Graph. + // inserted_time: Float UNIX timestamp with µs precision. This is + // always the wall time of when the row was inserted into the + // DB. It may be used as a hint for an archival job. + // graph_def: Contains the tf.GraphDef proto parts leftover which + // haven't been defined in SQL yet. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Graphs ( + rowid INTEGER PRIMARY KEY, + run_id INTEGER, + graph_id INTEGER NOT NULL, + inserted_time REAL, + graph_def BLOB + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS GraphIdIndex + ON Graphs (graph_id) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS GraphRunIndex + ON Graphs (run_id) + WHERE run_id IS NOT NULL + )sql")); + + // Nodes are the vertices in Graphs. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // graph_id: The Permanent ID of the associated Graph. + // node_id: ID for this node. This is more like a 0-index within + // the Graph. Please note indexes are allowed to be removed. + // node_name: Unique name for this Node within Graph. This is + // copied from the proto so it can be indexed. This is allowed + // to be NULL to save space on the index, in which case the + // node_def.name proto field must not be cleared. + // op: Copied from tf.NodeDef proto. + // device: Copied from tf.NodeDef proto. + // node_def: Contains the tf.NodeDef proto parts leftover which + // haven't been defined in SQL yet. + // + // TODO(jart): Make separate tables for op and device strings. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS Nodes ( + rowid INTEGER PRIMARY KEY, + graph_id INTEGER NOT NULL, + node_id INTEGER NOT NULL, + node_name TEXT, + op TEXT, + device TEXT, + node_def BLOB + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS NodeIdIndex + ON Nodes (graph_id, node_id) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS NodeNameIndex + ON Nodes (graph_id, node_name) + WHERE node_name IS NOT NULL + )sql")); + + // NodeInputs are directed edges between Nodes in Graphs. + // + // Fields: + // rowid: Ephemeral b-tree ID. + // graph_id: The Permanent ID of the associated Graph. + // node_id: Index of Node in question. This can be considered the + // 'to' vertex. + // idx: Used for ordering inputs on a given Node. + // input_node_id: Nodes.node_id of the corresponding input node. + // This can be considered the 'from' vertex. + // input_node_idx: Since a Node can output multiple Tensors, this + // is the integer index of which of those outputs is our input. + // NULL is treated as 0. + // is_control: If non-zero, indicates this input is a controlled + // dependency, which means this isn't an edge through which + // tensors flow. NULL means 0. + // + // TODO(jart): Rename to NodeEdges. + s.Update(Run(db, R"sql( + CREATE TABLE IF NOT EXISTS NodeInputs ( + rowid INTEGER PRIMARY KEY, + graph_id INTEGER NOT NULL, + node_id INTEGER NOT NULL, + idx INTEGER NOT NULL, + input_node_id INTEGER NOT NULL, + input_node_idx INTEGER, + is_control INTEGER + ) + )sql")); + + s.Update(Run(db, R"sql( + CREATE UNIQUE INDEX IF NOT EXISTS NodeInputsIndex + ON NodeInputs (graph_id, node_id, idx) + )sql")); + + return s; } } // namespace tensorflow diff --git a/tensorflow/contrib/tensorboard/db/schema.h b/tensorflow/contrib/tensorboard/db/schema.h index 900c10298ce0a69b92f7528db9742517243c3c51..3da450422523dbe4304446869a38d43981d76eb5 100644 --- a/tensorflow/contrib/tensorboard/db/schema.h +++ b/tensorflow/contrib/tensorboard/db/schema.h @@ -15,18 +15,18 @@ limitations under the License. #ifndef TENSORFLOW_CONTRIB_TENSORBOARD_DB_SCHEMA_H_ #define TENSORFLOW_CONTRIB_TENSORBOARD_DB_SCHEMA_H_ -#include - #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/db/sqlite.h" namespace tensorflow { +constexpr uint32 kTensorboardSqliteApplicationId = 0xfeedabee; + /// \brief Creates TensorBoard SQLite tables and indexes. /// /// If they are already created, this has no effect. If schema /// migrations are necessary, they will be performed with logging. -Status SetupTensorboardSqliteDb(std::shared_ptr db); +Status SetupTensorboardSqliteDb(Sqlite* db); } // namespace tensorflow diff --git a/tensorflow/contrib/tensorboard/db/schema_test.cc b/tensorflow/contrib/tensorboard/db/schema_test.cc index 463c4e59e7e76e6460b7ddfbd92262ac249aa9ed..4d3f2880bd02682ad00a90760f2a4478f1e6b2a2 100644 --- a/tensorflow/contrib/tensorboard/db/schema_test.cc +++ b/tensorflow/contrib/tensorboard/db/schema_test.cc @@ -23,7 +23,9 @@ namespace tensorflow { namespace { TEST(SchemaTest, SmokeTestTensorboardSchema) { - auto db = Sqlite::Open(":memory:").ValueOrDie(); + Sqlite* db; + TF_ASSERT_OK(Sqlite::Open(":memory:", SQLITE_OPEN_READWRITE, &db)); + core::ScopedUnref unref_db(db); TF_ASSERT_OK(SetupTensorboardSqliteDb(db)); } diff --git a/tensorflow/contrib/tensorboard/db/summary_converter.cc b/tensorflow/contrib/tensorboard/db/summary_converter.cc new file mode 100644 index 0000000000000000000000000000000000000000..93c1183072b4d791843e740f970234ba52857463 --- /dev/null +++ b/tensorflow/contrib/tensorboard/db/summary_converter.cc @@ -0,0 +1,322 @@ +/* Copyright 2017 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/tensorboard/db/summary_converter.h" + +#include "tensorflow/core/framework/register_types.h" +#include "tensorflow/core/framework/summary.pb.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/lib/histogram/histogram.h" +#include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/lib/png/png_io.h" +#include "tensorflow/core/lib/wav/wav_io.h" + +namespace tensorflow { +namespace { + +template +Status TensorValueAt(Tensor t, int64 i, T* out) { +#define CASE(I) \ + case DataTypeToEnum::value: \ + *out = static_cast(t.flat()(i)); \ + break; +#define COMPLEX_CASE(I) \ + case DataTypeToEnum::value: \ + *out = static_cast(t.flat()(i).real()); \ + break; + // clang-format off + switch (t.dtype()) { + TF_CALL_half(CASE) + TF_CALL_float(CASE) + TF_CALL_double(CASE) + TF_CALL_int8(CASE) + TF_CALL_int16(CASE) + TF_CALL_int32(CASE) + TF_CALL_int64(CASE) + TF_CALL_uint8(CASE) + TF_CALL_uint16(CASE) + TF_CALL_uint32(CASE) + TF_CALL_uint64(CASE) + TF_CALL_complex64(COMPLEX_CASE) + TF_CALL_complex128(COMPLEX_CASE) + default: + return errors::Unimplemented("SummaryFileWriter ", + DataTypeString(t.dtype()), + " not supported."); + } + // clang-format on + return Status::OK(); +#undef CASE +#undef COMPLEX_CASE +} + +typedef Eigen::Tensor Uint8Image; + +// Add the sequence of images specified by ith_image to the summary. +// +// Factoring this loop out into a helper function lets ith_image behave +// differently in the float and uint8 cases: the float case needs a temporary +// buffer which can be shared across calls to ith_image, but the uint8 case +// does not. +Status AddImages(const string& tag, int max_images, int batch_size, int w, + int h, int depth, + const std::function& ith_image, Summary* s) { + const int N = std::min(max_images, batch_size); + for (int i = 0; i < N; ++i) { + Summary::Value* v = s->add_value(); + // The tag depends on the number of requested images (not the number + // produced.) + // + // Note that later on avisu uses "/" to figure out a consistent naming + // convention for display, so we append "/image" to guarantee that the + // image(s) won't be displayed in the global scope with no name. + if (max_images > 1) { + v->set_tag(strings::StrCat(tag, "/image/", i)); + } else { + v->set_tag(strings::StrCat(tag, "/image")); + } + + const auto image = ith_image(i); + Summary::Image* si = v->mutable_image(); + si->set_height(h); + si->set_width(w); + si->set_colorspace(depth); + const int channel_bits = 8; + const int compression = -1; // Use zlib default + if (!png::WriteImageToBuffer(image.data(), w, h, w * depth, depth, + channel_bits, compression, + si->mutable_encoded_image_string(), nullptr)) { + return errors::Internal("PNG encoding failed"); + } + } + return Status::OK(); +} + +template +void NormalizeFloatImage(int hw, int depth, + typename TTypes::ConstMatrix values, + typename TTypes::ConstVec bad_color, + Uint8Image* image) { + if (!image->size()) return; // Nothing to do for empty images + + // Rescale the image to uint8 range. + // + // We are trying to generate an RGB image from a float/half tensor. We do + // not have any info about the expected range of values in the tensor + // but the generated image needs to have all RGB values within [0, 255]. + // + // We use two different algorithms to generate these values. If the + // tensor has only positive values we scale them all by 255/max(values). + // If the tensor has both negative and positive values we scale them by + // the max of their absolute values and center them around 127. + // + // This works for most cases, but does not respect the relative dynamic + // range across different instances of the tensor. + + // Compute min and max ignoring nonfinite pixels + float image_min = std::numeric_limits::infinity(); + float image_max = -image_min; + for (int i = 0; i < hw; i++) { + bool finite = true; + for (int j = 0; j < depth; j++) { + if (!Eigen::numext::isfinite(values(i, j))) { + finite = false; + break; + } + } + if (finite) { + for (int j = 0; j < depth; j++) { + float value(values(i, j)); + image_min = std::min(image_min, value); + image_max = std::max(image_max, value); + } + } + } + + // Pick an affine transform into uint8 + const float kZeroThreshold = 1e-6; + T scale, offset; + if (image_min < 0) { + const float max_val = std::max(std::abs(image_min), std::abs(image_max)); + scale = T(max_val < kZeroThreshold ? 0.0f : 127.0f / max_val); + offset = T(128.0f); + } else { + scale = T(image_max < kZeroThreshold ? 0.0f : 255.0f / image_max); + offset = T(0.0f); + } + + // Transform image, turning nonfinite values to bad_color + for (int i = 0; i < hw; i++) { + bool finite = true; + for (int j = 0; j < depth; j++) { + if (!Eigen::numext::isfinite(values(i, j))) { + finite = false; + break; + } + } + if (finite) { + image->chip<0>(i) = + (values.template chip<0>(i) * scale + offset).template cast(); + } else { + image->chip<0>(i) = bad_color; + } + } +} + +template +Status NormalizeAndAddImages(const Tensor& tensor, int max_images, int h, int w, + int hw, int depth, int batch_size, + const string& base_tag, Tensor bad_color_tensor, + Summary* s) { + // For float and half images, nans and infs are replaced with bad_color. + if (bad_color_tensor.dim_size(0) < depth) { + return errors::InvalidArgument( + "expected depth <= bad_color.size, got depth = ", depth, + ", bad_color.size = ", bad_color_tensor.dim_size(0)); + } + auto bad_color_full = bad_color_tensor.vec(); + typename TTypes::ConstVec bad_color(bad_color_full.data(), depth); + + // Float images must be scaled and translated. + Uint8Image image(hw, depth); + auto ith_image = [&tensor, &image, bad_color, batch_size, hw, depth](int i) { + auto tensor_eigen = tensor.template shaped({batch_size, hw, depth}); + typename TTypes::ConstMatrix values( + &tensor_eigen(i, 0, 0), Eigen::DSizes(hw, depth)); + NormalizeFloatImage(hw, depth, values, bad_color, &image); + return image; + }; + return AddImages(base_tag, max_images, batch_size, w, h, depth, ith_image, s); +} + +} // namespace + +Status AddTensorAsScalarToSummary(const Tensor& t, const string& tag, + Summary* s) { + Summary::Value* v = s->add_value(); + v->set_tag(tag); + float value; + TF_RETURN_IF_ERROR(TensorValueAt(t, 0, &value)); + v->set_simple_value(value); + return Status::OK(); +} + +Status AddTensorAsHistogramToSummary(const Tensor& t, const string& tag, + Summary* s) { + Summary::Value* v = s->add_value(); + v->set_tag(tag); + histogram::Histogram histo; + for (int64 i = 0; i < t.NumElements(); i++) { + double double_val; + TF_RETURN_IF_ERROR(TensorValueAt(t, i, &double_val)); + if (Eigen::numext::isnan(double_val)) { + return errors::InvalidArgument("Nan in summary histogram for: ", tag); + } else if (Eigen::numext::isinf(double_val)) { + return errors::InvalidArgument("Infinity in summary histogram for: ", + tag); + } + histo.Add(double_val); + } + histo.EncodeToProto(v->mutable_histo(), false /* Drop zero buckets */); + return Status::OK(); +} + +Status AddTensorAsImageToSummary(const Tensor& tensor, const string& tag, + int max_images, const Tensor& bad_color, + Summary* s) { + if (!(tensor.dims() == 4 && + (tensor.dim_size(3) == 1 || tensor.dim_size(3) == 3 || + tensor.dim_size(3) == 4))) { + return errors::InvalidArgument( + "Tensor must be 4-D with last dim 1, 3, or 4, not ", + tensor.shape().DebugString()); + } + if (!(tensor.dim_size(0) < (1LL << 31) && tensor.dim_size(1) < (1LL << 31) && + tensor.dim_size(2) < (1LL << 31) && + (tensor.dim_size(1) * tensor.dim_size(2)) < (1LL << 29))) { + return errors::InvalidArgument("Tensor too large for summary ", + tensor.shape().DebugString()); + } + // The casts and h * w cannot overflow because of the limits above. + const int batch_size = static_cast(tensor.dim_size(0)); + const int h = static_cast(tensor.dim_size(1)); + const int w = static_cast(tensor.dim_size(2)); + const int hw = h * w; // Compact these two dims for simplicity + const int depth = static_cast(tensor.dim_size(3)); + if (tensor.dtype() == DT_UINT8) { + // For uint8 input, no normalization is necessary + auto ith_image = [&tensor, batch_size, hw, depth](int i) { + auto values = tensor.shaped({batch_size, hw, depth}); + return typename TTypes::ConstMatrix( + &values(i, 0, 0), Eigen::DSizes(hw, depth)); + }; + TF_RETURN_IF_ERROR( + AddImages(tag, max_images, batch_size, w, h, depth, ith_image, s)); + } else if (tensor.dtype() == DT_HALF) { + TF_RETURN_IF_ERROR(NormalizeAndAddImages( + tensor, max_images, h, w, hw, depth, batch_size, tag, bad_color, s)); + } else if (tensor.dtype() == DT_FLOAT) { + TF_RETURN_IF_ERROR(NormalizeAndAddImages( + tensor, max_images, h, w, hw, depth, batch_size, tag, bad_color, s)); + } else { + return errors::InvalidArgument( + "Only DT_INT8, DT_HALF, and DT_FLOAT images are supported. Got ", + DataTypeString(tensor.dtype())); + } + return Status::OK(); +} + +Status AddTensorAsAudioToSummary(const Tensor& tensor, const string& tag, + int max_outputs, float sample_rate, + Summary* s) { + if (sample_rate <= 0.0f) { + return errors::InvalidArgument("sample_rate must be > 0"); + } + const int batch_size = tensor.dim_size(0); + const int64 length_frames = tensor.dim_size(1); + const int64 num_channels = + tensor.dims() == 2 ? 1 : tensor.dim_size(tensor.dims() - 1); + const int N = std::min(max_outputs, batch_size); + for (int i = 0; i < N; ++i) { + Summary::Value* v = s->add_value(); + if (max_outputs > 1) { + v->set_tag(strings::StrCat(tag, "/audio/", i)); + } else { + v->set_tag(strings::StrCat(tag, "/audio")); + } + + Summary::Audio* sa = v->mutable_audio(); + sa->set_sample_rate(sample_rate); + sa->set_num_channels(num_channels); + sa->set_length_frames(length_frames); + sa->set_content_type("audio/wav"); + + auto values = + tensor.shaped({batch_size, length_frames, num_channels}); + auto channels_by_frames = typename TTypes::ConstMatrix( + &values(i, 0, 0), + Eigen::DSizes(length_frames, num_channels)); + size_t sample_rate_truncated = lrintf(sample_rate); + if (sample_rate_truncated == 0) { + sample_rate_truncated = 1; + } + TF_RETURN_IF_ERROR(wav::EncodeAudioAsS16LEWav( + channels_by_frames.data(), sample_rate_truncated, num_channels, + length_frames, sa->mutable_encoded_audio_string())); + } + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/contrib/tensorboard/db/summary_converter.h b/tensorflow/contrib/tensorboard/db/summary_converter.h new file mode 100644 index 0000000000000000000000000000000000000000..329c7f9f2f9fe25cdff8d5ac2e52c25362f624c2 --- /dev/null +++ b/tensorflow/contrib/tensorboard/db/summary_converter.h @@ -0,0 +1,38 @@ +/* Copyright 2017 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_TENSORBOARD_DB_SUMMARY_CONVERTER_H_ +#define TENSORFLOW_CONTRIB_TENSORBOARD_DB_SUMMARY_CONVERTER_H_ + +#include "tensorflow/core/framework/summary.pb.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { + +// TODO(jart): Delete these methods in favor of new Python implementation. +Status AddTensorAsScalarToSummary(const Tensor& t, const string& tag, + Summary* s); +Status AddTensorAsHistogramToSummary(const Tensor& t, const string& tag, + Summary* s); +Status AddTensorAsImageToSummary(const Tensor& tensor, const string& tag, + int max_images, const Tensor& bad_color, + Summary* s); +Status AddTensorAsAudioToSummary(const Tensor& tensor, const string& tag, + int max_outputs, float sample_rate, + Summary* s); + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_TENSORBOARD_DB_SUMMARY_CONVERTER_H_ diff --git a/tensorflow/contrib/tensorboard/db/summary_db_writer.cc b/tensorflow/contrib/tensorboard/db/summary_db_writer.cc index 04b9c8e457bd52ff476ed8b13ff9608bdc8a933e..6590d6f7df4f35cad78db1fa9c4407bfb1270a2f 100644 --- a/tensorflow/contrib/tensorboard/db/summary_db_writer.cc +++ b/tensorflow/contrib/tensorboard/db/summary_db_writer.cc @@ -14,18 +14,38 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/tensorboard/db/summary_db_writer.h" -#include "tensorflow/contrib/tensorboard/db/schema.h" +#include "tensorflow/contrib/tensorboard/db/summary_converter.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/summary.pb.h" #include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/core/lib/db/sqlite.h" #include "tensorflow/core/lib/random/random.h" -#include "tensorflow/core/lib/strings/stringprintf.h" -#include "tensorflow/core/platform/fingerprint.h" -#include "tensorflow/core/platform/snappy.h" #include "tensorflow/core/util/event.pb.h" +// TODO(jart): Break this up into multiple files with excellent unit tests. +// TODO(jart): Make decision to write in separate op. +// TODO(jart): Add really good busy handling. + +// clang-format off +#define CALL_SUPPORTED_TYPES(m) \ + TF_CALL_string(m) \ + TF_CALL_half(m) \ + TF_CALL_float(m) \ + TF_CALL_double(m) \ + TF_CALL_complex64(m) \ + TF_CALL_complex128(m) \ + TF_CALL_int8(m) \ + TF_CALL_int16(m) \ + TF_CALL_int32(m) \ + TF_CALL_int64(m) \ + TF_CALL_uint8(m) \ + TF_CALL_uint16(m) \ + TF_CALL_uint32(m) \ + TF_CALL_uint64(m) +// clang-format on + namespace tensorflow { namespace { @@ -34,125 +54,145 @@ const uint64 kIdTiers[] = { 0x7fffffULL, // 23-bit (3 bytes on disk) 0x7fffffffULL, // 31-bit (4 bytes on disk) 0x7fffffffffffULL, // 47-bit (5 bytes on disk) - // Remaining bits reserved for future use. + // remaining bits for future use }; const int kMaxIdTier = sizeof(kIdTiers) / sizeof(uint64); const int kIdCollisionDelayMicros = 10; const int kMaxIdCollisions = 21; // sum(2**i*10µs for i in range(21))~=21s const int64 kAbsent = 0LL; -const int64 kReserved = 0x7fffffffffffffffLL; -double GetWallTime(Env* env) { +const char* kScalarPluginName = "scalars"; +const char* kImagePluginName = "images"; +const char* kAudioPluginName = "audio"; +const char* kHistogramPluginName = "histograms"; + +const int kScalarSlots = 10000; +const int kImageSlots = 10; +const int kAudioSlots = 10; +const int kHistogramSlots = 1; +const int kTensorSlots = 10; + +const int64 kReserveMinBytes = 32; +const double kReserveMultiplier = 1.5; + +// Flush is a misnomer because what we're actually doing is having lots +// of commits inside any SqliteTransaction that writes potentially +// hundreds of megs but doesn't need the transaction to maintain its +// invariants. This ensures the WAL read penalty is small and might +// allow writers in other processes a chance to schedule. +const uint64 kFlushBytes = 1024 * 1024; + +double DoubleTime(uint64 micros) { // TODO(@jart): Follow precise definitions for time laid out in schema. // TODO(@jart): Use monotonic clock from gRPC codebase. - return static_cast(env->NowMicros()) / 1.0e6; + return static_cast(micros) / 1.0e6; } -Status Serialize(const protobuf::MessageLite& proto, string* output) { - output->clear(); - if (!proto.SerializeToString(output)) { - return errors::DataLoss("SerializeToString failed"); +string StringifyShape(const TensorShape& shape) { + string result; + bool first = true; + for (const auto& dim : shape) { + if (first) { + first = false; + } else { + strings::StrAppend(&result, ","); + } + strings::StrAppend(&result, dim.size); } - return Status::OK(); + return result; } -Status Compress(const string& data, string* output) { - output->clear(); - if (!port::Snappy_Compress(data.data(), data.size(), output)) { - return errors::FailedPrecondition("TensorBase needs Snappy"); +Status CheckSupportedType(const Tensor& t) { +#define CASE(T) \ + case DataTypeToEnum::value: \ + break; + switch (t.dtype()) { + CALL_SUPPORTED_TYPES(CASE) + default: + return errors::Unimplemented(DataTypeString(t.dtype()), + " tensors unsupported on platform"); } return Status::OK(); +#undef CASE } -Status BindProto(SqliteStatement* stmt, int parameter, - const protobuf::MessageLite& proto) { - string serialized; - TF_RETURN_IF_ERROR(Serialize(proto, &serialized)); - string compressed; - TF_RETURN_IF_ERROR(Compress(serialized, &compressed)); - stmt->BindBlob(parameter, compressed); - return Status::OK(); +Tensor AsScalar(const Tensor& t) { + Tensor t2{t.dtype(), {}}; +#define CASE(T) \ + case DataTypeToEnum::value: \ + t2.scalar()() = t.flat()(0); \ + break; + switch (t.dtype()) { + CALL_SUPPORTED_TYPES(CASE) + default: + t2 = {DT_FLOAT, {}}; + t2.scalar()() = NAN; + break; + } + return t2; +#undef CASE } -Status BindTensor(SqliteStatement* stmt, int parameter, const Tensor& t) { - // TODO(@jart): Make portable between little and big endian systems. - // TODO(@jart): Use TensorChunks with minimal copying for big tensors. - // TODO(@jart): Add field to indicate encoding. - // TODO(@jart): Allow crunch tool to re-compress with zlib instead. - TensorProto p; - t.AsProtoTensorContent(&p); - return BindProto(stmt, parameter, p); +void PatchPluginName(SummaryMetadata* metadata, const char* name) { + if (metadata->plugin_data().plugin_name().empty()) { + metadata->mutable_plugin_data()->set_plugin_name(name); + } } -// Tries to fudge shape and dtype to something with smaller storage. -Status CoerceScalar(const Tensor& t, Tensor* out) { - switch (t.dtype()) { - case DT_DOUBLE: - *out = t; - break; - case DT_INT64: - *out = t; - break; - case DT_FLOAT: - *out = {DT_DOUBLE, {}}; - out->scalar()() = t.scalar()(); - break; - case DT_HALF: - *out = {DT_DOUBLE, {}}; - out->scalar()() = static_cast(t.scalar()()); - break; - case DT_INT32: - *out = {DT_INT64, {}}; - out->scalar()() = t.scalar()(); - break; - case DT_INT16: - *out = {DT_INT64, {}}; - out->scalar()() = t.scalar()(); - break; - case DT_INT8: - *out = {DT_INT64, {}}; - out->scalar()() = t.scalar()(); - break; - case DT_UINT32: - *out = {DT_INT64, {}}; - out->scalar()() = t.scalar()(); - break; - case DT_UINT16: - *out = {DT_INT64, {}}; - out->scalar()() = t.scalar()(); - break; - case DT_UINT8: - *out = {DT_INT64, {}}; - out->scalar()() = t.scalar()(); - break; - default: - return errors::Unimplemented("Scalar summary for dtype ", - DataTypeString(t.dtype()), - " is not supported."); +int GetSlots(const Tensor& t, const SummaryMetadata& metadata) { + if (metadata.plugin_data().plugin_name() == kScalarPluginName) { + return kScalarSlots; + } else if (metadata.plugin_data().plugin_name() == kImagePluginName) { + return kImageSlots; + } else if (metadata.plugin_data().plugin_name() == kAudioPluginName) { + return kAudioSlots; + } else if (metadata.plugin_data().plugin_name() == kHistogramPluginName) { + return kHistogramSlots; + } else if (t.dims() == 0 && t.dtype() != DT_STRING) { + return kScalarSlots; + } else { + return kTensorSlots; } - return Status::OK(); } -/// \brief Generates unique IDs randomly in the [1,2**63-2] range. +Status SetDescription(Sqlite* db, int64 id, const StringPiece& markdown) { + const char* sql = R"sql( + INSERT OR REPLACE INTO Descriptions (id, description) VALUES (?, ?) + )sql"; + SqliteStatement insert_desc; + TF_RETURN_IF_ERROR(db->Prepare(sql, &insert_desc)); + insert_desc.BindInt(1, id); + insert_desc.BindText(2, markdown); + return insert_desc.StepAndReset(); +} + +/// \brief Generates unique IDs randomly in the [1,2**63-1] range. /// /// This class starts off generating IDs in the [1,2**23-1] range, /// because it's human friendly and occupies 4 bytes max on disk with /// SQLite's zigzag varint encoding. Then, each time a collision /// happens, the random space is increased by 8 bits. /// -/// This class uses exponential back-off so writes will slow down as -/// the ID space becomes exhausted. +/// This class uses exponential back-off so writes gradually slow down +/// as IDs become exhausted but reads are still possible. +/// +/// This class is thread safe. class IdAllocator { public: - IdAllocator(Env* env, Sqlite* db) - : env_{env}, inserter_{db->Prepare("INSERT INTO Ids (id) VALUES (?)")} {} + IdAllocator(Env* env, Sqlite* db) : env_{env}, db_{db} { + DCHECK(env_ != nullptr); + DCHECK(db_ != nullptr); + } - Status CreateNewId(int64* id) { + Status CreateNewId(int64* id) LOCKS_EXCLUDED(mu_) { + mutex_lock lock(mu_); Status s; + SqliteStatement stmt; + TF_RETURN_IF_ERROR(db_->Prepare("INSERT INTO Ids (id) VALUES (?)", &stmt)); for (int i = 0; i < kMaxIdCollisions; ++i) { int64 tid = MakeRandomId(); - inserter_.BindInt(1, tid); - s = inserter_.StepAndReset(); + stmt.BindInt(1, tid); + s = stmt.StepAndReset(); if (s.ok()) { *id = tid; break; @@ -178,34 +218,38 @@ class IdAllocator { } private: - int64 MakeRandomId() { + int64 MakeRandomId() EXCLUSIVE_LOCKS_REQUIRED(mu_) { int64 id = static_cast(random::New64() & kIdTiers[tier_]); if (id == kAbsent) ++id; - if (id == kReserved) --id; return id; } - Env* env_; - SqliteStatement inserter_; - int tier_ = 0; + mutex mu_; + Env* const env_; + Sqlite* const db_; + int tier_ GUARDED_BY(mu_) = 0; + + TF_DISALLOW_COPY_AND_ASSIGN(IdAllocator); }; -class GraphSaver { +class GraphWriter { public: - static Status Save(Env* env, Sqlite* db, IdAllocator* id_allocator, - GraphDef* graph, int64* graph_id) { - TF_RETURN_IF_ERROR(id_allocator->CreateNewId(graph_id)); - GraphSaver saver{env, db, graph, *graph_id}; + static Status Save(Sqlite* db, SqliteTransaction* txn, IdAllocator* ids, + GraphDef* graph, uint64 now, int64 run_id, int64* graph_id) + SQLITE_EXCLUSIVE_TRANSACTIONS_REQUIRED(*db) { + TF_RETURN_IF_ERROR(ids->CreateNewId(graph_id)); + GraphWriter saver{db, txn, graph, now, *graph_id}; saver.MapNameToNodeId(); - TF_RETURN_IF_ERROR(saver.SaveNodeInputs()); - TF_RETURN_IF_ERROR(saver.SaveNodes()); - TF_RETURN_IF_ERROR(saver.SaveGraph()); + TF_RETURN_WITH_CONTEXT_IF_ERROR(saver.SaveNodeInputs(), "SaveNodeInputs"); + TF_RETURN_WITH_CONTEXT_IF_ERROR(saver.SaveNodes(), "SaveNodes"); + TF_RETURN_WITH_CONTEXT_IF_ERROR(saver.SaveGraph(run_id), "SaveGraph"); return Status::OK(); } private: - GraphSaver(Env* env, Sqlite* db, GraphDef* graph, int64 graph_id) - : env_(env), db_(db), graph_(graph), graph_id_(graph_id) {} + GraphWriter(Sqlite* db, SqliteTransaction* txn, GraphDef* graph, uint64 now, + int64 graph_id) + : db_(db), txn_(txn), graph_(graph), now_(now), graph_id_(graph_id) {} void MapNameToNodeId() { size_t toto = static_cast(graph_->node_size()); @@ -220,156 +264,193 @@ class GraphSaver { } Status SaveNodeInputs() { - auto insert = db_->Prepare(R"sql( - INSERT INTO NodeInputs (graph_id, node_id, idx, input_node_id, is_control) - VALUES (?, ?, ?, ?, ?) - )sql"); + const char* sql = R"sql( + INSERT INTO NodeInputs ( + graph_id, + node_id, + idx, + input_node_id, + input_node_idx, + is_control + ) VALUES (?, ?, ?, ?, ?, ?) + )sql"; + SqliteStatement insert; + TF_RETURN_IF_ERROR(db_->Prepare(sql, &insert)); for (int node_id = 0; node_id < graph_->node_size(); ++node_id) { const NodeDef& node = graph_->node(node_id); for (int idx = 0; idx < node.input_size(); ++idx) { StringPiece name = node.input(idx); - insert.BindInt(1, graph_id_); - insert.BindInt(2, node_id); - insert.BindInt(3, idx); + int64 input_node_id; + int64 input_node_idx = 0; + int64 is_control = 0; + size_t i = name.rfind(':'); + if (i != StringPiece::npos) { + if (!strings::safe_strto64(name.substr(i + 1, name.size() - i - 1), + &input_node_idx)) { + return errors::DataLoss("Bad NodeDef.input: ", name); + } + name.remove_suffix(name.size() - i); + } if (!name.empty() && name[0] == '^') { name.remove_prefix(1); - insert.BindInt(5, 1); + is_control = 1; } auto e = name_to_node_id_.find(name); if (e == name_to_node_id_.end()) { return errors::DataLoss("Could not find node: ", name); } - insert.BindInt(4, e->second); + input_node_id = e->second; + insert.BindInt(1, graph_id_); + insert.BindInt(2, node_id); + insert.BindInt(3, idx); + insert.BindInt(4, input_node_id); + insert.BindInt(5, input_node_idx); + insert.BindInt(6, is_control); + unflushed_bytes_ += insert.size(); TF_RETURN_WITH_CONTEXT_IF_ERROR(insert.StepAndReset(), node.name(), " -> ", name); + TF_RETURN_IF_ERROR(MaybeFlush()); } } return Status::OK(); } Status SaveNodes() { - auto insert = db_->Prepare(R"sql( - INSERT INTO Nodes (graph_id, node_id, node_name, op, device, node_def) + const char* sql = R"sql( + INSERT INTO Nodes ( + graph_id, + node_id, + node_name, + op, + device, + node_def) VALUES (?, ?, ?, ?, ?, ?) - )sql"); + )sql"; + SqliteStatement insert; + TF_RETURN_IF_ERROR(db_->Prepare(sql, &insert)); for (int node_id = 0; node_id < graph_->node_size(); ++node_id) { NodeDef* node = graph_->mutable_node(node_id); insert.BindInt(1, graph_id_); insert.BindInt(2, node_id); insert.BindText(3, node->name()); + insert.BindText(4, node->op()); + insert.BindText(5, node->device()); node->clear_name(); - if (!node->op().empty()) { - insert.BindText(4, node->op()); - node->clear_op(); - } - if (!node->device().empty()) { - insert.BindText(5, node->device()); - node->clear_device(); - } + node->clear_op(); + node->clear_device(); node->clear_input(); - TF_RETURN_IF_ERROR(BindProto(&insert, 6, *node)); + string node_def; + if (node->SerializeToString(&node_def)) { + insert.BindBlobUnsafe(6, node_def); + } + unflushed_bytes_ += insert.size(); TF_RETURN_WITH_CONTEXT_IF_ERROR(insert.StepAndReset(), node->name()); + TF_RETURN_IF_ERROR(MaybeFlush()); } return Status::OK(); } - Status SaveGraph() { - auto insert = db_->Prepare(R"sql( - INSERT INTO Graphs (graph_id, inserted_time, graph_def) - VALUES (?, ?, ?) - )sql"); - insert.BindInt(1, graph_id_); - insert.BindDouble(2, GetWallTime(env_)); + Status SaveGraph(int64 run_id) { + const char* sql = R"sql( + INSERT OR REPLACE INTO Graphs ( + run_id, + graph_id, + inserted_time, + graph_def + ) VALUES (?, ?, ?, ?) + )sql"; + SqliteStatement insert; + TF_RETURN_IF_ERROR(db_->Prepare(sql, &insert)); + if (run_id != kAbsent) insert.BindInt(1, run_id); + insert.BindInt(2, graph_id_); + insert.BindDouble(3, DoubleTime(now_)); graph_->clear_node(); - TF_RETURN_IF_ERROR(BindProto(&insert, 3, *graph_)); + string graph_def; + if (graph_->SerializeToString(&graph_def)) { + insert.BindBlobUnsafe(4, graph_def); + } return insert.StepAndReset(); } - Env* env_; - Sqlite* db_; - GraphDef* graph_; - int64 graph_id_; + Status MaybeFlush() { + if (unflushed_bytes_ >= kFlushBytes) { + TF_RETURN_WITH_CONTEXT_IF_ERROR(txn_->Commit(), "flushing ", + unflushed_bytes_, " bytes"); + unflushed_bytes_ = 0; + } + return Status::OK(); + } + + Sqlite* const db_; + SqliteTransaction* const txn_; + uint64 unflushed_bytes_ = 0; + GraphDef* const graph_; + const uint64 now_; + const int64 graph_id_; std::vector name_copies_; std::unordered_map name_to_node_id_; + + TF_DISALLOW_COPY_AND_ASSIGN(GraphWriter); }; -class RunWriter { +/// \brief Run metadata manager. +/// +/// This class gives us Tag IDs we can pass to SeriesWriter. In order +/// to do that, rows are created in the Ids, Tags, Runs, Experiments, +/// and Users tables. +/// +/// This class is thread safe. +class RunMetadata { public: - RunWriter(Env* env, std::shared_ptr db, const string& experiment_name, - const string& run_name, const string& user_name) - : env_{env}, - db_{std::move(db)}, - id_allocator_{env_, db_.get()}, + RunMetadata(IdAllocator* ids, const string& experiment_name, + const string& run_name, const string& user_name) + : ids_{ids}, experiment_name_{experiment_name}, run_name_{run_name}, - user_name_{user_name}, - insert_tensor_{db_->Prepare(R"sql( - INSERT OR REPLACE INTO Tensors (tag_id, step, computed_time, tensor) - VALUES (?, ?, ?, ?) - )sql")} {} - - ~RunWriter() { - if (run_id_ == kAbsent) return; - auto update = db_->Prepare(R"sql( - UPDATE Runs SET finished_time = ? WHERE run_id = ? - )sql"); - update.BindDouble(1, GetWallTime(env_)); - update.BindInt(2, run_id_); - Status s = update.StepAndReset(); - if (!s.ok()) { - LOG(ERROR) << "Failed to set Runs[" << run_id_ - << "].finish_time: " << s.ToString(); - } + user_name_{user_name} { + DCHECK(ids_ != nullptr); } - Status InsertTensor(int64 tag_id, int64 step, double computed_time, - Tensor t) { - insert_tensor_.BindInt(1, tag_id); - insert_tensor_.BindInt(2, step); - insert_tensor_.BindDouble(3, computed_time); - if (t.shape().dims() == 0 && t.dtype() == DT_INT64) { - insert_tensor_.BindInt(4, t.scalar()()); - } else if (t.shape().dims() == 0 && t.dtype() == DT_DOUBLE) { - insert_tensor_.BindDouble(4, t.scalar()()); - } else { - TF_RETURN_IF_ERROR(BindTensor(&insert_tensor_, 4, t)); - } - return insert_tensor_.StepAndReset(); + const string& experiment_name() { return experiment_name_; } + const string& run_name() { return run_name_; } + const string& user_name() { return user_name_; } + + int64 run_id() LOCKS_EXCLUDED(mu_) { + mutex_lock lock(mu_); + return run_id_; } - Status InsertGraph(std::unique_ptr g, double computed_time) { - TF_RETURN_IF_ERROR(InitializeRun(computed_time)); + Status SetGraph(Sqlite* db, uint64 now, double computed_time, + std::unique_ptr g) SQLITE_TRANSACTIONS_EXCLUDED(*db) + LOCKS_EXCLUDED(mu_) { + int64 run_id; + { + mutex_lock lock(mu_); + TF_RETURN_IF_ERROR(InitializeRun(db, now, computed_time)); + run_id = run_id_; + } int64 graph_id; + SqliteTransaction txn(*db); // only to increase performance TF_RETURN_IF_ERROR( - GraphSaver::Save(env_, db_.get(), &id_allocator_, g.get(), &graph_id)); - if (run_id_ != kAbsent) { - auto set = db_->Prepare("UPDATE Runs SET graph_id = ? WHERE run_id = ?"); - set.BindInt(1, graph_id); - set.BindInt(2, run_id_); - TF_RETURN_IF_ERROR(set.StepAndReset()); - } - return Status::OK(); + GraphWriter::Save(db, &txn, ids_, g.get(), now, run_id, &graph_id)); + return txn.Commit(); } - Status GetTagId(double computed_time, const string& tag_name, - const SummaryMetadata& metadata, int64* tag_id) { - TF_RETURN_IF_ERROR(InitializeRun(computed_time)); + Status GetTagId(Sqlite* db, uint64 now, double computed_time, + const string& tag_name, int64* tag_id, + const SummaryMetadata& metadata) LOCKS_EXCLUDED(mu_) { + mutex_lock lock(mu_); + TF_RETURN_IF_ERROR(InitializeRun(db, now, computed_time)); auto e = tag_ids_.find(tag_name); if (e != tag_ids_.end()) { *tag_id = e->second; return Status::OK(); } - TF_RETURN_IF_ERROR(id_allocator_.CreateNewId(tag_id)); + TF_RETURN_IF_ERROR(ids_->CreateNewId(tag_id)); tag_ids_[tag_name] = *tag_id; - if (!metadata.summary_description().empty()) { - SqliteStatement insert_description = db_->Prepare(R"sql( - INSERT INTO Descriptions (id, description) VALUES (?, ?) - )sql"); - insert_description.BindInt(1, *tag_id); - insert_description.BindText(2, metadata.summary_description()); - TF_RETURN_IF_ERROR(insert_description.StepAndReset()); - } - SqliteStatement insert = db_->Prepare(R"sql( + TF_RETURN_IF_ERROR( + SetDescription(db, *tag_id, metadata.summary_description())); + const char* sql = R"sql( INSERT INTO Tags ( run_id, tag_id, @@ -378,30 +459,54 @@ class RunWriter { display_name, plugin_name, plugin_data - ) VALUES (?, ?, ?, ?, ?, ?, ?) - )sql"); - if (run_id_ != kAbsent) insert.BindInt(1, run_id_); - insert.BindInt(2, *tag_id); - insert.BindText(3, tag_name); - insert.BindDouble(4, GetWallTime(env_)); - if (!metadata.display_name().empty()) { - insert.BindText(5, metadata.display_name()); - } - if (!metadata.plugin_data().plugin_name().empty()) { - insert.BindText(6, metadata.plugin_data().plugin_name()); - } - if (!metadata.plugin_data().content().empty()) { - insert.BindBlob(7, metadata.plugin_data().content()); - } + ) VALUES ( + :run_id, + :tag_id, + :tag_name, + :inserted_time, + :display_name, + :plugin_name, + :plugin_data + ) + )sql"; + SqliteStatement insert; + TF_RETURN_IF_ERROR(db->Prepare(sql, &insert)); + if (run_id_ != kAbsent) insert.BindInt(":run_id", run_id_); + insert.BindInt(":tag_id", *tag_id); + insert.BindTextUnsafe(":tag_name", tag_name); + insert.BindDouble(":inserted_time", DoubleTime(now)); + insert.BindTextUnsafe(":display_name", metadata.display_name()); + insert.BindTextUnsafe(":plugin_name", metadata.plugin_data().plugin_name()); + insert.BindBlobUnsafe(":plugin_data", metadata.plugin_data().content()); return insert.StepAndReset(); } + Status GetIsWatching(Sqlite* db, bool* is_watching) + SQLITE_TRANSACTIONS_EXCLUDED(*db) LOCKS_EXCLUDED(mu_) { + mutex_lock lock(mu_); + if (experiment_id_ == kAbsent) { + *is_watching = true; + return Status::OK(); + } + const char* sql = R"sql( + SELECT is_watching FROM Experiments WHERE experiment_id = ? + )sql"; + SqliteStatement stmt; + TF_RETURN_IF_ERROR(db->Prepare(sql, &stmt)); + stmt.BindInt(1, experiment_id_); + TF_RETURN_IF_ERROR(stmt.StepOnce()); + *is_watching = stmt.ColumnInt(0) != 0; + return Status::OK(); + } + private: - Status InitializeUser() { + Status InitializeUser(Sqlite* db, uint64 now) EXCLUSIVE_LOCKS_REQUIRED(mu_) { if (user_id_ != kAbsent || user_name_.empty()) return Status::OK(); - SqliteStatement get = db_->Prepare(R"sql( + const char* get_sql = R"sql( SELECT user_id FROM Users WHERE user_name = ? - )sql"); + )sql"; + SqliteStatement get; + TF_RETURN_IF_ERROR(db->Prepare(get_sql, &get)); get.BindText(1, user_name_); bool is_done; TF_RETURN_IF_ERROR(get.Step(&is_done)); @@ -409,22 +514,29 @@ class RunWriter { user_id_ = get.ColumnInt(0); return Status::OK(); } - TF_RETURN_IF_ERROR(id_allocator_.CreateNewId(&user_id_)); - SqliteStatement insert = db_->Prepare(R"sql( - INSERT INTO Users (user_id, user_name, inserted_time) VALUES (?, ?, ?) - )sql"); + TF_RETURN_IF_ERROR(ids_->CreateNewId(&user_id_)); + const char* insert_sql = R"sql( + INSERT INTO Users ( + user_id, + user_name, + inserted_time + ) VALUES (?, ?, ?) + )sql"; + SqliteStatement insert; + TF_RETURN_IF_ERROR(db->Prepare(insert_sql, &insert)); insert.BindInt(1, user_id_); insert.BindText(2, user_name_); - insert.BindDouble(3, GetWallTime(env_)); + insert.BindDouble(3, DoubleTime(now)); TF_RETURN_IF_ERROR(insert.StepAndReset()); return Status::OK(); } - Status InitializeExperiment(double computed_time) { + Status InitializeExperiment(Sqlite* db, uint64 now, double computed_time) + EXCLUSIVE_LOCKS_REQUIRED(mu_) { if (experiment_name_.empty()) return Status::OK(); if (experiment_id_ == kAbsent) { - TF_RETURN_IF_ERROR(InitializeUser()); - SqliteStatement get = db_->Prepare(R"sql( + TF_RETURN_IF_ERROR(InitializeUser(db, now)); + const char* get_sql = R"sql( SELECT experiment_id, started_time @@ -433,7 +545,9 @@ class RunWriter { WHERE user_id IS ? AND experiment_name = ? - )sql"); + )sql"; + SqliteStatement get; + TF_RETURN_IF_ERROR(db->Prepare(get_sql, &get)); if (user_id_ != kAbsent) get.BindInt(1, user_id_); get.BindText(2, experiment_name_); bool is_done; @@ -442,30 +556,41 @@ class RunWriter { experiment_id_ = get.ColumnInt(0); experiment_started_time_ = get.ColumnInt(1); } else { - TF_RETURN_IF_ERROR(id_allocator_.CreateNewId(&experiment_id_)); + TF_RETURN_IF_ERROR(ids_->CreateNewId(&experiment_id_)); experiment_started_time_ = computed_time; - SqliteStatement insert = db_->Prepare(R"sql( + const char* insert_sql = R"sql( INSERT INTO Experiments ( user_id, experiment_id, experiment_name, inserted_time, - started_time - ) VALUES (?, ?, ?, ?, ?) - )sql"); + started_time, + is_watching + ) VALUES (?, ?, ?, ?, ?, ?) + )sql"; + SqliteStatement insert; + TF_RETURN_IF_ERROR(db->Prepare(insert_sql, &insert)); if (user_id_ != kAbsent) insert.BindInt(1, user_id_); insert.BindInt(2, experiment_id_); insert.BindText(3, experiment_name_); - insert.BindDouble(4, GetWallTime(env_)); + insert.BindDouble(4, DoubleTime(now)); insert.BindDouble(5, computed_time); + insert.BindInt(6, 0); TF_RETURN_IF_ERROR(insert.StepAndReset()); } } if (computed_time < experiment_started_time_) { experiment_started_time_ = computed_time; - SqliteStatement update = db_->Prepare(R"sql( - UPDATE Experiments SET started_time = ? WHERE experiment_id = ? - )sql"); + const char* update_sql = R"sql( + UPDATE + Experiments + SET + started_time = ? + WHERE + experiment_id = ? + )sql"; + SqliteStatement update; + TF_RETURN_IF_ERROR(db->Prepare(update_sql, &update)); update.BindDouble(1, computed_time); update.BindInt(2, experiment_id_); TF_RETURN_IF_ERROR(update.StepAndReset()); @@ -473,13 +598,14 @@ class RunWriter { return Status::OK(); } - Status InitializeRun(double computed_time) { + Status InitializeRun(Sqlite* db, uint64 now, double computed_time) + EXCLUSIVE_LOCKS_REQUIRED(mu_) { if (run_name_.empty()) return Status::OK(); - TF_RETURN_IF_ERROR(InitializeExperiment(computed_time)); + TF_RETURN_IF_ERROR(InitializeExperiment(db, now, computed_time)); if (run_id_ == kAbsent) { - TF_RETURN_IF_ERROR(id_allocator_.CreateNewId(&run_id_)); + TF_RETURN_IF_ERROR(ids_->CreateNewId(&run_id_)); run_started_time_ = computed_time; - SqliteStatement insert = db_->Prepare(R"sql( + const char* insert_sql = R"sql( INSERT OR REPLACE INTO Runs ( experiment_id, run_id, @@ -487,19 +613,28 @@ class RunWriter { inserted_time, started_time ) VALUES (?, ?, ?, ?, ?) - )sql"); + )sql"; + SqliteStatement insert; + TF_RETURN_IF_ERROR(db->Prepare(insert_sql, &insert)); if (experiment_id_ != kAbsent) insert.BindInt(1, experiment_id_); insert.BindInt(2, run_id_); insert.BindText(3, run_name_); - insert.BindDouble(4, GetWallTime(env_)); + insert.BindDouble(4, DoubleTime(now)); insert.BindDouble(5, computed_time); TF_RETURN_IF_ERROR(insert.StepAndReset()); } if (computed_time < run_started_time_) { run_started_time_ = computed_time; - SqliteStatement update = db_->Prepare(R"sql( - UPDATE Runs SET started_time = ? WHERE run_id = ? - )sql"); + const char* update_sql = R"sql( + UPDATE + Runs + SET + started_time = ? + WHERE + run_id = ? + )sql"; + SqliteStatement update; + TF_RETURN_IF_ERROR(db->Prepare(update_sql, &update)); update.BindDouble(1, computed_time); update.BindInt(2, run_id_); TF_RETURN_IF_ERROR(update.StepAndReset()); @@ -507,137 +642,605 @@ class RunWriter { return Status::OK(); } - Env* env_; - std::shared_ptr db_; - IdAllocator id_allocator_; + mutex mu_; + IdAllocator* const ids_; const string experiment_name_; const string run_name_; const string user_name_; - int64 experiment_id_ = kAbsent; - int64 run_id_ = kAbsent; - int64 user_id_ = kAbsent; - std::unordered_map tag_ids_; - double experiment_started_time_ = 0.0; - double run_started_time_ = 0.0; - SqliteStatement insert_tensor_; + int64 experiment_id_ GUARDED_BY(mu_) = kAbsent; + int64 run_id_ GUARDED_BY(mu_) = kAbsent; + int64 user_id_ GUARDED_BY(mu_) = kAbsent; + double experiment_started_time_ GUARDED_BY(mu_) = 0.0; + double run_started_time_ GUARDED_BY(mu_) = 0.0; + std::unordered_map tag_ids_ GUARDED_BY(mu_); + + TF_DISALLOW_COPY_AND_ASSIGN(RunMetadata); }; +/// \brief Tensor writer for a single series, e.g. Tag. +/// +/// This class can be used to write an infinite stream of Tensors to the +/// database in a fixed block of contiguous disk space. This is +/// accomplished using Algorithm R reservoir sampling. +/// +/// The reservoir consists of a fixed number of rows, which are inserted +/// using ZEROBLOB upon receiving the first sample, which is used to +/// predict how big the other ones are likely to be. This is done +/// transactionally in a way that tries to be mindful of other processes +/// that might be trying to access the same DB. +/// +/// Once the reservoir fills up, rows are replaced at random, and writes +/// gradually become no-ops. This allows long training to go fast +/// without configuration. The exception is when someone is actually +/// looking at TensorBoard. When that happens, the "keep last" behavior +/// is turned on and Append() will always result in a write. +/// +/// If no one is watching training, this class still holds on to the +/// most recent "dangling" Tensor, so if Finish() is called, the most +/// recent training state can be written to disk. +/// +/// The randomly selected sampling points should be consistent across +/// multiple instances. +/// +/// This class is thread safe. +class SeriesWriter { + public: + SeriesWriter(int64 series, int slots, RunMetadata* meta) + : series_{series}, + slots_{slots}, + meta_{meta}, + rng_{std::mt19937_64::default_seed} { + DCHECK(series_ > 0); + DCHECK(slots_ > 0); + } + + Status Append(Sqlite* db, int64 step, uint64 now, double computed_time, + Tensor t) SQLITE_TRANSACTIONS_EXCLUDED(*db) + LOCKS_EXCLUDED(mu_) { + mutex_lock lock(mu_); + if (rowids_.empty()) { + Status s = Reserve(db, t); + if (!s.ok()) { + rowids_.clear(); + return s; + } + } + DCHECK(rowids_.size() == slots_); + int64 rowid; + size_t i = count_; + if (i < slots_) { + rowid = last_rowid_ = rowids_[i]; + } else { + i = rng_() % (i + 1); + if (i < slots_) { + rowid = last_rowid_ = rowids_[i]; + } else { + bool keep_last; + TF_RETURN_IF_ERROR(meta_->GetIsWatching(db, &keep_last)); + if (!keep_last) { + ++count_; + dangling_tensor_.reset(new Tensor(std::move(t))); + dangling_step_ = step; + dangling_computed_time_ = computed_time; + return Status::OK(); + } + rowid = last_rowid_; + } + } + Status s = Write(db, rowid, step, computed_time, t); + if (s.ok()) { + ++count_; + dangling_tensor_.reset(); + } + return s; + } + + Status Finish(Sqlite* db) SQLITE_TRANSACTIONS_EXCLUDED(*db) + LOCKS_EXCLUDED(mu_) { + mutex_lock lock(mu_); + // Short runs: Delete unused pre-allocated Tensors. + if (count_ < rowids_.size()) { + SqliteTransaction txn(*db); + const char* sql = R"sql( + DELETE FROM Tensors WHERE rowid = ? + )sql"; + SqliteStatement deleter; + TF_RETURN_IF_ERROR(db->Prepare(sql, &deleter)); + for (size_t i = count_; i < rowids_.size(); ++i) { + deleter.BindInt(1, rowids_[i]); + TF_RETURN_IF_ERROR(deleter.StepAndReset()); + } + TF_RETURN_IF_ERROR(txn.Commit()); + rowids_.clear(); + } + // Long runs: Make last sample be the very most recent one. + if (dangling_tensor_) { + DCHECK(last_rowid_ != kAbsent); + TF_RETURN_IF_ERROR(Write(db, last_rowid_, dangling_step_, + dangling_computed_time_, *dangling_tensor_)); + dangling_tensor_.reset(); + } + return Status::OK(); + } + + private: + Status Write(Sqlite* db, int64 rowid, int64 step, double computed_time, + const Tensor& t) SQLITE_TRANSACTIONS_EXCLUDED(*db) { + if (t.dtype() == DT_STRING) { + if (t.dims() == 0) { + return Update(db, step, computed_time, t, t.scalar()(), rowid); + } else { + SqliteTransaction txn(*db); + TF_RETURN_IF_ERROR( + Update(db, step, computed_time, t, StringPiece(), rowid)); + TF_RETURN_IF_ERROR(UpdateNdString(db, t, rowid)); + return txn.Commit(); + } + } else { + return Update(db, step, computed_time, t, t.tensor_data(), rowid); + } + } + + Status Update(Sqlite* db, int64 step, double computed_time, const Tensor& t, + const StringPiece& data, int64 rowid) { + // TODO(jart): How can we ensure reservoir fills on replace? + const char* sql = R"sql( + UPDATE OR REPLACE + Tensors + SET + step = ?, + computed_time = ?, + dtype = ?, + shape = ?, + data = ? + WHERE + rowid = ? + )sql"; + SqliteStatement stmt; + TF_RETURN_IF_ERROR(db->Prepare(sql, &stmt)); + stmt.BindInt(1, step); + stmt.BindDouble(2, computed_time); + stmt.BindInt(3, t.dtype()); + stmt.BindText(4, StringifyShape(t.shape())); + stmt.BindBlobUnsafe(5, data); + stmt.BindInt(6, rowid); + TF_RETURN_IF_ERROR(stmt.StepAndReset()); + return Status::OK(); + } + + Status UpdateNdString(Sqlite* db, const Tensor& t, int64 tensor_rowid) + SQLITE_EXCLUSIVE_TRANSACTIONS_REQUIRED(*db) { + DCHECK_EQ(t.dtype(), DT_STRING); + DCHECK_GT(t.dims(), 0); + const char* deleter_sql = R"sql( + DELETE FROM TensorStrings WHERE tensor_rowid = ? + )sql"; + SqliteStatement deleter; + TF_RETURN_IF_ERROR(db->Prepare(deleter_sql, &deleter)); + deleter.BindInt(1, tensor_rowid); + TF_RETURN_WITH_CONTEXT_IF_ERROR(deleter.StepAndReset(), tensor_rowid); + const char* inserter_sql = R"sql( + INSERT INTO TensorStrings ( + tensor_rowid, + idx, + data + ) VALUES (?, ?, ?) + )sql"; + SqliteStatement inserter; + TF_RETURN_IF_ERROR(db->Prepare(inserter_sql, &inserter)); + auto flat = t.flat(); + for (int64 i = 0; i < flat.size(); ++i) { + inserter.BindInt(1, tensor_rowid); + inserter.BindInt(2, i); + inserter.BindBlobUnsafe(3, flat(i)); + TF_RETURN_WITH_CONTEXT_IF_ERROR(inserter.StepAndReset(), "i=", i); + } + return Status::OK(); + } + + Status Reserve(Sqlite* db, const Tensor& t) SQLITE_TRANSACTIONS_EXCLUDED(*db) + EXCLUSIVE_LOCKS_REQUIRED(mu_) { + SqliteTransaction txn(*db); // only for performance + unflushed_bytes_ = 0; + if (t.dtype() == DT_STRING) { + if (t.dims() == 0) { + TF_RETURN_IF_ERROR(ReserveData(db, &txn, t.scalar()().size())); + } else { + TF_RETURN_IF_ERROR(ReserveTensors(db, &txn, kReserveMinBytes)); + } + } else { + TF_RETURN_IF_ERROR(ReserveData(db, &txn, t.tensor_data().size())); + } + return txn.Commit(); + } + + Status ReserveData(Sqlite* db, SqliteTransaction* txn, size_t size) + SQLITE_EXCLUSIVE_TRANSACTIONS_REQUIRED(*db) + EXCLUSIVE_LOCKS_REQUIRED(mu_) { + int64 space = + static_cast(static_cast(size) * kReserveMultiplier); + if (space < kReserveMinBytes) space = kReserveMinBytes; + return ReserveTensors(db, txn, space); + } + + Status ReserveTensors(Sqlite* db, SqliteTransaction* txn, + int64 reserved_bytes) + SQLITE_EXCLUSIVE_TRANSACTIONS_REQUIRED(*db) + EXCLUSIVE_LOCKS_REQUIRED(mu_) { + const char* sql = R"sql( + INSERT INTO Tensors ( + series, + data + ) VALUES (?, ZEROBLOB(?)) + )sql"; + SqliteStatement insert; + TF_RETURN_IF_ERROR(db->Prepare(sql, &insert)); + // TODO(jart): Maybe preallocate index pages by setting step. This + // is tricky because UPDATE OR REPLACE can have a side + // effect of deleting preallocated rows. + for (int64 i = 0; i < slots_; ++i) { + insert.BindInt(1, series_); + insert.BindInt(2, reserved_bytes); + TF_RETURN_WITH_CONTEXT_IF_ERROR(insert.StepAndReset(), "i=", i); + rowids_.push_back(db->last_insert_rowid()); + unflushed_bytes_ += reserved_bytes; + TF_RETURN_IF_ERROR(MaybeFlush(db, txn)); + } + return Status::OK(); + } + + Status MaybeFlush(Sqlite* db, SqliteTransaction* txn) + SQLITE_EXCLUSIVE_TRANSACTIONS_REQUIRED(*db) + EXCLUSIVE_LOCKS_REQUIRED(mu_) { + if (unflushed_bytes_ >= kFlushBytes) { + TF_RETURN_WITH_CONTEXT_IF_ERROR(txn->Commit(), "flushing ", + unflushed_bytes_, " bytes"); + unflushed_bytes_ = 0; + } + return Status::OK(); + } + + mutex mu_; + const int64 series_; + const int slots_; + RunMetadata* const meta_; + std::mt19937_64 rng_ GUARDED_BY(mu_); + uint64 count_ GUARDED_BY(mu_) = 0; + int64 last_rowid_ GUARDED_BY(mu_) = kAbsent; + std::vector rowids_ GUARDED_BY(mu_); + uint64 unflushed_bytes_ GUARDED_BY(mu_) = 0; + std::unique_ptr dangling_tensor_ GUARDED_BY(mu_); + int64 dangling_step_ GUARDED_BY(mu_) = 0; + double dangling_computed_time_ GUARDED_BY(mu_) = 0.0; + + TF_DISALLOW_COPY_AND_ASSIGN(SeriesWriter); +}; + +/// \brief Tensor writer for a single Run. +/// +/// This class farms out tensors to SeriesWriter instances. It also +/// keeps track of whether or not someone is watching the TensorBoard +/// GUI, so it can avoid writes when possible. +/// +/// This class is thread safe. +class RunWriter { + public: + explicit RunWriter(RunMetadata* meta) : meta_{meta} {} + + Status Append(Sqlite* db, int64 tag_id, int64 step, uint64 now, + double computed_time, Tensor t, int slots) + SQLITE_TRANSACTIONS_EXCLUDED(*db) LOCKS_EXCLUDED(mu_) { + SeriesWriter* writer = GetSeriesWriter(tag_id, slots); + return writer->Append(db, step, now, computed_time, std::move(t)); + } + + Status Finish(Sqlite* db) SQLITE_TRANSACTIONS_EXCLUDED(*db) + LOCKS_EXCLUDED(mu_) { + mutex_lock lock(mu_); + if (series_writers_.empty()) return Status::OK(); + for (auto i = series_writers_.begin(); i != series_writers_.end(); ++i) { + if (!i->second) continue; + TF_RETURN_WITH_CONTEXT_IF_ERROR(i->second->Finish(db), + "finish tag_id=", i->first); + i->second.reset(); + } + return Status::OK(); + } + + private: + SeriesWriter* GetSeriesWriter(int64 tag_id, int slots) LOCKS_EXCLUDED(mu_) { + mutex_lock sl(mu_); + auto spot = series_writers_.find(tag_id); + if (spot == series_writers_.end()) { + SeriesWriter* writer = new SeriesWriter(tag_id, slots, meta_); + series_writers_[tag_id].reset(writer); + return writer; + } else { + return spot->second.get(); + } + } + + mutex mu_; + RunMetadata* const meta_; + std::unordered_map> series_writers_ + GUARDED_BY(mu_); + + TF_DISALLOW_COPY_AND_ASSIGN(RunWriter); +}; + +/// \brief SQLite implementation of SummaryWriterInterface. +/// +/// This class is thread safe. class SummaryDbWriter : public SummaryWriterInterface { public: - SummaryDbWriter(Env* env, std::shared_ptr db, - const string& experiment_name, const string& run_name, - const string& user_name) + SummaryDbWriter(Env* env, Sqlite* db, const string& experiment_name, + const string& run_name, const string& user_name) : SummaryWriterInterface(), env_{env}, - run_writer_{env, std::move(db), experiment_name, run_name, user_name} {} - ~SummaryDbWriter() override {} + db_{db}, + ids_{env_, db_}, + meta_{&ids_, experiment_name, run_name, user_name}, + run_{&meta_} { + DCHECK(env_ != nullptr); + db_->Ref(); + } + + ~SummaryDbWriter() override { + core::ScopedUnref unref(db_); + Status s = run_.Finish(db_); + if (!s.ok()) { + // TODO(jart): Retry on transient errors here. + LOG(ERROR) << s.ToString(); + } + int64 run_id = meta_.run_id(); + if (run_id == kAbsent) return; + const char* sql = R"sql( + UPDATE Runs SET finished_time = ? WHERE run_id = ? + )sql"; + SqliteStatement update; + s = db_->Prepare(sql, &update); + if (s.ok()) { + update.BindDouble(1, DoubleTime(env_->NowMicros())); + update.BindInt(2, run_id); + s = update.StepAndReset(); + } + if (!s.ok()) { + LOG(ERROR) << "Failed to set Runs[" << run_id + << "].finish_time: " << s.ToString(); + } + } Status Flush() override { return Status::OK(); } Status WriteTensor(int64 global_step, Tensor t, const string& tag, const string& serialized_metadata) override { - mutex_lock ml(mu_); + TF_RETURN_IF_ERROR(CheckSupportedType(t)); SummaryMetadata metadata; - if (!serialized_metadata.empty()) { - metadata.ParseFromString(serialized_metadata); + if (!metadata.ParseFromString(serialized_metadata)) { + return errors::InvalidArgument("Bad serialized_metadata"); } - double now = GetWallTime(env_); - int64 tag_id; - TF_RETURN_IF_ERROR(run_writer_.GetTagId(now, tag, metadata, &tag_id)); - return run_writer_.InsertTensor(tag_id, global_step, now, t); + return Write(global_step, t, tag, metadata); } Status WriteScalar(int64 global_step, Tensor t, const string& tag) override { - Tensor t2; - TF_RETURN_IF_ERROR(CoerceScalar(t, &t2)); - // TODO(jart): Generate scalars plugin metadata on this value. - return WriteTensor(global_step, std::move(t2), tag, ""); + TF_RETURN_IF_ERROR(CheckSupportedType(t)); + SummaryMetadata metadata; + PatchPluginName(&metadata, kScalarPluginName); + return Write(global_step, AsScalar(t), tag, metadata); } Status WriteGraph(int64 global_step, std::unique_ptr g) override { - mutex_lock ml(mu_); - return run_writer_.InsertGraph(std::move(g), GetWallTime(env_)); + uint64 now = env_->NowMicros(); + return meta_.SetGraph(db_, now, DoubleTime(now), std::move(g)); } Status WriteEvent(std::unique_ptr e) override { - switch (e->what_case()) { - case Event::WhatCase::kSummary: { - mutex_lock ml(mu_); - Status s; - for (const auto& value : e->summary().value()) { - s.Update(WriteSummary(e.get(), value)); - } - return s; - } - case Event::WhatCase::kGraphDef: { - mutex_lock ml(mu_); - std::unique_ptr graph{new GraphDef}; - if (!ParseProtoUnlimited(graph.get(), e->graph_def())) { - return errors::DataLoss("parse event.graph_def failed"); - } - return run_writer_.InsertGraph(std::move(graph), e->wall_time()); - } - default: - // TODO(@jart): Handle other stuff. - return Status::OK(); - } + return MigrateEvent(std::move(e)); } Status WriteHistogram(int64 global_step, Tensor t, const string& tag) override { - return errors::Unimplemented( - "SummaryDbWriter::WriteHistogram not supported. Please use ", - "tensorboard.summary.histogram() instead."); + uint64 now = env_->NowMicros(); + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(DoubleTime(now)); + TF_RETURN_IF_ERROR( + AddTensorAsHistogramToSummary(t, tag, e->mutable_summary())); + return MigrateEvent(std::move(e)); } - Status WriteImage(int64 global_step, Tensor tensor, const string& tag, + Status WriteImage(int64 global_step, Tensor t, const string& tag, int max_images, Tensor bad_color) override { - return errors::Unimplemented( - "SummaryDbWriter::WriteImage not supported. Please use ", - "tensorboard.summary.image() instead."); + uint64 now = env_->NowMicros(); + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(DoubleTime(now)); + TF_RETURN_IF_ERROR(AddTensorAsImageToSummary(t, tag, max_images, bad_color, + e->mutable_summary())); + return MigrateEvent(std::move(e)); } - Status WriteAudio(int64 global_step, Tensor tensor, const string& tag, + Status WriteAudio(int64 global_step, Tensor t, const string& tag, int max_outputs, float sample_rate) override { - return errors::Unimplemented( - "SummaryDbWriter::WriteAudio not supported. Please use ", - "tensorboard.summary.audio() instead."); + uint64 now = env_->NowMicros(); + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(DoubleTime(now)); + TF_RETURN_IF_ERROR(AddTensorAsAudioToSummary( + t, tag, max_outputs, sample_rate, e->mutable_summary())); + return MigrateEvent(std::move(e)); } string DebugString() override { return "SummaryDbWriter"; } private: - Status WriteSummary(const Event* e, const Summary::Value& summary) - EXCLUSIVE_LOCKS_REQUIRED(mu_) { - switch (summary.value_case()) { - case Summary::Value::ValueCase::kSimpleValue: { - int64 tag_id; - TF_RETURN_IF_ERROR(run_writer_.GetTagId(e->wall_time(), summary.tag(), - summary.metadata(), &tag_id)); - Tensor t{DT_DOUBLE, {}}; - t.scalar()() = summary.simple_value(); - return run_writer_.InsertTensor(tag_id, e->step(), e->wall_time(), t); + Status Write(int64 step, const Tensor& t, const string& tag, + const SummaryMetadata& metadata) { + uint64 now = env_->NowMicros(); + double computed_time = DoubleTime(now); + int64 tag_id; + TF_RETURN_IF_ERROR( + meta_.GetTagId(db_, now, computed_time, tag, &tag_id, metadata)); + TF_RETURN_WITH_CONTEXT_IF_ERROR( + run_.Append(db_, tag_id, step, now, computed_time, t, + GetSlots(t, metadata)), + meta_.user_name(), "/", meta_.experiment_name(), "/", meta_.run_name(), + "/", tag, "@", step); + return Status::OK(); + } + + Status MigrateEvent(std::unique_ptr e) { + switch (e->what_case()) { + case Event::WhatCase::kSummary: { + uint64 now = env_->NowMicros(); + auto summaries = e->mutable_summary(); + for (int i = 0; i < summaries->value_size(); ++i) { + Summary::Value* value = summaries->mutable_value(i); + TF_RETURN_WITH_CONTEXT_IF_ERROR( + MigrateSummary(e.get(), value, now), meta_.user_name(), "/", + meta_.experiment_name(), "/", meta_.run_name(), "/", value->tag(), + "@", e->step()); + } + break; } + case Event::WhatCase::kGraphDef: + TF_RETURN_WITH_CONTEXT_IF_ERROR( + MigrateGraph(e.get(), e->graph_def()), meta_.user_name(), "/", + meta_.experiment_name(), "/", meta_.run_name(), "/__graph__@", + e->step()); + break; default: - // TODO(@jart): Handle the rest. - return Status::OK(); + // TODO(@jart): Handle other stuff. + break; } + return Status::OK(); } - mutex mu_; - Env* env_; - RunWriter run_writer_ GUARDED_BY(mu_); + Status MigrateGraph(const Event* e, const string& graph_def) { + uint64 now = env_->NowMicros(); + std::unique_ptr graph{new GraphDef}; + if (!ParseProtoUnlimited(graph.get(), graph_def)) { + return errors::InvalidArgument("bad proto"); + } + return meta_.SetGraph(db_, now, e->wall_time(), std::move(graph)); + } + + Status MigrateSummary(const Event* e, Summary::Value* s, uint64 now) { + switch (s->value_case()) { + case Summary::Value::ValueCase::kTensor: + TF_RETURN_WITH_CONTEXT_IF_ERROR(MigrateTensor(e, s, now), "tensor"); + break; + case Summary::Value::ValueCase::kSimpleValue: + TF_RETURN_WITH_CONTEXT_IF_ERROR(MigrateScalar(e, s, now), "scalar"); + break; + case Summary::Value::ValueCase::kHisto: + TF_RETURN_WITH_CONTEXT_IF_ERROR(MigrateHistogram(e, s, now), "histo"); + break; + case Summary::Value::ValueCase::kImage: + TF_RETURN_WITH_CONTEXT_IF_ERROR(MigrateImage(e, s, now), "image"); + break; + case Summary::Value::ValueCase::kAudio: + TF_RETURN_WITH_CONTEXT_IF_ERROR(MigrateAudio(e, s, now), "audio"); + break; + default: + break; + } + return Status::OK(); + } + + Status MigrateTensor(const Event* e, Summary::Value* s, uint64 now) { + Tensor t; + if (!t.FromProto(s->tensor())) return errors::InvalidArgument("bad proto"); + TF_RETURN_IF_ERROR(CheckSupportedType(t)); + int64 tag_id; + TF_RETURN_IF_ERROR(meta_.GetTagId(db_, now, e->wall_time(), s->tag(), + &tag_id, s->metadata())); + return run_.Append(db_, tag_id, e->step(), now, e->wall_time(), t, + GetSlots(t, s->metadata())); + } + + // TODO(jart): Refactor Summary -> Tensor logic into separate file. + + Status MigrateScalar(const Event* e, Summary::Value* s, uint64 now) { + // See tensorboard/plugins/scalar/summary.py and data_compat.py + Tensor t{DT_FLOAT, {}}; + t.scalar()() = s->simple_value(); + int64 tag_id; + PatchPluginName(s->mutable_metadata(), kScalarPluginName); + TF_RETURN_IF_ERROR(meta_.GetTagId(db_, now, e->wall_time(), s->tag(), + &tag_id, s->metadata())); + return run_.Append(db_, tag_id, e->step(), now, e->wall_time(), + std::move(t), kScalarSlots); + } + + Status MigrateHistogram(const Event* e, Summary::Value* s, uint64 now) { + const HistogramProto& histo = s->histo(); + int k = histo.bucket_size(); + if (k != histo.bucket_limit_size()) { + return errors::InvalidArgument("size mismatch"); + } + // See tensorboard/plugins/histogram/summary.py and data_compat.py + Tensor t{DT_DOUBLE, {k, 3}}; + auto data = t.flat(); + for (int i = 0; i < k; ++i) { + double left_edge = ((i - 1 >= 0) ? histo.bucket_limit(i - 1) + : std::numeric_limits::min()); + double right_edge = ((i + 1 < k) ? histo.bucket_limit(i + 1) + : std::numeric_limits::max()); + data(i + 0) = left_edge; + data(i + 1) = right_edge; + data(i + 2) = histo.bucket(i); + } + int64 tag_id; + PatchPluginName(s->mutable_metadata(), kHistogramPluginName); + TF_RETURN_IF_ERROR(meta_.GetTagId(db_, now, e->wall_time(), s->tag(), + &tag_id, s->metadata())); + return run_.Append(db_, tag_id, e->step(), now, e->wall_time(), + std::move(t), kHistogramSlots); + } + + Status MigrateImage(const Event* e, Summary::Value* s, uint64 now) { + // See tensorboard/plugins/image/summary.py and data_compat.py + Tensor t{DT_STRING, {3}}; + auto img = s->mutable_image(); + t.flat()(0) = strings::StrCat(img->width()); + t.flat()(1) = strings::StrCat(img->height()); + t.flat()(2) = std::move(*img->mutable_encoded_image_string()); + int64 tag_id; + PatchPluginName(s->mutable_metadata(), kImagePluginName); + TF_RETURN_IF_ERROR(meta_.GetTagId(db_, now, e->wall_time(), s->tag(), + &tag_id, s->metadata())); + return run_.Append(db_, tag_id, e->step(), now, e->wall_time(), + std::move(t), kImageSlots); + } + + Status MigrateAudio(const Event* e, Summary::Value* s, uint64 now) { + // See tensorboard/plugins/audio/summary.py and data_compat.py + Tensor t{DT_STRING, {1, 2}}; + auto wav = s->mutable_audio(); + t.flat()(0) = std::move(*wav->mutable_encoded_audio_string()); + t.flat()(1) = ""; + int64 tag_id; + PatchPluginName(s->mutable_metadata(), kAudioPluginName); + TF_RETURN_IF_ERROR(meta_.GetTagId(db_, now, e->wall_time(), s->tag(), + &tag_id, s->metadata())); + return run_.Append(db_, tag_id, e->step(), now, e->wall_time(), + std::move(t), kAudioSlots); + } + + Env* const env_; + Sqlite* const db_; + IdAllocator ids_; + RunMetadata meta_; + RunWriter run_; }; } // namespace -Status CreateSummaryDbWriter(std::shared_ptr db, - const string& experiment_name, +Status CreateSummaryDbWriter(Sqlite* db, const string& experiment_name, const string& run_name, const string& user_name, Env* env, SummaryWriterInterface** result) { - TF_RETURN_IF_ERROR(SetupTensorboardSqliteDb(db)); - *result = new SummaryDbWriter(env, std::move(db), experiment_name, run_name, - user_name); + *result = new SummaryDbWriter(env, db, experiment_name, run_name, user_name); return Status::OK(); } diff --git a/tensorflow/contrib/tensorboard/db/summary_db_writer.h b/tensorflow/contrib/tensorboard/db/summary_db_writer.h index 74f61e50b7cdf4b4151162a2e1e5e0af0d468be2..746da1533b157bf7b2be5c85ada8b61ba224cc3e 100644 --- a/tensorflow/contrib/tensorboard/db/summary_db_writer.h +++ b/tensorflow/contrib/tensorboard/db/summary_db_writer.h @@ -19,21 +19,21 @@ limitations under the License. #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/db/sqlite.h" #include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/types.h" namespace tensorflow { /// \brief Creates SQLite SummaryWriterInterface. /// /// This can be used to write tensors from the execution graph directly -/// to a database. The schema will be created automatically, but only -/// if necessary. Entries in the Users, Experiments, and Runs tables -/// will be created automatically if they don't already exist. +/// to a database. The schema must be created beforehand. Entries in +/// Users, Experiments, and Runs tables will be created automatically +/// if they don't already exist. /// /// Please note that the type signature of this function may change in /// the future if support for other DBs is added to core. -Status CreateSummaryDbWriter(std::shared_ptr db, - const string& experiment_name, +/// +/// The result holds a new reference to db. +Status CreateSummaryDbWriter(Sqlite* db, const string& experiment_name, const string& run_name, const string& user_name, Env* env, SummaryWriterInterface** result); diff --git a/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc b/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc index 5ea844b6685d15ac4c0549816770060c6f25ce38..29b8063218de72aac1a73bbfb440e75fcdd5013f 100644 --- a/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc +++ b/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc @@ -14,6 +14,8 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/tensorboard/db/summary_db_writer.h" +#include "tensorflow/contrib/tensorboard/db/schema.h" +#include "tensorflow/core/framework/function.pb.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/framework/summary.pb.h" @@ -27,8 +29,6 @@ limitations under the License. namespace tensorflow { namespace { -const float kTolerance = 1e-5; - Tensor MakeScalarInt64(int64 x) { Tensor t(DT_INT64, TensorShape({})); t.scalar()() = x; @@ -48,17 +48,22 @@ class FakeClockEnv : public EnvWrapper { class SummaryDbWriterTest : public ::testing::Test { protected: - void SetUp() override { db_ = Sqlite::Open(":memory:").ValueOrDie(); } + void SetUp() override { + TF_ASSERT_OK(Sqlite::Open(":memory:", SQLITE_OPEN_READWRITE, &db_)); + TF_ASSERT_OK(SetupTensorboardSqliteDb(db_)); + } void TearDown() override { if (writer_ != nullptr) { writer_->Unref(); writer_ = nullptr; } + db_->Unref(); + db_ = nullptr; } int64 QueryInt(const string& sql) { - SqliteStatement stmt = db_->Prepare(sql); + SqliteStatement stmt = db_->PrepareOrDie(sql); bool is_done; Status s = stmt.Step(&is_done); if (!s.ok() || is_done) { @@ -69,7 +74,7 @@ class SummaryDbWriterTest : public ::testing::Test { } double QueryDouble(const string& sql) { - SqliteStatement stmt = db_->Prepare(sql); + SqliteStatement stmt = db_->PrepareOrDie(sql); bool is_done; Status s = stmt.Step(&is_done); if (!s.ok() || is_done) { @@ -80,7 +85,7 @@ class SummaryDbWriterTest : public ::testing::Test { } string QueryString(const string& sql) { - SqliteStatement stmt = db_->Prepare(sql); + SqliteStatement stmt = db_->PrepareOrDie(sql); bool is_done; Status s = stmt.Step(&is_done); if (!s.ok() || is_done) { @@ -91,7 +96,7 @@ class SummaryDbWriterTest : public ::testing::Test { } FakeClockEnv env_; - std::shared_ptr db_; + Sqlite* db_ = nullptr; SummaryWriterInterface* writer_ = nullptr; }; @@ -134,7 +139,7 @@ TEST_F(SummaryDbWriterTest, TensorsWritten_RowsGetInitialized) { ASSERT_EQ(1LL, QueryInt("SELECT COUNT(*) FROM Experiments")); ASSERT_EQ(1LL, QueryInt("SELECT COUNT(*) FROM Runs")); ASSERT_EQ(1LL, QueryInt("SELECT COUNT(*) FROM Tags")); - ASSERT_EQ(2LL, QueryInt("SELECT COUNT(*) FROM Tensors")); + ASSERT_EQ(10000LL, QueryInt("SELECT COUNT(*) FROM Tensors")); int64 user_id = QueryInt("SELECT user_id FROM Users"); int64 experiment_id = QueryInt("SELECT experiment_id FROM Experiments"); @@ -166,17 +171,13 @@ TEST_F(SummaryDbWriterTest, TensorsWritten_RowsGetInitialized) { EXPECT_EQ("plugin_data", QueryString("SELECT plugin_data FROM Tags")); EXPECT_EQ("description", QueryString("SELECT description FROM Descriptions")); - EXPECT_EQ(tag_id, QueryInt("SELECT tag_id FROM Tensors WHERE step = 1")); + EXPECT_EQ(tag_id, QueryInt("SELECT series FROM Tensors WHERE step = 1")); EXPECT_EQ(0.023, QueryDouble("SELECT computed_time FROM Tensors WHERE step = 1")); - EXPECT_FALSE( - QueryString("SELECT tensor FROM Tensors WHERE step = 1").empty()); - EXPECT_EQ(tag_id, QueryInt("SELECT tag_id FROM Tensors WHERE step = 2")); + EXPECT_EQ(tag_id, QueryInt("SELECT series FROM Tensors WHERE step = 2")); EXPECT_EQ(0.046, QueryDouble("SELECT computed_time FROM Tensors WHERE step = 2")); - EXPECT_FALSE( - QueryString("SELECT tensor FROM Tensors WHERE step = 2").empty()); } TEST_F(SummaryDbWriterTest, EmptyParentNames_NoParentsCreated) { @@ -187,7 +188,7 @@ TEST_F(SummaryDbWriterTest, EmptyParentNames_NoParentsCreated) { ASSERT_EQ(0LL, QueryInt("SELECT COUNT(*) FROM Experiments")); ASSERT_EQ(0LL, QueryInt("SELECT COUNT(*) FROM Runs")); ASSERT_EQ(1LL, QueryInt("SELECT COUNT(*) FROM Tags")); - ASSERT_EQ(1LL, QueryInt("SELECT COUNT(*) FROM Tensors")); + ASSERT_EQ(10000LL, QueryInt("SELECT COUNT(*) FROM Tensors")); } TEST_F(SummaryDbWriterTest, WriteEvent_Scalar) { @@ -204,33 +205,24 @@ TEST_F(SummaryDbWriterTest, WriteEvent_Scalar) { TF_ASSERT_OK(writer_->WriteEvent(std::move(e))); TF_ASSERT_OK(writer_->Flush()); ASSERT_EQ(2LL, QueryInt("SELECT COUNT(*) FROM Tags")); - ASSERT_EQ(2LL, QueryInt("SELECT COUNT(*) FROM Tensors")); + ASSERT_EQ(20000LL, QueryInt("SELECT COUNT(*) FROM Tensors")); int64 tag1_id = QueryInt("SELECT tag_id FROM Tags WHERE tag_name = 'π'"); int64 tag2_id = QueryInt("SELECT tag_id FROM Tags WHERE tag_name = 'φ'"); EXPECT_GT(tag1_id, 0LL); EXPECT_GT(tag2_id, 0LL); EXPECT_EQ(123.456, QueryDouble(strings::StrCat( - "SELECT computed_time FROM Tensors WHERE tag_id = ", + "SELECT computed_time FROM Tensors WHERE series = ", tag1_id, " AND step = 7"))); EXPECT_EQ(123.456, QueryDouble(strings::StrCat( - "SELECT computed_time FROM Tensors WHERE tag_id = ", + "SELECT computed_time FROM Tensors WHERE series = ", tag2_id, " AND step = 7"))); - EXPECT_NEAR(3.14, - QueryDouble(strings::StrCat( - "SELECT tensor FROM Tensors WHERE tag_id = ", tag1_id, - " AND step = 7")), - kTolerance); // Summary::simple_value is float - EXPECT_NEAR(1.61, - QueryDouble(strings::StrCat( - "SELECT tensor FROM Tensors WHERE tag_id = ", tag2_id, - " AND step = 7")), - kTolerance); } TEST_F(SummaryDbWriterTest, WriteGraph) { TF_ASSERT_OK(CreateSummaryDbWriter(db_, "", "R", "", &env_, &writer_)); env_.AdvanceByMillis(23); GraphDef graph; + graph.mutable_library()->add_gradient()->set_function_name("funk"); NodeDef* node = graph.add_node(); node->set_name("x"); node->set_op("Placeholder"); @@ -256,11 +248,17 @@ TEST_F(SummaryDbWriterTest, WriteGraph) { ASSERT_EQ(4LL, QueryInt("SELECT COUNT(*) FROM Nodes")); ASSERT_EQ(3LL, QueryInt("SELECT COUNT(*) FROM NodeInputs")); + ASSERT_EQ(QueryInt("SELECT run_id FROM Runs"), + QueryInt("SELECT run_id FROM Graphs")); + int64 graph_id = QueryInt("SELECT graph_id FROM Graphs"); EXPECT_GT(graph_id, 0LL); - EXPECT_EQ(graph_id, QueryInt("SELECT graph_id FROM Runs")); EXPECT_EQ(0.023, QueryDouble("SELECT inserted_time FROM Graphs")); - EXPECT_FALSE(QueryString("SELECT graph_def FROM Graphs").empty()); + + GraphDef graph2; + graph2.ParseFromString(QueryString("SELECT graph_def FROM Graphs")); + EXPECT_EQ(0, graph2.node_size()); + EXPECT_EQ("funk", graph2.library().gradient(0).function_name()); EXPECT_EQ("x", QueryString("SELECT node_name FROM Nodes WHERE node_id = 0")); EXPECT_EQ("y", QueryString("SELECT node_name FROM Nodes WHERE node_id = 1")); @@ -303,33 +301,6 @@ TEST_F(SummaryDbWriterTest, WriteGraph) { EXPECT_EQ(1LL, QueryInt("SELECT is_control FROM NodeInputs WHERE idx = 2")); } -TEST_F(SummaryDbWriterTest, WriteScalarInt32_CoercesToInt64) { - TF_ASSERT_OK(CreateSummaryDbWriter(db_, "", "", "", &env_, &writer_)); - Tensor t(DT_INT32, {}); - t.scalar()() = -17; - TF_ASSERT_OK(writer_->WriteScalar(1, t, "t")); - TF_ASSERT_OK(writer_->Flush()); - ASSERT_EQ(-17LL, QueryInt("SELECT tensor FROM Tensors")); -} - -TEST_F(SummaryDbWriterTest, WriteScalarInt8_CoercesToInt64) { - TF_ASSERT_OK(CreateSummaryDbWriter(db_, "", "", "", &env_, &writer_)); - Tensor t(DT_INT8, {}); - t.scalar()() = static_cast(-17); - TF_ASSERT_OK(writer_->WriteScalar(1, t, "t")); - TF_ASSERT_OK(writer_->Flush()); - ASSERT_EQ(-17LL, QueryInt("SELECT tensor FROM Tensors")); -} - -TEST_F(SummaryDbWriterTest, WriteScalarUint8_CoercesToInt64) { - TF_ASSERT_OK(CreateSummaryDbWriter(db_, "", "", "", &env_, &writer_)); - Tensor t(DT_UINT8, {}); - t.scalar()() = static_cast(254); - TF_ASSERT_OK(writer_->WriteScalar(1, t, "t")); - TF_ASSERT_OK(writer_->Flush()); - ASSERT_EQ(254LL, QueryInt("SELECT tensor FROM Tensors")); -} - TEST_F(SummaryDbWriterTest, UsesIdsTable) { SummaryMetadata metadata; TF_ASSERT_OK(CreateSummaryDbWriter(db_, "mad-science", "train", "jart", &env_, diff --git a/tensorflow/contrib/tensorboard/db/summary_file_writer.cc b/tensorflow/contrib/tensorboard/db/summary_file_writer.cc new file mode 100644 index 0000000000000000000000000000000000000000..85b3e7231bcb433e9510522597c03c5f764f06cf --- /dev/null +++ b/tensorflow/contrib/tensorboard/db/summary_file_writer.cc @@ -0,0 +1,190 @@ +/* Copyright 2017 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/tensorboard/db/summary_file_writer.h" + +#include "tensorflow/contrib/tensorboard/db/summary_converter.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/resource_mgr.h" +#include "tensorflow/core/framework/summary.pb.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/util/events_writer.h" +#include "tensorflow/core/util/ptr_util.h" + +namespace tensorflow { +namespace { + +class SummaryFileWriter : public SummaryWriterInterface { + public: + SummaryFileWriter(int max_queue, int flush_millis, Env* env) + : SummaryWriterInterface(), + is_initialized_(false), + max_queue_(max_queue), + flush_millis_(flush_millis), + env_(env) {} + + Status Initialize(const string& logdir, const string& filename_suffix) { + const Status is_dir = env_->IsDirectory(logdir); + if (!is_dir.ok()) { + if (is_dir.code() != tensorflow::error::NOT_FOUND) { + return is_dir; + } + TF_RETURN_IF_ERROR(env_->RecursivelyCreateDir(logdir)); + } + mutex_lock ml(mu_); + events_writer_ = + tensorflow::MakeUnique(io::JoinPath(logdir, "events")); + TF_RETURN_WITH_CONTEXT_IF_ERROR( + events_writer_->InitWithSuffix(filename_suffix), + "Could not initialize events writer."); + last_flush_ = env_->NowMicros(); + is_initialized_ = true; + return Status::OK(); + } + + Status Flush() override { + mutex_lock ml(mu_); + if (!is_initialized_) { + return errors::FailedPrecondition("Class was not properly initialized."); + } + return InternalFlush(); + } + + ~SummaryFileWriter() override { + (void)Flush(); // Ignore errors. + } + + Status WriteTensor(int64 global_step, Tensor t, const string& tag, + const string& serialized_metadata) override { + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(GetWallTime()); + Summary::Value* v = e->mutable_summary()->add_value(); + t.AsProtoTensorContent(v->mutable_tensor()); + v->set_tag(tag); + if (!serialized_metadata.empty()) { + v->mutable_metadata()->ParseFromString(serialized_metadata); + } + return WriteEvent(std::move(e)); + } + + Status WriteScalar(int64 global_step, Tensor t, const string& tag) override { + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(GetWallTime()); + TF_RETURN_IF_ERROR( + AddTensorAsScalarToSummary(t, tag, e->mutable_summary())); + return WriteEvent(std::move(e)); + } + + Status WriteHistogram(int64 global_step, Tensor t, + const string& tag) override { + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(GetWallTime()); + TF_RETURN_IF_ERROR( + AddTensorAsHistogramToSummary(t, tag, e->mutable_summary())); + return WriteEvent(std::move(e)); + } + + Status WriteImage(int64 global_step, Tensor t, const string& tag, + int max_images, Tensor bad_color) override { + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(GetWallTime()); + TF_RETURN_IF_ERROR(AddTensorAsImageToSummary(t, tag, max_images, bad_color, + e->mutable_summary())); + return WriteEvent(std::move(e)); + } + + Status WriteAudio(int64 global_step, Tensor t, const string& tag, + int max_outputs, float sample_rate) override { + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(GetWallTime()); + TF_RETURN_IF_ERROR(AddTensorAsAudioToSummary( + t, tag, max_outputs, sample_rate, e->mutable_summary())); + return WriteEvent(std::move(e)); + } + + Status WriteGraph(int64 global_step, + std::unique_ptr graph) override { + std::unique_ptr e{new Event}; + e->set_step(global_step); + e->set_wall_time(GetWallTime()); + graph->SerializeToString(e->mutable_graph_def()); + return WriteEvent(std::move(e)); + } + + Status WriteEvent(std::unique_ptr event) override { + mutex_lock ml(mu_); + queue_.emplace_back(std::move(event)); + if (queue_.size() >= max_queue_ || + env_->NowMicros() - last_flush_ > 1000 * flush_millis_) { + return InternalFlush(); + } + return Status::OK(); + } + + string DebugString() override { return "SummaryFileWriter"; } + + private: + double GetWallTime() { + return static_cast(env_->NowMicros()) / 1.0e6; + } + + Status InternalFlush() EXCLUSIVE_LOCKS_REQUIRED(mu_) { + for (const std::unique_ptr& e : queue_) { + events_writer_->WriteEvent(*e); + } + queue_.clear(); + TF_RETURN_WITH_CONTEXT_IF_ERROR(events_writer_->Flush(), + "Could not flush events file."); + last_flush_ = env_->NowMicros(); + return Status::OK(); + } + + bool is_initialized_; + const int max_queue_; + const int flush_millis_; + uint64 last_flush_; + Env* env_; + mutex mu_; + std::vector> queue_ GUARDED_BY(mu_); + // A pointer to allow deferred construction. + std::unique_ptr events_writer_ GUARDED_BY(mu_); + std::vector> registered_summaries_ + GUARDED_BY(mu_); +}; + +} // namespace + +Status CreateSummaryFileWriter(int max_queue, int flush_millis, + const string& logdir, + const string& filename_suffix, Env* env, + SummaryWriterInterface** result) { + SummaryFileWriter* w = new SummaryFileWriter(max_queue, flush_millis, env); + const Status s = w->Initialize(logdir, filename_suffix); + if (!s.ok()) { + w->Unref(); + *result = nullptr; + return s; + } + *result = w; + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/contrib/tensorboard/db/summary_file_writer.h b/tensorflow/contrib/tensorboard/db/summary_file_writer.h new file mode 100644 index 0000000000000000000000000000000000000000..73b0a5542beabdc460c32156dd44aacc5f08610a --- /dev/null +++ b/tensorflow/contrib/tensorboard/db/summary_file_writer.h @@ -0,0 +1,43 @@ +/* Copyright 2017 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_TENSORBOARD_DB_SUMMARY_FILE_WRITER_H_ +#define TENSORFLOW_CONTRIB_TENSORBOARD_DB_SUMMARY_FILE_WRITER_H_ + +#include "tensorflow/core/kernels/summary_interface.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { + +/// \brief Creates SummaryWriterInterface which writes to a file. +/// +/// The file is an append-only records file of tf.Event protos. That +/// makes this summary writer suitable for file systems like GCS. +/// +/// It will enqueue up to max_queue summaries, and flush at least every +/// flush_millis milliseconds. The summaries will be written to the +/// directory specified by logdir and with the filename suffixed by +/// filename_suffix. The caller owns a reference to result if the +/// returned status is ok. The Env object must not be destroyed until +/// after the returned writer. +Status CreateSummaryFileWriter(int max_queue, int flush_millis, + const string& logdir, + const string& filename_suffix, Env* env, + SummaryWriterInterface** result); + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_TENSORBOARD_DB_SUMMARY_FILE_WRITER_H_ diff --git a/tensorflow/core/kernels/summary_interface_test.cc b/tensorflow/contrib/tensorboard/db/summary_file_writer_test.cc similarity index 93% rename from tensorflow/core/kernels/summary_interface_test.cc rename to tensorflow/contrib/tensorboard/db/summary_file_writer_test.cc index 58e021a0b3e889ce1efe1bb5c73bcc74e16db139..c61b4655961664a6c9c22a5f6d6f26a55c34bfcd 100644 --- a/tensorflow/core/kernels/summary_interface_test.cc +++ b/tensorflow/contrib/tensorboard/db/summary_file_writer_test.cc @@ -12,7 +12,7 @@ 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/kernels/summary_interface.h" +#include "tensorflow/contrib/tensorboard/db/summary_file_writer.h" #include "tensorflow/core/framework/summary.pb.h" #include "tensorflow/core/lib/core/errors.h" @@ -37,7 +37,7 @@ class FakeClockEnv : public EnvWrapper { uint64 current_millis_; }; -class SummaryInterfaceTest : public ::testing::Test { +class SummaryFileWriterTest : public ::testing::Test { protected: Status SummaryTestHelper( const string& test_name, @@ -47,8 +47,8 @@ class SummaryInterfaceTest : public ::testing::Test { CHECK(tests->insert(test_name).second) << ": " << test_name; SummaryWriterInterface* writer; - TF_CHECK_OK(CreateSummaryWriter(1, 1, testing::TmpDir(), test_name, &env_, - &writer)); + TF_CHECK_OK(CreateSummaryFileWriter(1, 1, testing::TmpDir(), test_name, + &env_, &writer)); core::ScopedUnref deleter(writer); TF_CHECK_OK(writer_fn(writer)); @@ -87,7 +87,7 @@ class SummaryInterfaceTest : public ::testing::Test { FakeClockEnv env_; }; -TEST_F(SummaryInterfaceTest, WriteTensor) { +TEST_F(SummaryFileWriterTest, WriteTensor) { TF_CHECK_OK(SummaryTestHelper("tensor_test", [](SummaryWriterInterface* writer) { Tensor one(DT_FLOAT, TensorShape({})); @@ -105,7 +105,7 @@ TEST_F(SummaryInterfaceTest, WriteTensor) { })); } -TEST_F(SummaryInterfaceTest, WriteScalar) { +TEST_F(SummaryFileWriterTest, WriteScalar) { TF_CHECK_OK(SummaryTestHelper( "scalar_test", [](SummaryWriterInterface* writer) { @@ -123,7 +123,7 @@ TEST_F(SummaryInterfaceTest, WriteScalar) { })); } -TEST_F(SummaryInterfaceTest, WriteHistogram) { +TEST_F(SummaryFileWriterTest, WriteHistogram) { TF_CHECK_OK(SummaryTestHelper("hist_test", [](SummaryWriterInterface* writer) { Tensor one(DT_FLOAT, TensorShape({})); @@ -141,7 +141,7 @@ TEST_F(SummaryInterfaceTest, WriteHistogram) { })); } -TEST_F(SummaryInterfaceTest, WriteImage) { +TEST_F(SummaryFileWriterTest, WriteImage) { TF_CHECK_OK(SummaryTestHelper( "image_test", [](SummaryWriterInterface* writer) { @@ -162,7 +162,7 @@ TEST_F(SummaryInterfaceTest, WriteImage) { })); } -TEST_F(SummaryInterfaceTest, WriteAudio) { +TEST_F(SummaryFileWriterTest, WriteAudio) { TF_CHECK_OK(SummaryTestHelper( "audio_test", [](SummaryWriterInterface* writer) { @@ -180,7 +180,7 @@ TEST_F(SummaryInterfaceTest, WriteAudio) { })); } -TEST_F(SummaryInterfaceTest, WriteEvent) { +TEST_F(SummaryFileWriterTest, WriteEvent) { TF_CHECK_OK( SummaryTestHelper("event_test", [](SummaryWriterInterface* writer) { @@ -198,7 +198,7 @@ TEST_F(SummaryInterfaceTest, WriteEvent) { })); } -TEST_F(SummaryInterfaceTest, WallTime) { +TEST_F(SummaryFileWriterTest, WallTime) { env_.AdvanceByMillis(7023); TF_CHECK_OK(SummaryTestHelper( "wall_time_test", diff --git a/tensorflow/contrib/tensorboard/db/vacuum.cc b/tensorflow/contrib/tensorboard/db/vacuum.cc new file mode 100644 index 0000000000000000000000000000000000000000..5febe63f0612046f96b89053811952e67d4c449b --- /dev/null +++ b/tensorflow/contrib/tensorboard/db/vacuum.cc @@ -0,0 +1,137 @@ +/* Copyright 2017 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 + +#include "tensorflow/core/lib/db/sqlite.h" +#include "tensorflow/core/platform/init_main.h" +#include "tensorflow/core/util/command_line_flags.h" + +namespace tensorflow { +namespace { + +void Vacuum(const char* path) { + LOG(INFO) << "Opening SQLite DB: " << path; + Sqlite* db; + TF_CHECK_OK(Sqlite::Open(path, SQLITE_OPEN_READWRITE, &db)); + core::ScopedUnref db_unref(db); + + // TODO(jart): Maybe defragment rowids on Tensors. + // TODO(jart): Maybe LIMIT deletes and incremental VACUUM. + + // clang-format off + + LOG(INFO) << "Deleting orphaned Experiments"; + db->PrepareOrDie(R"sql( + DELETE FROM + Experiments + WHERE + user_id IS NOT NULL + AND user_id NOT IN (SELECT user_id FROM Users) + )sql").StepAndResetOrDie(); + + LOG(INFO) << "Deleting orphaned Runs"; + db->PrepareOrDie(R"sql( + DELETE FROM + Runs + WHERE + experiment_id IS NOT NULL + AND experiment_id NOT IN (SELECT experiment_id FROM Experiments) + )sql").StepAndResetOrDie(); + + LOG(INFO) << "Deleting orphaned Tags"; + db->PrepareOrDie(R"sql( + DELETE FROM + Tags + WHERE + run_id IS NOT NULL + AND run_id NOT IN (SELECT run_id FROM Runs) + )sql").StepAndResetOrDie(); + + // TODO(jart): What should we do if plugins define non-tag tensor series? + LOG(INFO) << "Deleting orphaned Tensors"; + db->PrepareOrDie(R"sql( + DELETE FROM + Tensors + WHERE + series IS NOT NULL + AND series NOT IN (SELECT tag_id FROM Tags) + )sql").StepAndResetOrDie(); + + LOG(INFO) << "Deleting orphaned TensorStrings"; + db->PrepareOrDie(R"sql( + DELETE FROM + TensorStrings + WHERE + tensor_rowid NOT IN (SELECT rowid FROM Tensors) + )sql").StepAndResetOrDie(); + + LOG(INFO) << "Deleting orphaned Graphs"; + db->PrepareOrDie(R"sql( + DELETE FROM + Graphs + WHERE + run_id IS NOT NULL + AND run_id NOT IN (SELECT run_id FROM Runs) + )sql").StepAndResetOrDie(); + + LOG(INFO) << "Deleting orphaned Nodes"; + db->PrepareOrDie(R"sql( + DELETE FROM + Nodes + WHERE + graph_id NOT IN (SELECT graph_id FROM Graphs) + )sql").StepAndResetOrDie(); + + LOG(INFO) << "Deleting orphaned NodeInputs"; + db->PrepareOrDie(R"sql( + DELETE FROM + NodeInputs + WHERE + graph_id NOT IN (SELECT graph_id FROM Graphs) + )sql").StepAndResetOrDie(); + + LOG(INFO) << "Running VACUUM"; + db->PrepareOrDie("VACUUM").StepAndResetOrDie(); + + // clang-format on +} + +int main(int argc, char* argv[]) { + string usage = Flags::Usage(argv[0], {}); + bool parse_result = Flags::Parse(&argc, argv, {}); + if (!parse_result) { + std::cerr << "The vacuum tool rebuilds SQLite database files created by\n" + << "SummaryDbWriter, which makes them smaller.\n\n" + << "This means deleting orphaned rows and rebuilding b-tree\n" + << "pages so empty space from deleted rows is cleared. Any\n" + << "superfluous padding of Tensor BLOBs is also removed.\n\n" + << usage; + return -1; + } + port::InitMain(argv[0], &argc, &argv); + if (argc < 2 || argv[1][0] == '-') { + std::cerr << "Need at least one SQLite DB path.\n"; + return -1; + } + for (int i = 1; i < argc; ++i) { + Vacuum(argv[i]); + } + return 0; +} + +} // namespace +} // namespace tensorflow + +int main(int argc, char* argv[]) { return tensorflow::main(argc, argv); } diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..cf67c27b70f1a8c761b71074d3eb5cd962a68488 --- /dev/null +++ b/tensorflow/contrib/tensorrt/BUILD @@ -0,0 +1,246 @@ +# Description: +# Wrap NVIDIA TensorRT (http://developer.nvidia.com/tensorrt) with tensorflow +# and provide TensorRT operators and converter package. +# APIs are meant to change over time. + +package(default_visibility = ["//tensorflow:__subpackages__"]) + +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_test", + "tf_copts", + "tf_cuda_library", + "tf_custom_op_library", + "tf_custom_op_library_additional_deps", + "tf_gen_op_libs", + "tf_gen_op_wrapper_py", +) +load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test") +load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") +load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc") +load( + "@local_config_tensorrt//:build_defs.bzl", + "if_tensorrt", +) + +tf_cuda_cc_test( + name = "tensorrt_test_cc", + size = "small", + srcs = ["tensorrt_test.cc"], + tags = [ + "manual", + "notap", + ], + deps = [ + "//tensorflow/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ] + if_tensorrt([ + "@local_config_cuda//cuda:cuda_headers", + "@local_config_tensorrt//:nv_infer", + ]), +) + +tf_custom_op_library( + name = "python/ops/_trt_engine_op.so", + srcs = ["ops/trt_engine_op.cc"], + deps = [ + ":trt_engine_op_kernel", + ":trt_shape_function", + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]), +) + +tf_cuda_library( + name = "trt_shape_function", + srcs = ["shape_fn/trt_shfn.cc"], + hdrs = ["shape_fn/trt_shfn.h"], + visibility = ["//visibility:public"], + deps = [ + ":trt_logging", + ] + if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]) + tf_custom_op_library_additional_deps(), +) + +cc_library( + name = "trt_engine_op_kernel", + srcs = ["kernels/trt_engine_op.cc"], + hdrs = ["kernels/trt_engine_op.h"], + copts = tf_copts(), + deps = [ + ":trt_logging", + "//tensorflow/core:gpu_headers_lib", + "//tensorflow/core:lib_proto_parsing", + "//tensorflow/core:stream_executor_headers_lib", + ] + if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]) + tf_custom_op_library_additional_deps(), + alwayslink = 1, +) + +tf_gen_op_libs( + op_lib_names = ["trt_engine_op"], + deps = if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]), +) + +tf_cuda_library( + name = "trt_logging", + srcs = ["log/trt_logger.cc"], + hdrs = ["log/trt_logger.h"], + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]), +) + +tf_gen_op_wrapper_py( + name = "trt_engine_op", + deps = [ + ":trt_engine_op_op_lib", + ":trt_logging", + ":trt_shape_function", + ], +) + +tf_custom_op_py_library( + name = "trt_engine_op_loader", + srcs = ["python/ops/trt_engine_op.py"], + dso = [ + ":python/ops/_trt_engine_op.so", + ] + if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]), + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:resources", + ], +) + +py_library( + name = "init_py", + srcs = [ + "__init__.py", + "python/__init__.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":trt_convert_py", + ":trt_ops_py", + ], +) + +py_library( + name = "trt_ops_py", + srcs_version = "PY2AND3", + deps = [ + ":trt_engine_op", + ":trt_engine_op_loader", + ], +) + +py_library( + name = "trt_convert_py", + srcs = ["python/trt_convert.py"], + srcs_version = "PY2AND3", + deps = [ + ":wrap_conversion", + ], +) + +tf_py_wrap_cc( + name = "wrap_conversion", + srcs = ["trt_conversion.i"], + copts = tf_copts(), + deps = [ + ":trt_conversion", + "//tensorflow/core:framework_lite", + "//util/python:python_headers", + ], +) + +# Library for the node-level conversion portion of TensorRT operation creation +tf_cuda_library( + name = "trt_conversion", + srcs = [ + "convert/convert_graph.cc", + "convert/convert_nodes.cc", + ], + hdrs = [ + "convert/convert_graph.h", + "convert/convert_nodes.h", + ], + deps = [ + ":segment", + ":trt_logging", + "//tensorflow/core/grappler:grappler_item", + "//tensorflow/core/grappler:utils", + "//tensorflow/core:framework", + "//tensorflow/core:framework_lite", + "//tensorflow/core:graph", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core/grappler:devices", + "//tensorflow/core/grappler/clusters:virtual_cluster", + "//tensorflow/core/grappler/costs:graph_properties", + "//tensorflow/core/grappler/optimizers:constant_folding", + "//tensorflow/core/grappler/optimizers:layout_optimizer", + ] + if_tensorrt([ + "@local_config_tensorrt//:nv_infer", + ]) + tf_custom_op_library_additional_deps(), +) + +# Library for the segmenting portion of TensorRT operation creation +cc_library( + name = "segment", + srcs = ["segment/segment.cc"], + hdrs = [ + "segment/segment.h", + "segment/union_find.h", + ], + linkstatic = 1, + deps = [ + "//tensorflow/core:graph", + "//tensorflow/core:lib_proto_parsing", + "//tensorflow/core:protos_all_cc", + "@protobuf_archive//:protobuf_headers", + ], +) + +tf_cc_test( + name = "segment_test", + size = "small", + srcs = ["segment/segment_test.cc"], + deps = [ + ":segment", + "//tensorflow/c:c_api", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/contrib/tensorrt/README.md b/tensorflow/contrib/tensorrt/README.md new file mode 100644 index 0000000000000000000000000000000000000000..dfcce0fd00eedf3341850bbc23927dc3b2e2d2aa --- /dev/null +++ b/tensorflow/contrib/tensorrt/README.md @@ -0,0 +1,40 @@ +Using TensorRT in TensorFlow +============================ + +This module provides necessary bindings and introduces TRT_engine_op +operator that wraps a subgraph in TensorRT. + +Compilation +----------- + +In order to compile the module, you need to have a local TensorRT +installation (libnvinfer.so and respective include files). During the +configuration step, TensorRT should be enabled and installation path +should be set. If installed through package managers (deb,rpm), +configure script should find the necessary components from the system +automatically. If installed from tar packages, user has to set path to +location where the library is installed during configuration. + + +``` +bazel build --config=cuda --config=opt //tensorflow/tools/pip_package:build_pip_package +bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/ +``` + +After the installation of tensorflow package, TensorRT transformation +will be available. An example use is shown below. + +```python +import tensorflow as tf +import tensorflow.contrib.tensorrt as trt +#... create and train or load model +gdef = sess.graph.as_graph_def() +trt_gdef = trt.create_inference_graph( + gdef, #original graph_def + ["output"], #name of output node(s) + max_batch_size, #maximum batch size to run the inference + max_workspace_size_bytes) # max memory for TensorRT to use +tf.reset_default_graph() +tf.import_graph_def(graph_def=trt_gdef) +#...... run inference +``` diff --git a/tensorflow/contrib/tensorrt/__init__.py b/tensorflow/contrib/tensorrt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fd551d70b4385b14b84b7b98a6d16b0c03733d38 --- /dev/null +++ b/tensorflow/contrib/tensorrt/__init__.py @@ -0,0 +1,23 @@ +# 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. +# ============================================================================= +"""Exposes the python wrapper for TensorRT graph transforms.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# pylint: disable=unused-import,wildcard-import +from tensorflow.contrib.tensorrt.python import * +# pylint: enable=unused-import,wildcard-import diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc new file mode 100644 index 0000000000000000000000000000000000000000..899448004f917b36b35fb871a66a9d857736a338 --- /dev/null +++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc @@ -0,0 +1,273 @@ +/* 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/convert/convert_graph.h" + +#include +#include +#include +#include +#include + +#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" +#include "tensorflow/contrib/tensorrt/segment/segment.h" +#include "tensorflow/core/graph/algorithm.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/graph_constructor.h" +#include "tensorflow/core/grappler/clusters/virtual_cluster.h" +#include "tensorflow/core/grappler/costs/graph_properties.h" +#include "tensorflow/core/grappler/devices.h" +#include "tensorflow/core/grappler/grappler_item.h" +#include "tensorflow/core/grappler/optimizers/constant_folding.h" +#include "tensorflow/core/grappler/optimizers/layout_optimizer.h" +#include "tensorflow/core/grappler/utils.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/protobuf/device_properties.pb.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { +namespace convert { +namespace { + +static bool IsTensorRTCandidate(const tensorflow::NodeDef& node_def) { + // LINT.IfChange + // TODO(jie): Segmentation shouldn't associated with op name. + // Split it into a registration for each kernel. + static const std::set candidate_ops = { + "Identity", "Const", "Conv2D", "MaxPool", "BiasAdd", "Relu", + "Add", "Mul", "Sub", "Rsqrt", "Pad" // "Placeholder" ,"Mean" + }; + // LINT.ThenChange(//tensorflow/contrib/tensorrt/convert/convert_nodes.h) + return candidate_ops.count(node_def.op()); +} + +void GetSubGraphIncomingEdges(const tensorflow::Graph& graph, + const std::set& subgraph_node_ids, + tensorflow::EdgeSet* incoming_edges) { + for (int node_id : subgraph_node_ids) { + const tensorflow::Node* node = graph.FindNodeId(node_id); + for (const tensorflow::Edge* edge : node->in_edges()) { + if (!subgraph_node_ids.count(edge->src()->id()) && + !edge->src()->IsSource()) { + incoming_edges->insert(edge); + } + } + } +} + +void GetSubGraphOutgoingEdges(const tensorflow::Graph& graph, + const std::set& subgraph_node_ids, + tensorflow::EdgeSet* outgoing_edges) { + for (int node_id : subgraph_node_ids) { + const tensorflow::Node* node = graph.FindNodeId(node_id); + for (const tensorflow::Edge* edge : node->out_edges()) { + if (!subgraph_node_ids.count(edge->dst()->id()) && + !edge->dst()->IsSink()) { + outgoing_edges->insert(edge); + } + } + } +} + +std::pair ParseTensorName(string name, int default_idx = 0) { + int idx = default_idx; + size_t sep = name.find_last_of(':'); + if (sep != string::npos) { + name = name.substr(0, sep); + idx = std::stoi(name.substr(sep + 1)); + } + return std::make_pair(name, idx); +} + +std::unordered_map> BuildTensorNameMap( + const std::vector& tensor_names) { + std::unordered_map> result; + for (string const& tensor_name : tensor_names) { + string node_name; + int index; + std::tie(node_name, index) = ParseTensorName(tensor_name); + result[node_name].push_back(index); + } + return result; +} + +tensorflow::Status ConvertSubGraphToTensorRT( + const std::vector& output_names, + const std::set& subgraph_node_ids, + size_t max_batch_size, // Max batch size that engine will be created for + // Max amount of memory that engine will be allowed to consume, in bytes + size_t max_workspace_size_bytes, + const tensorflow::grappler::GraphProperties& graph_properties, + tensorflow::Graph* graph) { + tensorflow::EdgeSet subgraph_incoming_edges; + GetSubGraphIncomingEdges(*graph, subgraph_node_ids, &subgraph_incoming_edges); + + std::vector> subgraph_inputs; + + // Collect inputs by looking for incoming edges + for (const tensorflow::Edge* edge : subgraph_incoming_edges) { + subgraph_inputs.push_back({edge->src()->id(), edge->src_output()}); + } + std::set> subgraph_outputs_set; + // Collect outputs referenced from output_names + auto output_name_to_index_map = BuildTensorNameMap(output_names); + for (int node_id : subgraph_node_ids) { + tensorflow::Node* node = graph->FindNodeId(node_id); + if (output_name_to_index_map.count(node->name())) { + for (int index : output_name_to_index_map.at(node->name())) { + subgraph_outputs_set.insert({node_id, index}); + } + } + } + // Collect outputs referenced from outgoing edges + tensorflow::EdgeSet subgraph_outgoing_edges; + GetSubGraphOutgoingEdges(*graph, subgraph_node_ids, &subgraph_outgoing_edges); + for (const tensorflow::Edge* edge : subgraph_outgoing_edges) { + subgraph_outputs_set.insert({edge->src()->id(), edge->src_output()}); + } + // Impose an ordering on the outputs + std::vector> subgraph_outputs( + subgraph_outputs_set.begin(), subgraph_outputs_set.end()); + // Build TensorRT node and add it to the graph + tensorflow::NodeDef trt_node_def; + TF_RETURN_IF_ERROR(ConvertSubGraphToTensorRTNodeDef( + *graph, subgraph_node_ids, subgraph_inputs, subgraph_outputs, + max_batch_size, max_workspace_size_bytes, graph_properties, + &trt_node_def)); + tensorflow::Status status; + tensorflow::Node* trt_node = graph->AddNode(trt_node_def, &status); + TF_RETURN_IF_ERROR(status); + + // Re-map outgoing edges to use the new TRT node instead of the orig subgraph + std::map, int> subgraph_edge_to_output_map; + for (size_t i = 0; i < subgraph_outputs.size(); ++i) { + subgraph_edge_to_output_map.insert({subgraph_outputs.at(i), i}); + } + TF_RETURN_IF_ERROR(status); + for (const tensorflow::Edge* edge : subgraph_outgoing_edges) { + std::pair old_src = {edge->src()->id(), edge->src_output()}; + int new_src_output = subgraph_edge_to_output_map.at(old_src); + TF_RETURN_IF_ERROR(graph->UpdateEdge(trt_node, new_src_output, edge->dst(), + edge->dst_input())); + } + // Remove the original subgraph + for (int node_id : subgraph_node_ids) { + tensorflow::Node* node = graph->FindNodeId(node_id); + // Don't remove the input placeholders + if (node->type_string() == "Placeholder") { + continue; + } + graph->RemoveNode(node); + } + return tensorflow::Status::OK(); +} + +tensorflow::Status BuildNodeMap( + const tensorflow::Graph& graph, + std::unordered_map* node_map) { + for (auto* node : graph.op_nodes()) { + if (!node_map->insert({node->name(), node}).second) { + return tensorflow::errors::AlreadyExists( + "Node name is not unique in graph: " + node->name()); + } + } + return tensorflow::Status::OK(); +} + +} // namespace + +tensorflow::Status ConvertGraphDefToTensorRT( + const tensorflow::GraphDef& graph_def, + const std::vector& output_names, size_t max_batch_size, + size_t max_workspace_size_bytes, tensorflow::GraphDef* new_graph_def) { + // Optimization pass + tensorflow::grappler::GrapplerItem item; + item.fetch = output_names; + tensorflow::GraphDef gdef; + + // Layout optimization + item.graph = graph_def; + tensorflow::grappler::LayoutOptimizer optimizer; + tensorflow::grappler::Cluster* cluster; + + // Virtual cluster + tensorflow::DeviceProperties device_properties; + device_properties.set_type("GPU"); + device_properties.mutable_environment()->insert({"architecture", "6"}); + cluster = + new tensorflow::grappler::VirtualCluster({{"/GPU:0", device_properties}}); + + TF_RETURN_IF_ERROR(optimizer.Optimize(cluster, item, &gdef)); + + // Constant folding + item.graph = gdef; + tensorflow::grappler::ConstantFolding fold(nullptr); + TF_RETURN_IF_ERROR(fold.Optimize(nullptr, item, &gdef)); + + // AJ refactoring shape inference through grappler/GraphProperties. + tensorflow::grappler::GraphProperties static_graph_properties(item); + TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(false)); + + // Build full graph + tensorflow::FunctionLibraryDefinition flib(tensorflow::OpRegistry::Global(), + gdef.library()); + tensorflow::Graph graph(flib); + TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToGraph( + tensorflow::GraphConstructorOptions(), gdef, &graph)); + + // Segment the graph into subgraphs that can be converted to TensorRT + tensorflow::tensorrt::segment::SegmentOptions segment_options; + + // TODO(ben,jie,sami): exclude output nodes (DISCUSS IT) + for (auto node : output_names) { + segment_options.exclude_node_list.insert(node); + } + + // TODO(sami): this should be passed as a knob!!!! + segment_options.minimum_segment_size = 2; + tensorflow::tensorrt::segment::SegmentNodesVector segments; + TF_RETURN_IF_ERROR(tensorrt::segment::SegmentGraph( + gdef, IsTensorRTCandidate, segment_options, &segments)); + if (segments.size() > 1) { + VLOG(0) << "MULTIPLE tensorrt candidate conversion: " << segments.size(); + } + std::unordered_map node_map; + TF_RETURN_IF_ERROR(BuildNodeMap(graph, &node_map)); + for (const std::set& subgraph_node_names : segments) { + std::set subgraph_node_ids; + for (const string& node_name : subgraph_node_names) { + subgraph_node_ids.insert(node_map.at(node_name)->id()); + } + TF_RETURN_IF_ERROR(ConvertSubGraphToTensorRT( + output_names, subgraph_node_ids, max_batch_size, + max_workspace_size_bytes, static_graph_properties, &graph)); + } + graph.ToGraphDef(new_graph_def); + return tensorflow::Status::OK(); +} + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.h b/tensorflow/contrib/tensorrt/convert/convert_graph.h new file mode 100644 index 0000000000000000000000000000000000000000..154ad3f2e8fb0ae702448097fbdece510df30223 --- /dev/null +++ b/tensorflow/contrib/tensorrt/convert/convert_graph.h @@ -0,0 +1,47 @@ +/* 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_CONVERT_CONVERT_GRAPH_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_GRAPH_H_ + +#include + +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/types.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +namespace tensorflow { +namespace tensorrt { +namespace convert { + +// max_batch_size: maximum batch size which can be used for inference for +// optimization targets inference run with max batch size. +// max_workspace_size_bytes: The upper bound of memory allowence for +// engine building. +tensorflow::Status ConvertGraphDefToTensorRT( + const tensorflow::GraphDef& graph_def, + const std::vector& output_names, size_t max_batch_size, + size_t max_workspace_size_bytes, tensorflow::GraphDef* new_graph_def); + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_GRAPH_H_ diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc new file mode 100644 index 0000000000000000000000000000000000000000..9ee717dd7fb1eff4a11fb104cf5806ec8ab853d2 --- /dev/null +++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc @@ -0,0 +1,1601 @@ +/* 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/convert/convert_nodes.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tensorflow/core/framework/node_def_builder.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" // NOLINT +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/graph/algorithm.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/graph_constructor.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/tensor_coding.h" +#include "tensorflow/core/platform/types.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorflow/contrib/tensorrt/log/trt_logger.h" +#include "tensorrt/include/NvInfer.h" + +// Check if the types are equal. Cast to int first so that failure log message +// would work! +#define CHECK_EQ_TYPE(val1, val2) CHECK_EQ((int)val1, (int)val2) + +namespace tensorflow { +namespace tensorrt { +namespace convert { + +namespace { + +inline tensorflow::Status ConvertDType(tensorflow::DataType tf_dtype, + nvinfer1::DataType* trt_dtype) { + switch (tf_dtype) { + case tensorflow::DataType::DT_FLOAT: + *trt_dtype = nvinfer1::DataType::kFLOAT; + break; + case tensorflow::DataType::DT_INT8: + *trt_dtype = nvinfer1::DataType::kINT8; + break; + case tensorflow::DataType::DT_HALF: + *trt_dtype = nvinfer1::DataType::kHALF; + break; + default: + return tensorflow::errors::InvalidArgument("Unsupported data type"); + } + return tensorflow::Status::OK(); +} + +inline nvinfer1::Dims GetTensorShape(const tensorflow::Tensor& tensor) { + nvinfer1::Dims dims; + dims.nbDims = tensor.dims(); + for (int i = 0; i < dims.nbDims; i++) { + dims.d[i] = tensor.dim_size(i); + } + return dims; +} + +inline int64_t GetShapeSize(nvinfer1::Dims shape) { + // Returns total number of elements in shape + int64_t count = 1; + for (int d = 0; d < shape.nbDims; ++d) { + count *= shape.d[d]; + } + return count; +} + +static std::vector> CreateSamePadding( + const nvinfer1::DimsHW& stride, const nvinfer1::DimsHW& kernel, + const std::vector& input_dims) { + std::vector> padding(input_dims.size()); + CHECK_EQ((size_t)stride.nbDims, input_dims.size()); // TODO(jie): N+C? NC+? + + for (size_t i = 0; i < input_dims.size(); ++i) { + // Formula to calculate the padding + int p = ((input_dims[i] - 1) / stride.d[i]) * stride.d[i] + kernel.d[i] - + input_dims[i]; + p = (p > 0) ? p : 0; + + // Right precedence padding, like in TensorFlow + int left = p / 2; + int right = p - left; + + VLOG(2) << "PADDING_" << i << " pre: " << left << ", post: " << right + << "paras: " << input_dims[i] << ", " << stride.d[i] << ", " + << "kernel: " << kernel.d[i]; + padding[i] = {left, right}; + } + return padding; +} + +class TRT_ShapedWeights { + public: + TRT_ShapedWeights(tensorflow::DataType type, const void* values, + nvinfer1::Dims shape) + : shape_(shape), type_(type), values_(values), empty_weight_flag_(false) { + // Note: this->shape.type[] is not used + } + + explicit TRT_ShapedWeights(tensorflow::DataType type) + : shape_(), type_(type), values_(nullptr), empty_weight_flag_(true) {} + + TRT_ShapedWeights(const TRT_ShapedWeights& rhs) + : shape_(rhs.shape_), + type_(rhs.type_), + values_(rhs.values_), + empty_weight_flag_(rhs.empty_weight_flag_) {} + + int64_t count() const { + int64_t c = 1; + for (int i = 0; i < shape_.nbDims; i++) c *= shape_.d[i]; + return c; + } + + nvinfer1::Weights GetWeightsForTRT() const { + nvinfer1::DataType trt_type(nvinfer1::DataType::kFLOAT); + TF_CHECK_OK(ConvertDType(type_, &trt_type)); + if (empty_weight_flag_) return nvinfer1::Weights{trt_type, nullptr, 0}; + + // Note: this->shape.type[] is not used + return nvinfer1::Weights{trt_type, GetValues(), GetShapeSize(shape_)}; + } + + const void* GetValues() const { return values_; } + + void SetValues(const void* values) { values_ = values; } + + size_t size_bytes() const { + int type_size = tensorflow::DataTypeSize(this->type_); + return this->count() * type_size; + } + + // Default converter + operator nvinfer1::Weights() const { return GetWeightsForTRT(); } + + nvinfer1::Dims shape_; + tensorflow::DataType type_; + + private: + const void* values_; + bool empty_weight_flag_; +}; + +class TRT_TensorOrWeights { + public: + explicit TRT_TensorOrWeights(nvinfer1::ITensor* tensor) + : tensor_(tensor), weights_(DT_FLOAT), variant_(TRT_NODE_TENSOR) {} + explicit TRT_TensorOrWeights(const TRT_ShapedWeights& weights) + : tensor_(nullptr), weights_(weights), variant_(TRT_NODE_WEIGHTS) {} + TRT_TensorOrWeights(const TRT_TensorOrWeights& rhs) + : tensor_(rhs.tensor_), weights_(rhs.weights_), variant_(rhs.variant_) {} + ~TRT_TensorOrWeights() {} + + bool is_tensor() const { return variant_ == TRT_NODE_TENSOR; } + bool is_weights() const { return variant_ == TRT_NODE_WEIGHTS; } + + nvinfer1::ITensor* tensor() { + CHECK_EQ(is_tensor(), true); + return tensor_; + } + const nvinfer1::ITensor* tensor() const { + CHECK_EQ(is_tensor(), true); + return tensor_; + } + TRT_ShapedWeights& weights() { + CHECK_EQ(is_weights(), true); + return weights_; + } + const TRT_ShapedWeights& weights() const { + CHECK_EQ(is_weights(), true); + return weights_; + } + nvinfer1::Dims shape() const { + if (is_tensor()) { + return tensor()->getDimensions(); + } else { + return weights().shape_; + } + } + + private: + nvinfer1::ITensor* tensor_; + TRT_ShapedWeights weights_; + enum { TRT_NODE_TENSOR, TRT_NODE_WEIGHTS } variant_; +}; + +class TFAttrs { + public: + explicit TFAttrs(const tensorflow::NodeDef& tf_node) { + for (const auto& attr : tf_node.attr()) { + attrs_.insert({attr.first, &attr.second}); + } + } + bool count(string key) const { return attrs_.count(key); } + tensorflow::AttrValue const* at(string key) const { + if (!attrs_.count(key)) { + LOG(FATAL) << "Attribute not found: " << key; + } + return attrs_.at(key); + } + template + T get(string key) const; + template + T get(string key, const T& default_value) const { + return attrs_.count(key) ? this->get(key) : default_value; + } + + private: + typedef std::map AttrMap; + AttrMap attrs_; +}; + +template <> +string TFAttrs::get(string key) const { + return this->at(key)->s(); +} + +template <> +std::vector TFAttrs::get>(string key) const { + auto attr = this->at(key)->list().i(); + return std::vector(attr.begin(), attr.end()); +} + +template <> +nvinfer1::Dims TFAttrs::get(string key) const { + auto values = this->get>(key); + nvinfer1::Dims dims; + dims.nbDims = values.size(); + std::copy(values.begin(), values.end(), dims.d); + // Note: No dimension type information is included + return dims; +} + +template <> +nvinfer1::DataType TFAttrs::get(string key) const { + nvinfer1::DataType trt_dtype(nvinfer1::DataType::kFLOAT); + TF_CHECK_OK(ConvertDType(this->at(key)->type(), &trt_dtype)); + return trt_dtype; +} + +template <> +tensorflow::DataType TFAttrs::get(string key) const { + return this->at(key)->type(); +} + +template +void Reorder4(nvinfer1::DimsNCHW shape, const T* idata, + nvinfer1::DimsNCHW istrides, T* odata, + nvinfer1::DimsNCHW ostrides) { + for (int n = 0; n < shape.n(); ++n) { + for (int c = 0; c < shape.c(); ++c) { + for (int h = 0; h < shape.h(); ++h) { + for (int w = 0; w < shape.w(); ++w) { + odata[n * ostrides.n() + c * ostrides.c() + h * ostrides.h() + + w * ostrides.w()] = idata[n * istrides.n() + c * istrides.c() + + h * istrides.h() + w * istrides.w()]; + } + } + } + } +} + +void ReorderRSCKToKCRS(const TRT_ShapedWeights& iweights, + TRT_ShapedWeights* oweights) { + CHECK_EQ(iweights.type_, oweights->type_); + CHECK_EQ(iweights.size_bytes(), oweights->size_bytes()); + int r = iweights.shape_.d[0]; + int s = iweights.shape_.d[1]; + int c = iweights.shape_.d[2]; + int k = iweights.shape_.d[3]; + oweights->shape_.d[0] = k; + oweights->shape_.d[1] = c; + oweights->shape_.d[2] = r; + oweights->shape_.d[3] = s; + nvinfer1::DimsNCHW istrides = {1, k, s * k * c, c * k}; + nvinfer1::DimsNCHW ostrides = {c * r * s, r * s, s, 1}; + switch (iweights.type_) { + case tensorflow::DataType::DT_FLOAT: + Reorder4({k, c, r, s}, static_cast(iweights.GetValues()), + istrides, + static_cast(const_cast(oweights->GetValues())), + ostrides); + break; + default: + LOG(FATAL) << "!!!!!!!!!!!!!!!!!!!!!!!!broke!!!!!!!!!!!!"; + } +} + +struct InferDeleter { + template + void operator()(T* obj) const { + if (obj) { + obj->destroy(); + } + } +}; + +template +inline std::shared_ptr infer_object(T* obj) { + return std::shared_ptr(obj, InferDeleter()); +} + +// Logger for GIE info/warning/errors +class Converter; + +using OpConverter = + std::function const&, + std::vector*)>; + +class Converter { + std::unordered_map trt_tensors_; + std::unordered_map op_registry_; + nvinfer1::INetworkDefinition* trt_network_; + std::list> temp_bufs_; + + void register_op_converters(); + + std::vector get_inputs( + const tensorflow::NodeDef& node_def) { + std::vector inputs; + for (const auto& input_name : node_def.input()) { + VLOG(2) << "Retrieve input: " << input_name; + inputs.push_back(trt_tensors_.at(input_name)); + } + return inputs; + } + + public: + explicit Converter(nvinfer1::INetworkDefinition* trt_network) + : trt_network_(trt_network) { + this->register_op_converters(); + } + + TRT_ShapedWeights get_temp_weights(tensorflow::DataType type, + nvinfer1::Dims shape) { + TRT_ShapedWeights weights(type, nullptr, shape); + // TODO(jie): check weights size_bytes. 0 means type error + temp_bufs_.push_back(std::vector(weights.size_bytes())); + weights.SetValues(temp_bufs_.back().data()); + return weights; + } + + TRT_ShapedWeights get_temp_weights_like(const TRT_ShapedWeights& weights) { + return this->get_temp_weights(weights.type_, weights.shape_); + } + + tensorflow::Status convert_node(const tensorflow::NodeDef& node_def) { + std::vector inputs = this->get_inputs(node_def); + string op = node_def.op(); + if (!op_registry_.count(op)) { + return tensorflow::errors::Unimplemented( + "No converter registered for op: " + op); + } + OpConverter op_converter = op_registry_.at(op); + std::vector outputs; + TF_RETURN_IF_ERROR(op_converter(*this, node_def, inputs, &outputs)); + for (size_t i = 0; i < outputs.size(); ++i) { + TRT_TensorOrWeights output = outputs.at(i); + // TODO(jie): tf protobuf seems to be omitting the :0 suffix + string output_name = node_def.name(); + if (i != 0) output_name = output_name + ":" + std::to_string(i); + if (output.is_tensor()) { + output.tensor()->setName(output_name.c_str()); + } + VLOG(2) << "Write out tensor: " << output_name; + if (!trt_tensors_.insert({output_name, output}).second) { + return tensorflow::errors::AlreadyExists( + "Output tensor already exists for op: " + op); + } + } + return tensorflow::Status::OK(); + } + + nvinfer1::INetworkDefinition* network() { return trt_network_; } + + TRT_TensorOrWeights get_tensor(string name) { + if (!trt_tensors_.count(name)) { + return TRT_TensorOrWeights(nullptr); + } + return trt_tensors_.at(name); + } + + bool insert_input_tensor(string name, nvinfer1::ITensor* tensor) { + return trt_tensors_.insert({name, TRT_TensorOrWeights(tensor)}).second; + } + + nvinfer1::ITensor* TransposeTensor(nvinfer1::ITensor* input_tensor, + std::vector order) { + auto dims = input_tensor->getDimensions(); + + // TODO(jie): change the return to status and properly exit + if (order.size() - 1 != size_t(dims.nbDims)) + LOG(ERROR) << "Dimension does not match, fail gracefully"; + + nvinfer1::IShuffleLayer* layer = this->network()->addShuffle(*input_tensor); + nvinfer1::Permutation permutation; + for (int32_t i = 0; i < dims.nbDims; ++i) { + permutation.order[i] = order[i + 1] - 1; + } + layer->setFirstTranspose(permutation); + + nvinfer1::Dims reshape_dims; + reshape_dims.nbDims = dims.nbDims; + for (int32_t i = 0; i < reshape_dims.nbDims; ++i) { + reshape_dims.d[i] = 0; + reshape_dims.type[i] = dims.type[i]; + } + layer->setReshapeDimensions(reshape_dims); + return layer->getOutput(0); + } +}; + +// **************************************************************************** +// Constant folding functions +// TODO(jie): once optimizer kicks in, we should have done constant folding +// there. +//*****************************************************************************/ +struct LambdaFactory { + enum class OP_CATEGORY : int { RSQRT = 0, NEG, ADD, MUL, SUB }; + OP_CATEGORY op; + + template + std::function unary() { + switch (op) { + case OP_CATEGORY::RSQRT: { + VLOG(2) << "RSQRT GETS DONE"; + return [](T t) -> T { return 1.0 / std::sqrt(t); }; + } + case OP_CATEGORY::NEG: + return [](T t) -> T { return -t; }; + default: + VLOG(2) << "Not supported op for unary: " << static_cast(op); + return nullptr; + } + } + + template + std::function binary() { + switch (op) { + case OP_CATEGORY::ADD: + return [](T l, T r) -> T { return l + r; }; + case OP_CATEGORY::SUB: + return [](T l, T r) -> T { return l - r; }; + case OP_CATEGORY::MUL: + return [](T l, T r) -> T { return l * r; }; + default: + LOG(WARNING) << "Not supported op for binary: " << static_cast(op); + } + return [](T l, T r) -> T { + LOG(FATAL) << "Unsupported op type "; + return l; + }; + } + + template + std::function broadcast_r(T val) { + VLOG(2) << "LAMBDA VAL : " << val; + switch (op) { + case OP_CATEGORY::ADD: + return [val](T l) -> T { + VLOG(2) << "LAMBDA VAL : " << val; + return l + val; + }; + // Return [val](T l)-> T {return l+val;}; + case OP_CATEGORY::SUB: + return [val](T l) -> T { + VLOG(2) << "LAMBDA VAL : " << val; + return l - val; + }; + case OP_CATEGORY::MUL: + return [val](T l) -> T { + VLOG(2) << "LAMBDA VAL : " << val; + return l * val; + }; + default: + LOG(WARNING) << "Not supported op for binary: " << static_cast(op); + } + return [val](T l) -> T { + LOG(FATAL) << "Unsupported op type "; + return l; + }; + } + + template + std::function broadcast_l(T val) { + VLOG(2) << "LAMBDA VAL : " << val; + switch (op) { + case OP_CATEGORY::ADD: + return [val](T l) -> T { + VLOG(2) << "LAMBDA VAL : " << val; + return val + l; + }; + case OP_CATEGORY::SUB: + return [val](T l) -> T { + VLOG(2) << "LAMBDA VAL : " << val; + return val - l; + }; + case OP_CATEGORY::MUL: + return [val](T l) -> T { + VLOG(2) << "LAMBDA VAL : " << val; + return val * l; + }; + default: + LOG(ERROR) << "Not supported op for binary: " << static_cast(op); + } + return [val](T l) -> T { + LOG(FATAL) << "Unsupported op type "; + return l; + }; + } +}; + +tensorflow::Status UnaryCompute(const TRT_ShapedWeights& iweights, + TRT_ShapedWeights* oweights, + LambdaFactory unary_op) { + CHECK_EQ(iweights.type_, oweights->type_); + switch (iweights.type_) { + case tensorflow::DataType::DT_FLOAT: { + auto inp = static_cast(iweights.GetValues()); + auto oup = static_cast(const_cast(oweights->GetValues())); + std::transform(inp, inp + iweights.count(), oup, unary_op.unary()); + break; + } + default: + return tensorflow::errors::Unimplemented( + "Data type not supported: " + + tensorflow::DataTypeString(iweights.type_)); + } + return tensorflow::Status::OK(); +} + +tensorflow::Status BinaryCompute(const TRT_ShapedWeights& iweights_l, + const TRT_ShapedWeights& iweights_r, + TRT_ShapedWeights* oweights, + LambdaFactory binary_op) { + // Assume iweights_l.type == iweight_r.type + CHECK_EQ(iweights_l.type_, oweights->type_); + CHECK_EQ(iweights_r.type_, oweights->type_); + VLOG(2) << "SANITY CHECK!"; + + switch (iweights_l.type_) { + case tensorflow::DataType::DT_FLOAT: { + auto inp_l = static_cast(iweights_l.GetValues()); + auto inp_r = static_cast(iweights_r.GetValues()); + auto oup = static_cast(const_cast(oweights->GetValues())); + + if (iweights_l.count() != iweights_r.count()) { + // We only supports broadcast of RankZero + if (iweights_l.count() == 1) { + VLOG(2) << "I bet it is not working!" << (*inp_l); + std::transform(inp_r, inp_r + iweights_r.count(), oup, + binary_op.broadcast_l(*inp_l)); + } else if (iweights_r.count() == 1) { + VLOG(2) << "I bet it is not working!" << (*inp_r); + std::transform(inp_l, inp_l + iweights_l.count(), oup, + binary_op.broadcast_r(*inp_r)); + } else { + return tensorflow::errors::Unimplemented( + "Binary op with non-rankZero broadcast not supported"); + } + } else { + std::transform(inp_l, inp_l + iweights_l.count(), inp_r, oup, + binary_op.binary()); + } + break; + } + default: + return tensorflow::errors::Unimplemented( + "Data type not supported: " + + tensorflow::DataTypeString(iweights_l.type_)); + } + + return tensorflow::Status::OK(); +} + +tensorflow::Status ConstantFoldUnary( + Converter& ctx, const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + TRT_ShapedWeights weights_input = inputs.at(0).weights(); + + // Allocate output weights + TRT_ShapedWeights weights_output = ctx.get_temp_weights_like(weights_input); + + // FIXME assume type matches input weights + // Get trt type & shape + // Maybe this part has to be moved into the block of rsqrt later + // Check type consistency + CHECK_EQ(weights_input.type_, + TFAttrs(node_def).get("T")); + + // Maybe I should do a switch + LambdaFactory unary_op; + if (node_def.op() == "Rsqrt") { + // Compute rsqrt + unary_op.op = LambdaFactory::OP_CATEGORY::RSQRT; + auto ret = UnaryCompute(weights_input, &weights_output, unary_op); + // PAss the output + if (ret == tensorflow::Status::OK()) { + outputs->push_back(TRT_TensorOrWeights(weights_output)); + } + return ret; + } else { + return tensorflow::errors::Unimplemented("Binary op not supported: " + + node_def.op()); + } +} + +// TODO(jie,ben) broadcast is needed yet not implemented +// Let's get the simple stuff working first. Maybe we should fall bakc to TF +// approach for constant folding +tensorflow::Status ConstantFoldBinary( + Converter& ctx, const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + TRT_ShapedWeights weights_input_l = inputs.at(0).weights(); + TRT_ShapedWeights weights_input_r = inputs.at(1).weights(); + + // Check type consistency + CHECK_EQ(weights_input_l.type_, weights_input_r.type_); + + if (weights_input_l.shape_.nbDims != weights_input_r.shape_.nbDims) + return tensorflow::errors::Unimplemented( + "Binary op implicit broadcast not supported: " + node_def.op()); + + // TODO(jie): constant fold should really fall back to TF. + int nb_dims = weights_input_l.shape_.nbDims; + nvinfer1::Dims output_shape; + output_shape.nbDims = nb_dims; + VLOG(2) << "nb_dims: " << nb_dims + << ", the other: " << weights_input_r.shape_.nbDims; + for (int i = 0; i < nb_dims; i++) { + if (weights_input_l.shape_.d[i] == weights_input_r.shape_.d[i]) { + output_shape.d[i] = weights_input_l.shape_.d[i]; + } else if (weights_input_l.shape_.d[i] == 1 || + weights_input_r.shape_.d[i] == 1) { + output_shape.d[i] = + std::max(weights_input_l.shape_.d[i], weights_input_r.shape_.d[i]); + } else { + return tensorflow::errors::Unimplemented( + "Binary op with incompatible shape at, " + node_def.op()); + } + VLOG(2) << "left: " << weights_input_l.shape_.d[i] + << "right: " << weights_input_r.shape_.d[i] + << "output: " << output_shape.d[i]; + } + + // FIXME assume type matches input weights + // Get trt type & shape + TFAttrs attrs(node_def); + // Maybe this part has to be moved into the block of rsqrt later + tensorflow::DataType dtype = attrs.get("T"); + + // Allocate output weights + TRT_ShapedWeights weights_output = ctx.get_temp_weights(dtype, output_shape); + + // Maybe I should do a switch + LambdaFactory binary_op; + if (node_def.op() == "Sub") { + binary_op.op = LambdaFactory::OP_CATEGORY::SUB; + } else if (node_def.op() == "Mul") { + binary_op.op = LambdaFactory::OP_CATEGORY::MUL; + } else if (node_def.op() == "Add") { + binary_op.op = LambdaFactory::OP_CATEGORY::ADD; + } else { + return tensorflow::errors::Unimplemented("Binary op not supported: " + + node_def.op()); + } + auto ret = BinaryCompute(weights_input_l, weights_input_r, &weights_output, + binary_op); + + // Pass the output + if (ret == tensorflow::Status::OK()) { + outputs->push_back(TRT_TensorOrWeights(weights_output)); + } + + return ret; +} + +// TODO(jie): broadcast is needed yet not implemented. +// Only implemented channel wise for the time being +tensorflow::Status BinaryTensorOpWeight( + Converter& ctx, const tensorflow::NodeDef& node_def, + const nvinfer1::ITensor* tensor, TRT_ShapedWeights weights, + std::vector* outputs) { + // FIXME assume type matches input weights + // Get trt type & shape + // Maybe this part has to be moved into the block of rsqrt later + + // Check type consistency + auto dtype = TFAttrs(node_def).get("T"); + CHECK_EQ_TYPE(tensor->getType(), dtype); // Cast to int for error messages + nvinfer1::DataType ttype; + TF_CHECK_OK(ConvertDType(weights.type_, &ttype)); + CHECK_EQ_TYPE(ttype, dtype); // Cast to int for error message + + // Check scale mode + auto dims_w = weights.shape_; + auto dims_t = tensor->getDimensions(); + + // Default to channel-wise + auto scale_mode = nvinfer1::ScaleMode::kELEMENTWISE; + + if (weights.count() == 1) { + VLOG(2) << "UNIFORM"; + scale_mode = nvinfer1::ScaleMode::kUNIFORM; + } else { + // No broadcasting on Batch dimension; + assert(dims_w.d[0] == 1); + + // Broadcasting on Channel dimension only allowed in kUNIFORM + assert(dims_w.d[1] == dims_t.d[0]); + assert(dims_w.nbDims == dims_t.nbDims); + + // Default is element; + for (int i = 2; i < dims_w.nbDims; i++) { + if (dims_w.d[i] != dims_t.d[i - 1]) { + scale_mode = nvinfer1::ScaleMode::kCHANNEL; + break; + } + } + if (scale_mode == nvinfer1::ScaleMode::kELEMENTWISE) { + scale_mode = nvinfer1::ScaleMode::kELEMENTWISE; + for (int i = 2; i < dims_w.nbDims; i++) { + if (dims_w.d[i] != 1) + return tensorflow::errors::InvalidArgument( + "Weight shape not compatible at, " + node_def.name()); + } + } + } + + // Prepare weights + TRT_ShapedWeights shift_weights(weights.type_); + TRT_ShapedWeights scale_weights(weights.type_); + TRT_ShapedWeights power_weights(weights.type_); + + // Maybe I should do a switch + if (node_def.op() == "Sub") { + TRT_ShapedWeights neg_weights = ctx.get_temp_weights_like(weights); + LambdaFactory unary_op; + unary_op.op = LambdaFactory::OP_CATEGORY::NEG; + TF_RETURN_IF_ERROR(UnaryCompute(weights, &neg_weights, unary_op)); + shift_weights = neg_weights; + } else if (node_def.op() == "Mul") { + scale_weights = weights; + } else if (node_def.op() == "Add") { + shift_weights = weights; + } else { + return tensorflow::errors::Unimplemented("Binary op not supported: " + + node_def.op()); + } + + nvinfer1::IScaleLayer* layer = ctx.network()->addScale( + *const_cast(tensor), scale_mode, shift_weights, + scale_weights, power_weights); + + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + // Pass the output + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status BinaryTensorOpTensor( + Converter& ctx, const tensorflow::NodeDef& node_def, + const nvinfer1::ITensor* tensor_l, const nvinfer1::ITensor* tensor_r, + std::vector* outputs) { + static const std::unordered_map ops{ + {"Add", nvinfer1::ElementWiseOperation::kSUM}, + {"Mul", nvinfer1::ElementWiseOperation::kPROD}, + // {"max", nvinfer1::ElementWiseOperation::kMAX}, + // {"min", nvinfer1::ElementWiseOperation::kMIN}, + {"Sub", nvinfer1::ElementWiseOperation::kSUB}, + {"Div", nvinfer1::ElementWiseOperation::kDIV}, + }; + + // FIXME assume type matches input weights + // Get trt type & shape + TFAttrs attrs(node_def); + // Maybe this part has to be moved into the block of rsqrt later + nvinfer1::DataType dtype = attrs.get("T"); + + // Check type consistency + CHECK_EQ_TYPE(tensor_l->getType(), dtype); + 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()); + + nvinfer1::IElementWiseLayer* layer = ctx.network()->addElementWise( + *const_cast(tensor_l), + *const_cast(tensor_r), op_pair->second); + + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + // Pass the output + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertPlaceholder( + Converter& ctx, const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + VLOG(2) << "Placeholder should have been replace already"; + return tensorflow::errors::Unimplemented(", cannot convert Placeholder op"); + // OK this make sense since we are supposed to replace it with input + TFAttrs attrs(node_def); + nvinfer1::DataType dtype = attrs.get("dtype"); + nvinfer1::Dims dims = attrs.get("shape"); + + dims.nbDims--; + for (int i = 0; i < dims.nbDims; i++) dims.d[i] = dims.d[i + 1]; + + nvinfer1::ITensor* output = + ctx.network()->addInput(node_def.name().c_str(), dtype, dims); + if (!output) { + return tensorflow::errors::InvalidArgument("Failed to create Input layer"); + } + outputs->push_back(TRT_TensorOrWeights(output)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertConv2D(Converter& ctx, + const tensorflow::NodeDef& node_def, + const std::vector& inputs, + std::vector* outputs) { + nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + // TODO(jie): handle NHWC/NCHW transpose; + TRT_ShapedWeights weights_rsck = inputs.at(1).weights(); + TRT_ShapedWeights weights = ctx.get_temp_weights_like(weights_rsck); + ReorderRSCKToKCRS(weights_rsck, &weights); + TRT_ShapedWeights biases(weights.type_); + int noutput = weights.shape_.d[0]; + nvinfer1::DimsHW kernel_size; + kernel_size.h() = weights.shape_.d[2]; + kernel_size.w() = weights.shape_.d[3]; + TFAttrs attrs(node_def); + + int h_index = 2; + int w_index = 3; + auto data_format = attrs.get("data_format"); + if (data_format == "NHWC") { + tensor = ctx.TransposeTensor(const_cast(tensor), + {0, 3, 1, 2}); + h_index = 1; + w_index = 2; + // TODO(jie): transpose it + } + + // TODO(jie): stride. (NHWC/NCHW) + auto tf_stride = attrs.get>("strides"); + nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]); + + auto tensor_dim = tensor->getDimensions(); + std::vector> padding; + // TODO(jie): padding. + if (attrs.get("padding") == "SAME") { + // This is NCHW tensor with no batch dimension. + // 1 -> h + // 2 -> w + padding = CreateSamePadding( + stride, kernel_size, + {static_cast(tensor_dim.d[1]), static_cast(tensor_dim.d[2])}); + } else { + padding = {{0, 0}, {0, 0}}; + } + + if (padding[0].first != padding[0].second || + padding[1].first != padding[1].second) { + // TODO(jie): handle asymmetric padding + VLOG(2) << "Padding!!!: " << padding[0].first << padding[0].second + << padding[1].first << padding[1].second; + + auto dim_before = tensor->getDimensions(); + VLOG(2) << "TENSOR before: " << dim_before.d[0] << ", " << dim_before.d[1] + << dim_before.d[2] << ", " << dim_before.d[3]; + auto pad_layer = ctx.network()->addPadding( + *const_cast(tensor), + nvinfer1::DimsHW(padding[0].first, padding[1].first), + nvinfer1::DimsHW(padding[0].second, padding[1].second)); + padding = {{0, 0}, {0, 0}}; + tensor = pad_layer->getOutput(0); + auto dim_after = tensor->getDimensions(); + VLOG(2) << "TENSOR after: " << dim_after.d[0] << ", " << dim_after.d[1] + << dim_after.d[2] << ", " << dim_after.d[3]; + } + + nvinfer1::IConvolutionLayer* layer = + ctx.network()->addConvolution(*const_cast(tensor), + noutput, kernel_size, weights, biases); + + layer->setStride(stride); + layer->setPadding({padding[0].first, padding[1].first}); + layer->setName(node_def.name().c_str()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + auto dim_after = output_tensor->getDimensions(); + VLOG(2) << "TENSOR out: " << dim_after.d[0] << ", " << dim_after.d[1] + << dim_after.d[2] << ", " << dim_after.d[3]; + + if (data_format == "NHWC") { + // TODO(jie): transpose it back! + output_tensor = ctx.TransposeTensor(output_tensor, {0, 2, 3, 1}); + } else { + VLOG(2) << "NCHW !!!!"; + } + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertPool(Converter& ctx, + const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + TFAttrs attrs(node_def); + + int h_index = 2; + int w_index = 3; + auto data_format = attrs.get("data_format"); + if (data_format == "NHWC") { + h_index = 1; + w_index = 2; + tensor = ctx.TransposeTensor(const_cast(tensor), + {0, 3, 1, 2}); + } else { + VLOG(2) << "NCHW !!!!"; + } + nvinfer1::PoolingType type; + // TODO(jie): support other pooling type + if (node_def.op() == "MaxPool") + type = nvinfer1::PoolingType::kMAX; + else + return tensorflow::errors::Unimplemented("Only supports Max pool"); + + // TODO(jie): NCHW + auto tf_stride = attrs.get>("strides"); + nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]); + + auto tf_kernel = attrs.get>("ksize"); + nvinfer1::DimsHW ksize(tf_kernel[h_index], tf_kernel[w_index]); + + auto tensor_dim = tensor->getDimensions(); + std::vector> padding; + // TODO(jie): padding. + if (attrs.get("padding") == "SAME") { + // This is NCHW tensor with no batch dimension. + // 1 -> h + // 2 -> w + padding = CreateSamePadding( + stride, ksize, + {static_cast(tensor_dim.d[1]), static_cast(tensor_dim.d[2])}); + } else if (attrs.get("padding") == "VALID") { + // No padding for valid padding here + VLOG(2) << "No padding added for VALID padding in pool" << node_def.name(); + padding = {{0, 0}, {0, 0}}; + } else { + return tensorflow::errors::Unimplemented( + "Current MaxPool cannot support padding other than SAME"); + } + + if (padding[0].first != padding[0].second || + padding[1].first != padding[1].second) { + // TODO(jie): handle asymmetric padding + VLOG(2) << "Padding!!!: " << padding[0].first << padding[0].second + << padding[1].first << padding[1].second; + auto pad_layer = ctx.network()->addPadding( + *const_cast(tensor), + nvinfer1::DimsHW(padding[0].first, padding[1].first), + nvinfer1::DimsHW(padding[0].second, padding[1].second)); + padding = {{0, 0}, {0, 0}}; + tensor = pad_layer->getOutput(0); + } + + nvinfer1::IPoolingLayer* layer = ctx.network()->addPooling( + *const_cast(tensor), type, ksize); + + layer->setStride(stride); + layer->setPadding({padding[0].first, padding[1].first}); + layer->setName(node_def.name().c_str()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + if (data_format == "NHWC") { + // TODO(jie): transpose it back! + output_tensor = ctx.TransposeTensor(output_tensor, {0, 2, 3, 1}); + } else { + VLOG(2) << "NCHW !!!!"; + } + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertActivation( + Converter& ctx, const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + nvinfer1::IActivationLayer* layer = ctx.network()->addActivation( + *const_cast(tensor), nvinfer1::ActivationType::kRELU); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertScale(Converter& ctx, + const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + if (inputs.size() != 2 || !inputs.at(0).is_tensor() || + !inputs.at(1).is_weights()) + return tensorflow::errors::Unimplemented( + "Only supports tensor op weight for now, at " + node_def.name()); + // Implement tensor binaryOp weight [channel wise] for now; + nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + + // TODO(jie): handle NHWC/NCHW transpose; + TRT_ShapedWeights weights = inputs.at(1).weights(); + TRT_ShapedWeights empty_weights(weights.type_); + + TFAttrs attrs(node_def); + + // Transpose NHWC + auto data_format = attrs.get("data_format"); + if (data_format == "NHWC") { + tensor = ctx.TransposeTensor(const_cast(tensor), + {0, 3, 1, 2}); + // TODO(jie): transpose it + } else { + VLOG(2) << "NCHW !!!!"; + } + nvinfer1::IScaleLayer* layer = ctx.network()->addScale( + *const_cast(tensor), nvinfer1::ScaleMode::kCHANNEL, + weights, empty_weights, empty_weights); + + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + if (data_format == "NHWC") { + // TODO(jie): transpose it back! + output_tensor = ctx.TransposeTensor(output_tensor, {0, 2, 3, 1}); + } else { + VLOG(2) << "NCHW !!!!"; + } + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertConst(Converter& ctx, + const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + const auto& weights_tensor = node_def.attr().at("value").tensor(); + + // Get trt type & shape + TFAttrs attrs(node_def); + const tensorflow::DataType dtype = attrs.get("dtype"); + + // Create shaped weights as output + tensorflow::Tensor tensor; + if (!tensor.FromProto(weights_tensor)) + return tensorflow::errors::Internal("Cannot parse weight tensor proto: " + + node_def.name()); + + TRT_ShapedWeights weights(dtype); + if (!weights_tensor.float_val().empty()) { + VLOG(2) << "SCALAR!!!" << node_def.name(); + nvinfer1::Dims scalar_shape; + if (tensor.dims() > 0) { + VLOG(2) << "Dimensions: " << tensor.dims(); + weights = TRT_ShapedWeights(dtype, weights_tensor.float_val().data(), + GetTensorShape(tensor)); + } else { + VLOG(2) << "Dimensions: " << tensor.dims(); + scalar_shape.nbDims = 1; + scalar_shape.d[0] = 1; + scalar_shape.type[0] = nvinfer1::DimensionType::kSPATIAL; + for (int i = 1; i < nvinfer1::Dims::MAX_DIMS; i++) { + scalar_shape.d[i] = 0; + scalar_shape.type[i] = nvinfer1::DimensionType::kSPATIAL; + } + weights = TRT_ShapedWeights(dtype, weights_tensor.float_val().data(), + scalar_shape); + } + } else if (!weights_tensor.tensor_content().empty()) { + VLOG(2) << "TENSOR!!!" << node_def.name(); + const auto& content = weights_tensor.tensor_content(); + + weights = ctx.get_temp_weights(dtype, GetTensorShape(tensor)); + if (content.size() > 0) { + const int dtype_size = tensorflow::DataTypeSize(dtype); + CHECK_EQ(0, content.size() % dtype_size) + << "Tensor content size (" << content.size() + << ") is not a multiple of " << dtype_size; + port::CopyToArray( + content, static_cast(const_cast(weights.GetValues()))); + } + } else { + return tensorflow::errors::Unimplemented( + "Not supported constant type, at " + node_def.name()); + } + // Pass the output + outputs->push_back(TRT_TensorOrWeights(weights)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertIdentity( + Converter& ctx, const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + outputs->push_back(inputs.at(0)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertBinary(Converter& ctx, + const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + if (inputs.size() != 2) + return tensorflow::errors::FailedPrecondition( + "Binary ops require two tensor input, at " + node_def.name()); + + if (inputs.at(0).is_weights() && inputs.at(1).is_weights()) + return ConstantFoldBinary(ctx, node_def, inputs, outputs); + + if (inputs.at(0).is_tensor() && inputs.at(1).is_weights()) + return BinaryTensorOpWeight(ctx, node_def, inputs.at(0).tensor(), + inputs.at(1).weights(), outputs); + + if (inputs.at(0).is_weights() && inputs.at(1).is_tensor()) + return BinaryTensorOpWeight(ctx, node_def, inputs.at(1).tensor(), + inputs.at(0).weights(), outputs); + + if (inputs.at(0).is_tensor() && inputs.at(1).is_tensor()) + return BinaryTensorOpTensor(ctx, node_def, inputs.at(0).tensor(), + inputs.at(1).tensor(), outputs); + + return tensorflow::errors::Unknown("Binary op input error, at " + + node_def.name()); +} + +tensorflow::Status ConvertUnary(Converter& ctx, + const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + if (inputs.size() != 1) + return tensorflow::errors::FailedPrecondition( + "Unary ops require single tensor input, at " + node_def.name()); + + if (inputs.at(0).is_weights()) + return ConstantFoldUnary(ctx, node_def, inputs, outputs); + else if (inputs.at(0).is_tensor()) + return tensorflow::errors::Unimplemented( + "Unary op for tensor not supported, at " + node_def.name()); + + return tensorflow::errors::Unknown("Binary op input error, at " + + node_def.name()); +} + +tensorflow::Status ConvertReduce(Converter& ctx, + const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + if (inputs.size() != 2 || !inputs.at(0).is_tensor() || + !inputs.at(1).is_weights()) + return tensorflow::errors::InvalidArgument( + "Input expects tensor and weights, at" + node_def.name()); + + // Implement tensor binaryOp weight [channel wise] for now; + nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + auto dims = tensor->getDimensions(); + // Restore implicit batch dimension + int nb_dims = dims.nbDims + 1; + + TRT_ShapedWeights index_list = inputs.at(1).weights(); + + TFAttrs attrs(node_def); + // TODO(jie): handle data type. + // Index type here is done through TF type, so I can leverage their + // EnumToDataType for my cast + auto index_type = attrs.get("Tidx"); + + // Only expect to handle INT32 as attributes for now + if (index_type != tensorflow::DataType::DT_INT32) + return tensorflow::errors::Unimplemented("Tidx supports only DT_INT32"); + auto index_list_data = + static_cast(const_cast(index_list.GetValues())); + + // Hack warning: have to fall back to pool layer since reduce is not in public + // TRT yet. + if (nb_dims != 4) + return tensorflow::errors::InvalidArgument( + "TRT only support reduce on 4 dimensional tensors, at" + + node_def.name()); + if (index_list.count() > 2) + return tensorflow::errors::InvalidArgument( + "TRT cannot support reduce on more than 2 dimensions, at" + + node_def.name()); + + std::set idx_set; + // We cannot operate on Channel. permutation flag used to transpose tensor + int permuted_index = -1; + for (int i = 0; i < index_list.count(); i++) { + if (index_list_data[i] == 0) + return tensorflow::errors::InvalidArgument("TRT cannot reduce at 0, at" + + node_def.name()); + if (index_list_data[i] == 1) permuted_index = 1; + idx_set.emplace(index_list_data[i]); + } + + std::vector permutation_order(nb_dims); + nvinfer1::DimsHW pool_kernel; + if (permuted_index == 1) { + for (int i = 2; i < nb_dims; i++) { + if (idx_set.count(i)) { + permuted_index = i; + break; + } + } + for (int i = 0; i < nb_dims; i++) permutation_order[i] = i; + + permutation_order[permuted_index] = 1; + permutation_order[1] = permuted_index; + + // Apply permutation before extracting dimension for pool_kernel + tensor = ctx.TransposeTensor(const_cast(tensor), + permutation_order); + } + + // Apply permutation before extracting dimension for pool_kernel + pool_kernel.d[0] = (idx_set.count(2) || permuted_index == 2) ? dims.d[1] : 1; + pool_kernel.d[1] = (idx_set.count(3) || permuted_index == 3) ? dims.d[2] : 1; + + nvinfer1::ITensor* output_tensor; + + if (node_def.op() == "Mean") { + nvinfer1::IPoolingLayer* layer = + ctx.network()->addPooling(*const_cast(tensor), + nvinfer1::PoolingType::kAVERAGE, pool_kernel); + output_tensor = layer->getOutput(0); + } else { + return tensorflow::errors::Unimplemented( + "Op not supported " + node_def.op() + " , at " + node_def.name()); + } + if (permuted_index != -1) { + // Apply permutation before extracting dimension for pool_kernel + output_tensor = ctx.TransposeTensor( + const_cast(output_tensor), permutation_order); + } + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertPad(Converter& ctx, + const tensorflow::NodeDef& node_def, + std::vector const& inputs, + std::vector* outputs) { + if (inputs.size() != 2 || !inputs.at(0).is_tensor() || + !inputs.at(1).is_weights()) + return tensorflow::errors::InvalidArgument( + "Input expects tensor and weights, at" + node_def.name()); + + // Implement tensor binaryOp weight [channel wise] for now; + nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + auto dims = tensor->getDimensions(); + // Restore implicit batch dimension + int nb_dims = dims.nbDims + 1; + + TRT_ShapedWeights pads = inputs.at(1).weights(); + + TFAttrs attrs(node_def); + // Padding type here is done through TF type + // so I can leverage their EnumToDataType for my cast + auto padding_type = attrs.get("Tpaddings"); + // TODO(jie): handle data type conversion for TRT? + + if (pads.shape_.d[0] != nb_dims || pads.shape_.d[1] != 2) + return tensorflow::errors::InvalidArgument( + "Pad only supports explicit padding on 4 dimensional tensor, at " + + node_def.name()); + + // Only expect to handle INT32 as attributes for now + if (padding_type != tensorflow::DataType::DT_INT32) + return tensorflow::errors::Unimplemented( + "Tpaddings supports only DT_INT32"); + auto pad_data = static_cast(const_cast(pads.GetValues())); + + std::vector pad_index; + for (int i = 0; i < nb_dims; i++) { + if (pad_data[2 * i] != 0 || pad_data[2 * i + 1] != 0) + pad_index.push_back(i); + } + + // No padding at all, we should exit + if (pad_index.size() == 0) { + outputs->push_back(inputs.at(0)); + return tensorflow::Status::OK(); + } + + // Only supports padding on less than 2 axis GIE-2579 + if (pad_index.size() > 2) + return tensorflow::errors::InvalidArgument( + "Padding layer does not support padding on > 2"); + + // Padding on batch dimension is not supported + if (pad_index[0] == 0) + return tensorflow::errors::InvalidArgument( + "Padding layer does not support padding on batch dimension"); + + // Not doing the legit thing here. ignoring padding on dim 1 and 3; + // TODO(jie): implement pad as uff parser + if (pad_index.size() == 2 && pad_index[0] == 0 && pad_index[1] == 3) + return tensorflow::errors::Unimplemented( + "Padding layer does not support padding on dimension 1 and 3 yet"); + + bool legit_pad = true; + nvinfer1::DimsHW pre_padding(0, 0); + nvinfer1::DimsHW post_padding(0, 0); + + std::vector permuted_pad_index(pad_index); + if (pad_index[0] == 1) { + legit_pad = false; + tensor = ctx.TransposeTensor(const_cast(tensor), + {0, 3, 2, 1}); + permuted_pad_index[0] = 3; + } + + for (size_t i = 0; i < pad_index.size(); i++) { + int index = pad_index[i]; + if (permuted_pad_index[i] == 2) { + pre_padding.h() = pad_data[index * 2]; + post_padding.h() = pad_data[index * 2 + 1]; + } else if (permuted_pad_index[i] == 3) { + pre_padding.w() = pad_data[index * 2]; + post_padding.w() = pad_data[index * 2 + 1]; + } + } + + nvinfer1::IPaddingLayer* layer = ctx.network()->addPadding( + *const_cast(tensor), pre_padding, post_padding); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + if (!legit_pad) + output_tensor = ctx.TransposeTensor( + const_cast(output_tensor), {0, 3, 2, 1}); + + outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +void Converter::register_op_converters() { + // vgg_16 slim implementation + op_registry_["Placeholder"] = ConvertPlaceholder; + op_registry_["Conv2D"] = ConvertConv2D; + op_registry_["Relu"] = ConvertActivation; + op_registry_["MaxPool"] = ConvertPool; + // This could be really handled as ConvertBinary + op_registry_["BiasAdd"] = ConvertScale; + op_registry_["Const"] = ConvertConst; + // op_registry_["MatMul"] = ConvertFullyConnected; // Not used in vgg + // TODO(ben,jie): this is a temp hack. + op_registry_["Identity"] = ConvertIdentity; // Identity should be removed + // op_registry_["AvgPool"] = ConvertPool; + + // resnet_50_v1 slim implementation + op_registry_["Add"] = ConvertBinary; + op_registry_["Mul"] = ConvertBinary; + op_registry_["Sub"] = ConvertBinary; + op_registry_["Rsqrt"] = ConvertUnary; + op_registry_["Mean"] = ConvertReduce; + op_registry_["Pad"] = ConvertPad; + // TODO(ben,jie): Add more ops +} + +} // namespace + +tensorflow::Status ConvertSubGraphToTensorRTNodeDef( + const tensorflow::Graph& graph, const std::set& subgraph_node_ids, + const std::vector>& input_inds, + const std::vector>& output_inds, size_t max_batch_size, + size_t max_workspace_size_bytes, + const tensorflow::grappler::GraphProperties& graph_properties, + tensorflow::NodeDef* trt_node) { + // Visit nodes in reverse topological order and construct the TRT network. + + // Toposort + std::vector order_vec; + tensorflow::GetPostOrder(graph, &order_vec); + // Select just the subgraph + std::list order; + for (tensorflow::Node* node : order_vec) { + if (subgraph_node_ids.count(node->id())) { + // We want topological order to contstruct the + // network layer by layer + order.push_front(node); + } + } + // Topological order is needed to build TRT network + + tensorflow::tensorrt::Logger trt_logger; + + auto trt_builder = infer_object(nvinfer1::createInferBuilder(trt_logger)); + if (!trt_builder) { + return tensorflow::errors::Internal( + "Failed to create TensorRT builder object"); + } + + auto trt_network = infer_object(trt_builder->createNetwork()); + if (!trt_network) { + return tensorflow::errors::Internal( + "Failed to create TensorRT network object"); + } + + // Build the network + Converter converter(trt_network.get()); + + std::vector input_names; + std::vector input_dtypes; + for (std::pair const& input : input_inds) { + int node_id = input.first; + int output_idx = input.second; + tensorflow::Node* node = graph.FindNodeId(node_id); + auto node_name = node->name(); + input_names.push_back(node_name); // Insert original node name without port + // TODO(jie): alternative :) + if (!graph_properties.HasOutputProperties(node_name)) + return tensorflow::errors::Internal("Failed to find input node: " + + node_name); + + auto op_info_vec = graph_properties.GetOutputProperties(node_name); + if (static_cast(op_info_vec.size()) < output_idx) + return tensorflow::errors::Internal( + "Accessing output index of: " + std::to_string(output_idx) + + ", at node: " + node_name + " with output entry from shape_map: " + + std::to_string(op_info_vec.size())); + + auto op_info = op_info_vec.at(output_idx); + + tensorflow::DataType tf_dtype = op_info.dtype(); + input_dtypes.push_back(tf_dtype); + + nvinfer1::DataType dtype(nvinfer1::DataType::kFLOAT); + TF_CHECK_OK(ConvertDType(tf_dtype, &dtype)); + + VLOG(2) << "Accessing output index of: " << std::to_string(output_idx) + << ", at node: " << node_name + << " with output entry from shape_map: " + << std::to_string(op_info_vec.size()); + + // TODO(ben,jie): update TRT input format/dimension + 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_psuedo_chw.d[i - 1] = op_info.shape().dim(i).size(); + } + + // TODO(ben,jie): proper way to restore input tensor name? + auto input_tensor_name = node_name; + if (output_idx != 0) + 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_psuedo_chw); + + if (!input_tensor) + return tensorflow::errors::InvalidArgument( + "Failed to create Input layer"); + VLOG(2) << "Input tensor name :" << input_tensor_name; + + if (!converter.insert_input_tensor(input_tensor_name, input_tensor)) + return tensorflow::errors::AlreadyExists( + "Output tensor already exists for op: " + input_tensor_name); + } + + VLOG(2) << "Finished sorting"; + + for (const tensorflow::Node* node : order) { + const tensorflow::NodeDef& node_def = node->def(); + VLOG(2) << "Converting node: " << node_def.name() << " , " << node_def.op(); + TF_RETURN_IF_ERROR(converter.convert_node(node_def)); + } + + VLOG(2) << "Finished conversion"; + + // Gather output metadata + std::vector output_names; + std::vector output_dtypes; + for (std::pair const& output : output_inds) { + int node_id = output.first; + int output_idx = output.second; + tensorflow::Node* node = graph.FindNodeId(node_id); + string op_name = node->name(); + string tensor_name = op_name; + if (output_idx != 0) + tensor_name = tensor_name + ":" + std::to_string(output_idx); + VLOG(2) << "Output tensor name: " << tensor_name; + output_names.push_back(tensor_name); + auto tensor_or_weights = converter.get_tensor(tensor_name); + if (!tensor_or_weights.is_tensor()) { + return tensorflow::errors::InvalidArgument( + "Output node is weights not tensor"); + } + nvinfer1::ITensor* tensor = tensor_or_weights.tensor(); + if (!tensor) { + return tensorflow::errors::NotFound("Output tensor not found: " + + tensor_name); + } + converter.network()->markOutput(*tensor); + tensorflow::DataType tf_dtype = node->output_type(output_idx); + output_dtypes.push_back(tf_dtype); + nvinfer1::DataType trt_dtype = nvinfer1::DataType::kFLOAT; + TF_RETURN_IF_ERROR(ConvertDType(tf_dtype, &trt_dtype)); + tensor->setType(trt_dtype); + } + + VLOG(2) << "Finished output"; + // TODO(jie): static_id is not thread safe. + static int static_id = 0; + + // Build the engine + trt_builder->setMaxBatchSize(max_batch_size); + trt_builder->setMaxWorkspaceSize(max_workspace_size_bytes); + VLOG(0) << "Starting build engine " << static_id; + // TODO(ben,jie): half2 and int8 mode support + string engine_plan_string; + { + auto trt_engine = + infer_object(trt_builder->buildCudaEngine(*converter.network())); + VLOG(0) << "Built network"; + auto engine_plan = infer_object(trt_engine->serialize()); + VLOG(0) << "Serialized engine"; + const char* engine_plan_data = + static_cast(engine_plan->data()); + engine_plan_string = + string(engine_plan_data, engine_plan_data + engine_plan->size()); + } + + VLOG(0) << "Finished engine"; + + // Build the TRT op + // TODO(sami,ben,jie): proper naming! + tensorflow::NodeDefBuilder op_builder( + tensorflow::strings::StrCat("my_trt_op", static_id++), "TRTEngineOp"); + std::vector income_edges; + for (size_t i = 0; i < input_names.size(); ++i) { + int output_idx = input_inds.at(i).second; + // We wired up the input here already, it is redundant to do it again in + // ConvertSubGraphToTensorRT(convert_graph.cc) + auto incoming_edge = tensorflow::NodeDefBuilder::NodeOut( + input_names.at(i), output_idx, input_dtypes.at(i)); + income_edges.push_back(incoming_edge); + } + tensorflow::gtl::ArraySlice input_list( + income_edges); + op_builder.Input(input_list); + + VLOG(0) << "Finished op preparation"; + + auto status = op_builder.Attr("serialized_engine", engine_plan_string) + .Attr("input_nodes", input_names) + .Attr("output_nodes", output_names) + .Attr("OutT", output_dtypes) + .Finalize(trt_node); + + VLOG(0) << status.ToString() << " finished op building"; + + return tensorflow::Status::OK(); +} + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.h b/tensorflow/contrib/tensorrt/convert/convert_nodes.h new file mode 100644 index 0000000000000000000000000000000000000000..2e7fd19566e1ed3719b932c7443a9c3f652b2d3e --- /dev/null +++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.h @@ -0,0 +1,52 @@ +/* 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_CONVERT_CONVERT_NODES_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_NODES_H_ + +#include +#include +#include + +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/grappler/costs/graph_properties.h" +#include "tensorflow/core/lib/core/status.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +namespace tensorflow { +namespace tensorrt { +namespace convert { + +tensorflow::Status ConvertSubGraphToTensorRTNodeDef( + const tensorflow::Graph& graph, const std::set& subgraph_node_ids, + const std::vector>& + input_inds, // {node_id, output_idx} + const std::vector>& + output_inds, // {node_id, output_idx} + size_t max_batch_size, size_t max_workspace_size_bytes, + const tensorflow::grappler::GraphProperties& graph_prop, + tensorflow::NodeDef* trt_node); + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_NODES_H_ diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..8efdf63ebebc4d7a199c60635ca64348d2b30505 --- /dev/null +++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc @@ -0,0 +1,140 @@ +/* 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_engine_op.h" + +#include "tensorflow/contrib/tensorrt/log/trt_logger.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor.h" +#include "tensorflow/core/platform/types.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "cuda/include/cuda_runtime_api.h" + +namespace tensorflow { +namespace tensorrt { +static ::tensorflow::tensorrt::Logger logger; + +TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) : OpKernel(context) { + // read serialized_engine + string serialized_engine; + OP_REQUIRES_OK(context, + context->GetAttr("serialized_engine", &serialized_engine)); + + // register input output node name in trt_sub_graph + OP_REQUIRES_OK(context, context->GetAttr("input_nodes", &input_nodes_)); + OP_REQUIRES_OK(context, context->GetAttr("output_nodes", &output_nodes_)); + + // TODO(samikama) runtime should be taken from a resourcemanager as well. + // Only engine should be in the op and context and runtime should be taken + // from resourcemanager + nvinfer1::IRuntime* infer = nvinfer1::createInferRuntime(logger); + trt_engine_ptr_.reset(infer->deserializeCudaEngine( + serialized_engine.c_str(), serialized_engine.size(), nullptr)); + + trt_execution_context_ptr_.reset(trt_engine_ptr_->createExecutionContext()); + // Runtime is safe to delete after engine creation + infer->destroy(); +} + +void TRTEngineOp::Compute(OpKernelContext* context) { + int num_binding = context->num_inputs() + context->num_outputs(); + std::vector buffers(num_binding); + + size_t binding_index; + int num_batch = 0; + bool valid = true; + for (int i = 0; i < context->num_inputs(); i++) { + // Grab the input tensor + binding_index = trt_engine_ptr_->getBindingIndex(input_nodes_[i].c_str()); + + const Tensor& input_tensor = context->input(i); + const TensorShape& input_shape = input_tensor.shape(); + if (i == 0) { + num_batch = input_shape.dim_size(0); + } else if (num_batch != input_shape.dim_size(0)) { + valid = false; + break; + } + switch (trt_engine_ptr_->getBindingDataType(binding_index)) { + case nvinfer1::DataType::kFLOAT: + buffers[binding_index] = (void*)(input_tensor.flat().data()); + break; + case nvinfer1::DataType::kHALF: + LOG(FATAL) << "half size is not supported yet!"; + break; + case nvinfer1::DataType::kINT8: + LOG(FATAL) << "int8 is not supported yet!"; + break; + } + } + + // Might want a different way to inform the user of batch size inconsistency + if (!valid) LOG(WARNING) << "input data inconsistent batch size"; + + for (int i = 0; i < static_cast(output_nodes_.size()); i++) { + // This is bad that we have to reallocate output buffer every run. + // Create an output tensor + binding_index = trt_engine_ptr_->getBindingIndex(output_nodes_[i].c_str()); + Tensor* output_tensor = nullptr; + + TensorShape output_shape; + if (binding_index != -1) { + auto dims = trt_engine_ptr_->getBindingDimensions(binding_index); + std::vector trt_shape(dims.nbDims + 1); + trt_shape[0] = num_batch; + for (int j = 0; j < dims.nbDims; j++) trt_shape[j + 1] = dims.d[j]; + OP_REQUIRES_OK(context, + TensorShapeUtils::MakeShape( + trt_shape.data(), trt_shape.size(), &output_shape)); + } else { + LOG(FATAL) << "output node not found, at " << output_nodes_[i]; + break; + } + + OP_REQUIRES_OK(context, + context->allocate_output(i, output_shape, &output_tensor)); + switch (trt_engine_ptr_->getBindingDataType(binding_index)) { + case nvinfer1::DataType::kFLOAT: + buffers[binding_index] = + reinterpret_cast(output_tensor->flat().data()); + break; + case nvinfer1::DataType::kHALF: + LOG(FATAL) << "half size is not supported yet!"; + break; + case nvinfer1::DataType::kINT8: + LOG(FATAL) << "int8 is not supported yet!"; + break; + } + } + // copied from cuda_kernel_helper since it seems only valid in *.cu.cc files + const cudaStream_t* stream = CHECK_NOTNULL( + reinterpret_cast(context->op_device_context() + ->stream() + ->implementation() + ->CudaStreamMemberHack())); + + // execution handled by TF since we are getting stream from TF. + // it is safe for CPU pointer array (buffers) to go out of scope after enqueue + trt_execution_context_ptr_->enqueue(num_batch, &buffers[0], *stream, nullptr); +} + +REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp); + +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h new file mode 100644 index 0000000000000000000000000000000000000000..0964b4b18a781143fdd7884a2904321b9d14e354 --- /dev/null +++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h @@ -0,0 +1,62 @@ +/* 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_ENGINE_OP_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_ENGINE_OP_H_ + +#include +#include +#include + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "cuda/include/cuda_runtime_api.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { +class Logger; + +class TRTEngineOp : public OpKernel { + public: + explicit TRTEngineOp(OpKernelConstruction* context); + + void Compute(OpKernelContext* context) override; + + private: + template + struct Destroyer { + void operator()(T* d) { d->destroy(); } + }; + + template + using destroyed_ptr = std::unique_ptr>; + destroyed_ptr trt_engine_ptr_; + // TODO(samikama): context should go to a resource manager! + destroyed_ptr trt_execution_context_ptr_; + + std::vector input_nodes_; + std::vector output_nodes_; +}; + +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +#endif // TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_ENGINE_OP_H_ diff --git a/tensorflow/contrib/tensorrt/log/trt_logger.cc b/tensorflow/contrib/tensorrt/log/trt_logger.cc new file mode 100644 index 0000000000000000000000000000000000000000..7add8cb8b3d2a04206ee4174e79a1a4b86e05f30 --- /dev/null +++ b/tensorflow/contrib/tensorrt/log/trt_logger.cc @@ -0,0 +1,57 @@ +/* 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/log/trt_logger.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorflow/core/platform/logging.h" + +namespace tensorflow { +namespace tensorrt { + +// Use TF logging for TensorRT informations +void Logger::log(Severity severity, const char* msg) { + // Suppress info-level messages + switch (severity) { + case Severity::kINFO: { // Mark TRT info messages as debug! + VLOG(2) << msg; + break; + } + case Severity::kWARNING: { + LOG(WARNING) << msg; + break; + } + case Severity::kERROR: { + LOG(ERROR) << msg; + break; + } + case Severity::kINTERNAL_ERROR: { + LOG(FATAL) << msg; + break; + } + // This is useless for now. But would catch it in future if enum changes. It + // is always good to have default case! + default: { + LOG(FATAL) << name_ << "Got unknown severity level from TRT " << msg; + break; + } + } +} +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_CUDA +#endif // GOOGLE_TENSORRT diff --git a/tensorflow/contrib/tensorrt/log/trt_logger.h b/tensorflow/contrib/tensorrt/log/trt_logger.h new file mode 100644 index 0000000000000000000000000000000000000000..d71f66b933a8068a6276a7e070755e0075543bb5 --- /dev/null +++ b/tensorflow/contrib/tensorrt/log/trt_logger.h @@ -0,0 +1,42 @@ +/* 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_LOG_TRT_LOGGER_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_LOG_TRT_LOGGER_H_ + +#include "tensorflow/core/platform/types.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { + +// Logger for GIE info/warning/errors +class Logger : public nvinfer1::ILogger { + private: + void log(nvinfer1::ILogger::Severity severity, const char* msg) override; + + string name_; +}; + +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +#endif // TENSORFLOW_CONTRIB_TENSORRT_LOG_TRT_LOGGER_H_ diff --git a/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc b/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..079d73f7bec3f9a9740e455b31a259cec287f849 --- /dev/null +++ b/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc @@ -0,0 +1,43 @@ +/* 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. +==============================================================================*/ + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/shape_inference.h" +#include "tensorflow/core/framework/tensor_shape.h" + +namespace tensorflow { + +namespace shape_inference { +extern Status TRTEngineOpShapeInference(InferenceContext* c); +} + +REGISTER_OP("TRTEngineOp") + .Attr("serialized_engine: string") + .Attr("input_nodes: list(string)") + .Attr("output_nodes: list(string)") + .Attr("InT: list({float32})") + .Attr("OutT: list({float32})") + .Input("in_tensor: InT") + .Output("out_tensor: OutT") + .SetShapeFn(shape_inference::TRTEngineOpShapeInference); + +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/python/__init__.py b/tensorflow/contrib/tensorrt/python/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7e050a768ce97af1fc1d2c85cb52640b4c6a6a97 --- /dev/null +++ b/tensorflow/contrib/tensorrt/python/__init__.py @@ -0,0 +1,24 @@ +# 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. +# ============================================================================= +"""Exposes the python wrapper for TensorRT graph transforms.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# pylint: disable=unused-import,line-too-long +from tensorflow.contrib.tensorrt.python.ops import trt_engine_op +from tensorflow.contrib.tensorrt.python.trt_convert import create_inference_graph +# pylint: enable=unused-import,line-too-long diff --git a/tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py b/tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py new file mode 100644 index 0000000000000000000000000000000000000000..31a313182be9a2fca7457a539670dbc911ccabb1 --- /dev/null +++ b/tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py @@ -0,0 +1,34 @@ +# 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. +# ============================================================================= +"""Exposes the Python wrapper of TRTEngineOp.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import platform + +if platform.system() != "Windows": + # pylint: disable=wildcard-import,unused-import,g-import-not-at-top + from tensorflow.contrib.tensorrt.ops.gen_trt_engine_op import * + + from tensorflow.contrib.util import loader + from tensorflow.python.platform import resource_loader + # pylint: enable=wildcard-import,unused-import,g-import-not-at-top + + _trt_engine_op = loader.load_op_library( + resource_loader.get_path_to_datafile("_trt_engine_op.so")) +else: + raise RuntimeError("Windows platforms are not supported") diff --git a/tensorflow/contrib/tensorrt/python/trt_convert.py b/tensorflow/contrib/tensorrt/python/trt_convert.py new file mode 100644 index 0000000000000000000000000000000000000000..9454862f857ab743712ce409ff007de55e72a68e --- /dev/null +++ b/tensorflow/contrib/tensorrt/python/trt_convert.py @@ -0,0 +1,103 @@ +# 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. +# ============================================================================= +"""Exposes the Python wrapper conversion to trt_graph.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# pylint: disable=unused-import,line-too-long +import six as _six +from tensorflow.contrib.tensorrt.wrap_conversion import trt_convert +from tensorflow.core.framework import graph_pb2 +from tensorflow.python.framework import errors +from tensorflow.python.framework import errors_impl as _impl +from tensorflow.python.framework import ops + + +# TODO(skama): get outputs from session when implemented as c++ +# optimization pass +def create_inference_graph(input_graph_def, + outputs, + max_batch_size=1, + max_workspace_size_bytes=2 << 20): + """Python wrapper for the TRT transormation. + + + Args: + input_graph_def: GraphDef object containing a model to be transformed. + outputs: List of tensors or node names for the model outputs. + max_batch_size: max size for the input batch + max_workspace_size_bytes: parameter to control memory allocation (in Bytes) + + Returns: + New GraphDef with TRTEngineOps placed in graph replacing subgraphs. + + Raises: + RuntimeError: if the returned status message is malformed. + """ + + def py2bytes(inp): + return inp + + def py3bytes(inp): + return inp.encode("utf-8", errors="surrogateescape") + + def py2string(inp): + return inp + + def py3string(inp): + return inp.decode("utf-8") + + if _six.PY2: + to_bytes = py2bytes + to_string = py2string + else: + to_bytes = py3bytes + to_string = py3string + + out_names = [] + for i in outputs: + if isinstance(i, ops.Tensor): + out_names.append(to_bytes(i.name)) + else: + out_names.append(to_bytes(i)) + + input_graph_def_str = input_graph_def.SerializeToString() + + # TODO(sami): Fix this when we can return status from C++ library + # There is a problem with the TF internal library setup that doesn't + # allow us to return a status object from C++. Thus we return a + # pair or strings where first one is encoded status and the second + # one is the transformed graphs protobuf string. + out = trt_convert(input_graph_def_str, out_names, max_batch_size, + max_workspace_size_bytes) + status = to_string(out[0]) + output_graph_def_string = out[1] + del input_graph_def_str # Save some memory + if len(status) < 2: + raise _impl.UnknownError(None, None, status) + if status[:2] != "OK": + msg = status.split(";") + if len(msg) == 1: + raise RuntimeError("Status message is malformed {}".format(status)) + # pylint: disable=protected-access + raise _impl._make_specific_exception(None, None, ";".join(msg[1:]), + int(msg[0])) + # pylint: enable=protected-access + output_graph_def = graph_pb2.GraphDef() + output_graph_def.ParseFromString(output_graph_def_string) + del output_graph_def_string # Save some memory + return output_graph_def diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/contrib/tensorrt/segment/segment.cc new file mode 100644 index 0000000000000000000000000000000000000000..6193f0b0a13f6985d5fc8dd4c6fc09b15f72f139 --- /dev/null +++ b/tensorflow/contrib/tensorrt/segment/segment.cc @@ -0,0 +1,253 @@ +/* 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/segment/segment.h" + +#include +#include +#include + +#include "tensorflow/contrib/tensorrt/segment/union_find.h" +#include "tensorflow/core/graph/algorithm.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/graph_constructor.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +namespace tensorrt { +namespace segment { + +namespace { + +bool CanContractEdge(const tensorflow::Edge* edge, + const tensorflow::Graph& graph) { + const tensorflow::Node* src = edge->src(); + const tensorflow::Node* dst = edge->dst(); + + // Can't contract edge if doing so would cause a cycle in the + // graph. So, if there is a directed path from 'src' to 'dst', other + // than 'edge' (or any other direct edge from 'src' to 'dst'), then + // combining 'src' and 'dst' will cause a cycle along that path. + // + // In practice, to avoid modifying the graph and to take advantage + // of existing graph functions, we perform an equivalent. + // 1. Get all nodes incoming to 'dst', excluding 'src' + // 2. Reverse DFS from those nodes + // 3. If reverse DFS reaches 'src' then we have a cycle + std::vector dfs_start_nodes; + for (tensorflow::Node* node : dst->in_nodes()) { + if (node != src) { + dfs_start_nodes.push_back(node); + } + } + + bool is_cycle = false; + if (!dfs_start_nodes.empty()) { + tensorflow::ReverseDFSFrom(graph, dfs_start_nodes, {}, + [&is_cycle, src](tensorflow::Node* node) { + if (node == src) { + is_cycle = true; + } + }); + } + + return !is_cycle; +} + +void ContractEdge(tensorflow::Edge* edge, tensorflow::Graph* graph, + std::vector* remove_edges) { + // Transfer all inputs and outputs of 'dst' to 'src' except edges + // connecting the two. + tensorflow::Node* src = edge->src(); + tensorflow::Node* dst = edge->dst(); + + // We can use '0' for input/output index because we don't need them + // to be accurate for the way we are using the graph. + std::vector in_edges(dst->in_edges().begin(), + dst->in_edges().end()); + for (const tensorflow::Edge* in_edge : in_edges) { + if (in_edge->src() != src) { + tensorflow::Edge* e = const_cast(in_edge); + if (e->src() == graph->source_node()) { + graph->AddEdge(e->src(), e->src_output(), src, + tensorflow::Graph::kControlSlot); + } else { + graph->AddEdge(e->src(), e->src_output(), src, 0 /* input index */); + } + } + } + + std::vector out_edges(dst->out_edges().begin(), + dst->out_edges().end()); + for (const tensorflow::Edge* out_edge : out_edges) { + tensorflow::Edge* e = const_cast(out_edge); + if (e->dst() == graph->sink_node()) { + graph->AddEdge(src, tensorflow::Graph::kControlSlot, e->dst(), + e->dst_input()); + } else { + graph->AddEdge(src, 0 /* output index */, e->dst(), e->dst_input()); + } + } + + // Return the edges that must be removed to disconnect 'dst' from + // the graph. We don't actually remove 'dst' since the caller holds + // references to all the nodes. + for (const auto& in_edge : dst->in_edges()) { + remove_edges->push_back(in_edge); + } + for (const auto& out_edge : dst->out_edges()) { + remove_edges->push_back(out_edge); + } +} + +} // namespace + +tensorflow::Status SegmentGraph( + const tensorflow::GraphDef& gdef, + const std::function& candidate_fn, + const SegmentOptions& options, SegmentNodesVector* segments) { + // Create a Graph representation of the GraphDef. + tensorflow::FunctionLibraryDefinition flib(tensorflow::OpRegistry::Global(), + gdef.library()); + tensorflow::Graph graph(flib); + TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToGraph( + tensorflow::GraphConstructorOptions(), gdef, &graph)); + + // tensorflow::DumpGraph("Pre-Segment", &graph); + + // Use a union-find to collect the nodes that belong to the same + // segment. A node value of nullptr indicates that the node is not a + // candidate for TRT. + std::vector> node_segments; + for (int i = 0; i < graph.num_node_ids(); ++i) { + tensorflow::Node* node = graph.FindNodeId(i); + if (options.exclude_node_list.count(node->name()) != 0 || + !candidate_fn(node->def())) { + node = nullptr; + } + node_segments.emplace_back(node); + } + + // The segmentation algorithm below visits nodes in reverse + // topological order and attempts to merge nodes along output + // edges. That means that subgraphs grow from the output-side of the + // network towards the inputs. In general this is not guaranteed to + // produce a globally optimal segmentation. In the future if we have + // a measure of how beneficial it is to include a given node in a + // TRT subgraph then we can revisit this algorithm to take advantage + // of that information. + std::vector order; + tensorflow::GetPostOrder(graph, &order); + + for (const tensorflow::Node* node : order) { + // All output nodes of 'node' have been visited... + VLOG(2) << "Trying node " << node->name(); + + // 'node' must be a TRT candidate... + if (node_segments[node->id()].Value() == nullptr) { + VLOG(2) << "... not a TRT candidate"; + continue; + } + + // Contract output edges to combine 'node' with output + // nodes. Iterate since combining two nodes may unblock other + // combining. + while (true) { + std::set contract_edges; + for (const tensorflow::Edge* out_edge : node->out_edges()) { + VLOG(2) << "... out node " << out_edge->dst()->name(); + + // Out node must be TRT candidate... + if (node_segments[out_edge->dst()->id()].Value() == nullptr) { + VLOG(2) << "... ... not a TRT candidate"; + continue; + } + + if (CanContractEdge(out_edge, graph)) { + VLOG(2) << "... ... can contract"; + contract_edges.insert(out_edge); + } else { + VLOG(2) << "... ... cannot contract, would form cycle"; + } + } + + if (contract_edges.empty()) { + break; + } + + // Contract edges and collect the adjacent nodes into the same + // segment/subgraph. + while (!contract_edges.empty()) { + const tensorflow::Edge* contract_edge = *contract_edges.begin(); + const tensorflow::Node* src = contract_edge->src(); + const tensorflow::Node* dst = contract_edge->dst(); + + VLOG(2) << "Merge " << src->name() << " <- " << dst->name(); + node_segments[src->id()].Merge(&node_segments[dst->id()]); + + // Contracting the edge leaves disconnected graph edges. + // Remove these from the graph and from 'contract_edges' so we + // don't visit them again. + tensorflow::Edge* e = const_cast(contract_edge); + std::vector remove_edges; + ContractEdge(e, &graph, &remove_edges); + + for (const tensorflow::Edge* r : remove_edges) { + contract_edges.erase(r); + graph.RemoveEdge(r); + } + } + } + } + + // Collect the segments/subgraphs. Each subgraph is represented by a + // set of the names of the nodes in that subgraph. + std::unordered_map> sg_map; + for (auto& u : node_segments) { + if ((u.Value() != nullptr) && (u.ParentValue() != nullptr)) { + sg_map[u.ParentValue()->name()].insert(u.Value()->name()); + } + } + + // Convert the segments into the expected return format + for (const auto& itr : sg_map) { + const auto& segment_node_names = itr.second; + if (VLOG_IS_ON(1)) { + string s; + for (const auto& name : segment_node_names) { + s += " " + name; + } + VLOG(1) << "Segment " << segments->size() << ":" << s; + } + + // Don't use small segments. + if (static_cast(segment_node_names.size()) < + options.minimum_segment_size) { + VLOG(1) << "Segment " << segments->size() << " has only " + << segment_node_names.size() << " nodes, dropping"; + continue; + } + + segments->emplace_back(segment_node_names); + } + + return tensorflow::Status::OK(); +} + +} // namespace segment +} // namespace tensorrt +} // namespace tensorflow diff --git a/tensorflow/contrib/tensorrt/segment/segment.h b/tensorflow/contrib/tensorrt/segment/segment.h new file mode 100644 index 0000000000000000000000000000000000000000..ee6e2b3ed26cd1fabc0e952d882d549046cd9a30 --- /dev/null +++ b/tensorflow/contrib/tensorrt/segment/segment.h @@ -0,0 +1,56 @@ +/* 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_SEGMENT_SEGMENT_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_SEGMENT_H_ + +#include +#include + +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +namespace tensorrt { +namespace segment { + +using SegmentNodesVector = std::vector>; + +struct SegmentOptions { + // Segment must contain at least this many nodes. + int minimum_segment_size = 2; + std::set exclude_node_list; +}; + +// Get the subgraphs of a graph that can be handled by TensorRT. +// +// @param gdef The GraphDef describing the network +// @param candidate_fn A function that returns true for a NodeDef if +// that node can be handled by TensorRT. +// @param segments Returns the TensorRT segments/subgraphs. Each entry +// in the vector describes a subgraph by giving a set of the names of +// all the NodeDefs in that subgraph. +// @return the status. +tensorflow::Status SegmentGraph( + const tensorflow::GraphDef& gdef, + const std::function& candidate_fn, + const SegmentOptions& options, SegmentNodesVector* segments); + +} // namespace segment +} // namespace tensorrt +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_SEGMENT_H_ diff --git a/tensorflow/contrib/tensorrt/segment/segment_test.cc b/tensorflow/contrib/tensorrt/segment/segment_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..74cbc5f2b376b76324eed06d251767da6f928e3e --- /dev/null +++ b/tensorflow/contrib/tensorrt/segment/segment_test.cc @@ -0,0 +1,367 @@ +/* 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/segment/segment.h" +#include "tensorflow/c/c_api.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +namespace tensorrt { +namespace segment { +namespace test { + +class SegmentTest : public ::testing::Test { + public: + bool GetGraphDef(TF_Graph* graph, tensorflow::GraphDef* graph_def); + + TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s, const char* name); + TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name); + + std::function MakeCandidateFn( + const std::set& node_names); + + protected: + void PlaceholderHelper(TF_Graph* graph, TF_Status* s, const char* name, + TF_Operation** op); + void AddHelper(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name, TF_Operation** op, bool check); + + SegmentOptions default_options_; +}; + +bool SegmentTest::GetGraphDef(TF_Graph* graph, + tensorflow::GraphDef* graph_def) { + TF_Status* s = TF_NewStatus(); + TF_Buffer* buffer = TF_NewBuffer(); + TF_GraphToGraphDef(graph, buffer, s); + bool ret = TF_GetCode(s) == TF_OK; + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + if (ret) ret = graph_def->ParseFromArray(buffer->data, buffer->length); + TF_DeleteBuffer(buffer); + TF_DeleteStatus(s); + return ret; +} + +std::function SegmentTest::MakeCandidateFn( + const std::set& node_names) { + return [node_names](const NodeDef& node) -> bool { + return node_names.find(node.name()) != node_names.end(); + }; +} + +void SegmentTest::PlaceholderHelper(TF_Graph* graph, TF_Status* s, + const char* name, TF_Operation** op) { + TF_OperationDescription* desc = TF_NewOperation(graph, "Placeholder", name); + TF_SetAttrType(desc, "dtype", TF_INT32); + *op = TF_FinishOperation(desc, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); +} + +TF_Operation* SegmentTest::Placeholder(TF_Graph* graph, TF_Status* s, + const char* name) { + TF_Operation* op; + PlaceholderHelper(graph, s, name, &op); + return op; +} + +void SegmentTest::AddHelper(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name, TF_Operation** op, + bool check) { + TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name); + TF_Output add_inputs[2] = {{l, 0}, {r, 0}}; + TF_AddInputList(desc, add_inputs, 2); + *op = TF_FinishOperation(desc, s); + if (check) { + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); + } +} + +TF_Operation* SegmentTest::Add(TF_Operation* l, TF_Operation* r, + TF_Graph* graph, TF_Status* s, + const char* name) { + TF_Operation* op; + AddHelper(l, r, graph, s, name, &op, true); + return op; +} + +TEST_F(SegmentTest, Empty) { + TF_Graph* graph = TF_NewGraph(); + + GraphDef graph_def; + ASSERT_TRUE(GetGraphDef(graph, &graph_def)); + + SegmentNodesVector segments; + ASSERT_EQ( + SegmentGraph(graph_def, MakeCandidateFn({}), default_options_, &segments), + tensorflow::Status::OK()); + + // Expect no segments/subgraphs. + EXPECT_TRUE(segments.empty()); + TF_DeleteGraph(graph); +} + +TEST_F(SegmentTest, Simple) { + TF_Status* s = TF_NewStatus(); + TF_Graph* graph = TF_NewGraph(); + + // feed + // // || + // add0 add1 + // | | / + // | add2 + // | / || + // add3 add4 + // | / + // + // + TF_Operation* feed = Placeholder(graph, s, "feed"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("feed"), string(TF_OperationName(feed))); + + TF_Operation* add0 = Add(feed, feed, graph, s, "add0"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add1 = Add(feed, feed, graph, s, "add1"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add2 = Add(add0, add1, graph, s, "add2"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add3 = Add(add0, add2, graph, s, "add3"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("add3"), string(TF_OperationName(add3))); + TF_Operation* add4 = Add(add2, add2, graph, s, "add4"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("add4"), string(TF_OperationName(add4))); + + GraphDef graph_def; + ASSERT_TRUE(GetGraphDef(graph, &graph_def)); + + SegmentNodesVector segments; + ASSERT_EQ( + SegmentGraph(graph_def, + MakeCandidateFn({"add0", "add1", "add2", "add3", "add4"}), + default_options_, &segments), + tensorflow::Status::OK()); + + // Expect all Add operations to be collapsed into a single segment + ASSERT_EQ(segments.size(), 1); + std::vector expected{"add0", "add1", "add2", "add3", "add4"}; + for (const auto& ex : expected) { + EXPECT_TRUE(segments[0].find(ex) != segments[0].end()) + << "Missing expected node " << ex; + } + TF_DeleteGraph(graph); + TF_DeleteStatus(s); +} + +TEST_F(SegmentTest, AvoidCycle) { + TF_Status* s = TF_NewStatus(); + TF_Graph* graph = TF_NewGraph(); + + // add2 is not a TRT candidate so add0/add3 cannot be formed as a + // subgraph + // + // feed + // // || + // add0 add1 + // | | / + // | add2 + // | / || + // add3 add4 + // | / + // + // + TF_Operation* feed = Placeholder(graph, s, "feed"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("feed"), string(TF_OperationName(feed))); + + TF_Operation* add0 = Add(feed, feed, graph, s, "add0"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add1 = Add(feed, feed, graph, s, "add1"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add2 = Add(add0, add1, graph, s, "add2"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add3 = Add(add0, add2, graph, s, "add3"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("add3"), string(TF_OperationName(add3))); + TF_Operation* add4 = Add(add2, add2, graph, s, "add4"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("add4"), string(TF_OperationName(add4))); + + GraphDef graph_def; + ASSERT_TRUE(GetGraphDef(graph, &graph_def)); + + SegmentNodesVector segments; + ASSERT_EQ( + SegmentGraph(graph_def, MakeCandidateFn({"add0", "add1", "add3", "add4"}), + default_options_, &segments), + tensorflow::Status::OK()); + + // Expect no subgraphs + EXPECT_EQ(segments.size(), 0); + TF_DeleteGraph(graph); + TF_DeleteStatus(s); +} + +TEST_F(SegmentTest, Multiple) { + TF_Status* s = TF_NewStatus(); + TF_Graph* graph = TF_NewGraph(); + + // add5 is not a TRT candidate so two subgraphs should be formed + // + // feed + // // || || + // add0 add1 add7 + // | | / / || + // | add2-----add5 add8 + // | / | | | | + // add3 add4 add6 + // | | / + // + // + TF_Operation* feed = Placeholder(graph, s, "feed"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("feed"), string(TF_OperationName(feed))); + + TF_Operation* add0 = Add(feed, feed, graph, s, "add0"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add1 = Add(feed, feed, graph, s, "add1"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add7 = Add(feed, feed, graph, s, "add7"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add2 = Add(add0, add1, graph, s, "add2"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add5 = Add(add2, add7, graph, s, "add5"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add8 = Add(add7, add7, graph, s, "add8"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add3 = Add(add0, add2, graph, s, "add3"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("add3"), string(TF_OperationName(add3))); + TF_Operation* add4 = Add(add2, add5, graph, s, "add4"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("add4"), string(TF_OperationName(add4))); + TF_Operation* add6 = Add(add5, add8, graph, s, "add6"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("add6"), string(TF_OperationName(add6))); + + GraphDef graph_def; + ASSERT_TRUE(GetGraphDef(graph, &graph_def)); + + SegmentNodesVector segments; + ASSERT_EQ(SegmentGraph(graph_def, + MakeCandidateFn({"add0", "add1", "add2", "add3", + "add4", "add6", "add7", "add8"}), + default_options_, &segments), + tensorflow::Status::OK()); + + // Expect two subgraphs + EXPECT_EQ(segments.size(), 2); + + std::vector expected0{"add0", "add1", "add2", "add3"}; + for (const auto& ex : expected0) { + EXPECT_TRUE(segments[0].find(ex) != segments[0].end()) + << "Missing expected node " << ex; + } + + std::vector expected1{"add6", "add8"}; + for (const auto& ex : expected1) { + EXPECT_TRUE(segments[1].find(ex) != segments[1].end()) + << "Missing expected node " << ex; + } + TF_DeleteGraph(graph); + TF_DeleteStatus(s); +} + +TEST_F(SegmentTest, BigIfElse) { + TF_Status* s = TF_NewStatus(); + TF_Graph* graph = TF_NewGraph(); + + // add2 is not a TRT candidate + // + // feed + // || + // add0 + // // || + // add1 add4 + // || || + // add2 add5 + // || || + // add3 add6 + // || // + // add7 + // || + // + // + TF_Operation* feed = Placeholder(graph, s, "feed"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("feed"), string(TF_OperationName(feed))); + + TF_Operation* add0 = Add(feed, feed, graph, s, "add0"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add1 = Add(add0, add0, graph, s, "add1"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add2 = Add(add1, add1, graph, s, "add2"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add3 = Add(add2, add2, graph, s, "add3"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add4 = Add(add0, add0, graph, s, "add4"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add5 = Add(add4, add4, graph, s, "add5"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add6 = Add(add5, add5, graph, s, "add6"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_Operation* add7 = Add(add3, add6, graph, s, "add7"); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + EXPECT_EQ(string("add7"), string(TF_OperationName(add7))); + + GraphDef graph_def; + ASSERT_TRUE(GetGraphDef(graph, &graph_def)); + + SegmentNodesVector segments; + ASSERT_EQ(SegmentGraph(graph_def, + MakeCandidateFn({"add0", "add1", "add3", "add4", + "add5", "add6", "add7"}), + default_options_, &segments), + tensorflow::Status::OK()); + + // Expect 2 subgraphs + EXPECT_EQ(segments.size(), 2); + + std::vector expected0{"add3", "add4", "add5", "add6", "add7"}; + for (const auto& ex : expected0) { + EXPECT_TRUE(segments[0].find(ex) != segments[0].end()) + << "Missing expected node " << ex; + } + + std::vector expected1{"add0", "add1"}; + for (const auto& ex : expected1) { + EXPECT_TRUE(segments[1].find(ex) != segments[1].end()) + << "Missing expected node " << ex; + } + TF_DeleteGraph(graph); + TF_DeleteStatus(s); +} + +} // namespace test +} // namespace segment +} // namespace tensorrt +} // namespace tensorflow diff --git a/tensorflow/contrib/tensorrt/segment/union_find.h b/tensorflow/contrib/tensorrt/segment/union_find.h new file mode 100644 index 0000000000000000000000000000000000000000..1c64ebbb0ae532a4776ab8963515d19fd3b23b4c --- /dev/null +++ b/tensorflow/contrib/tensorrt/segment/union_find.h @@ -0,0 +1,79 @@ +/* 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_SEGMENT_UNION_FIND_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_UNION_FIND_H_ + +namespace tensorflow { +namespace tensorrt { +namespace segment { + +// Union-Find data structure. +// Each cluster has an associated value; when merging clusters we can control +// which value becomes the representative of the merged clusters. Values must be +// copyable. +template +class UnionFind { + public: + UnionFind() : size_(1), parent_(nullptr) {} + explicit UnionFind(const T& v) : size_(1), parent_(nullptr), value_(v) {} + + // Returns the number of elements in a cluster. + int Size() { return FindRoot()->size_; } + + // Merges this cluster with 'other'. This cluster's value becomes + // the value of the merged cluster; the value of 'other' is ignored. + void Merge(UnionFind* other); + + // Each cluster has an associated value. Retrieves the value associated + // with this cluster. + T& ParentValue() { return FindRoot()->value_; } + + // Get the original value of this node. + T& Value() { return value_; } + + private: + // Finds the root element of the cluster. Performs path compression. + UnionFind* FindRoot(); + + int size_; + UnionFind* parent_; + T value_; +}; + +template +void UnionFind::Merge(UnionFind* other) { + UnionFind* a = FindRoot(); + UnionFind* b = other->FindRoot(); + if (a == b) return; + + b->parent_ = a; + a->size_ += b->size_; +} + +template +UnionFind* UnionFind::FindRoot() { + if (!parent_) return this; + // Path compression: update intermediate nodes to point to the root of the + // equivalence class. + parent_ = parent_->FindRoot(); + return parent_; +} + +} // namespace segment +} // namespace tensorrt +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_UNION_FIND_H_ diff --git a/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.cc b/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.cc new file mode 100644 index 0000000000000000000000000000000000000000..8b475177bc670ddae2b26b6a494f758eba20b2c3 --- /dev/null +++ b/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.cc @@ -0,0 +1,89 @@ +/* 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/shape_fn/trt_shfn.h" + +#include +#include + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorflow/contrib/tensorrt/log/trt_logger.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace shape_inference { + +tensorflow::Status TRTEngineOpShapeInference(InferenceContext* context) { + tensorflow::tensorrt::Logger logger; + string serialized_engine; + TF_RETURN_IF_ERROR(context->GetAttr("serialized_engine", &serialized_engine)); + nvinfer1::IRuntime* infer = nvinfer1::createInferRuntime(logger); + nvinfer1::ICudaEngine* trt_engine = infer->deserializeCudaEngine( + serialized_engine.c_str(), serialized_engine.size(), nullptr); + + int num_batch = -1; + std::vector<::tensorflow::DataType> input_type; + TF_RETURN_IF_ERROR(context->GetAttr("InT", &input_type)); + for (size_t i = 0; i < context->num_inputs(); i++) { + // Check if input shape is legit + auto input_shape = context->input(i); + for (int j = 0; j < context->Rank(input_shape); j++) { + auto dim_handler = context->Dim(input_shape, j); + if (j == 0) { + if (i == 0) { + num_batch = context->Value(dim_handler); + } else if (num_batch != context->Value(dim_handler)) { + // TODO(jie): TensorRT engine requires consistent batch between inputs + // tensors. Segmenter should be aware of this. + LOG(FATAL) << "TensorRT engine requires consistent batch size"; + } + } + } + } + + // Arrange input here + std::vector input_nodes; + TF_RETURN_IF_ERROR(context->GetAttr("input_nodes", &input_nodes)); + + // Arrange output here + std::vector output_nodes; + TF_RETURN_IF_ERROR(context->GetAttr("output_nodes", &output_nodes)); + for (size_t i = 0; i < output_nodes.size(); i++) { + int binding_index = trt_engine->getBindingIndex(output_nodes[i].c_str()); + ShapeHandle output_shape; + std::vector dim_vec; + dim_vec.emplace_back(context->MakeDim(num_batch)); + if (binding_index != -1) { + auto dims = trt_engine->getBindingDimensions(binding_index); + for (int j = 0; j < dims.nbDims; j++) { + dim_vec.emplace_back(context->MakeDim(dims.d[j])); + } + } else { + LOG(FATAL) << "TensorRT engine cannot find binding: " << output_nodes[i]; + } + output_shape = context->MakeShape(dim_vec); + context->set_output(i, output_shape); + } + + return Status::OK(); +} + +} // namespace shape_inference +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.h b/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.h new file mode 100644 index 0000000000000000000000000000000000000000..4b50f66699f0965639e22169ee7d71e860314bf0 --- /dev/null +++ b/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.h @@ -0,0 +1,33 @@ +/* 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_SHAPE_FN_TRT_SHFN_H_ +#define TENSORFLOW_CONTRIB_TENSORRT_SHAPE_FN_TRT_SHFN_H_ + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorflow/core/framework/shape_inference.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { +namespace shape_inference { +Status TRTEngineOpShapeInference(InferenceContext* c); +} // namespace shape_inference +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +#endif // TENSORFLOW_CONTRIB_TENSORRT_SHAPE_FN_TRT_SHFN_H_ diff --git a/tensorflow/contrib/tensorrt/tensorrt_test.cc b/tensorflow/contrib/tensorrt/tensorrt_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..e11522ea5bda7f5a303d6ea332148dbd7b17f162 --- /dev/null +++ b/tensorflow/contrib/tensorrt/tensorrt_test.cc @@ -0,0 +1,159 @@ +/* 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/platform/logging.h" +#include "tensorflow/core/platform/test.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "cuda/include/cuda.h" +#include "cuda/include/cuda_runtime_api.h" +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace { + +class Logger : public nvinfer1::ILogger { + public: + void log(nvinfer1::ILogger::Severity severity, const char* msg) override { + switch (severity) { + case Severity::kINFO: + LOG(INFO) << msg; + break; + case Severity::kWARNING: + LOG(WARNING) << msg; + break; + case Severity::kINTERNAL_ERROR: + case Severity::kERROR: + LOG(ERROR) << msg; + break; + default: + break; + } + } +}; + +class ScopedWeights { + public: + ScopedWeights(float value) : value_(value) { + w.type = nvinfer1::DataType::kFLOAT; + w.values = &value_; + w.count = 1; + } + const nvinfer1::Weights& get() { return w; } + + private: + float value_; + nvinfer1::Weights w; +}; + +const char* kInputTensor = "input"; +const char* kOutputTensor = "output"; + +// Creates a network to compute y=2x+3. +nvinfer1::IHostMemory* CreateNetwork() { + Logger logger; + nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger); + ScopedWeights weights(2.0); + ScopedWeights bias(3.0); + + nvinfer1::INetworkDefinition* network = builder->createNetwork(); + // Add the input. + auto input = network->addInput(kInputTensor, nvinfer1::DataType::kFLOAT, + nvinfer1::DimsCHW{1, 1, 1}); + EXPECT_NE(input, nullptr); + // Add the hidden layer. + auto layer = network->addFullyConnected(*input, 1, weights.get(), bias.get()); + EXPECT_NE(layer, nullptr); + // Mark the output. + auto output = layer->getOutput(0); + output->setName(kOutputTensor); + network->markOutput(*output); + // Build the engine + builder->setMaxBatchSize(1); + builder->setMaxWorkspaceSize(1 << 10); + auto engine = builder->buildCudaEngine(*network); + EXPECT_NE(engine, nullptr); + // Serialize the engine to create a model, then close everything. + nvinfer1::IHostMemory* model = engine->serialize(); + network->destroy(); + engine->destroy(); + builder->destroy(); + return model; +} + +// Executes the network. +void Execute(nvinfer1::IExecutionContext& context, const float* input, + float* output) { + const nvinfer1::ICudaEngine& engine = context.getEngine(); + + // We have two bindings: input and output. + ASSERT_EQ(engine.getNbBindings(), 2); + const int input_index = engine.getBindingIndex(kInputTensor); + const int output_index = engine.getBindingIndex(kOutputTensor); + + // Create GPU buffers and a stream + void* buffers[2]; + ASSERT_EQ(0, cudaMalloc(&buffers[input_index], sizeof(float))); + ASSERT_EQ(0, cudaMalloc(&buffers[output_index], sizeof(float))); + cudaStream_t stream; + ASSERT_EQ(0, cudaStreamCreate(&stream)); + + // Copy the input to the GPU, execute the network, and copy the output back. + // + // Note that since the host buffer was not created as pinned memory, these + // async copies are turned into sync copies. So the following synchronization + // could be removed. + ASSERT_EQ(0, cudaMemcpyAsync(buffers[input_index], input, sizeof(float), + cudaMemcpyHostToDevice, stream)); + context.enqueue(1, buffers, stream, nullptr); + ASSERT_EQ(0, cudaMemcpyAsync(output, buffers[output_index], sizeof(float), + cudaMemcpyDeviceToHost, stream)); + cudaStreamSynchronize(stream); + + // Release the stream and the buffers + cudaStreamDestroy(stream); + ASSERT_EQ(0, cudaFree(buffers[input_index])); + ASSERT_EQ(0, cudaFree(buffers[output_index])); +} + +TEST(TensorrtTest, BasicFunctions) { + // Create the network model. + nvinfer1::IHostMemory* model = CreateNetwork(); + // Use the model to create an engine and then an execution context. + Logger logger; + nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger); + nvinfer1::ICudaEngine* engine = + runtime->deserializeCudaEngine(model->data(), model->size(), nullptr); + model->destroy(); + nvinfer1::IExecutionContext* context = engine->createExecutionContext(); + + // Execute the network. + float input = 1234; + float output; + Execute(*context, &input, &output); + EXPECT_EQ(output, input * 2 + 3); + + // Destroy the engine. + context->destroy(); + engine->destroy(); + runtime->destroy(); +} + +} // namespace +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/test/test_tftrt.py b/tensorflow/contrib/tensorrt/test/test_tftrt.py new file mode 100644 index 0000000000000000000000000000000000000000..18dba94acb3724cb2b5a1c53227bcf08bf9f8fcc --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/test_tftrt.py @@ -0,0 +1,88 @@ +# 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. +# ============================================================================== +"""Script to test TF-TensorRT integration.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +# normally we should do import tensorflow as tf and then +# tf.placeholder, tf.constant, tf.nn.conv2d etc but +# it looks like internal builds don't like it so +# importing every module individually + +from tensorflow.contrib import tensorrt as trt +from tensorflow.core.protobuf import config_pb2 as cpb2 +from tensorflow.python.client import session as csess +from tensorflow.python.framework import constant_op as cop +from tensorflow.python.framework import dtypes as dtypes +from tensorflow.python.framework import importer as importer +from tensorflow.python.framework import ops as ops +from tensorflow.python.ops import array_ops as aops +from tensorflow.python.ops import nn as nn +from tensorflow.python.ops import nn_ops as nn_ops + + +def get_simple_graph_def(): + """Create a simple graph and return its graph_def""" + g = ops.Graph() + with g.as_default(): + a = aops.placeholder( + dtype=dtypes.float32, shape=(None, 24, 24, 2), name="input") + e = cop.constant( + [[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]], + name="weights", + dtype=dtypes.float32) + conv = nn.conv2d( + input=a, filter=e, strides=[1, 2, 2, 1], padding="SAME", name="conv") + b = cop.constant( + [4., 1.5, 2., 3., 5., 7.], name="bias", dtype=dtypes.float32) + t = nn.bias_add(conv, b, name="biasAdd") + relu = nn.relu(t, "relu") + idty = aops.identity(relu, "ID") + v = nn_ops.max_pool( + idty, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool") + aops.squeeze(v, name="output") + return g.as_graph_def() + + +def run_graph(gdef, dumm_inp): + gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50) + ops.reset_default_graph() + g = ops.Graph() + with g.as_default(): + inp, out = importer.import_graph_def( + graph_def=gdef, return_elements=["input", "output"]) + inp = inp.outputs[0] + out = out.outputs[0] + with csess.Session( + config=cpb2.ConfigProto(gpu_options=gpu_options), graph=g) as sess: + val = sess.run(out, {inp: dumm_inp}) + return val + + +if "__main__" in __name__: + inp_dims = (100, 24, 24, 2) + dummy_input = np.random.random_sample(inp_dims) + gdef = get_simple_graph_def() + # Get optimized graph + trt_graph = trt.create_inference_graph(gdef, ["output"], inp_dims[0]) + o1 = run_graph(gdef, dummy_input) + o2 = run_graph(trt_graph, dummy_input) + o3 = run_graph(trt_graph, dummy_input) + assert np.array_equal(o1, o2) + assert np.array_equal(o3, o2) # sanity check + print("Pass") diff --git a/tensorflow/contrib/tensorrt/trt_conversion.i b/tensorflow/contrib/tensorrt/trt_conversion.i new file mode 100644 index 0000000000000000000000000000000000000000..d679945d569c1784448b6cb09c2f431b9cda56d7 --- /dev/null +++ b/tensorflow/contrib/tensorrt/trt_conversion.i @@ -0,0 +1,131 @@ +/* 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. +==============================================================================*/ + +/* Wrap trt_conversion */ +%{ +#define SWIG_FILE_WITH_INIT +%} +%include "std_pair.i" +%include "tensorflow/python/platform/base.i" + +%{ +PyObject* pair_helper(std::pair* in) { + PyObject *first(nullptr), *second(nullptr), *tuple(nullptr); + first = PyBytes_FromStringAndSize(in->first.data(), in->first.length()); + if (!first) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, "Pair conversion first argument failed"); + } + return NULL; + } + second = PyBytes_FromStringAndSize(in->second.data(), in->second.length()); + if (!second) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, + "Pair conversion second argument failed"); + } + return NULL; + } + tuple = Py_BuildValue("(OO)", first, second); + if (!tuple) { + if (!PyErr_Occurred()) { + PyErr_SetString(PyExc_TypeError, + "Tuple creation from pair failed!"); + } + return NULL; + } + return tuple; +} +%} +%typemap(out) std::pair { + PyObject *tuple = pair_helper(&$1); + if (!tuple) SWIG_fail; + $result = tuple; +} +%{ +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/util/stat_summarizer.h" +#include "tensorflow/contrib/tensorrt/convert/convert_graph.h" +%} + +%ignoreall +%unignore tensorflow; +%unignore trt_convert; + +%{ +std::pair trt_convert( + string graph_def_string, // The serialized GraphDef string. + std::vector output_names, + size_t max_batch_size, + size_t max_workspace_size_bytes + // Unfortunately we can't use TF_Status here since it + // is in c/c_api and brings in a lot of other libraries + // which in turn declare ops. These ops are included + // statically in our library and cause an abort when + // module is loaded due to double registration + // until Tensorflow properly exposes these headers + // we have to work around this by returning a string + // and converting it to exception on python side. + //,TF_Status* out_status) { +) { +#if GOOGLE_CUDA && GOOGLE_TENSORRT + string out_status; + + tensorflow::GraphDef graph_def; + if (!graph_def.ParseFromString(graph_def_string)) { + out_status = "InvalidArgument;Couldn't interpret input as a GraphDef"; + return std::pair{out_status, ""}; + } + + if (!output_names.size()) { + out_status = "InvalidArgument;Size of the output_names vector is 0"; + return std::pair{out_status, ""}; + // return ""; + } + tensorflow::GraphDef outGraph; + tensorflow::Status conversion_status = + tensorflow::tensorrt::convert::ConvertGraphDefToTensorRT( + graph_def, output_names, max_batch_size, max_workspace_size_bytes, + &outGraph); + if (!conversion_status.ok()) { + auto retCode = (int)conversion_status.code(); + char buff[2000]; + snprintf(buff, 2000, "%d;%s", retCode, + conversion_status.error_message().c_str()); + out_status = buff; + return std::pair{out_status, ""}; + } + string result; + if (!outGraph.SerializeToString(&result)) { + out_status = "InvalidArgument;Couldn't serialize output as a GraphDef"; + return std::pair{out_status, ""}; + } + out_status = "OK;All good!"; + return std::pair{out_status, result}; +#else + // Returns FAILED_PRECONDITION. + return std::pair{"9;TensorRT is not enabled!", ""}; +#endif // GOOGLE_CUDA && GOOGLE_TENSORRT +} +%} + +std::pair trt_convert(string graph_def_string, + std::vector output_names, + size_t max_batch_size, + size_t max_workspace_size_bytes); + + +%unignoreall diff --git a/tensorflow/contrib/timeseries/examples/data/multivariate_periods.csv b/tensorflow/contrib/timeseries/examples/data/multivariate_periods.csv index 02a60d1cf61765c7c916803fe918d8b7b186405e..b49a0662c29b1d810f4be31ca1f318f0571f533e 100644 --- a/tensorflow/contrib/timeseries/examples/data/multivariate_periods.csv +++ b/tensorflow/contrib/timeseries/examples/data/multivariate_periods.csv @@ -1,100 +1,100 @@ -0,0.926906299771,1.99107237682,2.56546245685,3.07914768197,4.04839057867 -1,0.108010001864,1.41645361423,2.1686839775,2.94963962176,4.1263503303 -2,-0.800567600028,1.0172132907,1.96434754116,2.99885333086,4.04300485864 -3,0.0607042871898,0.719540073421,1.9765012584,2.89265588817,4.0951014426 -4,0.933712200629,0.28052120776,1.41018552514,2.69232603996,4.06481164223 -5,-0.171730652974,0.260054421028,1.48770816369,2.62199129293,4.44572807842 -6,-1.00180162933,0.333045158863,1.50006392277,2.88888309683,4.24755865606 -7,0.0580061875336,0.688929398826,1.56543458772,2.99840358953,4.52726873347 -8,0.764139447412,1.24704875327,1.77649279698,3.13578593851,4.63238922951 -9,-0.230331874785,1.47903998963,2.03547545751,3.20624030377,4.77980005228 -10,-1.03846045211,2.01133000781,2.31977503972,3.67951536251,5.09716775897 -11,0.188643592253,2.23285349038,2.68338482249,3.49817168611,5.24928239634 -12,0.91207302309,2.24244446841,2.71362604985,3.96332587625,5.37802271594 -13,-0.296588665881,2.02594634141,3.07733910479,3.99698324956,5.56365901394 -14,-0.959961476551,1.45078629833,3.18996420137,4.3763059609,5.65356015609 -15,0.46313530679,1.01141441548,3.4980215948,4.20224896882,5.88842247449 -16,0.929354125798,0.626635305936,3.70508262244,4.51791573544,5.73945973251 -17,-0.519110731957,0.269249223148,3.39866823332,4.46802003061,5.82768174382 -18,-0.924330981367,0.349602834684,3.21762413294,4.72803587499,5.94918925767 -19,0.253239387885,0.345158023497,3.11071425333,4.79311566935,5.9489259713 -20,0.637408390225,0.698996675371,3.25232492145,4.73814732384,5.9612010251 -21,-0.407396859412,1.17456342803,2.49526823723,4.59323415742,5.82501686811 -22,-0.967485452118,1.66655933642,2.47284606244,4.58316034754,5.88721406681 -23,0.474480867904,1.95018556323,2.0228950072,4.48651142819,5.8255943735 -24,1.04309652155,2.23519892356,1.91924131572,4.19094661783,5.87457348436 -25,-0.517861513772,2.12501967336,1.70266619979,4.05280882887,5.72160912899 -26,-0.945301585146,1.65464653549,1.81567174251,3.92309850635,5.58270493814 -27,0.501153868974,1.40600764889,1.53991387719,3.72853247942,5.60169001727 -28,0.972859524418,1.00344321868,1.5175642828,3.64092376655,5.10567722582 -29,-0.70553406135,0.465306263885,1.7038540803,3.33236870312,5.09182481555 -30,-0.946093634916,0.294539309453,1.88052827037,2.93011492669,4.97354922696 -31,0.47922123231,0.308465865031,2.03445883031,2.90772899045,4.86241793548 -32,0.754030014252,0.549752241167,2.46115815089,2.95063349534,4.71834614627 -33,-0.64875949826,0.894615488148,2.5922463381,2.81269864022,4.43480095104 -34,-0.757829951086,1.39123914261,2.69258079904,2.61834837315,4.36580046156 -35,0.565653301088,1.72360022693,2.97794913834,2.80403840334,4.27327248459 -36,0.867440092372,2.21100730052,3.38648090792,2.84057515729,4.12210169576 -37,-0.894567758095,2.17549105818,3.45532493329,2.90446025717,4.00251740584 -38,-0.715442356893,2.15105389965,3.52041791902,3.03650393392,4.12809249577 -39,0.80671703672,1.81504564517,3.60463324866,3.00747789871,3.98440762467 -40,0.527014790142,1.31803513865,3.43842186337,3.3332594663,4.03232406566 -41,-0.795936862129,0.847809114454,3.09875133548,3.52863155938,3.94883924909 -42,-0.610245806946,0.425530441018,2.92581949152,3.77238736123,4.27287245021 -43,0.611662279431,0.178432049837,2.48128214822,3.73212087883,4.17319013831 -44,0.650866553108,0.220341648392,2.41694642022,4.2609098519,4.27271645905 -45,-0.774156982023,0.632667602331,2.05474356052,4.32889204886,4.18029723271 -46,-0.714058448409,0.924562377599,1.75706135146,4.52492718422,4.3972678094 -47,0.889627293379,1.46207968841,1.78299357672,4.64466731095,4.56317887554 -48,0.520140662861,1.8996333843,1.41377633823,4.48899091177,4.78805049769 -49,-1.03816935616,2.08997002059,1.51218375351,4.84167764204,4.93026048606 -50,-0.40772951362,2.30878972136,1.44144415128,4.76854460997,5.01538444629 -51,0.792730684781,1.91367048509,1.58887384677,4.71739397335,5.25690012199 -52,0.371311881576,1.67565079528,1.81688563053,4.60353107555,5.44265822961 -53,-0.814398070371,1.13374634126,1.80328814859,4.72264252878,5.52674761122 -54,-0.469017949323,0.601244136627,2.29690896736,4.49859178859,5.54126153454 -55,0.871044371426,0.407597593794,2.7499112487,4.19060637761,5.57693767301 -56,0.523764933017,0.247705192709,3.09002071379,4.02095509006,5.80510362182 -57,-0.881326403531,0.31513103164,3.11358205718,3.96079100808,5.81000652365 -58,-0.357928025339,0.486163915865,3.17884556771,3.72634990659,5.85693642011 -59,0.853038779822,1.04218094475,3.45835384454,3.36703969978,5.9585988449 -60,0.435311516013,1.59715085283,3.63313338588,3.11276729421,5.93643818229 -61,-1.02703719138,1.92205832542,3.47606111735,3.06247155999,6.02106646259 -62,-0.246661325557,2.14653802542,3.29446326567,2.89936259181,5.67531541272 -63,1.02554736569,2.25943737733,3.07031591528,2.78176218013,5.78206328989 -64,0.337814475969,2.07589147224,2.80356226089,2.55888206331,5.7094075496 -65,-1.12023369929,1.25333011618,2.56497288445,2.77361359194,5.50799418376 -66,-0.178980246554,1.11937139901,2.51598681313,2.91438309151,5.47469577206 -67,0.97550951531,0.60553823137,2.11657741073,2.88081098981,5.37034999502 -68,0.136653357206,0.365828836075,1.97386033165,3.13217903204,5.07254490219 -69,-1.05607596951,0.153152115069,1.52110743825,3.01308794192,5.08902539125 -70,-0.13095280331,0.337113974483,1.52703079853,3.16687131599,4.86649398514 -71,1.07081057754,0.714247566736,1.53761382634,3.45151989484,4.75892309166 -72,0.0153410376082,1.24631231847,1.61690939161,3.85481994498,4.35683752832 -73,-0.912801257303,1.60791309476,1.8729264524,4.03037260012,4.36072588913 -74,-0.0894895640338,2.02535207407,1.93484909619,4.09557485132,4.35327025188 -75,0.978646999652,2.20085086625,2.09003440427,4.27542353033,4.1805058388 -76,-0.113312642876,2.2444100761,2.50789248839,4.4151861502,4.03267168136 -77,-1.00215099149,1.84305628445,2.61691237246,4.45425147595,3.81203553766 -78,-0.0183234614205,1.49573923116,2.99308471214,4.71134960112,4.0273804959 -79,1.0823738177,1.12211589848,3.27079386925,4.94288270502,4.01851068083 -80,0.124370187893,0.616474412808,3.4284236674,4.76942168327,3.9749536483 -81,-0.929423379352,0.290977090976,3.34131726136,4.78590392707,4.10190661656 -82,0.23766302648,0.155302052254,3.49779513794,4.64605656795,4.15571321107 -83,1.03531486192,0.359702776204,3.4880725919,4.48167586667,4.21134561991 -84,-0.261234571382,0.713877760378,3.42756426614,4.426443869,4.25208300527 -85,-1.03572442277,1.25001113691,2.96908341113,4.25500915322,4.25723010649 -86,0.380034261243,1.70543355622,2.73605932518,4.16703432307,4.63700400788 -87,1.03734873488,1.97544410562,2.55586572141,3.84976673263,4.55282864289 -88,-0.177344253372,2.22614526325,2.09565864891,3.77378097953,4.82577400298 -89,-0.976821526892,2.18385079177,1.78522284118,3.67768223554,5.06302440873 -90,0.264820472091,1.86981946157,1.50048403865,3.43619796921,5.05651761669 -91,1.05642344868,1.47568646076,1.51347671977,3.20898518885,5.50149047462 -92,-0.311607433358,1.04226467636,1.52089650905,3.02291865417,5.4889046232 -93,-0.724285777937,0.553052311957,1.48573560173,2.7365973598,5.72549174225 -94,0.519859192905,0.226520626591,1.61543723167,2.84102086852,5.69330622288 -95,1.0323195039,0.260873217055,1.81913034804,2.83951143848,5.90325028086 -96,-0.53285682538,0.387695521405,1.70935609313,2.57977050631,5.79579213161 -97,-0.975127997215,0.920948771589,2.51292643636,2.71004616612,5.87016469227 -98,0.540246804099,1.36445470181,2.61949412896,2.98482553485,6.02447664937 -99,0.987764008058,1.85581989607,2.84685706149,2.94760204892,6.0212151724 +0,0.926906299771,1.99107237682,2.56546245685,3.07914768197,4.04839057867,1.,0. +1,0.108010001864,1.41645361423,2.1686839775,2.94963962176,4.1263503303,1.,0. +2,-0.800567600028,1.0172132907,1.96434754116,2.99885333086,4.04300485864,1.,0. +3,0.0607042871898,0.719540073421,1.9765012584,2.89265588817,4.0951014426,1.,0. +4,0.933712200629,0.28052120776,1.41018552514,2.69232603996,4.06481164223,1.,0. +5,-0.171730652974,0.260054421028,1.48770816369,2.62199129293,4.44572807842,1.,0. +6,-1.00180162933,0.333045158863,1.50006392277,2.88888309683,4.24755865606,1.,0. +7,0.0580061875336,0.688929398826,1.56543458772,2.99840358953,4.52726873347,1.,0. +8,0.764139447412,1.24704875327,1.77649279698,3.13578593851,4.63238922951,1.,0. +9,-0.230331874785,1.47903998963,2.03547545751,3.20624030377,4.77980005228,1.,0. +10,-1.03846045211,2.01133000781,2.31977503972,3.67951536251,5.09716775897,1.,0. +11,0.188643592253,2.23285349038,2.68338482249,3.49817168611,5.24928239634,1.,0. +12,0.91207302309,2.24244446841,2.71362604985,3.96332587625,5.37802271594,1.,0. +13,-0.296588665881,2.02594634141,3.07733910479,3.99698324956,5.56365901394,1.,0. +14,-0.959961476551,1.45078629833,3.18996420137,4.3763059609,5.65356015609,1.,0. +15,0.46313530679,1.01141441548,3.4980215948,4.20224896882,5.88842247449,1.,0. +16,0.929354125798,0.626635305936,3.70508262244,4.51791573544,5.73945973251,1.,0. +17,-0.519110731957,0.269249223148,3.39866823332,4.46802003061,5.82768174382,1.,0. +18,-0.924330981367,0.349602834684,3.21762413294,4.72803587499,5.94918925767,1.,0. +19,0.253239387885,0.345158023497,3.11071425333,4.79311566935,5.9489259713,1.,0. +20,0.637408390225,0.698996675371,3.25232492145,4.73814732384,5.9612010251,1.,0. +21,-0.407396859412,1.17456342803,2.49526823723,4.59323415742,5.82501686811,1.,0. +22,-0.967485452118,1.66655933642,2.47284606244,4.58316034754,5.88721406681,1.,0. +23,0.474480867904,1.95018556323,2.0228950072,4.48651142819,5.8255943735,1.,0. +24,1.04309652155,2.23519892356,1.91924131572,4.19094661783,5.87457348436,1.,0. +25,-0.517861513772,2.12501967336,1.70266619979,4.05280882887,5.72160912899,1.,0. +26,-0.945301585146,1.65464653549,1.81567174251,3.92309850635,5.58270493814,1.,0. +27,0.501153868974,1.40600764889,1.53991387719,3.72853247942,5.60169001727,1.,0. +28,0.972859524418,1.00344321868,1.5175642828,3.64092376655,5.10567722582,1.,0. +29,-0.70553406135,0.465306263885,1.7038540803,3.33236870312,5.09182481555,1.,0. +30,-0.946093634916,0.294539309453,1.88052827037,2.93011492669,4.97354922696,1.,0. +31,0.47922123231,0.308465865031,2.03445883031,2.90772899045,4.86241793548,1.,0. +32,0.754030014252,0.549752241167,2.46115815089,2.95063349534,4.71834614627,1.,0. +33,-0.64875949826,0.894615488148,2.5922463381,2.81269864022,4.43480095104,1.,0. +34,-0.757829951086,1.39123914261,2.69258079904,2.61834837315,4.36580046156,1.,0. +35,0.565653301088,1.72360022693,2.97794913834,2.80403840334,4.27327248459,1.,0. +36,0.867440092372,2.21100730052,3.38648090792,2.84057515729,4.12210169576,1.,0. +37,-0.894567758095,2.17549105818,3.45532493329,2.90446025717,4.00251740584,1.,0. +38,-0.715442356893,2.15105389965,3.52041791902,3.03650393392,4.12809249577,1.,0. +39,0.80671703672,1.81504564517,3.60463324866,3.00747789871,3.98440762467,1.,0. +40,0.527014790142,1.31803513865,3.43842186337,3.3332594663,4.03232406566,1.,0. +41,-0.795936862129,0.847809114454,3.09875133548,3.52863155938,3.94883924909,1.,0. +42,-0.610245806946,0.425530441018,2.92581949152,3.77238736123,4.27287245021,1.,0. +43,0.611662279431,0.178432049837,2.48128214822,3.73212087883,4.17319013831,1.,0. +44,0.650866553108,0.220341648392,2.41694642022,4.2609098519,4.27271645905,1.,0. +45,-0.774156982023,0.632667602331,2.05474356052,4.32889204886,4.18029723271,1.,0. +46,-0.714058448409,0.924562377599,1.75706135146,4.52492718422,4.3972678094,1.,0. +47,0.889627293379,1.46207968841,1.78299357672,4.64466731095,4.56317887554,1.,0. +48,0.520140662861,1.8996333843,1.41377633823,4.48899091177,4.78805049769,1.,0. +49,-1.03816935616,2.08997002059,1.51218375351,4.84167764204,4.93026048606,1.,0. +50,-0.40772951362,2.30878972136,1.44144415128,4.76854460997,5.01538444629,1.,0. +51,0.792730684781,1.91367048509,1.58887384677,4.71739397335,5.25690012199,1.,0. +52,0.371311881576,1.67565079528,1.81688563053,4.60353107555,5.44265822961,1.,0. +53,-0.814398070371,1.13374634126,1.80328814859,4.72264252878,5.52674761122,1.,0. +54,-0.469017949323,0.601244136627,2.29690896736,4.49859178859,5.54126153454,1.,0. +55,0.871044371426,0.407597593794,2.7499112487,4.19060637761,5.57693767301,1.,0. +56,0.523764933017,0.247705192709,3.09002071379,4.02095509006,5.80510362182,1.,0. +57,-0.881326403531,0.31513103164,3.11358205718,3.96079100808,5.81000652365,1.,0. +58,-0.357928025339,0.486163915865,3.17884556771,3.72634990659,5.85693642011,1.,0. +59,0.853038779822,1.04218094475,3.45835384454,3.36703969978,5.9585988449,1.,0. +60,0.435311516013,1.59715085283,3.63313338588,3.11276729421,5.93643818229,1.,0. +61,-1.02703719138,1.92205832542,3.47606111735,3.06247155999,6.02106646259,1.,0. +62,-0.246661325557,2.14653802542,3.29446326567,2.89936259181,5.67531541272,1.,0. +63,1.02554736569,2.25943737733,3.07031591528,2.78176218013,5.78206328989,1.,0. +64,0.337814475969,2.07589147224,2.80356226089,2.55888206331,5.7094075496,1.,0. +65,-1.12023369929,1.25333011618,2.56497288445,2.77361359194,5.50799418376,1.,0. +66,-0.178980246554,1.11937139901,2.51598681313,2.91438309151,5.47469577206,1.,0. +67,0.97550951531,0.60553823137,2.11657741073,2.88081098981,5.37034999502,1.,0. +68,0.136653357206,0.365828836075,1.97386033165,3.13217903204,5.07254490219,1.,0. +69,-1.05607596951,0.153152115069,1.52110743825,3.01308794192,5.08902539125,1.,0. +70,-0.13095280331,0.337113974483,1.52703079853,3.16687131599,4.86649398514,1.,0. +71,1.07081057754,0.714247566736,1.53761382634,3.45151989484,4.75892309166,1.,0. +72,0.0153410376082,1.24631231847,1.61690939161,3.85481994498,4.35683752832,1.,0. +73,-0.912801257303,1.60791309476,1.8729264524,4.03037260012,4.36072588913,1.,0. +74,-0.0894895640338,2.02535207407,1.93484909619,4.09557485132,4.35327025188,1.,0. +75,0.978646999652,2.20085086625,2.09003440427,4.27542353033,4.1805058388,1.,0. +76,-0.113312642876,2.2444100761,2.50789248839,4.4151861502,4.03267168136,1.,0. +77,-1.00215099149,1.84305628445,2.61691237246,4.45425147595,3.81203553766,1.,0. +78,-0.0183234614205,1.49573923116,2.99308471214,4.71134960112,4.0273804959,1.,0. +79,1.0823738177,1.12211589848,3.27079386925,4.94288270502,4.01851068083,1.,0. +80,0.124370187893,0.616474412808,3.4284236674,4.76942168327,3.9749536483,1.,0. +81,-0.929423379352,0.290977090976,3.34131726136,4.78590392707,4.10190661656,1.,0. +82,0.23766302648,0.155302052254,3.49779513794,4.64605656795,4.15571321107,1.,0. +83,1.03531486192,0.359702776204,3.4880725919,4.48167586667,4.21134561991,1.,0. +84,-0.261234571382,0.713877760378,3.42756426614,4.426443869,4.25208300527,1.,0. +85,-1.03572442277,1.25001113691,2.96908341113,4.25500915322,4.25723010649,1.,0. +86,0.380034261243,1.70543355622,2.73605932518,4.16703432307,4.63700400788,1.,0. +87,1.03734873488,1.97544410562,2.55586572141,3.84976673263,4.55282864289,1.,0. +88,-0.177344253372,2.22614526325,2.09565864891,3.77378097953,4.82577400298,1.,0. +89,-0.976821526892,2.18385079177,1.78522284118,3.67768223554,5.06302440873,1.,0. +90,0.264820472091,1.86981946157,1.50048403865,3.43619796921,5.05651761669,1.,0. +91,1.05642344868,1.47568646076,1.51347671977,3.20898518885,5.50149047462,1.,0. +92,-0.311607433358,1.04226467636,1.52089650905,3.02291865417,5.4889046232,1.,0. +93,-0.724285777937,0.553052311957,1.48573560173,2.7365973598,5.72549174225,1.,0. +94,0.519859192905,0.226520626591,1.61543723167,2.84102086852,5.69330622288,1.,0. +95,1.0323195039,0.260873217055,1.81913034804,2.83951143848,5.90325028086,1.,0. +96,-0.53285682538,0.387695521405,1.70935609313,2.57977050631,5.79579213161,1.,0. +97,-0.975127997215,0.920948771589,2.51292643636,2.71004616612,5.87016469227,1.,0. +98,0.540246804099,1.36445470181,2.61949412896,2.98482553485,6.02447664937,1.,0. +99,0.987764008058,1.85581989607,2.84685706149,2.94760204892,6.0212151724,1.,0. diff --git a/tensorflow/contrib/timeseries/examples/lstm.py b/tensorflow/contrib/timeseries/examples/lstm.py index c7193cef6915f9d0caf5b52fc084129cbc736994..f37cafcc502dc9415db0829b9b067b862f87dca7 100644 --- a/tensorflow/contrib/timeseries/examples/lstm.py +++ b/tensorflow/contrib/timeseries/examples/lstm.py @@ -18,13 +18,16 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools from os import path +import tempfile import numpy import tensorflow as tf from tensorflow.contrib.timeseries.python.timeseries import estimators as ts_estimators from tensorflow.contrib.timeseries.python.timeseries import model as ts_model +from tensorflow.contrib.timeseries.python.timeseries import state_management try: import matplotlib # pylint: disable=g-import-not-at-top @@ -45,7 +48,8 @@ _DATA_FILE = path.join(_MODULE_PATH, "data/multivariate_periods.csv") class _LSTMModel(ts_model.SequentialTimeSeriesModel): """A time series model-building example using an RNNCell.""" - def __init__(self, num_units, num_features, dtype=tf.float32): + def __init__(self, num_units, num_features, exogenous_feature_columns=None, + dtype=tf.float32): """Initialize/configure the model object. Note that we do not start graph building here. Rather, this object is a @@ -55,6 +59,10 @@ class _LSTMModel(ts_model.SequentialTimeSeriesModel): num_units: The number of units in the model's LSTMCell. num_features: The dimensionality of the time series (features per timestep). + exogenous_feature_columns: A list of tf.contrib.layers.FeatureColumn + objects representing features which are inputs to the model but are + not predicted by it. These must then be present for training, + evaluation, and prediction. dtype: The floating point data type to use. """ super(_LSTMModel, self).__init__( @@ -62,6 +70,7 @@ class _LSTMModel(ts_model.SequentialTimeSeriesModel): train_output_names=["mean"], predict_output_names=["mean"], num_features=num_features, + exogenous_feature_columns=exogenous_feature_columns, dtype=dtype) self._num_units = num_units # Filled in by initialize_graph() @@ -69,7 +78,7 @@ class _LSTMModel(ts_model.SequentialTimeSeriesModel): self._lstm_cell_run = None self._predict_from_lstm_output = None - def initialize_graph(self, input_statistics): + def initialize_graph(self, input_statistics=None): """Save templates for components, which can then be used repeatedly. This method is called every time a new graph is created. It's safe to start @@ -80,18 +89,19 @@ class _LSTMModel(ts_model.SequentialTimeSeriesModel): input_statistics: A math_utils.InputStatistics object. """ super(_LSTMModel, self).initialize_graph(input_statistics=input_statistics) - self._lstm_cell = tf.nn.rnn_cell.LSTMCell(num_units=self._num_units) - # Create templates so we don't have to worry about variable reuse. - self._lstm_cell_run = tf.make_template( - name_="lstm_cell", - func_=self._lstm_cell, - create_scope_now_=True) - # Transforms LSTM output into mean predictions. - self._predict_from_lstm_output = tf.make_template( - name_="predict_from_lstm_output", - func_= - lambda inputs: tf.layers.dense(inputs=inputs, units=self.num_features), - create_scope_now_=True) + with tf.variable_scope("", use_resource=True): + # Use ResourceVariables to avoid race conditions. + self._lstm_cell = tf.nn.rnn_cell.LSTMCell(num_units=self._num_units) + # Create templates so we don't have to worry about variable reuse. + self._lstm_cell_run = tf.make_template( + name_="lstm_cell", + func_=self._lstm_cell, + create_scope_now_=True) + # Transforms LSTM output into mean predictions. + self._predict_from_lstm_output = tf.make_template( + name_="predict_from_lstm_output", + func_=functools.partial(tf.layers.dense, units=self.num_features), + create_scope_now_=True) def get_start_state(self): """Return initial state for the time series model.""" @@ -100,6 +110,8 @@ class _LSTMModel(ts_model.SequentialTimeSeriesModel): tf.zeros([], dtype=tf.int64), # The previous observation or prediction. tf.zeros([self.num_features], dtype=self.dtype), + # The most recently seen exogenous features. + tf.zeros(self._get_exogenous_embedding_shape(), dtype=self.dtype), # The state of the RNNCell (batch dimension removed since this parent # class will broadcast). [tf.squeeze(state_element, axis=0) @@ -127,7 +139,7 @@ class _LSTMModel(ts_model.SequentialTimeSeriesModel): loss (note that we could also return other measures of goodness of fit, although only "loss" will be optimized). """ - state_from_time, prediction, lstm_state = state + state_from_time, prediction, exogenous, lstm_state = state with tf.control_dependencies( [tf.assert_equal(current_times, state_from_time)]): # Subtract the mean and divide by the variance of the series. Slightly @@ -139,16 +151,22 @@ class _LSTMModel(ts_model.SequentialTimeSeriesModel): (prediction - transformed_values) ** 2, axis=-1) # Keep track of the new observation in model state. It won't be run # through the LSTM until the next _imputation_step. - new_state_tuple = (current_times, transformed_values, lstm_state) + new_state_tuple = (current_times, transformed_values, + exogenous, lstm_state) return (new_state_tuple, predictions) def _prediction_step(self, current_times, state): """Advance the RNN state using a previous observation or prediction.""" - _, previous_observation_or_prediction, lstm_state = state + _, previous_observation_or_prediction, exogenous, lstm_state = state + # Update LSTM state based on the most recent exogenous and endogenous + # features. + inputs = tf.concat([previous_observation_or_prediction, exogenous], + axis=-1) lstm_output, new_lstm_state = self._lstm_cell_run( - inputs=previous_observation_or_prediction, state=lstm_state) + inputs=inputs, state=lstm_state) next_prediction = self._predict_from_lstm_output(lstm_output) - new_state_tuple = (current_times, next_prediction, new_lstm_state) + new_state_tuple = (current_times, next_prediction, + exogenous, new_lstm_state) return new_state_tuple, {"mean": self._scale_back_data(next_prediction)} def _imputation_step(self, current_times, state): @@ -160,36 +178,75 @@ class _LSTMModel(ts_model.SequentialTimeSeriesModel): def _exogenous_input_step( self, current_times, current_exogenous_regressors, state): - """Update model state based on exogenous regressors.""" - raise NotImplementedError( - "Exogenous inputs are not implemented for this example.") + """Save exogenous regressors in model state for use in _prediction_step.""" + state_from_time, prediction, _, lstm_state = state + return (state_from_time, prediction, + current_exogenous_regressors, lstm_state) def train_and_predict( - csv_file_name=_DATA_FILE, training_steps=200, estimator_config=None): + csv_file_name=_DATA_FILE, training_steps=200, estimator_config=None, + export_directory=None): """Train and predict using a custom time series model.""" # Construct an Estimator from our LSTM model. + exogenous_feature_columns = [ + # Exogenous features are not part of the loss, but can inform + # predictions. In this example the features have no extra information, but + # are included as an API example. + tf.contrib.layers.real_valued_column( + "2d_exogenous_feature", dimension=2)] estimator = ts_estimators.TimeSeriesRegressor( - model=_LSTMModel(num_features=5, num_units=128), - optimizer=tf.train.AdamOptimizer(0.001), config=estimator_config) + model=_LSTMModel(num_features=5, num_units=128, + exogenous_feature_columns=exogenous_feature_columns), + optimizer=tf.train.AdamOptimizer(0.001), config=estimator_config, + # Set state to be saved across windows. + state_manager=state_management.ChainingStateManager()) reader = tf.contrib.timeseries.CSVReader( csv_file_name, column_names=((tf.contrib.timeseries.TrainEvalFeatures.TIMES,) - + (tf.contrib.timeseries.TrainEvalFeatures.VALUES,) * 5)) + + (tf.contrib.timeseries.TrainEvalFeatures.VALUES,) * 5 + + ("2d_exogenous_feature",) * 2)) train_input_fn = tf.contrib.timeseries.RandomWindowInputFn( reader, batch_size=4, window_size=32) estimator.train(input_fn=train_input_fn, steps=training_steps) evaluation_input_fn = tf.contrib.timeseries.WholeDatasetInputFn(reader) evaluation = estimator.evaluate(input_fn=evaluation_input_fn, steps=1) # Predict starting after the evaluation + predict_exogenous_features = { + "2d_exogenous_feature": numpy.concatenate( + [numpy.ones([1, 100, 1]), numpy.zeros([1, 100, 1])], + axis=-1)} (predictions,) = tuple(estimator.predict( input_fn=tf.contrib.timeseries.predict_continuation_input_fn( - evaluation, steps=100))) + evaluation, steps=100, + exogenous_features=predict_exogenous_features))) times = evaluation["times"][0] observed = evaluation["observed"][0, :, :] predicted_mean = numpy.squeeze(numpy.concatenate( [evaluation["mean"][0], predictions["mean"]], axis=0)) all_times = numpy.concatenate([times, predictions["times"]], axis=0) + + # Export the model in SavedModel format. + if export_directory is None: + export_directory = tempfile.mkdtemp() + input_receiver_fn = estimator.build_raw_serving_input_receiver_fn() + export_location = estimator.export_savedmodel( + export_directory, input_receiver_fn) + # Predict using the SavedModel + with tf.Graph().as_default(): + with tf.Session() as session: + signatures = tf.saved_model.loader.load( + session, [tf.saved_model.tag_constants.SERVING], export_location) + saved_model_output = ( + tf.contrib.timeseries.saved_model_utils.predict_continuation( + continue_from=evaluation, signatures=signatures, + session=session, steps=100, + exogenous_features=predict_exogenous_features)) + # The exported model gives the same results as the Estimator.predict() + # call above. + numpy.testing.assert_allclose( + predictions["mean"], + numpy.squeeze(saved_model_output["mean"], axis=0)) return times, observed, all_times, predicted_mean diff --git a/tensorflow/contrib/timeseries/examples/lstm_test.py b/tensorflow/contrib/timeseries/examples/lstm_test.py index 3cace567266d497b12d836f44a335bbe5d916949..ca56e38ca079f71b38cf29605a295a50929945e8 100644 --- a/tensorflow/contrib/timeseries/examples/lstm_test.py +++ b/tensorflow/contrib/timeseries/examples/lstm_test.py @@ -36,7 +36,8 @@ class LSTMExampleTest(test.TestCase): def test_periodicity_learned(self): (observed_times, observed_values, all_times, predicted_values) = lstm.train_and_predict( - training_steps=100, estimator_config=_SeedRunConfig()) + training_steps=100, estimator_config=_SeedRunConfig(), + export_directory=self.get_temp_dir()) self.assertAllEqual([100], observed_times.shape) self.assertAllEqual([100, 5], observed_values.shape) self.assertAllEqual([200], all_times.shape) diff --git a/tensorflow/contrib/timeseries/python/timeseries/BUILD b/tensorflow/contrib/timeseries/python/timeseries/BUILD index 5f04eb2f5a4af031ad19662b05a8a2396299925d..fff972c1f3277ad5d83673a202a50d1e6f7df210 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/BUILD +++ b/tensorflow/contrib/timeseries/python/timeseries/BUILD @@ -296,6 +296,8 @@ py_test( ], srcs_version = "PY2AND3", tags = [ + "no_oss", # b/63709811 + "no_pip", # b/63709811 "no_pip_gpu", # b/63391119 ], deps = [ diff --git a/tensorflow/contrib/timeseries/python/timeseries/estimators.py b/tensorflow/contrib/timeseries/python/timeseries/estimators.py index 3738dfa154d4f39b9562446972443ed88f3fbe8b..f8355f366fe8e191ab570fd271bbe4a8bf71c73d 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/estimators.py +++ b/tensorflow/contrib/timeseries/python/timeseries/estimators.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib.layers.python.layers import feature_column + from tensorflow.contrib.timeseries.python.timeseries import ar_model from tensorflow.contrib.timeseries.python.timeseries import feature_keys from tensorflow.contrib.timeseries.python.timeseries import head as ts_head_lib @@ -72,15 +74,14 @@ class TimeSeriesRegressor(estimator_lib.Estimator): # tf.Example containing all features (times, values, any exogenous features) # and serialized model state (possibly also as a tf.Example). def build_raw_serving_input_receiver_fn(self, - exogenous_features=None, default_batch_size=None, default_series_length=None): """Build an input_receiver_fn for export_savedmodel which accepts arrays. + Automatically creates placeholders for exogenous `FeatureColumn`s passed to + the model. + Args: - exogenous_features: A dictionary mapping feature keys to exogenous - features (either Numpy arrays or Tensors). Used to determine the shapes - of placeholders for these features. default_batch_size: If specified, must be a scalar integer. Sets the batch size in the static shape information of all feature Tensors, which means only this batch size will be accepted by the exported model. If None @@ -94,9 +95,6 @@ class TimeSeriesRegressor(estimator_lib.Estimator): An input_receiver_fn which may be passed to the Estimator's export_savedmodel. """ - if exogenous_features is None: - exogenous_features = {} - def _serving_input_receiver_fn(): """A receiver function to be passed to export_savedmodel.""" placeholders = {} @@ -119,14 +117,22 @@ class TimeSeriesRegressor(estimator_lib.Estimator): dtype=self._model.dtype), shape=(default_batch_size, default_series_length, self._model.num_features))) - for feature_key, feature_value in exogenous_features.items(): - value_tensor = ops.convert_to_tensor(feature_value) - value_tensor.get_shape().with_rank_at_least(2) - feature_shape = value_tensor.get_shape().as_list() - feature_shape[0] = default_batch_size - feature_shape[1] = default_series_length + with ops.Graph().as_default(): + # Default placeholders have only an unknown batch dimension. Make them + # in a separate graph, then splice in the series length to the shapes + # and re-create them in the outer graph. + exogenous_feature_shapes = { + key: (value.get_shape(), value.dtype) for key, value + in feature_column.make_place_holder_tensors_for_base_features( + self._model.exogenous_feature_columns).items()} + for feature_key, (batch_only_feature_shape, value_dtype) in ( + exogenous_feature_shapes.items()): + batch_only_feature_shape = batch_only_feature_shape.with_rank_at_least( + 1).as_list() + feature_shape = ([default_batch_size, default_series_length] + + batch_only_feature_shape[1:]) placeholders[feature_key] = array_ops.placeholder( - dtype=value_tensor.dtype, name=feature_key, shape=feature_shape) + dtype=value_dtype, name=feature_key, shape=feature_shape) # Models may not know the shape of their state without creating some # variables/ops. Avoid polluting the default graph by making a new one. We # use only static metadata from the returned Tensors. diff --git a/tensorflow/contrib/timeseries/python/timeseries/model.py b/tensorflow/contrib/timeseries/python/timeseries/model.py index b32b5c5494ae14187954b900119678a5b53a3602..bac7d1ebf59b28d4688a3d1a69ecdc1fc12248e0 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/model.py +++ b/tensorflow/contrib/timeseries/python/timeseries/model.py @@ -22,6 +22,7 @@ import abc import collections from tensorflow.contrib import layers +from tensorflow.contrib.layers import feature_column from tensorflow.contrib.timeseries.python.timeseries import math_utils from tensorflow.contrib.timeseries.python.timeseries.feature_keys import PredictionFeatures @@ -83,6 +84,11 @@ class TimeSeriesModel(object): self._stats_means = None self._stats_sigmas = None + @property + def exogenous_feature_columns(self): + """`FeatureColumn` objects for features which are not predicted.""" + return self._exogenous_feature_columns + # TODO(allenl): Move more of the generic machinery for generating and # predicting into TimeSeriesModel, and possibly share it between generate() # and predict() @@ -250,6 +256,23 @@ class TimeSeriesModel(object): """ pass + def _get_exogenous_embedding_shape(self): + """Computes the shape of the vector returned by _process_exogenous_features. + + Returns: + The shape as a list. Does not include a batch dimension. + """ + if not self._exogenous_feature_columns: + return (0,) + with ops.Graph().as_default(): + placeholder_features = ( + feature_column.make_place_holder_tensors_for_base_features( + self._exogenous_feature_columns)) + embedded = layers.input_from_feature_columns( + columns_to_tensors=placeholder_features, + feature_columns=self._exogenous_feature_columns) + return embedded.get_shape().as_list()[1:] + def _process_exogenous_features(self, times, features): """Create a single vector from exogenous features. diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model_test.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model_test.py index ca57715e2b2e6bbadd276d641703c0a3b842652e..1fb4a3c121c8d7c1daf8fc4a3f59a8b8de38bf8f 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model_test.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model_test.py @@ -187,9 +187,7 @@ class StateSpaceEquivalenceTests(test.TestCase): estimator.train(combined_input_fn, steps=1) export_location = estimator.export_savedmodel( self.get_temp_dir(), - estimator.build_raw_serving_input_receiver_fn( - exogenous_features={ - "exogenous": numpy.zeros((0, 0), dtype=numpy.float32)})) + estimator.build_raw_serving_input_receiver_fn()) with ops.Graph().as_default() as graph: random_model.initialize_graph() with self.test_session(graph=graph) as session: @@ -209,7 +207,7 @@ class StateSpaceEquivalenceTests(test.TestCase): features={ feature_keys.FilteringFeatures.TIMES: [1, 2], feature_keys.FilteringFeatures.VALUES: [1., 2.], - "exogenous": [-1., -2.]}) + "exogenous": [[-1.], [-2.]]}) second_split_filtering = saved_model_utils.filter_continuation( continue_from=first_split_filtering, signatures=signatures, @@ -217,7 +215,7 @@ class StateSpaceEquivalenceTests(test.TestCase): features={ feature_keys.FilteringFeatures.TIMES: [3, 4], feature_keys.FilteringFeatures.VALUES: [3., 4.], - "exogenous": [-3., -4.] + "exogenous": [[-3.], [-4.]] }) combined_filtering = saved_model_utils.filter_continuation( continue_from={ @@ -227,7 +225,7 @@ class StateSpaceEquivalenceTests(test.TestCase): features={ feature_keys.FilteringFeatures.TIMES: [1, 2, 3, 4], feature_keys.FilteringFeatures.VALUES: [1., 2., 3., 4.], - "exogenous": [-1., -2., -3., -4.] + "exogenous": [[-1.], [-2.], [-3.], [-4.]] }) split_predict = saved_model_utils.predict_continuation( continue_from=second_split_filtering, @@ -235,14 +233,14 @@ class StateSpaceEquivalenceTests(test.TestCase): session=session, steps=1, exogenous_features={ - "exogenous": [[-5.]]}) + "exogenous": [[[-5.]]]}) combined_predict = saved_model_utils.predict_continuation( continue_from=combined_filtering, signatures=signatures, session=session, steps=1, exogenous_features={ - "exogenous": [[-5.]]}) + "exogenous": [[[-5.]]]}) for state_key, combined_state_value in combined_filtering.items(): if state_key == feature_keys.FilteringResults.TIMES: continue @@ -627,9 +625,8 @@ class UnknownShapeModel(TimeDependentStateSpaceModel): def get_observation_model(self, times): parent_model = super(UnknownShapeModel, self).get_observation_model(times) - parent_model._shape = tensor_shape.unknown_shape() - assert parent_model.get_shape().ndims is None - return parent_model + return array_ops.placeholder_with_default( + input=parent_model, shape=tensor_shape.unknown_shape()) class TimeDependentTests(test.TestCase): diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD index 43cb38f169c2bf2cd8b271cb332ee2cacab7df06..c48e84ddfaac8ac9c07e061847315eab3fd72152 100644 --- a/tensorflow/contrib/tpu/BUILD +++ b/tensorflow/contrib/tpu/BUILD @@ -28,6 +28,7 @@ cc_library( ":outfeed_ops_op_lib", ":replication_ops_op_lib", ":tpu_configuration_ops_op_lib", + ":tpu_embedding_ops_op_lib", ], ) @@ -35,13 +36,16 @@ py_library( name = "tpu_estimator", srcs = [ "python/tpu/tpu_config.py", + "python/tpu/tpu_context.py", "python/tpu/tpu_estimator.py", + "python/tpu/tpu_system_metadata.py", "python/tpu/util.py", ], srcs_version = "PY2AND3", deps = [ ":tpu_lib", ":tpu_py", + "//tensorflow/contrib/summary:summary_ops", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", @@ -69,9 +73,11 @@ tf_gen_op_libs( "outfeed_ops", "replication_ops", "tpu_configuration_ops", + "tpu_embedding_ops", ], deps = [ - "//tensorflow/core:lib", + "//tensorflow/contrib/tpu/proto:tpu_embedding_config_proto_cc", + "//tensorflow/core:lib_proto_parsing", ], ) @@ -83,6 +89,11 @@ tf_custom_op_library( "ops/outfeed_ops.cc", "ops/replication_ops.cc", "ops/tpu_configuration_ops.cc", + "ops/tpu_embedding_ops.cc", + ], + deps = [ + "//tensorflow/contrib/tpu/proto:tpu_embedding_config_proto_cc", + "//tensorflow/core:lib_proto_parsing", ], ) @@ -94,6 +105,7 @@ tf_gen_op_wrapper_py( ":outfeed_ops_op_lib", ":replication_ops_op_lib", ":tpu_configuration_ops_op_lib", + ":tpu_embedding_ops_op_lib", ], ) @@ -157,6 +169,7 @@ py_library( "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:control_flow_util", "//tensorflow/python:dtypes", "//tensorflow/python:framework", "//tensorflow/python:framework_ops", @@ -168,6 +181,19 @@ py_library( ], ) +tf_py_test( + name = "tpu_test", + size = "small", + srcs = ["python/tpu/tpu_test.py"], + additional_deps = [ + ":tpu", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework", + "//tensorflow/python:layers", + ], +) + tf_py_test( name = "tpu_sharding_test", size = "small", diff --git a/tensorflow/contrib/tpu/ops/infeed_ops.cc b/tensorflow/contrib/tpu/ops/infeed_ops.cc index 849c4a1102787870b372c35740cf0fe271efa078..efc546f9a6077de9cac5a5acefa3fc7206547fc6 100644 --- a/tensorflow/contrib/tpu/ops/infeed_ops.cc +++ b/tensorflow/contrib/tpu/ops/infeed_ops.cc @@ -41,6 +41,7 @@ REGISTER_OP("InfeedEnqueue") .Attr("dtype: type") .Attr("shape: shape = {}") .Attr("device_ordinal: int = -1") + .SetShapeFn(shape_inference::NoOutputs) .SetIsStateful() .Doc(R"doc( An op which feeds a single Tensor value into the computation. @@ -58,6 +59,7 @@ REGISTER_OP("InfeedEnqueueTuple") .Attr("dtypes: list(type)") .Attr("shapes: list(shape)") .Attr("device_ordinal: int = -1") + .SetShapeFn(shape_inference::NoOutputs) .SetIsStateful() .Doc(R"doc( An op which feeds multiple Tensor values into the computation as an XLA tuple. diff --git a/tensorflow/contrib/tpu/ops/tpu_configuration_ops.cc b/tensorflow/contrib/tpu/ops/tpu_configuration_ops.cc index 8c4fe5538d832f390845fe2d31aa6a08342b280b..f8de8baa65339383c7f92284ee274a434f12f8c2 100644 --- a/tensorflow/contrib/tpu/ops/tpu_configuration_ops.cc +++ b/tensorflow/contrib/tpu/ops/tpu_configuration_ops.cc @@ -26,29 +26,25 @@ using shape_inference::ShapeHandle; // Configuring a distributed TPU system is achieved by running // the following Ops: // -// 1 Run _DisconnectHostFromDistributedTPUSystem on the CPU of each -// host. This is needed in case the system had previously been -// configured. It returns, for each host, the number of TPU chips on -// the host. +// 1 Run _DisconnectHostFromDistributedTPUSystem on the TPU_SYSTEM of each +// host. This is needed in case the system had previously been configured. It +// returns, for each host, the number of TPU chips on the host. // -// 2 Run _ConfigureDistributedTPU on TPU_SYSTEM. Takes as input the -// number of chips on each host. Validates that all hosts have the -// same number of chips, and that the chips are consistent with the -// topology set by flags. Has a single output which is a proto -// describing the requested system configuration, which is sent to all -// hosts. +// 2 Run _ConfigureDistributedTPU on TPU_SYSTEM of worker 0. Takes as input the +// number of chips on each host. Validates that all hosts have the same number +// of chips, and that the chips are consistent with the topology set by +// flags. Has a single output which is a proto describing the requested system +// configuration, which is sent to all hosts. // -// 3 Run _InitializeHostForDistributedTPU on the CPU of each host, -// taking as input the output from ConfigureDistributedTPU. Has a -// single Tensor output which is a vector of int32 indicating, for -// each TPU on the host, what its global TPU system id is. +// 3 Run _InitializeHostForDistributedTPU on the TPU_SYSTEM of each host, taking +// as input the output from ConfigureDistributedTPU. Has a single Tensor output +// which is a vector of int32 indicating, for each TPU on the host, what its +// global TPU system id is. // // 4 Run _WaitForDistributedTPU on TPU_SYSTEM, taking as input the // outputs from all the _InitializeHostForDistributedTPU -// Ops. _WaitForDistributedTPU has an attr host_specs which is a -// vector giving the partial device spec for each host. These -// partial specs are combined in the Op with the outputs from the host -// initialization Ops to construct a mapping from full TPU device +// Ops. _These partial specs are combined in the Op with the outputs from +// the host initialization Ops to construct a mapping from full TPU device // specs to global TPU ids. Has a single Tensor output which is a // matrix of int32 indicating, for each host (outer dimension) and for // each TPU on the host (inner dimension) what that TPU's global id @@ -56,29 +52,28 @@ using shape_inference::ShapeHandle; // system to initialize fully, which may take several minutes for a // large system. // -// 5 Run _SetGlobalTPUArray on the CPU of each host, taking as input -// the output from _WaitForDistributedTPU. This Op tells each host the -// global Id of every TPU on every host. +// 5 Run _SetGlobalTPUArray on the TPU_SYSTEM of each host, taking as input the +// output from _WaitForDistributedTPU. This Op tells each host the global Id of +// every TPU on every host. // -// Most user code works by placing the ConfigureDistributedTPU Op on -// the desired TPU_SYSTEM device, and a graph rewrite replaces it by -// the subgraph described above. +// Most user code works by placing the ConfigureDistributedTPU Op on the desired +// TPU_SYSTEM device, and a graph rewrite replaces it by the subgraph described +// above. // // -// A distributed TPU system can be cleanly shut down by running -// the following Ops: +// A distributed TPU system can be cleanly shut down by running the following +// Ops: // -// 1 Run _DisconnectHostFromDistributedTPUSystem on the CPU of each -// host. +// 1 Run _DisconnectHostFromDistributedTPUSystem on the TPU_SYSTEM of each host. // // 2 Run _ShutdownDistributedTPU on the TPU_SYSTEM where -// _ConfigureDistributedTPU was run. The Op will return an error if no -// system is configured. +// _ConfigureDistributedTPU was run. The Op will return an error if no system is +// configured. // // -// Most user code works by placing the ShutdownDistributedTPU Op on -// the desired TPU_SYSTEM device, and a graph rewrite replaces it by -// the subgraph described above. +// Most user code works by placing the ShutdownDistributedTPU Op on the desired +// TPU_SYSTEM device, and a graph rewrite replaces it by the subgraph described +// above. REGISTER_OP("_ConfigureDistributedTPU") .Input("inputs: N * int32") @@ -108,7 +103,6 @@ in a host. REGISTER_OP("_WaitForDistributedTPU") .Input("inputs: N * int32") .Output("topology: string") - .Attr("host_specs: list(string)") .Attr("startup_timeout_sec: int = 20") .Attr("N: int") .SetIsStateful() @@ -196,6 +190,7 @@ chips on the host. REGISTER_OP("ConfigureDistributedTPU") .Output("topology: string") .Attr("embedding_config: string = ''") + .Attr("tpu_embedding_config: string = ''") .SetIsStateful() .SetShapeFn(shape_inference::UnknownShape) .Doc(R"doc( @@ -204,6 +199,9 @@ system. topology: A serialized tensorflow.tpu.TopologyProto that describes the TPU topology. +tpu_embedding_config: Serialized tensorflow.tpu.TPUEmbeddingConfiguration that +describes the embedding lookups of the program. +embedding_config: Reserved. Do not use. )doc"); REGISTER_OP("ShutdownDistributedTPU") @@ -214,4 +212,20 @@ An op that shuts down a running distributed TPU system. The Op returns an error if no system is running. )doc"); -} // namespace tensorflow +REGISTER_OP("SessionStatus") + .Input("fetch_start_timestamp: double") + .Output("status: string") + .SetShapeFn(shape_inference::ScalarShape) + .Doc(R"doc( +Not for public usage. + +Returns messages from the current session as a serialized SessionStatusProto. + +This includes the current state of the compiler, along with any critical +logging or warning messages. + +fetch_start_timestamp: any messages earlier than this will be excluded from the +returned proto. +)doc"); + +} // end namespace tensorflow diff --git a/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc b/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..cc32a265286951a1e4d59228da6b3ac83a75c5e9 --- /dev/null +++ b/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc @@ -0,0 +1,328 @@ +/* Copyright 2017 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/tpu/proto/tpu_embedding_config.pb.h" +#include "tensorflow/core/framework/common_shape_fns.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { + +// TPUs use a specialized mechanism for performing embedding lookups, +// necessitating differences in TF Graphs that use embeddings on TPUs relative +// to CPUs. Embedding lookups on TPU systems are achieved by including the +// following in the TF Graph. +// +// 0. Construct a TPUEmbeddingConfiguration, specifying the embedding tables +// in the model, the size of the TPU system to be used, and the optimizer to +// be used for each table. Some of this information is redundant with other +// pieces of the TF Graph. +// 1. Pass this TPUEmbeddingConfiguration to tpu.initialize_system() as the +// tpu_embedding_config parameter. +// 2. Use the TPUEmbeddingLoad Op to initialize the embedding tables in TPU +// memories, sharded across the memories attached to each Host. +// 3. Use TPUEmbeddingEnqueueSparseBatch to provide the TPU with embedding +// indices and aggregation weights. +// 4. TPUEmbeddingReceiveActivations returns a list of Tensors, containing the +// activations from each table specified in the configuration. +// 5. TPUEmbeddingActivations, when used with appropriate Python libraries, +// enables the automatic differentiation of models that use embeddings. +// 6. TPUEmbeddingSendGradients takes a list of Tensors (of the same shapes +// as those returned by TPUEmbeddingReceivActivations) containing gradients +// to use in updating the embedding tables. +// 7. Before saving a checkpoint, use the TPUEmbeddingRetrieve Op to update +// the Graph's embedding table Variables from the updated tables in the +// TPU memories. +// +// TPU Embeddings use dedicated ops to enforce Host/TPU consistency in the +// state of embedding table variables. Before beginning training or inference, +// the model must Load the optimizer parameters into the TPU memories. Before +// saving a checkpoint, the model must Retreieve the parameters back into the +// host CPU memory. + +REGISTER_OP("TPUEmbeddingLoadGradientDescentParameters") + .Input("parameters: float32") + .Attr("tpu_embedding_config: string") + .Attr("table_id: int >= 0") + .Attr("num_hosts: int >= 1") + .Attr("host_id: int >= 0") + .SetIsStateful() + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +Load an embedding table shard into TPU memory for use with GradientDescent. + +TPU embeddings use dedicated per-optimizer Ops for loading and retrieving +trainable variables and optimizer state from TPU memory. This op enables +functionality equivalent to GradientDescentOptimizer. + +parameters: The shard of the embedding table resident on the host executing this + op. For single-TPU models, this is the entire embedding table. +tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto. +table_id: The id of the table specified in the tpu_embedding_config. +num_hosts: The number of CPU hosts in the distributed training job. +host_id: Which CPU host in the distributed training job will execute this op. +)doc"); + +namespace tpu_embedding_config_util { + +Status GradientDescentShapes(shape_inference::InferenceContext *c) { + string config_string; + TF_RETURN_IF_ERROR(c->GetAttr("tpu_embedding_config", &config_string)); + tpu::TPUEmbeddingConfiguration config; + if (!config.ParseFromString(config_string)) { + return errors::InvalidArgument("Malformed tpu_embedding_config."); + } + + int table_id; + TF_RETURN_IF_ERROR(c->GetAttr("table_id", &table_id)); + int64 num_tables = config.table_config_size(); + if (table_id >= num_tables) { + return errors::InvalidArgument("Table id >= num_tables"); + } + int64 width = config.table_config(table_id).width(); + int64 num_rows = config.table_config(table_id).num_rows(); + + TF_RETURN_IF_ERROR(c->set_output("parameters", {c->Matrix(num_rows, width)})); + return Status::OK(); +} + +} // namespace tpu_embedding_config_util + +REGISTER_OP("TPUEmbeddingRetrieveGradientDescentParameters") + .Output("parameters: float32") + .Attr("tpu_embedding_config: string") + .Attr("table_id: int") + .Attr("num_hosts: int") + .Attr("host_id: int") + .SetIsStateful() + .SetShapeFn(tpu_embedding_config_util::GradientDescentShapes) + .Doc(R"doc( +Retrieve an embedding table shard from TPU memory. + +TPU embeddings use dedicated per-optimizer Ops for loading and retrieving +trainable variables and optimizer state from TPU memory. This op enables +functionality equivalent to GradientDescentOptimizer. + +tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto. +table_id: The id of the table specified in tpu_embedding_config. +num_hosts: The number of CPU hosts in the distributed training job. +host_id: Which CPU host in the distributed training job will execute this op. +)doc"); + +REGISTER_OP("TPUEmbeddingLoadAdagradParameters") + .Input("parameters: float32") + .Input("accumulators: float32") + .Attr("tpu_embedding_config: string") + .Attr("table_id: int >= 0") + .Attr("num_hosts: int >= 1") + .Attr("host_id: int >= 0") + .SetIsStateful() + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +Load an embedding table shard into TensorNode memories for use with Adagrad. + +TPU embeddings use dedicated per-optimizer Ops for loading and retrieving +trainable variables and optimizer state from TPU memory. This op enables +functionality equivalent to AdagradOptimizer. + +parameters: The shard of the embedding table resident on the host executing this + op. For single-TPU models, this is the entire embedding table. +accumulators: Shard of the Adagrad accumulators resident on the host executing + this op. +tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto. +table_id: The id of the table specified in the embedding_config. +num_hosts: The number of CPU hosts in the distributed training job. +host_id: Which CPU host in the distributed training job will execute this op. +)doc"); + +namespace tpu_embedding_config_util { + +Status AdagradShapes(shape_inference::InferenceContext *c) { + string config_string; + TF_RETURN_IF_ERROR(c->GetAttr("tpu_embedding_config", &config_string)); + tpu::TPUEmbeddingConfiguration config; + if (!config.ParseFromString(config_string)) { + return errors::InvalidArgument("Malformed tpu_embedding_config."); + } + + int table_id; + TF_RETURN_IF_ERROR(c->GetAttr("table_id", &table_id)); + int64 num_tables = config.table_config_size(); + if (table_id >= num_tables) { + return errors::InvalidArgument("Table id >= num_tables"); + } + int64 width = config.table_config(table_id).width(); + int64 num_rows = config.table_config(table_id).num_rows(); + + TF_RETURN_IF_ERROR(c->set_output("parameters", {c->Matrix(num_rows, width)})); + TF_RETURN_IF_ERROR( + c->set_output("accumulators", {c->Matrix(num_rows, width)})); + return Status::OK(); +} + +} // namespace tpu_embedding_config_util + +REGISTER_OP("TPUEmbeddingRetrieveAdagradParameters") + .Output("parameters: float32") + .Output("accumulators: float32") + .Attr("tpu_embedding_config: string") + .Attr("table_id: int >= 0") + .Attr("num_hosts: int >= 1") + .Attr("host_id: int >= 0") + .SetIsStateful() + .SetShapeFn(tpu_embedding_config_util::AdagradShapes) + .Doc(R"doc( +Retrieve an embedding table shard from TPU memory. + +TPU embeddings use dedicated per-optimizer Ops for loading and retrieving +trainable variables and optimizer state from TPU memory. This op enables +functionality equivalent to AdagradOptimizer. + +tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto. +table_id: The id of the table specified in the embedding_config_json. +num_hosts: The number of CPU hosts in the distributed training job. +host_id: Which CPU host in the distributed training job will execute this op. +)doc"); + +REGISTER_OP("TPUEmbeddingEnqueueSparseBatch") + .Input("sample_indices: num_tables * int32") + .Input("embedding_indices: num_tables * int32") + .Input("aggregation_weights: num_tables * float32") + .Attr("num_tables: int") + .Attr("device_ordinal: int = -1") + .SetIsStateful() + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +An op that feeds a batch of embedding indices and weights to the TPU. + +Embedding lookups are equivalent to sparse-dense matrix multiplications: the +sparse matrix contains nonzeros in column j in order to retrieve row j from the +embedding table. + +The three Tensor list arguments (sample_indices, embedding_indices, and +aggregation_weights) represent these sparse matrices in COO format. The Tensor +lists each have one entry for each embedding table specified in the model. +For the kth embedding table, the three Tensors at position k in the list +specify a COO-format sparse matrix. For the kth table, the row indices, +column indices, and nonzero values of the COO sparse matrix are specified by +sample_indices[k], embedding_indices[k], and aggregation_weights[k], +respectively. Entries must be sorted by row index, then by column index. + +There should be at most one TPUEmbeddingEnqueueSparseBatch op in a signle +training step per TPU shard. + +sample_indices: A list of rank 1 Tensors specifying row indices of the COO + sparse matrix representing the embedding lookups for each table. +embedding_indices: A list of rank 1 Tensors specifying column indices of the + COO sparse matrix representing the embedding lookups for each table. +aggregation_weights: A list of rank 1 Tensors specifying the nonzero values + of the COO sparse matrix representing the embedding lookups for each table. +device_ordinal: The TPU device to use. This should be -1 when the Op + is running on a TPU device, and >= 0 when the Op is running on the CPU + device. +)doc"); + +namespace tpu_embedding_config_util { + +Status ActivationShapes(shape_inference::InferenceContext *c) { + string config_string; + TF_RETURN_IF_ERROR(c->GetAttr("tpu_embedding_config", &config_string)); + tpu::TPUEmbeddingConfiguration config; + if (!config.ParseFromString(config_string)) { + return errors::InvalidArgument("Malformed tpu_embedding_config."); + } + int64 batch_size = config.batch_size(); + int64 num_tables = config.table_config_size(); + for (int table_id = 0; table_id < num_tables; ++table_id) { + int64 width = config.table_config(table_id).width(); + int64 num_features = config.table_config(table_id).num_features(); + c->set_output(table_id, c->Matrix(batch_size * num_features, width)); + } + return Status::OK(); +} + +} // namespace tpu_embedding_config_util + +REGISTER_OP("TPUEmbeddingReceiveActivations") + .Output("outputs: num_tables * float") + .Attr("num_tables: int >= 1") + .Attr("tpu_embedding_config: string") + .SetIsStateful() + .SetShapeFn(tpu_embedding_config_util::ActivationShapes) + .Doc(R"doc( +An op that receives embeddng activations on the TPU. + +The TPU system performs the embedding lookups and aggregations specified by +the arguments to TPUEmbeddingEnqueueSparseBatch. The results of these +aggregations are visible to the Tensorflow Graph as the outputs of a +TPUEmbeddingDequeueActivations Op. This op returns a list containing one +Tensor of activations per table specified in the model. There can be at most +one ReceieveActivations op in the TPU graph. + +outputs: A TensorList of embedding activations containing one Tensor per + embedding table in the model. +num_tables: The number of output activation tensors, equal to the number of + embedding tables in the model. +tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto. +)doc"); + +REGISTER_OP("TPUEmbeddingActivations") + .Input("embedding_variable: float32") + .Input("sliced_activations: float32") + .Output("output: float32") + .Attr("table_id: int >= 0") + .Attr("lookup_id: int >= 0") + .SetShapeFn([](shape_inference::InferenceContext *c) { + c->set_output(0, c->input(1)); + return Status::OK(); + }) + .Doc(R"doc( +An op enabling differentiation of TPU Embeddings. + +This op simply returns its first input, which is assumed to have been sliced +from the Tensors returnd by TPUEmbeddingDequeueActivations. The presence of this +op, and its first argument being a trainable Variable, enables automatic +differentiation of graphs containing embeddings via the TPU Embedding Python +libraries. + +embedding_variable: A trainable variable, enabling optimizers to find this op. +sliced_activations: The embedding activations Tensor to return. +table_id: The id of the table in the embedding layer configuration from which + these activations were computed. +lookup_id: Identifier of the set of embedding indices which produced these + activations. +)doc"); + +REGISTER_OP("TPUEmbeddingSendGradients") + .Input("gradients: num_tables * float32") + .Attr("num_tables: int >= 1") + .Attr("tpu_embedding_config: string") + .SetIsStateful() + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +An op that performs gradient updates of embedding tables. + +The TensorList argument has the same length and shapes as the return value of +TPUEmbeddingReceiveActivations, but contains gradients of the model's loss +with respect to the embedding activations. The embedding tables are updated +from these gradients via the optimizer specified in the configuration given +to tpu.initialize_system. + +gradients: A TensorList of gradients with which to update embedding tables. +tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto. +)doc"); + +} // namespace tensorflow diff --git a/tensorflow/contrib/tpu/profiler/BUILD b/tensorflow/contrib/tpu/profiler/BUILD index 0e1fca3d3c8b6f3a19b3e989dbee1863475796c5..198da0203a7d17249c4f50110713121b74d5ca4f 100644 --- a/tensorflow/contrib/tpu/profiler/BUILD +++ b/tensorflow/contrib/tpu/profiler/BUILD @@ -44,13 +44,22 @@ cc_library( ], ) +cc_library( + name = "version", + hdrs = ["version.h"], + visibility = ["//visibility:public"], +) + tf_cc_binary( name = "capture_tpu_profile", - srcs = ["capture_tpu_profile.cc"], - visibility = ["//tensorflow/contrib/tpu/profiler:__subpackages__"], + srcs = [ + "capture_tpu_profile.cc", + ], + visibility = ["//visibility:public"], deps = [ ":dump_tpu_profile", ":tpu_profiler_proto_cc", + ":version", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core/distributed_runtime/rpc:grpc_util", diff --git a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc index bff23a447f841339d9bf5bd3bf125d705bf1fee7..b1ef9fde37fe0647965f0818895be37d2d56d207 100644 --- a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc +++ b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc @@ -26,8 +26,10 @@ limitations under the License. #include "tensorflow/contrib/tpu/profiler/dump_tpu_profile.h" #include "tensorflow/contrib/tpu/profiler/tpu_profiler.grpc.pb.h" +#include "tensorflow/contrib/tpu/profiler/version.h" #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" #include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/platform/init_main.h" #include "tensorflow/core/util/command_line_flags.h" @@ -46,21 +48,40 @@ string GetCurrentTimeStampAsString() { return s; } -ProfileResponse Profile(const string& service_addr, int duration_ms) { +Status ValidateHostPortPair(const string& host_port) { + uint32 port; + std::vector parts = str_util::Split(host_port, ':'); + // Must be host:port, port must be a number, host must not contain a '/', + // host also must not be empty. + if (parts.size() != 2 || !strings::safe_strtou32(parts[1], &port) || + parts[0].find("/") != string::npos || parts[0].empty()) { + return errors::InvalidArgument("Could not interpret \"", host_port, + "\" as a host-port pair."); + } + return Status::OK(); +} + +ProfileResponse Profile(const string& service_addr, int duration_ms, + const ProfileOptions& opts) { ProfileRequest request; request.set_duration_ms(duration_ms); request.set_max_events(kMaxEvents); request.add_tools("input_pipeline"); + request.add_tools("overview_page"); + *request.mutable_opts() = opts; std::cout << "Limiting the number of trace events to " << kMaxEvents << std::endl; ::grpc::ClientContext context; ::grpc::ChannelArguments channel_args; // TODO(ioeric): use `SetMaxReceiveMessageSize` instead once it's available. + // TODO(qiuminxu): use `NewHostPortGrpcChannel` instead once their + // `ValidateHostPortPair` checks for empty host string case. channel_args.SetInt(GRPC_ARG_MAX_MESSAGE_LENGTH, std::numeric_limits::max()); std::unique_ptr stub = TPUProfiler::NewStub(::grpc::CreateCustomChannel( - service_addr, ::grpc::InsecureChannelCredentials(), channel_args)); + "dns:///" + service_addr, ::grpc::InsecureChannelCredentials(), + channel_args)); ProfileResponse response; TF_QCHECK_OK(FromGrpcStatus(stub->Profile(&context, request, &response))); return response; @@ -74,26 +95,69 @@ int main(int argc, char** argv) { tensorflow::string FLAGS_service_addr; tensorflow::string FLAGS_logdir; int FLAGS_duration_ms = 2000; + int FLAGS_num_tracing_attempts = 3; + bool FLAGS_include_dataset_ops = true; std::vector flag_list = { tensorflow::Flag("service_addr", &FLAGS_service_addr, "Address of TPU profiler service e.g. localhost:8466"), tensorflow::Flag("logdir", &FLAGS_logdir, - "Path of TensorBoard log directory e.g. /tmp/tb_log"), + "Path of TensorBoard log directory e.g. /tmp/tb_log, " + "gs://tb_bucket"), tensorflow::Flag("duration_ms", &FLAGS_duration_ms, "Duration of tracing in ms. Default is 2000ms."), + tensorflow::Flag("num_tracing_attempts", &FLAGS_num_tracing_attempts, + "Automatically retry N times when no trace event " + "is collected. Default is 3."), + tensorflow::Flag("include_dataset_ops", &FLAGS_include_dataset_ops, + "Set to false to profile longer TPU device traces."), }; + std::cout << "Welcome to the Cloud TPU Profiler v" << TPU_PROFILER_VERSION + << std::endl; + tensorflow::string usage = tensorflow::Flags::Usage(argv[0], flag_list); bool parse_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); if (!parse_ok || FLAGS_service_addr.empty() || FLAGS_logdir.empty()) { - std::printf("%s", usage.c_str()); + std::cout << usage.c_str() << std::endl; + return 2; + } + tensorflow::Status status = + tensorflow::tpu::ValidateHostPortPair(FLAGS_service_addr); + if (!status.ok()) { + std::cout << status.error_message() << std::endl; + std::cout << usage.c_str() << std::endl; return 2; } tensorflow::port::InitMain(argv[0], &argc, &argv); - int duration_ms = FLAGS_duration_ms; - tensorflow::ProfileResponse response = - tensorflow::tpu::Profile(FLAGS_service_addr, duration_ms); + // Sets the minimum duration_ms and tracing attempts to one. + int duration_ms = std::max(FLAGS_duration_ms, 1); + int remaining_attempts = std::max(FLAGS_num_tracing_attempts, 1); + tensorflow::ProfileOptions opts; + opts.set_include_dataset_ops(FLAGS_include_dataset_ops); + tensorflow::ProfileResponse response; + + while (true) { + std::cout << "Starting to profile TPU traces for " << duration_ms << " ms. " + << "Remaining attempt(s): " << remaining_attempts-- << std::endl; + response = tensorflow::tpu::Profile(FLAGS_service_addr, duration_ms, opts); + if (remaining_attempts <= 0 || !response.encoded_trace().empty()) break; + std::cout << "No trace event is collected. Automatically retrying." + << std::endl + << std::endl; + } + + if (response.encoded_trace().empty()) { + std::cout << "No trace event is collected after " + << FLAGS_num_tracing_attempts << " attempt(s). " + << "Perhaps, you want to try again (with more attempts?)." + << std::endl + << "Tip: increase number of attempts with --num_tracing_attempts." + << std::endl; + // Don't dump profile data if no trace is collected. + return 0; + } + // Use the current timestamp as the run name. tensorflow::string run = tensorflow::tpu::GetCurrentTimeStampAsString(); TF_CHECK_OK(tensorflow::tpu::WriteTensorboardTPUProfile( diff --git a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc index 120a38b6c2353deaf0b86d330cda999ba6be7dbf..ebd6185faad28ae7a22eb33f6b358eb2344c9c22 100644 --- a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc +++ b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc @@ -149,10 +149,9 @@ Status WriteTensorboardTPUProfile(const string& logdir, const string& run, // Dumps profile data to /plugins/profile//. string profile_run_dir = JoinPath(logdir, kProfilePluginDirectory, run); TF_RETURN_IF_ERROR(Env::Default()->RecursivelyCreateDir(profile_run_dir)); + // Ignore computation_graph for now. - if (response.encoded_trace().empty()) { - *os << "No trace event is collected." << std::endl; - } else { + if (!response.encoded_trace().empty()) { LOG(INFO) << "Converting trace events to TraceViewer JSON."; TF_RETURN_IF_ERROR( DumpTraceToLogDirectory(profile_run_dir, response.encoded_trace(), os)); @@ -163,13 +162,10 @@ Status WriteTensorboardTPUProfile(const string& logdir, const string& run, TF_RETURN_IF_ERROR(DumpOpProfileToLogDirectory(profile_run_dir, response.op_profile(), os)); } - if (!response.tool_data().empty()) { - for (const auto& tool_data : response.tool_data()) { - TF_RETURN_IF_ERROR( - DumpToolDataToLogDirectory(profile_run_dir, tool_data, os)); - } + for (const auto& tool_data : response.tool_data()) { + TF_RETURN_IF_ERROR( + DumpToolDataToLogDirectory(profile_run_dir, tool_data, os)); } - TF_RETURN_IF_ERROR(DumpGraphEvents(logdir, run, response, os)); return Status::OK(); } diff --git a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.h b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.h index 65b92aa41867ed9e2e8b06c9e34dd99068bb459c..29ef977bacfd61e163be49558c5b94277ed479c1 100644 --- a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.h +++ b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TPU_PROFILER_DUMP_TPU_PROFILE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TPU_PROFILER_DUMP_TPU_PROFILE_H_ +#ifndef TENSORFLOW_CONTRIB_TPU_PROFILER_DUMP_TPU_PROFILE_H_ +#define TENSORFLOW_CONTRIB_TPU_PROFILER_DUMP_TPU_PROFILE_H_ #include "tensorflow/contrib/tpu/profiler/tpu_profiler.grpc.pb.h" #include "tensorflow/core/lib/core/status.h" @@ -27,7 +27,10 @@ namespace tpu { // The following tools are supported: // - Trace viewer // - Op profile -// - HLO computation graph +// - Input pipeline analyzer +// - Overview page +// Note: this function creates a directory even when all fields in +// ProfileResponse are unset/empty. Status WriteTensorboardTPUProfile(const string& logdir, const string& run, const ProfileResponse& response, std::ostream* os); @@ -35,4 +38,4 @@ Status WriteTensorboardTPUProfile(const string& logdir, const string& run, } // namespace tpu } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TPU_PROFILER_DUMP_TPU_PROFILE_H_ +#endif // TENSORFLOW_CONTRIB_TPU_PROFILER_DUMP_TPU_PROFILE_H_ diff --git a/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py b/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py index 7970c20a2693cbbe91a136080240f676d29f2053..a730d6142d890cc41f72176cf617ac0b0434192c 100644 --- a/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py +++ b/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py @@ -17,6 +17,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl import flags import os import subprocess @@ -24,13 +25,36 @@ import sys import tensorflow as tf -tf.flags.DEFINE_string('service_addr', '', - 'Address of TPU profiler service e.g. localhost:8466') -tf.flags.DEFINE_string('logdir', '', - 'Path of TensorBoard log directory e.g. /tmp/tb_log') -tf.flags.DEFINE_integer('duration_ms', 2000, 'Duration of tracing in ms.') +# Cloud TPU Cluster Resolvers +flags.DEFINE_string( + 'gcp_project', None, + 'Project name for the Cloud TPU-enabled project. If not specified, we ' + 'will attempt to automatically detect the GCE project from metadata.') +flags.DEFINE_string( + 'tpu_zone', + None, + help='GCE zone where the Cloud TPU is located in. If not specified, we ' + 'will attempt to automatically detect the GCE project from metadata.') +flags.DEFINE_string('tpu_name', None, + 'Name of the Cloud TPU for Cluster Resolvers. You must ' + 'specify either this flag or --master.') -FLAGS = tf.flags.FLAGS +# Tool specific parameters +flags.DEFINE_string( + 'service_addr', None, 'Address of TPU profiler service e.g. ' + 'localhost:8466, you must specify either this flag or --tpu_name.') +flags.DEFINE_string('logdir', None, + 'Path of TensorBoard log directory e.g. /tmp/tb_log, ' + 'gs://tb_bucket') +flags.DEFINE_integer('duration_ms', 2000, 'Duration of tracing in ms.') +flags.DEFINE_integer('num_tracing_attempts', 3, + 'Automatically retry N times when no trace ' + 'event is collected.') +flags.DEFINE_boolean('include_dataset_ops', True, + 'Set to false to profile longer TPU ' + 'device traces.') + +FLAGS = flags.FLAGS EXECUTABLE = 'data/capture_tpu_profile' @@ -39,13 +63,35 @@ def run_main(): def main(unused_argv=None): - if not FLAGS.service_addr or not FLAGS.logdir: - sys.exit('service_addr and logdir must be provided.') + tf.logging.set_verbosity(tf.logging.INFO) + + if FLAGS.service_addr is None and FLAGS.tpu_name is None: + sys.exit('You must specify either --service_addr or --tpu_name.') + + if FLAGS.service_addr is not None: + if FLAGS.tpu_name is not None: + tf.logging.warn('Both --service_addr and --tpu_name are set. Ignoring ' + '--tpu_name and using --service_addr.') + service_addr = FLAGS.service_addr + else: + tpu_cluster_resolver = ( + tf.contrib.cluster_resolver.TPUClusterResolver( + tpu_names=[FLAGS.tpu_name], + zone=FLAGS.tpu_zone, + project=FLAGS.gcp_project)) + service_addr = tpu_cluster_resolver.get_master() + service_addr = service_addr.replace('grpc://', '').replace(':8470', ':8466') + + if not FLAGS.logdir: + sys.exit('logdir must be provided.') executable_path = os.path.join(os.path.dirname(__file__), EXECUTABLE) + logdir = os.path.expandvars(os.path.expanduser(FLAGS.logdir)) cmd = [executable_path] - cmd.append('--logdir='+FLAGS.logdir) - cmd.append('--service_addr='+FLAGS.service_addr) - cmd.append('--duration_ms='+str(FLAGS.duration_ms)) + cmd.append('--logdir=' + logdir) + cmd.append('--service_addr=' + service_addr) + cmd.append('--duration_ms=' + str(FLAGS.duration_ms)) + cmd.append('--num_tracing_attempts=' + str(FLAGS.num_tracing_attempts)) + cmd.append('--include_dataset_ops=' + str(FLAGS.include_dataset_ops).lower()) subprocess.call(cmd) diff --git a/tensorflow/contrib/tpu/profiler/pip_package/setup.py b/tensorflow/contrib/tpu/profiler/pip_package/setup.py index 179d29602b9f970fb450bc057332fa092066255c..8d99835b64152629c66607e6792495eb36319eb8 100644 --- a/tensorflow/contrib/tpu/profiler/pip_package/setup.py +++ b/tensorflow/contrib/tpu/profiler/pip_package/setup.py @@ -20,16 +20,12 @@ from __future__ import print_function from setuptools import setup -_VERSION = '1.3.0-a1' +_VERSION = '1.6.0-rc1' CONSOLE_SCRIPTS = [ 'capture_tpu_profile=cloud_tpu_profiler.main:run_main', ] -REQUIRED_PACKAGES = [ - 'tensorflow >= 1.2.0', -] - setup( name='cloud_tpu_profiler', version=_VERSION.replace('-', ''), @@ -45,27 +41,22 @@ setup( entry_points={ 'console_scripts': CONSOLE_SCRIPTS, }, - install_requires=REQUIRED_PACKAGES, classifiers=[ # How mature is this project? Common values are # 3 - Alpha # 4 - Beta # 5 - Production/Stable - 'Development Status :: 3 - Alpha', - + 'Development Status :: 4 - Beta', 'Intended Audience :: Developers', 'Intended Audience :: Education', 'Intended Audience :: Science/Research', - 'License :: OSI Approved :: Apache Software License', - 'Programming Language :: Python :: 2', 'Programming Language :: Python :: 2.7', 'Programming Language :: Python :: 3', 'Programming Language :: Python :: 3.4', 'Programming Language :: Python :: 3.5', 'Programming Language :: Python :: 3.6', - 'Topic :: Scientific/Engineering', 'Topic :: Scientific/Engineering :: Mathematics', 'Topic :: Scientific/Engineering :: Artificial Intelligence', @@ -74,4 +65,5 @@ setup( 'Topic :: Software Development :: Libraries :: Python Modules', ], license='Apache 2.0', - keywords='tensorflow performance tpu',) + keywords='tensorflow performance tpu', +) diff --git a/tensorflow/contrib/tpu/profiler/tf_op_stats.proto b/tensorflow/contrib/tpu/profiler/tf_op_stats.proto index 5440bbbfdd75207bd209c19d5cc42dc69504d39b..2094294baad63ae73712c8648b588accd4551ef8 100644 --- a/tensorflow/contrib/tpu/profiler/tf_op_stats.proto +++ b/tensorflow/contrib/tpu/profiler/tf_op_stats.proto @@ -61,6 +61,11 @@ message OpMetricsResult { message OpMetricsDbResult { // A bunch of OpMetricsResults. repeated OpMetricsResult metrics_db = 1; + // The total host infeed-enqueue duration in picoseconds. + optional uint64 total_host_infeed_enq_duration_ps = 2; + // The total of the difference between the start times of two + // consecutive infeed-enqueues (per host) in picoseconds. + optional uint64 total_host_infeed_enq_start_timestamp_ps_diff = 3; } // Result proto for StepInfo. diff --git a/tensorflow/contrib/tpu/profiler/tpu_profiler.proto b/tensorflow/contrib/tpu/profiler/tpu_profiler.proto index bf30d2ce091302eaf361a0018464d3b7de94ea6d..f3f3302ceb3d27dbb21bdce753aeb2d7fcd77448 100644 --- a/tensorflow/contrib/tpu/profiler/tpu_profiler.proto +++ b/tensorflow/contrib/tpu/profiler/tpu_profiler.proto @@ -13,6 +13,14 @@ service TPUProfiler { } } +message ProfileOptions { + // We don't collect the dataset ops by default for better trace-viewer + // scalability. The caller can mannually set this field to include the ops. + bool include_dataset_ops = 1; + + // next-field: 2 +} + message ProfileRequest { // In future, the caller will be able to customize when profiling starts and // stops. For now, it collects `duration_ms` milliseconds worth of data. @@ -25,10 +33,13 @@ message ProfileRequest { // required profiling tools name such as "input_pipeline_analyzer" etc repeated string tools = 3; + // Optional profiling options that control how a TF session will be profiled. + ProfileOptions opts = 4; + // In future, the caller will indicate which TF session is being profiled, and // only data relating to that program will be returned. For now, we assume // all activity during the profiling period is relevant. - // next-field: 4 + // next-field: 5 } message ProfileToolData { diff --git a/tensorflow/contrib/tpu/profiler/trace_events_to_json.h b/tensorflow/contrib/tpu/profiler/trace_events_to_json.h index 992eae43d903db495850ced7a59e38120d3fed34..3bd76dd01c7d0f35bad9386c11811743e1709fca 100644 --- a/tensorflow/contrib/tpu/profiler/trace_events_to_json.h +++ b/tensorflow/contrib/tpu/profiler/trace_events_to_json.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_TPU_PROFILER_TRACE_EVENTS_TO_JSON_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_TPU_PROFILER_TRACE_EVENTS_TO_JSON_H_ +#ifndef TENSORFLOW_CONTRIB_TPU_PROFILER_TRACE_EVENTS_TO_JSON_H_ +#define TENSORFLOW_CONTRIB_TPU_PROFILER_TRACE_EVENTS_TO_JSON_H_ #include "tensorflow/contrib/tpu/profiler/trace_events.pb.h" #include "tensorflow/core/platform/types.h" @@ -29,4 +29,4 @@ string TraceEventsToJson(const Trace &trace); } // namespace tpu } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_TPU_PROFILER_TRACE_EVENTS_TO_JSON_H_ +#endif // TENSORFLOW_CONTRIB_TPU_PROFILER_TRACE_EVENTS_TO_JSON_H_ diff --git a/tensorflow/contrib/tpu/profiler/version.h b/tensorflow/contrib/tpu/profiler/version.h new file mode 100644 index 0000000000000000000000000000000000000000..dc6a934891138018d32d511750120453bdf290cf --- /dev/null +++ b/tensorflow/contrib/tpu/profiler/version.h @@ -0,0 +1,21 @@ +/* Copyright 2017 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_TPU_PROFILER_VERSION_H_ +#define TENSORFLOW_CONTRIB_TPU_PROFILER_VERSION_H_ + +#define TPU_PROFILER_VERSION "1.5.0" + +#endif // TENSORFLOW_CONTRIB_TPU_PROFILER_VERSION_H_ diff --git a/tensorflow/contrib/tpu/proto/BUILD b/tensorflow/contrib/tpu/proto/BUILD index 79a79efb6b62d3e98127558e951ceefd276b580c..e1660985676e8c2efe3b01e32b48b211391885b7 100644 --- a/tensorflow/contrib/tpu/proto/BUILD +++ b/tensorflow/contrib/tpu/proto/BUILD @@ -15,6 +15,15 @@ filegroup( visibility = ["//tensorflow:__subpackages__"], ) +tf_proto_library( + name = "tpu_embedding_config_proto", + srcs = [ + "tpu_embedding_config.proto", + ], + cc_api_version = 2, + visibility = ["//visibility:public"], +) + tf_proto_library( name = "topology_proto", srcs = [ diff --git a/tensorflow/contrib/tpu/proto/tpu_embedding_config.proto b/tensorflow/contrib/tpu/proto/tpu_embedding_config.proto new file mode 100644 index 0000000000000000000000000000000000000000..b0ec968d3a401f1b80ed1bf6fd7a83a69c068fe2 --- /dev/null +++ b/tensorflow/contrib/tpu/proto/tpu_embedding_config.proto @@ -0,0 +1,76 @@ +syntax = "proto3"; + +package tensorflow.tpu; + +// The TPUEmbeddingConfiguration contains specification of TPU Embedding lookups +// and gradient updates separate from the TF Graph. +message TPUEmbeddingConfiguration { + // model_mode specifies whether the model is to be run in training or + // inference. In inference mode, gradient updates to embedding tables are not + // performed. + enum ModelMode { + INVALID = 0; + TRAINING = 1; + INFERENCE = 2; + } + + ModelMode model_mode = 1; + + // num_hosts is the number of host CPU systems in the training/inference job. + // Each embedding table must be sharded into num_hosts separate Variables, + // placed separately on the num_hosts CPU devices in the cluster. Sharding + // will be performed equivalently to the 'div' sharding_strategy option of + // embedding_lookup() and embedding_lookup_sparse(). + int32 num_hosts = 2; + + // The total number of TensorNodes. This is equal to num_hosts times the + // number of TensorNodes attached to each host. + int32 num_tensornodes = 3; + + // The number of training examples per TensorNode. + int32 batch_size = 4; + + message GradientDescentOptimizer { + float learning_rate = 1; + } + + message AdagradOptimizer { + float learning_rate = 1; + float initial_accumulator = 2; + } + + // Each Embedding + message TPUEmbeddingTable { + // Name of the embedding table. This will be used to name Variables in the + // Tensorflow Graph. + string name = 1; + + // Number of rows of the embedding table. The Variable created to hold the + // learned embedding table values will have shape (num_rows, width). + int32 num_rows = 3; + + // Width of the embedding table. The Variable created to hold the + // learned embedding table values will have shape (num_rows, width). + int32 width = 4; + + // Number of distinct embedding activation vectors per training example + // produced by lookups into this table during model evaluation. For each + // table, the Graph will receive an activations Tensor of shape + // (batch_size * table.num_features, table.width). + // For example, num_features = 1 produces equivalent behavior to a single + // tf.nn.embedding_lookup() call. In the case of 'multivalent' embeddings, + // (i.e. tf.nn.embedding_lookup_sparse()) which compute weighted averages of + // embedding table rows, num_features is the number of vectors produced + // after averaging. In sequence models num_features is typically equal + // to the sequence length, since each sequence element must be represented + // separately to the convolutional or recurrent network. + int32 num_features = 5; + + oneof optimizer { + GradientDescentOptimizer gradient_descent = 6; + AdagradOptimizer adagrad = 7; + } + } + + repeated TPUEmbeddingTable table_config = 5; +} diff --git a/tensorflow/contrib/tpu/python/ops/tpu_ops.py b/tensorflow/contrib/tpu/python/ops/tpu_ops.py index 33e47f674d798f622fb08121dabb67d7f45af15b..97876216793e0e6b20b7c072cac4f575b8fd48be 100644 --- a/tensorflow/contrib/tpu/python/ops/tpu_ops.py +++ b/tensorflow/contrib/tpu/python/ops/tpu_ops.py @@ -21,6 +21,7 @@ from __future__ import print_function import platform +from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops if platform.system() != "Windows": @@ -40,6 +41,64 @@ if platform.system() != "Windows": del op # Unused # The gradient of a cross replica sum is also a cross-replica sum. return gen_tpu_ops.cross_replica_sum(grad) + + # This extra type checking exists to give a more helpful error message in + # the common case that uint8 and int64 values are infed. Remove when both + # types are supported. + + _SUPPORTED_INFEED_DTYPES = set([ + dtypes.bool, dtypes.int32, dtypes.bfloat16, dtypes.float32, + dtypes.complex64 + ]) + + def infeed_dequeue(dtype, shape, name=None): + """A placeholder op for a value that will be fed into the computation. + + Args: + dtype: A `tf.DType`. The type of elements in the tensor. + shape: A `tf.TensorShape` or list of `ints`. The shape of the tensor. + name: A name for the operation (optional). + + Returns: + A `Tensor` of type `dtype`. + A tensor that will be provided using the infeed mechanism. + + Raises: + TypeError: If 'dtype` is not a supported infeed type. + """ + if dtype not in _SUPPORTED_INFEED_DTYPES: + raise TypeError( + "{} is not a supported TPU infeed type. Supported types are: " + "{}".format(dtype, list(_SUPPORTED_INFEED_DTYPES))) + + return gen_tpu_ops.infeed_dequeue(dtype, shape, name=name) + + # pylint: disable=redefined-outer-name + def infeed_dequeue_tuple(dtypes, shapes, name=None): + """A placeholder op for values fed into the TPU simultaneously as a tuple. + + Args: + dtypes: A list of `tf.DType`s that has length `>= 1`. + The element types of each element in `outputs`. + shapes: A list of shapes (each a `tf.TensorShape` or list of `ints`). + The shapes of each tensor in `outputs`. + name: A name for the operation (optional). + + Returns: + A list of `Tensor` objects of type `dtypes`. + A list of tensors that will be provided using the infeed mechanism. + + Raises: + TypeError: If a type in 'dtypes` is not a supported infeed type. + """ + for dtype in dtypes: + if dtype not in _SUPPORTED_INFEED_DTYPES: + raise TypeError( + "{} is not a supported TPU infeed type. Supported types are: " + "{}".format(dtype, list(_SUPPORTED_INFEED_DTYPES))) + return gen_tpu_ops.infeed_dequeue_tuple(dtypes, shapes, name=name) + # pylint: enable=redefined-outer-name + else: # We have already built the appropriate libraries into the binary via CMake # if we have built contrib, so we don't need this diff --git a/tensorflow/contrib/tpu/python/tpu/device_assignment.py b/tensorflow/contrib/tpu/python/tpu/device_assignment.py index ee202610a8a8a1406363b3010771e7806d5d84bf..bdd9b88af55fa4fb483ddbdbe5c51d7076cce675 100644 --- a/tensorflow/contrib/tpu/python/tpu/device_assignment.py +++ b/tensorflow/contrib/tpu/python/tpu/device_assignment.py @@ -87,6 +87,8 @@ class DeviceAssignment(object): core_assignment.shape)) self._core_assignment = core_assignment + self._task_and_cores_to_replicas = self._compute_task_and_cores_to_replicas( + self._core_assignment, self._topology_tasks) def _invert_topology(self, topology): """Inverts a [task,device,axis] topology to [x,y,z] -> task/device maps.""" @@ -100,6 +102,34 @@ class DeviceAssignment(object): devices[x, y, z] = device return tasks, devices + def _compute_task_and_cores_to_replicas(self, core_assignment, + topology_tasks): + """Computes a nested dict which maps task and logical core to replicas.""" + task_and_cores_to_replicas = {} + for replica in xrange(core_assignment.shape[0]): + for dx in xrange(core_assignment.shape[1]): + for dy in xrange(core_assignment.shape[2]): + for dz in xrange(core_assignment.shape[3]): + x, y, z = core_assignment[replica, dx, dy, dz, :] + task_id = topology_tasks[x, y, z] + if task_id not in task_and_cores_to_replicas: + task_and_cores_to_replicas[task_id] = {} + logical_core = (dx, dy, dz) + if logical_core not in task_and_cores_to_replicas[task_id]: + task_and_cores_to_replicas[task_id][logical_core] = set() + + task_and_cores_to_replicas[task_id][logical_core].add(replica) + + task_to_sorted_replica_id = {} + + for task, core_to_replicas in task_and_cores_to_replicas.items(): + core_to_sorted_replicas = {} + for core, replicas in core_to_replicas.items(): + core_to_sorted_replicas[core] = sorted(replicas) + + task_to_sorted_replica_id[task] = core_to_sorted_replicas + return task_to_sorted_replica_id + @property def topology(self): """A `Topology` that describes the TPU topology.""" @@ -119,6 +149,11 @@ class DeviceAssignment(object): """ return self._computation_shape + @property + def num_cores_per_replica(self): + """The number of cores per replica.""" + return np.prod(self.computation_shape) + @property def num_replicas(self): """The number of replicas of the computation.""" @@ -148,6 +183,26 @@ class DeviceAssignment(object): logical_offset = tuple([replica] + logical_core.tolist() + [slice(3)]) return tuple(self.core_assignment[logical_offset]) + def lookup_replicas(self, task_id, logical_core): + """Lookup replica ids by task number and logical core. + + Args: + task_id: TensorFlow task number. + logical_core: A tuple of three integers which represents a logical core. + Returns: + A sorted list of the replicas that are attached to that task and + loical_core. + Raises: + ValueError: If no replica exisis in the task which contains the logical + core. + """ + try: + return self._task_and_cores_to_replicas[task_id][logical_core] + except KeyError: + raise ValueError( + "Can not find any replica in task: {} contains logical_core: {} ". + format(task_id, logical_core)) + def tpu_ordinal(self, replica=0, logical_core=None): """Returns the ordinal of the TPU device assigned to a logical core.""" coordinates = self._coordinates(replica, logical_core) diff --git a/tensorflow/contrib/tpu/python/tpu/tpu.py b/tensorflow/contrib/tpu/python/tpu/tpu.py index 659b97e15898141019406b0b6ca4f59b543e6992..d5f54ff4fd278f0c84f79e0079bfb7a409dfba8d 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu.py @@ -30,6 +30,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import variable_scope from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.util import compat # Operations that indicate some error in the users graph, e.g. a placeholder @@ -105,7 +106,7 @@ def core(num): return "device:TPU_REPLICATED_CORE:{}".format(num) -class TPUReplicateContext(control_flow_ops.ControlFlowContext): +class TPUReplicateContext(control_flow_ops.XLAControlFlowContext): """A `ControlFlowContext` for nodes inside a TPU computation. The primary role of `TPUReplicateContext` is to mark operators inside a @@ -121,7 +122,7 @@ class TPUReplicateContext(control_flow_ops.ControlFlowContext): """ def __init__(self, name): - control_flow_ops.ControlFlowContext.__init__(self) + super(TPUReplicateContext, self).__init__() self._name = name self._unsupported_ops = [] @@ -141,8 +142,9 @@ class TPUReplicateContext(control_flow_ops.ControlFlowContext): def _AddOpInternal(self, op): # pylint: disable=protected-access if op.type in _BLACKLISTED_OPS: - raise ValueError("Operation of type %s (%s) is not supported on the TPU" % - (op.type, op.name)) + logging.error("Operation of type %s (%s) is not supported on the TPU. " + "Execution will fail if this op is used in the graph. " % + (op.type, op.name)) if op.type in _NOT_IMPLEMENTED_OPS: self._unsupported_ops.append(op) @@ -151,10 +153,11 @@ class TPUReplicateContext(control_flow_ops.ControlFlowContext): raise NotImplementedError( "Non-resource Variables are not supported inside TPU computations " "(operator name: %s)" % op.name) - # pylint: enable=protected-access if _TPU_REPLICATE_ATTR in op.node_def.attr: raise ValueError("TPU computations cannot be nested") - op.node_def.attr[_TPU_REPLICATE_ATTR].s = self._name + op._set_attr(_TPU_REPLICATE_ATTR, + attr_value_pb2.AttrValue(s=compat.as_bytes(self._name))) + # pylint: enable=protected-access op.graph.prevent_feeding(op) op.graph.prevent_fetching(op) diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_config.py b/tensorflow/contrib/tpu/python/tpu/tpu_config.py index 1b6ce2dfdf09a046038b293810115b3e6e2f05ab..644070218214643923b9ca3ee138615ec568e8b5 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_config.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_config.py @@ -23,6 +23,8 @@ import collections import json import os +import numpy as np + from tensorflow.contrib.tpu.python.tpu import util as util_lib from tensorflow.python.estimator import run_config as run_config_lib from tensorflow.python.platform import tf_logging as logging @@ -31,29 +33,44 @@ from tensorflow.python.platform import tf_logging as logging _TF_CONFIG_ENV = run_config_lib._TF_CONFIG_ENV _SERVICE_KEY = run_config_lib._SERVICE_KEY _TPU_WORKER_JOB_NAME = 'tpu_worker_job_name' +_NUM_CORES_PER_HOST = 8 # pylint: enable=protected-access +# TODO(b/72511246) Provide a simplified api to configure model parallelism. class TPUConfig( collections.namedtuple('TPUConfig', [ 'iterations_per_loop', 'num_shards', + 'computation_shape', 'per_host_input_for_training', 'tpu_job_name', 'initial_infeed_sleep_secs', ])): - """TPU related configuration required by `TPUEstimator`. + r"""TPU related configuration required by `TPUEstimator`. Args: - iterations_per_loop: This is the number of train steps runnining in TPU + iterations_per_loop: This is the number of train steps running in TPU system before returning to CPU host for each `Session.run`. This means global step is increased `iterations_per_loop` times in one `Session.run`. It is recommended to be set as number of global steps for next checkpoint. - num_shards: The number of TPU shards in the system. + num_shards: (Deprecated, ignored by TPUEstimator). + The number of model replicas in the system. For non-model-parallelism + case, this number equals the total number of TPU cores. For + model-parallelism, the total number of TPU cores equals + product(computation_shape) * num_shards. + computation_shape: Defaults to `None`, which disables model parallelism. A + list of size 3 which describes the shape of a model replica's block of + cores. This is required by model-parallelism which enables partitioning + the model to multiple cores. For example, [2, 2, 1] means the model is + partitioned across 4 cores which span two cores in both x and y + coordinates. Please refer to ${tf.contrib.tpu.TopologyProto} for the + geometry of a TPU mesh. per_host_input_for_training: If `True`, `input_fn` is invoked Per-Host rather than Per-Core. With Per-Host input pipeline deployment, `input_fn` - is invoked once on each host. To be precise, with a global batch size + is invoked once on each host. With Per-Core input pipeline deployment, it + is invoked once for each core. To be precise, with a global batch size `train_batch_size` in `TPUEstimator` constructor, the batch size for each shard is `train_batch_size` // #hosts. With Per-Core input pipeline deployment, the shard batch size is `train_batch_size` // #cores. @@ -64,11 +81,15 @@ class TPUConfig( initial_infeed_sleep_secs: The number of seconds the infeed thread should wait before enqueueing the first batch. This helps avoid timeouts for models that require a long compilation time. + + Raises: + ValueError: If `computation_shape` or `computation_shape` are invalid. """ def __new__(cls, iterations_per_loop=2, - num_shards=2, + num_shards=None, + computation_shape=None, per_host_input_for_training=True, tpu_job_name=None, initial_infeed_sleep_secs=None): @@ -78,7 +99,22 @@ class TPUConfig( 'TPUConfig iterations_per_loop') # Check num_shards. - util_lib.check_positive_integer(num_shards, 'TPUConfig num_shards') + if num_shards is not None: + util_lib.check_positive_integer(num_shards, 'TPUConfig num_shards') + + # Check computation_shape + if computation_shape is not None and len(computation_shape) != 3: + raise ValueError( + 'computation_shape must be a list with length 3 or None; got {}'. + format(str(computation_shape))) + + if computation_shape is not None: + computation_shape_array = np.asarray(computation_shape, dtype=np.int32) + # This prevents any computation being replicated across multiple hosts, so + # that each host feeds the same number of computations. + if any(computation_shape_array < 1) or any(computation_shape_array > 2): + raise ValueError('computation_shape elements can only be 1 or 2; got ' + 'computation_shape={}'.format(computation_shape)) # Check initial_infeed_sleep_secs. if initial_infeed_sleep_secs: @@ -91,6 +127,7 @@ class TPUConfig( cls, iterations_per_loop=iterations_per_loop, num_shards=num_shards, + computation_shape=computation_shape, per_host_input_for_training=per_host_input_for_training, tpu_job_name=tpu_job_name, initial_infeed_sleep_secs=initial_infeed_sleep_secs) @@ -99,7 +136,10 @@ class TPUConfig( class RunConfig(run_config_lib.RunConfig): """RunConfig with TPU support.""" - def __init__(self, tpu_config=None, evaluation_master=None, master='', + def __init__(self, + tpu_config=None, + evaluation_master=None, + master=None, **kwargs): """Constructs a RunConfig. @@ -108,16 +148,27 @@ class RunConfig(run_config_lib.RunConfig): evaluation_master: a string. The address of the master to use for eval. Defaults to master if not set. master: a string. The address of the master to use for training. - tf_random_seed: an int. Sets the TensorFlow random seed. Defaults to None, - which initializes it randomly based on the environment. + **kwargs: keyword config parameters. """ super(RunConfig, self).__init__(**kwargs) self._tpu_config = tpu_config or TPUConfig() - if evaluation_master is None: - self._evaluation_master = master - else: + + # If user sets master and/or evaluation_master explicilty, including empty + # string '', take it. Otherwise, take the values set by parent class. + if master is not None: + self._master = master + + if evaluation_master is not None: self._evaluation_master = evaluation_master - self._master = master + elif (not self._evaluation_master and + self.task_type != run_config_lib.TaskType.EVALUATOR): + # If the task type is EVALUATOR, it means some cluster manager sets the + # TF_CONFIG. In that case, we respect the configuration in TF_CONFIG. + # + # Otherwise, it means user executes the code without external cluster + # manager. For that, we optimize the user experience by setting + # evaluation_master to master, unless user overwrites it. + self._evaluation_master = self._master @property def evaluation_master(self): diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_config_test.py b/tensorflow/contrib/tpu/python/tpu/tpu_config_test.py index 618f2636184ac05e69376ddc581dde086aacfbe1..37ef3dbe1e66efe18b13ab9153ee346c08b9774a 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_config_test.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_config_test.py @@ -21,6 +21,7 @@ from __future__ import print_function import json from tensorflow.contrib.tpu.python.tpu import tpu_config as tpu_config_lib +from tensorflow.python.estimator import run_config as run_config_lib from tensorflow.python.platform import test @@ -42,6 +43,89 @@ class TPURunConfigTest(test.TestCase): tpu_config_lib.RunConfig( tpu_config=tpu_config_lib.TPUConfig(iterations_per_loop=0)) + def test_fail_with_invalid_computation_shape(self): + with self.assertRaisesRegexp(ValueError, + 'computation_shape must be a list with length' + ' 3 or None'): + tpu_config_lib.TPUConfig(computation_shape=[2, 1]) + + with self.assertRaisesRegexp(ValueError, + 'computation_shape elements can only be'): + tpu_config_lib.TPUConfig(computation_shape=[1, 3, 1]) + + +class TPURunConfigMasterTest(test.TestCase): + + def test_default_values(self): + run_config = tpu_config_lib.RunConfig() + self.assertEqual('', run_config.master) + self.assertEqual('', run_config.evaluation_master) + + def test_user_provided_master_and_evaluation_master(self): + run_config = tpu_config_lib.RunConfig( + master='_master_123', evaluation_master='_eval_master_123') + self.assertEqual('_master_123', run_config.master) + self.assertEqual('_eval_master_123', run_config.evaluation_master) + + def test_evaluation_master_defaults_to_master(self): + run_config = tpu_config_lib.RunConfig(master='_master_123') + self.assertEqual('_master_123', run_config.master) + self.assertEqual('_master_123', run_config.evaluation_master) + + def test_tf_config(self): + tf_config = { + 'session_master': '_master_123', + 'eval_session_master': '_eval_master_123' + } + with _set_tf_config_env_variable(tf_config): + run_config = tpu_config_lib.RunConfig() + self.assertEqual('_master_123', run_config.master) + self.assertEqual('_eval_master_123', run_config.evaluation_master) + + def test_evaluation_master_defaults_to_master_in_tf_config(self): + tf_config = { + 'session_master': '_master_123', + } + with _set_tf_config_env_variable(tf_config): + run_config = tpu_config_lib.RunConfig() + self.assertEqual('_master_123', run_config.master) + self.assertEqual('_master_123', run_config.evaluation_master) + + def test_respect_evaluation_master_in_tf_config(self): + tf_config = { + 'cluster': { + run_config_lib.TaskType.CHIEF: ['host0:0'], + }, + 'task': { + 'type': run_config_lib.TaskType.EVALUATOR, + 'index': 0 + }, + } + with _set_tf_config_env_variable(tf_config): + run_config = tpu_config_lib.RunConfig(master='_something') + self.assertEqual('', run_config.evaluation_master) + + def test_user_overwrites_tf_config(self): + tf_config = { + 'session_master': '_master_123', + 'eval_session_master': '_eval_master_123' + } + with _set_tf_config_env_variable(tf_config): + run_config = tpu_config_lib.RunConfig( + master='_new_master_123', evaluation_master='_new_eval_master_123') + self.assertEqual('_new_master_123', run_config.master) + self.assertEqual('_new_eval_master_123', run_config.evaluation_master) + + def test_user_overwrites_master_in_tf_config(self): + tf_config = { + 'session_master': '_master_123', + 'eval_session_master': '_eval_master_123' + } + with _set_tf_config_env_variable(tf_config): + run_config = tpu_config_lib.RunConfig(master='_new_master_123') + self.assertEqual('_new_master_123', run_config.master) + self.assertEqual('_eval_master_123', run_config.evaluation_master) + class TPUJobNameTest(test.TestCase): diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_context.py b/tensorflow/contrib/tpu/python/tpu/tpu_context.py new file mode 100644 index 0000000000000000000000000000000000000000..c5c46ea741ea64ca37089431f8ed66cad7bc31fb --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/tpu_context.py @@ -0,0 +1,510 @@ +# 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. +# =================================================================== +"""TPU system metdata and associated tooling.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from contextlib import contextmanager +import copy + +import numpy as np + +from tensorflow.contrib.tpu.python.tpu import device_assignment as tpu_device_assignment +from tensorflow.contrib.tpu.python.tpu import tpu_system_metadata as tpu_system_metadata_lib +from tensorflow.python.estimator import model_fn as model_fn_lib +from tensorflow.python.platform import tf_logging as logging + + +_DEFAULT_JOB_NAME = 'tpu_worker' +_DEFAULT_COORDINATOR_JOB_NAME = 'coordinator' +_LOCAL_MASTERS = ('', 'local') + + +class _TPUContext(object): + """A context holds immutable states of TPU computation. + + This immutable object holds TPUEstimator config, train/eval batch size, and + `TPUEstimator.use_tpu`, which is expected to be passed around. It also + provides utility functions, basded on the current state, to determine other + information commonly required by TPU computation, such as TPU device names, + TPU hosts, shard batch size, etc. + + N.B. As `mode` is not immutable state in Estimator, but essential to + distinguish between TPU training and evaluation, a common usage for + _TPUContext with `mode` is as follows: + ``` + with _ctx.with_mode(mode) as ctx: + if ctx.is_running_on_cpu(): + ... + ``` + """ + + def __init__(self, config, train_batch_size, eval_batch_size, + predict_batch_size, use_tpu): + self._config = config + self._train_batch_size = train_batch_size + self._eval_batch_size = eval_batch_size + self._predict_batch_size = predict_batch_size + self._use_tpu = use_tpu + self._model_parallelism_enabled = ( + use_tpu and config.tpu_config.computation_shape) + self._mode = None + + self._lazy_tpu_system_metadata_dict = {} # key by master address + self._lazy_device_assignment_dict = {} # key by master address + self._lazy_validation_dict = {} # key by ModeKeys + + def _assert_mode(self): + if self._mode is None: + raise RuntimeError( + '`mode` needs to be set via contextmanager `with_mode`.') + return self._mode + + @contextmanager + def with_mode(self, mode): + # NOTE(xiejw): Shallow copy is enough. It will share he lazy dictionaries, + # such as _lazy_tpu_system_metadata_dict between new copy and the original + # one. Note that all lazy states stored in properties _lazy_foo are sort of + # immutable as they should be same for the process lifetime. + new_ctx = copy.copy(self) + new_ctx._mode = mode # pylint: disable=protected-access + yield new_ctx + + @property + def mode(self): + return self._assert_mode() + + def _get_master_address(self): + mode = self._assert_mode() + config = self._config + master = ( + config.master + if mode != model_fn_lib.ModeKeys.EVAL else config.evaluation_master) + return master + + def _get_tpu_system_metadata(self): + """Gets the (maybe cached) TPU system metadata.""" + master = self._get_master_address() + tpu_system_metadata = self._lazy_tpu_system_metadata_dict.get(master) + if tpu_system_metadata is not None: + return tpu_system_metadata + + # pylint: disable=protected-access + tpu_system_metadata = ( + tpu_system_metadata_lib._query_tpu_system_metadata( + master, + run_config=self._config, + query_topology=self.model_parallelism_enabled)) + + self._lazy_tpu_system_metadata_dict[master] = tpu_system_metadata + return tpu_system_metadata + + def _get_device_assignment(self): + """Gets the (maybe cached) TPU device assignment.""" + master = self._get_master_address() + device_assignment = self._lazy_device_assignment_dict.get(master) + if device_assignment is not None: + return device_assignment + + tpu_system_metadata = self._get_tpu_system_metadata() + + device_assignment = tpu_device_assignment.device_assignment( + tpu_system_metadata.topology, + computation_shape=self._config.tpu_config.computation_shape, + num_replicas=self.num_replicas) + + logging.info('computation_shape: %s', + str(self._config.tpu_config.computation_shape)) + logging.info('num_replicas: %d', self.num_replicas) + logging.info('device_assignment.topology.device_coordinates: %s', + str(device_assignment.topology.device_coordinates)) + logging.info('device_assignment.core_assignment: %s', + str(device_assignment.core_assignment)) + + self._lazy_device_assignment_dict[master] = device_assignment + return device_assignment + + @property + def model_parallelism_enabled(self): + return self._model_parallelism_enabled + + @property + def device_assignment(self): + return (self._get_device_assignment() + if self._model_parallelism_enabled else None) + + @property + def num_of_cores_per_host(self): + metadata = self._get_tpu_system_metadata() + return metadata.num_of_cores_per_host + + @property + def num_cores(self): + metadata = self._get_tpu_system_metadata() + return metadata.num_cores + + @property + def num_of_replicas_per_host(self): + if self.model_parallelism_enabled: + return self.num_replicas // self.num_hosts + else: + return self.num_of_cores_per_host + + @property + def num_replicas(self): + num_cores_in_system = self.num_cores + + if self.model_parallelism_enabled: + computation_shape_array = np.asarray( + self._config.tpu_config.computation_shape, dtype=np.int32) + num_cores_per_replica = np.prod(computation_shape_array) + if num_cores_per_replica > num_cores_in_system: + raise ValueError( + 'The num of cores required by the model parallelism, specified by ' + 'TPUConfig.computation_shape, is larger than the total num of ' + 'TPU cores in the system. computation_shape: {}, num cores ' + 'in the system: {}'.format( + self._config.tpu_config.computation_shape, + num_cores_in_system)) + + if num_cores_in_system % num_cores_per_replica != 0: + raise RuntimeError( + 'The num of cores in the system ({}) is not divisible by the num ' + 'of cores ({}) required by the model parallelism, specified by ' + 'TPUConfig.computation_shape. This should never happen!'.format( + num_cores_in_system, num_cores_per_replica)) + + return num_cores_in_system // num_cores_per_replica + else: + return num_cores_in_system + + @property + def num_hosts(self): + metadata = self._get_tpu_system_metadata() + return metadata.num_hosts + + @property + def config(self): + return self._config + + def is_input_sharded_per_core(self): + """Return true if input_fn is invoked per-core (other than per-host).""" + mode = self._assert_mode() + return (mode == model_fn_lib.ModeKeys.TRAIN and + not self._config.tpu_config.per_host_input_for_training) + + def is_running_on_cpu(self, is_export_mode=False): + """Determines whether the input_fn and model_fn should be invoked on CPU. + + This API also validates user provided configuration, such as batch size, + according the lazy initialized TPU system metadata. + + Args: + is_export_mode: Indicates whether the current mode is for exporting the + model, when mode == PREDICT. Only with this bool, we could + tell whether user is calling the Estimator.predict or + Estimator.export_savedmodel, which are running on TPU and CPU + respectively. Parent class Estimator does not distingush these two. + + Returns: + bool, whether current input_fn or model_fn should be running on CPU. + + Raises: + ValueError: any configuration is invalid. + """ + + is_running_on_cpu = self._is_running_on_cpu(is_export_mode) + if not is_running_on_cpu: + self._validate_tpu_configuration() + return is_running_on_cpu + + def _is_running_on_cpu(self, is_export_mode): + """Determines whether the input_fn and model_fn should be invoked on CPU.""" + mode = self._assert_mode() + + if not self._use_tpu: + return True + + if mode != model_fn_lib.ModeKeys.PREDICT: + return False + + # There are actually 2 use cases when running with mode.PREDICT: prediction + # and saving the model. We run actual predictions on the TPU, but + # model export is run on the CPU. + if is_export_mode: + return True + + return False + + @property + def global_batch_size(self): + mode = self._assert_mode() + if mode == model_fn_lib.ModeKeys.TRAIN: + return self._train_batch_size + elif mode == model_fn_lib.ModeKeys.EVAL: + return self._eval_batch_size + elif mode == model_fn_lib.ModeKeys.PREDICT: + return self._predict_batch_size + else: + return None + + @property + def batch_size_for_input_fn(self): + """Returns the shard batch size for `input_fn`.""" + global_batch_size = self.global_batch_size + + if self.is_running_on_cpu(): + return global_batch_size + + # On TPU + if self.is_input_sharded_per_core(): + # We prohibit per core input sharding for the model parallelism case, + # therefore it is safe to use num_cores here. + return global_batch_size // self.num_cores + else: + return global_batch_size // self.num_hosts + + @property + def batch_size_for_model_fn(self): + """Returns the shard batch size for `model_fn`.""" + global_batch_size = self.global_batch_size + + if self.is_running_on_cpu(): + return global_batch_size + + # On TPU. always sharded per shard. + return global_batch_size // self.num_replicas + + @property + def master_job(self): + """Returns the job name to use to place TPU computations on. + + Returns: + A string containing the job name, or None if no job should be specified. + + Raises: + ValueError: If the user needs to specify a tpu_job_name, because we are + unable to infer the job name automatically, or if the user-specified job + names are inappropriate. + """ + run_config = self._config + # If the user specifies the tpu_job_name, use that. + if run_config.tpu_config.tpu_job_name: + return run_config.tpu_config.tpu_job_name + + # The tpu job is determined by the run_config. Right now, this method is + # required as tpu_config is not part of the RunConfig. + mode = self._assert_mode() + master = ( + run_config.evaluation_master + if mode == model_fn_lib.ModeKeys.EVAL else run_config.master) + if master in _LOCAL_MASTERS: + return None + + if (not run_config.session_config or + not run_config.session_config.cluster_def.job): + return _DEFAULT_JOB_NAME + cluster_def = run_config.session_config.cluster_def + job_names = set([job.name for job in cluster_def.job]) + if _DEFAULT_JOB_NAME in job_names: + # b/37868888 tracks allowing ClusterSpec propagation to reuse job names. + raise ValueError('Currently, tpu_worker is not an allowed job name.') + if len(job_names) == 1: + return cluster_def.job[0].name + if len(job_names) == 2: + if _DEFAULT_COORDINATOR_JOB_NAME in job_names: + job_names.remove(_DEFAULT_COORDINATOR_JOB_NAME) + return job_names.pop() + # TODO(b/67716447): Include more sophisticated heuristics. + raise ValueError( + 'Could not infer TPU job name. Please specify a tpu_job_name as part ' + 'of your TPUConfig.') + + @property + def tpu_host_placement_function(self): + """Returns the TPU host place function.""" + master = self.master_job + + def _placement_function(_sentinal=None, core_id=None, host_id=None): # pylint: disable=invalid-name + assert _sentinal is None + if core_id is not None and host_id is not None: + raise RuntimeError( + 'core_id and host_id can have only one non-None value.') + + if master is None: + return '/replica:0/task:0/device:CPU:0' + else: + if core_id is not None: + host_id = core_id / self.num_of_cores_per_host + return '/job:%s/task:%d/device:CPU:0' % (master, host_id) + + return _placement_function + + @property + def tpu_device_placement_function(self): + """Returns a TPU device placement Fn.""" + master = self.master_job + job_device = '' if master is None else ('/job:%s' % master) + + def _placement_function(i): + if self.model_parallelism_enabled: + return self.device_assignment.tpu_device(replica=i, job=master) + else: + num_of_cores_per_host = self.num_of_cores_per_host + host_id = i / num_of_cores_per_host + ordinal_id = i % num_of_cores_per_host + return '%s/task:%d/device:TPU:%d' % (job_device, host_id, ordinal_id) + + return _placement_function + + @property + def tpu_ordinal_function(self): + """Returns the TPU ordinal fn.""" + + def _tpu_ordinal_function(index): + """Return the TPU ordinal associated with a shard. + + Required because the enqueue ops are placed on CPU. + + Args: + index: the shard index + + Returns: + The ordinal of the TPU device the shard's infeed should be placed on. + """ + if self.model_parallelism_enabled: + return self.device_assignment.tpu_ordinal(replica=index) + else: + return index % self.num_of_cores_per_host + + return _tpu_ordinal_function + + def _validate_tpu_configuration(self): + """Validates the configuration based on the TPU system metadata.""" + mode = self._assert_mode() + if self._lazy_validation_dict.get(mode): + return + + # All following information is obtained from TPU system metadata. + num_cores = self.num_cores + num_replicas = self.num_replicas + num_hosts = self.num_hosts + + if not num_cores: + tpu_system_metadata = self._get_tpu_system_metadata() + raise RuntimeError( + 'Cannot find any TPU cores in the system. Please double check ' + 'Tensorflow master address and TPU worker(s). Available devices ' + 'are {}.'.format(tpu_system_metadata.devices)) + + if self._config.tpu_config.num_shards: + user_provided_num_replicas = self._config.tpu_config.num_shards + if user_provided_num_replicas != num_replicas: + message = ( + 'TPUConfig.num_shards is not set correctly. According to TPU ' + 'system metadata for Tensorflow master ({}): num_replicas should ' + 'be ({}), got ({}). For non-model-parallelism, num_replicas should ' + 'be the total num of TPU cores in the system. For ' + 'model-parallelism, the total number of TPU cores should be ' + 'product(computation_shape) * num_replicas. Please set it ' + 'accordingly or leave it as `None`'.format( + self._get_master_address(), num_replicas, + user_provided_num_replicas)) + + raise ValueError(message) + + if mode == model_fn_lib.ModeKeys.TRAIN: + if self._train_batch_size % num_replicas != 0: + raise ValueError( + 'train batch size {} must be divisible by number of replicas {}' + .format(self._train_batch_size, num_replicas)) + + elif mode == model_fn_lib.ModeKeys.EVAL: + if self._eval_batch_size is None: + raise ValueError( + 'eval_batch_size in TPUEstimator constructor cannot be `None`' + 'if .evaluate is running on TPU.') + if self._eval_batch_size % num_replicas != 0: + raise ValueError( + 'eval batch size {} must be divisible by number of replicas {}' + .format(self._eval_batch_size, num_replicas)) + if num_hosts > 1: + raise ValueError( + 'TPUEstimator.evaluate should be running on single TPU worker. ' + 'got {}.'.format(num_hosts)) + else: + assert mode == model_fn_lib.ModeKeys.PREDICT + if self._predict_batch_size is None: + raise ValueError( + 'predict_batch_size in TPUEstimator constructor should not be ' + '`None` if .predict is running on TPU.') + if self._predict_batch_size % num_replicas != 0: + raise ValueError( + 'predict batch size {} must be divisible by number of replicas {}' + .format(self._predict_batch_size, num_replicas)) + if num_hosts > 1: + raise ValueError( + 'TPUEstimator.predict should be running on single TPU worker. ' + 'got {}.'.format(num_hosts)) + + # Record the state "validated" into lazy dictionary. + self._lazy_validation_dict[mode] = True + + +class _OneCoreTPUContext(_TPUContext): + """Special _TPUContext for one core usage.""" + + def __init__(self, config, train_batch_size, eval_batch_size, + predict_batch_size, use_tpu): + + super(_OneCoreTPUContext, self).__init__( + config, train_batch_size, eval_batch_size, + predict_batch_size, use_tpu) + + def _get_tpu_system_metadata(self): + """Gets the (maybe cached) TPU system metadata.""" + master = self._get_master_address() + tpu_system_metadata = self._lazy_tpu_system_metadata_dict.get(master) + if tpu_system_metadata is not None: + return tpu_system_metadata + + tpu_system_metadata = ( + tpu_system_metadata_lib._TPUSystemMetadata( # pylint: disable=protected-access + num_cores=1, + num_hosts=1, + num_of_cores_per_host=1, + topology=None, + devices=[])) + + self._lazy_tpu_system_metadata_dict[master] = tpu_system_metadata + return tpu_system_metadata + + +def _get_tpu_context(config, train_batch_size, eval_batch_size, + predict_batch_size, use_tpu): + """Returns an instance of `_TPUContext`.""" + + if (config.tpu_config.num_shards == 1 and + config.tpu_config.computation_shape is None): + logging.warning( + 'Setting TPUConfig.num_shards==1 is an unsupported behavior. ' + 'Please fix as soon as possible (leaving num_shards as None.') + return _OneCoreTPUContext(config, train_batch_size, eval_batch_size, + predict_batch_size, use_tpu) + + return _TPUContext(config, train_batch_size, eval_batch_size, + predict_batch_size, use_tpu) diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py index 6bf11e1ae535f26cfa8326e224f2b8d4e0124b13..ff53fe4f5d0e219f56d77d3476640bb023c7535a 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # =================================================================== - """TPUEstimator class.""" from __future__ import absolute_import @@ -20,28 +19,33 @@ from __future__ import division from __future__ import print_function import collections -from contextlib import contextmanager import copy +import signal import threading import time +import traceback import six from six.moves import queue as Queue # pylint: disable=redefined-builtin +from six.moves import xrange # pylint: disable=redefined-builtin +from tensorflow.contrib.summary import summary_ops as contrib_summary from tensorflow.contrib.tpu.python.ops import tpu_ops from tensorflow.contrib.tpu.python.tpu import tpu from tensorflow.contrib.tpu.python.tpu import tpu_config +from tensorflow.contrib.tpu.python.tpu import tpu_context from tensorflow.contrib.tpu.python.tpu import tpu_feed from tensorflow.contrib.tpu.python.tpu import training_loop from tensorflow.contrib.tpu.python.tpu import util as util_lib - +from tensorflow.core.framework.summary_pb2 import Summary from tensorflow.core.protobuf import config_pb2 - +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.estimator import estimator as estimator_lib from tensorflow.python.estimator import model_fn as model_fn_lib from tensorflow.python.estimator import util from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -52,11 +56,12 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import tf_logging as logging from tensorflow.python.summary import summary +from tensorflow.python.training import basic_session_run_hooks from tensorflow.python.training import evaluation from tensorflow.python.training import session_run_hook from tensorflow.python.training import training from tensorflow.python.training import training_util - +from tensorflow.python.util import tf_inspect _INITIAL_LOSS = 1e7 _ZERO_LOSS = 0. @@ -64,9 +69,15 @@ _TPU_ESTIMATOR = 'tpu_estimator' _ITERATIONS_PER_LOOP_VAR = 'iterations_per_loop' _BATCH_SIZE_KEY = 'batch_size' _CROSS_REPLICA_SUM_OP = 'CrossReplicaSum' + _RESERVED_PARAMS_KEYS = [_BATCH_SIZE_KEY] -# TODO(b/65703635): Flip the value and remove all dead code. + +# TODO(b/65703635): Flip the value and remove all dead code. Currently, this is +# only used for per-core based deployments. For per-host based pipelines, if a +# user returns a Dataset instance it will be automatically wrapped in a +# tf.while_loop (This can be disabled by returning features and labels +# explicitly). _WRAP_INPUT_FN_INTO_WHILE_LOOP = False @@ -83,28 +94,28 @@ def _create_global_step(graph): initializer=init_ops.zeros_initializer(), trainable=False, use_resource=True, - collections=[ops.GraphKeys.GLOBAL_VARIABLES, - ops.GraphKeys.GLOBAL_STEP]) + collections=[ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.GLOBAL_STEP]) def _create_or_get_iterations_per_loop(): graph = ops.get_default_graph() - iter_vars = graph.get_collection(_TPU_ESTIMATOR) + collection_name = '{}_{}'.format(_TPU_ESTIMATOR, _ITERATIONS_PER_LOOP_VAR) + iter_vars = graph.get_collection(collection_name) if len(iter_vars) == 1: return iter_vars[0] elif len(iter_vars) > 1: raise RuntimeError('Multiple iterations_per_loop_var in collection.') with ops.colocate_with(training_util.get_global_step()): - with variable_scope.variable_scope(_TPU_ESTIMATOR, - reuse=variable_scope.AUTO_REUSE): + with variable_scope.variable_scope( + _TPU_ESTIMATOR, reuse=variable_scope.AUTO_REUSE): return variable_scope.get_variable( _ITERATIONS_PER_LOOP_VAR, initializer=init_ops.zeros_initializer(), shape=[], dtype=dtypes.int32, trainable=False, - collections=[_TPU_ESTIMATOR], + collections=[collection_name, ops.GraphKeys.LOCAL_VARIABLES], use_resource=True) @@ -136,218 +147,6 @@ def _increase_eval_step_op(iterations_per_loop): use_locking=True) -_DEFAULT_JOB_NAME = 'tpu_worker' -_DEFAULT_COORDINATOR_JOB_NAME = 'coordinator' -_LOCAL_MASTERS = ('', 'local') - - -class _TPUContext(object): - """A context holds immutable states of TPU computation. - - This immutable object holds TPUEstimator config, train/eval batch size, and - `TPUEstimator.use_tpu`, which is expected to be passed around. It also - provides utility functions, basded on the current state, to determine other - information commonly required by TPU computation, such as TPU device names, - TPU hosts, shard batch size, etc. - - N.B. As `mode` is not immutable state in Estimator, but essential to - distinguish between TPU training and evaluation, a common usage for - _TPUContext with `mode` is as follows: - ``` - with _ctx.with_mode(mode) as ctx: - if ctx.is_running_on_cpu(): - ... - ``` - """ - - def __init__(self, config, train_batch_size, eval_batch_size, use_tpu): - self._config = config - self._train_batch_size = train_batch_size - self._eval_batch_size = eval_batch_size - self._use_tpu = use_tpu - self._num_shards_or_none = self._config.tpu_config.num_shards - self._mode = None - - def _assert_mode(self): - if self._mode is None: - raise RuntimeError( - '`mode` needs to be set via contextmanager `with_mode`.') - return self._mode - - @property - def num_of_cores_per_host(self): - num_cores = self.num_cores - return min(num_cores, 8) - - @contextmanager - def with_mode(self, mode): - new_ctx = copy.copy(self) # Shallow copy is enough. - new_ctx._mode = mode # pylint: disable=protected-access - yield new_ctx - - @property - def mode(self): - return self._assert_mode() - - @property - def num_cores(self): - # TODO(xiejw): Adds lazy num_shards initialization. - return self._num_shards_or_none - - @property - def num_hosts(self): - return self.num_cores // self.num_of_cores_per_host - - @property - def config(self): - return self._config - - def is_input_sharded_per_core(self): - """Return true if input_fn is invoked per-core (other than per-host).""" - self._assert_mode() - return (self._mode == model_fn_lib.ModeKeys.TRAIN and - not self._config.tpu_config.per_host_input_for_training) - - def is_running_on_cpu(self): - """Determines whether the input_fn and model_fn should be invoked on CPU.""" - mode = self._assert_mode() - return ((not self._use_tpu) or mode == model_fn_lib.ModeKeys.PREDICT or - (mode == model_fn_lib.ModeKeys.EVAL and - self._eval_batch_size is None)) - - @property - def batch_size_for_input_fn(self): - """Returns the shard batch size for `input_fn`.""" - mode = self._assert_mode() - # Special case for eval. - if mode == model_fn_lib.ModeKeys.EVAL and self._eval_batch_size is None: - return None - if self.is_running_on_cpu(): - if mode == model_fn_lib.ModeKeys.TRAIN: - return self._train_batch_size - if mode == model_fn_lib.ModeKeys.EVAL: - return self._eval_batch_size - return None - - global_batch_size = (self._train_batch_size if - mode == model_fn_lib.ModeKeys.TRAIN - else self._eval_batch_size) - # On TPU - if self.is_input_sharded_per_core(): - return global_batch_size // self.num_cores - else: - return global_batch_size // self.num_hosts - - @property - def batch_size_for_model_fn(self): - """Returns the shard batch size for `model_fn`.""" - mode = self._assert_mode() - # Special case for eval. - if mode == model_fn_lib.ModeKeys.EVAL and self._eval_batch_size is None: - return None - if self.is_running_on_cpu(): - if mode == model_fn_lib.ModeKeys.TRAIN: - return self._train_batch_size - if mode == model_fn_lib.ModeKeys.EVAL: - return self._eval_batch_size - return None - - # On TPU. always sharded per core. - if mode == model_fn_lib.ModeKeys.TRAIN: - return self._train_batch_size // self.num_cores - else: - return self._eval_batch_size // self.num_cores - - @property - def master_job(self): - """Returns the job name to use to place TPU computations on. - - Returns: - A string containing the job name, or None if no job should be specified. - - Raises: - ValueError: If the user needs to specify a tpu_job_name, because we are - unable to infer the job name automatically, or if the user-specified job - names are inappropriate. - """ - run_config = self._config - # If the user specifies the tpu_job_name, use that. - if run_config.tpu_config.tpu_job_name: - return run_config.tpu_config.tpu_job_name - - # The tpu job is determined by the run_config. Right now, this method is - # required as tpu_config is not part of the RunConfig. - mode = self._assert_mode() - master = (run_config.evaluation_master if mode == model_fn_lib.ModeKeys.EVAL - else run_config.master) - if master in _LOCAL_MASTERS: - return None - - if (not run_config.session_config or - not run_config.session_config.cluster_def.job): - return _DEFAULT_JOB_NAME - cluster_def = run_config.session_config.cluster_def - job_names = set([job.name for job in cluster_def.job]) - if _DEFAULT_JOB_NAME in job_names: - # b/37868888 tracks allowing ClusterSpec propagation to reuse job names. - raise ValueError('Currently, tpu_worker is not an allowed job name.') - if len(job_names) == 1: - return cluster_def.job[0].name - if len(job_names) == 2: - if _DEFAULT_COORDINATOR_JOB_NAME in job_names: - job_names.remove(_DEFAULT_COORDINATOR_JOB_NAME) - return job_names.pop() - # TODO(b/67716447): Include more sophisticated heuristics. - raise ValueError( - 'Could not infer TPU job name. Please specify a tpu_job_name as part ' - 'of your TPUConfig.') - - @property - def tpu_host_placement_function(self): - """Returns the TPU host place function.""" - master = self.master_job - def _placement_function(_sentinal=None, core_id=None, host_id=None): # pylint: disable=invalid-name - assert _sentinal is None - if core_id is not None and host_id is not None: - raise RuntimeError( - 'core_id and host_id can have only one non-None value.') - - if master is None: - return '/replica:0/task:0/device:CPU:0' - else: - # This assumes that if using more than 8 shards, - # the job configuration varies 'task'. - if core_id is not None: - host_id = core_id / 8 - return '/job:%s/task:%d/device:CPU:0' % (master, host_id) - return _placement_function - - @property - def tpu_device_placement_function(self): - master = self.master_job - job_device = '' if master is None else ('/job:%s' % master) - def _placement_function(i): - return '%s/task:%d/device:TPU:%d' % (job_device, i / 8, i % 8) - return _placement_function - - @property - def tpu_ordinal_function(self): - """Returns the TPU ordinal fn.""" - def _tpu_ordinal_function(index): - """Return the TPU ordinal associated with a shard. - - Required because the enqueue ops are placed on CPU. - - Args: - index: the shard index - - Returns: - The ordinal of the TPU device the shard's infeed should be placed on. - """ - return index % 8 - return _tpu_ordinal_function - - class _SIGNAL(object): """Signal used to control the thread of infeed/outfeed. @@ -358,14 +157,17 @@ class _SIGNAL(object): STOP = -2 -class TPUEstimatorSpec(collections.namedtuple('TPUEstimatorSpec', [ - 'mode', - 'predictions', - 'loss', - 'train_op', - 'eval_metrics', - 'export_outputs', - 'scaffold_fn'])): +class TPUEstimatorSpec( + collections.namedtuple('TPUEstimatorSpec', [ + 'mode', + 'predictions', + 'loss', + 'train_op', + 'eval_metrics', + 'export_outputs', + 'scaffold_fn', + 'host_call' + ])): """Ops and objects returned from a `model_fn` and passed to `TPUEstimator`. See `EstimatorSpec` for `mode`, 'predictions, 'loss', 'train_op', and @@ -375,7 +177,7 @@ class TPUEstimatorSpec(collections.namedtuple('TPUEstimatorSpec', [ `metric_fn` runs on CPU to generate metrics and `tensors` represents the `Tensor`s transferred from TPU system to CPU host and passed to `metric_fn`. To be precise, TPU evaluation expects a slightly different signature from the - ${tf.estimator.Estimator}. While `EstimatorSpec.eval_metric_ops` expects a + @{tf.estimator.Estimator}. While `EstimatorSpec.eval_metric_ops` expects a dict, `TPUEstimatorSpec.eval_metrics` is a tuple of `metric_fn` and `tensors`. The `tensors` could be a list of `Tensor`s or dict of names to `Tensor`s. The `tensors` usually specify the model logits, which are transferred back from @@ -390,6 +192,15 @@ class TPUEstimatorSpec(collections.namedtuple('TPUEstimatorSpec', [ `scaffold_fn` is a function running on CPU to generate the `Scaffold`. This function should not capture any Tensors in `model_fn`. + + `host_call` is a tuple of a `function` and a list or dictionary of `tensors` + to pass to that function and returns a list of Tensors. `host_call` currently + works for train() and evaluate(). The Tensors returned by the function is + executed on the CPU on every step, so there is communication overhead when + sending tensors from TPU to CPU. To reduce the overhead, try reducing the + size of the tensors. The `tensors` are concatenated along their major (batch) + dimension, and so must be >= rank 1. The `host_call` is useful for writing + summaries with @{tf.contrib.summary.create_file_writer}. """ def __new__(cls, @@ -399,121 +210,176 @@ class TPUEstimatorSpec(collections.namedtuple('TPUEstimatorSpec', [ train_op=None, eval_metrics=None, export_outputs=None, - scaffold_fn=None): + scaffold_fn=None, + host_call=None): """Creates a validated `TPUEstimatorSpec` instance.""" + host_calls = {} if eval_metrics is not None: - _EvalMetrics.validate(eval_metrics) - return super(TPUEstimatorSpec, cls).__new__(cls, - mode=mode, - predictions=predictions, - loss=loss, - train_op=train_op, - eval_metrics=eval_metrics, - export_outputs=export_outputs, - scaffold_fn=scaffold_fn) + host_calls['eval_metrics'] = eval_metrics + if host_call is not None: + host_calls['host_call'] = host_call + _OutfeedHostCall.validate(host_calls) + return super(TPUEstimatorSpec, cls).__new__( + cls, + mode=mode, + predictions=predictions, + loss=loss, + train_op=train_op, + eval_metrics=eval_metrics, + export_outputs=export_outputs, + scaffold_fn=scaffold_fn, + host_call=host_call) def as_estimator_spec(self): """Creates an equivalent `EstimatorSpec` used by CPU train/eval.""" - eval_metric_ops = _EvalMetrics.to_metric_metric_ops_for_cpu( - self.eval_metrics) + host_calls = {} + if self.eval_metrics is not None: + host_calls['eval_metrics'] = self.eval_metrics + if self.host_call is not None: + host_calls['host_call'] = self.host_call + host_call_ret = _OutfeedHostCall.create_cpu_hostcall(host_calls) + eval_metric_ops = None + if self.eval_metrics is not None: + eval_metric_ops = host_call_ret['eval_metrics'] + hooks = None + if self.host_call is not None: + hooks = [_OutfeedHostCallHook(host_call_ret['host_call'])] scaffold = self.scaffold_fn() if self.scaffold_fn else None - return model_fn_lib.EstimatorSpec(mode=self.mode, - predictions=self.predictions, - loss=self.loss, - train_op=self.train_op, - eval_metric_ops=eval_metric_ops, - export_outputs=self.export_outputs, - scaffold=scaffold) + return model_fn_lib.EstimatorSpec( + mode=self.mode, + predictions=self.predictions, + loss=self.loss, + train_op=self.train_op, + eval_metric_ops=eval_metric_ops, + export_outputs=self.export_outputs, + scaffold=scaffold, + training_hooks=hooks, + evaluation_hooks=hooks, + prediction_hooks=hooks) + + +class _OpQueueContext(object): + """Manages work queue and thread for a infeed/outfeed thread.""" + + def __init__(self, name, target, args): + self._name = name + self._queue = Queue.Queue() + args = (self,) + args + self._thread = threading.Thread(name=name, target=target, args=args) + self._thread.daemon = True + self._thread.start() + + def stop(self): + self._queue.put(_SIGNAL.STOP) + + def send_next_batch_signal(self, iterations): + self._queue.put(iterations) + + def read_iteration_counts(self): + while True: + iterations = self._queue.get(block=True) + logging.debug('%s read iterations %s', self._name, iterations) + if iterations == _SIGNAL.STOP: + logging.info('%s received shutdown signal, stopping.', self._name) + return + yield iterations + def join(self): + logging.info('Shutting down %s thread.' % self._name) + self.stop() + self._thread.join() -class _InfeedOutfeedThreadBaseController(object): - """This wraps the infeed/outfeed thread and stops when Estimator finishes.""" - def __init__(self, thd): - self._signal_queue = Queue.Queue() - thd.daemon = True - thd.start() - self._thd = thd +class _OpSignalOnceQueueContext(_OpQueueContext): + """Manages work queue and thread for a infeed/outfeed thread. - def block_and_get_signal(self): - return self._signal_queue.get() + This subclass only signals once. + """ - def send_next_batch_signal(self, signal=_SIGNAL.NEXT_BATCH): - self._signal_queue.put(signal) + def __init__(self, name, target, args): + super(_OpSignalOnceQueueContext, self).__init__(name, target, args) + self._has_signaled = False - def join(self): - self._signal_queue.put(_SIGNAL.STOP) - self._thd.join() + def send_next_batch_signal(self, iterations): + if not self._has_signaled: + self._queue.put(iterations) + self._has_signaled = True -class _OutfeedThreadController(_InfeedOutfeedThreadBaseController): - """This wraps the outfeed thread and stops when Estimator finishes.""" +class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): + """A Session hook setting up the TPU initialization, infeed, and outfeed. - def __init__(self, session, dequeue_ops): - super(_OutfeedThreadController, self).__init__( - threading.Thread(target=self._execute_dequeue_ops, - args=(session, dequeue_ops))) + This hook does two major things: + 1. initialize and shutdown TPU system. + 2. launch and join the threads for infeed enqueue and (optional) outfeed + dequeue. + """ - def _execute_dequeue_ops(self, session, dequeue_ops): - count = 0 - while True: - signal = self.block_and_get_signal() - if signal == _SIGNAL.STOP: - logging.info('Stop outfeed thread.') - return + def __init__(self, + ctx, + enqueue_ops, + dequeue_ops, + run_infeed_loop_on_coordinator=True): + self._master_job = ctx.master_job + self._enqueue_ops = enqueue_ops + self._dequeue_ops = dequeue_ops - iterations = signal - for i in range(iterations): - logging.debug('Outfeed dequeue for iteration (%d, %d)', count, i) - session.run(dequeue_ops) - count += 1 + self._run_infeed_loop_on_coordinator = run_infeed_loop_on_coordinator + self._initial_infeed_sleep_secs = ( + ctx.config.tpu_config.initial_infeed_sleep_secs) - def join(self): - logging.info('Waiting for Outfeed Thread to exit.') - super(_OutfeedThreadController, self).join() - - -class _InfeedThreadController(_InfeedOutfeedThreadBaseController): - """This wraps the infeed thread and stops when Estimator finishes.""" - - def __init__(self, session, enqueue_ops, initial_infeed_sleep_secs): - super(_InfeedThreadController, self).__init__( - threading.Thread( - target=self._input_thread_fn_for_loading, - args=(session, enqueue_ops, initial_infeed_sleep_secs))) - - def _input_thread_fn_for_loading(self, session, enqueue_ops, - initial_infeed_sleep_secs): - count = 0 - if initial_infeed_sleep_secs: - logging.info('Infeed thread sleeping for %d seconds.', - initial_infeed_sleep_secs) - time.sleep(initial_infeed_sleep_secs) - logging.info('Infeed thread starting after sleep') - try: - while True: - signal = self._signal_queue.get() - if signal == _SIGNAL.STOP: - logging.info('Stop Infeed input thread.') - return - - if _WRAP_INPUT_FN_INTO_WHILE_LOOP: - # Enqueue batches for next loop. - session.run(enqueue_ops) - else: - iterations = signal - for i in range(iterations): - logging.debug('Infeed enqueue for iteration (%d, %d)', count, i) - session.run(enqueue_ops) - count += 1 + self._session_cancel_timer = None + + self._feed_error = None + self._finished = False - except Exception: # pylint: disable=broad-except + def begin(self): + logging.info('TPU job name %s', self._master_job) + self._iterations_per_loop_var = _create_or_get_iterations_per_loop() + self._init_ops = [tpu.initialize_system(job=self._master_job)] + self._finalize_ops = [tpu.shutdown_system(job=self._master_job)] + + summary_writer_init_ops = contrib_summary.summary_writer_initializer_op() + self._init_ops.extend(summary_writer_init_ops) + # Get all the writer resources from the initializer, so we know what to + # flush. + for op in summary_writer_init_ops: + self._finalize_ops.append(contrib_summary.flush(writer=op.inputs[0])) + + def _log_error(self, session, error): + """Log an infeed or outfeed error. + + This logs a short error message immediately, and schedules a timer to + emit the full stack trace and error message after a short period of time. + If the main session has terminated by the time the timer triggers, we + assume the real source of the error was from the main session and avoid + emitting a stack trace for the infeed. + + Args: + session: `tf.Session`, session to be terminated error: exception that + triggered logging. + error: the Exception to log. + """ + logging.warning( + '\n\n' + 'Error occurred during infeed/outfeed. This may be due to a compile ' + 'error in the main session. Waiting for a short time for the main ' + 'session to come back.\n\n%s', error) + + self._feed_error = traceback.format_exc() + + # If we've already encountered a feed error, don't schedule another + # cancellation op. + if self._session_cancel_timer: + return + + def _cancel_session(): # Close the session to avoid the main thread from hanging. If input # pipeline triggers any error, the infeed thread dies but the main thread # for TPU computation waits for the infeed enqueue forever. Close the # Session to cancel the main thread Session.run execution. # - # However, sleep for 2 minutes before explicit closing to give some time + # We sleep for a few seconds before closing to give some time # for the TPU compilation error, if any, propagating, from TPU to CPU # host. Compilation errors should be reported by the main thread so that # the program can be interrupted and users can take action. Due to a race @@ -522,80 +388,108 @@ class _InfeedThreadController(_InfeedOutfeedThreadBaseController): # exception in the main thread, instead of the expected compile error. # User code that depends on having the proper exception type will # therefore be confused. - logging.error( - 'Failed running infeed, closing session.\n' - 'You may see an exception from your main session after this. ' - 'Sleep for 2 minutes before close Session from infeed thread to ' - 'allow the main thread returning an error first, if any.', - exc_info=1 - ) - time.sleep(120) - logging.error('Closing the failed session.') - session.close() - - def join(self): - logging.info('Waiting for Infeed Thread to exit.') - super(_InfeedThreadController, self).join() + time.sleep(5) + + # If the main session is still running, the infeed/outfeed errors are + # legitimate, and should be logged. + if not self._finished and self._feed_error: + logging.error('Feed error: %s', self._feed_error) + logging.error('Closing session. A RuntimeError should follow.') + session.close() + + self._session_cancel_timer = threading.Thread(target=_cancel_session) + self._session_cancel_timer.daemon = True + self._session_cancel_timer.start() + + def _run_infeed(self, queue_ctx, session): + logging.info('Starting infeed thread controller.') + if self._initial_infeed_sleep_secs: + logging.info('%s thread sleeping for %d seconds.', self._name, + self._initial_infeed_sleep_secs) + time.sleep(self._initial_infeed_sleep_secs) + logging.info('%s thread starting after sleep', self._name) + try: + if self._run_infeed_loop_on_coordinator: + for count, steps in enumerate(queue_ctx.read_iteration_counts()): + for i in xrange(steps): + logging.debug('Infeed enqueue for iteration (%d, %d)', count, i) + session.run(self._enqueue_ops) + else: + for _ in queue_ctx.read_iteration_counts(): + session.run(self._enqueue_ops) + logging.info('Infeed thread finished, shutting down.') + except Exception as e: # pylint: disable=broad-except + self._log_error(session, e) + + def _run_outfeed(self, queue_ctx, session): + logging.info('Starting outfeed thread controller.') + try: + for count, steps in enumerate(queue_ctx.read_iteration_counts()): + for i in xrange(steps): + logging.debug('Outfeed dequeue for iteration (%d, %d)', count, i) + session.run(self._dequeue_ops) + logging.info('Outfeed thread finished, shutting down.') + except Exception as e: # pylint: disable=broad-except + self._log_error(session, e) -class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): - """A Session hook setting up the TPU initialization, infeed, and outfeed. - - This hook does two major things: - 1. initialize and shutdown TPU system. - 2. launch and join the threads for infeed enqueue and (optional) outfeed - dequeue. - """ - - def __init__(self, ctx, enqueue_ops, dequeue_ops=None): - self._master_job = ctx.master_job - self._enqueue_ops = enqueue_ops - self._dequeue_ops = dequeue_ops - self._initial_infeed_sleep_secs = ( - ctx.config.tpu_config.initial_infeed_sleep_secs) - - def begin(self): - logging.info('TPU job name %s', self._master_job) - self._iterations_per_loop_var = _create_or_get_iterations_per_loop() - self._init_op = [tpu.initialize_system(job=self._master_job)] - self._finalize_op = [tpu.shutdown_system(job=self._master_job)] + def _create_infeed_controller(self, name, target, args): + return _OpQueueContext(name=name, target=target, args=args) def after_create_session(self, session, coord): logging.info('Init TPU system') - session.run(self._init_op, - options=config_pb2.RunOptions(timeout_in_ms=5*60*1000)) + session.run(self._init_ops, + options=config_pb2.RunOptions(timeout_in_ms=5 * 60 * 1000)) logging.info('Start infeed thread controller') - self._infeed_thd_controller = _InfeedThreadController( - session, self._enqueue_ops, self._initial_infeed_sleep_secs) + self._infeed_controller = self._create_infeed_controller( + name='InfeedController', target=self._run_infeed, args=(session,)) - if self._dequeue_ops is not None: - logging.info('Start outfeed thread controller') - self._outfeed_thd_controller = _OutfeedThreadController( - session, self._dequeue_ops) + logging.info('Start outfeed thread controller') + self._outfeed_controller = _OpQueueContext( + name='OutfeedController', target=self._run_outfeed, args=(session,)) def before_run(self, run_context): + self._feed_error = None + + # Wait for the cancellation timer to complete before continuing. + if self._session_cancel_timer: + self._session_cancel_timer.join() + self._session_cancel_timer = None + iterations = run_context.session.run(self._iterations_per_loop_var) logging.info('Enqueue next (%d) batch(es) of data to infeed.', iterations) + self._infeed_controller.send_next_batch_signal(iterations) - self._infeed_thd_controller.send_next_batch_signal(iterations) - if self._dequeue_ops is not None: - # TODO(xiejw): Refactor the outfeed dequeue into tf.while_loop. - logging.info( - 'Dequeue next (%d) batch(es) of data from outfeed.', iterations) - self._outfeed_thd_controller.send_next_batch_signal(iterations) + logging.info('Dequeue next (%d) batch(es) of data from outfeed.', + iterations) + self._outfeed_controller.send_next_batch_signal(iterations) def end(self, session): + if self._session_cancel_timer: + logging.warning('Feed error occurred; waiting for message.') + self._session_cancel_timer.join() + + self._finished = True logging.info('Stop infeed thread controller') - self._infeed_thd_controller.join() + self._infeed_controller.join() - if self._dequeue_ops is not None: - logging.info('Stop output thread controller') - self._outfeed_thd_controller.join() + logging.info('Stop output thread controller') + self._outfeed_controller.join() logging.info('Shutdown TPU system.') - session.run(self._finalize_op) + session.run(self._finalize_ops) + + +class TPUInfeedOutfeedSessionHookForPrediction(TPUInfeedOutfeedSessionHook): + + def __init__(self, ctx, enqueue_ops, dequeue_ops): + super(TPUInfeedOutfeedSessionHookForPrediction, self).__init__( + ctx, enqueue_ops, dequeue_ops, run_infeed_loop_on_coordinator=False) + + def _create_infeed_controller(self, name, target, args): + return _OpSignalOnceQueueContext(name=name, target=target, args=args) class _TPUStopAtStepHook(session_run_hook.SessionRunHook): @@ -663,8 +557,8 @@ class _TPUStopAtStepHook(session_run_hook.SessionRunHook): run_context.request_stop() else: iterations = self._next_iterations(global_step, self._last_step) - self._iterations_per_loop_var.load(iterations, - session=run_context.session) + self._iterations_per_loop_var.load( + iterations, session=run_context.session) class _SetEvalIterationsHook(session_run_hook.SessionRunHook): @@ -685,8 +579,49 @@ class _SetEvalIterationsHook(session_run_hook.SessionRunHook): self._iterations_per_loop_var.load(self._num_steps, session=session) -def generate_per_core_enqueue_ops_fn_for_host( - ctx, input_fn, inputs_structure_recorder): +class _StoppingPredictHook(session_run_hook.SessionRunHook): + """Hook that requests stop according to the stopping signal in prediction.""" + + def __init__(self, scalar_stopping_signal): + self._scalar_stopping_signal = scalar_stopping_signal + + def begin(self): + self._iterations_per_loop_var = _create_or_get_iterations_per_loop() + + def after_create_session(self, session, coord): + # This is not necessary as we do not run infeed enqueue and outfeed dequeue + # in side threads for prediction model. But it makes the + # TPUInfeedOutfeedSessionHook prints nice message. + self._iterations_per_loop_var.load(1, session=session) + + def before_run(self, run_context): + return session_run_hook.SessionRunArgs(self._scalar_stopping_signal) + + def after_run(self, run_context, run_values): + _ = run_context + scalar_stopping_signal = run_values.results + if _StopSignals.should_stop(scalar_stopping_signal): + # NOTE(xiejw): In prediction, stopping signals are inserted for each + # batch. And we append one more batch to signal the system it should stop. + # The data flow might look like + # + # batch 0: images, labels, stop = 0 (user provideded) + # batch 1: images, labels, stop = 0 (user provideded) + # ... + # batch 99: images, labels, stop = 0 (user provideded) + # batch 100: images, labels, stop = 1 (TPUEstimator appended) + # + # where the final batch (id = 100) is appended by TPUEstimator, so we + # should drop it before returning the predictions to user. + # To achieve that, we throw the OutOfRangeError in after_run. Once + # Monitored Session sees this error in SessionRunHook.after_run, the + # "current" prediciton, i.e., batch with id=100, will be discarded + # immediately + raise errors.OutOfRangeError(None, None, 'Stopped by stopping signal.') + + +def generate_per_core_enqueue_ops_fn_for_host(ctx, input_fn, + inputs_structure_recorder): """Generates infeed enqueue ops for per-core input_fn on a single host.""" captured_infeed_queue = _CapturedObject() @@ -696,11 +631,14 @@ def generate_per_core_enqueue_ops_fn_for_host( per_host_sharded_inputs = [] for core_ordinal in range(num_cores_per_host): with ops.name_scope('ordinal_%d' % (core_ordinal)): - inputs = input_fn() - if isinstance(inputs, tuple): - features, labels = inputs - else: - features, labels = inputs, None + inputs = _Inputs.from_input_fn(input_fn()) + if inputs.is_dataset: + raise TypeError( + '`input_fn` returning `Dataset` is not yet supported in ' + 'per-Core input pipeline deployment yet. Please set ' + 'TPUConfig.per_host_input_for_training to True or return ' + '`features` and `labels` from `input_fn`') + features, labels = inputs.features_and_labels() inputs_structure_recorder.validate_and_record_structure( features, labels) @@ -716,44 +654,83 @@ def generate_per_core_enqueue_ops_fn_for_host( per_host_sharded_inputs) per_host_enqueue_ops = infeed_queue.generate_enqueue_ops( - per_host_sharded_inputs, - tpu_ordinal_function=ctx.tpu_ordinal_function) + per_host_sharded_inputs, tpu_ordinal_function=ctx.tpu_ordinal_function) return per_host_enqueue_ops + return enqueue_ops_fn, captured_infeed_queue def generate_per_host_enqueue_ops_fn_for_host( - ctx, input_fn, inputs_structure_recorder, batch_axis, device): + ctx, input_fn, inputs_structure_recorder, batch_axis, device, host_id): """Generates infeed enqueue ops for per-host input_fn on a single host.""" captured_infeed_queue = _CapturedObject() + hooks = [] + + with ops.device(device): + inputs = _Inputs.from_input_fn(input_fn()) + + is_dataset = inputs.is_dataset + if ctx.mode == model_fn_lib.ModeKeys.PREDICT: + if not is_dataset: + raise TypeError( + 'For mode PREDICT, `input_fn` must return `Dataset` instead of ' + '`features` and `labels`.') + inputs = _InputsWithStoppingSignals( + dataset=inputs.dataset, batch_size=ctx.batch_size_for_input_fn) + + if is_dataset: + hooks.append(inputs.dataset_initializer_hook()) + + # TODO(ylc): Refactoring the code to merge the tpu ordinal logic here and the + # _TPUContext.tpu_ordinal_function. We should either introduce another + # abstraction or a different helper method. + def _tpu_ordinal_function_impl(shard_index_in_host): + # We put both enqueue/dequeue op at tpu.core(0) in each replica. + replica = ctx.device_assignment.lookup_replicas( + host_id, (0, 0, 0))[shard_index_in_host] + return ctx.device_assignment.tpu_ordinal(replica=replica) + + if ctx.model_parallelism_enabled: + tpu_ordinal_function = _tpu_ordinal_function_impl + else: + tpu_ordinal_function = None + def enqueue_ops_fn(): with ops.device(device): - num_cores_per_host = ctx.num_of_cores_per_host - inputs = input_fn() - if isinstance(inputs, tuple): - features, labels = inputs - else: - features, labels = inputs, None + num_of_replicas_per_host = ctx.num_of_replicas_per_host + # Convert user input to features and labels. If the user returns a + # dataset, it is initialized and the features and labels extracted via + # `dataset.iterator.get_next()` + features, labels = inputs.features_and_labels() + signals = inputs.signals() + inputs_structure_recorder.validate_and_record_structure( - features, labels) + features, labels, signals) unsharded_tensor_list = ( inputs_structure_recorder.flatten_features_and_labels( - features, labels)) + features, labels, signals)) infeed_queue = tpu_feed.InfeedQueue( tuple_types=[t.dtype for t in unsharded_tensor_list], tuple_shapes=[t.shape for t in unsharded_tensor_list], shard_dimensions=batch_axis) captured_infeed_queue.capture(infeed_queue) - infeed_queue.set_number_of_shards(num_cores_per_host) - + infeed_queue.set_number_of_shards(num_of_replicas_per_host) per_host_enqueue_ops = ( infeed_queue.split_inputs_and_generate_enqueue_ops( unsharded_tensor_list, - placement_function=lambda x: device)) - return per_host_enqueue_ops - return enqueue_ops_fn, captured_infeed_queue + placement_function=lambda x: device, + tpu_ordinal_function=tpu_ordinal_function)) + if signals is None: + return per_host_enqueue_ops + else: + return { + 'ops': per_host_enqueue_ops, + 'signals': signals, + } + + return enqueue_ops_fn, captured_infeed_queue, hooks, is_dataset class _InputPipeline(object): @@ -793,6 +770,7 @@ class _InputPipeline(object): self._feature_names = [] self._label_names = [] self._has_labels = False + self._signals_helper = None # Internal state. self._initialized = False @@ -800,8 +778,9 @@ class _InputPipeline(object): def has_labels(self): return self._has_labels - def validate_and_record_structure(self, features, labels): + def validate_and_record_structure(self, features, labels, signals=None): """Validates and records the structure of features` and `labels`.""" + def _extract_key_names(tensor_or_dict): if tensor_or_dict is None: return [] @@ -812,6 +791,10 @@ class _InputPipeline(object): feature_names = _extract_key_names(features) label_names = _extract_key_names(labels) + if signals is not None and self._signals_helper is None: + # Record signals helper. + self._signals_helper = _SignalsHelper(signals) + if self._initialized: # Verify the structure is same. The following should never happen. assert feature_names == self._feature_names, 'feature keys mismatched' @@ -824,13 +807,13 @@ class _InputPipeline(object): self._label_names = label_names self._has_labels = has_labels - def flatten_features_and_labels(self, features, labels): + def flatten_features_and_labels(self, features, labels, signals=None): """Flattens the `features` and `labels` to a single tensor list.""" flattened_inputs = [] if self._feature_names: # We need a fixed ordering for enqueueing and dequeueing. - flattened_inputs.extend([features[name] - for name in self._feature_names]) + flattened_inputs.extend( + [features[name] for name in self._feature_names]) else: flattened_inputs.append(features) @@ -840,6 +823,9 @@ class _InputPipeline(object): flattened_inputs.extend([labels[name] for name in self._label_names]) else: flattened_inputs.append(labels) + + if signals is not None: + flattened_inputs.extend(_SignalsHelper.as_tensor_list(signals)) return flattened_inputs def unflatten_features_and_labels(self, flattened_inputs): @@ -857,15 +843,19 @@ class _InputPipeline(object): ValueError: If the number of expected tensors from `flattened_inputs` mismatches the recorded structure. """ - expected_num_features = (len(self._feature_names) if self._feature_names - else 1) + expected_num_features = ( + len(self._feature_names) if self._feature_names else 1) if self._has_labels: - expected_num_labels = (len(self._label_names) if self._label_names - else 1) + expected_num_labels = ( + len(self._label_names) if self._label_names else 1) else: expected_num_labels = 0 - expected_num_tensors = expected_num_features + expected_num_labels + expected_num_signals = ( + self._signals_helper.num_signals if self._signals_helper else 0) + + expected_num_tensors = ( + expected_num_features + expected_num_labels + expected_num_signals) if expected_num_tensors != len(flattened_inputs): raise ValueError( @@ -882,13 +872,20 @@ class _InputPipeline(object): if expected_num_labels == 0: unflattened_label = None elif self._label_names: - unflattened_label = dict(zip(self._label_names, - flattened_inputs[expected_num_features:])) + label_list = flattened_inputs[ + expected_num_features:expected_num_features + expected_num_labels] + unflattened_label = dict(zip(self._label_names, label_list)) else: # Single tensor case. unflattened_label = flattened_inputs[expected_num_features] - return unflattened_features, unflattened_label + signals = None + if expected_num_signals != 0: + tensor_list_for_signals = flattened_inputs[ + expected_num_features + expected_num_labels:] + signals = self._signals_helper.unflatten(tensor_list_for_signals) + + return _Inputs(unflattened_features, unflattened_label, signals=signals) def __init__(self, input_fn, batch_axis, ctx): """Constructor. @@ -916,25 +913,34 @@ class _InputPipeline(object): # While tf.while_loop is called, the body function, which invokes # `enqueue_fn` passed in, is called to construct the graph. So, input_fn # structure is recorded. - enqueue_ops = self._invoke_input_fn_and_record_structure() + enqueue_ops, all_hooks, run_infeed_loop_on_coordinator = ( + self._invoke_input_fn_and_record_structure()) self._validate_input_pipeline() def dequeue_fn(): """dequeue_fn is used by TPU to retrieve the tensors.""" - values = self._infeed_queue.generate_dequeue_op() + # In the model-parallel case, both the host-side and device-side + # computations must agree on the core on which infeed takes place. We + # choose to perform infeed on logical core 0 of each replica. + with ops.device(tpu.core(0)): + values = self._infeed_queue.generate_dequeue_op() # The unflatten process uses the structure information recorded above. return self._inputs_structure_recorder.unflatten_features_and_labels( values) - return (enqueue_ops, dequeue_fn) + return (enqueue_ops, dequeue_fn, all_hooks, run_infeed_loop_on_coordinator) def _invoke_input_fn_and_record_structure(self): """Deploys the input pipeline and record input structure.""" enqueue_ops = [] infeed_queues = [] + all_hooks = [] num_hosts = self._ctx.num_hosts tpu_host_placement_fn = self._ctx.tpu_host_placement_function + + run_infeed_loop_on_coordinator = True + if self._sharded_per_core: # Per-Core input pipeline deployment. # Invoke input pipeline for each core and placed on the corresponding @@ -948,8 +954,10 @@ class _InputPipeline(object): self._ctx, self._input_fn, self._inputs_structure_recorder)) if _WRAP_INPUT_FN_INTO_WHILE_LOOP: - enqueue_ops.append(_wrap_computation_in_while_loop( - device=host_device, op_fn=enqueue_ops_fn)) + run_infeed_loop_on_coordinator = False + enqueue_ops.append( + _wrap_computation_in_while_loop( + device=host_device, op_fn=enqueue_ops_fn)) else: enqueue_ops.append(enqueue_ops_fn()) # Infeed_queue_getter must be called after enqueue_ops_fn is called. @@ -960,14 +968,32 @@ class _InputPipeline(object): host_device = tpu_host_placement_fn(host_id=host_id) with ops.device(host_device): with ops.name_scope('input_pipeline_task%d' % (host_id)): - enqueue_ops_fn, captured_infeed_queue = ( + enqueue_ops_fn, captured_infeed_queue, hooks, is_dataset = ( generate_per_host_enqueue_ops_fn_for_host( self._ctx, self._input_fn, self._inputs_structure_recorder, - self._batch_axis, host_device)) - - if _WRAP_INPUT_FN_INTO_WHILE_LOOP: - enqueue_ops.append(_wrap_computation_in_while_loop( - device=host_device, op_fn=enqueue_ops_fn)) + self._batch_axis, host_device, host_id)) + all_hooks.extend(hooks) + + # NOTE(xiejw): We dispatch here based on the return type of the + # users `input_fn`. + # + # 1. If input_fn returns a Dataset instance, we initialize the + # iterator outside of tf.while_loop, and call the iterator.get_next + # inside tf.while_loop. This should be always safe. + # + # 2. If input_fn returns (features, labels), it is too late to wrap + # them inside tf.while_loop, as resource initialization cannot be + # handled in TF control flow properly. In this case, we will use + # python loop to enqueue the data into TPU system. This may be + # slow compared to the previous case. + if is_dataset: + run_infeed_loop_on_coordinator = False + wrap_fn = ( + _wrap_computation_in_while_loop + if self._ctx.mode != model_fn_lib.ModeKeys.PREDICT else + _wrap_computation_in_while_loop_with_stopping_signals) + enqueue_ops.append( + wrap_fn(device=host_device, op_fn=enqueue_ops_fn)) else: enqueue_ops.append(enqueue_ops_fn()) infeed_queues.append(captured_infeed_queue.get()) @@ -975,7 +1001,7 @@ class _InputPipeline(object): # dequeue is dtypes and types. So, any one can be used. Here, grab the # first one. self._infeed_queue = infeed_queues[0] - return enqueue_ops + return enqueue_ops, all_hooks, run_infeed_loop_on_coordinator def _validate_input_pipeline(self): # Perform some sanity checks to log user friendly information. We should @@ -1032,15 +1058,18 @@ class _ModelFnWrapper(object): infeed dequeue channel. Returns: - A Fn representing the train step for TPU. + A tuple of train_fn, host_calls, and captured scaffold_fn. The train_fn + representing the train step for TPU. """ + host_call = _OutfeedHostCall(self._ctx) captured_scaffold_fn = _CapturedObject() def train_step(loss): """Training step function for use inside a while loop.""" del loss # unused; required in function signature. - features, labels = dequeue_fn() + inputs = dequeue_fn() + features, labels = inputs.features_and_labels() estimator_spec = self._verify_estimator_spec( self._call_model_fn(features, labels)) @@ -1051,9 +1080,18 @@ class _ModelFnWrapper(object): else: captured_scaffold_fn.capture(None) + # We must run train_op to update the variables prior to running the + # outfeed. with ops.control_dependencies([train_op]): - return array_ops.identity(loss) - return train_step, captured_scaffold_fn + host_call_outfeed_ops = [] + if (isinstance(estimator_spec, TPUEstimatorSpec) and + estimator_spec.host_call is not None): + host_call.record({'host_call': estimator_spec.host_call}) + host_call_outfeed_ops = host_call.create_enqueue_op() + with ops.control_dependencies(host_call_outfeed_ops): + return array_ops.identity(loss) + + return train_step, host_call, captured_scaffold_fn def convert_to_single_tpu_eval_step(self, dequeue_fn): """Converts user provided model_fn` as a single eval step on TPU. @@ -1078,15 +1116,16 @@ class _ModelFnWrapper(object): infeed dequeue channel. Returns: - A tuple of eval_fn and eval_metrics. The eval_fn representing the eval - step for TPU. and eval_metrics is an `_EvalMetrics` instance. + A tuple of eval_fn, host_calls, and captured scaffold_fn. The eval_fn + representing the eval step for TPU. """ - eval_metrics = _EvalMetrics(self._ctx) + host_calls = _OutfeedHostCall(self._ctx) captured_scaffold_fn = _CapturedObject() def eval_step(total_loss): """Evaluation step function for use inside a while loop.""" - features, labels = dequeue_fn() + inputs = dequeue_fn() + features, labels = inputs.features_and_labels() tpu_estimator_spec = self._call_model_fn(features, labels) if not isinstance(tpu_estimator_spec, TPUEstimatorSpec): @@ -1096,14 +1135,68 @@ class _ModelFnWrapper(object): loss = tpu_estimator_spec.loss captured_scaffold_fn.capture(tpu_estimator_spec.scaffold_fn) - eval_metrics.record(tpu_estimator_spec) - outfeed_ops = tpu_ops.outfeed_enqueue_tuple(eval_metrics.outfeed_tensors) - - with ops.control_dependencies([outfeed_ops]): + to_record = {} + to_record['eval_metrics'] = tpu_estimator_spec.eval_metrics + if tpu_estimator_spec.host_call is not None: + # We assume that evaluate won't update global step, so we don't wrap + # this host_call. + to_record['host_call'] = tpu_estimator_spec.host_call + host_calls.record(to_record) + + with ops.control_dependencies(host_calls.create_enqueue_op()): return math_ops.add(total_loss, loss) - return eval_step, eval_metrics, captured_scaffold_fn - def _call_model_fn(self, features, labels): + return eval_step, host_calls, captured_scaffold_fn + + def convert_to_single_tpu_predict_step(self, dequeue_fn): + """Converts user provided model_fn` as a single predict step on TPU. + + Args: + dequeue_fn: The function to retrieve inputs, features and labels, from TPU + infeed dequeue channel. + + Returns: + A tuple of predict_fn, host_calls, and captured scaffold_fn. The + predict_fn representing the predict step for TPU. + """ + host_calls = _OutfeedHostCall(self._ctx) + captured_scaffold_fn = _CapturedObject() + + def predict_step(unused_scalar_stopping_signal): + """Evaluation step function for use inside a while loop.""" + inputs = dequeue_fn() + features, labels = inputs.features_and_labels() + stopping_signals = inputs.signals() + + assert stopping_signals is not None, ( + 'Internal Error: `signals` is missing.') + + tpu_estimator_spec = self._call_model_fn( + features, labels, is_export_mode=False) + if not isinstance(tpu_estimator_spec, TPUEstimatorSpec): + raise RuntimeError( + 'estimator_spec used by TPU prediction must have type' + '`TPUEstimatorSpec`. Got {}'.format(type(tpu_estimator_spec))) + + captured_scaffold_fn.capture(tpu_estimator_spec.scaffold_fn) + to_record = {} + identity_fn = lambda **kwargs: kwargs + # TODO(xiejw): Adds validation for prediction dictionrary. + # TODO(xiejw): Adds support for single tensor as predictions. + if not isinstance(tpu_estimator_spec.predictions, dict): + raise TypeError('TPUEstimatorSpec.predictions must be dict of Tensors.') + to_record['predictions'] = [identity_fn, tpu_estimator_spec.predictions] + to_record['signals'] = [identity_fn, stopping_signals] + if tpu_estimator_spec.host_call is not None: + to_record['host_call'] = tpu_estimator_spec.host_call + host_calls.record(to_record) + + with ops.control_dependencies(host_calls.create_enqueue_op()): + return _StopSignals.as_scalar_stopping_signal(stopping_signals) + + return predict_step, host_calls, captured_scaffold_fn + + def _call_model_fn(self, features, labels, is_export_mode=True): """Calls the model_fn with required parameters.""" model_fn_args = util.fn_args(self._model_fn) kwargs = {} @@ -1125,17 +1218,16 @@ class _ModelFnWrapper(object): kwargs['params'] = params if 'params' not in model_fn_args: - raise ValueError( - 'model_fn ({}) does not include params argument, ' - 'required by TPUEstimator to pass batch size as ' - 'params[\'batch_size\']'.format(self._model_fn)) + raise ValueError('model_fn ({}) does not include params argument, ' + 'required by TPUEstimator to pass batch size as ' + 'params[\'batch_size\']'.format(self._model_fn)) batch_size_for_model_fn = self._ctx.batch_size_for_model_fn if batch_size_for_model_fn is not None: params[_BATCH_SIZE_KEY] = batch_size_for_model_fn estimator_spec = self._model_fn(features=features, **kwargs) - if (self._ctx.is_running_on_cpu() and + if (self._ctx.is_running_on_cpu(is_export_mode) and isinstance(estimator_spec, TPUEstimatorSpec)): # The estimator_spec will be passed to `Estimator` directly, which expects # type `EstimatorSpec`. @@ -1162,158 +1254,255 @@ class _ModelFnWrapper(object): return estimator_spec -class _EvalMetrics(object): - """Class wraps TPUEstimator.eval_metrics.""" +class _OutfeedHostCall(object): + """Support for `eval_metrics` and `host_call` in TPUEstimatorSpec.""" def __init__(self, ctx): self._ctx = ctx - self._metric_fn = None - self._is_dict = False - self._tensor_keys = [] - self._tensors = [] - self._tensor_dtypes = [] - self._tensor_shapes = [] - self._recorded = False + self._names = [] + # All of these are dictionaries of lists keyed on the name. + self._host_fns = {} + self._tensor_keys = collections.defaultdict(list) + self._tensors = collections.defaultdict(list) + self._tensor_dtypes = collections.defaultdict(list) + self._tensor_shapes = collections.defaultdict(list) @staticmethod - def validate(eval_metrics): - """Validates the `eval_metrics` in `TPUEstimatorSpec`.""" - - if not isinstance(eval_metrics, (tuple, list)): - raise ValueError('eval_metrics should be tuple or list') - if len(eval_metrics) != 2: - raise ValueError('eval_metrics should have two elements.') - if not callable(eval_metrics[0]): - raise TypeError('eval_metrics[0] should be callable.') - if not isinstance(eval_metrics[1], (tuple, list, dict)): - raise ValueError('eval_metrics[1] should be tuple or list, or dict.') - - if isinstance(eval_metrics[1], (tuple, list)): - fn_args = util.fn_args(eval_metrics[0]) - if len(eval_metrics[1]) != len(fn_args): - raise RuntimeError( - 'In TPUEstimatorSpec.eval_metrics, length of tensors does not ' - 'match method args of metric_fn.') + def validate(host_calls): + """Validates the `eval_metrics` and `host_call` in `TPUEstimatorSpec`.""" + + for name, host_call in host_calls.items(): + if not isinstance(host_call, (tuple, list)): + raise ValueError('{} should be tuple or list'.format(name)) + if len(host_call) != 2: + raise ValueError('{} should have two elements.'.format(name)) + if not callable(host_call[0]): + raise TypeError('{}[0] should be callable.'.format(name)) + if not isinstance(host_call[1], (tuple, list, dict)): + raise ValueError('{}[1] should be tuple or list, or dict.'.format(name)) + + if isinstance(host_call[1], (tuple, list)): + fullargspec = tf_inspect.getfullargspec(host_call[0]) + fn_args = util.fn_args(host_call[0]) + # wrapped_hostcall_with_global_step uses varargs, so we allow that. + if fullargspec.varargs is None and len(host_call[1]) != len(fn_args): + raise RuntimeError( + 'In TPUEstimatorSpec.{}, length of tensors {} does not match ' + 'method args of the function, which takes {}.'.format( + name, len(host_call[1]), len(fn_args))) @staticmethod - def to_metric_metric_ops_for_cpu(eval_metrics): - """Converts `TPUEstimatorSpec.eval_metrics` to `eval_metric_ops` for CPU.""" - if not eval_metrics: - return None - - _EvalMetrics.validate(eval_metrics) + def create_cpu_hostcall(host_calls): + """Runs on the host_call on CPU instead of TPU when use_tpu=False.""" + + _OutfeedHostCall.validate(host_calls) + ret = {} + for name, host_call in host_calls.items(): + host_fn, tensors = host_call + if isinstance(tensors, (tuple, list)): + ret[name] = host_fn(*tensors) + else: + # Must be dict. + try: + ret[name] = host_fn(**tensors) + except TypeError as e: + logging.warning( + 'Exception while calling %s: %s. It is likely the tensors ' + '(%s[1]) do not match the ' + 'function\'s arguments', name, e, name) + raise e + return ret + + def record(self, host_calls): + """Records the host_call structure.""" + + for name, host_call in host_calls.items(): + host_fn, tensor_list_or_dict = host_call + self._names.append(name) + self._host_fns[name] = host_fn + + if isinstance(tensor_list_or_dict, dict): + for (key, tensor) in six.iteritems(tensor_list_or_dict): + self._tensor_keys[name].append(key) + self._tensors[name].append(tensor) + self._tensor_dtypes[name].append(tensor.dtype) + self._tensor_shapes[name].append(tensor.shape) + else: + # List or tuple. + self._tensor_keys[name] = None + for tensor in tensor_list_or_dict: + self._tensors[name].append(tensor) + self._tensor_dtypes[name].append(tensor.dtype) + self._tensor_shapes[name].append(tensor.shape) - metric_fn, tensors = eval_metrics + def create_enqueue_op(self): + """Create the op to enqueue the recorded host_calls. - if isinstance(tensors, (tuple, list)): - return metric_fn(*tensors) - else: - # Must be dict. - try: - return metric_fn(**tensors) - except TypeError as e: - logging.warning( - 'Exception while calling metric_fn for evalution: %s. ' - 'It is likely the tensors (eval_metrics[1]) do not match the ' - 'metric_fn arguments', e) - raise e - - def record(self, spec): - """Records the eval_metrics structure in `spec`.""" - if self._recorded: - raise RuntimeError('Eval metrics have been recorded already.') - - self._metric_fn, tensor_list_or_dict = spec.eval_metrics - - if isinstance(tensor_list_or_dict, dict): - self._is_dict = True - for (key, tensor) in six.iteritems(tensor_list_or_dict): - self._tensor_keys.append(key) - self._tensors.append(tensor) - self._tensor_dtypes.append(tensor.dtype) - self._tensor_shapes.append(tensor.shape) - else: - # List or tuple. - self._is_dict = False - self._tensors = tensor_list_or_dict - for tensor in tensor_list_or_dict: - self._tensor_dtypes.append(tensor.dtype) - self._tensor_shapes.append(tensor.shape) - self._recorded = True + Returns: + A list of enqueue ops, which is empty if there are no host calls. + """ + if not self._names: + return [] - @property - def outfeed_tensors(self): - if not self._recorded: - raise RuntimeError('Eval metrics have not been recorded yet') - return self._tensors + tensors = [] + # TODO(jhseu): Consider deduping tensors. + for name in self._names: + tensors.extend(self._tensors[name]) - def to_metric_metric_ops_for_tpu(self, dummy_update_op): - """Creates the eval_metric_ops now based on the TPU outfeed. + with ops.device(tpu.core(0)): + return [tpu_ops.outfeed_enqueue_tuple(tensors)] - `eval_metric_ops` is defined in `EstimatorSpec`. From all shards, tensors - are dequeued from outfeed and then concatenated (along batch size dimension) - to form global-like tensors. All global-like tensors are passed to the - metric fn. + def create_tpu_hostcall(self): + """Sends the tensors through outfeed and runs the host_fn on CPU. - Args: - dummy_update_op: A dummy update op. + The tensors are concatenated along dimension 0 to form a global tensor + across all shards. The concatenated function is passed to the host_fn and + executed on the first host. Returns: - A tuple of (`eval_metric_ops` and `update_ops`), where `update_ops` should - be invoked in Outfeed dequeue thread, which drive the outfeed dequeue and - update the state of metrics. + A dictionary mapping name to the return type of the host_call by that + name. Raises: RuntimeError: If outfeed tensor is scalar. """ + if not self._names: + return [] - num_cores = self._ctx.num_cores - + ret = {} # For each i, dequeue_ops[i] is a list containing the tensors from all # shards. This list is concatenated later. dequeue_ops = [] - for i in xrange(len(self._tensors)): - dequeue_ops.append([]) - - # Outfeed ops execute on each JF node. + tensor_dtypes = [] + tensor_shapes = [] + for name in self._names: + for _ in self._tensors[name]: + dequeue_ops.append([]) + for dtype in self._tensor_dtypes[name]: + tensor_dtypes.append(dtype) + for shape in self._tensor_shapes[name]: + tensor_shapes.append(shape) + + # Outfeed ops execute on each replica's first logical core. Note: we must + # constraint it such that we have at most one outfeed dequeue and enqueue + # per replica. tpu_device_placement_fn = self._ctx.tpu_device_placement_function - for i in xrange(num_cores): + for i in xrange(self._ctx.num_replicas): with ops.device(tpu_device_placement_fn(i)): outfeed_tensors = tpu_ops.outfeed_dequeue_tuple( - dtypes=self._tensor_dtypes, shapes=self._tensor_shapes) + dtypes=tensor_dtypes, shapes=tensor_shapes) for j, item in enumerate(outfeed_tensors): dequeue_ops[j].append(item) - # It is assumed evaluation always happends on single host TPU system. So, + # Deconstruct dequeue ops. + dequeue_ops_by_name = {} + pos = 0 + for name in self._names: + dequeue_ops_by_name[name] = dequeue_ops[pos:pos+len(self._tensors[name])] + pos += len(self._tensors[name]) + + # It is assumed evaluation always happens on single host TPU system. So, # place all ops on tpu host if possible. + # + # TODO(jhseu): Evaluate whether this is right for summaries. with ops.device(self._ctx.tpu_host_placement_function(core_id=0)): - for i, item in enumerate(dequeue_ops): - if dequeue_ops[i][0].shape.ndims == 0: - raise RuntimeError( - 'All tensors outfed from TPU should preseve batch size ' - 'dimension, but got scalar {}'.format(dequeue_ops[i][0])) - # TODO(xiejw): Allow users to specify the axis for batch size dimension. - dequeue_ops[i] = array_ops.concat(dequeue_ops[i], axis=0) + for name in self._names: + dequeue_ops = dequeue_ops_by_name[name] + for i, item in enumerate(dequeue_ops): + if dequeue_ops[i][0].shape.ndims == 0: + raise RuntimeError( + 'All tensors outfed from TPU should preserve batch size ' + 'dimension, but got scalar {}'.format(dequeue_ops[i][0])) + # TODO(xiejw): Allow users to specify the axis for batch size + # dimension. + dequeue_ops[i] = array_ops.concat(dequeue_ops[i], axis=0) + + if self._tensor_keys[name] is not None: + # The user-provided eval_metrics[1] is a dict. + dequeue_ops = dict(zip(self._tensor_keys[name], dequeue_ops)) + try: + ret[name] = self._host_fns[name](**dequeue_ops) + except TypeError as e: + logging.warning( + 'Exception while calling %s: %s. It is likely the tensors ' + '(%s[1]) do not match the ' + 'function\'s arguments', name, e, name) + raise e + else: + ret[name] = self._host_fns[name](*dequeue_ops) - if self._is_dict: - dequeue_ops = dict(zip(self._tensor_keys, dequeue_ops)) - try: - eval_metric_ops = self._metric_fn(**dequeue_ops) - except TypeError as e: - logging.warning( - 'Exception while calling metric_fn for evalution: %s. ' - 'It is likely the tensors (eval_metrics[1]) do not match the ' - 'metric_fn arguments', e) - raise e - else: - eval_metric_ops = self._metric_fn(*dequeue_ops) + return ret + + +class _OutfeedHostCallHook(session_run_hook.SessionRunHook): + """Hook to run host calls when use_tpu=False.""" + + def __init__(self, tensors): + self._tensors = tensors + + def begin(self): + # We duplicate this code from the TPUInfeedOutfeedSessionHook rather than + # create a separate hook to guarantee execution order, because summaries + # need to be initialized before the outfeed thread starts. + # TODO(jhseu): Make a wrapper hook instead? + self._init_ops = contrib_summary.summary_writer_initializer_op() + # Get all the writer resources from the initializer, so we know what to + # flush. + self._finalize_ops = [] + for op in self._init_ops: + self._finalize_ops.append(contrib_summary.flush(writer=op.inputs[0])) + + def after_create_session(self, session, coord): + session.run(self._init_ops) - eval_update_ops = [] - for k, v in eval_metric_ops.items(): - eval_metric_ops[k] = (v[0], dummy_update_op) - eval_update_ops.append(v[1]) + def before_run(self, run_context): + return basic_session_run_hooks.SessionRunArgs(self._tensors) + + def end(self, session): + session.run(self._finalize_ops) + + +class ExamplesPerSecondHook(basic_session_run_hooks.StepCounterHook): + """Count examples during runtime.""" + + def __init__(self, + batch_size, + every_n_steps=100, + every_n_secs=None, + output_dir=None, + summary_writer=None): + self._batch_size = batch_size + super(ExamplesPerSecondHook, self).__init__( + every_n_steps=every_n_steps, + every_n_secs=every_n_secs, + output_dir=output_dir, + summary_writer=summary_writer) + + def _log_and_record(self, elapsed_steps, elapsed_time, global_step): + examples_per_sec = self._batch_size * elapsed_steps / elapsed_time + if self._summary_writer is not None: + example_summary = Summary(value=[ + Summary.Value(tag='examples_sec', simple_value=examples_per_sec) + ]) + self._summary_writer.add_summary(example_summary, global_step) + logging.info('examples/sec: %g', examples_per_sec) + + +class InstallSignalHandlerHook(session_run_hook.SessionRunHook): + """Change SIGINT (CTRL^C) handler to force quit the process. + + The default behavior often results in hanging processes. + The original handler is restored after training/evaluation. + """ + + def __init__(self): + self._signal_fn = signal.getsignal(signal.SIGINT) + + def before_run(self, run_context): + signal.signal(signal.SIGINT, signal.SIG_DFL) - return eval_metric_ops, eval_update_ops + def end(self, session): + signal.signal(signal.SIGINT, self._signal_fn) class TPUEstimator(estimator_lib.Estimator): @@ -1323,30 +1512,28 @@ class TPUEstimator(estimator_lib.Estimator): replicating inputs and models for each core, and returning to host periodically to run hooks. - If `use_tpu` is false, all training, evaluation, and predict are executed on - CPU. - - For training, TPUEstimator transforms a global batch size in params to a - per-shard batch size when calling the `input_fn` and `model_fn`. Users should - specify `train_batch_size` in constructor, and then get the batch size for - each shard in `input_fn` and `model_fn` by `params['batch_size']`. If - `TPUConfig.per_host_input_for_training` is `True`, `input_fn` is invoked per - host rather than per core. In this case, a global batch size is transformed a - per-host batch size in params for `input_fn`, but `model_fn` still gets - per-core batch size. - - For evaluation, if `eval_batch_size` is None, it is executed on CPU, even if - `use_tpu` is `True`. If `eval_batch_size` is not `None`, it is executed on - TPU, which is an experimental feature. In this case, `model_fn` should return - `TPUEstimatorSpec` instead of `EstimatorSpec`, which expects the - `eval_metrics` for TPU evaluation. - + TPUEstimator transforms a global batch size in params to a per-shard batch + size when calling the `input_fn` and `model_fn`. Users should specify + global batch size in constructor, and then get the batch size for each shard + in `input_fn` and `model_fn` by `params['batch_size']`. + For training, `model_fn` gets per-core batch size; `input_fn` may get + per-core or per-host batch size depending on + `per_host_input_for_training` in `TPUConfig`. + For evaluation, `model_fn` gets per-core batch size and `input_fn` get + per-host batch size. + + `model_fn` should return `TPUEstimatorSpec`, which expects the `eval_metrics` + for TPU evaluation. `TPUEstimatorSpec.eval_metrics` is a tuple of `metric_fn` and `tensors`, where `tensors` could be a list of `Tensor`s or dict of names to `Tensor`s. (See `TPUEstimatorSpec` for details). `metric_fn` takes the `tensors` and returns a dict from metric string name to the result of calling a metric function, namely a `(metric_tensor, update_op)` tuple. + One can set `use_tpu` to `False` for testing. All training, evaluation, and + predict will be executed on CPU. `input_fn` and `model_fn` will receive + `train_batch_size` or `eval_batch_size` unmodified as `params['batch_size']`. + Current limitations: 1. TPU evaluation only works on single host. @@ -1401,6 +1588,7 @@ class TPUEstimator(estimator_lib.Estimator): use_tpu=True, train_batch_size=None, eval_batch_size=None, + predict_batch_size=None, batch_axis=None): """Constructs an `TPUEstimator` instance. @@ -1419,18 +1607,17 @@ class TPUEstimator(estimator_lib.Estimator): basic python types. There are reserved keys for `TPUEstimator`, including 'batch_size'. use_tpu: A bool indicating whether TPU support is enabled. Currently, - - TPU training respects this bit. - - If true, see `eval_batch_size` for evaluate support. + - TPU training and evaluation respect this bit. - Predict still happens on CPU. train_batch_size: An int representing the global training batch size. TPUEstimator transforms this global batch size to a per-shard batch size, as params['batch_size'], when calling `input_fn` and `model_fn`. - Cannot be `None` if `use_tpu` is `True`. Must be divisible by - `config.tpu_config.num_shards`. - eval_batch_size: An int representing the global training batch size. - Currently, if `None`, evaluation is still executed on CPU (even when - `use_tpu` is True). In near future, `use_tpu` will be the only option to - switch between TPU/CPU evaluation. + Cannot be `None` if `use_tpu` is `True`. + Must be divisible by total number of replicas. + eval_batch_size: An int representing evaluation batch size. + Must be divisible by total number of replicas. + predict_batch_size: An int representing the prediction batch size. + Must be divisible by total number of replicas. batch_axis: A python tuple of int values describing how each tensor produced by the Estimator `input_fn` should be split across the TPU compute shards. For example, if your input_fn produced (images, labels) @@ -1450,35 +1637,28 @@ class TPUEstimator(estimator_lib.Estimator): '`config` must be provided with type `tpu_config.RunConfig`') if params is not None and any(k in params for k in _RESERVED_PARAMS_KEYS): - raise ValueError( - '{} are reserved keys but existed in params {}.'.format( - _RESERVED_PARAMS_KEYS, params)) + raise ValueError('{} are reserved keys but existed in params {}.'.format( + _RESERVED_PARAMS_KEYS, params)) if use_tpu: + # Perform some very basic validations. More validations will be found in + # _TPUContext. if train_batch_size is None: raise ValueError('`train_batch_size` cannot be `None`') - if not isinstance(train_batch_size, int): - raise ValueError('`train_batch_size` must be an int') - if train_batch_size < 1: - raise ValueError('`train_batch_size` must be positive') - - # The specified batch size is the batch size for the entire computation. - # The input_fn and model_fn are called per-shard, so we want to calculate - # the per-shard batch size and pass that. - if train_batch_size % config.tpu_config.num_shards != 0: + util_lib.check_positive_integer(train_batch_size, 'train_batch_size') + + if (not config.tpu_config.per_host_input_for_training and + config.tpu_config.computation_shape): raise ValueError( - 'train batch size {} must be divisible by number of shards {}' - .format(train_batch_size, config.tpu_config.num_shards)) + 'Model parallelism only supports per host input for training. ' + 'Please adjust TPURunconfig.per_host_input_for_training.') if eval_batch_size is not None: - if config.tpu_config.num_shards > 8: - raise NotImplementedError( - 'TPU evaluation is only supported with one host.') + util_lib.check_positive_integer(eval_batch_size, 'eval_batch_size') - if eval_batch_size % config.tpu_config.num_shards != 0: - raise ValueError( - 'eval batch size {} must be divisible by number of shards {}' - .format(eval_batch_size, config.tpu_config.num_shards)) + if predict_batch_size is not None: + util_lib.check_positive_integer(predict_batch_size, + 'predict_batch_size') # Verifies the model_fn signature according to Estimator framework. estimator_lib._verify_model_fn_args(model_fn, params) # pylint: disable=protected-access @@ -1498,8 +1678,11 @@ class TPUEstimator(estimator_lib.Estimator): self._config.tpu_config.iterations_per_loop) # All properties passed to _TPUContext are immutable. - self._ctx = _TPUContext(self._config, train_batch_size, eval_batch_size, - use_tpu) + # pylint: disable=protected-access + self._ctx = tpu_context._get_tpu_context( + self._config, train_batch_size, + eval_batch_size, predict_batch_size, + use_tpu) def _create_global_step(self, graph): """Creates a global step suitable for TPUs. @@ -1533,8 +1716,9 @@ class TPUEstimator(estimator_lib.Estimator): if max_steps is not None: util_lib.check_positive_integer(max_steps, 'Train max_steps') - return [_TPUStopAtStepHook(self._iterations_per_training_loop, - steps, max_steps)] + return [ + _TPUStopAtStepHook(self._iterations_per_training_loop, steps, max_steps) + ] def _convert_eval_steps_to_hooks(self, steps): with self._ctx.with_mode(model_fn_lib.ModeKeys.EVAL) as ctx: @@ -1546,11 +1730,11 @@ class TPUEstimator(estimator_lib.Estimator): util_lib.check_positive_integer(steps, 'Eval steps') - hooks = [] - hooks.append(evaluation._StopAfterNEvalsHook( # pylint: disable=protected-access - num_evals=steps)) - hooks.append(_SetEvalIterationsHook(steps)) - return hooks + return [ + evaluation._StopAfterNEvalsHook( # pylint: disable=protected-access + num_evals=steps), + _SetEvalIterationsHook(steps) + ] def _call_input_fn(self, input_fn, mode): """Calls the input function. @@ -1586,7 +1770,9 @@ class TPUEstimator(estimator_lib.Estimator): if batch_size_for_input_fn is not None: kwargs['params'][_BATCH_SIZE_KEY] = batch_size_for_input_fn - if ctx.is_running_on_cpu(): + # For export_savedmodel, input_fn is never passed to Estimator. So, + # `is_export_mode` must be False. + if ctx.is_running_on_cpu(is_export_mode=False): with ops.device('/device:CPU:0'): return input_fn(**kwargs) @@ -1602,6 +1788,7 @@ class TPUEstimator(estimator_lib.Estimator): # `features` in `model_fn` signature. def _input_fn(): return input_fn(**kwargs) + return _input_fn def _augment_model_fn(self, model_fn, batch_axis): @@ -1612,8 +1799,13 @@ class TPUEstimator(estimator_lib.Estimator): with self._ctx.with_mode(mode) as ctx: model_fn_wrapper = _ModelFnWrapper(model_fn, config, params, ctx) - # TODO(jhseu): Move to PREDICT to TPU. - if ctx.is_running_on_cpu(): + # For export_savedmodel, input_fn is never passed to Estimator. So, + # if features is callable, it means it is the input_fn passed by + # TPUEstimator._call_input_fn. Then we can know if the mode == PREDICT, + # it implies, it is the .predict API, not export_savedmodel API. + is_export_mode = not callable(features) + + if ctx.is_running_on_cpu(is_export_mode=is_export_mode): logging.info('Running %s on CPU', mode) return model_fn_wrapper.call_without_tpu(features, labels) @@ -1623,19 +1815,31 @@ class TPUEstimator(estimator_lib.Estimator): input_fn = features input_holders = _InputPipeline(input_fn, batch_axis, ctx) - enqueue_ops, dequeue_fn = ( + enqueue_ops, dequeue_fn, input_hooks, run_infeed_loop_on_coordinator = ( input_holders.generate_infeed_enqueue_ops_and_dequeue_fn()) if mode == model_fn_lib.ModeKeys.TRAIN: - loss, scaffold = ( + loss, host_call, scaffold = ( _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn)) + host_ops = host_call.create_tpu_hostcall() + if host_ops is None: + host_ops = [] hooks = [ - TPUInfeedOutfeedSessionHook(ctx, enqueue_ops), + TPUInfeedOutfeedSessionHook( + ctx, + enqueue_ops, + host_ops, + run_infeed_loop_on_coordinator=( + run_infeed_loop_on_coordinator)), + ExamplesPerSecondHook(ctx.global_batch_size), + InstallSignalHandlerHook(), training.LoggingTensorHook( - {'loss': array_ops.identity(loss), - 'step': training.get_global_step()}, + { + 'loss': array_ops.identity(loss), + 'step': training.get_global_step() + }, every_n_secs=30) - ] + ] + input_hooks summary.scalar(model_fn_lib.LOSS_METRIC_KEY, loss) with ops.control_dependencies([loss]): update_ops = _sync_variables_ops() @@ -1650,92 +1854,198 @@ class TPUEstimator(estimator_lib.Estimator): train_op=control_flow_ops.group(*update_ops), scaffold=scaffold) - # Now eval. - total_loss, eval_metric_ops, scaffold = _eval_on_tpu_system( + if mode == model_fn_lib.ModeKeys.EVAL: + total_loss, host_calls, scaffold = _eval_on_tpu_system( + ctx, model_fn_wrapper, dequeue_fn) + iterations_per_loop_var = _create_or_get_iterations_per_loop() + mean_loss = math_ops.div(total_loss, + math_ops.cast( + iterations_per_loop_var, + dtype=total_loss.dtype)) + + # Creates a dummy metric update_op for all metrics. Estimator expects + # all metrics in eval_metric_ops have update_op and calls them one by + # one. The real metric update_ops are invoked in a separated thread. + # So, here give Estimator the dummy op for all metrics. + with ops.control_dependencies([mean_loss]): + # After TPU evaluation computation is done (the mean_loss tensor), + # reads all variables back from TPU and updates the eval step + # counter properly + internal_ops_to_run = _sync_variables_ops() + internal_ops_to_run.append( + _increase_eval_step_op(iterations_per_loop_var)) + with ops.control_dependencies(internal_ops_to_run): + dummy_update_op = control_flow_ops.no_op() + + host_call_ret = host_calls.create_tpu_hostcall() + eval_metric_ops = {} + eval_update_ops = [] + for k, v in host_call_ret['eval_metrics'].items(): + eval_metric_ops[k] = (v[0], dummy_update_op) + eval_update_ops.append(v[1]) + + if 'host_call' not in host_call_ret: + host_ops = [] + else: + host_ops = host_call_ret['host_call'] + hooks = [ + TPUInfeedOutfeedSessionHook( + ctx, + enqueue_ops, + eval_update_ops + host_ops, + run_infeed_loop_on_coordinator=( + run_infeed_loop_on_coordinator)), + ] + input_hooks + + return model_fn_lib.EstimatorSpec( + mode, + loss=mean_loss, + evaluation_hooks=hooks, + eval_metric_ops=eval_metric_ops, + scaffold=scaffold) + + # Predict + assert mode == model_fn_lib.ModeKeys.PREDICT + + dummy_predict_op, host_calls, scaffold = _predict_on_tpu_system( ctx, model_fn_wrapper, dequeue_fn) - iterations_per_loop_var = _create_or_get_iterations_per_loop() - mean_loss = math_ops.div( - total_loss, - math_ops.cast(iterations_per_loop_var, dtype=total_loss.dtype)) - - # Creates a dummy metric update_op for all metrics. Estimator expects - # all metrics in eval_metric_ops have update_op and calls them one by - # one. The real metric update_ops are invoked in a separated thread. So, - # here give Estimator the dummy op for all metrics. - with ops.control_dependencies([mean_loss]): - # After TPU evaluation computation is done (the mean_loss tensor), - # reads all variables back from TPU and updates the eval step counter - # properly + with ops.control_dependencies([dummy_predict_op]): internal_ops_to_run = _sync_variables_ops() - internal_ops_to_run.append( - _increase_eval_step_op(iterations_per_loop_var)) with ops.control_dependencies(internal_ops_to_run): - dummy_update_op = control_flow_ops.no_op() + dummy_predict_op = control_flow_ops.no_op() + + # In train and evaluation, the main TPU program is passed to monitored + # training session to run. Infeed enqueue and outfeed dequeue are + # executed in side threads. This is not the configuration for + # prediction mode. + # + # For prediction, the Estimator executes the EstimatorSpec.predictions + # directly and yield the element (via generator) to call site. So, the + # outfeed based prediction must be passed to MonitoredSession directly. + # Other parts of the TPU execution are organized as follows. + # + # 1. All outfeed based Tensors must be grouped with predictions Tensors + # to form a single invocation. This avoid the issue we might trigger + # multiple outfeeds incorrectly. To achieve this, `host_call` is + # placed in control_dependencies of `stopping_signals`, and + # `stopping_signals` is passed into _StoppingPredictHook, which sets + # the `stopping_signals` as SessionRunArgs. MonitoredSession merges + # all SessionRunArgs with the fetch in session.run together. + # + # 2. The TPU program (dummy_predict_op) and enqueue_ops (infeed Enqueue) + # are grouped together. They will be launched once and only once in + # side threads and they quit naturally according to the SAME stopping + # condition. + enqueue_ops.append(dummy_predict_op) + + host_call_ret = host_calls.create_tpu_hostcall() + if 'host_call' not in host_call_ret: + host_ops = [] + else: + host_ops = host_call_ret['host_call'] + + predictions = host_call_ret['predictions'] + stopping_signals = host_call_ret['signals'] + + with ops.control_dependencies(host_ops): + host_ops = [] # Empty, we do do not need it anymore. + scalar_stopping_signal = _StopSignals.as_scalar_stopping_signal( + stopping_signals) - eval_metric_ops, eval_update_ops = ( - eval_metric_ops.to_metric_metric_ops_for_tpu(dummy_update_op)) hooks = [ - TPUInfeedOutfeedSessionHook(ctx, enqueue_ops, eval_update_ops), - ] + _StoppingPredictHook(scalar_stopping_signal), + TPUInfeedOutfeedSessionHookForPrediction(ctx, enqueue_ops, + host_ops), + ] + input_hooks return model_fn_lib.EstimatorSpec( mode, - loss=mean_loss, - evaluation_hooks=hooks, - eval_metric_ops=eval_metric_ops, + prediction_hooks=hooks, + predictions=predictions, scaffold=scaffold) + return _model_fn def _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): """Executes `model_fn_wrapper` multiple times on all TPU shards.""" - num_cores = ctx.num_cores iterations_per_loop_var = _create_or_get_iterations_per_loop() - single_tpu_eval_step, eval_metric_ops, captured_scaffold_fn = ( + single_tpu_eval_step, host_calls, captured_scaffold_fn = ( model_fn_wrapper.convert_to_single_tpu_eval_step(dequeue_fn)) def multi_tpu_eval_steps_on_single_shard(): - return training_loop.repeat(iterations_per_loop_var, - single_tpu_eval_step, - [_ZERO_LOSS], - name='loop') + return training_loop.repeat( + iterations_per_loop_var, + single_tpu_eval_step, [_ZERO_LOSS], + name='loop') - (loss,) = tpu.shard(multi_tpu_eval_steps_on_single_shard, - inputs=[], - num_shards=num_cores, - outputs_from_all_shards=False) + (loss,) = tpu.shard( + multi_tpu_eval_steps_on_single_shard, + inputs=[], + num_shards=ctx.num_replicas, + outputs_from_all_shards=False, + device_assignment=ctx.device_assignment) scaffold = _get_scaffold(captured_scaffold_fn) - return loss, eval_metric_ops, scaffold + return loss, host_calls, scaffold def _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): """Executes `model_fn_wrapper` multiple times on all TPU shards.""" - num_cores = ctx.num_cores iterations_per_loop_var = _create_or_get_iterations_per_loop() - single_tpu_train_step, captured_scaffold_fn = ( + single_tpu_train_step, host_call, captured_scaffold_fn = ( model_fn_wrapper.convert_to_single_tpu_train_step(dequeue_fn)) def multi_tpu_train_steps_on_single_shard(): return training_loop.repeat( iterations_per_loop_var, - single_tpu_train_step, - [_INITIAL_LOSS], + single_tpu_train_step, [_INITIAL_LOSS], name=b'loop') - (loss,) = tpu.shard(multi_tpu_train_steps_on_single_shard, - inputs=[], - num_shards=num_cores, - outputs_from_all_shards=False) + (loss,) = tpu.shard( + multi_tpu_train_steps_on_single_shard, + inputs=[], + num_shards=ctx.num_replicas, + outputs_from_all_shards=False, + device_assignment=ctx.device_assignment) + + scaffold = _get_scaffold(captured_scaffold_fn) + return loss, host_call, scaffold + + +def _predict_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): + """Executes `model_fn_wrapper` multiple times on all TPU shards.""" + num_cores = ctx.num_cores + + single_tpu_predict_step, host_calls, captured_scaffold_fn = ( + model_fn_wrapper.convert_to_single_tpu_predict_step(dequeue_fn)) + + def multi_tpu_predict_steps_on_single_shard(): + + def cond(scalar_stopping_signal): + return math_ops.logical_not( + _StopSignals.should_stop(scalar_stopping_signal)) + + inputs = [_StopSignals.NON_STOPPING_SIGNAL] + outputs = training_loop.while_loop( + cond, single_tpu_predict_step, inputs=inputs, name=b'loop') + return outputs + + (dummy_predict_op,) = tpu.shard( + multi_tpu_predict_steps_on_single_shard, + inputs=[], + num_shards=num_cores, + outputs_from_all_shards=False) scaffold = _get_scaffold(captured_scaffold_fn) - return loss, scaffold + return dummy_predict_op, host_calls, scaffold def _wrap_computation_in_while_loop(device, op_fn): """Wraps the ops generated by `op_fn` in tf.while_loop.""" + def computation(i): with ops.control_dependencies(op_fn()): return i + 1 @@ -1747,7 +2057,31 @@ def _wrap_computation_in_while_loop(device, op_fn): iterations = array_ops.identity(iterations_per_loop_var) return control_flow_ops.while_loop( lambda i: i < iterations, - computation, [constant_op.constant(0)], parallel_iterations=1) + computation, [constant_op.constant(0)], + parallel_iterations=1) + + +def _wrap_computation_in_while_loop_with_stopping_signals(device, op_fn): + """Wraps the ops generated by `op_fn` in tf.while_loop.""" + + def cond(scalar_stopping_signal): + return math_ops.logical_not( + _StopSignals.should_stop(scalar_stopping_signal)) + + def computation(unused_scalar_stopping_signal): + return_value = op_fn() + execute_ops = return_value['ops'] + signals = return_value['signals'] + with ops.control_dependencies(execute_ops): + return _StopSignals.as_scalar_stopping_signal(signals) + + # By setting parallel_iterations=1, the parallel execution in while_loop is + # basically turned off. + with ops.device(device): + return control_flow_ops.while_loop( + cond, + computation, [_StopSignals.NON_STOPPING_SIGNAL], + parallel_iterations=1) def _validate_tpu_training_graph(): @@ -1760,8 +2094,9 @@ def _validate_tpu_training_graph(): # Check if there is atleast one CrossReplicaSum operation in the graph # This should be introduced by using the CrossShardOptimizer wrapper - cross_replica_sum_ops = [o for o in operations - if o.type == _CROSS_REPLICA_SUM_OP] + cross_replica_sum_ops = [ + o for o in operations if o.type == _CROSS_REPLICA_SUM_OP + ] if not cross_replica_sum_ops: raise ValueError( 'CrossShardOptimizer must be used for model training on TPUs.') @@ -1808,9 +2143,11 @@ def _get_scaffold(captured_scaffold_fn): if scaffold: wrapped_finalize = scaffold.finalize + def _finalize(): with _CapturingContext('Inside Scaffold.finalize'): wrapped_finalize() + scaffold.finalize = _finalize return scaffold @@ -1825,9 +2162,8 @@ class _CapturingContext(control_flow_ops.ControlFlowContext): def AddOp(self, op): # pylint: disable=invalid-name for c in op.inputs: if tpu._TPU_REPLICATE_ATTR in c.op.node_def.attr: # pylint: disable=protected-access - raise ValueError( - '{}: Op {} depends on TPU computation {}, ' - 'which is not allowed.'.format(self._message, op, c)) + raise ValueError('{}: Op {} depends on TPU computation {}, ' + 'which is not allowed.'.format(self._message, op, c)) def __enter__(self): # pylint: disable=protected-access @@ -1840,3 +2176,192 @@ class _CapturingContext(control_flow_ops.ControlFlowContext): self._g._set_control_flow_context(self._old) # pylint: disable=protected-access +class _Inputs(object): + """A data structure representing the input_fn returned values. + + This also supports the returned value from input_fn as `Dataset`. + """ + + def __init__(self, features=None, labels=None, dataset=None, signals=None): + if dataset is not None and (features is not None or labels is not None or + signals is not None): + raise RuntimeError('Internal Error: Either (features and labels) or ' + 'dataset should be provided, not both. Please file ' + 'bug') + + self._features = features + self._labels = labels + self._signals = signals + + self._dataset = dataset + self._iterator = None + + @staticmethod + def from_input_fn(return_values): + """Returns an `_Inputs` instance according to `input_fn` return value.""" + if isinstance(return_values, dataset_ops.Dataset): + dataset = return_values + return _Inputs(dataset=dataset) + + features, labels = _Inputs._parse_inputs(return_values) + return _Inputs(features, labels) + + @staticmethod + def _parse_inputs(return_values): + if isinstance(return_values, tuple): + features, labels = return_values + else: + features, labels = return_values, None + return features, labels + + @property + def is_dataset(self): + """Returns True if the return value from input_fn is Dataset.""" + return self._dataset is not None + + def dataset_initializer_hook(self): + """Returns a `SessionRunHook` to initialize this dataset. + + This must be called before `features_and_labels`. + """ + iterator = self._dataset.make_initializable_iterator() + # pylint: disable=protected-access + hook = estimator_lib._DatasetInitializerHook(iterator) + self._iterator = iterator + return hook + + def features_and_labels(self): + """Gets `features` and `labels`.""" + if self.is_dataset: + return _Inputs._parse_inputs(self._iterator.get_next()) + + return (self._features, self._labels) + + def signals(self): + return self._signals + + @property + def dataset(self): + return self._dataset + + +# TODO(xiejw): Extend this to support final partial batch. +class _InputsWithStoppingSignals(_Inputs): + """Inputs with `_StopSignals` inserted into the dataset.""" + + def __init__(self, dataset, batch_size): + + assert dataset is not None + + user_provided_dataset = dataset.map( + _InputsWithStoppingSignals.insert_stopping_signal( + stop=False, batch_size=batch_size)) + final_batch_dataset = dataset.take(1).map( + _InputsWithStoppingSignals.insert_stopping_signal( + stop=True, batch_size=batch_size)) + dataset = user_provided_dataset.concatenate(final_batch_dataset).prefetch(2) + + super(_InputsWithStoppingSignals, self).__init__(dataset=dataset) + self._current_inputs = None + + def features_and_labels(self): + if self._current_inputs is not None: + raise RuntimeError( + 'Internal Error: The previous inputs have not been properly ' + 'consumed. First call features_and_labels, then call signals.') + + inputs_with_signals = self._iterator.get_next() + features = inputs_with_signals['features'] + labels = inputs_with_signals.get('labels') + + self._current_inputs = inputs_with_signals + return features, labels + + def signals(self): + """Returns the `Signals` from `_Inputs`.""" + if self._current_inputs is None: + raise RuntimeError( + 'Internal Error: The current inputs have not been properly ' + 'generated. First call features_and_labels, then call signals.') + signals = self._current_inputs['signals'] + self._current_inputs = None + return signals + + @staticmethod + def insert_stopping_signal(stop, batch_size): + """Inserts stopping_signal into dataset via _map_fn. + + Here we change the data structure in the dataset, such that the return value + is a dictionary now and `features`, `labels`, and `signals` are three + distinguished keys in that dict. This provides a better structure, which + eases the process to decompose the inputs (see `features_and_labels`). + + Args: + stop: bool, state of current stopping signals. + batch_size: int, batch size. + + Returns: + A map_fn passed to dataset.map API. + """ + + def _map_fn(*args): + features, labels = _Inputs._parse_inputs(args) + new_input_dict = {} + new_input_dict['features'] = features + if labels is not None: + new_input_dict['labels'] = labels + new_input_dict['signals'] = _StopSignals( + stop=stop, batch_size=batch_size).as_dict() + return new_input_dict + + return _map_fn + + +class _StopSignals(object): + """Signals class holding all logic to handle TPU stopping condition.""" + + NON_STOPPING_SIGNAL = 0.0 + STOPPING_SIGNAL = 1.0 + + def __init__(self, stop, batch_size): + self._stop = stop + self._batch_size = batch_size + + def as_dict(self): + shape = [self._batch_size, 1] + dtype = dtypes.float32 + + if self._stop: + stopping = array_ops.ones(shape=shape, dtype=dtype) + else: + stopping = array_ops.zeros(shape=shape, dtype=dtype) + + return {'stopping': stopping} + + @staticmethod + def as_scalar_stopping_signal(signals): + return array_ops.identity(signals['stopping'][0][0]) + + @staticmethod + def should_stop(scalar_stopping_signal): + return scalar_stopping_signal >= _StopSignals.STOPPING_SIGNAL + + +class _SignalsHelper(object): + """A general helper class to handle common signals manipulation.""" + + def __init__(self, signals): + self._signal_keys = [] + for key in sorted(signals.iterkeys()): + self._signal_keys.append(key) + + @property + def num_signals(self): + return len(self._signal_keys) + + def unflatten(self, tensor_list): + return dict(zip(self._signal_keys, tensor_list)) + + @staticmethod + def as_tensor_list(signals): + return [signals[key] for key in sorted(signals.iterkeys())] diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_sharding.py b/tensorflow/contrib/tpu/python/tpu/tpu_sharding.py index f8ba7d45e20b2f48e1409427665878df40a6db02..f5af03f33ca8f13af517007672e9ce0e12be6205 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_sharding.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_sharding.py @@ -244,7 +244,8 @@ class ShardingPolicy(object): str(shapes), self.number_of_shards)) unsharded_shapes = [self._unshard_shape(s) for s in shapes] for i in xrange(self.number_of_shards - 1): - if unsharded_shapes[i] != unsharded_shapes[self.number_of_shards - 1]: + if not unsharded_shapes[i].is_compatible_with( + unsharded_shapes[self.number_of_shards - 1]): raise ValueError( "sharded shapes %s are not consistent shards of a full shape " "sharded %d ways along dimension %d" % ( diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py b/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..493d1848c072caa5254fc87c67badc2e99ec16ee --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py @@ -0,0 +1,155 @@ +# 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. +# =================================================================== +"""TPU system metadata and associated tooling.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import collections +import re + +from tensorflow.contrib.tpu.python.tpu import tpu +from tensorflow.core.protobuf import config_pb2 +from tensorflow.python.client import session as session_lib +from tensorflow.python.framework import errors +from tensorflow.python.framework import ops +from tensorflow.python.platform import tf_logging as logging + +_PINGING_MASTER_TIMEOUT_IN_MS = 60 * 1000 # 1 min +_RETRY_TIMES = 120 +_INITIAL_TPU_SYSTEM_TIMEOUT_IN_MS = 300 * 1000 # 5 mins + +_TPU_DEVICE_REG = re.compile(r'.*task:(\d+)/.*device:TPU:(\d+)$') + +# _TPUSystemMetadata is used by TPUEstimator to hold TPU configuration, +# including num_cores and num_hosts. +_TPUSystemMetadata = collections.namedtuple('_TPUSystemMetadata', [ + 'num_cores', + 'num_hosts', + 'num_of_cores_per_host', + 'topology', + 'devices', +]) + + +def _query_tpu_system_metadata(master_address, run_config, + query_topology=False): + """Automatically detects the TPU system metadata in the system.""" + tpu_core_count = 0 + devices = [] + device_dict = collections.defaultdict(list) + + retry_count = 1 + while True: + logging.info('Querying Tensorflow master (%s) for TPU system metadata.', + master_address) + try: + with ops.Graph().as_default(): + with session_lib.Session( + master_address, + config=_get_session_config_with_timeout( + _PINGING_MASTER_TIMEOUT_IN_MS, run_config)) as sess: + devices = sess.list_devices() + for device in devices: + match = _TPU_DEVICE_REG.match(device.name) + if match: + host_id = match.group(1) + core_id = match.group(2) + device_dict[host_id].append(core_id) + tpu_core_count += 1 + break + except errors.DeadlineExceededError: + msg = ('Fail to connect Tensorflow master. It could be the TPU worker is ' + 'not ready (still under scheduling) or Tensorflow ' + 'master address is correct: got (%s).' % + (master_address)) + + # TODO(xiejw): For local or grpc master we might not need retry logic + # here. + if retry_count <= _RETRY_TIMES: + logging.warning('%s', msg) + logging.warning('Retrying (%d/%d).', retry_count, _RETRY_TIMES) + retry_count += 1 + else: + raise ValueError(msg) + + num_of_cores_per_host = 0 + if tpu_core_count: + num_cores_per_host_set = set( + [len(core_ids) for core_ids in device_dict.values()]) + if len(num_cores_per_host_set) != 1: + raise RuntimeError( + 'TPU cores on each host is not same. This should not happen!. ' + 'devices: {}'.format(devices)) + num_of_cores_per_host = num_cores_per_host_set.pop() + + topology = None + if query_topology: + if not tpu_core_count: + raise RuntimeError( + 'Cannot find any TPU cores in the system (master address {}). ' + 'This usually means the master address is incorrect or the ' + 'TPU worker has some problems. Available devices: {}'.format( + master_address, devices)) + + topology = _obtain_topology(master_address, run_config) + + metadata = _TPUSystemMetadata( + num_cores=tpu_core_count, + num_hosts=len(device_dict), + num_of_cores_per_host=num_of_cores_per_host, + topology=topology, + devices=devices) + + if tpu_core_count: + logging.info('Found TPU system:') + logging.info('*** Num TPU Cores: %d', metadata.num_cores) + logging.info('*** Num TPU Workers: %d', metadata.num_hosts) + logging.info('*** Num TPU Cores Per Worker: %d', + metadata.num_of_cores_per_host) + logging.info('*** Available Devices: %s', metadata.devices) + else: + logging.info('Failed to find TPU: %s', metadata) + return metadata + + +def _obtain_topology(master_address, run_config): + try: + logging.info('Initializing TPU system (master: %s) to fetch topology ' + 'for model parallelism. This might take a while.', + master_address) + with ops.Graph().as_default(): + session_config = _get_session_config_with_timeout( + _INITIAL_TPU_SYSTEM_TIMEOUT_IN_MS, run_config) + with session_lib.Session( + master_address, config=session_config) as sess: + topology = sess.run(tpu.initialize_system()) + return topology + except errors.DeadlineExceededError: + raise ValueError( + 'Fail to initialize TPU system with master (%s). ' + 'Please double check the TPU system is functional.' % ( + master_address)) + + +def _get_session_config_with_timeout(timeout_in_secs, run_config): + cluster_def = None + if run_config.session_config and run_config.session_config.cluster_def.job: + cluster_def = run_config.session_config.cluster_def + + config = config_pb2.ConfigProto( + operation_timeout_in_ms=timeout_in_secs, cluster_def=cluster_def) + return config diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_test.py b/tensorflow/contrib/tpu/python/tpu/tpu_test.py new file mode 100644 index 0000000000000000000000000000000000000000..336d8260c3c8a5c30efa603e3faeabcc0944b8d0 --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/tpu_test.py @@ -0,0 +1,80 @@ +# Copyright 2017 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 tpu_function helpers.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.tpu.python.tpu import tpu +from tensorflow.contrib.tpu.python.tpu import tpu_feed +from tensorflow.contrib.tpu.python.tpu import training_loop + +from tensorflow.python.framework import dtypes +from tensorflow.python.layers import convolutional +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_util +from tensorflow.python.ops import math_ops + +from tensorflow.python.platform import test + + +class TPUContextTest(test.TestCase): + + def testIsInContext(self): + """Test that control_flow_util can check that we're in a TPU context.""" + z1 = array_ops.identity(1) + context = tpu.TPUReplicateContext(b"context") + context.Enter() + z2 = array_ops.identity(1) + context.Exit() + self.assertFalse(control_flow_util.IsInXLAContext(z1.op)) + self.assertTrue(control_flow_util.IsInXLAContext(z2.op)) + + +class TPULayerRewriteTest(test.TestCase): + + def testUsingInfeedQueueWithRegularizer(self): + """Test that Layer regularizers can reference data created in loops.""" + + def make_regularizer(scale): + return lambda inputs: scale * math_ops.reduce_sum(math_ops.square(inputs)) + + def training_step(inputs, scale): + outputs = convolutional.conv2d( + inputs, + filters=16, + kernel_size=(3, 3), + data_format="channels_first", + kernel_regularizer=make_regularizer(scale)) + loss = math_ops.reduce_mean(math_ops.square(outputs)) + return loss.op + + inputs = array_ops.zeros(shape=(128, 32, 32, 16)) + scale = array_ops.ones(shape=()) + infeed = tpu_feed.InfeedQueue( + tuple_types=[dtypes.float32, dtypes.float32], + tuple_shapes=[inputs.shape, scale.shape]) + + def loop(): + return training_loop.repeat(5, training_step, infeed_queue=infeed) + + # This should not throw an error. + tpu.rewrite(loop) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/tpu/tpu_estimator.md b/tensorflow/contrib/tpu/tpu_estimator.md new file mode 100644 index 0000000000000000000000000000000000000000..4ef8f9eebdb165e5fe221be8670276bf943159b3 --- /dev/null +++ b/tensorflow/contrib/tpu/tpu_estimator.md @@ -0,0 +1,241 @@ +# Using the Estimator API with TPUs + + +This document describes how to train a TensorFlow model on TPUs using the +Estimator API. If you are interested in the hardware itself, check out the +[Cloud TPU documentation](https://cloud.google.com/tpu/docs). + +The TPU Estimator simplifies running models on a Cloud TPU by automatically +handling numerous low-level hardware-specific details + +[TOC] + +## Introduction to Estimator + +[TensorFlow +tutorials](https://www.tensorflow.org/extend/estimators) cover the Estimator +API. At a high-level, the Estimator API provides: + +* `Estimator.train()` - train a model on a given input for a fixed number of + steps. +* `Estimator.evaluate()` - evaluate the model on a test set. +* `Estimator.predict()` - run inference using the trained model. +* `Estimator.export_savedmodel()` - export your model for serving. + +In addition, `Estimator` includes default behavior common to training jobs, +such as saving and restoring checkpoints, creating summaries for TensorBoard, +etc. + +`Estimator` requires you to write a `model_fn` and an `input_fn`, which +correspond to the model and input portions of your TensorFlow graph. + +The following code demonstrates using `TPUEstimator` with MNIST example to +handle training: + + def model_fn(features, labels, mode, params): + """A simple CNN.""" + del params # unused + + input_layer = tf.reshape(features, [-1, 28, 28, 1]) + conv1 = tf.layers.conv2d( + inputs=input_layer, filters=32, kernel_size=[5, 5], padding="same", + activation=tf.nn.relu) + pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2) + conv2 = tf.layers.conv2d( + inputs=pool1, filters=64, kernel_size=[5, 5], + padding="same", activation=tf.nn.relu) + pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2) + pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64]) + dense = tf.layers.dense(inputs=pool2_flat, units=128, activation=tf.nn.relu) + dropout = tf.layers.dropout( + inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN) + logits = tf.layers.dense(inputs=dropout, units=10) + onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10) + + loss = tf.losses.softmax_cross_entropy( + onehot_labels=onehot_labels, logits=logits) + + learning_rate = tf.train.exponential_decay( + FLAGS.learning_rate, tf.train.get_global_step(), 100000, 0.96) + + optimizer = tpu_optimizer.CrossShardOptimizer( + tf.train.GradientDescentOptimizer(learning_rate=learning_rate)) + + train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step()) + return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op) + + + def get_input_fn(filename): + """Returns an `input_fn` for train and eval.""" + + def input_fn(params): + """An input_fn to parse 28x28 images from filename using tf.data.""" + batch_size = params["batch_size"] + + def parser(serialized_example): + """Parses a single tf.Example into image and label tensors.""" + features = tf.parse_single_example( + serialized_example, + features={ + "image_raw": tf.FixedLenFeature([], tf.string), + "label": tf.FixedLenFeature([], tf.int64), + }) + image = tf.decode_raw(features["image_raw"], tf.uint8) + image.set_shape([28 * 28]) + # Normalize the values of the image from the range [0, 255] to [-0.5, 0.5] + image = tf.cast(image, tf.float32) * (1. / 255) - 0.5 + label = tf.cast(features["label"], tf.int32) + return image, label + + dataset = tf.contrib.data.TFRecordDataset( + filename, buffer_size=FLAGS.dataset_reader_buffer_size) + dataset = dataset.map(parser).cache().repeat().batch(batch_size) + images, labels = dataset.make_one_shot_iterator().get_next() + # set_shape to give inputs statically known shapes. + images.set_shape([batch_size, 28 * 28]) + labels.set_shape([batch_size]) + return images, labels + return input_fn + + + def main(unused_argv): + + tf.logging.set_verbosity(tf.logging.INFO) + + run_config = tpu_config.RunConfig( + master=FLAGS.master, + model_dir=FLAGS.model_dir, + session_config=tf.ConfigProto( + allow_soft_placement=True, log_device_placement=True), + tpu_config=tpu_config.TPUConfig(FLAGS.iterations, FLAGS.num_shards),) + + estimator = tpu_estimator.TPUEstimator( + model_fn=model_fn, + use_tpu=FLAGS.use_tpu, + train_batch_size=FLAGS.batch_size, + eval_batch_size=FLAGS.batch_size, + config=run_config) + + estimator.train(input_fn=get_input_fn(FLAGS.train_file), + max_steps=FLAGS.train_steps) + + +Although this code is quite simple by appearance, there are some new +concepts to learn for using `TPU`s. The next section will cover the most +important details. + +## New Concepts Related to TPU/TPUEstimator + +TF programs run with `TPU Estimator` use an [in-graph +replication](https://www.tensorflow.org/deploy/distributed) approach. + +In-graph replication (also known as single-session replication) differs from +the between-graph replication (also known as multi-session replication) +training typically used in distributed TensorFlow. The major +differences include: + +1. The TensorFlow Session master is not local anymore. The user python program + creates one single graph that is replicated across all the cores in the Cloud + TPU. The typical configuration today sets the TensorFlow session master to be + the first worker. + +1. The input pipeline is placed on remote hosts (instead of local) to ensure the + training examples can be fed as fast as possible to TPU system. All queue-based + input pipelines do not work effectively. Dataset (tf.data) is + required. + +1. Workers in the TPU system operate in synchronous fashion, and each perform + the same step at the same time. + +Regarding programming model, _"The programmer picks a (large) batch size B and +writes the program (and sets hyperparameters) based on that batch size. The +system distributes the computation across the available devices." + +To align these, `TPUEstimator` wraps the computation (the `model_fn`) and +distributes it to all available TPU chips. + +To summarize: + +- The `input_fn` models the input pipeline running on remote host CPU. Use + `tf.data` to program the input Ops. `input_fn` is expected to be invoked + multiple times when using TPU pods. Each handles one device's input of the + global batch. The shard batch size should be retrieved from + `params['batch_size']`. We plan to provide better abstraction about the + sharding mechanism for `tf.data` to remove the `params['batch_size']`. + +- The `model_fn` models the computation which will be replicated and distributed + to all TPU chips. It should only contains ops that are supported by TPUs. + +## Convert from Vanilla Estimator to TPUEstimator + +It is always recommended to port a small, simple model first to make sure that +you are familiar with the basic concepts of `TPUEstimator` and test end-to-end +behavior. Once your simple model runs, gradually add more functionality. +In addition, there are several sample models, available at +[github.com/tensorflow/tpu-demos](https://github.com/tensorflow/tpu-demos). + +To convert your code from the vanilla `Estimator` class to use TPUs, change the +following (note some of the details may change over time): + +- Switch from `tf.estimator.RunConfig` to `tf.contrib.tpu.RunConfig`. +- Set the `TPUConfig` (part of the `tf.contrib.tpu.RunConfig`) to specify the + `iterations_per_loop`, number of iterations to run on the TPU device for one + `session.run` call (per training loop), and `num_shards`, the number of shards + (typically the number of TPU cores you’re running on). TPUs run a number of + iterations of the training loop before returning to host. Until all iterations + on the TPU device are run, no checkpoints or summaries will be saved. In the + future, we’ll choose a reasonable default. +- In `model_fn`, use `tf.contrib.tpu.CrossShardOptimizer` to wrap your + optimizer. Example: + + optimizer = tpu_optimizer.CrossShardOptimizer( + tf.train.GradientDescentOptimizer(learning_rate=learning_rate)) + +- Switch from `tf.estimator.Estimator` to `tf.contrib.tpu.TPUEstimator`. + +The default `RunConfig` will save summaries for TensorBoard every 100 steps and +write checkpoints every 10 minutes. + + +## FAQ + +### Why `tf.data` is Required for the Input Pipeline + +There are two reasons: + +1. The user code runs on the client, while the TPU computation is executed on + the `worker`. Input pipeline ops must be placed on the remote worker for + good performance. Only `tf.data` (Dataset) supports this. + +1. In order to amortize the TPU launch cost, the model train step is wrapped in + a `tf.while_loop`, such that one `Session.run` actually runs many iterations + for one train loop. To remove network back and forth, the input pipeline + in the future will be wrapped in a `tf.while_loop` and be placed on the + corresponding `worker`. Withou this, unnecessary network latency becomes + the performance bottleneck for models with short training-step times, or in + environments where network latency is higher. Only `tf.data` can be wrapped + by a `tf.while_loop`. + + +### How to add other CPU Ops into Graph +As `model_fn` only allows TPU Ops for computation, the easier workaround to add +CPU Ops into Graph is: + +1. Create a [SessionRunHook](https://www.tensorflow.org/api_docs/python/tf/train/SessionRunHook). +1. Modify the graph in the `def begin(self)`, +1. Pass the hook to `TPUEstimator.train`. + +### Running On GCP Cloud TPUs +To run your models on GCP Cloud TPUs refer to the [Cloud Documentation](https://cloud.google.com/tpu/docs/tutorials/mnist). +Refer to this link for all [Cloud TPU documentation](https://cloud.google.com/tpu/docs). + + +### Profiling +You can profile the `worker` by using instructions as specified in the [Cloud TPU Tools](https://cloud.google.com/tpu/docs/cloud-tpu-tools). + + +### Is `int64` supported? +`int64` is not supported by TPU. Cast to int32 if applicable. The only exception +is global step, which relies on `assign_add`. `int64` support for global step +is added to ensure checkpoint compatibility between `TPUEstimator` and non-TPU +`Estimator`. diff --git a/tensorflow/contrib/training/BUILD b/tensorflow/contrib/training/BUILD index 6139c1d5838c24414549b4e2bc4722175f2d1925..6db373d2d5e20ea7da449530b2730403c3bb64cc 100644 --- a/tensorflow/contrib/training/BUILD +++ b/tensorflow/contrib/training/BUILD @@ -26,7 +26,7 @@ py_library( "python/training/resample.py", "python/training/sampling_ops.py", "python/training/sequence_queueing_state_saver.py", - "python/training/sgdr_learning_rate_decay.py", + "python/training/tensor_queue_dataset.py", "python/training/training.py", "python/training/tuner.py", ], @@ -286,6 +286,28 @@ py_test( ], ) +py_test( + name = "tensor_queue_dataset_test", + size = "large", + srcs = ["python/training/tensor_queue_dataset_test.py"], + srcs_version = "PY2AND3", + tags = ["notsan"], + deps = [ + ":training_py", + "//tensorflow/contrib/data/python/kernel_tests:dataset_serialization_test", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:gradients", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform", + "//tensorflow/python:random_seed", + "//tensorflow/python:training", + "//tensorflow/python:variables", + "//tensorflow/python/data", + "//third_party/py/numpy", + ], +) + filegroup( name = "all_files", srcs = glob( diff --git a/tensorflow/contrib/training/python/training/batch_sequences_with_states_test.py b/tensorflow/contrib/training/python/training/batch_sequences_with_states_test.py index 2a0ef0e6b3750b4f0464f1f4390819e1fc2c7872..dbdbb08a8252c799924812c83fff7f0631424761 100644 --- a/tensorflow/contrib/training/python/training/batch_sequences_with_states_test.py +++ b/tensorflow/contrib/training/python/training/batch_sequences_with_states_test.py @@ -53,7 +53,7 @@ class BatchSequencesWithStatesTest(test.TestCase): sp_tensor1 = sparse_tensor.SparseTensor( array_ops.constant(ind1, dtypes.int64), array_ops.constant(val1, dtypes.int64), - array_ops.constant(shape1, dtypes.int64)) + array_ops.placeholder_with_default(shape1, shape=[2])) ind2 = np.array([ [0, 0, 1], [0, 1, 0], @@ -68,7 +68,7 @@ class BatchSequencesWithStatesTest(test.TestCase): sp_tensor2 = sparse_tensor.SparseTensor( array_ops.constant(ind2, dtypes.int64), array_ops.constant(val2, dtypes.int64), - array_ops.constant(shape2, dtypes.int64)) + array_ops.placeholder_with_default(shape2, shape=[3])) sp_tensor3 = sparse_tensor.SparseTensor( array_ops.constant([[1, 9], [2, 2], [2, 10]], dtypes.int64), array_ops.constant([7, 15, 2], dtypes.int64), @@ -320,6 +320,18 @@ class BatchSequencesWithStatesTest(test.TestCase): def testNotAMultiple(self): num_unroll = 3 # Not a divisor of value_length - # so padding would have been necessary. + + # Use placeholder_with_default in sequences to make sure we get runtime + # error instead of shape inference error + sequences = { + "seq1": array_ops.placeholder_with_default(self.sequences["seq1"], + shape=(None, 5)), + "seq2": array_ops.placeholder_with_default(self.sequences["seq2"], + shape=(None, 4, 2)), + "seq3": self.sequences["seq3"], + "seq4": self.sequences["seq4"], + } + with self.test_session() as sess: with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, ".*should be a multiple of: 3, but saw " @@ -330,7 +342,7 @@ class BatchSequencesWithStatesTest(test.TestCase): with coord.stop_on_exception(): next_batch = sqss.batch_sequences_with_states( input_key=self.key, - input_sequences=self.sequences, + input_sequences=sequences, input_context=self.context, input_length=3, initial_states=self.initial_states, @@ -493,6 +505,18 @@ class BatchSequencesWithStatesTest(test.TestCase): expected_seq4_batch2=expected_seq4_batch2) +class BatchSequencesWithStatesTestWithCApi(BatchSequencesWithStatesTest): + + def setUp(self): + self._prev_value = ops._USE_C_API + ops._USE_C_API = True + super(BatchSequencesWithStatesTestWithCApi, self).setUp() + + def tearDown(self): + super(BatchSequencesWithStatesTestWithCApi, self).tearDown() + ops._USE_C_API = self._prev_value + + class PaddingTest(test.TestCase): def testPaddingInvalidLengths(self): diff --git a/tensorflow/contrib/training/python/training/hparam.py b/tensorflow/contrib/training/python/training/hparam.py index 80de0f6eb7e36a1c86f7d44e4053a9757b09f0ae..fdfd27d6a414933b0bec824bae512c45dac24d3c 100644 --- a/tensorflow/contrib/training/python/training/hparam.py +++ b/tensorflow/contrib/training/python/training/hparam.py @@ -40,7 +40,7 @@ PARAM_RE = re.compile(r""" ((?P[^,\[]*) # single value: "a" or None | \[(?P[^\]]*)\]) # list of values: None or "1,2,3" - ($|,)""", re.VERBOSE) + ($|,\s*)""", re.VERBOSE) def _parse_fail(name, var_type, value, values): diff --git a/tensorflow/contrib/training/python/training/hparam_test.py b/tensorflow/contrib/training/python/training/hparam_test.py index 28e4b4d01eda9bef07ff7929f74894e09a3e987c..16397622edd382bc6dcb12870de5fa22130a2c2b 100644 --- a/tensorflow/contrib/training/python/training/hparam_test.py +++ b/tensorflow/contrib/training/python/training/hparam_test.py @@ -55,7 +55,7 @@ class HParamsTest(test.TestCase): self.assertEqual(12, hparams.aaa) self.assertEqual(2.0, hparams.b) self.assertEqual('relu6', hparams.c_c) - hparams.parse('c_c=relu4,b=-2.0e10') + hparams.parse('c_c=relu4, b=-2.0e10') self.assertDictEqual({ 'aaa': 12, 'b': -2.0e10, diff --git a/tensorflow/contrib/training/python/training/sgdr_learning_rate_decay.py b/tensorflow/contrib/training/python/training/sgdr_learning_rate_decay.py deleted file mode 100644 index ed0f398e30a7f3c0b1b9378f8fc5d5bfbea1536a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/training/python/training/sgdr_learning_rate_decay.py +++ /dev/null @@ -1,187 +0,0 @@ -# Copyright 2015 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -"""SGDR learning rate decay function.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import math - -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import ops -from tensorflow.python.ops import math_ops, control_flow_ops - - -def sgdr_decay(learning_rate, global_step, initial_period_steps, - t_mul=2.0, m_mul=1.0, name=None): - """Implements Stochastic Gradient Descent with Warm Restarts (SGDR). - - As described in "SGDR: Stochastic Gradient Descent - with Warm Restarts" by Ilya Loshchilov & Frank Hutter, Proceedings of - ICLR'2017, available at https://arxiv.org/pdf/1608.03983.pdf - - The learning rate decreases according to cosine annealing: - - ```python - learning_rate * 0.5 * (1 + cos(x_val * pi)) # for x_val defined in [0, 1] - ``` - - Thus, at the beginning (when the restart index i = 0), - the learning rate decreases for `initial_period_steps` steps from the initial - learning rate `learning_rate` (when `x_val=0`, we get `cos(0)=1`) to - 0 (when `x_val=1`, we get `cos(pi)=-1`). - - The decrease within the i-th period takes `t_i` steps, - where `t_0` = `initial_period_steps` is the user-defined number of batch - iterations (not epochs as in the paper) to be performed before the first - restart is launched. - - Then, we perform the first restart (i=1) by setting the learning rate to - `learning_rate*(m_mul^i)`, where `m_mul in [0,1]` (set to 1 by default). - The i-th restart runs for `t_i=t_0*(t_mul^i)` steps, i.e., every new - restart runs `t_mul` times longer than the previous one. - - Importantly, when one has no access to a validation set, SGDR suggests - to report the best expected / recommended solution in the following way: - When we are within our initial run (i=0), every new solution represents - SGDR's recommended solution. Instead, when i>0, the recommended solution is - the one obtained at the end of each restart. - - Note that the minimum learning rate is set to 0 for simplicity, - you can adjust the code to deal with any positive minimum learning rate - as defined in the paper. - - `initial_period_steps` is the duration of the first period measured in terms - of number of minibatch updates. If one wants to use epochs, one should compute - the number of updates required for an epoch. - - For example, assume the following parameters and intention: - Minibatch size: 100 - Training dataset size: 10000 - If the user wants the first decay period to span across 5 epochs, then - `initial_period_steps` = 5 * 10000/100 = 500 - - Train for 10000 batch iterations with the initial learning rate set to - 0.1, then restart to run 2 times longer, i.e, for 20000 batch iterations - and with the initial learning rate 0.05, then restart again and again, - doubling the runtime of each new period and with two times smaller - initial learning rate. - - To accomplish the above, one would write: - - ```python - ... - global_step = tf.Variable(0, trainable=False) - starter_learning_rate = 0.1 - learning_rate = sgdr_decay(starter_learning_rate, global_step, - initial_period_steps=10000, t_mul=2, m_mul=0.5) - # Passing global_step to minimize() will increment it at each step. - learning_step = ( - tf.train.GradientDescentOptimizer(learning_rate) - .minimize(...my loss..., global_step=global_step) - ) - - # Step | 0 | 1000 | 5000 | 9000 | 9999 | 10000 | 11000 | - # LR | 0.1 | 0.097 | 0.05 | 0.002 | 0.00 | 0.05 | 0.0496 | - - # Step | 20000 | 29000 | 29999 | 30000 | - # LR | 0.025 | 0.0003 | 0.00 | 0.025 | - ``` - - Args: - learning_rate: A scalar `float32` or `float64` `Tensor` or a - Python number. The initial learning rate. - global_step: A scalar `int32` or `int64` `Tensor` or a Python number. - Global step to use for the decay computation. Must not be negative. - initial_period_steps: Duration of the first period measured as the number - of minibatch updates, if one wants to use epochs, one should compute - the number of updates required for an epoch. - t_mul: A scalar `float32` or `float64` `Tensor` or a Python number. - Must be positive. - Used to derive the number of iterations in the i-th period: - `initial_period_steps * (t_mul^i)`. Defaults to 2.0. - m_mul: A scalar `float32` or `float64` `Tensor` or a Python number. - Must be positive. - Used to derive the initial learning rate of the i-th period: - `learning_rate * (m_mul^i)`. Defaults to 1.0 - - Returns: - A scalar `Tensor` of the same type as `learning_rate`. - The learning rate for a provided global_step. - Raises: - ValueError: if `global_step` is not supplied. - """ - - if global_step is None: - raise ValueError("global_step is required for sgdr_decay.") - with ops.name_scope(name, "SGDRDecay", - [learning_rate, global_step, - initial_period_steps, t_mul, m_mul]) as name: - learning_rate = ops.convert_to_tensor(learning_rate, - name="initial_learning_rate") - dtype = learning_rate.dtype - global_step = math_ops.cast(global_step, dtype) - t_0 = math_ops.cast(initial_period_steps, dtype) - t_mul = math_ops.cast(t_mul, dtype) - m_mul = math_ops.cast(m_mul, dtype) - - c_one = math_ops.cast(constant_op.constant(1.0), dtype) - c_half = math_ops.cast(constant_op.constant(0.5), dtype) - c_pi = math_ops.cast(constant_op.constant(math.pi), dtype) - - # Find normalized value of the current step - x_val = math_ops.div(global_step, t_0) - - def compute_step(x_val, geometric=False): - if geometric: - # Consider geometric series where t_mul != 1 - # 1 + t_mul + t_mul^2 ... = (1 - t_mul^i_restart) / (1 - t_mul) - - # First find how many restarts were performed for a given x_val - # Find maximal integer i_restart value for which this equation holds - # x_val >= (1 - t_mul^i_restart) / (1 - t_mul) - # x_val * (1 - t_mul) <= (1 - t_mul^i_restart) - # t_mul^i_restart <= (1 - x_val * (1 - t_mul)) - - # tensorflow allows only log with base e - # i_restart <= log(1 - x_val * (1 - t_mul) / log(t_mul) - # Find how many restarts were performed - - i_restart = math_ops.floor( - math_ops.log(c_one - x_val * (c_one - t_mul)) / math_ops.log(t_mul)) - # Compute the sum of all restarts before the current one - sum_r = (c_one - t_mul ** i_restart) / (c_one - t_mul) - # Compute our position within the current restart - x_val = (x_val - sum_r) / t_mul ** i_restart - - else: - # Find how many restarts were performed - i_restart = math_ops.floor(x_val) - # Compute our position within the current restart - x_val = x_val - i_restart - return i_restart, x_val - - i_restart, x_val = control_flow_ops.cond( - math_ops.equal(t_mul, c_one), - lambda: compute_step(x_val, geometric=False), - lambda: compute_step(x_val, geometric=True)) - - # If m_mul < 1, then the initial learning rate of every new restart will be - # smaller, i.e., by a factor of m_mul ** i_restart at i_restart-th restart - m_fac = learning_rate * (m_mul ** i_restart) - - return math_ops.multiply(c_half * m_fac, - (math_ops.cos(x_val * c_pi) + c_one), name=name) diff --git a/tensorflow/contrib/training/python/training/sgdr_learning_rate_decay_test.py b/tensorflow/contrib/training/python/training/sgdr_learning_rate_decay_test.py deleted file mode 100644 index 4a46e9a49ef203384e36698f81d6cbe3a3881ef8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/training/python/training/sgdr_learning_rate_decay_test.py +++ /dev/null @@ -1,145 +0,0 @@ -# Copyright 2015 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -"""Functional test for sgdr learning rate decay.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import math - -from sgdr_learning_rate_decay import sgdr_decay -from tensorflow.python.platform import googletest -from tensorflow.python.framework import test_util -from tensorflow.python.framework import dtypes -from tensorflow import placeholder - - -class SGDRDecayTest(test_util.TensorFlowTestCase): - """Unit tests for SGDR learning rate decay.""" - - def get_original_values(self, lr, t_e, mult_factor, iter_per_epoch, epochs): - """Get an array with learning rate values from the consecutive steps using - the original implementation - (https://github.com/loshchil/SGDR/blob/master/SGDR_WRNs.py).""" - t0 = math.pi / 2.0 - tt = 0 - te_next = t_e - - lr_values = [] - sh_lr = lr - for epoch in range(epochs): - for _ in range(iter_per_epoch): - # In the original approach training function is executed here - lr_values.append(sh_lr) - dt = 2.0 * math.pi / float(2.0 * t_e) - tt = tt + float(dt) / iter_per_epoch - if tt >= math.pi: - tt = tt - math.pi - cur_t = t0 + tt - new_lr = lr * (1.0 + math.sin(cur_t)) / 2.0 # lr_min = 0, lr_max = lr - sh_lr = new_lr - if (epoch + 1) == te_next: # time to restart - sh_lr = lr - tt = 0 # by setting to 0 we set lr to lr_max, see above - t_e = t_e * mult_factor # change the period of restarts - te_next = te_next + t_e # note the next restart's epoch - - return lr_values - - def get_sgdr_values(self, lr, initial_period_steps, t_mul, iters): - """Get an array with learning rate values from the consecutive steps - using current tensorflow implementation.""" - with self.test_session(): - step = placeholder(dtypes.int32) - - decay = sgdr_decay(lr, step, initial_period_steps, t_mul) - lr_values = [] - for i in range(iters): - lr_values.append(decay.eval(feed_dict={step: i})) - - return lr_values - - def testCompareToOriginal(self): - """Compare values generated by tensorflow implementation to the values - generated by the original implementation - (https://github.com/loshchil/SGDR/blob/master/SGDR_WRNs.py).""" - with self.test_session(): - lr = 10.0 - init_steps = 2 - t_mul = 3 - iters = 10 - epochs = 50 - - org_lr = self.get_original_values(lr, init_steps, t_mul, iters, epochs) - sgdr_lr = self.get_sgdr_values(lr, init_steps*iters, t_mul, iters*epochs) - - for org, sgdr in zip(org_lr, sgdr_lr): - self.assertAllClose(org, sgdr) - - def testMDecay(self): - """Test m_mul argument. Check values for learning rate at the beginning - of the first, second, third and fourth period. """ - with self.test_session(): - step = placeholder(dtypes.int32) - - lr = 0.1 - t_e = 10 - t_mul = 3 - m_mul = 0.9 - - decay = sgdr_decay(lr, step, t_e, t_mul, m_mul) - - test_step = 0 - self.assertAllClose(decay.eval(feed_dict={step: test_step}), - lr) - - test_step = t_e - self.assertAllClose(decay.eval(feed_dict={step: test_step}), - lr * m_mul) - - test_step = t_e + t_e*t_mul - self.assertAllClose(decay.eval(feed_dict={step: test_step}), - lr * m_mul**2) - - test_step = t_e + t_e*t_mul + t_e * (t_mul**2) - self.assertAllClose(decay.eval(feed_dict={step: test_step}), - lr * (m_mul**3)) - - def testCos(self): - """Check learning rate values at the beginning, in the middle - and at the end of the period.""" - with self.test_session(): - step = placeholder(dtypes.int32) - lr = 0.2 - t_e = 1000 - t_mul = 1 - - decay = sgdr_decay(lr, step, t_e, t_mul) - - test_step = 0 - self.assertAllClose(decay.eval(feed_dict={step: test_step}), lr) - - test_step = t_e//2 - self.assertAllClose(decay.eval(feed_dict={step: test_step}), lr/2) - - test_step = t_e - self.assertAllClose(decay.eval(feed_dict={step: test_step}), lr) - - test_step = t_e*3//2 - self.assertAllClose(decay.eval(feed_dict={step: test_step}), lr/2) - -if __name__ == "__main__": - googletest.main() diff --git a/tensorflow/contrib/training/python/training/tensor_queue_dataset.py b/tensorflow/contrib/training/python/training/tensor_queue_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..409aba817c1ec37003eb98f000f6cf8918234c5d --- /dev/null +++ b/tensorflow/contrib/training/python/training/tensor_queue_dataset.py @@ -0,0 +1,200 @@ +# Copyright 2017 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. +# ============================================================================== +"""Python wrappers for Datasets and Iterators.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.data.util import nest +from tensorflow.python.data.util import sparse +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import tensor_util +from tensorflow.python.ops import gen_dataset_ops +from tensorflow.python.util import nest as tf_nest + + +class _PrependFromQueueAndPaddedBatchDataset(dataset_ops.Dataset): + """A `Dataset` that prepends a queue to another `Dataset`. + + A vector of handles to the queue is returned as the first component of + the associated iterator. This vector can be passed to + `enqueue_in_queue_dataset` to add new elements to the queue. + """ + + def __init__(self, input_dataset, batch_size, padded_shapes, padding_values): + """Initialize `PrependFromQueueAndPaddedBatchDataset`.""" + super(_PrependFromQueueAndPaddedBatchDataset, self).__init__() + if sparse.any_sparse(input_dataset.output_classes): + raise TypeError( + "Batching of padded sparse tensors is not currently supported") + self._input_dataset = input_dataset + self._batch_size = ops.convert_to_tensor( + batch_size, dtype=dtypes.int64, name="batch_size") + # pylint: disable=protected-access + if padded_shapes is None: + self._padded_shapes = nest.map_structure( + dataset_ops._partial_shape_to_tensor, input_dataset.output_shapes) + else: + self._padded_shapes = nest.map_structure_up_to( + input_dataset.output_shapes, dataset_ops._partial_shape_to_tensor, + padded_shapes) + padding_values = ( + padding_values if padding_values is not None else + dataset_ops._default_padding(input_dataset)) + self._padding_values = nest.map_structure_up_to( + input_dataset.output_shapes, dataset_ops._padding_value_to_tensor, + padding_values, input_dataset.output_types) + # pylint: enable=protected-access + + def _as_variant_tensor(self): + # pylint: disable=protected-access + return gen_dataset_ops.prepend_from_queue_and_padded_batch_dataset( + self._input_dataset._as_variant_tensor(), + batch_size=self._batch_size, + padded_shapes=[ + ops.convert_to_tensor(s, dtype=dtypes.int64) + for s in nest.flatten(self._padded_shapes) + ], + padding_values=nest.flatten(self._padding_values), + output_shapes=nest.flatten( + sparse.as_dense_shapes(self.output_shapes, self.output_classes))) + # pylint: enable=protected-access + + @property + def output_classes(self): + return (ops.Tensor, self._input_dataset.output_classes) + + def _as_batch_shape(self, shape_like): + return tensor_shape.vector(None).concatenate( + tensor_util.constant_value_as_shape(shape_like)) + + @property + def output_shapes(self): + # First output is a variant representing the Queue + return (tensor_shape.vector(None), + nest.map_structure(self._as_batch_shape, self._padded_shapes)) + + @property + def output_types(self): + # First output is a variant representing the Queue + return (dtypes.variant, self._input_dataset.output_types) + + +def prepend_from_queue_and_padded_batch_dataset(batch_size, + padding_values=None, + padded_shapes=None): + """A transformation that prepends a queue to a `Dataset` and batches results. + + A vector of handles to the queue is returned as the first component of the + associated iterator. This vector can be passed to `enqueue_in_queue_dataset` + to add new elements to the queue. + + Below is an example of how this dataset might be used to split incoming + variable-length sequences into "head" and "rest" parts, where "rest" parts + are re-enqueued back into the dataset. A more realistic example would + perform some calculation on the "head" and modify some components of "rest" + with the result (before re-enqueueing). + + ```python + dataset = tf.data.Dataset.from_tensor_slices([2*x for x in range(10)]) + # Make a dataset of variable-length vectors and their lengths. + dataset = dataset.map(lambda count: (count, tf.ones((count,)))) + # Emit a queue we can prepend to, and counts/values as padded batch. + dataset = dataset.apply( + tf.contrib.training.prepend_from_queue_and_padded_batch_dataset( + batch_size=10)) + dataset = dataset.prefetch(1) + + iterator = dataset.make_one_shot_iterator() + queue, (count, padded_value) = iterator.get_next() + + # Split the padded_value into two pieces: head and rest + rest_indices = tf.squeeze(tf.where(count > 3), axis=1) + bound = tf.minimum(3, tf.reduce_max(count)) + value_head = padded_value[:, :bound] + count_rest = tf.gather(count - 3, rest_indices) + value_rest = tf.gather(padded_value[:, bound:], rest_indices) + queue_rest = tf.gather(queue, rest_indices) + enqueue_rest_op = tf.contrib.training.enqueue_in_queue_dataset( + queue_rest, (count_rest, value_rest)) + with tf.control_dependencies([enqueue_rest_op]): + calculation = fn(value_head) + + while True: # Will raise OutOfRange when finished with all pieces. + session.run(calculation) + ``` + + Args: + batch_size: `int64` scalar tensor. The batch size to use when performing + padded batching. + padding_values: (optional) Nested tuple of scalar tensors. If provided, + the structure and dtypes of padding_values should match that of + incoming dataset's `output_types`. + padded_shapes: (optional) Nested tuple of `int64` vector tensors. + If provided, the structure must match that of the incoming dataset's + `output_types`. If not provided, the incoming dataset's `output_shapes` + is used. Any unknown (`None` or `-1`) dimensions in the shapes are + treated as being unique per-batch: for each batch time, an unknown + dimension is replaced with the maximum given value of this dimension + across all tensors for the given component in the batch. + + Returns: + A `Dataset` transformation function, which can be passed to + @{tf.data.Dataset.apply}. + """ + + def _apply_fn(dataset): + return _PrependFromQueueAndPaddedBatchDataset( + dataset, + batch_size=batch_size, + padding_values=padding_values, + padded_shapes=padded_shapes) + + return _apply_fn + + +def enqueue_in_queue_dataset(queue, components): + """Enqueue components into queue from `PrependFromQueueAndPaddedBatchDataset`. + + The components' dtypes and shapes must be compatible with the `output_shapes` + attribute of the `dataset` created by + `prepend_from_queue_and_padded_batch_dataset`. This operation supports both + non-batched and batched modes. + + For more details, see the example in the docstring for + `prepend_from_queue_and_padded_batch_dataset`. + + Args: + queue: `variant` scalar or vector tensor. + The tensor emitted by the first component of the iterator associated with + `prepend_from_queue_and_padded_batch_dataset`. If this is a scalar, + then the `components` input tensors should not have a prepended batch + dimension. + components: Nested tuple of tensors, each with a leading batch dimension + if `queue` is a vector. The structure, dtypes, and shapes + (excluding batch dimension) must match the nested tuples + `dataset.output_types[1]` and `dataset.output_shapes[1]` (the non-queue + output types and shapes) of the `dataset` emitted by + the original `prepend_from_queue_and_padded_batch_dataset` call. + + Returns: + An `Operation` that enqueues `components` into the dataset(s) associated + with entries of `queue`. + """ + return gen_dataset_ops.enqueue_in_queue_dataset( + queue=queue, components=tf_nest.flatten(components)) diff --git a/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py b/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py new file mode 100644 index 0000000000000000000000000000000000000000..0338f409a203c232e63e99534a8f6d6a43fa661e --- /dev/null +++ b/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py @@ -0,0 +1,355 @@ +# Copyright 2016 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 TensorQueueDataset.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base +from tensorflow.contrib.training.python.training import tensor_queue_dataset as tqd +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors +from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import string_ops +from tensorflow.python.platform import test + + +class PrependFromQueueAndPaddedBatchDatasetTest(test.TestCase): + + def testNoEnqueue(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) + self.assertEqual((dtypes.variant, dtypes.int32), dataset.output_types) + self.assertAllEqual(([None],) * 2, + [x.as_list() for x in dataset.output_shapes]) + iterator = dataset.make_one_shot_iterator() + _, value = iterator.get_next() + self.assertEqual([0], self.evaluate(value)) + self.assertEqual([1], self.evaluate(value)) + self.assertEqual([2], self.evaluate(value)) + with self.assertRaisesOpError("End of sequence"): + self.evaluate(value) + + def testBatchedNoEnqueue(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=2)) + iterator = dataset.make_one_shot_iterator() + _, value = iterator.get_next() + self.assertAllEqual([0, 1], self.evaluate(value)) + self.assertAllEqual([2], self.evaluate(value)) + with self.assertRaisesOpError("End of sequence"): + self.evaluate(value) + + def testBatchedWithBiggerPaddingNoEnqueue(self): + dataset = dataset_ops.Dataset.from_tensor_slices([[0], [1], [2]]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset( + batch_size=2, padded_shapes=[3])) + iterator = dataset.make_one_shot_iterator() + _, value = iterator.get_next() + self.assertAllEqual([[0, 0, 0], [1, 0, 0]], self.evaluate(value)) + self.assertAllEqual([[2, 0, 0]], self.evaluate(value)) + with self.assertRaisesOpError("End of sequence"): + self.evaluate(value) + + def testBatchedWithBiggerPaddingOneEnqueue(self): + dataset = dataset_ops.Dataset.from_tensor_slices([[0], [1], [2]]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset( + batch_size=1, padded_shapes=[3])) + iterator = dataset.make_one_shot_iterator() + queue_handle, value = iterator.get_next() + enqueue_negative = tqd.enqueue_in_queue_dataset(queue_handle, -value) + with self.test_session() as sess: + self.assertAllEqual([[0, 0, 0]], sess.run(value)) + value_1, _ = sess.run([value, enqueue_negative]) + self.assertAllEqual([[1, 0, 0]], value_1) + value_2, _ = sess.run([value, enqueue_negative]) + self.assertAllEqual([[-1, 0, 0]], value_2) + value_3 = sess.run(value) + self.assertAllEqual([[1, 0, 0]], value_3) + value_4, _ = sess.run([value, enqueue_negative]) + self.assertAllEqual([[2, 0, 0]], value_4) + value_5 = sess.run(value) + self.assertAllEqual([[-2, 0, 0]], value_5) + with self.assertRaisesOpError("End of sequence"): + sess.run(value) + + def testOneEnqueue(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) + iterator = dataset.make_one_shot_iterator() + queue_handle, value = iterator.get_next() + enqueue_negative = tqd.enqueue_in_queue_dataset(queue_handle, -value) + with self.test_session() as sess: + self.assertEqual([0], sess.run(value)) + value_1, _ = sess.run([value, enqueue_negative]) + self.assertEqual([1], value_1) + value_2, _ = sess.run([value, enqueue_negative]) + self.assertEqual([-1], value_2) + value_3 = sess.run(value) + self.assertEqual([1], value_3) + value_4, _ = sess.run([value, enqueue_negative]) + self.assertEqual([2], value_4) + value_5 = sess.run(value) + self.assertEqual([-2], value_5) + with self.assertRaisesOpError("End of sequence"): + sess.run(value) + + def testBatchedOneEnqueue(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=2)) + iterator = dataset.make_one_shot_iterator() + queue_handle, value = iterator.get_next() + enqueue_negative = tqd.enqueue_in_queue_dataset(queue_handle, -value) + enqueue_zeroth = tqd.enqueue_in_queue_dataset([queue_handle[0]], + array_ops.expand_dims( + value[0], axis=0)) + with self.test_session() as sess: + value_0, _ = sess.run([value, enqueue_negative]) + self.assertAllEqual([0, 1], value_0) + value_1, _ = sess.run([value, enqueue_zeroth]) + self.assertAllEqual([0, -1], value_1) + value_2, _ = sess.run([value, enqueue_negative]) + self.assertAllEqual([0, 2], value_2) + self.assertAllEqual([0, -2], sess.run(value)) + with self.assertRaisesOpError("End of sequence"): + sess.run(value) + + def testManyEnqueue(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0, 1]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) + iterator = dataset.make_one_shot_iterator() + queue_handle, value = iterator.get_next() + enqueue_many_more = [ + tqd.enqueue_in_queue_dataset(queue_handle, value + 100 + i) + for i in range(1000) + ] + with self.test_session() as sess: + value_0, _ = sess.run((value, enqueue_many_more)) + self.assertEqual([0], value_0) + rest = [] + for _ in range(1000): + rest.append(sess.run(value)) + self.assertEquals([[100 + i] for i in range(1000)], sorted(rest)) + # Going back to the original input. + value_1, _ = sess.run((value, enqueue_many_more)) + self.assertEqual(1, value_1) + rest = [] + for _ in range(1000): + rest.append(sess.run(value)) + self.assertEquals([[100 + i + 1] for i in range(1000)], sorted(rest)) + with self.assertRaisesOpError("End of sequence"): + sess.run(value) + + def testEnqueueWithPrefetch(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) + # Prefetching will request additional values before they are + # available to the queue. + dataset = dataset.prefetch(buffer_size=3) + iterator = dataset.make_one_shot_iterator() + queue_handle, value = iterator.get_next() + enqueue = tqd.enqueue_in_queue_dataset(queue_handle, value + 1) + with self.test_session() as sess: + i = 0 + while i < 4: + received, _ = sess.run((value, enqueue)) + if received.size > 0: + self.assertAllEqual([i], received) + i += 1 + received_last = False + while True: + try: + received = sess.run(value) + if received.size > 0: + self.assertAllEqual([4], received) + received_last = True + except errors.OutOfRangeError: + break + self.assertTrue(received_last) + + def testDatasetWithPaddedShapeSmallerThanInputFails(self): + dataset = dataset_ops.Dataset.from_tensor_slices([[0, 0, 0]]).repeat(None) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset( + batch_size=1, padded_shapes=[2])) + iterator = dataset.make_one_shot_iterator() + _, value = iterator.get_next() + with self.test_session() as sess: + with self.assertRaisesOpError( + r"Incompatible input shapes at component 0 between " + r"input dataset this dataset: \[3\] vs. \[2\]"): + sess.run(value) + + def testEnqueueWithIncompatibleInputsFailsWithInformativeError(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0]).repeat(None) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) + iterator = dataset.make_one_shot_iterator() + queue_handle, value = iterator.get_next() + + enqueue_bad_structure = tqd.enqueue_in_queue_dataset( + queue_handle, (value, value)) + enqueue_bad_dtype = tqd.enqueue_in_queue_dataset(queue_handle, + np.array( + [1.0], + dtype=np.float32)) + enqueue_bad_shape_no_batch_dim = tqd.enqueue_in_queue_dataset( + queue_handle, ([1],)) + enqueue_bad_shape = tqd.enqueue_in_queue_dataset(queue_handle, + np.array( + [[1]], dtype=np.int32)) + + with self.test_session() as sess: + with self.assertRaisesOpError( + "mismatched number of tensors. Queue expects 1 tensors but " + "tried to insert 2"): + sess.run(enqueue_bad_structure) + with self.assertRaisesOpError(r"Expected component 0 to have batched " + r"shape \[1,...\], but saw shape: \[\]"): + sess.run(enqueue_bad_shape_no_batch_dim) + with self.assertRaisesOpError( + r"mismatched shapes at component 0. Attempted to insert tensor " + r"with shape \[1\] but queue expected shape: \[\]"): + sess.run(enqueue_bad_shape) + with self.assertRaisesOpError( + r"mismatched dtypes at component 0. Attempted to insert tensor " + r"of type float but queue expected type: int32"): + sess.run(enqueue_bad_dtype) + + def testEnqueueWithPaddedBatchFailsWithInformativeError(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) + with self.assertRaisesRegexp( + TypeError, r"Unable to create padding for field of type 'variant'"): + dataset.padded_batch(batch_size=10, padded_shapes=[1]) + + def testOneEnqueueWithPadding(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0, 2, 4, 6]) + # Make a dataset of variable-length vectors and their lengths. + dataset = dataset.map( + lambda c: (c, c * array_ops.ones((c,), dtype=c.dtype))) + # Emit a queue we can prepend to, and counts/values as padded + # batch. + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=3)) + + iterator = dataset.make_one_shot_iterator() + queue, (count, padded_value) = iterator.get_next() + + # Split the padded_value into two pieces: head and rest + rest_indices = array_ops.squeeze(array_ops.where(count > 2), axis=1) + bound = math_ops.minimum(2, math_ops.reduce_max(count)) + value_head = padded_value[:, :bound] + count_rest = array_ops.gather(count - 2, rest_indices) + value_rest = array_ops.gather(padded_value, rest_indices)[:, bound:] + queue_rest = array_ops.gather(queue, rest_indices) + enqueue_rest_op = tqd.enqueue_in_queue_dataset(queue_rest, + (count_rest, value_rest)) + with ops.control_dependencies([enqueue_rest_op]): + calc = array_ops.identity(value_head) + + with self.test_session() as sess: + self.assertAllEqual([[0, 0], [2, 2], [4, 4]], sess.run(calc)) + self.assertAllEqual([[4, 4], [6, 6]], sess.run(calc)) + self.assertAllEqual([[6, 6]], sess.run(calc)) + self.assertAllEqual([[6, 6]], sess.run(calc)) + # Get some final batches due to prefetching. + for _ in range(3): + try: + self.assertAllEqual( + np.empty(shape=(0, 0), dtype=np.int32), sess.run(calc)) + except errors.OutOfRangeError as e: + self.assertTrue(str(e).startswith("End of sequence")) + + def testNonstandardPadding(self): + dataset = dataset_ops.Dataset.from_tensor_slices([0, 2, 4, 6]) + # Make a dataset of variable-length vectors and their lengths. + dataset = dataset.map( + lambda c: (c, c * array_ops.ones((c,), dtype=c.dtype))) + # Emit a queue we can prepend to, and counts/values as padded + # batch. + dataset = dataset.apply( + tqd.prepend_from_queue_and_padded_batch_dataset( + batch_size=3, padding_values=( + 0, + -1, + ))) + + iterator = dataset.make_one_shot_iterator() + _, (unused_count, padded_value) = iterator.get_next() + + with self.test_session() as sess: + self.assertAllEqual([[-1, -1, -1, -1], [2, 2, -1, -1], [4, 4, 4, 4]], + sess.run(padded_value)) + self.assertAllEqual([[6] * 6], sess.run(padded_value)) + with self.assertRaisesOpError("End of sequence"): + sess.run(padded_value) + + +# TODO(ebrevdo): Figure out how to use run_core_tests to test state +# saving of an iterator that's had some tensors enqueued into its queue. +class PrependFromQueueAndPaddedBatchDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): + + def testPrependFromQueueAndPaddedBatch(self): + + def build_dataset(seq_lens): + return dataset_ops.Dataset.from_tensor_slices(seq_lens).map( + lambda x: array_ops.fill([x], x)).apply( + tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=4)) + + seq_lens1 = np.random.randint(1, 20, size=(32,)).astype(np.int32) + seq_lens2 = np.random.randint(21, 40, size=(32,)).astype(np.int32) + self.run_core_tests(lambda: build_dataset(seq_lens1), + lambda: build_dataset(seq_lens2), 8) + + def testPrependFromQueueAndPaddedBatchNonDefaultPadding(self): + + def build_dataset(seq_lens): + + def fill_tuple(x): + filled = array_ops.fill([x], x) + return (filled, string_ops.as_string(filled)) + + padded_shape = [-1] + return dataset_ops.Dataset.from_tensor_slices(seq_lens).map( + fill_tuple).apply( + tqd.prepend_from_queue_and_padded_batch_dataset( + batch_size=4, + padded_shapes=(padded_shape, padded_shape), + padding_values=(-1, ""))) + + seq_lens1 = np.random.randint(1, 20, size=(32,)).astype(np.int32) + seq_lens2 = np.random.randint(21, 40, size=(32,)).astype(np.int32) + self.run_core_tests(lambda: build_dataset(seq_lens1), + lambda: build_dataset(seq_lens2), 8) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/util/convert_graphdef_memmapped_format_lib.cc b/tensorflow/contrib/util/convert_graphdef_memmapped_format_lib.cc index 2992a61ea8186caada394208e9c27ddffe896dd1..9675428e56e93c9669753371dbca47d56325b0c4 100644 --- a/tensorflow/contrib/util/convert_graphdef_memmapped_format_lib.cc +++ b/tensorflow/contrib/util/convert_graphdef_memmapped_format_lib.cc @@ -142,9 +142,9 @@ Status ConvertConstantsToImmutable(const string& in_graph_filename, const auto load_graph_status = ReadBinaryProto(default_env, in_graph_filename, &graph_def); if (!load_graph_status.ok()) { - return tensorflow::errors::NotFound("Failed to load graph at '", - in_graph_filename, "' : ", - load_graph_status.error_message()); + return tensorflow::errors::NotFound( + "Failed to load graph at '", in_graph_filename, + "' : ", load_graph_status.error_message()); } NodeConverter node_converter; diff --git a/tensorflow/contrib/util/convert_graphdef_memmapped_format_lib.h b/tensorflow/contrib/util/convert_graphdef_memmapped_format_lib.h index 6518e7a10f587a687f2ee3258f5399d74d87364e..61fc6f36f7e5211e43c279506faf09624086d167 100644 --- a/tensorflow/contrib/util/convert_graphdef_memmapped_format_lib.h +++ b/tensorflow/contrib/util/convert_graphdef_memmapped_format_lib.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_UTIL_CONVERT_GRAPHDEF_MEMMAPPED_FORMAT_LIB_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_UTIL_CONVERT_GRAPHDEF_MEMMAPPED_FORMAT_LIB_H_ +#ifndef TENSORFLOW_CONTRIB_UTIL_CONVERT_GRAPHDEF_MEMMAPPED_FORMAT_LIB_H_ +#define TENSORFLOW_CONTRIB_UTIL_CONVERT_GRAPHDEF_MEMMAPPED_FORMAT_LIB_H_ #include @@ -31,4 +31,4 @@ Status ConvertConstantsToImmutable(const string& in_graph_filename, } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_UTIL_CONVERT_GRAPHDEF_MEMMAPPED_FORMAT_LIB_H_ +#endif // TENSORFLOW_CONTRIB_UTIL_CONVERT_GRAPHDEF_MEMMAPPED_FORMAT_LIB_H_ diff --git a/tensorflow/contrib/util/inspect_checkpoint.cc b/tensorflow/contrib/util/inspect_checkpoint.cc index 39088aeaad68e26344b2e89ce10ae6da8026e481..9b578ceb07548b8d198f64bc859d31c92774a286 100644 --- a/tensorflow/contrib/util/inspect_checkpoint.cc +++ b/tensorflow/contrib/util/inspect_checkpoint.cc @@ -13,10 +13,10 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/core/platform/init_main.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/init_main.h" #include "tensorflow/core/util/tensor_slice_reader.h" namespace tensorflow { diff --git a/tensorflow/contrib/verbs/BUILD b/tensorflow/contrib/verbs/BUILD index 38a84ffb10e594568a18dbd06debf32545cb2229..80a5d07ea43531ed2532443b6ff9327b9ece6df7 100644 --- a/tensorflow/contrib/verbs/BUILD +++ b/tensorflow/contrib/verbs/BUILD @@ -99,7 +99,7 @@ cc_library( alwayslink = 1, ) -tf_cuda_library( +cc_library( name = "rdma_rendezvous_mgr", srcs = ["rdma_rendezvous_mgr.cc"], hdrs = ["rdma_rendezvous_mgr.h"], @@ -114,7 +114,7 @@ tf_cuda_library( ], ) -cc_library( +tf_cuda_library( name = "rdma_mgr", srcs = ["rdma_mgr.cc"], hdrs = ["rdma_mgr.h"], @@ -141,6 +141,8 @@ tf_cuda_library( "//conditions:default": [], }), deps = [ + ":grpc_verbs_client", + ":verbs_service_proto_cc", ":verbs_util", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", diff --git a/tensorflow/contrib/verbs/README.md b/tensorflow/contrib/verbs/README.md index 7c1c8ea45912be8c471efbe42f43e083639e91fc..58fed4e5cb4c24b0f21dfe9b99cf4c665d2591c7 100644 --- a/tensorflow/contrib/verbs/README.md +++ b/tensorflow/contrib/verbs/README.md @@ -24,66 +24,144 @@ The design is based on TensorFlow r1.0. An RDMA path is added between servers fo During the server setup, an RDMA manager is created to manage low-level RDMA components such as RDMA channel and RDMA adapter, an RDMA rendezvous manager is created to oversee send/recv operations between servers. Following the distributed TensorFlow design philosophy, the send operation is passive, i.e. merely placing a tensor in the local out-going table. It is the receive operation that actually initiates the tensor transfer. -TensorFlow dynamically allocates memory for tensors that are to be sent or received. This causes difficulty for RDMA operations where pinned memory is required. Two remedies are possible, either the memory is pinned, transfer, then unpinned for each and every tensor to be transferred, or a buffer is pre-allocated and pinned for each tensor. The former incurs significant operation overhead since pinning and unpinning memory for each dynamically generated tensor is slow. The latter incurs large memory overhead and extra copying from the tensor to its pinned buffer, but may still be faster than the former. The second approach is adopted in this design. Each RDMA channel, representing a RDMA connection to a peer, contains a table of pinned buffers for all the seen tensors that requires transfer. It is assumed that the tensor size rarely changes across different steps. So only one buffer is created for the same tensor across all the steps. In the rare case when the tensor size does increases, the old buffer is discarded and new buffer of larger size is created and pinned. +TensorFlow dynamically allocates memory for tensors that are to be sent or received. This causes difficulty for RDMA operations where pinned memory is required. Few remedies are possible: +1. The memory is pinned, transferred, then unpinned for each and every tensor to be transferred. This incurs significant operation overhead since pinning and unpinning memory for each dynamically generated tensor is slow. +2. Buffer is pre-allocated and pinned for each tensor. This incurs large memory overhead and extra copying from the tensor to its pinned buffer, but may still be faster than the former. +3. Following HKUST research on the use of GPU direct, and their [GDR implementation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/gdr/README.md), there is a smart way to benefit from the TensorFlow allocation theme which is mostly pool based, i.e allocators pre-allocate a large memory block, and allocate the tensors from there. By attaching a custom Visitor to relevant allocators, we can do a single registration of the entire memory block, which zeros the registration overhead. Once the block is registered, each new tensor allocated will be at a registered address, which will allow us to do direct RDMA writes to it. -When a tensor is prepared for transfer, it is first converted to TensorProto, then the proto is serialized to byte array and copied to the pinned buffer. The content of the buffer is transferred to the remote node via RDMA write. On the remote side, the process is reversed. This is illustrated in the diagram below. The conversion of TensorProto is introduced to simplify transfer of string-tensors. Also since the TensorProto lives in host memory, even if the origin tensor lives in the device, the pinned buffers are all allocated in the host memory. -![TensorFlow RDMA path](./design_diagram.png) +For best performance, we will adopt HKUST 0 copies approach in our solution. This means: + +1. Tensor writes will be done directly from the source tensor to the **result** tensor, with no memory copies in between. This should be done for all DMAable tensors which are located either on CPU or on a RDMA compatible GPU device (GPU direct). +2. Non DMAable tensors (CanMemCopy == false) will be serialized to a TensorProto on the sender side, RDMA written to a registered buffer on the receiver side, and then deserialized by the receiver. +3. Tensors which are located on a non-RDMA-compatible GPU, will be RDMA written to a registered CPU **proxy** buffer on the receiver side, and then copied to GPU by the receiver. -The following improvements can be made in the future. First, conversion to TensorProto and serialization can be avoided for numeric (float/int) tensors since their internal buffer can be access directly as byte array. Second, the pinned buffer may be allocated on device if the tensor is located in the device. This avoids extra device-to-host copy at the expense of extra device memory consumption. ## Design details -### RDMA components +### Terminology -* **RDMA adapter:** The base for RDMA communications. It may contain multiple channels and buffers. It is responsible for handling various incoming RDMA messages. -* **RDMA channel:** Responsible for RDMA connection to a particular node. It manages multiple buffers. A channel has a callback table which stores all the callbacks for the requested tensors. -* **RDMA buffer:** Responsible for sending or receiving data. It has a fixed size memory to store the data. It has a queue to store the pending jobs. There are three types of buffers, message buffer, ACK buffer and tensor buffer. A channel has two message buffers, two ack buffers and many tensor buffers. -* **RDMA manager:** Manages the adapter and channels, including channel creation, channel setup via GRPC service, channel lookup, etc. -* **RDMA rendezvous manager:** manages multiple rdma rendezvous. -* **RDMA rendezvous:** a derived class of BaseRemoteRendezvous. This class is the back end for "send" and "recv" ops. When the sendrecv_op wants to send or receive a tensor, it calls the rendezvous' "send" and "recv" functions respectively. Rendezvous are identified by "step_id", a random number, so that tensors for different iterations don't get mixed up. +* **Sender** - The node which sends the tensor. +* **Receiver** - The node which receives the tensor. +* **Result tensor** - The destination tensor, allocated on its appropriate device. +* **Proxy tensor** - A CPU allocated tensor, which will be used in the case where the result tensor cannot be RDMA written to directly (GPU direct is disabled or not available). The RDMA write will therefore be done to the proxy tensor, and afterwards we will do a manual local copy from it to the result tensor. -### The SEND operation +### Messages -In TensorFlow, when rendezvous sends a tensor, it merely puts a tensor in a local table in the corresponding rendezvous. If the tensor has been requested, a callback exists in the table. "send" will activate the callback, which tries to send the tensor across the node. +* RDMA_MESSAGE_TENSOR_REQUEST +* RDMA_MESSAGE_META_DATA_RESPONSE +* RDMA_MESSAGE_TENSOR_RE_REQUEST +### Transport protocol -### The RECV operation +The tensor transfer process is initiated when the receiver requests a tensor. In code it is done by calling **Rendezvous::Recv()** or **Rendezvous::RecvAsync()**. The TensorFlow base implementation handles the case where the requested tensor is located on the same node. The more interesting case where the requested tensor is located on a remote node (receiver != sender) is to be handled in a derivation of the pure virtual **BaseRemoteRendezvous::RecvFromRemoteAsync()**. TensorFlow provides a default GRPC based implementation which comes in the vanilla version but suffers in scalability when running large models. Our RDMA based implementation presumes to be more scalable. HKUST's contrib GDR implementation is more scalable than GRPC, and less scalable than ours, only because we did our evolution based on it. -When a tensor is requested, rendezvous' recv function is called. The function first places a callback in the channel's callback table, which will be activated once the tensor is sent from the source. In the next step, a message is sent to notify the source of the requested tensor. Once the source receives the message, it will check locally for the tensor, if not found, a callback is placed in the table, otherwise, the tensor id will be placed at corresponding RDMA buffer's job queue for future transmission. When a tensor is scheduled to be transmitted, the RDMA buffer needs to have the memory allocated and initialized (registered with the remote buffer info). If the memory is not ready, the transmission is deferred, a message is sent to the destination to establish the memory first. The other case a transmission can be deferred is when the buffer is still being used by an on-going transmission. +Our entry point is the implementation of **RdmaRemoteRendezvous::RecvFromRemoteAsync()**, located in rdma_rendezvous_mgr.cc. The implementation creates a new **RdmaTensorRequest** object, keyed by request index (uint32_t), stores it in a list of pending requests, and calls its **Start()** method. The **Start()** method basically does 2 things: -### Three types of RDMA buffers +1. Allocate the result tensor (and the proxy tensor if required). +2. Send a **RDMA_MESSAGE_TENSOR_REQUEST** to the sender, containing the address of the destination tensor (result/proxy) for RDMA write. -* **Message buffer:** responsible for sending message only. -* **Ack buffer:** once a message is sent, the recipient needs to send an ack via the ack buffer to free up the message buffer. An ack buffer is exclusively for its coupled message buffer. -* **Tensor buffer:** responsible for sending tensors. The recipient needs to send back a message to free up the sending buffer. +In order to allocate the result and proxy tensors, we need to know the tensor's meta-data, i.e. shape and data-type for DMAable tensors, and proto-size for serialized tensors. Unfortunately, this information is only available on the sender side which complicates manners. In order to avoid sending extra messages for querying the meta-data at each step, we store a local meta-data cache per tensor, which will only be update upon changes. Based on the assumption that the meta-data of a tensor rarely changes between steps, we expect that on most times the cache will only be updated once. The sender is responsible to detect changes in the meta-data, and update the receiver. In order for the sender to know that the meta-data had changed, each **RDMA_MESSAGE_TENSOR_REQUEST** will contain the meta-data that the receiver had grabbed from the local cache. The sender will then compare the meta-data from the message to the tensor's new meta-data. -### RDMA packet format +When the sender receives an **RDMA_MESSAGE_TENSOR_REQUEST**, it will create a new **RdmaTensorResponse** object for the given request message, store it in a list of pending responses, and will invoke its **Start()** method. The **Start()** method does the following: -|type|name_size|name|step_id|buffer_size|remote_addr|rkey|is_dead|data_type|tensor_shape|tensor_bytes|tensor_buffer| +1. Grab the source tensor from the local table (In code, **RecvLocalAsync()**). +2. If the source tensor is not DMAable, serialize it to a TensorProto. +3. If the source tensor is located on a device which cannot be DMA written from, copy it to CPU. +4. If it is the first time this tensor is requested, or if the tensor's meta-data changed: + 1. Clone the tensor's data to be sent later. + 2. Send a **RDMA_MESSAGE_META_DATA_RESPONSE** containing the new meta-data. +5. Otherwise: + 1. RDMA write the tensor (or TensorProto) to the destination address and rkey specified in the request message. The immediate value for the write will be the request index. -### Six types of RDMA messages -* RDMA_MESSAGE_ACK -* RDMA_MESSAGE_BUFFER_IDLE -* RDMA_MESSAGE_BUFFER_REQUEST -* RDMA_MESSAGE_BUFFER_RESPONSE -* RDMA_MESSAGE_TENSOR_REQUEST -* RDMA_MESSAGE_TENSOR_WRITE - -### Actions upon receiving RDMA messages -* RDMA_MESSAGE_ACK - * sender: mark local ack buffer idle. - * receiver: mark remote message buffer idle, send next item. -* RDMA_MESSAGE_BUFFER_IDLE - * sender: mark local message buffer idle, send next item. - * receiver: send ack, set remote tensor buffer idle, send next item. -* RDMA_MESSAGE_BUFFER_REQUEST - * sender: mark local message buffer idle, send next item. - * receiver: send ack, find or create tensor buffer, send BUFFER_RESPONSE. -* RDMA_MESSAGE_BUFFER_RESPONSE - * sender: mark local message buffer idle, send next item. - * receiver: send ack, set remote buffer info, set local and remote buffer idle, send next item. -* RDMA_MESSAGE_TENSOR_REQUEST - * sender: mark local message buffer idle, send next item. - * receiver: send ack, find or create tensor buffer, enqueue tensor id, send next item. -* RDMA_MESSAGE_TENSOR_WRITE - * sender: mark local message buffer idle, send next item. - * receiver: run callback. + +When the receiver receives the **RDMA_MESSAGE_META_DATA_RESPONSE**, it will locate the relevant **RdmaTensorRequest** using the request index specified in the message, and invoke its **RecvTensorMetaData()** which does the following: + +1. Update the local meta-data cache. +2. Reallocate the result/proxy tensors. +3. Re-send the tensor request. For traceability, the new message has a different name: **RDMA_MESSAGE_TENSOR_RE_REQUEST**. + +When the sender receives a **RDMA_MESSAGE_TENSOR_RE_REQUEST**, it will locate the relevant **RdmaTensorResponse** using the request index specified in the message, and invoke its **Resume()** method, which will RDMA write the contents of the tensor that was cloned earlier, to the new remote address specified in the re-request. + +When the receiver receives the RDMA write, it will locate the relevant **RdmaTensorRequest** using the request index which is the immediate value. It will then invoke its **RecvTensorContent()** which does the following: + +1. Proxy copy/deserialize if required. +2. Invoke the done callback. +3. Deallocate the result/proxy tensors and remove the request from the pending list. + +![alt text](verbs_with_0_copies.png "Transport protocol") + +### Additional design notes + +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. + 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. +5. ACK writes are empty writes (hence they require no buffer) with immediate value 0xFFFFFFFE. Message writes have the immediate value 0xFFFFFFFF. All other writes are tensor-content writes whose immediate value is the request-index. + +### RDMA components + +* **enum RdmaImmDataType** - Immediate types to distinguish between different RDMA writes on the remote side. Ack writes and control-message writes have a fixed immediate value. The rest of the writes are tensor writes and the immediate value is the relevant request index. +* **enum RdmaWriteIDType** - Types to distinguish between different RDMA write-complete events: Ack, control message and tensor writes. +* **class RdmaWriteID** - Context for RDMA write complete events. Holds the RdmaWriteIDType and additional data. +* **class RdmaTensorMetaData** - Meta-data for a tensor (type, shape, is_dead, proto_size). +* **class RdmaMemoryMgr** - Manages the meta-data cache, and the registered memory regions. +* **class RdmaTensorRequest** - Holds and manages information for a single tensor request throughout the entire receive cycle. API: + * **Start()** - Start the request sequence. + * Allocate the result tensor (and proxy tensor if required). + * Send RDMA_MESSAGE_TENSOR_REQUEST to the remote side. + * **RecvTensorMetaData()** - Receive meta-data from the remote side. + * Update the local meta-data cache. + * Reallocate the result tensor (and proxy tensor if required). + * Re-send the request to the remote side. + * **RecvTensorContent()** - Receive tensor content from the remote side (RDMA write was completed). + * Decode proto if required and/or move to GPU if the content was not written to it directly (GPU direct is not available). + * Invoke the done callback. +* **class RdmaTensorResponse** - Holds and manages information for a single tensor response throughout the entire send cycle. API: + * **Start()** - Start the response sequence. + * Find the tensor in the local tag-match table. + * Compare the tensor's meta-data to the meta-data in the message (taken from the requester's local cache). + * If meta-data changed: + * Clone the tensor to be sent later. + * Send a meta-data update message and wait for re-request. + * Else: + * Send the tensor's content (using direct RDMA write). + * **Resume()** - Resume the response sequence after a re-request. Send the tensor's content that was cloned earlier. + * **Destroy()** - Destroy the response's resources and remove it form the pending list. +* **class RdmaAdapter** - The base for RDMA communications. It may contain multiple channels and buffers. It is responsible for handling various incoming RDMA messages. +* **class RdmaChannel** - Responsible for RDMA connection to a particular node. It manages messagee buffers. A channel has a request table which stores all the pending tensor requests. +* **class RdmaMessageBuffer** - Responsible for sending or receiving messages. It has a fixed size memory to store the data. It has a queue to store the pending jobs. A channel has two message buffers one for tx and one for rx. +* **class RdmaMgr** - Manages the adapter and channels, including channel creation, channel setup via GRPC service, channel lookup, etc. +* **class RdmaRendezvousMgr** - Manages multiple rdma rendezvous. +* **class RdmaRemoteRendezvous** - A derived class of BaseRemoteRendezvous. This class is the back end for "send" and "recv" ops. When the sendrecv_op wants to send or receive a tensor, it calls the rendezvous' "send" and "recv" functions respectively. Rendezvous are identified by "step_id", a random number, so that tensors for different iterations don't get mixed up. + +### Message structure: + +| type | name_size | name | step_id | request_index | remote_addr/checksum | rkey | is_dead | data_type | tensor_shape | tensor_bytes | error_status | +|------|---------- |------|---------|---------------|----------------------|------|---------|-----------|--------------|--------------|-----------------------| +| 1B | 2B | 512 | 8B | 8B | 8B | 4B | 1B | XB | XB | 8B | Size - 4B, proto - XB | + +* **RDMA_MESSAGE_TENSOR_REQUEST** - (receiver ==> sender) The original tensor request. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * remote_addr/rkey - Address/rkey of the result/proxy tensor. Irrelevant for first-time request. + * is_dead/data_type/tensor_shape/tensor_bytes - The current meta-data as stored in the receiver local cache. The sender will use that information to know if the receiver's cache requires updating. +* **RDMA_MESSAGE_META_DATA_RESPONSE** - (sender ==> receiver) The meta-data update message in case meta-data had changed (or if it is the first time the tensor is requested). + * type - The message type. + * request_index - Request index. + * is_dead/data_type/tensor_shape/tensor_bytes - The up-to-date meta-data. + * checksum - In data validation mode, this will hold the checksum of the source tensor. +* **RDMA_MESSAGE_TENSOR_RE_REQUEST** - (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. + * request_index - Request index. + * remote_addr/rkey - Address/rkey of the reallocated result/proxy tensor. +* **RDMA_MESSAGE_ERROR_STATUS** - (sender ==> receiver) Notify the receiver that an error had occured on the sender side, so it can propagate it to the upper levels. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * error_status - The error status (code, message, details). diff --git a/tensorflow/contrib/verbs/grpc_verbs_client.h b/tensorflow/contrib/verbs/grpc_verbs_client.h index 358977f92543e1a38b594cf45cdbff34f89277be..2cfaa4986cb0923d9687cb77b8e1116a937594a1 100644 --- a/tensorflow/contrib/verbs/grpc_verbs_client.h +++ b/tensorflow/contrib/verbs/grpc_verbs_client.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_GRPC_VERBS_CLIENT_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_GRPC_VERBS_CLIENT_H_ +#ifndef TENSORFLOW_CONTRIB_GRPC_VERBS_CLIENT_H_ +#define TENSORFLOW_CONTRIB_GRPC_VERBS_CLIENT_H_ #include "tensorflow/contrib/verbs/grpc_verbs_service_impl.h" #include "tensorflow/contrib/verbs/verbs_service.pb.h" @@ -47,4 +47,4 @@ class GrpcVerbsClient { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_GRPC_VERBS_CLIENT_H_ +#endif // TENSORFLOW_CONTRIB_GRPC_VERBS_CLIENT_H_ diff --git a/tensorflow/contrib/verbs/grpc_verbs_service.cc b/tensorflow/contrib/verbs/grpc_verbs_service.cc index f2af6b79fba6a480afbfe88fcbefcbf8a6670ce6..742f946c9536973eb8a6a11afda1b32ae4a7726b 100644 --- a/tensorflow/contrib/verbs/grpc_verbs_service.cc +++ b/tensorflow/contrib/verbs/grpc_verbs_service.cc @@ -122,17 +122,15 @@ Status GrpcVerbsService::GetRemoteAddressSync( rc->SetRemoteAddress(ra, false); rc->Connect(); int i = 0; - int idx[] = {1, 0, 3, 2}; - std::vector mb(rc->message_buffers()); - CHECK_EQ(request->mr_size(), 4); + int idx[] = {1, 0}; + std::vector mb(rc->message_buffers()); + CHECK_EQ(request->mr_size(), RdmaChannel::kNumMessageBuffers); for (const auto& mr : request->mr()) { // the connections are crossed, i.e. // local tx_message_buffer <---> remote rx_message_buffer_ // local rx_message_buffer <---> remote tx_message_buffer_ - // local tx_ack_buffer <---> remote rx_ack_buffer_ - // local rx_ack_buffer <---> remote tx_ack_buffer_ - // hence idx[] = {1, 0, 3, 2}. - RdmaBuffer* rb = mb[idx[i]]; + // hence idx[] = {1, 0}. + RdmaMessageBuffer* rb = mb[idx[i]]; RemoteMR rmr; rmr.remote_addr = mr.remote_addr(); rmr.rkey = mr.rkey(); diff --git a/tensorflow/contrib/verbs/grpc_verbs_service.h b/tensorflow/contrib/verbs/grpc_verbs_service.h index aa509602b51e7749547f1ff8eb5193acd1a3ec65..444c863b942ef8bce8d54d59765563b12eb6087e 100644 --- a/tensorflow/contrib/verbs/grpc_verbs_service.h +++ b/tensorflow/contrib/verbs/grpc_verbs_service.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_GRPC_VERBS_SERVICE_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_GRPC_VERBS_SERVICE_H_ +#ifndef TENSORFLOW_CONTRIB_VERBS_GRPC_VERBS_SERVICE_H_ +#define TENSORFLOW_CONTRIB_VERBS_GRPC_VERBS_SERVICE_H_ #ifdef TENSORFLOW_USE_VERBS @@ -69,4 +69,4 @@ void SetNewVerbsService(GrpcVerbsService** handle, const WorkerEnv* worker_env, } // namespace tensorflow #endif // TENSORFLOW_USE_VERBS -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_GRPC_VERBS_SERVICE_H_ +#endif // TENSORFLOW_CONTRIB_VERBS_GRPC_VERBS_SERVICE_H_ diff --git a/tensorflow/contrib/verbs/grpc_verbs_service_impl.h b/tensorflow/contrib/verbs/grpc_verbs_service_impl.h index 86431ca030c38c56155801202714ee4a49b764df..1f0f10517e98a32ae882c027330091928f1a6ee2 100644 --- a/tensorflow/contrib/verbs/grpc_verbs_service_impl.h +++ b/tensorflow/contrib/verbs/grpc_verbs_service_impl.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_GRPC_VERBS_SERVICE_IMPL_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_GRPC_VERBS_SERVICE_IMPL_H_ +#ifndef TENSORFLOW_CONTRIB_GRPC_VERBS_SERVICE_IMPL_H_ +#define TENSORFLOW_CONTRIB_GRPC_VERBS_SERVICE_IMPL_H_ #include "grpc++/impl/codegen/async_stream.h" #include "grpc++/impl/codegen/async_unary_call.h" @@ -86,4 +86,4 @@ class VerbsService GRPC_FINAL { } // namespace tensorflow -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_GRPC_VERBS_SERVICE_IMPL_H_ +#endif // TENSORFLOW_CONTRIB_GRPC_VERBS_SERVICE_IMPL_H_ diff --git a/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md b/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md new file mode 100644 index 0000000000000000000000000000000000000000..956b8f2147cf8154b6f1ade006d7bff194864c9b --- /dev/null +++ b/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md @@ -0,0 +1,87 @@ +## Verbs implementation to use direct tensor writes (0 copies) + +### Motivation: + +Following HKUST research on the use of GPU direct, and their [GDR implementation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/gdr/README.md), we wish to adopt the 0 copies approach and apply it to the current verbs implementation, while keeping the current implementation advantages, such as configurability and the use of RDMA for control messages. + +### Performance: + +Compared with the current GRPC, verbs and GDR implementation, the result implementation gave the best performance for every model, with any number of nodes. For VGG16 on 8 nodes with 4 P100 GPUs each, the prototype beat the second place by over 15%. + +### Implementation requirements: + +1. Tensor writes need to be done directly from the source Tensor to the destination Tensor, with no memory copies in between. This should be done for all DMAble tensors which are located either on CPU or on a RDMA compatible GPU device (GPU direct). +2. Non DMAble tensors (CanMemCopy == false) will be serialized to proto on the sender side, RDMA written to a registered buffer on the receiver side, and then deserialized by the receiver. +3. Tensors which are located on a non-RDMA-compatible GPU, will be RDMA written to a registered CPU proxy buffer on the receiver side, and then copied to GPU by the receiver. + +### Implementation constrains: + +For best stability and proof of correctness, we will divide the implementation to two stages: +1. At first stage we will keep changes to the current implementation to the minimum possible. The expense will be that we may have unused or unnecessary code leftovers, which may also affect performance. +2. At second stage, we will re-iterate over the code and remove irrelevant code parts. +The design of the solution aims that we will achieve both stages with relative ease. + +### Design guidelines: + +1. Since we do not want to do any unnecessary memory copying, we will no longer allocate a fixed CPU buffer as the destination for the RDMA write. Instead we will do the writing directly to the result tensor, or if the result tensor is on a device which does not support RDMA, we will do the writing to a proxy CPU tensor and then copy its content to the result tensor. +2. The address of the destination Tensor needs to be sent to the sender side for writing, meaning that the result/proxy tensor should be pre-allocated on the receiver side, prior to sending the tensor request. In order to do that, we need to know its meta-data, i.e. shape and data-type for DMAble tensors, and proto-size for serialized tensors. Unfortunately, this information is only available on the sender side which complicates manners. In order to avoid sending extra messages for querying the meta-data on each step, we store a local meta-data cache per tensor. Based on the assumption that the meta-data of a tensor rarely changes between steps, we expect that on most times the cache will only be updated once. When the sender receives a request for a tensor, if it is the first time this tensor is requested, or in the rare case that the meta-data did change, the sender will first send a meta-data response, on which the receiver will update the local cache, and reallocate the result/proxy tensors if required. When the receiver sends the tensor request, it will contain also the meta-data currently stored in its local cache, so the sender can compare it to see if there was a change. +3. When the sender writes the tensor content to the result tensor, no additional data is being written with it. That means we need to reside on ibverbs immediate (uint32_t) to indicate which request we are responding to (in order to trigger the receive callback). The easiest and most elegant way is to key the recv callback with a unique request_index (uint32_t), instead of the current key_with_step_id (string). +4. Since the sender no longer writes the tensor from/to fixed buffers, we no longer need to schedule the writes using the local/remote status. In addition we no longer rely on the RmdaTensorBuffer members as the source/destination addresses and rkey/lkey. Instead, each RdmaTensorBuffer will hold multiple "Response" objects (one per step-id), from which we derive destination address and rkey. The source address and lkey are always the ones of the source Tensor. +5. With the addition of tensor pre-allocation, we noticed there is a large code similarity between sending the first tensor request and re-sending the request in case of meta-data changes. After implementing a common method for tensor pre-allocation, it turned out that implementation becomes much simpler by encapsulating the process of request sending/re-sending, meta-data response callback and content response callback, all in a single "Request" class. The request class holds all the relevant request information, which reduces excessive parameter passing and lambda capturing. This decision is purely for elegance and code simplicity, and we decided to implement it in first stage because it makes the implementation much easier. + +### New types/classes: + +* **enum RdmaImmDataType** - Immediate types to distinguish between different RDMA writes on the remote side. Ack writes and control-message writes have a fixed immediate value. The rest of the writes are tensor writes and the immediate value is the relevant request index. +* **enum RdmaWriteIDType** - Types to distinguish between different RDMA write-complete events: Ack, control message, tensor DMA write and tensor proto write. +* **class RdmaWriteID** - Context for RDMA write complete events. Holds the RdmaWriteIDType and additional data. +* **class RemoteAddressContext** - Remote address information (address + mr). Will be passed as write context for tensor proto writes. +* **class RdmaTensorMetaData** - Meta-data for a tensor (type, shape, is_dead, proto_size). +* **class RdmaMemoryMgr** - Manages the meta-data cache, and the registered memory regions. +* **class RdmaTensorRequest** - Holds and manages information for a single tensor request throughout the entire receive cycle. API: + * Start() - Start the request. + * RecvTensorMetaData() - Receive meta-data from the remote side. + * RecvTensorContent() - Receive tensor content from the remote side and invoke the done() callback. +* **class RdmaTensorResponse** - Holds information for a single tensor response, such as destination address and rkey. + +### Protocol changes: + +The protocol messages themselves will remain mostly unchanged at the first stage, but will be used differently, as described below. The current messages structures already have most of the required fields for the new implementation. The only change is the "buffer_size" field which is no longer used since we are no longer sending additional information with the tensor, and thus it is now always equal to the "tensor_bytes" field. Instead, we use that field to pass the "request_index". + +### Message structure: + +| type | name_size | name | step_id | request_index | remote_addr | rkey | is_dead | data_type | tensor_shape | tensor_bytes | +|------|---------- |------|---------|---------------|-------------|------|---------|-----------|--------------|--------------| +| 1B | 2B | 512 | 8B | 8B | 8B | 4B | 1B | XB | XB | 8B | + +* **RDMA_MESSAGE_TENSOR_REQUEST** - (receiver ==> sender) The original tensor request. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * remote_addr/rkey - Address/rkey of the result/proxy tensor. Irrelevant for first-time request. + * is_dead/data_type/tensor_shape/tensor_bytes - The current meta-data as stored in the receiver local cache. The sender will use that information to know if the receiver's cache requires updating. +* **RDMA_MESSAGE_BUFFER_REQUEST** - (sender ==> receiver) The meta-data update message in case meta-data had changed (or if it is the first time the tensor is requested). + * 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. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * remote_addr/rkey - Address/rkey of the reallocated result/proxy tensor. + * is_dead/data_type/tensor_shape/tensor_bytes - The new meta-data. Will be removed in the next phase. +* **RDMA_MESSAGE_TENSOR_WRITE** - (sender ==> receiver) No longer sent. There is only a direct write of the tensor content to the result/proxy tensor. Request index passed as the immediate value of the write. +* **RDMA_MESSAGE_TENSOR_IDLE** - (receiver ==> sender) No longer sent. + +![alt text](verbs_with_0_copies_phase1_protocol.jpg "Phase 1 message protocol") + +### Second stage optimizations: +1. Remove unused code leftovers. +2. Remove the ACK buffer completely, since we can rely completely on its immediate value. + +### Future optimizations: +1. Map the tensor names to indexes, to significantly reduce the request message size. +2. Understand the purpose of empty tensors and if we can skip remote fetching for them. +3. Consider concatenating multiple requests and/or using multiple message buffers. +4. Consider a no-request architecture. diff --git a/tensorflow/contrib/verbs/rdma.cc b/tensorflow/contrib/verbs/rdma.cc index ae9a384565a6ad0e63a6cf3acf07c591c65f0637..86350a08e57e5050f18d019fe80d70f6381c1f7d 100644 --- a/tensorflow/contrib/verbs/rdma.cc +++ b/tensorflow/contrib/verbs/rdma.cc @@ -15,58 +15,48 @@ limitations under the License. #ifdef TENSORFLOW_USE_VERBS -#include "tensorflow/contrib/verbs/rdma.h" #include #include -#include -#include "tensorflow/contrib/verbs/verbs_util.h" + +#include "tensorflow/contrib/verbs/rdma.h" +#include "tensorflow/contrib/verbs/verbs_service.pb.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/dma_helper.h" +#include "tensorflow/core/common_runtime/process_util.h" #if GOOGLE_CUDA #include "tensorflow/core/common_runtime/gpu/gpu_util.h" #include "tensorflow/core/common_runtime/gpu/process_state.h" #endif #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" #include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/lib/core/threadpool.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/random/random.h" -#include "tensorflow/core/lib/core/threadpool.h" namespace tensorflow { #define RoCE_V2 "RoCE v2" namespace { -// hash name to 32-bit integer -uint32_t NameHash(const string& name) { - return Hash32(name.data(), name.size(), 0x1234ABCD); -} // convenience function for printing message string MessageTypeToString(RdmaMessageType rmt) { switch (rmt) { - case RDMA_MESSAGE_ACK: - return "RDMA_MESSAGE_ACK"; - break; - case RDMA_MESSAGE_BUFFER_IDLE: - return "RDMA_MESSAGE_BUFFER_IDLE"; + case RDMA_MESSAGE_META_DATA_UPDATE: + return "RDMA_MESSAGE_META_DATA_UPDATE"; break; - case RDMA_MESSAGE_BUFFER_REQUEST: - return "RDMA_MESSAGE_BUFFER_REQUEST"; - break; - case RDMA_MESSAGE_BUFFER_RESPONSE: - return "RDMA_MESSAGE_BUFFER_RESPONSE"; + case RDMA_MESSAGE_TENSOR_RE_REQUEST: + return "RDMA_MESSAGE_TENSOR_RE_REQUEST"; break; case RDMA_MESSAGE_TENSOR_REQUEST: return "RDMA_MESSAGE_TENSOR_REQUEST"; break; - case RDMA_MESSAGE_TENSOR_WRITE: - return "RDMA_MESSAGE_TENSOR_WRITE"; - break; default: return "UNKNOWN MESSAGE"; } @@ -347,7 +337,7 @@ uint32_t set_param(uint32_t default_val, const char* env_param) { enum ibv_mtu set_mtu(uint8_t port_num, ibv_context* context) { ibv_port_attr port_attr; - enum ibv_mtu mtu; + enum ibv_mtu mtu = IBV_MTU_512; string mtu_s; int rc, mtu_i; @@ -459,106 +449,79 @@ void RdmaAdapter::Process_CQ() { CHECK_GE(ne, 0); for (int i = 0; i < ne; ++i) { CHECK(wc_[i].status == IBV_WC_SUCCESS) - << "Failed status \n" << ibv_wc_status_str(wc_[i].status) << " " - << wc_[i].status << " " << static_cast(wc_[i].wr_id) << " " - << wc_[i].vendor_err; + << "Failed status \n" + << ibv_wc_status_str(wc_[i].status) << " " << wc_[i].status << " " + << static_cast(wc_[i].wr_id) << " " << wc_[i].vendor_err; if (wc_[i].opcode == IBV_WC_RECV_RDMA_WITH_IMM) { RdmaChannel* rc = reinterpret_cast(wc_[i].wr_id); // put back a recv wr. rc->Recv(); // imm_data is the index of RX buffer in the buffer table. uint32_t imm_data = wc_[i].imm_data; - RdmaBuffer* rb = rc->FindBuffer(imm_data); + RdmaMessageBuffer* rb; RdmaMessage rm; - RdmaMessage::ParseMessage(rm, rb->buffer_); - VLOG(2) << "recv RDMA message: " << MessageTypeToString(rm.type_); - if (rm.type_ == RDMA_MESSAGE_ACK) { + if (imm_data == RDMA_IMM_DATA_ACK) { // receive an ack to a message rb = rc->tx_message_buffer_; rb->SetBufferStatus(remote, idle); rb->SendNextItem(); - } else if (rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) { - // received a request-for-tensor message - // send ack to release remote tx message buffer - RdmaBuffer* ab = rc->tx_ack_buffer_; - ab->SendNextItem(); - // find or create buffer - RdmaBuffer* tb = rc->FindOrCreateBuffer(rm.name_); - string key_with_step_id = - VerbsUtil::AppendStepidToKey(rm.name_, rm.step_id_); - tb->EnqueueItem(key_with_step_id); - // send the next tensor - worker_env_->compute_pool->Schedule([tb]() { tb->SendNextItem(); }); - } else if (rm.type_ == RDMA_MESSAGE_BUFFER_IDLE) { - // receive tensor-buffer-ready message - // send ack to release remote tx message buffer - RdmaBuffer* ab = rc->tx_ack_buffer_; - ab->SendNextItem(); - // find buffer - RdmaTensorBuffer* tb = - reinterpret_cast(rc->FindBuffer(rm.name_)); - tb->SetBufferStatus(remote, idle); - worker_env_->compute_pool->Schedule([tb]() { tb->ReSendNextItem(); }); - } else if (rm.type_ == RDMA_MESSAGE_BUFFER_REQUEST) { - // remote host requests to create a tensor buffer; - // send ack to release remote tx message buffer - RdmaBuffer* ab = rc->tx_ack_buffer_; - ab->SendNextItem(); - // find or create the buffer - RdmaBuffer* tb = rc->FindOrCreateBuffer(rm.name_, TENSOR); - RemoteMR rmr; - rmr.remote_addr = rm.remote_addr_; - rmr.rkey = rm.rkey_; - tb->SetRemoteMR(rmr, true); - tb->CreateCPUBuffer(rm.buffer_size_); - // create RDMA_MESSAGE_BUFFER_RESPONSE message - RdmaMessage br; - br.type_ = RDMA_MESSAGE_BUFFER_RESPONSE; - br.name_size_ = rm.name_.size(); - br.name_ = rm.name_; - br.buffer_size_ = rm.buffer_size_; - br.remote_addr_ = reinterpret_cast(tb->buffer_); - br.rkey_ = tb->self_->rkey; - string message = RdmaMessage::CreateMessage(br); - RdmaBuffer* mb = rc->tx_message_buffer_; - mb->EnqueueItem(message); - mb->SendNextItem(); - } else if (rm.type_ == RDMA_MESSAGE_BUFFER_RESPONSE) { - // remote creates a buffer and responds - // send ack to release remote tx message buffer - RdmaBuffer* ab = rc->tx_ack_buffer_; - ab->SendNextItem(); - // find buffer - RdmaTensorBuffer* tb = - reinterpret_cast(rc->FindBuffer(rm.name_)); - CHECK(rm.buffer_size_ == tb->size_) - << "rm.buffer_size = " << rm.buffer_size_ - << "tb->size_ = " << tb->size_ << "rm.name_ = " << rm.name_; - RemoteMR rmr; - rmr.remote_addr = rm.remote_addr_; - rmr.rkey = rm.rkey_; - tb->SetRemoteMR(rmr, true); - tb->SetBufferStatus(local, idle); - tb->SetBufferStatus(remote, idle); - worker_env_->compute_pool->Schedule([tb]() { tb->ReSendNextItem(); }); - } else if (rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) { - // tensor RDMA write completed - worker_env_->compute_pool->Schedule([rm, rc]() { - string key_with_step_id = - VerbsUtil::AppendStepidToKey(rm.name_, rm.step_id_); - rc->RunRecvCallback(key_with_step_id); - }); + continue; } - } else if (wc_[i].opcode == IBV_WC_RDMA_WRITE) { - RdmaBuffer* rb = reinterpret_cast(wc_[i].wr_id); - rb->SetBufferStatus(local, idle); - RdmaMessage rm; + + if (imm_data <= RDMA_IMM_MAX_REQUEST_ID) { + // receive a tensor RDMA write + uint32_t request_index = imm_data; + RdmaTensorRequest* request = rc->GetTensorRequest(request_index); + request->RecvTensorContent(); + continue; + } + + // receive a control message + rb = rc->rx_message_buffer_; RdmaMessage::ParseMessage(rm, rb->buffer_); - VLOG(2) << "sent RDMA message: " << MessageTypeToString(rm.type_); - if (rm.type_ != RDMA_MESSAGE_ACK) { - worker_env_->compute_pool->Schedule([rb]() { rb->SendNextItem(); }); + RdmaMessageBuffer::SendAck(rc); + RDMA_LOG(1) << "Step 0x" << std::hex << rm.step_id_ << std::dec + << ": Received " << MessageTypeToString(rm.type_) << " " + << "#" << rm.request_index_ << ": " << rm.name_; + + if (rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) { + RdmaTensorResponse* response = rc->AddTensorResponse(rm); + response->Start(); + } else if (rm.type_ == RDMA_MESSAGE_META_DATA_UPDATE) { + RdmaTensorRequest* request = rc->GetTensorRequest(rm.request_index_); + request->RecvTensorMetaData(rm.data_type_, rm.tensor_shape_, + rm.is_dead_, rm.tensor_bytes_); +#ifdef RDMA_DATA_VALIDATION + request->RecvTensorChecksum(rm.checksum_); +#endif + } else if (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST) { + RdmaTensorResponse* response = rc->UpdateTensorResponse(rm); + response->Resume(); + } else if (rm.type_ == RDMA_MESSAGE_ERROR_STATUS) { + RdmaTensorRequest* request = rc->GetTensorRequest(rm.request_index_); + request->RecvErrorStatus(rm.status_); } + } else if (wc_[i].opcode == IBV_WC_RDMA_WRITE) { + RdmaWriteID* wr_id = reinterpret_cast(wc_[i].wr_id); + RDMA_LOG(2) << "Write complete of type " << wr_id->write_type; + switch (wr_id->write_type) { + case RDMA_WRITE_ID_ACK: + break; + case RDMA_WRITE_ID_MESSAGE: { + RdmaMessageBuffer* rb = + reinterpret_cast(wr_id->write_context); + rb->SetBufferStatus(local, idle); + rb->SendNextItem(); + break; + } + case RDMA_WRITE_ID_TENSOR_WRITE: { + RdmaTensorResponse* response = + reinterpret_cast(wr_id->write_context); + response->Destroy(); + } + } + delete wr_id; } } } @@ -577,7 +540,7 @@ int RdmaChannel::PingPostRecv() { int RdmaChannel::PingPostSend() { struct ibv_send_wr wr, *bad_wr; memset(&wr, 0, sizeof(wr)); - wr.wr_id = (uint64_t) this; + wr.wr_id = (uint64_t)this; wr.sg_list = &ping_sge_list_; wr.num_sge = 1; wr.opcode = IBV_WR_SEND; @@ -588,8 +551,10 @@ int RdmaChannel::PingPostSend() { RdmaChannel::RdmaChannel(const RdmaAdapter* adapter, const string local_name, const string remote_name) - : adapter_(adapter), local_name_(local_name), remote_name_(remote_name) { - + : adapter_(adapter), + local_name_(local_name), + remote_name_(remote_name), + request_serial_(0) { struct ibv_sge list; mr_ = ibv_reg_mr(adapter_->pd_, ping_buff_, kPingBuffSize, @@ -651,29 +616,15 @@ RdmaChannel::RdmaChannel(const RdmaAdapter* adapter, const string local_name, // create message and ack buffers, then initialize the tables. { - const string buffer_names[] = {"tx_message_buffer", "rx_message_buffer", - "tx_ack_buffer", "rx_ack_buffer"}; + const string buffer_names[] = {"tx_message_buffer", "rx_message_buffer"}; tx_message_buffer_ = new RdmaMessageBuffer(this, buffer_names[0]); rx_message_buffer_ = new RdmaMessageBuffer(this, buffer_names[1]); - tx_ack_buffer_ = new RdmaAckBuffer(this, buffer_names[2]); - rx_ack_buffer_ = new RdmaAckBuffer(this, buffer_names[3]); message_buffers_.reserve(kNumMessageBuffers); message_buffers_.push_back(tx_message_buffer_); message_buffers_.push_back(rx_message_buffer_); - message_buffers_.push_back(tx_ack_buffer_); - message_buffers_.push_back(rx_ack_buffer_); // create buffer on host tx_message_buffer_->CreateCPUBuffer(RdmaMessage::kRdmaMessageBufferSize); rx_message_buffer_->CreateCPUBuffer(RdmaMessage::kRdmaMessageBufferSize); - tx_ack_buffer_->CreateCPUBuffer(RdmaMessage::kRdmaAckBufferSize); - rx_ack_buffer_->CreateCPUBuffer(RdmaMessage::kRdmaAckBufferSize); - // bt_mu_.lock() is not used in constructor. - for (int i = 0; i < kNumMessageBuffers; i++) { - uint32_t index = NameHash(buffer_names[i]); - buffer_table_.insert({index, message_buffers_[i]}); - buffer_index_name_table_.insert({index, buffer_names[i]}); - buffer_name_index_table_.insert({buffer_names[i], index}); - } } CHECK(PingPostRecv() == 0) << "Couldn't post receive from " << remote_name_ << " with error " << std::strerror(errno); @@ -684,8 +635,6 @@ RdmaChannel::~RdmaChannel() { CHECK(!ibv_destroy_qp(qp_)) << "Failed to destroy QP"; delete tx_message_buffer_; delete rx_message_buffer_; - delete tx_ack_buffer_; - delete rx_ack_buffer_; } void RdmaChannel::SetRemoteAddress(const RdmaAddress& ra, bool override) { @@ -711,119 +660,36 @@ void RdmaChannel::SetRemoteAddress(const RdmaAddress& ra, bool override) { void RdmaChannel::Recv() { struct ibv_recv_wr wr; memset(&wr, 0, sizeof(wr)); - wr.wr_id = (uint64_t) this; + wr.wr_id = (uint64_t)this; struct ibv_recv_wr* bad_wr; CHECK(!ibv_post_recv(qp_, &wr, &bad_wr)) << "Failed to post recv"; } -// Lookup 32-bit buffer index from buffer name -// Args: -// buffer_name: name of the buffer -// Returns: -// 32-bit index -uint32_t RdmaChannel::LookupBufferIndex(const string& buffer_name) { - mutex_lock lock{bt_mu_}; - BufferNameIndexTable::iterator iter = - buffer_name_index_table_.find(buffer_name); - CHECK(iter != buffer_name_index_table_.end()); - return iter->second; -} - -// Find a buffer by its 32-bit index -// Args: -// index: 32-bit hash code of the tensor buffer name -// Returns: -// name of the tensor buffer -RdmaBuffer* RdmaChannel::FindBuffer(const uint32_t index) { - mutex_lock lock{bt_mu_}; - BufferTable::iterator iter = buffer_table_.find(index); - CHECK(iter != buffer_table_.end()); - return iter->second; -} - -// Find a buffer by its name -// Args: -// name: name of the buffer -// Returns: -// the named rdma buffer -RdmaBuffer* RdmaChannel::FindBuffer(const string& name) { - uint32_t index = LookupBufferIndex(name); - return FindBuffer(index); -} - -// Find a buffer if it exists, otherwise create one. -// The memory inside the created buffer is not allocated. -// Args: -// name: the name of the buffer -// buffer_type: TENSOR, MESSAGE or ACK. -// Returns: -// the named buffer -RdmaBuffer* RdmaChannel::FindOrCreateBuffer(const string& name, - BufferType buffer_type) { - mutex_lock lock{bt_mu_}; - RdmaBuffer* rb; - // find index - BufferNameIndexTable::iterator iter = buffer_name_index_table_.find(name); - if (iter != buffer_name_index_table_.end()) { - uint32_t index = iter->second; - // find buffer - BufferTable::iterator iter = buffer_table_.find(index); - CHECK(iter != buffer_table_.end()); - rb = iter->second; - } else { - uint32_t index = NameHash(name); - if (buffer_type == TENSOR) { - rb = new RdmaTensorBuffer(this, name); - } else if (buffer_type == MESSAGE) { - rb = new RdmaMessageBuffer(this, name); - } else if (buffer_type == ACK) { - rb = new RdmaAckBuffer(this, name); - } - buffer_name_index_table_.insert({name, index}); - buffer_index_name_table_.insert({index, name}); - buffer_table_.insert({index, rb}); +RdmaTensorRequest* RdmaChannel::InsertTensorRequest( + const string& key, int64 step_id, Device* dst_dev, + const Rendezvous::Args recv_args, + const RdmaTensorRequest::RecvDoneCallback& done) { + mutex_lock lock{ct_mu_}; + uint32_t request_index = request_serial_++; + if (request_serial_ > RDMA_IMM_MAX_REQUEST_ID) { + request_serial_ = 0; } - CHECK(rb); - return rb; + RdmaTensorRequest request(request_index, key, step_id, this, dst_dev, + recv_args, done); + auto it = request_table_.emplace(request_index, request); + return &it.first->second; } -// Insert callback to the callback_table. -// The callback is activated when the corresponding tensor is received. -// Arg: -// key: the name of the tensor -// recv_done: the callback associated with the tensor. -// Returns: -// None -void RdmaChannel::InsertRecvCallback(const string& key, - std::function recv_done) { +void RdmaChannel::RemoveTensorRequest(uint32_t request_index) { mutex_lock lock{ct_mu_}; - callback_table_.insert({key, recv_done}); + request_table_.erase(request_index); } -// Remove callback from the callback_table. -// Arg: -// key: the name of the tensor -// Returns: -// None -void RdmaChannel::RemoveRecvCallback(const string& key) { +RdmaTensorRequest* RdmaChannel::GetTensorRequest(uint32_t request_index) { mutex_lock lock{ct_mu_}; - callback_table_.erase(key); -} - -// Run named callback in the callback_table. -// Arg: -// key: the name of the tensor -// Returns: -// None -void RdmaChannel::RunRecvCallback(const string& key) { - std::function recv_done; - { - mutex_lock lock{ct_mu_}; - CallbackTable::iterator iter = callback_table_.find(key); - CHECK(iter != callback_table_.end()); - recv_done = iter->second; - } - recv_done(); + RequestTable::iterator iter = request_table_.find(request_index); + CHECK(iter != request_table_.end()); + return &iter->second; } void RdmaChannel::Connect() { @@ -865,11 +731,11 @@ void RdmaChannel::Connect(const RdmaAddress& remoteAddr) { attr.ah_attr.grh.traffic_class = adapter_->params_.traffic_class; int r; - CHECK(!(r = ibv_modify_qp(qp_, &attr, IBV_QP_STATE | IBV_QP_AV | - IBV_QP_PATH_MTU | - IBV_QP_DEST_QPN | IBV_QP_RQ_PSN | - IBV_QP_MAX_DEST_RD_ATOMIC | - IBV_QP_MIN_RNR_TIMER))) + CHECK(!(r = ibv_modify_qp(qp_, &attr, + IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | + IBV_QP_DEST_QPN | IBV_QP_RQ_PSN | + IBV_QP_MAX_DEST_RD_ATOMIC | + IBV_QP_MIN_RNR_TIMER))) << "QP to Ready to Receive " << r; memset(&attr, 0, sizeof(ibv_qp_attr)); @@ -880,33 +746,30 @@ void RdmaChannel::Connect(const RdmaAddress& remoteAddr) { attr.rnr_retry = 7; /* infinite */ attr.max_rd_atomic = 1; - CHECK(!(r = ibv_modify_qp(qp_, &attr, IBV_QP_STATE | IBV_QP_TIMEOUT | - IBV_QP_RETRY_CNT | - IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN | - IBV_QP_MAX_QP_RD_ATOMIC))) + CHECK(!(r = ibv_modify_qp(qp_, &attr, + IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | + IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN | + IBV_QP_MAX_QP_RD_ATOMIC))) << "QP to Ready to Send " << r; connected_ = true; } else { - LOG(INFO) << "channel already connected"; + RDMA_LOG(2) << "channel already connected"; } } -RdmaBuffer::RdmaBuffer(RdmaChannel* channel, string name) +RdmaMessageBuffer::RdmaMessageBuffer(RdmaChannel* channel, string name) : channel_(channel), name_(name) {} -RdmaBuffer::~RdmaBuffer() { +RdmaMessageBuffer::~RdmaMessageBuffer() { CHECK(!ibv_dereg_mr(self_)) << "ibv_dereg_mr failed"; FreeBuffer(); } -void RdmaBuffer::FreeBuffer() { +void RdmaMessageBuffer::FreeBuffer() { if ((buffer_ != nullptr) && buffer_on_host_) { free(buffer_); } - // TODO - // release buffer if it is on device. - // We don't support RDMABuffer on device at this moment. } // Allocate CPU memory for the Rdma buffer @@ -915,7 +778,7 @@ void RdmaBuffer::FreeBuffer() { // lock: whether or not mutex_lock the process to protect concurrency. // Returns: // None -void RdmaBuffer::CreateCPUBuffer(size_t size, bool lock) { +void RdmaMessageBuffer::CreateCPUBuffer(size_t size, bool lock) { CHECK(size > 0); if (lock) { mu_.lock(); @@ -943,7 +806,7 @@ void RdmaBuffer::CreateCPUBuffer(size_t size, bool lock) { // override: whether override existing information // Returns: // None -void RdmaBuffer::SetRemoteMR(RemoteMR rmr, bool override) { +void RdmaMessageBuffer::SetRemoteMR(RemoteMR rmr, bool override) { mutex_lock lock{mu_}; if ((override) || (remote_status_ == none)) { remote_.remote_addr = rmr.remote_addr; @@ -956,63 +819,51 @@ void RdmaBuffer::SetRemoteMR(RemoteMR rmr, bool override) { } // Put a task in the buffer's job queue -void RdmaBuffer::EnqueueItem(string item) { +void RdmaMessageBuffer::EnqueueItem(string item) { mutex_lock lock{mu_}; queue_.push(item); } // Rdma-Write the content of the buffer -void RdmaBuffer::Write(uint32_t imm_data, size_t buffer_size) { +void RdmaMessageBuffer::Write(uint32_t imm_data, size_t buffer_size) { + Write(channel_, imm_data, buffer_size, (uint64_t)buffer_, self_->lkey, + remote_.remote_addr, remote_.rkey, RDMA_WRITE_ID_MESSAGE, this); +} + +// Generalized Write method +void RdmaMessageBuffer::Write(const RdmaChannel* channel, uint32_t imm_data, + size_t buffer_size, uint64_t src_addr, + uint32_t lkey, uint64_t remote_addr, + uint32_t rkey, RdmaWriteIDType write_type, + void* write_context) { struct ibv_sge list; - list.addr = (uint64_t)buffer_; + list.addr = src_addr; list.length = buffer_size; - list.lkey = self_->lkey; + list.lkey = lkey; struct ibv_send_wr wr; memset(&wr, 0, sizeof(wr)); - wr.wr_id = (uint64_t) this; + wr.wr_id = (uint64_t) new RdmaWriteID(write_type, write_context); wr.sg_list = &list; wr.num_sge = 1; wr.opcode = IBV_WR_RDMA_WRITE_WITH_IMM; wr.send_flags = IBV_SEND_SIGNALED; wr.imm_data = imm_data; - wr.wr.rdma.remote_addr = (uint64_t)remote_.remote_addr; - wr.wr.rdma.rkey = remote_.rkey; + wr.wr.rdma.remote_addr = remote_addr; + wr.wr.rdma.rkey = rkey; struct ibv_send_wr* bad_wr; - CHECK(!ibv_post_send(channel_->qp_, &wr, &bad_wr)) << "Failed to post send"; -} - -RdmaAckBuffer::RdmaAckBuffer(RdmaChannel* channel, string name) - : RdmaBuffer(channel, name) {} - -RdmaMessageBuffer::RdmaMessageBuffer(RdmaChannel* channel, string name) - : RdmaBuffer(channel, name) {} - -RdmaTensorBuffer::RdmaTensorBuffer(RdmaChannel* channel, string name) - : RdmaBuffer(channel, name) {} - -RdmaTensorBuffer::~RdmaTensorBuffer() { - for (Itable it = retable.begin(); it != retable.end(); ++it) { - delete (it->second); - } + CHECK(!ibv_post_send(channel->qp_, &wr, &bad_wr)) << "Failed to post send"; } // Send the next ack from the buffer's job queue. -void RdmaAckBuffer::SendNextItem() { - uint32_t imm_data = LookupBufferIndex("rx_ack_buffer"); - RdmaMessage rm; - rm.name_ = "rx_ack_buffer"; - rm.type_ = RDMA_MESSAGE_ACK; - rm.name_size_ = rm.name_.size(); - string message = RdmaMessage::CreateMessage(rm); - memcpy(buffer_, message.data(), message.size()); - Write(imm_data, message.size()); +void RdmaMessageBuffer::SendAck(const RdmaChannel* channel) { + Write(channel, RDMA_IMM_DATA_ACK, 0, 0, 0, 0, 0, RDMA_WRITE_ID_ACK, nullptr); } // Send the next message from the buffer's job queue. void RdmaMessageBuffer::SendNextItem() { - uint32_t imm_data = LookupBufferIndex("rx_message_buffer"); + uint32_t imm_data = RDMA_IMM_DATA_MESSAGE; mu_.lock(); if (!queue_.empty() && (local_status_ == idle) && (remote_status_ == idle)) { local_status_ = busy; @@ -1029,244 +880,390 @@ void RdmaMessageBuffer::SendNextItem() { } } -Rendezvous::DoneCallback RdmaTensorBuffer::getRecvTensorCallback( - const string& key_with_step_id, const string& key, int64 step_id, - const Rendezvous::ParsedKey& parsed) { - Rendezvous::DoneCallback cb = [this, key_with_step_id, key, step_id, parsed]( - const Status& status, const Rendezvous::Args& send_args, - const Rendezvous::Args& recv_args, const Tensor& in, bool is_dead) { - CHECK(status.ok()) << "RecvLocalAsync was not ok, key" << key_with_step_id - << " error message: " << status.error_message(); - size_t buffer_size = RdmaMessage::kMessageTotalBytes; - size_t tensor_bytes = 0; - // Figures out which device the tensor is hosted on. - Device* src_dev = nullptr; - Status s = channel_->adapter_->worker_env_->device_mgr->LookupDevice( - parsed.src_device, &src_dev); - CHECK(s.ok()) << "src device not found"; - // Does the device have the right incarnation number we expect? - CHECK(src_dev->attributes().incarnation() == parsed.src_incarnation) - << "RecvTensor expects a different device incarnation: " - << parsed.src_incarnation << " vs. " - << src_dev->attributes().incarnation() - << ". Your worker job was probably restarted. Check your " - << "worker job for the reason why it was restarted."; - Device* dst_dev = nullptr; - // destination is on CPU. - s = channel_->adapter_->worker_env_->device_mgr->LookupDevice("CPU:0", - &dst_dev); - CHECK(s.ok()) << "dst device not found"; - AllocatorAttributes dst_alloc_attr; - dst_alloc_attr.set_on_host(true); - - bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); - // string tensor needs to be serialized - Tensor copy; - TensorProto proto; - if (src_dev->tensorflow_gpu_device_info() && - (!send_args.alloc_attrs.on_host())) { #if GOOGLE_CUDA - CHECK(send_args.device_context) << "send dev name: " << src_dev->name() - << " gpu_info: " - << src_dev->tensorflow_gpu_device_info(); - - if (can_memcpy) { - AllocatorAttributes host_alloc_attrs; - host_alloc_attrs.set_gpu_compatible(true); - host_alloc_attrs.set_on_host(true); - Allocator* alloc = ProcessState::singleton()->GetCUDAHostAllocator(0); - copy = Tensor(alloc, in.dtype(), in.shape()); - tensor_bytes = in.TotalBytes(); - buffer_size += tensor_bytes; - GPUUtil::CopyGPUTensorToCPU( - src_dev, send_args.device_context, &in, ©, - [this, copy, tensor_bytes, buffer_size, key, in, step_id, - key_with_step_id, is_dead, send_args, recv_args](const Status& s) { - CHECK(s.ok()) << "copy tensor from gpu sync"; - StringPiece copy_buf; - copy_buf = copy.tensor_data(); - PostCopyOperations(true, buffer_size, tensor_bytes, key, in, - step_id, is_dead, key_with_step_id, ©, - NULL, ©_buf, send_args, recv_args); - }); - } else { - // "val" is on a GPU. No longer uses GPUUtil to fill the proto, use - // aync instead - GPUUtil::SetProtoFromGPU( - in, src_dev, send_args.device_context, &proto, is_dead, - [this, proto, buffer_size, key, in, step_id, key_with_step_id, - is_dead, send_args, recv_args](const Status& s) mutable { - CHECK(s.ok()) << "copy proto from gpu sync"; - auto tensor_bytes = proto.ByteSize(); - buffer_size += tensor_bytes; - PostCopyOperations(false, buffer_size, tensor_bytes, key, in, - step_id, is_dead, key_with_step_id, NULL, - &proto, NULL, send_args, recv_args); - }); - } +static void CountCopies(const std::string& key, void* src_addr, void* dst_addr, + size_t tensor_bytes, bool is_gpu_to_cpu) { +#ifdef RDMA_COUNT_COPIES + static uint64_t numGPUToCPUCopies = 0; + static uint64_t numGPUToCPUCopiedBytes = 0; + static uint64_t numCPUToGPUCopies = 0; + static uint64_t numCPUToGPUCopiedBytes = 0; + static uint64_t numTotalCopies = 0; + + if (is_gpu_to_cpu) { + ++numGPUToCPUCopies; + numGPUToCPUCopiedBytes += tensor_bytes; + } else { + ++numCPUToGPUCopies; + numCPUToGPUCopiedBytes += tensor_bytes; + } + if ((++numTotalCopies % 0x400) == 0) { + RDMA_LOG(0) << "Tensor copies:" + << " GPU to CPU: " << numGPUToCPUCopies << " (" + << numGPUToCPUCopiedBytes << " Bytes)" + << " CPU to GPU: " << numCPUToGPUCopies << " (" + << numCPUToGPUCopiedBytes << " Bytes)"; + } + RDMA_LOG(2) << "Copying tensor " << key << " From: " << src_addr + << " To: " << dst_addr; +#endif // RDMA_COUNT_COPIES +} #endif // GOOGLE_CUDA - } else { - // tensor is in CPU memory. - StringPiece copy_buf; - if (can_memcpy) { - copy_buf = in.tensor_data(); - tensor_bytes = in.TotalBytes(); - } else { - in.AsProtoTensorContent(&proto); - tensor_bytes = proto.ByteSize(); - } - buffer_size += tensor_bytes; - PostCopyOperations(can_memcpy, buffer_size, tensor_bytes, key, in, - step_id, is_dead, key_with_step_id, ©, &proto, - ©_buf, send_args, recv_args); + +#ifdef RDMA_DATA_VALIDATION +static uint64_t Checksum(Device* device, const DeviceContext* device_context, + const Tensor& in) { + uint64 checksum = 0; + if (DataTypeCanUseMemcpy(in.dtype())) { +#if GOOGLE_CUDA + if (in.TotalBytes() == 0) { + return 0; } - }; - return cb; + checksum = (device_context != nullptr) + ? GPUUtil::Checksum(device, device_context, in) + : GPUUtil::Checksum(in); +#endif // GOOGLE_CUDA + } else { + string s = in.SummarizeValue(999999); + checksum = Hash64(s.c_str(), s.size(), 0); + } + return checksum; } -// Send the next tensor from the buffer's job queue. -void RdmaTensorBuffer::SendNextItem() { - // get the key - string key_with_step_id = ""; - { - mutex_lock lock{mu_}; - if (!queue_.empty()) { - key_with_step_id = queue_.front(); - queue_.pop(); +static void ValidateChecksum(uint64_t expected, uint64_t actual, + const Tensor& in, uint32_t request_index, + const std::string& key, const std::string& msg) { + RDMA_LOG(2) << "Request #" << request_index << ": " << key + << ": Checksum: " << std::hex << " Expected = 0x" << expected + << ". Actual = 0x" << actual << "."; + + if (expected != actual) { + // Checksum failed. There is one case where this is allowed - if the + // tensor is an AssignAdd of the global step. Since the data-validation + // always postpones the Tensor response in order to send a checksum message, + // it is possible that the global-step was updated while the response was + // still in queue. + if ((in.TotalBytes() == 8) && (in.dtype() == DT_INT64)) { + int64_t prev_val = *(int64_t*)DMAHelper::base(&in) - 1; + actual = Hash64((const char*)&prev_val, 8, 0); + } + if (expected != actual) { + LOG(FATAL) << "[" << msg << "]: Checksum validation failed for request #" + << request_index << ": " << key << std::hex << " " + << DataTypeString(in.dtype()) << " " + << in.shape().DebugString() << " (0x" << in.TotalBytes() + << " bytes): " + << " Expected 0x" << expected << ". Got 0x" << actual << "."; } } +} +#endif // RDMA_DATA_VALIDATION + +#if GOOGLE_CUDA +// Sync the 'done' operation on the GPU stream, but without all the data +// copying. +static void StreamGPUOp(Device* gpu_device, const DeviceContext* device_context, + StatusCallback done) { + Tensor dummy1, dummy2; + GPUUtil::CopyGPUTensorToCPU(gpu_device, device_context, &dummy1, &dummy2, + done); +} +#endif // GOOGLE_CUDA + +RdmaTensorResponse* RdmaChannel::AddTensorResponse(const RdmaMessage& rm) { + mutex_lock lock{mu_}; + auto it = + responses_table_.emplace(rm.request_index_, RdmaTensorResponse(this, rm)); + CHECK(it.second) << "Response with the ID " << rm.request_index_ + << " already exists."; + return &it.first->second; +} + +RdmaTensorResponse* RdmaChannel::UpdateTensorResponse(const RdmaMessage& rm) { + mutex_lock lock{mu_}; + auto it = responses_table_.find(rm.request_index_); + CHECK(it != responses_table_.end()) << "No response found."; + RdmaTensorResponse* response = &it->second; + response->Update(rm); + return response; +} + +void RdmaChannel::RemoveTensorResponse(uint32_t request_index) { + mutex_lock lock{mu_}; + responses_table_.erase(request_index); +} + +void RdmaTensorResponse::Start() { + Rendezvous::ParsedKey parsed; + Status s = Rendezvous::ParseKey(rm_.name_, &parsed); + if (!s.ok()) { + SendErrorStatus(s); + return; + } - // send the tensor if a key is acquired. - if (key_with_step_id != "") { - VLOG(2) << "try to send tensor: " << key_with_step_id; - string key; - int64 step_id; - VerbsUtil::GetKeyAndStepId(key_with_step_id, key, step_id); - CHECK(key.compare(name_) == 0); - Rendezvous::ParsedKey parsed; - Rendezvous::ParseKey(key, &parsed); - Rendezvous::DoneCallback cb = - getRecvTensorCallback(key_with_step_id, key, step_id, parsed); - channel_->adapter_->worker_env_->rendezvous_mgr->RecvLocalAsync(step_id, - parsed, cb); + channel_->adapter_->worker_env_->rendezvous_mgr->RecvLocalAsync( + rm_.step_id_, parsed, + [this, parsed](const Status& status, const Rendezvous::Args& send_args, + const Rendezvous::Args& recv_args, const Tensor& in, + bool is_dead) { + CHECK(status.ok()) << "RecvLocalAsync was not ok." + << " error message: " << status.error_message(); + RecvHandler(parsed, send_args, recv_args, in, is_dead); + }); +} + +void RdmaTensorResponse::Resume() { SendContent(*tensor_, *proto_, is_dead_); } + +// Helper for RecvTensor. Validates "key" and returns the source +// device in "*src_dev". +Status RdmaTensorResponse::PrepareRecvTensor( + const Rendezvous::ParsedKey& parsed, Device** src_dev) { + // Figures out which device the tensor is hosted on. + string local_name = DeviceNameUtils::LocalName(parsed.src_device); + TF_RETURN_IF_ERROR(channel_->adapter_->worker_env_->device_mgr->LookupDevice( + local_name, src_dev)); + + // Does the device have the right incarnation number we expect? + if ((*src_dev)->attributes().incarnation() != parsed.src_incarnation) { + return errors::Aborted( + "RecvTensor expects a different device incarnation: ", + parsed.src_incarnation, " vs. ", (*src_dev)->attributes().incarnation(), + ". Your worker job was probably restarted. Check your " + "worker job for the reason why it was restarted."); } + + return Status::OK(); } -void RdmaTensorBuffer::ReSendNextItem() { - // get the key - string key_with_step_id = ""; - { - mutex_lock lock{mu_}; - if (!requeue.empty()) { - key_with_step_id = requeue.front(); - requeue.pop(); - } +void RdmaTensorResponse::RecvHandler(Rendezvous::ParsedKey parsed, + const Rendezvous::Args& send_args, + const Rendezvous::Args& recv_args, + const Tensor& in, bool is_dead) { + Status s = PrepareRecvTensor(parsed, &src_dev_); + if (!s.ok()) { + SendErrorStatus(s); + return; } - // send the tensor if a key is acquired. - if (key_with_step_id != "") { - VLOG(2) << "try to send tensor: " << key_with_step_id; - string key; - int64 step_id; - VerbsUtil::GetKeyAndStepId(key_with_step_id, key, step_id); - CHECK(key.compare(name_) == 0); - Rendezvous::ParsedKey parsed; - Rendezvous::ParseKey(key, &parsed); - Rendezvous::DoneCallback cb = - getRecvTensorCallback(key_with_step_id, key, step_id, parsed); - ReItem* item; - { - mutex_lock lock{mu_}; - Itable it = retable.find(key_with_step_id); - CHECK(it != retable.end()) << "Could not find dup-recv context"; - item = it->second; - retable.erase(it); + meta_data_changed_ = TensorMetaDataChanged(in, is_dead); +#ifdef RDMA_DATA_VALIDATION + // Always send a meta data message with the source checksum + meta_data_changed_ = rm_.type_ == RDMA_MESSAGE_TENSOR_REQUEST; + checksum_ = Checksum(src_dev_, send_args.device_context, in); +#endif + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + // string tensor needs to be serialized + Tensor copy; + TensorProto proto; + const bool on_host = send_args.alloc_attrs.on_host(); + if (src_dev_->tensorflow_gpu_device_info() && !on_host) { +#if GOOGLE_CUDA + DeviceContext* send_dev_context = send_args.device_context; + CHECK(send_dev_context) + << "send dev name: " << src_dev_->name() + << " gpu_info: " << src_dev_->tensorflow_gpu_device_info(); + + if (can_memcpy) { + // If the tensor is located on a GDR compatible GPU, there is no need to + // copy it. We can send directly from the source, just need to make sure + // we are in sync with the GPU stream. + // If the tensor's meta-data changed however, we will need to clone it, + // so anyway we'll have to copy it from GPU to CPU first. If at some + // point in time Clone() is changed to only save a shallow copy, we can + // skip the copy here as well. + if ((in.TotalBytes() > 0) && !meta_data_changed_ && + (RdmaMemoryMgr::Singleton().FindMemoryRegion( + (void*)DMAHelper::base(&in), in.TotalBytes()) != nullptr)) { + StreamGPUOp(src_dev_, send_dev_context, + [this, in, proto, is_dead](const Status& s) { + Send(in, proto, is_dead, s); + }); + return; + } + + // The tensor must be copied from GPU to CPU, because either: + // 1. The tensor is located on a non GDR compatible GPU. + // 2. The tensor's meta-data has changed. + Allocator* alloc = ProcessState::singleton()->GetCUDAHostAllocator(0); + copy = Tensor(alloc, in.dtype(), in.shape()); + CountCopies(rm_.name_, (void*)DMAHelper::base(&in), + (void*)DMAHelper::base(©), in.TotalBytes(), true); + GPUUtil::CopyGPUTensorToCPU( + src_dev_, send_dev_context, &in, ©, + [this, copy, proto, is_dead](const Status& s) { + Send(copy, proto, is_dead, s); + }); + } else { + GPUUtil::SetProtoFromGPU( + in, src_dev_, send_args.device_context, &proto, is_dead, + [this, in, proto, is_dead](const Status& s) mutable { + Send(in, proto, is_dead, s); + }); + } +#else + SendErrorStatus(errors::Internal("No GPU device in process")); +#endif // GOOGLE_CUDA + } else { + // tensor is in CPU memory. + if (!can_memcpy) { + in.AsProtoTensorContent(&proto); } - cb(Status::OK(), item->send_args, item->recv_args, item->in, item->is_dead); - delete (item); + Send(in, proto, is_dead, Status::OK()); + } +} + +void RdmaTensorResponse::Send(const Tensor& in, const TensorProto& proto, + bool is_dead, const Status& status) { + if (!status.ok()) { + SendErrorStatus(status); + return; + } + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + bool proto_size_changed = + (!can_memcpy) && (proto.ByteSize() != rm_.tensor_bytes_); + if (meta_data_changed_ || proto_size_changed) { + Clone(in, proto, is_dead); + SendMetaData(in, proto, is_dead); + } else { + SendContent(in, proto, is_dead); } } -void RdmaTensorBuffer::PostCopyOperations( - bool can_memcpy, size_t buffer_size, size_t tensor_bytes, const string& key, - const Tensor& in, int64 step_id, bool is_dead, - const string& key_with_step_id, const Tensor* copy, - const TensorProto* proto, const StringPiece* copy_buf, - const Rendezvous::Args& send_args, const Rendezvous::Args& recv_args) { - // prepare message +bool RdmaTensorResponse::TensorMetaDataChanged(const Tensor& in, bool is_dead) { + return (rm_.data_type_ != in.dtype()) || (rm_.tensor_shape_ != in.shape()) || + (rm_.is_dead_ != is_dead); +} + +void RdmaTensorResponse::Clone(const Tensor& in, const TensorProto& proto, + bool is_dead) { + // Clone the data to be sent later. For simplicity, we clone the tensor's + // data even if it is already a copy. Performance is less of a concern here + // since the meta-data hardly ever changes. The reason we create a copy, is + // that some tensors share their buffer between different step-ids, so the + // tensor content may change before re-request was completed. + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + if (can_memcpy && (in.TotalBytes() > 0)) { + AllocatorAttributes host_alloc_attrs; + host_alloc_attrs.set_nic_compatible(true); + host_alloc_attrs.set_on_host(true); + Allocator* allocator = src_dev_->GetAllocator(host_alloc_attrs); + tensor_ = new Tensor(allocator, in.dtype(), in.shape()); + memcpy(DMAHelper::base(tensor_), DMAHelper::base(&in), in.TotalBytes()); + } else { + tensor_ = new Tensor(in.dtype(), in.shape()); + } + if (!can_memcpy) { + proto_ = new TensorProto(proto); + } + is_dead_ = is_dead; +} + +void RdmaTensorResponse::SendMetaData(const Tensor& in, + const TensorProto& proto, bool is_dead) { + RDMA_LOG(2) << "Request #" << rm_.request_index_ + << ": Meta data changed: " << rm_.name_; + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + size_t tensor_bytes = (can_memcpy) ? in.TotalBytes() : proto.ByteSize(); + + // Send meta-data update: RdmaMessage rm; - rm.name_size_ = key.size(); - rm.name_ = key; + rm.type_ = RDMA_MESSAGE_META_DATA_UPDATE; + rm.name_size_ = rm_.name_.size(); + rm.name_ = rm_.name_; rm.tensor_shape_ = in.shape(); rm.data_type_ = in.dtype(); - rm.step_id_ = step_id; + rm.step_id_ = rm_.step_id_; rm.is_dead_ = is_dead; rm.tensor_bytes_ = tensor_bytes; - rm.buffer_size_ = buffer_size; - mu_.lock(); - if (local_status_ == none || (buffer_size > size_ && local_status_ == idle && - remote_status_ == idle)) { - if ((local_status_ != none) && (buffer_size > size_)) { - VLOG(2) << "Extend RDMA buffer from " << size_ << " to " << buffer_size; - } - CreateCPUBuffer(buffer_size, false); - // Need to be received again, put into the re-recv queue and the table - requeue.push(key_with_step_id); - ReItem* item = new ReItem(send_args, recv_args, in, is_dead); - retable.insert(std::pair(key_with_step_id, item)); - mu_.unlock(); - // no longer used: put back the key since it is not sent; - // ask the remote to create the same buffer - rm.type_ = RDMA_MESSAGE_BUFFER_REQUEST; - rm.remote_addr_ = reinterpret_cast(buffer_); - rm.rkey_ = self_->rkey; - string message = RdmaMessage::CreateMessage(rm); - channel_->tx_message_buffer_->EnqueueItem(message); - channel_->tx_message_buffer_->SendNextItem(); - } else if ((local_status_ == idle) && (remote_status_ == idle)) { - // both buffers are ready, send the tensor - local_status_ = busy; - remote_status_ = busy; - // local/remote_status_ won't be set back to idle - // unitl Write() is successful - mu_.unlock(); - if (!((buffer_size == size_ && rm.data_type_ != DT_STRING) || - (buffer_size <= size_ && rm.data_type_ == DT_STRING))) { - VLOG(2) << "Tensor and buffer size do not agree," - << " buffer_size = " << size_ - << " requested tensor size = " << buffer_size << in.DebugString(); - } - uint32_t imm_data = LookupBufferIndex(key); - rm.type_ = RDMA_MESSAGE_TENSOR_WRITE; - string message = RdmaMessage::CreateMessage(rm); - memcpy(buffer_, message.data(), message.size()); - if (!is_dead) { - // copy the tensor buffer content - void* output = static_cast(static_cast(buffer_) + - RdmaMessage::kTensorBufferStartIndex); - CHECK(tensor_bytes + RdmaMessage::kTensorBufferStartIndex <= size_); - if (can_memcpy) { - CHECK(copy != NULL) << "callback missing pointer to copy tensor"; - CHECK(copy_buf != NULL) << "callback missing pointer to copy buffer"; - CHECK(copy_buf->size() == tensor_bytes) - << "unexpected tensor size: " << copy_buf->size() - << " != " << tensor_bytes; - memcpy(output, copy_buf->data(), tensor_bytes); - } else { - CHECK(proto != NULL) << "callback missing pointer to proto tensor"; - proto->SerializeToArray(output, tensor_bytes); + rm.request_index_ = rm_.request_index_; +#ifdef RDMA_DATA_VALIDATION + rm.checksum_ = checksum_; +#endif + RDMA_LOG(1) << "Step 0x" << std::hex << rm.step_id_ << std::dec + << ": Sending RDMA_MESSAGE_META_DATA_UPDATE #" + << rm.request_index_ << ": " << rm.name_ + << " (shape = " << rm.tensor_shape_.DebugString() << "." + << " data-type = " << DataTypeString(rm.data_type_) << "." + << " is-dead = " << rm.is_dead_ << ")"; + + string message = RdmaMessage::CreateMessage(rm); + channel_->tx_message_buffer_->EnqueueItem(message); + channel_->tx_message_buffer_->SendNextItem(); +} + +void RdmaTensorResponse::SendContent(const Tensor& in, const TensorProto& proto, + bool is_dead) { + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + size_t tensor_bytes = (can_memcpy) ? in.TotalBytes() : proto.ByteSize(); + uint32_t imm_data = rm_.request_index_; + if (!is_dead) { + if (can_memcpy) { + src_buffer_ = const_cast(DMAHelper::buffer(&in)); + if (src_buffer_ != nullptr) { + src_buffer_->Ref(); // Keep buffer alive until write is complete + src_addr_ = src_buffer_->data(); + mr_ = RdmaMemoryMgr::Singleton().FindMemoryRegion(src_addr_, + tensor_bytes); } } else { - buffer_size = RdmaMessage::kMessageTotalBytes; + RDMA_LOG(2) << "Encoding proto: " << rm_.name_ + << " (Size: " << tensor_bytes << ") " << in.DebugString(); + src_addr_ = malloc(tensor_bytes); + mr_ = ibv_reg_mr(channel_->adapter_->pd_, src_addr_, tensor_bytes, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + proto.SerializeToArray(src_addr_, tensor_bytes); } - Write(imm_data, buffer_size); } else { - // Need to be received again, put into the re-recv queue and the table - requeue.push(key_with_step_id); - ReItem* item = new ReItem(send_args, recv_args, in, is_dead); - retable.insert(std::pair(key_with_step_id, item)); - mu_.unlock(); + tensor_bytes = 0; + } + + uint32_t lkey = (mr_ == nullptr) ? 0 : mr_->lkey; + RDMA_LOG(1) << "Step 0x" << std::hex << rm_.step_id_ << std::dec + << ": Sending tensor content #" << rm_.request_index_ << " from " + << std::hex << src_addr_ << " (0x" << lkey << ")" + << " to " << rm_.remote_addr_ << " (0x" << rm_.rkey_ + << "): " << rm_.name_ << " (size: 0x" << std::hex << tensor_bytes + << ")"; + + RdmaMessageBuffer::Write(channel_, imm_data, tensor_bytes, + (uint64_t)src_addr_, lkey, rm_.remote_addr_, + rm_.rkey_, RDMA_WRITE_ID_TENSOR_WRITE, this); +} + +void RdmaTensorResponse::SendErrorStatus(const Status& status) { + RdmaMessage rm; + rm.type_ = RDMA_MESSAGE_ERROR_STATUS; + rm.name_size_ = rm_.name_.size(); + rm.name_ = rm_.name_; + rm.step_id_ = rm_.step_id_; + rm.request_index_ = rm_.request_index_; + rm.status_ = status; + LOG(ERROR) << "Step 0x" << std::hex << rm.step_id_ << std::dec + << ": Sending RDMA_MESSAGE_ERROR_STATUS #" << rm.request_index_ + << ": " << rm.name_ << ". Status: " << status.ToString(); + + string message = RdmaMessage::CreateMessage(rm); + channel_->tx_message_buffer_->EnqueueItem(message); + channel_->tx_message_buffer_->SendNextItem(); + + // Destroy the response. + Destroy(); +} + +void RdmaTensorResponse::Destroy() { + if (src_buffer_ != nullptr) { + src_buffer_->Unref(); + } + if (tensor_ != nullptr) { + delete tensor_; + } + if (proto_ != nullptr) { + ibv_dereg_mr(mr_); + free(src_addr_); + delete proto_; } + // Remove response from the pending list: + channel_->RemoveTensorResponse(rm_.request_index_); } // Create a RdmaMessage according to the pre-defined format @@ -1276,43 +1273,46 @@ void RdmaTensorBuffer::PostCopyOperations( // message in string format string RdmaMessage::CreateMessage(const RdmaMessage& rm) { // Rdma Message format - // type|name_size|name|step_id|buffer_size|remote_addr|rkey|is_dead|... - // 1B| 2B | 512| 8B | 8B | 8B | 4B | 1B |... - // ...|data_type|tensor_shape|tensor_bytes|tensor_buffer - // ...| XB | XB | 8B |... + // type|name_size|name|step_id|request_index|remote_addr|rkey|is_dead|... + // 1B| 2B | 512| 8B | 8B | 8B | 4B | 1B |... + // ...|data_type|tensor_shape|tensor_bytes|error_status | + // ...| XB | XB | 8B |size - 4B, proto - XB | // - // ACK: type|13|"rx_ack_buffer" - // TENSOR_REQUEST: type|name_size|tensor_name|step_id - // TENSOR_WRITE: type|name_size|tensor_name|step_id|...|is_dead - // |data_type|tensor_shape|tensor_bytes - // BUFFER_IDLE: type|name_size|buffer_name - // BUFFER_REQUEST: - // type|name_size|buffer_name|...|buffer_size|remote_addr|rkey| - // BUFFER_RESPONSE: - // type|name_size|buffer_name|...|buffer_size|remote_addr|rkey| - char message[kMessageTotalBytes]; + // ACK: Imm-type: ACK + // TENSOR_REQUEST: Imm-type: MESSAGE + // Fields: type, request_index, name, step_id, remote_addr, + // rkey, is_dead, data_type, tensor_shape, tensor_bytes + // META_DATA_UPDATE: Imm-type: MESSAGE + // Fields: type, request_index, is_dead, data_type, + // tensor_shape, tensor_bytes + // TENSOR_RE_REQUST: Imm-type: MESSAGE + // Fields: type, request_index, name, step_id, remote_addr, + // rkey, is_dead, data_type, tensor_shape, tensor_bytes + // ERROR_STATUS: Imm-type: MESSAGE + // Fields: type, request_index, name, step_id, error_status + // Tensor content: Imm-type: request_index + size_t message_size = kMessageTotalBytes; + char message[kMessageTotalBytes + kErrorStatusMaxSize]; // type message[kTypeStartIndex] = static_cast(rm.type_) & 0xff; - // size of name - memcpy(&message[kNameSizeStartIndex], &rm.name_size_, sizeof(rm.name_size_)); - // name - memcpy(&message[kNameStartIndex], rm.name_.data(), rm.name_.size()); - // buffer_size, remote_addr, rkey - if ((rm.type_ == RDMA_MESSAGE_BUFFER_REQUEST) || - (rm.type_ == RDMA_MESSAGE_BUFFER_RESPONSE)) { - memcpy(&message[kBufferSizeStartIndex], &rm.buffer_size_, - sizeof(rm.buffer_size_)); + // request index + memcpy(&message[kRequestIndexStartIndex], &rm.request_index_, + sizeof(rm.request_index_)); + // name, step_id, remote_addr, rkey + if ((rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) || + (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST)) { + memcpy(&message[kNameSizeStartIndex], &rm.name_size_, + sizeof(rm.name_size_)); + memcpy(&message[kNameStartIndex], rm.name_.data(), rm.name_.size()); memcpy(&message[kRemoteAddrStartIndex], &rm.remote_addr_, sizeof(rm.remote_addr_)); memcpy(&message[kRkeyStartIndex], &rm.rkey_, sizeof(rm.rkey_)); - } - // step_id - if ((rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) || - (rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST)) { memcpy(&message[kStepIdStartIndex], &rm.step_id_, sizeof(rm.step_id_)); } // is_dead, data_type, tensor_shape, tensor_bytes - if (rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) { + if ((rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) || + (rm.type_ == RDMA_MESSAGE_META_DATA_UPDATE) || + (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST)) { memcpy(&message[kIsDeadStartIndex], &rm.is_dead_, sizeof(rm.is_dead_)); memcpy(&message[kDataTypeStartIndex], &rm.data_type_, @@ -1322,7 +1322,30 @@ string RdmaMessage::CreateMessage(const RdmaMessage& rm) { memcpy(&message[kTensorBytesStartIndex], &rm.tensor_bytes_, sizeof(rm.tensor_bytes_)); } - return string(message, kMessageTotalBytes); + // checksum +#ifdef RDMA_DATA_VALIDATION + memcpy(&message[kChecksumStartIndex], &rm.checksum_, sizeof(rm.checksum_)); +#endif + // error status + if (rm.type_ == RDMA_MESSAGE_ERROR_STATUS) { + ::grpc::Status gs = ToGrpcStatus(rm.status_); + ErrorStatusProto gsProto; + gsProto.set_error_code(gs.error_code()); + gsProto.set_error_message(gs.error_message()); + gsProto.set_error_details(gs.error_details()); + uint32_t gsProtoSize = gsProto.ByteSize(); + if (gsProtoSize + 4 > kErrorStatusMaxSize) { + LOG(ERROR) << "Error status (" << gsProtoSize + 4 << " bytes) " + << "is too big to fit in RDMA message (" << kErrorStatusMaxSize + << " bytes). Truncated."; + gsProtoSize = kErrorStatusMaxSize - 4; + } + uint32_t* proto_size = (uint32_t*)&message[kErrorStatusStartIndex]; + *proto_size = gsProtoSize; + gsProto.SerializeToArray(&message[kErrorStatusStartIndex + 4], gsProtoSize); + message_size += gsProtoSize + 4; + } + return string(message, message_size); } // Parse a RdmaMessage according to the pre-defined format @@ -1335,26 +1358,24 @@ void RdmaMessage::ParseMessage(RdmaMessage& rm, void* buffer) { char* message = static_cast(buffer); // type rm.type_ = static_cast(message[kTypeStartIndex]); - // name_size_ - memcpy(&rm.name_size_, &message[kNameSizeStartIndex], sizeof(rm.name_size_)); - // name - rm.name_ = string(&message[kNameStartIndex], rm.name_size_); - // buffer_size, remote_addr, rkey - if ((rm.type_ == RDMA_MESSAGE_BUFFER_REQUEST) || - (rm.type_ == RDMA_MESSAGE_BUFFER_RESPONSE)) { - memcpy(&rm.buffer_size_, &message[kBufferSizeStartIndex], - sizeof(rm.buffer_size_)); + // request index + memcpy(&rm.request_index_, &message[kRequestIndexStartIndex], + sizeof(rm.request_index_)); + // name, step_id, remote_addr, rkey + if ((rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) || + (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST)) { + memcpy(&rm.name_size_, &message[kNameSizeStartIndex], + sizeof(rm.name_size_)); + rm.name_ = string(&message[kNameStartIndex], rm.name_size_); memcpy(&rm.remote_addr_, &message[kRemoteAddrStartIndex], sizeof(rm.remote_addr_)); memcpy(&rm.rkey_, &message[kRkeyStartIndex], sizeof(rm.rkey_)); - } - // step_id - if ((rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) || - (rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST)) { memcpy(&rm.step_id_, &message[kStepIdStartIndex], sizeof(rm.step_id_)); } // data_type, tensor_bytes, tensor_shape, is_dead - if (rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) { + if ((rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) || + (rm.type_ == RDMA_MESSAGE_META_DATA_UPDATE) || + (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST)) { memcpy(&rm.is_dead_, &message[kIsDeadStartIndex], sizeof(rm.is_dead_)); memcpy(&rm.data_type_, &message[kDataTypeStartIndex], sizeof(rm.data_type_)); @@ -1363,6 +1384,291 @@ void RdmaMessage::ParseMessage(RdmaMessage& rm, void* buffer) { memcpy(&rm.tensor_bytes_, &message[kTensorBytesStartIndex], sizeof(rm.tensor_bytes_)); } + // checksum +#ifdef RDMA_DATA_VALIDATION + memcpy(&rm.checksum_, &message[kChecksumStartIndex], sizeof(rm.checksum_)); +#endif + // error status + if (rm.type_ == RDMA_MESSAGE_ERROR_STATUS) { + ErrorStatusProto gsProto; + uint32_t gsProtoSize = *(uint32_t*)&message[kErrorStatusStartIndex]; + CHECK(ParseProtoUnlimited(&gsProto, &message[kErrorStatusStartIndex + 4], + gsProtoSize)) + << "Failed to parse error status proto from message. Aborting."; + ::grpc::Status gs((::grpc::StatusCode)gsProto.error_code(), + gsProto.error_message(), gsProto.error_details()); + rm.status_ = FromGrpcStatus(gs); + } +} + +//***************************************************************************** +// RdmaMemoryMgr +//***************************************************************************** + +ibv_mr* RdmaMemoryMgr::FindMemoryRegion(void* addr, size_t length) { + mutex_lock l(mrs_mu_); + auto iter = std::upper_bound(mrs_.begin(), mrs_.end(), addr, &Comparator); + if (iter == std::end(mrs_) || iter->get()->addr > addr) { + return nullptr; + } else { + return iter->get(); + } +} + +void RdmaMemoryMgr::InsertMemoryRegion(void* addr, size_t length, + const std::string& allocator_name) { + if (length == 0) return; + ibv_mr* mr = ibv_reg_mr(pd_, addr, length, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + RDMA_LOG(1) << "Insert memory region 0x" << std::hex << mr->rkey << ". [" + << addr << "-" << (void*)((uint64_t)addr + length - 1) << "]" + << " SIZE: 0x" << length << " (" << allocator_name << ")."; + if (mr != nullptr) { + mutex_lock l(mrs_mu_); + auto iter = std::upper_bound(mrs_.begin(), mrs_.end(), addr, &Comparator); + mrs_.insert(iter, {mr, &MRDeleter}); + } else { + LOG(WARNING) << "Cannot register memory region"; + } +} + +void RdmaMemoryMgr::EvictMemoryRegion(void* addr, size_t length) { + if (length == 0) return; + mutex_lock l(mrs_mu_); + auto iter = std::upper_bound(mrs_.begin(), mrs_.end(), addr, &Comparator); + if (iter != std::end(mrs_) && iter->get()->addr == addr) { + mrs_.erase(iter); + RDMA_LOG(1) << "Evict memory region 0x" << std::hex << iter->get()->rkey; + + } else { + LOG(WARNING) << "Failed to de-register memory region"; + } +} + +const TensorMetaData* RdmaMemoryMgr::GetTensorMetaData( + const std::string& tensor_name) { + mutex_lock l(tensor_meta_data_mu_); + auto it = tensors_meta_data_.find(tensor_name); + if (it == tensors_meta_data_.end()) { + return nullptr; + } + return &it->second; +} + +const TensorMetaData* RdmaMemoryMgr::SetTensorMetaData( + const std::string& tensor_name, DataType dtype, const TensorShape& shape, + bool is_dead, size_t proto_size) { + mutex_lock l(tensor_meta_data_mu_); + TensorMetaData& meta_data = tensors_meta_data_[tensor_name]; + meta_data.data_type_ = dtype; + meta_data.tensor_shape_ = shape; + meta_data.proto_size_ = proto_size; + meta_data.is_dead_ = is_dead; + return &meta_data; +} + +//***************************************************************************** +// RdmaTensorRequest +//***************************************************************************** + +RdmaTensorRequest::RdmaTensorRequest( + uint32_t index, const string& key, int64 step_id, RdmaChannel* channel, + Device* dst_dev, const Rendezvous::Args recv_args, + const RdmaTensorRequest::RecvDoneCallback& done) + : index_(index), + key_(key), + step_id_(step_id), + channel_(channel), + dst_dev_(dst_dev), + recv_args_(recv_args), + meta_data_(RdmaMemoryMgr::Singleton().GetTensorMetaData(key)), + result_tensor_(nullptr), + proxy_tensor_(nullptr), + rdma_addr_(nullptr), + mr_(nullptr), + done_(done) {} + +RdmaTensorRequest::~RdmaTensorRequest() { DeallocateTensors(); } + +void RdmaTensorRequest::Done(const Status& s) { + Tensor val = std::move(*result_tensor_); + +#ifdef RDMA_DATA_VALIDATION + // Validate checksum + // Unfortunately we can't always do a Checksum directly on the result tensor. + // If the result tensor is on GPU, then we need to copy it back to CPU. If + // we happen to be in the midst of a proxy callback, then the copying will + // get stuck. + uint64_t checksum = (proxy_tensor_ != nullptr) + ? Checksum(nullptr, nullptr, *proxy_tensor_) + : Checksum(dst_dev_, recv_args_.device_context, val); + ValidateChecksum(checksum_, checksum, val, index_, key_, "RDMA"); +#endif + + Rendezvous::Args recv_args = std::move(recv_args_); + bool is_dead = (meta_data_ == nullptr) ? false : meta_data_->is_dead_; + RecvDoneCallback done = done_; + DeallocateTensors(); + channel_->RemoveTensorRequest(index_); + done(s, Rendezvous::Args(), recv_args, val, is_dead); +} + +void RdmaTensorRequest::DeallocateTensors() { + if (result_tensor_ != nullptr) { + delete result_tensor_; + result_tensor_ = nullptr; + } + if (proxy_tensor_ != nullptr) { + delete proxy_tensor_; + proxy_tensor_ = nullptr; + } +} + +bool RdmaTensorRequest::AllocateTensors() { + result_tensor_ = + new Tensor(dst_dev_->GetAllocator(recv_args_.alloc_attrs), + meta_data_->data_type_, meta_data_->tensor_shape_); + + size_t tensor_size = result_tensor_->TotalBytes(); + bool can_memcpy = DataTypeCanUseMemcpy(result_tensor_->dtype()); + if (can_memcpy) { + if (tensor_size == 0) { + return true; + } + rdma_addr_ = DMAHelper::base(result_tensor_); + mr_ = RdmaMemoryMgr::Singleton().FindMemoryRegion(rdma_addr_, tensor_size); +#if GOOGLE_CUDA + if (mr_ == nullptr) { + // Can't RDMA directly to result. Use a proxy. + proxy_tensor_ = + new Tensor(ProcessState::singleton()->GetCUDAHostAllocator(0), + result_tensor_->dtype(), result_tensor_->shape()); + rdma_addr_ = DMAHelper::base(proxy_tensor_); + mr_ = + RdmaMemoryMgr::Singleton().FindMemoryRegion(rdma_addr_, tensor_size); + } +#endif + } else { + uint32_t proto_size = meta_data_->proto_size_; + rdma_addr_ = malloc(proto_size); + mr_ = ibv_reg_mr(RdmaMemoryMgr::Singleton().pd_, rdma_addr_, proto_size, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + } + CHECK(mr_ != nullptr) << " No memory region found for address " << rdma_addr_ + << ": " << key_; + return true; +} + +void RdmaTensorRequest::AllocateTensorsAsync(StatusCallback done) { + AllocateTensors(); + bool on_host = recv_args_.alloc_attrs.on_host(); + if (dst_dev_->tensorflow_gpu_device_info() && !on_host && + (proxy_tensor_ == nullptr)) { +#if GOOGLE_CUDA + // We need to sync the memory allocation on the GPU: + StreamGPUOp(dst_dev_, recv_args_.device_context, done); +#endif + } else { + done(Status::OK()); + } +} + +void RdmaTensorRequest::Send(RdmaMessageType message_type) { + RdmaMessageBuffer* rb = channel_->tx_message_buffer_; + RdmaMessage rm; + rm.type_ = message_type; + rm.request_index_ = index_; + rm.name_size_ = key_.size(); + rm.name_ = key_; + rm.step_id_ = step_id_; + rm.remote_addr_ = (uint64_t)rdma_addr_; + if (meta_data_ != nullptr) { + rm.data_type_ = meta_data_->data_type_; + rm.tensor_shape_ = meta_data_->tensor_shape_; + rm.is_dead_ = meta_data_->is_dead_; + rm.tensor_bytes_ = meta_data_->proto_size_; + } else { + rm.data_type_ = DT_INVALID; + } + rm.rkey_ = (mr_ == nullptr) ? 0 : mr_->rkey; + + RDMA_LOG(1) << "Step 0x" << std::hex << rm.step_id_ << std::dec + << ": Sending " << MessageTypeToString(message_type) << " #" + << index_ << ": " << rm.name_ << " on " << rdma_addr_ + << " (rkey: 0x" << std::hex << rm.rkey_ << ")"; + + string message = RdmaMessage::CreateMessage(rm); + rb->EnqueueItem(message); + rb->SendNextItem(); +} + +void RdmaTensorRequest::RecvTensorMetaData(DataType dtype, TensorShape shape, + bool is_dead, size_t proto_size) { + meta_data_ = RdmaMemoryMgr::Singleton().SetTensorMetaData( + key_, dtype, shape, is_dead, proto_size); + + DeallocateTensors(); + AllocateTensorsAsync( + [this](const Status& s) { Send(RDMA_MESSAGE_TENSOR_RE_REQUEST); }); +} + +void RdmaTensorRequest::RecvTensorContent() { + bool can_memcpy = DataTypeCanUseMemcpy(meta_data_->data_type_); + size_t message_size = + can_memcpy ? result_tensor_->TotalBytes() : meta_data_->proto_size_; + RDMA_LOG(1) << "Step 0x" << std::hex << step_id_ << std::dec + << ": Received tensor content #" << index_ << ": " << key_ + << " (Size: 0x" << std::hex << message_size << ")"; + + Tensor val; + +#if GOOGLE_CUDA + if (proxy_tensor_ != nullptr) { + CountCopies(key_, (void*)DMAHelper::base(proxy_tensor_), + (void*)DMAHelper::base(result_tensor_), + result_tensor_->TotalBytes(), false); + GPUUtil::CopyCPUTensorToGPU(proxy_tensor_, recv_args_.device_context, + dst_dev_, result_tensor_, + [this](const Status& s) { + CHECK(s.ok()) << "copy tensor to gpu sync"; + Done(s); + }); + return; + } +#endif + + if (can_memcpy) { + Done(Status::OK()); + } else { + RDMA_LOG(2) << "Decoding proto: " << key_ + << " (Size: " << meta_data_->proto_size_ << ")"; + TensorProto proto; + CHECK(ParseProtoUnlimited(&proto, rdma_addr_, meta_data_->proto_size_)) + << "fail to parse proto from array"; + ibv_dereg_mr(mr_); + free(rdma_addr_); + Status s = dst_dev_->MakeTensorFromProto(proto, recv_args_.alloc_attrs, + result_tensor_); + Done(s); + } +} + +void RdmaTensorRequest::RecvErrorStatus(const Status& status) { + if (result_tensor_ == nullptr) { + result_tensor_ = new Tensor(); + } + LOG(ERROR) << "Received RDMA_MESSAGE_ERROR_STATUS: " << status.ToString(); + Done(status); +} + +void RdmaTensorRequest::Start() { + meta_data_ = RdmaMemoryMgr::Singleton().GetTensorMetaData(key_); + if (meta_data_ != nullptr) { + AllocateTensorsAsync( + [this](const Status& s) { Send(RDMA_MESSAGE_TENSOR_REQUEST); }); + } else { + Send(RDMA_MESSAGE_TENSOR_REQUEST); + } } } // end namespace tensorflow diff --git a/tensorflow/contrib/verbs/rdma.h b/tensorflow/contrib/verbs/rdma.h index fea2327d77ffff67c4b3c45835a81f790bbd1574..94203ee2b3654bffe82d203cde8780a64f63ba2a 100644 --- a/tensorflow/contrib/verbs/rdma.h +++ b/tensorflow/contrib/verbs/rdma.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_H_ +#ifndef TENSORFLOW_CONTRIB_VERBS_RDMA_H_ +#define TENSORFLOW_CONTRIB_VERBS_RDMA_H_ #ifdef TENSORFLOW_USE_VERBS @@ -27,6 +27,7 @@ limitations under the License. #include #include +#include "tensorflow/contrib/verbs/verbs_util.h" #include "tensorflow/core/distributed_runtime/worker_env.h" #include "tensorflow/core/framework/rendezvous.h" #include "tensorflow/core/framework/tensor.h" @@ -43,6 +44,11 @@ namespace tensorflow { #define SL_DEFAULT 0 #define TRAFFIC_CLASS 0 +#define RDMA_LOG_0 LOG(INFO) +#define RDMA_LOG_1 VLOG(1) +#define RDMA_LOG_2 VLOG(2) +#define RDMA_LOG(LEVEL) RDMA_LOG_##LEVEL + struct RdmaParams { uint8_t port_num; uint8_t sgid_index; @@ -67,38 +73,305 @@ struct RemoteMR { uint64_t remote_addr; uint32_t rkey; }; -enum BufferStatus { - none, - idle, - busy +enum BufferStatus { none, idle, busy }; +enum Location { local, remote }; + +enum RdmaMessageType { + RDMA_MESSAGE_META_DATA_UPDATE, + RDMA_MESSAGE_TENSOR_RE_REQUEST, + RDMA_MESSAGE_TENSOR_REQUEST, + RDMA_MESSAGE_ERROR_STATUS, +}; + +struct RdmaMessage { + RdmaMessageType type_; + uint16_t name_size_; + string name_; + int64 step_id_; + uint64_t request_index_; + union { + uint64_t remote_addr_; +#ifdef RDMA_DATA_VALIDATION + uint64_t checksum_; +#endif + }; + uint32_t rkey_; + bool is_dead_; + DataType data_type_; + TensorShape tensor_shape_; + size_t tensor_bytes_; + + // For error status: + Status status_; + + // type|name_size|name|step_id|request_index|remote_addr/checksum|rkey|... + // 1B| 2B | 512| 8B | 8B | 8B | 4B |... + // ...|is_dead|data_type|tensor_shape|tensor_bytes|error_status | + // ...| 1B | XB | XB | 8B |size - 4B, proto - XB | + static const size_t kNameCapacity = 512; + static const size_t kTypeStartIndex = 0; + static const size_t kNameSizeStartIndex = kTypeStartIndex + sizeof(type_); + static const size_t kNameStartIndex = + kNameSizeStartIndex + sizeof(name_size_); + static const size_t kStepIdStartIndex = kNameStartIndex + kNameCapacity; + static const size_t kRequestIndexStartIndex = + kStepIdStartIndex + sizeof(step_id_); + static const size_t kRemoteAddrStartIndex = + kRequestIndexStartIndex + sizeof(request_index_); + static const size_t kChecksumStartIndex = kRemoteAddrStartIndex; + static const size_t kRkeyStartIndex = + kRemoteAddrStartIndex + sizeof(remote_addr_); + static const size_t kIsDeadStartIndex = kRkeyStartIndex + sizeof(rkey_); + static const size_t kDataTypeStartIndex = + kIsDeadStartIndex + sizeof(is_dead_); + static const size_t kTensorShapeStartIndex = + kDataTypeStartIndex + sizeof(data_type_); + static const size_t kTensorBytesStartIndex = + kTensorShapeStartIndex + sizeof(TensorShape); + static const size_t kErrorStatusStartIndex = + kTensorBytesStartIndex + sizeof(tensor_bytes_); + static const size_t kErrorStatusMaxSize = 4096; + + static const size_t kMessageTotalBytes = kErrorStatusStartIndex; + static const size_t kRdmaMessageBufferSize = + kMessageTotalBytes + kErrorStatusMaxSize; + static string CreateMessage(const RdmaMessage& rm); + static void ParseMessage(RdmaMessage& rm, void* buffer); +}; + +// Immediate types for RDMA write +enum RdmaImmDataType { + RDMA_IMM_MAX_REQUEST_ID = 0xFFFFFFFD, + RDMA_IMM_DATA_ACK = 0xFFFFFFFE, + RDMA_IMM_DATA_MESSAGE = 0xFFFFFFFF }; -enum Location { - local, - remote + +// Write types for RDMA write-complete events +enum RdmaWriteIDType { + RDMA_WRITE_ID_ACK, + RDMA_WRITE_ID_MESSAGE, + RDMA_WRITE_ID_TENSOR_WRITE }; -enum BufferType { - ACK, - MESSAGE, - TENSOR + +// Context for RDMA write-complete events +class RdmaWriteID { + public: + RdmaWriteID(RdmaWriteIDType write_type, void* write_context) + : write_type(write_type), write_context(write_context) {} + + RdmaWriteIDType write_type; + void* write_context; }; -enum RdmaMessageType { - RDMA_MESSAGE_ACK, - RDMA_MESSAGE_BUFFER_IDLE, - RDMA_MESSAGE_BUFFER_REQUEST, - RDMA_MESSAGE_BUFFER_RESPONSE, - RDMA_MESSAGE_TENSOR_REQUEST, - RDMA_MESSAGE_TENSOR_WRITE + +// Tensor meta-data +class TensorMetaData { + public: + TensorShape tensor_shape_; + DataType data_type_; + size_t proto_size_; + bool is_dead_; + + std::ostream& print(std::ostream& out) const { + out << "Dtype = " << DataTypeString(data_type_) + << ", Shape = " << tensor_shape_.DebugString() << ", Proto size = 0x" + << std::hex << proto_size_ << ", Is dead = " << is_dead_; + return out; + } +}; + +inline std::ostream& operator<<(std::ostream& out, + const TensorMetaData& meta_data) { + return meta_data.print(out); +} + +class RdmaChannel; + +void MRDeleter(ibv_mr* mr); +using MemoryRegionPtr = std::unique_ptr; + +// RdmaMemoryMgr +// Manages the local meta-data cache, and the registered RDMA memory regions. +class RdmaMemoryMgr { + public: + static RdmaMemoryMgr& Singleton() { + static RdmaMemoryMgr instance; + return instance; + } + + // Memory regions + ibv_mr* FindMemoryRegion(void* addr, size_t length); + void InsertMemoryRegion(void* addr, size_t length, + const std::string& allocator_name); + void EvictMemoryRegion(void* addr, size_t length); + + // Tensor meta-data cache + const TensorMetaData* GetTensorMetaData(const std::string& tensor_name); + const TensorMetaData* SetTensorMetaData(const std::string& tensor_name, + DataType dtype, + const TensorShape& shape, + bool is_dead, size_t proto_size); + + struct ibv_pd* pd_; + + protected: + RdmaMemoryMgr() : pd_(nullptr) {} + + static bool Comparator(const void* ptr, const MemoryRegionPtr& other) { + return ptr < reinterpret_cast(other->addr) + other->length; + } + + private: + mutex tensor_meta_data_mu_; + std::unordered_map tensors_meta_data_; + + // Managed memory regions + mutex mrs_mu_; + std::vector mrs_ GUARDED_BY(mrs_mu_); }; -class RdmaBuffer; + +// RdmaTensorRequest +// Represents a single tensor request. +class RdmaTensorRequest { + public: + typedef Rendezvous::DoneCallback RecvDoneCallback; + + // Creates a tensor request identified by index. + RdmaTensorRequest(uint32_t index, const string& key, int64 step_id, + RdmaChannel* channel, Device* dst_dev, + const Rendezvous::Args recv_args, + const RecvDoneCallback& done); + ~RdmaTensorRequest(); + + // Request unique index. + uint32_t index() { return index_; } + + // Start the tensor request sequence. + // + // 1. Allocate the result tensor (and proxy tensor if required). + // 2. Send RDMA_MESSAGE_TENSOR_REQUEST to the remote side. + void Start(); + + // Receive tensor meta-data. + // + // 1. Update the local meta-data cache. + // 2. Reallocate the result tensor (and proxy tensor if required). + // 3. Re-send the request to the remote side. + void RecvTensorMetaData(DataType dtype, TensorShape shape, bool is_dead, + size_t proto_size); + + // Receive tensor content (RDMA write was completed). + // + // Decode proto if required and/or move to GPU if the content was not + // written to it directly (GPU direct is not avaliable). Afterwards, + // invoke Done(). + void RecvTensorContent(); + + // Receive error status (in case of a remote error). + // Invoke Done() with the status code. + void RecvErrorStatus(const Status& status); + +#ifdef RDMA_DATA_VALIDATION + // Receive tensor checksum + // + // For validation: Get and store the Tensor's expected checksum for the + // current request. Compare the result Tensor's checksum with the stored + // checksum right before invoking Done(). + void RecvTensorChecksum(uint64_t checksum) { checksum_ = checksum; } +#endif + + private: + void Done(const Status& s); + void Send(RdmaMessageType message_type); + bool AllocateTensors(); + void AllocateTensorsAsync(StatusCallback done); + void DeallocateTensors(); + + uint32_t index_; + string key_; + int64 step_id_; + RdmaChannel* channel_; + Device* dst_dev_; + Rendezvous::Args recv_args_; + const TensorMetaData* meta_data_; + Tensor* result_tensor_; + Tensor* proxy_tensor_; + void* rdma_addr_; + ibv_mr* mr_; + RecvDoneCallback done_; +#ifdef RDMA_DATA_VALIDATION + uint64_t checksum_; +#endif +}; + +// RdmaTensorResponse +// Represents a single tensor response. +class RdmaTensorResponse { + public: + // Creates a response for request message. + RdmaTensorResponse(RdmaChannel* channel, const RdmaMessage& rm) + : channel_(channel), rm_(rm) {} + + void Update(const RdmaMessage& rm) { rm_ = rm; } + + // Start the tensor response sequence. + // + // 1. Find the tensor in the local tag-match table and invoke RecvHandler. + // (Using RecvLocalAsync()). + // 2. Compare the tensor's meta-data to the meta-data in the message (taken + // from the requester's local cache). + // If meta-data changed: + // a. Clone the tensor to be sent later. + // b. Send a meta-data update message and wait for re-request. + // Else: + // a. Send the tensor's content (using direct RDMA write). + void Start(); + + // Resume the response sequence, after a re-request. + // + // 1. Send the tensor's content that was cloned earlier. + void Resume(); + + // Destroy the response's resources and remove it from the pending list. + void Destroy(); + + private: + void RecvHandler(Rendezvous::ParsedKey parsed, + const Rendezvous::Args& send_args, + const Rendezvous::Args& recv_args, const Tensor& in, + bool is_dead); + void Clone(const Tensor& in, const TensorProto& proto, bool is_dead); + void Send(const Tensor& in, const TensorProto& proto, bool is_dead, + const Status& status); + bool TensorMetaDataChanged(const Tensor& in, bool is_dead); + Status PrepareRecvTensor(const Rendezvous::ParsedKey& parsed, + Device** src_dev); + void SendMetaData(const Tensor& in, const TensorProto& proto, bool is_dead); + void SendContent(const Tensor& in, const TensorProto& proto, bool is_dead); + void SendErrorStatus(const Status& status); + + RdmaChannel* channel_; + RdmaMessage rm_; // The request message + Device* src_dev_ = nullptr; + TensorBuffer* src_buffer_ = nullptr; + void* src_addr_ = nullptr; + ibv_mr* mr_ = nullptr; + uint64_t checksum_ = 0; + bool meta_data_changed_ = false; + + // Re-item: + TensorProto* proto_ = nullptr; + Tensor* tensor_ = nullptr; + bool is_dead_ = false; +}; + +class RdmaMessageBuffer; // Class that represents the Rdma Adapter. // Responsible for creation of the completion queue, and handling // of work completions. class RdmaAdapter { friend class RdmaChannel; - friend class RdmaBuffer; - friend class RdmaAckBuffer; friend class RdmaMessageBuffer; - friend class RdmaTensorBuffer; + friend class RdmaTensorResponse; friend class RdmaMgr; friend class RdmaRemoteRendezvous; @@ -133,10 +406,10 @@ class RdmaAdapter { // Responsible for connecting queue pairs. class RdmaChannel { friend class RdmaAdapter; - friend class RdmaBuffer; - friend class RdmaAckBuffer; friend class RdmaMessageBuffer; friend class RdmaTensorBuffer; + friend class RdmaTensorRequest; + friend class RdmaTensorResponse; friend class RdmaMgr; friend class RdmaRemoteRendezvous; @@ -146,22 +419,28 @@ class RdmaChannel { ~RdmaChannel(); inline const RdmaAddress& self() { return self_; } RdmaAddress address() const; - inline const std::vector& message_buffers() const { + inline const std::vector& message_buffers() const { return message_buffers_; } void Connect(const RdmaAddress& remoteAddr); void Connect(); void Recv(); - RdmaBuffer* FindBuffer(const uint32_t index); - RdmaBuffer* FindBuffer(const string& name); - RdmaBuffer* FindOrCreateBuffer(const string& name, - BufferType buffer_type = TENSOR); - uint32_t LookupBufferIndex(const string& buffer_name); void SetRemoteAddress(const RdmaAddress& ra, bool override); - void InsertRecvCallback(const string& key, std::function recv_done); - void RemoveRecvCallback(const string& key); - void RunRecvCallback(const string& key); - static const int kNumMessageBuffers = 4; + + // Requests: + RdmaTensorRequest* InsertTensorRequest( + const string& key, int64 step_id, Device* dst_dev, + const Rendezvous::Args recv_args, + const RdmaTensorRequest::RecvDoneCallback& done); + void RemoveTensorRequest(uint32_t request_index); + RdmaTensorRequest* GetTensorRequest(uint32_t request_index); + + // Responses: + RdmaTensorResponse* AddTensorResponse(const RdmaMessage& rm); + RdmaTensorResponse* UpdateTensorResponse(const RdmaMessage& rm); + void RemoveTensorResponse(uint32_t request_index); + + static const int kNumMessageBuffers = 2; static const int kPingRecvWrid = 0; private: @@ -179,36 +458,31 @@ class RdmaChannel { string remote_name_; ibv_qp* qp_; mutex mu_; - bool connected_ GUARDED_BY(bt_mu_) = false; - RdmaAddress remote_ GUARDED_BY(bt_mu_); - bool remote_set_ GUARDED_BY(bt_mu_) = false; + bool connected_ GUARDED_BY(mu_) = false; + RdmaAddress remote_ GUARDED_BY(mu_); + bool remote_set_ GUARDED_BY(mu_) = false; mutex ct_mu_; - typedef std::unordered_map > CallbackTable; - CallbackTable callback_table_ GUARDED_BY(ct_mu_); - mutex bt_mu_; - typedef std::unordered_map BufferTable; - BufferTable buffer_table_ GUARDED_BY(bt_mu_); - typedef std::unordered_map BufferIndexNameTable; - BufferIndexNameTable buffer_index_name_table_ GUARDED_BY(bt_mu_); - typedef std::unordered_map BufferNameIndexTable; - BufferNameIndexTable buffer_name_index_table_ GUARDED_BY(bt_mu_); - RdmaBuffer* tx_message_buffer_; - RdmaBuffer* rx_message_buffer_; - RdmaBuffer* tx_ack_buffer_; - RdmaBuffer* rx_ack_buffer_; - std::vector message_buffers_; + typedef std::unordered_map RequestTable; + RequestTable request_table_ GUARDED_BY(ct_mu_); + uint32_t request_serial_ GUARDED_BY(ct_mu_); + mutex responses_mu_; + typedef std::unordered_map ResponsesTable; + ResponsesTable responses_table_ GUARDED_BY(responses_mu_); + RdmaMessageBuffer* tx_message_buffer_; + RdmaMessageBuffer* rx_message_buffer_; + std::vector message_buffers_; }; -// Class that represents a buffer for Rdma writes and reads. -class RdmaBuffer { +// Class that represents a buffer for Rdma message sending. +class RdmaMessageBuffer { friend class RdmaChannel; friend class RdmaAdapter; friend class RdmaMgr; friend class RdmaRemoteRendezvous; public: - explicit RdmaBuffer(RdmaChannel* channel, string name); - virtual ~RdmaBuffer(); + explicit RdmaMessageBuffer(RdmaChannel* channel, string name); + ~RdmaMessageBuffer(); inline void* buffer() const { return buffer_; } inline ibv_mr* self() const { return self_; } @@ -223,13 +497,15 @@ class RdmaBuffer { } void FreeBuffer(); void EnqueueItem(string Item); - virtual void SendNextItem() {}; + void SendNextItem(); void CreateCPUBuffer(size_t size, bool lock = true); void SetRemoteMR(RemoteMR rmi, bool override); - uint32_t LookupBufferIndex(const string& buffer_name) { - return const_cast(channel_)->LookupBufferIndex(buffer_name); - } void Write(uint32_t imm_data, size_t buffer_size); + static void Write(const RdmaChannel* channel, uint32_t imm_data, + size_t buffer_size, uint64_t src_addr, uint32_t lkey, + uint64_t remote_addr, uint32_t rkey, + RdmaWriteIDType write_type, void* write_context); + static void SendAck(const RdmaChannel* channel); protected: const RdmaChannel* channel_; @@ -245,126 +521,7 @@ class RdmaBuffer { BufferStatus remote_status_ GUARDED_BY(mu_) = none; }; -class RdmaAckBuffer : public RdmaBuffer { - public: - explicit RdmaAckBuffer(RdmaChannel* channel, string name); - virtual ~RdmaAckBuffer() override {} - void SendNextItem() override; -}; - -class RdmaMessageBuffer : public RdmaBuffer { - friend class RdmaChannel; - friend class RdmaAapater; - - public: - explicit RdmaMessageBuffer(RdmaChannel* channel, string name); - virtual ~RdmaMessageBuffer() override {} - void SendNextItem() override; -}; - -class RdmaTensorBuffer : public RdmaBuffer { - public: - explicit RdmaTensorBuffer(RdmaChannel* channel, string name); - virtual ~RdmaTensorBuffer() override; - void SendNextItem() override; - void PostCopyOperations(bool can_memcpy, size_t buffer_size, - size_t tensor_bytes, const string& key, - const Tensor& in, int64 step_id, bool is_dead, - const string& key_with_step_id, const Tensor* copy, - const TensorProto* proto, const StringPiece* copy_buf, - const Rendezvous::Args& send_args, - const Rendezvous::Args& recv_args); - - void ReSendNextItem(); - - private: - Rendezvous::DoneCallback getRecvTensorCallback( - const string& key_with_step_id, const string& key, int64 step_id, - const Rendezvous::ParsedKey& parsed); - - struct ReItem { - Rendezvous::Args send_args; - Rendezvous::Args recv_args; - Tensor in; - bool is_dead; - - ReItem(const Rendezvous::Args& send_args_, - const Rendezvous::Args& recv_args_, const Tensor& in_, bool is_dead_) - : send_args(send_args_), - recv_args(recv_args_), - in(in_), - is_dead(is_dead_) { - if (send_args.device_context) { - send_args.device_context->Ref(); - } - if (recv_args.device_context) { - recv_args.device_context->Ref(); - } - } - - ~ReItem() { - if (send_args.device_context) { - send_args.device_context->Unref(); - } - if (recv_args.device_context) { - recv_args.device_context->Unref(); - } - } - }; - typedef std::map Table; - typedef Table::iterator Itable; - - std::queue requeue GUARDED_BY(mu_); - Table retable GUARDED_BY(mu_); -}; - -struct RdmaMessage { - RdmaMessageType type_; - uint16_t name_size_; - string name_; - int64 step_id_; - uint64_t buffer_size_; - uint64_t remote_addr_; - uint32_t rkey_; - bool is_dead_; - DataType data_type_; - TensorShape tensor_shape_; - size_t tensor_bytes_; - - // type|name_size|name|step_id|buffer_size|remote_addr|rkey|is_dead|... - // 1B| 2B | 512| 8B | 8B | 8B | 4B | 1B |... - // ...|data_type|tensor_shape|tensor_bytes|tensor_buffer - // ...| XB | XB | 8B |... - // - static const size_t kNameCapacity = 512; - static const size_t kTypeStartIndex = 0; - static const size_t kNameSizeStartIndex = kTypeStartIndex + sizeof(type_); - static const size_t kNameStartIndex = - kNameSizeStartIndex + sizeof(name_size_); - static const size_t kStepIdStartIndex = kNameStartIndex + kNameCapacity; - static const size_t kBufferSizeStartIndex = - kStepIdStartIndex + sizeof(step_id_); - static const size_t kRemoteAddrStartIndex = - kBufferSizeStartIndex + sizeof(buffer_size_); - static const size_t kRkeyStartIndex = - kRemoteAddrStartIndex + sizeof(remote_addr_); - static const size_t kIsDeadStartIndex = kRkeyStartIndex + sizeof(rkey_); - static const size_t kDataTypeStartIndex = - kIsDeadStartIndex + sizeof(is_dead_); - static const size_t kTensorShapeStartIndex = - kDataTypeStartIndex + sizeof(data_type_); - static const size_t kTensorBytesStartIndex = - kTensorShapeStartIndex + sizeof(TensorShape); - static const size_t kTensorBufferStartIndex = - kTensorBytesStartIndex + sizeof(tensor_bytes_); - static const size_t kMessageTotalBytes = kTensorBufferStartIndex; - static const size_t kRdmaMessageBufferSize = kMessageTotalBytes; - static const size_t kRdmaAckBufferSize = kMessageTotalBytes; - static string CreateMessage(const RdmaMessage& rm); - static void ParseMessage(RdmaMessage& rm, void* buffer); -}; - } // namespace tensorflow #endif // TENSORFLOW_USE_VERBS -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_H_ +#endif // TENSORFLOW_CONTRIB_VERBS_RDMA_H_ diff --git a/tensorflow/contrib/verbs/rdma_mgr.cc b/tensorflow/contrib/verbs/rdma_mgr.cc index 9cb307bcfa06cfdf5ecb9b4faa1d3710e5701080..369bd986df5313955bc22d6e5c6d38815908ada3 100644 --- a/tensorflow/contrib/verbs/rdma_mgr.cc +++ b/tensorflow/contrib/verbs/rdma_mgr.cc @@ -16,11 +16,16 @@ limitations under the License. #ifdef TENSORFLOW_USE_VERBS #include "tensorflow/contrib/verbs/rdma_mgr.h" +#include #include #include "tensorflow/contrib/verbs/grpc_verbs_client.h" #include "tensorflow/contrib/verbs/verbs_service.pb.h" +#include "tensorflow/core/common_runtime/bfc_allocator.h" +#include "tensorflow/core/common_runtime/gpu/gpu_util.h" +#include "tensorflow/core/common_runtime/gpu/process_state.h" #include "tensorflow/core/distributed_runtime/rpc/grpc_worker_cache.h" #include "tensorflow/core/distributed_runtime/session_mgr.h" +#include "tensorflow/core/framework/allocator_registry.h" #include "tensorflow/core/lib/core/status.h" namespace tensorflow { @@ -53,7 +58,7 @@ RdmaMgr::RdmaMgr(const WorkerEnv* const worker_env, void RdmaMgr::SetupChannels() { for (const auto& p : channel_table_) { string worker_name = p.first; - LOG(INFO) << "connecting to remote node " << worker_name; + RDMA_LOG(2) << "Connecting to remote node " << worker_name; RdmaChannel* rc = p.second; GetRemoteAddressRequest req; GetRemoteAddressResponse resp; @@ -78,39 +83,49 @@ void RdmaMgr::SetupChannels() { mr->set_rkey(rc->message_buffers_[i]->self_->rkey); } // synchronous call - Status s = client->GetRemoteAddress(&req, &resp); - // save obtained remote addresses - // connect to the remote channel - if (s.ok()) { - CHECK(worker_name.compare(resp.host_name()) == 0); - RdmaAddress ra; - ra.lid = resp.channel().lid(); - ra.qpn = resp.channel().qpn(); - ra.psn = resp.channel().psn(); - ra.snp = resp.channel().snp(); - ra.iid = resp.channel().iid(); - rc->SetRemoteAddress(ra, false); - rc->Connect(); - int i = 0; - int idx[] = {1, 0, 3, 2}; - for (const auto& mr : resp.mr()) { - // the connections are crossed, i.e. - // local tx_message_buffer <---> remote rx_message_buffer_ - // local rx_message_buffer <---> remote tx_message_buffer_ - // local tx_ack_buffer <---> remote rx_ack_buffer_ - // local rx_ack_buffer <---> remote tx_ack_buffer_ - // hence idx[] = {1, 0, 3, 2}. - RdmaBuffer* rb = rc->message_buffers_[idx[i]]; - RemoteMR rmr; - rmr.remote_addr = mr.remote_addr(); - rmr.rkey = mr.rkey(); - rb->SetRemoteMR(rmr, false); - i++; + Status s; + int attempts = 0; + static const int max_num_attempts = 5; + do { + s = client->GetRemoteAddress(&req, &resp); + // save obtained remote addresses + // connect to the remote channel + if (s.ok()) { + CHECK(worker_name.compare(resp.host_name()) == 0); + RdmaAddress ra; + ra.lid = resp.channel().lid(); + ra.qpn = resp.channel().qpn(); + ra.psn = resp.channel().psn(); + ra.snp = resp.channel().snp(); + ra.iid = resp.channel().iid(); + rc->SetRemoteAddress(ra, false); + rc->Connect(); + int i = 0; + int idx[] = {1, 0}; + for (const auto& mr : resp.mr()) { + // the connections are crossed, i.e. + // local tx_message_buffer <---> remote rx_message_buffer_ + // local rx_message_buffer <---> remote tx_message_buffer_ + // hence idx[] = {1, 0}. + RdmaMessageBuffer* rb = rc->message_buffers_[idx[i]]; + RemoteMR rmr; + rmr.remote_addr = mr.remote_addr(); + rmr.rkey = mr.rkey(); + rb->SetRemoteMR(rmr, false); + i++; + } + CHECK(i == RdmaChannel::kNumMessageBuffers); + } else { + LOG(ERROR) << "Connecting to " << worker_name << ": Got " + << s.error_message() << ". Retrying (" << (attempts + 1) + << "/" << max_num_attempts << ")..."; + if (++attempts == max_num_attempts) { + break; + } + worker_env_->env->SleepForMicroseconds(2000000); } - CHECK(i == RdmaChannel::kNumMessageBuffers); - } else { - LOG(ERROR) << s.error_message(); - } + } while (!s.ok()); + RDMA_LOG(0) << "Connected to remote node " << worker_name; delete client; } } @@ -144,19 +159,17 @@ bool RdmaMgr::ConnectivityCheck() { ibv_wc_status s = rdma_adapter_->wc_[i].status; // recv complete if ((int)rdma_adapter_->wc_[i].wr_id == RdmaChannel::kPingRecvWrid) { - CHECK(s == IBV_WC_SUCCESS) << ": " << ibv_wc_status_str( - rdma_adapter_->wc_[i].status) - << "(" << rdma_adapter_->wc_[i].status - << ") for PING_RECV_WRID"; + CHECK(s == IBV_WC_SUCCESS) + << ": " << ibv_wc_status_str(rdma_adapter_->wc_[i].status) << "(" + << rdma_adapter_->wc_[i].status << ") for PING_RECV_WRID"; ++rcnt; // send complete } else { RdmaChannel* rc = reinterpret_cast(rdma_adapter_->wc_[i].wr_id); - CHECK(s == IBV_WC_SUCCESS) << ": " << ibv_wc_status_str( - rdma_adapter_->wc_[i].status) - << "(" << rdma_adapter_->wc_[i].status - << ") to " << rc->remote_name_; + CHECK(s == IBV_WC_SUCCESS) + << ": " << ibv_wc_status_str(rdma_adapter_->wc_[i].status) << "(" + << rdma_adapter_->wc_[i].status << ") to " << rc->remote_name_; ++scnt; } } // for @@ -183,6 +196,139 @@ RdmaChannel* RdmaMgr::FindChannel(const string& name) { return iter->second; } +bool IsGDRAvailable() { +#if defined(__APPLE__) + return false; +#elif defined(PLATFORM_WINDOWS) + return false; +#else + std::ifstream ifs("/proc/modules"); + string line; + while (std::getline(ifs, line)) { + auto sep = line.find(' '); + CHECK_NE(sep, std::string::npos); + if (line.substr(0, sep) == "nv_peer_mem") { + return true; + } + } + return false; +#endif +} + +int TryToReadNumaNode(ibv_device* device) { +#if defined(__APPLE__) + LOG(INFO) << "OS X does not support NUMA - returning NUMA node 0"; + return 0; +#elif defined(PLATFORM_WINDOWS) + // Windows support for NUMA is not currently implemented. Return node 0. + return 0; +#else + VLOG(2) << "Trying to read NUMA node for device: " << device->name; + static const int kUnknownNumaNode = -1; + + auto filename = string(device->ibdev_path) + "/device/numa_node"; + + std::ifstream ifs(filename.c_str()); + string content; + CHECK(std::getline(ifs, content)); + + int32 value; + if (strings::safe_strto32(content, &value)) { + if (value < 0) { + LOG(INFO) << "Successful NUMA node read from SysFS had negative value (" + << value + << "), but there must be at least one NUMA node" + ", so returning NUMA node zero"; + return 0; + } + LOG(INFO) << "NUMA node for device: " << device->name << " is " << value; + return value; + } + return kUnknownNumaNode; +#endif +} + +void MRDeleter(ibv_mr* mr) { + if (mr) { + ibv_dereg_mr(mr); + } +} + +// TODO(byronyi): remove this class duplicated from the one in +// common/runtime/gpu/pool_allocator.h when it is available in common_runtime +class BasicCPUAllocator : public SubAllocator { + public: + ~BasicCPUAllocator() override {} + + void* Alloc(size_t alignment, size_t num_bytes) override { + return port::AlignedMalloc(num_bytes, alignment); + } + void Free(void* ptr, size_t) override { port::AlignedFree(ptr); } +}; + +// TODO(byronyi): remove this class and its registration when the default +// cpu_allocator() returns visitable allocator +class BFCRdmaAllocator : public BFCAllocator { + public: + BFCRdmaAllocator() + : BFCAllocator(new BasicCPUAllocator(), 1LL << 36, true, "cpu_rdma_bfc") { + } +}; + +REGISTER_MEM_ALLOCATOR("BFCRdmaAllocator", 101, BFCRdmaAllocator); + +void RdmaMgr::InitAllocators() { + RdmaMemoryMgr::Singleton().pd_ = rdma_adapter_->pd_; + + Allocator* allocators[] = { +#if GOOGLE_CUDA + ProcessState::singleton()->GetCUDAHostAllocator(0), + ProcessState::singleton()->GetCPUAllocator(0), +#endif // GOOGLE_CUDA + cpu_allocator(), + }; + + using namespace std::placeholders; + + std::set instrumented_; + + // Host memory allocators + for (Allocator* allocator : allocators) { + VisitableAllocator::Visitor alloc_visitor = + std::bind(&RdmaMemoryMgr::InsertMemoryRegion, + &RdmaMemoryMgr::Singleton(), _1, _2, allocator->Name()); + VisitableAllocator::Visitor free_visitor = std::bind( + &RdmaMemoryMgr::EvictMemoryRegion, &RdmaMemoryMgr::Singleton(), _1, _2); + + auto* visitable_allocator = dynamic_cast(allocator); + CHECK(visitable_allocator) + << "is not visitable for instrumentation" << allocator->Name(); + // Make sure we don't instrument the same allocator twice + if (instrumented_.find(allocator) == std::end(instrumented_)) { + visitable_allocator->AddAllocVisitor(alloc_visitor); + visitable_allocator->AddFreeVisitor(free_visitor); + instrumented_.insert(allocator); + LOG(INFO) << "Instrumenting CPU allocator " << allocator->Name(); + } + } + +#if GOOGLE_CUDA + if (IsGDRAvailable()) { + // Note we don't free allocated GPU memory so there is no free visitor + int32_t bus_id = TryToReadNumaNode(rdma_adapter_->context_->device) + 1; + + char buf[8]; + sprintf(buf, "gpu"); + VisitableAllocator::Visitor cuda_alloc_visitor = + std::bind(&RdmaMemoryMgr::InsertMemoryRegion, + &RdmaMemoryMgr::Singleton(), _1, _2, std::string(buf)); + + ProcessState::singleton()->AddGPUAllocVisitor(bus_id, cuda_alloc_visitor); + LOG(INFO) << "Instrumenting GPU allocator with bus_id " << bus_id; + } +#endif // GOOGLE_CUDA +} + } // end namespace tensorflow #endif diff --git a/tensorflow/contrib/verbs/rdma_mgr.h b/tensorflow/contrib/verbs/rdma_mgr.h index e711e604788b12ff0c1a0977a90db21f9f8fa50e..9fffc335bbe2bf47a626736f6d3073f52b32a9c2 100644 --- a/tensorflow/contrib/verbs/rdma_mgr.h +++ b/tensorflow/contrib/verbs/rdma_mgr.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_MGR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_MGR_H_ +#ifndef TENSORFLOW_CONTRIB_VERBS_RDMA_MGR_H_ +#define TENSORFLOW_CONTRIB_VERBS_RDMA_MGR_H_ #ifdef TENSORFLOW_USE_VERBS @@ -38,6 +38,7 @@ class RdmaMgr { RdmaChannel* FindChannel(const string& key); void SetupChannels(); bool ConnectivityCheck(); + void InitAllocators(); const string& local_worker() { return local_worker_; } private: @@ -54,4 +55,4 @@ class RdmaMgr { } // namespace tensorflow #endif // TENSORFLOW_USE_VERBS -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_MGR_H_ +#endif // TENSORFLOW_CONTRIB_VERBS_RDMA_MGR_H_ diff --git a/tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc b/tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc index 74f6681af3c29f370d6cdb37d64e10a30cbb7b84..ad3dce17844c5a43237372fb7fe074416e8b7117 100644 --- a/tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc +++ b/tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc @@ -21,10 +21,6 @@ limitations under the License. #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/dma_helper.h" -#if GOOGLE_CUDA -#include "tensorflow/core/common_runtime/gpu/gpu_util.h" -#include "tensorflow/core/common_runtime/gpu/process_state.h" -#endif // GOOGLE_CUDA #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/strings/numbers.h" #include "tensorflow/core/lib/strings/str_util.h" @@ -36,11 +32,6 @@ class RdmaRemoteRendezvous : public BaseRemoteRendezvous { RdmaRemoteRendezvous(const WorkerEnv* env, int64 step_id, RdmaMgr* rdma_mgr) : BaseRemoteRendezvous(env, step_id), rdma_mgr_(rdma_mgr) {} - void RecvPostCopyOps(const string& key, const string& key_with_step_id, - const Rendezvous::Args& recv_args, - const DoneCallback& done, const RdmaMessage& rm, - RdmaChannel* rc, Tensor& val, const Status& s); - protected: void RecvFromRemoteAsync(const Rendezvous::ParsedKey& parsed, const Rendezvous::Args& args, @@ -74,101 +65,18 @@ void RdmaRemoteRendezvous::RecvFromRemoteAsync( RdmaChannel* rc = rdma_mgr_->FindChannel(src_name); string key(std::move(parsed.FullKey().ToString())); string key_with_step_id = VerbsUtil::AppendStepidToKey(key, step_id_); - // insert callback - rc->InsertRecvCallback(key_with_step_id, [this, key, key_with_step_id, rc, - recv_args, parsed, done]() { - Status src_s, dst_s, s; - Device* src_dev, *dst_dev; - src_s = env_->device_mgr->LookupDevice("CPU:0", &src_dev); - dst_s = env_->device_mgr->LookupDevice(parsed.dst_device, &dst_dev); - if (!src_s.ok() || !dst_s.ok()) { - s = src_s.ok() ? dst_s : src_s; - LOG(ERROR) << "s is not ok, error code " << s.error_message(); - done(s, Args(), recv_args, Tensor(), true); - return; - } - RdmaBuffer* rb = rc->FindBuffer(key); - RdmaMessage rm; - CHECK(rb->size_ >= RdmaMessage::kMessageTotalBytes); - RdmaMessage::ParseMessage(rm, rb->buffer_); - CHECK(rm.type_ == RDMA_MESSAGE_TENSOR_WRITE); - Tensor val; - if (!rm.is_dead_) { - void* input = static_cast(rb->buffer_) + - RdmaMessage::kTensorBufferStartIndex; - bool can_memcpy = DataTypeCanUseMemcpy(rm.data_type_); - if (can_memcpy) { - if (dst_dev->tensorflow_gpu_device_info() && - (!recv_args.alloc_attrs.on_host())) { -#if GOOGLE_CUDA - CHECK(recv_args.device_context) - << "send dev name: " << src_dev->name() - << " gpu_info: " << src_dev->tensorflow_gpu_device_info(); - Allocator* alloc = ProcessState::singleton()->GetCUDAHostAllocator(0); - Tensor copy(alloc, rm.data_type_, rm.tensor_shape_); - memcpy(DMAHelper::base(©), input, rm.tensor_bytes_); - - Allocator* dst_alloc = dst_dev->GetAllocator(recv_args.alloc_attrs); - Tensor gpu_copy(dst_alloc, rm.data_type_, rm.tensor_shape_); - - GPUUtil::CopyCPUTensorToGPU( - ©, recv_args.device_context, dst_dev, &gpu_copy, - [this, gpu_copy, key, key_with_step_id, recv_args, done, rm, rc]( - const Status& s) { - CHECK(s.ok()) << "copy tensor to gpu sync"; - Tensor val; - val = std::move(gpu_copy); - RecvPostCopyOps(key, key_with_step_id, recv_args, done, rm, rc, - val, s); - }); -#endif // GOOGLE_CUDA - return; - } else { - AllocatorAttributes host_alloc_attrs; - host_alloc_attrs.set_gpu_compatible(true); - host_alloc_attrs.set_on_host(true); - Allocator* alloc = dst_dev->GetAllocator(host_alloc_attrs); - Tensor copy(alloc, rm.data_type_, rm.tensor_shape_); - memcpy(DMAHelper::base(©), input, rm.tensor_bytes_); - val = std::move(copy); - } - } else { - TensorProto proto; - CHECK(rm.tensor_bytes_ + RdmaMessage::kTensorBufferStartIndex <= - rb->size_); - CHECK(ParseProtoUnlimited(&proto, input, rm.tensor_bytes_)) - << "fail to parse proto from array"; - s = dst_dev->MakeTensorFromProto(proto, recv_args.alloc_attrs, &val); - } - } - RecvPostCopyOps(key, key_with_step_id, recv_args, done, rm, rc, val, s); - }); - // append key to message queue - RdmaBuffer* rb = rc->tx_message_buffer_; - RdmaMessage rm; - rm.type_ = RDMA_MESSAGE_TENSOR_REQUEST; - rm.name_size_ = key.size(); - rm.name_ = key; - rm.step_id_ = step_id_; - string message = RdmaMessage::CreateMessage(rm); - rb->EnqueueItem(message); - rb->SendNextItem(); -} -void RdmaRemoteRendezvous::RecvPostCopyOps( - const string& key, const string& key_with_step_id, - const Rendezvous::Args& recv_args, const DoneCallback& done, - const RdmaMessage& rm, RdmaChannel* rc, Tensor& val, const Status& s) { - rc->RemoveRecvCallback(key_with_step_id); - RdmaMessage br; - br.type_ = RDMA_MESSAGE_BUFFER_IDLE; - br.name_size_ = key.size(); - br.name_ = key; - string message = RdmaMessage::CreateMessage(br); - RdmaBuffer* tb = rc->tx_message_buffer_; - tb->EnqueueItem(message); - tb->SendNextItem(); - done(s, Args(), recv_args, val, rm.is_dead_); + Device* dst_dev; + s = env_->device_mgr->LookupDevice(parsed.dst_device, &dst_dev); + CHECK(s.ok()) << "s is not ok, error code " << s.error_message(); + if (!s.ok()) { + done(s, Args(), recv_args, Tensor(), true); + return; + } + + RdmaTensorRequest* request = + rc->InsertTensorRequest(key, step_id_, dst_dev, recv_args, done); + request->Start(); } RdmaRendezvousMgr::RdmaRendezvousMgr(const WorkerEnv* env) diff --git a/tensorflow/contrib/verbs/rdma_rendezvous_mgr.h b/tensorflow/contrib/verbs/rdma_rendezvous_mgr.h index 2dedd6c48f96a6ecf2b69c757f525ac1bfd6f2d0..c0d6f59c4842e28e37b2a3b45e955f8d92712dd7 100644 --- a/tensorflow/contrib/verbs/rdma_rendezvous_mgr.h +++ b/tensorflow/contrib/verbs/rdma_rendezvous_mgr.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_RENDEZVOUS_MGR_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_RENDEZVOUS_MGR_H_ +#ifndef TENSORFLOW_CONTRIB_VERBS_RDMA_RENDEZVOUS_MGR_H_ +#define TENSORFLOW_CONTRIB_VERBS_RDMA_RENDEZVOUS_MGR_H_ #ifdef TENSORFLOW_USE_VERBS @@ -60,4 +60,4 @@ class RdmaRendezvousMgr : public BaseRendezvousMgr { } // end namespace tensorflow #endif // TENSORFLOW_USE_VERBS -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_RDMA_RENDEZVOUS_MGR_H_ +#endif // TENSORFLOW_CONTRIB_VERBS_RDMA_RENDEZVOUS_MGR_H_ diff --git a/tensorflow/contrib/verbs/verbs_server_lib.cc b/tensorflow/contrib/verbs/verbs_server_lib.cc index a606ef75a42069b3c32eb13a69e981a5c4c8f83c..1a0b5028febb7b11f979abd179a3227a2615252d 100644 --- a/tensorflow/contrib/verbs/verbs_server_lib.cc +++ b/tensorflow/contrib/verbs/verbs_server_lib.cc @@ -49,8 +49,8 @@ VerbsServer::~VerbsServer() { Status VerbsServer::ChannelCacheFactory(const ServerDef& server_def, GrpcChannelCache** channel_cache) { string name_prefix = - strings::StrCat("/job:", server_def.job_name(), "/replica:0", "/task:", - server_def.task_index()); + strings::StrCat("/job:", server_def.job_name(), "/replica:0", + "/task:", server_def.task_index()); GrpcChannelSpec channel_spec; TF_RETURN_IF_ERROR(ParseChannelSpec(server_def, &channel_spec)); @@ -104,6 +104,7 @@ Status VerbsServer::Start() { [this] { verbs_service_->HandleRPCsLoop(); })); rdma_mgr_->SetupChannels(); CHECK(rdma_mgr_->ConnectivityCheck()) << "Connectivity check failed!"; + rdma_mgr_->InitAllocators(); verbs_state_ = CONNECTED; } } diff --git a/tensorflow/contrib/verbs/verbs_server_lib.h b/tensorflow/contrib/verbs/verbs_server_lib.h index 855380129f21bd8162cdf28a4d88c098db7ddc55..54ce8c1d47737f4da742925f99e3d1cd73160ffb 100644 --- a/tensorflow/contrib/verbs/verbs_server_lib.h +++ b/tensorflow/contrib/verbs/verbs_server_lib.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_VERBS_SERVER_LIB_H_ -#define THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_VERBS_SERVER_LIB_H_ +#ifndef TENSORFLOW_CONTRIB_VERBS_VERBS_SERVER_LIB_H_ +#define TENSORFLOW_CONTRIB_VERBS_VERBS_SERVER_LIB_H_ #ifdef TENSORFLOW_USE_VERBS @@ -63,4 +63,4 @@ class VerbsServer : public GrpcServer { } // namespace tensorflow #endif // TENSORFLOW_USE_VERBS -#endif // THIRD_PARTY_TENSORFLOW_CONTRIB_VERBS_VERBS_SERVER_LIB_H_ +#endif // TENSORFLOW_CONTRIB_VERBS_VERBS_SERVER_LIB_H_ diff --git a/tensorflow/contrib/verbs/verbs_service.proto b/tensorflow/contrib/verbs/verbs_service.proto index 0df1fed4b9de81d7d99be3de9fba4be8b88ad404..abdae1d84f74b076bb5f457d0cf6f74bf07d75b4 100644 --- a/tensorflow/contrib/verbs/verbs_service.proto +++ b/tensorflow/contrib/verbs/verbs_service.proto @@ -50,6 +50,12 @@ message GetRemoteAddressResponse { repeated MemoryRegion mr = 3; } +message ErrorStatusProto { + int32 error_code = 1; + string error_message = 2; + string error_details = 3; +} + //////////////////////////////////////////////////////////////////////////////// // // VerbsService diff --git a/tensorflow/contrib/verbs/verbs_with_0_copies.png b/tensorflow/contrib/verbs/verbs_with_0_copies.png new file mode 100644 index 0000000000000000000000000000000000000000..0641e2fd50da3738e3b8113f4324156063bf52a2 Binary files /dev/null and b/tensorflow/contrib/verbs/verbs_with_0_copies.png differ diff --git a/tensorflow/contrib/verbs/verbs_with_0_copies.xml b/tensorflow/contrib/verbs/verbs_with_0_copies.xml new file mode 100644 index 0000000000000000000000000000000000000000..16130a961ba5185c415463ccb9636d010fe05e68 --- /dev/null +++ b/tensorflow/contrib/verbs/verbs_with_0_copies.xml @@ -0,0 +1 @@ +7Vxtc9o4EP41zKQfmsGW3/hIgPQ60/RyIZ1rPzHClsFXY1FZEOivP8mW8ZsAB2yXtHQ6jb2SJXl3n0e7K6cdMFhsPhC4nD9gB/kdtetsOmDYUVVFUw32g0u2scRUu7FgRjxHdEoFY+8nEsKk28pzUJjrSDH2qbfMC20cBMimORkkBL/ku7nYz8+6hDNUEoxt6Jel/3oOnQup0u2mDX8hbzYXU1u6aJhC+/uM4FUg5uuowI3+xM0LmIwl+odz6OCXjAiMOmBAMKbx1WIzQD7XbaK2+Ln7Pa27dRMU0CoP6CB+Yg39FUqWHC2MbhNlRK+D+APdDrh7mXsUjZfQ5q0vzPxMNqcLn90p7NL1fH+AfUzYfYAD1ulOzIAIRZu9y1R2L8+cCuEFomTLumx2mo8fEf5kiduX1DhWIptn7GIkQigcYrYbOlUKuxB6kevIkqjI8NkMd463zqnK+LHihrtjL0rfQ9+bBR3QZz185NK0lV3NxM9olHAJg0Q2ppDQm3dJE1tatjUjjqbOS+tfTSKbEskK2lqYcsua+r6PbUgRJwIUhJiEt03OqfI5xyhwbp5Hn8d/P02eRv98GY2f37VhAam2c7PVBVDGTg5Elmtzu1OCv6NMi2FbaOru5iuBVQLp/fgFefwqRhnAalcC4B3jngPghGxrR/ATstfPkTs+IAqHkMKbC/GQ2oD3ZenEsGNah+/ZNeTbLrTnqHkAPiHYLuxlHAjKVCBjg8q0+XsBGahtAlkxjkcryGGRnLjFhM7xDAfQH6XSu7yWMxr9D1G6FcEoXFHMROkInzBein579RjiFbGTEFIsjW3oM5R002IZX+NBbRPkQ+qt89HoWZpTG6LAiCQGBMUgJShc4iBshRvs9SfGDX4/3AZ2s7QbUcAAL7dRGsKvH79wyCPisWd/8hf/6EZv/2PlEeTsffs3B3dT+aX7tv4r0M20RbZfszff+GC3Or/dePRr7u6bmKgaJ2hlTkhS5fo4QTz6iD22lJ0pe728KU2jYKF4UeKp1Eh9QuA2023JO4QH5jELO4RZyECP9E9SttRH4hWkHrPTSTXm00rMd++RkD57Cw7cjjngf1UDLjjggmm4LPJGjN4kwhvMYTBDDmMccF8V53O8mK7C4xh3GHvY1MOMjYbMbzjCLgP3LW/zvePbfDiHS37p+mjT5xWfCKyOuBzaPgxDzz5Ioa5lI9vOIr5bRnxJzVNL1/TKiLckQYBaEfAZXesSVSeyM3kBFCI6tcgL8euUeKE0kNbt3jK/MUxLUSxD10wjP65SjW9OgHjiHm35y5k+Id0FuhflFIbSu1WtHlAVy1KApqs5U2pFU1Z1kcPDgoo70ikeUi5zfkNhyUl4OJh3gbypRUFTUuMUMeTQZoZHTH7Hudbj4aloWHiOE8Unsi0gH7M0wd+gzN+axH3UOqK28ob7Gf++qraK8U6bqpblw00VQqJM75GFyfI0r61yMM/+5MFadgVPwwc+xAvxorz0Jq7SdaITI8qM/e61S3/zqZsh8cvmQjigH9+S28zlRMK2i+2q7tWqKdmrW8rYrKIFtWr74/6M7Zwd1CwZdFcaztDBm4eJ1mqFA1Q4fn0jmU6yn+WQYlZESjtRrVpIdTTzxDi2OJBe9IWaaimleZR6ayOgQj29EfeTlNbOdT+j7H7gstzvcPZjgkaSqtIXEPUlVaC8JdR9rDqIo7Xf6VRVUlqMI2uCN9soQOXnDPcOyh4veJWOF0oDyyLj6PTkY7BmYGMXDs+p+IGu7/NPl45Fxa9S0ZEz1aHkdJf7aeBE77rAayReGoX0tEzjzQfxxePWdoP4JN6UAHyuVIKAyNFlCEeMSEnGUHzEPXY6vVbAXMX2ghkT6Ondc5QevFf3GZ05HnH9aHebe46DgibKBoqp5yzbKxt299Fb1rDFHOAku6pN2ZV/J/EnW9U0jlq115RRZamEYO0iL7o4CiDsHWmldgSu2+1y8ihBdvjQnzyMxuP+h9Ek/1Vcxt7xyCUWnjbgBRdWBwRWnqoVyTeq0uiyjkKgVq65Nmf8h9FzfzLss3+eRuPHvz+PR1cHkDgA0Np0AFm9rXH0XwnggP21Vglg/0lALfYv1s+vFpdY3NDbtHhT2XfNSXUi+LhYxP2ruCF3Qv47M8VRBQO9yvsOJYiyVLLm9773kO+E+VfPLpBulyzPHaSp7sRjXrqJRQFciMaQouWE2V70XGCKJtBxiBB8R9v4in+mPYk/0z6SGZ9w6nUMuTwvNqaGbnTKNUDXVaMa4KVh2MhjWJV86QRkGbZeB4ab/tWihjAsg1m31PvPBbpUPweBfoXtebAFkmCrOdjKvk+8wvYPhO3r9ucrsk9Att7mhpyMcUV2ceqC818+viueVF1xtwd3hqR2XRfu2G36XxzEJ8/p/yMBRv8D \ No newline at end of file diff --git a/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.jpg b/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.jpg new file mode 100644 index 0000000000000000000000000000000000000000..8bc69b889d68f0a6a8faa64f08eab22637646842 Binary files /dev/null and b/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.jpg differ diff --git a/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.xml b/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.xml new file mode 100644 index 0000000000000000000000000000000000000000..484e7c78ae863b543dee360969e2ef1c8ae6114d --- /dev/null +++ b/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.xml @@ -0,0 +1 @@ +7Vxbc5s4FP41nuk+pMMd8ujYTrYzTZqNk9nmyaOAbNhgRIWc2P31K4G4yzZ1AHtaZzJjOBLS8TnnOzeRDNTRcn2DQejeIgf6A0Vy1gN1PFAUWVMM+sEom4RiGpcJYYE9h0/KCVPvJ+REiVNXngOj0kSCkE+8sEy0URBAm5RoAGP0Xp42R3551xAsYI0wtYFfp/7rOcTlVFmS8oG/obdw+daWzgdegP26wGgV8P0GijqPf5LhJUjX4vMjFzjovUBSJwN1hBEiydVyPYI+k20qtuS56y2jGd8YBqTJA1bywBvwVzDl2PDpo1eO98b4IxsuE+PHijF1ReCaXADfWwQDdUhn+HBO8lF6teCf8SpRCIKUNiUAk09/pUOUqeJogRxvXaa2z01Ke8ECDnpkDCxDehG8ROzjgs4c+j6yAYGPMIgQjkoC64WBKQycT4+Tu+m3h9nD5J+nyfSxax62a6K0m1LaRYlxBpklS3T43fUInIbAZqPv1C9RmkuWPr2Ts6ffIKZMbQWLnEGQujaIlpDgDZ3CH7A4aLlTkw1+/567CCX1EG7BO2RuA3C3tMiWzqFJLzg6xUhNPUbrUH2A9ltia7eQgDEgoHOTa6juNnZjBj2GoE9M1UHc/X4xZq+erq8nDLPT+29308nWTU8LRqrSJ4xkQwCjikCgQ5MBfoswcdECBcCf5NSrssgK4vkPErLh+QxYEURJ+QpfEQpLYmQb7RYi5QutsJ3O4rzSQLqA6TRNLGwMfUC8t/L6H5Kc0pEDivHiON/wU+hQyDzAKERBBLsHDfN8XylM/WG0Cezu9/syj/XyY+VhZjvDPgP7gKGsJRL7LiMUbm7unx7R6F6UXT2dUp7XqSCmEHuUsZ/wqN/45HMnQztq8qQfw8lT0eDN9+LNM1vss85u1x75Xrp75hsdGBq0emhIqvAPhAb+6D3y6M6ZKi+VsipNo6KhhAf+VK6kIcZgU5gWsglR831Us1LL7plvWLvnW9rO+fQi4Ti3sEyGzQKmVguY1x6OyKO3hMyZmCP2W/MyBbeQYDdNy0cuCBbQoX5GvW6KchctX1bRfoQ7NCTZxEPU2YypWTFEdoH6nnO9y/25XuSCkF3Ofbgess5RDFWHX45tH0SRZ5eFNfd8f4R8hOMl0gZPAe9SHe+HodoS5HuKWOIFieoCgaa0D2K/mrwrVewn3NewX1tI1aXPpiwZpiXLlqFrplFeV27mUw6AZWoEfVlFi/5cGhxR9bpx+VmxLlVFtixZ1XSlpDCtqrCmhrB7WbVhbDnEDtSaHTzDqGYKLA8rKzoiGL3CVNUBCmBF+5zEk7exTfUMKf2KuVKPlRt8YOk5TpxpiJxzOfvowherdV+sCcxHaSP/qofCO/T7itqSjihqUYOjq+KKFUD3KBSW7H2VQBfbUqgiAw/jW7bCO6bap5+fkr7cID5BIlSrv8r4iVVThsDAusurVH1/BO2zvOI1VJZwHRx0FVMQdLsposyqBrVmgW57EfWRUGjWFLqlIXdidq/12kVQ6xlDTSKnXU+kAaZk4KZY5P1klXKloNDMA/PI6kJ6VeMtdSVq+8jtdg3UBgcUnRiZoEl1oJEZdSNTj2pku2sMU+2kdEnbSR2ULmrdX7d9FjxK8qLf6ShY0Fo78FSmvyOWETtiubl/aqSHftgaw0h45LGbq/hJWqzte+TEEm3rmHl2muwIYO7KjYDA62ERziF1p7ggdbbiFqEfXpfTUsr2ggUl6PndY5zBXyjbNIgoY3M/jmQuLdth0E2JXlLsZUO9VrP0g9QqakC2olb2FsifrNRqedCrVkXFAQ9vVSc3R3EaYWfpWK5ImphJEuOxBtnx7XB2O5lOhzeTWfntvILCk5VrLvWlAzM4sZ5b1mNLD5gtgfJFOWYbTTet3t/sTvnZa15n5W9Tvqr1qXxRO6xz5Sfv+J21L9C+1iv0t/fbW9F+loLHK1T71tloVe1na8hydr1Pa+iqMm+54E4JX5bLZH4TE2UGyv6Spboq902/ZH27akDRAUzL3/vag74Tlb96kUGyCeFAGfHOAIzIzKNWuk5IAVjywYjAcEZ1z2cuEYEz4DiYE17hJrmidgRmDiBgb/F7wOHTPuRSzTnGi6EbA1EXULHtE8SwXMawInhvSBVl8nobGO76j6I6wrAIZlJt9p8LdKF8dgL9DNuPwVYVJGLdwVb0tt8Ztn8gbD8Un7e+kXvG/i9hX+8zZKdrnLFf3boCj9P3AA2PAs+424I7Q9D0bgt39Db/1wTJuXX+/x/Uyf8= \ No newline at end of file diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index ae38025942b711c4ff9f635066f4429d965a1984..30ac270109dbbb77cd6a400d1feaa1ac116456c1 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -79,6 +79,7 @@ load( "if_linux_x86_64", "if_mobile", "if_not_mobile", + "if_windows", "if_not_windows", "tf_copts", "tf_cc_test", @@ -135,6 +136,8 @@ load( "tf_nano_proto_library", "tf_protos_all", "tf_protos_all_impl", + "tf_protos_grappler", + "tf_protos_grappler_impl", ) load( "//tensorflow/core:platform/default/build_config_root.bzl", @@ -190,6 +193,7 @@ CORE_PROTO_SRCS = [ "protobuf/rewriter_config.proto", "protobuf/tensor_bundle.proto", "protobuf/saver.proto", + "util/event.proto", "util/memmapped_file_system.proto", "util/saved_tensor_slice.proto", ] @@ -202,11 +206,12 @@ CORE_PROTO_SRCS = [ ADDITIONAL_CORE_PROTO_SRCS = [ "example/example_parser_configuration.proto", "protobuf/control_flow.proto", + # TODO(ebrevdo): Re-enable once CriticalSection is in core. + # "protobuf/critical_section.proto", "protobuf/meta_graph.proto", "protobuf/named_tensor.proto", "protobuf/saved_model.proto", "protobuf/tensorflow_server.proto", - "util/event.proto", "util/test_log.proto", ] @@ -274,7 +279,9 @@ cc_library( "platform/platform.h", "platform/protobuf.h", "platform/types.h", - ] + glob(tf_additional_proto_hdrs()) + glob(tf_env_time_hdrs()), + "platform/windows/cpu_info.h", + "lib/bfloat16/bfloat16.h", + ] + tf_additional_proto_hdrs() + glob(tf_env_time_hdrs()), copts = tf_copts(), deps = tf_lib_proto_parsing_deps(), ) @@ -286,6 +293,7 @@ cc_library( cc_library( name = "lib", hdrs = [ + "lib/bfloat16/bfloat16.h", "lib/core/arena.h", "lib/core/bitmap.h", "lib/core/bits.h", @@ -362,6 +370,33 @@ cc_library( ], ) +cc_library( + name = "abi", + srcs = ["platform/abi.cc"], + hdrs = ["platform/abi.h"], +) + +cc_library( + name = "session_message", + srcs = ["util/session_message.cc"], + hdrs = ["util/session_message.h"], + deps = [ + ":framework", + ":lib", + ":protos_all_cc", + ], +) + +cc_library( + name = "stacktrace_handler", + srcs = ["platform/stacktrace_handler.cc"], + hdrs = ["platform/stacktrace_handler.h"], + deps = [ + ":lib", + ":lib_platform", + ], +) + # Test support library needed for all tests # This is currently public, but may be made internal in the # future. Try to avoid depending on it. @@ -408,6 +443,8 @@ tf_cuda_library( "framework/cancellation.h", "framework/common_shape_fns.h", "framework/control_flow.h", # TODO(josh11b): Make internal? + "framework/dataset.h", + "framework/dataset_stateful_op_whitelist.h", "framework/device_base.h", "framework/function.h", "framework/graph_def_util.h", @@ -428,6 +465,7 @@ tf_cuda_library( "framework/reader_interface.h", "framework/reader_op_kernel.h", "framework/register_types.h", + "framework/register_types_traits.h", "framework/resource_mgr.h", "framework/resource_op_kernel.h", "framework/selective_registration.h", @@ -511,14 +549,6 @@ cc_library( ], ) -tf_proto_library_cc( - name = "op_gen_overrides_proto", - srcs = ["framework/op_gen_overrides.proto"], - cc_api_version = 2, - protodeps = tf_additional_all_protos(), - visibility = ["//visibility:public"], -) - cc_library( name = "op_gen_lib", srcs = ["framework/op_gen_lib.cc"], @@ -527,7 +557,6 @@ cc_library( deps = [ ":lib", ":lib_internal", - ":op_gen_overrides_proto_cc", ":protos_all_cc", ], ) @@ -549,6 +578,7 @@ cc_library( "framework/numeric_types.h", "framework/tensor_types.h", "framework/type_traits.h", + "lib/bfloat16/bfloat16.h", "platform/default/dynamic_annotations.h", "platform/default/integral_types.h", "platform/default/logging.h", @@ -562,7 +592,7 @@ cc_library( "platform/prefetch.h", "platform/thread_annotations.h", "platform/types.h", - ], + ] + if_windows(["platform/windows/integral_types.h"]), visibility = ["//visibility:public"], deps = [ @@ -577,6 +607,7 @@ cc_library( tf_gen_op_libs( is_external = False, op_lib_names = [ + "batch_ops", "bitwise_ops", "candidate_sampling_ops", "checkpoint_ops", @@ -589,8 +620,10 @@ tf_gen_op_libs( "image_ops", "io_ops", "linalg_ops", + "list_ops", "lookup_ops", "logging_ops", + "manip_ops", "math_ops", "nn_ops", "no_op", @@ -657,6 +690,7 @@ cc_library( deps = [ ":array_ops_op_lib", ":audio_ops_op_lib", + ":batch_ops_op_lib", ":bitwise_ops_op_lib", ":candidate_sampling_ops_op_lib", ":checkpoint_ops_op_lib", @@ -669,8 +703,10 @@ cc_library( ":image_ops_op_lib", ":io_ops_op_lib", ":linalg_ops_op_lib", + ":list_ops_op_lib", ":logging_ops_op_lib", ":lookup_ops_op_lib", + ":manip_ops_op_lib", ":math_ops_op_lib", ":nn_ops_op_lib", ":no_op_op_lib", @@ -762,6 +798,7 @@ tf_cuda_library( "graph/graph.h", "graph/graph_constructor.h", "graph/graph_def_builder.h", + "graph/graph_def_builder_util.h", "graph/node_builder.h", "graph/validate.h", "graph/while_context.h", @@ -791,6 +828,7 @@ cc_library( deps = [ "//tensorflow/core/kernels:array", "//tensorflow/core/kernels:audio", + "//tensorflow/core/kernels:batch_kernels", "//tensorflow/core/kernels:bincount_op", "//tensorflow/core/kernels:candidate_sampler_ops", "//tensorflow/core/kernels:checkpoint_ops", @@ -800,12 +838,15 @@ cc_library( "//tensorflow/core/kernels:dataset_ops", "//tensorflow/core/kernels:fake_quant_ops", "//tensorflow/core/kernels:function_ops", + "//tensorflow/core/kernels:functional_ops", "//tensorflow/core/kernels:histogram_op", "//tensorflow/core/kernels:image", "//tensorflow/core/kernels:io", "//tensorflow/core/kernels:linalg", + "//tensorflow/core/kernels:list_kernels", "//tensorflow/core/kernels:lookup", "//tensorflow/core/kernels:logging", + "//tensorflow/core/kernels:manip", "//tensorflow/core/kernels:math", "//tensorflow/core/kernels:multinomial_op", "//tensorflow/core/kernels:nn", @@ -842,6 +883,7 @@ cc_library( "//tensorflow/core/kernels:mkl_pooling_ops", "//tensorflow/core/kernels:mkl_relu_op", "//tensorflow/core/kernels:mkl_reshape_op", + "//tensorflow/core/kernels:mkl_softmax_op", "//tensorflow/core/kernels:mkl_tfconv_op", "//tensorflow/core/kernels:mkl_aggregate_ops", ]), @@ -1039,18 +1081,31 @@ cc_library( name = "ios_tensorflow_lib", srcs = if_ios([ ":android_op_registrations_and_gradients", - "//tensorflow/core:android_srcs", "//tensorflow/core/kernels:android_core_ops", "//tensorflow/core/kernels:android_extended_ops", ]), copts = tf_copts() + ["-Os"] + ["-std=c++11"], visibility = ["//visibility:public"], deps = [ + ":ios_tensorflow_lib_lite", ":protos_all_cc_impl", "//third_party/eigen3", "//third_party/fft2d:fft2d_headers", - "@fft2d//:fft2d", - "@gemmlowp//:gemmlowp", + "@fft2d", + "@gemmlowp", + "@protobuf_archive//:protobuf", + ], + alwayslink = 1, +) + +cc_library( + name = "ios_tensorflow_lib_lite", + srcs = if_ios(["//tensorflow/core:android_srcs"]), + copts = tf_copts() + ["-Os"] + ["-std=c++11"], + visibility = ["//visibility:public"], + deps = [ + ":protos_all_cc_impl", + "//third_party/eigen3", "@nsync//:nsync_cpp", "@protobuf_archive//:protobuf", ], @@ -1114,6 +1169,7 @@ cc_library( deps = [ ":protos_all_cc_impl", "//third_party/eigen3", + "@nsync//:nsync_cpp", "@protobuf_archive//:protobuf", ], alwayslink = 1, @@ -1279,6 +1335,13 @@ tf_pyclif_proto_library( visibility = ["//visibility:public"], ) +tf_pyclif_proto_library( + name = "framework/kernel_def_pyclif", + proto_lib = ":protos_all_cc", + proto_srcfile = "framework/kernel_def.proto", + visibility = ["//visibility:public"], +) + tf_pyclif_proto_library( name = "framework/node_def_pyclif", proto_lib = ":protos_all_cc", @@ -1286,6 +1349,13 @@ tf_pyclif_proto_library( visibility = ["//visibility:public"], ) +tf_pyclif_proto_library( + name = "framework/function_pyclif", + proto_lib = ":protos_all_cc", + proto_srcfile = "framework/function.proto", + visibility = ["//visibility:public"], +) + tf_pyclif_proto_library( name = "framework/graph_pyclif", proto_lib = ":protos_all_cc", @@ -1307,6 +1377,13 @@ tf_pyclif_proto_library( visibility = ["//visibility:public"], ) +tf_pyclif_proto_library( + name = "protobuf/device_properties_pyclif", + proto_lib = ":protos_all_cc", + proto_srcfile = "protobuf/device_properties.proto", + visibility = ["//visibility:public"], +) + # ----------------------------------------------------------------------------- # Internal targets @@ -1497,7 +1574,7 @@ cc_library( "@snappy", "@zlib_archive//:zlib", "@protobuf_archive//:protobuf", - ] + tf_protos_all_impl(), + ] + tf_protos_all_impl() + tf_protos_grappler_impl(), ) # File compiled with extra flags to get cpu-specific acceleration. @@ -1562,6 +1639,7 @@ cc_library( "platform/jpeg.h", ]), hdrs = [ + "lib/bfloat16/bfloat16.h", "lib/core/stringpiece.h", "lib/jpeg/jpeg_handle.h", "lib/jpeg/jpeg_mem.h", @@ -1589,6 +1667,7 @@ cc_library( "platform/gif.h", ]), hdrs = [ + "lib/bfloat16/bfloat16.h", "lib/core/stringpiece.h", "lib/gif/gif_io.h", "lib/gtl/cleanup.h", @@ -1616,6 +1695,7 @@ cc_library( "platform/png.h", ]), hdrs = [ + "lib/bfloat16/bfloat16.h", "lib/core/casts.h", "lib/core/stringpiece.h", "lib/png/png_io.h", @@ -1663,6 +1743,9 @@ FRAMEWORK_INTERNAL_PRIVATE_HEADERS = [ "platform/variant_coding.h", "graph/edgeset.h", "graph/graph.h", + "graph/graph_def_builder.h", + "graph/node_builder.h", + "graph/tensor_id.h", ] + glob( [ "example/**/*.h", @@ -1680,6 +1763,7 @@ FRAMEWORK_INTERNAL_PRIVATE_HEADERS = [ "framework/reader_base.*", "util/memmapped_file_system.*", "util/memmapped_file_system_writer.*", + "util/session_message.*", "util/version_info.cc", ], ) + select({ @@ -1749,6 +1833,9 @@ tf_cuda_library( ] + [ "graph/edgeset.cc", "graph/graph.cc", + "graph/graph_def_builder.cc", + "graph/node_builder.cc", + "graph/tensor_id.cc", "graph/while_context.h", "graph/while_context.cc", ], @@ -1763,6 +1850,7 @@ tf_cuda_library( "framework/resource_handle.cc", "util/memmapped_file_system.*", "util/memmapped_file_system_writer.*", + "util/session_message.cc", "util/version_info.cc", ], ) + select({ @@ -1848,6 +1936,13 @@ cc_library( ], ) +tf_cuda_library( + name = "cuda_device_functions", + hdrs = ["util/cuda_device_functions.h"], + visibility = ["//visibility:public"], + deps = [":framework_lite"], +) + # TODO(josh11b): Is this needed, or can we just use ":protos_all_cc"? cc_library( name = "protos_cc", @@ -1870,6 +1965,7 @@ GRAPH_HDRS = [ "graph/graph.h", "graph/graph_constructor.h", # NOTE(mrry): Don't include the .cc since it depends on common_runtime. "graph/graph_def_builder.h", + "graph/graph_def_builder_util.h", "graph/graph_partition.h", "graph/mkl_layout_pass.h", "graph/mkl_tfconversion_pass.h", @@ -1890,12 +1986,9 @@ tf_cuda_library( "graph/colors.cc", "graph/control_flow.cc", "graph/costmodel.cc", - "graph/graph_def_builder.cc", "graph/graph_partition.cc", - "graph/node_builder.cc", "graph/optimizer_cse.cc", "graph/subgraph.cc", - "graph/tensor_id.cc", "graph/validate.cc", ], hdrs = GRAPH_HDRS, @@ -1924,6 +2017,7 @@ tf_cuda_library( "common_runtime/shape_refiner.h", "framework/versions.h", "graph/graph_constructor.cc", # Depends on common_runtime. + "graph/graph_def_builder_util.cc", # Depends on common_runtime. "public/session.h", "public/session_options.h", "public/version.h", @@ -2059,7 +2153,7 @@ tf_cuda_library( ":core_cpu_base", ":proto_text", "//tensorflow/core/grappler:grappler_item", - ] + if_static([":core_cpu_impl"]) + tf_protos_all(), + ] + if_static([":core_cpu_impl"]) + tf_protos_all() + tf_protos_grappler(), ) tf_cuda_library( @@ -2168,12 +2262,25 @@ tf_cuda_library( ] + tf_additional_device_tracer_deps(), ) +cc_library( + name = "gpu_id", + srcs = ["common_runtime/gpu/gpu_id_manager.cc"], + hdrs = [ + "common_runtime/gpu/gpu_id.h", + "common_runtime/gpu/gpu_id_manager.h", + ], + deps = [ + ":lib", + ], +) + GPU_RUNTIME_HEADERS = [ "common_runtime/gpu/gpu_bfc_allocator.h", "common_runtime/gpu/gpu_cudamalloc_allocator.h", "common_runtime/gpu/gpu_debug_allocator.h", "common_runtime/gpu/gpu_device.h", "common_runtime/gpu/gpu_id.h", + "common_runtime/gpu/gpu_id_manager.h", "common_runtime/gpu/gpu_id_utils.h", "common_runtime/gpu/gpu_init.h", "common_runtime/gpu/gpu_managed_allocator.h", @@ -2192,7 +2299,6 @@ tf_cuda_library( "common_runtime/gpu/gpu_debug_allocator.cc", "common_runtime/gpu/gpu_device.cc", "common_runtime/gpu/gpu_device_factory.cc", - "common_runtime/gpu/gpu_id_utils.cc", "common_runtime/gpu/gpu_managed_allocator.cc", "common_runtime/gpu/gpu_stream_util.cc", "common_runtime/gpu/gpu_util.cc", @@ -2207,6 +2313,7 @@ tf_cuda_library( ":core_cpu_lib", ":framework", ":framework_internal", + ":gpu_id", ":gpu_init_impl", ":gpu_lib", ":graph", @@ -2296,7 +2403,7 @@ cc_library( ":lib_internal", ":proto_text", "//third_party/eigen3", - "@local_config_sycl//sycl:sycl", + "@local_config_sycl//sycl", ], alwayslink = 0, ) @@ -2428,6 +2535,7 @@ tf_cc_tests( "platform/net_test.cc", "platform/port_test.cc", "platform/profile_utils/cpu_utils_test.cc", + "platform/stacktrace_handler_test.cc", "platform/subprocess_test.cc", ], deps = [ @@ -2783,6 +2891,7 @@ tf_cc_test_mkl( "//tensorflow/core/kernels:mkl_pooling_ops", "//tensorflow/core/kernels:mkl_relu_op", "//tensorflow/core/kernels:mkl_reshape_op", + "//tensorflow/core/kernels:mkl_softmax_op", "//tensorflow/core/kernels:mkl_tfconv_op", ]), ) @@ -2794,6 +2903,7 @@ tf_cc_tests_gpu( linkstatic = tf_kernel_tests_linkstatic(), deps = [ ":gpu_headers_lib", + ":gpu_id", ":gpu_runtime", ":test", ], @@ -2805,7 +2915,7 @@ tf_cc_tests_gpu( srcs = glob(["user_ops/**/*_test.cc"]) + [ "common_runtime/gpu/gpu_bfc_allocator_test.cc", "common_runtime/gpu/gpu_device_test.cc", - "common_runtime/gpu/gpu_id_utils_test.cc", + "common_runtime/gpu/gpu_id_manager_test.cc", "common_runtime/gpu/gpu_event_mgr_test.cc", "common_runtime/gpu/pool_allocator_test.cc", ], @@ -2817,6 +2927,7 @@ tf_cc_tests_gpu( ":direct_session", ":framework", ":framework_internal", + ":gpu_id", ":gpu_runtime", ":lib", ":lib_internal", @@ -3212,6 +3323,7 @@ tf_cc_test_gpu( ":direct_session", ":framework", ":framework_internal", + ":gpu_id", ":gpu_runtime", ":lib", ":lib_internal", @@ -3421,37 +3533,6 @@ tf_cc_test( ], ) -filegroup( - name = "base_api_def", - srcs = glob(["api_def/base_api/*"]), -) - -filegroup( - name = "python_api_def", - srcs = glob(["api_def/python_api/*"]), -) - -tf_cc_test( - name = "api_test", - srcs = ["api_def/api_test.cc"], - data = [ - ":base_api_def", - "//tensorflow/cc:ops/op_gen_overrides.pbtxt", - ], - deps = [ - ":framework", - ":framework_internal", - ":lib", - ":lib_internal", - ":lib_test_internal", - ":op_gen_lib", - ":op_gen_overrides_proto_cc", - ":ops", - ":protos_all_cc", - ":test", - ], -) - tf_cc_test_gpu( name = "device_tracer_test", size = "small", diff --git a/tensorflow/core/api_def/BUILD b/tensorflow/core/api_def/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..58dbac4e8edac7079d315fbfcdafbd136793df0b --- /dev/null +++ b/tensorflow/core/api_def/BUILD @@ -0,0 +1,114 @@ +# Description: +# Provides ApiDef access and ApiDef validation for TensorFlow. +# +# The following targets can be used to access ApiDefs: +# :base_api_def +# :python_api_def + +package( + default_visibility = ["//visibility:private"], +) + +licenses(["notice"]) # Apache 2.0 + +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_binary", + "tf_cc_test", +) + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + +filegroup( + name = "base_api_def", + srcs = glob(["base_api/*"]), + visibility = ["//tensorflow:internal"], +) + +filegroup( + name = "python_api_def", + srcs = glob(["python_api/*"]), + visibility = ["//tensorflow:internal"], +) + +cc_library( + name = "excluded_ops_lib", + srcs = ["excluded_ops.cc"], + hdrs = ["excluded_ops.h"], +) + +cc_library( + name = "update_api_def_lib", + srcs = ["update_api_def.cc"], + hdrs = ["update_api_def.h"], + deps = [ + ":excluded_ops_lib", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:op_gen_lib", + "//tensorflow/core:ops", + "//tensorflow/core:protos_all_cc", + ], +) + +tf_cc_test( + name = "update_api_def_test", + srcs = ["update_api_def_test.cc"], + deps = [ + ":update_api_def_lib", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +tf_cc_binary( + name = "update_api_def", + srcs = [ + "update_api_def_main.cc", + ], + data = [ + ":base_api_def", + ], + deps = [ + ":update_api_def_lib", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "api_test", + srcs = ["api_test.cc"], + data = [ + ":base_api_def", + ":python_api_def", + ], + deps = [ + ":excluded_ops_lib", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:lib_test_internal", + "//tensorflow/core:op_gen_lib", + "//tensorflow/core:ops", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) diff --git a/tensorflow/core/api_def/api_test.cc b/tensorflow/core/api_def/api_test.cc index 2cdc14843f61a2585b61e214527e0a0b5bdea446..477a0b670e49f8aa4ee8c250d4957886eb865ed5 100644 --- a/tensorflow/core/api_def/api_test.cc +++ b/tensorflow/core/api_def/api_test.cc @@ -13,9 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// Test that verifies tensorflow/core/api_def/base_api/api_def*.pbtxt files -// are correct. If api_def*.pbtxt do not match expected contents, run -// tensorflow/core/api_def/base_api/update_api_def.sh script to update them. +// Test that validates tensorflow/core/api_def/base_api/api_def*.pbtxt files. #include #include @@ -23,12 +21,11 @@ limitations under the License. #include #include +#include "tensorflow/core/api_def/excluded_ops.h" #include "tensorflow/core/framework/api_def.pb.h" -#include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_def.pb.h" #include "tensorflow/core/framework/op_gen_lib.h" -#include "tensorflow/core/framework/op_gen_overrides.pb.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/lib/io/path.h" @@ -44,309 +41,235 @@ namespace tensorflow { namespace { constexpr char kDefaultApiDefDir[] = "tensorflow/core/api_def/base_api"; -constexpr char kOverridesFilePath[] = - "tensorflow/cc/ops/op_gen_overrides.pbtxt"; -constexpr char kApiDefFileFormat[] = "api_def_%s.pbtxt"; +constexpr char kPythonApiDefDir[] = + "tensorflow/core/api_def/python_api"; constexpr char kApiDefFilePattern[] = "api_def_*.pbtxt"; -void FillBaseApiDef(ApiDef* api_def, const OpDef& op) { - api_def->set_graph_op_name(op.name()); - // Add arg docs - for (auto& input_arg : op.input_arg()) { - if (!input_arg.description().empty()) { - auto* api_def_in_arg = api_def->add_in_arg(); - api_def_in_arg->set_name(input_arg.name()); - api_def_in_arg->set_description(input_arg.description()); - } - } - for (auto& output_arg : op.output_arg()) { - if (!output_arg.description().empty()) { - auto* api_def_out_arg = api_def->add_out_arg(); - api_def_out_arg->set_name(output_arg.name()); - api_def_out_arg->set_description(output_arg.description()); - } - } - // Add attr docs - for (auto& attr : op.attr()) { - if (!attr.description().empty()) { - auto* api_def_attr = api_def->add_attr(); - api_def_attr->set_name(attr.name()); - api_def_attr->set_description(attr.description()); - } +// Reads golden ApiDef files and returns a map from file name to ApiDef file +// contents. +void GetGoldenApiDefs(Env* env, const string& api_files_dir, + std::unordered_map* name_to_api_def) { + std::vector matching_paths; + TF_CHECK_OK(env->GetMatchingPaths( + io::JoinPath(api_files_dir, kApiDefFilePattern), &matching_paths)); + + for (auto& file_path : matching_paths) { + string file_contents; + TF_CHECK_OK(ReadFileToString(env, file_path, &file_contents)); + file_contents = PBTxtFromMultiline(file_contents); + + ApiDefs api_defs; + CHECK(tensorflow::protobuf::TextFormat::ParseFromString(file_contents, + &api_defs)) + << "Failed to load " << file_path; + CHECK_EQ(api_defs.op_size(), 1); + (*name_to_api_def)[api_defs.op(0).graph_op_name()] = api_defs.op(0); } - // Add docs - api_def->set_summary(op.summary()); - api_def->set_description(op.description()); } -// Checks if arg1 should be before arg2 according to ordering in args. -bool CheckArgBefore(const ApiDef::Arg* arg1, const ApiDef::Arg* arg2, - const protobuf::RepeatedPtrField& args) { - for (auto& arg : args) { - if (arg.name() == arg2->name()) { - return false; - } else if (arg.name() == arg1->name()) { - return true; - } +void TestAllApiDefsHaveCorrespondingOp( + const OpList& ops, const std::unordered_map& api_defs_map) { + std::unordered_set op_names; + for (const auto& op : ops.op()) { + op_names.insert(op.name()); + } + for (const auto& name_and_api_def : api_defs_map) { + ASSERT_TRUE(op_names.find(name_and_api_def.first) != op_names.end()) + << name_and_api_def.first << " op has ApiDef but missing from ops. " + << "Does api_def_" << name_and_api_def.first << " need to be deleted?"; } - return false; } -// Checks if attr1 should be before attr2 according to ordering in op_def. -bool CheckAttrBefore(const ApiDef::Attr* attr1, const ApiDef::Attr* attr2, - const OpDef& op_def) { - for (auto& attr : op_def.attr()) { - if (attr.name() == attr2->name()) { - return false; - } else if (attr.name() == attr1->name()) { - return true; +void TestAllApiDefInputArgsAreValid( + const OpList& ops, const std::unordered_map& api_defs_map) { + for (const auto& op : ops.op()) { + const auto api_def_iter = api_defs_map.find(op.name()); + if (api_def_iter == api_defs_map.end()) { + continue; + } + const auto& api_def = api_def_iter->second; + for (const auto& api_def_arg : api_def.in_arg()) { + bool found_arg = false; + for (const auto& op_arg : op.input_arg()) { + if (api_def_arg.name() == op_arg.name()) { + found_arg = true; + break; + } + } + ASSERT_TRUE(found_arg) + << "Input argument " << api_def_arg.name() + << " (overwritten in api_def_" << op.name() + << ".pbtxt) is not defined in OpDef for " << op.name(); } } - return false; } -// Applies renames to args. -void ApplyArgOverrides( - protobuf::RepeatedPtrField* args, - const protobuf::RepeatedPtrField& renames, - const protobuf::RepeatedPtrField& op_args, - const string& op_name) { - for (auto& rename : renames) { - // First check if rename is valid. - bool valid = false; - for (const auto& op_arg : op_args) { - if (op_arg.name() == rename.from()) { - valid = true; - } +void TestAllApiDefOutputArgsAreValid( + const OpList& ops, const std::unordered_map& api_defs_map) { + for (const auto& op : ops.op()) { + const auto api_def_iter = api_defs_map.find(op.name()); + if (api_def_iter == api_defs_map.end()) { + continue; } - QCHECK(valid) << rename.from() << " is not a valid argument for " - << op_name; - bool found_arg = false; - // If Arg is already in ApiDef, just update it. - for (int i = 0; i < args->size(); ++i) { - auto* arg = args->Mutable(i); - if (arg->name() == rename.from()) { - arg->set_rename_to(rename.to()); - found_arg = true; - break; + const auto& api_def = api_def_iter->second; + for (const auto& api_def_arg : api_def.out_arg()) { + bool found_arg = false; + for (const auto& op_arg : op.output_arg()) { + if (api_def_arg.name() == op_arg.name()) { + found_arg = true; + break; + } } - } - if (!found_arg) { // not in ApiDef, add a new arg. - auto* new_arg = args->Add(); - new_arg->set_name(rename.from()); - new_arg->set_rename_to(rename.to()); + ASSERT_TRUE(found_arg) + << "Output argument " << api_def_arg.name() + << " (overwritten in api_def_" << op.name() + << ".pbtxt) is not defined in OpDef for " << op.name(); } } - // We don't really need a specific order here right now. - // However, it is clearer if order follows OpDef. - std::sort(args->pointer_begin(), args->pointer_end(), - [&](ApiDef::Arg* arg1, ApiDef::Arg* arg2) { - return CheckArgBefore(arg1, arg2, op_args); - }); } -// Returns existing attribute with the given name if such -// attribute exists. Otherwise, adds a new attribute and returns it. -ApiDef::Attr* FindOrAddAttr(ApiDef* api_def, const string attr_name) { - // If Attr is already in ApiDef, just update it. - for (int i = 0; i < api_def->attr_size(); ++i) { - auto* attr = api_def->mutable_attr(i); - if (attr->name() == attr_name) { - return attr; +void TestAllApiDefAttributeNamesAreValid( + const OpList& ops, const std::unordered_map& api_defs_map) { + for (const auto& op : ops.op()) { + const auto api_def_iter = api_defs_map.find(op.name()); + if (api_def_iter == api_defs_map.end()) { + continue; + } + const auto& api_def = api_def_iter->second; + for (const auto& api_def_attr : api_def.attr()) { + bool found_attr = false; + for (const auto& op_attr : op.attr()) { + if (api_def_attr.name() == op_attr.name()) { + found_attr = true; + } + } + ASSERT_TRUE(found_attr) + << "Attribute " << api_def_attr.name() << " (overwritten in api_def_" + << op.name() << ".pbtxt) is not defined in OpDef for " << op.name(); } } - // Add a new Attr. - auto* new_attr = api_def->add_attr(); - new_attr->set_name(attr_name); - return new_attr; } +} // namespace -// Applies renames and default values to attributes. -void ApplyAttrOverrides(ApiDef* api_def, const OpGenOverride& op_override, - const OpDef& op_def) { - for (auto& attr_rename : op_override.attr_rename()) { - auto* attr = FindOrAddAttr(api_def, attr_rename.from()); - attr->set_rename_to(attr_rename.to()); - } - - for (auto& attr_default : op_override.attr_default()) { - auto* attr = FindOrAddAttr(api_def, attr_default.name()); - *(attr->mutable_default_value()) = attr_default.value(); - } - // We don't really need a specific order here right now. - // However, it is clearer if order follows OpDef. - std::sort(api_def->mutable_attr()->pointer_begin(), - api_def->mutable_attr()->pointer_end(), - [&](ApiDef::Attr* attr1, ApiDef::Attr* attr2) { - return CheckAttrBefore(attr1, attr2, op_def); - }); -} +class BaseApiTest : public ::testing::Test { + protected: + BaseApiTest() { + OpRegistry::Global()->Export(false, &ops_); + const std::vector multi_line_fields = {"description"}; -void ApplyOverridesToApiDef(ApiDef* api_def, const OpDef& op, - const OpGenOverride& op_override) { - // Fill ApiDef with data based on op and op_override. - // Set visibility - if (op_override.skip()) { - api_def->set_visibility(ApiDef_Visibility_SKIP); - } else if (op_override.hide()) { - api_def->set_visibility(ApiDef_Visibility_HIDDEN); + Env* env = Env::Default(); + GetGoldenApiDefs(env, kDefaultApiDefDir, &api_defs_map_); } - // Add endpoints - if (!op_override.rename_to().empty()) { - api_def->add_endpoint()->set_name(op_override.rename_to()); - } else if (!op_override.alias().empty()) { - api_def->add_endpoint()->set_name(op.name()); - } - - for (auto& alias : op_override.alias()) { - auto* endpoint = api_def->add_endpoint(); - endpoint->set_name(alias); + OpList ops_; + std::unordered_map api_defs_map_; +}; + +// Check that all ops have an ApiDef. +TEST_F(BaseApiTest, AllOpsAreInApiDef) { + auto* excluded_ops = GetExcludedOps(); + for (const auto& op : ops_.op()) { + if (excluded_ops->find(op.name()) != excluded_ops->end()) { + continue; + } + ASSERT_TRUE(api_defs_map_.find(op.name()) != api_defs_map_.end()) + << op.name() << " op does not have api_def_*.pbtxt file. " + << "Please add api_def_" << op.name() << ".pbtxt file " + << "under tensorflow/core/api_def/base_api/ directory."; } - - ApplyArgOverrides(api_def->mutable_in_arg(), op_override.input_rename(), - op.input_arg(), api_def->graph_op_name()); - ApplyArgOverrides(api_def->mutable_out_arg(), op_override.output_rename(), - op.output_arg(), api_def->graph_op_name()); - ApplyAttrOverrides(api_def, op_override, op); } -// Get map from ApiDef file path to corresponding ApiDefs proto. -std::unordered_map GenerateApiDef( - const string& api_def_dir, const OpList& ops, - const OpGenOverrides& overrides) { - std::unordered_map name_to_override; - for (const auto& op_override : overrides.op()) { - name_to_override[op_override.name()] = op_override; - } +// Check that ApiDefs have a corresponding op. +TEST_F(BaseApiTest, AllApiDefsHaveCorrespondingOp) { + TestAllApiDefsHaveCorrespondingOp(ops_, api_defs_map_); +} - std::unordered_map api_defs_map; +string GetOpDefHasDocStringError(const string& op_name) { + return strings::Printf( + "OpDef for %s has a doc string. " + "Doc strings must be defined in ApiDef instead of OpDef. " + "Please, add summary and descriptions in api_def_%s" + ".pbtxt file instead", + op_name.c_str(), op_name.c_str()); +} - // These ops are included in OpList only if TF_NEED_GCP - // is set to true. So, we skip them for now so that this test passes - // whether TF_NEED_GCP is set or not. - const std::unordered_set ops_to_exclude = { - "BigQueryReader", "GenerateBigQueryReaderPartitions"}; - for (const auto& op : ops.op()) { - CHECK(!op.name().empty()) - << "Encountered empty op name: %s" << op.DebugString(); - if (ops_to_exclude.find(op.name()) != ops_to_exclude.end()) { - LOG(INFO) << "Skipping " << op.name(); +// Check that OpDef's do not have descriptions and summaries. +// Descriptions and summaries must be in corresponding ApiDefs. +TEST_F(BaseApiTest, OpDefsShouldNotHaveDocs) { + auto* excluded_ops = GetExcludedOps(); + for (const auto& op : ops_.op()) { + if (excluded_ops->find(op.name()) != excluded_ops->end()) { continue; } - string file_path = io::JoinPath(api_def_dir, kApiDefFileFormat); - file_path = strings::Printf(file_path.c_str(), op.name().c_str()); - ApiDef* api_def = api_defs_map[file_path].add_op(); - FillBaseApiDef(api_def, op); - - if (name_to_override.find(op.name()) != name_to_override.end()) { - ApplyOverridesToApiDef(api_def, op, name_to_override[op.name()]); + ASSERT_TRUE(op.summary().empty()) << GetOpDefHasDocStringError(op.name()); + ASSERT_TRUE(op.description().empty()) + << GetOpDefHasDocStringError(op.name()); + for (const auto& arg : op.input_arg()) { + ASSERT_TRUE(arg.description().empty()) + << GetOpDefHasDocStringError(op.name()); + } + for (const auto& arg : op.output_arg()) { + ASSERT_TRUE(arg.description().empty()) + << GetOpDefHasDocStringError(op.name()); + } + for (const auto& attr : op.attr()) { + ASSERT_TRUE(attr.description().empty()) + << GetOpDefHasDocStringError(op.name()); } } - return api_defs_map; } -// Reads golden ApiDef files and returns a map from file name to ApiDef file -// contents. -std::unordered_map GetGoldenApiDefs( - Env* env, const string& api_files_dir) { - std::vector matching_paths; - TF_CHECK_OK(env->GetMatchingPaths( - io::JoinPath(api_files_dir, kApiDefFilePattern), &matching_paths)); - - std::unordered_map file_path_to_api_def; - for (auto& file_path : matching_paths) { - string file_contents; - TF_CHECK_OK(ReadFileToString(env, file_path, &file_contents)); - file_path_to_api_def[file_path] = file_contents; - } - return file_path_to_api_def; +// Checks that input arg names in an ApiDef match input +// arg names in corresponding OpDef. +TEST_F(BaseApiTest, AllApiDefInputArgsAreValid) { + TestAllApiDefInputArgsAreValid(ops_, api_defs_map_); } -void RunApiTest(bool update_api_def, const string& api_files_dir) { - // Read C++ overrides file - OpGenOverrides overrides; - Env* env = Env::Default(); - TF_EXPECT_OK(ReadTextProto(env, kOverridesFilePath, &overrides)); - - // Read all ops - OpList ops; - OpRegistry::Global()->Export(false, &ops); - const std::vector multi_line_fields = {"description"}; +// Checks that output arg names in an ApiDef match output +// arg names in corresponding OpDef. +TEST_F(BaseApiTest, AllApiDefOutputArgsAreValid) { + TestAllApiDefOutputArgsAreValid(ops_, api_defs_map_); +} - // Get expected ApiDefs - const auto new_api_defs_map = GenerateApiDef(api_files_dir, ops, overrides); +// Checks that attribute names in an ApiDef match attribute +// names in corresponding OpDef. +TEST_F(BaseApiTest, AllApiDefAttributeNamesAreValid) { + TestAllApiDefAttributeNamesAreValid(ops_, api_defs_map_); +} - bool updated_at_least_one_file = false; - const auto golden_api_defs_map = GetGoldenApiDefs(env, api_files_dir); +class PythonApiTest : public ::testing::Test { + protected: + PythonApiTest() { + OpRegistry::Global()->Export(false, &ops_); + const std::vector multi_line_fields = {"description"}; - for (auto new_api_entry : new_api_defs_map) { - const auto& file_path = new_api_entry.first; - std::string golden_api_defs_str = ""; - if (golden_api_defs_map.find(file_path) != golden_api_defs_map.end()) { - golden_api_defs_str = golden_api_defs_map.at(file_path); - } - string new_api_defs_str = new_api_entry.second.DebugString(); - new_api_defs_str = PBTxtToMultiline(new_api_defs_str, multi_line_fields); - if (golden_api_defs_str == new_api_defs_str) { - continue; - } - if (update_api_def) { - std::cout << "Updating " << file_path << "..." << std::endl; - TF_EXPECT_OK(WriteStringToFile(env, file_path, new_api_defs_str)); - updated_at_least_one_file = true; - } else { - EXPECT_EQ(golden_api_defs_str, new_api_defs_str) - << "To update golden API files, run " - << "tensorflow/core/api_def/update_api_def.sh."; - } + Env* env = Env::Default(); + GetGoldenApiDefs(env, kPythonApiDefDir, &api_defs_map_); } + OpList ops_; + std::unordered_map api_defs_map_; +}; - for (const auto& golden_api_entry : golden_api_defs_map) { - const auto& file_path = golden_api_entry.first; - if (new_api_defs_map.find(file_path) == new_api_defs_map.end()) { - if (update_api_def) { - std::cout << "Deleting " << file_path << "..." << std::endl; - TF_EXPECT_OK(env->DeleteFile(file_path)); - updated_at_least_one_file = true; - } else { - EXPECT_EQ("", golden_api_entry.second) - << "To update golden API files, run " - << "tensorflow/core/api_def/update_api_def.sh."; - } - } - } +// Check that ApiDefs have a corresponding op. +TEST_F(PythonApiTest, AllApiDefsHaveCorrespondingOp) { + TestAllApiDefsHaveCorrespondingOp(ops_, api_defs_map_); +} - if (update_api_def && !updated_at_least_one_file) { - std::cout << "Api def files are already up to date." << std::endl; - } +// Checks that input arg names in an ApiDef match input +// arg names in corresponding OpDef. +TEST_F(PythonApiTest, AllApiDefInputArgsAreValid) { + TestAllApiDefInputArgsAreValid(ops_, api_defs_map_); } -TEST(ApiTest, GenerateBaseAPIDef) { RunApiTest(false, kDefaultApiDefDir); } -} // namespace -} // namespace tensorflow +// Checks that output arg names in an ApiDef match output +// arg names in corresponding OpDef. +TEST_F(PythonApiTest, AllApiDefOutputArgsAreValid) { + TestAllApiDefOutputArgsAreValid(ops_, api_defs_map_); +} -int main(int argc, char** argv) { - bool update_api_def = false; - tensorflow::string api_files_dir = tensorflow::kDefaultApiDefDir; - std::vector flag_list = { - tensorflow::Flag( - "update_api_def", &update_api_def, - "Whether to update tensorflow/core/api_def/base_api/api_def*.pbtxt " - "files if they differ from expected API."), - tensorflow::Flag("api_def_dir", &api_files_dir, - "Base directory of api_def*.pbtxt files.")}; - std::string usage = tensorflow::Flags::Usage(argv[0], flag_list); - bool parsed_values_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); - if (!parsed_values_ok) { - std::cerr << usage << std::endl; - return 2; - } - if (update_api_def) { - tensorflow::port::InitMain(argv[0], &argc, &argv); - tensorflow::RunApiTest(update_api_def, api_files_dir); - return 0; - } - testing::InitGoogleTest(&argc, argv); - // Run tests - return RUN_ALL_TESTS(); +// Checks that attribute names in an ApiDef match attribute +// names in corresponding OpDef. +TEST_F(PythonApiTest, AllApiDefAttributeNamesAreValid) { + TestAllApiDefAttributeNamesAreValid(ops_, api_defs_map_); } + +} // namespace tensorflow diff --git a/tensorflow/core/api_def/base_api/api_def_AssignAddVariableOp.pbtxt b/tensorflow/core/api_def/base_api/api_def_AssignAddVariableOp.pbtxt index 5d21d7bab699ff481c65ed44eb9bf66ec14ea387..ac05b54eea95f70e4a6db843aab13adf7b94602c 100644 --- a/tensorflow/core/api_def/base_api/api_def_AssignAddVariableOp.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_AssignAddVariableOp.pbtxt @@ -20,10 +20,7 @@ END } summary: "Adds a value to the current value of a variable." description: <